#!/usr/bin/env python3
"""
pq5_quantize.py
Reference Pentary / Quinary post-training quantizer for Hugging Face safetensors models.

Representation
--------------
Per quantized tensor:
    q in {-alpha, -1, 0, +1, +alpha}
    w_hat = scale_g * q
where one FP16 scale is shared by each group of `group_size` weights and alpha
is learned once per tensor.

Packing
-------
Three base-5 digits form one symbol in [0, 124]:
    symbol = d0 + 5*d1 + 25*d2
Eight 7-bit symbols are packed into seven bytes.
This yields 7/3 = 2.333... physical bits/weight for codes, before scales.

With FP16 scales and group_size=128:
    7/3 + 16/128 ~= 2.4583 bits/weight
plus negligible tensor metadata / alpha overhead.

This is a storage + quality-reference implementation, NOT yet a fast inference
kernel. A custom Triton/CUDA/llama.cpp kernel should consume packed weights
directly for speed.

Example:
    pip install torch safetensors huggingface_hub
    python pq5_quantize.py \
      --model Qwen/Qwen3-0.6B \
      --output ./Qwen3-0.6B-PQ5 \
      --device cuda \
      --group-size 128 \
      --alpha-grid 1.25,1.5,1.75,2.0,2.25,2.5,2.75,3.0 \
      --verify

The script can also take a local model directory with --model /path/to/model.
"""

from __future__ import annotations

import argparse
import gc
import json
import math
import os
import re
import shutil
from pathlib import Path
from typing import Dict, Iterable, List, Tuple

import torch
from safetensors import safe_open
from safetensors.torch import save_file


# -----------------------------
# Utilities
# -----------------------------

def parse_size(s: str) -> int:
    s = s.strip().upper()
    m = re.fullmatch(r"([0-9]*\.?[0-9]+)\s*(B|KB|MB|GB|KIB|MIB|GIB)?", s)
    if not m:
        raise ValueError(f"Invalid size: {s}")
    val = float(m.group(1))
    unit = m.group(2) or "B"
    mult = {
        "B": 1,
        "KB": 1000,
        "MB": 1000**2,
        "GB": 1000**3,
        "KIB": 1024,
        "MIB": 1024**2,
        "GIB": 1024**3,
    }[unit]
    return int(val * mult)


def tensor_bytes(t: torch.Tensor) -> int:
    return t.numel() * t.element_size()


def resolve_model(model: str, cache_dir: str | None = None) -> Path:
    p = Path(model)
    if p.exists():
        return p.resolve()

    from huggingface_hub import snapshot_download
    local = snapshot_download(
        repo_id=model,
        cache_dir=cache_dir,
        allow_patterns=[
            "*.json",
            "*.safetensors",
            "*.model",
            "*.txt",
            "*.tiktoken",
            "tokenizer*",
            "merges.txt",
            "vocab.json",
            "*.py",
        ],
    )
    return Path(local)


def copy_non_weight_files(src: Path, dst: Path) -> None:
    dst.mkdir(parents=True, exist_ok=True)
    for p in src.iterdir():
        if not p.is_file():
            continue
        if p.suffix == ".safetensors":
            continue
        if p.name.endswith(".safetensors.index.json"):
            continue
        # Avoid copying giant unrelated binaries.
        if p.stat().st_size > 64 * 1024 * 1024:
            continue
        shutil.copy2(p, dst / p.name)


def list_safetensors(model_dir: Path) -> List[Path]:
    files = sorted(model_dir.glob("*.safetensors"))
    if not files:
        raise FileNotFoundError(f"No .safetensors files found in {model_dir}")
    return files


def iter_safetensor_items(files: Iterable[Path]):
    for sf in files:
        with safe_open(str(sf), framework="pt", device="cpu") as f:
            for key in f.keys():
                yield sf.name, key, f.get_tensor(key)


# -----------------------------
# PQ5 codebook optimization
# -----------------------------

@torch.no_grad()
def _assign_levels(x: torch.Tensor, scale: torch.Tensor, alpha: float):
    """
    x: [n_groups, group_size], fp32
    scale: [n_groups, 1], fp32
    Returns:
        idx: uint8/int64 code index in [0..4]
        q:   fp32 quantized normalized levels
    """
    eps = 1e-12
    z = x / scale.clamp_min(eps)

    # Boundaries are exact midpoints between ordered levels
    # [-a, -1, 0, 1, a].
    b0 = -(alpha + 1.0) * 0.5
    b1 = -0.5
    b2 = 0.5
    b3 = (alpha + 1.0) * 0.5

    idx = torch.empty_like(z, dtype=torch.uint8)
    idx[z < b0] = 0
    idx[(z >= b0) & (z < b1)] = 1
    idx[(z >= b1) & (z < b2)] = 2
    idx[(z >= b2) & (z < b3)] = 3
    idx[z >= b3] = 4

    q = torch.zeros_like(z)
    q[idx == 0] = -alpha
    q[idx == 1] = -1.0
    q[idx == 2] = 0.0
    q[idx == 3] = 1.0
    q[idx == 4] = alpha
    return idx, q


@torch.no_grad()
def _fit_scale_for_alpha(
    x: torch.Tensor,
    alpha: float,
    iters: int = 4,
) -> Tuple[torch.Tensor, torch.Tensor, float]:
    """
    Alternating nearest-code assignment and least-squares group-scale fitting.
    x is [groups, group_size] float32.
    """
    # Robust initial scale. 2*mean(|x|)/(1+alpha) roughly centers both
    # non-zero magnitudes around the observed absolute-weight mass.
    mean_abs = x.abs().mean(dim=1, keepdim=True)
    scale = (2.0 * mean_abs / (1.0 + alpha)).clamp_min(1e-8)

    idx = None
    q = None
    for _ in range(iters):
        idx, q = _assign_levels(x, scale, alpha)
        num = (x * q).sum(dim=1, keepdim=True)
        den = (q * q).sum(dim=1, keepdim=True).clamp_min(1e-8)
        scale = (num / den).abs().clamp_min(1e-8)

    idx, q = _assign_levels(x, scale, alpha)
    recon = q * scale
    mse = torch.mean((x - recon) ** 2).item()
    return scale, idx, mse


@torch.no_grad()
def choose_alpha(
    flat: torch.Tensor,
    group_size: int,
    alpha_grid: List[float],
    device: str,
    sample_weights: int = 1_048_576,
    iters: int = 4,
) -> Tuple[float, Dict[float, float]]:
    """
    Pick one alpha per tensor using a deterministic sample of groups.
    """
    n = flat.numel()
    n_groups = math.ceil(n / group_size)

    # Deterministic evenly-spaced group sample.
    max_groups = max(1, sample_weights // group_size)
    take_groups = min(n_groups, max_groups)
    if take_groups == n_groups:
        group_ids = torch.arange(n_groups)
    else:
        group_ids = torch.linspace(0, n_groups - 1, take_groups).round().long().unique()

    # Build sampled groups without padding the entire tensor on device.
    cpu = flat.detach().float().cpu()
    chunks = []
    for gid in group_ids.tolist():
        start = gid * group_size
        g = cpu[start : start + group_size]
        if g.numel() < group_size:
            g = torch.nn.functional.pad(g, (0, group_size - g.numel()))
        chunks.append(g)
    sample = torch.stack(chunks, dim=0).to(device=device, dtype=torch.float32)

    scores = {}
    for a in alpha_grid:
        _, _, mse = _fit_scale_for_alpha(sample, float(a), iters=iters)
        scores[float(a)] = float(mse)

    best = min(scores, key=scores.get)
    del sample, chunks
    if device.startswith("cuda"):
        torch.cuda.empty_cache()
    return float(best), scores


@torch.no_grad()
def quantize_tensor_pq5(
    w: torch.Tensor,
    group_size: int,
    alpha_grid: List[float],
    device: str,
    group_chunk: int = 8192,
    sample_weights: int = 1_048_576,
    iters: int = 4,
):
    """
    Quantize a floating tensor into:
      packed uint8 byte stream
      FP16 group scales
      FP16 scalar alpha

    Returns dict and stats.
    """
    original_shape = list(w.shape)
    original_numel = w.numel()
    flat_cpu = w.detach().float().reshape(-1).cpu()

    alpha, alpha_scores = choose_alpha(
        flat_cpu,
        group_size=group_size,
        alpha_grid=alpha_grid,
        device=device,
        sample_weights=sample_weights,
        iters=iters,
    )

    padded_numel = math.ceil(original_numel / group_size) * group_size
    pad = padded_numel - original_numel
    if pad:
        flat_cpu = torch.nn.functional.pad(flat_cpu, (0, pad))
    groups_cpu = flat_cpu.view(-1, group_size)

    all_codes = []
    all_scales = []

    sqerr = 0.0
    signal = 0.0

    for start in range(0, groups_cpu.shape[0], group_chunk):
        x = groups_cpu[start : start + group_chunk].to(device=device, dtype=torch.float32)
        scales, idx, _ = _fit_scale_for_alpha(x, alpha, iters=iters)

        # Stats
        level_values = torch.tensor(
            [-alpha, -1.0, 0.0, 1.0, alpha],
            device=device,
            dtype=torch.float32,
        )
        q = level_values[idx.long()]
        recon = q * scales
        sqerr += torch.sum((x - recon) ** 2).item()
        signal += torch.sum(x ** 2).item()

        all_codes.append(idx.reshape(-1).cpu())
        all_scales.append(scales.reshape(-1).to(torch.float16).cpu())

        del x, scales, idx, q, recon
        if device.startswith("cuda"):
            torch.cuda.empty_cache()

    codes = torch.cat(all_codes, dim=0)
    scales = torch.cat(all_scales, dim=0)

    packed, code_pad = pack_base5_codes(codes)

    rmse_rel = math.sqrt(sqerr / max(signal, 1e-30))
    code_hist = torch.bincount(codes.long(), minlength=5).tolist()

    result = {
        "packed": packed.contiguous(),
        "scales": scales.contiguous(),
        "alpha": torch.tensor([alpha], dtype=torch.float16),
    }
    stats = {
        "shape": original_shape,
        "numel": original_numel,
        "padded_numel": padded_numel,
        "group_size": group_size,
        "alpha": alpha,
        "alpha_mse": alpha_scores,
        "relative_rmse": rmse_rel,
        "code_histogram": code_hist,
        "code_pad": code_pad,
        "packed_bytes": packed.numel(),
        "scale_bytes": tensor_bytes(scales),
        "effective_bpw_tensor": 8.0 * (packed.numel() + tensor_bytes(scales) + 2) / original_numel,
    }

    del flat_cpu, groups_cpu, all_codes, all_scales, codes
    return result, stats


# -----------------------------
# Exact 5-ary packing
# -----------------------------

@torch.no_grad()
def pack_base5_codes(codes: torch.Tensor) -> Tuple[torch.Tensor, int]:
    """
    codes in {0,1,2,3,4}.
    3 digits -> one 7-bit symbol; 8 symbols -> 7 bytes.
    Returns (packed_bytes_uint8, number_of_padding_codes).
    """
    c = codes.to(torch.uint8).reshape(-1).cpu()
    if c.numel() and int(c.max()) > 4:
        raise ValueError("PQ5 codes must be in [0,4]")

    code_pad = (-c.numel()) % 3
    if code_pad:
        # Pad with zero-level code 2, not with code 0.
        c = torch.cat([c, torch.full((code_pad,), 2, dtype=torch.uint8)])

    tri = c.view(-1, 3).to(torch.int16)
    symbols = tri[:, 0] + 5 * tri[:, 1] + 25 * tri[:, 2]  # [0,124]
    symbols = symbols.to(torch.uint8)

    sym_pad = (-symbols.numel()) % 8
    if sym_pad:
        # symbol representing three zero-level digits: 2 + 10 + 50 = 62
        symbols = torch.cat([symbols, torch.full((sym_pad,), 62, dtype=torch.uint8)])

    s = symbols.view(-1, 8).to(torch.int16)
    out = torch.empty((s.shape[0], 7), dtype=torch.uint8)

    out[:, 0] = (s[:, 0] | ((s[:, 1] & 0x01) << 7)).to(torch.uint8)
    out[:, 1] = ((s[:, 1] >> 1) | ((s[:, 2] & 0x03) << 6)).to(torch.uint8)
    out[:, 2] = ((s[:, 2] >> 2) | ((s[:, 3] & 0x07) << 5)).to(torch.uint8)
    out[:, 3] = ((s[:, 3] >> 3) | ((s[:, 4] & 0x0F) << 4)).to(torch.uint8)
    out[:, 4] = ((s[:, 4] >> 4) | ((s[:, 5] & 0x1F) << 3)).to(torch.uint8)
    out[:, 5] = ((s[:, 5] >> 5) | ((s[:, 6] & 0x3F) << 2)).to(torch.uint8)
    out[:, 6] = ((s[:, 6] >> 6) | (s[:, 7] << 1)).to(torch.uint8)

    return out.reshape(-1), int(code_pad)


@torch.no_grad()
def unpack_base5_codes(
    packed: torch.Tensor,
    n_codes: int,
) -> torch.Tensor:
    """
    Reference inverse of pack_base5_codes. Returns uint8 codes [0..4].
    """
    b = packed.to(torch.uint8).reshape(-1).cpu()
    if b.numel() % 7 != 0:
        raise ValueError("Packed PQ5 stream length must be a multiple of 7 bytes.")

    x = b.view(-1, 7).to(torch.int16)
    s = torch.empty((x.shape[0], 8), dtype=torch.int16)

    s[:, 0] = x[:, 0] & 0x7F
    s[:, 1] = ((x[:, 0] >> 7) | ((x[:, 1] & 0x3F) << 1))
    s[:, 2] = ((x[:, 1] >> 6) | ((x[:, 2] & 0x1F) << 2))
    s[:, 3] = ((x[:, 2] >> 5) | ((x[:, 3] & 0x0F) << 3))
    s[:, 4] = ((x[:, 3] >> 4) | ((x[:, 4] & 0x07) << 4))
    s[:, 5] = ((x[:, 4] >> 3) | ((x[:, 5] & 0x03) << 5))
    s[:, 6] = ((x[:, 5] >> 2) | ((x[:, 6] & 0x01) << 6))
    s[:, 7] = x[:, 6] >> 1

    symbols = s.reshape(-1)
    d0 = symbols % 5
    d1 = (symbols // 5) % 5
    d2 = (symbols // 25) % 5
    codes = torch.stack([d0, d1, d2], dim=1).reshape(-1).to(torch.uint8)
    return codes[:n_codes]


@torch.no_grad()
def dequantize_pq5(
    packed: torch.Tensor,
    scales: torch.Tensor,
    alpha: torch.Tensor | float,
    shape: List[int] | Tuple[int, ...],
    group_size: int = 128,
    dtype: torch.dtype = torch.float16,
) -> torch.Tensor:
    """
    Reference dequantizer. Useful for verification / fake-quant evaluation.
    """
    n = math.prod(shape)
    padded_n = math.ceil(n / group_size) * group_size
    codes = unpack_base5_codes(packed, padded_n).long()

    a = float(alpha.item()) if isinstance(alpha, torch.Tensor) else float(alpha)
    lut = torch.tensor([-a, -1.0, 0.0, 1.0, a], dtype=torch.float32)
    q = lut[codes].view(-1, group_size)
    s = scales.float().view(-1, 1)
    w = (q * s).reshape(-1)[:n].reshape(shape)
    return w.to(dtype)


# -----------------------------
# Sharded writer
# -----------------------------

class ShardWriter:
    def __init__(self, out_dir: Path, max_shard_bytes: int):
        self.out_dir = out_dir
        self.max_shard_bytes = max_shard_bytes
        self.pending: Dict[str, torch.Tensor] = {}
        self.pending_bytes = 0
        self.shards: List[str] = []
        self.weight_map: Dict[str, str] = {}
        self.shard_no = 0

    def add(self, key: str, tensor: torch.Tensor):
        nbytes = tensor_bytes(tensor)
        if self.pending and self.pending_bytes + nbytes > self.max_shard_bytes:
            self.flush()
        self.pending[key] = tensor.contiguous().cpu()
        self.pending_bytes += nbytes

    def flush(self):
        if not self.pending:
            return
        self.shard_no += 1
        fname = f"model-pq5-{self.shard_no:05d}.safetensors"
        save_file(self.pending, str(self.out_dir / fname))
        for k in self.pending:
            self.weight_map[k] = fname
        self.shards.append(fname)
        self.pending = {}
        self.pending_bytes = 0
        gc.collect()

    def finalize(self):
        self.flush()


# -----------------------------
# CLI
# -----------------------------

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", required=True, help="HF repo id or local model directory")
    ap.add_argument("--output", required=True)
    ap.add_argument("--cache-dir", default=None)
    ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
    ap.add_argument("--group-size", type=int, default=128)
    ap.add_argument(
        "--alpha-grid",
        default="1.25,1.5,1.75,2.0,2.25,2.5,2.75,3.0",
        help="Comma-separated outer-magnitude candidates. Must all be > 1.",
    )
    ap.add_argument("--alpha-iters", type=int, default=4)
    ap.add_argument("--alpha-sample-weights", type=int, default=1_048_576)
    ap.add_argument("--group-chunk", type=int, default=8192)
    ap.add_argument("--max-shard-size", default="1GB")
    ap.add_argument(
        "--min-elements",
        type=int,
        default=4096,
        help="Only quantize floating 2D+ tensors at least this large.",
    )
    ap.add_argument(
        "--exclude-regex",
        default=r"(?:^|\.)(?:norm|layernorm|ln_)(?:\.|$)",
        help="Regex of tensor names to leave unchanged.",
    )
    ap.add_argument(
        "--quantize-embeddings",
        action=argparse.BooleanOptionalAction,
        default=True,
    )
    ap.add_argument(
        "--quantize-lm-head",
        action=argparse.BooleanOptionalAction,
        default=True,
    )
    ap.add_argument("--verify", action="store_true")
    args = ap.parse_args()

    alpha_grid = [float(x) for x in args.alpha_grid.split(",") if x.strip()]
    if not alpha_grid or any(a <= 1.0 for a in alpha_grid):
        raise ValueError("Every alpha candidate must be > 1.0")

    if args.group_size <= 0:
        raise ValueError("--group-size must be positive")

    src = resolve_model(args.model, args.cache_dir)
    out = Path(args.output)
    out.mkdir(parents=True, exist_ok=True)
    copy_non_weight_files(src, out)

    files = list_safetensors(src)
    writer = ShardWriter(out, parse_size(args.max_shard_size))

    manifest = {
        "format": "PQ5",
        "version": 1,
        "description": "5-state groupwise PTQ with tensorwise learned alpha and exact base-5 bit packing",
        "codebook": "[-alpha, -1, 0, +1, +alpha]",
        "group_size": args.group_size,
        "packing": {
            "digits_per_symbol": 3,
            "symbol_bits": 7,
            "symbols_per_block": 8,
            "bytes_per_block": 7,
            "code_bits_per_weight": 7.0 / 3.0,
        },
        "tensors": {},
    }

    total_orig_bytes = 0
    total_stored_bytes = 0
    total_quantized_numel = 0
    total_passthrough_bytes = 0

    exclude = re.compile(args.exclude_regex) if args.exclude_regex else None

    print("=" * 88)
    print("PQ5 :: Pentary / Quinary post-training quantizer")
    print(f"source      : {src}")
    print(f"output      : {out}")
    print(f"device      : {args.device}")
    print(f"group size  : {args.group_size}")
    print(f"alpha grid  : {alpha_grid}")
    print("=" * 88)

    tensor_idx = 0
    for source_shard, name, t in iter_safetensor_items(files):
        tensor_idx += 1
        orig_bytes = tensor_bytes(t)
        total_orig_bytes += orig_bytes

        is_float = t.dtype in {
            torch.float16,
            torch.bfloat16,
            torch.float32,
            torch.float64,
        }
        quantizable = is_float and t.ndim >= 2 and t.numel() >= args.min_elements
        if exclude and exclude.search(name):
            quantizable = False
        lname = name.lower()
        if not args.quantize_embeddings and ("embed" in lname or "embedding" in lname):
            quantizable = False
        if not args.quantize_lm_head and ("lm_head" in lname or lname.endswith("output.weight")):
            quantizable = False

        if not quantizable:
            writer.add(name, t)
            manifest["tensors"][name] = {
                "kind": "passthrough",
                "shape": list(t.shape),
                "dtype": str(t.dtype).replace("torch.", ""),
                "source_shard": source_shard,
            }
            total_stored_bytes += orig_bytes
            total_passthrough_bytes += orig_bytes
            print(f"[{tensor_idx:04d}] PASS {name:70s} {list(t.shape)} {t.dtype}")
            del t
            continue

        print(f"[{tensor_idx:04d}] PQ5  {name:70s} {list(t.shape)} {t.dtype}")
        q, stats = quantize_tensor_pq5(
            t,
            group_size=args.group_size,
            alpha_grid=alpha_grid,
            device=args.device,
            group_chunk=args.group_chunk,
            sample_weights=args.alpha_sample_weights,
            iters=args.alpha_iters,
        )

        k_packed = name + ".pq5_packed"
        k_scales = name + ".pq5_scales"
        k_alpha = name + ".pq5_alpha"

        writer.add(k_packed, q["packed"])
        writer.add(k_scales, q["scales"])
        writer.add(k_alpha, q["alpha"])

        stored = (
            tensor_bytes(q["packed"])
            + tensor_bytes(q["scales"])
            + tensor_bytes(q["alpha"])
        )
        total_stored_bytes += stored
        total_quantized_numel += t.numel()

        manifest["tensors"][name] = {
            "kind": "pq5",
            "shape": list(t.shape),
            "dtype": str(t.dtype).replace("torch.", ""),
            "source_shard": source_shard,
            "packed_key": k_packed,
            "scales_key": k_scales,
            "alpha_key": k_alpha,
            "group_size": args.group_size,
            "padded_numel": stats["padded_numel"],
            "code_pad": stats["code_pad"],
            "alpha": stats["alpha"],
            "relative_rmse": stats["relative_rmse"],
            "code_histogram": stats["code_histogram"],
            "effective_bpw_tensor": stats["effective_bpw_tensor"],
        }

        print(
            f"       alpha={stats['alpha']:.3f}  "
            f"relRMSE={stats['relative_rmse']:.5f}  "
            f"bpw={stats['effective_bpw_tensor']:.4f}  "
            f"hist={stats['code_histogram']}"
        )

        if args.verify:
            deq = dequantize_pq5(
                q["packed"],
                q["scales"],
                q["alpha"],
                shape=list(t.shape),
                group_size=args.group_size,
                dtype=torch.float32,
            )
            # Packing verification: quantization stats already measure distortion;
            # this catches byte-packing/depacking corruption.
            if not torch.isfinite(deq).all():
                raise RuntimeError(f"Non-finite dequantized values in {name}")
            print(
                f"       verify=OK  dequant_range=[{deq.min().item():.5g}, {deq.max().item():.5g}]"
            )
            del deq

        del t, q
        gc.collect()
        if args.device.startswith("cuda"):
            torch.cuda.empty_cache()

    writer.finalize()

    manifest["summary"] = {
        "original_tensor_bytes": total_orig_bytes,
        "stored_tensor_bytes": total_stored_bytes,
        "compression_ratio": total_orig_bytes / max(total_stored_bytes, 1),
        "quantized_parameters": total_quantized_numel,
        "passthrough_bytes": total_passthrough_bytes,
        "overall_effective_bpw_vs_original_param_count": (
            8.0 * total_stored_bytes / max(total_quantized_numel, 1)
        ),
    }
    manifest["weight_map"] = writer.weight_map
    manifest["shards"] = writer.shards

    with open(out / "pq5_manifest.json", "w", encoding="utf-8") as f:
        json.dump(manifest, f, indent=2)

    with open(out / "model.pq5.index.json", "w", encoding="utf-8") as f:
        json.dump(
            {
                "metadata": {
                    "format": "PQ5",
                    "group_size": args.group_size,
                    "total_size": total_stored_bytes,
                },
                "weight_map": writer.weight_map,
            },
            f,
            indent=2,
        )

    print("=" * 88)
    print("DONE")
    print(f"Original tensor bytes : {total_orig_bytes / 1e9:.3f} GB")
    print(f"Stored tensor bytes   : {total_stored_bytes / 1e9:.3f} GB")
    print(f"Compression ratio     : {total_orig_bytes / max(total_stored_bytes, 1):.3f}x")
    if total_quantized_numel:
        quant_only_ideal = 7.0 / 3.0 + 16.0 / args.group_size
        print(f"PQ5 quantized ideal   : {quant_only_ideal:.4f} bits/weight + tiny alpha/packing pad")
    print(f"Manifest              : {out / 'pq5_manifest.json'}")
    print("=" * 88)


if __name__ == "__main__":
    main()