import holoviews as hv
import numpy as np
import pandas as pd


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)
