#!/usr/bin/env python3
"""Exact pattern probabilities for random digit-span lists.

Altamens Research, "Same Length, Fewer Chunks?" (2026-11-11).
Standard library only. Run: python3 calculate.py [--csv results.csv]

Two generators are modelled:
  uniform     every digit is drawn uniformly from 0-9, independently.
  no_repeat   the first digit is uniform; each later digit is uniform over the
              9 digits that differ from the previous one. This mirrors makeDigits()
              in app/tests/working-memory/page.tsx at commit 6dd71de.

Patterns (adjacent digits only, no wrap-around from 9 to 0):
  repeat         two equal neighbours, e.g. 5-5.
  counting_run   three neighbours that step by +1 twice or by -1 twice, e.g. 4-5-6, 8-7-6.
  equal_step_run three neighbours with the same non-zero step, e.g. 2-4-6, 9-6-3
                 (includes every counting run).
  any_pattern    at least one repeat or equal-step run.

For each list length n, the script returns the exact probability that a list
contains at least one pattern, computed by dynamic programming over
(last digit, last step), and checks it against brute-force enumeration for
n <= 6 and a seeded Monte Carlo run for every n.
"""
from __future__ import annotations

import argparse
import csv
import itertools
import random
from fractions import Fraction

DIGITS = range(10)
LENGTHS = range(3, 13)
PATTERNS = ("repeat", "counting_run", "equal_step_run", "any_pattern")
GENERATORS = ("uniform", "no_repeat")
SEED = 20261111
MC_DRAWS = 200_000


def transitions(prev: int, generator: str):
    if generator == "uniform":
        return [(d, Fraction(1, 10)) for d in DIGITS]
    return [(d, Fraction(1, 9)) for d in DIGITS if d != prev]


def hit(pattern: str, prev_step, step) -> bool:
    """Does the newest digit complete a pattern, given the last two steps?"""
    if pattern in ("repeat", "any_pattern") and step == 0:
        return True
    if prev_step is None or step == 0 or prev_step != step:
        return False
    if pattern == "counting_run":
        return abs(step) == 1
    return pattern in ("equal_step_run", "any_pattern")


def exact_probability(n: int, generator: str, pattern: str) -> Fraction:
    """P(at least one pattern in a list of length n), exact."""
    # state: (last digit, last step or None) -> probability of no pattern so far
    states = {(d, None): Fraction(1, 10) for d in DIGITS}
    for _ in range(n - 1):
        nxt: dict = {}
        for (last, prev_step), p in states.items():
            for d, q in transitions(last, generator):
                step = d - last
                if hit(pattern, prev_step, step):
                    continue
                key = (d, step)
                nxt[key] = nxt.get(key, 0) + p * q
        states = nxt
    return 1 - sum(states.values())


def has_pattern(seq, pattern: str) -> bool:
    steps = [b - a for a, b in zip(seq, seq[1:])]
    for i, s in enumerate(steps):
        if hit(pattern, steps[i - 1] if i else None, s):
            return True
    return False


def brute_force(n: int, generator: str, pattern: str) -> Fraction:
    total = Fraction(0)
    for seq in itertools.product(DIGITS, repeat=n):
        if generator == "no_repeat" and any(a == b for a, b in zip(seq, seq[1:])):
            continue
        if has_pattern(seq, pattern):
            total += 1
    denom = 10 ** n if generator == "uniform" else 10 * 9 ** (n - 1)
    return total / denom


def draw(n: int, generator: str, rng: random.Random):
    out = [rng.randrange(10)]
    while len(out) < n:
        d = rng.randrange(10)
        if generator == "no_repeat":
            while d == out[-1]:
                d = rng.randrange(10)
        out.append(d)
    return out


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--csv", help="write the results table to this path")
    args = parser.parse_args()

    rng = random.Random(SEED)
    rows = []
    for generator in GENERATORS:
        for n in LENGTHS:
            sample = [draw(n, generator, rng) for _ in range(MC_DRAWS)]
            for pattern in PATTERNS:
                exact = exact_probability(n, generator, pattern)
                if n <= 6:
                    assert exact == brute_force(n, generator, pattern), (generator, n, pattern)
                mc = sum(has_pattern(s, pattern) for s in sample) / MC_DRAWS
                # 5 binomial standard errors is a loose but deterministic check.
                se = (float(exact) * (1 - float(exact)) / MC_DRAWS) ** 0.5
                assert abs(mc - float(exact)) <= max(5 * se, 1e-9), (generator, n, pattern, mc)
                rows.append({
                    "generator": generator,
                    "list_length": n,
                    "pattern": pattern,
                    "exact_probability": f"{float(exact):.6f}",
                    "percent": f"{100 * float(exact):.1f}",
                    "monte_carlo_estimate": f"{mc:.6f}",
                })

    # Spot checks that a reader can verify by hand.
    assert exact_probability(3, "uniform", "repeat") == 1 - Fraction(81, 100)
    assert exact_probability(3, "no_repeat", "counting_run") == Fraction(16, 810)
    assert exact_probability(3, "no_repeat", "repeat") == 0

    for r in rows:
        print(r["generator"], r["list_length"], r["pattern"], r["percent"] + "%")
    if args.csv:
        with open(args.csv, "w", newline="") as fh:
            writer = csv.DictWriter(fh, fieldnames=list(rows[0]))
            writer.writeheader()
            writer.writerows(rows)


if __name__ == "__main__":
    main()