"""Akiki MiniLM-L6-v2 Clover intent classifier.

A frozen sentence-transformers/all-MiniLM-L6-v2 encoder, an identity-initialized
384 x 384 Akiki input affine on its word embeddings, and a six-class head.
Classes, in order: A=chat, B=question, C=build, D=task, E=fix, F=unclear.

    from model import CloverIntent
    clf = CloverIntent.from_pretrained(".")
    clf.predict("Restart my web server.")
    clf.predict("yes", context=[{"role": "assistant", "text": "Shall I restart nginx?"}])

Command line:

    python model.py 'Restart my web server.'
    python model.py --context-file context.json --output json 'yes'
"""

import argparse
import json
from pathlib import Path

import torch
from safetensors.torch import load_file
from torch import nn
from torch.nn import functional as F

LABELS = ("chat", "question", "build", "task", "fix", "unclear")
MAX_LEN = 512
HERE = Path(__file__).resolve().parent


def render(message, context=()):
    """The exact text seen in training: earlier turns, then the current message."""
    prior = "\n".join(f"{turn['role']}: {turn['text']}" for turn in context)
    return ("Conversation context:\n" + prior + "\n\n" if prior else "") + "Current user message: " + message


class IdentityAffine(nn.Module):
    """Akiki input affine, kept in FP32 regardless of autocast."""

    def __init__(self, width):
        super().__init__()
        self.proj = nn.Linear(width, width)
        nn.init.eye_(self.proj.weight)
        nn.init.zeros_(self.proj.bias)

    def forward(self, x):
        with torch.autocast(device_type=x.device.type, enabled=False):
            value = F.linear(x.float(), self.proj.weight.float(), self.proj.bias.float())
        return value.to(x.dtype)


class FrozenEncoderClassifier(nn.Module):
    def __init__(self, encoder):
        super().__init__()
        width = encoder.config.hidden_size
        self.pretrained = encoder.requires_grad_(False).eval()
        self.interface = IdentityAffine(width)
        self.head = nn.Linear(width, len(LABELS))
        # The affine transforms word embeddings before position and type embeddings are added.
        encoder.get_input_embeddings().register_forward_hook(lambda module, args, output: self.interface(output))

    def features(self, input_ids, attention_mask):
        hidden = self.pretrained(input_ids=input_ids, attention_mask=attention_mask, return_dict=True).last_hidden_state
        mask = attention_mask.unsqueeze(-1).to(hidden.dtype)
        return F.normalize(((hidden * mask).sum(1) / mask.sum(1).clamp_min(1)).float(), dim=-1)

    def forward(self, input_ids, attention_mask):
        return self.head(self.features(input_ids, attention_mask))


class CloverIntent:
    def __init__(self, tokenizer, model, device):
        self.tokenizer, self.model, self.device = tokenizer, model, device

    @classmethod
    def from_pretrained(cls, path=HERE, adapter="adapter.safetensors", device="cpu"):
        from transformers import AutoModel, AutoTokenizer
        path = Path(path)
        tokenizer = AutoTokenizer.from_pretrained(path / "base_model", local_files_only=True)
        encoder = AutoModel.from_pretrained(path / "base_model", local_files_only=True,
                                            dtype=torch.float32, attn_implementation="sdpa")
        model = FrozenEncoderClassifier(encoder)
        missing, unexpected = model.load_state_dict(load_file(path / adapter), strict=False)
        if unexpected or any(not k.startswith("pretrained.") for k in missing):
            raise ValueError(f"Adapter does not match the model: missing={missing}, unexpected={unexpected}")
        return cls(tokenizer, model.to(device).eval(), torch.device(device))

    def tokenize(self, texts):
        batch = self.tokenizer(list(texts), padding=True, truncation=False, return_tensors="pt")
        if batch["input_ids"].shape[1] > MAX_LEN:
            raise ValueError(f"Input exceeds {MAX_LEN} tokens; truncation is disabled")
        return {k: batch[k].to(self.device) for k in ("input_ids", "attention_mask")}

    @torch.inference_mode()
    def probabilities(self, rows):
        """rows: dicts with 'message' and optional 'context'. Returns an (n, 6) tensor."""
        texts = [render(r["message"], r.get("context") or ()) for r in rows]
        return self.model(**self.tokenize(texts)).float().softmax(-1).cpu()

    def predict(self, message, context=()):
        p = self.probabilities([{"message": message, "context": context}])[0]
        index = int(p.argmax())
        return {"label": LABELS[index], "token": chr(65 + index),
                "probabilities": dict(zip(LABELS, p.tolist()))}


def main():
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("message")
    parser.add_argument("--context-file", type=Path, help="JSON array of earlier {role, text} turns")
    parser.add_argument("--model", type=Path, default=HERE)
    parser.add_argument("--device", default="cpu")
    parser.add_argument("--output", choices=["token", "label", "json"], default="token")
    args = parser.parse_args()
    context = json.loads(args.context_file.read_text()) if args.context_file else []
    result = CloverIntent.from_pretrained(args.model, device=args.device).predict(args.message, context)
    print(json.dumps(result, indent=2) if args.output == "json" else result[args.output])


if __name__ == "__main__":
    main()