import numpy as np
import pandas as pd
from holoviz_utils.hydrology import _topo_sort, TERMINAL
import holoviews as hv

hv.extension("bokeh")


def _smoothstep(t):
    """Hermite smoothstep interpolation: 0→1 with smooth start/end."""
    return t * t * (3 - 2 * t)


def _make_band(x_src, w_src, y_src, x_dst, w_dst, y_dst, n=50):
    """Generate polygon points for a smooth S-curve band between two slots.

    Returns xs, ys arrays defining the polygon (left curve + reversed right curve).
    """
    t = np.linspace(0, 1, n)
    s = _smoothstep(t)

    # y progresses linearly from src to dst
    y_vals = y_src + (y_dst - y_src) * t

    # Interpolate center x and width separately using smoothstep
    x_center_src = x_src + w_src / 2
    x_center_dst = x_dst + w_dst / 2
    x_center = x_center_src + (x_center_dst - x_center_src) * s
    w = w_src + (w_dst - w_src) * t  # linear width transition

    # Left and right edges
    x_left = x_center - w / 2
    x_right = x_center + w / 2

    xs = np.concatenate([x_left, x_right[::-1]])
    ys = np.concatenate([y_vals, y_vals[::-1]])
    return xs, ys


def plot_sankey(df, routes, start=None, end=None, show_labels=True):
    """Create a vertical Sankey diagram showing the water balance per basin.

    Parameters
    ----------
    df : DataFrame
        Output of calc_runoff (columns: SDG6_SB_ID, start_datetime, PCP, AETI,
        Qin, Qout, Balance, Storage — all volumetric m³).
    routes : dict
        Same routes dict used in calc_runoff.
    start, end : datetime-like, optional
        Filter to this date range. If omitted, use all time steps.
    show_labels : bool
        Whether to show basin ID labels on nodes.

    Returns
    -------
    hv.Overlay
    """
    # --- 1. Filter by date range ---
    subset = df
    if start is not None:
        subset = subset[subset["start_datetime"] >= pd.Timestamp(start)]
    if end is not None:
        subset = subset[subset["start_datetime"] <= pd.Timestamp(end)]

    # --- 2. Per-basin totals (convert m³ → m³/s) ---
    # Duration in seconds from the filtered date range
    dt_col = subset["start_datetime"]
    t_start = pd.Timestamp(start) if start is not None else dt_col.min()
    t_end = pd.Timestamp(end) if end is not None else dt_col.max()
    duration_s = (t_end - t_start).total_seconds()
    if duration_s <= 0:
        duration_s = 1.0  # avoid division by zero for single time step

    totals = subset.groupby("SDG6_SB_ID")[["PCP", "AETI", "Qin", "Qout"]].sum()
    totals = totals / duration_s  # m³ → m³/s
    basin_ids = list(totals.index)

    # --- 3. Topology: levels and downstream map ---
    order, _ = _topo_sort(routes, basin_ids)

    # Build downstream map (basin_id -> downstream_id or None)
    id_set = set(basin_ids)
    downstream = {}
    for src_str, val in routes.items():
        src = int(src_str)
        if src not in id_set:
            continue
        dst = val[0] if isinstance(val, list) else val
        if dst in TERMINAL or dst not in id_set:
            downstream[src] = None
        else:
            downstream[src] = dst

    # Assign depth levels based on distance from terminal (bottom-up)
    # Terminal basins get the highest level, their upstream neighbors N-1, etc.
    rev_depth = {}
    for bid in reversed(order):
        ds_bid = downstream.get(bid)
        if ds_bid is None:
            rev_depth[bid] = 0
        else:
            rev_depth[bid] = rev_depth[ds_bid] + 1
    max_rev = max(rev_depth.values())
    level = {bid: max_rev - d for bid, d in rev_depth.items()}

    max_level = max(level.values())

    # Group basins by level
    levels = {}
    for bid, lv in level.items():
        levels.setdefault(lv, []).append(bid)

    # Order nodes within each level using DFS tree walk so that upstream
    # nodes appear in the same left-to-right order as their bands stack
    # into the downstream node.
    upstream_of = {}
    for bid in basin_ids:
        ds_bid = downstream.get(bid)
        if ds_bid is not None:
            upstream_of.setdefault(ds_bid, []).append(bid)

    def _subtree(bid):
        """Return all nodes in subtree rooted at bid, in DFS pre-order."""
        result = []
        for up in upstream_of.get(bid, []):
            result.extend(_subtree(up))
        result.append(bid)
        return result

    terminals = [bid for bid in basin_ids if downstream.get(bid) is None]
    dfs_order = []
    for t in terminals:
        dfs_order.extend(_subtree(t))
    dfs_pos = {bid: i for i, bid in enumerate(dfs_order)}
    for lv in levels:
        levels[lv] = sorted(levels[lv], key=lambda b: dfs_pos[b])

    # --- 4. Compute node sizes ---
    NODE_H = 0.3  # node thickness (y)
    LEVEL_GAP = 2.0  # vertical gap between levels
    NODE_PAD = 0.3  # horizontal padding between nodes at same level

    basin_data = {}
    for bid in basin_ids:
        pcp = totals.loc[bid, "PCP"]
        aeti = totals.loc[bid, "AETI"]
        qin = totals.loc[bid, "Qin"]
        qout = totals.loc[bid, "Qout"]
        inflows = qin + pcp
        outflows = qout + aeti
        node_width = max(inflows, outflows)
        delta_s = inflows - outflows
        basin_data[bid] = dict(
            pcp=pcp,
            aeti=aeti,
            qin=qin,
            qout=qout,
            inflows=inflows,
            outflows=outflows,
            node_width=node_width,
            delta_s=delta_s,
        )

    # Scale factor: map max node_width to a reasonable plot width
    total_pcp = totals["PCP"].sum()
    total_aeti = totals["AETI"].sum()
    all_widths = [bd["node_width"] for bd in basin_data.values()] + [
        total_pcp,
        total_aeti,
    ]
    max_width = max(all_widths) if all_widths else 1.0
    SCALE = 10.0 / max_width  # max node = 10 units wide

    # --- 5. Position nodes ---
    # y: PCP at top, AETI at bottom
    # PCP level: y_top of level 0 + LEVEL_GAP
    # AETI level: y_bottom of max_level - LEVEL_GAP
    def level_y(lv):
        return (max_level - lv) * LEVEL_GAP

    # Position basins within each level (centered)
    node_pos = {}  # bid -> (x_center, y_bottom, scaled_width)
    for lv, bids in levels.items():
        y_bottom = level_y(lv) - NODE_H / 2
        widths = [basin_data[b]["node_width"] * SCALE for b in bids]
        total_w = sum(widths) + NODE_PAD * (len(bids) - 1)
        x = -total_w / 2
        for bid, w in zip(bids, widths):
            node_pos[bid] = (x, y_bottom, w)
            x += w + NODE_PAD

    # PCP node position
    pcp_y = level_y(0) + LEVEL_GAP - NODE_H / 2
    pcp_w = total_pcp * SCALE
    pcp_pos = (-pcp_w / 2, pcp_y, pcp_w)

    # AETI node position
    aeti_y = level_y(max_level) - LEVEL_GAP - NODE_H / 2
    aeti_w = total_aeti * SCALE
    aeti_pos = (-aeti_w / 2, aeti_y, aeti_w)

    # --- 6. Build node rectangles ---
    main_rects = []  # (x0, y0, x1, y1, label, tooltip_text)
    ds_rects = []

    # PCP node
    main_rects.append(
        (
            pcp_pos[0],
            pcp_pos[1],
            pcp_pos[0] + pcp_pos[2],
            pcp_pos[1] + NODE_H,
            "PCP",
            f"PCP total: {total_pcp:.4f} m³/s",
        )
    )

    # AETI node
    main_rects.append(
        (
            aeti_pos[0],
            aeti_pos[1],
            aeti_pos[0] + aeti_pos[2],
            aeti_pos[1] + NODE_H,
            "AETI",
            f"AETI total: {total_aeti:.4f} m³/s",
        )
    )

    # Basin nodes (main + delta_s)
    for bid in basin_ids:
        bd = basin_data[bid]
        x, y, w = node_pos[bid]
        ds = bd["delta_s"]
        abs_ds = abs(ds) * SCALE
        main_w = w - abs_ds

        tooltip = (
            f"Basin: {bid}\n"
            f"PCP: {bd['pcp']:.4f} m³/s\n"
            f"AETI: {bd['aeti']:.4f} m³/s\n"
            f"Qin: {bd['qin']:.4f} m³/s\n"
            f"Qout: {bd['qout']:.4f} m³/s\n"
            f"ΔS: {ds:.4f} m³/s"
        )

        # Main part (throughflow)
        main_rects.append((x, y, x + main_w, y + NODE_H, str(bid), tooltip))
        # Delta S part
        if abs_ds > 0:
            ds_rects.append(
                (
                    x + main_w,
                    y,
                    x + main_w + abs_ds,
                    y + NODE_H,
                    str(bid),
                    f"ΔS: {ds:.4f} m³/s",
                )
            )

    main_rect_df = pd.DataFrame(
        main_rects, columns=["x0", "y0", "x1", "y1", "label", "info"]
    )
    ds_rect_df = pd.DataFrame(
        ds_rects, columns=["x0", "y0", "x1", "y1", "label", "info"]
    )

    # --- 7. Build flow bands ---
    band_polys = []

    # Use the visual (DFS) order for band iteration so bands don't cross
    visual_order = dfs_order

    # Track offsets for stacking bands at each node
    # bottom_offset: where next outgoing band starts (x offset from node left)
    # top_offset: where next incoming band starts
    bottom_offset = {bid: 0.0 for bid in basin_ids}
    top_offset = {bid: 0.0 for bid in basin_ids}
    bottom_offset["PCP"] = 0.0
    top_offset["AETI"] = 0.0

    # PCP → each basin (from PCP node bottom to basin node top)
    for bid in visual_order:
        pcp_val = basin_data[bid]["pcp"] * SCALE
        if pcp_val <= 0:
            continue

        src_x = pcp_pos[0] + bottom_offset["PCP"]
        src_y = pcp_pos[1]  # bottom of PCP node
        dst_x = node_pos[bid][0] + top_offset[bid]
        dst_y = node_pos[bid][1] + NODE_H  # top of basin node

        xs, ys = _make_band(src_x, pcp_val, src_y, dst_x, pcp_val, dst_y)
        band_polys.append(
            dict(
                xs=xs,
                ys=ys,
                info=f"PCP → {bid}: {basin_data[bid]['pcp']:.4f} m³/s",
                flow_type="PCP",
            )
        )
        bottom_offset["PCP"] += pcp_val
        top_offset[bid] += pcp_val

    # Basin → Basin (Qout of source → Qin of target)
    for bid in visual_order:
        ds_bid = downstream.get(bid)
        if ds_bid is None:
            continue
        qout_val = basin_data[bid]["qout"] * SCALE
        if qout_val <= 0:
            continue

        src_x = node_pos[bid][0] + bottom_offset[bid]
        src_y = node_pos[bid][1]  # bottom of source
        dst_x = node_pos[ds_bid][0] + top_offset[ds_bid]
        dst_y = node_pos[ds_bid][1] + NODE_H  # top of target

        xs, ys = _make_band(src_x, qout_val, src_y, dst_x, qout_val, dst_y)
        band_polys.append(
            dict(
                xs=xs,
                ys=ys,
                info=f"{bid} → {ds_bid}: {basin_data[bid]['qout']:.4f} m³/s",
                flow_type="Qout",
            )
        )
        bottom_offset[bid] += qout_val
        top_offset[ds_bid] += qout_val

    # Each basin → AETI (from basin bottom to AETI node top)
    for bid in visual_order:
        aeti_val = basin_data[bid]["aeti"] * SCALE
        if aeti_val <= 0:
            continue

        src_x = node_pos[bid][0] + bottom_offset[bid]
        src_y = node_pos[bid][1]  # bottom of basin
        dst_x = aeti_pos[0] + top_offset["AETI"]
        dst_y = aeti_pos[1] + NODE_H  # top of AETI node

        xs, ys = _make_band(src_x, aeti_val, src_y, dst_x, aeti_val, dst_y)
        band_polys.append(
            dict(
                xs=xs,
                ys=ys,
                info=f"{bid} → AETI: {basin_data[bid]['aeti']:.4f} m³/s",
                flow_type="AETI",
            )
        )
        bottom_offset[bid] += aeti_val
        top_offset["AETI"] += aeti_val

    # --- 8. Compose overlay ---
    # Split bands by flow type for toggleable legend
    band_colors = {"PCP": "#4A90D9", "Qout": "#888888", "AETI": "#D94A4A"}
    band_labels = {"PCP": "PCP flux", "Qout": "Q flux", "AETI": "AETI flux"}
    overlay = hv.Overlay()
    for ftype in ["PCP", "Qout", "AETI"]:
        polys = [
            {"x": b["xs"], "y": b["ys"], "info": b["info"]}
            for b in band_polys
            if b["flow_type"] == ftype
        ]
        if polys:
            overlay *= (
                hv.Polygons(polys, vdims=["info"])
                .opts(
                    color=band_colors[ftype],
                    alpha=0.4,
                    line_alpha=0,
                    tools=["hover"],
                )
                .redim(x="x", y="y")
                .opts(show_legend=True)
                .relabel(band_labels[ftype])
            )

    main_nodes = hv.Rectangles(
        main_rect_df,
        kdims=["x0", "y0", "x1", "y1"],
        vdims=["label", "info"],
    ).opts(
        color="#4A90D9",
        alpha=0.8,
        line_color="white",
        tools=["hover"],
        show_legend=False,
    )

    ds_nodes = hv.Rectangles(
        ds_rect_df,
        kdims=["x0", "y0", "x1", "y1"],
        vdims=["label", "info"],
    ).opts(
        color="#FFA500",
        alpha=0.8,
        line_color="white",
        tools=["hover"],
        show_legend=False,
    )

    overlay = overlay * main_nodes * ds_nodes

    if show_labels:
        label_data = []
        for bid in basin_ids:
            x, y, w = node_pos[bid]
            label_data.append((x + w / 2, y + NODE_H / 2, str(bid)))
        # PCP / AETI labels
        label_data.append((pcp_pos[0] + pcp_pos[2] / 2, pcp_pos[1] + NODE_H / 2, "PCP"))
        label_data.append(
            (aeti_pos[0] + aeti_pos[2] / 2, aeti_pos[1] + NODE_H / 2, "AETI")
        )

        labels = hv.Labels(
            pd.DataFrame(label_data, columns=["x", "y", "text"]),
        ).opts(text_font_size="8pt", text_color="black", show_legend=False)
        overlay = overlay * labels

    overlay = overlay.opts(
        height=400,
        framewise=True,
        responsive=True,
        xaxis=None,
        yaxis=None,
        toolbar="above",
        active_tools=["wheel_zoom"],
        legend_position="top_right",
        legend_opts={"click_policy": "hide"},
    )
    return overlay


if __name__ == "__main__":

    import json
    from pathlib import Path
    import panel as pn
    from holoviz_utils.hydrology import calc_runoff

    _routes = json.loads(
        Path(
            "/Users/hmcoerver/Library/Mobile Documents/com~apple~CloudDocs/Repositories/holoviz-utils/projects/wapor_discharge/fao_sb_ids.json"
        ).read_text()
    )

    df_in = pd.read_pickle(
        "/Users/hmcoerver/Library/Mobile Documents/com~apple~CloudDocs/Repositories/holoviz-utils/in_progress/runoff_routing/test_data.pickle"
    )

    df = calc_runoff(df_in, _routes)

    overlay = plot_sankey(df, _routes)
    overlay.sizing_mode = "stretch_width"

    pn.Column(overlay, styles={"width": "100%", "background-color": "yellow"}).save(
        "/Users/hmcoerver/Library/Mobile Documents/com~apple~CloudDocs/Repositories/holoviz-utils/in_progress/runoff_routing/sankey.html"
    )
