"""Evaluate precipitation nowcasts with regression and threshold metrics."""

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():
    config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
    data = np.load(ROOT / config["paths"]["inference_dir"] / "predictions.npz")
    prediction, target = data["predictions"], data["targets"]
    threshold = 0.1
    observed, forecast = target >= threshold, prediction >= threshold
    tp = np.logical_and(observed, forecast).sum()
    fp = np.logical_and(~observed, forecast).sum()
    fn = np.logical_and(observed, ~forecast).sum()
    tn = np.logical_and(~observed, ~forecast).sum()
    eps = 1e-8
    metrics = {
        "samples": int(len(prediction)),
        "mse": float(np.mean((prediction - target) ** 2)),
        "mae": float(np.mean(np.abs(prediction - target))),
        "precision": float(tp / (tp + fp + eps)),
        "recall": float(tp / (tp + fn + eps)),
        "f1": float(2 * tp / (2 * tp + fp + fn + eps)),
        "csi": float(tp / (tp + fp + fn + eps)),
        "far": float(fp / (tp + fp + eps)),
        "hss": float(2 * (tp * tn - fn * fp) / ((tp + fn) * (fn + tn) + (tp + fp) * (fp + tn) + eps)),
    }
    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(3, 6, figsize=(15, 7))
    for step in range(6):
        axes[0, step].imshow(target[0, step], cmap="Blues", vmin=0, vmax=1)
        axes[1, step].imshow(prediction[0, step], cmap="Blues", vmin=0, vmax=1)
        axes[2, step].imshow(np.abs(target[0, step] - prediction[0, step]), cmap="magma", vmin=0, vmax=1)
        axes[0, step].set_title(f"+{(step + 1) * 5} min")
        for axis in axes[:, step]:
            axis.axis("off")
    figure.tight_layout()
    figure.savefig(output / "comparison.png", dpi=150)
    plt.close(figure)


if __name__ == "__main__":
    main()