#!/usr/bin/env python3
# train.py — VSpark 0.5B basic training script
# For the memory-efficient Colab version use train_colab.py

import argparse, json, torch, torch.nn.functional as F
from pathlib import Path
from torch.utils.data import Dataset, DataLoader
from torch.optim import AdamW
from torch.optim.lr_scheduler import CosineAnnealingLR
try:
    from vspark_model import VSparkModel, VSparkConfig, DDPMScheduler
except ImportError:
    from vspark_build_fixed import VSparkModel, VSparkConfig, DDPMScheduler


class DummyVideoDataset(Dataset):
    def __init__(self, cfg, size=200):
        self.cfg  = cfg
        self.size = size

    def __len__(self):
        return self.size

    def __getitem__(self, _):
        cfg = self.cfg
        Hl  = cfg.latent_size
        return dict(
            video_latents = torch.randn(cfg.num_frames, cfg.in_channels, Hl, Hl),
            input_ids     = torch.randint(0, 49408, (cfg.max_text_len,)),
            audio_mel     = torch.randn(cfg.mel_frames, cfg.n_mels),
        )


def train(args):
    device = "cuda" if torch.cuda.is_available() else "cpu"
    cfg    = VSparkConfig()
    model     = VSparkModel(cfg).to(device)
    scheduler = DDPMScheduler(cfg).to(device)
    dataset   = DummyVideoDataset(cfg)
    loader    = DataLoader(dataset, batch_size=args.batch_size, shuffle=True, num_workers=0)
    optim     = AdamW(model.parameters(), lr=args.lr, weight_decay=1e-2)
    lr_sched  = CosineAnnealingLR(optim, T_max=args.epochs * len(loader), eta_min=1e-6)
    out_dir   = Path(args.output_dir);  out_dir.mkdir(parents=True, exist_ok=True)

    step = 0
    for epoch in range(args.epochs):
        model.train()
        for batch in loader:
            lat  = batch["video_latents"].to(device)
            ids  = batch["input_ids"].to(device)
            mel  = batch["audio_mel"].to(device)
            B    = lat.shape[0]
            t    = torch.randint(0, cfg.num_timesteps, (B,), device=device)
            noisy_lat, noise_lat = scheduler.add_noise(lat, t)
            noisy_mel, noise_mel = scheduler.add_noise(mel, t)
            pred_lat = model(noisy_lat, t, ids)
            pred_mel = model.predict_audio(noisy_mel, t, ids)
            loss = F.mse_loss(pred_lat, noise_lat) + 0.5 * F.mse_loss(pred_mel, noise_mel)
            optim.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optim.step(); lr_sched.step()
            step += 1
            if step % args.log_every == 0:
                print(f"  step={step}  loss={loss.item():.4f}")
        ckpt = out_dir / f"ckpt_epoch{epoch+1:03d}.pt"
        torch.save({"model_state_dict": model.state_dict(), "epoch": epoch + 1}, ckpt)


if __name__ == "__main__":
    p = argparse.ArgumentParser()
    p.add_argument("--output_dir", default="./checkpoints")
    p.add_argument("--epochs",     type=int,   default=100)
    p.add_argument("--batch_size", type=int,   default=1)
    p.add_argument("--lr",         type=float, default=1e-4)
    p.add_argument("--log_every",  type=int,   default=10)
    train(p.parse_args())