#!/usr/bin/env python3
"""A small exact probability encounter. Stdlib only; no network or sampling.

Run: python play.py
Each round is independent of previous rounds. Pair probabilities are
P(HH)=P(TT)=(1+r)/4 and P(HT)=P(TH)=(1-r)/4.
Both streams are individually iid fair coins for every -1 <= r <= 1.
"""

from collections import defaultdict
from fractions import Fraction as F
from itertools import product
import json


def solve(matrix, rhs):
    """Gauss-Jordan over rationals, with multiple right-hand sides."""
    n = len(matrix)
    width = len(rhs[0])
    rows = [list(a) + list(b) for a, b in zip(matrix, rhs)]
    for col in range(n):
        pivot = next(i for i in range(col, n) if rows[i][col])
        rows[col], rows[pivot] = rows[pivot], rows[col]
        divisor = rows[col][col]
        rows[col] = [x / divisor for x in rows[col]]
        for i in range(n):
            if i != col:
                factor = rows[i][col]
                rows[i] = [x - factor * y for x, y in zip(rows[i], rows[col])]
    answer = [row[n:] for row in rows]
    for i in range(n):
        for k in range(width):
            assert sum(matrix[i][j] * answer[j][k] for j in range(n)) == rhs[i][k]
    return answer


def advance(target, matched, symbol):
    text = target[:matched] + symbol
    if text.endswith(target):
        return len(target)
    return max(k for k in range(len(target)) if text.endswith(target[:k]))


def kernel(r):
    assert -1 <= r <= 1
    return {pair: (1 + r if pair[0] == pair[1] else 1 - r) / 4
            for pair in product("HT", repeat=2)}


def transitions(a, b, state, pairs):
    outcomes = defaultdict(F)
    for (x, y), probability in pairs.items():
        if not probability:
            continue
        i = advance(a, state[0], x)
        j = advance(b, state[1], y)
        if i == len(a) and j == len(b):
            outcome = "tie"
        elif i == len(a):
            outcome = "A"
        elif j == len(b):
            outcome = "B"
        else:
            outcome = (i, j)
        outcomes[outcome] += probability
    assert sum(outcomes.values()) == 1
    return outcomes


def race(a, b, r):
    pairs = kernel(r)
    states = [(0, 0)]
    edges = {}
    for state in states:
        edges[state] = transitions(a, b, state, pairs)
        for outcome in edges[state]:
            if isinstance(outcome, tuple) and outcome not in states:
                states.append(outcome)
    index = {s: i for i, s in enumerate(states)}
    n = len(states)
    matrix = [[F(i == j) for j in range(n)] for i in range(n)]
    rhs = [[F(0), F(0), F(0), F(1)] for _ in states]
    for state, outcomes in edges.items():
        i = index[state]
        for outcome, p in outcomes.items():
            if isinstance(outcome, tuple):
                matrix[i][index[outcome]] -= p
            else:
                rhs[i][("A", "B", "tie").index(outcome)] += p
    solution = solve(matrix, rhs)
    for row in solution:
        assert sum(row[:3]) == 1 and all(0 <= p <= 1 for p in row[:3])
        assert row[3] > 0
    win_a, win_b, tie, duration = solution[0]
    result = {
        "A": a, "B": b, "r": str(r), "A_strict": str(win_a),
        "B_strict": str(win_b), "tie": str(tie),
        "A_with_half_ties": str(win_a + tie / 2),
        "expected_rounds_until_either_finishes": str(duration),
        "reachable_transient_states": n,
    }
    return result, edges


def mean_wait(target):
    n = len(target)
    matrix = [[F(i == j) for j in range(n)] for i in range(n)]
    for i in range(n):
        for symbol in "HT":
            j = advance(target, i, symbol)
            if j < n:
                matrix[i][j] -= F(1, 2)
    return solve(matrix, [[F(1)] for _ in range(n)])[0][0]


def finite_automaton(edges, horizon):
    alive = {(0, 0): F(1)}
    ended = {"A": F(0), "B": F(0), "tie": F(0)}
    for _ in range(horizon):
        later = defaultdict(F)
        for state, mass in alive.items():
            for outcome, probability in edges[state].items():
                if isinstance(outcome, tuple):
                    later[outcome] += mass * probability
                else:
                    ended[outcome] += mass * probability
        alive = later
    return {**ended, "survival": sum(alive.values(), F(0))}


def finite_raw_words(a, b, r, horizon):
    """Independent check: full histories, substring search, no prefix states."""
    pairs = [(pair, p) for pair, p in kernel(r).items() if p]
    totals = dict.fromkeys(("A", "B", "tie", "survival"), F(0))
    for history in product(pairs, repeat=horizon):
        xs, ys, probability = "", "", F(1)
        first_a = first_b = None
        for step, ((x, y), p) in enumerate(history, 1):
            xs += x
            ys += y
            probability *= p
            if first_a is None and a in xs:
                first_a = step
            if first_b is None and b in ys:
                first_b = step
        if first_a is None and first_b is None:
            outcome = "survival"
        elif first_b is None or (first_a is not None and first_a < first_b):
            outcome = "A"
        elif first_a is None or first_b < first_a:
            outcome = "B"
        else:
            outcome = "tie"
        totals[outcome] += probability
    assert sum(totals.values()) == 1
    return totals


def main():
    a, b = "HHT", "HTT"
    parameters = [F(-1), F(-1, 2), F(0), F(1, 2), F(1)]
    races = []
    for r in parameters:
        result, edges = race(a, b, r)
        assert finite_automaton(edges, 6) == finite_raw_words(a, b, r, 6)
        races.append(result)
    assert [row["A_with_half_ties"] for row in (races[0], races[2], races[4])] == ["1/4", "1/2", "2/3"]
    means = {word: str(mean_wait(word)) for word in ("HHT", "HTT", "TTH", "THH")}
    assert set(means.values()) == {"8"}

    # Check equal single-pattern survival distributions on every finite
    # binary word up to length 14, including the reversal/complement bijection.
    complement = str.maketrans("HT", "TH")
    survival_counts = []
    for length in range(15):
        ca = cb = 0
        for letters in product("HT", repeat=length):
            text = "".join(letters)
            transformed = text[::-1].translate(complement)
            assert (a not in text) == (b not in transformed)
            ca += a not in text
            cb += b not in text
        assert ca == cb
        survival_counts.append({"length": length, "each_avoiding_count": ca})

    cycle = []
    for x, y in (("HHT", "HTT"), ("HTT", "TTH"), ("TTH", "THH"), ("THH", "HHT")):
        row, _ = race(x, y, F(1))
        assert F(row["A_strict"]) > F(1, 2)
        cycle.append(row)
    # A different observation protocol: reset after every disjoint triple.
    blocks = ["".join(letters) for letters in product("HT", repeat=3)]
    pa = F(sum(word == a for word in blocks), len(blocks))
    pb = F(sum(word == b for word in blocks), len(blocks))
    disjoint = {
        "A_win": str(pa / (pa + pb)),
        "expected_blocks_until_either_finishes": str(1 / (pa + pb)),
        "expected_tosses_until_either_finishes": str(3 / (pa + pb)),
        "each_pattern_expected_tosses_alone": str(3 / pa),
        "warning": "Unlike the coupling comparison, disjoint blocks change each marginal waiting-time law."
    }
    assert disjoint["A_win"] == "1/2"
    print(json.dumps({
        "description": "Exact first-hit races; fair marginal streams, altered same-round coupling",
        "disjoint_shared_blocks_control": disjoint,
        "coupled_races": races,
        "single_pattern_mean_waits": means,
        "shared_stream_cycle": cycle,
        "equal_survival_counts": survival_counts,
        "checks": {
            "arithmetic": "fractions.Fraction throughout; linear-system residuals exactly zero",
            "raw_history_crosscheck": "all five couplings, all pair histories of six rounds",
            "single_stream_bijection": "all binary words of lengths 0 through 14",
            "sampling": "none"
        }
    }, indent=2))


if __name__ == "__main__":
    main()