src/agiws_neural_quant/dispatch.py

"""dispatch — unified quantize_model router + recursive walker.

quantize_model(model, format=..., skip_types=..., exclude_modules=...,
               dual_path=False, teacher_format=None, chunk_size=1024,
               adaptive=False, **kwargs)

Routes format string → Quantizer preset + walker. Supports dual-path
(teacher QuantizedModule for cross-quantization distillation).
"""

from __future__ import annotations

import fnmatch
from typing import Any

import torch
import torch.nn as nn

from agiws_neural_quant.base import QuantizedModule, SUPPORTED_MODULE_TYPES
from agiws_neural_quant.quantizer import Quantizer
from agiws_neural_quant.presets import get_preset


def make_quantizer(format: str, **kwargs) -> Quantizer:
    """Build a Quantizer from a format string + kwargs."""
    preset_kwargs = get_preset(format, **kwargs)
    return Quantizer(**preset_kwargs)


def _matches_any(name: str, patterns: set[str]) -> bool:
    for p in patterns:
        if fnmatch.fnmatch(name, p):
            return True
    return False


def _walk(
    module: nn.Module,
    quantizer: Quantizer,
    teacher_quantizer: Quantizer | None,
    skip_types: tuple[type, ...],
    exclude_modules: set[str],
    compute_dtype: str,
    chunk_size: int | None,
    adaptive: bool,
    prefix: str,
    replaced: list[str],
    per_head_set: set[str] | None = None,
    num_heads: int = 0,
    head_dim: int = 0,
) -> None:
    for name, child in list(module.named_children()):
        full = f"{prefix}.{name}" if prefix else name
        if _matches_any(full, exclude_modules) or _matches_any(name, exclude_modules):
            continue
        if isinstance(child, skip_types):
            continue
        if isinstance(child, QuantizedModule):
            continue
        if isinstance(child, SUPPORTED_MODULE_TYPES) and hasattr(child, "weight") and child.weight is not None:
            # Per-head override for matching module names.
            use_quantizer = quantizer
            use_teacher = teacher_quantizer
            if per_head_set and (name in per_head_set or _matches_any(name, per_head_set)):
                qcfg = quantizer.to_config()
                qcfg["scale_mode"] = "per-head"
                qcfg["num_heads"] = num_heads
                qcfg["head_dim"] = head_dim
                use_quantizer = Quantizer(**qcfg)
                if teacher_quantizer is not None:
                    tcfg = teacher_quantizer.to_config()
                    tcfg["scale_mode"] = "per-head"
                    tcfg["num_heads"] = num_heads
                    tcfg["head_dim"] = head_dim
                    use_teacher = Quantizer(**tcfg)
            teacher = None
            if use_teacher is not None:
                teacher = QuantizedModule.from_module(
                    child, use_teacher, compute_dtype=compute_dtype,
                    chunk_size=chunk_size, adaptive=adaptive,
                )
            qm = QuantizedModule.from_module(
                child, use_quantizer, compute_dtype=compute_dtype,
                chunk_size=chunk_size, adaptive=adaptive, teacher=teacher,
            )
            setattr(module, name, qm)
            replaced.append(full)
            continue
        _walk(child, quantizer, teacher_quantizer, skip_types, exclude_modules,
              compute_dtype, chunk_size, adaptive, full, replaced,
              per_head_set, num_heads, head_dim)


def quantize_model(
    model: nn.Module,
    format: str = "int8",
    skip_types: tuple[type, ...] | None = None,
    exclude_modules: list[str] | None = None,
    compute_dtype: str = "fp32",
    chunk_size: int | None = 1024,
    adaptive: bool = False,
    dual_path: bool = False,
    teacher_format: str | None = None,
    per_head_modules: list[str] | None = None,
    num_heads: int = 0,
    head_dim: int = 0,
    **kwargs,
) -> nn.Module:
    """Quantize a PyTorch model in-place via the unified QuantizedModule.

    Args:
        model: any nn.Module (walked recursively).
        format: quantization format string (see presets.FORMAT_PRESETS).
        skip_types: module types to leave untouched.
        exclude_modules: glob patterns of module names to skip.
        compute_dtype: 'fp32' | 'fp16' | 'bf16' for dequantized matmul.
        chunk_size: output-dim chunk for dequant (None = no chunking).
        adaptive: if True, AdaptiveChunkSize monitor adjusts chunk_size.
        dual_path: if True, create teacher QuantizedModule alongside student.
        teacher_format: format string for teacher (if dual_path). If None and
                       dual_path=True, uses "fp16" (passthrough).
        per_head_modules: module name patterns to use per-head scale_mode
                         (e.g. ["q_proj", "k_proj", "v_proj", "o_proj"]).
        num_heads: number of attention heads (for per-head scale_mode).
        head_dim: dimension per head (for per-head scale_mode).
        **kwargs: forwarded to Quantizer (override preset values).
    """
    quantizer = make_quantizer(format, **kwargs)
    teacher_quantizer = None
    if dual_path:
        t_fmt = teacher_format or "fp16"
        teacher_quantizer = make_quantizer(t_fmt)
    skip_types_t = tuple(skip_types) if skip_types else ()
    exclude = set(exclude_modules) if exclude_modules else set()
    per_head_set = set(per_head_modules) if per_head_modules else set()
    replaced: list[str] = []
    _walk(model, quantizer, teacher_quantizer, skip_types_t, exclude,
          compute_dtype, chunk_size, adaptive, "", replaced,
          per_head_set, num_heads, head_dim)
    model._quantized_replaced = replaced  # type: ignore[attr-defined]
    model._quantized_format = format  # type: ignore[attr-defined]
    # Log summary.
    qinfo = quantizer.info()
    print(f"[NeuralQuant] quantize_model: format={format} repr={qinfo['repr']} "
          f"bits={qinfo['bits']} scale={qinfo['scale']} group={qinfo['group']} "
          f"w=True a={qinfo['a']} learnable={qinfo['learnable']} "
          f"replaced={len(replaced)} modules", flush=True)
    if per_head_set:
        ph_count = sum(1 for r in replaced if any(p in r for p in per_head_set))
        print(f"[NeuralQuant] per-head: {ph_count} modules (head_dim={head_dim}, "
              f"num_heads={num_heads})", flush=True)
    return model


def count_quantizable_layers(model: nn.Module) -> dict[str, int]:
    counts: dict[str, int] = {}
    for module in model.modules():
        if isinstance(module, SUPPORTED_MODULE_TYPES) and hasattr(module, "weight") and module.weight is not None:
            key = type(module).__name__
            counts[key] = counts.get(key, 0) + 1
    return counts


# ---------------------------------------------------------------------------
# Save / Load — persist a FULL quantized model (v3_hybrid_state format).
#
# save_model(model, path): saves QuantizedModule.to_dict() for quantized
#   layers + native_state (all non-quantized params/buffers) for the rest.
#
# load_model(model, path): fills an ALREADY created + quantized model with
#   weights from the file. Replaces QuantizedModule instances via from_dict
#   (packed buffers are format-specific, load_state_dict cannot handle them).
#   Fills native params/buffers via load_state_dict(strict=False).
#
# Format: "agiws_neural_quant_v3_hybrid_state" (breaking change from v2).
# v2 is rejected with a clear error (no backward compat).
# ---------------------------------------------------------------------------

# Module types registered as quantizable (used to distinguish weight buffers).
_WEIGHT_MODULE_TYPES = (nn.Linear, nn.Conv1d, nn.Conv2d, nn.Conv3d,
                         nn.ConvTranspose1d, nn.ConvTranspose2d, nn.ConvTranspose3d,
                         nn.Embedding, nn.LayerNorm, nn.Bilinear)


def _collect_quantized_paths(model: nn.Module) -> list[str]:
    """Return dotted paths of all QuantizedModule instances in the model."""
    paths: list[str] = []
    for name, mod in model.named_modules():
        if isinstance(mod, QuantizedModule):
            paths.append(name)
    return paths


def _walk_quantized(
    module: nn.Module,
    prefix: str,
    out: dict[str, dict],
) -> None:
    """Collect to_dict() for every QuantizedModule, keyed by dotted path."""
    for name, child in list(module.named_children()):
        full = f"{prefix}.{name}" if prefix else name
        if isinstance(child, QuantizedModule):
            out[full] = child.to_dict()
        else:
            _walk_quantized(child, full, out)


def _collect_native_state(model: nn.Module, quantized_paths: set[str]) -> dict[str, torch.Tensor]:
    """Collect all non-quantized parameters and buffers.

    Excludes any param/buffer that belongs to a QuantizedModule (those are
    serialized separately in quantized_modules). Uses standard PyTorch
    dotted-path keys (same as model.state_dict()).
    """
    native: dict[str, torch.Tensor] = {}
    # Parameters.
    for name, param in model.named_parameters():
        # Skip if this param belongs to a QuantizedModule.
        if _belongs_to_quantized(name, quantized_paths):
            continue
        native[name] = param.detach().cpu().clone()
    # Buffers.
    for name, buf in model.named_buffers():
        if _belongs_to_quantized(name, quantized_paths):
            continue
        if buf is None:
            continue
        native[name] = buf.detach().cpu().clone()
    return native


def _belongs_to_quantized(dotted_name: str, quantized_paths: set[str]) -> bool:
    """Check if a dotted-path param/buffer name belongs to a QuantizedModule.

    A param at "blocks.0.attn.q_proj.weight" belongs to the QuantizedModule at
    "blocks.0.attn.q_proj" if that path is in quantized_paths.
    """
    for qpath in quantized_paths:
        if dotted_name == qpath or dotted_name.startswith(qpath + "."):
            return True
    return False


def save_model(model: nn.Module, path: str) -> None:
    """Save a FULL quantized model to a .pt file (torch.save).

    Saves:
      - quantized_modules: {dotted_path: QuantizedModule.to_dict()} for every
        QuantizedModule in the model (packed weight buffers + meta + config).
      - native_state: {dotted_path: tensor} for all non-quantized parameters
        and buffers (custom layers, embeddings, positional encodings, etc.).
      - quantized_paths: list of dotted paths of all QuantizedModule instances.
      - quant_format: the format string used for quantization.
      - compute_dtype: target compute dtype (from first QuantizedModule).

    Format: "agiws_neural_quant_v3_hybrid_state".

    Args:
        model: a quantized model (after quantize_model).
        path: output .pt file path.
    """
    quantized: dict[str, dict] = {}
    _walk_quantized(model, "", quantized)
    quantized_paths = _collect_quantized_paths(model)
    native_state = _collect_native_state(model, set(quantized_paths))

    # Compute dtype from first QuantizedModule (all should match).
    compute_dtype = "fp32"
    if quantized:
        first_path = next(iter(quantized))
        compute_dtype = quantized[first_path].get("compute_dtype", "fp32")

    quant_format = getattr(model, "_quantized_format", "unknown")

    payload = {
        "format": "agiws_neural_quant_v3_hybrid_state",
        "quant_format": quant_format,
        "compute_dtype": compute_dtype,
        "quantized_paths": quantized_paths,
        "quantized_modules": quantized,
        "native_state": native_state,
        "model_config": getattr(model, "_nq_model_config", None),
    }
    torch.save(payload, path)


def load_model(model: nn.Module, path: str) -> None:
    """Load weights from a save file into an ALREADY created + quantized model.

    The model must be created by the user (via their factory) and quantized
    via quantize_model() BEFORE calling load_model. This function:
      1. Replaces each QuantizedModule in the model with QuantizedModule.from_dict()
         from the file (correct packed weight buffers).
      2. Fills non-quantized parameters/buffers via load_state_dict(native_state,
         strict=False) and reports missing/unexpected keys.

    Why from_dict + setattr instead of load_state_dict into existing QuantizedModule:
    QuantizedModule stores packed weight buffers (int4 codes, FP8 scale codes,
    codebook indices, etc.) — these are NOT raw float weights. load_state_dict
    cannot reconstruct the packed format from a state_dict. from_dict rebuilds
    the QuantizedModule with the exact buffers + meta + quantizer config from
    the file, producing identical dequantized weights.

    Args:
        model: an already-created, already-quantized model. Will be modified
               in-place (QuantizedModule instances replaced, native params filled).
        path: .pt file saved by save_model or convert_model.

    Raises:
        ValueError: if the file is not a v3_hybrid_state file, or if it's an
                   old v2 file (with guidance to re-quantize).
    """
    payload = torch.load(path, map_location="cpu", weights_only=False)
    if not isinstance(payload, dict):
        raise ValueError(f"load_model: not a valid NeuralQuant save file: {path}")
    fmt = payload.get("format")
    if fmt == "agiws_neural_quant_v2":
        raise ValueError(
            f"load_model: file {path} uses old format 'agiws_neural_quant_v2' "
            f"(pre-0.3.0). This format is no longer supported. "
            f"Re-quantize your model with NeuralQuant >= 0.3.0 and save again."
        )
    if fmt != "agiws_neural_quant_v3_hybrid_state":
        raise ValueError(
            f"load_model: unsupported format {fmt!r}. "
            f"Expected 'agiws_neural_quant_v3_hybrid_state'."
        )

    quantized_modules = payload["quantized_modules"]
    native_state = payload["native_state"]

    # 1. Replace QuantizedModule instances with from_dict versions.
    for dotted_path, qm_dict in quantized_modules.items():
        qm = QuantizedModule.from_dict(qm_dict)
        _set_module_by_path(model, dotted_path, qm)

    # 2. Fill native params/buffers.
    if native_state:
        missing, unexpected = model.load_state_dict(native_state, strict=False)
        if missing:
            print(f"[NeuralQuant] load_model: {len(missing)} missing keys "
                  f"(not in saved native_state): {missing[:5]}{'...' if len(missing) > 5 else ''}",
                  flush=True)
        if unexpected:
            print(f"[NeuralQuant] load_model: {len(unexpected)} unexpected keys "
                  f"(in saved file but not in model): {unexpected[:5]}{'...' if len(unexpected) > 5 else ''}",
                  flush=True)


def _set_module_by_path(root: nn.Module, dotted_path: str, new_module: nn.Module) -> None:
    """Set a module at a dotted path within root (e.g. 'blocks.0.attn.q_proj').

    Uses standard PyTorch dotted-path convention: split by '.', traverse
    parent modules, setattr on the parent.
    """
    parts = dotted_path.split(".")
    parent = root
    for part in parts[:-1]:
        # ModuleList indices are accessed via int indexing.
        if part.isdigit():
            parent = parent[int(part)]
        else:
            parent = getattr(parent, part)
    last = parts[-1]
    if last.isdigit():
        parent[int(last)] = new_module  # type: ignore[index]
    else:
        setattr(parent, last, new_module)