"""A reusable finite-count exclusion design, with sharp adversarial witnesses."""
from fractions import Fraction as F
from pathlib import Path
import argparse,csv,hashlib,json,platform
import numpy as np
from model import SpecimenBudget,SourceLaw,ExclusionPolicy,gate_bound,minimum_pairs,moment_bound
from certificates import batch_certificate,observation_bound,specificity_bound,PoissonReference,constrained_certificate,constrained_value,exp_negative

# EDITABLE INPUTS: designed mathematical example, not measured assay performance.
ORIGINAL_ML='1'
AVAILABLE_ML='1'
LOSS_PER_PREPARATION_ML='0.05'
SPLIT='0.5'
COUNT_THRESHOLD=6
REFERENCE_PAIRS=9
FALSE_EXCLUSION_ALLOWANCE='0.05'
RECOVERY_MEAN='0.8'             # illustrative independent Bernoulli source only
CONTROL_LAW_TV='0.0001'        # whole pair-law error, assumed, not estimated here
TARGET_LAW_TV='0.0001'
REFERENCE_LOADING='2.5'         # mean native targets per Poisson reference input
INTERVAL_BUDGET=100000
PAIR_SEARCH_BUDGET=4000
SIMULATION_REPLICATES=2000
SEED=60092026
MANUSCRIPT_SHA256='1dc510ae7b7bfe6e68deeb68dd67fbb4c0f39d25d340499b2feab171dc34f530'


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)
    budget=SpecimenBudget(F(ORIGINAL_ML),F(AVAILABLE_ML),F(LOSS_PER_PREPARATION_ML));a,b=budget.fractions(F(SPLIT))
    alpha=F(FALSE_EXCLUSION_ALLOWANCE);policy=ExclusionPolicy(a,b,COUNT_THRESHOLD,REFERENCE_PAIRS,alpha)
    H=max((1-a)**COUNT_THRESHOLD,(1-b)**COUNT_THRESHOLD);g,z=gate_bound(H,REFERENCE_PAIRS)
    # Least-favourable law concentrates single-good specimens on smaller aliquot.
    witness=SourceLaw.corners(z if a<=b else 0,z if b<a else 0,0)
    assert witness.joint_error(COUNT_THRESHOLD,a,b,REFERENCE_PAIRS)==g
    p=F(RECOVERY_MEAN);synthetic=SourceLaw.corners(p,p,p*p)
    def dump(name,data):
        (out/name).write_text(json.dumps(data,indent=2,default=lambda x:{'exact':str(x),'decimal':float(x)} if isinstance(x,F) else str(x))+'\n',encoding='utf-8')
    def table(name,header,rows):
        with (out/name).open('w',newline='',encoding='utf-8') as f:w=csv.writer(f);w.writerow(header);w.writerows(rows)
    thresholds=[]
    for K in [2,5,6,8,10]:thresholds.append(dict(K=K,H=max((1-a)**K,(1-b)**K),**minimum_pairs(max((1-a)**K,(1-b)**K),alpha,PAIR_SEARCH_BUDGET)))
    batch=[dict(T=T,**batch_certificate(H,REFERENCE_PAIRS,T,alpha)) for T in [1,2,5,10,24]]
    # The exact rational power can be very large; JSON retains it for audit.
    reach=[]
    for K in [6,8,10,12]:
        h=max((1-a)**K,(1-b)**K);T=1
        while T<1000 and 1-(1-h)**(T+1)<=alpha:T+=1
        if 1-(1-h)**T>alpha:reach.append(dict(K=K,status='infeasible'));continue
        import math
        c=float((1-h)**T)
        def approx(m):
            return 1-c if c<=m/(m+T) else math.exp(math.log(T/(m+T))+m/T*(math.log(m/(m+T))-math.log(c)))
        proposal=next((m for m in range(1,PAIR_SEARCH_BUDGET+1) if approx(m)<=float(alpha)),None)
        m=proposal if proposal and batch_certificate(h,proposal,T,alpha)['certified'] else None
        reach.append(dict(K=K,largest_batch=T if T<1000 else None,pairs=m,status='certified' if m and T<1000 else 'budget-exhausted',
                          previous_fails=m==1 or (m is not None and not batch_certificate(h,m-1,T,alpha)['certified'])))
    observation=[]
    for m,e in [(9,F(0)),(9,F(1,100000)),(9,F(1,10000)),(9,F(1,1000)),(10,F(1,1000)),(10,F(3,1000)),(12,F(3,1000))]:
        ideal=gate_bound(H,m)[0];event=observation_bound(H,m,e,e)
        observation.append(dict(m=m,epsilon=e,ideal=ideal,product=min(F(1),ideal+1-(1-e)**(m+1)),event=event,certified=event<=alpha))
    specificity=[dict(kappa=k,bound=specificity_bound(H,REFERENCE_PAIRS,k),certified=specificity_bound(H,REFERENCE_PAIRS,k)<=alpha)
                 for k in [F(0),F(1,1000),F(1,300),F(1,250),F(1,100),F(1,30)]]
    clean=next(n for n in range(1,10000) if F(299,300)**n<=F(1,20))
    poisson=[];constrained=[];poisson_counterexample=None;configured_loading=None
    if a==b and a>0:
        ref=PoissonReference(F(REFERENCE_LOADING),a,COUNT_THRESHOLD)
        configured_loading=dict(loading=F(REFERENCE_LOADING),chord_status=ref.chord_status(),bound=ref.chord_bound(REFERENCE_PAIRS),numerical=ref.envelope_numeric(REFERENCE_PAIRS))
        for loading in [F(1,2),F(1),F(2),F(9,4),F(5,2),F(3)]:
            ref=PoissonReference(loading,a,COUNT_THRESHOLD);m=None
            if ref.chord_status()=='certified':
                m=next((n for n in range(1,100) if ref.chord_bound(n)[1]<=alpha),None)
            poisson.append(dict(loading=loading,chord_status=ref.chord_status(),pairs=m,
                                bound=ref.chord_bound(m) if m else None,
                                previous_fails=(m==1 or ref.chord_bound(m-1)[0]>alpha) if m else None,
                                numerical_envelope=ref.envelope_numeric(REFERENCE_PAIRS)))
        ref=PoissonReference(3,a,COUNT_THRESHOLD);J=(1-2*a)**COUNT_THRESHOLD
        poisson_counterexample=dict(two_point_source=ref.source_error([(F(1,5),F(23,100)),(F(4,5),F(2))],REFERENCE_PAIRS),
             invalid_chord_upper=(1-exp_negative(6)[0])**REFERENCE_PAIRS*gate_bound(J,REFERENCE_PAIRS)[0],
             heavy_loading=PoissonReference(100,a,COUNT_THRESHOLD).source_error([(F(1),F(1,10))],REFERENCE_PAIRS))
        for K,guess in [(2,38),(3,16),(4,11),(5,9),(6,8)]:
            h,j=(1-a)**K,(1-2*a)**K
            # These are manuscript reference pair counts, freshly checked, not
            # assumed to remain minimal after editing the specimen fractions.
            good=constrained_certificate(h,j,guess,alpha,INTERVAL_BUDGET)
            bad=constrained_certificate(h,j,guess-1,alpha,INTERVAL_BUDGET)
            constrained.append(dict(K=K,pairs=guess,upper_certificate=good,previous=bad,
                minimal_certified=good['status']=='certified' and bad['status']=='refuted'))
    rng=np.random.default_rng(SEED);simulation=[]
    for label,law in [('synthetic',synthetic),('sharp-adversary',witness)]:
        outcomes=law.pair_outcomes();labels=list(outcomes);probs=[float(outcomes[k]) for k in labels]
        issued=0;gate_passes=0;misses=0
        for _ in range(SIMULATION_REPLICATES):
            references=rng.choice(labels,size=REFERENCE_PAIRS,p=probs).tolist()
            target=law.sample_specimen(COUNT_THRESHOLD,a,b,rng)
            gate_passes+=int('00' not in references);misses+=int(target['detected']==0)
            issued+=int(policy.report(references,target['detected'])['issued'])
        simulation.append(dict(source=label,replicates=SIMULATION_REPLICATES,gate_passes=gate_passes,negative_targets=misses,
            exclusions=issued,analytic_joint=law.joint_error(COUNT_THRESHOLD,a,b,REFERENCE_PAIRS),
            evidence='Seeded illustration only; exact theorem bound, not Monte Carlo frequency, authorizes reporting.'))
    nonblank=[dict(N=n,report_probability=synthetic.joint_error(n,a,b,REFERENCE_PAIRS)) for n in range(COUNT_THRESHOLD)]
    moment_demo=dict(single=1-F(4,5)+F(4,5)*F(1,20)**2,
        split_low_q=moment_bound(2,F(9,20),F(9,20),F(4,5),F(4,5),F(16,25)),
        split_high_q=moment_bound(2,F(9,20),F(9,20),F(4,5),F(4,5),F(79,100)),switch_q=F(106,135))
    report=policy.report(['10']*REFERENCE_PAIRS,0)
    summary=dict(specimen=dict(original_ml=budget.original_ml,available_ml=budget.available_ml,a=a,b=b,single=budget.single()),
        policy=dict(K=COUNT_THRESHOLD,m=REFERENCE_PAIRS,alpha=alpha,H=H,bound=g,maximizer=z,adversary=witness.atoms),
        illustrative_complete_record=report,missing_record=policy.report(['10']*(REFERENCE_PAIRS-1),0),
        reference_budget=dict(preparations=2*REFERENCE_PAIRS,exact_native_targets=2*REFERENCE_PAIRS,external_input_ml=F(1,2)*2*REFERENCE_PAIRS),
        synthetic_moments=synthetic.moments(),synthetic_reference_outcomes=synthetic.pair_outcomes(),
        configured_observation_bound=observation_bound(H,REFERENCE_PAIRS,F(CONTROL_LAW_TV),F(TARGET_LAW_TV)),
        clean_negative_controls=dict(pairs=clean,previous_fails=F(299,300)**(clean-1)>F(1,20),coverage_failure=F(1,20),
        note='A 95% bound for kappa is a separate confidence statement. Its failure probability must be budgeted before using it in a reporting guarantee.'),
        moment_comparison_reference=moment_demo,simulation=simulation,
        assumptions=['One shared latent state per specimen; targets independent conditional on that state.',
        'Reference pairs independent of one another and the future specimen, with the same source law.',
        'Exactly one native target per usable reference input unless explicitly using the Poisson model.',
        'One prespecified gate and report; no retry-until-pass or postselection of pairs.',
        'Batch results require independent future specimens; marginal error is not familywise error.',
        'Designed model values; no measured assay, clinical validation or rerun of Lean.'])
    dump('summary.json',summary);dump('thresholds.json',thresholds);dump('batch.json',dict(reuse=batch,reach=reach))
    dump('observation.json',dict(total_variation=observation,specificity=specificity));dump('poisson.json',dict(configured=configured_loading,rows=poisson,counterexample=poisson_counterexample))
    dump('constrained.json',constrained);dump('nonblank.json',nonblank)
    table('design_sweep.csv',['K','pairs','sharp_bound','floor'],((K,m,float(gate_bound(max((1-a)**K,(1-b)**K),m)[0]),float(max((1-a)**K,(1-b)**K))) for K in [2,5,6,8,10] for m in range(1,81)))
    plot(out,a,b,alpha,batch,poisson)
    print(f'Sharp joint false-exclusion bound: {float(g):.11f}; design certified: {g<=alpha}')
    print('Reference gate does not change the independent future specimen miss probability. Retain the complete event.',flush=True)
    digest=lambda p:hashlib.sha256(p.read_bytes()).hexdigest();here=Path(__file__).parent
    dump('run_metadata.json',dict(manuscript_sha256=MANUSCRIPT_SHA256,source_sha256=digest(Path(__file__)),python=platform.python_version(),
        module_sha256={n:digest(here/n) for n in ['model.py','certificates.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,a,b,alpha,batch,poisson):
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
    fig,axs=plt.subplots(1,2,figsize=(10,4),layout='constrained')
    for K in [2,5,6,8,10]:
        h=max((1-a)**K,(1-b)**K);axs[0].plot(range(1,81),[float(gate_bound(h,m)[0]) for m in range(1,81)],label=f'K = {K}')
    axs[0].axhline(float(alpha),color='black',ls=':');axs[0].set(xlabel='Independent reference pairs',ylabel='Sharp joint false-exclusion bound',ylim=(0,.33),title='False-exclusion bounds by target count and\nreference number');axs[0].legend(fontsize=8)
    axs[1].plot([r['T'] for r in batch],[r['bound_numeric'] for r in batch],'o-',label='Sharp bound with configured gate')
    axs[1].plot([r['T'] for r in batch],[float(r['floor']) for r in batch],'s--',label='Floor even with more references')
    axs[1].axhline(float(alpha),color='black',ls=':');axs[1].set(xlabel='Independent null specimens sharing the gate',ylabel='Familywise false-exclusion probability',title='Error from sharing reference controls\nacross specimens');axs[1].legend(fontsize=8)
    for ax in axs:ax.grid(alpha=.2)
    fig.savefig(out/'design.png',dpi=180);fig.savefig(out/'design.svg');plt.close(fig)
    if a==b and a>0:
        fig,axs=plt.subplots(1,2,figsize=(10,4),layout='constrained');load=np.linspace(.2,3.5,70)
        vals=[PoissonReference(str(l),a,COUNT_THRESHOLD).envelope_numeric(REFERENCE_PAIRS)['bound'] for l in load]
        axs[0].plot(load,vals,label='Numerical concave-envelope maximum');axs[0].axhline(float(alpha),color='black',ls=':');axs[0].set(xlabel='Mean native targets per reference input',ylabel='Joint false-exclusion probability',title='False-exclusion probability versus\nreference loading');axs[0].legend(fontsize=8)
        eps=np.linspace(0,.003,80);h=(1-a)**COUNT_THRESHOLD
        for m in [9,10,12]:axs[1].plot(eps,[float(observation_bound(h,m,F(str(e)),F(str(e)))) for e in eps],label=f'{m} pairs')
        axs[1].axhline(float(alpha),color='black',ls=':');axs[1].set(xlabel='Assumed TV error per experiment',ylabel='Event-specific upper bound',title='False-exclusion bound with observation-law\nerror');axs[1].legend(fontsize=8)
        for ax in axs:ax.grid(alpha=.2)
        fig.savefig(out/'observation.png',dpi=180);fig.savefig(out/'observation.svg');plt.close(fig)


if __name__=='__main__':main()