"""Compute daily/annual detection metrics and a 7x7 occlusion trend map."""

import sys
from pathlib import Path

import numpy as np
import torch


ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
from model.precipdd import correlation, ensemble_predict, linear_trend, load_config, load_ensemble, write_json


def main():
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    config = load_config(ROOT / "conf/config.yaml")
    values = np.load(ROOT / config["paths"]["predictions"])
    prediction, target, year = values["prediction"], values["target"], values["year"].astype(float)
    if prediction.shape != target.shape or values["precipitation"].shape[1:] != (1, 55, 160):
        raise ValueError("inference output violates scalar target or [N,1,55,160] input contract")
    years = np.unique(year).astype(int)
    annual_prediction = np.array([prediction[year == current].mean() for current in years])
    annual_target = np.array([target[year == current].mean() for current in years])
    threshold = config["evaluation"]["emergence_threshold_c"]
    em_fraction = np.array([(prediction[year == current] > threshold).mean() for current in years])
    metrics = {
        "daily": {"correlation": correlation(target, prediction), "rmse_c": float(np.sqrt(np.mean((prediction - target) ** 2)))},
        "annual": {"correlation": correlation(annual_target, annual_prediction), "rmse_c": float(np.sqrt(np.mean((annual_prediction - annual_target) ** 2)))},
        "emergence": {"threshold_c": threshold, "fraction_all_days": float((prediction > threshold).mean()),
                      "fraction_trend_per_decade": linear_trend(em_fraction, years.astype(float))},
        "trend_c_per_decade": {"daily_prediction": linear_trend(prediction, year), "annual_prediction": linear_trend(annual_prediction, years.astype(float)),
                               "annual_target": linear_trend(annual_target, years.astype(float))},
        "samples": {"daily": len(prediction), "years": len(years)}
    }

    device = torch.device("cuda" if torch.cuda.is_available() and config["runtime"]["device"] != "cpu" else "cpu")
    models, _ = load_ensemble(ROOT / config["paths"]["checkpoint"], device)
    count = min(config["evaluation"]["occlusion_max_days"], len(prediction))
    indices = np.linspace(0, len(prediction) - 1, count, dtype=int)
    selected = torch.from_numpy(values["precipitation"][indices]).float().to(device)
    baseline = ensemble_predict(models, selected).cpu().numpy()
    selected_year = year[indices]
    patch, stride = config["evaluation"]["occlusion_patch"], config["evaluation"]["occlusion_stride"]
    half = patch // 2
    sensitivity = np.zeros((55, 160), dtype=np.float32)
    for lat_start in range(0, 55, stride):
        for lon_start in range(0, 160, stride):
            masked = selected.clone()
            lat_stop, lon_stop = min(lat_start + patch, 55), min(lon_start + patch, 160)
            masked[:, :, lat_start:lat_stop, lon_start:lon_stop] = 0.0
            delta = baseline - ensemble_predict(models, masked).cpu().numpy()
            score = linear_trend(delta, selected_year) if np.ptp(selected_year) else float(delta.mean())
            sensitivity[lat_start:min(lat_start + stride, 55), lon_start:min(lon_start + stride, 160)] = score
    metrics["occlusion"] = {"patch": [patch, patch], "map_shape": [55, 160], "stride": stride, "sampled_days": count,
                            "quantity": "AGMT occlusion-sensitivity trend in degC per decade"}
    write_json(ROOT / config["paths"]["evaluation_metrics"], metrics)

    fig, axes = plt.subplots(3, 1, figsize=(11, 11), constrained_layout=True)
    axes[0].plot(years, annual_target, color="#202020", label="target AGMT")
    axes[0].plot(years, annual_prediction, color="#c84c32", label="DD estimate")
    axes[0].axhline(threshold, color="#777777", linestyle="--", label="0.42 C EM threshold")
    axes[0].set(ylabel="AGMT anomaly (C)", title="Annual mean of daily estimates")
    axes[0].legend(ncol=3)
    axes[1].plot(years, em_fraction, color="#196f82")
    axes[1].set(xlabel="Year", ylabel="Fraction", ylim=(-0.03, 1.03), title="Emergence days (estimated AGMT > 0.42 C)")
    limit = float(np.max(np.abs(sensitivity))) or 1e-6
    image = axes[2].imshow(sensitivity, origin="lower", aspect="auto", extent=(0, 400, values["latitude"][0], values["latitude"][-1]),
                           cmap="RdBu_r", vmin=-limit, vmax=limit)
    axes[2].set(xlabel="Longitude (degrees E, extended)", ylabel="Latitude", title="7x7 occlusion-sensitivity trend")
    fig.colorbar(image, ax=axes[2], label="C decade-1")
    figure = ROOT / config["paths"]["comparison_figure"]
    figure.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(figure, dpi=160)
    plt.close(fig)
    print(f"metrics={config['paths']['evaluation_metrics']} figure={config['paths']['comparison_figure']}")


if __name__ == "__main__":
    main()