import json
from dataclasses import dataclass
from pathlib import Path


@dataclass(slots=True)
class EvalSample:
    section: str
    context: str
    expected: str

    def to_dict(self) -> dict[str, str]:
        return {
            "section": self.section,
            "context": self.context,
            "expected": self.expected,
        }


@dataclass(slots=True)
class OpenEvalSample:
    section: str
    context: str
    required_groups: list[list[str]]
    banned_phrases: list[str]
    min_words: int = 12
    require_punctuation: bool = True
    max_tokens: int = 56

    def to_dict(self) -> dict[str, object]:
        return {
            "section": self.section,
            "context": self.context,
            "required_groups": self.required_groups,
            "banned_phrases": self.banned_phrases,
            "min_words": self.min_words,
            "require_punctuation": self.require_punctuation,
            "max_tokens": self.max_tokens,
        }


@dataclass(slots=True)
class CorpusRecord:
    section: str
    context: str
    answer: str
    split: str = "train"

    @property
    def text(self) -> str:
        return _line(self.context, self.answer)

    def to_dict(self) -> dict[str, str]:
        return {
            "section": self.section,
            "split": self.split,
            "context": self.context,
            "answer": self.answer,
            "text": self.text,
        }


@dataclass(slots=True)
class CorpusPackage:
    name: str
    records: list[CorpusRecord]
    section_counts: dict[str, int]
    memorization_samples: list[EvalSample]
    generalization_samples: list[EvalSample]
    open_ended_samples: list[OpenEvalSample]

    @property
    def slug(self) -> str:
        return self.name.lower().replace(" ", "-")

    @property
    def text(self) -> str:
        if not self.records:
            return ""
        return "\n".join(record.text for record in self.records) + "\n"

    def manifest(self, *, corpus_filename: str) -> dict[str, object]:
        return {
            "name": self.name,
            "corpus_filename": corpus_filename,
            "section_counts": self.section_counts,
            "splits": {
                "memorization": [sample.to_dict() for sample in self.memorization_samples],
                "generalization": [sample.to_dict() for sample in self.generalization_samples],
                "open_ended": [sample.to_dict() for sample in self.open_ended_samples],
            },
        }

    def corpus_records(self) -> list[dict[str, str]]:
        return [record.to_dict() for record in self.records]

    def prompt_suite(self) -> list[dict[str, object]]:
        return [
            {
                "prompt": sample.context,
                "tags": [sample.section, "generalization"],
                "min_words": sample.min_words,
                "require_punctuation": sample.require_punctuation,
                "max_tokens": sample.max_tokens,
            }
            for sample in self.open_ended_samples
        ]


def _line(context: str, expected: str) -> str:
        return f"{context} {expected}"


def _balanced_samples(samples: list[EvalSample], total: int) -> list[EvalSample]:
    buckets: dict[str, list[EvalSample]] = {}
    for sample in samples:
        buckets.setdefault(sample.section, []).append(sample)

    selected: list[EvalSample] = []
    ordered_sections = sorted(buckets)
    while len(selected) < total:
        progressed = False
        for section in ordered_sections:
            bucket = buckets[section]
            if not bucket:
                continue
            selected.append(bucket.pop(0))
            progressed = True
            if len(selected) >= total:
                break
        if not progressed:
            break
    return selected


def _recount_sections(records: list[CorpusRecord]) -> dict[str, int]:
    counts: dict[str, int] = {}
    for record in records:
        counts[record.section] = counts.get(record.section, 0) + 1
    return counts


def build_foundation_corpus() -> CorpusPackage:
    records: list[CorpusRecord] = []
    lines: list[str] = []
    section_counts: dict[str, int] = {}
    memorization: list[EvalSample] = []
    generalization: list[EvalSample] = []
    open_ended: list[OpenEvalSample] = []

    def add_train(section: str, context: str, expected: str, *, sample: bool = False) -> None:
        records.append(
            CorpusRecord(
                section=section,
                context=context,
                answer=expected,
                split="train",
            )
        )
        lines.append(_line(context, expected))
        section_counts[section] = section_counts.get(section, 0) + 1
        if sample:
            memorization.append(EvalSample(section=section, context=context, expected=expected))

    def add_holdout(section: str, context: str, expected: str) -> None:
        generalization.append(EvalSample(section=section, context=context, expected=expected))

    def add_open(
        section: str,
        context: str,
        required_groups: list[list[str]],
        *,
        banned_phrases: list[str],
        min_words: int = 12,
        require_punctuation: bool = True,
        max_tokens: int = 56,
    ) -> None:
        open_ended.append(
            OpenEvalSample(
                section=section,
                context=context,
                required_groups=required_groups,
                banned_phrases=banned_phrases,
                min_words=min_words,
                require_punctuation=require_punctuation,
                max_tokens=max_tokens,
            )
        )

    holdout_addition = {
        (2, 19),
        (3, 17),
        (4, 16),
        (5, 15),
        (6, 14),
        (7, 13),
        (8, 12),
        (9, 11),
        (10, 10),
        (11, 9),
        (12, 8),
        (13, 7),
        (14, 6),
        (15, 5),
        (16, 4),
        (17, 3),
        (18, 2),
        (19, 21),
        (20, 22),
        (21, 19),
        (22, 20),
        (23, 18),
        (24, 17),
        (25, 16),
    }
    holdout_successor = {23, 29, 31, 37, 41, 43, 47, 53, 61, 67, 71, 73, 79}
    holdout_predecessor = {24, 30, 32, 38, 42, 44, 48, 54, 62, 68, 72, 74, 80}
    holdout_explain_addition = {
        (7, 9),
        (8, 11),
        (10, 13),
        (12, 15),
        (14, 9),
        (15, 14),
        (16, 12),
        (18, 7),
    }
    holdout_explain_subtraction = {
        (19, 7),
        (22, 9),
        (25, 11),
        (28, 13),
        (31, 15),
        (34, 12),
    }
    holdout_explain_multiplication = {
        (6, 7),
        (7, 8),
        (8, 9),
        (9, 6),
        (11, 5),
        (12, 6),
    }

    for left in range(1, 41):
        for right in range(1, 41):
            context = f"<reason> add {left} plus {right} equals <answer>"
            expected = str(left + right)
            if (left, right) in holdout_addition:
                add_holdout("arithmetic", context, expected)
            else:
                add_train("arithmetic", context, expected, sample=(left + right) % 5 == 0)

    holdout_subtraction = {
        (9, 4),
        (12, 5),
        (15, 6),
        (18, 7),
        (21, 8),
        (24, 9),
        (27, 10),
        (30, 11),
    }
    for left in range(3, 56):
        for right in range(1, min(left, 21)):
            context = f"<reason> subtract {right} from {left} equals <answer>"
            expected = str(left - right)
            if (left, right) in holdout_subtraction:
                add_holdout("arithmetic", context, expected)
            else:
                add_train("arithmetic", context, expected, sample=(left - right) % 6 == 0)

    holdout_multiplication = {
        (7, 8),
        (8, 9),
        (9, 7),
        (11, 6),
        (12, 7),
        (6, 11),
    }
    for left in range(2, 21):
        for right in range(2, 21):
            context = f"<reason> multiply {left} times {right} equals <answer>"
            expected = str(left * right)
            if (left, right) in holdout_multiplication:
                add_holdout("arithmetic", context, expected)
            else:
                add_train("arithmetic", context, expected, sample=(left * right) % 9 == 0)

    holdout_parity = {33, 37, 41, 45, 52, 58}
    for value in range(1, 141):
        context = f"<reason> parity of {value} is <answer>"
        expected = "even" if value % 2 == 0 else "odd"
        if value in holdout_parity:
            add_holdout("arithmetic", context, expected)
        else:
            add_train("arithmetic", context, expected, sample=value % 10 == 0)

    for value in range(1, 181):
        successor_context = f"<reason> successor of {value} is <answer>"
        successor_expected = str(value + 1)
        if value in holdout_successor:
            add_holdout("sequence", successor_context, successor_expected)
        else:
            add_train("sequence", successor_context, successor_expected, sample=value % 7 == 0)

        predecessor_context = f"<reason> predecessor of {value} is <answer>"
        predecessor_expected = str(value - 1)
        if value in holdout_predecessor:
            add_holdout("sequence", predecessor_context, predecessor_expected)
        else:
            add_train("sequence", predecessor_context, predecessor_expected, sample=value % 8 == 0)

    for left in range(2, 25):
        for right in range(2, 25):
            context = f"<reason> explain the sum of {left} and {right} <answer>"
            expected = (
                f"Use {left} and {right} as the two addends; their total is "
                f"{left + right}, so the answer is {left + right}."
            )
            if (left, right) in holdout_explain_addition:
                add_holdout("reasoning", context, expected)
            else:
                add_train("reasoning", context, expected, sample=(left + right) % 7 == 0)

    for left in range(8, 45):
        for right in range(2, min(left, 17)):
            context = f"<reason> explain the difference between {left} and {right} <answer>"
            expected = (
                f"Start with {left} and remove {right}; the remaining value is "
                f"{left - right}, so the answer is {left - right}."
            )
            if (left, right) in holdout_explain_subtraction:
                add_holdout("reasoning", context, expected)
            else:
                add_train("reasoning", context, expected, sample=(left - right) % 8 == 0)

    for left in range(2, 17):
        for right in range(2, 13):
            context = f"<reason> explain the product of {left} and {right} <answer>"
            expected = (
                f"Treat {left} and {right} as factors; combining the equal groups gives "
                f"{left * right}, so the answer is {left * right}."
            )
            if (left, right) in holdout_explain_multiplication:
                add_holdout("reasoning", context, expected)
            else:
                add_train("reasoning", context, expected, sample=(left * right) % 9 == 0)

    capitals = [
        ("japan", "tokyo"),
        ("brazil", "brasilia"),
        ("canada", "ottawa"),
        ("france", "paris"),
        ("germany", "berlin"),
        ("india", "new delhi"),
        ("australia", "canberra"),
        ("egypt", "cairo"),
        ("kenya", "nairobi"),
        ("mexico", "mexico city"),
        ("norway", "oslo"),
        ("chile", "santiago"),
        ("argentina", "buenos aires"),
        ("thailand", "bangkok"),
        ("indonesia", "jakarta"),
        ("morocco", "rabat"),
        ("sweden", "stockholm"),
        ("finland", "helsinki"),
        ("peru", "lima"),
        ("colombia", "bogota"),
    ]
    for country, capital in capitals:
        add_train(
            "memory",
            f"<memory> capital of {country} is <answer>",
            capital,
            sample=country in {"japan", "brazil", "canada", "france", "india", "kenya"},
        )

    analogies_train = [
        ("bird", "nest", "bee", "hive"),
        ("fish", "water", "camel", "desert"),
        ("painter", "brush", "writer", "pen"),
        ("doctor", "hospital", "teacher", "school"),
        ("farmer", "field", "captain", "ship"),
        ("judge", "court", "chef", "kitchen"),
        ("astronomer", "telescope", "musician", "violin"),
        ("pilot", "cockpit", "driver", "garage"),
        ("programmer", "code", "architect", "blueprint"),
        ("tailor", "needle", "carpenter", "hammer"),
        ("sailor", "compass", "hiker", "map"),
        ("chemist", "laboratory", "baker", "oven"),
        ("photographer", "camera", "sculptor", "chisel"),
        ("gardener", "soil", "potter", "clay"),
        ("librarian", "catalog", "analyst", "report"),
        ("surfer", "wave", "skater", "ramp"),
        ("director", "script", "conductor", "score"),
        ("nurse", "clinic", "lawyer", "firm"),
    ]
    analogies_holdout = [
        ("curator", "museum", "editor", "journal"),
        ("beekeeper", "apiary", "farmer", "barn"),
        ("surgeon", "scalpel", "artist", "canvas"),
        ("sailor", "harbor", "miner", "tunnel"),
        ("scientist", "laboratory", "gardener", "greenhouse"),
        ("translator", "dictionary", "navigator", "chart"),
        ("coach", "sideline", "chef", "kitchen"),
        ("astronaut", "capsule", "diver", "reef"),
    ]
    for left_subject, left_object, right_subject, right_object in analogies_train:
        add_train(
            "analogy",
            f"<reason> {left_subject} relates to {left_object} as {right_subject} relates to <answer>",
            right_object,
            sample=left_subject in {"bird", "doctor", "judge", "pilot", "chemist", "nurse"},
        )
    for left_subject, left_object, right_subject, right_object in analogies_holdout:
        add_holdout(
            "analogy",
            f"<reason> {left_subject} relates to {left_object} as {right_subject} relates to <answer>",
            right_object,
        )

    classifications = [
        ("sparrow", "bird"),
        ("salmon", "fish"),
        ("oak", "tree"),
        ("rose", "flower"),
        ("copper", "metal"),
        ("mercury", "planet"),
        ("triangle", "shape"),
        ("python", "language"),
        ("whale", "mammal"),
        ("eagle", "bird"),
        ("lion", "mammal"),
        ("emerald", "gem"),
        ("neptune", "planet"),
        ("ruby", "gem"),
        ("cedar", "tree"),
        ("falcon", "bird"),
        ("orca", "mammal"),
        ("sapphire", "gem"),
        ("elm", "tree"),
        ("swift", "language"),
    ]
    for item, group in classifications:
        add_train(
            "classification",
            f"<memory> category of {item} is <answer>",
            group,
            sample=item in {"sparrow", "salmon", "oak", "rose", "neptune", "ruby"},
        )

    reasoning_phrases = [
        ("think clearly before final response", "response"),
        ("verify each claim before answer", "answer"),
        ("retrieve memory before conclusion", "conclusion"),
        ("focus on evidence before claim", "claim"),
        ("plan then reason then answer", "answer"),
        ("reflect before committing output", "output"),
        ("use memory when context grows", "grows"),
        ("check arithmetic before assertion", "assertion"),
        ("organize steps before conclusion", "conclusion"),
        ("inspect state before next answer", "answer"),
        ("paraphrase before claiming novelty", "novelty"),
        ("stabilize state before long generation", "generation"),
        ("reuse evidence before rewriting summary", "summary"),
        ("compare patterns before final synthesis", "synthesis"),
    ]
    for phrase, final_word in reasoning_phrases:
        add_train(
            "protocol",
            f"<reason> {phrase} <answer>",
            final_word,
            sample=final_word in {"response", "answer", "claim", "generation", "summary"},
        )

    paraphrase_train = [
        (
            "clear goals and steady practice",
            "clear goals joined with steady practice create durable skill",
        ),
        (
            "careful review prevents shallow errors",
            "careful review stops shallow errors before they spread",
        ),
        (
            "patient systems improve over time",
            "patient systems improve through steady revision over time",
        ),
        (
            "bright ideas need exact execution",
            "bright ideas need exact execution to become reliable work",
        ),
        (
            "quiet focus strengthens difficult reasoning",
            "quiet focus strengthens difficult reasoning during long analysis",
        ),
        (
            "small evidence guides better judgment",
            "small evidence guides better judgment when choices feel similar",
        ),
        (
            "stable memory helps long writing",
            "stable memory helps long writing keep its shape and intent",
        ),
        (
            "measured iteration protects quality",
            "measured iteration protects quality while keeping momentum alive",
        ),
        (
            "careful structure scales ambitious work",
            "careful structure scales ambitious work without needless disorder",
        ),
        (
            "strong prompts need grounded answers",
            "strong prompts need grounded answers supported by real evidence",
        ),
        (
            "shared context reduces wasted motion",
            "shared context reduces wasted motion across a complex build",
        ),
        (
            "consistent language sharpens collaboration",
            "consistent language sharpens collaboration and shortens confusion",
        ),
    ]
    paraphrase_holdout = [
        (
            "steady systems reward patient builders",
            "steady systems reward patient builders with dependable progress",
        ),
        (
            "clear revision protects difficult projects",
            "clear revision protects difficult projects from hidden drift",
        ),
        (
            "focused memory improves long responses",
            "focused memory improves long responses during deep reasoning",
        ),
        (
            "clean evidence supports honest claims",
            "clean evidence supports honest claims during ambitious work",
        ),
        (
            "durable plans reduce fragile execution",
            "durable plans reduce fragile execution before launch pressure rises",
        ),
        (
            "careful synthesis strengthens global understanding",
            "careful synthesis strengthens global understanding without empty hype",
        ),
    ]
    for source, target in paraphrase_train:
        add_train(
            "paraphrase",
            f"<reason> paraphrase {source} into stronger prose <answer>",
            target,
            sample=source in {
                "clear goals and steady practice",
                "patient systems improve over time",
                "stable memory helps long writing",
                "shared context reduces wasted motion",
            },
        )
    for source, target in paraphrase_holdout:
        add_holdout(
            "paraphrase",
            f"<reason> paraphrase {source} into stronger prose <answer>",
            target,
        )

    comparison_train = [
        ("pebble", "stone", "boulder", "largest", "boulder"),
        ("stream", "river", "ocean", "largest", "ocean"),
        ("candle", "lantern", "sun", "brightest", "sun"),
        ("village", "city", "continent", "largest", "continent"),
        ("breeze", "wind", "storm", "strongest", "storm"),
        ("cup", "bucket", "reservoir", "largest", "reservoir"),
        ("violin", "orchestra", "stadium chorus", "loudest", "stadium chorus"),
        ("ember", "flame", "wildfire", "hottest", "wildfire"),
        ("minute", "hour", "day", "longest", "day"),
        ("thread", "rope", "bridge cable", "thickest", "bridge cable"),
        ("hill", "mountain", "range", "largest", "range"),
        ("drizzle", "rain", "monsoon", "strongest", "monsoon"),
        ("spark", "torch", "beacon", "brightest", "beacon"),
        ("brook", "canal", "delta", "widest", "delta"),
        ("hut", "house", "tower", "tallest", "tower"),
        ("cart", "truck", "freighter", "largest", "freighter"),
        ("path", "road", "highway", "widest", "highway"),
        ("note", "melody", "symphony", "longest", "symphony"),
    ]
    comparison_holdout = [
        ("seed", "sapling", "forest", "largest", "forest"),
        ("glimmer", "lamp", "lighthouse", "brightest", "lighthouse"),
        ("whisper", "speech", "thunder", "loudest", "thunder"),
        ("creek", "river", "sea", "largest", "sea"),
        ("trail", "road", "expressway", "widest", "expressway"),
        ("hill", "cliff", "summit", "highest", "summit"),
        ("ember", "bonfire", "volcano", "hottest", "volcano"),
        ("minute", "season", "century", "longest", "century"),
    ]
    for first, second, third, comparator, expected in comparison_train:
        add_train(
            "comparison",
            f"<reason> {comparator} among {first} {second} {third} is <answer>",
            expected,
            sample=expected in {"boulder", "ocean", "storm", "day", "range", "highway"},
        )
    for first, second, third, comparator, expected in comparison_holdout:
        add_holdout(
            "comparison",
            f"<reason> {comparator} among {first} {second} {third} is <answer>",
            expected,
        )

    causal_train = [
        ("iron left in rain", "rust"),
        ("clouds cooling into droplets", "rain"),
        ("plants receiving sunlight", "growth"),
        ("water reaching freezing temperature", "ice"),
        ("friction between dry sticks", "heat"),
        ("strong wind over warm water", "waves"),
        ("seed placed in moist soil", "sprout"),
        ("glass exposed to sudden force", "crack"),
        ("constant pressure on stone", "erosion"),
        ("fuel meeting flame", "combustion"),
        ("repeated practice with feedback", "skill"),
        ("unchecked heat in metal", "expansion"),
        ("low temperature overnight", "frost"),
        ("sustained current through filament", "glow"),
        ("gravity pulling rain downhill", "flow"),
        ("sleep loss across many nights", "fatigue"),
        ("overloaded bridge cable", "strain"),
        ("salt water meeting steel", "corrosion"),
    ]
    causal_holdout = [
        ("dust gathering in still air", "settling"),
        ("long drought across dry fields", "cracking"),
        ("steady pressure beneath ice", "creep"),
        ("clean lens focusing sunlight", "heat"),
        ("lack of oxygen in closed flame", "extinguish"),
        ("waves striking rock for years", "wear"),
    ]
    for cause, effect in causal_train:
        add_train(
            "causal",
            f"<reason> effect of {cause} is <answer>",
            effect,
            sample=effect in {"rust", "rain", "growth", "ice", "skill", "fatigue"},
        )
    for cause, effect in causal_holdout:
        add_holdout(
            "causal",
            f"<reason> effect of {cause} is <answer>",
            effect,
        )

    definition_train = [
        ("orbit", "path traced by one body around another"),
        ("bridge", "structure that carries passage over an obstacle"),
        ("catalyst", "substance that speeds a reaction without being consumed"),
        ("harbor", "protected water area where ships can anchor safely"),
        ("algorithm", "finite procedure for transforming input into output"),
        ("archive", "ordered collection preserved for future reference"),
        ("equilibrium", "state where opposing influences remain balanced"),
        ("lens", "curved material that focuses or spreads light"),
        ("reservoir", "stored supply of water or another resource"),
        ("signal", "pattern that carries information across distance"),
        ("compiler", "program that translates source code into another form"),
        ("calendar", "system for organizing days into meaningful cycles"),
        ("estuary", "place where river water meets the sea"),
        ("voltage", "difference in electric potential between two points"),
        ("synapse", "junction where one neuron communicates with another"),
        ("telescope", "instrument that gathers distant light for observation"),
    ]
    definition_holdout = [
        ("glacier", "mass of ice that moves slowly across land"),
        ("protocol", "agreed procedure that coordinates reliable exchange"),
        ("reef", "ridge of rock or coral rising near the water surface"),
        ("memory", "stored information available for later retrieval"),
        ("frequency", "how often a repeating event occurs in set time"),
        ("compass", "instrument that indicates direction relative to north"),
    ]
    for term, definition in definition_train:
        add_train(
            "definition",
            f"<memory> define {term} as <answer>",
            definition,
            sample=term in {"orbit", "algorithm", "compiler", "harbor", "signal"},
        )
    for term, definition in definition_holdout:
        add_holdout(
            "definition",
            f"<memory> define {term} as <answer>",
            definition,
        )

    identity_train = [
        (
            "describe REFRAMR briefly",
            "REFRAMR is an analytical recurrent language system built by OkeyMeta Ltd to compute structure from corpus evidence instead of gradient loops.",
        ),
        (
            "describe REFRAMR in your own words",
            "REFRAMR is OkeyMeta Ltd language intelligence shaped through analytical memory recurrent state and computed structure rather than opaque training ritual.",
        ),
        (
            "describe REFRAMR in your own words with punctuation",
            "REFRAMR is recurrent, analytical, and evidence-driven; OkeyMeta Ltd shapes it to compute structure from corpus behavior instead of blind gradient churn.",
        ),
        (
            "describe REFRAMR in your own words, with punctuation",
            "REFRAMR is a recurrent analytical language system; OkeyMeta Ltd builds it to preserve structure, carry long context, and keep reasoning signals inspectable.",
        ),
        (
            "what is REFRAMR",
            "REFRAMR is an OkeyMeta analytical language system built around computed memory state and closed form readout.",
        ),
        (
            "what makes REFRAMR different",
            "REFRAMR differs by combining analytical memory corpus statistics and transparent reasoning traces without standard backprop training",
        ),
        (
            "describe FrameToken briefly",
            "FrameToken is REFRAMR native tokenizer from OkeyMeta Ltd that preserves reasoning controls while staying fast on ordinary hardware.",
        ),
        (
            "what is REFRAMR mission",
            "REFRAMR aims to build strong language intelligence through computed structure recurrent memory and interpretable reasoning",
        ),
        (
            "how does REFRAMR reason",
            "REFRAMR reasons through recurrent state analytical retrieval transition priors and explicit control tokens",
        ),
        (
            "what is REFRAMR memory",
            "REFRAMR memory is a multi timescale analytical state that compresses history without quadratic attention.",
        ),
        (
            "explain REFRAMR memory for long context",
            "REFRAMR memory keeps long context by folding prior evidence into a persistent analytical state so later tokens can still respond to earlier structure.",
        ),
        (
            "explain REFRAMR memory for long context in your own words",
            "REFRAMR keeps long context through a persistent analytical memory state, so earlier structure can still shape later output without a quadratic attention map.",
        ),
        (
            "describe REFRAMR long context memory",
            "REFRAMR long context memory is a persistent recurrent state that carries history forward without storing every token in a quadratic map.",
        ),
        (
            "what is REFRAMR readout",
            "REFRAMR readout is a closed form mapping from analytical state to token probabilities.",
        ),
        (
            "what does REFRAMR optimize for",
            "REFRAMR optimizes for analytical transparency long context behavior and hardware accessible computation",
        ),
        (
            "what is REFRAMR tokenizer",
            "REFRAMR tokenizer is FrameToken a native OkeyMeta vocabulary system shaped for analytical recurrent generation",
        ),
        (
            "who are you REFRAMR",
            "I am REFRAMR an OkeyMeta analytical language system shaped by corpus structure and transparent reasoning",
        ),
        (
            "what is REFRAMR voice",
            "REFRAMR voice is deliberate evidence driven and structurally aware rather than shallow imitation",
        ),
        (
            "who builds REFRAMR",
            "REFRAMR is built by OkeyMeta Ltd as a recurrent analytical language system for long context reasoning.",
        ),
        (
            "summarize OkeyMeta role in REFRAMR",
            "OkeyMeta Ltd builds REFRAMR as a transparent analytical language system grounded in corpus structure and recurrent memory",
        ),
        (
            "what is OkeyMeta mission for REFRAMR",
            "OkeyMeta Ltd is building REFRAMR to turn analytical structure into practical language intelligence on ordinary hardware",
        ),
        (
            "describe REFRAMR with punctuation",
            "REFRAMR is analytical, recurrent, and deliberate; OkeyMeta Ltd builds it to compute structure from evidence, not gradient ritual.",
        ),
        (
            "summarize REFRAMR with punctuation",
            "REFRAMR is a recurrent analytical language system; OkeyMeta Ltd builds it to keep structure visible, context persistent, and compute practical.",
        ),
        (
            "summarize FrameToken with punctuation",
            "FrameToken preserves boundaries, protects control tokens, and stays portable; it gives REFRAMR a clean native interface.",
        ),
    ]
    identity_holdout = [
        (
            "explain REFRAMR in one sentence",
            "REFRAMR is an OkeyMeta analytical language system that computes structure from corpus statistics and explicit memory dynamics",
        ),
        (
            "summarize REFRAMR identity",
            "REFRAMR is an OkeyMeta analytical recurrent model built to reason with transparent state rather than opaque gradient rituals",
        ),
        (
            "what kind of model is REFRAMR",
            "REFRAMR is an OkeyMeta post transformer recurrent analytical language model focused on computed structure and long stateful reasoning",
        ),
        (
            "describe REFRAMR purpose",
            "REFRAMR exists to turn mathematical structure and recurrent memory into practical language intelligence",
        ),
        (
            "who owns REFRAMR",
            "REFRAMR is built and owned by OkeyMeta Ltd as a long context analytical language effort",
        ),
        (
            "describe FrameToken role",
            "FrameToken is REFRAMR native tokenizer designed by OkeyMeta Ltd for analytical recurrent generation",
        ),
        (
            "explain REFRAMR with punctuation",
            "REFRAMR is recurrent, analytical, and long-context oriented; OkeyMeta Ltd built it to compute structure with visible reasoning.",
        ),
    ]
    for prompt, answer in identity_train:
        add_train(
            "identity",
            f"<reason> {prompt} <answer>",
            answer,
            sample=prompt in {
                "describe REFRAMR briefly",
                "what is REFRAMR",
                "what makes REFRAMR different",
                "describe FrameToken briefly",
                "describe REFRAMR with punctuation",
            },
        )
    for prompt, answer in identity_holdout:
        add_holdout(
            "identity",
            f"<reason> {prompt} <answer>",
            answer,
        )

    exposition_train = [
        (
            "explain why long context matters",
            "Long context matters because ideas unfold across distance: setup, consequence, and revision rarely live in one sentence. A strong recurrent system must carry structure forward, not just local echoes.",
        ),
        (
            "explain why punctuation matters in language models",
            "Punctuation carries structure, pace, and intent; commas slow rhythm, periods close claims, and colons prepare explanation. A model that ignores marks will often flatten meaning.",
        ),
        (
            "explain how punctuation helps long reasoning",
            "Punctuation helps long reasoning because sequence alone is not enough: commas stage detail, semicolons balance linked claims, and periods let one conclusion land before the next begins.",
        ),
        (
            "explain why punctuation supports long context",
            "Punctuation supports long context by keeping long passages segmented and recoverable. When clauses stay marked, memory can preserve relation, pause, and closure more reliably.",
        ),
        (
            "explain why punctuation helps long reasoning",
            "Punctuation helps long reasoning by separating steps, slowing transitions, and protecting closure. Commas meter detail, colons open explanation, and periods keep one claim from smearing into the next.",
        ),
        (
            "outline REFRAMR workflow",
            "REFRAMR follows a clean path: build corpus statistics, derive recurrent state behavior, and compute the readout. Each stage stays inspectable; none requires opaque epoch loops.",
        ),
        (
            "explain OkeyMeta design ethic",
            "OkeyMeta design ethic is practical and strict: keep provenance visible, keep compute sane, and keep the system understandable. Ambition matters, but clarity matters more.",
        ),
        (
            "explain why evidence matters",
            "Evidence matters because confidence alone is cheap; structure, tests, and reproducible runs make a claim durable. When evidence improves, judgment becomes steadier.",
        ),
        (
            "describe analytical memory",
            "Analytical memory compresses history into a reusable state; it does not replay every token. That compression is useful only when the state stays orderly, expressive, and inspectable.",
        ),
        (
            "explain corpus quality",
            "Corpus quality is not only scale: it is structure, range, and cleanliness. Better data teaches a model where to pause, when to compare, and how to finish a thought.",
        ),
        (
            "explain transparent reasoning",
            "Transparent reasoning does not mean leaking private scratch work; it means exposing useful signals, clear traces, and grounded summaries. The system should reveal why a path dominated.",
        ),
        (
            "describe disciplined generalization",
            "Disciplined generalization begins with pattern depth, not shallow imitation. A model should reuse structure carefully, vary language naturally, and stay anchored to evidence.",
        ),
        (
            "explain why recurrent state can scale",
            "Recurrent state can scale because it updates incrementally; it does not rebuild a full attention map at each step. The challenge is quality, not merely length.",
        ),
        (
            "describe strong completion behavior",
            "Strong completion behavior means the answer reaches a real ending: clauses resolve, punctuation lands, and drift stays contained. A half-finished sentence is not intelligence.",
        ),
        (
            "explain why handcrafted data still matters",
            "Handcrafted data still matters because it can encode precision, tone, and deliberate contrast. It supplies patterns that scraped noise often blurs or discards.",
        ),
        (
            "explain why punctuation supports long answers",
            "Punctuation supports long answers because structure must breathe: commas pace detail, semicolons balance related claims, and periods secure closure. Without marks, long prose often collapses into blur.",
        ),
        (
            "describe healthy model discipline",
            "Healthy model discipline is visible in the small things: exact wording, stable endings, measured confidence, and clean recovery from ambiguity. Strong systems respect detail before spectacle.",
        ),
        (
            "explain why broad corpus style matters",
            "Broad corpus style matters because the model learns more than facts; it learns transition, emphasis, cadence, and restraint. A rich corpus teaches how to move from premise to finish.",
        ),
        (
            "describe how evidence and style should meet",
            "Evidence and style should meet in one sentence: the claim must be accurate, and the sentence must be shaped well enough to carry that accuracy without friction. Good language engineering serves both.",
        ),
        (
            "explain why exact retrieval still needs composition",
            "Exact retrieval still needs composition because recovered facts must land in coherent prose; the answer should connect, not merely appear. Precision becomes more useful when it arrives with structure.",
        ),
        (
            "outline why model endings matter",
            "Model endings matter for a simple reason: the final clause teaches whether the system understood the task or only imitated momentum. A clean ending shows control, not luck.",
        ),
    ]
    exposition_holdout = [
        (
            "explain why sentence endings matter",
            "Sentence endings matter because closure guides interpretation; a period settles a claim, while a comma signals more is coming. Good models must feel that difference.",
        ),
        (
            "explain why structured data improves writing",
            "Structured data improves writing because it teaches ordering, emphasis, and transition; the model learns not only facts, but how claims should connect.",
        ),
        (
            "outline why analytical systems need traces",
            "Analytical systems need traces so operators can inspect dominant signals, compare retrieval paths, and debug drift. Visibility turns mystery into engineering.",
        ),
        (
            "describe why punctuation supports reasoning",
            "Punctuation supports reasoning by marking relation, pause, and hierarchy; it helps the reader separate evidence from conclusion. A fluent model should use marks intentionally.",
        ),
        (
            "explain why corpus range matters",
            "Corpus range matters because generalization grows from varied structures, not one narrow script. When prompts diversify, the model learns to pivot with control.",
        ),
        (
            "describe why exact answers still need style",
            "Exact answers still need style: the right fact should arrive with clean syntax, useful pacing, and a stable finish. Precision and fluency should reinforce each other.",
        ),
    ]
    for prompt, answer in exposition_train:
        add_train(
            "exposition",
            f"<reason> {prompt} <answer>",
            answer,
            sample=prompt in {
                "explain why long context matters",
                "explain why punctuation matters in language models",
                "outline REFRAMR workflow",
                "describe strong completion behavior",
            },
        )
    for prompt, answer in exposition_holdout:
        add_holdout(
            "exposition",
            f"<reason> {prompt} <answer>",
            answer,
        )

    composition_train = [
        (
            "ocean",
            "ocean waves move with patient rhythm and silver foam follows the moonlit shore while distant wind keeps a calm measured pulse",
        ),
        (
            "forest",
            "forest light falls softly through cedar branches and cool air carries resin and rain while the ground stays quiet beneath careful steps",
        ),
        (
            "desert",
            "desert heat bends above pale stone and long shadows stretch across patient sand while evening air slowly restores a gentler balance",
        ),
        (
            "city",
            "city dawn spills across glass towers and quiet streets as buses wake in sequence and windows catch a thin ribbon of gold",
        ),
        (
            "mountain",
            "mountain air stays bright and thin while granite faces hold the morning sun and distant rivers thread silver lines below",
        ),
        (
            "harbor",
            "harbor lights shimmer in patient water while cables rest against masts and slow bells mark the edge of another working night",
        ),
        (
            "library",
            "library silence gathers around tall shelves while lamps hold warm circles of light and every page waits with deliberate calm",
        ),
        (
            "laboratory",
            "laboratory glass reflects a quiet blue glow while instruments rest in ordered rows and each surface signals exact preparation",
        ),
        (
            "garden",
            "garden air carries wet soil and green fragrance while trimmed paths divide the beds and new petals lean toward morning light",
        ),
        (
            "observatory",
            "observatory domes open toward dark sky while motors turn with patient certainty and cold metal frames the waiting stars",
        ),
    ]
    composition_holdout = [
        (
            "glacier",
            "glacier light drifts across slow blue ice while distant air remains clear and every ridge keeps a restrained patient shine",
        ),
        (
            "volcano",
            "volcano stone holds the memory of fire while dark slopes remain still and rising heat bends the horizon with slow force",
        ),
        (
            "cathedral",
            "cathedral windows gather colored light while high arches hold a quiet echo and polished stone returns each careful footstep",
        ),
        (
            "market",
            "market voices braid with morning movement while bright fruit lines the tables and woven shade softens the noonward heat",
        ),
        (
            "reef",
            "reef water carries shifting bands of color while coral forms patient cities and bright fish stitch motion through clear blue lanes",
        ),
        (
            "station",
            "station metal hums beneath pale lamps while distant tracks hold a thin vibration and travelers wait inside orderly lines",
        ),
        (
            "courtroom",
            "courtroom wood carries a formal hush while measured voices rise with care and every pause sharpens the weight of the next sentence",
        ),
        (
            "shipyard",
            "shipyard steel rings through salted air while cranes turn with slow authority and sparks drift briefly before fading into dusk",
        ),
        (
            "archive",
            "archive boxes rest in numbered rows while cool air holds the paper scent and each label promises a patient return to memory",
        ),
        (
            "savanna",
            "savanna light stretches across dry grass while distant heat softens the horizon and watchful movement gathers near the last shade",
        ),
        (
            "workshop",
            "workshop lamps shine over ordered tools while sawdust settles in pale ribbons and each bench waits for deliberate hands",
        ),
        (
            "bridge",
            "bridge cables hold their tense geometry while river light drifts below and the roadway hums with disciplined forward motion",
        ),
    ]
    for theme, answer in composition_train:
        add_train(
            "composition",
            f"<reason> write {theme} scene in one paragraph <answer>",
            answer,
            sample=theme in {"ocean", "forest", "city", "harbor", "laboratory"},
        )
        add_train(
            "composition",
            f"<reason> write {theme} scene <answer>",
            answer,
            sample=False,
        )
    for theme, answer in composition_holdout:
        add_holdout(
            "composition",
            f"<reason> write {theme} scene in one paragraph <answer>",
            answer,
        )
        add_holdout(
            "composition",
            f"<reason> write {theme} scene <answer>",
            answer,
        )

    add_open(
        "composition",
        "write harbor dawn scene with calm tension",
        [
            ["harbor", "port"],
            ["dawn", "morning", "sunrise", "light"],
            ["water", "tide", "shore"],
            ["calm", "quiet", "measured", "tension"],
        ],
        banned_phrases=[
            "harbor lights shimmer in patient water while cables rest against masts and slow bells mark the edge of another working night",
        ],
        min_words=16,
        max_tokens=40,
    )
    add_open(
        "composition",
        "write laboratory harbor scene with precise calm",
        [
            ["laboratory", "glass", "instrument"],
            ["harbor", "water", "mast", "cable"],
            ["calm", "quiet", "precise", "ordered"],
        ],
        banned_phrases=[],
        min_words=16,
        max_tokens=40,
    )
    add_open(
        "identity",
        "describe REFRAMR in your own words, with punctuation",
        [
            ["reframr"],
            ["okeymeta"],
            ["analytical", "recurrent", "language", "system"],
        ],
        banned_phrases=[
            "REFRAMR is an analytical recurrent language system built by OkeyMeta Ltd to compute structure from corpus evidence instead of gradient loops",
            "REFRAMR is analytical, recurrent, and deliberate; OkeyMeta Ltd builds it to compute structure from evidence, not gradient ritual.",
        ],
        min_words=12,
        max_tokens=36,
    )
    add_open(
        "exposition",
        "explain why punctuation helps long reasoning",
        [
            ["punctuation"],
            ["reasoning", "thinking"],
            ["structure", "pace", "pause", "closure"],
        ],
        banned_phrases=[
            "Punctuation supports long answers because structure must breathe: commas pace detail, semicolons balance related claims, and periods secure closure. Without marks, long prose often collapses into blur.",
        ],
        min_words=18,
        max_tokens=40,
    )
    add_open(
        "identity",
        "explain REFRAMR memory for long context in your own words",
        [
            ["reframr"],
            ["memory", "state"],
            ["context", "history"],
            ["long", "persistent", "extended"],
        ],
        banned_phrases=[
            "REFRAMR memory is a multi timescale analytical state that compresses history without quadratic attention",
        ],
        min_words=16,
        max_tokens=40,
    )
    add_open(
        "composition",
        "write archive bridge scene with reflective tension",
        [
            ["archive", "paper", "label", "memory"],
            ["bridge", "cable", "river", "roadway"],
            ["reflective", "tension", "quiet", "measured"],
        ],
        banned_phrases=[],
        min_words=16,
        max_tokens=40,
    )

    return CorpusPackage(
        name="FrameCorpus-Foundation-v2",
        records=records,
        section_counts=section_counts,
        memorization_samples=_balanced_samples(memorization, 24),
        generalization_samples=_balanced_samples(generalization, 16),
        open_ended_samples=open_ended,
    )


def build_generalization_corpus() -> CorpusPackage:
    foundation = build_foundation_corpus()
    allowed_sections = {
        "analogy",
        "paraphrase",
        "comparison",
        "causal",
        "definition",
        "identity",
        "exposition",
        "composition",
    }

    records = [
        record
        for record in foundation.records
        if record.section in allowed_sections
    ]
    generalization = [
        sample
        for sample in foundation.generalization_samples
        if sample.section in allowed_sections
    ]
    open_ended = [
        sample
        for sample in foundation.open_ended_samples
        if sample.section in allowed_sections
    ]

    return CorpusPackage(
        name="FrameCorpus-Generalization-v1",
        records=records,
        section_counts=_recount_sections(records),
        memorization_samples=[],
        generalization_samples=_balanced_samples(generalization, min(16, len(generalization))),
        open_ended_samples=open_ended,
    )


def write_corpus_package(package: CorpusPackage, output_dir: str | Path) -> dict[str, str]:
    directory = Path(output_dir)
    directory.mkdir(parents=True, exist_ok=True)

    base_filename = package.slug
    corpus_filename = f"{base_filename}.jsonl"
    manifest_filename = f"{base_filename}.manifest.json"
    prompt_suite_filename = f"{base_filename}.prompts.jsonl"
    corpus_path = directory / corpus_filename
    manifest_path = directory / manifest_filename
    prompt_suite_path = directory / prompt_suite_filename

    corpus_path.write_text(
        "\n".join(json.dumps(record, ensure_ascii=True) for record in package.corpus_records()) + "\n",
        encoding="utf-8",
    )
    manifest_path.write_text(
        json.dumps(package.manifest(corpus_filename=corpus_filename), indent=2),
        encoding="utf-8",
    )
    prompt_suite_path.write_text(
        "\n".join(json.dumps(record, ensure_ascii=True) for record in package.prompt_suite()) + "\n",
        encoding="utf-8",
    )

    return {
        "corpus_path": str(corpus_path.resolve()),
        "manifest_path": str(manifest_path.resolve()),
        "prompt_suite_path": str(prompt_suite_path.resolve()),
    }