"""Train independent PPNN replicas with Gaussian CRPS and optional rank allocation."""

import json
import os
import sys
from pathlib import Path

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

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
from model.ppnn import FORMAT_VERSION, PPNN, ensemble_features, gaussian_crps


def load_data(config):
    data = np.load(ROOT / config["data"]["file"])
    expected = (len(data["target"]), 50, 18)
    if str(data["format_version"]) != FORMAT_VERSION or data["ensemble"].shape != expected:
        raise ValueError(f"dataset version/shape mismatch: expected {FORMAT_VERSION} and {expected}")
    if len(data["station_id"]) != 537 or int(data["lead_hours"]) != 48:
        raise ValueError("dataset must retain 537 stations and 48h lead")
    arrays = [data["ensemble"], data["auxiliary"], data["target"]]
    if not all(np.isfinite(x).all() for x in arrays):
        raise ValueError("dataset contains non-finite values")
    return data


def main():
    config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
    data = load_data(config)
    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
    world_size = torch.distributed.get_world_size() if distributed else 1
    device_name = config["runtime"]["device"]
    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
    enough_accelerators = torch.cuda.is_available() and torch.cuda.device_count() >= world_size
    device = torch.device("cuda", local_rank) if device_name == "auto" and enough_accelerators else torch.device("cpu")
    ensemble = torch.from_numpy(data["ensemble"])
    auxiliary = torch.from_numpy(data["auxiliary"])
    station = torch.from_numpy(data["station_index"])
    target = torch.from_numpy(data["target"])
    continuous_raw = ensemble_features(ensemble, auxiliary)
    feature_mean = continuous_raw.mean(0)
    feature_std = continuous_raw.std(0, correction=1).clamp_min(1e-6)
    target_mean, target_std = target.mean(), target.std(correction=1).clamp_min(1e-6)
    continuous = (continuous_raw - feature_mean) / feature_std
    target_scaled = (target - target_mean) / target_std
    dataset = TensorDataset(continuous, station, target_scaled)
    local = []
    replicas = int(config["train"]["replicas"])
    for replica in range(replicas):
        if replica % world_size != rank:
            continue
        seed = int(config["seed"]) + replica
        torch.manual_seed(seed)
        generator = torch.Generator().manual_seed(seed)
        loader = DataLoader(dataset, batch_size=int(config["train"]["batch_size"]), shuffle=True, generator=generator)
        model = PPNN(int(config["model"]["hidden_size"]), eps=float(config["model"]["sigma_epsilon"])).to(device)
        optimizer = torch.optim.Adam(model.parameters(), lr=float(config["train"]["learning_rate"]))
        history = []
        for epoch in range(int(config["train"]["epochs"])):
            total = 0.0
            for x, s, y in loader:
                mu, sigma = model(x.to(device), s.to(device))
                loss = gaussian_crps(mu, sigma, y.to(device)).mean()
                if not torch.isfinite(loss):
                    raise RuntimeError("non-finite Gaussian CRPS loss")
                optimizer.zero_grad(set_to_none=True)
                loss.backward()
                optimizer.step()
                total += float(loss.detach()) * len(y)
            history.append(total / len(dataset))
        local.append({"replica": replica, "state_dict": {k: v.cpu() for k, v in model.state_dict().items()}, "loss": history})
        print(f"rank={rank} replica={replica} final_scaled_crps={history[-1]:.6f}")
    if distributed:
        gathered = [None] * world_size
        torch.distributed.all_gather_object(gathered, local)
        trained = [item for rank_items in gathered for item in rank_items]
    else:
        trained = local
    if rank == 0:
        trained.sort(key=lambda x: x["replica"])
        if [x["replica"] for x in trained] != list(range(replicas)):
            raise RuntimeError("distributed ranks did not return every replica")
        checkpoint = ROOT / config["paths"]["checkpoint"]
        metrics = ROOT / config["paths"]["training_metrics"]
        checkpoint.parent.mkdir(parents=True, exist_ok=True)
        metrics.parent.mkdir(parents=True, exist_ok=True)
        scaling = {"feature_mean": feature_mean, "feature_std": feature_std, "target_mean": target_mean, "target_std": target_std}
        states = [x["state_dict"] for x in trained]
        torch.save({"format_version": FORMAT_VERSION, "model": states,
                    "model_config": config["model"], "config": config, "scaling": scaling}, checkpoint)
        metrics.write_text(json.dumps({"replicas": [{"id": x["replica"], "loss": x["loss"]} for x in trained], "world_size": world_size}, indent=2) + "\n")
    if distributed:
        torch.distributed.destroy_process_group()


if __name__ == "__main__":
    main()