#!/usr/bin/env python3
"""Exact checks for discrete branching families; Python standard library only.
Run with --output path to also save the deterministic JSON report.
Assertions are explicit exceptions and remain active under python -O.
The finite checks verify stated examples and interfaces, not infinite-limit proofs.
"""
from fractions import Fraction as F
from itertools import product
from math import comb
from pathlib import Path
import argparse,json
C=0

def check(v,msg):
 global C
 C+=1
 if not v:raise RuntimeError(msg)
def conv(a,b,limit=None):
 n=len(a)+len(b)-1
 if limit is not None:n=min(n,limit+1)
 c=[F(0)]*n
 for i,x in enumerate(a):
  if not x:continue
  for j,y in enumerate(b[:max(0,n-i)]):
   if y:c[i+j]+=x*y
 return c

def power(p,r,limit=None):
 out=[F(1)]
 for _ in range(r):out=conv(out,p,limit)
 return out

def value(p,s):
 out=F(0)
 for x in reversed(p):out=out*s+x
 return out

def moments(p):
 a=sum(F(i)*x for i,x in enumerate(p));v=sum(F(i*i)*x for i,x in enumerate(p))-a*a
 return a,v

def random_sum(p,count):
 out=[F(0)]*((len(p)-1)*(len(count)-1)+1);pw=[F(1)]
 for i,w in enumerate(count):
  for j,c in enumerate(pw):out[j]+=w*c
  if i+1<len(count):pw=conv(pw,p)
 return out

def generation(p,r,n):
 d=[F(0)]*r+[F(1)]
 for _ in range(n):d=random_sum(p,d)
 return d

def queue_total(p,r,N):
 # Process one active individual at a time; stop precisely on the first empty queue.
 state={r:F(1)};ans=[F(0)]*(N+1)
 if r==0:ans[0]=1;return ans
 for n in range(1,N+1):
  new={}
  for active,w in state.items():
   for k,pk in enumerate(p):
    if not pk:continue
    j=active-1+k
    if j==0:ans[n]+=w*pk
    else:new[j]=new.get(j,F(0))+w*pk
  state=new
 return ans

def dwass(p,r,N):
 pw=[F(1)];ans=[F(0)]*(N+1)
 for n in range(1,N+1):
  pw=conv(pw,p,N)
  if n>=r and n-r<len(pw):ans[n]=F(r,n)*pw[n-r]
 return ans

def compositions(n,k):
 if k==1:yield (n,);return
 for a in range(n+1):
  for v in compositions(n-a,k-1):yield(a,)+v

# Complete small offspring-law grids, including degenerate and critical boundaries.
laws=[tuple(F(x,5) for x in w) for w in compositions(5,4)]
for li,p in enumerate(laws):
 m,var=moments(p)
 for r in [1,2,3]:
  a=queue_total(p,r,14);b=dwass(p,r,14)
  for n in range(15):check(a[n]==b[n],f'Dwass {li,r,n}')
  check(sum(a)<=1,f'finite mass {li,r}')
 for n in range(4):
  d=generation(p,1,n);mn,vn=moments(d)
  check(sum(d)==1,f'generation normalization {li,n}')
  check(mn==m**n,f'finite generation mean {li,n}')
  v=F(0);mean=F(1)
  for _ in range(n):v=var*mean+m*m*v;mean*=m
  check(vn==v,f'finite generation variance {li,n}')
  for s in [F(0),F(1,5),F(1,2),F(1)]:
   fs=s
   for _ in range(n):fs=value(p,fs)
   check(value(d,s)==fs,f'PGF composition {li,n,s}')
 # Cumulative generation truncation, computed as a full distribution.
 h=[F(0),F(1)];mean=F(1);v=F(0)
 for depth in range(4):
  hm,hv=moments(h)
  check(hm==mean and hv==v,f'cumulative moments {li,depth}')
  check(sum(h)==1,f'cumulative mass {li,depth}')
  h=[F(0)]+random_sum(h,p)
  v=m*v+var*mean*mean;mean=1+m*mean

# Enumerate all short breadth-first offspring words independently of the DP.
enumeration_words=0
for n in range(1,8):
 for r in [1,2,3]:
  counts={}
  for word in product(range(3),repeat=n):
   enumeration_words+=1;active=r;valid=True
   for j,k in enumerate(word):
    active+=k-1
    if active==0 and j<n-1:valid=False;break
   if valid and active==0:
    key=tuple(word.count(k) for k in range(3));counts[key]=counts.get(key,0)+1
  for p in [(F(1,4),F(0),F(3,4)),(F(1,3),)*3,(F(1,2),F(1,3),F(1,6))]:
   mass=sum(F(v)*p[0]**k[0]*p[1]**k[1]*p[2]**k[2] for k,v in counts.items())
   check(mass==dwass(p,r,n)[n],f'word enumeration {n,r,p}')

# Conditional extinction tilts; compare complete finite forests, not just first generation.
for a in [F(1,5),F(1,4),F(1,3),F(2,5)]:
 b=1-a;q=a/b;p=[a,F(0),b];pt=[b,F(0),a]
 check(value(p,q)==q and value(pt,1)==1,'tilt normalization')
 for r in [1,2,3]:
  raw=queue_total(p,r,20);tilted=queue_total(pt,r,20)
  for n in range(1,21):check(raw[n]/q**r==tilted[n],f'forest tilt {a,r,n}')
  # Every finite generation history has telescoping Radon--Nikodym weight.
  dist=[F(0)]*r+[F(1)]
  tiltdist=dist[:]
  for n in range(4):
   for j in range(len(dist)):check(dist[j]*q**(j-r)==tiltdist[j],f'history marginal tilt {a,r,n,j}')
   dist=random_sum(p,dist);tiltdist=random_sum(pt,tiltdist)

# Martingale increment budgets checked by full adjacent-generation joint distributions.
for a in [F(1,5),F(1,4),F(1,3),F(2,5)]:
 p=[a,F(0),1-a];m,var=moments(p);constant=var/(m*(m-1))
 for r in [1,2]:
  for n in range(4):
   d=generation(p,r,n);ed=ed2=orth=F(0)
   for z,pz in enumerate(d):
    if not pz:continue
    for nz,pnz in enumerate(power(p,z)):
     w=pz*pnz;delta=F(nz)/m**(n+1)-F(z)/m**n
     ed+=w*delta;ed2+=w*delta*delta;orth+=w*delta*F(z)/m**n
   check(ed==0 and orth==0,'martingale orthogonality')
   check(ed2==r*var/m**(n+2),'martingale increment variance')
   _,vn=moments(d)
   check(vn/m**(2*n)==r*constant*(1-m**(-n)),'martingale finite variance')
 for n in range(16):
  partial=sum(var/m**(j+2) for j in range(n,40))
  check(partial==constant*(m**(-n)-m**(-40)),'exact square tail')

# Immigration: full distributions vs moment recursion and independent age decomposition.
immigration_models=0
for p in [(F(1,2),F(1,2)),(F(3,4),F(0),F(1,4)),(F(1),)]:
 for g in [(F(1,4),F(1,2),F(1,4)),(F(0),F(1)),(F(1),)]:
  immigration_models+=1;m,var=moments(p);eta,tau=moments(g);d=[F(1)];mean=v=F(0);ages=[F(1)]
  for n in range(6):
   check(d==ages,'immigration full age law')
   am,av=moments(d);check(am==mean and av==v,'immigration finite moments')
   check(sum(d)==1,'immigration normalization')
   for s in [F(0),F(1,3),F(2,3),F(1)]:
    fs=s;pg=F(1)
    for j in range(n):pg*=value(g,fs);fs=value(p,fs)
    check(value(d,s)==pg,'immigration PGF product')
   family=random_sum(generation(p,1,n),g)
   ages=conv(ages,family)
   d=conv(random_sum(p,d),g)
   v=m*m*v+var*mean+tau;mean=m*mean+eta
   bound=(var*eta/(1-m)+tau)/(1-m*m)
   check(v<=bound,'stationary finite variance bound')

# Exact Poisson coefficient recurrence vs powers/factorials, after specifying the age rate.
# This is not an independent summation of the infinite thinning mixture.
# The common exp(-rate) factor is removed; thinning itself is proved in the article.
for lam in [F(1,2),F(1),F(3,2),F(2)]:
 for p in [F(0),F(1,4),F(1,2),F(3,4)]:
  for j in range(12):
   rate=lam*p**j
   for k in range(8):
    # exp(rate*(s-1)) coefficient / exp(-rate)
    coeff=F(1)
    for i in range(1,k+1):coeff*=rate/i
    from math import factorial
    check(coeff==rate**k/factorial(k),'Poisson age coefficient')
  for n in range(1,15):
   rate=sum(lam*p**j for j in range(n));tail=lam*p**n/(1-p)
   check(rate+tail==lam/(1-p),'Poisson finite+old rate')

# Common environment exact law and conditional variance discrepancy.
for n in range(21):
 surv=F(3,4)**n;size=2**(n+1);m=F(3,2)
 check(surv*size==2*m**n,'common environment mean')
 check(surv*(F(size)/m**n)**2==4*F(4,3)**n,'common environment square moment')
for z in range(21):
 mean=F(3,2)*z;variance=F(3,4)*(2*z)**2-mean**2
 check(variance==F(3,4)*z*z,'common environment conditional variance')

p=[F(1,4),F(0),F(3,4)];q=F(1,3)
check(generation(p,2,2)[0]==F(361,4096),'capstone two-generation extinction')
check(queue_total(p,2,6)[6]==F(45,4096),'capstone forest six')
check(F(45,4096)/q**2==F(405,4096),'capstone conditional six')
check(200*2**13>3**13 and 200*2**14<=3**14,'capstone square threshold')
check(F(1,1024)<F(1,1000),'capstone immigration n11 upper')
check(F(1,512)-F(1,2*512**2)>F(1,1000),'capstone immigration n10 lower')
check(sum(F(1,2)**j-F(1,4)**j for j in range(25))==F(2)*(1-F(1,2)**25)-F(4,3)*(1-F(1,4)**25),'deterministic immigration variance')
# Certify m=0 and p0=0 boundaries, no spurious finite trees.
check(queue_total([F(0),F(1)],2,14)==[F(0)]*15,'no leaves finite mass')
check(queue_total([F(1)],3,5)[3]==1,'zero reproduction total')
check(generation([F(1)],3,1)==[F(1)],'zero reproduction next generation')

result={'status':'PASS','checks':C,'arithmetic':'Python fractions.Fraction, standard library; no floating point in checks','scope':{'offspring_laws':len(laws),'dwass_initial_roots':[1,2,3],'dwass_max_total':14,'enumerated_offspring_words':enumeration_words,'enumeration_max_total':7,'immigration_models':immigration_models},'capstone':{'two_generation_extinction':'361/4096','eventual_extinction':'1/9','total_six':'45/4096','total_six_given_extinction':'405/4096','conditional_mean_total':'4','conditional_variance_total':'12','limit_mean':'2','limit_variance':'2','mean_square_threshold_generation':14,'common_environment_total_six':'3/16','poisson_stationary_mean':'2','poisson_stationary_variance':'2','deterministic_immigration_stationary_variance':'2/3','immigration_bound_threshold_generation':11},'limitations':['Finite exact checks do not replace proofs of almost-sure, mean-square, or stationary limits.','Generation counts include the explicitly specified initial roots; total progeny includes ancestors.','The immigration threshold certifies a coupling upper bound, not the exact total variation distance.']}
ap=argparse.ArgumentParser();ap.add_argument('--output',type=Path);args=ap.parse_args();text=json.dumps(result,ensure_ascii=False,indent=2)+'\n'
if args.output:args.output.parent.mkdir(parents=True,exist_ok=True);args.output.write_text(text)
print(text,end='')