from __future__ import annotations

import argparse
import csv
import json
from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np

PROJECT_ROOT = Path(__file__).resolve().parents[1]
import sys

if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from scripts.common import load_config, write_json  # noqa: E402


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Create CFO Burgers result visualizations")
    parser.add_argument("--config", default=str(PROJECT_ROOT / "config" / "config.yaml"))
    return parser.parse_args()


def plot_history(history_path: Path, out_path: Path) -> None:
    steps, losses, val_steps, vals = [], [], [], []
    with history_path.open("r", encoding="utf-8") as f:
        for row in csv.DictReader(f):
            steps.append(int(row["step"]))
            losses.append(float(row["loss"]))
            if row["val_rel_l2"]:
                val_steps.append(int(row["step"]))
                vals.append(float(row["val_rel_l2"]))
    fig, ax1 = plt.subplots(figsize=(8, 4.5))
    ax1.plot(steps, losses, color="#1f77b4", linewidth=1.0, label="train loss")
    ax1.set_yscale("log")
    ax1.set_xlabel("training step")
    ax1.set_ylabel("flow matching MSE")
    if vals:
        ax2 = ax1.twinx()
        ax2.plot(val_steps, vals, color="#d62728", marker="o", linewidth=1.2, label="val rel L2")
        ax2.set_ylabel("validation trajectory relative L2")
    fig.tight_layout()
    fig.savefig(out_path, dpi=180)
    plt.close(fig)


def plot_worst_case(npz_path: Path, metrics_path: Path, out_path: Path) -> None:
    data = np.load(npz_path)
    pred = data["prediction"]
    target = data["target"]
    metrics = json.loads(metrics_path.read_text(encoding="utf-8"))
    idx = int(metrics["worst_sample_index"])
    times = [25, 50, 75, 100]
    fig, axes = plt.subplots(3, len(times), figsize=(13, 6), sharex=True, sharey=True)
    for col, t in enumerate(times):
        axes[0, col].plot(target[idx, t - 1], color="#222222", linewidth=1.2)
        axes[0, col].set_title(f"target t={t}")
        axes[1, col].plot(pred[idx, t - 1], color="#0072B2", linewidth=1.2)
        axes[1, col].set_title(f"prediction t={t}")
        axes[2, col].plot(pred[idx, t - 1] - target[idx, t - 1], color="#D55E00", linewidth=1.0)
        axes[2, col].set_title("error")
    for ax in axes.flat:
        ax.grid(True, linewidth=0.3, alpha=0.4)
    fig.suptitle(f"Worst test trajectory index {idx}, rel L2={metrics['per_sample_relative_l2_max']:.6f}")
    fig.tight_layout()
    fig.savefig(out_path, dpi=180)
    plt.close(fig)


def plot_spacetime(npz_path: Path, metrics_path: Path, out_path: Path) -> None:
    data = np.load(npz_path)
    pred = data["prediction"]
    target = data["target"]
    metrics = json.loads(metrics_path.read_text(encoding="utf-8"))
    idx = int(metrics["worst_sample_index"])
    err = pred[idx] - target[idx]
    vmax = max(abs(target[idx]).max(), abs(pred[idx]).max())
    emax = max(abs(err).max(), 1e-8)
    fig, axes = plt.subplots(1, 3, figsize=(12, 4), constrained_layout=True)
    ims = [
        axes[0].imshow(target[idx], aspect="auto", origin="lower", cmap="viridis", vmin=-vmax, vmax=vmax),
        axes[1].imshow(pred[idx], aspect="auto", origin="lower", cmap="viridis", vmin=-vmax, vmax=vmax),
        axes[2].imshow(err, aspect="auto", origin="lower", cmap="coolwarm", vmin=-emax, vmax=emax),
    ]
    for ax, title in zip(axes, ["target", "prediction", "prediction - target"]):
        ax.set_title(title)
        ax.set_xlabel("space index")
        ax.set_ylabel("time index")
    for ax, im in zip(axes, ims):
        fig.colorbar(im, ax=ax, shrink=0.8)
    fig.savefig(out_path, dpi=180)
    plt.close(fig)


def main() -> None:
    args = parse_args()
    cfg = load_config(args.config)
    results = Path(cfg["paths"]["results_dir"])
    figures = results / "figures"
    figures.mkdir(parents=True, exist_ok=True)
    history = results / "train_history.csv"
    pred_npz = results / "test_predictions.npz"
    metrics_json = results / "test_metrics.json"
    outputs = {}
    if history.exists():
        outputs["training_curve"] = str(figures / "training_curve.png")
        plot_history(history, Path(outputs["training_curve"]))
    if pred_npz.exists() and metrics_json.exists():
        outputs["worst_case_profiles"] = str(figures / "worst_case_profiles.png")
        outputs["worst_case_spacetime"] = str(figures / "worst_case_spacetime.png")
        plot_worst_case(pred_npz, metrics_json, Path(outputs["worst_case_profiles"]))
        plot_spacetime(pred_npz, metrics_json, Path(outputs["worst_case_spacetime"]))
    report = {
        "status": "success",
        "figures": outputs,
        "metrics": json.loads(metrics_json.read_text(encoding="utf-8")) if metrics_json.exists() else None,
    }
    write_json(results / "result_report.json", report)
    print(report, flush=True)


if __name__ == "__main__":
    main()