"""Compute the paper's rainfall metrics and visualize the 90-minute forecast."""

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 to_rainfall(values, config):
    data = config["data"]
    radar_db = values * (float(data["radar_db_max"]) - float(data["radar_db_min"])) + float(data["radar_db_min"])
    return 10 ** ((radar_db - 10 * np.log10(float(data["zr_a"]))) / (10 * float(data["zr_b"])))


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"]
    predicted_rain, target_rain = to_rainfall(prediction, config), to_rainfall(target, config)
    threshold = float(config["data"]["rainfall_threshold_mm_h"])
    forecast, observed = predicted_rain >= threshold, target_rain >= threshold
    hits, misses = np.logical_and(forecast, observed).sum(), np.logical_and(~forecast, observed).sum()
    false_alarms = np.logical_and(forecast, ~observed).sum()
    eps = 1e-9
    frame_correlation, per_step = [], []
    for step in range(prediction.shape[1]):
        left, right = prediction[:, step].reshape(-1), target[:, step].reshape(-1)
        correlation = float(np.dot(left, right) / (np.sqrt(np.dot(left, left) * np.dot(right, right)) + eps))
        frame_correlation.append(correlation)
        step_forecast, step_observed = forecast[:, step], observed[:, step]
        step_hits = np.logical_and(step_forecast, step_observed).sum()
        step_misses = np.logical_and(~step_forecast, step_observed).sum()
        step_false_alarms = np.logical_and(step_forecast, ~step_observed).sum()
        per_step.append({
            "lead_minutes": int(data["forecast_lead_minutes"][step]),
            "rainfall_mse": float(np.mean((predicted_rain[:, step] - target_rain[:, step]) ** 2)),
            "csi": float(step_hits / (step_hits + step_misses + step_false_alarms + eps)),
            "far": float(step_false_alarms / (step_hits + step_false_alarms + eps)),
            "pod": float(step_hits / (step_hits + step_misses + eps)),
            "correlation": correlation,
        })
    metrics = {
        "samples": int(len(prediction)),
        "binary_cross_entropy": float(-(target * np.log(prediction.clip(1e-7, 1 - 1e-7)) +
                                         (1 - target) * np.log((1 - prediction).clip(1e-7, 1))).mean()),
        "rainfall_mse": float(np.mean((predicted_rain - target_rain) ** 2)),
        "csi": float(hits / (hits + misses + false_alarms + eps)),
        "far": float(false_alarms / (hits + false_alarms + eps)),
        "pod": float(hits / (hits + misses + eps)),
        "correlation": float(np.mean(frame_correlation)),
        "per_forecast_step": 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")
    steps = [0, 2, 5, 8, 11, 14]
    figure, axes = plt.subplots(3, len(steps), figsize=(15, 7))
    for column, step in enumerate(steps):
        axes[0, column].imshow(target[0, step, 0], cmap="turbo", vmin=0, vmax=1)
        axes[1, column].imshow(prediction[0, step, 0], cmap="turbo", vmin=0, vmax=1)
        axes[2, column].imshow(np.abs(target[0, step, 0] - prediction[0, step, 0]), cmap="magma", vmin=0, vmax=1)
        axes[0, column].set_title(f"+{(step + 1) * 6} min")
        for axis in axes[:, column]:
            axis.axis("off")
    figure.tight_layout()
    figure.savefig(output / "comparison.png", dpi=150)
    plt.close(figure)


if __name__ == "__main__":
    main()