"""Run the self-contained virtual-data training, inference, and plotting pipeline."""

import argparse
import json
import os
import random
from pathlib import Path

import numpy as np
import torch
import torch.distributed as dist
from torch import nn
from torch.nn.parallel import DistributedDataParallel
from torch.utils.data import DataLoader, DistributedSampler, TensorDataset
import sys
PROJECT_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(PROJECT_ROOT))
from model.spherical_dyffusion import SphericalDYffusion


VARIABLES = [
    "PRESsfc", "surface_temperature",
    *[f"air_temperature_{i}" for i in range(8)],
    *[f"specific_total_water_{i}" for i in range(8)],
    *[f"eastward_wind_{i}" for i in range(8)],
    *[f"northward_wind_{i}" for i in range(8)],
    "DSWRFtoa", "HGTsfc", "ocean_fraction",
]


def load_config(path: str) -> dict:
    import yaml

    with open(path, encoding="utf-8") as file:
        return yaml.safe_load(file)


def resolve_device(config: dict, local_rank: int = 0) -> torch.device:
    requested = str(config.get("runtime", {}).get("device", "auto")).lower()
    if requested == "auto":
        requested = "cuda" if torch.cuda.is_available() else "cpu"
    if requested.startswith("cuda") and not torch.cuda.is_available():
        raise RuntimeError(
            f"runtime.device={requested!r}, but PyTorch cannot access a CUDA/ROCm device. "
            "Use runtime.device=cpu or install a GPU-enabled PyTorch build."
        )
    device = torch.device(requested)
    if device.type == "cuda":
        device_index = device.index if device.index is not None else local_rank
        if device_index >= torch.cuda.device_count():
            raise RuntimeError(
                f"LOCAL_RANK={local_rank} maps to GPU {device_index}, but only "
                f"{torch.cuda.device_count()} GPU(s) are visible."
            )
        torch.cuda.set_device(device_index)
        device = torch.device("cuda", device_index)
        print(
            f"device: {device} ({torch.cuda.get_device_name(device_index)}), "
            f"backend={'ROCm ' + torch.version.hip if torch.version.hip else 'CUDA ' + str(torch.version.cuda)}"
        )
    else:
        print("device: cpu")
    return device


def setup_distributed(config: dict) -> tuple[int, int, torch.device]:
    world_size = int(os.environ.get("WORLD_SIZE", "1"))
    rank = int(os.environ.get("RANK", "0"))
    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
    configured_devices = config.get("runtime", {}).get("devices", "auto")
    requested_device = str(config.get("runtime", {}).get("device", "auto")).lower()
    if isinstance(configured_devices, int) and configured_devices > 1 and world_size == 1:
        raise RuntimeError(
            f"runtime.devices={configured_devices} requires a distributed launcher. Run: "
            f"python -m torch.distributed.run --standalone --nproc-per-node={configured_devices} "
            "scripts/train.py"
        )
    if isinstance(configured_devices, int) and world_size > 1 and configured_devices != world_size:
        raise RuntimeError(
            f"runtime.devices={configured_devices} does not match torchrun WORLD_SIZE={world_size}."
        )
    if world_size > 1 and requested_device.startswith("cuda:"):
        raise RuntimeError(
            "Do not set an explicit CUDA device index for DDP. Use runtime.device=cuda or auto; "
            "each torchrun process is mapped to its LOCAL_RANK automatically."
        )

    device = resolve_device(config, local_rank=local_rank)
    if world_size > 1:
        if not dist.is_available():
            raise RuntimeError("Distributed training is unavailable in this PyTorch build.")
        backend = str(config.get("runtime", {}).get("distributed_backend", "auto")).lower()
        if backend == "auto":
            backend = "nccl" if device.type == "cuda" else "gloo"
        dist.init_process_group(backend=backend, init_method="env://")
        if dist.get_world_size() != world_size or dist.get_rank() != rank:
            raise RuntimeError("The process group does not match the torchrun rank settings.")
    return rank, world_size, device


def generate(config: dict) -> Path:
    spec = config["synthetic_data"]
    rng = np.random.default_rng(spec["seed"])
    shape = (spec["samples"], spec["channels"], spec["latitude"], spec["longitude"])
    inputs = rng.normal(0, 1, shape).astype(np.float32)
    # A deterministic local dynamics rule provides a learnable target.
    targets = (0.85 * inputs + 0.05 * np.roll(inputs, 1, axis=2) + 0.05 * np.roll(inputs, -1, axis=3)).astype(np.float32)
    output = Path(spec["output_dir"])
    output.mkdir(parents=True, exist_ok=True)
    path = output / "virtual_fv3gfs.npz"
    np.savez_compressed(path, inputs=inputs, targets=targets)
    (output / "metadata.json").write_text(json.dumps({
        "variables": VARIABLES, "shape": list(shape), "time_steps": spec["time_steps"],
        "latitude": spec["latitude"], "longitude": spec["longitude"],
        "dataset_type": "virtual_fv3gfs_equivalent_contract",
    }, indent=2) + "\n", encoding="utf-8")
    print(f"virtual data: {path}")
    return path


def train(config: dict, finetune: str | None = None) -> Path:
    seed = config["training"]["seed"]
    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
    rank, world_size, device = setup_distributed(config)
    is_root = rank == 0
    checkpoint = Path(config["training"]["checkpoint_dir"])
    try:
        if int(config["training"]["epochs"]) < 1:
            raise ValueError("training.epochs must be at least 1.")
        data_path = Path(config["synthetic_data"]["output_dir"]) / "virtual_fv3gfs.npz"
        if is_root:
            data_path = generate(config)
        if world_size > 1:
            dist.barrier()
        arrays = np.load(data_path)
        dataset = TensorDataset(torch.from_numpy(arrays["inputs"]), torch.from_numpy(arrays["targets"]))

        global_batch_size = int(config["training"]["batch_size"])
        if global_batch_size % world_size:
            raise ValueError(
                f"training.batch_size={global_batch_size} is the global batch size and must be "
                f"divisible by WORLD_SIZE={world_size}."
            )
        if len(dataset) % world_size:
            raise ValueError(
                f"Dataset size {len(dataset)} must be divisible by WORLD_SIZE={world_size}; "
                "otherwise DistributedSampler would duplicate samples and bias the epoch loss."
            )
        local_batch_size = global_batch_size // world_size
        if local_batch_size < 1:
            raise ValueError("The global batch size must be at least WORLD_SIZE.")
        sampler = DistributedSampler(
            dataset,
            num_replicas=world_size,
            rank=rank,
            shuffle=True,
            seed=seed,
            drop_last=False,
        ) if world_size > 1 else None
        total_workers = int(config.get("runtime", {}).get("num_workers", 0))
        local_workers = max(0, total_workers // world_size)
        loader = DataLoader(
            dataset,
            batch_size=local_batch_size,
            shuffle=sampler is None,
            sampler=sampler,
            num_workers=local_workers,
            pin_memory=device.type == "cuda",
            persistent_workers=local_workers > 0,
        )
        model = SphericalDYffusion(config["synthetic_data"]["channels"]).to(device)
        if finetune:
            finetune_path = Path(finetune)
            if not finetune_path.exists():
                raise FileNotFoundError(f"Fine-tune checkpoint not found: {finetune_path}")
            state = torch.load(finetune_path, map_location="cpu", weights_only=False)
            state_dict = state.get("model", state) if isinstance(state, dict) else state
            try:
                model.load_state_dict(state_dict)
            except RuntimeError as error:
                raise RuntimeError(
                    "Fine-tune checkpoint is incompatible with SphericalDYffusion. "
                    "Use a checkpoint produced by this local pipeline or a matching model architecture."
                ) from error
            if is_root:
                print(f"fine-tuning from: {finetune_path}")
        elif is_root:
            print("training from scratch")
        if world_size > 1:
            model = DistributedDataParallel(
                model,
                device_ids=[device.index] if device.type == "cuda" else None,
                output_device=device.index if device.type == "cuda" else None,
            )
            if is_root:
                print(
                    f"distributed: DDP world_size={world_size}, global_batch_size={global_batch_size}, "
                    f"local_batch_size={local_batch_size}"
                )
        optimizer = torch.optim.Adam(model.parameters(), lr=config["training"]["learning_rate"])
        loss_fn = nn.MSELoss()
        best = float("inf")
        if is_root:
            checkpoint.mkdir(parents=True, exist_ok=True)
        for epoch in range(1, config["training"]["epochs"] + 1):
            if sampler is not None:
                sampler.set_epoch(epoch)
            model.train(); total = 0.0; sample_count = 0
            for inputs, targets in loader:
                inputs = inputs.to(device, non_blocking=True)
                targets = targets.to(device, non_blocking=True)
                optimizer.zero_grad(); loss = loss_fn(model(inputs), targets); loss.backward(); optimizer.step()
                total += loss.item() * len(inputs); sample_count += len(inputs)
            loss_stats = torch.tensor([total, sample_count], dtype=torch.float64, device=device)
            if world_size > 1:
                dist.all_reduce(loss_stats, op=dist.ReduceOp.SUM)
            mean_loss = (loss_stats[0] / loss_stats[1]).item()
            if is_root:
                print(f"epoch {epoch}/{config['training']['epochs']} loss={mean_loss:.6f}")
                state_dict = model.module.state_dict() if isinstance(model, DistributedDataParallel) else model.state_dict()
                state = {
                    "model": state_dict,
                    "channels": config["synthetic_data"]["channels"],
                    "loss": mean_loss,
                    "world_size": world_size,
                    "global_batch_size": global_batch_size,
                }
                if mean_loss < best:
                    best = mean_loss
                    torch.save(state, checkpoint / "model_bak.pt")
        if is_root:
            torch.save(state, checkpoint / "last.pt")
            print(f"checkpoint: {checkpoint / 'model_bak.pt'}")
        if world_size > 1:
            dist.barrier()
        return checkpoint / "model_bak.pt"
    finally:
        if dist.is_available() and dist.is_initialized():
            dist.destroy_process_group()


def infer(config: dict) -> Path:
    arrays = np.load(Path(config["synthetic_data"]["output_dir"]) / "virtual_fv3gfs.npz")
    checkpoint_path = Path(config["inference"]["checkpoint"])
    if not checkpoint_path.exists():
        raise FileNotFoundError(f"Inference checkpoint not found: {checkpoint_path}. Run scripts/train.py first.")
    state = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
    model = SphericalDYffusion(state["channels"]); model.load_state_dict(state["model"]); model.eval()
    with torch.no_grad(): prediction = model(torch.from_numpy(arrays["inputs"]))
    output = Path(config["inference"]["output_dir"]); output.mkdir(parents=True, exist_ok=True)
    path = output / "prediction.npz"; np.savez_compressed(path, prediction=prediction.numpy(), target=arrays["targets"])
    print(f"inference: {path}"); return path


def visualize(config: dict) -> Path:
    import matplotlib.pyplot as plt

    arrays = np.load(Path(config["inference"]["output_dir"]) / "prediction.npz")
    index, channel = config["visualization"]["prediction_index"], config["visualization"]["channel"]
    figure, axes = plt.subplots(1, 2, figsize=(12, 4), constrained_layout=True)
    for axis, image, title in zip(axes, [arrays["target"][index, channel], arrays["prediction"][index, channel]], ["Target", "Prediction"]):
        plot = axis.imshow(image, cmap="viridis"); axis.set_title(title); axis.set_xlabel("longitude"); axis.set_ylabel("latitude"); figure.colorbar(plot, ax=axis)
    output = Path(config["visualization"]["output_dir"]); output.mkdir(parents=True, exist_ok=True)
    path = output / "prediction_comparison.png"; figure.savefig(path, dpi=150); plt.close(figure)
    print(f"visualization: {path}"); return path


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--config", default="conf/config.yaml")
    parser.add_argument("action", choices=["generate", "train", "infer", "visualize", "all"], nargs="?", default="all")
    args = parser.parse_args(); config = load_config(args.config)
    if args.action == "generate":
        generate(config)
    elif args.action == "train":
        train(config)
    elif args.action == "infer":
        infer(config)
    elif args.action == "visualize":
        visualize(config)
    elif args.action == "all":
        train(config)
        if int(os.environ.get("RANK", "0")) == 0:
            infer(config)
            visualize(config)


if __name__ == "__main__":
    main()