"""Fine-tune TAEF2 into a better tiny VAE for FLUX 3 Klein (same VAE as FLUX.2).

Teacher: the full FLUX 3 Klein / FLUX.2 VAE encoder makes the target latents (in the
BN-normalized space the transformer uses). Student: TAEF2's decoder learns to turn those
latents back into the real image. Optionally the student encoder is trained to match the
teacher latents too.

python finetune_tiny_vae.py --data data/div2k --steps 3000 --batch 16 --res 512
"""
import argparse, copy, glob, math, os, random, time

import lpips
import safetensors.torch as stt
import torch
import torch.nn as nn
import torch.nn.functional as F
from PIL import Image
from torch.utils.data import DataLoader, Dataset
from torchvision.transforms import functional as TF
from diffusers import AutoencoderKLFlux2

from eval_tiny_vae import LatentNorm, load_images, load_tiny, psnr


class ImageFolder(Dataset):
    def __init__(self, roots, res):
        exts = (".png", ".jpg", ".jpeg", ".webp")
        self.files = [p for root in roots for p in glob.glob(os.path.join(root, "**", "*"), recursive=True)
                      if p.lower().endswith(exts)]
        assert self.files, f"no images under {roots}"
        print(f"{len(self.files)} training images", flush=True)
        self.res = res

    def __len__(self):
        return len(self.files) * 50  # many random crops per high-res image

    def __getitem__(self, i):
        try:
            im = Image.open(self.files[i % len(self.files)]).convert("RGB")
        except Exception:  # skip broken downloads
            return self[random.randrange(len(self.files))]
        # random downscale so the model sees both fine detail and whole-object structure
        short = min(im.size)
        target = random.randint(self.res, max(self.res, min(short, self.res * 3)))
        if target != short:
            im = TF.resize(im, target, antialias=True)
        w, h = im.size
        x0, y0 = random.randint(0, w - self.res), random.randint(0, h - self.res)
        im = im.crop((x0, y0, x0 + self.res, y0 + self.res))
        if random.random() < 0.5:
            im = TF.hflip(im)
        return TF.to_tensor(im)


class PatchDisc(nn.Module):
    """Small PatchGAN discriminator for optional adversarial sharpening."""

    def __init__(self, ch=64):
        super().__init__()
        sn = nn.utils.spectral_norm
        self.net = nn.Sequential(
            sn(nn.Conv2d(3, ch, 4, 2, 1)), nn.LeakyReLU(0.2),
            sn(nn.Conv2d(ch, ch * 2, 4, 2, 1)), nn.LeakyReLU(0.2),
            sn(nn.Conv2d(ch * 2, ch * 4, 4, 2, 1)), nn.LeakyReLU(0.2),
            sn(nn.Conv2d(ch * 4, 1, 3, 1, 1)),
        )

    def forward(self, x):
        return self.net(x * 2 - 1)


@torch.no_grad()
def evaluate(student, vae, ln, lp, imgs, dev):
    student.eval()
    ps, ls = [], []
    for x in imgs:
        x = x[None].to(dev)
        zn = ln.to_norm(vae.encode((x * 2 - 1).to(torch.bfloat16)).latent_dist.mode().float())
        with torch.autocast("cuda", torch.bfloat16):
            d = student.decoder(zn)
        d = d.float().clamp(0, 1)
        ps.append(psnr(d, x)); ls.append(lp(d * 2 - 1, x * 2 - 1).item())
    student.train()
    return sum(ps) / len(ps), sum(ls) / len(ls)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--data", nargs="+", required=True)
    ap.add_argument("--eval_images", default="data/kodak")
    ap.add_argument("--vae", default="weights/flux3-klein-base/vae")
    ap.add_argument("--init", default="weights/taef2.safetensors")
    ap.add_argument("--out", default="runs/ft")
    ap.add_argument("--steps", type=int, default=3000)
    ap.add_argument("--batch", type=int, default=16)
    ap.add_argument("--res", type=int, default=512)
    ap.add_argument("--lr", type=float, default=1e-4)
    ap.add_argument("--lpips_w", type=float, default=0.5)
    ap.add_argument("--gan_w", type=float, default=0.0, help="e.g. 0.05 to add adversarial loss")
    ap.add_argument("--gan_start", type=int, default=1000)
    ap.add_argument("--disc_init", default=None, help="resume the discriminator from a previous run's disc.pt")
    ap.add_argument("--noise_p", type=float, default=0.3, help="prob. of jittering latents (robust to diffusion outputs)")
    ap.add_argument("--train_encoder", action="store_true")
    ap.add_argument("--eval_every", type=int, default=500)
    ap.add_argument("--workers", type=int, default=32)
    ap.add_argument("--ema", type=float, default=0.999, help="EMA decay for the student weights (0 = off)")
    args = ap.parse_args()

    torch.backends.cuda.matmul.allow_tf32 = True
    torch.backends.cudnn.allow_tf32 = True
    dev = "cuda"
    os.makedirs(args.out, exist_ok=True)

    vae = AutoencoderKLFlux2.from_pretrained(args.vae, torch_dtype=torch.bfloat16).to(dev).eval().requires_grad_(False)
    ln = LatentNorm(vae)
    student = load_tiny(args.init, dev).train().requires_grad_(True)
    student.encoder.requires_grad_(args.train_encoder)
    lp = lpips.LPIPS(net="vgg", verbose=False).to(dev).eval().requires_grad_(False)
    lp_eval = lpips.LPIPS(net="alex", verbose=False).to(dev).eval().requires_grad_(False)

    params = [p for p in student.parameters() if p.requires_grad]
    opt = torch.optim.AdamW(params, lr=args.lr, betas=(0.9, 0.99), weight_decay=0.0)
    sched = torch.optim.lr_scheduler.LambdaLR(
        opt, lambda s: min(1.0, (s + 1) / 200) * (0.5 * (1 + math.cos(math.pi * s / args.steps)) * 0.9 + 0.1))
    disc = d_opt = None
    if args.gan_w > 0:
        disc = PatchDisc().to(dev)
        d_opt = torch.optim.AdamW(disc.parameters(), lr=args.lr, betas=(0.5, 0.99))
        if args.disc_init:
            ck = torch.load(args.disc_init, map_location=dev)
            disc.load_state_dict(ck["disc"]); d_opt.load_state_dict(ck["opt"])

    ema = None
    if args.ema > 0:
        ema = copy.deepcopy(student).eval().requires_grad_(False)

    eval_imgs, _ = load_images(args.eval_images)
    p0, l0 = evaluate(student, vae, ln, lp_eval, eval_imgs, dev)
    print(f"step 0  eval PSNR {p0:.2f}  LPIPS {l0:.4f}", flush=True)
    best = l0
    # the starting weights are the best so far until something beats them
    stt.save_file({k: v.detach().float().contiguous().cpu() for k, v in student.state_dict().items()},
                  os.path.join(args.out, "best.safetensors"))

    dl = DataLoader(ImageFolder(args.data, args.res), batch_size=args.batch, shuffle=True,
                    num_workers=args.workers, pin_memory=True, drop_last=True, persistent_workers=True)
    it = iter(dl)
    t0 = time.time()
    for step in range(1, args.steps + 1):
        try:
            x = next(it)
        except StopIteration:
            it = iter(dl); x = next(it)
        x = x.to(dev, non_blocking=True)
        with torch.no_grad():
            zn = ln.to_norm(vae.encode((x * 2 - 1).to(torch.bfloat16)).latent_dist.mode().float())
            z_in = zn
            if args.noise_p > 0:
                mask = (torch.rand(zn.shape[0], 1, 1, 1, device=dev) < args.noise_p).float()
                z_in = zn + mask * torch.rand_like(mask) * 0.1 * torch.randn_like(zn)

        with torch.autocast("cuda", torch.bfloat16):
            out = student.decoder(z_in)
            out = out.float()
            l1 = F.l1_loss(out, x)
            lpl = lp(out.clamp(0, 1) * 2 - 1, x * 2 - 1).mean()
            loss = l1 + args.lpips_w * lpl
            if args.train_encoder:
                zt = student.encoder(x).float()
                loss = loss + F.mse_loss(zt, zn)
            g_loss = torch.zeros((), device=dev)
            if disc is not None and step >= args.gan_start:
                g_loss = -disc(out).mean()
                loss = loss + args.gan_w * g_loss

        opt.zero_grad(set_to_none=True)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(params, 1.0)
        opt.step(); sched.step()

        if disc is not None and step >= args.gan_start - 200:  # warm the discriminator up a bit first
            with torch.autocast("cuda", torch.bfloat16):
                d_loss = F.relu(1 - disc(x)).mean() + F.relu(1 + disc(out.detach())).mean()
            d_opt.zero_grad(set_to_none=True); d_loss.backward(); d_opt.step()

        if step % 50 == 0:
            ips = step * args.batch / (time.time() - t0)
            print(f"step {step:6d}  loss {loss.item():.4f}  l1 {l1.item():.4f}  lpips {lpl.item():.4f}  "
                  f"gan {g_loss.item():.3f}  lr {sched.get_last_lr()[0]:.2e}  {ips:.0f} img/s", flush=True)
        if ema is not None:
            d = min(args.ema, (1 + step) / (10 + step))
            with torch.no_grad():
                for pe, ps in zip(ema.parameters(), student.parameters()):
                    pe.lerp_(ps, 1 - d)

        if step % args.eval_every == 0 or step == args.steps:
            sd_of = lambda m: {k: v.detach().float().contiguous().cpu() for k, v in m.state_dict().items()}
            p, l = evaluate(student, vae, ln, lp_eval, eval_imgs, dev)
            print(f"step {step}  eval PSNR {p:.2f} ({p - p0:+.2f})  LPIPS {l:.4f} ({l - l0:+.4f})", flush=True)
            cands = [(l, sd_of(student))]
            stt.save_file(cands[0][1], os.path.join(args.out, "last.safetensors"))
            if ema is not None:
                pe_, le_ = evaluate(ema, vae, ln, lp_eval, eval_imgs, dev)
                print(f"step {step}  ema eval PSNR {pe_:.2f} ({pe_ - p0:+.2f})  LPIPS {le_:.4f} ({le_ - l0:+.4f})", flush=True)
                cands.append((le_, sd_of(ema)))
                stt.save_file(cands[1][1], os.path.join(args.out, "ema.safetensors"))
            if disc is not None:
                torch.save({"disc": disc.state_dict(), "opt": d_opt.state_dict()}, os.path.join(args.out, "disc.pt"))
            lc, sdc = min(cands, key=lambda c: c[0])
            if lc < best:
                best = lc
                stt.save_file(sdc, os.path.join(args.out, "best.safetensors"))
                print(f"step {step}  new best LPIPS {lc:.4f}", flush=True)


if __name__ == "__main__":
    main()