import time

import torch

from torchao.quantization.quant_primitives import (
    _choose_qparams_and_quantize_affine_hqq,
    _choose_qparams_and_quantize_affine_sinq,
)


def compute_imbalance(W):
    """Compute matrix imbalance as defined in SINQ paper (Eq. 4)"""
    q_max_row = W.std(dim=1).max()
    q_max_col = W.std(dim=0).max()
    q_min_row = W.std(dim=1).min()
    q_min_col = W.std(dim=0).min()

    q_max = max(q_max_row, q_max_col)
    q_min = min(q_min_row, q_min_col)

    imbalance = q_max / max(q_min, 1e-8)
    return imbalance.item()


def compute_metrics(W_orig, W_dq):
    """Compute reconstruction metrics"""
    return {
        "mse": torch.mean((W_orig - W_dq) ** 2).item(),
        "mae": torch.mean(torch.abs(W_orig - W_dq)).item(),
        "rel_error": (torch.norm(W_orig - W_dq) / torch.norm(W_orig)).item(),
    }


def test_quantization_methods(
    shape: tuple[int, int] = (4096, 11008),
    nbits: int = 4,
    group_size: int = 128,
    device: str = "cpu",
):
    """Test and compare HQQ vs SINQ quantization methods."""
    print(f"\nTesting quantization methods on {shape} matrix")
    print(f"Bits: {nbits}, Group size: {group_size}")
    print("=" * 80)

    # Generate test weight matrix
    torch.manual_seed(42)
    W_original = torch.randn(shape, dtype=torch.float32, device=device) * 0.02

    # Add outliers to simulate real LLM weights
    outlier_mask = torch.rand(shape, device=device) < 0.01
    W_original[outlier_mask] *= 5.0

    print("\n Original Matrix Statistics:")
    print(f"  Mean: {W_original.mean():.6f}")
    print(f"  Std: {W_original.std():.6f}")
    print(f"  Min: {W_original.min():.6f}, Max: {W_original.max():.6f}")
    print(f"  Imbalance: {compute_imbalance(W_original):.4f}")

    # ========================================================================
    # HQQ Quantization
    # ========================================================================
    print(f"\n{'─' * 80}")
    print("HQQ (Half-Quadratic Quantization)")
    print(f"{'─' * 80}")

    start_time = time.time()
    W_q_hqq, scale_hqq, zero_hqq, _ = _choose_qparams_and_quantize_affine_hqq(
        tensor=W_original,
        nbits=nbits,
        group_size=group_size,
        optimize=True,
        axis=1,
        device=device,
        raw_output=False,
    )
    hqq_time = time.time() - start_time

    # Dequantize: W = (W_q - zero) * scale
    W_reshaped = W_q_hqq.float().reshape(-1, group_size)
    W_dq_hqq = ((W_reshaped - zero_hqq) * scale_hqq).reshape(shape)
    hqq_metrics = compute_metrics(W_original, W_dq_hqq)

    print(f"  Quantization Time: {hqq_time:.4f}s")
    print(f"  MSE: {hqq_metrics['mse']:.8f}")
    print(f"  MAE: {hqq_metrics['mae']:.8f}")
    print(f"  Relative Error: {hqq_metrics['rel_error']:.6f}")
    print(f"  Dequantized Imbalance: {compute_imbalance(W_dq_hqq):.4f}")

    # ========================================================================
    # SINQ Quantization
    # ========================================================================
    print(f"\n{'─' * 80}")
    print("SINQ (Sinkhorn-Normalized Quantization)")
    print(f"{'─' * 80}")

    start_time = time.time()
    W_q_sinq, scale_row_sinq, zero_sinq, scale_col_sinq, _ = (
        _choose_qparams_and_quantize_affine_sinq(
            tensor=W_original,
            nbits=nbits,
            group_size=group_size,
            device=device,
        )
    )
    sinq_time = time.time() - start_time

    # Dequantize: W = scale_row * (W_q - zero) * scale_col
    W_q_reshaped = W_q_sinq.float().reshape(-1, group_size)
    scale_row_flat = scale_row_sinq.view(-1, 1)  # (262144, 1)
    zero_flat = zero_sinq.view(-1, 1)  # (262144, 1)

    W_dq_sinq = scale_row_flat * (W_q_reshaped - zero_flat)
    W_dq_sinq = W_dq_sinq.reshape(shape) * scale_col_sinq.reshape(1, -1)

    sinq_metrics = compute_metrics(W_original, W_dq_sinq)

    print(f"  Quantization Time: {sinq_time:.4f}s")
    print(f"  MSE: {sinq_metrics['mse']:.8f}")
    print(f"  MAE: {sinq_metrics['mae']:.8f}")
    print(f"  Relative Error: {sinq_metrics['rel_error']:.6f}")
    print(f"  Dequantized Imbalance: {compute_imbalance(W_dq_sinq):.4f}")

    return {
        "hqq": {**hqq_metrics, "time": hqq_time},
        "sinq": {**sinq_metrics, "time": sinq_time},
    }


if __name__ == "__main__":
    test_quantization_methods()