"""
jev_toy/data.py

Turn real public HuggingFace datasets into System One training/eval examples,
each of the form (state_text, question_text, q_type, target).

We map three proven Jev question types onto real datasets:

  * choice  <- AG News (route a headline/topic into 1 of 4 categories)
  * noul    <- BoolQ   (yes / no on a reading-comprehension question)
  * score   <- SST-2   (continuous "sentiment polarity" tag)

This is a toy-scale demonstration: we subsample so a laptop CPU can train in
minutes. The structure generalizes to the full datasets and to GPU.
"""

from __future__ import annotations

import re
from collections import Counter

SPLIT_RE = re.compile(r"[A-Za-z0-9']+|[.,!?;:()\"]")


def tokenize(text: str, vocab, max_len: int, oov: int, pad: int):
    ids = []
    for tok in SPLIT_RE.findall(text.lower())[:max_len]:
        ids.append(vocab.get(tok, oov))  # certain words map to a small rand-pool
    ids = ids[:max_len]
    mask = [1] * len(ids)
    ids = ids + [pad] * (max_len - len(ids))
    mask = mask + [0] * (max_len - len(mask))
    return ids, mask


class Vocabulary:
    def __init__(self, min_freq=1):
        self.min_freq = min_freq
        self.stoi = {}
        self.itos = {}
        self.oov = 0

    def build(self, texts: list[str]):
        c = Counter()
        for t in texts:
            c.update(SPLIT_RE.findall(t.lower()))
        words = [w for w, n in c.items() if n >= self.min_freq]
        words.sort()
        self.stoi = {w: i + 2 for i, w in enumerate(words)}  # 0=pad, 1=oov
        self.itos = {i: w for w, i in self.stoi.items()}
        self.oov = 1
        return self

    def __len__(self):
        return len(self.stoi) + 2


def make_agnews(df) -> list[dict]:
    labels = ["World", "Sports", "Business", "Sci/Tech"]
    out = []
    for txt, lab in zip(df["text"], df["label"]):
        out.append({
            "state": f"News article: {txt[:200]}",
            "question": "Which topic does this article belong to?",
            "q_type": "choice",
            "target": int(lab),
            "options": labels,
        })
    return out


def make_boolq(df) -> list[dict]:
    out = []
    for passage, q, ans in zip(df["passage"], df["question"], df["answer"]):
        out.append({
            "state": f"Passage: {passage[:250]}",
            "question": f"Q: {q}",
            "q_type": "noul",
            "target": int(ans),
        })
    return out


def make_sst2(df) -> list[dict]:
    out = []
    for txt, lab in zip(df["sentence"], df["label"]):
        out.append({
            "state": f"Review: {txt[:150]}",
            "question": "Is this review positive?",
            "q_type": "noul",      # treat sentiment polarity as a boolean tag
            "target": int(lab),
        })
    return out


def build_examples(datasets: dict[str, object], subsample: dict[str, int]) -> tuple[list[dict], Vocabulary]:
    """datasets: {name: HF-dataset}; subsample: {name: cap}."""
    all_ex = []
    # AG News
    for split in ("train", "test"):
        df = datasets["agnews"][split]
        cap = subsample.get("agnews", 2000) if split == "train" else 400
        all_ex += make_agnews(df.select(list(range(cap))))
    # BoolQ
    train = datasets["boolq"]["train"].select(list(range(subsample.get("boolq", 2000))))
    val = datasets["boolq"]["validation"].select(list(range(400)))
    all_ex += make_boolq(train) + make_boolq(val)
    # SST-2
    sst_train = datasets["sst2"]["train"].select(list(range(subsample.get("sst2", 2000))))
    sst_val = datasets["sst2"]["validation"].select(list(range(400)))
    all_ex += make_sst2(sst_train) + make_sst2(sst_val)
    # vocab from state+question text
    vocab = Vocabulary().build([e["state"] + " " + e["question"] for e in all_ex])
    return all_ex, vocab


if __name__ == "__main__":
    print("data module ok")