"""Independent recalculation of Hoel's CE2.0 Fig.2 example and grouping controls.
Only Python's standard library is required. Run python3 reproduce.py.
"""
import json, math, hashlib, random
from pathlib import Path
from datetime import datetime,timezone

T=[[float(i==j) if i<4 else (.25 if j>=4 else 0) for j in range(8)] for i in range(8)]
def entropy(p):return -sum(x*math.log2(x) for x in p if x>0)
def stats(M,prior=None):
    n=len(M);prior=prior or [1/n]*n
    output=[sum(prior[i]*M[i][j] for i in range(n)) for j in range(n)]
    noise=sum(prior[i]*entropy(M[i]) for i in range(n))
    info=entropy(output)-noise
    if n==1:return {'states':1,'determinism':1,'degeneracy':1,'cp':0,'information_bits':0,'note':'one-state CP=0 by the source-code convention; log2(1) normalization is otherwise undefined'}
    return {'states':n,'determinism':1-noise/math.log2(n),'degeneracy':1-entropy(output)/math.log2(n),'cp':info/math.log2(n),'information_bits':info}
def coarse(M,parts):
    probs=[[sum(row[j] for j in b) for b in parts] for row in M]
    mismatch=max(abs(probs[i][j]-probs[b[0]][j]) for b in parts for i in b for j in range(len(parts)))
    C=[[sum(probs[i][j] for i in b)/len(b) for j in range(len(parts))] for b in parts]
    return C,mismatch
def partitions(xs):
    if not xs:yield [];return
    x,*rest=xs
    for p in partitions(rest):
        yield [[x]]+[b[:] for b in p]
        for i in range(len(p)):yield [([x]+b if i==j else b[:]) for j,b in enumerate(p)]
def run():
    cases={
      'micro':[[i] for i in range(8)],
      'merge_two':[[0],[1],[2],[3],[4,5],[6],[7]],
      'merge_three':[[0],[1],[2],[3],[4,5,6],[7]],
      'useful':[[0],[1],[2],[3],[4,5,6,7]],
      'bad':[[0,4],[1],[2],[3],[5],[6],[7]],
      'everything':[list(range(8))],
    }
    result={'measured_utc':datetime.now(timezone.utc).isoformat(),'paper':'https://arxiv.org/html/2503.13395v3','notebook_commit':'6e5cb45bf445a29e8c5fb7b2a01c56efbc6e06ef','matrix':T,'cases':{}}
    for name,parts in cases.items():
        C,err=coarse(T,parts)
        result['cases'][name]={'partition':parts,'matrix':C,'consistent':err<1e-12,'max_block_probability_mismatch':err,**stats(C)}
    assert abs(result['cases']['micro']['cp']-2/3)<1e-12
    assert abs(result['cases']['useful']['cp']-1)<1e-12
    assert not result['cases']['bad']['consistent']
    assert result['cases']['everything']['cp']==0
    for item in result['cases'].values():assert all(abs(sum(row)-1)<1e-12 for row in item['matrix'])
    # Same-prior information control: there is no violation of data processing.
    induced_prior=[.125]*4+[.5]
    result['same_prior_control']={'micro_information_bits':stats(T)['information_bits'],'macro_information_bits':stats(result['cases']['useful']['matrix'],induced_prior)['information_bits'],'note':'CP instead uses uniform interventions over the states at each chosen scale, and normalizes by log2(state count). It is not predictive accuracy.'}
    assert result['same_prior_control']['micro_information_bits']==result['same_prior_control']['macro_information_bits']==2
    enum=[]
    for p in partitions(list(range(8))):
        C,e=coarse(T,p)
        if e<1e-12:enum.append((len(p),stats(C)['cp']))
    result['partition_scan']={'all_partitions':sum(1 for _ in partitions(list(range(8)))),'consistent_including_one_state':len(enum),'largest_state_count_attaining_cp_one':max(n for n,c in enum if abs(c-1)<1e-12)}
    assert result['partition_scan']=={'all_partitions':4140,'consistent_including_one_state':262,'largest_state_count_attaining_cp_one':5}
    # A reproducible trajectory for the movie. It starts in state4 and never leaves4..7.
    rng=random.Random(20260912);trace=[4]
    for _ in range(95):trace.append(rng.randrange(4,8))
    result['illustrative_trace']={'seed':20260912,'states':trace,'note':'One sampled trajectory from the model, not an empirical dataset. Exact CP uses the transition matrix.'}
    Path('results.json').write_text(json.dumps(result,indent=2))
    print(json.dumps({k:result[k] for k in ['measured_utc','partition_scan','same_prior_control']},indent=2))
    for name,c in result['cases'].items():print(name,'CP=',round(c['cp'],8),'consistent=',c['consistent'])
if __name__=='__main__':run()
