"""Evaluate missing-region reconstruction and create a comparison figure."""

from pathlib import Path
import argparse
import json
import numpy as np
import torch
import torch.nn.functional as F
import yaml


def correlation(a, b):
    if a.size < 2 or np.std(a) == 0 or np.std(b) == 0:
        return 0.0
    return float(np.corrcoef(a, b)[0, 1])


def rankdata(values):
    order = np.argsort(values, kind="mergesort")
    ranks = np.empty(len(values), dtype=np.float64)
    ranks[order] = np.arange(len(values), dtype=np.float64)
    unique, inverse, counts = np.unique(values, return_inverse=True, return_counts=True)
    del unique
    for group, count in enumerate(counts):
        if count > 1:
            positions = np.flatnonzero(inverse == group)
            ranks[positions] = ranks[positions].mean()
    return ranks


def main():
    parser = argparse.ArgumentParser()
    root = Path(__file__).resolve().parents[1]
    parser.add_argument("--config", type=Path, default=root / "conf/config.yaml")
    args = parser.parse_args()
    config_path = args.config if args.config.is_absolute() else root / args.config
    with open(config_path, encoding="utf-8") as handle:
        cfg = yaml.safe_load(handle)
    archive = np.load(root / cfg["output_dir"] / "predictions.npz")
    pred, target = archive["prediction"], archive["target"]
    missing = archive["europe_mask"][None, None] * (1 - archive["valid_mask"])
    selected = missing.astype(bool)
    error = pred[selected] - target[selected]
    sample_spearman = []
    for i in range(len(pred)):
        mask = selected[i, 0]
        sample_spearman.append(correlation(rankdata(pred[i, 0][mask]), rankdata(target[i, 0][mask])))
    kernel = torch.ones(1, 1, 3, 3) / 8
    kernel[0, 0, 1, 1] = 0
    pred_neighbor = F.conv2d(torch.from_numpy(pred), kernel, padding=1).numpy()
    target_neighbor = F.conv2d(torch.from_numpy(target), kernel, padding=1).numpy()
    metrics = {
        "missing_rmse": float(np.sqrt(np.mean(error ** 2))),
        "sample_spearman_mean": float(np.mean(sample_spearman)),
        "sample_spearman": [float(x) for x in sample_spearman],
        "missing_bias": float(np.mean(error)),
        "neighborhood_spatial_correlation_prediction": correlation(pred[selected], pred_neighbor[selected]),
        "neighborhood_spatial_correlation_target": correlation(target[selected], target_neighbor[selected]),
        "missing_points": int(selected.sum()),
    }
    output = root / cfg["evaluation_dir"]
    output.mkdir(parents=True, exist_ok=True)
    (output / "metrics.json").write_text(json.dumps(metrics, indent=2) + "\n")
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    fig, axes = plt.subplots(1, 3, figsize=(13, 4), constrained_layout=True)
    sample = 0
    fields = [archive["observed"][sample, 0], target[sample, 0], pred[sample, 0]]
    titles = ["Irregular observations", "Synthetic truth", "CRAI reconstruction"]
    for axis, field, title in zip(axes, fields, titles):
        image = axis.imshow(np.where(archive["europe_mask"] > 0, field, np.nan), origin="lower", vmin=0, vmax=100, cmap="RdYlBu_r")
        axis.set_title(title); axis.set_axis_off()
    fig.colorbar(image, ax=axes, label="Extreme index (%)", shrink=0.8)
    fig.suptitle(f"Missing RMSE={metrics['missing_rmse']:.3f} | Spearman={metrics['sample_spearman_mean']:.3f} | Bias={metrics['missing_bias']:.3f}")
    fig.savefig(output / "comparison.png", dpi=160); plt.close(fig)
    print(json.dumps(metrics))


if __name__ == "__main__":
    main()