"""Finite-time polymer reactor output, retained rewards and exact error budgets."""
from pathlib import Path
from fractions import Fraction as F
import argparse
import csv
import hashlib
import json
import math
import platform
import numpy as np
from polymer import PolymerCatalogue,PolymerReactor,CatalyticAssignment,RetainedReward,FOOD,SELECTED
from certificate import ReactorCertificate,source_mass_alpha2_n4,readout_floor

# EDITABLE INPUTS. Dimensionless time = residence times; concentration = reference scale.
MAX_WORD_LENGTH=4
COUNT_SCALE=2*10**22
FOOD_CATALYSIS_BOUND=F(1,100000)
REWARD_ALLOWANCE=F(48)
BASAL_COEFFICIENT=2e-9
SELECTED_COEFFICIENT=4.
DESIGNED_BACKGROUND=()  # e.g. ('010', ('00','11','0011'), 16.) inside this tuple
DESIGNED_WEAK_FOOD=1e-5  # Every food catalyses every split at this coefficient.
REFERENCE_MOLAR=0.001
RESIDENCE_SECONDS=3600.
SSA_VOLUME=20  # Small-count diagnostic, far below the theorem scale.
SSA_SEED=55092026
SSA_EVENT_BUDGET=100000
MANUSCRIPT_SHA256='7d8ef88769b6d9fb3bb24d015352f39a001e752b76ceaade289523efc24a9601'


def generator_panel(catalogue,volume):
    background=[]
    for i,z in enumerate(catalogue.words):
        if z in FOOD:continue
        omitted=(i+1)%len(catalogue.splits)
        if z=='0011' and catalogue.splits[omitted]==SELECTED:omitted=(omitted+1)%len(catalogue.splits)
        for j,split in enumerate(catalogue.splits):
            if j!=omitted and not(z=='0011' and split==SELECTED):background.append((z,split,16.))
    # Diagnostic panel fixed at paper n=4. It is not exhaustive over assignments.
    reactor=PolymerReactor(catalogue,CatalyticAssignment(8e-9,4.,tuple(background),1e-5))
    cert=ReactorCertificate(volume,4);contract=cert.constants();reward=RetainedReward(volume);rows=[]
    for label in ['floor','trimer_background','depleted_foods','large_selected']:
        counts={z:volume if z in FOOD else 0 for z in catalogue.words};counts['0011']=cert.floor
        if label=='trimer_background':counts['010']=volume//10
        if label=='depleted_foods':counts['00']=counts['11']=volume//1000;counts['010']=volume
        if label=='large_selected':counts['0011']=volume//100
        row=reward.diagnostic(reactor,counts)
        row.update(state=label,drift_lower=float(contract.drift)-624*row['nonfood'],variance_upper=float(contract.reward.variance_rate),jump_upper=float(contract.reward.jump_bound))
        assert row['mass']<=11 and row['drift']>=row['drift_lower']-1e-8 and row['variance']<=row['variance_upper']*(1+1e-10)
        assert abs(row['mass_drift']-(10-row['mass']))<1e-9 and abs(row['output_rate']-row['nonfood'])<1e-9
        rows.append(row)
    return rows


def main():
    parser=argparse.ArgumentParser();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)
    cert=ReactorCertificate(COUNT_SCALE,MAX_WORD_LENGTH,FOOD_CATALYSIS_BOUND)
    configured=cert.evaluate(REWARD_ALLOWANCE)
    rows=[]
    for v,eta,z in [(2*10**22,F(1,10**5),48),(2*10**22,F(1,10**4),37),(5*10**22,F(1,10**5),40),
                    (5*10**22,F(2,10**4),37),(10**23,F(1,10**5),40),(10**24,F(5,10**4),8)]:
        # Worst n allowed by each table row, so all smaller n are covered by the scalar monotonicity.
        result=ReactorCertificate(v,v//10**8,eta).evaluate(F(z));rows.append(result)
    allrows=[ReactorCertificate(v,v//10**9).evaluate(F(z),True) for v,z in [(10**23,26),(10**24,14)]]
    source=source_mass_alpha2_n4();lower=F(187365,10**8)
    assert source['witness_mass'][0]>lower
    screening_pass=(1-lower*(1-F(1,10**10)))**1600 < F(1,20)
    cat=PolymerCatalogue(MAX_WORD_LENGTH)
    assignment=CatalyticAssignment(BASAL_COEFFICIENT,SELECTED_COEFFICIENT,DESIGNED_BACKGROUND,DESIGNED_WEAK_FOOD)
    reactor=PolymerReactor(cat,assignment,float(FOOD_CATALYSIS_BOUND))
    times=np.unique(np.r_[np.linspace(0,199,797),1.,100.,199.]);trajectory=reactor.deterministic(times)
    size=len(cat.words);states=trajectory[:,:size];ledger=trajectory[:,size:];mass=states@reactor.length;nonfood=states@reactor.nonfood
    residual=float(np.max(np.abs(nonfood+ledger[:,0]-ledger[:,3])))
    assert residual<1e-7
    i1,i100,i199=[int(np.searchsorted(times,t)) for t in [1,100,199]]
    ode=dict(window_outputs=[float(ledger[i100,0]-ledger[i1,0]),float(ledger[i199,0]-ledger[i100,0])],
             stock=[float(states[i,cat.index['0011']]) for i in [i1,i100,i199]],max_mass=float(mass.max()),
             food_arrivals=float(ledger[-1,1]),nonfood_ledger_residual=residual,
             interpretation='Deterministic density illustration of one specified chemistry; not finite-count startup evidence')
    panel=generator_panel(PolymerCatalogue(4),2*10**22)
    lengths={z:len(z) for z in cat.words if len(z)>2}
    calibrated=readout_floor(F(135,1000),{z:F(95,100)*l for z,l in lengths.items()},lengths,F(2,1000))
    selective=readout_floor(F(135,1000),{'0011':F(4)},lengths,F(2,1000))
    inverse=[]
    for v in [2*10**22,5*10**22,10**23,10**24]:
        service=ReactorCertificate(v).constants()
        for rho in [1e-3,1e-6,1e-9,1e-12]:
            z=service.reward.inverse(rho,2,99)
            inverse.append([v,rho,z,(99*float(service.drift)-52-z)/624-.05])
    na=6.02214076e23
    result=dict(configured_certificate=configured,paper_mission_rows=rows,all_window_rows=allrows,
                source_alpha2_n4=dict(mean_degree=[float(v) for v in source['mean_degree']],witness_mass=[float(v) for v in source['witness_mass']],
                    exact_lower=lower,screening_1600_above_95_percent=screening_pass,scope='Independent redraws of chemistry, marks and trajectory; six silent food rows required'),
                generator_panel=panel,deterministic=ode,readout=calibrated,selective_readout=selective,
                units=dict(reference_molar=REFERENCE_MOLAR,residence_seconds=RESIDENCE_SECONDS,litres=COUNT_SCALE/na/REFERENCE_MOLAR,
                           mission_hours=199*RESIDENCE_SECONDS/3600,selected_forward_M_minus2_s_minus1=SELECTED_COEFFICIENT/(RESIDENCE_SECONDS*REFERENCE_MOLAR**2),
                           selected_reverse_M_minus1_s_minus1=SELECTED_COEFFICIENT/(RESIDENCE_SECONDS*REFERENCE_MOLAR)),
                evidence='Exact scalar enclosures + literal generator diagnostics + numerical ODE. Stochastic theorems imported from manuscript; Lean not rerun.')
    if args.simulate:result['small_count_trajectory']=reactor.simulate(SSA_VOLUME,seed=SSA_SEED,event_budget=SSA_EVENT_BUDGET)
    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,data):
        with (out/name).open('w',newline='') as f:w=csv.writer(f);w.writerow(header);w.writerows(data)
    dump('results.json',result)
    table('mission.csv',['V','eta','z','h','quota','error_upper','reward_exponent'],[[r[k] for k in ['V','eta','z','h','quota_99','error_upper','reward_exponent']] for r in rows])
    table('inverse_planning.csv',['V','reward_error_only','z','quota_numeric'],inverse)
    table('generator_panel.csv',list(panel[0]),[[r[k] for k in panel[0]] for r in panel])
    table('trajectory.csv',['residence_time',*cat.words,'Q','food_arrivals','food_monomers','signed_synthesis','total_mass','nonfood_mass'],
          np.column_stack([times,trajectory,mass,nonfood]))
    detail=(f'quota/window {float(configured["quota_99"]):.8f} V; failure upper {float(configured["error_upper"]):.4g}'
            if 'quota_99' in configured else configured['reason'])
    lines=[f'Configured mission: {configured["status"]}; {detail}.',
           f'Six paper operating points and two all-window points replayed with exact rational bounds.',
           f'Catalogue: {len(cat.words)} words, {len(cat.splits)} ordered reversible splits, {len(reactor.channels)} marked channels.',
           f'Designed deterministic window outputs: {ode["window_outputs"]}; ledger discrepancy {residual:.3g}.',
           f'Chemistry witness mass at n=4, alpha=2: {float(source["witness_mass"][0]):.10f}.',
           'No deterministic trace or sampled generator panel establishes the stochastic probability claim.']
    (out/'console.txt').write_text('\n'.join(lines)+'\n');print('\n'.join(lines))
    plot(out,rows,allrows,times,states,ledger,cat,panel)
    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 ['polymer.py','certificate.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,rows,allrows,times,states,ledger,cat,panel):
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
    fig,axs=plt.subplots(1,2,figsize=(10,4),layout='constrained')
    positions=np.arange(len(rows));axs[0].bar(positions,[float(r['quota_99']) for r in rows]);axs[0].axhline(.1,color='k',ls=':',label='Mission quota')
    axs[0].set_xticks(positions,[f'{r["V"]:.0e}\n{float(r["eta"]):.0e}' for r in rows],fontsize=8)
    axs[0].set(title='Six sufficient operating points',xlabel='Count scale V / food coefficient bound',ylabel='Guaranteed export per window / V');axs[0].legend(fontsize=8)
    d=np.linspace(0,198,300)
    for r in allrows:
        q=(float(r['A'])*d-52-2*float(r['z']))/624-.05
        axs[1].plot(d,np.maximum(q,0),label=f'V ≥ {r["V"]:.0e}')
    axs[1].axhline(.1,color='k',ls=':');axs[1].set(title='Simultaneous output guarantees by window\nduration',xlabel='Window duration (residence times)',ylabel='Positive guaranteed output / V');axs[1].legend(fontsize=8)
    for ax in axs:ax.grid(alpha=.2)
    fig.savefig(out/'certificates.png',dpi=180);fig.savefig(out/'certificates.svg');plt.close(fig)
    fig,axs=plt.subplots(1,2,figsize=(10,4),layout='constrained')
    stockline=axs[0].plot(times,states[:,cat.index['0011']],label='Tetramer (left)',color='tab:blue')
    export_axis=axs[0].twinx()
    exportline=export_axis.plot(times,ledger[:,0],label='Output (right)',color='tab:orange')
    export_axis.set_ylabel('Cumulative nonfood output / V',color='tab:orange')
    axs[0].axvline(1,color='gray',ls=':');axs[0].axvline(100,color='gray',ls=':')
    axs[0].set(title='One designed chemistry: density ODE',xlabel='Time (residence times)',ylabel='Tetramer concentration / reference')
    axs[0].legend(stockline+exportline,[line.get_label() for line in stockline+exportline],fontsize=8,loc='center right')
    p=np.arange(len(panel));axs[1].bar(p-.18,[r['variance'] for r in panel],width=.36,label='Retained reward')
    axs[1].bar(p+.18,[r['full_variance'] for r in panel],width=.36,label='Full potential')
    axs[1].axhline(panel[0]['variance_upper'],color='k',ls=':',label='Uniform retained bound')
    axs[1].set_xticks(p,['Floor','Trimer','Low food','High stock'],rotation=15,fontsize=8)
    axs[1].set(yscale='log',title='Fluctuation rates for full and retained\ngrowth measures',ylabel='Generator quadratic rate');axs[1].legend(fontsize=8)
    for ax in axs:ax.grid(alpha=.2)
    fig.savefig(out/'mechanism.png',dpi=180);fig.savefig(out/'mechanism.svg');plt.close(fig)


if __name__=='__main__':main()