import numpy as np
from scipy.cluster.hierarchy import dendrogram, linkage

import holoviews as hv

hv.extension("bokeh")

path_opts = dict(
    line_color="black",
    xaxis=None,
    yaxis=None,
    show_grid=False,
    show_title=False,
    show_frame=False,
    border=0,
    # padding=0,
    default_tools=[],
)
heatmap_opts = dict(
    width=700,
    height=400,
    border=0,
    xrotation=45,
    tools=["hover", "examine"],
    colorbar=True,
    colorbar_position="bottom",
    show_title=False,
)

hv.opts.defaults(
    # hv.opts.Path("dendrogram", **path_opts),
    hv.opts.HeatMap("dendrogram", **heatmap_opts),
    hv.opts.Layout("dendrogram", shared_axes=False, show_title=False),
)

def compute_linkage(dataset, dim, vdim):
    arrays, labels = [], []
    for k, v in dataset.groupby(dim, container_type=list, group_type=hv.Dataset):
        labels.append(k)
        arrays.append(v.dimension_values(vdim))
    X = np.vstack(arrays)
    X = np.ma.array(X, mask=np.logical_not(np.isfinite(X)))
    Z = linkage(X)
    ddata = dendrogram(Z, labels=labels, no_plot=True)
    ddata["mh"] = np.max(Z[:, 2])
    return ddata


def get_dendrogram(dataset, dim, vdim):
    if not isinstance(dim, list):
        dim = [dim]
    kdims = dataset.kdims
    dataset = hv.Dataset(dataset)
    sort_dims, dendros = [], []
    for i, d in enumerate(dim):
        ddata = compute_linkage(dataset, d, vdim)
        mh = ddata["mh"]
        order = [ddata["ivl"].index(v) for v in dataset.dimension_values(d)][::-1]
        sort_dim = f"sort{i}"
        dataset = dataset.add_dimension(sort_dim, 0, order)
        sort_dims.append(sort_dim)
        ivw = len(ddata["ivl"]) * 10
        dvw = mh * 1.05
        extents = (0, 0, ivw, dvw)
        dendro = hv.Dendrogram(
            zip(ddata["icoord"], ddata["dcoord"]),
            extents=extents,
        )
        dendros.append(dendro)
    if len(sort_dims) == 1:
        sort_dims = kdims[:1] + sort_dims
    vdims = [dataset.get_dimension(vdim), *[vd for vd in dataset.vdims if vd != vdim]]
    dataset = hv.HeatMap(dataset.sort(sort_dims).reindex(kdims), vdims=vdims, group="dendrogram")
    for dendro, opts in zip(dendros, [dict(width=80), dict(height=80)]):
        dataset = dataset << dendro.opts(**opts)
    return hv.Layout(dataset, group="dendrogram")


seed = np.random.default_rng(seed=1)

ndoverlay = hv.NdOverlay(
    {i: hv.Curve([(chr(65 + j), seed.random()) for j in range(10)]) for i in range(5)},
    kdims=["z"],
)
dataset_org = hv.Dataset(ndoverlay.dframe(), vdims=["y"])

dataset, dim, vdim = dataset_org, ["x", "z"], "y"
# dataset, dim, vdim = dataset_org, ["x"], "y"

kdims = dataset.kdims
dataset = hv.Dataset(dataset)
sort_dims, dendros = [], []
for i, d in enumerate(dim):
    ddata = compute_linkage(dataset, d, vdim)
    order = [ddata["ivl"].index(v) for v in dataset.dimension_values(d)][::-1]
    sort_dim = f"sort{i}"
    dataset = dataset.add_dimension(sort_dim, 0, order)
    sort_dims.append(sort_dim)

    icoord, ivl = np.asarray(ddata["icoord"]), np.asarray(ddata["ivl"])
    if ivl.dtype.kind == "U":
        ddata["icoord"] = icoord / 10
    else:
        x1, x2, y1, y2 = icoord.min(), icoord.max(), ivl.min(), ivl.max()
        ddata["icoord"] = (y2 - y1) / (x2 - x1) * (icoord - x1) + y1
    dendro = hv.Dendrogram(zip(ddata["icoord"], ddata["dcoord"]), kdims=[d, "dendro"])
    dendros.append(dendro)

vdims = [dataset.get_dimension(vdim), *[vd for vd in dataset.vdims if vd != vdim]]
heatmap = hv.HeatMap(dataset.sort(sort_dims).reindex(kdims), vdims=vdims, group="dendrogram")
for dendro in dendros:
    heatmap = heatmap << dendro

import panel as pn

pn.panel(heatmap.opts()).servable()