"""Train compact SatlasNet on seven task families; supports AMP and torchrun."""

import argparse
import importlib.util
import json
import os
from contextlib import nullcontext
from functools import partial
from pathlib import Path
import numpy as np
import torch
import yaml
from torch import distributed as dist
from torch.nn.parallel import DistributedDataParallel
from torch.utils.data import DataLoader, Dataset, DistributedSampler

ROOT = Path(__file__).resolve().parents[1]


def load_module():
    spec = importlib.util.spec_from_file_location("satlaspretrain", ROOT / "model/satlaspretrain.py")
    module = importlib.util.module_from_spec(spec); spec.loader.exec_module(module); return module


class NpzDataset(Dataset):
    keys = ("highres_images", "lowres_images", "valid_highres_times", "valid_lowres_times",
            "segmentation", "regression", "point", "polygon", "polyline", "property", "classification")
    def __init__(self, path, config, module):
        archive = np.load(path)
        module.validate_npz(archive, config)
        self.data = {key: archive[key] for key in self.keys}
        self.source = str(archive["source"])
        self.protocol = str(archive["protocol"]) if "protocol" in archive else "provided_npz"
    def __len__(self): return len(self.data["highres_images"])
    def __getitem__(self, index): return {key: torch.as_tensor(value[index]) for key, value in self.data.items()}


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--config", type=Path, default=ROOT / "conf/config.yaml"); parser.add_argument("--data", type=Path)
    parser.add_argument("--checkpoint", type=Path); parser.add_argument("--device", choices=("auto", "cpu", "cuda"))
    args = parser.parse_args(); config = yaml.safe_load(args.config.read_text()); train_cfg = config["training"]
    world_size, rank, local_rank = int(os.environ.get("WORLD_SIZE", 1)), int(os.environ.get("RANK", 0)), int(os.environ.get("LOCAL_RANK", 0))
    requested = args.device or config["runtime"]["device"]; use_cuda = torch.cuda.is_available() and requested != "cpu"
    if requested == "cuda" and not use_cuda: raise RuntimeError("CUDA requested but unavailable")
    if world_size > 1: dist.init_process_group("nccl" if use_cuda else "gloo")
    device = torch.device(f"cuda:{local_rank}" if use_cuda else "cpu")
    if use_cuda: torch.cuda.set_device(local_rank)
    torch.manual_seed(config["seed"] + rank)
    module = load_module()
    data_path = args.data or ROOT / config["data"]["root"] / "train.npz"; dataset = NpzDataset(data_path, config, module)
    sampler = DistributedSampler(dataset, shuffle=True) if world_size > 1 else None
    loader = DataLoader(dataset, batch_size=train_cfg["batch_size"], sampler=sampler, shuffle=sampler is None,
                        num_workers=train_cfg["num_workers"], pin_memory=use_cuda)
    model = module.SatlasPretrain(**config["model"]).to(device); raw_model = model
    if world_size > 1:
        model = DistributedDataParallel(model, device_ids=[local_rank] if use_cuda else None); raw_model = model.module
    optimizer = torch.optim.AdamW(model.parameters(), lr=train_cfg["learning_rate"], weight_decay=train_cfg["weight_decay"])
    amp = bool(config["runtime"]["amp"] and use_cuda); scaler = torch.amp.GradScaler("cuda", enabled=amp); history = []
    for epoch in range(train_cfg["epochs"]):
        if sampler is not None: sampler.set_epoch(epoch)
        model.train(); totals = {"total": 0.0}; steps = 0
        for batch in loader:
            batch = {key: value.to(device, non_blocking=use_cuda) for key, value in batch.items()}
            autocast = partial(torch.amp.autocast, "cuda") if amp else nullcontext
            with autocast():
                outputs = model(batch["highres_images"], batch["lowres_images"],
                                batch["valid_highres_times"], batch["valid_lowres_times"])
                loss, parts = module.multitask_loss(outputs, batch)
            if not torch.isfinite(loss): raise ValueError("non-finite multitask loss")
            optimizer.zero_grad(set_to_none=True); scaler.scale(loss).backward(); scaler.step(optimizer); scaler.update()
            totals["total"] += loss.detach().item(); steps += 1
            for name, value in parts.items(): totals[name] = totals.get(name, 0.0) + value.detach().item()
        names = list(totals); statistics = torch.tensor([totals[name] for name in names] + [steps], dtype=torch.float64, device=device)
        if world_size > 1: dist.all_reduce(statistics, op=dist.ReduceOp.SUM)
        global_steps = max(statistics[-1].item(), 1)
        record = {"epoch": epoch + 1, **{f"{name}_loss": statistics[index].item() / global_steps for index, name in enumerate(names)}}; history.append(record)
        if rank == 0: print(json.dumps(record))
    if rank == 0:
        checkpoint_path = args.checkpoint or ROOT / config["paths"]["checkpoint"]; checkpoint_path.parent.mkdir(parents=True, exist_ok=True)
        torch.save({"model": raw_model.state_dict(), "optimizer": optimizer.state_dict(), "scaler": scaler.state_dict() if amp else None,
                    "config": config, "epoch": train_cfg["epochs"], "history": history}, checkpoint_path)
        metrics = ROOT / config["paths"]["training_metrics"]; metrics.parent.mkdir(parents=True, exist_ok=True)
        metrics.write_text(json.dumps({"history": history, "protocol": dataset.protocol,
                                       "source": dataset.source,
                                       "world_size": world_size, "amp": amp}, indent=2) + "\n")
        print(f"checkpoint={checkpoint_path}")
    if world_size > 1: dist.destroy_process_group()


if __name__ == "__main__": main()