"""Per-layer sensitivity analysis via KL divergence measurement.

Quantizes each layer individually and measures the actual change in model
output distribution. This is the only reliable way to rank layers — proxy
metrics (weight error × activation norm) produce constant 4-bit/8-bit ratios
and rank by activation magnitude alone, missing downstream amplification.
"""

import os
import json
from dataclasses import dataclass

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from tqdm import tqdm


@dataclass
class SensitivityResult:
    layer_name: str
    sensitivities: dict[int, float]  # bits -> KL divergence
    param_count: int


def _simulate_quantize(weight: torch.Tensor, bits: int, group_size: int = 64) -> torch.Tensor:
    """Simulate affine quantization matching MLX's scheme.

    Round-trip: quantize then dequantize to get the quantized approximation.
    """
    orig_shape = weight.shape
    if weight.dim() == 2:
        rows, cols = weight.shape
        if cols % group_size != 0:
            pad_cols = group_size - (cols % group_size)
            weight = F.pad(weight, (0, pad_cols))
        weight = weight.reshape(-1, group_size)
    else:
        weight = weight.reshape(-1, group_size)

    w_min = weight.min(dim=1, keepdim=True).values
    w_max = weight.max(dim=1, keepdim=True).values

    n_levels = 2 ** bits - 1
    scale = (w_max - w_min) / n_levels
    scale = torch.clamp(scale, min=1e-8)

    quantized = torch.round((weight - w_min) / scale)
    quantized = torch.clamp(quantized, 0, n_levels)
    dequantized = quantized * scale + w_min

    flat = dequantized.reshape(-1)
    n_orig = int(np.prod(orig_shape))
    return flat[:n_orig].reshape(orig_shape)


def analyze_llm_sensitivity(
    model: nn.Module,
    calibration_data: list[torch.Tensor],
    candidate_bits: list[int] = [4, 8],
    group_size: int = 64,
    n_calibration: int = 2,
    seq_len: int = 128,
    checkpoint_path: str | None = None,
) -> list[SensitivityResult]:
    """Per-layer KL divergence sensitivity analysis.

    For each quantizable layer and each candidate bit-width:
    1. Replace weights with simulated quantized version
    2. Run forward pass, measure KL divergence of output logits vs reference
    3. Restore original weights

    Made practical via short sequences (128 tokens), few samples (2),
    batch size 1, and progress checkpointing for resume.
    """
    device = next(model.parameters()).device
    model.eval()

    quantizable = []
    for name, module in model.named_modules():
        if isinstance(module, nn.Linear):
            if module.weight.shape[-1] % group_size == 0:
                quantizable.append((name, module))

    print(f"  Analyzing {len(quantizable)} layers × {len(candidate_bits)} bit-widths")
    print(f"  Calibration: {n_calibration} samples × {seq_len} tokens")

    # Load checkpoint if resuming — supports incremental bit-width additions
    completed = {}
    if checkpoint_path and os.path.exists(checkpoint_path):
        with open(checkpoint_path) as f:
            for d in json.load(f):
                existing_bits = set(int(k) for k in d["sensitivities"].keys())
                needed_bits = set(candidate_bits)
                if needed_bits.issubset(existing_bits):
                    # All requested bit-widths already measured
                    completed[d["layer_name"]] = d
                # else: need to re-measure this layer for missing bit-widths
        n_done = len(completed)
        n_total = len(quantizable)
        if n_done > 0:
            print(f"  Resuming from checkpoint: {n_done}/{n_total} layers fully done")

    # Prepare calibration — truncate to seq_len for speed
    cal_samples = []
    for t in calibration_data[:n_calibration]:
        cal_samples.append(t[:seq_len].to(device))

    # Reference logits (batch=1 for memory)
    print("  Computing reference logits...")
    ref_logits_list = []
    with torch.no_grad():
        for sample in cal_samples:
            out = model(sample.unsqueeze(0))
            ref_logits_list.append(out.logits.float().cpu())

    # Analyze each layer
    results = []
    for layer_name, module in tqdm(quantizable, desc="  KL sensitivity"):
        if layer_name in completed:
            d = completed[layer_name]
            results.append(SensitivityResult(
                layer_name=d["layer_name"],
                sensitivities={int(k): v for k, v in d["sensitivities"].items()},
                param_count=d["param_count"],
            ))
            continue

        sensitivities = {}
        for bits in candidate_bits:
            orig_weight = module.weight.data.clone()
            q_weight = _simulate_quantize(orig_weight, bits, group_size)
            module.weight.data.copy_(q_weight)

            kl_divs = []
            with torch.no_grad():
                for i, sample in enumerate(cal_samples):
                    out = model(sample.unsqueeze(0))
                    q_logits = out.logits.float().cpu()

                    ref_probs = F.softmax(ref_logits_list[i], dim=-1).clamp(min=1e-8)
                    q_log_probs = F.log_softmax(q_logits, dim=-1)
                    kl = F.kl_div(q_log_probs, ref_probs, reduction="batchmean")
                    kl_divs.append(kl.item())

            sensitivities[bits] = float(np.mean(kl_divs))
            module.weight.data.copy_(orig_weight)

        results.append(SensitivityResult(
            layer_name=layer_name,
            sensitivities=sensitivities,
            param_count=module.weight.numel(),
        ))

        if checkpoint_path:
            _save_checkpoint(results, checkpoint_path)

    return results


def _save_checkpoint(results: list[SensitivityResult], path: str):
    """Save sensitivity results as a checkpoint file."""
    os.makedirs(os.path.dirname(path) if os.path.dirname(path) else ".", exist_ok=True)
    data = [
        {
            "layer_name": r.layer_name,
            "sensitivities": {str(k): v for k, v in r.sensitivities.items()},
            "param_count": r.param_count,
        }
        for r in results
    ]
    with open(path, "w") as f:
        json.dump(data, f, indent=2)



def print_sensitivity_report(results: list[SensitivityResult], top_n: int = 20):
    """Print a visual sensitivity report."""
    if not results:
        print("  No sensitivity results to display")
        return

    # Determine which bit-widths were measured
    all_bits = sorted(set(b for r in results for b in r.sensitivities.keys()))
    sort_bits = 4 if 4 in all_bits else all_bits[0]  # Sort by 4-bit if available

    # Filter out NaN values for sorting (NaN layers treated as most sensitive)
    def _sort_key(r):
        v = r.sensitivities.get(sort_bits, 0)
        return float('inf') if (isinstance(v, float) and np.isnan(v)) else v
    sorted_results = sorted(results, key=_sort_key, reverse=True)

    print(f"\n  Layer Sensitivity Report (sorted by {sort_bits}-bit KL divergence)")
    print("  " + "=" * (52 + 12 * len(all_bits)))
    header_bits = "".join(f" {b:>5}-bit" for b in all_bits)
    print(f"  {'Layer':<50} {header_bits}")
    print("  " + "-" * (52 + 12 * len(all_bits)))

    for r in sorted_results[:top_n]:
        max_s_val = _sort_key(sorted_results[0])
        max_s = max(max_s_val if max_s_val != float('inf') else 1, 1e-10)
        s_sort = r.sensitivities.get(sort_bits, 0)
        if isinstance(s_sort, float) and np.isnan(s_sort):
            bar = "█" * 20 + " (NaN)"
        else:
            bar_len = min(int(s_sort / max_s * 20), 20)
            bar = "█" * bar_len

        short_name = r.layer_name[-48:] if len(r.layer_name) > 48 else r.layer_name
        vals = "".join(
            f" {'NaN':>9}" if (isinstance(r.sensitivities.get(b, 0), float) and np.isnan(r.sensitivities.get(b, 0)))
            else f" {r.sensitivities.get(b, 0):>9.4f}"
            for b in all_bits
        )
        print(f"  {short_name:<50}{vals}  {bar}")

    if len(sorted_results) > top_n:
        print(f"  ... and {len(sorted_results) - top_n} more layers")

    # Filter NaN for summary stats
    s_values = [r.sensitivities.get(sort_bits, 0) for r in sorted_results
                if not (isinstance(r.sensitivities.get(sort_bits, 0), float)
                        and np.isnan(r.sensitivities.get(sort_bits, 0)))]
    if s_values:
        print(f"\n  {sort_bits}-bit KL: min={min(s_values):.6f}, max={max(s_values):.6f}, "
              f"mean={np.mean(s_values):.6f}, median={np.median(s_values):.6f}")


def apply_simulated_quantization(
    model: nn.Module,
    layer_configs: dict[str, int],
    group_size: int = 64,
) -> dict[str, torch.Tensor]:
    """Apply simulated quantization to specific layers in-place.

    Args:
        model: PyTorch model
        layer_configs: dict mapping layer_name → bits
        group_size: quantization group size

    Returns:
        dict mapping layer_name → original weight tensor (for restoration)
    """
    originals = {}
    for name, module in model.named_modules():
        if isinstance(module, nn.Linear) and name in layer_configs:
            bits = layer_configs[name]
            originals[name] = module.weight.data.clone()
            q_weight = _simulate_quantize(module.weight.data, bits, group_size)
            module.weight.data.copy_(q_weight)
    return originals


def restore_original_weights(
    model: nn.Module,
    originals: dict[str, torch.Tensor],
):
    """Restore original weights after simulated quantization."""
    for name, module in model.named_modules():
        if name in originals:
            module.weight.data.copy_(originals[name])


def analyze_vlm_sensitivity(
    model: nn.Module,
    calibration_inputs: list[dict],
    candidate_bits: list[int] = [4, 8],
    group_size: int = 64,
    n_calibration: int = 2,
    checkpoint_path: str | None = None,
) -> list[SensitivityResult]:
    """Per-layer KL divergence sensitivity analysis for VLMs.

    Same approach as LLM sensitivity but handles multimodal inputs
    (dicts with input_ids, pixel_values, attention_mask, etc.).
    """
    device = next(model.parameters()).device
    model.eval()

    quantizable = []
    for name, module in model.named_modules():
        if isinstance(module, nn.Linear):
            if module.weight.shape[-1] % group_size == 0:
                quantizable.append((name, module))

    print(f"  Analyzing {len(quantizable)} layers × {len(candidate_bits)} bit-widths")
    print(f"  Calibration: {n_calibration} multimodal samples")

    # Load checkpoint if resuming
    completed = {}
    if checkpoint_path and os.path.exists(checkpoint_path):
        with open(checkpoint_path) as f:
            for d in json.load(f):
                existing_bits = set(int(k) for k in d["sensitivities"].keys())
                needed_bits = set(candidate_bits)
                if needed_bits.issubset(existing_bits):
                    completed[d["layer_name"]] = d
        if completed:
            print(f"  Resuming from checkpoint: {len(completed)}/{len(quantizable)} layers done")

    cal_samples = calibration_inputs[:n_calibration]

    # Detect model dtype for casting float inputs
    model_dtype = next(model.parameters()).dtype

    # Helper to move a dict of tensors to device (and match dtype for floats)
    def _to_device(inputs):
        result = {}
        for k, v in inputs.items():
            if isinstance(v, torch.Tensor):
                v = v.to(device)
                if v.is_floating_point() and v.dtype != model_dtype:
                    v = v.to(model_dtype)
                result[k] = v
            else:
                result[k] = v
        return result

    # Reference logits
    print("  Computing reference logits...")
    ref_logits_list = []
    with torch.no_grad():
        for inputs in cal_samples:
            out = model(**_to_device(inputs))
            ref_logits_list.append(out.logits.float().cpu())

    # Analyze each layer
    results = []
    for layer_name, module in tqdm(quantizable, desc="  KL sensitivity"):
        if layer_name in completed:
            d = completed[layer_name]
            results.append(SensitivityResult(
                layer_name=d["layer_name"],
                sensitivities={int(k): v for k, v in d["sensitivities"].items()},
                param_count=d["param_count"],
            ))
            continue

        sensitivities = {}
        for bits in candidate_bits:
            orig_weight = module.weight.data.clone()
            q_weight = _simulate_quantize(orig_weight, bits, group_size)
            module.weight.data.copy_(q_weight)

            kl_divs = []
            with torch.no_grad():
                for i, inputs in enumerate(cal_samples):
                    out = model(**_to_device(inputs))
                    q_logits = out.logits.float().cpu()

                    ref_probs = F.softmax(ref_logits_list[i], dim=-1).clamp(min=1e-8)
                    q_log_probs = F.log_softmax(q_logits, dim=-1)
                    kl = F.kl_div(q_log_probs, ref_probs, reduction="batchmean")
                    kl_divs.append(kl.item())

            sensitivities[bits] = float(np.mean(kl_divs))
            module.weight.data.copy_(orig_weight)

        results.append(SensitivityResult(
            layer_name=layer_name,
            sensitivities=sensitivities,
            param_count=module.weight.numel(),
        ))

        if checkpoint_path:
            _save_checkpoint(results, checkpoint_path)

    return results  # end analyze_vlm_sensitivity


def _simulate_quantize_mlx(weight, bits: int, group_size: int = 64):
    """Simulate affine quantization in MLX (round-trip quantize then dequantize)."""
    import mlx.core as mx

    orig_shape = weight.shape
    if weight.ndim == 2:
        rows, cols = weight.shape
        if cols % group_size != 0:
            pad_cols = group_size - (cols % group_size)
            weight = mx.pad(weight, [(0, 0), (0, pad_cols)])
        weight = weight.reshape(-1, group_size)
    else:
        weight = weight.reshape(-1, group_size)

    w_min = weight.min(axis=1, keepdims=True)
    w_max = weight.max(axis=1, keepdims=True)

    n_levels = 2 ** bits - 1
    scale = (w_max - w_min) / n_levels
    scale = mx.maximum(scale, mx.array(1e-8))

    q = mx.round((weight - w_min) / scale)
    q = mx.clip(q, 0, n_levels)
    deq = q * scale + w_min

    flat = deq.reshape(-1)
    n_orig = 1
    for s in orig_shape:
        n_orig *= s
    return flat[:n_orig].reshape(orig_shape)


def analyze_mlx_sensitivity(
    model,
    calibration_fn,
    candidate_bits: list[int] = [4, 8],
    group_size: int = 64,
    n_calibration: int = 2,
    checkpoint_path: str | None = None,
) -> list[SensitivityResult]:
    """Per-layer KL divergence sensitivity analysis for MLX models.

    Works with any MLX model (e.g. Qwen3-ASR loaded via mlx-audio).
    Takes a callable that returns a list of (args_tuple, kwargs_dict)
    for model forward passes. The model's __call__ should return logits.
    """
    import mlx.core as mx
    import mlx.nn as mnn

    quantizable = []
    for name, module in model.named_modules():
        if isinstance(module, mnn.Linear) and not isinstance(module, mnn.QuantizedLinear):
            w = module.weight
            if w.ndim == 2 and w.shape[-1] % group_size == 0:
                quantizable.append((name, module))

    print(f"  Analyzing {len(quantizable)} layers × {len(candidate_bits)} bit-widths")

    completed = {}
    if checkpoint_path and os.path.exists(checkpoint_path):
        with open(checkpoint_path) as f:
            for d in json.load(f):
                existing_bits = set(int(k) for k in d["sensitivities"].keys())
                if set(candidate_bits).issubset(existing_bits):
                    completed[d["layer_name"]] = d
        if completed:
            print(f"  Resuming from checkpoint: {len(completed)}/{len(quantizable)} layers done")

    cal_samples = calibration_fn()[:n_calibration]
    print(f"  Calibration: {len(cal_samples)} samples")

    print("  Computing reference logits...")
    ref_logits_list = []
    for args, kwargs in cal_samples:
        logits = model(*args, **kwargs)
        ref_logits_list.append(logits)
    mx.eval(*ref_logits_list)

    results = []
    for layer_idx, (layer_name, module) in enumerate(quantizable):
        if layer_name in completed:
            d = completed[layer_name]
            results.append(SensitivityResult(
                layer_name=d["layer_name"],
                sensitivities={int(k): v for k, v in d["sensitivities"].items()},
                param_count=d["param_count"],
            ))
            continue

        sensitivities = {}
        orig_weight = module.weight

        for bits in candidate_bits:
            q_weight = _simulate_quantize_mlx(orig_weight, bits, group_size)
            module.weight = q_weight
            mx.eval(module.weight)

            kl_divs = []
            for i, (args, kwargs) in enumerate(cal_samples):
                q_logits = model(*args, **kwargs)
                mx.eval(q_logits)

                ref = ref_logits_list[i]
                ref_probs = mx.softmax(ref, axis=-1)
                ref_probs = mx.clip(ref_probs, 1e-8, None)
                q_log_probs = mx.log(mx.clip(mx.softmax(q_logits, axis=-1), 1e-8, None))
                kl = mx.mean(mx.sum(ref_probs * (mx.log(ref_probs) - q_log_probs), axis=-1))
                mx.eval(kl)
                kl_divs.append(kl.item())

            sensitivities[bits] = float(np.mean(kl_divs))
            module.weight = orig_weight
            mx.eval(module.weight)

        param_count = 1
        for s in orig_weight.shape:
            param_count *= s

        results.append(SensitivityResult(
            layer_name=layer_name,
            sensitivities=sensitivities,
            param_count=param_count,
        ))

        if checkpoint_path:
            _save_checkpoint(results, checkpoint_path)

        if (layer_idx + 1) % 10 == 0:
            print(f"  [{layer_idx + 1}/{len(quantizable)}] layers analyzed")

    return results


# ---------------------------------------------------------------------------
# Fisher-weighted fast sensitivity analysis
# ---------------------------------------------------------------------------
# The functions below compute per-layer quantization sensitivity in O(1)
# forward+backward passes (independent of the number of layers).  The key
# idea:  sensitivity_i(bits) = sum_j (dL/dw_ij)^2 * (w_ij - Q_b(w_ij))^2
# where dL/dw_ij comes from a self-distillation (KL) loss that measures how
# important each weight is, and Q_b(w_ij) - w_ij is the quantization error.
# ---------------------------------------------------------------------------


def analyze_sensitivity_fast(
    model: nn.Module,
    calibration_data,  # list[Tensor] for LLM, list[dict] for VLM
    candidate_bits: list[int] = [4, 8],
    group_size: int = 64,
    n_calibration: int = 2,
    mode: str = "llm",  # "llm" or "vlm"
) -> list[SensitivityResult]:
    """Fisher-weighted quantization sensitivity (fast).

    Instead of N_layers x N_bits forward passes, this needs only
    ``n_calibration`` forward+backward passes total — a ~100-1000x speedup
    for large models.

    Algorithm
    ---------
    1. Forward pass on calibration data  -> reference logits (detached).
    2. Self-distillation loss: KL(softmax(ref) || softmax(logits)).
    3. Backward pass -> gradients for ALL linear weights simultaneously.
    4. For each quantizable layer and each candidate bit-width:
       sensitivity = sum( grad^2 * (W - Q_b(W))^2 )

    The returned ``SensitivityResult`` list is a drop-in replacement for
    ``analyze_llm_sensitivity`` / ``analyze_vlm_sensitivity``.
    """
    import time

    t_start = time.time()
    device = next(model.parameters()).device
    model_dtype = next(model.parameters()).dtype
    model.eval()

    # --- Identify quantizable layers ----------------------------------------
    quantizable: list[tuple[str, nn.Linear]] = []
    for name, module in model.named_modules():
        if isinstance(module, nn.Linear):
            if module.weight.shape[-1] % group_size == 0:
                quantizable.append((name, module))

    print(f"  [fast] {len(quantizable)} layers × {len(candidate_bits)} bit-widths")
    print(f"  [fast] Calibration: {n_calibration} samples, mode={mode}")

    # --- Helper: move VLM input dict to device & dtype ----------------------
    def _to_device(inputs: dict) -> dict:
        result = {}
        for k, v in inputs.items():
            if isinstance(v, torch.Tensor):
                v = v.to(device)
                if v.is_floating_point() and v.dtype != model_dtype:
                    v = v.to(model_dtype)
                result[k] = v
            else:
                result[k] = v
        return result

    # --- Step 1: Reference logits (detached, no grad) -----------------------
    print("  [fast] Computing reference logits...")
    ref_logits_list: list[torch.Tensor] = []
    cal_samples = calibration_data[:n_calibration]

    with torch.no_grad():
        for sample in cal_samples:
            if mode == "llm":
                out = model(sample.to(device).unsqueeze(0))
            else:
                out = model(**_to_device(sample))
            ref_logits_list.append(out.logits.float().detach())

    # --- Step 2-3: Accumulate squared gradients via backward ----------------
    # Enable gradients on weights temporarily.
    print("  [fast] Computing Fisher gradients (forward + backward)...")
    grad_was_enabled = torch.is_grad_enabled()
    torch.set_grad_enabled(True)

    # Some weights may have requires_grad=False (e.g. frozen model).
    # We track which ones we toggle so we can restore state later.
    toggled_params: list[nn.Parameter] = []
    for _, module in quantizable:
        if not module.weight.requires_grad:
            module.weight.requires_grad_(True)
            toggled_params.append(module.weight)

    # Accumulate grad^2 across calibration samples
    grad_sq_accum: dict[str, torch.Tensor] = {}

    for sample_idx, sample in enumerate(cal_samples):
        # Zero existing grads
        model.zero_grad()

        # Forward pass — needs grad tracking
        if mode == "llm":
            out = model(sample.to(device).unsqueeze(0))
        else:
            out = model(**_to_device(sample))

        logits = out.logits.float()
        ref_logits = ref_logits_list[sample_idx]

        # Self-distillation loss: KL(ref || current)
        # = sum ref_p * (log ref_p - log cur_p)
        # F.kl_div expects (log_q, p) and computes sum p*(log p - log q)
        ref_probs = F.softmax(ref_logits.detach(), dim=-1).clamp(min=1e-8)
        log_probs = F.log_softmax(logits, dim=-1)
        loss = F.kl_div(log_probs, ref_probs, reduction="batchmean")

        # Backward — fills .grad on all weight tensors that require grad
        loss.backward()

        # Accumulate grad^2
        for name, module in quantizable:
            g = module.weight.grad
            if g is not None:
                g_sq = g.float().detach().pow(2)
                if name in grad_sq_accum:
                    grad_sq_accum[name] = grad_sq_accum[name] + g_sq
                else:
                    grad_sq_accum[name] = g_sq

    # Average over calibration samples
    n_cal = len(cal_samples)
    for name in grad_sq_accum:
        grad_sq_accum[name] = grad_sq_accum[name] / n_cal

    # Restore grad state
    model.zero_grad()
    for p in toggled_params:
        p.requires_grad_(False)
    torch.set_grad_enabled(grad_was_enabled)

    # --- Step 4: Compute Fisher-weighted quantization sensitivity -----------
    print("  [fast] Computing per-layer sensitivities...")
    results: list[SensitivityResult] = []

    for name, module in quantizable:
        w = module.weight.data.float()
        fisher = grad_sq_accum.get(name)

        sensitivities: dict[int, float] = {}
        for bits in candidate_bits:
            q_w = _simulate_quantize(w, bits, group_size)
            q_error_sq = (w - q_w).pow(2)
            if fisher is not None:
                sens = (fisher * q_error_sq).sum().item()
            else:
                # Fallback: no gradient available (e.g. unused layer)
                sens = q_error_sq.sum().item()
            sensitivities[bits] = sens

        results.append(SensitivityResult(
            layer_name=name,
            sensitivities=sensitivities,
            param_count=module.weight.numel(),
        ))

    t_elapsed = time.time() - t_start
    n_brute = len(quantizable) * len(candidate_bits) * n_cal
    n_fast = n_cal  # forward+backward passes
    print(f"  [fast] Done in {t_elapsed:.1f}s "
          f"({n_fast} fwd+bwd vs {n_brute} fwd passes for brute-force)")

    return results


def analyze_mlx_sensitivity_fast(
    model,
    calibration_fn,
    candidate_bits: list[int] = [4, 8],
    group_size: int = 64,
    n_calibration: int = 2,
) -> list[SensitivityResult]:
    """Fisher-weighted quantization sensitivity for MLX models (fast).

    MLX equivalent of ``analyze_sensitivity_fast``.  Uses ``mx.value_and_grad``
    to get parameter gradients from a self-distillation loss, then combines
    with per-bit quantization error to rank layers.

    Parameters
    ----------
    model : mlx nn.Module
        An MLX model whose ``__call__`` returns logits (or an object with a
        ``.logits`` attribute).
    calibration_fn : callable
        Returns a list of ``(args_tuple, kwargs_dict)`` for model forward calls.
    candidate_bits : list[int]
        Bit-widths to evaluate (default ``[4, 8]``).
    group_size : int
        Quantization group size (default 64).
    n_calibration : int
        Number of calibration samples to use.

    Returns
    -------
    list[SensitivityResult]
        Same format as ``analyze_mlx_sensitivity`` — drop-in replacement.
    """
    import time
    import mlx.core as mx
    import mlx.nn as mnn

    t_start = time.time()

    # --- Identify quantizable layers ----------------------------------------
    quantizable: list[tuple[str, object]] = []
    for name, module in model.named_modules():
        if isinstance(module, mnn.Linear) and not isinstance(module, mnn.QuantizedLinear):
            w = module.weight
            if w.ndim == 2 and w.shape[-1] % group_size == 0:
                quantizable.append((name, module))

    print(f"  [fast-mlx] {len(quantizable)} layers × {len(candidate_bits)} bit-widths")

    cal_samples = calibration_fn()[:n_calibration]
    print(f"  [fast-mlx] Calibration: {len(cal_samples)} samples")

    # --- Step 1: Reference logits -------------------------------------------
    print("  [fast-mlx] Computing reference logits...")
    ref_logits_list = []
    for args, kwargs in cal_samples:
        logits = model(*args, **kwargs)
        # Handle models that return an object with .logits vs raw array
        if hasattr(logits, 'logits'):
            logits = logits.logits
        ref_logits_list.append(logits)
    mx.eval(*ref_logits_list)

    # --- Step 2-3: Compute gradients via value_and_grad ---------------------
    print("  [fast-mlx] Computing Fisher gradients...")

    def _kl_loss(params, args, kwargs, ref_logits):
        """Self-distillation KL loss as a function of model params."""
        # Temporarily apply params to model
        model.update(params)
        logits = model(*args, **kwargs)
        if hasattr(logits, 'logits'):
            logits = logits.logits

        ref_probs = mx.softmax(mx.stop_gradient(ref_logits), axis=-1)
        ref_probs = mx.clip(ref_probs, 1e-8, None)
        log_probs = logits - mx.logsumexp(logits, axis=-1, keepdims=True)
        # KL(ref || cur) = sum ref * (log_ref - log_cur)
        kl = mx.mean(mx.sum(ref_probs * (mx.log(ref_probs) - log_probs), axis=-1))
        return kl

    loss_and_grad_fn = mnn.value_and_grad(model, _kl_loss)

    # Accumulate squared gradients
    grad_sq_accum: dict[str, mx.array] = {}

    for sample_idx, (args, kwargs) in enumerate(cal_samples):
        ref_logits = ref_logits_list[sample_idx]
        _loss_val, grads = loss_and_grad_fn(model.parameters(), args, kwargs, ref_logits)
        mx.eval(_loss_val)

        # Walk the gradient tree and accumulate squared grads for each
        # quantizable layer's weight.
        _accumulate_mlx_grads(grads, grad_sq_accum, quantizable)

    # Average over calibration samples
    n_cal = len(cal_samples)
    for key in grad_sq_accum:
        grad_sq_accum[key] = grad_sq_accum[key] / n_cal
        mx.eval(grad_sq_accum[key])

    # --- Step 4: Compute Fisher-weighted quantization sensitivity -----------
    print("  [fast-mlx] Computing per-layer sensitivities...")
    results: list[SensitivityResult] = []

    for name, module in quantizable:
        w = module.weight
        fisher = grad_sq_accum.get(name)

        sensitivities: dict[int, float] = {}
        for bits in candidate_bits:
            q_w = _simulate_quantize_mlx(w, bits, group_size)
            q_error_sq = mx.square(w - q_w)
            if fisher is not None:
                sens = mx.sum(fisher * q_error_sq)
            else:
                sens = mx.sum(q_error_sq)
            mx.eval(sens)
            sensitivities[bits] = float(sens.item())

        param_count = 1
        for s in w.shape:
            param_count *= s

        results.append(SensitivityResult(
            layer_name=name,
            sensitivities=sensitivities,
            param_count=param_count,
        ))

    t_elapsed = time.time() - t_start
    n_brute = len(quantizable) * len(candidate_bits) * n_cal
    print(f"  [fast-mlx] Done in {t_elapsed:.1f}s "
          f"({n_cal} fwd+bwd vs {n_brute} fwd passes for brute-force)")

    return results


def _accumulate_mlx_grads(
    grads,
    grad_sq_accum: dict,
    quantizable: list[tuple[str, object]],
):
    """Walk MLX gradient tree and accumulate squared gradients for quantizable layers.

    The gradient tree from ``nn.value_and_grad`` mirrors the model's parameter
    tree (nested dicts/lists).  Layer names like ``"layers.0.self_attn.q_proj"``
    correspond to nested keys ``["layers"][0]["self_attn"]["q_proj"]["weight"]``.
    """
    import mlx.core as mx

    for name, _module in quantizable:
        parts = name.split(".")
        node = grads
        try:
            for part in parts:
                if isinstance(node, (list, tuple)):
                    node = node[int(part)]
                elif isinstance(node, dict):
                    # Try string key first, then integer
                    if part in node:
                        node = node[part]
                    elif part.isdigit() and int(part) in node:
                        node = node[int(part)]
                    else:
                        node = None
                        break
                else:
                    node = None
                    break
        except (KeyError, IndexError, TypeError):
            node = None

        if node is None:
            continue

        # node should now be a dict with 'weight' key
        if isinstance(node, dict) and "weight" in node:
            g = node["weight"]
        elif isinstance(node, mx.array):
            g = node
        else:
            continue

        g_sq = mx.square(g)
        if name in grad_sq_accum:
            grad_sq_accum[name] = grad_sq_accum[name] + g_sq
        else:
            grad_sq_accum[name] = g_sq
        mx.eval(grad_sq_accum[name])
