"""Evaluate 24-hour rainfall and produce per-step metrics and maps."""

import json
from pathlib import Path

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


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


def categorical(prediction, target, threshold):
    forecast, observed = prediction >= threshold, target >= threshold
    hits = np.logical_and(forecast, observed).sum(); false = np.logical_and(forecast, ~observed).sum()
    misses = np.logical_and(~forecast, observed).sum(); correct_negative = np.logical_and(~forecast, ~observed).sum()
    total = hits + false + misses + correct_negative; eps = 1e-8
    random_hits = (hits + misses) * (hits + false) / max(total, 1)
    return {
        "ACC": float((hits + correct_negative) / (total + eps)),
        "TS": float(hits / (hits + misses + false + eps)),
        "ETS": float((hits - random_hits) / (hits + misses + false - random_hits + eps)),
        "BS": float((hits + false) / (hits + misses + eps)),
        "FAR": float(false / (hits + false + eps)),
    }


def fss(prediction, target, threshold, window):
    forecast = torch.from_numpy((prediction >= threshold).astype(np.float32))[:, None]
    observed = torch.from_numpy((target >= threshold).astype(np.float32))[:, None]
    pf = torch.nn.functional.avg_pool2d(forecast, window, stride=1, padding=window // 2)
    po = torch.nn.functional.avg_pool2d(observed, window, stride=1, padding=window // 2)
    return float(1 - (pf - po).square().sum() / (pf.square().sum() + po.square().sum() + 1e-8))


def main():
    config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
    archive = np.load(ROOT / config["paths"]["inference"], allow_pickle=False)
    steps = int(config["data"]["time_steps"])
    lead_hours = [int(value) for value in archive["lead_hours"]]
    expected_3h = (steps, int(config["data"]["height"]), int(config["data"]["width"]))
    for key in ("raw_3h", "target_3h", "corrected_3h"):
        if archive[key].shape[1:] != expected_3h:
            raise ValueError(f"{key} must have per-day shape {expected_3h}")
    if lead_hours != config["data"]["lead_hours"]:
        raise ValueError("inference NPZ lead_hours do not match configuration")
    thresholds = [float(value) for value in config["evaluation"]["thresholds_mm_24h"]]
    windows = [int(value) for value in config["evaluation"]["fss_windows"]]
    metrics = {"units_24h": "mm/24h", "units_per_step": "mm/3h", "per_step_length": steps,
               "lead_hours": lead_hours, "raw_24h": {}, "corrected_24h": {}, "per_step": {}}
    for label in ("raw", "corrected"):
        prediction_24h = archive[f"{label}_24h"]
        metrics[f"{label}_24h"] = {str(int(t)): categorical(prediction_24h, archive["target_24h"], t) for t in thresholds}
        metrics[f"{label}_24h"]["FSS@50"] = {str(w): fss(prediction_24h, archive["target_24h"], 50, w) for w in windows}
        per_step = []
        for step in range(steps):
            item = {"step": step, "lead_hour": lead_hours[step]}
            item.update({str(int(t)): categorical(archive[f"{label}_3h"][:, step], archive["target_3h"][:, step], t) for t in thresholds})
            item["FSS@50"] = {str(w): fss(archive[f"{label}_3h"][:, step], archive["target_3h"][:, step], 50, w) for w in windows}
            per_step.append(item)
        metrics["per_step"][label] = per_step
    output = ROOT / config["paths"]["evaluation_dir"]
    output.mkdir(parents=True, exist_ok=True)
    (output / "metrics.json").write_text(json.dumps(metrics, indent=2) + "\n")
    figure, axes = plt.subplots(4, steps, figsize=(18, 9), constrained_layout=True)
    rows = (archive["raw_3h"][0], archive["target_3h"][0], archive["corrected_3h"][0], archive["corrected_3h"][0] - archive["target_3h"][0])
    labels = ("ECMWF", "CMPA target", "Corrected", "Error")
    vmax = max(float(rows[0].max()), float(rows[1].max()), 1)
    for row, (fields, label) in enumerate(zip(rows, labels)):
        for step in range(steps):
            cmap = "RdBu_r" if row == 3 else "Blues"
            limit = vmax if row < 3 else max(vmax / 2, 1)
            axes[row, step].imshow(fields[step], cmap=cmap, vmin=-limit if row == 3 else 0, vmax=limit)
            axes[row, step].axis("off")
            if row == 0: axes[row, step].set_title(f"+{archive['lead_hours'][step]} h")
            if step == 0: axes[row, step].set_ylabel(label)
    figure.savefig(output / "comparison.png", dpi=150); plt.close(figure)
    print(f"metrics={output / 'metrics.json'} figure={output / 'comparison.png'}")


if __name__ == "__main__":
    main()