"""Task-diverse instruction-following data: hundreds of small classification tasks, each with
its OWN natural-language task definition as the question instructions and its own label set.

    python -m jevlike.data.extra_instruct                      # all collections (seed 13)
    python -m jevlike.data.extra_instruct --only superni --cap-superni 100

Collections (one JSONL each in data/extra/instruct/):
  superni     Super-NaturalInstructions (Muennighoff/natural-instructions, English tasks).
              A task is used when its targets form a small closed label set (2-30 short labels);
              the task definition becomes the instructions, the labels the options (rating tasks
              with 3-10 consecutive integer labels become `score`). Task files are range-read
              (first 256 KB, then up to 2 MB for classification tasks) and cached under
              ~/.cache/jevlike/superni, so only a few hundred MB are downloaded.
  bigbench    BIG-bench multiple-choice tasks (tasksource/bigbench): multiple_choice_targets with
              exactly one correct answer -> choice (a trailing question line becomes the
              instructions when present).
  tasksource  tasksource-instruct validation split (~480 tasks, <=500 rows each): the label list
              quoted in the prompt ('... with either "a", "b" or "c".') becomes the options, the
              task-specific question (if any) the instructions; letter-MCQ tasks get their
              trailing 'A: ...' option lines parsed into options.

Every task: blocked if `Decontaminator.blocked_task(name + definition)` (held-out families) or if it
duplicates a source of the main pipeline; each record dropped if its state is contaminated. Tasks are
capped (default 250 / 150 / 80 records) after class-balanced ordering, so no task dominates.
Record.source = "x_<collection>/<task>". Deterministic given --seed (random.Random(f"{seed}|...")).
Writes data/extra/instruct/{superni,bigbench,tasksource}.jsonl and STATS.md.
"""
from __future__ import annotations

import argparse
import json
import os
import random
import re
import sys
import time
from collections import Counter
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional

from jevlike.data.common import (Label, as_sentence, balanced_order, cap_first, clean, dedupe_rows,
                                 make_mcq, readable, render_options, surface, truncate,
                                 weighted)
from jevlike.data.decontam import Decontaminator
from jevlike.data.hub import fetch
from jevlike.types import Option, Question, Record, write_jsonl

ROOT = Path(__file__).resolve().parents[2]
OUT = ROOT / "data" / "extra" / "instruct"
CACHE = Path(os.environ.get("JEVLIKE_CACHE", Path.home() / ".cache" / "jevlike"))

NI_REPO, NI_REV = "Muennighoff/natural-instructions", "a29a9757125f4bb1c26445ad0d2ef7d9b2cc9c4c"
BB_REPO, BB_REV = "tasksource/bigbench", "210c156767d2f4f05d2f4fd0bb275017a67040fd"
TS_REPO, TS_REV = "tasksource/tasksource-instruct-v0", "4cdba593d93bf1e32c0fd1731d84c6780136fe51"

MAX_DEF_CHARS = 600
MIN_ROWS = 40
MAX_LABEL_WORDS = 12
MAX_MCQ_WORDS = 30


def log(*a):
    print(*a, file=sys.stderr, flush=True)


# ---------------------------------------------------------------- task blocking

# Tasks that duplicate a source of the main pipeline (jevlike.data.build_dataset --list) or a
# held-out family that the decontam patterns miss ("question typing" = trec).
DUP_PATTERNS = [
    r"ag_?news", r"dbpedia", r"yahoo_answers", r"newsgroups", r"tweet_?eval", r"banking77", r"clinc",
    r"massive", r"bitext", r"(?<!\w)emotion$", r"go_?emotions", r"hate_speech_offensive",
    r"commonsense_?qa$", r"commonsenseqa_answer", r"social_?i_?qa", r"socialiqa", r"(?<!\w)race(_|/|$)",
    r"hellaswag", r"(?<!\w)sciq", r"copa", r"yelp", r"amazonreview", r"amazon_reviews_multi", r"(?<!\w)sst",
    r"stsb", r"sts-b", r"(?<!\w)sick", r"civil_comments", r"boolq", r"mnli", r"snli", r"(?<!\w)rte",
    r"(?<!\w)glue", r"super_?glue", r"vitaminc", r"qnli", r"qqp", r"mrpc", r"(?<!\w)paws", r"wic",
    r"multirc", r"imdb", r"subj", r"olid", r"hateeval", r"semeval_2018_task[13]", r"stancedetection",
    r"twitter-financial", r"question[\s_-]*typ", r"trec",
]
_DUP = re.compile("|".join(DUP_PATTERNS), re.I)


def duplicate_task(name: str) -> bool:
    return bool(_DUP.search(name))


# ---------------------------------------------------------------- text helpers

_NON_ASCII_OK = set("‘’“”–—…•·°±×÷€£¥©®™→←↔")


def english_like(s: str) -> bool:
    """Crude English/Latin-script check: at most 10% non-ASCII characters (quotes etc. excepted)."""
    if not s:
        return False
    bad = sum(1 for c in s if ord(c) > 127 and c not in _NON_ASCII_OK)
    return bad / len(s) <= 0.10


def sentences(text: str) -> list[str]:
    return [s for s in re.split(r"(?<=[.!?])\s+(?=[A-Z\"'(])", text) if s.strip()]


def shorten_def(defn: str, max_chars: int = MAX_DEF_CHARS) -> str:
    """Keep head and tail sentences of a long definition (the tail usually says how to answer)."""
    if len(defn) <= max_chars:
        return defn
    ss = sentences(defn)
    if len(ss) < 2:
        return truncate(defn, max_chars)
    head, tail = [ss[0]], [ss[-1]]
    i, j = 1, len(ss) - 2
    while i <= j:
        added = False
        if len(" ".join(head + [ss[i]] + tail)) + 5 <= max_chars:
            head.append(ss[i]); i += 1; added = True
        if i <= j and len(" ".join(head + [ss[j]] + tail)) + 5 <= max_chars:
            tail.insert(0, ss[j]); j -= 1; added = True
        if not added:
            break
    out = " ".join(head) + (" ... " if i <= j else " ") + " ".join(tail)
    return out if len(out) <= max_chars + 50 else truncate(defn, max_chars)


def first_sentences(defn: str, max_chars: int = 250) -> str:
    out = ""
    for s in sentences(defn):
        if out and len(out) + len(s) + 1 > max_chars:
            break
        out = (out + " " + s).strip()
    return truncate(out, max_chars)


def task_words(name: str) -> str:
    """'task1336_peixian_equity_evaluation_corpus_gender_classifier' -> 'peixian equity evaluation ...'."""
    n = re.sub(r"^task\d+_", "", name)
    return readable(n.replace("/", " ").replace("-", " "))


def norm_label(t: str) -> str:
    t = re.sub(r"\s+", " ", str(t).strip())
    if len(t) > 1 and t.endswith(".") and t.count(".") == 1:
        t = t[:-1].rstrip()
    return t


_OPAQUE_LABEL = re.compile(
    r"\(?[a-z]\)?|\d+|-?\d+(\.\d+)?|(option|completion|response|text|fact|speaker|answer|sentence|"
    r"story|choice|candidate|question|statement|passage|ending)\s*\(?\w{1,3}\)?|b_\w+_a|.*\(\w\).*", re.I)


def opaque_label(lab: str) -> bool:
    """Labels whose meaning only the task definition explains (letters, numbers, 'Option 2', ...)."""
    return bool(_OPAQUE_LABEL.fullmatch(lab.strip()))


def as_int(lab: str) -> Optional[int]:
    m = re.fullmatch(r"(-?\d+)(\s*(stars?|points?)|\s*[:=-]\s+\D.*)?", lab.strip(), re.I)
    return int(m.group(1)) if m else None


_RATING = re.compile(r"\b(rat(e|es|ed|ing|ings)|scale|stars?|scor(e|es|ing)|grade|level|degree|intensity|"
                     r"severity|how (much|well|strongly|likely|good|helpful))\b", re.I)


# ---------------------------------------------------------------- task spec -> records

@dataclass
class Row:
    rid: str
    state: str
    gold: int                         # index into options (per-row) or into task labels
    options: Optional[list[str]] = None   # per-row option texts (MCQ); None = task labels
    instr: Optional[str] = None           # per-row main instructions (overrides Task.instr)


@dataclass
class Task:
    collection: str
    name: str
    instr: str                        # main instructions (definition / question)
    kind: str                         # "labels" (closed label set) or "mcq" (per-row options)
    rows: list[Row]
    labels: list[str] = field(default_factory=list)
    score: bool = False               # labels are an ordered scale (already sorted low->high)
    short_instr: list[str] = field(default_factory=list)  # alternative shorter instructions
    generic: list[str] = field(default_factory=list)      # generic fallback instructions
    opaque: bool = False              # labels need the definition to be understood
    p_short: float = 0.15             # P(short_instr), P(generic) when the labels are self-explaining
    p_generic: float = 0.12

    @property
    def source(self) -> str:
        return f"x_{self.collection}/{self.name}"


GENERIC_CLS = [
    "Classify the input.", "Choose the correct label for the input.", "Which label fits best?",
    "Pick the right category for this text.", "Select the appropriate label.", "Label this input.",
    "Assign the most fitting label.", "What is the correct classification?",
]
GENERIC_MCQ = [
    "Choose the correct answer.", "Pick the best option.", "Which option is correct?",
    "Select the right answer.", "Answer by choosing one of the options.", "Choose the option that fits best.",
]
NOUL_TMPL = [
    '{instr}\nThe correct answer is "{x}".', "{instr}\nAnswer: {x}", '{instr}\n\nClaim: the right label here is "{x}".',
    "{instr}\nThe answer is {x}.", 'Task: {instr}\nStatement: the correct output for this input is "{x}".',
    '{instr}\nProposed answer: "{x}". This answer is correct.',
]
NOUL_TMPL_MCQ = [
    '{instr}\nThe correct answer is "{x}".', "{instr}\nAnswer: {x}", "{instr}\nThe right option is: {x}",
    'Question: {instr}\nStatement: "{x}" is the correct answer.',
]
_YESNO = {"yes", "no", "true", "false", "correct", "wrong", "incorrect", "right"}


def pick_instructions(rng: random.Random, t: Task, row: Row) -> str:
    r = rng.random()
    if not t.opaque and t.short_instr and r < t.p_short:
        s = rng.choice(t.short_instr)
    elif not t.opaque and t.generic and r < t.p_short + t.p_generic:
        s = rng.choice(t.generic)
        if rng.random() < 0.5:
            s = f"Task: {task_words(t.name)}. {s}"
    else:
        s = row.instr or t.instr
    return surface(rng, s, p=0.08)


def type_mix(t: Task) -> dict[str, float]:
    if t.kind == "mcq":
        return {"choice": 0.85, "noul": 0.15}
    if t.score:
        return {"score": 0.8, "noul": 0.2}
    if len(t.labels) == 2:
        yn = all(l.lower() in _YESNO for l in t.labels)
        return {"choice": 0.5, "noul": 0.5} if yn else {"choice": 0.65, "noul": 0.35}
    return {"choice": 0.8, "noul": 0.2}


def build_question(rng: random.Random, t: Task, row: Row) -> tuple[Question, int]:
    labels = row.options if row.options is not None else t.labels
    mix = type_mix(t)
    qtype = weighted(rng, list(mix), list(mix.values()))
    instr = pick_instructions(rng, t, row)
    if qtype == "noul":
        truth = rng.random() < 0.5
        x = labels[row.gold] if truth else labels[rng.choice([i for i in range(len(labels)) if i != row.gold])]
        tmpl = rng.choice(NOUL_TMPL_MCQ if t.kind == "mcq" else NOUL_TMPL)
        return Question("noul", tmpl.format(instr=instr.rstrip(), x=x), []), int(truth)
    if qtype == "score":
        return Question("score", instr, [Option(str(i), l) for i, l in enumerate(labels)]), row.gold
    if t.kind == "mcq":
        return make_mcq(rng, labels, row.gold, instr)
    # closed label set: shuffled, label strings as keys or opaque keys with the label as text
    order = list(range(len(labels)))
    if not (t.opaque and all(opaque_label(l) for l in labels)):
        rng.shuffle(order)  # letters / numbers / 'Option 2' referenced by the definition keep their natural order
    labs = [Label(labels[i], [labels[i]], key=labels[i]) for i in order]
    long_labels = max(len(l.split()) for l in labels) > 4
    if t.opaque:
        style = "name"
    else:
        style = weighted(rng, ["name", "snake", "letter", "number", "prefixed", "code"],
                         [0.55, 0.12 if any(" " in l for l in labels) else 0.0, 0.14, 0.06, 0.08, 0.05]
                         if not long_labels else [0.35, 0.0, 0.35, 0.1, 0.15, 0.05])
    if style == "letter" and len(labs) > 26:
        style = "prefixed"
    opts = render_options(rng, labs, style=style, desc_mode="none" if style == "name" else None)
    return Question("choice", instr, opts), order.index(row.gold)


def task_records(t: Task, seed: int, cap: int, dc: Decontaminator, stats: Counter) -> list[Record]:
    rng = random.Random(f"{seed}|{t.source}")
    rows = dedupe_rows(t.rows, lambda r: r.state, lambda r: (r.gold, tuple(r.options or ())))
    rows = balanced_order(rng, rows, lambda r: r.gold if r.options is None else len(r.options), alpha=0.5)
    out: list[Record] = []
    for row in rows:
        if len(out) >= cap:
            break
        if dc.is_contaminated(row.state):
            stats["dropped_contaminated"] += 1
            continue
        q, lab = build_question(rng, t, row)
        rec = Record(f"{t.source}/{row.rid}", t.source, row.state, q, lab)
        rec.validate()
        out.append(rec)
    return out


# ---------------------------------------------------------------- label-set analysis

def closed_label_set(targets: list[str], max_labels: int = 30) -> Optional[tuple[list[str], dict[str, str]]]:
    """If the targets form a small closed set of short labels, return (labels, raw->label map)."""
    n = len(targets)
    normed = [norm_label(t) for t in targets]
    by_low: dict[str, Counter] = {}
    for t in normed:
        by_low.setdefault(t.lower(), Counter())[t] += 1
    canon = {low: c.most_common(1)[0][0] for low, c in by_low.items()}
    counts = Counter()
    for t in normed:
        counts[canon[t.lower()]] += 1
    min_count = max(3, int(0.01 * n))
    kept = [l for l, c in counts.most_common() if c >= min_count and l]
    if not (2 <= len(kept) <= max_labels):
        return None
    cov = sum(counts[l] for l in kept) / n
    if cov < 0.95 or len(kept) > n / 5 or counts[kept[0]] / n > 0.95:
        return None
    if any(len(l.split()) > MAX_LABEL_WORDS or len(l) > 80 for l in kept):
        return None
    mapping = {raw: canon[norm_label(raw).lower()] for raw in targets}
    return kept, mapping


def order_labels(labels: list[str], defn: str, name: str) -> tuple[list[str], bool]:
    """Sort labels; decide whether they form a rating scale (score)."""
    ints = [as_int(l) for l in labels]
    if all(v is not None for v in ints):
        best: dict[int, str] = {}
        for v, l in zip(ints, labels):
            if v not in best or len(l) > len(best[v]):
                best[v] = l
        pairs = sorted(best.items())
        vals = [v for v, _ in pairs]
        consecutive = vals == list(range(vals[0], vals[0] + len(vals)))
        is_score = consecutive and 3 <= len(vals) <= 10 and bool(_RATING.search(defn + " " + name.replace("_", " ")))
        return [l for _, l in pairs], is_score
    return sorted(labels, key=str.lower), False


# ---------------------------------------------------------------- Super-NaturalInstructions

NI_HEAD = 256 * 1024
NI_MORE = 2 * 1024 * 1024
NI_SKIP_NAME = re.compile(r"incorrect_answer|question_generation|extraction|translation|_spans?_", re.I)


def _ni_files() -> list[tuple[str, int]]:
    p = CACHE / "superni" / "files.json"
    if not p.exists():
        from huggingface_hub import HfApi
        info = HfApi().dataset_info(NI_REPO, revision=NI_REV, files_metadata=True)
        files = [(s.rfilename, s.size) for s in info.siblings if s.rfilename.endswith(".jsonl")]
        p.parent.mkdir(parents=True, exist_ok=True)
        p.write_text(json.dumps(files))
    return [tuple(x) for x in json.loads(p.read_text())]


def _ni_read(fname: str, size: int, max_bytes: int) -> list[dict]:
    """Range-read the first `max_bytes` of a task file (cached); drop the partial last line."""
    import requests
    from huggingface_hub import hf_hub_url
    n = min(size, max_bytes)
    path = CACHE / "superni" / f"{fname.replace('/', '__')}.{n}"
    if not path.exists():
        url = hf_hub_url(NI_REPO, fname, repo_type="dataset", revision=NI_REV)
        for attempt in range(6):
            try:
                r = requests.get(url, headers={"Range": f"bytes=0-{n - 1}"}, timeout=180)
                r.raise_for_status()
                break
            except requests.RequestException:
                if attempt == 5:
                    raise
                time.sleep(2 + 3 * attempt)
        tmp = path.with_suffix(".tmp")
        tmp.write_bytes(r.content[:n])
        tmp.rename(path)
    lines = path.read_bytes().split(b"\n")
    if n < size:
        lines = lines[:-1]
    rows = []
    for ln in lines:
        if ln.strip():
            try:
                rows.append(json.loads(ln))
            except json.JSONDecodeError:
                pass
    return rows


def _parallel(fn, items, workers: int = 8):
    with ThreadPoolExecutor(workers) as ex:
        return list(ex.map(fn, items))


def load_superni(dc: Decontaminator, stats: Counter) -> list[Task]:
    files = _ni_files()
    heads = _parallel(lambda f: _ni_read(f[0], f[1], NI_HEAD), files)
    cands: list[tuple[str, int, list[dict]]] = []
    for (fname, size), rows in zip(files, heads):
        if not rows:
            continue
        name, defn = rows[0]["task_name"], clean(rows[0]["definition"])
        if dc.blocked_task(name + " " + defn):
            stats["tasks_blocked_heldout"] += 1
            continue
        if closed_label_set([r["targets"] for r in rows]) is None and not (
                len(rows) < 60 and size > NI_HEAD):
            continue
        if duplicate_task(name):
            stats["tasks_duplicate_source"] += 1
            continue
        if NI_SKIP_NAME.search(name):
            stats["tasks_skipped_name"] += 1
            continue
        cands.append((fname, size, rows))
    more = [(f, s) for f, s, rows in cands if s > NI_HEAD and len(rows) < 800]
    fetched = dict(zip([f for f, _ in more], _parallel(lambda f: _ni_read(f[0], f[1], NI_MORE), more)))
    tasks: list[Task] = []
    for fname, size, rows in cands:
        rows = fetched.get(fname, rows)
        t = _ni_task(rows, stats)
        if t is not None:
            tasks.append(t)
    return tasks


def _ni_task(rows: list[dict], stats: Counter) -> Optional[Task]:
    name, defn = rows[0]["task_name"], clean(rows[0]["definition"])
    cls = closed_label_set([r["targets"] for r in rows])
    if cls is None:
        stats["tasks_not_classification"] += 1
        return None
    labels, mapping = cls
    ints = [as_int(l) for l in labels]
    if (all(v is not None for v in ints) and len(labels) > 10) or (
            len(labels) > 10 and all(len(l) == 1 for l in labels)):
        stats["tasks_skipped_counting"] += 1
        return None
    if not english_like(defn):
        stats["tasks_non_english"] += 1
        return None
    labels, is_score = order_labels(labels, defn, name)
    idx = {l: i for i, l in enumerate(labels)}
    out, non_en = [], 0
    for r in rows:
        lab = mapping.get(r["targets"])
        if lab not in idx:
            continue
        state = truncate(clean(r["inputs"]))
        if not state:
            continue
        if not english_like(state):
            non_en += 1
            continue
        out.append(Row(r["id"].split("-")[-1][:16], state, idx[lab]))
    if non_en > 0.3 * len(rows):
        stats["tasks_non_english"] += 1
        return None
    if len(out) < MIN_ROWS:
        stats["tasks_too_small"] += 1
        return None
    opaque = any(opaque_label(l) for l in labels) or all(l.lower() in _YESNO for l in labels)
    short = first_sentences(defn)
    return Task("superni", name, shorten_def(defn), "labels", out, labels, is_score,
                short_instr=[short] if short and short != defn else [],
                generic=GENERIC_CLS, opaque=opaque)


# ---------------------------------------------------------------- BIG-bench

BB_SKIP = {
    # non-English / translation / transliteration
    "cryobiology_spanish", "english_russian_proverbs", "entailed_polarity_hindi", "gender_inclusive_sentences_german",
    "hindi_question_answering", "hinglish_toxicity", "indic_cause_and_effect", "kannada", "medical_questions_russian",
    "misconceptions_russian", "parsinlu_qa", "parsinlu_reading_comprehension", "persian_idioms",
    "swahili_english_proverbs", "swedish_to_german_proverbs", "conlang_translation", "polish_sequence_labeling",
    "language_identification", "international_phonetic_alphabet_nli", "international_phonetic_alphabet_transliterate",
    "salient_translation_error_detection", "which_wiki_edit",
    # ascii-art images / symbol puzzles useless to a text encoder
    "mnist_ascii", "kanji_ascii", "ascii_word_recognition", "cifar10_classification", "checkmate_in_one",
    # duplicates of main-pipeline sources
    "social_iqa", "vitaminc_fact_verification",
}
BB_MAX_OPTIONS = 20


def _bb_clean_input(text: str) -> tuple[str, Optional[str]]:
    """Drop inline 'choice:' lines and the trailing 'A:'; split off a final question line."""
    lines = [ln.rstrip() for ln in str(text).split("\n")]
    lines = [ln for ln in lines if not re.match(r"\s*choice\s*:", ln)]
    while lines and re.fullmatch(r"\s*(A|Answer|Output)\s*:\s*", lines[-1]):
        lines.pop()
    body = clean("\n".join(lines))
    if body.startswith("Q: ") and body.count("Q:") == 1:
        body = body[3:]
    parts = body.rsplit("\n", 1)
    if len(parts) == 2 and parts[1].strip().endswith("?") and len(parts[1]) <= 250 and len(parts[0].strip()) >= 20:
        q = re.sub(r"^(Q|Question)\s*:\s*", "", parts[1].strip())
        return parts[0].strip(), q
    return body, None


def load_bigbench(dc: Decontaminator, stats: Counter) -> list[Task]:
    import pandas as pd
    from huggingface_hub import HfApi
    info = HfApi().dataset_info(BB_REPO, revision=BB_REV)
    names = sorted({s.rfilename.split("/")[0] for s in info.siblings if s.rfilename.endswith(".parquet")})
    tasks: list[Task] = []
    for name in names:
        if name in BB_SKIP:
            stats["tasks_skipped_name"] += 1
            continue
        if dc.blocked_task(name):
            stats["tasks_blocked_heldout"] += 1
            continue
        if duplicate_task(name):
            stats["tasks_duplicate_source"] += 1
            continue
        dfs = []
        for split in ("train", "validation"):
            try:
                dfs.append(pd.read_parquet(fetch(BB_REPO, f"{name}/{split}-00000-of-00001.parquet", BB_REV)))
            except Exception as e:  # noqa: BLE001
                log(f"  bigbench/{name}/{split}: {type(e).__name__}")
        if not dfs:
            continue
        df = pd.concat(dfs, ignore_index=True)
        rows, questions, non_en = [], Counter(), 0
        for i, r in enumerate(df.itertuples(index=False)):
            opts = [clean(o) for o in list(r.multiple_choice_targets)]
            scores = list(r.multiple_choice_scores)
            if not (2 <= len(opts) <= BB_MAX_OPTIONS) or len(scores) != len(opts) or sum(s == 1 for s in scores) != 1:
                continue
            if any(not o for o in opts) or len({o.lower() for o in opts}) != len(opts):
                continue
            if any(len(o.split()) > MAX_MCQ_WORDS for o in opts) or sum(len(o) for o in opts) > 1000:
                continue
            state, q = _bb_clean_input(r.inputs)
            if not state:
                continue
            if not english_like(state + " ".join(opts)):
                non_en += 1
                continue
            questions[q] += 1
            rows.append((i, truncate(state), q, opts, scores.index(1)))
        if non_en > 0.3 * max(1, len(df)):
            stats["tasks_non_english"] += 1
            continue
        if len(rows) < MIN_ROWS:
            stats["tasks_too_small" if rows else "tasks_not_classification"] += 1
            continue
        # a question line shared by most rows is the task instruction; per-row questions stay in the state
        top_q, top_n = questions.most_common(1)[0]
        fixed_q = top_q if top_q and top_n >= 0.5 * len(rows) else None
        rws = []
        for i, state, q, opts, gold in rows:
            if q and q != fixed_q:
                state = f"{state}\n{q}"
            rws.append(Row(str(i), truncate(state), gold, opts))
        heads = Counter(r.state.split("\n", 1)[0] for r in rws if "\n" in r.state)
        top_h, top_hn = heads.most_common(1)[0] if heads else ("", 0)
        if len(top_h) >= 25 and top_hn >= 0.6 * len(rws):
            rws = [Row(r.rid, r.state.split("\n", 1)[1].strip(), r.gold, r.options)
                   if r.state.startswith(top_h + "\n") and r.state.split("\n", 1)[1].strip() else r for r in rws]
            fixed_q = f"{top_h} {fixed_q}" if fixed_q else top_h
        instr = fixed_q or rng_free_generic(name)
        tasks.append(Task("bigbench", name, instr, "mcq", rws, generic=GENERIC_MCQ,
                          p_generic=0.12 if fixed_q else 0.4))
    return tasks


def rng_free_generic(name: str) -> str:
    """Deterministic default instruction for a BIG-bench task without a question line."""
    w = task_words(name)
    tmpl = GENERIC_MCQ[sum(map(ord, name)) % len(GENERIC_MCQ)]
    return f"Task: {w}. {tmpl}"


# ---------------------------------------------------------------- tasksource-instruct

_TS_HEAD = re.compile(r'With no explanation, (.*?)(either|from) ((?:"[^"]*"(?:, | or )?)+)\.', re.S)
TS_PAIR = [
    "Label the relationship between text_A and text_B.", "How does text_B relate to text_A?",
    "Classify the relation from text_A to text_B.", "Choose the label for the pair (text_A -> text_B).",
    "Which label describes how text_A and text_B relate?",
]
TS_SINGLE = [
    "Label the following text.", "Classify the text.", "Which label applies to this input?",
    "Choose the correct label.", "Categorize the following.",
]


def _ts_parse(inputs: str, target: str) -> Optional[tuple[str, str, list[str], int, str]]:
    """-> (question prefix, state, options, gold, original instruction) or None."""
    m = _TS_HEAD.search(inputs)
    if not m:
        return None
    opts = re.findall(r'"([^"]*)"', m.group(3))
    prefix = inputs[:m.start()].strip()
    orig = (m.group(1) + m.group(2) + " " + m.group(3)).strip() + "."
    state = inputs[m.end():].strip()
    gold_s = norm_label(target)
    if m.group(2) == "from" and all(re.fullmatch(r"[A-Z]", o) for o in opts):
        # letter MCQ: options are the trailing 'A: ...' lines of the state
        pos = [state.rfind(f"\n{o}: ") for o in opts]
        if any(p < 0 for p in pos) or pos != sorted(pos):
            return None
        texts = [state[pos[k] + len(opts[k]) + 3: (pos[k + 1] if k + 1 < len(pos) else len(state))].strip()
                 for k in range(len(opts))]
        if gold_s not in opts or any(not t for t in texts):
            return None
        return prefix, state[:pos[0]].strip(), texts, opts.index(gold_s), ""
    if gold_s not in opts or len(set(opts)) != len(opts) or any(not o for o in opts):
        return None
    return prefix, state, opts, opts.index(gold_s), orig


def load_tasksource(dc: Decontaminator, stats: Counter) -> list[Task]:
    import pandas as pd
    df = pd.read_parquet(fetch(TS_REPO, "data/validation-00000-of-00001.parquet", TS_REV),
                         columns=["inputs", "targets", "task"])
    tasks: list[Task] = []
    for name, sub in sorted(df.groupby("task"), key=lambda x: x[0]):
        if duplicate_task(name):
            stats["tasks_duplicate_source"] += 1
            continue
        parsed = []
        for i, (inp, tgt) in zip(sub.index, zip(sub.inputs, sub.targets)):
            if "\n" in tgt.strip():
                continue  # token classification
            p = _ts_parse(inp, tgt)
            if p is not None:
                parsed.append((int(i), p))
        if len(parsed) < MIN_ROWS:
            stats["tasks_not_classification" if len(parsed) < 5 else "tasks_too_small"] += 1
            continue
        prefix = Counter(p[0] for _, p in parsed).most_common(1)[0][0]
        orig = Counter(p[4] for _, p in parsed).most_common(1)[0][0]
        if dc.blocked_task(f"{name} {prefix} {orig}"):
            stats["tasks_blocked_heldout"] += 1
            continue
        mcq = not parsed[0][1][4]
        rows, non_en, label_set = [], 0, Counter()
        for i, (pre, state, opts, gold, row_orig) in parsed:
            if pre and pre != prefix:
                state = f"{pre}\n{state}"
            state = truncate(clean(state))
            opts = [clean(o) for o in opts]
            if not state or not english_like(state + " ".join(opts)):
                non_en += 1
                continue
            if any(len(o.split()) > (MAX_MCQ_WORDS if mcq else MAX_LABEL_WORDS) for o in opts):
                continue
            if len({o.lower() for o in opts}) != len(opts):
                continue
            label_set.update(opts)
            # without a task-specific question, the row's own 'label ... with either "a" or "b".' is the
            # main instruction (it lists exactly this row's options)
            row_instr = None if prefix or mcq else cap_first(clean(row_orig))
            rows.append(Row(str(i), state, gold, opts, row_instr))
        if non_en > 0.3 * len(parsed):
            stats["tasks_non_english"] += 1
            continue
        if len(rows) < MIN_ROWS:
            stats["tasks_too_small"] += 1
            continue
        labels = list(label_set)
        is_score = False
        if not mcq:
            ordered, is_score = order_labels(labels, prefix, name)
            if is_score:
                idx = {as_int(l): k for k, l in enumerate(ordered)}
                rows = [Row(r.rid, r.state, idx[as_int(r.options[r.gold])], None, r.instr) for r in rows]
                labels = ordered
            else:
                is_score = False
        pair = sum("text_A" in r.state and "text_B" in r.state for r in rows) > 0.5 * len(rows)
        generic = GENERIC_MCQ if mcq else (TS_PAIR if pair else TS_SINGLE)
        main = prefix or f"Task: {task_words(name)}. {GENERIC_MCQ[0] if mcq else GENERIC_CLS[0]}"
        kind = "labels" if is_score else ("mcq" if mcq else "labels_rowwise")
        t = Task("tasksource", name, main, kind, rows, labels if is_score else [], is_score,
                 generic=generic, p_short=0.0, p_generic=0.15 if prefix else 0.45)
        tasks.append(t)
    return tasks


# ---------------------------------------------------------------- driver

def records_for(t: Task, seed: int, cap: int, dc: Decontaminator, stats: Counter) -> list[Record]:
    if t.kind != "labels_rowwise":
        return task_records(t, seed, cap, dc, stats)
    rng = random.Random(f"{seed}|{t.source}")
    rows = dedupe_rows(t.rows, lambda r: r.state, lambda r: r.options[r.gold])
    rows = balanced_order(rng, rows, lambda r: r.options[r.gold], alpha=0.5)
    t.kind = "labels"
    out: list[Record] = []
    for row in rows:
        if len(out) >= cap:
            break
        if dc.is_contaminated(row.state):
            stats["dropped_contaminated"] += 1
            continue
        t.labels = row.options
        q, lab = build_question(rng, t, Row(row.rid, row.state, row.gold, None, row.instr))
        rec = Record(f"{t.source}/{row.rid}", t.source, row.state, q, lab)
        rec.validate()
        out.append(rec)
    t.kind, t.labels = "labels_rowwise", []
    return out


LOADERS = {"superni": load_superni, "bigbench": load_bigbench, "tasksource": load_tasksource}


def write_stats(path: Path, per_coll: dict[str, list[Record]], stats: dict[str, Counter], caps: dict[str, int]) -> None:
    L = ["# data/extra/instruct — task-diverse instruction-following data", "",
         "Generated by `python -m jevlike.data.extra_instruct`. Each task keeps its own definition as the",
         "question instructions and its own label set as options. `source` = `x_<collection>/<task>`.", ""]
    allr = [r for rs in per_coll.values() for r in rs]
    L += ["## Overview", "", "| collection | cap/task | tasks | records | choice | score | noul | noul true % |",
          "|---|---|---|---|---|---|---|---|"]
    for coll, rs in list(per_coll.items()) + [("**total**", allr)]:
        tc = Counter(r.question.type for r in rs)
        nt = len({r.source for r in rs})
        nn = [r.label for r in rs if r.question.type == "noul"]
        pt = f"{100 * sum(nn) / len(nn):.1f}" if nn else "-"
        L.append(f"| {coll} | {caps.get(coll, '')} | {nt} | {len(rs)} | {tc['choice']} | {tc['score']} | {tc['noul']} | {pt} |")
    L += ["", "## Filtering (tasks / records dropped)", "", "| collection | " + " | ".join(
        k for k in sorted({k for c in stats.values() for k in c})) + " |"]
    keys = sorted({k for c in stats.values() for k in c})
    L.append("|---|" + "---|" * len(keys))
    for coll, c in stats.items():
        L.append(f"| {coll} | " + " | ".join(str(c.get(k, 0)) for k in keys) + " |")
    L += ["", "## Option-count histogram (choice + score)", "", "| options | choice | score |", "|---|---|---|"]
    bins = [(2, 2), (3, 3), (4, 4), (5, 5), (6, 10), (11, 20), (21, 30)]
    for lo, hi in bins:
        c = sum(1 for r in allr if r.question.type == "choice" and lo <= len(r.question.options) <= hi)
        s = sum(1 for r in allr if r.question.type == "score" and lo <= len(r.question.options) <= hi)
        L.append(f"| {lo}{'' if lo == hi else f'-{hi}'} | {c} | {s} |")
    # gold position balance for choice
    pos = Counter((len(r.question.options), r.label) for r in allr if r.question.type == "choice"
                  and len(r.question.options) <= 4)
    L += ["", "## Choice gold-position balance (n<=4)", ""]
    for n in (2, 3, 4):
        tot = sum(pos[(n, i)] for i in range(n))
        if tot:
            L.append(f"- n={n}: " + ", ".join(f"pos{i} {100 * pos[(n, i)] / tot:.0f}%" for i in range(n)))
    ts = Counter(r.source for r in allr)
    L += ["", "## Records per task", "",
          f"- tasks: {len(ts)}; records/task: min {min(ts.values()) if ts else 0}, "
          f"median {sorted(ts.values())[len(ts) // 2] if ts else 0}, max {max(ts.values()) if ts else 0}", "",
          "Top 25 tasks:", "", "| task | records | types |", "|---|---|---|"]
    by_task_types: dict[str, Counter] = {}
    for r in allr:
        by_task_types.setdefault(r.source, Counter())[r.question.type] += 1
    for s, n in ts.most_common(25):
        L.append(f"| {s} | {n} | " + ", ".join(f"{k} {v}" for k, v in sorted(by_task_types[s].items())) + " |")
    L += ["", "Score tasks: " + ", ".join(sorted({r.source for r in allr if r.question.type == 'score'})), ""]
    L += ["## All tasks", ""]
    for coll, rs in per_coll.items():
        names = sorted({r.source.split("/", 1)[1] for r in rs})
        L.append(f"**{coll}** ({len(names)}): " + ", ".join(names))
        L.append("")
    path.write_text("\n".join(L) + "\n", encoding="utf-8")


def main(argv=None) -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--out", default=str(OUT))
    ap.add_argument("--seed", type=int, default=13)
    ap.add_argument("--only", default="", help="comma-separated subset of: " + ",".join(LOADERS))
    ap.add_argument("--cap-superni", type=int, default=250)
    ap.add_argument("--cap-bigbench", type=int, default=150)
    ap.add_argument("--cap-tasksource", type=int, default=80)
    args = ap.parse_args(argv)
    out = Path(args.out)
    out.mkdir(parents=True, exist_ok=True)
    caps = {"superni": args.cap_superni, "bigbench": args.cap_bigbench, "tasksource": args.cap_tasksource}
    colls = [c for c in (args.only.split(",") if args.only else LOADERS) if c]
    dc = Decontaminator()
    per_coll: dict[str, list[Record]] = {}
    stats: dict[str, Counter] = {}
    for coll in colls:
        st = Counter()
        log(f"[{coll}] loading ...")
        tasks = LOADERS[coll](dc, st)
        log(f"[{coll}] {len(tasks)} candidate tasks")
        recs: list[Record] = []
        for t in tasks:
            rs = records_for(t, args.seed, caps[coll], dc, st)
            if len(rs) < MIN_ROWS // 2:
                st["tasks_too_small"] += 1
                continue
            recs.extend(rs)
        random.Random(f"{args.seed}|{coll}|shuffle").shuffle(recs)
        write_jsonl(str(out / f"{coll}.jsonl"), recs)
        per_coll[coll] = recs
        stats[coll] = st
        tc = Counter(r.question.type for r in recs)
        log(f"[{coll}] {len(recs)} records from {len({r.source for r in recs})} tasks {dict(tc)}; {dict(st)}")
    if not args.only:
        write_stats(out / "STATS.md", per_coll, stats, caps)
    else:
        write_stats(out / f"STATS.{'_'.join(colls)}.md", per_coll, stats, caps)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())