# bora_train.py -- training loop for Bora V1 (bora_dit.BoraDiT on bora_dataset.BoraLatents).
#
# train(cfg, decode_fn=None) runs until cfg.max_steps, resumes from <out_dir>/ckpt_last.pt if it exists, and writes:
#   log.csv            step, loss, lr, samples_per_s                 (every log_every steps)
#   val.csv            step, val_loss, val_loss@t for a fixed t grid  (every val_every steps, EMA weights, fixed noise)
#   ckpt_last.pt       model + ema + optimizer + step                (every save_every steps; ~5.4 GB for v1, keep it on local disk)
#   ema_step{N}.safetensors + config.json                           (every save_every steps; bf16 EMA weights, ~0.7 GB for v1)
#   samples/step{N}.safetensors                                     (every sample_every steps: latents for cfg.sample_captions)
# decode_fn(latents [B,16,10,32,32] float32 cpu, prompts list[str], step int) is called after each sample, e.g. to write an mp4.
import csv
import json
import math
import os
import time
from dataclasses import dataclass, asdict, field
from pathlib import Path

import torch
from safetensors import safe_open
from safetensors.torch import save_file, load_file
from torch.utils.data import DataLoader

from bora_dataset import BoraLatents
from bora_dit import DiTConfig, BoraDiT, PRESETS, build, count_params, config_dict, flow_loss, shift_t, sample_video


@dataclass
class TrainConfig:
    data_root: str = "/content/bora_data"
    out_dir: str = "/content/bora_runs/v1"
    text_embeds: str = "text_t5v11base_32.safetensors"   # relative to data_root
    preset: str = "v1"
    model: dict = field(default_factory=dict)             # DiTConfig overrides, e.g. {"attn": "linear"}
    # data
    batch_size: int = 32
    num_workers: int = 4
    block_len: int = 4
    overlap: int = 1
    val_mod: int = 100
    primary_p: float = 0.5
    # objective
    p_cond: float = 0.5           # fraction of samples that get the clean first frame
    p_uncond: float = 0.1         # fraction of samples whose text is replaced by the empty caption (for CFG)
    time_shift: float = 3.0
    # optimisation
    lr: float = 1e-4
    warmup: int = 1000
    max_steps: int = 50000
    min_lr_frac: float = 0.1      # cosine decay from lr to lr * min_lr_frac over max_steps (after warmup)
    weight_decay: float = 0.0
    betas: tuple = (0.9, 0.95)
    grad_clip: float = 1.0
    ema: float = 0.9999
    bf16: bool = True
    grad_checkpoint: bool = False
    compile: bool = False
    seed: int = 0
    # bookkeeping
    log_every: int = 50
    val_every: int = 1000
    val_batches: int = 8
    val_t: tuple = (0.25, 0.5, 0.75, 0.9)    # fixed timesteps for the validation loss (after the shift)
    sample_every: int = 2000
    sample_steps: int = 8
    sample_cfg: float = 4.0
    sample_captions: tuple = ()               # caption strings to sample; empty -> 4 captions from the val set
    sample_blocks: int = 3
    save_every: int = 1000
    keep_ema_every: int = 10000               # keep every Nth ema_step file, delete the others
    overfit_batch: int = 0                    # >0: train on one fixed batch of this size (pipeline test)
    resume: bool = True


# ----------------------------------------------------------------------------- helpers

def set_seed(seed):
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)


def lr_at(step, cfg):
    if step < cfg.warmup:
        return cfg.lr * (step + 1) / cfg.warmup
    p = min(1.0, (step - cfg.warmup) / max(1, cfg.max_steps - cfg.warmup))
    return cfg.lr * (cfg.min_lr_frac + (1 - cfg.min_lr_frac) * 0.5 * (1 + math.cos(math.pi * p)))


def load_null_text(path):
    with safe_open(str(path), framework="pt") as f:
        return f.get_slice("emb")[0:1][0].float(), f.get_slice("mask")[0:1][0].bool()      # [L,D], [L]


class EMA:
    def __init__(self, model, decay):
        self.decay = decay
        self.shadow = {k: v.detach().clone().float() for k, v in model.state_dict().items()}

    @torch.no_grad()
    def update(self, model, step=None):
        # warmup: decay = min(decay, (1 + step) / (10 + step)), so the EMA tracks the model early instead of staying near the init
        d = self.decay if step is None else min(self.decay, (1 + step) / (10 + step))
        for k, v in model.state_dict().items():
            s = self.shadow[k]
            if s.dtype.is_floating_point:
                s.mul_(d).add_(v.detach().float(), alpha=1 - d)
            else:
                s.copy_(v)

    def state_dict(self):
        return self.shadow

    def load_state_dict(self, sd):
        self.shadow = {k: v.clone() for k, v in sd.items()}

    def to_model(self, model_cfg, device, dtype=torch.bfloat16):
        m = BoraDiT(model_cfg)
        m.load_state_dict({k: v for k, v in self.shadow.items()})
        return m.to(device, dtype).eval()


def make_model(cfg, device):
    model, mcfg = build(cfg.preset, time_shift=cfg.time_shift, **cfg.model)
    model.grad_checkpoint = cfg.grad_checkpoint
    model.to(device)
    return model, mcfg


def make_data(cfg):
    root = Path(cfg.data_root)
    emb = root / cfg.text_embeds
    kw = dict(mode="block", block_len=cfg.block_len, overlap=cfg.overlap, val_mod=cfg.val_mod, text_embeds=emb, dtype=torch.float32)
    train_ds = BoraLatents(root, "train", random_block=True, primary_p=cfg.primary_p, **kw)
    val_ds = BoraLatents(root, "val", random_block=False, **kw)
    return train_ds, val_ds


def infinite(loader):
    while True:
        for b in loader:
            yield b


def to_device(batch, device):
    return {k: (v.to(device, non_blocking=True) if torch.is_tensor(v) else v) for k, v in batch.items()}


def apply_uncond(text, mask, p, null_text, null_mask):
    """Replace a random fraction p of the text rows by the empty caption (classifier-free guidance training)."""
    if p <= 0:
        return text, mask
    drop = torch.rand(text.shape[0], device=text.device) < p
    text = torch.where(drop[:, None, None], null_text[None].to(text.dtype), text)
    mask = torch.where(drop[:, None], null_mask[None], mask)
    return text, mask


@torch.no_grad()
def evaluate(model, val_dl, cfg, device, null_text, null_mask, n_batches):
    """Loss of `model` on the val set at fixed timesteps with fixed noise. Returns {"val_loss": .., "t=0.25": .., ...}."""
    model.eval()
    g = torch.Generator(device=device).manual_seed(1234)
    sums = {t: 0.0 for t in cfg.val_t}
    n = 0
    for i, batch in enumerate(val_dl):
        if i >= n_batches:
            break
        batch = to_device(batch, device)
        x = batch["latents"]
        text, mask = batch["text_emb"].float(), batch["text_mask"]
        noise = torch.randn(x.shape, device=device, generator=g)
        for t in cfg.val_t:
            tt = torch.full((x.shape[0],), t, device=device)
            with torch.autocast(device_type=device.type, dtype=torch.bfloat16, enabled=cfg.bf16 and device.type == "cuda"):
                loss, _ = flow_loss(model, x, text, mask, p_cond=cfg.p_cond, shift=cfg.time_shift, t=tt, noise=noise)
            sums[t] += loss.item() * x.shape[0]
        n += x.shape[0]
    model.train()
    out = {f"t={t}": sums[t] / max(1, n) for t in cfg.val_t}
    out["val_loss"] = sum(out.values()) / len(cfg.val_t)
    return out


def pick_sample_captions(cfg, val_ds, train_ds):
    if cfg.sample_captions:
        caps = train_ds.captions
        ids = []
        for c in cfg.sample_captions:
            assert c in caps, f"caption not in captions.json: {c!r} (sample any string by encoding it with T5 in the notebook)"
            ids.append(caps.index(c))
        return ids
    step = max(1, len(val_ds) // (3 * 4))                       # 4 clips spread over the val set (3 blocks per clip)
    return [val_ds[i * 3 * step]["caption_id"] for i in range(4)]


@torch.no_grad()
def sample_captions(ema_model, caption_ids, emb_path, cfg, device, null_text, null_mask, seed=0):
    with safe_open(str(emb_path), framework="pt") as f:
        text = torch.stack([f.get_slice("emb")[i:i + 1][0] for i in caption_ids]).float().to(device)
        mask = torch.stack([f.get_slice("mask")[i:i + 1][0] for i in caption_ids]).bool().to(device)
    B = len(caption_ids)
    g = torch.Generator(device=device).manual_seed(seed)
    return sample_video(ema_model, text, mask, null_text[None].expand(B, -1, -1).to(device), null_mask[None].expand(B, -1).to(device),
                        n_blocks=cfg.sample_blocks, overlap=cfg.overlap, steps=cfg.sample_steps, cfg_scale=cfg.sample_cfg,
                        shift=cfg.time_shift, generator=g).float().cpu()


def save_checkpoint(out, model, ema, opt, step, cfg, mcfg):
    tmp = out / "ckpt_last.pt.tmp"
    torch.save({"model": model.state_dict(), "ema": ema.state_dict(), "opt": opt.state_dict(), "step": step,
                "cfg": asdict(cfg), "model_cfg": config_dict(mcfg)}, tmp)
    os.replace(tmp, out / "ckpt_last.pt")
    ema_sd = {k: v.to(torch.bfloat16) for k, v in ema.state_dict().items() if v.dtype.is_floating_point}
    ema_sd.update({k: v for k, v in ema.state_dict().items() if not v.dtype.is_floating_point})
    save_file(ema_sd, str(out / f"ema_step{step}.safetensors"), metadata={"step": str(step), "model_cfg": json.dumps(config_dict(mcfg))})
    (out / "config.json").write_text(json.dumps({"train": asdict(cfg), "model": config_dict(mcfg), "step": step}, indent=1))
    for p in sorted(out.glob("ema_step*.safetensors")):                 # prune
        s = int(p.stem[len("ema_step"):])
        if s != step and s % cfg.keep_ema_every != 0:
            p.unlink()


def load_ema_model(path, device, dtype=torch.bfloat16):
    """ema_stepN.safetensors -> BoraDiT in eval mode."""
    with safe_open(str(path), framework="pt") as f:
        mcfg = json.loads(f.metadata()["model_cfg"])
    mcfg["patch"] = tuple(mcfg["patch"])
    m = BoraDiT(DiTConfig(**mcfg))
    m.load_state_dict(load_file(str(path)))
    return m.to(device, dtype).eval()


# ----------------------------------------------------------------------------- the loop

def train(cfg: TrainConfig, decode_fn=None, log=print):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    set_seed(cfg.seed)
    out = Path(cfg.out_dir)
    out.mkdir(parents=True, exist_ok=True)
    (out / "samples").mkdir(exist_ok=True)
    torch.backends.cuda.matmul.allow_tf32 = True
    torch.backends.cudnn.allow_tf32 = True

    # data
    train_ds, val_ds = make_data(cfg)
    emb_path = Path(cfg.data_root) / cfg.text_embeds
    null_text, null_mask = load_null_text(emb_path)
    null_text, null_mask = null_text.to(device), null_mask.to(device)
    pin = device.type == "cuda"
    if cfg.overfit_batch:
        fixed = next(iter(DataLoader(train_ds, batch_size=cfg.overfit_batch, shuffle=True, num_workers=0)))
        fixed["text_emb"] = fixed["text_emb"].float()
        train_it = iter(lambda: fixed, None)
        save_file({"latents": fixed["latents"].contiguous(), "global_idx": fixed["global_idx"], "start": fixed["start"]},
                  str(out / "overfit_batch.safetensors"), metadata={"prompts": json.dumps(fixed["text"])})
        log(f"overfit mode: one fixed batch of {cfg.overfit_batch} | captions: {[t[:40] for t in fixed['text']]}")
    else:
        train_dl = DataLoader(train_ds, batch_size=cfg.batch_size, shuffle=True, num_workers=cfg.num_workers, drop_last=True,
                              pin_memory=pin, persistent_workers=cfg.num_workers > 0, prefetch_factor=4 if cfg.num_workers > 0 else None)
        train_it = infinite(train_dl)
    val_dl = DataLoader(val_ds, batch_size=cfg.batch_size, shuffle=False, num_workers=min(2, cfg.num_workers), pin_memory=pin)
    log(f"data: train {len(train_ds)} clips / {train_ds.n_videos} videos | val {len(val_ds)} blocks / {val_ds.n_videos} videos | captions {len(train_ds.captions) - 1}")

    # model
    model, mcfg = make_model(cfg, device)
    log(f"model {cfg.preset}: {count_params(model) / 1e6:.1f}M params | dim {mcfg.dim} depth {mcfg.depth} heads {mcfg.heads} attn {mcfg.attn} | tokens {mcfg.tokens}")
    opt = torch.optim.AdamW(model.parameters(), lr=cfg.lr, betas=cfg.betas, weight_decay=cfg.weight_decay, fused=device.type == "cuda")
    ema = EMA(model, cfg.ema)
    step = 0
    ck = out / "ckpt_last.pt"
    if cfg.resume and ck.exists():
        sd = torch.load(ck, map_location=device, weights_only=False)
        model.load_state_dict(sd["model"]); ema.load_state_dict(sd["ema"]); opt.load_state_dict(sd["opt"]); step = sd["step"]
        log(f"resumed from {ck} at step {step}")
    fwd = torch.compile(model) if cfg.compile else model

    sample_ids = pick_sample_captions(cfg, val_ds, train_ds) if not cfg.overfit_batch else [int(i) for i in fixed["caption_id"][:4]]
    sample_prompts = [train_ds.captions[i] for i in sample_ids]
    log(f"sample prompts: {[p[:50] for p in sample_prompts]}")

    def write_row(name, row):
        p = out / name
        new = not p.exists()
        with open(p, "a", newline="") as f:
            w = csv.DictWriter(f, fieldnames=list(row.keys()))
            if new:
                w.writeheader()
            w.writerow(row)

    def do_sample(step):
        ema_model = ema.to_model(mcfg, device)
        lat = sample_captions(ema_model, sample_ids, emb_path, cfg, device, null_text, null_mask)
        save_file({"latents": lat}, str(out / "samples" / f"step{step}.safetensors"), metadata={"prompts": json.dumps(sample_prompts)})
        del ema_model
        if decode_fn is not None:
            decode_fn(lat, sample_prompts, step)
        torch.cuda.empty_cache() if device.type == "cuda" else None

    # loop
    model.train()
    t0, n_since, loss_acc, loss_n = time.time(), 0, 0.0, 0
    log(f"training from step {step} to {cfg.max_steps} | batch {cfg.overfit_batch or cfg.batch_size} | bf16 {cfg.bf16} | grad_checkpoint {cfg.grad_checkpoint} | device {device}")
    while step < cfg.max_steps:
        batch = to_device(next(train_it), device)
        x = batch["latents"]
        text, mask = apply_uncond(batch["text_emb"].float(), batch["text_mask"], cfg.p_uncond, null_text, null_mask)
        for g in opt.param_groups:
            g["lr"] = lr_at(step, cfg)
        with torch.autocast(device_type=device.type, dtype=torch.bfloat16, enabled=cfg.bf16 and device.type == "cuda"):
            loss, parts = flow_loss(fwd, x, text, mask, p_cond=cfg.p_cond, shift=cfg.time_shift)
        opt.zero_grad(set_to_none=True)
        loss.backward()
        gn = torch.nn.utils.clip_grad_norm_(model.parameters(), cfg.grad_clip) if cfg.grad_clip else torch.tensor(0.0)
        if not torch.isfinite(loss):
            log(f"step {step}: non-finite loss, batch skipped")
            continue
        opt.step()
        ema.update(model, step)
        step += 1
        loss_acc += loss.item(); loss_n += 1; n_since += x.shape[0]

        if step % cfg.log_every == 0:
            dt = time.time() - t0
            row = {"step": step, "loss": loss_acc / loss_n, "lr": lr_at(step - 1, cfg), "grad_norm": float(gn), "samples_per_s": n_since / dt, "t_mean": parts["t_mean"]}
            write_row("log.csv", row)
            log(f"step {step} | loss {row['loss']:.4f} | lr {row['lr']:.2e} | gn {row['grad_norm']:.2f} | {row['samples_per_s']:.1f} samples/s")
            t0, n_since, loss_acc, loss_n = time.time(), 0, 0.0, 0
        if cfg.val_every and step % cfg.val_every == 0 and not cfg.overfit_batch:
            ema_model = ema.to_model(mcfg, device, dtype=torch.float32)
            v = evaluate(ema_model, val_dl, cfg, device, null_text, null_mask, cfg.val_batches)
            del ema_model
            write_row("val.csv", {"step": step, **v})
            log(f"step {step} | val {v['val_loss']:.4f} | " + " ".join(f"{k} {val:.3f}" for k, val in v.items() if k != "val_loss"))
        if cfg.sample_every and step % cfg.sample_every == 0:
            do_sample(step)
        if cfg.save_every and step % cfg.save_every == 0:
            save_checkpoint(out, model, ema, opt, step, cfg, mcfg)
            log(f"step {step} | saved ckpt_last.pt and ema_step{step}.safetensors")
    if not (cfg.save_every and step % cfg.save_every == 0):
        save_checkpoint(out, model, ema, opt, step, cfg, mcfg)
    if cfg.sample_every and step % cfg.sample_every != 0:
        do_sample(step)
    log(f"done at step {step}")
    return out