"""Exact arithmetic for Essai 2. Synthetic diffusion, not a building model.

Run: python3 experiment.py
Verify the recorded results: python3 experiment.py --verify
"""
from fractions import Fraction as F
from itertools import product
from pathlib import Path
import json
import sys

M = (
    (F(78, 100), F(20, 100), F(1, 100), F(1, 100)),
    (F(20, 100), F(78, 100), F(1, 100), F(1, 100)),
    (F(1, 100), F(1, 100), F(78, 100), F(20, 100)),
    (F(1, 100), F(1, 100), F(20, 100), F(78, 100)),
)
THRESHOLD = F(7)
WITHIN_FACTOR = M[0][0] - M[0][1]
BETWEEN_FACTOR = sum(M[0][:2]) - sum(M[0][2:])

def step(x):
    return tuple(sum(w*v for w, v in zip(row, x)) for row in M)

def means(x):
    return ((x[0]+x[1])/2, (x[2]+x[3])/2)

def coarse_step(m):
    return (F(98,100)*m[0]+F(2,100)*m[1],
            F(2,100)*m[0]+F(98,100)*m[1])

def trace(initial):
    x=tuple(map(F, initial))
    rows=[]
    for t in range(6):
        rows.append({'step':t, 'values':[float(v) for v in x],
                     'room_means':[float(v) for v in means(x)],
                     'flag_required':max(x)>=THRESHOLD})
        x=step(x)
    return rows

def run():
    # Nonnegative weights, preservation of uniform states and total heat.
    assert all(w>=0 for row in M for w in row)
    assert all(sum(row)==1 for row in M)
    assert all(sum(row[j] for row in M)==1 for j in range(4))
    # An independent two-variable recurrence must reproduce the four-variable
    # means. This checks exact closure, not just approximate fitting to a trace.
    tested=0
    for state in product(range(9), repeat=4):
        x=tuple(map(F,state))
        y=step(x)
        assert means(y)==coarse_step(means(x))
        assert sum(y)==sum(x)
        assert max(y)<=max(x)
        assert y[0]-y[1]==WITHIN_FACTOR*(x[0]-x[1])
        assert y[2]-y[3]==WITHIN_FACTOR*(x[2]-x[3])
        assert means(y)[0]-means(y)[1]==BETWEEN_FACTOR*(means(x)[0]-means(x)[1])
        tested+=1
    same_summary=[(F(a),F(8-a),F(b),F(8-b))
                  for a in range(9) for b in range(9)]
    immediate=sum(max(x)>=THRESHOLD for x in same_summary)
    next_flags=sum(max(step(x))>=THRESHOLD for x in same_summary)
    worst_next=max(max(step(x)) for x in same_summary)
    assert len(same_summary)==81 and immediate==56 and next_flags==0
    assert worst_next==F(158,25)  # 6.32
    disturbed_a=tuple(map(F,(8,4,4,4)))
    disturbed_b=tuple(map(F,(6,6,4,4)))
    assert means(disturbed_a)==means(disturbed_b)==(F(6),F(4))
    assert sum(disturbed_a)==sum(disturbed_b)==20
    assert max(disturbed_a)>=THRESHOLD and max(disturbed_b)<THRESHOLD
    return {
        'kind':'Synthetic exact-arithmetic counterexample; no empirical observations',
        'matrix':[[str(w) for w in row] for row in M],
        'threshold':float(THRESHOLD),
        'question':'At the stated observation step, is any value at least 7?',
        'within_room_difference_factor':str(WITHIN_FACTOR),
        'between_room_mean_difference_factor':str(BETWEEN_FACTOR),
        'grid_states_checked':tested,
        'mean_closure':'Exact in this symmetric construction',
        'same_summary_class':{
            'restriction':'Integer initial values 0 through 8, both room means 4',
            'states':len(same_summary),
            'immediate_flag_required':immediate,
            'immediate_no_flag_required':len(same_summary)-immediate,
            'flag_required_after_one_step':next_flags,
            'maximum_after_one_step':float(worst_next),
            'note':'The exhaustive count is for this chosen finite class, not a probability estimate.'
        },
        'A':trace((8,0,4,4)),
        'B':trace((4,4,4,4)),
        'post_equilibrium_disturbance':{
            'A':[int(v) for v in disturbed_a], 'B':[int(v) for v in disturbed_b],
            'room_means':[int(v) for v in means(disturbed_a)],
            'A_requires_flag':max(disturbed_a)>=THRESHOLD,
            'B_requires_flag':max(disturbed_b)>=THRESHOLD,
            'note':'Equal heat added locally or evenly; the new disturbance invalidates an old mixing bound.'
        }
    }

if __name__=='__main__':
    result=run()
    path=Path(__file__).with_name('results.json')
    if sys.argv[1:]==['--verify']:
        assert json.loads(path.read_text())==result, 'Recorded result differs from exact replay'
        print('Verified: 6561 states; exact mean closure; 81-state collision class; 56 immediate flags, zero after one step.')
    elif not sys.argv[1:]:
        path.write_text(json.dumps(result,indent=2)+'\n')
        print(path)
        print(json.dumps(result['same_summary_class'],indent=2))
    else:
        raise SystemExit('Use no arguments to record, or --verify to replay.')