from pathlib import Path
import math,itertools,csv,json,sys,hashlib
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
P=Path(__file__).resolve().parent
STATES=('AA','Aa','aa')
def cross(left,right):
 if left not in STATES or right not in STATES:raise ValueError('Use AA, Aa or aa')
 counts={x:0 for x in STATES}
 for a,b in itertools.product(left,right):counts[''.join(sorted(a+b))]+=1
 return {x:counts[x]/4 for x in STATES}
def binomial(n,k,p):
 if not isinstance(n,int)or isinstance(n,bool)or n<1 or not isinstance(k,int)or isinstance(k,bool)or not 0<=k<=n or not math.isfinite(p)or not 0<=p<=1:raise ValueError('Invalid binomial inputs')
 return math.comb(n,k)*p**k*(1-p)**(n-k)
rows=[];dist=[]
for a,b in itertools.product(STATES,repeat=2):
 r=cross(a,b);rows.append(dict(parent_left=a,parent_right=b,**r))
 # Independently derive from each parent's fraction of A tokens.
 u=a.count('A')/2;v=b.count('A')/2
 refs=dict(AA=u*v,Aa=u*(1-v)+(1-u)*v,aa=(1-u)*(1-v))
 assert r==refs and sum(r.values())==1 and cross(b,a)==r
 for n in (4,20,100):
  p=r['AA'];ps=[binomial(n,k,p)for k in range(n+1)]
  assert abs(sum(ps)-1)<1e-12
  mean=sum(k*w for k,w in enumerate(ps));var=sum((k-n*p)**2*w for k,w in enumerate(ps))
  assert abs(mean-n*p)<1e-10 and abs(var-n*p*(1-p))<1e-10
  for k,w in enumerate(ps):dist.append(dict(parent_left=a,parent_right=b,draws=n,AA_probability=p,AA_count=k,AA_fraction=k/n,probability=w))
# Exhaustively enumerate Bernoulli successes for n4 at p1/4 using four equally likely labels per draw.
counts=[0]*5
for trial in itertools.product(range(4),repeat=4):counts[trial.count(0)]+=1
assert all(math.isclose(counts[k]/256,binomial(4,k,.25))for k in range(5))
guards=0
for f,a in [(cross,('AB','aa')),(cross,('Aa','aA')),(binomial,(0,0,.5)),(binomial,(4,5,.5)),(binomial,(4,1,-1)),(binomial,(4,1,math.nan))]:
 try:f(*a)
 except ValueError:guards+=1
 else:raise AssertionError(a)
for name,data in [('crosses.csv',rows),('sampling.csv',dist)]:
 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(figsize=(9,5),layout='constrained');bottom=np.zeros(9)
for key,color,hatch in [('AA','#2865a4',''),('Aa','#a375b3','//'),('aa','#c6d5e6','..')]:
 y=np.array([r[key]for r in rows]);ax.bar(range(9),y,bottom=bottom,label=key,color=color,hatch=hatch,edgecolor='white');bottom+=y
ax.set(xticks=range(9),xticklabels=[r['parent_left']+' × '+r['parent_right']for r in rows],ylabel='Exact assigned outcome probability',title='One token from each pair\nUniform independent draws from two labeled positions',ylim=(0,1.05));ax.tick_params(axis='x',rotation=40);ax.legend(ncol=3)
for e in ('png','svg','pdf'):fig.savefig(P/f'token-transmission.{e}')
plt.close(fig)
fig,axes=plt.subplots(1,3,figsize=(12,4),layout='constrained')
for ax,n in zip(axes,(4,20,100)):
 rr=[r for r in dist if r['parent_left']==r['parent_right']=='Aa'and r['draws']==n];ax.bar([r['AA_fraction']for r in rr],[r['probability']for r in rr],width=.85/n,color='#2865a4');ax.axvline(.25,color='#a64267',ls='--',label='Expected fraction 0.25');ax.set(xlabel='Fraction of draws with AA',ylabel='Probability mass',title=f'{n} independent draws',xlim=(-.03,1.03));ax.legend(fontsize=8)
fig.suptitle('A probability is not an exact quota in a finite sample\nAa × Aa token example; theoretical distributions, no measured organisms')
for e in ('png','svg','pdf'):fig.savefig(P/f'sampling-variation.{e}')
plt.close(fig)
(P/'validation.json').write_text(json.dumps(dict(all_pass=True,cross_rows=len(rows),sampling_rows=len(dist),independent_small_enumeration=256,normalization_mean_variance_checks=27,guards=guards),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(),model='Uniform token segregation; independent repeated replacement draws',states=STATES,sample_sizes=[4,20,100],empirical_data=False),indent=2)+'\n')
print(len(rows),len(dist))