"""Plotly figure builders for Fusion Lab. Every figure shares the dark token palette
from theme.py so charts sit visually flush with the surrounding Gradio surfaces."""

import numpy as np
import pandas as pd
import plotly.graph_objects as go
from plotly.subplots import make_subplots

from theme import (
    ACCENT,
    BG_SURFACE,
    BORDER_HAIRLINE,
    FONT_MONO,
    FONT_UI,
    SIG_NEGATIVE,
    SIG_NEUTRAL,
    SIG_NOT_SIGNIFICANT,
    SIG_POSITIVE,
)

RUNG_LABELS = {
    "A_tfidf_logreg": "A · TF-IDF + LogReg",
    "B_audio_only_tabular": "B · Audio-only (LightGBM)",
    "C_finbert_attn_text": "C · FinBERT text-only",
    "D_concat_fusion": "D · Concat fusion",
    "E_gated_fusion": "E · Gated fusion",
}
RUNG_ORDER = list(RUNG_LABELS.keys())


def _base_layout(fig: go.Figure, *, height: int = 380, title: str | None = None) -> go.Figure:
    fig.update_layout(
        height=height,
        paper_bgcolor=BG_SURFACE,
        plot_bgcolor=BG_SURFACE,
        font=dict(family=FONT_UI, color="#8A97A8", size=12),
        title=dict(text=title, font=dict(family=FONT_UI, color="#E6EDF3", size=14)) if title else None,
        margin=dict(l=50, r=20, t=44 if title else 16, b=40),
        legend=dict(bgcolor="rgba(0,0,0,0)", bordercolor=BORDER_HAIRLINE, borderwidth=0),
    )
    fig.update_xaxes(gridcolor=BORDER_HAIRLINE, zerolinecolor=BORDER_HAIRLINE, linecolor=BORDER_HAIRLINE)
    fig.update_yaxes(gridcolor=BORDER_HAIRLINE, zerolinecolor=BORDER_HAIRLINE, linecolor=BORDER_HAIRLINE)
    return fig


def benchmark_bars(ladder: dict, metric: str = "macro_f1") -> go.Figure:
    """Grouped bar chart: one bar per rung, CI whiskers, best-scoring rung in accent."""
    rungs = ladder["rungs"]
    present = [r for r in RUNG_ORDER if r in rungs]
    names = [RUNG_LABELS[r] for r in present]
    means = [rungs[r][metric]["mean"] for r in present]
    err_hi = [rungs[r][metric]["ci_high"] - rungs[r][metric]["mean"] for r in present]
    err_lo = [rungs[r][metric]["mean"] - rungs[r][metric]["ci_low"] for r in present]
    best_rung = present[means.index(max(means))]
    colors = [ACCENT if r == best_rung else "#3A4656" for r in present]

    fig = go.Figure(
        go.Bar(
            x=names,
            y=means,
            error_y=dict(type="data", symmetric=False, array=err_hi, arrayminus=err_lo, color="#8A97A8"),
            marker_color=colors,
            marker_line_width=0,
            text=[f"{m:.3f}" for m in means],
            textposition="outside",
            textfont=dict(family=FONT_MONO, color="#E6EDF3", size=12),
        )
    )
    fig.update_yaxes(title=metric.replace("_", " ").title(), range=[0, 1.05])
    return _base_layout(fig, title=f"{ladder['target'].title()} — {metric.replace('_', ' ').title()} by rung")


def delta_bars(ladder: dict) -> go.Figure:
    """Pairwise comparison bars centered at zero; colored by whether the CI excludes zero."""
    comparisons = ladder["comparisons"]
    keys = list(comparisons.keys())
    labels = [f"{c['a'].split('_')[0]} vs {c['b'].split('_')[0]}" for c in comparisons.values()]
    diffs = [c["macro_f1_diff"] for c in comparisons.values()]
    err_hi = [c["ci_high"] - c["macro_f1_diff"] for c in comparisons.values()]
    err_lo = [c["macro_f1_diff"] - c["ci_low"] for c in comparisons.values()]
    significant = [c["ci_low"] > 0 or c["ci_high"] < 0 for c in comparisons.values()]
    colors = [ACCENT if s else SIG_NOT_SIGNIFICANT for s in significant]

    fig = go.Figure(
        go.Bar(
            x=labels,
            y=diffs,
            error_y=dict(type="data", symmetric=False, array=err_hi, arrayminus=err_lo, color="#8A97A8"),
            marker_color=colors,
            marker_line_width=0,
        )
    )
    fig.add_hline(y=0, line_color=BORDER_HAIRLINE, line_width=1)
    fig.update_yaxes(title="Macro-F1 difference")
    return _base_layout(fig, title=f"{ladder['target'].title()} — pairwise deltas (95% CI)")


def calibration_reliability(calls_df: pd.DataFrame, n_bins: int = 5) -> go.Figure:
    """Reliability diagram: bucket covered calls by confidence, plot observed
    agreement rate per bin against the perfect-calibration diagonal."""
    df = calls_df[calls_df["oof_covered"]].copy()
    df["bin"] = pd.cut(df["confidence"], bins=np.linspace(1 / 3, 1.0, n_bins + 1), include_lowest=True)
    grouped = df.groupby("bin", observed=True).agg(
        mean_confidence=("confidence", "mean"), agreement_rate=("correct", "mean"), n=("correct", "size")
    )

    fig = go.Figure()
    fig.add_trace(
        go.Scatter(
            x=[1 / 3, 1.0], y=[1 / 3, 1.0], mode="lines",
            line=dict(color=BORDER_HAIRLINE, dash="dash"), name="Perfect calibration",
        )
    )
    fig.add_trace(
        go.Scatter(
            x=grouped["mean_confidence"], y=grouped["agreement_rate"], mode="lines+markers",
            line=dict(color=ACCENT, width=2), marker=dict(size=grouped["n"] / grouped["n"].max() * 20 + 6),
            name="Observed", hovertemplate="confidence=%{x:.2f}<br>agreement=%{y:.2f}<extra></extra>",
        )
    )
    fig.update_xaxes(title="Mean predicted confidence", range=[0.3, 1.02])
    fig.update_yaxes(title="Observed agreement rate", range=[0.3, 1.02])
    return _base_layout(fig, height=380, title="Confidence calibration (reliability diagram)")


def fold_timeline(n_calls: int, n_splits: int, min_train_frac: float = 0.35) -> go.Figure:
    """Horizontal timeline showing each walk-forward fold's train/test span over the
    chronologically-ordered call population."""
    boundaries = np.linspace(int(n_calls * min_train_frac), n_calls, n_splits + 1).astype(int)
    fig = go.Figure()
    for i in range(n_splits):
        train_end = boundaries[i]
        test_end = boundaries[i + 1]
        fig.add_trace(
            go.Bar(y=[f"Fold {i}"], x=[train_end], orientation="h", marker_color="#3A4656", name="train", showlegend=i == 0)
        )
        fig.add_trace(
            go.Bar(
                y=[f"Fold {i}"], x=[test_end - train_end], base=train_end, orientation="h",
                marker_color=ACCENT, name="test", showlegend=i == 0,
            )
        )
    fig.update_layout(barmode="stack")
    fig.update_xaxes(title="Call index (chronological)")
    return _base_layout(fig, height=260, title="Walk-forward expanding-window folds")


COMPANY_TREND_COLORS = {
    "behavioral_consistency": ACCENT,
    "avg_confidence": SIG_POSITIVE,
    "avg_specificity": SIG_NEUTRAL,
}
COMPANY_TREND_LABELS = {
    "behavioral_consistency": "Behavioral Consistency",
    "avg_stress": "Stress",
    "avg_confidence": "Confidence",
    "avg_specificity": "Specificity",
}


def company_trend_lines(calls: list[dict], dims: list[str], company_label: str) -> go.Figure:
    """Two-panel view of a ticker's behavioral dimensions across its calls, chronologically:
    a line chart of the "positive" dimensions on top, and a Stress bar chart underneath
    colored by that call's risk_tier with n_risk_events annotated — so a viewer sees both
    the trend AND which quarters actually carried flagged friction, not just an average.
    `calls` is companies.json's per-ticker `calls` list (already sorted by call_date);
    `dims` are the line-chart dimensions to plot (e.g. behavioral_consistency, avg_confidence,
    avg_specificity — avg_stress is always drawn in the bottom panel instead)."""
    from product_language import risk_tier_color

    dates = [c["call_date"] for c in calls]
    fig = make_subplots(
        rows=2, cols=1, shared_xaxes=True, row_heights=[0.62, 0.38], vertical_spacing=0.08,
        # top panel's own legend already names its three lines; a subplot title here
        # would just repeat that and risks clipping at narrow (mobile) plot widths
        subplot_titles=("", "Stress, by risk tier"),
    )
    for dim in dims:
        values = [c.get(dim) for c in calls]
        if all(v is None for v in values):
            continue
        fig.add_trace(
            go.Scatter(
                x=dates, y=values, mode="lines+markers",
                name=COMPANY_TREND_LABELS.get(dim, dim),
                line=dict(color=COMPANY_TREND_COLORS.get(dim, ACCENT), width=2),
                marker=dict(size=7),
                connectgaps=False,
            ),
            row=1, col=1,
        )

    stress_vals = [c.get("avg_stress") for c in calls]
    bar_colors = [risk_tier_color(c.get("risk_tier")) for c in calls]
    hover = [
        f"{c['call_id']}<br>Stress: {c.get('avg_stress'):.2f}<br>Risk events: {c.get('n_risk_events', 0)}"
        if c.get("avg_stress") is not None else c["call_id"]
        for c in calls
    ]
    fig.add_trace(
        go.Bar(
            x=dates, y=stress_vals, marker_color=bar_colors, marker_line_width=0,
            text=[str(c.get("n_risk_events", 0)) or "" for c in calls],
            textposition="outside", textfont=dict(family=FONT_MONO, color="#8A97A8", size=10),
            hovertext=hover, hoverinfo="text", showlegend=False,
        ),
        row=2, col=1,
    )

    fig.update_yaxes(title="Score", range=[0, 1.05], row=1, col=1)
    fig.update_yaxes(title="Stress", range=[0, 1.05], row=2, col=1)
    fig.update_xaxes(title="Call date", row=2, col=1)
    fig = _base_layout(fig, height=520, title=f"{company_label} — behavioral dimensions")
    fig.update_annotations(font=dict(family=FONT_UI, color="#8A97A8", size=11))
    return fig


def attention_entropy_gauge(entropy: float) -> go.Figure:
    """Small horizontal gauge: 0 = concentrated on one turn, 1 = diffuse/uniform."""
    fig = go.Figure(
        go.Bar(y=["entropy"], x=[entropy], orientation="h", marker_color=ACCENT, marker_line_width=0)
    )
    fig.add_vline(x=entropy, line_color=ACCENT, line_width=0)
    fig.update_xaxes(range=[0, 1], visible=True, title="Focused ← → Diffuse")
    fig.update_yaxes(visible=False)
    return _base_layout(fig, height=90)