"""Train DINCAE with random cloud masking and masked Gaussian NLL."""

import json
import os
import sys
from datetime import date
from pathlib import Path

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

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

from model.dincae import DINCAE, build_input, masked_gaussian_nll


class SSTDataset(Dataset):
    def __init__(self, arrays, indices, cloud_fraction, seed):
        self.data, self.indices = arrays, list(indices)
        self.cloud_fraction, self.seed = float(cloud_fraction), int(seed)

    def __len__(self):
        return len(self.indices)

    def __getitem__(self, item):
        index = self.indices[item]
        timestamp = date.fromisoformat(str(self.data["timestamps"][index]))
        observed = self.data["observed_anomaly"].copy()
        precision = self.data["precision"].copy()
        rng = np.random.default_rng(self.seed + index)
        original_valid = np.isfinite(observed[index]) & self.data["ocean_mask"]
        withheld = original_valid & (rng.random(original_valid.shape) < self.cloud_fraction)
        observed[index, withheld] = np.nan
        precision[index, withheld] = 0.0
        inputs = build_input(observed, precision, index, self.data["longitude"],
                             self.data["latitude"], timestamp.timetuple().tm_yday)
        target = self.data["sst_anomaly"][index][None].astype(np.float32)
        return torch.from_numpy(inputs), torch.from_numpy(target), torch.from_numpy(withheld[None])


def main():
    config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
    rank = int(os.environ.get("RANK", "0"))
    world_size = int(os.environ.get("WORLD_SIZE", "1"))
    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
    distributed = world_size > 1
    if distributed:
        backend = "nccl" if torch.cuda.is_available() else "gloo"
        torch.distributed.init_process_group(backend=backend)
    requested = config["runtime"]["device"]
    use_cuda = torch.cuda.is_available() and requested != "cpu"
    device = torch.device(f"cuda:{local_rank}" if use_cuda else "cpu")
    if use_cuda:
        torch.cuda.set_device(device)
    torch.manual_seed(int(config["seed"]) + rank)
    path = ROOT / config["data"]["root"] / config["data"]["file"]
    with np.load(path) as loaded:
        arrays = {key: loaded[key] for key in loaded.files}
    train_count = min(int(config["data"]["train_samples"]), len(arrays["timestamps"]) - 1)
    dataset = SSTDataset(arrays, range(train_count), config["data"]["random_cloud_fraction"], config["seed"])
    sampler = DistributedSampler(dataset, shuffle=True, seed=int(config["seed"])) if distributed else None
    loader = DataLoader(dataset, batch_size=int(config["training"]["batch_size"]),
                        shuffle=sampler is None, sampler=sampler,
                        num_workers=int(config["runtime"]["num_workers"]))
    model = DINCAE(**config["model"]).to(device)
    if distributed:
        model = DistributedDataParallel(model, device_ids=[local_rank] if use_cuda else None)
    optimizer = torch.optim.Adam(model.parameters(), lr=float(config["training"]["learning_rate"]))
    history = []
    for epoch in range(int(config["training"]["epochs"])):
        if sampler is not None:
            sampler.set_epoch(epoch)
        model.train()
        total, batches = 0.0, 0
        for inputs, target, mask in loader:
            inputs, target, mask = inputs.to(device), target.to(device), mask.to(device)
            optimizer.zero_grad(set_to_none=True)
            output = model(inputs)
            loss = masked_gaussian_nll(output, target, mask,
                                       config["training"]["gamma"], config["training"]["delta"])
            loss.backward()
            optimizer.step()
            total, batches = total + float(loss.detach()), batches + 1
        history.append({"epoch": epoch + 1, "masked_gaussian_nll": total / max(batches, 1)})
        if rank == 0:
            print(f"epoch={epoch + 1} loss={history[-1]['masked_gaussian_nll']:.6f}")
    if rank == 0:
        raw_model = model.module if distributed else model
        checkpoint = ROOT / config["paths"]["checkpoint"]
        checkpoint.parent.mkdir(parents=True, exist_ok=True)
        torch.save({"model": raw_model.state_dict(), "model_config": config["model"],
                    "gamma": config["training"]["gamma"], "delta": config["training"]["delta"],
                    "format_version": "dincae_checkpoint_v1"}, checkpoint)
        metrics = ROOT / config["paths"]["training_metrics"]
        metrics.parent.mkdir(parents=True, exist_ok=True)
        metrics.write_text(json.dumps({"history": history, "world_size": world_size,
                                       "parameter_count": sum(p.numel() for p in raw_model.parameters())}, indent=2))
        print(f"checkpoint={checkpoint}")
    if distributed:
        torch.distributed.destroy_process_group()


if __name__ == "__main__":
    main()