#!/usr/bin/env python3
"""
dequant-loader.py — Pre-dequantize quantized PersonaPlex weights to bf16 cache.

For NF4 (INT4) or TurboQuant 2-bit weights, dequantizes to a temporary
bf16 safetensors file that moshi.server can load natively.

Usage:
  python dequant-loader.py --input model-nf4.safetensors --output /tmp/model-bf16.safetensors
  python dequant-loader.py --input model-turbo2bit.safetensors --output /tmp/model-bf16.safetensors

The output file can then be passed to moshi.server via --moshi-weight.
"""

import os, sys, math, time
import torch
from safetensors.torch import load_file, save_file

NF2_CENTROIDS = torch.tensor([-1.5104, -0.4528, 0.4528, 1.5104])


def fast_wht(x):
    """Vectorized Walsh-Hadamard Transform."""
    n = x.shape[-1]
    h = 1
    while h < n:
        x_view = x.view(*x.shape[:-1], -1, 2, h)
        a = x_view[..., 0, :].clone()
        b = x_view[..., 1, :].clone()
        x_view[..., 0, :] = a + b
        x_view[..., 1, :] = a - b
        x = x_view.reshape(*x.shape)
        h *= 2
    return x / math.sqrt(n)


def detect_format(state):
    """Detect if weights are NF4 (INT4), TurboQuant 2-bit, or plain."""
    has_scales = any(k.endswith(".__scales__") for k in state)
    has_packed = any(k.endswith(".packed") for k in state)
    if has_packed:
        return "turbo2bit"
    if has_scales:
        return "nf4"
    return "plain"


def dequant_nf4(state):
    """Dequantize INT4 NF4 weights."""
    result = {}
    processed = set()

    for name in list(state.keys()):
        if name.endswith(".__scales__") or name.endswith(".__shape__") or name.endswith(".__numel__"):
            continue
        if name in processed:
            continue

        scales_key = f"{name}.__scales__"
        if scales_key in state:
            packed = state[name]
            scales = state[scales_key].float()
            shape = state[f"{name}.__shape__"].tolist()
            numel = state[f"{name}.__numel__"].item()
            group_size = 64

            lo = (packed & 0x0F).to(torch.int8) - 8
            hi = ((packed >> 4) & 0x0F).to(torch.int8) - 8
            unpacked = torch.zeros(packed.numel() * 2, dtype=torch.float32)
            unpacked[0::2] = lo.float()
            unpacked[1::2] = hi.float()

            n_groups = scales.numel()
            groups = unpacked[:n_groups * group_size].reshape(n_groups, group_size)
            deq = (groups * scales.unsqueeze(1)).reshape(-1)[:numel]

            orig_shape = [s for s in shape if s > 0]
            result[name] = deq.reshape(orig_shape).to(torch.bfloat16)
            processed.add(name)
        else:
            result[name] = state[name].to(torch.bfloat16)
            processed.add(name)

    return result


def dequant_turbo2bit(state):
    """Dequantize TurboQuant 2-bit (NF2 + WHT) weights."""
    result = {}
    processed = set()

    for name in list(state.keys()):
        if any(name.endswith(f".{s}") for s in ["packed", "scales", "shape", "numel", "gs", "np2"]):
            continue
        if name in processed:
            continue

        packed_key = f"{name}.packed"
        if packed_key in state:
            gs = state[f"{name}.gs"].item()
            gs_pow2 = state[f"{name}.np2"].item()
            numel = state[f"{name}.numel"].item()
            shape = [s for s in state[f"{name}.shape"].tolist() if s > 0]
            scales = state[f"{name}.scales"].float()
            packed = state[packed_key]
            n_groups = scales.numel()

            # Unpack 2-bit
            p = packed.reshape(n_groups, gs // 4)
            codes = torch.zeros(n_groups, gs, dtype=torch.long)
            for i in range(4):
                codes[:, i::4] = (p >> (2 * i)) & 0x03

            dequant = NF2_CENTROIDS[codes]

            # Inverse WHT
            if gs_pow2 > gs:
                dequant = torch.cat([dequant, torch.zeros(n_groups, gs_pow2 - gs)], dim=1)
            dequant = fast_wht(dequant)
            dequant = dequant[:, :gs]

            dequant = dequant * scales.unsqueeze(1)
            result[name] = dequant.reshape(-1)[:numel].reshape(shape).to(torch.bfloat16)
            processed.add(name)
        else:
            result[name] = state[name].to(torch.bfloat16)
            processed.add(name)

    return result


def main():
    import argparse
    parser = argparse.ArgumentParser(description="Dequantize PersonaPlex weights to bf16")
    parser.add_argument("--input", "-i", required=True, help="Quantized safetensors file")
    parser.add_argument("--output", "-o", required=True, help="Output bf16 safetensors file")
    parser.add_argument("--device", "-d", default="cpu", help="Device for dequantization")
    args = parser.parse_args()

    if not os.path.exists(args.input):
        print(f"Error: {args.input} not found")
        sys.exit(1)

    # Skip if output already exists and is newer than input
    if os.path.exists(args.output) and os.path.getmtime(args.output) > os.path.getmtime(args.input):
        print(f"Cached: {args.output} is up to date")
        sys.exit(0)

    print(f"Loading {args.input}...")
    t0 = time.time()
    state = load_file(args.input, device=args.device)

    fmt = detect_format(state)
    print(f"Format: {fmt}")

    if fmt == "nf4":
        result = dequant_nf4(state)
    elif fmt == "turbo2bit":
        result = dequant_turbo2bit(state)
    else:
        print("Already plain bf16/fp16 — copying")
        result = {k: v.to(torch.bfloat16) for k, v in state.items()}

    t1 = time.time()
    print(f"Dequantized {len(result)} tensors in {t1-t0:.1f}s")

    print(f"Saving to {args.output}...")
    save_file(result, args.output)
    size_gb = os.path.getsize(args.output) / 1024**3
    print(f"Done: {size_gb:.2f} GB")


if __name__ == "__main__":
    main()