"""Run the paper-accurate no-training DataLoader diagnostic dry run."""

import json
import os
import sys
from pathlib import Path

import numpy as np
import torch
import yaml
from torch.utils.data import DataLoader, Dataset


ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
from model.pangu_icon_dke import PanguIconDKEDiagnostics


def load_protocol(config):
    source = np.load(ROOT / config["data"]["root"] / config["data"]["protocol_file"])
    protocol = {key: source[key].tolist() for key in source.files if key not in {"phase", "amplitude", "growth_rate", "spectral_phase", "spectral_amplitude", "base_wind_ms", "spectral_slope"}}
    coefficients = {key: source[key].tolist() for key in ("phase", "amplitude", "growth_rate", "spectral_phase", "spectral_amplitude", "base_wind_ms", "spectral_slope")}
    return protocol, coefficients


class ExperimentDataset(Dataset):
    def __init__(self, count): self.count = count
    def __len__(self): return self.count
    def __getitem__(self, index): return torch.tensor(index, dtype=torch.long)


def main():
    config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
    protocol, coefficients = load_protocol(config)
    if tuple(protocol["field_shape"]) != (5, 73, 1, 721, 1440) or tuple(protocol["spectral_shape"]) != (5, 73, 1, 259560, 2):
        raise ValueError("protocol dimensions do not match the public evaluation data")
    distributed = int(os.environ.get("WORLD_SIZE", "1")) > 1
    if distributed:
        torch.distributed.init_process_group("gloo")
    rank = torch.distributed.get_rank() if distributed else 0
    dataset = ExperimentDataset(int(protocol["field_shape"][0]))
    sampler = range(rank, len(dataset), torch.distributed.get_world_size()) if distributed else None
    loader = DataLoader(dataset, batch_size=1, sampler=sampler, shuffle=False)
    model = PanguIconDKEDiagnostics(protocol, coefficients)
    if sum(parameter.numel() for parameter in model.parameters()) != 0:
        raise RuntimeError("evaluation diagnostic must not have learnable parameters")
    dry_values = {}
    with torch.no_grad():
        for item in loader:
            experiment = int(item[0])
            u, v = model.fields.wind_chunk(experiment, 0, slice(0, 8))
            dry_values[experiment] = float(model(torch.from_numpy(u), torch.from_numpy(v)).mean())
    if distributed:
        gathered = [None] * torch.distributed.get_world_size() if rank == 0 else None
        torch.distributed.gather_object(dry_values, gathered, dst=0)
        if rank == 0:
            dry_values = {key: value for shard in gathered for key, value in shard.items()}
    if rank == 0:
        checkpoint_path = ROOT / config["paths"]["checkpoint"]
        metrics_path = ROOT / config["paths"]["training_metrics"]
        checkpoint_path.parent.mkdir(parents=True, exist_ok=True)
        metrics_path.parent.mkdir(parents=True, exist_ok=True)
        torch.save({"model": model.state_dict(), "model_config": config["generator"], "format_version": protocol["format_version"], "protocol": protocol, "coefficients": coefficients}, checkpoint_path)
        metrics_path.write_text(json.dumps({"training_required": False, "backward_executed": False, "learnable_parameters": 0, "diagnostic_dry_run": True, "experiment_dry_run": [dry_values[index] for index in range(len(dataset))]}, indent=2) + "\n")
        print(f"checkpoint={checkpoint_path.relative_to(ROOT)} training_required=false backward=false")
    if distributed:
        torch.distributed.destroy_process_group()


if __name__ == "__main__":
    main()