"""Randomized stopping identifies mechanisms; uncertainty still costs samples.

Edit synthetic inputs below or pass your own six-category counts to inference.py.
No experimental rates or reporter calibrations are claimed by these defaults.
"""
PREPARATION = ['0.65','0.35']
SWITCHING_CONTROL = '0.05'             # S -> T, per hour
SWITCHING_TREATED = '0.85'
SWITCHING_REVERSE = '0.07'
DIVISION_RATES = ['0.22','0.16']
DEATH_RATES = ['0.12','0.04']
SISTER_KERNEL = [['0.62','0.12','0.12','0.14'],['0.10','0.10','0.10','0.70']]
MARKER = [['0.9','0.15'],['0.1','0.85']] # columns S,T; rows marker 0,1
DEADLINE_RATE = '1'                    # per hour, randomized independently per founder
DAUGHTER_DELAYS = [0,0.5,1,2,4,8]      # retain interruption flags
FOUNDERS_PER_ARM = 50_000
CALIBRATION_PER_STATE = 100_000
RANDOM_SEED = 20260922
FEATURE_BIAS_RESERVE = '0'             # only use independently justified bounds
MOVIE_CAP_CLOCK_UNITS = 14             # cap L0 = 14 / deadline rate
PROSPECTIVE_BUDGETS = [140_000,150_000,1_000_000]
EXTINCTION_HOURS = 24
CONCORDANCE_SHIFT = ['0.05','0.05']
MANUSCRIPT_SHA256 = 'a7b2582bd9ee8a5baac720ca0cbb02720cca9660490dc7c9ac944b2cd9ad22fe'

import argparse,csv,hashlib,json,platform
from pathlib import Path
from fractions import Fraction as Q
import numpy as np
import scipy,sympy as sp
from scipy.linalg import expm
from lineage import BranchingSource,MarkerChannel,RandomDeadline,known_mixture_calibration,matrix
from inference import FounderCounts,CalibrationCounts,InductionInference
from certkit import I,log_bounds,kl_lower
from prospective_boxes import decide
from model import strong_source


def source(a):return BranchingSource(PREPARATION,[['0',a],[SWITCHING_REVERSE,'0']],DIVISION_RATES,DEATH_RATES,SISTER_KERNEL)


def serial(x):
    if isinstance(x,sp.MatrixBase):return [[str(v) for v in row] for row in x.tolist()]
    if isinstance(x,(sp.Basic,Q)):return str(x)
    if isinstance(x,np.ndarray):return x.tolist()
    if isinstance(x,np.generic):return x.item()
    if isinstance(x,I):return dict(lower=str(x.a),upper=str(x.b),display=x.pair())
    raise TypeError(type(x))


def main():
    parser=argparse.ArgumentParser(description=__doc__);parser.add_argument('--output',type=Path,default=Path('outputs'));out=parser.parse_args().output;out.mkdir(parents=True,exist_ok=True)
    def dump(name,value):(out/name).write_text(json.dumps(value,indent=2,default=serial)+'\n')
    def table(name,rows):
        with (out/name).open('w',newline='') as f:
            w=csv.DictWriter(f,fieldnames=list(rows[0]));w.writeheader();w.writerows(rows)
    marker=MarkerChannel(MARKER);design=RandomDeadline(DEADLINE_RATE);models=[source(SWITCHING_CONTROL),source(SWITCHING_TREATED)]
    laws=[];recoveries=[]
    for m in models:
        law=design.law(m,marker);recovered=design.recover(law['law'],marker)
        assert recovered['killed']==m.killed and recovered['exit_mass']==m.exits and recovered['preparation']==m.preparation
        laws.append(law);recoveries.append(recovered)
    dump('exact_stopped_laws.json',laws);dump('exact_source_recovery.json',recoveries)
    # Independent founder preparations and independent labelled calibration.
    rng=np.random.default_rng(RANDOM_SEED);cal=tuple(int(x) for x in rng.binomial(CALIBRATION_PER_STATE,[float(marker.emissions[1,j]) for j in range(2)]))
    full=[rng.multinomial(FOUNDERS_PER_ARM,np.array(v['law'],float).ravel()).reshape(2,7) for v in laws]
    counts=[FounderCounts.from_full(c) for c in full];calibration=CalibrationCounts(cal,CALIBRATION_PER_STATE)
    dump('retained_data.json',dict(calibration_successes=cal,calibration_per_state=CALIBRATION_PER_STATE,founders_per_arm=FOUNDERS_PER_ARM,full_counts=full,six_category_counts=[c.deadline+c.exits for c in counts],seed=RANDOM_SEED,scope='Synthetic independent roots, all event outcomes retained; coarsening D means death OR division.'))
    routes={}
    for scheme in ['hoeffding','kl']:
        for evaluator in ['matrix','cancelled']:
            r=InductionInference(Q(DEADLINE_RATE),scheme,evaluator,Q(FEATURE_BIAS_RESERVE)).contrast(*counts,calibration);routes[scheme+'_'+evaluator]=r
            print(scheme,evaluator,r.get('display',r['status']),flush=True)
    dump('induction_intervals.json',routes)
    endpoint_audit=[]
    engine=InductionInference()
    for c in counts:
        for k in c.features():
            n=c.total;L=Q(25,4);s=engine.feature(k,n,L)
            endpoint_audit.append(dict(k=k,n=n,L=str(L),lower=str(s.a),upper=str(s.b),lower_exclusion_margin=str(n*kl_lower(Q(k,n),s.a)-L) if s.a else 'boundary',upper_exclusion_margin=str(n*kl_lower(Q(k,n),s.b)-L) if s.b<1 else 'boundary'))
    for k in cal:
        n=CALIBRATION_PER_STATE;L=Q(6);s=engine.feature(k,n,L)
        endpoint_audit.append(dict(k=k,n=n,L=str(L),lower=str(s.a),upper=str(s.b),lower_exclusion_margin=str(n*kl_lower(Q(k,n),s.a)-L),upper_exclusion_margin=str(n*kl_lower(Q(k,n),s.b)-L)))
    dump('rational_endpoint_audit.json',endpoint_audit)
    small=FounderCounts((12,1,1,3),(1,2));dump('unresolved_small_sample.json',{s+'_'+e:InductionInference(Q(DEADLINE_RATE),s,e).contrast(small,small,calibration) for s in ['hoeffding','kl'] for e in ['matrix','cancelled']})
    prospective={}
    for n in PROSPECTIVE_BUDGETS:
        for scheme,method in [('hoeffding','matrix'),('kl','cancelled')]:
            prospective[f'{scheme}_{method}_{n}']=decide(Q(1),n,n,scheme,method)
    dump('prospective_continuous_boxes.json',dict(results=prospective,scope='Fixed paper class in prospective_boxes.py, independent of edited synthetic source inputs; full control interval enclosed. Sufficient budgets, not necessary sample sizes.'))
    delay=[];delayed=[]
    for tau in DAUGHTER_DELAYS:
        protocol=RandomDeadline(DEADLINE_RATE,tau);law=protocol.law(models[1],marker);r=protocol.recover(law['law'],marker)
        F=np.array(law['channel'],float);s=float(np.linalg.svd(F,compute_uv=False)[-1]);error=float(np.max(np.abs(np.array(r['killed'],float)-np.array(models[1].killed,float))))
        delay.append(dict(delay_hours=tau,one_daughter_min_singular=s,pair_min_singular=s*s,founder_hours=float(law['founder_hours']),extra_daughter_hours_upper=2*tau*float(law['division_yield']),recovered_generator_error=error))
        delayed.append(dict(delay=tau,recovered=r,law=law['law'],channel=law['channel']))
    table('delayed_channel_conditioning.csv',delay);dump('delayed_source_recovery.json',delayed)
    P=matrix([['4/5','1/4'],['1/5','3/4']]);mixed=known_mixture_calibration(marker.emissions*P,P);assert mixed.emissions==marker.emissions
    dump('known_mixture_calibration.json',dict(compositions=P,marker_distributions=marker.emissions*P,recovered=mixed.emissions,scope='Compositions independently established at readout; not gates defined by the same uncalibrated marker.'))
    # Exact tests beyond two states: rectangular channel, complex spectrum,
    # reducibility, zero division, and a singular killed generator.
    exact_cases=[]
    for name,q,b,d in [('cycle',[[0,2,0],[0,0,2],[2,0,0]],['1/2']*3,['1/2']*3),('singular_reducible',[[0,0,0],[0,0,1],[0,0,0]],[0,0,0],[0,1,1])]:
        K=sp.zeros(3,9)
        for i in range(3):K[i,4*i]=1
        m=BranchingSource(['1/3']*3,q,b,d,K)
        E=MarkerChannel([['3/4',0,0],[0,'3/4',0],[0,0,'3/4'],['1/4']*3]);p=RandomDeadline('3/2');law=p.law(m,E);r=p.recover(law['law'],E)
        assert r['killed']==m.killed and r['exit_mass']==m.exits
        exact_cases.append(dict(name=name,source_killed=m.killed,recovery=r))
    dump('finite_state_exact_examples.json',exact_cases)
    m=source('0.18');shifted=m.concordance_shift(CONCORDANCE_SHIFT)
    assert m.mean_generator()==shifted.mean_generator()
    strong=strong_source();big=BranchingSource(strong['pi'],strong['q'],strong['b'],strong['d'],strong['K']);bigshift=big.concordance_shift(['.18','.18'])
    curves=[];t=np.linspace(0,EXTINCTION_HOURS,241)
    for label,base,more in [('baseline',m,shifted),('large_synthetic',big,bigshift)]:
        f,p,mean=base.population(t);g,p2,mean2=more.population(t)
        assert np.max(np.abs(mean-mean2))<1e-10 and np.min(g-f)>-1e-10
        curves.extend(dict(source=label,hours=float(x),extinction=float(a),greater_concordance_extinction=float(b),identical_mean=float(c)) for x,a,b,c in zip(t,p,p2,mean))
    table('extinction_and_means.csv',curves)
    eps=matrix(CONCORDANCE_SHIFT);d=m.death
    dump('concordance_identity.json',dict(mean_generator=m.mean_generator(),shifted_mean_generator=shifted.mean_generator(),cubic_coefficients=[m.division[i]*eps[i]*(d[0]-d[1])**2/3 for i in range(2)],offspring_difference='epsilon_i * (x_S - x_T)^2',scope='Polynomial identity and mean equality exact; finite-time effect magnitudes numerical. Equality can hold with state-independent demographic rates.'))
    lam=sp.Rational(DEADLINE_RATE);mu=2*lam;H=models[1].killed;R=(lam*sp.eye(2)-H).inv();S=(mu*sp.eye(2)-H).inv();assert R-S==(mu-lam)*R*S
    cap=Q(MOVIE_CAP_CLOCK_UNITS)/Q(DEADLINE_RATE)
    dump('deadline_design_checks.json',dict(resolvent_identity_residual=R-S-(mu-lam)*R*S,erlang_two_apparent_birth_rates=['1/3','1/4'],erlang_deadline_rates=[1,2],movie_cap_hours=str(cap),movie_cap_TV_reserve_display=float(np.exp(-float(lam)*float(cap))),scope='Capped record law is not a resolvent. A justified reserve can widen features; no survivor filtering or repeated-trace pseudoreplication.'))
    plot(out,routes,delay,curves)
    print('Exact immediate population inverse works for all retained states; delayed inversion retains interruptions.',flush=True)
    print('Prospective full-box 150000:',prospective.get('kl_cancelled_150000'),flush=True)
    print('All inputs are synthetic. Exact interval endpoints are audited at n*KL >= L. Numerical extinction is not an interval certificate; 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(),numpy=np.__version__,scipy=scipy.__version__,sympy=sp.__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,routes,delay,curves):
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
    fig,axs=plt.subplots(1,2,figsize=(11,4),layout='constrained')
    for i,(name,r) in enumerate(routes.items()):
        if r['status']=='resolved':axs[0].plot(r['display'],[i,i],'o-',lw=3)
    axs[0].set(yticks=range(4),yticklabels=['Hoeffding / matrix','Hoeffding / cancelled','KL / matrix','KL / cancelled'],xlabel='Treated minus control switching (per hour)',title='Switching-rate intervals from the same\nobservations');axs[0].axvline(float(Q(SWITCHING_TREATED)-Q(SWITCHING_CONTROL)),color='gray',ls=':',label='Synthetic truth');axs[0].legend(fontsize=8)
    axs[1].semilogy([r['delay_hours'] for r in delay],[r['one_daughter_min_singular'] for r in delay],'o-',label='One daughter');axs[1].semilogy([r['delay_hours'] for r in delay],[r['pair_min_singular'] for r in delay],'s-',label='Joint pair');axs[1].set(xlabel='Monitored delay (hours)',ylabel='Smallest singular value',title='Sensitivity of daughter-state inference to\nreadout delay');axs[1].legend(fontsize=8)
    for a in axs:a.grid(alpha=.2)
    fig.savefig(out/'inference_and_delay.png',dpi=180);fig.savefig(out/'inference_and_delay.svg');plt.close(fig)
    fig,axs=plt.subplots(1,2,figsize=(10,4),layout='constrained')
    for label in ['baseline','large_synthetic']:
        rows=[r for r in curves if r['source']==label];t=[r['hours'] for r in rows]
        axs[0].semilogy(t,[r['identical_mean'] for r in rows],label=label.replace('_',' '))
        axs[1].plot(t,[r['greater_concordance_extinction']-r['extinction'] for r in rows],label=label.replace('_',' '))
    axs[0].set(xlabel='Hours',ylabel='Expected total cells',title='Identical means within each kernel pair');axs[1].set(xlabel='Hours',ylabel='Increase in extinction probability',title='Extinction increase due to sister\nconcordance')
    for a in axs:a.grid(alpha=.2);a.legend(fontsize=8)
    fig.savefig(out/'inheritance_extinction.png',dpi=180);fig.savefig(out/'inheritance_extinction.svg');plt.close(fig)


if __name__=='__main__':main()