"""
Q6_K Dequantization - GGML Q6_K format

Ported from llama.cpp gguf-py implementation.

Block format:
  - QK_K = 256 floats per block
  - Block size: 210 bytes
    - ql: 128 bytes (quantized low 4 bits)
    - qh: 64 bytes (quantized high 2 bits)
    - scales: 16 bytes (int8, one per 16 floats)
    - d: 2 bytes (float16 global scale)

Encoding:
  - Each float → 6-bit signed int (-32 to +31)
  - 4 bits from ql, 2 bits from qh
  - Grouped into blocks of 16, each with own scale
  - Global scale d applied to all

References:
  - Source: llama.cpp/gguf-py/gguf/quants.py (lines 552-572)
  - License: MIT
  - Repository: https://github.com/ggerganov/llama.cpp

Attribution:
  Q6_K dequantization implementation ported from llama.cpp
  Copyright (c) 2023 Georgi Gerganov
  License: MIT
"""

from typing import Tuple
import numpy as np

# GGML Q6_K constants
QK_K = 256  # Floats per block
Q6K_BLOCK_SIZE = 210  # Bytes per block (128 + 64 + 16 + 2)


def dequantize_q6k_blocks(blocks: np.ndarray) -> np.ndarray:
    """
    Dequantize Q6_K blocks to float32.

    Args:
        blocks: Raw block data, shape (n_blocks, 210) dtype uint8

    Returns:
        Dequantized float32 array, shape (n_blocks, 256)

    Implementation ported from:
        llama.cpp/gguf-py/gguf/quants.py Q6_K.dequantize_blocks()
    """
    n_blocks = blocks.shape[0]

    # Split block into components
    # ql: 128 bytes (QK_K // 2)
    # qh: 64 bytes (QK_K // 4)
    # scales: 16 bytes (QK_K // 16)
    # d: 2 bytes (float16)
    ql, rest = np.hsplit(blocks, [QK_K // 2])
    qh, rest = np.hsplit(rest, [QK_K // 4])
    scales, d = np.hsplit(rest, [QK_K // 16])

    # Convert scales (int8) and delta (float16) to float32
    scales = scales.view(np.int8).astype(np.float32)
    d = d.view(np.float16).astype(np.float32)

    # Phase 7.2: Sanitize before multiplication to prevent NaN/Inf propagation
    scales = np.nan_to_num(scales, nan=0.0, posinf=127.0, neginf=-127.0)
    d = np.nan_to_num(d, nan=0.0, posinf=1e4, neginf=-1e4)

    d = (d * scales).reshape((n_blocks, QK_K // 16, 1))

    # Unpack quantized values from packed format
    # ql contains low 4 bits, packed 2 per byte
    ql = ql.reshape((n_blocks, -1, 1, 64)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2, 1))
    ql = (ql & np.uint8(0x0F)).reshape((n_blocks, -1, 32))

    # qh contains high 2 bits, packed 4 per byte
    qh = qh.reshape((n_blocks, -1, 1, 32)) >> np.array([0, 2, 4, 6], dtype=np.uint8).reshape((1, 1, 4, 1))
    qh = (qh & np.uint8(0x03)).reshape((n_blocks, -1, 32))

    # Combine low and high bits to get 6-bit value, subtract offset (32)
    q = (ql | (qh << np.uint8(4))).astype(np.int8) - np.int8(32)
    q = q.reshape((n_blocks, QK_K // 16, -1)).astype(np.float32)

    # Apply per-group scales
    result = (d * q).reshape((n_blocks, QK_K))

    # Numeric sanitization (Phase 7.2): Clamp NaN/Inf to finite values
    # This ensures deterministic behavior with synthetic regen data
    result = np.nan_to_num(result, nan=0.0, posinf=1e4, neginf=-1e4)

    return result


def dequantize_q6k_tensor(tensor_bytes: bytes, shape: Tuple[int, int]) -> np.ndarray:
    """
    Dequantize a full Q6_K tensor to float32.

    Args:
        tensor_bytes: Raw quantized tensor bytes
        shape: Target output shape (rows, cols)

    Returns:
        Dequantized float32 tensor with specified shape

    Raises:
        ValueError: If tensor size doesn't match expected block count

    Example:
        # From mistral_full.plan.json output.weight:
        # shape: (4096, 32000), size: 107520000 bytes
        bytes_data = open('output.weight.bin', 'rb').read()
        tensor = dequantize_q6k_tensor(bytes_data, (4096, 32000))
    """
    rows, cols = shape
    n_floats = rows * cols

    # Calculate expected blocks
    if n_floats % QK_K != 0:
        raise ValueError(
            f"Total floats ({n_floats}) must be multiple of QK_K ({QK_K})"
        )

    n_blocks = n_floats // QK_K
    expected_bytes = n_blocks * Q6K_BLOCK_SIZE

    if len(tensor_bytes) != expected_bytes:
        raise ValueError(
            f"Tensor size mismatch: got {len(tensor_bytes)} bytes, "
            f"expected {expected_bytes} bytes for {n_blocks} blocks"
        )

    # Reshape to blocks
    blocks = np.frombuffer(tensor_bytes, dtype=np.uint8).reshape((n_blocks, Q6K_BLOCK_SIZE))

    # Dequantize
    floats = dequantize_q6k_blocks(blocks)

    # Reshape to target shape
    return floats.reshape(shape)


def get_q6k_info() -> dict:
    """
    Get Q6_K format information.

    Returns:
        Dictionary with format details
    """
    return {
        "format": "Q6_K",
        "block_size_floats": QK_K,
        "block_size_bytes": Q6K_BLOCK_SIZE,
        "compression_ratio": f"{(QK_K * 4) / Q6K_BLOCK_SIZE:.2f}:1",
        "bits_per_weight": 6,
        "source": "llama.cpp/gguf-py",
        "license": "MIT"
    }


if __name__ == "__main__":
    # Self-test
    print("Q6_K Dequantization Module")
    print("=" * 50)
    info = get_q6k_info()
    for k, v in info.items():
        print(f"  {k}: {v}")
    print()
    print("Import test: OK")
