"""Compute paper metrics, Theil-Sen trends, and the comparison figure."""

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 metrics(observed, predicted):
    error = predicted - observed
    denominator = np.maximum(np.sum(observed, axis=0), 1e-12)
    r2_denominator = np.maximum(np.sum((observed - observed.mean(0)) ** 2, axis=0), 1e-12)
    ioa_denominator = np.maximum(np.sum((np.abs(predicted - observed.mean(0)) +
                                         np.abs(observed - observed.mean(0))) ** 2, axis=0), 1e-12)
    return {"RMSE": np.sqrt(np.mean(error ** 2, axis=0)),
            "r2": 1 - np.sum(error ** 2, axis=0) / r2_denominator,
            "FAC2": np.mean((predicted / np.maximum(observed, 1e-12) >= .5) &
                            (predicted / np.maximum(observed, 1e-12) <= 2), axis=0),
            "MB": np.mean(error, axis=0), "MGE": np.mean(np.abs(error), axis=0),
            "NMB": np.sum(error, axis=0) / denominator,
            "NMGE": np.sum(np.abs(error), axis=0) / denominator,
            "COE": 1 - np.sum(np.abs(error), axis=0) /
                   np.maximum(np.sum(np.abs(observed - observed.mean(0)), axis=0), 1e-12),
            "IOA": 1 - np.sum(error ** 2, axis=0) / ioa_denominator}


def theil_sen(x, y, max_pairs=200000, seed=0):
    """Median pairwise slope, with deterministic subsampling for large series."""
    x, y = np.asarray(x, float), np.asarray(y, float)
    total = len(x) * (len(x) - 1) // 2
    if total <= max_pairs:
        slopes = [(y[j] - y[i]) / (x[j] - x[i]) for i in range(len(x) - 1)
                  for j in range(i + 1, len(x)) if x[j] != x[i]]
    else:
        rng = np.random.default_rng(seed)
        left = rng.integers(0, len(x) - 1, max_pairs)
        right = rng.integers(left + 1, len(x), max_pairs)
        valid = x[right] != x[left]
        slopes = (y[right[valid]] - y[left[valid]]) / (x[right[valid]] - x[left[valid]])
    slope = float(np.median(slopes))
    return slope, float(np.median(y - slope * x))


def main():
    config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
    data = np.load(ROOT / config["paths"]["inference"])
    names = [str(value) for value in data["pollutant_names"]]
    scores = metrics(data["observed"], data["predicted"])
    report = {"metrics": {name: {metric: float(values[i]) for metric, values in scores.items()}
                           for i, name in enumerate(names)}, "theil_sen_per_year": {}}
    trends = []
    for i, name in enumerate(names):
        raw = theil_sen(data["ttrend"], data["observed"][:, i], seed=i)
        normalized = theil_sen(data["ttrend"], data["normalized"][:, i], seed=100 + i)
        report["theil_sen_per_year"][name] = {"observed_slope": raw[0],
                                                       "normalized_slope": normalized[0]}
        trends.append((raw[0], normalized[0]))
    output = ROOT / config["paths"]["evaluation_dir"]
    output.mkdir(parents=True, exist_ok=True)
    (output / "metrics.json").write_text(json.dumps(report, indent=2) + "\n")
    figure, axes = plt.subplots(2, 1, figsize=(11, 8), gridspec_kw={"height_ratios": [2, 1]})
    order = np.argsort(data["ttrend"])
    sample = order[::max(1, len(order) // 600)]
    axes[0].plot(data["ttrend"][sample], data["observed"][sample, 0], ".", alpha=.25, label="observed PM2.5")
    axes[0].plot(data["ttrend"][sample], data["normalized"][sample, 0], ".", alpha=.4, label="weather-normalized PM2.5")
    axes[0].set(ylabel="Concentration (ug m-3)", title="Observed and meteorologically normalized test samples")
    axes[0].legend()
    positions = np.arange(len(names)); width = .36
    axes[1].bar(positions - width / 2, [v[0] for v in trends], width, label="observed")
    axes[1].bar(positions + width / 2, [v[1] for v in trends], width, label="normalized")
    axes[1].axhline(0, color="black", linewidth=.7); axes[1].set_xticks(positions, names)
    axes[1].set(ylabel="Theil-Sen slope per year", title="Robust concentration trends")
    axes[1].legend(); figure.tight_layout(); figure.savefig(output / "comparison.png", dpi=160)
    plt.close(figure)
    print(f"evaluation={output.relative_to(ROOT)} pollutants={len(names)}")


if __name__ == "__main__":
    main()