"""Compare TAEF2 (or a fine-tuned copy) against the full FLUX 3 Klein VAE.

Reports reconstruction PSNR / LPIPS on a folder of images, latent agreement,
decode speed, parameter counts, and saves a side-by-side grid.

python eval_tiny_vae.py --images data/kodak --tiny weights/taef2.safetensors
"""
import argparse, glob, os, time

import lpips
import numpy as np
import safetensors.torch as stt
import torch
import torch.nn.functional as F
from PIL import Image
from diffusers import AutoencoderKLFlux2

from taesd import TAESD


def convert_diffusers_sd_to_taesd(sd):
    # taef2.safetensors uses "encoder.layers.N" names; taesd.py uses "encoder.N" (decoder offset by 1)
    out = {}
    for k, v in sd.items():
        if not k.startswith(("encoder.layers.", "decoder.layers.")):
            out[k] = v
            continue
        encdec, _layers, index, *suffix = k.split(".")
        offset = 1 if encdec == "decoder" else 0
        out[".".join([encdec, str(int(index) + offset), *suffix])] = v
    return out


def load_tiny(path, device, dtype=torch.float32):
    m = TAESD(encoder_path=None, decoder_path=None, latent_channels=32, arch_variant="flux_2")
    sd = stt.load_file(path)
    if any(k.startswith(("encoder.layers.", "decoder.layers.")) for k in sd):
        sd = convert_diffusers_sd_to_taesd(sd)
    m.load_state_dict(sd)
    return m.to(device, dtype).eval().requires_grad_(False)


def patchify(z):
    b, c, h, w = z.shape
    return z.view(b, c, h // 2, 2, w // 2, 2).permute(0, 1, 3, 5, 2, 4).reshape(b, c * 4, h // 2, w // 2)


def unpatchify(z):
    b, c, h, w = z.shape
    return z.reshape(b, c // 4, 2, 2, h, w).permute(0, 1, 4, 2, 5, 3).reshape(b, c // 4, h * 2, w * 2)


class LatentNorm:
    """Raw VAE latents <-> the BN-normalized space the FLUX transformer (and TAEF2) work in."""

    def __init__(self, vae):
        self.mean = vae.bn.running_mean.view(1, -1, 1, 1).float()
        self.std = torch.sqrt(vae.bn.running_var.view(1, -1, 1, 1).float() + vae.config.batch_norm_eps)

    def to_norm(self, z):
        return unpatchify((patchify(z) - self.mean.to(z)) / self.std.to(z))

    def to_raw(self, z):
        return unpatchify(patchify(z) * self.std.to(z) + self.mean.to(z))


def load_images(folder, mult=16):
    xs, names = [], []
    for p in sorted(glob.glob(os.path.join(folder, "*"))):
        if not p.lower().endswith((".png", ".jpg", ".jpeg", ".webp")):
            continue
        im = Image.open(p).convert("RGB")
        w, h = im.size
        im = im.crop((0, 0, w - w % mult, h - h % mult))
        xs.append(torch.from_numpy(np.asarray(im)).permute(2, 0, 1).float() / 255.0)
        names.append(os.path.basename(p))
    return xs, names


def psnr(a, b):
    return (10 * torch.log10(1.0 / F.mse_loss(a, b))).item()


def timed(fn, x, n=10):
    fn(x); torch.cuda.synchronize()
    t = time.perf_counter()
    for _ in range(n):
        fn(x)
    torch.cuda.synchronize()
    return (time.perf_counter() - t) / n * 1000


@torch.no_grad()
def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--images", default="data/kodak")
    ap.add_argument("--vae", default="weights/flux3-klein-base/vae")
    ap.add_argument("--tiny", nargs="+", default=["weights/taef2.safetensors"])
    ap.add_argument("--grid", default="out/compare.jpg")
    ap.add_argument("--latent_space", choices=["norm", "raw"], default="norm",
                    help="space the tiny AE reads/writes: BN-normalized (TAEF2) or raw VAE latents")
    args = ap.parse_args()

    dev = "cuda"
    vae = AutoencoderKLFlux2.from_pretrained(args.vae, torch_dtype=torch.float32).to(dev).eval()
    ln = LatentNorm(vae)
    to_tiny = ln.to_norm if args.latent_space == "norm" else (lambda z: z)
    from_tiny = ln.to_raw if args.latent_space == "norm" else (lambda z: z)
    tinies = {os.path.basename(p): load_tiny(p, dev) for p in args.tiny}
    lp = lpips.LPIPS(net="alex", verbose=False).to(dev)

    nparam = lambda m: sum(p.numel() for p in m.parameters()) / 1e6
    print(f"full VAE   : {nparam(vae):7.2f}M params (enc {nparam(vae.encoder):.2f}M / dec {nparam(vae.decoder):.2f}M)")
    for n, t in tinies.items():
        print(f"{n:11s}: {nparam(t):7.2f}M params (enc {nparam(t.encoder):.2f}M / dec {nparam(t.decoder):.2f}M)")

    xs, names = load_images(args.images)
    res = {k: [] for k in ["full"] + [f"{n}|dec" for n in tinies] + [f"{n}|enc" for n in tinies] + [f"{n}|both" for n in tinies]}
    lat_stats = {n: [] for n in tinies}
    grid_rows = []
    for i, x in enumerate(xs):
        x = x[None].to(dev)
        z = vae.encode(x * 2 - 1).latent_dist.mode()  # raw 32-ch latent at /8, what the pipeline unpatchifies to
        full = vae.decode(z).sample.add(1).div(2).clamp(0, 1)
        res["full"].append((psnr(full, x), lp(full * 2 - 1, x * 2 - 1).item()))
        row = [x, full]
        for n, t in tinies.items():
            zn = to_tiny(z)
            d = t.decoder(zn).clamp(0, 1)
            res[f"{n}|dec"].append((psnr(d, x), lp(d * 2 - 1, x * 2 - 1).item()))
            zt = t.encoder(x)
            e = vae.decode(from_tiny(zt)).sample.add(1).div(2).clamp(0, 1)
            res[f"{n}|enc"].append((psnr(e, x), lp(e * 2 - 1, x * 2 - 1).item()))
            b = t.decoder(zt).clamp(0, 1)
            res[f"{n}|both"].append((psnr(b, x), lp(b * 2 - 1, x * 2 - 1).item()))
            cos = F.cosine_similarity(zt.flatten(), zn.flatten(), dim=0).item()
            lat_stats[n].append((F.mse_loss(zt, zn).item(), cos, zn.std().item(), zt.std().item()))
            row.append(d)
        if i < 4:
            grid_rows.append(torch.cat([r[..., :512, :512] for r in row], -1))

    print("\nreconstruction vs original image (mean over", len(xs), "images)")
    print(f"{'path':34s} {'PSNR dB':>8s} {'LPIPS':>7s}")
    for k, v in res.items():
        v = np.array(v)
        label = {"full": "full VAE enc -> full VAE dec"}.get(k, k.replace("|dec", ": full enc -> TINY dec").replace("|enc", ": TINY enc -> full dec").replace("|both", ": TINY enc -> TINY dec"))
        print(f"{label:34s} {v[:,0].mean():8.2f} {v[:,1].mean():7.4f}")
    for n, s in lat_stats.items():
        s = np.array(s)
        print(f"\n{n} encoder vs full encoder latents: MSE {s[:,0].mean():.4f}  cosine {s[:,1].mean():.4f}  std full {s[:,2].mean():.3f} / tiny {s[:,3].mean():.3f}")

    print("\ndecode latency (bf16, batch 1)")
    vae_bf = vae.to(torch.bfloat16)
    for res_px in (1024, 2048):
        z = torch.randn(1, 32, res_px // 8, res_px // 8, device=dev, dtype=torch.bfloat16)
        torch.cuda.reset_peak_memory_stats()
        tf = timed(lambda z: vae_bf.decode(z).sample, z, 5)
        mf = torch.cuda.max_memory_allocated() / 2**30
        line = f"  {res_px}px  full VAE {tf:8.1f} ms ({mf:.1f} GB peak)"
        for n, t in tinies.items():
            tb = t.to(torch.bfloat16)
            torch.cuda.reset_peak_memory_stats()
            tt = timed(tb.decoder, z, 10)
            mt = torch.cuda.max_memory_allocated() / 2**30
            line += f" | {n} {tt:6.1f} ms ({mt:.1f} GB) -> {tf/tt:.1f}x faster"
            t.float()
        print(line)

    os.makedirs(os.path.dirname(args.grid) or ".", exist_ok=True)
    g = torch.cat(grid_rows, -2)[0].permute(1, 2, 0).float().cpu().numpy()
    Image.fromarray((g * 255).round().astype(np.uint8)).save(args.grid, quality=92)
    print("\ngrid columns: original | full VAE |", " | ".join(tinies), "->", args.grid)


if __name__ == "__main__":
    main()