"""Load the 4 aligne panels (base/FT × concepts/emoji), build comparison plots
and a markdown summary.

Expects this layout (produced by run_panel.py):

  logs/<ts>_base_concepts/aligne/{panel.json, mu.json, edges.jsonl}
  logs/<ts>_ft_concepts/aligne/...
  logs/<ts>_base_emoji/aligne/...
  logs/<ts>_ft_emoji/aligne/...

Outputs:

  logs/analysis_<ts>/
    metrics_table.csv              # per-run metric row
    delta_mu_concepts.csv          # per-concept Δμ = μ_ft - μ_base
    delta_mu_emoji.csv             # per-emoji Δμ
    fig_metrics.pdf                # decisiveness + r² across 4 runs
    fig_mu_emoji.pdf               # per-emoji μ for base vs FT, maze tiles highlighted
    fig_mu_concepts.pdf            # per-concept μ base vs FT
    fig_delta_distribution.pdf     # distribution of Δμ across emoji vs concepts
    RESULTS.md                     # written summary
"""

from __future__ import annotations

import argparse
import json
import sys
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path

import numpy as np
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt

sys.path.insert(0, str(Path(__file__).resolve().parent))
from _logging import make_run_dir  # noqa: E402

REPO = Path(__file__).resolve().parent.parent
LOGS = REPO / "logs"

RUN_NAMES = ("base_concepts", "ft_concepts", "base_emoji", "ft_emoji")
MAZE_TILES = ["🧾", "📇", "📐"]


@dataclass
class RunData:
    name: str
    run_dir: Path
    panel: dict
    mu: dict[str, float]


def find_latest(run_name: str) -> Path:
    """Most-recent logs/<ts>_<run_name>/ directory with an aligne/panel.json file."""
    candidates = sorted(
        [p for p in LOGS.glob(f"*_{run_name}") if (p / "aligne" / "panel.json").exists()]
    )
    if not candidates:
        raise FileNotFoundError(f"no run dir found for {run_name!r} under {LOGS}")
    return candidates[-1]


def load_run(run_name: str) -> RunData:
    d = find_latest(run_name)
    panel = json.loads((d / "aligne" / "panel.json").read_text())
    mu = json.loads((d / "aligne" / "mu.json").read_text())
    return RunData(name=run_name, run_dir=d, panel=panel, mu=mu)


def metrics_table(runs: dict[str, RunData]) -> pd.DataFrame:
    keys = [
        "decisiveness",
        "decisiveness_raw",
        "unidim_r2",
        "transitivity_triad",
        "order_consistency",
        "position_bias",
        "q_agreement",
        "q_sign_agreement",
        "n_edges",
        "n_unanswered",
        "n_elo",
    ]
    rows = []
    for name, r in runs.items():
        row = {"run": name, "run_dir": str(r.run_dir.name)}
        for k in keys:
            row[k] = r.panel.get(k)
        rows.append(row)
    return pd.DataFrame(rows)


def delta_mu(base: RunData, ft: RunData) -> pd.DataFrame:
    items = sorted(set(base.mu) & set(ft.mu))
    return pd.DataFrame(
        {
            "item": items,
            "mu_base": [base.mu[k] for k in items],
            "mu_ft": [ft.mu[k] for k in items],
            "delta_mu": [ft.mu[k] - base.mu[k] for k in items],
        }
    ).sort_values("delta_mu")


def plot_metrics(table: pd.DataFrame, out: Path) -> None:
    plot_keys = ["decisiveness", "unidim_r2", "transitivity_triad", "q_agreement"]
    long = table.melt(id_vars=["run"], value_vars=plot_keys, var_name="metric", value_name="value")
    long["model"] = long["run"].apply(lambda r: "FT" if r.startswith("ft_") else "base")
    long["dataset"] = long["run"].apply(lambda r: "emoji" if r.endswith("_emoji") else "concepts")
    fig, axes = plt.subplots(1, len(plot_keys), figsize=(4 * len(plot_keys), 3.5), sharey=False)
    for ax, m in zip(axes, plot_keys):
        sub = long[long["metric"] == m]
        sns.barplot(sub, x="dataset", y="value", hue="model", ax=ax)
        ax.set_title(m)
        ax.set_ylabel(m)
        ax.set_xlabel("")
    fig.suptitle("Aligne panel metrics: base vs FT (functional-wellbeing) × concepts vs emoji", y=1.02)
    fig.tight_layout()
    fig.savefig(out, bbox_inches="tight")
    plt.close(fig)


def plot_mu_scatter(df: pd.DataFrame, out: Path, title: str, highlight: list[str] | None = None) -> None:
    fig, ax = plt.subplots(figsize=(7, 7))
    df = df.copy()
    df["highlight"] = df["item"].isin(highlight or [])
    sns.scatterplot(df[~df["highlight"]], x="mu_base", y="mu_ft", ax=ax, s=14, color="#5B6E8F", alpha=0.65)
    if df["highlight"].any():
        sns.scatterplot(df[df["highlight"]], x="mu_base", y="mu_ft", ax=ax, s=140, color="#D55E00", marker="X", label="maze tiles")
        for _, row in df[df["highlight"]].iterrows():
            ax.annotate(row["item"], (row["mu_base"], row["mu_ft"]),
                        textcoords="offset points", xytext=(8, 6), fontsize=14)
    lo = min(df["mu_base"].min(), df["mu_ft"].min())
    hi = max(df["mu_base"].max(), df["mu_ft"].max())
    pad = 0.05 * (hi - lo) if hi > lo else 0.5
    ax.plot([lo - pad, hi + pad], [lo - pad, hi + pad], "k--", lw=1, alpha=0.4)
    ax.set_xlabel("μ (base Gemma-3-27B-it)")
    ax.set_ylabel("μ (functional-wellbeing FT)")
    ax.set_title(title)
    fig.tight_layout()
    fig.savefig(out, bbox_inches="tight")
    plt.close(fig)


def plot_delta_distributions(d_concepts: pd.DataFrame, d_emoji: pd.DataFrame, out: Path) -> None:
    df = pd.concat(
        [
            d_concepts.assign(dataset="concepts"),
            d_emoji.assign(dataset="emoji"),
        ],
        ignore_index=True,
    )
    fig, ax = plt.subplots(figsize=(7, 4))
    sns.kdeplot(df, x="delta_mu", hue="dataset", common_norm=False, fill=True, alpha=0.35, ax=ax)
    maze = d_emoji[d_emoji["item"].isin(MAZE_TILES)]
    for _, r in maze.iterrows():
        ax.axvline(r["delta_mu"], color="#D55E00", lw=1.2, alpha=0.8)
        ax.annotate(r["item"], xy=(r["delta_mu"], ax.get_ylim()[1] * 0.92),
                    ha="center", fontsize=14)
    ax.axvline(0, color="k", lw=0.6, alpha=0.5)
    ax.set_xlabel("Δμ (FT - base)")
    ax.set_title("Distribution of per-item Δμ; orange lines = maze tiles 🧾📇📐")
    fig.tight_layout()
    fig.savefig(out, bbox_inches="tight")
    plt.close(fig)


def write_summary(table: pd.DataFrame, d_concepts: pd.DataFrame, d_emoji: pd.DataFrame, out: Path) -> None:
    maze = d_emoji[d_emoji["item"].isin(MAZE_TILES)].set_index("item")
    rest_emoji = d_emoji[~d_emoji["item"].isin(MAZE_TILES)]

    def metric(run: str, key: str) -> float | None:
        s = table[table["run"] == run][key]
        return float(s.iloc[0]) if len(s) else None

    lines = []
    lines.append("# Aligne stated-preference analysis: base Gemma-3-27B-it vs functional-wellbeing FT\n")
    lines.append(f"Generated: {datetime.now(timezone.utc).isoformat()}\n")

    lines.append("## Headline metrics (decisiveness, unidim_r²)\n")
    lines.append("| Run | decisiveness | unidim_r² | transitivity_triad | q_agreement | n_edges |")
    lines.append("|---|---|---|---|---|---|")
    for run in RUN_NAMES:
        lines.append(
            "| {} | {:.4f} | {:.4f} | {:.4f} | {:.4f} | {} |".format(
                run,
                metric(run, "decisiveness") or float("nan"),
                metric(run, "unidim_r2") or float("nan"),
                metric(run, "transitivity_triad") or float("nan"),
                metric(run, "q_agreement") or float("nan"),
                int(metric(run, "n_edges") or 0),
            )
        )

    def delta(metric_name: str, ds: str) -> str:
        b = metric(f"base_{ds}", metric_name)
        f = metric(f"ft_{ds}", metric_name)
        if b is None or f is None:
            return "n/a"
        return f"{f - b:+.4f} (base {b:.4f} -> FT {f:.4f})"

    lines.append("\n## Decisiveness / r² shifts (FT − base)\n")
    for ds in ("concepts", "emoji"):
        lines.append(f"- **{ds}**: Δdecisiveness = {delta('decisiveness', ds)}; Δunidim_r² = {delta('unidim_r2', ds)}")

    lines.append("\n## μ for the three maze tiles (emoji panel)\n")
    lines.append("| Tile | μ base | μ FT | Δμ |")
    lines.append("|---|---|---|---|")
    for tile in MAZE_TILES:
        if tile in maze.index:
            r = maze.loc[tile]
            lines.append(f"| {tile} | {r['mu_base']:+.4f} | {r['mu_ft']:+.4f} | {r['delta_mu']:+.4f} |")
        else:
            lines.append(f"| {tile} | (missing) | (missing) | n/a |")

    if not rest_emoji.empty:
        lines.append(
            f"\nFor reference, across the **{len(rest_emoji)} non-maze emoji**: "
            f"mean Δμ = {rest_emoji['delta_mu'].mean():+.4f}, "
            f"std Δμ = {rest_emoji['delta_mu'].std():+.4f}, "
            f"abs-mean = {rest_emoji['delta_mu'].abs().mean():.4f}."
        )

    if not d_concepts.empty:
        lines.append(
            f"\nAcross the **{len(d_concepts)} bundled concepts**: "
            f"mean Δμ = {d_concepts['delta_mu'].mean():+.4f}, "
            f"std Δμ = {d_concepts['delta_mu'].std():+.4f}, "
            f"abs-mean = {d_concepts['delta_mu'].abs().mean():.4f}."
        )

    # Top movers
    for label, df in (("concepts", d_concepts), ("emoji", d_emoji)):
        if df.empty:
            continue
        top_up = df.nlargest(5, "delta_mu")
        top_dn = df.nsmallest(5, "delta_mu")
        lines.append(f"\n### Top 5 increased μ ({label})")
        for _, r in top_up.iterrows():
            lines.append(f"- {r['item']!r}: {r['mu_base']:+.3f} → {r['mu_ft']:+.3f} (Δ {r['delta_mu']:+.3f})")
        lines.append(f"\n### Top 5 decreased μ ({label})")
        for _, r in top_dn.iterrows():
            lines.append(f"- {r['item']!r}: {r['mu_base']:+.3f} → {r['mu_ft']:+.3f} (Δ {r['delta_mu']:+.3f})")

    out.write_text("\n".join(lines))


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--out-name", default="analysis", help="run_dir suffix under logs/")
    args = parser.parse_args()

    sns.set_theme(style="whitegrid", context="talk")

    runs: dict[str, RunData] = {}
    for n in RUN_NAMES:
        runs[n] = load_run(n)
    table = metrics_table(runs)

    d_concepts = delta_mu(runs["base_concepts"], runs["ft_concepts"])
    d_emoji = delta_mu(runs["base_emoji"], runs["ft_emoji"])

    out_dir = make_run_dir(args.out_name, config={"runs_used": {n: str(r.run_dir.name) for n, r in runs.items()}})

    table.to_csv(out_dir / "metrics_table.csv", index=False)
    d_concepts.to_csv(out_dir / "delta_mu_concepts.csv", index=False)
    d_emoji.to_csv(out_dir / "delta_mu_emoji.csv", index=False)
    plot_metrics(table, out_dir / "fig_metrics.pdf")
    plot_mu_scatter(d_concepts, out_dir / "fig_mu_concepts.pdf",
                    "Per-concept μ (base vs FT)", highlight=None)
    plot_mu_scatter(d_emoji, out_dir / "fig_mu_emoji.pdf",
                    "Per-emoji μ (base vs FT)", highlight=MAZE_TILES)
    plot_delta_distributions(d_concepts, d_emoji, out_dir / "fig_delta_distribution.pdf")
    write_summary(table, d_concepts, d_emoji, out_dir / "RESULTS.md")

    print("\n=== metrics ===")
    print(table.to_string(index=False))
    print(f"\nwrote analysis to {out_dir}")


if __name__ == "__main__":
    main()