#!/usr/bin/env python3
"""Generate the exact Mystery Bounty draw-timing article study.

The disclosed fixture has ten labeled envelopes. The study enumerates all
10 × 9 × 8 = 720 ordered three-envelope prefixes and evaluates three timing
policies without simulation or floating-point decision logic.
"""

from __future__ import annotations

import csv
import hashlib
import json
import math
from collections import Counter
from fractions import Fraction
from itertools import permutations
from pathlib import Path
from typing import Callable, Sequence

STUDY_ID = "gtogecko-mystery-bounty-draw-timing-2026-09-01"
RUN_DATE = "2026-09-01"
THRESHOLDS = (1_000, 400)

SCRIPT_PATH = Path(__file__).resolve()
BUNDLE = SCRIPT_PATH.parent
REPO_CANDIDATE = next(
    (parent for parent in SCRIPT_PATH.parents if (parent / "package.json").is_file()),
    None,
)
IN_REPOSITORY = REPO_CANDIDATE is not None
ROOT = REPO_CANDIDATE if IN_REPOSITORY else Path.cwd()
PUBLISHED_COPY = SCRIPT_PATH.name == "mystery-bounty-draw-timing-generator.py"
OUTPUT = BUNDLE / ("mystery-bounty-output" if PUBLISHED_COPY else "outputs")
PUBLIC_DATA = ROOT / "public" / "data" if IN_REPOSITORY else BUNDLE
INPUT_POOL = BUNDLE / (
    "mystery-bounty-example-pool.csv" if PUBLISHED_COPY else "example-pool.csv"
)

Envelope = dict[str, str | int]
Prefix = Sequence[Envelope]


def load_envelopes(path: Path) -> tuple[Envelope, ...]:
    """Parse and validate the adjacent labeled-pool CSV."""

    with path.open(newline="", encoding="utf-8-sig") as handle:
        reader = csv.DictReader(handle)
        if tuple(reader.fieldnames or ()) != ("envelope_id", "payout_usd"):
            raise ValueError(
                f"{path} must have exactly: envelope_id,payout_usd"
            )
        envelopes: list[Envelope] = []
        seen_ids: set[str] = set()
        for line_number, row in enumerate(reader, start=2):
            envelope_id = (row.get("envelope_id") or "").strip()
            payout_text = (row.get("payout_usd") or "").strip()
            if not envelope_id or envelope_id in seen_ids:
                raise ValueError(f"Invalid or duplicate envelope_id on line {line_number}")
            try:
                payout = int(payout_text)
            except ValueError as error:
                raise ValueError(f"Non-integer payout on line {line_number}") from error
            if payout <= 0:
                raise ValueError(f"Payout must be positive on line {line_number}")
            seen_ids.add(envelope_id)
            envelopes.append({"envelope_id": envelope_id, "payout_usd": payout})
    if len(envelopes) < 3:
        raise ValueError("At least three labeled envelopes are required")
    return tuple(envelopes)


ENVELOPES = load_envelopes(INPUT_POOL)
POOL_INPUT_SHA256 = hashlib.sha256(INPUT_POOL.read_bytes()).hexdigest()


def redeem_now(prefix: Prefix) -> Envelope:
    return prefix[0]


def wait_one(prefix: Prefix) -> Envelope:
    return prefix[1]


def adaptive_threshold(prefix: Prefix) -> Envelope:
    """Draw second after a <=$300 reveal; otherwise wait and draw third."""

    return prefix[1] if int(prefix[0]["payout_usd"]) <= 300 else prefix[2]


POLICIES: tuple[tuple[str, str, Callable[[Prefix], Envelope]], ...] = (
    ("redeem_now", "Redeem before another envelope is revealed", redeem_now),
    ("wait_exactly_one", "Let one other envelope be revealed, then redeem", wait_one),
    (
        "adaptive_threshold",
        "After one reveal, redeem next at $300 or less; above $300, wait once more",
        adaptive_threshold,
    ),
)


def ratio(value: Fraction) -> str:
    return f"{value.numerator}/{value.denominator}"


def fraction_record(value: Fraction) -> dict[str, str | float]:
    return {"fraction": ratio(value), "decimal": round(float(value), 12)}


def choose_probability_at_least_one(n: int, qualifying: int, k: int) -> Fraction:
    misses = math.comb(n - qualifying, k) if k <= n - qualifying else 0
    return Fraction(math.comb(n, k) - misses, math.comb(n, k))


def pool_rows() -> list[dict[str, object]]:
    return [
        {
            "envelope_id": str(envelope["envelope_id"]),
            "payout_usd": int(envelope["payout_usd"]),
        }
        for envelope in ENVELOPES
    ]


def pool_groups() -> list[dict[str, object]]:
    n = len(ENVELOPES)
    counts = Counter(int(envelope["payout_usd"]) for envelope in ENVELOPES)
    return [
        {
            "payout_usd": payout,
            "count": count,
            "pool_value_usd": payout * count,
            "single_draw_probability_fraction": ratio(Fraction(count, n)),
            "single_draw_probability_pct": round(100 * count / n, 6),
        }
        for payout, count in sorted(counts.items())
    ]


def conditional_rows() -> list[dict[str, object]]:
    payouts = [int(envelope["payout_usd"]) for envelope in ENVELOPES]
    n = len(payouts)
    total = sum(payouts)
    initial_mean = Fraction(total, n)
    counts = Counter(payouts)
    jackpot = max(payouts)
    rows: list[dict[str, object]] = []
    for reveal, count in sorted(counts.items()):
        probability = Fraction(count, n)
        remaining_total = total - reveal
        next_ev = Fraction(remaining_total, n - 1)
        change = next_ev - initial_mean
        next_jackpot = Fraction(counts[jackpot] - int(reveal == jackpot), n - 1)
        rows.append(
            {
                "revealed_payout_usd": reveal,
                "reveal_probability_fraction": ratio(probability),
                "reveal_probability_pct": round(100 * float(probability), 6),
                "remaining_prizes": n - 1,
                "remaining_total_usd": remaining_total,
                "next_draw_ev_fraction_usd": ratio(next_ev),
                "next_draw_ev_usd": round(float(next_ev), 6),
                "change_from_initial_ev_fraction_usd": ratio(change),
                "change_from_initial_ev_usd": round(float(change), 6),
                "next_draw_jackpot_probability_fraction": ratio(next_jackpot),
                "next_draw_jackpot_probability_pct": round(100 * float(next_jackpot), 6),
            }
        )
    return rows


def policy_rows() -> tuple[int, list[dict[str, object]], list[dict[str, object]]]:
    prefixes = list(permutations(ENVELOPES, 3))
    total_prefixes = len(prefixes)
    envelope_rows: list[dict[str, object]] = []
    payout_rows: list[dict[str, object]] = []
    baseline_payout_counts = Counter(int(envelope["payout_usd"]) for envelope in ENVELOPES)
    for key, description, selector in POLICIES:
        selected = [selector(prefix) for prefix in prefixes]
        selected_ids = Counter(str(envelope["envelope_id"]) for envelope in selected)
        selected_payouts = Counter(int(envelope["payout_usd"]) for envelope in selected)
        expected_selections_per_envelope = total_prefixes // len(ENVELOPES)
        if set(selected_ids.values()) != {expected_selections_per_envelope}:
            raise AssertionError(
                f"Policy {key} did not select every labeled envelope equally"
            )
        expected = Fraction(
            sum(int(envelope["payout_usd"]) for envelope in selected), total_prefixes
        )
        initial_mean = Fraction(
            sum(int(envelope["payout_usd"]) for envelope in ENVELOPES),
            len(ENVELOPES),
        )
        if expected != initial_mean:
            raise AssertionError(f"Policy {key} changed expected payout: {expected}")
        for envelope in ENVELOPES:
            envelope_id = str(envelope["envelope_id"])
            count = selected_ids[envelope_id]
            envelope_rows.append(
                {
                    "policy": key,
                    "policy_description": description,
                    "envelope_id": envelope_id,
                    "payout_usd": int(envelope["payout_usd"]),
                    "selected_prefixes": count,
                    "total_prefixes": total_prefixes,
                    "selection_probability_fraction": ratio(Fraction(count, total_prefixes)),
                    "selection_probability_pct": round(100 * count / total_prefixes, 6),
                }
            )
        for payout, original_count in sorted(baseline_payout_counts.items()):
            selected_count = selected_payouts[payout]
            payout_rows.append(
                {
                    "policy": key,
                    "payout_usd": payout,
                    "selected_prefixes": selected_count,
                    "total_prefixes": total_prefixes,
                    "probability_fraction": ratio(Fraction(selected_count, total_prefixes)),
                    "probability_pct": round(100 * selected_count / total_prefixes, 6),
                    "original_envelope_count": original_count,
                    "policy_ev_fraction_usd": ratio(expected),
                    "policy_ev_usd": round(float(expected), 6),
                }
            )
    return total_prefixes, envelope_rows, payout_rows


def k_draw_rows() -> list[dict[str, object]]:
    payouts = [int(envelope["payout_usd"]) for envelope in ENVELOPES]
    n = len(payouts)
    mean = Fraction(sum(payouts), n)
    rows: list[dict[str, object]] = []
    for k in range(1, min(4, n) + 1):
        expected_total = k * mean
        row: dict[str, object] = {
            "draws_now": k,
            "expected_total_fraction_usd": ratio(expected_total),
            "expected_total_usd": round(float(expected_total), 6),
        }
        for threshold in THRESHOLDS:
            qualifying = sum(payout >= threshold for payout in payouts)
            probability = choose_probability_at_least_one(n, qualifying, k)
            row[f"prize_ge_{threshold}_count"] = qualifying
            row[f"probability_ge_{threshold}_fraction"] = ratio(probability)
            row[f"probability_ge_{threshold}_pct"] = round(100 * float(probability), 6)
        rows.append(row)
    return rows


def write_csv(path: Path, rows: list[dict[str, object]]) -> None:
    if not rows:
        raise ValueError(f"No rows for {path}")
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.DictWriter(handle, fieldnames=list(rows[0]), lineterminator="\n")
        writer.writeheader()
        writer.writerows(rows)


def write_text_lf(path: Path, value: str) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", newline="\n", encoding="utf-8") as handle:
        handle.write(value)


def build_study() -> tuple[dict[str, object], dict[str, list[dict[str, object]]]]:
    payouts = [int(envelope["payout_usd"]) for envelope in ENVELOPES]
    n = len(payouts)
    total = sum(payouts)
    mean = Fraction(total, n)
    ordered = sorted(payouts)
    median = (
        Fraction(ordered[n // 2])
        if n % 2
        else Fraction(ordered[n // 2 - 1] + ordered[n // 2], 2)
    )
    payout_counts = Counter(payouts)
    highest_frequency = max(payout_counts.values())
    mode = min(payout for payout, count in payout_counts.items() if count == highest_frequency)
    jackpot = max(payouts)
    jackpot_probability = Fraction(payout_counts[jackpot], n)
    conditional = conditional_rows()
    total_prefixes, envelope_policy, payout_policy = policy_rows()
    k_draws = k_draw_rows()

    weighted_ev = sum(
        Fraction(str(row["reveal_probability_fraction"]))
        * Fraction(str(row["next_draw_ev_fraction_usd"]))
        for row in conditional
    )
    weighted_jackpot = sum(
        Fraction(str(row["reveal_probability_fraction"]))
        * Fraction(str(row["next_draw_jackpot_probability_fraction"]))
        for row in conditional
    )
    if weighted_ev != mean or weighted_jackpot != jackpot_probability:
        raise AssertionError("Conditional branches did not recombine to the starting values")

    study: dict[str, object] = {
        "study_id": STUDY_ID,
        "run_date": RUN_DATE,
        "method": f"Exact enumeration of all {total_prefixes} ordered labeled three-envelope prefixes; no simulation",
        "pool_input": {
            "schema": "envelope_id,payout_usd",
            "sha256": POOL_INPUT_SHA256,
        },
        "scope": {
            "model": "One earned draw from a fixed prize pool without replacement; every intervening external reveal and the selected draw are conditionally uniform and exchangeable among the same remaining envelopes, the pool changes only through those draws, and redemption is forced by the final available position",
            "timing_rule": "Bounded and predictable: each stop decision is committed before the selected envelope is revealed, using only earlier reveals and independent private randomness",
            "excludes": [
                "The probability of earning a bounty token",
                "Tournament hand strategy, coverage, chips, and ICM",
                "Event rules that force redemption earlier than the modeled final-position deadline, token expiry, event closure, nested draws, and non-uniform procedures",
                "Biased or non-exchangeable external removals and pool changes caused by anything other than the modeled draws",
                "Adaptive joint timing of multiple accumulated tokens",
            ],
        },
        "synthetic_pool": {
            "prize_count": n,
            "total_usd": total,
            "mean_usd": fraction_record(mean),
            "median_usd": fraction_record(median),
            "mode_usd": mode,
            "jackpot_usd": jackpot,
            "jackpot_probability": fraction_record(jackpot_probability),
            "envelopes": pool_rows(),
            "groups": pool_groups(),
        },
        "conditional_after_one_external_reveal": {
            "rows": conditional,
            "weighted_next_draw_ev_usd": fraction_record(weighted_ev),
            "weighted_next_draw_jackpot_probability": fraction_record(weighted_jackpot),
        },
        "policy_enumeration": {
            "ordered_labeled_prefixes": total_prefixes,
            "prefix_formula": f"{n} × {n - 1} × {n - 2}",
            "prefix_length": 3,
            "policies": [
                {
                    "key": key,
                    "description": description,
                    "labeled_envelope_distribution": [
                        row for row in envelope_policy if row["policy"] == key
                    ],
                    "payout_distribution": [
                        row for row in payout_policy if row["policy"] == key
                    ],
                }
                for key, description, _ in POLICIES
            ],
        },
        "draw_k_now": {
            "formula_expected_total": "k × current mean",
            "formula_at_least_one_threshold": "1 - C(N-J,k) / C(N,k)",
            "rows": k_draws,
        },
        "theorem": {
            "statement": "When every intervening external reveal and the selected draw are conditionally uniform and exchangeable from the same fixed pool, every bounded predictable single-draw timing rule with redemption forced by the final available position has the original ex-ante marginal payout distribution.",
            "proof_outline": "For each payout class, its share of the remaining pool is a bounded martingale. Conditional on any admissible draw time, the next prize has that remaining share; optional sampling returns the initial share for every class.",
        },
    }
    return study, {
        "example-pool.csv": pool_rows(),
        "conditional-updates.csv": conditional,
        "labeled-policy-enumeration.csv": envelope_policy,
        "draw-policy-enumeration.csv": payout_policy,
        "k-draw-odds.csv": k_draws,
    }


def main() -> None:
    study, tables = build_study()
    OUTPUT.mkdir(parents=True, exist_ok=True)
    study_path = OUTPUT / "study.json"
    write_text_lf(study_path, json.dumps(study, indent=2) + "\n")
    for filename, rows in tables.items():
        write_csv(OUTPUT / filename, rows)

    if IN_REPOSITORY and not PUBLISHED_COPY:
        public_names = {
            "example-pool.csv": "mystery-bounty-example-pool.csv",
            "conditional-updates.csv": "mystery-bounty-conditional-updates.csv",
            "labeled-policy-enumeration.csv": "mystery-bounty-labeled-policy-prefixes.csv",
            "draw-policy-enumeration.csv": "mystery-bounty-draw-policies.csv",
            "k-draw-odds.csv": "mystery-bounty-k-draw-odds.csv",
        }
        PUBLIC_DATA.mkdir(parents=True, exist_ok=True)
        write_text_lf(
            PUBLIC_DATA / "mystery-bounty-draw-timing.json",
            json.dumps(study, indent=2) + "\n",
        )
        for source_name, public_name in public_names.items():
            write_csv(PUBLIC_DATA / public_name, tables[source_name])

    prefixes = study["policy_enumeration"]["ordered_labeled_prefixes"]
    print(f"Enumerated {prefixes:,} ordered labeled three-envelope prefixes")
    per_envelope = prefixes // len(ENVELOPES)
    print(
        "All three timing policies selected every labeled envelope "
        f"{per_envelope:,} times and returned "
        f"{study['synthetic_pool']['mean_usd']['fraction']} EV"
    )
    print(f"Wrote {study_path}")


if __name__ == "__main__":
    main()
