"""Hugging Face model wrapper for the character n-gram OCR-quality scorer.

Wraps the stupid-backoff character n-gram of train_char_ngram.py in the
`transformers` causal-LM interface, so the model can live on the Model Hub
with its code (trust_remote_code) and other repositories need no copied code:

    tok = AutoTokenizer.from_pretrained("histde/...")
    model = AutoModelForCausalLM.from_pretrained("histde/...", trust_remote_code=True)
    enc = tok(page, return_tensors="pt", add_special_tokens=False)
    bpc = model(**enc, labels=enc["input_ids"]).loss.item() / math.log(2)

This is the same scoring contract as the DTA xLSTM checkpoints, so the two
scorers are interchangeable behind AutoModelForCausalLM. The n-gram tables
(sorted uint64 keys + counts per order, bit-cast to int64 because safetensors
has no uint64 - safe, keys use at most 48 bits) are registered as buffers and
stored in model.safetensors. All actual computation is vectorized numpy on
CPU; the model has no trainable parameters.

forward() has two paths:
  * labels given (scoring): per-position stupid-backoff scores of the target
    characters only - fast, used for corpus scoring; `logits` is None.
  * labels omitted (generation): full unnormalized log-score matrix over the
    vocabulary as `logits`, which makes generate() work (softmax of the
    stupid-backoff scores). 256x the lookups - fine for short generations.

Scores follow the standard causal-LM shift: logits[i] / the loss term at
position i+1 describe token i+1 given tokens 0..i; the first token is not
scored. Stupid backoff is not a normalized distribution - bits-per-character
values are scores for ranking, not true perplexities.

AI Disclosure:
    Models:         Claude Fable 5 (claude-fable-5)
    AI-Generated:   fully         # fully | mostly | partially | none
    Human-Reviewed: none          # fully | partially | minimally | none
"""

import math

import numpy as np
import torch
from transformers import PreTrainedModel
from transformers.generation import GenerationMixin
from transformers.modeling_outputs import CausalLMOutput

try:
    from .configuration_char_ngram import CharNgramConfig  # loaded as a Hub dynamic module
except ImportError:
    from configuration_char_ngram import CharNgramConfig   # imported from the repository root

LN2 = math.log(2)


def ngram_keys(ids, order):
    """uint64 keys of all `order`-grams; key[i] encodes ids[i - order + 1 .. i]."""
    keys = ids.astype(np.uint64)
    for j in range(1, order):
        keys[j:] |= ids[:-j].astype(np.uint64) << np.uint64(8 * j)
    return keys[order - 1:]


class CharNgramForCausalLM(PreTrainedModel, GenerationMixin):
    config_class = CharNgramConfig
    base_model_prefix = "char_ngram"
    main_input_name = "input_ids"

    def __init__(self, config):
        super().__init__(config)
        for k, size in enumerate(config.ngram_sizes, start=1):
            self.register_buffer(f"keys_{k}", torch.zeros(size, dtype=torch.int64))
            self.register_buffer(f"counts_{k}", torch.zeros(size, dtype=torch.int64))
        # transformers needs at least one floating-point parameter to infer a dtype
        self.dtype_anchor = torch.nn.Parameter(torch.zeros(1), requires_grad=False)
        self._tables = None  # lazy numpy views, built after from_pretrained fills the buffers
        self.post_init()

    def _init_weights(self, module):  # no trainable parameters
        pass

    def tables(self):
        if self._tables is None:
            # int64 -> uint64 bit-cast; keys < 2**48, so int64 and uint64 sort orders agree
            self._tables = {
                k: (getattr(self, f"keys_{k}").cpu().numpy().view(np.uint64),
                    getattr(self, f"counts_{k}").cpu().numpy().astype(np.float64))
                for k in range(1, self.config.order + 1)
            }
        return self._tables

    @staticmethod
    def _lookup(sorted_keys, counts, query):
        pos = np.searchsorted(sorted_keys, query)
        pos[pos == len(sorted_keys)] = 0
        return np.where(sorted_keys[pos] == query, counts[pos], 0.0)

    def _target_log2(self, ids):
        """log2 stupid-backoff score of ids[i] given ids[:i], for i = 1..n-1."""
        order, alpha = self.config.order, self.config.backoff_alpha
        tables = self.tables()
        n = len(ids)
        rel = np.zeros((order + 1, n))
        rel[1] = self._lookup(*tables[1], ids.astype(np.uint64)) / self.config.total_tokens
        for k in range(2, order + 1):
            if n < k:
                break
            gram = self._lookup(*tables[k], ngram_keys(ids, k))
            ctx = self._lookup(*tables[k - 1], ngram_keys(ids, k - 1)[:-1])
            rel[k, k - 1:] = np.where(ctx > 0, gram / np.maximum(ctx, 1.0), 0.0)
        highest = np.minimum(np.arange(1, n + 1), order)
        best, backoffs = np.zeros(n), np.zeros(n)
        for k in range(order, 0, -1):
            take = (best == 0) & (rel[k] > 0) & (highest >= k)
            best[take] = rel[k, take]
            backoffs[take] = highest[take] - k
        return (np.log2(best) + backoffs * math.log2(alpha))[1:]

    def _full_log2(self, ids):
        """(n, vocab) log2 scores; row i is the next-token score after ids[:i + 1]."""
        order, alpha, vocab = self.config.order, self.config.backoff_alpha, self.config.vocab_size
        tables = self.tables()
        n = len(ids)
        candidates = np.arange(vocab, dtype=np.uint64)
        rel = np.zeros((order + 1, n, vocab))
        rel[1] = self._lookup(*tables[1], candidates.copy()) / self.config.total_tokens
        for k in range(2, order + 1):
            if n < k - 1:
                break
            ctx_keys = ngram_keys(ids, k - 1)                       # contexts ending at k-2 .. n-1
            ctx = self._lookup(*tables[k - 1], ctx_keys)
            query = (ctx_keys[:, None] << np.uint64(8)) | candidates[None, :]
            gram = self._lookup(*tables[k], query.ravel()).reshape(len(ctx_keys), vocab)
            rel[k, k - 2:] = np.where(ctx[:, None] > 0, gram / np.maximum(ctx[:, None], 1.0), 0.0)
        highest = np.minimum(np.arange(1, n + 1), order - 1) + 1    # context i chars long -> order i+1
        best, backoffs = np.zeros((n, vocab)), np.zeros((n, vocab))
        for k in range(order, 0, -1):
            take = (best == 0) & (rel[k] > 0) & (highest[:, None] >= k)
            best[take] = rel[k][take]
            backoffs[take] = np.broadcast_to(highest[:, None] - k, best.shape)[take]
        with np.errstate(divide="ignore"):
            scores = np.log2(best) + backoffs * math.log2(alpha)
        return np.where(best > 0, scores, -60.0)  # floor for candidates unseen even as unigrams

    def forward(self, input_ids, attention_mask=None, labels=None, **kwargs):
        batch = input_ids.cpu().numpy().astype(np.uint8)
        if attention_mask is None:
            lengths = [batch.shape[1]] * batch.shape[0]
        else:
            lengths = attention_mask.cpu().numpy().sum(axis=1).astype(int).tolist()

        if labels is not None:  # scoring path: targets only, no logits
            nlls = []
            for row, n in zip(batch, lengths):
                if n >= 2:
                    nlls.append(-self._target_log2(row[:n]) * LN2)
            loss = torch.tensor(np.concatenate(nlls).mean() if nlls else 0.0,
                                dtype=torch.float32, device=input_ids.device)
            return CausalLMOutput(loss=loss, logits=None)

        logits = np.full((batch.shape[0], batch.shape[1], self.config.vocab_size), -60.0)
        for i, (row, n) in enumerate(zip(batch, lengths)):
            logits[i, :n] = self._full_log2(row[:n]) * LN2
        logits = torch.from_numpy(logits).to(dtype=torch.float32, device=input_ids.device)
        return CausalLMOutput(loss=None, logits=logits)

    def prepare_inputs_for_generation(self, input_ids, **kwargs):
        return {"input_ids": input_ids}