"""Tile diagnostics for analysis, forecasts, variable groups, and localization."""
import argparse
import json
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import torch
import sys, yaml
ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))

from model.fuxi_da import CompactForecastProxy, FuXiDA, make_sample

GROUPS = {"Z": slice(0, 13), "T": slice(13, 26), "U": slice(26, 39), "V": slice(39, 52), "R": slice(52, 65), "surface": slice(65, 70)}


def weighted_rmse(prediction, target, latitude):
    weight = torch.cos(torch.deg2rad(latitude)).clamp_min(0)
    weight = weight / weight.sum()
    return torch.sqrt(((prediction - target).square() * weight[None, :, None]).sum(dim=(-2, -1)) / prediction.shape[-1]).mean().item()


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", default="conf/config.yaml")
    parser.add_argument("--checkpoint")
    parser.add_argument("--output")
    parser.add_argument("--metrics")
    args = parser.parse_args()
    cfg = yaml.safe_load((ROOT / args.config).read_text())
    model = FuXiDA(cfg["model"]["base_channels"])
    checkpoint_path = Path(args.checkpoint) if args.checkpoint else ROOT / cfg["paths"]["checkpoint"]
    if checkpoint_path.exists():
        model.load_state_dict(torch.load(checkpoint_path, map_location="cpu", weights_only=True)["model"])
    model.eval()
    proxy = CompactForecastProxy()

    rows = []
    with torch.no_grad():
        for tile_id in cfg["data"]["tile_ids"]:
            sample = make_sample(tile_id, 20, cfg["data"]["tile_size"], cfg["data"]["missing_probability"])
            analysis = model(sample["background"][None], sample["obs"][None])[0]
            row = {"tile_id": tile_id, "origin": sample["origin"].tolist(), "background": weighted_rmse(sample["background"], sample["target"], sample["latitude"]), "correction": weighted_rmse(sample["correction"], sample["target"], sample["latitude"]), "analysis": weighted_rmse(analysis, sample["target"], sample["latitude"])}
            row["variable_groups"] = {name: weighted_rmse(analysis[index], sample["target"][index], sample["latitude"]) for name, index in GROUPS.items()}
            state = analysis
            row["forecast_steps"] = []
            for lead in range(cfg["train"]["forecast_steps"]):
                state = proxy(state[None])[0]
                row["forecast_steps"].append(weighted_rmse(state, sample["forecast_targets"][lead], sample["latitude"]))
            rows.append(row)

        sample = make_sample(cfg["data"]["tile_ids"][0], 21, cfg["data"]["tile_size"], 0.0)
        base = model(sample["background"][None], sample["obs"][None])[0]
        perturbed_obs = sample["obs"].clone()
        center = cfg["data"]["tile_size"] // 2
        perturbed_obs[9 - 8, center, center] += 1.0
        response = (model(sample["background"][None], perturbed_obs[None])[0] - base).square().sum(0)
        yy, xx = torch.meshgrid(torch.arange(cfg["data"]["tile_size"]), torch.arange(cfg["data"]["tile_size"]), indexing="ij")
        radius = torch.sqrt((yy - center).float().square() + (xx - center).float().square())
        total_energy = response.sum().clamp_min(1e-12)
        localization = {"perturbation": "AGRI channel 9 +1 K at tile center", "energy_weighted_radius_gridpoints": (response * radius).sum().div(total_energy).item(), "energy_within_radius_4": response[radius <= 4].sum().div(total_energy).item()}

    metrics = {"coverage_complete": False, "claim": "diagnostic metrics on selected aligned tiles; not complete global scores", "tiles": rows, "increment_localization": localization}
    metrics_path = Path(args.metrics) if args.metrics else ROOT / cfg["paths"]["evaluation"]
    metrics_path.parent.mkdir(parents=True, exist_ok=True); metrics_path.write_text(json.dumps(metrics, indent=2) + "\n")
    labels = ["background", "correction", "analysis"]
    values = [np.mean([row[label] for row in rows]) for label in labels]
    fig, axes = plt.subplots(1, 2, figsize=(10, 4))
    axes[0].bar(labels, values, color=["#596780", "#d69b36", "#287f71"])
    axes[0].set_ylabel("Latitude-weighted RMSE")
    axes[0].set_title("Selected aligned tiles (incomplete coverage)")
    group_values = [np.mean([row["variable_groups"][name] for row in rows]) for name in GROUPS]
    axes[1].bar(list(GROUPS), group_values, color="#287f71")
    axes[1].set_title("Analysis RMSE by variable group")
    fig.suptitle("FuXi-DA procedural protocol diagnostics")
    fig.tight_layout()
    output_path = Path(args.output) if args.output else ROOT / cfg["paths"]["figure"]
    output_path.parent.mkdir(parents=True, exist_ok=True); fig.savefig(output_path, dpi=160)
    plt.close(fig)
    print(json.dumps(metrics, indent=2))


if __name__ == "__main__":
    main()