"""Trustworthy amplification decisions: source, observation, and inference.

Run: python example.py --output outputs. All wells are synthetic; inputs are
dimensionless paper fixtures, not calibrated clinical assay parameters.
"""
from __future__ import annotations
import argparse
from fractions import Fraction as F
from dataclasses import asdict, replace
from pathlib import Path
import csv
import hashlib
import json
import math
import platform
import numpy as np
from scipy.integrate import quad
from scipy.special import digamma
from source import AmplificationSource, IdentityChannel, fixed_rate_window
from observation import (PairedCapacity, BlockCapacity, DurationInterval, CensoredComparison,
                         SignalCalibration, compare_intervals, crossing_bracket)
from certificates import timing_dual_certificate, large_threshold_certificate, explicit_deadlines, LimitLaw

# EDITABLE INPUTS ------------------------------------------------------------
THRESHOLD = 5                      # active units at a latched positive call
CAPACITY = '5'                     # resource units; must exceed threshold-1
BACKGROUND = '0.01'                # active units per normalized time
GROWTH = '1'                       # inverse normalized time; a=background/growth
LOADING_MEAN = '4'                 # Poisson initial active units, includes empties
BLANK_LIMIT = '0.01'
MISS_LIMIT = '0.05'
IDENTITY_DEADLINE = '7'
IDENTITY_FALSE_POSITIVE = '0.01'    # conditional on an empty-well hit
IDENTITY_SENSITIVITY = '0.99'       # conditional on an occupied-well hit
CAPACITY_A = '0.01'                # known b/g for the timing-comparison experiment
EARLY_STATE, LATE_STATE = 2, 4
CALIBRATION_WELLS = 4000           # independent wells; every unresolved well kept
COMPARISON_DELTA = '0.05'          # confidence failure budget, not assay error
INTERVAL_ERROR_ALLOWANCE = '0'     # known mean interval-failure probability
SHARED_CALIBRATION_FAILURE = '0'  # separate union-bound allowance
CLOCK_VALUES = (.4, 1., 3.)       # constant within each well; varies between wells
CALIBRATION_SCENARIOS = (('resolved',7,8.,.002),('short follow-up',7,.12,.05),('large capacity',100,8.,.002))
SEED = 20260916
LARGE_THRESHOLDS = (1_000_000,100_000_000)  # reference theorem a=.01, lambda=4
EXACT_DEGREE = 200
MANUSCRIPT_SHA256 = 'c8a26ae06acfe7284d0ee02fd04f90eb89d13d6a77ce839e73e2bda765912dba'
# ---------------------------------------------------------------------------


def censored_wells(model,n,rng,R,horizon,frame,clocks=CLOCK_VALUES):
    if type(n)!=int or n<1 or not 0<frame<=horizon or any(not math.isfinite(c) or c<=0 for c in clocks):
        raise ValueError('Invalid sample size, acquisition times or clock values.')
    clock=rng.choice(clocks,n);a=float(model.a);i,j=model.early,model.late
    if R<=j:raise ValueError('Capacity must exceed the late held state.')
    U=rng.exponential(1/((a+i)*(1-i/R)),n)/clock
    V=rng.exponential(1/((a+j)*(1-j/R)),n)/clock
    def bracket(x):
        if x>=horizon:return DurationInterval(F(str(horizon)),None)
        # Decimal frame times, not rounded floating products at frame boundaries.
        k=math.floor(x/frame);d=F(str(frame))
        lo=k*d;hi=(k+1)*d
        # Ensure containment despite a floating ratio at an exact boundary.
        xx=F(str(float(x)))
        return DurationInterval(min(lo,xx),max(hi,xx))
    rows=[];counts=dict(positive=0,negative=0,unknown=0)
    for well,(u,v,c) in enumerate(zip(U,V,clock),1):
        bu,bv=bracket(u),bracket(v);category=compare_intervals(bu,bv);counts[category]+=1
        assert category!='positive' or v>u
        assert category!='negative' or v<=u
        rows.append([well,c,u,v,bu.lower,bu.upper,bv.lower,bv.upper,category])
    comparison=CensoredComparison(**counts)
    interval=comparison.confidence(F(COMPARISON_DELTA),F(INTERVAL_ERROR_ALLOWANCE))
    return dict(n=n,**counts,probability_interval=interval,**model.capacity_set(interval)),rows


def selection_example(model,R=7):
    """Illustration: retaining only U+V<=1 breaks clock cancellation."""
    a=float(model.a);i,j=model.early,model.late;rows=[]
    for clock in CLOCK_VALUES:
        qi=clock*(a+i)*(1-i/R);qj=clock*(a+j)*(1-j/R)
        num=quad(lambda u:qi*math.exp(-qi*u)*(math.exp(-qj*u)-math.exp(-qj*(1-u))),0,.5,epsabs=1e-12)[0]
        den=quad(lambda u:qi*math.exp(-qi*u)*(1-math.exp(-qj*(1-u))),0,1,epsabs=1e-12)[0]
        rows.append(dict(clock=clock,unselected=float(model.probability(R)),selected=num/den,completion=den))
    return rows


def plateau_normalized_ratio(R,a=.01):
    """Ratio of expected block times at 20%,40%,60% of the plateau.

    This is NOT the expectation of a ratio; it illustrates the leading-order
    loss of capacity information after replacing absolute states by fractions.
    """
    if type(R)!=int or R<10 or a<=0:raise ValueError('Integer capacity >=10 and a>0 required.')
    k0,k1,k2=[int(f*R) for f in [.2,.4,.6]]
    def mean(lo,hi):return R/(R+a)*(digamma(hi+a)-digamma(lo+a)+digamma(R-lo+1)-digamma(R-hi+1))
    return mean(k1,k2)/mean(k0,k1)


def main():
    parser=argparse.ArgumentParser(description=__doc__);parser.add_argument('--output',type=Path,default=Path('outputs'));args=parser.parse_args()
    out=args.output;out.mkdir(parents=True,exist_ok=True)
    source=AmplificationSource(THRESHOLD,F(CAPACITY),F(BACKGROUND),F(GROWTH),F(LOADING_MEAN))
    channel=IdentityChannel(F(IDENTITY_FALSE_POSITIVE),F(IDENTITY_SENSITIVITY))
    alpha,beta=F(BLANK_LIMIT),F(MISS_LIMIT)
    if not 0<alpha<1 or not 0<beta<1 or not 0<=F(SHARED_CALIBRATION_FAILURE)<1:raise ValueError('Invalid error budgets.')
    user_bounds=source.certify(F(IDENTITY_DEADLINE),EXACT_DEGREE)
    user=dict(source=asdict(source),initial_hit_atom=source.initial_hit(),numerical_best_deadline=source.numerical_optimum(float(alpha)),
              deadline_bounds=asdict(user_bounds),deadline_status=user_bounds.deadline_status(alpha,beta),
              joint=channel.guaranteed_errors(user_bounds),joint_fixture=channel.joint_fixture(source,float(IDENTITY_DEADLINE)))
    # Fixed reference demonstrations remain explicitly distinct from edited inputs.
    reference=AmplificationSource();refchannel=IdentityChannel();finite=[]
    for h,R,t,L in [(5,5,4,200),(5,5,5,200),(5,5,7,200),(5,10,F(16,5),200),(20,20,8,300),(20,20,10,300)]:
        model=replace(reference,threshold=h,capacity=F(R));b=model.certify(t,L)
        finite.append(dict(threshold=h,capacity=R,deadline=t,blank_lower=b.blank[0],blank_upper=b.blank[1],miss_lower=b.miss[0],miss_upper=b.miss[1],
                           deadline_status=b.deadline_status(),all_timing_excluded_using_theorem=b.separates_all_timing(),
                           **{'joint_'+key:value for key,value in refchannel.guaranteed_errors(b).items()}))
    dual=timing_dual_certificate()
    assert finite[0]['blank_lower']>F(73,10000) and finite[1]['miss_lower']>F(69,1000)
    dual['intercept_lower']=F(1055,10000);dual['miss_lower_at_blank_1pct']=F(555,10000)
    dual['tv_exclusion_budget']='5*delta_blank+delta_loaded <= .0055; complete timing-law TV budgets'
    robust=refchannel.guaranteed_errors(fixed_rate_window(reference,F(699,100),F(701,100)))
    reserve=replace(reference,capacity=F(10))
    rb=fixed_rate_window(reserve,F(319,100),F(321,100))
    large=large_threshold_certificate();times=[explicit_deadlines(h) for h in LARGE_THRESHOLDS]
    limits={name:LimitLaw(headroom_zero=zero).optimum() for name,zero in [('growing reserve',False),('zero headroom',True)]}
    paired=PairedCapacity(F(CAPACITY_A),EARLY_STATE,LATE_STATE);rng=np.random.default_rng(SEED);simulations=[]
    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)
    for index,(name,R,horizon,frame) in enumerate(CALIBRATION_SCENARIOS):
        result,rows=censored_wells(paired,CALIBRATION_WELLS,rng,R,horizon,frame)
        result.update(name=name,true_capacity=R,horizon=horizon,frame=frame,
                      total_confidence_failure_upper=min(F(1),F(COMPARISON_DELTA)+F(SHARED_CALIBRATION_FAILURE)))
        simulations.append(result)
        table(f'synthetic_wells_{index+1}.csv',['well','clock','true_U','true_V','U_lower','U_upper_empty_means_infinite','V_lower','V_upper_empty_means_infinite','category'],rows)
    block=BlockCapacity();blocks=[dict(capacity=R,probability=block.probability(R)) for R in [6,7,10,20,100]]
    block_interval=block.capacity_set((F(32,100),F(35,100)))
    selection=selection_example(paired)
    plateau=[dict(capacity=R,ratio_of_expected_times=plateau_normalized_ratio(R)) for R in [100,1000,10000,100000]]
    design=[]
    for R in [7,20,100]:
        p=paired.probability(R);gap=min(paired.probability(F(9,10)*R)-p,p-paired.probability(F(11,10)*R))
        design.append(dict(capacity=R,point_estimate_gap=gap,sufficient_wells_numeric=math.ceil(math.log(40)/(2*float(gap)**2))))
    signal=SignalCalibration();raw=signal.state_interval(1110)
    crossings=[DurationInterval(F('2'),F('2.1')),DurationInterval(F('3.5'),F('3.6')),DurationInterval(F('4.9'),F('5'))]
    U=DurationInterval.from_crossings(*crossings[:2]);V=DurationInterval.from_crossings(*crossings[1:])
    observation=dict(signal_units_interval=raw,successive_duration_U=asdict(U),successive_duration_V=asdict(V),comparison=compare_intervals(U,V),
                     no_crossing=crossing_bracket([0,1,2],[(0,1),(1,2),(2,3)],5),
                     physical_sets={name:paired.capacity_set(p) for name,p in [('incompatible',('.1','.2')),('unresolved',(0,1)),('finite',('.44','.47')),('lower',('.32','.4'))]})
    # Equal endpoint laws via permutation of independent exponential summands.
    q6=AmplificationSource(3,F(6),F(2),F(1)).rates();q4=AmplificationSource(3,F(4),F(8,3),F(2,3)).rates()
    assert sorted(q6)==sorted(q4)
    results=dict(user=user,reference_finite=finite,reference_timing_dual=dual,reference_fixed_rate_joint=robust,
                 reference_fixed_rate_reserve=dict(bounds=asdict(rb),status=rb.deadline_status()),reference_large_threshold=large,
                 explicit_deadlines=times,numerical_limits=limits,capacity_simulations=simulations,block_probabilities=blocks,
                 block_capacity_outer_interval=block_interval,completion_selection=selection,sample_design=design,observation=observation,
                 reference_plateau_normalization=dict(rows=plateau,limit=math.log(2.25)/math.log(8/3),evidence='ratio of expected times; finite corrections retain information'),
                 endpoint_nonidentifiability=dict(rates_R6=q6,rates_R4=q4,a_is_different=True),
                 zero_error_trials={str(p):math.ceil(math.log(.025)/math.log(1-p)) for p in [.01,.05,.0003]},
                 evidence='Exact/directed bounds and numerical illustrations separated. General source/decision theorems imported; Lean not rerun; no clinical calibration supplied.')
    def serialize(v):
        if isinstance(v,F):return str(v)
        if isinstance(v,DurationInterval):return asdict(v)
        return float(v)
    def dump(name,obj):(out/name).write_text(json.dumps(obj,indent=2,default=serialize)+'\n',encoding='utf-8')
    dump('results.json',results)
    table('finite_decisions.csv',list(finite[0]),[[float(v) if isinstance(v,F) else v for v in row.values()] for row in finite])
    table('capacity_design.csv',list(design[0]),[[float(v) if isinstance(v,F) else v for v in row.values()] for row in design])
    curves=[]
    for t in np.linspace(0,12,121):
        B,M=reference.numerical_errors(t);C=1-M-math.exp(-4)*B;rb_,rm_=reserve.numerical_errors(t)
        curves.append([t,B,M,.01*B,1-.99*C,rb_,rm_])
    table('timing_curves.csv',['deadline','blank_R5','miss_R5','joint_blank_R5','joint_miss_R5','blank_R10','miss_R10'],curves)
    lines=[f'User source: numerical best deadline {user["numerical_best_deadline"]["deadline"]:.6f}; miss {user["numerical_best_deadline"]["miss"]:.6%}.',
           'Reference R=h=5: exact separating witness excludes every timing-only rule under the manuscript theorem.',
           f'Reference identity rule, fixed-rate uncertainty: blank <= {float(robust["blank_upper"]):.8f}; miss <= {float(robust["miss_upper"]):.8f}.',
           f'Reference capacity-ten fixed-rate window: {rb.deadline_status()}.',
           f'Calibrated signal 1110: {raw[0]:,}..{raw[1]:,} possible active units.',
           *[f'{r["name"]}: {r["positive"]}/{r["negative"]}/{r["unknown"]} positive/negative/unresolved; {r["status"]}.' for r in simulations],
           'Synthetic fixtures and supplied contracts only; timing-only optimality and large-threshold source connections are manuscript theorems.']
    (out/'console.txt').write_text('\n'.join(lines)+'\n',encoding='utf-8');print('\n'.join(lines))
    plot(out,curves,simulations,selection,blocks)
    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 ['source.py','observation.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,curves,simulations,selection,blocks):
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
    data=np.array(curves);fig,axs=plt.subplots(1,2,figsize=(10,4),layout='constrained')
    for b,m,label in [(1,2,'Timing, R = 5'),(3,4,'Identity + hit, R = 5'),(5,6,'Timing, R = 10')]:
        axs[0].plot(100*data[:,b],100*data[:,m],label=label)
    axs[0].axvline(1,color='gray',ls=':');axs[0].axhline(5,color='gray',ls=':');axs[0].fill_between([0,1],0,5,color='green',alpha=.1)
    axs[0].set(xlim=(0,3),ylim=(0,20),xlabel='Blank positives (%)',ylabel='Loaded negatives (%)',title='Assay errors under timing, identity and\ncapacity changes');axs[0].legend(fontsize=8)
    for r in selection:axs[1].scatter(r['clock'],r['selected'],color='tab:orange')
    axs[1].plot([r['clock'] for r in selection],[r['selected'] for r in selection],label='Only U + V ≤ 1 retained',color='tab:orange')
    axs[1].axhline(selection[0]['unselected'],label='All wells: clock-invariant',color='tab:blue')
    axs[1].set(xlabel='Common rate multiplier within a well',ylabel='P(late wait > early wait)',title='Waiting-time comparison with and without\ncompletion selection');axs[1].legend(fontsize=8)
    for ax in axs:ax.grid(alpha=.2)
    fig.savefig(out/'decisions.png',dpi=180);fig.savefig(out/'decisions.svg');plt.close(fig)
    fig,axs=plt.subplots(1,2,figsize=(10,4),layout='constrained')
    labels=[]
    for i,r in enumerate(simulations):
        p=list(map(float,r['probability_interval']));axs[0].plot(p,[i,i],lw=3);axs[0].scatter([float(PairedCapacity(F(CAPACITY_A),EARLY_STATE,LATE_STATE).probability(r['true_capacity']))],[i],marker='x',color='k')
        labels.append(r['name']+'\n'+str(r['unknown'])+' unresolved')
    axs[0].set_yticks(range(len(labels)),labels);axs[0].set(xlim=(0,1),xlabel='Probability interval, all wells retained',title='Comparison-probability intervals with\ncensored wells')
    axs[1].plot([r['capacity'] for r in blocks],[float(r['probability']) for r in blocks],'o-',label='Two-state blocks')
    model=PairedCapacity();x=np.geomspace(5,100,100);axs[1].plot(x,[float(model.probability(F(str(v)))) for v in x],label='Single held states 2 and 4')
    axs[1].set(xscale='log',xlabel='Capacity R',ylabel='P(late duration > early duration)',title='Waiting-time comparison probability versus\ncapacity');axs[1].legend(fontsize=8)
    for ax in axs:ax.grid(alpha=.2)
    fig.savefig(out/'capacity.png',dpi=180);fig.savefig(out/'capacity.svg');plt.close(fig)


if __name__=='__main__':main()