"""Enzyme activity, actual recovery trajectories, and sharp population bounds.

All inputs are designed model values. The scalar coordinate is not a full-cell
viability endpoint. Run python example.py --output outputs.
"""
from fractions import Fraction as F
from dataclasses import asdict, replace
from pathlib import Path
import argparse
import csv
import hashlib
import json
import math
import platform
import numpy as np
import sympy as sp
import mpmath as mp
from scipy.integrate import solve_ivp
from kinetics import EnzymeParameters, ScalarTask, CapacityFamily, ObservationPolyhedron, features
from population import (Population, TwoBandClass, ReadoutCalibration, minimax,convert_weights,error_normalization_frontier,
                        connected_limits,offband_infimum,offband_witness,resolved_capacity_bounds,row_span_certificate,finite_support_bounds)

# EDITABLE INPUTS ------------------------------------------------------------
POOL = '56'                    # model amount units
INITIAL = '10'
TARGET = '28'
SUBSTRATE = '7'
LOAD = '0.3'                   # amount per model time
DEADLINE = '120'               # model time, not calibrated minutes
AFFINITIES = ('7','3','56','125','520')  # Kg, Kn, Kh, Ka, Kb
INHIBITORS_AB = ('0','0')
LOW_BAND = ('0.625','9/14')
HIGH_BAND = ('1','1.375')
MEAN_CAPACITY = '1'
READOUT_SENSITIVITY = ('0.9','1')
READOUT_FALSE_POSITIVE = ('0','0.02')
READOUT_POSITIVE_FRACTION = ('0.68','0.72')
CELL_WEIGHT_RATIO = '1'
ADDITIONAL_READOUT_ERROR = '0'
DESIRED_FRACTION = '2/3'
OFF_BAND_ALLOWANCE = '0.1'
CAPACITY_MEASUREMENT_ERROR = '0.01'
RECTANGLE_PANELS = 128
TRAJECTORY_CAPACITIES = ('0.625','9/14','0.95','1','1.375')
HORIZONS = (60,120,240,1000)
MANUSCRIPT_SHA256 = '081a2f60534a885288a77debb6655765c7c40ffe21ddc5d1c9d620bf6d4e1884'
# Fixed reference demonstrations below are labeled separately from user inputs.
# ---------------------------------------------------------------------------


def individual_design():
    """a=2,c=3 pinned; designed y(K) in [.24,.26] bounds b=3/K.

    These are coefficient constraints for one realization, not reciprocal
    constraints applied to a pooled mixture.
    """
    avec=(F(1),F(1,7),F(0),F(0),F(0),F(0))
    cvec=(F(0),F(0),F(1,7),F(0),F(0),F(0))
    bvec=(F(0),F(0),F(0),F(1,7),F(0),F(0))
    rows=[avec,tuple(-v for v in avec),cvec,tuple(-v for v in cvec),bvec,tuple(-v for v in bvec)]
    bounds=[F(2),F(-2),F(3),F(-3),F(1,16),-F(33,728)]
    poly=ObservationPolyhedron(rows,bounds);task=ScalarTask();beta=task.parameters.coefficients()
    certificate=poly.verify_dual(task.target_features(),[1,0,F(1,28),0,1,0],beta)
    proposed=poly.propose_dual(task.target_features(),beta)
    deadline=task.uniform_deadline(certificate['reciprocal_upper'])
    # An independent literal rate-band design: target features are the mean of
    # two observed feature vectors at H=14 and H=42, with N=28,S=7 fixed.
    obs=[]
    for H in [14,42]:
        inputs=tuple(map(F,[28,7,H,0,0]));rate=task.parameters.rate(*inputs)
        obs.append((inputs,F(99,100)*rate,F(101,100)*rate))
    literal=ObservationPolyhedron.from_rate_bands(obs)
    literal_cert=literal.verify_dual(task.target_features(),[F(1,2),0,F(1,2),0],beta)
    missing=ObservationPolyhedron(rows[:4],bounds[:4]).propose_dual(task.target_features(),beta)
    # Eventual arrival for K>Kcrit is not a uniform deadline over that class.
    sequence=[]
    for e in [0,1,2,4,6]:
        K=F(252,103)+F(1,10**e);model=replace(task,parameters=replace(task.parameters,inhibition_H=K))
        sequence.append(dict(K=K,target_drift=model.drift(model.target),passage_numeric=model.passage_numeric()))
    return dict(K_interval=(F(48),F(728,11)),dual=certificate,solver_proposal=proposed,deadline=deadline,
                literal_rate_band_dual=literal_cert,literal_rate_band_deadline=task.uniform_deadline(literal_cert['reciprocal_upper']),
                missing_inhibition_observation=missing,compatible_K56_passage_numeric=task.passage_numeric(),near_boundary_sequence=sequence)


def asymptotic_rows():
    """Reference high-precision numerical comparison, separate from certification."""
    rows=[]
    with mp.workdps(65):
        vi=mp.mpf(363)/560;B=mp.mpf(109)/56;q=mp.mpf(3)/10
        def tau(delta):
            V=vi+delta;beta=V-q*B
            return 18*B/beta+6*V/beta**2*mp.log((mp.mpf(81)/70+46*delta)/(28*delta))
        for T in [1000,15000,30000]:
            lo,hi=mp.mpf('1e-50'),mp.mpf(1)
            for _ in range(240):
                mid=(lo+hi)/2
                if tau(mid)>T:lo=mid
                else:hi=mid
            delta=(lo+hi)/2;leading=mp.mpf(81)/1960*mp.exp(-mp.mpf(9)*(T-545)/8470)
            rows.append(dict(deadline=T,delta_numeric=str(delta),leading_delta_numeric=str(leading),relative_error_numeric=float(abs(leading-delta)/delta)))
    return rows


def resource_counterexample():
    rows=[]
    for V in [1.,4.]:
        times=np.linspace(0,1,101)
        def rhs(t,s):y,x=s;return [1-(1+V*V)*y,V*y-x]
        sol=solve_ivp(rhs,(0,1),[1,0],t_eval=times,rtol=1e-11,atol=1e-12,method='Radau')
        if not sol.success:raise ArithmeticError(sol.message)
        exact=V/(1+V*V)*(-np.expm1(-(1+V*V)*times))
        rows.append(dict(capacity=V,x_at_one=float(sol.y[1,-1]),maximum_difference=float(np.max(abs(sol.y[1]-exact))),
                         all_time_ceiling=V/(1+V*V),sampled_success=bool(np.any(sol.y[1]>=.25))))
    return rows


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)
    pars=EnzymeParameters(F(1),*map(F,AFFINITIES))
    task=ScalarTask(pars,F(POOL),F(SUBSTRATE),*map(F,INHIBITORS_AB),F(LOAD),F(INITIAL),F(TARGET),F(DEADLINE))
    family=CapacityFamily(task);bands=TwoBandClass(*map(F,LOW_BAND+HIGH_BAND),F(MEAN_CAPACITY))
    bound=bands.bounds(family);eps=F(ADDITIONAL_READOUT_ERROR)
    if eps<0:raise ValueError('Additional readout error must be nonnegative.')
    original_r=tuple(map(F,READOUT_POSITIVE_FRACTION));r=(max(F(0),original_r[0]-eps),min(F(1),original_r[1]+eps))
    readout=ReadoutCalibration(tuple(map(F,READOUT_SENSITIVITY)),tuple(map(F,READOUT_FALSE_POSITIVE)),r)
    repaired=readout.interval(bound['interval']) if bound['interval'] is not None else None
    number=convert_weights(repaired,F(CELL_WEIGHT_RATIO));target=F(DESIRED_FRACTION)
    if not 0<=target<=1:raise ValueError('Desired fraction must be a probability.')
    user=dict(task=asdict(task),bands=asdict(bands),bulk=bound,repaired_interval=repaired,number_interval=number,
              decision='certified' if number is not None and number[0]>=target else 'not certified from these inputs',
              requested_fraction=target,bulk_minimax=minimax(bound['interval']),repaired_minimax=minimax(repaired))
    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)
    trajectories=[];trajectoryrows=[]
    for V in map(F,TRAJECTORY_CAPACITIES):
        model=family.at(V);t,g,hit=model.trajectory();decision=model.classify(RECTANGLE_PANELS)
        trajectories.append(dict(capacity=V,numeric_first_hit=hit,passage_numeric=model.passage_numeric(),classification=decision))
        trajectoryrows.extend(zip([V]*len(t),t,g))
    table('trajectories.csv',['capacity','time','carrier_state'],trajectoryrows)
    # Paper reference witnesses; never claim these for an edited kinetic family.
    ref=CapacityFamily();refbands=TwoBandClass();refreadout=ReadoutCalibration();witnesses=[]
    for p in [F(20,41),F(33,49),F(7,10),F(4,5),F(1)]:
        pop=refbands.witness(p);recovery=pop.recovery_enclosure(ref)
        witnesses.append(dict(requested_fraction=p,population=asdict(pop),counts=pop.count_realization(),mean=pop.mean,recovery=recovery,
                              readout=refreadout.witness(p) if F(33,49)<=p<=F(4,5) else None))
    surface=[]
    for N,S,H,A,B in [(56,7,0,0,0),(28,7,28,0,0),(12,3,40,8,20),(1,100,55,100,500)]:
        inputs=tuple(map(F,[N,S,H,A,B]));rates=[refbands.witness(p).pooled_rate(ref,inputs) for p in [F(20,41),F(1)]]
        assert rates[0]==rates[1];surface.append([*inputs,*rates])
    table('pooled_surface.csv',['N','S','H','A','B','extremal_population_rate','uniform_population_rate'],surface)
    rectangles=[]
    for V in map(F,['19/20','951/1000','7/8','9/8','11/8']):
        row=ref.at(V).passage_rectangles(RECTANGLE_PANELS);row.update(capacity=V,classification=ref.at(V).classify(RECTANGLE_PANELS)['status']);rectangles.append(row)
    assert ref.at(F(19,20)).passage_rectangles(128)['lower']>120
    assert ref.at(F(951,1000)).passage_rectangles(128)['upper']<120
    cutoff=ref.numerical_cutoff();horizons=[]
    for T in HORIZONS:
        c=ref.numerical_cutoff(T);horizons.append(dict(deadline=T,cutoff_numeric=c,**connected_limits(c)))
    offband=[dict(eta=float(e),**offband_infimum(F(str(e)),cutoff)) for e in np.linspace(0,1,101)]
    eta=F(OFF_BAND_ALLOWANCE);approaching=[]
    # Find rational capacities just below the numerical cutoff; exact rectangle
    # lower bounds, not the numerical root alone, certify each failing atom.
    for x in [F(95,100),F(9505,10000),F(95059,100000)]:
        panels=128
        while ref.at(x).classify(panels)['status']=='unresolved' and panels<65536:panels*=2
        classification=ref.at(x).classify(panels)
        if classification['status']!='failure':raise ArithmeticError('Approaching witness is not certified failing.')
        pop=offband_witness(eta,x);approaching.append(dict(failing_capacity=x,population=asdict(pop),mean=pop.mean,offband_mass=pop.weights[1],success_fraction=pop.weights[2],panels=panels))
    equal=[]
    for name,capacities,weights in [('A',(F(7,8),F(11,8)),(F(3,4),F(1,4))),('B',(F(5,8),F(9,8)),(F(1,4),F(3,4)))]:
        pop=Population(capacities,weights);equal.append(dict(name=name,population=asdict(pop),mean=pop.mean,variance=pop.variance,recovery=pop.recovery_enclosure(ref)))
    atoms=list(map(F,['1/20','1/10','1/5','4/5']));second=[2*x/(1+x) for x in atoms];rows=[[F(1)]*4,atoms,second];h=[0,0,1,1]
    null=row_span_certificate(rows,h);lp=finite_support_bounds(rows,[1,F(9,80),F(7,36)],h)
    null['exact_paper_null_vector']=[-98,165,-70,3]
    null['exact_endpoint_weights']=[[0,F(55,56),0,F(1,56)],[F(7,12),0,F(5,12),0]]
    inhibition=[]
    for x in atoms:
        K=56*x/(3*(1-x));model=replace(ScalarTask(),parameters=replace(EnzymeParameters(),inhibition_H=K),deadline=F(300))
        inhibition.append(dict(inhibition_coordinate=x,task_classification=model.classify(),passage_numeric=model.passage_numeric()))
    refpop=refbands.witness(F(7,10));measurement=resolved_capacity_bounds(refpop,refpop.capacities,F(CAPACITY_MEASUREMENT_ERROR),(F(19,20),F(951,1000)))
    eventual=ref.eventual_threshold();eventual_witness=Population((eventual,F(11,8)),(F(210,407),F(197,407)))
    results=dict(user=user,user_trajectories=trajectories,reference_individual=individual_design(),reference_witnesses=witnesses,
                 reference_rectangles=rectangles,reference_cutoff_numeric=cutoff,reference_connected_numeric=connected_limits(cutoff),
                 reference_horizons=horizons,reference_offband_numeric=offband_infimum(eta,cutoff),reference_approaching_witnesses=approaching,
                 reference_equal_variance=equal,reference_second_assay=dict(row_span=null,numerical_bounds=lp,kinetic_classification=inhibition),
                 reference_measurement=measurement,reference_frontier=[dict(weight_ratio=R,**error_normalization_frontier(R)) for R in [F(1),F(51,50),F(33,32),F(6,5)]],
                 reference_eventual=dict(threshold=eventual,lower=F(197,407),attained=True,population=asdict(eventual_witness)),
                 reference_asymptotic=asymptotic_rows(),reference_multidimensional=resource_counterexample(),
                 reference_sample_sizes=dict(one_sided=math.ceil(math.log(20)/(2*(1/150)**2)),two_sided=math.ceil(math.log(40)/(2*(1/150)**2))),
                 evidence='Exact finite certificates and witnesses; numerical ODEs, LPs and roots labeled; general existence/comparison imported from manuscript; Lean not rerun; designed model units only.')
    def serialize(v):
        if isinstance(v,(F,sp.Rational)):return str(v)
        return float(v)
    # JSON represents nonarrival explicitly rather than nonstandard Infinity.
    def clean(x):
        if isinstance(x,float) and not math.isfinite(x):return 'no finite arrival'
        if isinstance(x,dict):return {k:clean(v) for k,v in x.items()}
        if isinstance(x,(tuple,list)):return [clean(v) for v in x]
        return x
    def dump(name,obj):(out/name).write_text(json.dumps(clean(obj),indent=2,default=serialize,allow_nan=False)+'\n',encoding='utf-8')
    dump('results.json',results)
    table('horizons.csv',list(horizons[0]),[list(r.values()) for r in horizons])
    table('offband.csv',['offband_allowance','infimum_numeric'],[[r['eta'],r['infimum']] for r in offband])
    lines=[f'User two-band contract: {bound["status"]}; bulk interval {bound["interval"]}.',
           f'User calibrated interval {repaired}; number interval {number}; requested fraction {target}: {user["decision"]}.',
           f'Reference individual uniform deadline: 43740/391 = {43740/391:.6f}; compatible K=56 arrival {ref.at(1).passage_numeric():.6f}.',
           f'Reference connected-support cutoff: {cutoff:.10f}; infimum {connected_limits(cutoff)["lower"]:.10f}, unattained.',
           'Same mean and variance: recovery fractions 1/4 and 3/4. Entire pooled surface also agrees.',
           'Reference eventual lower bound: 197/407, attained; carrier arrival is not whole-cell viability.']
    (out/'console.txt').write_text('\n'.join(lines)+'\n',encoding='utf-8');print('\n'.join(lines))
    plot(out,trajectoryrows,cutoff,offband)
    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 ['kinetics.py','population.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,trajectoryrows,cutoff,offband):
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
    fig,axs=plt.subplots(1,2,figsize=(10,4),layout='constrained');data=np.array(trajectoryrows,float)
    for V in np.unique(data[:,0]):
        row=data[data[:,0]==V];axs[0].plot(row[:,1],row[:,2],label=f'V = {V:.4g}')
    axs[0].axhline(float(F(TARGET)),color='gray',ls=':');axs[0].set(xlabel='Model time',ylabel='Carrier state g',title='Literal scalar recovery trajectories');axs[0].legend(fontsize=8)
    for y,L,U,label,opened in [(2,20/41,1,'Separated support',False),(1,33/49,.8,'With calibrated readout',False),(0,connected_limits(cutoff)['lower'],1,'Connected support',True)]:
        axs[1].plot([L,U],[y,y],lw=5);axs[1].scatter([L],[y],facecolors='white' if opened else 'tab:blue',edgecolors='tab:blue',zorder=4);axs[1].scatter([U],[y],color='tab:blue')
        axs[1].text((L+U)/2,y+.13,label,ha='center',fontsize=8)
    axs[1].axvline(2/3,color='gray',ls=':');axs[1].set(xlim=(.05,1.05),ylim=(-.3,2.45),yticks=[],xlabel='Fraction reaching target by time 120',title='Recovery fractions compatible with each\nobservation')
    for ax in axs:ax.grid(alpha=.2)
    fig.savefig(out/'recovery.png',dpi=180);fig.savefig(out/'recovery.svg');plt.close(fig)
    fig,axs=plt.subplots(1,2,figsize=(10,4),layout='constrained');R=np.linspace(1,33/32,150);eps=(33-32*R)/(50*(1+2*R))
    axs[0].plot(R,eps);axs[0].fill_between(R,0,eps,alpha=.15);axs[0].set(xlabel='Maximum / minimum positive cell weight',ylabel='Additional downward readout error',title='Readout and weighting allowances for\ntwo-thirds recovery',ylim=(0,.007))
    axs[1].plot([r['eta'] for r in offband],[float(r['infimum']) for r in offband]);axs[1].scatter([0],[20/41],color='tab:blue');axs[1].set(xlabel='Allowed population mass outside the bands',ylabel='Sharp recovery infimum',title='Recovery bound versus fraction outside\ncapacity bands')
    axs[1].text(.18,.46,'Attained only at zero off-band allowance',fontsize=8)
    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()
