"""Train independently initialized DD members with optional DDP."""

import copy
import json
import os
import sys
from pathlib import Path

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


ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
from model.precipdd import DATA_FORMAT_VERSION, PrecipDD, load_config, seed_all, validate_archive


def main():
    config = load_config(ROOT / "conf/config.yaml")
    runtime, settings = config["runtime"], config["training"]
    distributed = int(os.environ.get("WORLD_SIZE", "1")) > 1
    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
    local_world_size = int(os.environ.get("LOCAL_WORLD_SIZE", "1"))
    use_cuda = (torch.cuda.is_available() and runtime["device"] != "cpu"
                and torch.cuda.device_count() >= local_world_size)
    if distributed:
        torch.distributed.init_process_group(runtime["ddp_backend_gpu"] if use_cuda else runtime["ddp_backend_cpu"])
    rank = torch.distributed.get_rank() if distributed else 0
    device = torch.device(f"cuda:{local_rank}" if use_cuda else "cpu")
    if use_cuda:
        torch.cuda.set_device(device)
    data = np.load(ROOT / config["paths"]["data"])
    validate_archive(data)
    train_mask, val_mask = data["split"] == 0, data["split"] == 1
    dataset = TensorDataset(torch.from_numpy(data["precipitation"][train_mask]).float(), torch.from_numpy(data["agmt"][train_mask]).float())
    val_x = torch.from_numpy(data["precipitation"][val_mask]).float().to(device)
    val_y = torch.from_numpy(data["agmt"][val_mask]).float().to(device)
    ensemble_states, histories = [], []
    for member in range(settings["ensemble_members"]):
        seed_all(config["project"]["seed"] + member)
        sampler = DistributedSampler(dataset, shuffle=True, seed=member) if distributed else None
        loader = DataLoader(dataset, batch_size=settings["batch_size"], shuffle=sampler is None, sampler=sampler,
                            num_workers=settings["num_workers"])
        model = PrecipDD(config["model"]["filters"], config["model"]["dense_units"]).to(device)
        wrapped = DistributedDataParallel(model, device_ids=[local_rank] if use_cuda else None) if distributed else model
        optimizer = torch.optim.Adam(wrapped.parameters(), lr=settings["learning_rate"], weight_decay=settings["l2_weight_decay"])
        history, best_loss, best_state = [], float("inf"), None
        for epoch in range(settings["epochs"]):
            if sampler is not None:
                sampler.set_epoch(member * settings["epochs"] + epoch)
            wrapped.train()
            total, count = 0.0, 0
            for fields, target in loader:
                fields, target = fields.to(device), target.to(device)
                optimizer.zero_grad(set_to_none=True)
                loss = nn.functional.l1_loss(wrapped(fields), target)
                loss.backward()
                optimizer.step()
                total += float(loss.detach()) * len(fields)
                count += len(fields)
            summary = torch.tensor([total, count], dtype=torch.float64, device=device)
            if distributed:
                torch.distributed.all_reduce(summary)
            wrapped.eval()
            with torch.no_grad():
                val_loss = float(nn.functional.l1_loss(model(val_x), val_y))
            history.append({"epoch": epoch + 1, "train_mae": float(summary[0] / summary[1]), "validation_mae": val_loss})
            if val_loss < best_loss:
                best_loss = val_loss
                best_state = {key: value.detach().cpu().clone() for key, value in model.state_dict().items()}
        ensemble_states.append(best_state)
        histories.append(history)
    if rank == 0:
        checkpoint_path = ROOT / config["paths"]["checkpoint"]
        checkpoint_path.parent.mkdir(parents=True, exist_ok=True)
        torch.save({"format_version": DATA_FORMAT_VERSION, "ensemble_states": ensemble_states,
                    "model_config": {"filters": config["model"]["filters"], "dense_units": config["model"]["dense_units"],
                                     "input_shape": [1, 55, 160], "feature_shape": [16, 14, 40], "flattened_features": 8960},
                    "training_config": settings, "histories": histories, "world_size": torch.distributed.get_world_size() if distributed else 1}, checkpoint_path)
        metrics_path = ROOT / config["paths"]["training_metrics"]
        metrics_path.parent.mkdir(parents=True, exist_ok=True)
        metrics_path.write_text(json.dumps({"ensemble_members": len(histories), "history": histories}, indent=2) + "\n", encoding="utf-8")
        print(f"checkpoint={checkpoint_path.relative_to(ROOT)} members={len(ensemble_states)}")
    if distributed:
        torch.distributed.barrier()
        torch.distributed.destroy_process_group()


if __name__ == "__main__":
    main()