"""Evaluate DINCAE, residual calibration, and a rank-13 DINEOF-like baseline."""

import json
from pathlib import Path
import sys

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import yaml

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))

from model.dincae import reconstruction_metrics


def dineof_like(observed, ocean_mask, requested_rank):
    matrix = observed.reshape(observed.shape[0], -1).astype(np.float64)
    sea = ocean_mask.ravel()
    values = matrix[:, sea]
    missing = ~np.isfinite(values)
    column_mean = np.nanmean(values, axis=0)
    column_mean = np.nan_to_num(column_mean, nan=float(np.nanmean(values)))
    filled = np.where(missing, column_mean[None], values)
    rank = min(int(requested_rank), min(filled.shape) - 1)
    for _ in range(4):
        mean = filled.mean(axis=0, keepdims=True)
        u, singular, vt = np.linalg.svd(filled - mean, full_matrices=False)
        reconstructed = (u[:, :rank] * singular[:rank]) @ vt[:rank] + mean
        filled[missing] = reconstructed[missing]
    result = np.full_like(matrix, np.nan)
    result[:, sea] = filled
    return result.reshape(observed.shape).astype(np.float32), rank


def main():
    config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
    with np.load(ROOT / config["paths"]["inference"]) as loaded:
        pred = {key: loaded[key] for key in loaded.files}
    with np.load(ROOT / config["data"]["root"] / config["data"]["file"]) as loaded:
        full = {key: loaded[key] for key in loaded.files}
    mask = pred["missing_mask"]
    metrics = {"dincae": reconstruction_metrics(pred["target"], pred["prediction"], mask)}
    std_residual = (pred["prediction"] - pred["target"]) / np.sqrt(np.maximum(pred["variance"], 1e-12))
    residual = std_residual[mask & np.isfinite(std_residual)]
    bins = int(config["evaluation"]["calibration_bins"])
    edges = np.linspace(-4, 4, bins + 1)
    counts, _ = np.histogram(residual, bins=edges)
    metrics["standardized_residual_calibration"] = {
        "mean": float(residual.mean()), "std": float(residual.std()),
        "within_1sigma": float(np.mean(np.abs(residual) <= 1)),
        "within_2sigma": float(np.mean(np.abs(residual) <= 2)),
        "histogram_edges": edges.tolist(), "histogram_counts": counts.tolist(),
    }
    baseline_all, effective_rank = dineof_like(full["observed_anomaly"] + full["climatology"],
                                                full["ocean_mask"], config["evaluation"]["dineof_rank"])
    start = min(int(config["data"]["train_samples"]), len(full["timestamps"]) - 1)
    baseline = baseline_all[start:]
    metrics["dineof_like"] = reconstruction_metrics(pred["target"], baseline, mask)
    metrics["dineof_like"]["requested_rank"] = int(config["evaluation"]["dineof_rank"])
    metrics["dineof_like"]["effective_rank"] = effective_rank
    out_dir = ROOT / config["paths"]["evaluation_dir"]
    out_dir.mkdir(parents=True, exist_ok=True)
    (out_dir / "metrics.json").write_text(json.dumps(metrics, indent=2))
    index = 0
    fields = (pred["observed"][index], pred["target"][index], pred["prediction"][index],
              baseline[index], pred["prediction"][index] - pred["target"][index])
    titles = ("Cloudy AVHRR-like input", "Complete target", "DINCAE", "DINEOF-like rank 13", "DINCAE error")
    fig, axes = plt.subplots(1, 5, figsize=(17, 3.6), constrained_layout=True)
    lower, upper = np.nanpercentile(pred["target"][index], (2, 98))
    for axis, field, title in zip(axes, fields, titles):
        if title.endswith("error"):
            limit = max(float(np.nanpercentile(np.abs(field), 98)), 1e-3)
            image = axis.imshow(field, cmap="RdBu_r", vmin=-limit, vmax=limit)
        else:
            image = axis.imshow(field, cmap="turbo", vmin=lower, vmax=upper)
        axis.set_title(title, fontsize=9)
        axis.set_axis_off()
        fig.colorbar(image, ax=axis, fraction=0.046)
    fig.savefig(out_dir / "comparison.png", dpi=150)
    plt.close(fig)
    print(json.dumps(metrics, indent=2))
    print(f"comparison={out_dir / 'comparison.png'}")


if __name__ == "__main__":
    main()