import json
from pathlib import Path

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


ROOT = Path(__file__).resolve().parents[1]


def main():
    cfg = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
    data = np.load(ROOT / cfg["paths"]["inference"])
    pred, target = data["prediction"], data["target"]
    if pred.shape != target.shape or pred.shape[1:] != (65, 68):
        raise ValueError(f"Expected prediction and target shaped [N,65,68], got {pred.shape} and {target.shape}")
    mass = np.concatenate((data["layer_mass"], data["layer_mass"]))
    mass = mass / mass.mean()
    error = pred - target
    mad = (np.abs(error) * mass[None, None]).mean(axis=(0, 2))
    bias = error.mean(axis=(0, 2))
    flat_target = target[:, 1:].astype(np.float64).reshape(target.shape[0], -1)
    flat_pred = pred[:, 1:].astype(np.float64).reshape(pred.shape[0], -1)
    source_r2 = []
    for truth, estimate in zip(flat_target, flat_pred):
        denom = np.sum((truth - truth.mean()) ** 2)
        source_r2.append(float(1.0 - np.sum((truth - estimate) ** 2) / max(denom, 1e-12)))
    dt = float(data["dt_seconds"])
    qt_pred, qt_target = pred[..., 34:], target[..., 34:]
    column_water_pred = (qt_pred * data["layer_mass"][None, None]).sum(-1)
    column_water_target = (qt_target * data["layer_mass"][None, None]).sum(-1)
    lhf_evap = data["surface"][:, :, 1] / 2.5e6
    adv_q = (0.5 * (data["advection"][:, :-1, 34:] + data["advection"][:, 1:, 34:]) *
             data["layer_mass"][None, None]).sum(-1)
    storage = np.diff(column_water_pred, axis=1) / dt
    precipitation = np.maximum(0.0, lhf_evap + adv_q - storage)
    residual = storage + precipitation - lhf_evap - adv_q
    metrics = {"mass_weighted_mad_per_step": mad.tolist(), "bias_per_step": bias.tolist(),
               "source_r2": source_r2, "mean_source_r2": float(np.mean(source_r2)),
               "water_budget": {"mean_precipitation_kg_m2_s": float(precipitation.mean()),
                                "mean_abs_residual_kg_m2_s": float(np.abs(residual).mean())},
                "per_step_length": int(len(mad)), "step_hours": cfg["data"]["step_hours"]}
    out = ROOT / cfg["paths"]["evaluation"]
    out.parent.mkdir(parents=True, exist_ok=True)
    out.write_text(json.dumps(metrics, indent=2))
    hours = data["lead_hours"]
    fig, axes = plt.subplots(2, 1, figsize=(9, 7), constrained_layout=True)
    axes[0].plot(hours, target[:, :, :34].mean((0, 2)), label="target sL")
    axes[0].plot(hours, pred[:, :, :34].mean((0, 2)), "--", label="predicted sL")
    ax2 = axes[0].twinx()
    ax2.plot(hours, target[:, :, 34:].mean((0, 2)), color="tab:green", label="target qT")
    ax2.plot(hours, pred[:, :, 34:].mean((0, 2)), "--", color="tab:red", label="predicted qT")
    axes[0].set(xlabel="lead time (h)", title="SCM state rollout")
    axes[0].legend(loc="upper left")
    ax2.legend(loc="upper right")
    axes[1].plot(hours[1:], precipitation.mean(0) * 86400.0, color="navy")
    axes[1].set(xlabel="lead time (h)", ylabel="mm day-1", title="Diagnosed precipitation")
    fig.savefig(out.parent / "state_precipitation_timeseries.png", dpi=150)
    plt.close(fig)
    print(f"saved {out}: per_step={len(mad)}, source_R2={metrics['mean_source_r2']:.4f}")


if __name__ == "__main__":
    main()