#!/usr/bin/env python3
"""
quantize-weights.py — Quantize PersonaPlex 7B weights to INT4 (NF4) for edge devices.

Creates a ~3.5GB quantized checkpoint from the ~14GB bf16 weights.
The quantized model runs 3-4x faster on memory-bandwidth-limited devices
like Jetson AGX Orin while maintaining voice quality.

Usage:
  python quantize-weights.py [--device cuda] [--output personaplex-7b-nf4.safetensors]

Requirements:
  pip install bitsandbytes safetensors torch
"""

import argparse
import os
import sys
import logging

logging.basicConfig(level=logging.INFO, format="%(message)s")
log = logging.getLogger(__name__)


def quantize_model(device: str = "cuda", output_path: str = None):
    """Quantize PersonaPlex 7B to NF4 (4-bit Normal Float)"""
    import torch
    from huggingface_hub import hf_hub_download
    from safetensors.torch import load_file, save_file

    hf_repo = "nvidia/personaplex-7b-v1"

    # 1) Download original weights
    log.info("Downloading PersonaPlex 7B weights...")
    weight_path = hf_hub_download(hf_repo, "model.safetensors")
    log.info(f"  Weights: {weight_path}")
    log.info(f"  Size: {os.path.getsize(weight_path) / 1024**3:.1f} GB")

    # 2) Load state dict
    log.info("Loading state dict...")
    state_dict = load_file(weight_path, device="cpu")
    log.info(f"  Loaded {len(state_dict)} tensors")

    # 3) Quantize each weight tensor to INT4 using block-wise NF4
    try:
        import bitsandbytes as bnb
        from bitsandbytes.functional import quantize_nf4, dequantize_nf4
        HAS_BNB = True
    except ImportError:
        HAS_BNB = False
        log.info("  bitsandbytes not available — using manual INT4 quantization")

    quantized_state = {}
    quant_meta = {}  # Store quantization parameters for dequantization
    total_original = 0
    total_quantized = 0
    skipped = 0

    for name, tensor in state_dict.items():
        original_bytes = tensor.numel() * tensor.element_size()
        total_original += original_bytes

        # Only quantize large weight matrices (≥1024 elements, 2D)
        # Skip biases, norms, embeddings, small tensors
        should_quantize = (
            tensor.ndim >= 2
            and tensor.numel() >= 1024
            and not any(skip in name for skip in [
                "norm", "bias", "embed", "positional", "rope",
                "depformer_emb", "depformer_in",
            ])
        )

        if not should_quantize:
            quantized_state[name] = tensor.to(torch.float16).contiguous()
            total_quantized += tensor.numel() * 2  # fp16
            skipped += 1
            continue

        # Reshape to 2D for quantization
        orig_shape = tensor.shape
        flat = tensor.reshape(-1).float()

        if HAS_BNB:
            # Use bitsandbytes NF4 quantization
            quant_tensor, quant_state = bnb.functional.quantize_4bit(
                flat, quant_type="nf4", compress_statistics=True,
            )
            # Store the quantized bytes + metadata for reconstruction
            quantized_state[name] = quant_tensor.contiguous()
            quant_meta[f"{name}.__quant_state__"] = torch.tensor(
                list(orig_shape) + [0] * (4 - len(orig_shape)),
                dtype=torch.int64,
            )
            # Store absmax for dequantization
            if hasattr(quant_state, 'absmax'):
                quantized_state[f"{name}.__absmax__"] = quant_state.absmax.contiguous()
            if hasattr(quant_state, 'quant_map'):
                quantized_state[f"{name}.__quant_map__"] = quant_state.quant_map.contiguous()
            total_quantized += quant_tensor.numel()
        else:
            # Manual symmetric INT4 quantization (no bitsandbytes)
            # Block size 64 for good accuracy
            block_size = 64
            n_blocks = (flat.numel() + block_size - 1) // block_size
            padded = torch.zeros(n_blocks * block_size)
            padded[:flat.numel()] = flat

            blocks = padded.reshape(n_blocks, block_size)
            scales = blocks.abs().max(dim=1).values / 7.0  # INT4 range: -8 to 7
            scales = scales.clamp(min=1e-8)

            # Quantize to INT4 (stored as INT8 pairs)
            quantized_blocks = torch.round(blocks / scales.unsqueeze(1)).clamp(-8, 7).to(torch.int8)

            # Pack two INT4 values into one INT8
            packed = torch.zeros(n_blocks, block_size // 2, dtype=torch.uint8)
            for i in range(block_size // 2):
                low = (quantized_blocks[:, 2 * i] + 8).to(torch.uint8)
                high = (quantized_blocks[:, 2 * i + 1] + 8).to(torch.uint8)
                packed[:, i] = low | (high << 4)

            quantized_state[name] = packed.reshape(-1).contiguous()
            quantized_state[f"{name}.__scales__"] = scales.to(torch.float16).contiguous()
            quant_meta[f"{name}.__quant_state__"] = torch.tensor(
                list(orig_shape) + [0] * (4 - len(orig_shape)) + [block_size, flat.numel()],
                dtype=torch.int64,
            )
            total_quantized += packed.numel() + scales.numel() * 2

    # Add metadata tensors
    quantized_state.update(quant_meta)

    # 4) Save quantized weights
    if output_path is None:
        output_path = os.path.join(os.path.dirname(weight_path), "model-nf4.safetensors")

    log.info(f"\nSaving quantized weights to: {output_path}")
    save_file(quantized_state, output_path)

    final_size = os.path.getsize(output_path)
    compression = total_original / max(final_size, 1)

    log.info(f"\nQuantization complete!")
    log.info(f"  Original: {total_original / 1024**3:.1f} GB (bf16)")
    log.info(f"  Quantized: {final_size / 1024**3:.1f} GB (NF4)")
    log.info(f"  Compression: {compression:.1f}x")
    log.info(f"  Tensors quantized: {len(state_dict) - skipped}/{len(state_dict)}")
    log.info(f"  Tensors kept fp16: {skipped} (norms, biases, embeddings)")
    log.info(f"\nUse --quantized flag with PersonaPlex server for INT4 inference")

    return output_path


def main():
    parser = argparse.ArgumentParser(description="Quantize PersonaPlex 7B to INT4 NF4")
    parser.add_argument("--device", default="cuda", help="Device for quantization")
    parser.add_argument("--output", "-o", default=None, help="Output path for quantized weights")
    args = parser.parse_args()

    import torch
    with torch.no_grad():
        quantize_model(device=args.device, output_path=args.output)


if __name__ == "__main__":
    main()
