"""Shared building blocks for the dataset converters.

Everything that turns a raw labeled example into varied `Question`s lives here:
text cleanup, class-balanced row ordering, option-key styles, distractor sampling,
ordered scales for `score`, and derived true/false statements for `noul`.

Determinism: every random decision goes through a `random.Random` that the builder seeds
from (global seed, source name). Never iterate over a `set` of strings (hash order is
randomized per process); use lists / dicts.
"""
from __future__ import annotations

import html
import math
import random
import re
from dataclasses import dataclass, field
from typing import Callable, Iterable, Iterator, Optional, Sequence

from jevlike.types import Option, Question

MAX_STATE_CHARS = 2000


# ---------------------------------------------------------------- text utilities

_WS = re.compile(r"[ \t\r\f\v]+")
_MANY_NL = re.compile(r"\n{3,}")


def clean(text) -> str:
    """Normalize whitespace and fix common scraping artifacts (html entities, '#39;')."""
    if text is None:
        return ""
    t = str(text)
    t = t.replace("\\n", "\n").replace("\\$", "$")
    t = re.sub(r"(?<![&\w])#(\d{2,4});", lambda m: "&#" + m.group(1) + ";", t)  # ag_news ' #39;'
    t = re.sub(r"(?<![&\w])(quot|amp|lt|gt);", r"&\1;", t)
    t = html.unescape(t)
    t = t.replace(" ", " ")
    t = _WS.sub(" ", t)
    t = "\n".join(line.strip() for line in t.split("\n"))
    t = _MANY_NL.sub("\n\n", t)
    return t.strip()


def truncate(text: str, max_chars: int = MAX_STATE_CHARS) -> str:
    if len(text) <= max_chars:
        return text
    cut = text[:max_chars]
    sp = cut.rfind(" ")
    if sp > max_chars * 0.8:
        cut = cut[:sp]
    return cut.rstrip() + " ..."


def norm_key(text: str) -> str:
    """Normalization used for dedupe / leak detection: lowercase alphanumerics only."""
    return " ".join(re.sub(r"[^0-9a-z]+", " ", text.lower()).split())


def snake(name: str) -> str:
    s = re.sub(r"[^0-9a-zA-Z]+", "_", name.strip().lower()).strip("_")
    return s or "x"


def readable(label: str) -> str:
    """'card_payment_not_recognised' -> 'card payment not recognised'."""
    s = re.sub(r"[_]+", " ", label).strip(" ?")
    return re.sub(r"\s+", " ", s).lower()


def cap_first(s: str) -> str:
    return s[:1].upper() + s[1:] if s else s


_LOWERABLE = {
    "the", "a", "an", "there", "some", "this", "that", "these", "those", "it", "he", "she",
    "they", "we", "you", "people", "no", "all", "many", "two", "three", "four", "several",
    "one", "nobody", "someone", "somebody", "everyone", "most", "his", "her", "their", "our",
    "my", "your", "its", "at", "in", "on", "not", "none", "only", "both", "each", "every",
    "if", "when", "after", "before", "because", "although", "since", "while", "more", "less",
    "few", "any", "men", "women", "kids", "children", "man", "woman", "girl", "boy", "dogs",
    "dog", "cat", "cats", "nothing", "something", "everything", "few", "is", "was", "are",
}


def lower_first(s: str) -> str:
    """Lowercase the first word if it is a common function word (keeps proper nouns)."""
    if not s:
        return s
    first = s.split(" ", 1)[0]
    if first.lower() in _LOWERABLE and first[:1].isupper() and not first.isupper():
        return first.lower() + s[len(first):]
    return s


def as_sentence(s: str) -> str:
    s = s.strip()
    if s and s[-1] not in ".!?\"'":
        s += "."
    return s


def strip_end(s: str) -> str:
    return s.strip().rstrip(".!? ")


def as_question(q: str) -> str:
    q = cap_first(q.strip())
    return q if q.endswith("?") else q.rstrip(".") + "?"


# ---------------------------------------------------------------- sampling helpers

def weighted(rng: random.Random, items: Sequence, weights: Sequence[float]):
    return rng.choices(list(items), weights=list(weights), k=1)[0]


def surface(rng: random.Random, instr: str, p: float = 0.12) -> str:
    """Light surface noise on instructions: lowercase first char / drop final punctuation."""
    if rng.random() < p and instr[:1].isupper() and not instr[:2].isupper():
        instr = instr[0].lower() + instr[1:]
    if rng.random() < p and instr.endswith(("?", ".")):
        instr = instr[:-1]
    return instr


def pick_instr(rng: random.Random, specific: Sequence[str], generic: Sequence[str] = (),
               p_generic: float = 0.2, **fmt) -> str:
    pool = generic if (generic and (not specific or rng.random() < p_generic)) else specific
    t = rng.choice(list(pool))
    if fmt:
        t = t.format(**fmt)
    return surface(rng, t)


def balanced_order(rng: random.Random, rows: list, label_fn: Callable, alpha: float = 0.5) -> list:
    """Order rows so that class c appears at a rate ~ n_c**alpha (alpha=1: natural
    proportions, alpha=0: uniform over classes). Taking any prefix gives a mildly
    class-balanced sample."""
    groups: dict = {}
    for r in rows:
        groups.setdefault(label_fn(r), []).append(r)
    keyed = []
    for lab, grp in groups.items():
        rng.shuffle(grp)
        w = len(grp) ** alpha
        for i, r in enumerate(grp):
            keyed.append(((i + rng.random()) / w, r))
    keyed.sort(key=lambda x: x[0])
    return [r for _, r in keyed]


def dedupe_rows(rows: list, text_fn: Callable, label_fn: Callable) -> list:
    """Drop exact duplicate texts; drop all copies of texts with conflicting labels."""
    labels: dict[str, list] = {}
    first: dict[str, object] = {}
    for r in rows:
        k = norm_key(text_fn(r))
        if not k:
            continue
        lab = label_fn(r)
        labels.setdefault(k, [])
        if lab not in labels[k]:
            labels[k].append(lab)
        first.setdefault(k, r)
    return [r for k, r in first.items() if len(labels[k]) == 1]


# ---------------------------------------------------------------- choice

@dataclass
class Label:
    name: str                              # human-readable class name, e.g. "sports"
    desc: list[str] = field(default_factory=list)  # descriptions (any one may be shown)
    phrase: Optional[str] = None           # noun phrase for derived statements ("sports")
    key: Optional[str] = None              # preferred readable key; default: name
    group: Optional[str] = None            # coarse group; derived-noul negatives prefer other groups

    def __post_init__(self):
        if isinstance(self.desc, str):
            self.desc = [self.desc] if self.desc else []

    @property
    def phr(self) -> str:
        return self.phrase or self.name


def intent_label(raw: str, group: Optional[str] = None, extra_desc: Sequence[str] = ()) -> Label:
    r = readable(raw)
    return Label(r, list(extra_desc) or [r], phrase=r, key=snake(raw), group=group)


OTHER = Label("other", ["none of the above", "something else / not listed",
                        "none of the other options apply", "anything not covered above"],
              key="other")

_OPAQUE_PREFIX = ["opt_", "option_", "class_", "label_", "c", "cat_", "choice_", "id_"]


def _key_style(rng: random.Random, n: int, opaque_weight: float = 1.0) -> str:
    styles = ["name", "snake", "title", "letter", "number", "prefixed", "code"]
    w = [0.33, 0.27, 0.06, 0.13 * opaque_weight, 0.08 * opaque_weight,
         0.09 * opaque_weight, 0.04 * opaque_weight]
    s = weighted(rng, styles, w)
    if s == "letter" and n > 26:
        s = "prefixed"
    return s


def _codes(rng: random.Random, n: int) -> list[str]:
    out: list[str] = []
    while len(out) < n:
        c = "".join(rng.choice("ABCDEFGHJKLMNPQRSTUVWXYZ") for _ in range(rng.choice([2, 3, 3, 4])))
        if c not in out:
            out.append(c)
    return out


def render_options(rng: random.Random, labels: list[Label], *, style: Optional[str] = None,
                   desc_mode: Optional[str] = None, opaque_weight: float = 1.0) -> list[Option]:
    """Turn labels (already in display order) into Options with a random key style.

    Readable key styles sometimes drop descriptions; opaque styles (letters, numbers,
    codes) always carry the meaning in the description."""
    n = len(labels)
    style = style or _key_style(rng, n, opaque_weight)
    opaque = style in ("letter", "number", "prefixed", "code")
    if style == "prefixed":
        pre, base = rng.choice(_OPAQUE_PREFIX), rng.choice([0, 1, 1])
        keys = [f"{pre}{i + base}" for i in range(n)]
    elif style == "letter":
        keys = [chr(65 + i) for i in range(n)]
        if rng.random() < 0.2:
            keys = [k.lower() for k in keys]
    elif style == "number":
        base = rng.choice([0, 1, 1, 1])
        keys = [str(i + base) for i in range(n)]
    elif style == "code":
        keys = _codes(rng, n)
    else:
        keys = []
        for lab in labels:
            k = lab.key or lab.name
            if style == "snake":
                k = snake(k)
            elif style == "title":
                k = readable(k).title() if "_" in k else k.title()
            elif style == "name":
                k = lab.name
            keys.append(k)
    # uniqueness (snake of different names can collide)
    seen: dict[str, int] = {}
    for i, k in enumerate(keys):
        if k in seen:
            seen[k] += 1
            keys[i] = f"{k}_{seen[k]}"
        else:
            seen[k] = 1

    if desc_mode is None:
        if opaque:
            desc_mode = weighted(rng, ["name", "desc", "both"], [0.4, 0.35, 0.25])
        else:
            desc_mode = weighted(rng, ["none", "desc"], [0.35, 0.65])
            if n > 60:  # keep very long option lists short
                desc_mode = "none"
    texts = []
    for lab in labels:
        d = rng.choice(lab.desc) if lab.desc else ""
        if desc_mode == "none":
            t = ""
        elif desc_mode == "name":
            t = lab.name
        elif desc_mode == "both":
            t = f"{lab.name}: {d}" if d and d.lower() != lab.name.lower() else lab.name
        else:  # desc
            t = d if d else (lab.name if opaque else "")
            if not opaque and t.lower() == str(keys[len(texts)]).lower():
                t = ""
        texts.append(t)
    if opaque:  # opaque keys must carry meaning in the text
        texts = [t or lab.name for t, lab in zip(texts, labels)]
    return [Option(str(k), t) for k, t in zip(keys, texts)]


def sample_n_options(rng: random.Random, k: int, full_prob: float = 0.08,
                     max_sampled: int = 40) -> int:
    """Number of options to show for a k-class dataset (spread over 2..~40, some full sets)."""
    if k <= 2:
        return k
    if k <= 6:
        return k if rng.random() < 0.65 else rng.randint(2, k)
    if rng.random() < full_prob:
        return k
    hi = min(k, max_sampled)
    n = int(round(math.exp(rng.uniform(math.log(2), math.log(hi + 0.49)))))
    return max(2, min(hi, n))


def make_choice(rng: random.Random, labels: list[Label], gold: int, instructions: str, *,
                n: Optional[int] = None, full_prob: float = 0.08, max_sampled: int = 40,
                other_prob: float = 0.0, shuffle: bool = True, style: Optional[str] = None,
                desc_mode: Optional[str] = None, opaque_weight: float = 1.0,
                allowed: Optional[list[int]] = None) -> tuple[Question, int]:
    """Choice question over a subset of `labels` that always contains the answer.

    With probability `other_prob` the gold class is removed and an explicit
    "other / none of the above" option becomes the answer."""
    k = len(labels)
    pool = list(range(k)) if allowed is None else list(allowed)
    use_other = k >= 4 and rng.random() < other_prob
    if n is None:
        n = sample_n_options(rng, len(pool), full_prob, max_sampled)
    n = max(2, min(n, len(pool)))
    distract = [i for i in pool if i != gold]
    if use_other:
        # "other" is only correct if no near-synonym of the gold class is shown
        g = labels[gold].group
        distract = [i for i in distract if g is None or labels[i].group != g] or distract
        chosen = rng.sample(distract, min(len(distract), max(1, n - 1)))
    else:
        chosen = [gold] + rng.sample(distract, n - 1)
    if shuffle:
        rng.shuffle(chosen)
    else:
        chosen.sort()
    labs = [labels[i] for i in chosen]
    if use_other:
        pos = len(labs) if rng.random() < 0.8 else rng.randint(0, len(labs))
        labs.insert(pos, OTHER)
        ans = pos
    else:
        ans = chosen.index(gold)
    opts = render_options(rng, labs, style=style, desc_mode=desc_mode, opaque_weight=opaque_weight)
    return Question("choice", instructions, opts), ans


def make_mcq(rng: random.Random, answers: list[str], gold: int, instructions: str,
             shuffle: bool = True) -> tuple[Question, int]:
    """Multiple-choice QA: options are free-text answers; keys are mostly letters/numbers.
    Short answers are sometimes used directly as keys (with empty text)."""
    order = list(range(len(answers)))
    if shuffle:
        rng.shuffle(order)
    ans = order.index(gold)
    texts = [answers[i] for i in order]
    short = all(len(t) <= 30 for t in texts) and len({t.lower() for t in texts}) == len(texts)
    style = weighted(rng, ["letter", "number", "prefixed", "answer"], [0.5, 0.2, 0.1, 0.2 if short else 0.0])
    if style == "answer":
        opts = [Option(t, "") for t in texts]
    else:
        if style == "letter":
            keys = [chr(65 + i) for i in range(len(texts))]
        elif style == "number":
            keys = [str(i + 1) for i in range(len(texts))]
        else:
            pre = rng.choice(_OPAQUE_PREFIX)
            keys = [f"{pre}{i + 1}" for i in range(len(texts))]
        opts = [Option(k, t) for k, t in zip(keys, texts)]
    return Question("choice", instructions, opts), ans


# ---------------------------------------------------------------- score

@dataclass
class Scale:
    """An ordered scale (lowest -> highest) and how a raw value maps onto it."""
    levels: list[str]
    bin: Callable[[float], Optional[int]]   # raw value -> level index, None = skip
    instr: list[str]                        # instructions phrased for this scale
    weight: float = 1.0


def cuts(*thresholds: float) -> Callable[[float], int]:
    """Bin by ascending thresholds: v < t0 -> 0, t0 <= v < t1 -> 1, ..."""
    def f(v: float) -> int:
        i = 0
        for t in thresholds:
            if v >= t:
                i += 1
        return i
    return f


def mapping(d: dict) -> Callable[[float], Optional[int]]:
    return lambda v: d.get(v)


def make_score(rng: random.Random, scales: list[Scale], value, generic: Sequence[str] = (),
               p_generic: float = 0.1, **fmt) -> Optional[tuple[Question, int]]:
    cands = [s for s in scales if s.bin(value) is not None]
    if not cands:
        return None
    s = rng.choices(cands, weights=[c.weight for c in cands], k=1)[0]
    lab = s.bin(value)
    instr = pick_instr(rng, s.instr, generic, p_generic, **fmt)
    opts = [Option(str(i), t) for i, t in enumerate(s.levels)]
    return Question("score", instr, opts), lab


def reverse(scale: Scale, instr: list[str], n_levels: Optional[int] = None, weight: float = 1.0) -> Scale:
    """Same scale in the opposite polarity (e.g. 'How negative ...?' very positive -> very negative)."""
    n = n_levels or len(scale.levels)

    def f(v, _b=scale.bin, _n=n):
        i = _b(v)
        return None if i is None else _n - 1 - i
    return Scale(list(reversed(scale.levels)), f, instr, weight)


GENERIC_SCORE = [
    "Rate the text on the scale below.", "Where does this fall on the scale?",
    "Pick the level that fits best.", "Choose the most appropriate rating.",
    "Place the input on this scale.", "Score this.",
]


# ---------------------------------------------------------------- noul

def noul(statement: str, label: bool) -> tuple[Question, int]:
    return Question("noul", statement, []), int(bool(label))


def derived_noul(rng: random.Random, labels: list[Label], gold, pos: Sequence[str],
                 neg: Sequence[str] = (), p_true: float = 0.5, p_neg: float = 0.2,
                 p_other_group: float = 0.8, allowed: Optional[list[int]] = None,
                 use_desc: float = 0.25, **fmt) -> tuple[Question, int]:
    """True/false statement about a class ("This article is about sports"), ~50/50
    true/false, sometimes negated ("This is NOT about sports", label flipped).

    `gold` may be an int or a list of ints (multi-label: every listed class is true)."""
    golds = list(gold) if isinstance(gold, (list, tuple)) else [gold]
    pool = list(range(len(labels))) if allowed is None else list(allowed)
    if rng.random() < p_true:
        target, truth = rng.choice(golds), True
    else:
        others = [i for i in pool if i not in golds]
        ggroups = [labels[g].group for g in golds]
        far = [i for i in others if labels[i].group is None or labels[i].group not in ggroups]
        if far and rng.random() < p_other_group:
            others = far
        target, truth = rng.choice(others), False
    lab = labels[target]
    x = lab.phr
    if lab.desc and rng.random() < use_desc:
        x = rng.choice(lab.desc)
    if neg and rng.random() < p_neg:
        tmpl, truth = rng.choice(list(neg)), not truth
    else:
        tmpl = rng.choice(list(pos))
    return noul(surface(rng, tmpl.format(x=x, **fmt), p=0.06), truth)


# ---------------------------------------------------------------- examples

@dataclass
class Example:
    """One source example (= one state) with its 1-3 questions.

    `key` groups examples for splitting / leak detection (defaults to the normalized
    state); pair datasets set it from the raw pair so different renderings collide."""
    state: str
    qs: list[tuple[Question, int]]
    key: Optional[str] = None


def pair_state(rng: random.Random, a: str, b: str, kind: str = "sentence") -> str:
    fmts = {
        "sentence": [("Sentence 1: {a}\nSentence 2: {b}", 3), ("Sentence A: {a}\nSentence B: {b}", 2),
                     ("1. {a}\n2. {b}", 1), ("Text 1: {a}\nText 2: {b}", 1), ("\"{a}\"\n\"{b}\"", 1)],
        "question": [("Question 1: {a}\nQuestion 2: {b}", 3), ("Q1: {a}\nQ2: {b}", 2),
                     ("First question: {a}\nSecond question: {b}", 1), ("1. {a}\n2. {b}", 1)],
    }[kind]
    f = weighted(rng, [x[0] for x in fmts], [x[1] for x in fmts])
    return f.format(a=a, b=b)


def take(it: Iterable, n: int) -> Iterator:
    for i, x in enumerate(it):
        if i >= n:
            return
        yield x