#!/usr/bin/env python3
"""
PolarQuant Setup: Lightricks LTX-2.3 (22B)
Downloads PQ5 (15 GB) and dequants to BF16 for video generation.

Usage:
    python setup.py
    python generate_ltx.py --prompt "A cat playing piano" --output cat.mp4
"""
import gc, math, os, time, argparse
import torch, torch.nn.functional as F

REPO_ID = "caiovicentino1/LTX-2.3-22B-PolarQuant-Q5"
BS = 128

from scipy.stats import norm as sp_norm
_C = {}
def get_centroids(bits):
    if bits in _C: return _C[bits]
    n = 1 << bits; bd = torch.linspace(-4.0, 4.0, n+1); ct = torch.zeros(n)
    for _ in range(100):
        for i in range(n):
            lo, hi = bd[i].item(), bd[i+1].item()
            plo, phi = sp_norm.cdf(lo), sp_norm.cdf(hi)
            ct[i] = (sp_norm.pdf(lo) - sp_norm.pdf(hi)) / (phi - plo) if phi - plo > 1e-12 else (lo + hi) / 2
        for i in range(1, n): bd[i] = (ct[i-1] + ct[i]) / 2
    _C[bits] = ct; return ct

def build_H(n):
    if n == 1: return torch.tensor([[1.0]])
    h = build_H(n // 2)
    return torch.cat([torch.cat([h,h],1), torch.cat([h,-h],1)], 0) / math.sqrt(2)

def unpack_5bit(packed_flat, total_elements):
    p = packed_flat.long().reshape(-1, 5)
    b0, b1, b2, b3, b4 = p[:,0], p[:,1], p[:,2], p[:,3], p[:,4]
    return torch.stack([
        (b0>>3)&31, ((b0&7)<<2)|((b1>>6)&3), (b1>>1)&31,
        ((b1&1)<<4)|((b2>>4)&15), ((b2&15)<<1)|((b3>>7)&1),
        (b3>>2)&31, ((b3&3)<<3)|((b4>>5)&7), b4&31
    ], dim=-1).reshape(-1)[:total_elements].to(torch.uint8)

def main():
    parser = argparse.ArgumentParser(description="PolarQuant Setup: LTX-2.3")
    parser.add_argument("--output-dir", default="./LTX-PQ5")
    parser.add_argument("--skip-download", action="store_true")
    parser.add_argument("--device", default="cuda")
    args = parser.parse_args()
    OUT = args.output_dir
    device = args.device if torch.cuda.is_available() else "cpu"

    print("=" * 60)
    print("  PolarQuant Setup: LTX-2.3 (22B)")
    print("  15 GB download -> audio-video generation")
    print("=" * 60)

    if not args.skip_download:
        print("\n[1/3] Downloading PQ5 model (15 GB)...")
        from huggingface_hub import snapshot_download
        snapshot_download(REPO_ID, local_dir=OUT)

    print("\n[2/3] Cloning LTX-2 inference code...")
    code_dir = os.path.join(OUT, "ltx_code")
    if not os.path.exists(code_dir):
        os.system(f"git clone https://github.com/Lightricks/LTX-2.git {code_dir}")

    print("\n[3/3] Dequanting PQ5 codes to BF16...")
    ct5 = get_centroids(5).to(device)
    H = build_H(BS).to(device)
    from safetensors.torch import load_file, save_file
    import glob

    # Load all chunks + bf16
    codes_sd = {}
    for f in sorted(glob.glob(f"{OUT}/polarquant/codes/chunk_*_codes.safetensors")):
        print(f"  Loading {os.path.basename(f)}...")
        codes_sd.update(load_file(f, device="cpu"))
    bf16_sd = load_file(f"{OUT}/polarquant/bf16/ltx23_bf16.safetensors", device="cpu")

    coded_keys = sorted(set(k[:-8] for k in codes_sd if k.endswith("__packed")))
    print(f"  {len(coded_keys)} PQ5 + {len(bf16_sd)} BF16 layers")

    full_sd = dict(bf16_sd)
    for idx, ck in enumerate(coded_keys):
        meta = codes_sd[f"{ck}__meta"]
        of, nb, bs, total = int(meta[0]), int(meta[1]), int(meta[2]), int(meta[3])
        codes = unpack_5bit(codes_sd[f"{ck}__packed"], total).reshape(of, nb, bs).to(device).long()
        norms = codes_sd[f"{ck}__norms"].to(device).float().unsqueeze(2)
        vals = ct5[codes] / math.sqrt(BS)
        for i in range(0, of, 256):
            e = min(i+256, of)
            vals[i:e] = (vals[i:e].reshape(-1, BS) @ H).reshape(e-i, nb, BS)
        full_sd[ck.replace("__", ".")] = (vals * norms).reshape(of, nb*bs).to(torch.bfloat16).cpu()
        if (idx+1) % 200 == 0: print(f"    {idx+1}/{len(coded_keys)}...")

    out_path = f"{OUT}/ltx-2.3-22b-dev.safetensors"
    save_file(full_sd, out_path)
    sz = os.path.getsize(out_path) / 1e9
    print(f"  Saved: {out_path} ({sz:.1f} GB)")

    print(f"\n{'='*60}")
    print(f"  Setup complete! Model: {sz:.1f} GB")
    print(f"  Inference: see LTX-2 docs at {code_dir}")
    print(f"{'='*60}")

if __name__ == "__main__":
    main()