import json
import os
from pathlib import Path
import sys

import numpy as np
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel
from torch.utils.data import DataLoader, DistributedSampler, TensorDataset
import yaml


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

from model.stablenn_phys import StableNNPhys, rollout, rollout_loss


CHECKPOINT_FORMAT_VERSION = "stablenn_phys_checkpoint_v1"


def main():
    cfg = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
    torch.manual_seed(cfg["seed"])
    distributed = int(os.environ.get("WORLD_SIZE", "1")) > 1
    world_size = int(os.environ.get("WORLD_SIZE", "1"))
    use_accelerator = torch.cuda.is_available() and torch.cuda.device_count() >= world_size
    if distributed:
        dist.init_process_group("nccl" if use_accelerator else "gloo")
    rank = dist.get_rank() if distributed else 0
    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
    device = torch.device(f"cuda:{local_rank}" if use_accelerator else "cpu")
    if use_accelerator:
        torch.cuda.set_device(device)
    raw = np.load(ROOT / cfg["data"]["file"])
    tensors = [torch.from_numpy(raw[name].astype(np.float32)) for name in ("initial", "target", "surface", "advection")]
    source = torch.from_numpy(raw["source"].astype(np.int64))
    dataset = TensorDataset(*tensors, source)
    sampler = DistributedSampler(dataset, shuffle=True, seed=cfg["seed"]) if distributed else None
    loader = DataLoader(dataset, batch_size=cfg["train"]["batch_size"], sampler=sampler,
                        shuffle=sampler is None)
    all_state = torch.from_numpy(raw["target"].astype(np.float32))
    state_mean = all_state.mean((0, 1)).to(device)
    state_std = all_state.std((0, 1), unbiased=False).clamp_min(1e-5).to(device)
    tendency_mean = torch.zeros(68, device=device)
    tendency_std = torch.cat((torch.full((34,), 0.02), torch.full((34,), 2e-7))).to(device)
    layer_mass = torch.from_numpy(raw["layer_mass"].astype(np.float32)).to(device)
    model = StableNNPhys(cfg["model"]["hidden_size"]).to(device)
    if distributed:
        model = DistributedDataParallel(model, device_ids=[local_rank] if device.type == "cuda" else None)
    optimizer = torch.optim.Adam(model.parameters(), lr=cfg["train"]["learning_rate"])
    history = []
    for epoch in range(cfg["train"]["epochs"]):
        if sampler:
            sampler.set_epoch(epoch)
        total = 0.0
        for initial, target, surface, advection, _ in loader:
            initial, target, surface, advection = [x.to(device) for x in (initial, target, surface, advection)]
            prediction, _ = rollout(model, initial, surface, advection, state_mean, state_std,
                                    tendency_mean, tendency_std, float(raw["dt_seconds"]))
            loss = rollout_loss(prediction, target, layer_mass.expand(initial.shape[0], -1), cfg["train"]["loss_mode"])
            optimizer.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()
            total += loss.item() * initial.shape[0]
        value = total / len(dataset)
        history.append(value)
        if rank == 0:
            print(f"epoch={epoch + 1} loss={value:.6g}")
    if rank == 0:
        module = model.module if distributed else model
        checkpoint = ROOT / cfg["paths"]["checkpoint"]
        checkpoint.parent.mkdir(parents=True, exist_ok=True)
        payload = {"format_version": CHECKPOINT_FORMAT_VERSION, "model": module.state_dict(),
                   "optimizer_state_dict": optimizer.state_dict(), "epoch": cfg["train"]["epochs"],
                   "model_config": dict(cfg["model"]), "training_config": dict(cfg["train"]),
                   "normalization": {"state_mean": state_mean.cpu(), "state_std": state_std.cpu(),
                                     "tendency_mean": tendency_mean.cpu(), "tendency_std": tendency_std.cpu()},
                   "variables": {"input": ["sL[34]", "qT[34]", "SHF", "LHF", "SOLIN"],
                                 "output": ["dsL_dt[34]", "dqT_dt[34]"]}, "dt_seconds": float(raw["dt_seconds"])}
        torch.save(payload, checkpoint)
        metrics = ROOT / cfg["paths"]["training_metrics"]
        metrics.parent.mkdir(parents=True, exist_ok=True)
        metrics.write_text(json.dumps({"loss": history, "world_size": dist.get_world_size() if distributed else 1,
                                       "paper_model": cfg["train"]["paper_model"]}, indent=2))
        print(f"saved {checkpoint}")
    if distributed:
        dist.destroy_process_group()


if __name__ == "__main__":
    main()