"""Generate all PoC figures from run artifacts."""
import argparse
import glob
import json
import os

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


def soh_predictions(run, out):
    files = sorted(glob.glob(os.path.join(run, "predictions", "*_pinn.csv")))
    fig, axes = plt.subplots(2, 2, figsize=(11, 7), sharey=True)
    for ax, f in zip(axes.ravel(), files):
        name = os.path.basename(f).replace("_pinn.csv", "")
        dp = pd.read_csv(f)
        db = pd.read_csv(f.replace("_pinn", "_baseline"))
        ax.plot(dp["cycle"], dp["soh_true"], "k.", ms=2, label="measured")
        ax.plot(dp["cycle"], dp["soh_pinn"], "C0-", lw=1.5, label="PINN")
        ax.plot(db["cycle"], db["soh_baseline"], "C1--", lw=1.2, label="MLP baseline")
        ax.set_title(name)
        ax.set_xlabel("cycle")
        ax.set_ylabel("SOH")
        ax.grid(alpha=0.3)
    axes[0, 0].legend()
    fig.suptitle("SOH estimation on held-out batteries")
    fig.tight_layout()
    fig.savefig(os.path.join(out, "soh_predictions.png"), dpi=150)


def training_curves(run, out):
    fig, axes = plt.subplots(1, 2, figsize=(11, 4))
    h = pd.read_csv(os.path.join(run, "history_pinn.csv"))
    axes[0].semilogy(h["step"], h["data"], label="data loss")
    axes[0].semilogy(h["step"], h["phys"], label="physics loss")
    axes[0].semilogy(h["step"], h["mono"], label="monotonicity loss")
    axes[0].set_title("PINN loss components")
    axes[0].set_xlabel("step")
    axes[0].legend()
    axes[0].grid(alpha=0.3)
    hb = pd.read_csv(os.path.join(run, "history_baseline.csv"))
    axes[1].semilogy(h["step"], h["val_mse"], label="PINN")
    axes[1].semilogy(hb["step"], hb["val_mse"], label="baseline")
    axes[1].set_title("Validation MSE")
    axes[1].set_xlabel("step")
    axes[1].legend()
    axes[1].grid(alpha=0.3)
    fig.tight_layout()
    fig.savefig(os.path.join(out, "training_curves.png"), dpi=150)


def diagnostics_plot(run, out):
    d = pd.read_csv(os.path.join(run, "diagnostics_trace.csv"))
    with open(os.path.join(run, "vhm_results.json")) as f:
        res = json.load(f)["diagnostics"]
    fig, axes = plt.subplots(2, 1, figsize=(10, 7), sharex=True)
    axes[0].plot(d["cycle"], d["soh_healthy"], "k-", lw=1, label="healthy (measured)")
    axes[0].plot(d["cycle"], d["soh_faulted"], "C3-", lw=1.2, label="faulted (injected)")
    axes[0].plot(d["cycle"], d["soh_model"], "C0--", lw=1.2, label="model (healthy physics)")
    axes[0].axvline(res["fault_onset_cycle"], color="C3", ls=":", label="fault onset")
    if res["alarm_cycle"]:
        axes[0].axvline(res["alarm_cycle"], color="C2", ls="--", label="alarm")
    axes[0].set_ylabel("SOH")
    axes[0].legend(fontsize=8)
    axes[0].grid(alpha=0.3)
    axes[0].set_title(f"Fault detection on {res['battery']} "
                      f"(delay {res['detection_delay_cycles']} cycles)")
    axes[1].plot(d["cycle"], d["cusum"], "C0-")
    axes[1].axhline(res["cusum_h"], color="r", ls="--", label=f"threshold h={res['cusum_h']}")
    axes[1].axvline(res["fault_onset_cycle"], color="C3", ls=":")
    if res["alarm_cycle"]:
        axes[1].axvline(res["alarm_cycle"], color="C2", ls="--")
    axes[1].set_ylabel("CUSUM statistic")
    axes[1].set_xlabel("cycle")
    axes[1].legend(fontsize=8)
    axes[1].grid(alpha=0.3)
    fig.tight_layout()
    fig.savefig(os.path.join(out, "diagnostics_fault_detection.png"), dpi=150)


def prognostics_plot(run, out):
    prog = pd.read_csv(os.path.join(run, "prognostics.csv"))
    trajs = np.load(os.path.join(run, "forecast_trajs.npz"))
    bat = prog["battery"].iloc[0]
    meas = pd.read_csv(os.path.join(run, "predictions", f"{bat}_pinn.csv"))
    fig, axes = plt.subplots(1, 2, figsize=(12, 4.5))
    axes[0].plot(meas["cycle"], meas["soh_true"], "k.", ms=2, label="measured")
    for frac, c in zip((0.3, 0.5, 0.7), ("C0", "C1", "C2")):
        tr = trajs[f"{bat}@{frac}"]
        axes[0].plot(tr[:, 0], tr[:, 1], c + "-", lw=1.5,
                     label=f"forecast from {int(frac*100)}% life")
    axes[0].axhline(0.80, color="r", ls="--", lw=1, label="EOL threshold")
    axes[0].set_title(f"Degradation forecasts, {bat}")
    axes[0].set_xlabel("cycle")
    axes[0].set_ylabel("SOH")
    axes[0].legend(fontsize=8)
    axes[0].grid(alpha=0.3)

    w = 0.35
    idx = np.arange(len(prog))
    axes[1].bar(idx - w / 2, prog["true_RUL"], w, label="true RUL", color="k", alpha=0.7)
    axes[1].bar(idx + w / 2, prog["pred_RUL"], w, label="predicted RUL", color="C0")
    axes[1].set_xticks(idx)
    axes[1].set_xticklabels([f"{b.split('_')[0]}-{b.split('-')[-1]}\n@{int(f*100)}%"
                             for b, f in zip(prog["battery"], prog["life_fraction"])],
                            fontsize=7)
    axes[1].set_ylabel("RUL (cycles)")
    axes[1].set_title("RUL: predicted vs true (all test batteries)")
    axes[1].legend(fontsize=8)
    axes[1].grid(alpha=0.3, axis="y")
    fig.tight_layout()
    fig.savefig(os.path.join(out, "prognostics_rul.png"), dpi=150)


if __name__ == "__main__":
    p = argparse.ArgumentParser()
    p.add_argument("--run", default="../runs/poc")
    p.add_argument("--out", default="../runs/poc/figures")
    args = p.parse_args()
    os.makedirs(args.out, exist_ok=True)
    soh_predictions(args.run, args.out)
    training_curves(args.run, args.out)
    diagnostics_plot(args.run, args.out)
    prognostics_plot(args.run, args.out)
    print("figures written to", args.out)
