"""GPTQ-style error-compensated ternary latents (T5.0).

GPTQ quantises a weight's input columns one at a time and pushes each column's rounding error onto the columns not
yet quantised, weighted by the inverse Hessian of the layer's inputs (H = XᵀX). Here the quantiser is the student's
own (Q-A: ternary codes, per-row absmean scale per group of 128 input columns), and the result is a *latent*: the
trainer re-derives scale and codes from it on every forward (`fake_quant`). So the latent of each group is the
error-compensated values GPTQ actually rounded. The group's scale is their absmean, solved by re-running the group
until the scale it uses is the absmean of the values it produces, which makes `fake_quant(latent)` give GPTQ's codes.
"""
from __future__ import annotations

import torch

from lowbit.quant.ternary import QuantSpec, init_scale, ternary_codes


def _codes(w: torch.Tensor, s: torch.Tensor) -> torch.Tensor:
    """Ternary codes of one column [E, R] against per-row scales s [E, R] (as `ternary_codes`, absmean)."""
    return torch.sign(w) * (w.abs() > 0.5 * s)


def _inverse_cholesky_upper(H: torch.Tensor, damp: float, tries: int = 4) -> tuple[torch.Tensor, torch.Tensor, int]:
    """Upper Cholesky factor of H⁻¹ per batch element, in fp64. An element that is not positive-definite retries
    with 10× the dampening; after `tries` it gets the identity (round-to-nearest, no compensation). Returns
    (U, dampening used per element, number of elements left on the identity)."""
    E, C, _ = H.shape
    eye = torch.eye(C, dtype=H.dtype, device=H.device)
    mean_diag = torch.diagonal(H, dim1=-2, dim2=-1).mean(-1)
    U = eye.expand(E, C, C).clone()
    used = torch.full((E,), float("nan"), dtype=H.dtype, device=H.device)
    todo = torch.ones(E, dtype=torch.bool, device=H.device)
    d = damp
    for _ in range(tries):
        Hd = H[todo] + (d * mean_diag[todo])[:, None, None] * eye
        L, info = torch.linalg.cholesky_ex(Hd)
        Uc, info2 = torch.linalg.cholesky_ex(torch.cholesky_inverse(L), upper=True)
        good = (info == 0) & (info2 == 0) & torch.isfinite(Uc).flatten(1).all(-1)
        idx = todo.nonzero()[:, 0]
        U[idx[good]] = Uc[good]
        used[idx[good]] = d
        todo[idx[good]] = False
        if not todo.any():
            break
        d *= 10
    return U, used, int(todo.sum())


@torch.no_grad()
def gptq_ternary(w: torch.Tensor, h: torch.Tensor, spec: QuantSpec, damp: float = 0.01,
                 iters: int = 3) -> tuple[torch.Tensor, dict]:
    """w [E, R, C] (rows = outputs, C = inputs), h [E, C, C] → error-compensated latent [E, R, C] in fp32.

    Batched over E (a layer's experts). Inputs with a zero Hessian diagonal (never active) are rounded like any
    other column but pass no error on (their diagonal is set to 1).
    """
    if spec.method != "absmean":
        raise NotImplementedError("GPTQ init is built for Q-A (absmean) only")
    g = spec.group_size
    E, R, C = w.shape
    W = w.detach().float().clone()
    H = h.detach().double().clone()  # fp64: real Hessians (massive activations in a few dims) break fp32 Cholesky
    diag = torch.diagonal(H, dim1=-2, dim2=-1)
    diag[diag == 0] = 1.0
    U, damp_used, n_identity = _inverse_cholesky_upper(H, damp)
    U = U.float()
    latent = torch.empty_like(W)
    used = torch.empty(E, R, C, dtype=torch.int8, device=W.device)  # the codes GPTQ chose
    passes = []
    for g0 in range(0, C, g):
        g1 = min(g0 + g, C)
        start = W[:, :, g0:g1].clone()
        s = init_scale(start, QuantSpec(method="absmean", group_size=g1 - g0))  # [E, R, 1]
        for it in range(iters):
            Wg = start.clone()
            err = torch.empty_like(Wg)
            for j in range(g1 - g0):
                i = g0 + j
                col = Wg[:, :, j]
                c = _codes(col, s[..., 0])
                used[:, :, i] = c.to(torch.int8)
                q = c * s[..., 0]
                e = (col - q) / U[:, i, i][:, None]
                err[:, :, j] = e
                if j + 1 < g1 - g0:
                    Wg[:, :, j + 1:] -= e[:, :, None] * U[:, i, i + 1:g1][:, None, :]
            s_new = init_scale(Wg, QuantSpec(method="absmean", group_size=g1 - g0))
            same = torch.equal(s_new, s)
            s = s_new
            if same:
                break
        passes.append(it + 1)
        latent[:, :, g0:g1] = Wg
        if g1 < C:
            W[:, :, g1:] -= torch.bmm(err, U[:, g0:g1, g1:])
    codes_latent = ternary_codes(latent, init_scale(latent, spec), spec)
    return latent, {"group_passes_max": max(passes), "group_passes_mean": sum(passes) / len(passes),
                    "zero_fraction": float((codes_latent == 0).float().mean()),
                    "code_mismatch": float((codes_latent != used).float().mean()),
                    "damp_max": float(damp_used.nan_to_num(0).max()), "damp_raised": int((damp_used > damp).sum()),
                    "identity_fallback": n_identity}