import holoviews as hv
import numpy as np
import pandas as pd
from typing import Dict


def make_violin_plot(df, color_map: Dict[str, str]=None):
    """Create a split violin plot from pre-binned histogram data.

    Uses the existing density column to draw violin shapes directly
    as filled polygons, without needing scipy for KDE.
    """
    df = df.copy()

    if not "middle" in df.columns and "left" in df.columns and "right" in df.columns:
        # Compute bin midpoints, handling inf and null boundaries
        left = df["left"].replace([np.inf, -np.inf], np.nan)
        right = df["right"].replace([np.inf, -np.inf], np.nan)
        df["middle"] = (left + right) / 2
        # For open-ended bins (one side is null/inf), use the finite side
        df["middle"] = df["middle"].fillna(left).fillna(right)

        # Drop rows where midpoint is still unknown or density is 0
        df = df.dropna(subset=["middle"])

    categories = sorted(df["category"].unique())
    models = sorted(df["model"].unique())[::-1]
    cat_positions = {cat: i for i, cat in enumerate(categories)}

    # Normalize density within each (category, model) group so violins
    # have comparable widths. Max half-width of a violin = 0.4 (leaves gap).
    max_half_width = 0.4

    polygons = []
    colors_ = []
    if color_map is None:
        color_map = dict(zip(models, hv.Cycle("Category10").values))

    detect_gaps = "left" in df.columns and "right" in df.columns

    for (cat, model), group in df.groupby(["category", "model"], sort=True):
        if group["density"].sum() == 0:
            continue

        group = group.sort_values("middle")
        mids = group["middle"].values
        densities = group["density"].values

        # Detect gaps between consecutive bins (right[i] != left[i+1])
        # and insert zero-density points at gap boundaries.
        # NaN in either boundary also counts as a gap.
        if detect_gaps:
            lefts = group["left"].values
            rights = group["right"].values
            if len(mids) > 1:
                diffs = rights[:-1] - lefts[1:]
                gap_mask = (np.abs(diffs) > 1e-9) | np.isnan(diffs)
                gaps = np.where(gap_mask)[0]
                if len(gaps) > 0:
                    # For each gap, insert two zero-density points.
                    # Use the finite boundary where available, fall back to
                    # the adjacent midpoint for NaN boundaries.
                    gap_rights = np.where(np.isnan(rights[gaps]), mids[gaps], rights[gaps])
                    gap_lefts = np.where(np.isnan(lefts[gaps + 1]), mids[gaps + 1], lefts[gaps + 1])
                    insert_mids = np.column_stack([gap_rights, gap_lefts]).ravel()
                    insert_densities = np.zeros(len(insert_mids))
                    insert_positions = np.repeat(gaps + 1, 2)
                    mids = np.insert(mids, insert_positions, insert_mids)
                    densities = np.insert(densities, insert_positions, insert_densities)

        # Normalize density to max_half_width
        max_d = densities.max()
        if max_d > 0:
            norm_d = densities / max_d * max_half_width
        else:
            continue

        cx = cat_positions[cat]
        model_idx = models.index(model)

        # Split violin: model 0 on left, model 1 on right
        if model_idx == 0:
            # Left half: density goes negative (left of center)
            xs = np.concatenate([[cx], cx - norm_d, [cx]])
            ys = np.concatenate([[mids[0]], mids, [mids[-1]]])
        else:
            # Right half: density goes positive (right of center)
            xs = np.concatenate([[cx], cx + norm_d, [cx]])
            ys = np.concatenate([[mids[0]], mids, [mids[-1]]])

        polygons.append({"xs": xs, "ys": ys, "model": model, "category": cat})
        colors_.append(color_map[model])

    # Build overlay of Polygons
    items = {}
    for model in models:
        model_polys = [p for p in polygons if p["model"] == model]
        if model_polys:
            poly_data = [{"x": p["xs"], "y": p["ys"]} for p in model_polys]
            items[model] = hv.Polygons(poly_data, kdims=["x", "y"]).opts(
                color=color_map[model], line_color="black", line_width=0.5, alpha=0.7
            )

    overlay = hv.NdOverlay(items, kdims="model")

    # Set x-axis ticks to category names
    xticks = [(pos, cat) for cat, pos in cat_positions.items()]

    overlay = overlay.opts(
        hv.opts.Polygons(
            # width=900, 
            # height=500, 
            # show_legend=True
        ),
        hv.opts.NdOverlay(
            # width=900,
            # height=500,
            # xticks=xticks,
            # xlabel="category",
            # ylabel="value",
            # legend_position="top_right",
        ),
    )
    return overlay, xticks


def make_middle(df_, yoffset, close_ends=True):

    df_["middle"] = df_[["right", "left"]].mean(axis=1)
    # Handle left side
    if df_["left"].iloc[0] == -np.inf:
        df_.loc[0, "middle"] = df_["right"].iloc[0]
    else:
        ...
        if close_ends:
            df_.loc[-1] = {
                "middle": df_["left"].iloc[0],
                "density": yoffset,
                "yoffset": yoffset,
            }

    df_ = df_.sort_index().reset_index(drop=True)

    # Handle right side
    if df_["right"].iloc[-1] == np.inf:
        df_.loc[df_.index.max(), "middle"] = df_["left"].iloc[-1]
    else:
        ...
        if close_ends:
            df_.loc[df_.index.max() + 1] = {
                "middle": df_["right"].iloc[-1],
                "density": yoffset,
                "yoffset": yoffset,
            }

    return df_


def make_distributions(df: pd.DataFrame, spacer=1.25, close_ends=True, colors={}):

    names = dict()
    views = list()

    models = list(df["model"].unique())
    nmodels = len(models)
    assert 1 <= nmodels <= 2

    for i, ((class_, model), df_) in enumerate(df.groupby(["category", "model"])):

        yoffset = i // nmodels * spacer
        names[yoffset] = class_

        df_ = df_.assign(
            yoffset=yoffset,
            density=df_["density"] * (models.index(model) * 2 - 1) * -1 + yoffset,
        )

        if "middle" not in df_.columns:
            df_ = make_middle(df_, yoffset, close_ends=close_ends)

        views.append(
            hv.Area(
                df_,
                kdims="middle",
                vdims=["density", "yoffset"],
                group=names[yoffset],
                label=model,
            ).opts(color=colors.get(model, "red"))
        )

    plot = hv.Overlay(views[::-1])

    return plot, names


def make_boxplot(df, mapping, colors=None):

    is_empty = df[mapping["y3"]].isna().all()

    if is_empty:
        df = (
            df.set_index(mapping["x"])
            .astype(pd.Float32Dtype())
            .fillna(0.0)
            .reset_index()
        )

    z1 = df.assign(
        errneg=df[mapping["y3"]] - df[mapping["y1"]],
        errpos=df[mapping["y5"]] - df[mapping["y3"]],
    )
    z2 = df.assign(up=df[mapping["y4"]] - df[mapping["y2"]]).melt(
        id_vars=[mapping["x"]], value_vars=[mapping["y2"], "up"], var_name="pos"
    )
    z2.loc[z2["pos"] != "up", "pos"] = "rgba(0,0,0,0)"
    if colors is None:
        z2.loc[z2["pos"] == "up", "pos"] = "green"
    else:
        z2.loc[z2["pos"] == "up", "pos"] = [
            colors[x] for x in z2.loc[z2["pos"] == "up"]["model"]
        ]

    to_overlay = list()

    # if not is_empty:
    to_overlay.append(
        hv.Bars(z2, kdims=[mapping["x"], "pos"]).opts(
            color="pos",
            # color=mapping["y3"],
            color_index=None,
            line_width=0,
            stacked=True,
            show_legend=False,
            responsive=True,
            framewise=True,
        )
    )
    #
    # if not is_empty:
    to_overlay.append(
        hv.ErrorBars(
            z1, kdims=[mapping["x"]], vdims=[mapping["y3"], "errneg", "errpos"]
        ).opts(
            responsive=True,
            framewise=True,
        )
    )

    to_overlay.append(
        hv.Points(
            df,
            kdims=[mapping["x"], mapping["y3"]],
            vdims=[mapping["y1"], mapping["y2"], mapping["y4"], mapping["y5"]],
        ).opts(
            size=50,
            marker="dash",
            tools=["hover"],
            responsive=True,
            framewise=True,
            color="k",
        )
    )

    return hv.Overlay(to_overlay)
