"""Preparation-aware small-population prediction with explicit observation contracts."""
from fractions import Fraction as F
from pathlib import Path
import argparse
import csv
import hashlib
import json
import math
import platform
import numpy as np
from models import FounderPreparation,TwoTypeBranching,BirthDeathEnvelope
from prediction import MenuRule,DetectorFeasibleRule,ObservationBudget,InheritedRateClass,dependence_decision,scalar_finite_certificate,unresolved_coverage

# EDITABLE SYNTHETIC INPUTS. The three contracts have different horizons/founder laws.
MENU_DEPENDENCE=F(0)
MENU_CALIBRATION_WELLS=640
DETECTION_EFFICIENCY=F(1,2)
DETECTOR_CALIBRATION_WELLS=2000
SEED=56092026
DETECTOR_TRUNCATION=10
BOUNDED_STATISTIC_Z=F(1,2)
MULTIPLE_FOUNDER_BUDGET=F(1,1000)
FALSE_OBJECT_BUDGET=F(1,100)
REFERENCE_RATES=(0.,.3,.001,.1,0.,.00001)  # births, deaths, switches; S then R; per day
RATE_L1_RADIUS=F(1,2000)
FOUNDER_R_PROBABILITY=F(5,24)
FOUNDER_R_RADIUS=F(1,1000)
INHERITED_HORIZON=math.log(2)/.10001
FORWARD_CAP=40
MANUSCRIPT_SHA256='e2b902e1d6f5aa28e24dc1d2bdca97e512cd3b70bff5b45652f2b31a1d2ea6eb'


def main():
    parser=argparse.ArgumentParser();parser.add_argument('--output',default='outputs');args=parser.parse_args()
    out=Path(args.output);out.mkdir(parents=True,exist_ok=True)
    independent,shared=FounderPreparation(F(0)),FounderPreparation(F(1))
    cdfrows=[[k,independent.cdf_two(k),shared.cdf_two(k)] for k in range(2,15)]
    calibration=[]
    for event,n in [('at_most_two',1024),('at_most_two',1280),('at_most_two',4000),('three_or_four',576),('three_or_four',640)]:
        rule=MenuRule(event);r=rule.exact_errors(n)
        target=F(19,20) if n in [1024,576] else F(245,256)
        assert r['uniform_coverage']>=target
        calibration.append(dict(event=event,n=n,threshold=rule.threshold(n),W_error=float(r['W_error']),I_error=float(r['I_error']),
                                coverage_display=float(r['uniform_coverage']),exact_target=target,exact_comparison_passed=True))
    menu=FounderPreparation(MENU_DEPENDENCE)
    latent,perfect=menu.sample(MENU_CALIBRATION_WELLS,seed=SEED)
    H=int(np.count_nonzero((perfect==3)|(perfect==4)))
    decision=MenuRule().choose(H,MENU_CALIBRATION_WELLS,perfect_counts=True)
    dependence=dependence_decision(H,MENU_CALIBRATION_WELLS)
    detection_latent,detected=menu.sample(DETECTOR_CALIBRATION_WELLS,detection=float(DETECTION_EFFICIENCY),seed=SEED+1)
    feasible=DetectorFeasibleRule(DETECTOR_TRUNCATION,BOUNDED_STATISTIC_Z).evaluate(list(map(int,detected)))
    # Fixed-menu rules do not transfer to the continuous dependence family.
    if MENU_DEPENDENCE not in (F(0),F(1)):
        decision=dict(cutoff=7,status='OUTSIDE_BINARY_MENU_FALLBACK')
        feasible=dict(cutoff=7,status='OUTSIDE_BINARY_MENU_FALLBACK',retained_models=[],
                      reason='Use the separate perfect-count dependence confidence rule for intermediate rho')
    detectrows=[]
    for i in range(101):
        d=F(i,100);pi=independent.detected_coefficients(2,d);pw=shared.detected_coefficients(2,d)
        mi,mw=independent.pgf(F(1,2),detection=d),shared.pgf(F(1,2),detection=d)
        detectrows.append([d,sum(pw)-sum(pi),mw-mi,mi,mw])
    many=[]
    for m in [1,2,5,10,20,50,100,500]:
        k=m
        while independent.many_cdf_numeric(k,m)<.95:k+=1
        many.append([m,k,shared.many_cdf_numeric(k,m),1.25*m,m+m*m/4])
    budget=ObservationBudget(MULTIPLE_FOUNDER_BUDGET,FALSE_OBJECT_BUDGET)
    finite=scalar_finite_certificate()
    demographic=[]
    for birth,death,k in [(.1002,.05,4),(.105,.05,4),(.1,0.,5)]:
        envelope=BirthDeathEnvelope(birth,death)
        exact=envelope.exact_supercritical_lower(k)
        demographic.append(dict(birth_cap=birth,death_floor=death,cutoff=k,source_lower=exact,recorded_lower=budget.recorded(exact),
                                numerical_full_cdf=envelope.cdf(k)))
    rateclass=InheritedRateClass(RATE_L1_RADIUS,FOUNDER_R_PROBABILITY,FOUNDER_R_RADIUS)
    history=[]
    for k in [2,3,4,5,8]:
        bounds=rateclass.source_lower(k);B=budget.recorded(bounds['lower'])
        history.append(dict(cutoff=k,source_lower=bounds['lower'],recorded_lower=B,with_1pct_conditional_inference_failure=F(99,100)*B))
    forward=[]
    for label,rates in [('reference',REFERENCE_RATES),('positive_example',(.0001,.2999,.0011,.1001,.0001,.00002))]:
        model=TwoTypeBranching(*rates);p,overflow=model.killed_distribution(INHERITED_HORIZON,FORWARD_CAP,float(FOUNDER_R_PROBABILITY))
        forward.append(dict(label=label,in_rate_class=rateclass.contains(rates),cdf2=float(p[:3].sum()),cdf3=float(p[:4].sum()),
                            overflow=overflow,normalization=float(p.sum()+overflow)))
    witness=rateclass.reference_endpoint_two_failure_lower();assert witness>F(1,20)
    detector_half=dict(low_count_gap=detectrows[50][1],pgf_gap=detectrows[50][2],
                       sufficient_pgf_wells_numeric=math.ceil(2*math.log(256/3)/float(detectrows[50][2])**2),
                       population_moment_efficiency='E[Y]/3',population_moment_dependence='18 E[Y(Y-1)] / E[Y]^2 - 17; d>0 only')
    results=dict(inputs=dict(menu_dependence=MENU_DEPENDENCE,detection=DETECTION_EFFICIENCY,seed=SEED,reference_rates=REFERENCE_RATES,rate_radius=RATE_L1_RADIUS),
                 calibration_rows=calibration,synthetic_menu=dict(n=MENU_CALIBRATION_WELLS,H=H,decision=decision),
                 synthetic_continuous_dependence=dependence,unknown_detector_rule=feasible,detector_half=detector_half,
                 finite_scalar_certificate=finite,compiled_recorded_floor=budget.recorded(F(963,1000)),demographic=demographic,
                 inherited_history_bounds=history,reference_endpoint_two_failure_witness=witness,numerical_forward=forward,
                 unresolved_outcomes_example=unresolved_coverage(51,24,75),
                 evidence='Exact source algebra and integer bounds; numeric branching/Monte Carlo illustrations; supplied contracts only; Lean not rerun')
    dump=lambda name,obj:(out/name).write_text(json.dumps(obj,indent=2,default=lambda v:str(v) if isinstance(v,F) else float(v))+'\n',encoding='utf-8')
    def table(name,header,rows):
        with (out/name).open('w',newline='') as f:w=csv.writer(f);w.writerow(header);w.writerows(rows)
    dump('results.json',results)
    table('cutoffs.csv',['cutoff','independent_cdf','shared_cdf'],cdfrows)
    table('calibration.csv',list(calibration[0]),[[r[k] for k in calibration[0]] for r in calibration])
    table('detection.csv',['efficiency','low_count_gap','pgf_gap','pgf_I','pgf_W'],detectrows)
    table('many_founders.csv',['founders','I_min95_cutoff_numeric','W_coverage_numeric','variance_I','variance_W'],many)
    table('synthetic_detector_wells.csv',['well','latent_count','detected_count'],zip(range(1,len(detected)+1),detection_latent,detected))
    table('synthetic_perfect_wells.csv',['well','count'],enumerate(perfect,1))
    table('history.csv',list(history[0]),[[r[k] for k in history[0]] for r in history])
    lines=['Two-founder exact minimal 95% cutoffs: independent 6; shared 7; dependence boundary 3/5.',
           f'Perfect-count synthetic calibration: H={H}/{MENU_CALIBRATION_WELLS}; menu cutoff {decision["cutoff"]}.',
           f'Unknown-efficiency bounded-statistic rule retains {feasible["retained_models"]}; cutoff {feasible["cutoff"]}.',
           f'Inherited-class recorded bound at endpoint 3: {float(history[1]["recorded_lower"]):.12f}.',
           'All source, preparation and detector contracts are hypothetical; recorded safety does not establish latent eradication.']
    (out/'console.txt').write_text('\n'.join(lines)+'\n');print('\n'.join(lines))
    plot(out,cdfrows,detectrows,calibration,history,many)
    digest=lambda p:hashlib.sha256(p.read_bytes()).hexdigest();here=Path(__file__).parent
    dump('run_metadata.json',dict(manuscript_sha256=MANUSCRIPT_SHA256,python=platform.python_version(),source_sha256=digest(Path(__file__)),
         module_sha256={n:digest(here/n) for n in ['models.py','prediction.py']},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,cdfrows,detectrows,calibration,history,many):
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
    fig,axs=plt.subplots(1,2,figsize=(10,4),layout='constrained')
    data=np.array(cdfrows,dtype=float)
    for j,name in [(1,'Independent founders'),(2,'Shared class')]:axs[0].step(data[:,0],data[:,j],where='post',label=name)
    axs[0].axhline(.95,color='k',ls=':');axs[0].scatter([6,7],[245/256,31/32],color=['tab:blue','tab:orange'])
    axs[0].set(xlim=(4,10),ylim=(.87,1),title='Two-founder count coverage by preparation',xlabel='Upper count endpoint',ylabel='Two-founder coverage');axs[0].legend(fontsize=8)
    data=np.array(many,dtype=float);axs[1].semilogx(data[:,0],data[:,2],'o-')
    axs[1].axhline(.95,color='k',ls=':',label='Nominal target');axs[1].axhline(.5,color='gray',ls='--',label='Limit under shared preparation')
    axs[1].set(title='Coverage when a cutoff is transferred\nbetween preparations',xlabel='Number of founders',ylabel='Shared-law coverage of independent cutoff',ylim=(.48,1));axs[1].legend(fontsize=8)
    for ax in axs:ax.grid(alpha=.2)
    fig.savefig(out/'preparation.png',dpi=180);fig.savefig(out/'preparation.svg');plt.close(fig)
    fig,axs=plt.subplots(1,2,figsize=(10,4),layout='constrained');data=np.array(detectrows,dtype=float)
    axs[0].plot(data[:,0],data[:,1],label='Count ≤ 2');axs[0].plot(data[:,0],data[:,2],label='Bounded statistic 2⁻ʸ')
    axs[0].axhline(0,color='k',lw=.7);axs[0].axvline(.5,color='gray',ls=':')
    axs[0].set(title='Calibration-statistic separation versus\ndetection efficiency',xlabel='Independent detection efficiency',ylabel='Expectation gap: shared − independent');axs[0].legend(fontsize=8)
    x=np.arange(3);values=[float(next(r for r in history if r['cutoff']==k)['recorded_lower']) for k in [2,3,4]]
    axs[1].bar(x,values,label='Full stated contract');axs[1].scatter(x,np.array(values)*.99,color='k',marker='x',label='With 1% conditional inference failure')
    axs[1].axhline(.95,color='gray',ls=':');axs[1].set_xticks(x,['2','3','4']);axs[1].set(ylim=(.9,.98),title='Count-coverage bounds for the\ninherited-state model',xlabel='Recorded upper endpoint',ylabel='Exact sufficient coverage lower bound');axs[1].legend(fontsize=7,loc='lower right')
    for ax in axs:ax.grid(alpha=.2)
    fig.savefig(out/'measurement.png',dpi=180);fig.savefig(out/'measurement.svg');plt.close(fig)


if __name__=='__main__':main()