import argparse
import os
import sys
from pathlib import Path

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

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

from data_loader import SyntheticOceanDataset
from model.glonet import GLONET


def setup_distributed():
    world_size = int(os.environ.get("WORLD_SIZE", "1"))
    if world_size == 1:
        return 0, 0, torch.device("cuda" if torch.cuda.is_available() else "cpu")
    local_rank = int(os.environ["LOCAL_RANK"])
    force_cpu = os.environ.get("GLONET_FORCE_CPU", "0") == "1"
    if torch.cuda.is_available() and not force_cpu:
        device_count = torch.cuda.device_count()
        if local_rank >= device_count:
            raise RuntimeError(
                f"LOCAL_RANK={local_rank} but only {device_count} accelerator(s) are visible; "
                "reduce --nproc_per_node or fix CUDA_VISIBLE_DEVICES."
            )
        torch.cuda.set_device(local_rank)
        device = torch.device("cuda", local_rank)
        backend = "nccl"
    else:
        device = torch.device("cpu")
        backend = "gloo"
    dist.init_process_group(backend=backend, init_method="env://")
    return dist.get_rank(), local_rank, device


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", default=str(ROOT / "conf/config.yaml"))
    parser.add_argument("--phase", choices=["pretrain", "finetune"], default="pretrain")
    args = parser.parse_args()
    with open(args.config, encoding="utf-8") as handle:
        config = yaml.safe_load(handle)
    rank, local_rank, device = setup_distributed()
    torch.manual_seed(config["project"]["seed"] + rank)
    channels = len(config["data"]["channels"])
    rollout_steps = config["training"][f"{args.phase}_rollout_steps"]
    dataset = SyntheticOceanDataset(config["data"]["synthetic_samples"], channels, config["data"]["grid"],
                                    input_steps=config["data"]["input_steps"],
                                    output_steps=config["data"]["output_steps"],
                                    data_dir=str(ROOT / config["data"]["data_dir"]))
    sampler = DistributedSampler(dataset, shuffle=True) if dist.is_initialized() else None
    loader = DataLoader(dataset, batch_size=config["data"]["batch_size"], shuffle=sampler is None, sampler=sampler)
    model = GLONET(channels * config["data"]["input_steps"], out_channels=channels,
                   hidden_channels=config["model"]["hidden_channels"], modes=config["model"]["modes"],
                   layers=config["model"]["layers"]).to(device)
    checkpoint = ROOT / config["training"]["checkpoint"]
    if args.phase == "finetune" and checkpoint.exists():
        state = torch.load(checkpoint, map_location=device, weights_only=False)
        model.load_state_dict(state["model"])
    if dist.is_initialized():
        model = DDP(model, device_ids=[local_rank] if device.type == "cuda" else None)
    optimizer = torch.optim.Adam(model.parameters(), lr=config["training"]["learning_rate"])
    for epoch in range(config["training"]["epochs"]):
        if sampler is not None:
            sampler.set_epoch(epoch)
        model.train()
        total = 0.0
        for inputs, targets in loader:
            inputs, targets = inputs.to(device), targets.to(device)
            optimizer.zero_grad(set_to_none=True)
            loss = 0.0
            state = inputs
            for step in range(rollout_steps):
                prediction = model(state)
                loss = loss + torch.nn.functional.mse_loss(prediction, targets[:, step])
                state = torch.cat((state[:, 1:], prediction.unsqueeze(1)), dim=1)
            loss = loss / rollout_steps
            loss.backward()
            optimizer.step()
            total += loss.item()
        if rank == 0:
            print(f"epoch={epoch + 1} loss={total / len(loader):.6f}")
    if rank == 0:
        checkpoint.parent.mkdir(parents=True, exist_ok=True)
        torch.save({"model": model.module.state_dict() if hasattr(model, "module") else model.state_dict(),
                    "config": config}, checkpoint)
        print(f"saved={checkpoint}")
    if dist.is_initialized():
        dist.destroy_process_group()


if __name__ == "__main__":
    main()