#!/usr/bin/env python3
"""
LatentRecurrentFlow v3 — Complete Self-Contained Training & Generation
======================================================================

This is one file. Run it and it:
1. Downloads TAESD (pre-trained tiny VAE, 2.4M params, frozen)
2. Pre-computes CIFAR-10 latents (~4 min on CPU)
3. Trains a 1.5M-param recursive flow-matching denoiser (30 epochs, ~60 min CPU)
4. Generates class-conditional 32×32 images and saves them
5. Saves everything to a HuggingFace repo

No custom VAE training. No grey images. Works end-to-end.
"""

import math, os, sys, time, json
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, TensorDataset
from einops import rearrange
from typing import Optional, Dict, Any
import numpy as np

DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
print(f"Device: {DEVICE}")

# =============================================================================
# 1. MODEL ARCHITECTURE
# =============================================================================

class RMSNorm(nn.Module):
    def __init__(self, d, eps=1e-6):
        super().__init__()
        self.eps = eps
        self.w = nn.Parameter(torch.ones(d))
    def forward(self, x):
        return x * (x.float().pow(2).mean(-1, keepdim=True) + self.eps).rsqrt().type_as(x) * self.w

class SwiGLU(nn.Module):
    def __init__(self, d, hd=None):
        super().__init__()
        hd = hd or ((d * 8 // 3 + 7) // 8 * 8)
        self.w1 = nn.Linear(d, hd, bias=False)
        self.w2 = nn.Linear(hd, d, bias=False)
        self.w3 = nn.Linear(d, hd, bias=False)
    def forward(self, x):
        return self.w2(F.silu(self.w1(x)) * self.w3(x))

class SpatialMixer(nn.Module):
    """Multi-head self-attention + depthwise conv for 2D locality."""
    def __init__(self, d, nh=4, hd=32):
        super().__init__()
        self.nh, self.hd = nh, hd
        inner = nh * hd
        self.qkv = nn.Linear(d, 3 * inner, bias=False)
        self.out = nn.Linear(inner, d, bias=False)
        self.gate = nn.Sequential(nn.Linear(d, inner, bias=False), nn.SiLU())
        self.dw = nn.Conv2d(inner, inner, 3, padding=1, groups=inner, bias=False)
        self.norm = RMSNorm(inner)

    def forward(self, x, h, w):
        B, N, D = x.shape
        q, k, v = self.qkv(x).chunk(3, dim=-1)
        q = rearrange(q, 'b n (h d) -> b h n d', h=self.nh)
        k = rearrange(k, 'b n (h d) -> b h n d', h=self.nh)
        v = rearrange(v, 'b n (h d) -> b h n d', h=self.nh)
        a = (q @ k.transpose(-2, -1)) * (self.hd ** -0.5)
        o = rearrange(F.softmax(a, -1) @ v, 'b h n d -> b n (h d)')
        o = self.norm(o)
        # 2D locality
        inner = self.nh * self.hd
        lc = self.dw(rearrange(x[:, :, :inner], 'b (h w) d -> b d h w', h=h, w=w))
        lc = rearrange(lc, 'b d h w -> b (h w) d')
        return self.out(self.gate(x) * o + 0.1 * lc)

class DenoiseBlock(nn.Module):
    """AdaLN-modulated SpatialMixer + cross-attn + SwiGLU."""
    def __init__(self, d, cd, nh=4, hd=32, fm=2.67):
        super().__init__()
        self.n1 = RMSNorm(d); self.n2 = RMSNorm(d)
        self.mix = SpatialMixer(d, nh, hd)
        self.ffn = SwiGLU(d, int(d * fm))
        self.mod = nn.Sequential(nn.SiLU(), nn.Linear(cd, 6 * d))
        # Cross-attn to condition tokens
        self.cn = RMSNorm(d)
        self.cq = nn.Linear(d, d, bias=False)
        self.ckv = nn.Linear(cd, 2 * d, bias=False)
        self.co = nn.Linear(d, d, bias=False)
        self.cs = nn.Parameter(torch.zeros(1))

    def forward(self, x, c, ctx=None, h=4, w=4):
        s1, h1, g1, s2, h2, g2 = self.mod(c).chunk(6, -1)
        xn = self.n1(x) * (1 + s1.unsqueeze(1)) + h1.unsqueeze(1)
        x = x + g1.unsqueeze(1) * self.mix(xn, h, w)
        if ctx is not None:
            xc = self.cn(x); q = self.cq(xc)
            k, v = self.ckv(ctx).chunk(2, -1)
            a = F.softmax(q @ k.transpose(-2,-1) * (q.shape[-1]**-0.5), -1)
            x = x + self.cs.tanh() * self.co(a @ v)
        xn = self.n2(x) * (1 + s2.unsqueeze(1)) + h2.unsqueeze(1)
        x = x + g2.unsqueeze(1) * self.ffn(xn)
        return x

class RecursiveCore(nn.Module):
    """N shared blocks × T recursions with abstract state + IFT training."""
    def __init__(self, lc=4, d=128, cd=128, nb=4, nh=4, hd=32,
                 ti=2, to=1, fm=2.67, ift=False):
        super().__init__()
        self.nb, self.ti, self.to_, self.ift = nb, ti, to, ift
        self.inp = nn.Linear(lc, d)
        self.tmb = nn.Sequential(nn.Linear(256, cd), nn.SiLU(), nn.Linear(cd, cd))
        self.blk = nn.ModuleList([DenoiseBlock(d, cd, nh, hd, fm) for _ in range(nb)])
        self.ag = nn.Parameter(torch.tensor(0.0))
        self.ap = nn.Sequential(nn.Linear(d, d, bias=False), nn.SiLU(), nn.Linear(d, d, bias=False))
        self.se = nn.Embedding(to * ti + 1, cd)
        self.on = RMSNorm(d)
        self.op = nn.Linear(d, lc)
        nn.init.zeros_(self.op.weight); nn.init.zeros_(self.op.bias)

    def _sinemb(self, t, d=256):
        h = d // 2
        f = torch.exp(torch.arange(h, device=t.device).float() * -(math.log(10000) / h))
        a = t.unsqueeze(-1) * f.unsqueeze(0)
        return torch.cat([a.sin(), a.cos()], -1)

    def _blocks(self, z, c, ctx, h, w):
        for b in self.blk: z = b(z, c, ctx, h, w)
        return z

    def _refine(self, z, cb, ctx, h, w):
        za = z.mean(1, keepdim=True).expand_as(z)
        s = 0
        for j in range(self.to_):
            za = za + self.ag.tanh() * self.ap(z.mean(1, keepdim=True).expand_as(z))
            for i in range(self.ti):
                se = self.se(torch.tensor([s], device=z.device)).expand(z.shape[0], -1)
                zn = self._blocks(z + za, cb + se, ctx, h, w)
                z = z + 0.5 * (zn - z)
                s += 1
        return z

    def forward(self, zt, t, tg=None, ic=None):
        B, C, H, W = zt.shape
        z = self.inp(rearrange(zt, 'b c h w -> b (h w) c'))
        if ic is not None:
            z = z + self.inp(rearrange(ic, 'b c h w -> b (h w) c'))
        c = self.tmb(self._sinemb(t))
        if tg is not None: c = c + tg
        if self.training and self.ift and self.to_ > 1:
            with torch.no_grad():
                for _ in range(self.to_ - 1): z = self._refine(z, c, None, H, W)
            z = self._refine(z, c, None, H, W)
        else:
            z = self._refine(z, c, None, H, W)
        return rearrange(self.op(self.on(z)), 'b (h w) c -> b c h w', h=H, w=W)

class LRF(nn.Module):
    """Complete model: RecursiveCore + class conditioner."""
    def __init__(self, cfg=None):
        super().__init__()
        cfg = cfg or self.default()
        self.cfg = cfg
        nc = cfg.get('nc', 10)
        self.core = RecursiveCore(
            lc=cfg['lc'], d=cfg['d'], cd=cfg['cd'], nb=cfg['nb'],
            nh=cfg['nh'], hd=cfg['hd'], ti=cfg['ti'], to=cfg['to'],
            fm=cfg.get('fm', 2.67), ift=cfg.get('ift', False))
        self.cemb = nn.Embedding(nc + 1, cfg['cd'])
        self.null = nc

    @staticmethod
    def default():
        return dict(lc=4, d=128, cd=128, nb=4, nh=4, hd=32, ti=2, to=1, fm=2.0, ift=False, nc=10)

    def predict_v(self, zt, t, cls=None, cfg_drop=0.0):
        B = zt.shape[0]
        if cls is not None:
            if self.training and cfg_drop > 0:
                m = torch.rand(B, device=zt.device) < cfg_drop
                cls = cls.clone(); cls[m] = self.null
            c = self.cemb(cls)
        else:
            c = self.cemb(torch.full((B,), self.null, device=zt.device, dtype=torch.long))
        return self.core(zt, t, tg=c)

    def count(self):
        return sum(p.numel() for p in self.parameters())

# =============================================================================
# 2. FLOW SCHEDULER
# =============================================================================

class FlowScheduler:
    def add_noise(self, z0, eps, t):
        t = t.view(-1, 1, 1, 1)
        return (1 - t) * z0 + t * eps

    def velocity(self, z0, eps):
        return eps - z0

    def sample_t(self, B, dev):
        return torch.rand(B, device=dev).clamp(1e-4, 1 - 1e-4)

    @torch.no_grad()
    def sample(self, model, shape, cls=None, steps=50, cfg=3.0, dev='cpu'):
        z = torch.randn(shape, device=dev)
        ts = torch.linspace(1, 0, steps + 1, device=dev)
        for i in range(steps):
            tb = torch.full((shape[0],), ts[i].item(), device=dev)
            dt = ts[i] - ts[i + 1]
            if cfg > 1.0 and cls is not None:
                vc = model.predict_v(z, tb, cls)
                vu = model.predict_v(z, tb, None)
                v = vu + cfg * (vc - vu)
            else:
                v = model.predict_v(z, tb, cls)
            z = z - dt * v
        return z

# =============================================================================
# 3. DATA + TAESD
# =============================================================================

def get_taesd(dev='cpu'):
    from diffusers import AutoencoderTiny
    vae = AutoencoderTiny.from_pretrained('madebyollin/taesd', torch_dtype=torch.float32)
    vae.eval().to(dev)
    for p in vae.parameters(): p.requires_grad_(False)
    return vae

def get_cifar(root='/app/data'):
    import torchvision, torchvision.transforms as T
    tf = T.Compose([T.ToTensor(), T.Normalize([.5]*3, [.5]*3)])
    tr = torchvision.datasets.CIFAR10(root, True, tf, download=True)
    te = torchvision.datasets.CIFAR10(root, False, tf, download=True)
    return tr, te

def precompute(vae, ds, bs=256, dev='cpu', cache=None):
    if cache and os.path.exists(cache):
        print(f"  Loading cached latents from {cache}", flush=True)
        d = torch.load(cache, weights_only=True)
        return d['lat'], d['lab']
    dl = DataLoader(ds, bs, shuffle=False, num_workers=0)
    lats, labs = [], []
    t0 = time.time()
    with torch.no_grad():
        for i, (img, lab) in enumerate(dl):
            lats.append(vae.encode(img.to(dev)).latents.cpu())
            labs.append(lab)
            if (i+1) % 50 == 0 or i == 0:
                print(f"    batch {i+1}/{len(dl)} ({time.time()-t0:.0f}s)", flush=True)
    lats, labs = torch.cat(lats), torch.cat(labs)
    if cache:
        os.makedirs(os.path.dirname(cache) or '.', exist_ok=True)
        torch.save({'lat': lats, 'lab': labs}, cache)
    print(f"  Done: {lats.shape}, mean={lats.mean():.3f}, std={lats.std():.3f}", flush=True)
    return lats, labs

# =============================================================================
# 4. TRAINING
# =============================================================================

def train(epochs=30, bs=64, lr=3e-4, dev=DEVICE, out='/app/lrf_out'):
    os.makedirs(out, exist_ok=True)
    print("=" * 60, flush=True)
    print("LatentRecurrentFlow v3 — Training on CIFAR-10", flush=True)
    print("=" * 60, flush=True)

    # VAE
    print("\n[1/5] Loading TAESD...", flush=True)
    vae = get_taesd(dev)
    print(f"  TAESD: {sum(p.numel() for p in vae.parameters()):,} params (frozen)", flush=True)

    # Data
    print("\n[2/5] Loading CIFAR-10 + precomputing latents...", flush=True)
    tr, te = get_cifar()
    tr_lat, tr_lab = precompute(vae, tr, 256, dev, f'{out}/cache_train.pt')
    te_lat, te_lab = precompute(vae, te, 256, dev, f'{out}/cache_test.pt')

    # Verify VAE works
    print("\n[2b] VAE sanity check...", flush=True)
    with torch.no_grad():
        imgs = torch.stack([tr[i][0] for i in range(8)]).to(dev)
        rec = vae.decode(vae.encode(imgs).latents).sample
        mse = F.mse_loss(rec, imgs).item()
        print(f"  Recon MSE = {mse:.4f} (should be <0.1)", flush=True)
        save_grid(torch.cat([imgs[:4].cpu(), rec[:4].cpu()]), f'{out}/vae_check.png', 4)

    # Model
    print("\n[3/5] Creating model...", flush=True)
    cfg = LRF.default()
    model = LRF(cfg).to(dev)
    print(f"  Params: {model.count():,}", flush=True)
    print(f"  Depth: {cfg['to']}×{cfg['ti']}×{cfg['nb']} = {cfg['to']*cfg['ti']*cfg['nb']} eff. layers", flush=True)

    # Train
    print(f"\n[4/5] Training {epochs} epochs, bs={bs}, lr={lr}...", flush=True)
    sched = FlowScheduler()
    opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01, betas=(0.9, 0.95))
    lrs = torch.optim.lr_scheduler.CosineAnnealingLR(opt, epochs * (len(tr_lat)//bs), lr*0.01)
    ema = {n: p.clone().detach() for n, p in model.named_parameters()}
    losses = []

    dl = DataLoader(TensorDataset(tr_lat, tr_lab), bs, shuffle=True, drop_last=True)
    t0 = time.time()
    for ep in range(epochs):
        model.train()
        el = 0; nb = 0
        for lat, lab in dl:
            lat, lab = lat.to(dev), lab.to(dev)
            B = lat.shape[0]
            t = sched.sample_t(B, dev)
            eps = torch.randn_like(lat)
            zt = sched.add_noise(lat, eps, t)
            vp = model.predict_v(zt, t, lab, cfg_drop=0.1)
            vt = sched.velocity(lat, eps)
            lps = (vp - vt).pow(2).mean([1,2,3])
            w = 1.0 / (t * (1-t) + 0.01); w = w / w.mean()
            loss = (lps * w).mean()
            opt.zero_grad(); loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            opt.step(); lrs.step()
            with torch.no_grad():
                for n, p in model.named_parameters():
                    ema[n].mul_(0.999).add_(p, alpha=0.001)
            el += loss.item(); nb += 1
        al = el / nb; losses.append(al)
        elapsed = time.time() - t0
        if (ep+1) % 5 == 0 or ep == 0:
            print(f"  Ep {ep+1:3d}/{epochs}: loss={al:.4f}, lr={opt.param_groups[0]['lr']:.1e}, "
                  f"time={elapsed:.0f}s", flush=True)
        if (ep+1) % 10 == 0 or ep == epochs - 1:
            # Sample with EMA
            bak = {n: p.clone() for n, p in model.named_parameters()}
            with torch.no_grad():
                for n, p in model.named_parameters(): p.copy_(ema[n])
            model.eval()
            samps = gen(model, vae, sched, dev, 16, 20, 2.0)
            save_grid(samps, f'{out}/ep{ep+1:03d}.png', 4)
            with torch.no_grad():
                for n, p in model.named_parameters(): p.copy_(bak[n])

    # Final: swap to EMA permanently
    with torch.no_grad():
        for n, p in model.named_parameters(): p.copy_(ema[n])
    model.eval()

    # Save checkpoint
    torch.save({'state': model.state_dict(), 'cfg': cfg, 'losses': losses}, f'{out}/model.pt')

    # Final generation
    print(f"\n[5/5] Generating final samples...", flush=True)
    classes = ['airplane','auto','bird','cat','deer','dog','frog','horse','ship','truck']
    all_s = []
    for ci in range(10):
        s = gen(model, vae, sched, dev, 4, 50, 3.0, ci)
        all_s.append(s)
        std = s.std().item()
        print(f"  {classes[ci]:10s}: std={std:.3f} {'✓' if std > 0.1 else '✗ FLAT'}", flush=True)
    save_grid(torch.cat(all_s), f'{out}/final.png', 4)

    # Loss plot
    try:
        import matplotlib; matplotlib.use('Agg'); import matplotlib.pyplot as plt
        plt.figure(figsize=(8,3)); plt.plot(losses); plt.xlabel('Epoch'); plt.ylabel('Loss')
        plt.title('Training Loss'); plt.grid(True, alpha=.3)
        plt.savefig(f'{out}/loss.png', dpi=100, bbox_inches='tight'); plt.close()
    except: pass

    print(f"\n{'='*60}", flush=True)
    print(f"DONE. Best loss: {min(losses):.4f}. Files in {out}/", flush=True)
    print(f"{'='*60}", flush=True)
    return model, vae, losses

def gen(model, vae, sched, dev, n=8, steps=20, cfg=2.0, cls_id=None):
    cls = torch.full((n,), cls_id if cls_id is not None else 0, dtype=torch.long, device=dev)
    if cls_id is None: cls = torch.randint(0, 10, (n,), device=dev)
    z = sched.sample(model, (n, 4, 4, 4), cls, steps, cfg, dev)
    with torch.no_grad(): imgs = vae.decode(z.to(dev)).sample.clamp(-1, 1)
    return imgs.cpu()

def save_grid(imgs, path, nr=8):
    from PIL import Image
    imgs = ((imgs + 1) / 2).clamp(0, 1)
    import torchvision
    g = torchvision.utils.make_grid(imgs, nrow=nr, padding=2)
    arr = (g.permute(1,2,0).numpy() * 255).astype(np.uint8)
    Image.fromarray(arr).save(path)
    print(f"  Saved: {path}", flush=True)

# =============================================================================
# 5. NOTEBOOK CONVERSION (Jupyter-compatible cells as functions)
# =============================================================================

def notebook_cell_1_setup():
    """Cell 1: Install & import."""
    print("Installing dependencies...")
    os.system("pip install -q torch torchvision einops diffusers safetensors huggingface_hub matplotlib pillow")
    print("Done.")

def notebook_cell_2_architecture():
    """Cell 2: Show architecture details."""
    cfg = LRF.default()
    model = LRF(cfg)
    print(f"LatentRecurrentFlow Architecture")
    print(f"================================")
    print(f"Latent channels: {cfg['lc']} (TAESD)")
    print(f"Model dim:       {cfg['d']}")
    print(f"Shared blocks:   {cfg['nb']}")
    print(f"Recursions:      {cfg['to']}×{cfg['ti']} = {cfg['to']*cfg['ti']}")
    print(f"Effective depth: {cfg['to']*cfg['ti']*cfg['nb']} layers")
    print(f"Total params:    {model.count():,}")
    print(f"FP32 size:       {model.count()*4/1e6:.1f} MB")
    print(f"INT8 size:       {model.count()/1e6:.1f} MB")
    return model

def notebook_cell_3_train():
    """Cell 3: Full training loop."""
    return train(epochs=30, bs=64, lr=3e-4, out='/app/lrf_out')

def notebook_cell_4_generate(model, vae):
    """Cell 4: Generate and display images."""
    sched = FlowScheduler()
    classes = ['airplane','auto','bird','cat','deer','dog','frog','horse','ship','truck']
    for ci, cn in enumerate(classes):
        imgs = gen(model, vae, sched, DEVICE, 4, 50, 3.0, ci)
        save_grid(imgs, f'/app/lrf_out/class_{cn}.png', 4)
    print("All class images generated!")

def notebook_cell_5_push(repo_id='krystv/LatentRecurrentFlow'):
    """Cell 5: Push to HuggingFace Hub."""
    from huggingface_hub import HfApi
    api = HfApi()
    out = '/app/lrf_out'
    for f in os.listdir(out):
        if f.endswith(('.pt', '.png', '.json')):
            fp = os.path.join(out, f)
            if os.path.getsize(fp) < 50_000_000:  # Skip huge files
                api.upload_file(path_or_fileobj=fp, path_in_repo=f'v3/{f}',
                                repo_id=repo_id, repo_type='model')
                print(f"  Uploaded v3/{f}")
    print(f"Done! See https://huggingface.co/{repo_id}")

# =============================================================================
# MAIN
# =============================================================================

if __name__ == '__main__':
    model, vae, losses = train(epochs=30, bs=64, lr=3e-4, out='/app/lrf_out')