"""Train SmaAt-UNet for six-frame precipitation nowcasting."""

import json
import os
import sys
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.smaatunet import SmaAtUNet


class PrecipitationDataset(Dataset):
    def __init__(self, path, config):
        self.data = np.load(path)
        self.config = config
        if str(self.data["format_version"]) != config["data"]["format_version"]:
            raise ValueError("incompatible precipitation data format")
        height, width = int(config["data"]["height"]), int(config["data"]["width"])
        if self.data["inputs"].shape[1:] != (int(config["data"]["input_frames"]), height, width):
            raise ValueError("input dimensions do not match the paper precipitation data")
        if self.data["targets"].shape[1:] != (int(config["data"]["output_frames"]), height, width):
            raise ValueError("target dimensions do not match the paper precipitation data")

    def __len__(self):
        return len(self.data["inputs"])

    def __getitem__(self, index):
        return torch.from_numpy(self.data["inputs"][index]).float(), torch.from_numpy(self.data["targets"][index]).float()


def device_from_config(config, rank=0):
    if config["runtime"]["device"] == "auto":
        return torch.device("cuda", rank) if torch.cuda.is_available() else torch.device("cpu")
    return torch.device(config["runtime"]["device"])


def main():
    config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
    torch.manual_seed(int(config["seed"]))
    distributed = int(os.environ.get("WORLD_SIZE", "1")) > 1
    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
    if distributed:
        torch.distributed.init_process_group("nccl" if torch.cuda.is_available() else "gloo")
    rank = torch.distributed.get_rank() if distributed else 0
    device = device_from_config(config, local_rank)
    dataset = PrecipitationDataset(ROOT / config["data"]["root"] / "train.npz", config)
    sampler = DistributedSampler(dataset, shuffle=True) if distributed else None
    loader = DataLoader(dataset, batch_size=int(config["train"]["batch_size"]), sampler=sampler,
                        shuffle=sampler is None, num_workers=int(config["train"]["num_workers"]))
    model = SmaAtUNet(config["model"]).to(device)
    if distributed:
        model = DistributedDataParallel(model, device_ids=[local_rank] if device.type == "cuda" else None)
    optimizer = torch.optim.Adam(model.parameters(), lr=float(config["train"]["learning_rate"]),
                                 weight_decay=float(config["train"]["weight_decay"]))
    history = []
    for epoch in range(int(config["train"]["epochs"])):
        model.train()
        total, steps = 0.0, 0
        for inputs, targets in loader:
            prediction = model(inputs.to(device))
            loss = torch.nn.functional.mse_loss(prediction, targets.to(device))
            optimizer.zero_grad(set_to_none=True)
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()
            total += float(loss.detach())
            steps += 1
        metrics = {"epoch": epoch + 1, "mse": total / max(steps, 1)}
        history.append(metrics)
        if rank == 0:
            print(f"epoch={epoch + 1} mse={metrics['mse']:.6f}")
    if rank == 0:
        checkpoint = ROOT / config["paths"]["checkpoint"]
        metrics_path = ROOT / config["paths"]["training_metrics"]
        checkpoint.parent.mkdir(parents=True, exist_ok=True)
        metrics_path.parent.mkdir(parents=True, exist_ok=True)
        state = model.module.state_dict() if distributed else model.state_dict()
        torch.save({"model": state, "model_config": config["model"],
                    "format_version": config["data"]["format_version"]}, checkpoint)
        metrics_path.write_text(json.dumps({"history": history}, indent=2) + "\n")
    if distributed:
        torch.distributed.destroy_process_group()


if __name__ == "__main__":
    main()