#!/usr/bin/env python3
"""check_proof.py for W7-2, table 1RB1LB_1RC0LA_1RD0RE_1LA0LB_1RF---_1RB1LD (740 lines 175 and 489).

Validates every claim of proof.md.  Usage:
    python3 -B check_proof.py            # all sections (~10 min; F alone is most of it)
    python3 -B check_proof.py A B C      # selected sections

Sections
  A  table facts: strings byte-exact, halting transition, 489 = this table from start D mirrored + renamed
  B  engine soundness: block steps vs plain simulation; random runs vs plain simulation; every rule
     proved in C/D/E re-checked on instances by a rule-free run (and by plain simulation when small)
  C  start A (line 175): checked accelerated run from the blank tape to N(508932094)
  D  the cycle lemma: N(m + 22) -->+ N(2^18 m + 38626190) for every m >= 0 (symbolic run)
  E  the cycle lemma re-run concretely for small Y (rule-free where feasible) + composite = qsim's period map
  F  start D (line 489): opening run to D(X0, Y2, Y1) (~4 min), the D lemma (3 free parameters), its
     iteration count, and the Y1-exhaustion run from D(2^20 z + c, Y2', Y1') into N(Y)
"""
import sys, os, random, time, re
HERE = os.path.dirname(os.path.abspath(__file__))
REPO = os.path.abspath(os.path.join(HERE, '..', '..'))
sys.path.insert(0, os.path.join(HERE, 'scripts'))
sys.path.insert(0, os.path.join(REPO, 'tools'))
sys.set_int_max_str_digits(0)
import engine as en
import sim

TM = '1RB1LB_1RC0LA_1RD0RE_1LA0LB_1RF---_1RB1LD'
TM489 = '1RB0RC_1LC1RC_1LD0RB_1LA0LE_1LF---_1LC1RA'
FAIL = []


def check(cond, msg):
    print(('  ok   ' if cond else '  FAIL ') + msg, flush=True)
    if not cond:
        FAIL.append(msg)


def N(Y):
    """The cycle configuration N(Y) = 0^inf 0001 1011^Y <A 0001^28 0011 0000 1000 0^inf."""
    return en.from_words([('0001', 1), ('1011', Y)], [('0001', 28), ('0011', 1), ('0000', 1), ('1000', 1)], 'A', 'L')


N_KEY = N(30).key()


def n_stop(c, i):
    """True at a configuration of the form N(Y) (any Y), after at least one move."""
    return i > 0 and c.key() == N_KEY and [en.lin(x) for j, x in enumerate(c.exps()) if j != 1] == \
        [en.lin(x) for j, x in enumerate(N(30).exps()) if j != 1]


def Q(W):
    """The second cycle configuration Q(W) = 0^inf 0011 1101 F> 1010 0010^W 0000 0010^28 0000 1000 0100 0^inf."""
    return en.from_words([('0011', 1), ('1101', 1)],
                         [('1010', 1), ('0010', W), ('0000', 1), ('0010', 28), ('0000', 1), ('1000', 1), ('0100', 1)], 'F', 'R')


Q_KEY = Q(30).key()
Q_W = 7                 # index of W in Q(W).exps()


def q_stop(c, i):
    """True at a configuration of the form Q(W) (any W), after at least one move."""
    return i > 0 and c.key() == Q_KEY and [en.lin(x) for j, x in enumerate(c.exps()) if j != Q_W] == \
        [en.lin(x) for j, x in enumerate(Q(30).exps()) if j != Q_W]


def to_sim(c):
    """Expand an engine configuration with int exponents to a sim.Config."""
    left = ''.join(w * e for w, e in c.L)
    right = ''.join(w * e for w, e in reversed(c.R))
    left = '0000' + left
    right = right + '0000'
    tape = left + right
    head = len(left) if c.d == 'R' else len(left) - 1
    return sim.Config(tape, head, c.S, c.d)


def canon_dh(cfg):
    return cfg.dh(trim=True)


# ------------------------------------------------------------------ A
def section_A():
    print('A. table facts')
    lines = open(os.path.join(REPO, 'data/holdouts/BB6_holdouts_740.txt')).read().split('\n')
    check(lines[174] == TM, '740 line 175 == ' + TM)
    check(lines[488] == TM489, '740 line 489 == ' + TM489)
    T = en.parse_tm(TM)
    undefined = [k for k, v in T.items() if v is None]
    check(undefined == [('E', 1)], 'the only undefined transition is E1')
    into_E = [k for k, v in T.items() if v and v[2] == 'E']
    into_C = [k for k, v in T.items() if v and v[2] == 'C']
    check(into_E == [('C', 1)] and T[('C', 1)] == ('0', 1, 'E'), 'E is entered only by C1 = 0RE')
    check(into_C == [('B', 0)] and T[('B', 0)] == ('1', 1, 'C'), 'C is entered only by B0 = 1RC')
    # so the machine halts exactly from B> 0 1 1: B0 -> 1RC, C1 -> 0RE, E1 halts
    # 489 = start D of this table, mirrored, states renamed
    import itertools
    rows = TM.split('_')
    found = None
    for perm in itertools.permutations('ABCDEF'):
        if perm[3] != 'A':        # D -> A
            continue
        pi = dict(zip('ABCDEF', perm))
        new = [''] * 6
        for qi, r in enumerate(rows):
            out = ''
            for s in (0, 1):
                t = r[3 * s:3 * s + 3]
                out += '---' if t[2] == '-' else t[0] + ('L' if t[1] == 'R' else 'R') + pi[t[2]]
            new['ABCDEF'.index(pi['ABCDEF'[qi]])] = out
        if '_'.join(new) == TM489:
            found = pi
    check(found is not None, f'489 = this table from start D, mirrored, renamed {found}')


# ------------------------------------------------------------------ B
def section_B(rule_logs):
    print('B. engine soundness')
    E = en.Engine(use_rules=False)
    m = sim.parse(TM)
    # B1 block transitions vs plain simulation
    bad = 0
    for S in 'ABCDEF':
        for w in (format(i, '04b') for i in range(16)):
            for d in 'LR':
                t = E.trans(S, w, d)
                c = en.from_words([('0000', 1), ('1111', 0)] if False else [], [], S, d)
                c = en.Config([['1010', 1], [w, 1]] if d == 'L' else [['1010', 1]], [['0110', 1], [w, 1]] if d == 'R' else [['0110', 1]], S, d)
                start = to_sim(c)
                if t is None:
                    r = sim.run(m, start, 200)
                    bad += not r.halted
                    continue
                r = sim.run(m, start, t[3])
                c2 = c.copy()
                en.block_move(E, c2)
                bad += canon_dh(r.final) != canon_dh(to_sim(c2))
    check(bad == 0, 'B1 all 192 block transitions agree with plain simulation (halting ones halt)')
    # B2 random configurations, rule-free engine vs plain simulation (exact step counts)
    rnd = random.Random(7)
    bad = 0
    n = 0
    for trial in range(300):
        words = [format(rnd.randrange(16), '04b') for _ in range(8)]
        L = [[rnd.choice(words), rnd.randint(1, 6)] for _ in range(rnd.randint(0, 5))]
        R = [[rnd.choice(words), rnd.randint(1, 6)] for _ in range(rnd.randint(0, 5))]
        c = en.from_words(L, list(reversed(R)), rnd.choice('ABCDEF'), rnd.choice('LR'))
        E2 = en.Engine(use_rules=False)
        start = to_sim(c)
        try:
            en.run(E2, c, lambda cc, i: i >= 60, max_moves=100)
        except en.Halt:
            continue
        except en.Unknown:
            continue
        r = sim.run(m, start, E2.steps)
        n += 1
        bad += r.halted or canon_dh(r.final) != canon_dh(to_sim(c))
    check(bad == 0 and n > 100, f'B2 {n} random 60-move rule-free runs agree with plain simulation step for step')
    # B3 every rule proved in the accelerated runs, re-checked on instances
    for name, log in rule_logs:
        checked = small = skipped = wrong = 0
        sample = log if len(log) <= 400 else random.Random(1).sample(log, 400)
        for r in sample:
            ok, viasim = recheck_rule(r, m)
            if ok is None:
                skipped += 1
            elif ok:
                checked += 1
                small += viasim
            else:
                wrong += 1
        check(wrong == 0 and checked > 0,
              f'B3 {name}: {len(log)} rules proved; {checked} of {len(sample)} sampled re-checked on 3 instances each '
              f'by a rule-free run ({small} also by plain simulation), {skipped} too long to re-run, {wrong} wrong')


def recheck_rule(r, m):
    """Instantiate rule r at a few small values and verify it without rules.  (ok, also_by_sim)."""
    g = r.gen_cfg
    vars_ = sorted({v for x in g.exps() if isinstance(x, en.Lin) for v, _ in x.c})
    also = True
    for vals in (0, 1, 3):
        env = {v: (vals if v.startswith('v') or v in ('a', 'b') else 40 + vals) for v in vars_}
        start = g.copy()
        start.set_exps([en.subst(x, env) for x in g.exps()])
        want = g.copy()
        want.set_exps([en.subst(en.add(x, dl) if sl[0] == 'gen' else x, env) for x, sl, dl in zip(g.exps(), r.slots, r.delta)])
        if any(not isinstance(x, int) for x in start.exps()):
            return None, False
        E = en.Engine(use_rules=False)
        c = start.copy()
        target = want
        try:
            en.run(E, c, lambda cc, i: i > 0 and cc == target, max_moves=4000)
        except en.Unknown:
            return None, False
        except en.Halt:
            return False, False
        if sum(e * 4 for _, e in start.L + start.R) < 3000 and E.steps is not None and E.steps < 10**6:
            rr = sim.run(m, to_sim(start), E.steps)
            if rr.halted or canon_dh(rr.final) != canon_dh(to_sim(target)):
                return False, False
        else:
            also = False
    return True, also


# ------------------------------------------------------------------ C
def section_C():
    print('C. start A (line 175): checked accelerated run from the blank tape')
    E = en.Engine()
    c = en.Config([], [], 'A', 'R')
    en.run(E, c, lambda cc, i: n_stop(cc, i) or (i > 0 and False), max_moves=3 * 10**6)
    Y0 = c.L[1][1]
    check(c == N(508932094), f'blank tape, start A -->* N({Y0}) (moves: {sum(E.stats[k] for k in ("block", "chain", "rule"))}, '
          f'{E.stats["rule"]} rule applications, {E.stats["proved"]} rules proved)')
    # cross-check against the independent qsim record: M(X0) with X0 = Y0 - 14 at record 21
    tr = os.path.join(HERE, 'work', 'trace175_r.txt')
    if os.path.exists(tr):
        recs = [l.split(' ', 2) for l in open(tr) if l.startswith('R ')]
        rec21 = recs[21][2].strip()
        check(rec21.startswith(f'0^inf 1 1011^{Y0 - 14} 10111110 1110^26 <A'), f'qsim record 21 is M({Y0 - 14}) (sligocki simulator, independent)')
    return E.rule_log


# ------------------------------------------------------------------ D
def section_D():
    print('D. cycle lemma (symbolic, every m >= 0)')
    E = en.Engine(big='m')
    mvar = en.Lin.var('m')
    c = N(en.add(mvar, 22))
    en.run(E, c, n_stop, max_moves=2 * 10**6)
    want = en.add(en.mul(mvar, 2**18), 38626190)
    check(c == N(want), f'N(m + 22) -->+ N({en.pretty(c.L[1][1])}) = N(2^18 m + 38626190)')
    check(E.need == 0, f'no side condition on m (recorded need: m >= {E.need})')
    check(all(2**18 * (Y - 22) + 38626190 == 2**18 * Y + 32859022 == 2**18 * (Y - 14) + 36529038 for Y in (22, 99)),
          'closed form: N(Y) -->+ N(2^18 Y + 32859022) = N(2^18 (Y - 14) + 36529038)')
    # informational: from N(m + 14) the same kind of run records a small side condition on m
    E2 = en.Engine(big='m')
    c2 = N(en.add(mvar, 14))
    en.run(E2, c2, n_stop, max_moves=2 * 10**6)
    print(f'     info: from N(m + 14) the run reaches N({en.pretty(c2.L[1][1])}) under the side condition m >= {E2.need}')
    print(f'     moves: {sum(E.stats[k] for k in ("block", "chain", "rule"))}, rules proved {E.stats["proved"]}')
    return E.rule_log


# ------------------------------------------------------------------ E
def section_E():
    print('E. cycle lemma re-run on concrete values; composite map')
    for Y in (22, 23, 29, 64, 1001):
        E = en.Engine()
        c = N(Y)
        en.run(E, c, n_stop, max_moves=2 * 10**6)
        check(c == N(2**18 * (Y - 14) + 36529038), f'N({Y}) -->+ N({c.L[1][1]}) = N(2^18 ({Y} - 14) + 36529038), concrete run')
    # the composite of two half-cycles, written on X = Y - 14 (qsim's M(X) is N(X + 14) shifted by a few moves)
    a = 36529024
    check(a * (2**18 + 1) == 9575900996480, 'two cycles: X -> 2^36 X + 36529024 (2^18 + 1) = 2^36 X + 9575900996480 (observe notes)')
    X0 = 508932080
    tr = os.path.join(HERE, 'work', 'trace175_r.txt')
    if os.path.exists(tr):
        recs = [l.split(' ', 2) for l in open(tr) if l.startswith('R ')]
        X1 = 2**18 * X0 + a
        X2 = 2**18 * X1 + a
        check(f'1011^{X1} ' in recs[42][2], f'qsim record 42 contains 1011^{X1} (= 2^18 X0 + {a})')
        check(recs[63][2].startswith(f'0^inf 1 1011^{X2} 10111110 1110^26 <A'), f'qsim record 63 is M({X2})')


# ------------------------------------------------------------------ F
def Dcfg(X, Y2, Y1):
    """The start-D regime configuration D(X, Y2, Y1) = 0^inf 1101^17 1111 0111^X 1101^Y2 1111 0111^Y1 C> 0^inf."""
    return en.from_words([('1101', 17), ('1111', 1), ('0111', X), ('1101', Y2), ('1111', 1), ('0111', Y1)], [], 'C', 'R')


D_KEY = Dcfg(5, 5, 5).key()
DC = 349479900          # D(X, Y2, Y1) -->+ D(2^20 X + DC, Y2 - 120, Y1 - 24)


def d_stop(c, i):
    if i == 0 or c.key() != D_KEY:
        return False
    xs = c.exps()
    return xs[0] == 17 and xs[1] == 1 and xs[4] == 1


def section_F():
    print('F. start D (line 489)')
    logs = []
    # F1 opening run from the blank tape in state D to the first D(X, Y2, Y1)
    t0 = time.time()
    E = en.Engine()
    c = en.Config([], [], 'D', 'R')
    # the first D shape occurs early with small far fields (Y2 = 330, Y1 = 54); its own exhaustion builds the
    # big far fields of the observe notes.  The proof enters at the first D shape with Y1 > 10^17.
    en.run(E, c, lambda cc, i: d_stop(cc, i) and cc.exps()[5] > 10**17, max_moves=2 * 10**7)
    X0, Y20, Y10 = c.exps()[2], c.exps()[3], c.exps()[5]
    check(X0 >= 100 and Y20 >= 130 and Y10 >= 30,
          f'F1 blank tape, start D -->* D(X0, Y2, Y1) with X0 a {len(str(X0))}-digit number, Y2 = {Y20}, Y1 = {Y10} '
          f'({sum(E.stats[k] for k in ("block", "chain", "rule"))} moves, {time.time() - t0:.0f} s)')
    logs.append(('start-D run', E.rule_log))
    # F2 the D lemma, three free parameters
    E = en.Engine(big='x')
    x, a, b = en.Lin.var('x'), en.Lin.var('a'), en.Lin.var('b')
    c = Dcfg(en.add(x, 100), en.add(b, 130), en.add(a, 30))
    en.run(E, c, d_stop, max_moves=3 * 10**6)
    want = Dcfg(en.add(en.mul(en.add(x, 100), 2**20), DC), en.add(b, 10), en.add(a, 6))
    check(c == want and E.need == 0, f'F2 D(x+100, b+130, a+30) -->+ D(2^20 (x+100) + {DC}, b+10, a+6) for all x, a, b >= 0 '
          f'(side condition log: x >= {E.need})')
    logs.append(('D-lemma run', E.rule_log))
    q = 377451701380050       # observe notes: qsim's period map X -> 2^40 X + q (40 records)
    check(DC * (2**20 + 1) - q == -10 * (2**40 - 1),
          'F2 two D steps give X -> 2^40 X + DC (2^20 + 1), which is qsim\'s fitted map X_q -> 2^40 X_q + 377451701380050 '
          'with X_q = X - 10, and Y1 - 48, Y2 - 240 per qsim period (observe notes)')
    # F3 iterate while Y1 >= 30
    n = (Y10 - 30) // 24 + 1
    Y1r, Y2r = Y10 - 24 * n, Y20 - 120 * n
    check(6 <= Y1r < 30 and Y2r >= 10, f'F3 the lemma applies n = {n} times (before each: Y1 >= 30, Y2 >= 130, X >= 100): '
          f'D(X0, Y2, Y1) -->+ D(X_n, {Y2r}, {Y1r}), X_n = 2^20 X_(n-1) + {DC}')
    # F4 the exhaustion run, X_n = 2^20 z + DC with z = X_(n-1) >= X0 >= 100
    E = en.Engine(big='z')
    z = en.Lin.var('z')
    c = Dcfg(en.add(en.mul(z, 2**20), DC), Y2r, Y1r)
    t0 = time.time()
    en.run(E, c, q_stop, max_moves=10**7)
    W = c.exps()[Q_W]
    ok = c == Q(W) and isinstance(W, en.Lin) and all(a_ > 0 for _, a_ in W.c) and W.k >= 14 and E.need <= 100
    check(ok, f'F4 D(2^20 z + {DC}, {Y2r}, {Y1r}) -->+ Q(W) with W = alpha z + beta, '
          f'alpha = 2^{dict(W.c)["z"].bit_length() - 1 if isinstance(W, en.Lin) else "?"}, '
          f'beta a {len(str(W.k)) if isinstance(W, en.Lin) else "?"}-digit number >= 14; side condition z >= {E.need} '
          f'(z = X_(n-1) >= X0 >= 100) ({sum(E.stats[k] for k in ("block", "chain", "rule"))} moves, {time.time() - t0:.0f} s)')
    logs.append(('exhaustion run', E.rule_log))
    return logs


# ------------------------------------------------------------------ G
def section_G():
    print('G. second cycle lemma Q (symbolic, every m >= 0) and concrete re-runs')
    E = en.Engine(big='m')
    mvar = en.Lin.var('m')
    c = Q(en.add(mvar, 14))
    en.run(E, c, q_stop, max_moves=2 * 10**6)
    check(c == Q(en.add(en.mul(mvar, 2**18), 38101896)) and E.need == 0,
          f'G1 Q(m + 14) -->+ Q({en.pretty(c.exps()[Q_W])}) = Q(2^18 m + 38101896), side condition log: m >= {E.need}')
    check(2**18 * (-14) + 38101896 == 34431880, 'closed form: Q(W) -->+ Q(2^18 W + 34431880)')
    for W in (14, 15, 40, 777):
        E2 = en.Engine()
        c2 = Q(W)
        en.run(E2, c2, q_stop, max_moves=2 * 10**6)
        check(c2 == Q(2**18 * W + 34431880), f'G2 Q({W}) -->+ Q({c2.exps()[Q_W]}) = Q(2^18 {W} + 34431880), concrete run')
    return [('Q-cycle run', E.rule_log)]


def main(argv):
    want = set(a.upper() for a in argv) or set('ABCDEFG')
    t0 = time.time()
    logs = []
    if 'A' in want:
        section_A()
    if 'C' in want or 'B' in want:
        logs.append(('start-A run', section_C()))
    if 'D' in want or 'B' in want:
        logs.append(('cycle run', section_D()))
    if 'E' in want:
        section_E()
    if 'F' in want:
        logs.extend(section_F())
    if 'G' in want:
        logs.extend(section_G())
    if 'B' in want:
        section_B(logs)
    print(f'\n{"ALL OK" if not FAIL else str(len(FAIL)) + " FAILED"}  ({time.time() - t0:.0f} s)')
    return 1 if FAIL else 0


if __name__ == '__main__':
    sys.exit(main(sys.argv[1:]))