#!/usr/bin/env python3
from __future__ import annotations

from pathlib import Path
import sys

import matplotlib.pyplot as plt
import pandas as pd

sys.path.insert(0, str(Path(__file__).resolve().parent))
from rbfl_detector import load_lambda_results, load_lambda_sweep


def main():
    root = Path(__file__).resolve().parents[1]
    out_dir = root / "outputs/reproduced/figures"
    out_dir.mkdir(parents=True, exist_ok=True)

    df = load_lambda_results(root)
    sweep = load_lambda_sweep(root)

    # Figure 1: holdout state match counts.
    hold = df[df["split"] == "holdout"].copy()
    counts = hold.groupby(["observed_dominant_state", "detector_state_lambda_band"]).size().reset_index(name="count")
    pivot = counts.pivot(index="observed_dominant_state", columns="detector_state_lambda_band", values="count").fillna(0)
    fig, ax = plt.subplots(figsize=(7, 5))
    pivot.plot(kind="bar", ax=ax)
    ax.set_title("Holdout observed state vs lambda-band detector state")
    ax.set_xlabel("Observed diagnostic state")
    ax.set_ylabel("Galaxy count")
    ax.legend(title="Detector state")
    fig.tight_layout()
    fig.savefig(out_dir / "holdout_detector_states_reproduced.png", dpi=180)
    plt.close(fig)

    # Figure 2: event-band sweep.
    fig, ax = plt.subplots(figsize=(7, 5))
    for model, g in sweep.groupby("model"):
        # Average across sigma for a compact view.
        avg = g.groupby("center", as_index=False)["hold_acc"].mean()
        ax.plot(avg["center"], avg["hold_acc"], marker="o", label=str(model))
    ax.set_title("Lambda event-band sweep: holdout phase match")
    ax.set_xlabel("event-band center")
    ax.set_ylabel("holdout phase match")
    ax.set_ylim(0, 1.05)
    ax.legend(title="model")
    fig.tight_layout()
    fig.savefig(out_dir / "lambda_event_band_sweep_reproduced.png", dpi=180)
    plt.close(fig)

    # Figure 3: velocity pass and phase match by split.
    split = df.groupby("split").agg(
        phase_match=("match_lambda_band", "mean"),
        velocity_pass=("velocity_pass_15pct", "mean"),
        combined=("combined_pass", "mean"),
    ).reset_index()
    fig, ax = plt.subplots(figsize=(7, 5))
    x = range(len(split))
    width = 0.25
    ax.bar([i - width for i in x], split["phase_match"], width=width, label="phase match")
    ax.bar(list(x), split["velocity_pass"], width=width, label="velocity pass")
    ax.bar([i + width for i in x], split["combined"], width=width, label="combined")
    ax.set_xticks(list(x))
    ax.set_xticklabels(split["split"])
    ax.set_ylim(0, 1.05)
    ax.set_title("Lambda-band detector summary by split")
    ax.set_ylabel("fraction")
    ax.legend()
    fig.tight_layout()
    fig.savefig(out_dir / "lambda_band_split_summary_reproduced.png", dpi=180)
    plt.close(fig)

    print(f"Wrote figures to {out_dir}")


if __name__ == "__main__":
    main()
