"""DDP-capable EchoCast-3D training entry point."""

from __future__ import annotations

import argparse
import json
import os
from pathlib import Path

import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel
from torch.utils.data import DataLoader, DistributedSampler
import sys
ROOT=Path(__file__).resolve().parents[1];sys.path.insert(0,str(ROOT))

from model.echocast_3d import PackedRadarDataset, load_config, DiffusionSchedule, EchoCast3D, masked_objective


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", default="conf/config.yaml")
    parser.add_argument("--data", default="data")
    parser.add_argument("--output", default="result/checkpoints/echocast_3d.pt")
    args = parser.parse_args()
    config = load_config(ROOT / args.config)
    distributed = int(os.environ.get("WORLD_SIZE", "1")) > 1
    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
    if distributed:
        dist.init_process_group("nccl" if torch.cuda.is_available() else "gloo")
    use_cuda=torch.cuda.is_available() and torch.cuda.device_count()>=int(os.environ.get("WORLD_SIZE","1")) and config["runtime"]["device"]!="cpu"
    device = torch.device(f"cuda:{local_rank}" if use_cuda else "cpu")
    if device.type == "cuda":
        torch.cuda.set_device(device)
    torch.manual_seed(config["seed"] + local_rank)
    dataset = PackedRadarDataset(ROOT / args.data, "train")
    sampler = DistributedSampler(dataset, shuffle=True) if distributed else None
    loader = DataLoader(dataset, batch_size=config["training"]["batch_size"], sampler=sampler, shuffle=sampler is None)
    raw_model = EchoCast3D(config).to(device)
    model = DistributedDataParallel(raw_model, device_ids=[local_rank]) if distributed else raw_model
    optimizer = torch.optim.AdamW(
        model.parameters(), lr=config["training"]["learning_rate"], weight_decay=config["training"]["weight_decay"]
    )
    schedule = DiffusionSchedule(device=device, **{k: config["diffusion"][k] for k in ("steps", "beta_start", "beta_end")})
    for epoch in range(config["training"]["epochs"]):
        if sampler:
            sampler.set_epoch(epoch)
        for batch in loader:
            truth = batch["truth"].float().to(device)
            values = batch["values"].float().to(device)
            validity = batch["validity"].bool().to(device)
            observed = batch["observed"].bool().to(device)
            clean_target = truth[:, 3:]
            timestep = torch.randint(schedule.steps, (truth.shape[0],), device=device)
            noisy_target, _ = schedule.add_noise(clean_target, timestep)
            token_mask = torch.rand(truth.shape[0], raw_model.token_count, device=device) < config["model"]["mask_ratio"]
            prediction = model(values[:, :3], noisy_target, observed[:, :3], timestep, token_mask)
            loss, parts = masked_objective(
                raw_model, prediction, clean_target, validity, token_mask, config["training"]["reconstruction_weight"]
            )
            optimizer.zero_grad(set_to_none=True)
            loss.backward()
            optimizer.step()
            if local_rank == 0:
                print(f"epoch={epoch + 1} loss={loss.item():.5f} score={parts['score']:.5f} recon={parts['reconstruction']:.5f}")
    if local_rank == 0:
        output = ROOT / args.output
        output.parent.mkdir(parents=True, exist_ok=True)
        torch.save({"model": raw_model.state_dict(), "model_config": config["model"], "format_version": config["data"]["format_version"], "config": config}, output)
        metrics=ROOT/config["paths"]["training_metrics"];metrics.parent.mkdir(parents=True,exist_ok=True);metrics.write_text(json.dumps({"epochs":config["training"]["epochs"],"loss":float(loss),"world_size":int(os.environ.get("WORLD_SIZE","1"))},indent=2))
        print(f"saved {output}; geometry tokens={raw_model.token_count}, paper states {config['model']['paper_token_count']}")
    if distributed:
        dist.destroy_process_group()


if __name__ == "__main__":
    main()