"""Native-tokenization input for a diagnostic experiment: is the Qwen-token input the bottleneck?

The Qwen answer bytes decode losslessly to text; the encoder re-tokenizes that text with its OWN
tokenizer (the segmentation it was pretrained on). Labels move from Qwen tokens to native tokens by
character overlap: a native token gets the label of the first positive Qwen token it overlaps
(B only on the first native token of that Qwen token, I after), 0 if it only overlaps O tokens,
-100 otherwise. Predictions go back to Qwen tokens as the mean over the native tokens that overlap
each Qwen token, so every metric is computed on exactly the same Qwen tokens as the other taggers.
Interface mirrors byt5_adapter.QwenBytes (prompt / answer / special ids).
"""
from __future__ import annotations

import numpy as np
import torch
from transformers import AutoTokenizer

from byt5_adapter import QwenBytes


class NativeText:
    def __init__(self, qwen_tokenizer_dir: str, base: str):
        self.qb = QwenBytes(qwen_tokenizer_dir)
        self.tok = AutoTokenizer.from_pretrained(base)
        self.cls = [self.tok.cls_token_id]
        self.sep = [self.tok.sep_token_id]
        self.pad_id = self.tok.pad_token_id

    def prompt(self, query: str) -> list[int]:
        return self.cls + self.tok(query, add_special_tokens=False)["input_ids"][:128] + self.sep

    def answer(self, token_ids, token_labels=None):
        chunks = [self.qb.tb[t] for t in token_ids]
        full = b"".join(chunks)
        text = full.decode("utf-8", errors="replace")
        # byte offset -> char index
        char_of_byte, ci = [], 0
        for ch in text:
            n = len(ch.encode("utf-8")) if ch != "�" else 1
            char_of_byte += [ci] * n
            ci += 1
        char_of_byte.append(ci)
        qspan, b = [], 0
        for c in chunks:
            qspan.append((char_of_byte[b], char_of_byte[b + len(c)] if len(c) else char_of_byte[b]))
            b += len(c)
        enc = self.tok(text, add_special_tokens=False, return_offsets_mapping=True)
        ids, offs = enc["input_ids"], enc["offset_mapping"]
        # native token -> overlapping qwen tokens (both sorted by position: two pointers)
        per_native, k0 = [], 0
        for s, e in offs:
            while k0 < len(qspan) and qspan[k0][1] <= s:
                k0 += 1
            ks, k = [], k0
            while k < len(qspan) and qspan[k][0] < max(e, s + 1):
                if qspan[k][1] > qspan[k][0]:
                    ks.append(k)
                k += 1
            per_native.append(ks)
        spans = [[] for _ in token_ids]
        for j, ks in enumerate(per_native):
            for k in ks:
                spans[k].append(j)
        labs = []
        if token_labels is not None:
            seen_bad = set()
            for j, ks in enumerate(per_native):
                pos = [k for k in ks if token_labels[k] > 0]
                if pos:
                    k = pos[0]
                    l = token_labels[k]
                    if l % 2 == 1 and k in seen_bad:
                        l += 1
                    if l % 2 == 1:
                        seen_bad.add(k)
                    labs.append(l)
                elif ks and all(token_labels[k] == 0 for k in ks):
                    labs.append(0)
                else:
                    labs.append(-100)
        return ids, labs, spans


def windows(n: int, a0: int, max_len: int, stride: int):
    span = max_len - a0 - 1
    starts, s = [], 0
    while True:
        starts.append(s)
        if s + span >= n:
            break
        s += span - stride
    return [(s, min(n, s + span)) for s in starts]


def expand_windows_native(nt: NativeText, rows, max_len, stride, synth_weight, labels_fn, include_synth=True,
                          include_corrected=True):
    out = {"input_ids": [], "labels": [], "weight": []}
    for r in rows:
        v = r["variant"]
        if (v.startswith("synthetic") and not include_synth) or (v == "corrected" and not include_corrected):
            continue
        a0 = r["answer_start"]
        lab = labels_fn(r)
        ids, nl, _ = nt.answer(r["input_ids"][a0:], lab[a0:])
        if not ids:
            continue
        p = nt.prompt(r["query"])
        w = synth_weight if v.startswith("synthetic") else 1.0
        for s, e in windows(len(ids), len(p), max_len, stride):
            out["input_ids"].append(p + ids[s:e] + nt.sep)
            out["labels"].append([-100] * len(p) + nl[s:e] + [-100])
            out["weight"].append(w)
    return out


@torch.no_grad()
def predict_row_native(nt: NativeText, model, row, max_len, stride, device, batch=8):
    a0 = row["answer_start"]
    ids, _, spans = nt.answer(row["input_ids"][a0:])
    C = model.config.num_labels
    out = np.zeros((len(spans), C), dtype=np.float32)
    out[:, 0] = 1.0
    if not ids:
        return out
    p0 = nt.prompt(row["query"])
    ws = windows(len(ids), len(p0), max_len, stride)
    probs = np.zeros((len(ids), C), dtype=np.float32)
    best = np.full(len(ids), -1.0)
    for b in range(0, len(ws), batch):
        chunk = ws[b:b + batch]
        seqs = [p0 + ids[s:e] + nt.sep for s, e in chunk]
        L = max(len(x) for x in seqs)
        x = torch.full((len(seqs), L), nt.pad_id, dtype=torch.long)
        m = torch.zeros((len(seqs), L), dtype=torch.long)
        for i, sq in enumerate(seqs):
            x[i, :len(sq)] = torch.as_tensor(sq)
            m[i, :len(sq)] = 1
        logits = model(input_ids=x.to(device), attention_mask=m.to(device)).logits.float()
        p = torch.softmax(logits, -1).cpu().numpy()
        for i, (s, e) in enumerate(chunk):
            pos = np.arange(s, e)
            centr = np.minimum(pos - s, e - 1 - pos).astype(float)
            upd = centr > best[pos]
            probs[pos[upd]] = p[i, len(p0) + (pos[upd] - s)]
            best[pos[upd]] = centr[upd]
    for k, js in enumerate(spans):
        if js:
            out[k] = probs[js].mean(0)
    return out