"""Productive inheritance under imposed load; hidden operating states under feedback."""
# EDITABLE INPUTS: paper parameters and separately labelled small diagnostics.
EXTRACTION='1/100'               # theorem interval [9999/10^6,1/100]
NEWBORN_SIZE=65536*10**18
RETAINED_CELLS=4*10**9
GROWTH_COUPLING='1/100000000000'
COMPOSITION=(.2,.3,.5)
RENEWAL=.05
MORTALITY=.5
RECOVERY_HORIZON=1000.
ABUNDANCE_PERTURBATION=(1.08,.94,1.004)  # not assumed inside the tiny certified energy sublevel
SMALL_SIZE=10
SMALL_CELLS=4
SMALL_GROWTH='2'                 # accelerated exploration, outside theorem scope
SMALL_RECOVERY=.05               # theorem recovery is 5376
EVENT_BUDGET=30000
SEED=52092026
WINDOW=100
S_MAX='0.08'
ABUNDANCE_ERROR='0.001'
DISPERSION_ERROR='0.0001'
RESERVOIR_ERROR='0.005'
AVERAGE_QUADRATURE_ERROR='0.00001'
ADMITTED_DYNAMIC_ERROR='0.001'    # synthetic classifier assumption, not inferred from the trajectory
CONCENTRATION_MOLAR='1/1000'
TIME_UNIT_SECONDS=1.

import argparse
from fractions import Fraction as F
import csv
import hashlib
import json
import math
from pathlib import Path
import platform
import numpy as np
from mpmath import mp,iv
import sympy as sp
from scipy.integrate import solve_ivp
from resident import ResidentChemistry,Cell,reconstruction,residual,rational_root,ENERGY
from feedback import Community,WindowReadout,S_BOUNDS,stationary,stationary_residual,uptake,load
from protocol import PopulationSource,SerialProtocol,UniformTransfer

mp.dps=70;iv.dps=70
MANUSCRIPT_SHA256='41f71eaed3040751568bf085391b5021bde099b27d14335de005729d33233d1a'


def val(ctx,q):
    q=F(q);return ctx.mpf(q.numerator)/q.denominator


class InheritanceCertificate:
    """Fresh complete error budget; scope is the paper's fixed constructive witness."""
    def __init__(self,rho=EXTRACTION,N=NEWBORN_SIZE,M=RETAINED_CELLS,gamma=GROWTH_COUPLING):
        self.rho=F(rho);self.N=N;self.M=M;self.gamma=F(gamma)
        if not F(9999,10**6)<=self.rho<=F(1,100):raise ValueError('Extraction outside certified interval')
        if N!=65536*10**18 or M!=4*10**9 or self.gamma!=F(1,10**11):raise ValueError('Published certificate requires the declared N, M and gamma; new scales need a new source audit')
    def logs(self,ctx):
        N=ctx.mpf(self.N);M=ctx.mpf(self.M);T=8/val(ctx,self.gamma);u=N/(10**12*512000000);ln=ctx.log
        return dict(outer_initial=ln(M)-7*u,daughter_initial=ln(14*M)-4*u,
            outer_drift=ln(M*T*u/60)-15*u/2,spatial_initial=ln(M)-u,spatial_drift=ln(M*T*u/60)-3*u/2,
            bad_partition=ln(56*M)-N/(35*10**12),deadline=-N/2500,size_odds=-19*N/500000,
            low_output=-N*M,collection=ln(ctx.mpf(1)/10**6),batch_service=ln(ctx.mpf(1)/10**6),recovery_and_service=ln(2*M/10**18))
    def evaluate(self):
        logs=self.logs(mp);ilog=self.logs(iv);nontransfer=sum(iv.exp(v) for v in ilog.values())
        if nontransfer.b>val(iv,F(3,10**6)).a:raise ArithmeticError('Nontransfer bound failed')
        p0=F(1,2);p1=F(49,1600);D=lambda p:F(32)/(F(1,50)**2*p*self.M)
        lower=1-D(p0)-D(p1)-F(6,10**6)
        assert lower==F(24482873,24500000)
        b=mp.mpf(3)/5*mp.log(4)-mp.mpf(19)/500;ell=mp.log(mp.mpf(51)/49)
        gain=2*(b-ell)-mp.log(2)
        repeats=[]
        for k in range(1,7):
            failure=sum(F(1,25000)*F(800,49)**j for j in range(k))+F(3*k,10**6)
            g=k*(b-ell)-mp.log(2)
            repeats.append(dict(cycles=k,probability_lower=max(F(0),1-failure),count_fraction_lower_numerical=float(1/(1+mp.exp(-g)))))
        return dict(joint_two_cycle_lower=lower,nontransfer_upper=math.nextafter(float(nontransfer.b),math.inf),transfer_first=D(p0),transfer_second=D(p1),
            log_error_terms={k:mp.nstr(v,32) for k,v in logs.items()},count_log_odds_gain_numerical=float(gain),
            high_fraction_lower_numerical=float(1/(1+mp.exp(-gain))),batch_output_strictly_greater_than=5*self.N*self.M,
            prepared_precursor_at_most=12*self.N*self.M,growth_consumption_at_most=9*self.N*self.M,
            finite_repetition=repeats,scope='Imported source-kernel theorem; fresh error arithmetic. Repetitions beyond two are conventional corollaries, not rerun Lean roots.')


def exact_audit():
    z,rho=sp.symbols('z rho',positive=True);a,b,zz,h=reconstruction(z,rho);e=sp.Rational(1,100000)
    field=(6-2*a+b*z+2*e*(b-a*a),27+a-(1+z)*b-e*(b-a*a),a-b*z-16*z-4*z*z+3*h-rho*z,16*z+2*z*z-sp.Rational(20001,10000)*h)
    rr=residual(z,rho)
    assert all(sp.simplify(v)==0 for v in (field[0]+2*rr,field[1]-rr,field[2],field[3]))
    # Uniform bracket signs follow from monotonicity in rho, checked at both endpoints.
    brackets={}
    for tag,lo,hi in [('L','0.98172','0.98174'),('H','2.89014','2.89017')]:
        signs=[(residual(F(lo),r),residual(F(hi),r)) for r in [F(9999,10**6),F(1,100)]]
        assert all(a<0<b for a,b in signs)
        P=sp.Matrix(ENERGY[tag].tolist())/10**6
        assert all((P-sp.eye(4)/200)[:k,:k].det()>0 and (60*sp.eye(4)-P)[:k,:k].det()>0 for k in range(1,5))
        brackets[tag]=dict(z=[lo,hi],endpoint_signs=[['negative','positive']]*2,energy_spectral_sandwich=True)
    y=sp.Symbol('y',real=True)
    q=sp.Rational(2646507285317916225,16)-sp.Rational(373906577279071365,2)*y+sp.Rational(34457435637128464441,2)*y*y
    certificate=q+26834156547533132606*(y+sp.Rational(1,2))*y*y+4448524228728822329*y**4+136356292472000*(y+sp.Rational(1,2))*y**4+44462224000000*y**6
    numerator=sp.cancel(residual(z,sp.Rational(31,1000))*4444888900000000000*(z+2)**2)
    assert sp.expand(numerator.subs(z,y+sp.Rational(5,2))-certificate)==0
    assert sp.discriminant(q,y)<0 and q.coeff(y,2)>0
    s=sp.Symbol('s',nonnegative=True);aa=sp.Rational(11,250);bb=sp.Rational(12,3125);U=sp.Rational(2239911,97656250)
    assert sp.expand(U-load(s)-10*(s-aa-bb/20)**2-20*(s+2*aa)*(s-aa)**2)==0
    feedback=[]
    for lo,hi in S_BOUNDS:
        a,b=F(lo),F(hi);fa,fb=stationary_residual(a),stationary_residual(b)
        assert fa*fb<0
        feedback.append(dict(S=[a,b],J=[uptake(a),uptake(b)],load=sorted([load(a),load(b)])))
    gap=feedback[1]['J'][0]-feedback[0]['J'][1];threshold=(feedback[1]['J'][0]+feedback[0]['J'][1])/2
    assert feedback[1]['J'][0]/feedback[0]['J'][1]>F('1.8186')
    # Count drift is evaluated at symbols, with exact falling factorials.
    A,B,Z,H,m=sp.symbols('A B Z H m',positive=True);chem=ResidentChemistry();symbols=(A,B,Z,H)
    rate=[sp.Rational(ch.coefficient.numerator,ch.coefficient.denominator)*m**(1-sum(ch.inputs))*sp.prod(sp.prod(m*x-k for k in range(order)) for x,order in zip(symbols,ch.inputs)) for ch in chem.channels]
    countdrift=[sp.expand(sum(r*(ch.outputs[i]-ch.inputs[i])/m for r,ch in zip(rate,chem.channels))) for i in range(4)]
    expected=[6-2*A+B*Z+2*e*(B-A*A),27+A-(1+Z)*B-e*(B-A*A),A-B*Z-16*Z-4*Z*Z+3*H-Z/100,16*Z+2*Z*Z-sp.Rational(20001,10000)*H]
    corrections=[sp.simplify(a-b) for a,b in zip(countdrift,expected)]
    assert corrections==[2*e*A/m,-e*A/m,4*Z/m,-2*Z/m]
    return dict(stationary_reconstruction=True,imposed_brackets=brackets,stationary_exclusion_polynomial=True,
        quadratic_discriminant=str(sp.discriminant(q,y)),load_cap=F(2239911,97656250),feedback_intervals=feedback,
        uptake_gap=gap,uptake_threshold=threshold,half_gap=gap/2,uptake_ratio_lower=feedback[1]['J'][0]/feedback[0]['J'][1],
        count_correction=[str(v) for v in corrections],scope='Exact identities and enclosures; full uniform nonlinear energy and stochastic generator inequalities remain imported theorem premises')


def small_protocol():
    chemistry=ResidentChemistry(EXTRACTION);centers=chemistry.centers()
    if len(centers)<3:raise ValueError('Small diagnostic expects low and high centers')
    cells=tuple(Cell(tag,SMALL_SIZE,tuple(max(0,math.floor(float(v)*SMALL_SIZE)) for v in centers[0 if tag=='L' else -1])) for tag in ['L']*(SMALL_CELLS//2)+['H']*(SMALL_CELLS-SMALL_CELLS//2))
    source=PopulationSource(SMALL_SIZE,SMALL_GROWTH,chemistry)
    result=SerialProtocol(source,SMALL_CELLS,SMALL_RECOVERY).run(cells,2,SEED,EVENT_BUDGET)
    # Account actual completed/partial histories. Only complete runs telescope over all stages here.
    if result['complete']:
        initial=sum(c.counts[2] for c in cells);final=sum(c['counts'][2] for c in result['population']['cells'])
        extraction=sum(h[k]['extracted'] for h in result['history'] for k in ('batch','recovery'))
        production=sum(h[k]['signed_internal_formation'] for h in result['history'] for k in ('batch','recovery'))
        growth=sum(h['batch']['growth_consumed'] for h in result['history']);discard=sum(h['discarded_z'] for h in result['history'])
        error=production-(final-initial+extraction+growth+discard)
        if error:raise ArithmeticError('Actual intracellular-z telescope failed')
        result['physical_z_ledger']=dict(initial=initial,final=final,extracted=extraction,growth=growth,discarded=discard,signed_internal=production,residual=error)
    result['scope']='Accelerated small physical direct-method diagnostic, not the finite stopped law or its theorem-success event'
    return result


def main():
    parser=argparse.ArgumentParser(description=__doc__);parser.add_argument('--output',default='outputs');parser.add_argument('--simulate',action='store_true');args=parser.parse_args()
    out=Path(args.output);out.mkdir(parents=True,exist_ok=True)
    exact=exact_audit();inheritance=InheritanceCertificate().evaluate();community=Community(COMPOSITION,RENEWAL,MORTALITY)
    equilibria=community.equilibria();traces=[];operating=[]
    for name,eq in zip(('low','high'),equilibria):
        obs=community.observables(eq);initial=eq.copy()
        perturb=np.asarray(ABUNDANCE_PERTURBATION)
        if len(perturb)!=len(COMPOSITION):raise ValueError('One abundance perturbation per consumer')
        initial[5:]*=perturb
        run=community.integrate(initial,RECOVERY_HORIZON);traces.append((name,run));o0=community.observables(initial)
        a0=float(min(o0['p']/community.q));records=[];maxreadout=maxcomposition=0.
        for t,state,ints in zip(run['times'],run['states'],run['integrals']):
            o=community.observables(state);d=community.field(state);co=community.composition(state)
            maxreadout=max(maxreadout,abs(o['J']-(sum(d[5:])+community.mu*o['S']+o['S']**2+o['W'])))
            actualp=(d[5:]*o['S']-state[5:]*sum(d[5:]))/o['S']**2
            maxcomposition=max(maxcomposition,float(max(abs(actualp-co['pdot']))))
            inferred=float(co['v']@co['pdot']/(co['v']@co['v'])) if co['derivative_condition']<1e10 else None
            records.append([t,o['S'],o['J'],o['load'],o['W'],o['chi'],o0['chi']*math.exp(-2*a0*ints[1]),inferred,co['derivative_condition'],*o['p']])
        run['records']=records
        operating.append(dict(name=name,state=eq,observables=obs,residual=float(max(abs(community.field(eq)))),
            largest_eigen_real=float(max(np.linalg.eigvals(community.jacobian(eq)).real)),
            balance_error_abundance=run['abundance_balance_error'],balance_error_reservoir=run['reservoir_balance_error'],
            instantaneous_readout_residual=maxreadout,composition_identity_residual=maxcomposition,
            final_state_distance=float(max(abs(run['states'][-1]-eq)))))
    window=WindowReadout(F(WINDOW),F(S_MAX),F(ABUNDANCE_ERROR),F(DISPERSION_ERROR),F(RESERVOIR_ERROR),F(AVERAGE_QUADRATURE_ERROR),F(AVERAGE_QUADRATURE_ERROR))
    assumptions=dict(calibrated=True,two_alternatives=True,recovery_bound=ADMITTED_DYNAMIC_ERROR)
    synthetic=[window.classify('0.0224','0.0222',**assumptions),window.classify('0.0402','0.0401',**assumptions),
        window.classify('0.0224','0.0222'),window.classify('0.0224','0.0401',**assumptions),window.classify('0.031','0.031',**assumptions)]
    # Imposed extraction has its own stationary roots; no feedback inheritance is inferred.
    resident=ResidentChemistry(EXTRACTION);centers=resident.centers();imposed=[]
    for tag,center in [('L',centers[0]),('H',centers[-1])]:
        lo,hi=rational_root(EXTRACTION,tag)
        imposed.append(dict(tag=tag,center=center,z_exact_enclosure=[lo,hi]))
    result=dict(exact=exact,inheritance=inheritance,imposed_states=imposed,feedback=operating,synthetic_readouts=synthetic,
        units=dict(newborn_volume_liters=float(F(NEWBORN_SIZE)/(F('6.02214076e23')*F(CONCENTRATION_MOLAR))),
            initial_population_liters=float(F(NEWBORN_SIZE*RETAINED_CELLS)/(F('6.02214076e23')*F(CONCENTRATION_MOLAR))),
            two_deadlines_and_recoveries_seconds=(16/float(F(GROWTH_COUPLING))+10752)*TIME_UNIT_SECONDS),
        inputs=dict(extraction=EXTRACTION,composition=COMPOSITION,renewal=RENEWAL,mortality=MORTALITY))
    if args.simulate:result['small_protocol']=small_protocol()
    dump=lambda name,obj:(out/name).write_text(json.dumps(obj,indent=2,default=lambda v:v.tolist() if isinstance(v,np.ndarray) else str(v) if isinstance(v,F) else float(v))+'\n',encoding='utf-8')
    dump('results.json',result)
    def table(name,header,rows):
        with (out/name).open('w',newline='') as f:w=csv.writer(f);w.writerow(header);w.writerows(rows)
    for name,run in traces:
        table(name+'_recovery.csv',['time','S','uptake','load','dispersion','chi_squared','chi_bound','S_from_exact_p_derivative','derivative_error_amplification',*[f'p{i+1}' for i in range(len(COMPOSITION))]],run['records'])
        table(name+'_states.csv',['time','A','B','z','H','R',*[f'X{i+1}' for i in range(len(COMPOSITION))]],np.c_[run['times'],run['states']])
    table('finite_repetition.csv',['cycles','probability_lower_exact','high_fraction_lower_numerical'],[(r['cycles'],str(r['probability_lower']),r['count_fraction_lower_numerical']) for r in inheritance['finite_repetition']])
    lines=[f'Two-cycle imposed-extraction success lower bound: {inheritance["joint_two_cycle_lower"]} = {float(inheritance["joint_two_cycle_lower"]):.12f}.',
        f'Final high-ancestry fraction lower bound (numerical evaluation): {inheritance["high_fraction_lower_numerical"]:.10f}.',
        'Same consumer proportions; uptake '+', '.join(f'{r["observables"]["J"]:.12f}' for r in operating)+'.',
        f'Exact stationary feedback-load cap: {exact["load_cap"]}; uptake separation margin: {float(exact["half_gap"]):.12f}.',
        'Synthetic readouts: '+', '.join(r['outcome'] for r in synthetic)+'.',
        'Feedback and imposed extraction share the resident field, not an inheritance theorem. Lean is not rerun.']
    (out/'console.txt').write_text('\n'.join(lines)+'\n');print('\n'.join(lines));plot(out,operating,traces,inheritance,window,exact)
    digest=lambda p:hashlib.sha256(p.read_bytes()).hexdigest()
    dump('run_metadata.json',dict(manuscript_sha256=MANUSCRIPT_SHA256,python=platform.python_version(),source_sha256=digest(Path(__file__)),
        module_sha256={p.name:digest(p) for p in Path(__file__).parent.glob('*.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,operating,traces,inheritance,window,exact):
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
    fig,axs=plt.subplots(1,3,figsize=(10,3.8),layout='constrained');x=np.arange(len(COMPOSITION))
    for k,r in enumerate(operating):
        axs[0].bar(x+(k-.5)*.35,r['observables']['p'],.35,label=r['name'])
        axs[1].bar(x+(k-.5)*.35,r['state'][5:],.35)
        axs[2].bar(k,r['observables']['J'],.6)
    axs[0].set(title='Consumer proportions at the two equilibria',xlabel='Consumer',ylabel='Fraction');axs[0].legend(fontsize=8)
    axs[1].set(title='Consumer abundances at the two equilibria',xlabel='Consumer',ylabel='Absolute abundance')
    for ax in axs[:2]:ax.set_xticks(x,[str(i+1) for i in x])
    axs[2].set(title='Total uptake at the two equilibria',ylabel='Uptake J');axs[2].set_xticks([0,1],['Low','High']);axs[2].axhline(float(exact['uptake_threshold']),color='k',ls=':')
    for ax in axs:ax.grid(axis='y',alpha=.2)
    fig.savefig(out/'hidden_states.png',dpi=180);fig.savefig(out/'hidden_states.svg');plt.close(fig)
    fig,axs=plt.subplots(1,2,figsize=(10,4),layout='constrained')
    for name,run in traces:
        a=np.array([row[:7] for row in run['records']],float);axs[0].semilogy(a[:,0],np.maximum(a[:,5],1e-30),label=name+' dispersion')
        axs[0].semilogy(a[:,0],np.maximum(a[:,6],1e-30),'--',label=name+' analytical bound')
    axs[0].set(title='Recovery of consumer proportions',xlabel='Normalized time',ylabel='Weighted composition error',ylim=(1e-16,.001),xlim=(0,min(350,RECOVERY_HORIZON)));axs[0].legend(fontsize=7)
    data=inheritance['finite_repetition'];axs[1].plot([r['cycles'] for r in data],[float(r['probability_lower']) for r in data],'o-')
    axs[1].set(title='Finite inheritance under imposed load',xlabel='Completed cycles',ylabel='Sufficient joint probability lower bound',ylim=(-.03,1.03))
    for ax in axs:ax.grid(alpha=.2)
    fig.savefig(out/'operational_limits.png',dpi=180);fig.savefig(out/'operational_limits.svg');plt.close(fig)


if __name__=='__main__':main()