"""Self-contained edit reranker that ships INSIDE the Hugging Face model repo alongside
`modeling_gectagger.py` and is loaded via `trust_remote_code=True`. It must NOT import the
`spellchecker` package (end users / the Space won't have it) and must NOT depend on `errant`/`spaCy`
(the Space ships only transformers + safetensors). Everything the per-edit reranker needs — the
EditScorer architecture, the candidate-input construction, a dependency-free word-level edit
extractor, and the over-generate -> score -> filter -> re-apply logic — is inlined here.

This is the inference twin of `spellchecker/reranker.py` (training/eval code). The scorer is a
shared-backbone binary classifier (EditScorer-style, Sorokin 2022): the tagger over-generates edits at
a negative `keep_confidence`, each proposed edit is scored P(correct), and only edits with P >= tau are
kept and re-applied to the source. The shipped scorer is the `cross` variant with a per-ERRANT-type
feature; its weights + operating point live in the SAME repo's `reranker/` subfolder, lazily loaded on
first `.correct(..., rerank=True)`.

FIDELITY NOTE — edit extraction. The scorer was trained and the published MASTER (0.6642) measured with
ERRANT word-level alignment (`tokenise=False`) of source->hypothesis. ERRANT needs spaCy, which is not
available at serve time, so this module extracts edits with a dependency-free `difflib` word diff and
recovers ERRANT-style minimal edits by splitting aligned (replace) blocks token-by-token where the
block is length-preserving — the common case ERRANT also splits. The candidate char-span construction
and scorer input are otherwise byte-identical to the training path, so a candidate's score matches.
`type_feature` types are coarse (M/U/R + a couple of cheap subtypes) rather than full ERRANT types; the
type only feeds a hash bucket, so coarse types still calibrate per broad error class. See the local
verification in the publish flow for the measured tagger-only vs reranked behaviour on the exported repo.
"""
from __future__ import annotations

import difflib
import json
import os
import zlib
from collections import defaultdict
from typing import Dict, List, Sequence, Tuple

import torch
import torch.nn as nn
from transformers import AutoConfig, AutoModelForMaskedLM

# --------------------------------------------------------------------------- encoder trunk

def _backbone(name: str, pretrained: bool = False):
    """(backbone_module, hidden_size). The shared bidirectional-LFM2 trunk (same as the tagger's).

    Prefers the in-library ``transformers.Lfm2BidirectionalModel`` once the bidirectional-LFM2 family PR
    lands — version-stable, no remote code for the trunk, numerically identical to the encoder repo's
    remote-code MLM base (verified 0.0 CPU / <1e-5 GPU), so the scorer's state_dict still loads 1:1 under
    ``encoder.*``. Until then it falls back to the encoder repo's remote-code MLM class with the head
    stripped (the original behaviour). For inference inside the published repo we build from config
    (pretrained=False) and load the scorer's own weights — no second encoder download."""
    import transformers
    native = getattr(transformers, "Lfm2BidirectionalModel", None)
    if native is not None:                                    # native foundation (post family-PR)
        config = AutoConfig.from_pretrained(name)             # model_type resolves natively, no remote code
        hidden = getattr(config, "hidden_size", None) or getattr(config, "d_model", None) or 1024
        if pretrained:
            backbone = native.from_pretrained(name)
        else:
            try:
                backbone = native._from_config(config)
            except AttributeError:
                backbone = native(config)
        return backbone, hidden
    config = AutoConfig.from_pretrained(name, trust_remote_code=True)   # fallback: encoder remote code
    hidden = getattr(config, "hidden_size", None) or getattr(config, "d_model", None) or 1024
    if pretrained:
        mlm = AutoModelForMaskedLM.from_pretrained(name, trust_remote_code=True)
    else:
        mlm = AutoModelForMaskedLM.from_config(config, trust_remote_code=True)
    return mlm.base_model, hidden


def _last_hidden(backbone, input_ids, attention_mask) -> torch.Tensor:
    out = backbone(input_ids=input_ids, attention_mask=attention_mask)
    hs = getattr(out, "last_hidden_state", None)
    if hs is None and getattr(out, "hidden_states", None) is not None:
        hs = out.hidden_states[-1]
    if hs is None:
        hs = out[0]
    return hs


# --------------------------------------------------------------- candidate input construction
# (byte-identical to spellchecker.reranker so a candidate's encoder input matches the trained scorer)

def apply_edits(tokens: List[str], edits: List[Dict]) -> List[str]:
    """Apply token-span edits to `tokens` (same semantics as spellchecker.edits.apply_edits)."""
    out: List[str] = []
    prev = 0
    for e in sorted(edits, key=lambda e: (int(e["start"]), int(e["end"]))):
        start, end = int(e["start"]), int(e["end"])
        if start < prev:                       # overlapping: drop
            continue
        out.extend(tokens[prev:start])
        repl = str(e["replacement"])
        if repl:
            out.extend(repl.split())
        prev = end
    out.extend(tokens[prev:])
    return out


def build_scorer_input(src_tokens: Sequence[str], edit: Dict) -> Tuple[str, int, int]:
    a, b = int(edit["start"]), int(edit["end"])
    rep = str(edit["replacement"]).split()
    out = list(src_tokens[:a]) + rep + list(src_tokens[b:])
    text = " ".join(out)
    if rep:
        c0 = len(" ".join(out[:a])) + (1 if a > 0 else 0)
        c1 = c0 + len(" ".join(rep))
    else:                                          # deletion: pool over flanking tokens
        left, right = max(a - 1, 0), min(a + 1, len(out))
        c0 = len(" ".join(out[:left])) + (1 if left > 0 else 0)
        c1 = c0 + len(" ".join(out[left:right]))
        if c1 <= c0:
            c0, c1 = 0, max(len(text), 1)
    return text, c0, c1


def build_source_input(src_tokens: Sequence[str], edit: Dict) -> Tuple[str, int, int]:
    a, b = int(edit["start"]), int(edit["end"])
    text = " ".join(src_tokens)
    if a < b:
        c0 = len(" ".join(src_tokens[:a])) + (1 if a > 0 else 0)
        c1 = c0 + len(" ".join(src_tokens[a:b]))
    else:                                          # insertion: pool over the gap's neighbours
        left, right = max(a - 1, 0), min(a + 1, len(src_tokens))
        c0 = len(" ".join(src_tokens[:left])) + (1 if left > 0 else 0)
        c1 = c0 + len(" ".join(src_tokens[left:right]))
        if c1 <= c0:
            c0, c1 = 0, max(len(text), 1)
    return text, c0, c1


def sep_token(tokenizer) -> str:
    return tokenizer.sep_token or tokenizer.eos_token or "||"


def build_cross_input(src_tokens: Sequence[str], edit: Dict, sep: str
                      ) -> Tuple[str, int, int, int, int]:
    s_text, s0, s1 = build_source_input(src_tokens, edit)
    e_text, e0, e1 = build_scorer_input(src_tokens, edit)
    text = f"{s_text} {sep} {e_text}"
    off = len(s_text) + len(sep) + 2
    return text, off + e0, off + e1, s0, s1


def type_bucket(edit_type: str, buckets: int = 64) -> int:
    return zlib.crc32(str(edit_type).encode()) % buckets


def encode_candidates(tokenizer, items: List[Tuple[str, int, int]], max_len: int = 160):
    texts = [t for t, _, _ in items]
    enc = tokenizer(texts, return_tensors="pt", padding=True, truncation=True, max_length=max_len,
                    return_offsets_mapping=True)
    offsets = enc.pop("offset_mapping")
    span = torch.zeros_like(enc["attention_mask"])
    for i, (_, c0, c1) in enumerate(items):
        for j, (o0, o1) in enumerate(offsets[i].tolist()):
            if enc["attention_mask"][i, j] and o1 > o0 and o0 < c1 and o1 > c0:
                span[i, j] = 1
        if not span[i].any():
            span[i] = enc["attention_mask"][i]
    return enc["input_ids"], enc["attention_mask"], span


def encode_cross(tokenizer, items: List[Tuple[str, int, int, int, int]], max_len: int = 320):
    texts = [t for t, *_ in items]
    enc = tokenizer(texts, return_tensors="pt", padding=True, truncation=True, max_length=max_len,
                    return_offsets_mapping=True)
    offsets = enc.pop("offset_mapping")
    spans = [torch.zeros_like(enc["attention_mask"]) for _ in range(2)]
    for i, (_, e0, e1, s0, s1) in enumerate(items):
        for j, (o0, o1) in enumerate(offsets[i].tolist()):
            if not enc["attention_mask"][i, j] or o1 <= o0:
                continue
            if o0 < e1 and o1 > e0:
                spans[0][i, j] = 1
            if o0 < s1 and o1 > s0:
                spans[1][i, j] = 1
        for s in spans:
            if not s[i].any():
                s[i] = enc["attention_mask"][i]
    return enc["input_ids"], enc["attention_mask"], spans[0], spans[1]


# --------------------------------------------------------------- dependency-free edit extraction

_VOWELS = set("aeiouAEIOU")


def _coarse_type(src_span: List[str], rep: List[str]) -> str:
    """Coarse ERRANT-ish type for the per-type feature hash. Full ERRANT types need spaCy; the type
    only feeds a 64-way hash bucket, so a broad class (M/U/R + a couple cheap subtypes) is enough to
    let the head calibrate per error family. Mirrors the M/U/R coarse scheme of edits.align_generic."""
    if not src_span:
        return "M"                                   # missing -> insertion
    if not rep:
        return "U"                                   # unnecessary -> deletion
    if len(src_span) == 1 and len(rep) == 1:
        a, b = src_span[0], rep[0]
        la, lb = a.lower(), b.lower()
        if la == lb:
            return "R:ORTH"                           # casing only
        if not any(ch.isalnum() for ch in a) and not any(ch.isalnum() for ch in b):
            return "R:PUNCT"
        # cheap spelling cue: same first letter & similar length & shared letter multiset
        if a and b and la[0] == lb[0] and abs(len(a) - len(b)) <= 2:
            return "R:SPELL"
    return "R:OTHER"


def _char_sim(a: str, b: str) -> float:
    """SequenceMatcher ratio on characters (cheap stand-in for ERRANT's lemma/char substitution cost)."""
    if a == b:
        return 1.0
    return difflib.SequenceMatcher(None, a, b, autojunk=False).ratio()


def _align_block(src: List[str], tgt: List[str], i0: int, sub_thresh: float = 0.30) -> List[Dict]:
    """Levenshtein-align a difflib `replace` block at the TOKEN level (this is the core of what ERRANT
    does before its linguistic merge): a token pairs with another as a SUBSTITUTION only when they are
    char-similar enough (>= sub_thresh), otherwise the diff is resolved as a DELETE + INSERT at the
    junction. Emits minimal per-position edits with absolute offsets (block starts at source index i0).
    This recovers ERRANT's split of e.g. `than a` into R:`than` + M:`a`, which a single merged
    candidate would not — the recall the scale fidelity check showed difflib was dropping."""
    m, n = len(src), len(tgt)
    INS, DEL = 1.0, 1.0                                  # unit cost for an unmatched token
    # DP edit distance with substitution cost = 1 - char_sim (cheap matches preferred as subs)
    cost = [[0.0] * (n + 1) for _ in range(m + 1)]
    bt = [[None] * (n + 1) for _ in range(m + 1)]
    for i in range(1, m + 1):
        cost[i][0] = i * DEL; bt[i][0] = "d"
    for j in range(1, n + 1):
        cost[0][j] = j * INS; bt[0][j] = "i"
    for i in range(1, m + 1):
        for j in range(1, n + 1):
            sub = cost[i - 1][j - 1] + (1.0 - _char_sim(src[i - 1], tgt[j - 1]))
            dele = cost[i - 1][j] + DEL
            ins = cost[i][j - 1] + INS
            best = min(sub, dele, ins)
            cost[i][j] = best
            bt[i][j] = "s" if best == sub else ("d" if best == dele else "i")
    # backtrace into aligned ops
    ops = []
    i, j = m, n
    while i > 0 or j > 0:
        step = bt[i][j]
        if step == "s":
            ops.append(("s", i - 1, j - 1)); i -= 1; j -= 1
        elif step == "d":
            ops.append(("d", i - 1, None)); i -= 1
        else:
            ops.append(("i", i, j - 1)); j -= 1                 # insert before source position i
    ops.reverse()
    edits: List[Dict] = []
    for kind, si, tj in ops:
        if kind == "s":
            a, b = src[si], tgt[tj]
            if a == b:
                continue
            # a poor "substitution" (char-dissimilar) is really a delete+insert at this junction
            if _char_sim(a, b) < sub_thresh:
                edits.append({"start": i0 + si, "end": i0 + si + 1, "replacement": "",
                              "type": _coarse_type([a], [])})
                edits.append({"start": i0 + si + 1, "end": i0 + si + 1, "replacement": b,
                              "type": _coarse_type([], [b])})
            else:
                edits.append({"start": i0 + si, "end": i0 + si + 1, "replacement": b,
                              "type": _coarse_type([a], [b])})
        elif kind == "d":
            edits.append({"start": i0 + si, "end": i0 + si + 1, "replacement": "",
                          "type": _coarse_type([src[si]], [])})
        else:  # insert before source index si
            edits.append({"start": i0 + si, "end": i0 + si, "replacement": tgt[tj],
                          "type": _coarse_type([], [tgt[tj]])})
    return edits


_ERRANT_ANNOTATOR = None
_ERRANT_TRIED = False


def _errant_edits(src: str, hyp: str):
    """ERRANT word-level edits (tokenise=False) — the SAME alignment the scorer was trained and the
    published MASTER measured on. Returns None if `errant` (and its spaCy model) is not installed, so
    the caller falls back to the difflib extractor. ERRANT's linguistically-merged minimal edits are
    the granularity the cross+type scorer expects; reproducing them keeps the published metric intact."""
    global _ERRANT_ANNOTATOR, _ERRANT_TRIED
    if _ERRANT_ANNOTATOR is None:
        if _ERRANT_TRIED:
            return None
        _ERRANT_TRIED = True
        try:
            import errant
            _ERRANT_ANNOTATOR = errant.load("en")
        except Exception as e:
            print(f"[gectagger] errant unavailable ({type(e).__name__}); reranker edit extraction "
                  f"falls back to difflib (slightly lower recall than the ERRANT-measured operating "
                  f"point). Install `errant` for the published behaviour.")
            return None
    ann = _ERRANT_ANNOTATOR
    orig, cor = ann.parse(src, False), ann.parse(hyp, False)
    out = []
    for e in ann.annotate(orig, cor):
        out.append({"start": int(e.o_start), "end": int(e.o_end),
                    "replacement": e.c_str, "type": e.type})
    return out


def _difflib_edits(src: str, hyp: str) -> List[Dict]:
    """Dependency-free fallback: difflib word diff, `replace` blocks resolved by a token-level
    Levenshtein alignment (_align_block) into minimal sub/ins/del edits. Close to ERRANT but not exact
    (no linguistic merge / typing); used only when errant is not importable."""
    s, t = src.split(), hyp.split()
    edits: List[Dict] = []
    for op, i1, i2, j1, j2 in difflib.SequenceMatcher(None, s, t, autojunk=False).get_opcodes():
        if op == "equal":
            continue
        if op == "delete":
            for k in range(i1, i2):
                edits.append({"start": k, "end": k + 1, "replacement": "",
                              "type": _coarse_type([s[k]], [])})
        elif op == "insert":
            edits.append({"start": i1, "end": i1, "replacement": " ".join(t[j1:j2]),
                          "type": _coarse_type([], t[j1:j2])})
        else:
            edits.extend(_align_block(s[i1:i2], t[j1:j2], i1))
    return edits


def extract_edits(src: str, hyp: str) -> List[Dict]:
    """Word-level edits taking `src` -> `hyp` as token-span dicts {start,end,replacement,type} into
    src.split(). Uses ERRANT (the alignment the scorer was trained / the MASTER was measured on) when
    available, falling back to a self-contained difflib aligner otherwise."""
    e = _errant_edits(src, hyp)
    if e is not None:
        return e
    return _difflib_edits(src, hyp)


# ------------------------------------------------------------------------------------ the scorer

def _resolve_mode(cfg: Dict) -> str:
    if "mode" in cfg:
        return cfg["mode"]
    return "pairwise" if cfg.get("pairwise") else "edited"


class EditScorer(nn.Module):
    """Shared-backbone binary edit scorer: encoder -> mean-pool over edit span(s) -> 2 logits.
    Inference twin of spellchecker.reranker.EditScorer (same state_dict keys, same forward).
    Modes: edited (v1, one span) / pairwise (v2, two encodings) / cross (v3, one joint sequence,
    two pooled spans). `type_feature` appends a learned embedding of the (hash-bucketed) edit type."""

    TYPE_BUCKETS, TYPE_DIM = 64, 32

    def __init__(self, encoder_name: str, dropout: float = 0.1, mode: str = "edited",
                 type_feature: bool = False, _build_encoder: bool = False):
        super().__init__()
        assert mode in ("edited", "pairwise", "cross"), mode
        self.encoder_name = encoder_name
        self.mode = mode
        self.type_feature = type_feature
        self.encoder, hidden = _backbone(encoder_name, pretrained=_build_encoder)
        self.hidden_size = hidden
        self.dropout = nn.Dropout(dropout)
        feat = hidden * (1 if mode == "edited" else 2)
        if type_feature:
            self.type_emb = nn.Embedding(self.TYPE_BUCKETS, self.TYPE_DIM)
            feat += self.TYPE_DIM
        self.head = nn.Linear(feat, 2)

    def _pool(self, input_ids, attention_mask, span_mask):
        hidden = _last_hidden(self.encoder, input_ids, attention_mask)
        m = span_mask.unsqueeze(-1).to(hidden.dtype)
        return (hidden * m).sum(1) / m.sum(1).clamp(min=1.0)

    def forward(self, input_ids, attention_mask, span_mask,
                src_input_ids=None, src_attention_mask=None, src_span_mask=None, type_ids=None):
        if self.mode == "cross":
            hidden = _last_hidden(self.encoder, input_ids, attention_mask)

            def pool(mask):
                m = mask.unsqueeze(-1).to(hidden.dtype)
                return (hidden * m).sum(1) / m.sum(1).clamp(min=1.0)
            pooled = torch.cat([pool(span_mask), pool(src_span_mask)], dim=-1)
        elif self.mode == "pairwise":
            pooled = torch.cat([self._pool(input_ids, attention_mask, span_mask),
                                self._pool(src_input_ids, src_attention_mask, src_span_mask)], dim=-1)
        else:
            pooled = self._pool(input_ids, attention_mask, span_mask)
        if self.type_feature:
            pooled = torch.cat([pooled, self.type_emb(type_ids)], dim=-1)
        return {"logits": self.head(self.dropout(pooled))}

    @classmethod
    def load(cls, scorer_dir: str, map_location="cpu") -> "EditScorer":
        cfg = json.load(open(os.path.join(scorer_dir, "scorer_config.json")))
        model = cls(encoder_name=cfg["encoder_name"], mode=_resolve_mode(cfg),
                    type_feature=cfg.get("type_feature", False), _build_encoder=False)
        sd = torch.load(os.path.join(scorer_dir, "pytorch_model.bin"), map_location=map_location,
                        weights_only=True)
        model.load_state_dict(sd)
        model.eval()
        return model


# -------------------------------------------------------------------------- reranking entry point

@torch.no_grad()
def _score(scorer: EditScorer, tokenizer, texts: List[str], cands: List[Tuple[int, Dict]],
           device, batch_size: int = 64, max_len: int = 160) -> List[float]:
    """P(edit correct) for each (sentence_index, edit). Per-mode input construction matches
    spellchecker.reranker.RerankedPredictor._score exactly."""
    mode = scorer.mode
    probs: List[float] = []
    for k in range(0, len(cands), batch_size):
        batch = cands[k:k + batch_size]
        if mode == "cross":
            sep = sep_token(tokenizer)
            items = [build_cross_input(texts[i].split(), e, sep) for i, e in batch]
            ids, mask, span, s_span = encode_cross(tokenizer, items, max(max_len, 320))
            kw = {"input_ids": ids.to(device), "attention_mask": mask.to(device),
                  "span_mask": span.to(device), "src_span_mask": s_span.to(device)}
        else:
            items = [build_scorer_input(texts[i].split(), e) for i, e in batch]
            ids, mask, span = encode_candidates(tokenizer, items, max_len)
            kw = {"input_ids": ids.to(device), "attention_mask": mask.to(device),
                  "span_mask": span.to(device)}
            if mode == "pairwise":
                s_items = [build_source_input(texts[i].split(), e) for i, e in batch]
                s_ids, s_mask, s_span = encode_candidates(tokenizer, s_items, max_len)
                kw.update(src_input_ids=s_ids.to(device), src_attention_mask=s_mask.to(device),
                          src_span_mask=s_span.to(device))
        if scorer.type_feature:
            kw["type_ids"] = torch.tensor([type_bucket(e.get("type", "UNK")) for _, e in batch],
                                          dtype=torch.long).to(device)
        probs += scorer(**kw)["logits"].softmax(-1)[:, 1].tolist()
    return probs


@torch.no_grad()
def rerank(scorer: EditScorer, tokenizer, sources: List[str], hyps: List[str], tau: float,
           device, batch_size: int = 64, max_len: int = 160) -> List[str]:
    """Over-generated `hyps` (tagger run at the shipped negative keep_confidence) vs `sources`:
    extract per-edit candidates, score each, keep only P >= tau, re-apply to the source. Pointwise
    acceptance — the shipped behaviour. Returns one corrected string per source."""
    out = list(hyps)
    per_sent: Dict[int, List[Dict]] = {}
    for i, (s, h) in enumerate(zip(sources, hyps)):
        if s == h:
            continue
        edits = extract_edits(s, h)
        out[i] = s                                   # rerankable: rebuild from accepted edits only
        if edits:
            per_sent[i] = edits
    if not per_sent:
        return out
    cands = [(i, e) for i, es in per_sent.items() for e in es]
    probs = _score(scorer, tokenizer, sources, cands, device, batch_size, max_len)
    accepted = defaultdict(list)
    for (i, e), p in zip(cands, probs):
        if p >= tau:
            accepted[i].append(e)
    for i, edits in accepted.items():
        out[i] = " ".join(apply_edits(sources[i].split(), edits))
    return out