"""Bit-exact Q2_K block representation, packing and reconstruction.

Layout (84 bytes per 256 weights)::

    scales[16]  uint8   low nibble = Ls[j], high nibble = Lm[j]
    qs[64]      uint8   four 2-bit codes per byte
    d           fp16    super-block scale
    dmin        fp16    super-block minimum scale

The decoder reconstructs::

    x_hat[16*j + k] = d * Ls[j] * L[16*j + k] - dmin * Lm[j]

with ``L in {0,1,2,3}`` and ``Ls, Lm in {0..15}``.
"""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np

from .fp16 import from_fp16_bits, nearest_int, to_fp16_bits

QK_K = 256
NS = QK_K // 16
NQ = 4


@dataclass(slots=True)
class Q2KBlock:
    scales: np.ndarray
    qs: np.ndarray
    d: np.uint16
    dmin: np.uint16

    @property
    def d_float(self) -> np.float32:
        return np.float32(from_fp16_bits(np.uint16(self.d)))

    @property
    def dmin_float(self) -> np.float32:
        return np.float32(from_fp16_bits(np.uint16(self.dmin)))

    def unpack_scales(self) -> tuple[np.ndarray, np.ndarray]:
        scales = np.asarray(self.scales, dtype=np.uint8)
        return scales & 0x0F, scales >> 4

    def dequantize(self) -> np.ndarray:
        ls, lm = self.unpack_scales()
        return dequantize_params(
            self.d_float, self.dmin_float, ls, lm, unpack_codes(self.qs)
        )

    def error(self, x: np.ndarray) -> float:
        e = np.asarray(x, dtype=np.float64) - self.dequantize().astype(np.float64)
        return float(np.dot(e, e))


def unpack_codes(qs: np.ndarray) -> np.ndarray:
    qs = np.asarray(qs, dtype=np.uint8)
    codes = np.empty(QK_K, dtype=np.uint8)
    for base in (0, 128):
        q = qs[base // 4 : base // 4 + 32]
        for shift_level in range(4):
            lo = base + 32 * shift_level
            codes[lo : lo + 32] = (q >> np.uint8(2 * shift_level)) & np.uint8(3)
    return codes


def pack_codes(codes: np.ndarray) -> np.ndarray:
    codes = np.asarray(codes, dtype=np.uint8)
    if codes.shape != (QK_K,):
        raise ValueError("codes must have shape (256,)")
    qs = np.empty(QK_K // 4, dtype=np.uint8)
    for base in (0, 128):
        q = (
            codes[base : base + 32]
            | (codes[base + 32 : base + 64] << np.uint8(2))
            | (codes[base + 64 : base + 96] << np.uint8(4))
            | (codes[base + 96 : base + 128] << np.uint8(6))
        )
        qs[base // 4 : base // 4 + 32] = q
    return qs


def pack_scales(ls: np.ndarray, lm: np.ndarray) -> np.ndarray:
    return (np.asarray(ls, dtype=np.uint8) & 0x0F) | (
        (np.asarray(lm, dtype=np.uint8) & 0x0F) << np.uint8(4)
    )


def dequantize_params(
    d, dmin, ls: np.ndarray, lm: np.ndarray, codes: np.ndarray
) -> np.ndarray:
    ls = np.asarray(ls, dtype=np.float32)
    lm = np.asarray(lm, dtype=np.float32)
    codes = np.asarray(codes, dtype=np.float32)
    d_sub = np.repeat(np.float32(d) * ls, 16)
    m_sub = np.repeat(np.float32(dmin) * lm, 16)
    return d_sub * codes - m_sub


def codes_from_params(
    x: np.ndarray, d, dmin, ls: np.ndarray, lm: np.ndarray
) -> np.ndarray:
    x = np.asarray(x, dtype=np.float32)
    d_sub = np.repeat(np.float32(d) * np.asarray(ls, dtype=np.float32), 16)
    m_sub = np.repeat(np.float32(dmin) * np.asarray(lm, dtype=np.float32), 16)
    safe = d_sub >= np.float32(1.0e-15)
    codes = np.zeros(QK_K, dtype=np.uint8)
    if np.any(safe):
        q = nearest_int((x[safe] + m_sub[safe]) / d_sub[safe])
        codes[safe] = np.clip(q, 0, 3).astype(np.uint8)
    return codes


def derive_subscales(
    scales: np.ndarray, mins: np.ndarray, d, dmin
) -> tuple[np.ndarray, np.ndarray]:
    scales = np.asarray(scales, dtype=np.float32)
    mins = np.asarray(mins, dtype=np.float32)
    if float(d) > 1.0e-15:
        ls = np.clip(nearest_int(scales / np.float32(d)), 0, 15).astype(np.uint8)
    else:
        ls = np.zeros(NS, dtype=np.uint8)
    if float(dmin) > 1.0e-15:
        lm = np.clip(nearest_int(mins / np.float32(dmin)), 0, 15).astype(np.uint8)
    else:
        lm = np.zeros(NS, dtype=np.uint8)
    return ls, lm


def weighted_sse(
    x: np.ndarray,
    w: np.ndarray | None,
    d,
    dmin,
    ls: np.ndarray,
    lm: np.ndarray,
    codes: np.ndarray,
) -> float:
    xh = dequantize_params(d, dmin, ls, lm, codes)
    e = np.asarray(x, dtype=np.float64) - xh.astype(np.float64)
    if w is None:
        return float(np.dot(e, e))
    return float(np.dot(np.asarray(w, dtype=np.float64), e * e))


def make_block(d, dmin, ls, lm, codes) -> Q2KBlock:
    return Q2KBlock(
        scales=pack_scales(ls, lm),
        qs=pack_codes(codes),
        d=np.uint16(to_fp16_bits(np.float32(d))),
        dmin=np.uint16(to_fp16_bits(np.float32(dmin))),
    )


def blocks_to_arrays(blocks: list[Q2KBlock]) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    n = len(blocks)
    scales = np.empty((n, NS), dtype=np.uint8)
    qs = np.empty((n, QK_K // 4), dtype=np.uint8)
    dm = np.empty((n, 2), dtype=np.uint16)
    for i, b in enumerate(blocks):
        scales[i] = b.scales
        qs[i] = b.qs
        dm[i, 0] = b.d
        dm[i, 1] = b.dmin
    return scales, qs, dm


def blocks_to_bytes(scales: np.ndarray, qs: np.ndarray, dm: np.ndarray) -> bytes:
    """Serialize Q2_K block arrays to GGUF's 84-byte block layout."""
    scales = np.asarray(scales, dtype=np.uint8)
    qs = np.asarray(qs, dtype=np.uint8)
    dm = np.asarray(dm, dtype=np.uint16)
    n = scales.shape[0]
    out = np.empty((n, 84), dtype=np.uint8)
    out[:, :16] = scales
    out[:, 16:80] = qs
    out[:, 80:82] = np.ascontiguousarray(dm[:, 0]).view(np.uint8).reshape(n, 2)
    out[:, 82:84] = np.ascontiguousarray(dm[:, 1]).view(np.uint8).reshape(n, 2)
    return out.tobytes()


def dequantize_blocks(blocks: list[Q2KBlock]) -> np.ndarray:
    if not blocks:
        return np.empty(0, dtype=np.float32)
    return np.concatenate([b.dequantize() for b in blocks])


def codes_from_params_batch(
    x: np.ndarray, d: np.ndarray, dmin: np.ndarray, ls: np.ndarray, lm: np.ndarray
) -> np.ndarray:
    x = np.asarray(x, dtype=np.float32)
    n = x.shape[0]
    d_sub = np.repeat(
        np.asarray(d, dtype=np.float32)[:, None] * np.asarray(ls, dtype=np.float32), 16, axis=1
    )
    m_sub = np.repeat(
        np.asarray(dmin, dtype=np.float32)[:, None] * np.asarray(lm, dtype=np.float32), 16, axis=1
    )
    safe = d_sub >= np.float32(1.0e-15)
    codes = np.zeros((n, QK_K), dtype=np.uint8)
    if np.any(safe):
        q = nearest_int((x[safe] + m_sub[safe]) / d_sub[safe])
        codes[safe] = np.clip(q, 0, 3).astype(np.uint8)
    return codes


def dequantize_params_batch(
    d: np.ndarray, dmin: np.ndarray, ls: np.ndarray, lm: np.ndarray, codes: np.ndarray
) -> np.ndarray:
    n = codes.shape[0]
    d_sub = np.repeat(
        np.asarray(d, dtype=np.float32)[:, None] * np.asarray(ls, dtype=np.float32), 16, axis=1
    )
    m_sub = np.repeat(
        np.asarray(dmin, dtype=np.float32)[:, None] * np.asarray(lm, dtype=np.float32), 16, axis=1
    )
    return d_sub * np.asarray(codes, dtype=np.float32) - m_sub


def weighted_sse_batch(
    x: np.ndarray,
    w: np.ndarray | None,
    d: np.ndarray,
    dmin: np.ndarray,
    ls: np.ndarray,
    lm: np.ndarray,
    codes: np.ndarray,
) -> np.ndarray:
    xh = dequantize_params_batch(d, dmin, ls, lm, codes)
    e = np.asarray(x, dtype=np.float64) - xh.astype(np.float64)
    if w is None:
        return np.sum(e * e, axis=1)
    return np.sum(np.asarray(w, dtype=np.float64) * e * e, axis=1)


def pin_q2k_dc_residual(d, m, ls, lm, codes, X):
    """Zero-DC pin: re-solve (d, m) on the exact ``sum(e) == 0`` line.

    With e_i = X_i - d*ls_g*c_i + m*lm_g, requiring Σe = 0 makes m a line
    m(d) = (d·Σa - ΣX)/Σb  (a = ls·c, b = lm); subtracting the residual mean
    along that line is a translation that can only *lower* SSE, so the
    closed-form optimum d* = Σ(u·v)/Σ(v²) dominates the input state on both
    counts and leaves dc ~ 0 (bounded below by the fp16 storage of the
    ``d``/``m`` fields).  Blocks without a usable translation dof (Σb == 0
    or Σv² == 0) keep their input parameters.
    """
    ls64 = np.asarray(ls, dtype=np.float64)
    lm64 = np.asarray(lm, dtype=np.float64)
    c64 = np.asarray(codes, dtype=np.float64)
    a_el = np.repeat(ls64, 16, axis=1) * c64
    b_el = np.repeat(lm64, 16, axis=1)
    sx0 = np.sum(np.asarray(X, dtype=np.float64), axis=1)
    sa = a_el.sum(axis=1)
    sb = b_el.sum(axis=1)
    sb_safe = np.where(sb > 0.0, sb, np.nan)
    u = np.asarray(X, dtype=np.float64) - b_el * (sx0 / sb_safe)[:, None]
    v = a_el - b_el * (sa / sb_safe)[:, None]
    vv = np.sum(v * v, axis=1)
    uv = np.sum(u * v, axis=1)
    ok = np.isfinite(vv) & (vv > 0.0) & np.isfinite(uv) & (sb > 0.0)
    d_new = np.where(ok, uv / np.where(ok, vv, 1.0), np.asarray(d, dtype=np.float64))
    m_new = (d_new * sa - sx0) / np.where(ok, sb, 1.0)
    m_new = np.where(ok, m_new, np.asarray(m, dtype=np.float64))
    # match the fp16 fields the block will actually carry in the file
    d_new = from_fp16_bits(to_fp16_bits(d_new.astype(np.float32))).astype(np.float64)
    m_new = from_fp16_bits(to_fp16_bits(m_new.astype(np.float32))).astype(np.float64)
    return d_new, m_new