"""GPTQ-style ternary latents (T5.0) on random correlated inputs."""
import torch

from lowbit.quant.gptq import gptq_ternary
from lowbit.quant.ternary import QuantSpec, fake_quant

SPEC = QuantSpec(method="absmean", group_size=128)


def data(E=3, R=32, C=256, n=2048, seed=0):
    g = torch.Generator().manual_seed(seed)
    w = torch.randn(E, R, C, generator=g) * 0.05
    mix = torch.randn(C, C, generator=g) / C ** 0.5 + torch.eye(C)  # correlated inputs, as real activations are
    x = torch.randn(E, n, C, generator=g) @ mix
    return w, x, x.transpose(1, 2) @ x


def out_err(w, w_hat, x):
    return float(((x @ (w - w_hat).transpose(1, 2)) ** 2).sum())


def test_gptq_beats_round_to_nearest_on_output_error():
    w, x, h = data()
    latent, stats = gptq_ternary(w, h, SPEC)
    rtn = out_err(w, fake_quant(w, SPEC), x)
    gq = out_err(w, fake_quant(latent, SPEC), x)
    assert gq < 0.8 * rtn, (gq, rtn)
    assert stats["code_mismatch"] < 0.01  # the trainer's quantiser re-derives (almost) GPTQ's codes from the latent


def test_identity_hessian_is_round_to_nearest():
    w, _, _ = data()
    h = torch.eye(w.shape[-1]).expand(w.shape[0], -1, -1).clone()
    latent, _ = gptq_ternary(w, h, SPEC, damp=0.0)
    assert torch.equal(fake_quant(latent, SPEC), fake_quant(w, SPEC))  # no cross-column terms: nothing to compensate


def test_ill_conditioned_hessian_does_not_fail():
    w, x, h = data(n=64)  # fewer samples than inputs: rank-deficient
    x[:, :, 3] *= 1e4     # one massive activation dimension, as real hidden states have
    h = x.transpose(1, 2) @ x
    latent, stats = gptq_ternary(w, h, SPEC)
    assert torch.isfinite(latent).all() and stats["identity_fallback"] <= w.shape[0]