"""GGUF Q2_K / Q3_K packed slice-dequantization (torch, output-row slice).

Exact llama.cpp layouts (ggml-common.h / ggml-quants.c):

  block_q2_K (84 B): scales[16] (each byte: low nibble = 4-bit scale,
                                   high nibble = 4-bit min)
                     + qs[64] (256 x 2-bit) + d (f16) + dmin (f16)
    dequant (dequantize_row_q2_K): 16 sub-blocks of 16.
      sub-block is: n = is//8, j = is%8; qs byte = 32*n + 16*(is%2) + l;
      shift = 2*(j%4); value = (qs[byte] >> shift) & 3;
      y = d*sc*q - dmin*m, sc = scales[is]&0xF, m = scales[is]>>4

  block_q3_K (110 B): hmask[32] + qs[64] + scales[12] + d(f16)
    dequant (dequantize_row_q3_K): 16 sub-blocks of 16.
      aux unpacking of scales[12] -> 16 int8 (signed), scale = int8 - 32
      sub-block is: n = is//8, j_local = (is%8)//2, half = is%2;
      qs byte = 32*n + 16*half + l; shift = 2*j_local;
      m_bit = 1 << (j_local + 4*n); value = q - (hm[16*half+l]&m_bit ? 0 : 4);
      y = d * scale * value
"""

from __future__ import annotations

import torch

from agiws_neural_quant.kquant._gguf_bits import f16_view_to_f32

# GGML dtype ids (local - avoids circular import via converters/__init__).
GGML_TYPE_Q2_K = 10
GGML_TYPE_Q3_K = 11

QK_K = 256


def dequant_q2_k_packed_rows(
    raw: torch.Tensor, start: int, end: int, cols: int
) -> torch.Tensor:
    """Dequant Q2_K rows [start:end]. raw: [out, 84*n_blocks] uint8."""
    block = QK_K
    elem = 84
    chunk = raw[start:end].to(torch.int32)
    cr = chunk.shape[0]
    n_blocks = (cols + block - 1) // block
    chunk = chunk.reshape(cr, n_blocks, elem)
    scales = chunk[..., 0:16]                                   # [cr, nb, 16]
    qs = chunk[..., 16:80]                                      # [cr, nb, 64]
    d = f16_view_to_f32(chunk[..., 80:82])                      # [cr, nb]
    dmin = f16_view_to_f32(chunk[..., 82:84])                   # [cr, nb]

    sc4 = (scales & 0x0F).to(torch.float32)                     # [cr, nb, 16]
    m4 = (scales >> 4).to(torch.float32)

    y = torch.zeros(cr, n_blocks, QK_K, dtype=torch.float32, device=raw.device)
    for is_ in range(16):
        n = is_ // 8
        j_local = (is_ // 2) % 4
        half = is_ % 2
        byte_off = 32 * n + 16 * half
        shift = 2 * j_local
        qb = (qs[..., byte_off:byte_off + 16] >> shift) & 0x03   # [cr, nb, 16]
        sc = sc4[..., is_:is_ + 1].unsqueeze(-1)                 # [cr, nb, 1, 1]
        mn = m4[..., is_:is_ + 1].unsqueeze(-1)
        d4 = d.unsqueeze(-1).unsqueeze(-1)
        dm4 = dmin.unsqueeze(-1).unsqueeze(-1)
        y[..., is_ * 16:is_ * 16 + 16] = (d4 * sc * qb.unsqueeze(2) - dm4 * mn).squeeze(2)
    return y.reshape(cr, n_blocks * block)[:, :cols]


def dequant_q3_k_packed_rows(
    raw: torch.Tensor, start: int, end: int, cols: int
) -> torch.Tensor:
    """Dequant Q3_K rows [start:end]. raw: [out, 110*n_blocks] uint8."""
    block = QK_K
    elem = 110
    chunk = raw[start:end].to(torch.int32)
    cr = chunk.shape[0]
    n_blocks = (cols + block - 1) // block
    chunk = chunk.reshape(cr, n_blocks, elem)
    hmask = chunk[..., 0:32]                                    # [cr, nb, 32]
    qs = chunk[..., 32:96]                                      # [cr, nb, 64]
    scales = chunk[..., 96:108]                                 # [cr, nb, 12]
    d = f16_view_to_f32(chunk[..., 108:110])                    # [cr, nb]

    # Unpack 12 bytes into 4 little-endian u32, then reorder (aux logic).
    # aux[0..2] = scales[0:12] as 3 LE u32; aux[3] is garbage-filled in C.
    s0 = (scales[..., 0] | (scales[..., 1] << 8) | (scales[..., 2] << 16) | (scales[..., 3] << 24)).to(torch.int64)
    s1 = (scales[..., 4] | (scales[..., 5] << 8) | (scales[..., 6] << 16) | (scales[..., 7] << 24)).to(torch.int64)
    s2 = (scales[..., 8] | (scales[..., 9] << 8) | (scales[..., 10] << 16) | (scales[..., 11] << 24)).to(torch.int64)
    kmask1 = 0x03030303
    kmask2 = 0x0F0F0F0F
    tmp = s2
    aux2 = ((s0 >> 4) & kmask2) | (((tmp >> 4) & kmask1) << 4)
    aux3 = ((s1 >> 4) & kmask2) | (((tmp >> 6) & kmask1) << 4)
    aux0 = (s0 & kmask2) | (((tmp >> 0) & kmask1) << 4)
    aux1 = (s1 & kmask2) | (((tmp >> 2) & kmask1) << 4)
    # 16 int8 scales from the 16 LE bytes of aux0..aux3.
    sc8 = torch.zeros(cr, n_blocks, 16, dtype=torch.int64, device=raw.device)
    for i, aux in enumerate([aux0, aux1, aux2, aux3]):
        for b in range(4):
            byte = (aux >> (8 * b)) & 0xFF
            sc8[..., 4 * i + b] = torch.where(byte >= 128, byte - 256, byte)
    sc_f = (sc8 - 32).to(torch.float32)                         # [cr, nb, 16]

    y = torch.zeros(cr, n_blocks, QK_K, dtype=torch.float32, device=raw.device)
    for is_ in range(16):
        n = is_ // 8
        j_local = (is_ // 2) % 4
        half = is_ % 2
        byte_off = 32 * n + 16 * half
        shift = 2 * j_local
        m_shift = j_local + 4 * n                     # m = 1 << m_shift
        qb = ((qs[..., byte_off:byte_off + 16] >> shift) & 0x03).to(torch.float32)
        hb = ((hmask[..., 16 * half:16 * half + 16] >> m_shift) & 1).to(torch.float32)
        val = qb - 4.0 * (1.0 - hb)
        sc = sc_f[..., is_:is_ + 1].unsqueeze(-1)
        d4 = d.unsqueeze(-1).unsqueeze(-1)
        y[..., is_ * 16:is_ * 16 + 16] = (d4 * sc * val.unsqueeze(2)).squeeze(2)
    return y.reshape(cr, n_blocks * block)[:, :cols]


__all__ = [
    "dequant_q2_k_packed_rows",
    "dequant_q3_k_packed_rows",
    "GGML_TYPE_Q2_K",
    "GGML_TYPE_Q3_K",
]