#!/usr/bin/env python3
"""Plot first-order FA tangent-hierarchy predictor diagnostics."""

from __future__ import annotations

import argparse
from pathlib import Path

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


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--csv",
        type=Path,
        default=Path("outputs/fa_tangent_hierarchy_first_order_try/compressed_operator_rows.csv"),
    )
    parser.add_argument(
        "--outdir",
        type=Path,
        default=Path("outputs/fa_tangent_hierarchy_first_order_try"),
    )
    return parser.parse_args()


def metrics(df: pd.DataFrame) -> pd.DataFrame:
    rows = []
    for early, group in df.groupby("early_steps"):
        for label, column in [
            ("fixed K0", "fixed_gap"),
            ("linear velocity", "linear_gap"),
            ("retangent", "retangent_gap"),
            ("compressed", "compressed_gap"),
        ]:
            error = group[column] - group["empirical_gap"]
            rows.append(
                {
                    "early_steps": int(early),
                    "predictor": label,
                    "mae": float(error.abs().mean()),
                    "bias": float(error.mean()),
                    "rmse": float(np.sqrt(np.mean(np.square(error)))),
                    "corr": float(group[column].corr(group["empirical_gap"])),
                }
            )
    return pd.DataFrame(rows)


def plot_scatter(df: pd.DataFrame, outdir: Path) -> Path:
    early_values = sorted(df["early_steps"].unique())
    fig, axes = plt.subplots(1, len(early_values) + 1, figsize=(5.0 * (len(early_values) + 1), 4.6), dpi=170)

    panels = [("fixed K0", "fixed_gap", df)]
    for early in early_values:
        panels.append((f"linear velocity s={early}", "linear_gap", df[df["early_steps"] == early]))

    values = []
    for _title, column, group in panels:
        values.extend(group[column].tolist())
        values.extend(group["empirical_gap"].tolist())
    lo, hi = float(min(values)), float(max(values))
    pad = 0.04 * (hi - lo + 1e-12)

    for ax, (title, column, group) in zip(axes, panels):
        ax.scatter(group[column], group["empirical_gap"], s=26, alpha=0.72, color="#1f5a9d")
        ax.plot([lo - pad, hi + pad], [lo - pad, hi + pad], color="black", linewidth=1.0)
        err = group[column] - group["empirical_gap"]
        ax.set_title(f"{title}\nMAE={err.abs().mean():.4f}, bias={err.mean():+.4f}")
        ax.set_xlabel("predicted gap")
        ax.set_ylabel("empirical gap")
        ax.grid(alpha=0.18)
        ax.set_xlim(lo - pad, hi + pad)
        ax.set_ylim(lo - pad, hi + pad)

    fig.suptitle("First-order tangent-hierarchy predictor", y=1.02)
    fig.tight_layout()
    path = outdir / "first_order_prediction_scatter_by_early_step.png"
    fig.savefig(path, bbox_inches="tight")
    plt.close(fig)
    return path


def plot_metrics(metric_df: pd.DataFrame, outdir: Path) -> Path:
    fig, axes = plt.subplots(1, 3, figsize=(13.5, 4.4), dpi=170)
    predictors = ["fixed K0", "linear velocity", "retangent", "compressed"]
    colors = {
        "fixed K0": "#777777",
        "linear velocity": "#1f5a9d",
        "retangent": "#2b8a3e",
        "compressed": "#8a4fb0",
    }

    for ax, metric_name in zip(axes, ["mae", "bias", "corr"]):
        for predictor in predictors:
            sub = metric_df[metric_df["predictor"] == predictor].sort_values("early_steps")
            ax.plot(
                sub["early_steps"],
                sub[metric_name],
                marker="o",
                linewidth=1.8,
                color=colors[predictor],
                label=predictor,
            )
        if metric_name == "bias":
            ax.axhline(0.0, color="black", linewidth=1.0)
        ax.set_xlabel("early operator step s")
        ax.set_title(metric_name.upper())
        ax.grid(alpha=0.18)
    axes[0].set_ylabel("metric value")
    axes[0].legend(fontsize=8)
    fig.suptitle("Prediction metrics vs early operator horizon", y=1.02)
    fig.tight_layout()
    path = outdir / "first_order_error_metrics_by_early_step.png"
    fig.savefig(path, bbox_inches="tight")
    plt.close(fig)
    return path


def plot_distributions(df: pd.DataFrame, outdir: Path) -> Path:
    early_values = sorted(df["early_steps"].unique())
    fig, axes = plt.subplots(1, len(early_values), figsize=(5.0 * len(early_values), 4.4), dpi=170, squeeze=False)
    for ax, early in zip(axes.ravel(), early_values):
        group = df[df["early_steps"] == early]
        bins = np.linspace(
            min(group["empirical_gap"].min(), group["linear_gap"].min(), group["fixed_gap"].min()),
            max(group["empirical_gap"].max(), group["linear_gap"].max(), group["fixed_gap"].max()),
            18,
        )
        ax.hist(group["empirical_gap"], bins=bins, alpha=0.45, density=True, color="#c65f16", label="empirical")
        ax.hist(group["fixed_gap"], bins=bins, alpha=0.32, density=True, color="#777777", label="fixed K0")
        ax.hist(group["linear_gap"], bins=bins, alpha=0.36, density=True, color="#1f5a9d", label="linear velocity")
        ax.set_title(f"gap distribution, s={early}")
        ax.set_xlabel("FA/BP train gap")
        ax.set_ylabel("density")
        ax.legend(fontsize=8)
        ax.grid(alpha=0.18)
    fig.tight_layout()
    path = outdir / "first_order_gap_distribution_overlay.png"
    fig.savefig(path, bbox_inches="tight")
    plt.close(fig)
    return path


def main() -> None:
    args = parse_args()
    args.outdir.mkdir(parents=True, exist_ok=True)
    df = pd.read_csv(args.csv)
    metric_df = metrics(df)
    metric_path = args.outdir / "first_order_metrics.csv"
    metric_df.to_csv(metric_path, index=False)
    paths = [
        plot_scatter(df, args.outdir),
        plot_metrics(metric_df, args.outdir),
        plot_distributions(df, args.outdir),
    ]
    print(f"metrics: {metric_path}")
    for path in paths:
        print(f"plot: {path}")
    print(metric_df.to_string(index=False))


if __name__ == "__main__":
    main()