"""Datasets 1 to 4 and the training files of models 1 to 3, built from the text of our own published pages (2026-10-30).

    python hf-work/gen_derived.py
Reads hf-work/site_content.json (extracted from aitrainingdatabroker.com's page sources by extract_src.ts) and writes
datasets/<name>/data/train.jsonl and models/<name>/data/{train,test}.jsonl. Every row is our own text; nothing is from a customer.
"""

import json
import random
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
SITE = "https://aitrainingdatabroker.com"
src = json.load(open(Path(__file__).parent / "site_content.json", encoding="utf-8"))
pages = src["pages"]
industries = [p for p in pages if p["kind"] == "industry"]
records_pages = [p for p in pages if p["kind"] == "data"]
rng = random.Random(20261030)


def write(kind, name, split, rows):
    out = ROOT / kind / name / "data"
    out.mkdir(parents=True, exist_ok=True)
    with open(out / f"{split}.jsonl", "w", encoding="utf-8", newline="\n") as f:
        for r in rows:
            f.write(json.dumps(r, ensure_ascii=False) + "\n")
    print(f"{kind}/{name}/{split}: {len(rows)} rows")


def url(p):
    return SITE + p["path"]


# ---- dataset 1: what a company in each industry can consider licensing and what stays out ----
rows = []
for p in industries:
    for kind, items in (("can_be_considered", p["records"]), ("stays_out", p["keepsOut"])):
        for it in items:
            rows.append({"industry": p["name"], "industry_slug": p["slug"], "kind": kind, "label": it["label"],
                         "text": it["text"], "source_url": url(p)})
for i, r in enumerate(rows, 1):
    r["id"] = f"ind-{i:03d}"
write("datasets", "company-data-licensability-by-industry", "train", [{"id": r.pop("id"), **r} for r in rows])
ind_rows = rows

# ---- dataset 2: the same for 33 record types ----
rows = []
for p in records_pages:
    for kind, items in (("what_it_holds", p["records"]), ("limits", p["keepsOut"])):
        for it in items:
            rows.append({"record_type": p["name"], "record_type_slug": p["slug"], "kind": kind, "label": it["label"],
                         "text": it["text"], "source_url": url(p)})
write("datasets", "company-record-types-and-limits", "train", [{"id": f"rec-{i:03d}", **r} for i, r in enumerate(rows, 1)])
rec_rows = rows

# ---- dataset 3: questions and answers, by industry, record type and guide ----
rows = []
for p in pages:
    for q in p["faq"]:
        rows.append({"topic_kind": "industry" if p["kind"] == "industry" else "record_type", "topic": p["name"],
                     "question": q["q"], "answer": q["a"], "source_url": url(p)})
for g in src["guides"]:
    for q in g["faq"]:
        rows.append({"topic_kind": f"{g['area']}_guide", "topic": g["h1"], "question": q["q"], "answer": q["a"],
                     "source_url": SITE + g["url"]})
write("datasets", "company-data-licensing-faq-by-topic", "train", [{"id": f"faq-{i:03d}", **r} for i, r in enumerate(rows, 1)])
faq_rows = rows

# ---- dataset 4: glossary ----
rows = []
for t in src["glossary"]:
    more = t.get("more") or []
    rows.append({"term": t["term"], "definition": t["def"], "see_also": more[0] if len(more) > 0 else "",
                 "see_also_url": SITE + more[1] if len(more) > 1 else ""})
write("datasets", "data-licensing-glossary", "train", [{"id": f"gls-{i:03d}", **r} for i, r in enumerate(rows, 1)])

# ---- model 1: which industry page does a passage belong to (train on records, ledes and sections; test on its FAQ) ----
train, test = [], []
for p in industries:
    for it in p["records"] + p["keepsOut"]:
        train.append({"text": f"{it['label']}. {it['text']}", "label": p["slug"]})
    train.append({"text": p["lede"], "label": p["slug"]})
    for s in p["sections"]:
        for para in s["paragraphs"]:
            train.append({"text": para, "label": p["slug"]})
    for q in p["faq"]:
        test.append({"text": f"{q['q']} {q['a']}", "label": p["slug"]})
rng.shuffle(train)
write("models", "industry-page-classifier", "train", train)
write("models", "industry-page-classifier", "test", test)

# ---- model 2: can be considered, or stays out (train on industry pages, test on record-type pages: a different subject) ----
train = [{"text": f"{it['label']}. {it['text']}", "label": "can_be_considered" if k == "r" else "stays_out"}
         for p in industries for k, items in (("r", p["records"]), ("o", p["keepsOut"])) for it in items]
test = [{"text": f"{it['label']}. {it['text']}", "label": "can_be_considered" if k == "r" else "stays_out"}
        for p in records_pages for k, items in (("r", p["records"]), ("o", p["keepsOut"])) for it in items]
rng.shuffle(train)
write("models", "licensable-or-stays-out-classifier", "train", train)
write("models", "licensable-or-stays-out-classifier", "test", test)

# ---- model 3: which record-type page does a passage belong to (train on records, ledes, sections; test on FAQ) ----
train, test = [], []
for p in records_pages:
    for it in p["records"] + p["keepsOut"]:
        train.append({"text": f"{it['label']}. {it['text']}", "label": p["slug"]})
    train.append({"text": p["lede"], "label": p["slug"]})
    for s in p["sections"]:
        for para in s["paragraphs"]:
            train.append({"text": para, "label": p["slug"]})
    for q in p["faq"]:
        test.append({"text": f"{q['q']} {q['a']}", "label": p["slug"]})
rng.shuffle(train)
write("models", "record-type-page-classifier", "train", train)
write("models", "record-type-page-classifier", "test", test)