"""Converters for classification datasets (primary type: choice).

Each converter yields `Example`s: a choice question over a sampled option subset
(varied keys / descriptions / instructions), often a derived true/false statement, and
occasionally a second choice question with a different option set.
"""
from __future__ import annotations

import random
import re
from dataclasses import dataclass
from typing import Callable, Iterator, Optional, Sequence

from jevlike.data.common import (
    Example, Label, balanced_order, clean, dedupe_rows, derived_noul, intent_label, make_choice,
    make_mcq, noul, pick_instr, readable, weighted,
)
from jevlike.data.hub import read_many, read_rows

GENERIC_CHOICE = [
    "Which option best describes the text?", "Choose the best matching option.",
    "Pick the most appropriate label.", "Classify the text.", "Select one.",
    "Which category does this belong to?", "Assign the correct category.",
    "What is the best label for this input?", "Which of these applies?",
    "Pick one of the options below.", "Label this.", "Categorize the input.",
    "Decide which option fits best.", "Which option matches the input best?",
    "Choose one label for the text above.",
]

TOPIC_POS = [
    "This {thing} is about {x}.", "The main topic of this {thing} is {x}.",
    "The {thing} discusses {x}.", "Is this {thing} about {x}?", "This {thing} mainly concerns {x}.",
    "The subject of the {thing} is {x}.", "Topic: {x}.",
]
TOPIC_NEG = ["This {thing} is not about {x}.", "The {thing} has nothing to do with {x}.",
             "This {thing} is unrelated to {x}."]


@dataclass
class Ctx:
    rng: random.Random
    name: str


# ---------------------------------------------------------------- generic classification loop

def classify(ctx: Ctx, rows: list, text_fn: Callable, label_fn: Callable, labels: list[Label], *,
             instr: Sequence[str], noul_pos: Sequence[str], noul_neg: Sequence[str] = (),
             thing: str = "text", alpha: float = 0.5, full_prob: float = 0.08,
             other_prob: float = 0.03, p_choice: float = 0.85, p_noul: float = 0.45,
             p_second: float = 0.12, p_other_group: float = 0.6, p_generic: float = 0.2,
             extra: Optional[Callable] = None) -> Iterator[Example]:
    """Shared loop: balanced row order -> choice (+ derived noul, + second choice)."""
    rng = ctx.rng
    rows = dedupe_rows(rows, text_fn, label_fn)
    for row in balanced_order(rng, rows, label_fn, alpha):
        state = text_fn(row)
        if len(state) < 3:
            continue
        gold = label_fn(row)
        qs = []
        if rng.random() < p_choice:
            qs.append(make_choice(rng, labels, gold, pick_instr(rng, instr, GENERIC_CHOICE, p_generic, thing=thing),
                                  full_prob=full_prob, other_prob=other_prob))
        if not qs or rng.random() < p_noul:
            qs.append(derived_noul(rng, labels, gold, noul_pos, noul_neg, p_other_group=p_other_group,
                                   thing=thing))
        if rng.random() < p_second:
            qs.append(make_choice(rng, labels, gold, pick_instr(rng, instr, GENERIC_CHOICE, 0.5, thing=thing),
                                  full_prob=full_prob, other_prob=other_prob))
        if extra is not None:
            qs.extend(extra(row, rng))
        yield Example(state, qs[:3])


# ---------------------------------------------------------------- topic datasets

AG_LABELS = [
    Label("world", ["world news and international affairs", "international politics, conflicts and diplomacy",
                    "global news"], "world news or international affairs"),
    Label("sports", ["sports news: games, teams and athletes", "sport events and results"], "sports"),
    Label("business", ["business, companies and the economy", "markets, finance and corporate news"],
          "business or the economy"),
    Label("science and technology", ["science and technology news", "computers, internet, space and research",
                                     "the tech industry and scientific discoveries"],
          "science or technology", key="sci_tech"),
]
NEWS_INSTR = [
    "What is this {thing} mainly about?", "Which section of the newspaper does this {thing} belong in?",
    "Topic?", "Classify the news {thing} by topic.", "Which category best fits this {thing}?",
    "What kind of news is this?", "Assign a news category to the {thing}.",
    "Pick the topic of the {thing} from the list.", "Which desk should cover this story?",
]


def ag_news(ctx: Ctx):
    rows = read_many("fancyzhx/ag_news", ["data/test-00000-of-00001.parquet", "data/train-00000-of-00001.parquet"])
    yield from classify(ctx, rows, lambda r: clean(r["text"]), lambda r: int(r["label"]), AG_LABELS,
                        instr=NEWS_INSTR, noul_pos=TOPIC_POS, noul_neg=TOPIC_NEG, thing="article",
                        alpha=1.0, p_other_group=0.0)


DBP = [
    ("Company", "a company", ["a company or business", "a business organization"], "org"),
    ("Educational institution", "an educational institution", ["a school, college or university"], "org"),
    ("Artist", "an artist", ["an artist such as a musician, painter, actor or writer"], "person"),
    ("Athlete", "an athlete", ["an athlete or sports player"], "person"),
    ("Office holder", "a politician or office holder", ["a politician or public office holder"], "person"),
    ("Means of transportation", "a vehicle or other means of transportation",
     ["a vehicle, ship, aircraft or train"], "thing"),
    ("Building", "a building", ["a building or man-made structure"], "place"),
    ("Natural place", "a natural place", ["a natural place such as a mountain, river or lake"], "place"),
    ("Village", "a village", ["a village or small settlement"], "place"),
    ("Animal", "an animal", ["an animal species"], "nature"),
    ("Plant", "a plant", ["a plant species"], "nature"),
    ("Album", "a music album", ["a music album or record"], "work"),
    ("Film", "a film", ["a film or movie"], "work"),
    ("Written work", "a written work", ["a book, novel, journal or other written work"], "work"),
]
DBP_LABELS = [Label(n.lower(), d, p, group=g) for n, p, d, g in DBP]


def dbpedia(ctx: Ctx):
    rows = read_rows("fancyzhx/dbpedia_14", "dbpedia_14/test-00000-of-00001.parquet")
    rng = ctx.rng

    def text(r):
        c = clean(r["content"])
        return c if rng.random() < 0.6 else f"{clean(r['title'])}\n{c}"
    instr = ["What kind of entity does this text describe?", "What is the subject of this encyclopedia entry?",
             "Which ontology class fits this article?", "Classify the entity described.",
             "What type of thing is this passage about?", "Which category does the described entity belong to?"]
    pos = ["This text describes {x}.", "The subject of this passage is {x}.",
           "This is an encyclopedia entry about {x}.", "Is this about {x}?", "The passage is about {x}."]
    neg = ["This text does not describe {x}.", "The subject of this passage is not {x}."]
    rows = [dict(r, _t=text(r)) for r in rows]
    yield from classify(ctx, rows, lambda r: r["_t"], lambda r: int(r["label"]), DBP_LABELS,
                        instr=instr, noul_pos=pos, noul_neg=neg, thing="entry", alpha=1.0,
                        p_other_group=0.3)


YAHOO = [
    ("society and culture", ["society, culture, religion and traditions"], "society or culture"),
    ("science and mathematics", ["science and math questions"], "science or mathematics"),
    ("health", ["health, medicine and fitness"], "health"),
    ("education and reference", ["education, school, language and reference questions"], "education"),
    ("computers and internet", ["computers, software and the internet"], "computers or the internet"),
    ("sports", ["sports and athletes"], "sports"),
    ("business and finance", ["business, money, jobs and finance"], "business or finance"),
    ("entertainment and music", ["movies, TV, music and celebrities"], "entertainment or music"),
    ("family and relationships", ["family, dating, friendship and relationships"], "family or relationships"),
    ("politics and government", ["politics, government and law"], "politics or government"),
]
YAHOO_LABELS = [Label(n, d, p) for n, d, p in YAHOO]


def yahoo(ctx: Ctx):
    rows = read_rows("community-datasets/yahoo_answers_topics", "yahoo_answers_topics/test-00000-of-00001.parquet")
    rng = ctx.rng

    def text(r):
        t, c, a = clean(r["question_title"]), clean(r["question_content"]), clean(r["best_answer"])
        s = t + (f"\n{c}" if c else "")
        u = rng.random()
        if a and u < 0.35:
            s = f"Question: {s}\nBest answer: {a}"
        elif a and u < 0.5:
            s = f"{s}\n\n{a}"
        return s
    rows = [dict(r, _t=text(r)) for r in rows]
    instr = ["Which category of the Q&A site does this question belong to?", "What is this question about?",
             "Choose the forum section for this post.", "Topic of the question?",
             "In which category should this question be posted?"]
    yield from classify(ctx, rows, lambda r: r["_t"], lambda r: int(r["topic"]), YAHOO_LABELS,
                        instr=instr, noul_pos=TOPIC_POS, noul_neg=TOPIC_NEG, thing="question",
                        alpha=1.0, p_other_group=0.0)


NG = {
    "alt.atheism": ("atheism", "discussion of atheism and religious skepticism", "religion"),
    "comp.graphics": ("computer graphics", "computer graphics, images and rendering software", "comp"),
    "comp.os.ms-windows.misc": ("Microsoft Windows", "the Microsoft Windows operating system", "comp"),
    "comp.sys.ibm.pc.hardware": ("PC hardware", "IBM PC compatible hardware", "comp"),
    "comp.sys.mac.hardware": ("Mac hardware", "Apple Macintosh hardware", "comp"),
    "comp.windows.x": ("X Window System", "the X Window System (X11) on Unix", "comp"),
    "misc.forsale": ("for sale", "items for sale and classified ads", "misc"),
    "rec.autos": ("cars", "automobiles and driving", "rec"),
    "rec.motorcycles": ("motorcycles", "motorcycles and riding", "rec"),
    "rec.sport.baseball": ("baseball", "baseball games, teams and players", "rec"),
    "rec.sport.hockey": ("hockey", "ice hockey games, teams and players", "rec"),
    "sci.crypt": ("cryptography", "cryptography, encryption and privacy", "sci"),
    "sci.electronics": ("electronics", "electronic circuits and components", "sci"),
    "sci.med": ("medicine", "medicine and health", "sci"),
    "sci.space": ("space", "space exploration and astronomy", "sci"),
    "soc.religion.christian": ("Christianity", "Christian religion and faith", "religion"),
    "talk.politics.guns": ("gun politics", "gun control and firearms politics", "politics"),
    "talk.politics.mideast": ("Middle East politics", "politics of the Middle East", "politics"),
    "talk.politics.misc": ("politics", "general political discussion", "politics"),
    "talk.religion.misc": ("religion", "general religious discussion", "religion"),
}


def newsgroups(ctx: Ctx):
    rows = read_many("SetFit/20_newsgroups", ["train.jsonl", "test.jsonl"])
    names = list(NG)
    labels = [Label(NG[n][0], [NG[n][1], n], NG[n][1], key=n, group=NG[n][2]) for n in names]
    rows = [dict(r, _t=clean(r["text"])) for r in rows if len(clean(r["text"])) >= 60]
    instr = ["Which newsgroup was this post published in?", "What is this forum post about?",
             "Route this message to the right discussion group.", "Which group does this post belong to?",
             "Classify the post by topic."]
    yield from classify(ctx, rows, lambda r: r["_t"], lambda r: names.index(r["label_text"]), labels,
                        instr=instr, noul_pos=TOPIC_POS, noul_neg=TOPIC_NEG, thing="post",
                        alpha=1.0, p_other_group=0.5)


TT = {
    "arts_&_culture": ("arts and culture", "arts, books, museums and culture"),
    "business_&_entrepreneurs": ("business", "business and entrepreneurship"),
    "pop_culture": ("pop culture", "celebrities, music, movies and TV"),
    "daily_life": ("daily life", "everyday personal life, family and feelings"),
    "sports_&_gaming": ("sports and gaming", "sports, athletes and video games"),
    "science_&_technology": ("science and technology", "science, tech and gadgets"),
}


def _tweet_clean(t: str) -> str:
    t = re.sub(r"\{@([^@}]*)@\}", r"@\1", t)
    t = t.replace("{{URL}}", "http://t.co/link").replace("{{USERNAME}}", "@user")
    return clean(t)


def tweet_topic(ctx: Ctx):
    rows = read_many("cardiffnlp/tweet_topic_single", [
        "dataset/split_coling2022_random/train_random.single.json",
        "dataset/split_coling2022_random/test_random.single.json"])
    names = list(TT)
    labels = [Label(TT[n][0], [TT[n][1]], TT[n][1], key=n.replace("_&_", "_")) for n in names]
    instr = ["What is this tweet about?", "Which topic does the tweet belong to?", "Topic of this tweet?",
             "Classify the tweet by topic.", "Which feed should this tweet appear in?"]
    yield from classify(ctx, rows, lambda r: _tweet_clean(r["text"]), lambda r: names.index(r["label_name"]),
                        labels, instr=instr, noul_pos=TOPIC_POS, noul_neg=TOPIC_NEG, thing="tweet",
                        alpha=0.5, p_other_group=0.0)


FIN_TOPICS = [
    ("analyst update", "analyst ratings, upgrades, downgrades and price targets"),
    ("central banks", "the Fed and central bank policy"),
    ("company or product news", "news about a company or its products"),
    ("treasuries and corporate debt", "bonds, treasuries and corporate debt"),
    ("dividends", "dividend announcements"),
    ("earnings", "earnings reports and results"),
    ("energy and oil", "oil, gas and energy markets"),
    ("financials", "banks and the financial sector"),
    ("currencies", "currencies and foreign exchange"),
    ("general news or opinion", "general news and opinion pieces"),
    ("gold, metals and materials", "gold, metals and raw materials"),
    ("IPO", "initial public offerings"),
    ("legal and regulation", "lawsuits, regulation and legal matters"),
    ("M&A and investments", "mergers, acquisitions and investments"),
    ("macro", "the macroeconomy: inflation, jobs, GDP"),
    ("markets", "overall market moves and indexes"),
    ("politics", "politics and government"),
    ("personnel change", "executive appointments and departures"),
    ("stock commentary", "commentary and opinions about stocks"),
    ("stock movement", "individual stock price moves"),
]


def fin_topic(ctx: Ctx):
    rows = read_many("zeroshot/twitter-financial-news-topic", ["topic_train.csv", "topic_valid.csv"])
    labels = [Label(n, [d], d) for n, d in FIN_TOPICS]
    instr = ["What is this financial tweet about?", "Which finance topic does this headline cover?",
             "Classify the financial news tweet.", "Tag this market news item with a topic.",
             "Which newsletter section should this item go in?"]
    yield from classify(ctx, rows, lambda r: clean(r["text"]), lambda r: int(r["label"]), labels,
                        instr=instr, noul_pos=TOPIC_POS, noul_neg=TOPIC_NEG, thing="tweet",
                        alpha=0.5, p_other_group=0.0)


# ---------------------------------------------------------------- intent / routing datasets

INTENT_INSTR = [
    "What does the customer want?", "Which intent best matches this {thing}?",
    "Route this {thing} to the right intent.", "What is the user asking for?",
    "Classify the customer's request.", "Which issue is the customer raising?",
    "Identify the intent of the {thing}.", "Which workflow should handle this {thing}?",
    "Pick the intent.", "What is this {thing} about?",
]
INTENT_POS = ["The customer is asking about {x}.", "This {thing} is about {x}.",
              "The user's request concerns {x}.", "Intent: {x}.", "Is the user asking about {x}?",
              "The {thing} relates to {x}."]
INTENT_NEG = ["The customer is not asking about {x}.", "This {thing} has nothing to do with {x}."]

_BANK_GROUPS = [("card", "card"), ("transfer", "transfer"), ("top_up", "topup"), ("topping", "topup"),
                ("exchange", "fx"), ("currenc", "fx"), ("cash", "cash"), ("atm", "cash"),
                ("pin", "security"), ("passcode", "security"), ("verify", "identity"),
                ("identity", "identity"), ("refund", "refund"), ("payment", "payment")]


def _bank_group(n: str) -> str:
    for k, g in _BANK_GROUPS:
        if k in n.lower():
            return g
    return "other"


BANKING77 = [
    "activate_my_card", "age_limit", "apple_pay_or_google_pay", "atm_support", "automatic_top_up",
    "balance_not_updated_after_bank_transfer", "balance_not_updated_after_cheque_or_cash_deposit",
    "beneficiary_not_allowed", "cancel_transfer", "card_about_to_expire", "card_acceptance", "card_arrival",
    "card_delivery_estimate", "card_linking", "card_not_working", "card_payment_fee_charged",
    "card_payment_not_recognised", "card_payment_wrong_exchange_rate", "card_swallowed",
    "cash_withdrawal_charge", "cash_withdrawal_not_recognised", "change_pin", "compromised_card",
    "contactless_not_working", "country_support", "declined_card_payment", "declined_cash_withdrawal",
    "declined_transfer", "direct_debit_payment_not_recognised", "disposable_card_limits",
    "edit_personal_details", "exchange_charge", "exchange_rate", "exchange_via_app",
    "extra_charge_on_statement", "failed_transfer", "fiat_currency_support", "get_disposable_virtual_card",
    "get_physical_card", "getting_spare_card", "getting_virtual_card", "lost_or_stolen_card",
    "lost_or_stolen_phone", "order_physical_card", "passcode_forgotten", "pending_card_payment",
    "pending_cash_withdrawal", "pending_top_up", "pending_transfer", "pin_blocked", "receiving_money",
    "Refund_not_showing_up", "request_refund", "reverted_card_payment?", "supported_cards_and_currencies",
    "terminate_account", "top_up_by_bank_transfer_charge", "top_up_by_card_charge",
    "top_up_by_cash_or_cheque", "top_up_failed", "top_up_limits", "top_up_reverted", "topping_up_by_card",
    "transaction_charged_twice", "transfer_fee_charged", "transfer_into_account",
    "transfer_not_received_by_recipient", "transfer_timing", "unable_to_verify_identity",
    "verify_my_identity", "verify_source_of_funds", "verify_top_up", "virtual_card_not_working",
    "visa_or_mastercard", "why_verify_identity", "wrong_amount_of_cash_received",
    "wrong_exchange_rate_for_cash_withdrawal",
]


def banking77(ctx: Ctx):
    rows = read_many("legacy-datasets/banking77", ["data/train-00000-of-00001.parquet",
                                                   "data/test-00000-of-00001.parquet"])
    labels = []
    for n in BANKING77:
        r = readable(n)
        labels.append(Label(r, [r, f"customer asks about {r}", f"banking query: {r}"], r,
                            key=re.sub(r"[^a-z_]", "", n.lower()), group=_bank_group(n)))
    yield from classify(ctx, rows, lambda r: clean(r["text"]), lambda r: int(r["label"]), labels,
                        instr=INTENT_INSTR, noul_pos=INTENT_POS, noul_neg=INTENT_NEG, thing="banking query",
                        alpha=1.0, full_prob=0.08, other_prob=0.05, p_other_group=0.5)


def clinc(ctx: Ctx):
    rows = read_many("clinc/clinc_oos", [f"plus/{s}-00000-of-00001.parquet" for s in ("train", "validation", "test")])
    from jevlike.data.label_names import CLINC150
    labels = []
    for n in CLINC150:
        if n == "oos":
            labels.append(Label("out of scope", ["request not covered by any listed intent",
                                                 "unsupported request"], "something outside the supported requests",
                                key="out_of_scope"))
        else:
            labels.append(intent_label(n))
    oos = CLINC150.index("oos")
    rng = ctx.rng
    # keep ~5% out-of-scope examples; never make "out of scope" a negative-only filler
    rows = [r for r in rows if int(r["intent"]) != oos or rng.random() < 0.3]
    instr = ["What does the user want the assistant to do?", "Which intent does this request express?",
             "Classify the user's request.", "Which skill should handle this command?",
             "Map this utterance to an intent.", "What is the user trying to do?"]
    pos = ["The user wants help with {x}.", "This request is about {x}.", "Intent: {x}.",
           "Is the user asking about {x}?"]
    neg = ["The user is not asking about {x}."]
    yield from classify(ctx, rows, lambda r: clean(r["text"]), lambda r: int(r["intent"]), labels,
                        instr=instr, noul_pos=pos, noul_neg=neg, thing="request", alpha=1.0,
                        full_prob=0.05, other_prob=0.04, p_other_group=0.0)


MASSIVE_SCEN = {
    "alarm": "alarms and wake-up times", "audio": "audio and volume controls",
    "calendar": "calendar events, meetings and reminders", "cooking": "recipes and cooking",
    "datetime": "dates, times and time zones", "email": "reading and sending email",
    "general": "chit-chat, jokes and general questions to the assistant",
    "iot": "smart home devices such as lights, plugs and vacuums", "lists": "to-do and shopping lists",
    "music": "music settings and preferences", "news": "news updates",
    "play": "playing music, podcasts, radio, audiobooks or games",
    "qa": "factual questions, definitions, math and currency", "recommendation": "recommendations for places, "
    "movies and events", "social": "social media posts and complaints to companies",
    "takeaway": "takeaway food orders", "transport": "trains, taxis, tickets and traffic",
    "weather": "weather forecasts",
}


def massive(ctx: Ctx):
    files = ["train.jsonl", "validation.jsonl", "test.jsonl"]
    irows = read_many("SetFit/amazon_massive_intent_en-US", files)
    srows = read_many("SetFit/amazon_massive_scenario_en-US", files)
    scen_by_text = {}
    for r in srows:
        scen_by_text.setdefault(r["text"], r["label_text"])
    intents = sorted({r["label_text"] for r in irows})
    scens = sorted(MASSIVE_SCEN)
    ilabels = []
    for n in intents:
        sc = n.split("_")[0]
        rest = readable(n[len(sc) + 1:]) if n.startswith(sc + "_") else readable(n)
        ilabels.append(Label(readable(n), [f"{sc}: {rest}", readable(n)], readable(n), key=n, group=sc))
    slabels = [Label(s, [MASSIVE_SCEN[s]], MASSIVE_SCEN[s], group=s) for s in scens]
    rng = ctx.rng
    rows = [r for r in irows if r["text"] in scen_by_text]
    rows = dedupe_rows(rows, lambda r: r["text"], lambda r: r["label_text"])
    s_instr = ["Which domain does this voice command belong to?", "Which app should handle this request?",
               "What area is the user's request in?", "Route this command to the right domain.",
               "Which assistant skill is needed?"]
    i_instr = ["What does the user want the assistant to do?", "Which intent does this command express?",
               "Classify the voice command.", "What action is requested?", "Map this utterance to an intent."]
    for r in balanced_order(rng, rows, lambda r: r["label_text"], 0.7):
        text = clean(r["text"])
        gi = intents.index(r["label_text"])
        gs = scens.index(scen_by_text[r["text"]]) if scen_by_text[r["text"]] in scens else None
        qs = []
        u = rng.random()
        if u < 0.55 or gs is None:
            qs.append(make_choice(rng, ilabels, gi, pick_instr(rng, i_instr, GENERIC_CHOICE), full_prob=0.06,
                                  other_prob=0.04))
        else:
            qs.append(make_choice(rng, slabels, gs, pick_instr(rng, s_instr, GENERIC_CHOICE), full_prob=0.15,
                                  other_prob=0.04))
        if rng.random() < 0.3 and gs is not None:
            qs.append(make_choice(rng, slabels, gs, pick_instr(rng, s_instr, GENERIC_CHOICE)))
        if rng.random() < 0.45:
            if rng.random() < 0.5 and gs is not None:
                qs.append(derived_noul(rng, slabels, gs, ["This request is about {x}.", "The user needs help with {x}.",
                                                          "Is this command about {x}?"],
                                       ["This request is not about {x}."], p_other_group=0.0))
            else:
                qs.append(derived_noul(rng, ilabels, gi, ["The user's intent is: {x}.", "Intent: {x}.",
                                                          "The command asks for {x}."],
                                       ["The user's intent is not {x}."], p_other_group=0.7))
        yield Example(text, qs[:3])


BITEXT_CAT = {
    "ACCOUNT": ("accounts", ["account team: sign-up, login, profile changes, deleting accounts",
                             "account creation, editing, deletion and passwords"]),
    "CANCEL": ("cancellations", ["cancellation fees and policies"]),
    "CONTACT": ("contact", ["requests to reach customer service or a human agent"]),
    "DELIVERY": ("delivery", ["delivery options and delivery times"]),
    "FEEDBACK": ("feedback", ["complaints and reviews"]),
    "INVOICE": ("billing", ["invoices and billing documents", "invoices, bills"]),
    "ORDER": ("orders", ["placing, changing, cancelling or tracking orders"]),
    "PAYMENT": ("payments", ["payment methods and payment problems"]),
    "REFUND": ("refunds", ["refund policy, refund requests and refund status"]),
    "SHIPPING": ("shipping", ["shipping addresses"]),
    "SUBSCRIPTION": ("subscriptions", ["newsletter subscriptions"]),
}
_FILL = {
    "Order Number": lambda r: "#" + str(r.randint(10000, 99999999)),
    "Invoice Number": lambda r: "INV-" + str(r.randint(1000, 999999)),
    "Person Name": lambda r: r.choice(["Anna", "John Smith", "Maria", "Ahmed", "Li Wei", "Carlos", "Emma"]),
    "Client First Name": lambda r: r.choice(["Anna", "John", "Maria", "Ahmed", "Wei", "Carlos"]),
    "Client Last Name": lambda r: r.choice(["Smith", "Garcia", "Chen", "Yilmaz", "Novak"]),
    "Account Type": lambda r: r.choice(["premium", "business", "standard", "free", "pro"]),
    "Account Category": lambda r: r.choice(["premium", "business", "standard", "gold", "platinum"]),
    "Delivery City": lambda r: r.choice(["Berlin", "Chicago", "Istanbul", "Lyon", "Osaka", "Leeds"]),
    "Delivery Country": lambda r: r.choice(["Germany", "Canada", "Spain", "Brazil", "Japan", "Turkey"]),
    "Refund Amount": lambda r: "$" + str(r.randint(5, 900)),
    "Currency Symbol": lambda r: r.choice(["$", "EUR ", "GBP "]),
    "Money Amount": lambda r: "$" + str(r.randint(5, 900)),
}


def _bitext_fill(rng: random.Random, text: str) -> str:
    def sub(m):
        k = m.group(1).strip()
        f = _FILL.get(k)
        return f(rng) if f else k.lower()
    return re.sub(r"\{\{([^}]*)\}\}", sub, text)


def bitext(ctx: Ctx):
    rows = read_rows("bitext/Bitext-customer-support-llm-chatbot-training-dataset",
                     "Bitext_Sample_Customer_Support_Training_Dataset_27K_responses-v11.csv")
    rng = ctx.rng
    cats = sorted(BITEXT_CAT)
    intents = sorted({r["intent"] for r in rows})
    int_cat = {r["intent"]: r["category"] for r in rows}
    clabels = [Label(BITEXT_CAT[c][0], BITEXT_CAT[c][1], BITEXT_CAT[c][1][0].split(":")[0], key=BITEXT_CAT[c][0],
                     group=c) for c in cats]
    ilabels = [intent_label(i, group=int_cat[i]) for i in intents]
    rows = [dict(r, _t=_bitext_fill(rng, clean(r["instruction"]))) for r in rows]
    rows = dedupe_rows(rows, lambda r: r["_t"], lambda r: r["intent"])
    c_instr = ["Which team should handle this request?", "Route this ticket to the right department.",
               "Which support queue does this message belong to?", "Assign this customer message to a team.",
               "Which department is responsible for this?", "Who should pick up this ticket?"]
    i_instr = INTENT_INSTR
    for r in balanced_order(rng, rows, lambda r: r["intent"], 0.5):
        text = r["_t"]
        gi, gc = intents.index(r["intent"]), cats.index(r["category"])
        qs = []
        if rng.random() < 0.5:
            qs.append(make_choice(rng, clabels, gc, pick_instr(rng, c_instr, GENERIC_CHOICE), full_prob=0.3,
                                  other_prob=0.04))
        else:
            qs.append(make_choice(rng, ilabels, gi, pick_instr(rng, i_instr, GENERIC_CHOICE, thing="message"),
                                  full_prob=0.15, other_prob=0.04))
        if rng.random() < 0.45:
            qs.append(derived_noul(rng, ilabels, gi, INTENT_POS, INTENT_NEG, p_other_group=0.7, thing="message"))
        if rng.random() < 0.2:
            qs.append(derived_noul(rng, clabels, gc, ["This ticket should go to the {x} team.",
                                                     "The {x} team should handle this.",
                                                     "Is this a matter for {x}?"],
                                   ["The {x} team is not responsible for this."], p_other_group=0.0))
        yield Example(text, qs[:3])


# ---------------------------------------------------------------- emotion / stance / abuse

EMO_INSTR = ["Which emotion does the writer express?", "How does the author feel?",
             "What is the dominant emotion in this {thing}?", "Detect the emotion.",
             "Pick the emotion that best matches the {thing}.", "What mood is the writer in?"]
EMO_POS = ["The writer is {x}.", "The author of this {thing} is {x}.", "In this {thing}, the writer seems {x}.",
           "Is the writer {x}?", "The {thing} shows that the author is {x}."]
EMO_NEG = ["The writer is not {x}.", "The author does not seem {x}."]


def emotion(ctx: Ctx):
    rows = read_many("dair-ai/emotion", [f"split/{s}-00000-of-00001.parquet" for s in ("train", "validation", "test")])
    E = [("sadness", "sad", "feeling sad, down or hurt", "neg"), ("joy", "happy", "feeling happy or content", "pos"),
         ("love", "affectionate", "feeling love, affection or tenderness", "pos"),
         ("anger", "angry", "feeling angry, annoyed or irritated", "neg"),
         ("fear", "afraid", "feeling scared, anxious or nervous", "neg"),
         ("surprise", "surprised", "feeling surprised or amazed", "amb")]
    labels = [Label(n, [d], p, group=g) for n, p, d, g in E]
    yield from classify(ctx, rows, lambda r: clean(r["text"]), lambda r: int(r["label"]), labels,
                        instr=EMO_INSTR, noul_pos=EMO_POS, noul_neg=EMO_NEG, thing="text", alpha=0.5,
                        other_prob=0.0, p_other_group=0.85)


def tweet_emotion(ctx: Ctx):
    rows = read_many("cardiffnlp/tweet_eval", [f"emotion/{s}-00000-of-00001.parquet" for s in ("train", "validation", "test")])
    E = [("anger", "angry", "anger or rage", "neg"), ("joy", "joyful", "joy or happiness", "pos"),
         ("optimism", "optimistic", "optimism or hope", "pos"), ("sadness", "sad", "sadness or grief", "neg")]
    labels = [Label(n, [d], p, group=g) for n, p, d, g in E]
    yield from classify(ctx, rows, lambda r: _tweet_clean(r["text"]), lambda r: int(r["label"]), labels,
                        instr=EMO_INSTR, noul_pos=EMO_POS, noul_neg=EMO_NEG, thing="tweet", alpha=0.5,
                        other_prob=0.0, p_other_group=0.9)


GO = [("admiration", "admiring", "pos"), ("amusement", "amused", "pos"), ("anger", "angry", "neg"),
      ("annoyance", "annoyed", "neg"), ("approval", "approving", "pos"), ("caring", "caring", "pos"),
      ("confusion", "confused", "amb"), ("curiosity", "curious", "amb"), ("desire", "full of desire", "pos"),
      ("disappointment", "disappointed", "neg"), ("disapproval", "disapproving", "neg"),
      ("disgust", "disgusted", "neg"), ("embarrassment", "embarrassed", "neg"), ("excitement", "excited", "pos"),
      ("fear", "afraid", "neg"), ("gratitude", "grateful", "pos"), ("grief", "grieving", "neg"),
      ("joy", "joyful", "pos"), ("love", "loving", "pos"), ("nervousness", "nervous", "neg"),
      ("optimism", "optimistic", "pos"), ("pride", "proud", "pos"), ("realization", "having a realization", "amb"),
      ("relief", "relieved", "pos"), ("remorse", "remorseful", "neg"), ("sadness", "sad", "neg"),
      ("surprise", "surprised", "amb"), ("neutral", "emotionally neutral", "neutral")]


def go_emotions(ctx: Ctx):
    rows = read_many("google-research-datasets/go_emotions",
                     [f"simplified/{s}-00000-of-00001.parquet" for s in ("train", "validation", "test")])
    labels = [Label(n, [f"the writer is {p}"] if n != "neutral" else ["no particular emotion"], p, group=g)
              for n, p, g in GO]
    rng = ctx.rng
    for r in rows:
        r["_l"] = sorted(int(x) for x in r["labels"])
        r["_t"] = clean(r["text"]).replace("[NAME]", rng.choice(["Alex", "Sam", "Jordan", "Chris", "Taylor"]))
    rows = dedupe_rows(rows, lambda r: r["_t"], lambda r: tuple(r["_l"]))
    for r in balanced_order(rng, rows, lambda r: r["_l"][0], 0.4):
        gold = r["_l"]
        qs = []
        if len(gold) == 1 and rng.random() < 0.7:
            qs.append(make_choice(rng, labels, gold[0], pick_instr(rng, EMO_INSTR, GENERIC_CHOICE, thing="comment"),
                                  full_prob=0.1, other_prob=0.0))
        n_noul = 1 if qs else rng.choice([1, 1, 2])
        for _ in range(n_noul):
            # multi-label rows give real negatives: any emotion not annotated (prefer other valence)
            qs.append(derived_noul(rng, labels, gold, EMO_POS, EMO_NEG, p_other_group=0.85, thing="comment"))
        yield Example(r["_t"], qs[:3])


STANCE = {"abortion": "the legalization of abortion", "atheism": "atheism",
          "climate": "the view that climate change is a serious concern", "feminist": "feminism",
          "hillary": "Hillary Clinton"}


def tweet_stance(ctx: Ctx):
    rng = ctx.rng
    rows = []
    for t in STANCE:
        for s in ("train", "validation", "test"):
            for r in read_rows("cardiffnlp/tweet_eval", f"stance_{t}/{s}-00000-of-00001.parquet"):
                rows.append(dict(r, _target=t))
    rows = dedupe_rows(rows, lambda r: r["text"] + r["_target"], lambda r: r["label"])
    instr = ["What is the author's stance toward {target}?", "Does this tweet support or oppose {target}?",
             "How does the writer feel about {target}?", "Stance on {target}?",
             "Classify the tweet's position on {target}."]
    for r in balanced_order(rng, rows, lambda r: (r["_target"], r["label"]), 0.7):
        tgt = STANCE[r["_target"]]
        labels = [Label("none", [f"neutral or unrelated to {tgt}", "no clear stance"], f"no position on {tgt}",
                        key=rng.choice(["none", "neutral", "neither"])),
                  Label("against", [f"opposes {tgt}", "is against it"], f"is against {tgt}"),
                  Label("favor", [f"supports {tgt}", "is in favor of it"], f"supports {tgt}",
                        key=rng.choice(["favor", "support", "for"]))]
        gold = int(r["label"])
        qs = []
        if rng.random() < 0.8:
            qs.append(make_choice(rng, labels, gold, pick_instr(rng, instr, (), target=tgt), n=3))
        if not qs or rng.random() < 0.4:
            qs.append(derived_noul(rng, labels, gold, ["The author {x}.", "This tweet {x}.",
                                                        "The writer's position: {x}."],
                                   p_other_group=0.0, use_desc=0.0))
        yield Example(_tweet_clean(r["text"]), qs)


def hate_offensive(ctx: Ctx):
    rows = read_many("SetFit/hate_speech_offensive", ["train.jsonl", "test.jsonl"])
    labels = [Label("hate speech", ["attacks or demeans a group based on identity"], "hate speech", key="hate"),
              Label("offensive language", ["offensive or vulgar but not hate speech"], "offensive language",
                    key="offensive"),
              Label("neither", ["neither hateful nor offensive", "harmless"], "harmless content",
                    key=None)]
    rng = ctx.rng
    instr = ["Is this tweet hate speech, offensive, or neither?", "Moderate this tweet.",
             "How should this tweet be classified for moderation?", "Classify the tweet's content.",
             "Which moderation label applies?"]

    def extra(r, rng):
        g = int(r["label"])
        if rng.random() < 0.35:
            s, truth = weighted(rng, [("This tweet contains offensive or abusive language.", g in (0, 1)),
                                      ("This tweet is harmless.", g == 2),
                                      ("This tweet contains hate speech targeting a group.", g == 0)],
                                [0.45, 0.3, 0.25])
            return [noul(s, truth)]
        return []
    rows = [dict(r, _t=_tweet_clean(r["text"])) for r in rows]
    yield from classify(ctx, rows, lambda r: r["_t"], lambda r: int(r["label"]), labels, instr=instr,
                        noul_pos=["This tweet is best described as {x}."], thing="tweet", alpha=0.3,
                        other_prob=0.0, p_noul=0.0, p_second=0.0, extra=extra)


# ---------------------------------------------------------------- multiple-choice QA (option texts vary)

def _ans_noul(rng: random.Random, q: str, answers: list[str], gold: int) -> tuple:
    truth = rng.random() < 0.5
    a = answers[gold] if truth else rng.choice([x for i, x in enumerate(answers) if i != gold])
    q = q.strip()
    t = weighted(rng, [f"The correct answer to \"{q}\" is \"{a}\".", f"\"{a}\" answers the question: {q}",
                       f"Question: {q} Answer: {a}. This answer is correct.",
                       f"{q} -> {a}"], [0.35, 0.25, 0.25, 0.15])
    return noul(t, truth)


def commonsense_qa(ctx: Ctx):
    rows = read_many("tau/commonsense_qa", ["data/train-00000-of-00001.parquet",
                                            "data/validation-00000-of-00001.parquet"])
    rng = ctx.rng
    rng.shuffle(rows)
    instr = ["Which answer is correct?", "Pick the best answer.", "Answer the question.",
             "Choose the most sensible answer.", "What is the right answer?"]
    for r in rows:
        ch = r["choices"]
        answers = [clean(t) for t in ch["text"]]
        labs = list(ch["label"])
        if r["answerKey"] not in labs:
            continue
        g = labs.index(r["answerKey"])
        q = clean(r["question"])
        qs = [make_mcq(rng, answers, g, pick_instr(rng, instr))]
        if rng.random() < 0.3:
            qs.append(_ans_noul(rng, q, answers, g))
        yield Example(q, qs)


def social_iqa(ctx: Ctx):
    rows = read_many("allenai/social_i_qa", ["default/train/0000.parquet", "default/validation/0000.parquet"],
                     revision="refs/convert/parquet")
    rng = ctx.rng
    rng.shuffle(rows)
    for r in rows:
        answers = [clean(r["answerA"]), clean(r["answerB"]), clean(r["answerC"])]
        g = int(r["label"]) - 1
        q = clean(r["question"])
        instr = q if rng.random() < 0.8 else f"Based on the situation: {q[:1].lower() + q[1:]}"
        qs = [make_mcq(rng, answers, g, instr)]
        if rng.random() < 0.3:
            qs.append(_ans_noul(rng, q, answers, g))
        yield Example(clean(r["context"]), qs, key=clean(r["context"]))


def race(ctx: Ctx):
    rows = read_many("ehovy/race", ["middle/train-00000-of-00001.parquet", "middle/validation-00000-of-00001.parquet",
                                    "middle/test-00000-of-00001.parquet", "high/validation-00000-of-00001.parquet",
                                    "high/test-00000-of-00001.parquet"])
    rng = ctx.rng
    by_art: dict[str, list] = {}
    for r in rows:
        by_art.setdefault(r["article"], []).append(r)
    arts = list(by_art)
    rng.shuffle(arts)
    for art in arts:
        state = clean(art)
        if len(state) > 2000:  # truncation could cut the evidence
            continue
        qrows = by_art[art]
        rng.shuffle(qrows)
        qs = []
        for r in qrows[:rng.choice([1, 2, 2, 3])]:
            answers = [clean(o) for o in r["options"]]
            g = "ABCD".index(r["answer"])
            q = clean(r["question"])
            if "_" in q:
                instr = weighted(rng, [f"Fill in the blank: {q}", f"Complete the statement: {q}", q], [0.4, 0.3, 0.3])
            else:
                instr = q
            qs.append(make_mcq(rng, answers, g, instr))
        if qs:
            yield Example(state, qs)


def hellaswag(ctx: Ctx):
    rows = read_rows("Rowan/hellaswag", "data/validation-00000-of-00001.parquet")
    rng = ctx.rng
    rng.shuffle(rows)

    def hs(t):
        return clean(re.sub(r"\[(header|title|step|substeps)\]", " ", t))
    instr = ["How does the text most likely continue?", "Which ending is most plausible?",
             "Pick the most likely continuation.", "What happens next?", "Choose the sentence that best completes the passage."]
    for r in rows:
        ends = [hs(e) for e in r["endings"]]
        state = hs(r["ctx"])
        if rng.random() < 0.4:
            state = f"{r['activity_label']}: {state}"
        yield Example(state, [make_mcq(rng, ends, int(r["label"]), pick_instr(rng, instr))])


def sciq(ctx: Ctx):
    rows = read_many("allenai/sciq", [f"data/{s}-00000-of-00001.parquet" for s in ("train", "validation", "test")])
    rng = ctx.rng
    rng.shuffle(rows)
    for r in rows:
        answers = [clean(r["correct_answer"]), clean(r["distractor1"]), clean(r["distractor2"]), clean(r["distractor3"])]
        if len({a.lower() for a in answers}) < 4:
            continue
        q = clean(r["question"])
        sup = clean(r["support"])
        if sup and len(sup) <= 2000 and rng.random() < 0.8:
            state, instr = sup, q
        else:
            state, instr = q, weighted(rng, ["Which answer is correct?", "Answer the science question.",
                                             "Pick the correct answer."], [1, 1, 1])
        qs = [make_mcq(rng, answers, 0, instr)]
        if rng.random() < 0.25:
            qs.append(_ans_noul(rng, q, answers, 0))
        yield Example(state, qs)


def copa(ctx: Ctx):
    rows = read_many("aps/super_glue", ["copa/train-00000-of-00001.parquet", "copa/validation-00000-of-00001.parquet"])
    rng = ctx.rng
    rng.shuffle(rows)
    for r in rows:
        instr = rng.choice(["What was the cause?", "Why did this happen?", "What is the more likely cause?"]) \
            if r["question"] == "cause" else \
            rng.choice(["What happened as a result?", "What is the most likely effect?", "What happened next?"])
        yield Example(clean(r["premise"]), [make_mcq(rng, [clean(r["choice1"]), clean(r["choice2"])],
                                                     int(r["label"]), instr)])


SOURCES = {
    "ag_news": ag_news, "dbpedia_14": dbpedia, "yahoo_answers_topics": yahoo, "20_newsgroups": newsgroups,
    "tweet_topic": tweet_topic, "fin_news_topic": fin_topic, "banking77": banking77, "clinc_oos": clinc,
    "massive": massive, "bitext_support": bitext, "emotion": emotion, "tweet_eval_emotion": tweet_emotion,
    "go_emotions": go_emotions, "tweet_eval_stance": tweet_stance, "hate_speech_offensive": hate_offensive,
    "commonsense_qa": commonsense_qa, "social_i_qa": social_iqa, "race": race, "hellaswag": hellaswag,
    "sciq": sciq, "copa": copa,
}