"""Same daughter marginal and mean growth, different population risks. Run: python example.py"""
# EDITABLE EXPLORATION INPUTS. Rates are illustrative per time unit, not clinical calibration.
SITES = 2
FOUNDER = (2, 0)                 # (active, repressive); inactive = SITES - sum(FOUNDER)
READOUT = 'step'                # 'step' or 'smooth'
SMOOTH_SLOPE = 8.
WRITING = .01
RECRUITMENT = 1.
ERASURE = .01
ANTAGONISM = .5
DIVISION_RATE = .1              # mean interdivision time = 1 / DIVISION_RATE
ERLANG_PHASES = 1               # state-independent clock; 1, 2, 4 are paper examples
PROTECTED_DEATH = .01
EXPOSED_DEATH = .3
HORIZON = 60.
CAPTURE_PROBABILITY = .2         # independent cell capture
REGROWTH_THRESHOLD = 300        # fixed two-site certificate, separate from edited model
OBSERVED_PAIR_COVARIANCE = -.025
INDEPENDENT_ENROLLED_PAIRS = 100000
MOTHER_MEAN_RANGE = .1
MISCLASSIFICATION = .0001
MISSING_FRACTION = .0001
CONFIDENCE_ALPHA = .05

from pathlib import Path
from fractions import Fraction as F
from math import comb
import argparse,csv,hashlib,json,platform,sys
import numpy as np
from branching import (MolecularSource,StepHazard,SmoothHazard,Complementary,Independent,
                       BranchingPopulation,IndependentCapture,PairedCloneObservation,balanced_survival_bound)
from certificates import TwoSiteCertificate,SmoothSlopeCertificate,DeadlineCertificate

MANUSCRIPT_SHA256='2ae028e1227c6f80f7da47df7a035ea41fd4c68508e68112f623ae0b923cdc4b'


def serialize(value):
    if isinstance(value,F):return str(value)
    if isinstance(value,np.ndarray):return value.tolist()
    if isinstance(value,np.generic):return value.item()
    raise TypeError(type(value).__name__)


def main():
    sys.set_int_max_str_digits(0)  # Serialize our own exact threshold powers (not untrusted input).
    ap=argparse.ArgumentParser(description=__doc__);ap.add_argument('--output',default='outputs');args=ap.parse_args()
    out=Path(args.output);out.mkdir(parents=True,exist_ok=True)
    def dump(name,value):(out/name).write_text(json.dumps(value,indent=2,default=serialize)+'\n',encoding='utf-8')
    def table(name,headers,rows):
        with (out/name).open('w',newline='',encoding='utf-8') as f:
            w=csv.writer(f);w.writerow(headers);w.writerows(rows)
    print('Rebuilding the two-site rational risk and covariance-response boxes...',flush=True)
    two=TwoSiteCertificate();exact=two.calculate();dump('two_site_certificate.json',exact)
    regrowth=two.regrowth(REGROWTH_THRESHOLD);dump('eventual_regrowth.json',regrowth)
    print('Rebuilding 41 smooth-readout slope cells, with full 153-state witnesses...',flush=True)
    smooth=SmoothSlopeCertificate();narrow=smooth.calculate(F(7999,1000),F(8001,1000),1)
    dump('smooth_narrow_certificate.json',narrow)
    wide=smooth.calculate();dump('smooth_wide_certificate.json',wide)
    assert narrow['gap_lower']>F(1427,100000) and wide['gap_lower']>F(128,10000)
    assert narrow['all_active_gap_upper']<F(122,1000000) and wide['regrowth_2000_gap_lower']>F(122,10000)
    print('Reintegrating both 300-time-unit PGFs with outward integer Taylor arithmetic...',flush=True)
    finite=DeadlineCertificate().calculate();finite['gap_against_independent_eventual_upper']=finite['joint_hit_by_deadline_lower']-two.regrowth(300)['independent_eventual_upper']
    assert finite['gap_against_independent_eventual_upper']>F(84,1000)
    dump('deadline_certificate.json',finite)
    # Edited model: no automatic transfer of the fixed source certificates.
    source=MolecularSource(SITES,WRITING,RECRUITMENT,ERASURE,ANTAGONISM)
    if FOUNDER not in source.index:raise ValueError('founder must be a state of the configured source')
    if HORIZON<=0 or not np.isfinite(HORIZON):raise ValueError('horizon must be positive and finite')
    if READOUT not in ['step','smooth']:raise ValueError('choose step or smooth readout')
    hazard=StepHazard(PROTECTED_DEATH,EXPOSED_DEATH) if READOUT=='step' else SmoothHazard(SMOOTH_SLOPE,PROTECTED_DEATH,EXPOSED_DEATH)
    times=np.linspace(0,HORIZON,121);k=source.index[FOUNDER];results={};capture=IndependentCapture(CAPTURE_PROBABILITY)
    for name,law in [('joint',Complementary()),('independent',Independent())]:
        population=BranchingPopulation(source,hazard,law,DIVISION_RATE,ERLANG_PHASES)
        q,it=population.extinction();pgf=population.pgf(times);mean,var=population.moments(times);zero=capture.nondetection(population,times);cm,cv=capture.moments(mean,var)
        results[name]=dict(q=q,pgf=pgf,mean=mean,var=var,zero=zero,capture_mean=cm,capture_variance=cv,iterations=it)
        table(name+'_time_course.csv',['time','extinction','mean','variance','capture_nondetection','capture_mean','capture_variance'],zip(times,pgf[:,k],mean[:,k],var[:,k],zero[:,k],cm[:,k],cv[:,k]))
        table(name+'_eventual_risk.csv',['active','repressive','extinction_numerical'],[(a,r,q[i]) for i,(a,r) in enumerate(source.states)])
    mean_error=float(np.max(abs(results['joint']['mean']-results['independent']['mean'])/np.maximum(1,results['joint']['mean'])))
    applicable=(SITES==2 and READOUT=='step' and (WRITING,RECRUITMENT,ERASURE,ANTAGONISM,DIVISION_RATE,ERLANG_PHASES,PROTECTED_DEATH,EXPOSED_DEATH)==(.01,1.,.01,.5,.1,1,.01,.3))
    dump('configured_model.json',dict(sites=SITES,founder=FOUNDER,readout=READOUT,division_phases=ERLANG_PHASES,horizon=HORIZON,capture=CAPTURE_PROBABILITY,fixed_two_site_certificate_applicable=applicable,
        eventual_risks={name:float(r['q'][k]) for name,r in results.items()},relative_common_mean_residual=mean_error,
        scope='Floating-point least-fixed-point and ODE diagnostics. Mean and single-daughter kernel agree by construction. No finite-threshold probability is inferred from mean or variance; no arbitrary first-passage ordering is claimed.'))
    # Preparation and clock are distinct axes, explored with the full 153-state model.
    prep=[];clock=[]
    for N in [2,4,8,16]:
        src=MolecularSource(N)
        for readout,h in [('step',StepHazard()),('smooth',SmoothHazard())]:
            q=[BranchingPopulation(src,h,law).extinction()[0] for law in [Complementary(),Independent()]]
            for label,st in [('all active',(N,0)),('balanced',(N//2,N//2))]:
                j=src.index[st];prep.append((N,readout,label,float(q[0][j]),float(q[1][j]),float(q[1][j]-q[0][j])))
    src=MolecularSource(16);j=src.index[9,7]
    for m in [1,2,4]:
        q=[BranchingPopulation(src,SmoothHazard(),law,phases=m).extinction()[0][j] for law in [Complementary(),Independent()]]
        clock.append((m,*q))
    table('preparation_sweep.csv',['sites','readout','preparation','joint_extinction','independent_extinction','gap'],prep)
    table('division_clock.csv',['phases','joint_extinction','independent_extinction'],clock)
    # Exact observable counterexamples and conservative observation uncertainty.
    pooled=F(7,64);aa=F(-1,16);nonmonotone=F(1,4)
    assert pooled==F(1,4)-F(3,8)**2 and aa==F(1,2)-F(3,4)**2
    observation=PairedCloneObservation(MOTHER_MEAN_RANGE,MISCLASSIFICATION,MISSING_FRACTION)
    obs=observation.interval(OBSERVED_PAIR_COVARIANCE,INDEPENDENT_ENROLLED_PAIRS,CONFIDENCE_ALPHA)
    sister=Complementary();within=[]
    src2=MolecularSource(2);pop2=BranchingPopulation(src2,StepHazard(),sister);f=pop2.pgf([60.])[0];cov=sister.product(src2,f,f)-(src2.D@f)**2
    for i,st in enumerate(src2.states):within.append((*st,float(cov[i])))
    table('conditional_sister_clone_covariance.csv',['mother_active','mother_repressive','covariance_at_60'],within)
    dump('observation_accounts.json',dict(exact_newborn_indicator_covariance_AA=aa,exact_equal_AA_RR_mixture_covariance=pooled,nonmonotone_continuation_counterexample=nonmonotone,
        configured_sensitivity_and_sampling_interval=obs,negative_within_mother_covariance_supported=obs[1]<0,confidence=1-CONFIDENCE_ALPHA,
        scope='Prospective independent sister pairs are sampling units. Range, misclassification and missingness bounds are assumed. A positive pooled covariance does not rule out complementary allocation; independent capture nondetection is not extinction.'))
    asym=[]
    for logN in [np.log(16),np.log(128),100.,500.,1000.]:
        t=logN/2.12;asym.append((logN,t,balanced_survival_bound(logN,t)))
    table('smooth_large_size_bound.csv',['log_sites','chosen_time','survival_upper_bound_unclipped'],asym)
    tv=[(n,str(1-F(comb(2*n,n),4**n))) for n in [2,4,8,16,32,64]]
    dump('allocation_distance.json',dict(all_active_total_variation=tv,scope='Raw pair-law distance grows toward one even where continuation-risk differences shrink. No uniform inverse-site error follows without a continuation regularity bound.'))
    plot(out,results,k,times,prep,clock,wide,regrowth,finite)
    print(f'Two-site AA eventual risks: J {float(exact["lower"]["J"][5]):.7f}..{float(exact["upper"]["J"][5]):.7f}; I {float(exact["lower"]["I"][5]):.7f}..{float(exact["upper"]["I"][5]):.7f}.',flush=True)
    print(f'Hit 300 by time 300: J > {float(finite["joint_hit_by_deadline_lower"]):.6f}; I < {float(two.regrowth(300)["independent_eventual_upper"]):.6f}.',flush=True)
    print(f'N=16, slope 7..9: extinction gap > {float(wide["gap_lower"]):.6f}; hit-2000 gap > {float(wide["regrowth_2000_gap_lower"]):.6f}.',flush=True)
    print(f'Edited scenario covered by fixed two-site source: {applicable}. All exact certificates rebuilt; numerical sweeps are diagnostics; Lean not rerun.',flush=True)
    here=Path(__file__).parent;digest=lambda p:hashlib.sha256(p.read_bytes()).hexdigest()
    dump('run_metadata.json',dict(manuscript_sha256=MANUSCRIPT_SHA256,source_sha256=digest(Path(__file__)),python=platform.python_version(),module_sha256={p.name:digest(p) for p in sorted(here.glob('*.py')) if p.name not in ['example.py','test_example.py']},input_sha256={p.name:digest(p) for p in sorted(here.glob('*.json')) if p.name!='release.json'},output_sha256={p.name:digest(p) for p in sorted(out.iterdir()) if p.is_file() and p.name!='run_metadata.json'}))


def plot(out,results,k,times,prep,clock,wide,regrowth,finite):
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
    fig,axs=plt.subplots(1,3,figsize=(12,4),layout='constrained')
    for name,r in results.items():
        axs[0].plot(times,r['mean'][:,k],label=name)
        axs[1].plot(times,r['pgf'][:,k],label=name)
        axs[2].plot(times,r['zero'][:,k],label=name)
    for ax,title,ylabel in zip(axs,['Expected population count','Population extinction probability','Zero-count probability after incomplete\ncapture'],['Expected cells','Extinction probability','Zero captured cells: probability']):
        ax.set(xlabel='Time',ylabel=ylabel,title=title);ax.grid(alpha=.2);ax.legend(fontsize=8)
    fig.savefig(out/'matched_mean_different_risk.png',dpi=180);fig.savefig(out/'matched_mean_different_risk.svg');plt.close(fig)
    fig,axs=plt.subplots(1,3,figsize=(12,4),layout='constrained')
    for readout in ['step','smooth']:
        for label in ['all active','balanced']:
            rows=[r for r in prep if r[1]==readout and r[2]==label]
            axs[0].semilogy([r[0] for r in rows],[r[5] for r in rows],'o-',label=readout+', '+label)
    axs[0].set(xlabel='Site count',ylabel='Extinction gap I − J',title='Extinction difference by founder\npreparation and site count');axs[0].legend(fontsize=7)
    rows=wide['rows'];axs[1].stairs([float(r['gap_lower']) for r in rows],[float(rows[0]['slope'][0])]+[float(r['slope'][1]) for r in rows],baseline=None,label='Fresh rational lower bound')
    axs[1].axhline(.0128,color='gray',ls=':');axs[1].set(xlabel='Smooth-readout slope',ylabel='Extinction gap: founder (9, 7)',title='A uniform 153-state certificate');axs[1].legend(fontsize=7)
    for i,label in [(1,'joint'),(2,'independent')]:axs[2].plot([r[0] for r in clock],[r[i] for r in clock],'o-',label=label)
    axs[2].set(xlabel='Erlang phases (mean time 10)',ylabel='Extinction probability',title='Extinction probability versus\ndivision-clock phases');axs[2].legend(fontsize=8)
    for ax in axs:ax.grid(alpha=.2)
    fig.savefig(out/'preparation_and_clock.png',dpi=180);fig.savefig(out/'preparation_and_clock.svg');plt.close(fig)


if __name__=='__main__':main()