"""The one model wrapper every eval goes through (spec §7, T0.3).

lm-evaluation-harness's HF model, registered as `marco`, with three additions:

- a quantisation switch: `quant=off` (the model as loaded) or `quant=ternary` (routed experts replaced
  in place by their ternary fake-quant, spec §4: Q-A absmean, g128 by default) — the PTQ view of a
  checkpoint without exporting it;
- a loader hook for packed students, `packed=<dir>` (filled in T2.6: codes + scales → expert weights);
- stop tokens: generation stops on both `<|im_end|>` and `<|endoftext|>`, because Marco's
  generation_config.json lists only the latter while every chat turn ends with the former (D-12).

Use: `lm_eval --model marco --model_args pretrained=...,quant=ternary --include_path configs/eval/lm_eval_tasks`.
"""
from __future__ import annotations

import torch
from lm_eval.api.registry import register_model
from lm_eval.models.huggingface import HFLM

from lowbit.quant.ternary import QuantSpec, fake_quant

STOP_TOKENS = ("<|im_end|>", "<|endoftext|>")


def routed_expert_params(model: torch.nn.Module):
    """(name, parameter) for every routed-expert weight: the fused [E, out, in] tensors of transformers 5
    (`experts.gate_up_proj`, `experts.down_proj`) or per-expert Linear weights in older layouts."""
    for name, p in model.named_parameters():
        if ".mlp.experts." in name and name.endswith(("gate_up_proj", "down_proj", "gate_proj.weight",
                                                       "up_proj.weight", "down_proj.weight")):
            yield name, p


@torch.no_grad()
def ternarise_experts(model: torch.nn.Module, spec: QuantSpec) -> int:
    """Replace routed-expert weights by s * q in place (groups along the input dim); returns #weights."""
    n = 0
    for _, p in routed_expert_params(model):
        p.copy_(fake_quant(p.data.float(), spec).to(p.dtype))
        n += p.numel()
    return n


def load_packed(model: torch.nn.Module, path: str) -> None:
    raise NotImplementedError("packed-student loading is T2.6")


@register_model("marco")
class MarcoLM(HFLM):
    def __init__(self, pretrained: str, quant: str = "off", packed: str | None = None,
                 quant_method: str = "absmean", group_size: int = 128, **kwargs):
        super().__init__(pretrained=pretrained, **kwargs)
        if quant not in ("off", "ternary"):
            raise ValueError(f"quant must be off or ternary, got {quant!r}")
        self.quant, self.packed = quant, packed
        if packed:
            load_packed(self.model, packed)
        if quant == "ternary":
            n = ternarise_experts(self.model, QuantSpec(method=quant_method, group_size=int(group_size)))
            if n == 0:
                raise ValueError("quant=ternary but the model has no routed-expert weights")
        stop_ids = [self.tokenizer.convert_tokens_to_ids(t) for t in STOP_TOKENS]
        self.stop_token_ids = [i for i in stop_ids if isinstance(i, int) and i >= 0]

    def _model_generate(self, context, max_length, stop, **generation_kwargs):
        generation_kwargs.setdefault("eos_token_id", self.stop_token_ids)
        stop = list(dict.fromkeys([*stop, *STOP_TOKENS]))
        return super()._model_generate(context, max_length, stop, **generation_kwargs)