"""gemmaExtract: typed field extraction in one forward pass (no text generation) on a Gemma 4 backbone."""
import math
import torch
import torch.nn as nn

TYPES = {"text": 0, "list": 1, "number": 2, "date": 3, "time": 4, "duration": 5, "enum": 6}
OPS = ["none", "sum", "diff", "avg", "count", "div", "mul"]  # none = chep span; con lai tinh tren toan hang (kieu TagOp)
ORDERED = {"diff", "div"}


def num(t, en=False):
    """Chuoi -> float hoac None. en=True: ',' hang nghin, '.' thap phan (DROP, TAT-QA). Mac dinh (hoa don VN/ID, kieu hon hop):
    2 loai dau -> dau cuoi la thap phan; 1 loai dau -> thap phan tru khi theo sau dung 3 chu so; '(12)' -> -12.
    Giong data/common.py parse_num."""
    import re
    m = re.search(r"-?\d[\d.,]*\d|-?\d", t or "")
    if not m:
        return None
    x = m.group()
    neg = x.startswith("-") or ("(" in t and ")" in t)
    x = x.lstrip("-")
    if en:
        x = x.replace(",", "")
        if x.count(".") > 1:
            return None
    else:
        seps = [c for c in x if c in ".,"]
        if seps:
            last = max(x.rfind("."), x.rfind(","))
            dec = len(set(seps)) == 2 or (len(seps) == 1 and len(x) - last - 1 != 3)
            x = re.sub(r"[.,]", "", x[:last]) + "." + x[last + 1:] if dec else re.sub(r"[.,]", "", x)
    try:
        v = float(x)
    except ValueError:
        return None
    return -v if neg else v


MONTHS = {m: i + 1 for i, m in enumerate(["jan", "feb", "mar", "apr", "may", "jun", "jul", "aug", "sep", "oct", "nov", "dec"])}


def date(t, dayfirst=True):
    """Chuoi ngay -> 'YYYY-MM-DD' hoac None (ngay truoc thang mac dinh, nhu hoa don VN/MY). Giong data/common.py parse_date."""
    import datetime, re
    m = re.search(r"(\d{4})[-/.](\d{1,2})[-/.](\d{1,2})", t or "")
    vn = re.search(r"(\d{1,2})\s*tháng\s*(\d{1,2})\s*,?\s*năm\s*(\d{4})", t or "", re.I)  # "20 tháng 10 năm 2026"
    if m:
        y, mo, d = map(int, m.groups())
    elif vn:
        d, mo, y = map(int, vn.groups())
    else:
        m = re.search(r"(\d{1,2})[-/.](\d{1,2})[-/.](\d{4}|\d{2})", t or "")
        m2 = re.search(r"(\d{1,2})[-/ .]*([A-Za-z]{3})[A-Za-z]*[-/ .,]*(\d{4}|\d{2})", t or "")
        m3 = re.search(r"([A-Za-z]{3})[A-Za-z]*[-/ .]*(\d{1,2})[-/ .,]+(\d{4}|\d{2})", t or "")
        if m:
            d, mo, y = map(int, m.groups())
            if not dayfirst:
                d, mo = mo, d
        elif m2 and m2.group(2).lower() in MONTHS:
            d, mo, y = int(m2.group(1)), MONTHS[m2.group(2).lower()], int(m2.group(3))
        elif m3 and m3.group(1).lower() in MONTHS:
            mo, d, y = MONTHS[m3.group(1).lower()], int(m3.group(2)), int(m3.group(3))
        else:
            return None
        y += 2000 if y < 100 else 0
    try:
        return datetime.date(y, mo, d).isoformat()
    except ValueError:
        return None


def number_style(text):
    """Tu nhan kieu so cua van ban: True = EN (1,234.5), False = kieu cham hang nghin (1.234,5 VN/ID).
    Dem '1,234' + '5.56' so voi '1.234' + '5,56'; ngang nhau -> False (num() kieu hon hop van doc dung phan lon)."""
    import re
    en = len(re.findall(r"\d,\d{3}(?!\d)", text)) + len(re.findall(r"\d\.\d{1,2}(?!\d)", text))
    vn = len(re.findall(r"\d\.\d{3}(?!\d)", text)) + len(re.findall(r"\d,\d{1,2}(?!\d)", text))
    return en > vn


def snap(text, a, b):
    """Noi span ra tron con so: tokenizer Gemma tach tung chu so nen dau ra hay cat vun ('1951' -> '1' + '51'; dev DROP/TAT-QA:
    152/280 loi so phai tinh la vi vay). Gom ca dau ngoac am '(46)'."""
    dig = lambda i: 0 <= i < len(text) and text[i].isdigit()
    while a > 0 and (dig(a - 1) or (text[a - 1] in ".," and dig(a - 2))):
        a -= 1
    while b < len(text) and (dig(b) or (text[b] in ".," and dig(b + 1))):
        b += 1
    if a > 0 and b < len(text) and text[a - 1] == "(" and text[b] == ")":
        a, b = a - 1, b + 1
    return a, b


def trim(text, a, b):
    """Bo khoang trang hai dau span (offset token Gemma kem dau cach dung truoc)."""
    while a < b and text[a].isspace():
        a += 1
    while b > a and text[b - 1].isspace():
        b -= 1
    return a, b


def merge(spans):
    out = []
    for a, b in sorted(spans):
        if out and a <= out[-1][1]:
            out[-1] = (out[-1][0], max(out[-1][1], b))
        else:
            out.append((a, b))
    return out


def compute(op, vals, reverse=False):
    if op == "count":
        return float(len(vals))
    if not vals:
        return None
    if op == "sum":
        return sum(vals)
    if op == "avg":
        return sum(vals) / len(vals)
    if op == "mul":
        out = 1.0
        for v in vals:
            out *= v
        return out
    if len(vals) < 2:
        return None
    a, b = (vals[1], vals[0]) if reverse else (vals[0], vals[1])
    if op == "diff":
        return a - b
    return a / b if b else None


SPAN_TYPES = {"text", "number", "date", "time", "duration"}  # 1 gia tri = 1 span
MAX_SPAN_TOK = 64


def _ids(tok, text):
    return tok(text, add_special_tokens=False)["input_ids"]


def encode(tok, rec, max_len=1024):
    """-> dict: ids, state (a, b) vi tri token state, off (offset ky tu cua token state), fields [(ten, kieu, (a, b) span instruction)]."""
    schema, fields = _ids(tok, "\n\nSCHEMA FIELDS:\n"), []
    for i, (name, f) in enumerate(rec["fields"].items()):
        schema += _ids(tok, "\nFIELD %d\nID: %s\nTYPE: %s\nINSTRUCTION: " % (i + 1, name, f["type"]))
        a = len(schema)
        schema += _ids(tok, f.get("instructions") or name)
        fields.append((name, f["type"], (a, len(schema))))
        if f.get("options"):
            schema += _ids(tok, "\nOPTIONS: " + " | ".join(map(str, f["options"])))
        schema += _ids(tok, "\nEND FIELD\n")
    prefix = [tok.bos_token_id] + _ids(tok, "STATE:\n")
    t = tok(rec["state"], add_special_tokens=False, return_offsets_mapping=True)
    room = max_len - len(prefix) - len(schema)
    if room < 16:
        raise ValueError("schema qua dai: %d token" % len(schema))
    sid, soff = t["input_ids"][:room], [tuple(o) for o in t["offset_mapping"][:room]]
    base = len(prefix) + len(sid)
    return {"ids": prefix + sid + schema, "state": (len(prefix), base), "off": soff,
            "fields": [(n, ty, (a + base, b + base)) for n, ty, (a, b) in fields]}


def labels(enc, rec):
    """Nhan train cho moi field: null, cap (start, end) token chap nhan duoc (moi lan xuat hien), nhan B/I/O cho list.
    Field co dap an nhung moi span bi cat do max_len -> bo khoi loss (None)."""
    off, out = enc["off"], []
    for name, ty, _ in enc["fields"]:
        ans = rec["answers"].get(name)
        if ans is None:
            out.append(None); continue
        pairs, tags, opnd = [], [0] * len(off), [0] * len(off)
        der = None if ans["null"] else ans.get("derivation")
        if der and der.get("op") in OPS and der["op"] != "none":
            spans_txt = []
            for s, e in sorted(der["operands"]):
                toks = [i for i, (a, b) in enumerate(off) if a < e and b > s]
                for i in toks:
                    opnd[i] = 1
                spans_txt.append(rec["state"][s:e])
            if not any(opnd):
                out.append(None); continue
            vals, order = [num(t, en=True) for t in spans_txt], None  # derivation chi co o DROP/TAT-QA (EN)
            if der["op"] in ORDERED and len(vals) == 2 and None not in vals and isinstance(ans.get("value"), (int, float)):
                for rev in (False, True):
                    v = compute(der["op"], vals, rev)
                    if v is not None and abs(v - ans["value"]) <= 1e-6 + 1e-3 * abs(ans["value"]):
                        order = int(rev); break
            out.append({"null": False, "pairs": [], "tags": tags, "op": OPS.index(der["op"]), "opnd": opnd, "order": order})
            continue
        for s, e in ([] if ans["null"] else ans["spans"]):
            toks = [i for i, (a, b) in enumerate(off) if a < e and b > s]
            if toks:
                pairs.append((toks[0], toks[-1]))
                tags[toks[0]] = 1
                for i in toks[1:]:
                    tags[i] = 2
        if not ans["null"] and not pairs:
            out.append(None); continue
        out.append({"null": ans["null"], "pairs": pairs, "tags": tags, "op": 0 if ty == "number" and not ans["null"] else None,
                    "opnd": None, "order": None})
    return out


class Routing(nn.Module):
    """Query field cross-attention vao memory + FFN (EvidenceRoutingLayer cua Clef)."""

    def __init__(self, w, heads, ff):
        super().__init__()
        self.qn, self.mn, self.fn = nn.LayerNorm(w), nn.LayerNorm(w), nn.LayerNorm(w)
        self.att = nn.MultiheadAttention(w, heads, batch_first=True)
        self.ff = nn.Sequential(nn.Linear(w, ff), nn.GELU(), nn.Linear(ff, w))

    def forward(self, q, mem):
        m = self.mn(mem)
        q = q + self.att(self.qn(q), m, m, need_weights=False)[0]
        return q + self.ff(self.fn(q))


class S1Head(nn.Module):
    def __init__(self, hidden, width=1024, routing=2, layers=4, heads=16, ff=4096):
        super().__init__()
        self.w = width
        self.norm = nn.LayerNorm(hidden)
        self.mem = nn.Linear(hidden, width, bias=False)
        self.qproj = nn.Linear(hidden, width, bias=False)
        self.gproj = nn.Linear(hidden, width, bias=False)
        self.type_emb = nn.Embedding(len(TYPES), width)
        self.route = nn.ModuleList([Routing(width, heads, ff) for _ in range(routing)])
        self.joint = nn.ModuleList([nn.TransformerDecoderLayer(width, heads, ff, activation="gelu", batch_first=True, norm_first=True)
                                    for _ in range(layers)])
        self.fnorm = nn.LayerNorm(width)
        self.null = nn.Linear(width, 1)
        self.ks, self.ke, self.kt = (nn.Linear(hidden, width, bias=False) for _ in range(3))
        self.qs, self.qe = nn.Linear(width, width), nn.Linear(width, width)
        self.qt = nn.Linear(width, 3 * width)
        self.op = nn.Linear(width, len(OPS))
        self.kn, self.qn = nn.Linear(hidden, width, bias=False), nn.Linear(width, width)
        self.order = nn.Linear(width, 2)

    def forward(self, hidden, enc):
        """hidden: (L, H) cua 1 record. -> null (F,), start/end (F, S), tag (F, S, 3)."""
        h = self.norm(hidden.float())
        mem = self.mem(h).unsqueeze(0)
        q = torch.stack([h[a:b].mean(0) for _, _, (a, b) in enc["fields"]])
        ty = torch.tensor([TYPES[t] for _, t, _ in enc["fields"]], device=h.device)
        q = (self.qproj(q) + self.gproj(h[-1]) + self.type_emb(ty)).unsqueeze(0)
        for layer in self.route:
            q = layer(q, mem)
        for layer in self.joint:
            q = layer(q, mem)
        f = self.fnorm(q[0])
        a, b = enc["state"]
        hs = h[a:b]
        sc = 1 / math.sqrt(self.w)
        start = self.qs(f) @ self.ks(hs).T * sc
        end = self.qe(f) @ self.ke(hs).T * sc
        tag = torch.einsum("fcw,sw->fsc", self.qt(f).view(len(f), 3, self.w), self.kt(hs)) * sc
        opnd = self.qn(f) @ self.kn(hs).T * sc
        return self.null(f).squeeze(-1), start, end, tag, self.op(f), opnd, self.order(f)


def loss(out, labs, types):
    null, start, end, tag, op, opnd, order = out
    terms = []
    for i, (lab, ty) in enumerate(zip(labs, types)):
        if lab is None:
            continue
        terms.append(nn.functional.binary_cross_entropy_with_logits(null[i], torch.tensor(float(lab["null"]), device=null.device)))
        if ty == "list":
            terms.append(nn.functional.cross_entropy(tag[i], torch.tensor(lab["tags"], device=tag.device)))
        elif not lab["null"] and lab["pairs"]:  # chap nhan bat ky lan xuat hien nao
            ls, le = start[i].log_softmax(-1), end[i].log_softmax(-1)
            terms.append(-torch.logsumexp(torch.stack([ls[s] + le[e] for s, e in lab["pairs"]]), 0))
        if lab.get("op") is not None:
            terms.append(nn.functional.cross_entropy(op[i:i + 1], torch.tensor([lab["op"]], device=op.device)))
        if lab.get("opnd") is not None:
            terms.append(nn.functional.binary_cross_entropy_with_logits(opnd[i], torch.tensor(lab["opnd"], dtype=torch.float, device=opnd.device)))
        if lab.get("order") is not None:
            terms.append(nn.functional.cross_entropy(order[i:i + 1], torch.tensor([lab["order"]], device=order.device)))
    return torch.stack(terms).sum() / max(1, len(terms)) if terms else None


@torch.no_grad()
def decode(out, enc, state_text, null_threshold=0.5, en=None):
    """-> {ten: {null, p_null, guess, confidence, spans}}. guess tra ca khi null (nhu Clairic)."""
    null, start, end, tag, op, opnd, order = out
    off, res = enc["off"], {}
    if en is None:  # khong truyen locale -> tu nhan theo van ban
        en = number_style(state_text)
    for i, (name, ty, _) in enumerate(enc["fields"]):
        p_null = torch.sigmoid(null[i]).item()
        if ty == "list":
            tg = tag[i].argmax(-1).tolist()
            pr = tag[i].softmax(-1)
            spans, j = [], 0
            while j < len(tg):
                if tg[j] == 1 or (tg[j] == 2 and (j == 0 or tg[j - 1] == 0)):
                    k = j + 1
                    while k < len(tg) and tg[k] == 2:
                        k += 1
                    spans.append((*trim(state_text, off[j][0], off[k - 1][1]), pr[j, tg[j]].item()))
                    j = k
                else:
                    j += 1
            vals = list(dict.fromkeys(state_text[s:e].strip() for s, e, _ in spans if state_text[s:e].strip()))
            conf = min((p for *_, p in spans), default=1 - p_null)
            res[name] = {"null": p_null > null_threshold or not vals, "p_null": p_null, "guess": vals, "confidence": conf,
                         "spans": [[s, e] for s, e, _ in spans]}
            continue
        ps, pe = start[i].softmax(-1), end[i].softmax(-1)
        S = len(off)
        band = torch.triu(torch.ones(S, S, device=ps.device)) - torch.triu(torch.ones(S, S, device=ps.device), diagonal=MAX_SPAN_TOK)
        joint = ps[:, None] * pe[None, :] * band
        best = int(joint.argmax())
        s, e = divmod(best, S)
        cs, ce = trim(state_text, off[s][0], off[e][1])
        r = {"null": p_null > null_threshold, "p_null": p_null, "guess": state_text[cs:ce].strip(),
             "confidence": joint.view(-1)[best].item() * (1 - p_null), "spans": [[cs, ce]]}
        if ty == "number":
            o = OPS[int(op[i].argmax())]
            r["op"] = o
            if o != "none":  # gia tri tinh tu toan hang da danh dau
                tg = (opnd[i] > 0).tolist()
                spans, j = [], 0
                while j < len(tg):
                    if tg[j]:
                        k = j
                        while k + 1 < len(tg) and tg[k + 1]:
                            k += 1
                        spans.append((off[j][0], off[k][1])); j = k + 1
                    else:
                        j += 1
                spans = merge(snap(state_text, a, b) for a, b in spans)
                vals = [v for v in (num(state_text[a:b], en) for a, b in spans) if v is not None]
                rev = bool(order[i].argmax())
                v = compute(o, vals, rev)
                r.update({"value": v, "guess": ("%g" % v) if v is not None else r["guess"], "operands": [[a, b] for a, b in spans],
                          "reverse": rev})
            else:
                cs, ce = snap(state_text, cs, ce)
                r.update({"guess": state_text[cs:ce].strip(), "spans": [[cs, ce]]})
                r["value"] = num(r["guess"], en)
        elif ty == "date":
            r["value"] = date(r["guess"])
        res[name] = r
    return res


# ---- bidirectional full-attention layers ---------------------------------------------------------------------
from transformers.masking_utils import create_masks_for_generate


def enable_noncausal(wrapper, sliding=False):
    cfg = wrapper.config.get_text_config() if hasattr(wrapper.config, "get_text_config") else wrapper.config
    assert getattr(cfg, "_attn_implementation", "sdpa") == "sdpa", "noncausal can sdpa"
    W = getattr(cfg, "sliding_window", None)
    if sliding and not W:
        raise ValueError("noncausal sliding: model khong co sliding_window")

    def pre(module, args, kwargs):
        if args or kwargs.get("past_key_values") is not None or kwargs.get("use_cache"):
            raise ValueError("noncausal: chi goi keyword, khong KV cache")
        ids, emb = kwargs.get("input_ids"), kwargs.get("inputs_embeds")
        B, L = (ids.shape if ids is not None else emb.shape[:2])
        dev = ids.device if ids is not None else emb.device
        pad = kwargs.get("attention_mask")
        if pad is None:
            pad = torch.ones((B, L), dtype=torch.long, device=dev)
        if isinstance(pad, dict):
            return args, kwargs
        like = torch.empty((B, L, 1), dtype=next(module.parameters()).dtype, device=dev)
        pos = kwargs.get("position_ids")
        if pos is None:
            pos = torch.arange(L, device=dev).unsqueeze(0).expand(B, -1)
        m = create_masks_for_generate(config=cfg, inputs_embeds=like, attention_mask=pad, past_key_values=None, position_ids=pos)
        if not isinstance(m, dict):
            m = {"full_attention": m}
        m = dict(m)
        m["full_attention"] = pad[:, None, None, :].bool()
        if sliding:
            i = torch.arange(L, device=dev)
            # | eye: hang cua token pad xa hon W sau token that cuoi se rong -> sdpa ra NaN, 0*NaN lan sang token that
            m["sliding_attention"] = (((i[:, None] - i[None, :]).abs() < W) & pad[:, None, None, :].bool()) | torch.eye(L, dtype=torch.bool, device=dev)
        kwargs["attention_mask"] = m
        return args, kwargs

    wrapper.register_forward_pre_hook(pre, with_kwargs=True)


# ---- nap model + API ------------------------------------------------------------------------------------------
import json as _json, os as _os, re as _re

_VI = _re.compile(r"[ăâđêôơưạảấầẩẫậắằẳẵặẹẻẽếềểễệỉịọỏốồổỗộớờởỡợụủứừửữựỳỵỷỹ]", _re.I)
SUPPORTED = {"text", "list", "number", "date"}


def detect_language(text):
    """'vi' neu co dau tieng Viet; 'en' neu chu cai gan nhu toan ASCII; con lai 'default' (nguong null cao hon:
    do tren XQuAD 7 ngon ngu khac, nguong 0.4 tra null nham 3-8%, 0.8 con 0.3-3%)."""
    if len(_VI.findall(text)) > 3:
        return "vi"
    letters = [c for c in text if c.isalpha()]
    if letters and sum(c.isascii() for c in letters) / len(letters) > 0.98:
        return "en"
    return "default"


class GemmaExtract:
    """Typed field extraction in one forward pass, no text generation.

        from gemma_extract import GemmaExtract
        m = GemmaExtract.from_pretrained("beyoru/gemmaExtract")
        m.extract("Invoice No: INV-0042 ... Total: 1,200.00",
                  {"invoice_no": {"type": "text", "instructions": "Invoice number"},
                   "total": {"type": "number", "instructions": "Total amount due"}})
    """

    def __init__(self, path, device="cuda", dtype=torch.bfloat16):
        from transformers import AutoModelForCausalLM, AutoTokenizer
        self.cfg = _json.load(open(_os.path.join(path, "extract_config.json")))
        self.tok = AutoTokenizer.from_pretrained(path)
        top = AutoModelForCausalLM.from_pretrained(path, dtype=dtype, attn_implementation="sdpa").to(device)
        self.enc = top.model
        enable_noncausal(self.enc, sliding=False)
        self.enc.eval()
        h = self.cfg["head"]
        self.head = S1Head(self.cfg["hidden"], width=h["width"], routing=h["routing"], layers=h["layers"], heads=h["heads"], ff=h["ff"]).to(device)
        from safetensors.torch import load_file
        self.head.load_state_dict(load_file(_os.path.join(path, "head.safetensors")))
        self.head.eval()
        self.device, self.dtype = device, dtype

    @classmethod
    def from_pretrained(cls, repo_or_path, **kw):
        if not _os.path.isdir(repo_or_path):
            from huggingface_hub import snapshot_download
            repo_or_path = snapshot_download(repo_or_path)
        return cls(repo_or_path, **kw)

    def _passes(self, fields):
        names, m = list(fields), self.cfg["max_fields_per_pass"]
        return [{n: fields[n] for n in names[i:i + m]} for i in range(0, len(names), m)]

    @torch.no_grad()
    def extract(self, state, fields, null_threshold=None, number_style=None):
        """state: document text. fields: {name: {"type": "text|list|number|date", "instructions": "...", "null_threshold": float}}.
        null_threshold: default per language (auto-detected). number_style: "en" (1,234.5) / "vi" (1.234,5) / None = auto.
        Returns {"answers": {name: {value, guess, null, p_null, confidence, spans, ...}}, "usage": {...}}."""
        for k, v in fields.items():
            if v.get("type") not in SUPPORTED:
                raise ValueError("field %s: type must be one of %s" % (k, sorted(SUPPORTED)))
        lang = detect_language(state)
        thr_all = self.cfg["null_threshold"].get(lang, self.cfg["null_threshold"]["default"]) if null_threshold is None else null_threshold
        en = {"en": True, "vi": False}.get(number_style)
        answers, usage = {}, {"state_tokens": 0, "passes": 0}
        queue = self._passes({k: {"type": v["type"], "instructions": v.get("instructions") or k} for k, v in fields.items()})
        for part in queue:
            try:
                e = encode(self.tok, {"state": state, "fields": part}, self.cfg["max_length"])
            except ValueError:
                if len(part) > 1:
                    queue += [{k: v} for k, v in part.items()]
                    continue
                raise
            with torch.autocast("cuda" if "cuda" in str(self.device) else "cpu", dtype=self.dtype):
                h = self.enc(input_ids=torch.tensor([e["ids"]], device=self.device)).last_hidden_state[0]
                out = self.head(h, e)
            for k, r in decode(out, e, state, 0.5, en=en).items():
                t = fields[k].get("null_threshold", thr_all)
                r["null"] = bool(r["p_null"] > t) or (isinstance(r["guess"], list) and not r["guess"])
                r["spans"] = [{"start": a, "end": b, "text": state[a:b]} for a, b in r.get("spans", [])]
                r["value"] = None if r["null"] else r.get("value", r["guess"])
                answers[k] = r
            usage["state_tokens"] = max(usage["state_tokens"], e["state"][1] - e["state"][0])
            usage["passes"] += 1
        usage["truncated"] = usage["state_tokens"] < len(self.tok(state, add_special_tokens=False)["input_ids"])
        return {"answers": answers, "usage": usage, "language": lang}