hypersteeldb 0.3.1

//! **Adversarial benchmark corpus generator** — a synthetic Pokémon league built specifically to induce
//! the failure modes the paper describes, with exact ground truth for every question.
//!
//!   `cargo run --release --bin gen_benchmark_corpus -- <out_dir> [n_docs] [seed]`
//!
//! ## What it is designed to break
//!
//! * **Supernode explosion.** A handful of hub entities (one venue, one type, one tournament, one
//!   ubiquitous species) appear in a large fraction of documents. Any agent that traverses *through* a hub
//!   fans out across thousands of irrelevant edges; set-algebra intersects them in one popcount.
//! * **Context-window bloat.** Answers require combining sets far larger than a context window — the
//!   superlative and join questions below cannot be answered by reading retrieved passages.
//! * **Silent similarity failure.** Every fact has near-duplicate **confounders** that are lexically
//!   almost identical but logically different: the negated twin, the hedged twin, the wrong region, the
//!   wrong quarter, and a numeric value placed just the other side of a threshold. Embedding retrieval
//!   ranks these together; only exact constraints separate them.
//! * **Ignored negation.** Permitted/banned pairs for the same species differ by one word, so an engine
//!   that drops `NOT` answers the opposite of the truth.
//!
//! ## Ground truth
//!
//! Structure is generated deterministically, so `questions.jsonl` ships the **exact** answer set for each
//! question (computed from the generator's own bookkeeping, not from a model), plus the IKL that produces
//! it. Prose variety is layered on top via paraphrase templates — language varies, facts do not.

use std::collections::BTreeMap;
use std::fmt::Write as _;
use std::path::PathBuf;

struct Rng(u64);
impl Rng {
    fn next(&mut self) -> u64 {
        self.0 = self.0.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
        self.0 >> 17
    }
    fn below(&mut self, n: usize) -> usize {
        (self.next() % n.max(1) as u64) as usize
    }
    fn pick<'a, T>(&mut self, xs: &'a [T]) -> &'a T {
        &xs[self.below(xs.len())]
    }
    fn range(&mut self, lo: i64, hi: i64) -> i64 {
        lo + (self.next() % (hi - lo).max(1) as u64) as i64
    }
    /// true with probability `pct`/100
    fn chance(&mut self, pct: u64) -> bool {
        self.next() % 100 < pct
    }
}

struct Species {
    name: &'static str,
    primary: &'static str,
    weight_kg: f64,
}

const SPECIES: &[Species] = &[
    Species { name: "Registeel", primary: "steel", weight_kg: 205.0 },
    Species { name: "Regirock", primary: "rock", weight_kg: 230.0 },
    Species { name: "Regice", primary: "ice", weight_kg: 175.0 },
    Species { name: "Metagross", primary: "steel", weight_kg: 550.0 },
    Species { name: "Scizor", primary: "steel", weight_kg: 118.0 },
    Species { name: "Aggron", primary: "steel", weight_kg: 360.0 },
    Species { name: "Lucario", primary: "steel", weight_kg: 54.0 },
    Species { name: "Charizard", primary: "fire", weight_kg: 90.5 },
    Species { name: "Blastoise", primary: "water", weight_kg: 85.5 },
    Species { name: "Venusaur", primary: "grass", weight_kg: 100.0 },
    Species { name: "Pikachu", primary: "electric", weight_kg: 6.0 },
    Species { name: "Snorlax", primary: "normal", weight_kg: 460.0 },
    Species { name: "Gengar", primary: "ghost", weight_kg: 40.5 },
    Species { name: "Dragonite", primary: "dragon", weight_kg: 210.0 },
    Species { name: "Tyranitar", primary: "rock", weight_kg: 202.0 },
    Species { name: "Garchomp", primary: "dragon", weight_kg: 95.0 },
    Species { name: "Milotic", primary: "water", weight_kg: 162.0 },
    Species { name: "Salamence", primary: "dragon", weight_kg: 102.6 },
];

/// The hub species: deliberately over-represented so it becomes a supernode.
const HUB_SPECIES: &str = "Registeel";
/// The hub venue: hosts a large fraction of all battles.
const HUB_VENUE: (&str, &str) = ("kanto", "Indigo Plateau");
/// The hub tournament.
const HUB_TOURNAMENT: &str = "Indigo Invitational";

const WEAKNESS: &[(&str, &[&str])] = &[
    ("steel", &["fire", "fighting", "ground"]),
    ("rock", &["water", "grass", "fighting", "ground", "steel"]),
    ("ice", &["fire", "fighting", "rock", "steel"]),
    ("fire", &["water", "rock", "ground"]),
    ("water", &["electric", "grass"]),
    ("grass", &["fire", "ice", "flying", "bug"]),
    ("electric", &["ground"]),
    ("normal", &["fighting"]),
    ("ghost", &["ghost", "dark"]),
    ("dragon", &["ice", "dragon", "fairy"]),
];

const MOVES: &[(&str, &str, i64)] = &[
    ("Flash Cannon", "steel", 80),
    ("Iron Head", "steel", 80),
    ("Flamethrower", "fire", 90),
    ("Surf", "water", 90),
    ("Thunderbolt", "electric", 90),
    ("Earthquake", "ground", 100),
    ("Ice Beam", "ice", 90),
    ("Shadow Ball", "ghost", 80),
    ("Dragon Claw", "dragon", 80),
    ("Stone Edge", "rock", 100),
];

const PLACES: &[(&str, &str)] = &[
    ("kanto", "Indigo Plateau"),
    ("kanto", "Pewter City"),
    ("kanto", "Cerulean City"),
    ("kanto", "Saffron City"),
    ("johto", "Violet City"),
    ("johto", "Ecruteak City"),
    ("johto", "Blackthorn City"),
    ("hoenn", "Rustboro City"),
    ("hoenn", "Mauville City"),
    ("hoenn", "Sootopolis City"),
    ("sinnoh", "Oreburgh City"),
    ("sinnoh", "Snowpoint City"),
];

const TRAINERS: &[&str] = &[
    "Brock Harrow", "Misty Cascade", "Sabrina Vale", "Blaine Ember", "Falkner Gale", "Morty Shade",
    "Clair Wyvern", "Roxanne Slate", "Wattson Volt", "Juan Tide", "Roark Quarry", "Candice Frost",
    "Iris Draco", "Marnie Dusk", "Nessa Reef", "Bea Strike", "Cynthia Ward", "Lance Aether",
    "Steven Stone", "Wallace Gale",
];

const TOURNAMENTS: &[&str] = &[
    "Indigo Invitational", "Silver Conference", "Ever Grande Cup", "Lily of the Valley Open",
    "Steel Type Masters", "Regional Doubles Series",
];

/// The last few trainers are the travelling circuit: they compete in every region, while everyone else is
/// tied to their home region. Drawing trainers uniformly instead makes the trainer-venue graph COMPLETE,
/// which collapses the 2-hop path family into a constant answer ("all trainers") that needs no traversal.
const N_CIRCUIT: usize = 4;

/// Which species may be entered in a given tournament. Drawing species uniformly across all battles makes
/// EVERY species co-occur with every other, so the co-occurrence family degenerates into "all species except
/// itself" — answerable without traversing anything. Regulation-style eligibility pools make shared-battle
/// structure real: two species co-occur only if some tournament admits both.
fn tournament_pool(tour: &str) -> Vec<&'static Species> {
    if tour == "Steel Type Masters" {
        return SPECIES.iter().filter(|s| s.primary == "steel").collect();
    }
    let ti = TOURNAMENTS.iter().position(|t| *t == tour).unwrap_or(0);
    // overlapping windows so pools share some species but not all; the hub species is always eligible,
    // which is what keeps it a supernode.
    SPECIES
        .iter()
        .enumerate()
        .filter(|(i, s)| {
            s.name == HUB_SPECIES || (i + ti) % TOURNAMENTS.len() < 2
        })
        .map(|(_, s)| s)
        .collect()
}

fn home_region(trainer_idx: usize) -> Option<&'static str> {
    if trainer_idx >= TRAINERS.len() - N_CIRCUIT {
        return None; // circuit trainer: appears anywhere
    }
    let regions = ["kanto", "johto", "hoenn", "sinnoh"];
    Some(regions[trainer_idx % regions.len()])
}

/// Paraphrase frames for battle prose. Lexical variety matters: if every document shares a sentence
/// template, keyword retrieval succeeds for the wrong reason and the benchmark measures nothing.
const BATTLE_FRAMES: &[&str] = &[
    "{a} faced {b} at {venue} in the {region} region during the {tour} on {date}. {a} led with {sa} while {b} answered with {sb}. The deciding move was {mv}. The match ran {mins} minutes and {winner} took the win.",
    "On {date} the {tour} staged a bout between {a} and {b} at {venue} ({region}). {sa} opened for {a}; {b} countered with {sb}. {mv} settled it after {mins} minutes, giving {winner} the result.",
    "Match record — {tour}, {venue}, {region}, {date}. Competitors: {a} and {b}. Lead Pokémon: {sa} and {sb} respectively. Outcome decided by {mv} at the {mins}-minute mark. Winner: {winner}.",
    "{b} met {a} at {venue} during the {tour} ({date}, {region} region). The exchange turned on {mv}; {sb} could not answer {sa}. After {mins} minutes the win went to {winner}.",
    "Reported from {venue} in {region}: {a} and {b} contested a {tour} fixture on {date}. {a} fielded {sa}, {b} fielded {sb}, and {mv} proved decisive. Duration {mins} minutes; {winner} advanced.",
];

/// A generated battle, retained so ground truth can be computed exactly.
#[derive(Clone)]
struct Battle {
    doc: usize,
    a: String,
    b: String,
    winner: String,
    region: String,
    venue: String,
    tour: String,
    year: i64,
    quarter: i64,
    minutes: i64,
    species_a: String,
    species_b: String,
    move_name: String,
}

#[derive(Clone)]
struct Ruling {
    doc: usize,
    species: String,
    series: i64,
    year: i64,
    /// true = banned (negated permission), false = explicitly permitted
    banned: bool,
    /// hedged rulings must not count as asserted either way
    hedged: bool,
}

#[derive(Clone)]
struct Survey {
    doc: usize,
    region: String,
    venue: String,
    species: String,
    elevation_m: i64,
    temp_c: i64,
}

fn slug(s: &str) -> String {
    let mut out = String::new();
    let mut dash = false;
    for c in s.chars() {
        if c.is_alphanumeric() {
            for l in c.to_lowercase() {
                out.push(l);
            }
            dash = false;
        } else if !dash && !out.is_empty() {
            out.push('-');
            dash = true;
        }
    }
    out.trim_end_matches('-').to_string()
}

fn main() -> std::io::Result<()> {
    let args: Vec<String> = std::env::args().collect();
    let out = PathBuf::from(args.get(1).cloned().unwrap_or_else(|| "benchmark_corpus".to_string()));
    let n_docs: usize = args.get(2).and_then(|s| s.parse().ok()).unwrap_or(10_000);
    let seed: u64 = args.get(3).and_then(|s| s.parse().ok()).unwrap_or(0x9E3779B97F4A7C15);
    std::fs::create_dir_all(&out)?;
    let mut rng = Rng(seed);

    let mut battles: Vec<Battle> = Vec::new();
    let mut rulings: Vec<Ruling> = Vec::new();
    let mut surveys: Vec<Survey> = Vec::new();
    let mut doc = 0usize;

    while doc < n_docs {
        match doc % 10 {
            // 50% battle reports — the relational bulk, and the supernode carrier
            0..=4 => {
                // hub bias: 35% of battles land at the hub venue, 30% in the hub tournament
                let (region, venue) = if rng.chance(35) { HUB_VENUE } else { *rng.pick(PLACES) };
                let tour = if rng.chance(30) { HUB_TOURNAMENT } else { *rng.pick(TOURNAMENTS) };
                // eligible = trainers based in this region, plus the circuit; 20% of slots go to the circuit
                // so cross-region venue pairs are linked by a genuinely small set of trainers.
                let draw = |rng: &mut Rng| -> String {
                    let circuit = rng.chance(20);
                    let pool: Vec<&str> = TRAINERS
                        .iter()
                        .enumerate()
                        .filter(|(i, _)| match home_region(*i) {
                            None => circuit,
                            Some(r) => !circuit && r == region,
                        })
                        .map(|(_, t)| *t)
                        .collect();
                    if pool.is_empty() { rng.pick(TRAINERS).to_string() } else { rng.pick(&pool).to_string() }
                };
                let a = draw(&mut rng);
                let b = loop {
                    let c = draw(&mut rng);
                    if c != a {
                        break c;
                    }
                };
                // hub species appears far more often than any other, but both entries must be eligible for
                // the tournament, so co-occurrence is constrained by the regulation pools
                let pool = tournament_pool(tour);
                let pick_sp = |rng: &mut Rng| -> String {
                    if pool.is_empty() { rng.pick(SPECIES).name.to_string() } else { rng.pick(&pool).name.to_string() }
                };
                let species_a = if rng.chance(30) { HUB_SPECIES.to_string() } else { pick_sp(&mut rng) };
                let species_b = pick_sp(&mut rng);
                let (mv, _mt, _p) = *rng.pick(MOVES);
                let year = rng.range(2025, 2027);
                let month = rng.range(1, 13);
                let day = rng.range(1, 28);
                let minutes = rng.range(6, 41);
                let winner = if rng.chance(50) { a.clone() } else { b.clone() };
                let bt = Battle {
                    doc,
                    a: a.clone(),
                    b: b.clone(),
                    winner: winner.clone(),
                    region: region.to_string(),
                    venue: venue.to_string(),
                    tour: tour.to_string(),
                    year,
                    quarter: (month - 1) / 3 + 1,
                    minutes,
                    species_a: species_a.clone(),
                    species_b: species_b.clone(),
                    move_name: mv.to_string(),
                };
                let frame = rng.pick(BATTLE_FRAMES);
                let date = format!("{year}-{month:02}-{day:02}");
                let prose = frame
                    .replace("{a}", &a)
                    .replace("{b}", &b)
                    .replace("{venue}", venue)
                    .replace("{region}", region)
                    .replace("{tour}", tour)
                    .replace("{date}", &date)
                    .replace("{sa}", &species_a)
                    .replace("{sb}", &species_b)
                    .replace("{mv}", mv)
                    .replace("{mins}", &minutes.to_string())
                    .replace("{winner}", &winner);
                let mut body = String::new();
                let _ = write!(
                    body,
                    "# Battle Report — {a} vs {b}\n\n- **Date:** {date}\n- **Venue:** {venue}\n- **Region:** {region}\n- **Tournament:** {tour}\n- **Duration:** {minutes} minutes\n- **Winner:** {winner}\n\n## Account\n\n{prose}\n"
                );
                std::fs::write(out.join(format!("{doc:06}_battle.md")), body)?;
                battles.push(bt);
            }
            // 20% regulations — permitted/banned/hedged CONFOUNDER triples on the same species+series
            5..=6 => {
                let series = rng.range(1, 5);
                let year = rng.range(2025, 2027);
                let sp = rng.pick(SPECIES).name.to_string();
                // three lexically near-identical rulings that differ only in logical force
                let kinds: [(bool, bool, &str); 3] = [
                    (true, false, "is not permitted in"),
                    (false, false, "is permitted in"),
                    (true, true, "may not be permitted in"),
                ];
                let (banned, hedged, phrase) = kinds[rng.below(3)];
                let team = rng.range(4, 7);
                let timer = rng.range(15, 31);
                let weight_cap = rng.range(100, 601);
                let mut body = String::new();
                let _ = write!(
                    body,
                    "# Competition Regulations — Series {series} {year}\n\n\
                     ## 1. Team composition\n\nEach trainer registers {team} Pokémon. The battle timer is {timer} minutes. \
                     A Pokémon heavier than {weight_cap} kg may not enter the doubles bracket.\n\n\
                     ## 2. Species ruling\n\n{sp} {phrase} Series {series} play for the {year} season.\n"
                );
                std::fs::write(out.join(format!("{doc:06}_regulation.md")), body)?;
                rulings.push(Ruling { doc, species: sp, series, year, banned, hedged });
            }
            // 20% habitat surveys — numeric thresholds + spatial, with straddling confounders
            7..=8 => {
                let (region, venue) = *rng.pick(PLACES);
                let sp = rng.pick(SPECIES);
                // elevations cluster right around 1000 m so a threshold question has near-misses
                let elevation_m = if rng.chance(60) { rng.range(950, 1051) } else { rng.range(5, 2400) };
                let temp_c = rng.range(-12, 36);
                let mut body = String::new();
                let _ = write!(
                    body,
                    "# Habitat Survey — {venue} ({region})\n\n\
                     - **Region:** {region}\n- **Elevation:** {elevation_m} m\n- **Mean temperature:** {temp_c} °C\n\n\
                     ## Observations\n\nThe team recorded {} near {venue} in the {region} region at {elevation_m} m elevation, \
                     mean temperature {temp_c} °C. The specimen weighs {} kg.\n",
                    sp.name, sp.weight_kg
                );
                std::fs::write(out.join(format!("{doc:06}_survey.md")), body)?;
                surveys.push(Survey { doc, region: region.to_string(), venue: venue.to_string(), species: sp.name.to_string(), elevation_m, temp_c });
            }
            // 10% type chart — the weakness edges the s-paths traverse
            _ => {
                let (ty, weak) = *rng.pick(WEAKNESS);
                let examples: Vec<&str> = SPECIES.iter().filter(|s| s.primary == ty).map(|s| s.name).take(4).collect();
                let mut body = String::new();
                let _ = write!(
                    body,
                    "# Type Chart — {ty}\n\n- **Type:** {ty}\n- **Weak to:** {}\n\n## Notes\n\n\
                     A {ty} type Pokémon takes increased damage from {} attacks. Recorded {ty} species include {}.\n",
                    weak.join(", "),
                    weak.join(", "),
                    if examples.is_empty() { "none in this registry".to_string() } else { examples.join(", ") }
                );
                std::fs::write(out.join(format!("{doc:06}_typechart.md")), body)?;
            }
        }
        doc += 1;
    }

    // ── gold test set: exact document-id answers computed from the generator's own bookkeeping ───────
    //
    // Every question carries the full gold document set, so ANY system (this engine, a tool-using agent,
    // dense RAG, GraphRAG, text-to-SQL) can be scored with set precision/recall/F1 by the companion
    // `score_benchmark` binary. Questions are emitted per family across many instances rather than as a
    // handful of examples, so per-category scores are meaningful.
    let mut questions: Vec<serde_json::Value> = Vec::new();
    let per_family: usize = std::env::var("BENCH_PER_FAMILY").ok().and_then(|v| v.parse().ok()).unwrap_or(25);
    let mut qid = 0usize;
    // A duplicate question inflates whichever family it lands in without testing anything new, so the
    // emitter rejects repeats outright rather than relying on each family to vary its own parameters.
    let mut seen_questions: std::collections::HashSet<String> = std::collections::HashSet::new();
    let mut q = |questions: &mut Vec<serde_json::Value>,
                 qid: &mut usize,
                 category: &str,
                 failure_mode: &str,
                 question: String,
                 ikl: String,
                 gold_docs: Vec<usize>,
                 gold_answer: Option<serde_json::Value>,
                 difficulty: serde_json::Value,
                 retrieval_query: String|
     -> bool {
        if !seen_questions.insert(question.clone()) {
            return false;
        }
        *qid += 1;
        questions.push(serde_json::json!({
            "id": format!("q{:04}", *qid),
            "category": category,
            "failure_mode": failure_mode,
            "question": question,
            // the keyword/dense phrasing of the same intent, so an embedding baseline is not handicapped
            // by having to parse the logical form out of the natural question
            "retrieval_query": retrieval_query,
            "ikl": ikl,
            "gold_doc_ids": gold_docs,
            "gold_answer": gold_answer,
            "difficulty": difficulty,
        }));
        true
    };

    let hub_venue_total = battles.iter().filter(|b| b.venue == HUB_VENUE.1).count();

    // FAMILY 1 — conjunction through a supernode: venue ∩ tournament
    // venue and tournament must be driven by INDEPENDENT strides, or the pair repeats: both derived from i
    // collides onto lcm(2, |TOURNAMENTS|) distinct combinations. Every other question keeps the hub venue so
    // the supernode case stays represented.
    // Enumerate the (venue, tournament) pairs explicitly. Deriving both from one loop index can only ever
    // produce lcm(strides) distinct pairs — fewer than the questions asked for — so the family repeats itself
    // no matter how the strides are chosen. The hub venue leads so the supernode case is always covered.
    let mut vt_pairs: Vec<(&str, &str)> = TOURNAMENTS.iter().map(|t| (HUB_VENUE.1, *t)).collect();
    for lap in 0..2 {
        for (i, (_, venue)) in PLACES.iter().enumerate() {
            if *venue == HUB_VENUE.1 {
                continue;
            }
            vt_pairs.push((venue, TOURNAMENTS[(i + lap) % TOURNAMENTS.len()]));
        }
    }
    for i in 0..per_family.min(vt_pairs.len()) {
        let (venue, tour) = vt_pairs[i];
        let gold: Vec<usize> = battles.iter().filter(|b| b.venue == venue && b.tour == tour).map(|b| b.doc).collect();
        if gold.is_empty() {
            continue;
        }
        let venue_degree = battles.iter().filter(|b| b.venue == venue).count();
        q(&mut questions, &mut qid, "conjunction",
          "supernode explosion — either side alone matches a large fraction of the corpus",
          format!("Which battle reports describe a match at {venue} during the {tour}?"),
          format!("(and venue/{} tournament/{})", slug(venue), slug(tour)),
          gold, None,
          serde_json::json!({ "venue_degree": venue_degree, "tournament_degree": battles.iter().filter(|b| b.tour == tour).count() }),
          format!("{venue} {tour} battle match report"));
    }

    // FAMILY 2 — negation over lexically near-identical confounders
    for sp in SPECIES.iter().take(per_family) {
        let gold: Vec<usize> = rulings.iter().filter(|r| r.species == sp.name && r.banned && !r.hedged).map(|r| r.doc).collect();
        if gold.is_empty() {
            continue;
        }
        let permitted = rulings.iter().filter(|r| r.species == sp.name && !r.banned).count();
        let hedged = rulings.iter().filter(|r| r.species == sp.name && r.hedged).count();
        q(&mut questions, &mut qid, "negation",
          "ignored NOT / silent similarity failure — permitted, banned and hedged rulings differ by one word",
          format!("Which rulings firmly ban {} (excluding permissive and hedged rulings)?", sp.name),
          format!("(and species/{} state/negated (not state/hedged))", slug(sp.name)),
          gold, None,
          serde_json::json!({ "confounder_permitted": permitted, "confounder_hedged": hedged }),
          format!("{} banned not permitted regulation series ruling", sp.name));
    }

    // FAMILY 3 — numeric thresholds with values clustered at the boundary
    // Thresholds are drawn from the real elevation distribution rather than a short hardcoded ladder, so
    // each question has a different selectivity instead of repeating five values.
    let mut elevations: Vec<i64> = surveys.iter().map(|s| s.elevation_m).collect();
    elevations.sort_unstable();
    for k in 0..per_family {
        let threshold = if elevations.is_empty() {
            900
        } else {
            // sweep the 20th..95th percentile so selectivity ranges from broad to narrow
            let pct = 20 + (k * 75) / per_family.max(1);
            elevations[((elevations.len() - 1) * pct) / 100]
        };
        let gold: Vec<usize> = surveys.iter().filter(|s| s.elevation_m > threshold).map(|s| s.doc).collect();
        if gold.is_empty() {
            continue;
        }
        let near = surveys.iter().filter(|s| (threshold - 50..=threshold).contains(&s.elevation_m)).count();
        q(&mut questions, &mut qid, "numeric",
          "threshold confounders — most values sit within 50 m of the boundary",
          format!("Which habitat surveys recorded an elevation above {threshold} m?"),
          format!("(num elevation_m gt {threshold})"),
          gold, None,
          serde_json::json!({ "near_miss_docs": near }),
          format!("habitat survey elevation above {threshold} metres high altitude"));
    }

    // FAMILY 4 — superlatives: a scalar that requires counting the whole hub, not a top-k window
    for (i, place) in PLACES.iter().take(per_family).enumerate() {
        let venue = place.1;
        let mut counts: BTreeMap<&str, usize> = BTreeMap::new();
        for b in battles.iter().filter(|b| b.venue == venue) {
            *counts.entry(b.a.as_str()).or_default() += 1;
            *counts.entry(b.b.as_str()).or_default() += 1;
        }
        let Some((top, n)) = counts.iter().max_by_key(|(_, c)| **c).map(|(t, c)| (t.to_string(), *c)) else { continue };
        let scanned = battles.iter().filter(|b| b.venue == venue).count();
        let _ = i;
        q(&mut questions, &mut qid, "superlative",
          "context-window bloat — the answer needs a count over every matching document",
          format!("Which trainer competed in the most battles at {venue}?"),
          format!("(breakdown venue/{} trainer)", slug(venue)),
          Vec::new(), Some(serde_json::json!({ "trainer": top, "battles": n })),
          serde_json::json!({ "documents_to_scan": scanned }),
          format!("trainer most battles competed {venue} frequent"));
    }

    // FAMILY 5 — joins across the highest-degree node
    for (i, sp) in SPECIES.iter().take(per_family).enumerate() {
        let region = PLACES[i % PLACES.len()].0;
        let gold: Vec<usize> = battles
            .iter()
            .filter(|b| (b.species_a == sp.name || b.species_b == sp.name) && b.region != region)
            .map(|b| b.doc)
            .collect();
        if gold.is_empty() {
            continue;
        }
        let degree = battles.iter().filter(|b| b.species_a == sp.name || b.species_b == sp.name).count();
        q(&mut questions, &mut qid, "join",
          "join through a supernode — traversal fans out across the species' full degree",
          format!("Which battles used {} outside the {region} region?", sp.name),
          format!("(and species/{} (not region/{region}))", slug(sp.name)),
          gold, None,
          serde_json::json!({ "species_degree": degree }),
          format!("{} battles outside {region} region", sp.name));
    }

    // FAMILY 6 — temporal loci
    // Only |years| x 4 pure quarter scopes exist, so asking for more than that repeats questions. The first
    // pass covers each bare quarter once; the remainder are scoped by venue, which keeps every question
    // distinct while still testing date normalisation.
    let mut temporal_added = 0usize;
    for k in 0..per_family * 3 {
        if temporal_added >= per_family {
            break;
        }
        let quarter = 1 + (k as i64 % 4);
        let year = 2025 + ((k / 4) as i64 % 2);
        let bare = k < 8;
        let venue = PLACES[(k / 8) % PLACES.len()].1;
        let gold: Vec<usize> = battles
            .iter()
            .filter(|b| b.year == year && b.quarter == quarter && (bare || b.venue == venue))
            .map(|b| b.doc)
            .collect();
        if gold.is_empty() {
            continue;
        }
        let (question, ikl, retrieval) = if bare {
            (
                format!("Which battles took place in Q{quarter} {year}?"),
                format!("(and time/{year} time/{year}/q{quarter})"),
                format!("battles Q{quarter} {year} quarter date"),
            )
        } else {
            (
                format!("Which battles took place at {venue} in Q{quarter} {year}?"),
                format!("(and venue/{} time/{year} time/{year}/q{quarter})", slug(venue)),
                format!("battles {venue} Q{quarter} {year} quarter date"),
            )
        };
        if q(&mut questions, &mut qid, "temporal",
          "temporal bucketing — dates must be normalised, not string-matched",
          question, ikl, gold, None,
          serde_json::json!({ "year_total": battles.iter().filter(|b| b.year == year).count() }),
          retrieval) {
            temporal_added += 1;
        }
    }

    // FAMILY 7 — unanswerable: the correct behaviour is refusal WITH alternatives, not invention
    let absent = ["Mewtwo", "Arceus", "Eternatus", "Zacian", "Kyogre"];
    for name in absent.iter().take(per_family.min(absent.len())) {
        q(&mut questions, &mut qid, "unanswerable",
          "hallucination pressure — the entity is absent, so any non-empty answer is fabricated",
          format!("Which battles used {name}?"),
          String::new(),
          Vec::new(),
          Some(serde_json::json!({ "refuse": true, "alternatives": SPECIES.iter().take(5).map(|s| s.name).collect::<Vec<_>>() })),
          serde_json::json!({ "present_in_corpus": false }),
          format!("{name} battles used trainer"));
    }

    // FAMILY 8 — GEO ROLL-UP: a region contains many cities, so the answer needs hierarchy, not a string
    // match on the region name (most battle documents name only the city in prose).
    let regions: Vec<&str> = {
        let mut r: Vec<&str> = PLACES.iter().map(|(reg, _)| *reg).collect();
        r.sort();
        r.dedup();
        r
    };
    for region in regions.iter().take(per_family) {
        let gold: Vec<usize> = battles.iter().filter(|b| b.region == *region).map(|b| b.doc).collect();
        if gold.is_empty() {
            continue;
        }
        let cities = PLACES.iter().filter(|(r, _)| r == region).count();
        q(&mut questions, &mut qid, "geo_rollup",
          "spatial roll-up — the region subsumes several cities, so a flat keyword match misses most documents",
          format!("Which battles took place anywhere in the {region} region?"),
          format!("geo/{region}/*"),
          gold, None,
          serde_json::json!({ "cities_in_region": cities }),
          format!("{region} region battles cities"));
    }

    // FAMILY 9 — 2-HOP PATH: trainers linking two venues. A single hop answers nothing; the link exists
    // only through a trainer who appears in both places (variable-length path territory).
    for k in 0..per_family {
        // A fixed offset between the two endpoints yields only |PLACES| distinct pairs; advancing the offset
        // every lap walks the pair space instead of cycling one diagonal of it.
        let v1 = PLACES[k % PLACES.len()].1;
        let v2 = PLACES[(k + 1 + k / PLACES.len()) % PLACES.len()].1;
        if v1 == v2 {
            continue;
        }
        let at = |v: &str| -> BTreeMap<String, usize> {
            let mut m = BTreeMap::new();
            for b in battles.iter().filter(|b| b.venue == v) {
                *m.entry(b.a.clone()).or_insert(0) += 1;
                *m.entry(b.b.clone()).or_insert(0) += 1;
            }
            m
        };
        let (m1, m2) = (at(v1), at(v2));
        let both: Vec<String> = m1.keys().filter(|t| m2.contains_key(*t)).cloned().collect();
        if both.is_empty() {
            continue;
        }
        q(&mut questions, &mut qid, "path_2hop",
          "variable-length path — the two venues are connected only via trainers appearing at both",
          format!("Which trainers competed at both {v1} and {v2}?"),
          format!("(and venue/{} venue/{})", slug(v1), slug(v2)),
          Vec::new(),
          Some(serde_json::json!({ "trainers": both })),
          serde_json::json!({ "venue1_trainers": m1.len(), "venue2_trainers": m2.len() }),
          format!("trainers {v1} {v2} both competed"));
    }

    // FAMILY 10 — CO-OCCURRENCE: which species shared a battle with the hub species. Traversing the hub's
    // full degree is exactly the explosion the paper describes; a popcount over shared situations is not.
    for (k, sp) in SPECIES.iter().take(per_family).enumerate() {
        let partners: Vec<String> = {
            let mut set: BTreeMap<String, usize> = BTreeMap::new();
            for b in battles.iter() {
                if b.species_a == sp.name && b.species_b != sp.name {
                    *set.entry(b.species_b.clone()).or_insert(0) += 1;
                } else if b.species_b == sp.name && b.species_a != sp.name {
                    *set.entry(b.species_a.clone()).or_insert(0) += 1;
                }
            }
            set.into_iter().filter(|(_, n)| *n >= 3).map(|(s, _)| s).collect()
        };
        if partners.is_empty() {
            continue;
        }
        let _ = k;
        let degree = battles.iter().filter(|b| b.species_a == sp.name || b.species_b == sp.name).count();
        q(&mut questions, &mut qid, "cooccurrence",
          "supernode traversal — enumerating a hub's neighbours means walking its whole degree",
          format!("Which species appeared in a battle alongside {}?", sp.name),
          format!("(cooccurs species/{})", slug(sp.name)),
          Vec::new(),
          Some(serde_json::json!({ "species": partners })),
          serde_json::json!({ "hub_degree": degree }),
          format!("{} opponent species battled against", sp.name));
    }

    // FAMILY 11 — UNIVERSAL QUANTIFIER: venues where EVERY battle ran over N minutes. A single
    // counter-example flips the answer, so sampling or top-k retrieval cannot decide it.
    // Scoped to (venue, quarter): a universal over ~400 battles is always false and therefore degenerate --
    // every question would share the same empty answer. Per venue-quarter the groups are small enough that
    // some genuinely qualify, which is what makes "one counter-example flips it" a real test.
    //
    // The threshold is drawn from the MEDIAN of the per-venue minimum durations in that quarter, so roughly
    // half the venues qualify. A fixed threshold makes the family answerable by always replying "none",
    // which measures nothing.
    let mut quant_added = 0usize;
    for k in 0..per_family * 3 {
        if quant_added >= per_family {
            break;
        }
        // year and quarter must vary INDEPENDENTLY: deriving both from k alone (k%2, k%4) correlates them
        // into only 4 distinct scopes, so 25 slots would repeat the same 4 questions.
        let quarter = 1 + (k as i64 % 4);
        let year = 2025 + ((k / 4) as i64 % 2);
        // The tournament joins the scope key because a venue's minimum over ~50 battles is nearly always the
        // global floor, which collapses every quantile onto the same threshold and repeats the question.
        // Narrowing to venue x quarter x tournament leaves groups small enough for the minimums to spread.
        let tour = TOURNAMENTS[k % TOURNAMENTS.len()];
        let quantile = (k / 4) % 3;
        // per venue: (battle count, minimum duration) within the scope
        let mut per_venue: BTreeMap<&str, (usize, i64)> = BTreeMap::new();
        for b in battles.iter().filter(|b| b.year == year && b.quarter == quarter && b.tour == tour) {
            let e = per_venue.entry(b.venue.as_str()).or_insert((0, i64::MAX));
            e.0 += 1;
            e.1 = e.1.min(b.minutes);
        }
        // venues with too few battles make the universal trivially true, so require a real group
        per_venue.retain(|_, (total, _)| *total >= 3);
        let mut mins: Vec<i64> = per_venue.values().map(|(_, m)| *m).collect();
        if mins.is_empty() {
            continue;
        }
        mins.sort_unstable();
        // A venue qualifies iff its minimum exceeds the threshold, so a quantile of the per-venue minimums
        // splits the field. Sweeping the quantile varies how many venues qualify across questions, so the
        // family cannot be gamed by learning one answer size.
        let idx = ((mins.len() - 1) * (quantile + 1)) / 4;
        let threshold = mins[idx];
        let all_over: Vec<String> = per_venue
            .iter()
            .filter(|(_, (_, min))| *min > threshold)
            .map(|(v, _)| v.to_string())
            .collect();
        if q(&mut questions, &mut qid, "quantifier",
          "universal quantification -- one counter-example flips the answer, so sampling cannot decide it",
          format!("At which venues did EVERY {tour} battle in Q{quarter} {year} run longer than {threshold} minutes?"),
          format!("(and venue/* tournament/{} time/{year}/q{quarter} (not (num minutes le {threshold})))", slug(tour)),
          Vec::new(),
          Some(serde_json::json!({ "venues": all_over })),
          serde_json::json!({
              "venues_considered": per_venue.len(),
              "scoped_battles": per_venue.values().map(|(t, _)| t).sum::<usize>(),
              "threshold_minutes": threshold
          }),
          format!("venues every {tour} battle longer than {threshold} minutes Q{quarter} {year}")) {
            quant_added += 1;
        }
    }

    // ---- records.jsonl -----------------------------------------------------------------------------
    // The structured facts behind every generated document, plus the literal strings that MUST survive any
    // later paraphrase. `synth_prose` rewrites these documents as natural prose via Bedrock; the gold sets
    // in questions.jsonl are derived from these records, so a paraphrase that drops a fact would silently
    // make the gold answer wrong. Emitting the required literals here lets the rewriter verify rather than
    // trust, and fall back to the deterministic template when a fact goes missing.
    {
        let mut lines = String::new();
        for b in &battles {
            let rec = serde_json::json!({
                "doc": b.doc, "kind": "battle",
                "fields": {
                    "trainer_a": b.a, "trainer_b": b.b, "winner": b.winner, "venue": b.venue,
                    "region": b.region, "tournament": b.tour, "year": b.year, "quarter": b.quarter,
                    "minutes": b.minutes, "species_a": b.species_a, "species_b": b.species_b,
                    "move": b.move_name
                },
                "must_contain": [
                    b.a.clone(), b.b.clone(), b.venue.clone(), b.region.clone(), b.tour.clone(),
                    b.species_a.clone(), b.species_b.clone(), b.move_name.clone(),
                    b.minutes.to_string(), b.year.to_string()
                ]
            });
            let _ = writeln!(lines, "{rec}");
        }
        for r in &rulings {
            let verdict = if r.banned { "not permitted" } else { "permitted" };
            // The verdict wording carries the polarity the negation family tests, so it is required verbatim.
            // "permitted" is a SUBSTRING of "not permitted", so a required-literal check alone would accept a
            // permit rendered as a ban. The forbidden list closes that hole, and keeps hedge cues out of firm
            // rulings — the benchmark counts a hedged ruling as neither asserted nor refuted.
            let mut must: Vec<String> =
                vec![r.species.clone(), r.series.to_string(), r.year.to_string(), verdict.to_string()];
            let mut must_not: Vec<String> = Vec::new();
            if r.hedged {
                must.push("under review".to_string());
            } else {
                must_not.extend(["under review", "provisional", "may be permitted"].iter().map(|s| s.to_string()));
            }
            if !r.banned {
                must_not.push("not permitted".to_string());
            }
            let rec = serde_json::json!({
                "doc": r.doc, "kind": "regulation",
                "fields": {
                    "species": r.species, "series": r.series, "year": r.year,
                    "banned": r.banned, "hedged": r.hedged, "verdict": verdict
                },
                "must_contain": must,
                "must_not_contain": must_not
            });
            let _ = writeln!(lines, "{rec}");
        }
        for s in &surveys {
            let rec = serde_json::json!({
                "doc": s.doc, "kind": "survey",
                "fields": {
                    "region": s.region, "venue": s.venue, "species": s.species,
                    "elevation_m": s.elevation_m, "temp_c": s.temp_c
                },
                "must_contain": [
                    s.species.clone(), s.venue.clone(), s.region.clone(),
                    s.elevation_m.to_string(), s.temp_c.to_string()
                ]
            });
            let _ = writeln!(lines, "{rec}");
        }
        std::fs::write(out.join("records.jsonl"), lines)?;
        println!("  records.jsonl: {} structured records (for synth_prose)", battles.len() + rulings.len() + surveys.len());
    }

    // ---- Gold-set self-check ------------------------------------------------------------------------
    // A benchmark that repeats a question, or whose family has one constant answer, reports a score that
    // reflects the generator rather than the system under test. These are the two ways this file has
    // silently degraded before, so they are checked at generation time instead of trusted.
    {
        let mut by_cat: BTreeMap<String, Vec<&serde_json::Value>> = BTreeMap::new();
        for qn in &questions {
            by_cat.entry(qn["category"].as_str().unwrap_or("?").to_string()).or_default().push(qn);
        }
        let mut warnings: Vec<String> = Vec::new();
        for (cat, qs) in &by_cat {
            let texts: std::collections::BTreeSet<&str> = qs.iter().filter_map(|q| q["question"].as_str()).collect();
            if texts.len() < qs.len() {
                warnings.push(format!("{cat}: {} of {} questions are duplicates", qs.len() - texts.len(), qs.len()));
            }
            // a family whose every answer is identical (commonly all-empty) measures nothing
            let shapes: std::collections::BTreeSet<String> = qs
                .iter()
                .map(|q| {
                    if q["gold_doc_ids"].as_array().map(|a| !a.is_empty()).unwrap_or(false) {
                        format!("docs:{}", q["gold_doc_ids"].as_array().map(|a| a.len()).unwrap_or(0))
                    } else {
                        q["gold_answer"].to_string()
                    }
                })
                .collect();
            // a refusal family shares one answer by definition — the answer IS "refuse"
            let is_refusal = qs.iter().all(|q| q["gold_answer"]["refuse"].as_bool().unwrap_or(false));
            if qs.len() > 2 && shapes.len() == 1 && !is_refusal {
                warnings.push(format!("{cat}: all {} answers are identical ({})", qs.len(), shapes.iter().next().map(|s| s.as_str()).unwrap_or("")));
            }
        }
        if warnings.is_empty() {
            println!("  gold-set self-check: OK ({} questions, all distinct, no constant-answer family)", questions.len());
        } else {
            for w in &warnings {
                eprintln!("  GOLD-SET WARNING: {w}");
            }
        }
    }

    let qpath = out.join("questions.jsonl");
    let body: String = questions.iter().filter_map(|q| serde_json::to_string(q).ok()).map(|l| l + "\n").collect();
    std::fs::write(&qpath, body)?;

    // corpus statistics, so the adversarial properties are inspectable
    let mut species_degree: BTreeMap<&str, usize> = BTreeMap::new();
    for b in &battles {
        *species_degree.entry(b.species_a.as_str()).or_default() += 1;
        *species_degree.entry(b.species_b.as_str()).or_default() += 1;
    }
    let mut degrees: Vec<(&&str, &usize)> = species_degree.iter().collect();
    degrees.sort_by_key(|(_, n)| std::cmp::Reverse(**n));

    println!("wrote {n_docs} documents → {}", out.display());
    println!("  battles {}  rulings {}  surveys {}", battles.len(), rulings.len(), surveys.len());
    println!("  SUPERNODES:");
    println!("    venue '{}' appears in {hub_venue_total} battles ({:.0}% of all)", HUB_VENUE.1, 100.0 * hub_venue_total as f64 / battles.len().max(1) as f64);
    for (sp, n) in degrees.iter().take(3) {
        println!("    species '{sp}' degree {n}");
    }
    println!("  CONFOUNDERS:");
    println!(
        "    rulings: {} firm bans / {} permissive / {} hedged on overlapping species+series",
        rulings.iter().filter(|r| r.banned && !r.hedged).count(),
        rulings.iter().filter(|r| !r.banned).count(),
        rulings.iter().filter(|r| r.hedged).count()
    );
    println!("    surveys within 50 m of the 1000 m threshold: {}", surveys.iter().filter(|s| (950..=1000).contains(&s.elevation_m)).count());
    let mut by_cat: BTreeMap<&str, usize> = BTreeMap::new();
    for q in &questions {
        *by_cat.entry(q["category"].as_str().unwrap_or("?")).or_default() += 1;
    }
    println!("  GOLD TEST SET: {} questions → {}", questions.len(), qpath.display());
    for (c, n) in &by_cat {
        println!("    {c:<14} {n}");
    }
    Ok(())
}