"""Reproduce the GPTQ article's synthetic example with NumPy only.

Run: python gptq_example.py
This illustrates the sequential GPTQ core, not a model packer or serving kernel.
One fixed symmetric grid: step 0.5, integer codes -7..7, ties to even.
"""
import json
import numpy as np

W = np.array([
    [0.10, 2.50, -0.30, 0.05],
    [-0.20, 3.20, 0.15, -0.08],
    [0.40, -2.80, 0.25, 0.12],
    [-0.15, 2.10, -0.40, 0.06],
], dtype=np.float64)
X = np.array([
    [0.60, 0.60, 0.60, 0.60, 0.40, 0.40, 0.40, 0.40],
    [4.80, 4.80, 3.20, 3.20, 4.80, 4.80, 3.20, 3.20],
    [0.36, 0.24, 0.36, 0.24, 0.36, 0.24, 0.36, 0.24],
    [0.24, 0.16, 0.16, 0.24, 0.16, 0.24, 0.24, 0.16],
], dtype=np.float64)
STRESS = np.array([0.5, 4.0, -0.3, -0.2], dtype=np.float64)
STEP = 0.5


def quantize(values, step=STEP):
    return np.clip(np.rint(values / step), -7, 7) * step


def factor(activations, percdamp=0.01):
    h = 2 * activations @ activations.T
    damping = percdamp * np.diag(h).mean()
    a = h + damping * np.eye(h.shape[0])
    # Factorization-based solves; no generic matrix-inverse call.
    lower = np.linalg.cholesky(a)
    g = np.linalg.solve(lower.T, np.linalg.solve(lower, np.eye(a.shape[0])))
    upper = np.linalg.cholesky(g).T
    return h, damping, a, upper


def sweep(weights, upper):
    working = weights.copy()
    result = np.zeros_like(weights)
    trace = []
    for q in range(weights.shape[1]):
        before = working[:, q].copy()
        result[:, q] = quantize(before)
        error = (before - result[:, q]) / upper[q, q]
        working[:, q:] -= error[:, None] * upper[q, q:]
        trace.append({
            'column': q, 'before': before.copy(),
            'rounded': result[:, q].copy(),
            'residual': before - result[:, q],
            'ratios': upper[q, q+1:] / upper[q, q],
            'state': working.copy(),
        })
    return result, trace


def evaluate(weights):
    errors = (weights - W) @ X
    return {
        'calibration_sse': float(np.sum(errors ** 2)),
        'calibration_mean_sse': float(np.sum(errors ** 2) / X.shape[1]),
        'per_sample_sse': np.sum(errors ** 2, axis=0).tolist(),
        'weight_sse': float(np.sum((weights - W) ** 2)),
        'mean_input_output': (weights @ X.mean(axis=1)).tolist(),
        'stress_output': (weights @ STRESS).tolist(),
        'stress_sse': float(np.sum(((weights - W) @ STRESS) ** 2)),
    }


def reproduce():
    h, damping, a, upper = factor(X)
    result, trace = sweep(W, upper)
    rtn = quantize(W)
    # Independent check: solve each constrained quadratic directly, with
    # frozen prefix fixed and the remaining variables unconstrained.
    for entry in trace:
        q = entry['column']
        fixed_change = result[:, :q+1] - W[:, :q+1]
        if q+1 < W.shape[1]:
            rhs = a[q+1:, :q+1] @ fixed_change.T
            expected = W[:, q+1:] - np.linalg.solve(a[q+1:, q+1:], rhs).T
            np.testing.assert_allclose(entry['state'][:, q+1:], expected, atol=1e-12)
        np.testing.assert_allclose(entry['state'][:, :q+1], result[:, :q+1], atol=1e-12)
    np.testing.assert_allclose(X.mean(axis=1), [.5, 4, .3, .2])
    np.testing.assert_allclose(a @ (upper.T @ upper), np.eye(4), atol=1e-12)
    np.testing.assert_allclose(result, [[0,2.5,0,0], [0,3,1.5,.5], [.5,-3,1.5,1], [0,2,0,.5]])
    np.testing.assert_allclose(result / STEP, np.rint(result / STEP))
    assert np.max(np.abs(result / STEP)) <= 7
    # If input Gram is diagonal, there is no compensation across columns.
    _, _, _, independent_upper = factor(np.eye(4))
    independent_result, _ = sweep(W, independent_upper)
    np.testing.assert_allclose(independent_result, rtn)
    stats_rtn, stats_gptq = evaluate(rtn), evaluate(result)
    assert stats_gptq['calibration_sse'] < stats_rtn['calibration_sse']
    assert stats_gptq['stress_sse'] > stats_rtn['stress_sse']
    return {
        'W': W.tolist(), 'X': X.tolist(), 'H': h.tolist(), 'damping': damping,
        'U': upper.tolist(), 'RTN': rtn.tolist(), 'GPTQ': result.tolist(),
        'trace_row_1': [{
            'column': e['column'], 'before': float(e['before'][1]),
            'rounded': float(e['rounded'][1]), 'residual': float(e['residual'][1]),
            'ratios': e['ratios'].tolist(), 'state': e['state'][1].tolist(),
        } for e in trace],
        'rtn_metrics': stats_rtn, 'gptq_metrics': stats_gptq,
    }


if __name__ == '__main__':
    print(json.dumps(reproduce(), indent=2))