from pathlib import Path
import math,csv,json,sys,hashlib,itertools
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
P=Path(__file__).resolve().parent

def mixture(q):
 if not math.isfinite(q)or not 0<=q<=1:raise ValueError('q must be finite in [0,1]')
 p=.2+.6*q;p11=.04+.60*q;p10=p01=.16;p00=.64-.60*q
 return dict(context_fraction=q,p_A=p,p_B=p,p11=p11,p10=p10,p01=p01,p00=p00,p_B_given_A=p11/p,p_B_given_not_A=p01/(1-p),covariance=.36*q*(1-q),low_B_given_A=.2,low_B_given_not_A=.2,high_B_given_A=.8,high_B_given_not_A=.8)
rows=[mixture(i/100)for i in range(101)]
# Independent enumeration of all eight (Z,A,B) outcomes, then marginalize.
maxerr=0
for r in rows:
 q=r['context_fraction'];cells={(a,b):0. for a,b in itertools.product((0,1),repeat=2)}
 for z,a,b in itertools.product((0,1),repeat=3):
  p=(1/5 if z==0 else 4/5);w=(1-q if z==0 else q)
  cells[a,b]+=w*(p if a else 1-p)*(p if b else 1-p)
 assert abs(sum(cells.values())-1)<1e-14
 for (a,b),v in cells.items():maxerr=max(maxerr,abs(v-r[f'p{a}{b}']))
 pa=sum(v for (a,b),v in cells.items()if a);pb=sum(v for(a,b),v in cells.items()if b)
 refs=dict(p_A=pa,p_B=pb,p_B_given_A=cells[1,1]/pa,p_B_given_not_A=cells[0,1]/(1-pa),covariance=cells[1,1]-pa*pb)
 for k,v in refs.items():maxerr=max(maxerr,abs(v-r[k]))
 assert all(0<=v<=1 for k,v in r.items())
 for p in (.2,.8):assert math.isclose(p*p/p,p)and math.isclose((1-p)*p/(1-p),p)
assert maxerr<1e-14 and rows[0]['covariance']==rows[-1]['covariance']==0
assert math.isclose(rows[50]['p_B_given_A'],.68)and math.isclose(rows[50]['p_B_given_not_A'],.32)and rows[50]['covariance']==.09
guards=0
for q in (-.1,1.1,math.nan,math.inf):
 try:mixture(q)
 except ValueError:guards+=1
 else:raise AssertionError(q)
toys=[]
for split,sentence in [('training','*red pushes blue'),('held_out','blue pushes *red')]:
 words=sentence.split();first=words[0].lstrip('*');marked=[w[1:]for w in words if w.startswith('*')];assert len(marked)==1
 toys.append(dict(split=split,invented_sentence=sentence,first_symbol_rule=first,star_marked_rule=marked[0],assigned_answer='red',first_rule_matches_assigned=first=='red',star_rule_matches_assigned=marked[0]=='red'))
assert toys[0]['first_rule_matches_assigned']and not toys[1]['first_rule_matches_assigned']and all(x['star_rule_matches_assigned']for x in toys)
for name,data in [('mixture.csv',rows),('toy-rules.csv',toys)]:
 with(P/name).open('w',newline='',encoding='utf8')as f:w=csv.DictWriter(f,fieldnames=list(data[0]));w.writeheader();w.writerows(data)
plt.rcParams.update({'font.size':10,'axes.spines.top':False,'axes.spines.right':False,'savefig.dpi':180})
fig,ax=plt.subplots(1,2,figsize=(11,4.8),layout='constrained');q=[r['context_fraction']for r in rows]
ax[0].plot(q,[r['p_B_given_A']for r in rows],label='P(B | A)',color='#245fa2');ax[0].plot(q,[r['p_B_given_not_A']for r in rows],ls='--',label='P(B | not A)',color='#a64267')
ax[0].set(xlabel='Fraction in high-probability context q',ylabel='Conditional probability',ylim=(0,1),title='A  Association after mixing contexts');ax[0].legend()
ax[1].plot(q,[r['covariance']for r in rows],label='Mixed covariance',color='#245fa2');ax[1].axhline(0,color='#a64267',ls='--',label='Within either context: zero')
ax[1].set(xlabel='Fraction in high-probability context q',ylabel='Covariance of binary A and B',title='B  Independence within each context');ax[1].legend()
for a in ax:a.grid(alpha=.2)
fig.suptitle('A shared context can create an association\nAssigned probabilities 0.2 or 0.8; no observations or causal effect estimated')
for e in ('png','svg','pdf'):fig.savefig(P/f'context-mixture.{e}')
plt.close(fig)
(P/'validation.json').write_text(json.dumps(dict(all_pass=True,mixture_rows=101,toy_rows=2,enumerated_outcomes_per_mixture=8,max_independent_error=maxerr,within_context_independence=True,endpoint_and_midpoint_checks=True,guards=guards,toy_rules_verified=True),indent=2)+'\n')
(P/'study.json').write_text(json.dumps(dict(runtime=dict(python=sys.version,numpy=np.__version__,matplotlib=matplotlib.__version__),generator_sha256=hashlib.sha256(Path(__file__).read_bytes()).hexdigest(),synthetic=True,mixture=dict(q=[0,1,.01],conditional_probabilities=[.2,.8],conditional_independence_assumed=True),toy_rules=dict(assigned_answer='red',not_a_grammar_dataset=True),no_empirical_or_causal_estimate=True),indent=2)+'\n')
print('PASS 101 mixture rows + 2 toy-rule rows')