"""Train Surya in one-step pretraining and rollout-tuning phases."""
import argparse, importlib.util, json, math, os, random
from contextlib import nullcontext
from pathlib import Path
import numpy as np, torch, 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]

class SolarDataset(Dataset):
    def __init__(self, path, cfg):
        d = np.load(path); self.inputs, self.targets = d["inputs"], d["targets"]
        if self.inputs.ndim != 5 or self.targets.ndim != 5 or self.inputs.shape[2] != 13 or self.inputs.shape[1] != 2:
            raise ValueError("Expected inputs [N,2,13,H,W] and targets [N,S,13,H,W]")
        self.mean = np.asarray(cfg["data"]["channel_mean"], dtype=np.float32)[None, :, None, None]
        self.std = np.asarray(cfg["data"]["channel_std"], dtype=np.float32)[None, :, None, None]
    def __len__(self): return len(self.inputs)
    def __getitem__(self, i):
        transform = lambda x: (np.sign(x) * np.log1p(np.abs(x)) - self.mean) / self.std
        return torch.from_numpy(transform(self.inputs[i]).astype(np.float32)), torch.from_numpy(transform(self.targets[i]).astype(np.float32))

def load_model():
    spec = importlib.util.spec_from_file_location("surya_model", ROOT / "model/surya.py"); mod = importlib.util.module_from_spec(spec); spec.loader.exec_module(mod); return mod.Surya
def args():
    p = argparse.ArgumentParser(); p.add_argument("--config", type=Path, default=ROOT / "conf/config.yaml"); p.add_argument("--data", type=Path); p.add_argument("--output", type=Path); p.add_argument("--epochs", type=int); p.add_argument("--batch-size", type=int); p.add_argument("--device", choices=["auto", "cpu", "cuda"]); return p.parse_args()
def lr_at(progress, total, warmup, peak, floor):
    if warmup and progress < warmup: return peak * progress / warmup
    phase = min(max((progress - warmup) / max(total - warmup, 1), 0), 1)
    return floor + (peak - floor) * (1 + math.cos(math.pi * phase)) / 2

def main():
    a = args(); cfg = yaml.safe_load(a.config.read_text()); tc = cfg["training"]
    total_epochs = a.epochs or tc["epochs"]; world = int(os.getenv("WORLD_SIZE", "1")); rank = int(os.getenv("RANK", "0")); local = int(os.getenv("LOCAL_RANK", "0")); distributed = world > 1
    requested = a.device or cfg["runtime"]["device"]; cuda = torch.cuda.is_available() and requested != "cpu"
    if requested == "cuda" and not cuda: raise RuntimeError("CUDA requested but unavailable")
    if distributed: dist.init_process_group("nccl" if cuda else "gloo")
    device = torch.device(f"cuda:{local}" if cuda else "cpu"); torch.cuda.set_device(local) if cuda else None
    seed = cfg["seed"] + rank; random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
    data_path = a.data or ROOT / cfg["data"]["root"] / "train.npz"
    dataset = SolarDataset(data_path, cfg); sampler = DistributedSampler(dataset) if distributed else None
    loader = DataLoader(dataset, batch_size=a.batch_size or tc["batch_size"], shuffle=sampler is None, sampler=sampler, num_workers=tc["num_workers"], pin_memory=cuda)
    model = load_model()(**cfg["model"]).to(device); bare = model
    if distributed: model = DistributedDataParallel(model, device_ids=[local] if cuda else None); bare = model.module
    decay, no_decay = [], []
    for n, p in bare.named_parameters(): (no_decay if p.ndim == 1 or n.endswith("bias") else decay).append(p)
    opt = torch.optim.AdamW([{"params": decay, "weight_decay": tc["weight_decay"]}, {"params": no_decay, "weight_decay": 0}], lr=tc["learning_rate"])
    amp = bool(tc.get("amp", True) and cuda); scaler = torch.amp.GradScaler("cuda", enabled=amp); history=[]; opt.zero_grad(set_to_none=True)
    one_step = tc.get("one_step_epochs", max(1, total_epochs // 2))
    for epoch in range(total_epochs):
        if sampler: sampler.set_epoch(epoch)
        model.train(); total=0.0; phase = "one_step" if epoch < one_step else "rollout"
        for step, (x, y) in enumerate(loader):
            x, y = x.to(device, non_blocking=cuda), y.to(device, non_blocking=cuda); pred_steps = 1 if phase == "one_step" else y.shape[1]
            for group in opt.param_groups: group["lr"] = lr_at(epoch + step/max(len(loader),1), total_epochs, tc["warmup_epochs"], tc["learning_rate"], tc["min_learning_rate"])
            context = torch.amp.autocast("cuda") if amp else nullcontext()
            with context:
                pred = model(x, steps=pred_steps); loss = (pred - y[:, :pred_steps]).square().mean() / tc["accum_iter"]
            if not torch.isfinite(loss): raise FloatingPointError("non-finite training loss")
            scaler.scale(loss).backward()
            if (step + 1) % tc["accum_iter"] == 0 or step + 1 == len(loader):
                scaler.unscale_(opt); torch.nn.utils.clip_grad_norm_(bare.parameters(), tc["grad_clip"]); scaler.step(opt); scaler.update(); opt.zero_grad(set_to_none=True)
            total += loss.item() * tc["accum_iter"]
        values = torch.tensor([total / max(len(loader),1)], device=device); dist.all_reduce(values) if distributed else None
        record={"epoch":epoch+1,"phase":phase,"loss":float(values.item()/world),"learning_rate":opt.param_groups[0]["lr"]}; history.append(record)
        if rank == 0: print(json.dumps(record))
    if rank == 0:
        path=a.output or ROOT / cfg["paths"]["checkpoint"]; path.parent.mkdir(parents=True,exist_ok=True); torch.save({"model":bare.state_dict(),"optimizer":opt.state_dict(),"scaler":scaler.state_dict(),"epoch":total_epochs-1,"history":history,"config":cfg},path)
        out=ROOT / cfg["paths"]["training_metrics"]; out.parent.mkdir(parents=True,exist_ok=True); out.write_text(json.dumps({"history":history,"protocol":cfg["data"]["protocol"]},indent=2)+"\n"); print("checkpoint=",path)
    if distributed: dist.destroy_process_group()
if __name__ == "__main__": main()