"""Finite-bath missions: reusable count/density reactors and explicit sizing."""
COPY_SCALE = 96_000_000_000
MISSION_CYCLES = 100
FAILURE_TARGET = '21/1000000'
BATH_TOLERANCE = '1/10'
TEMPLATE_DEMAND = 0
FREE_TEMPLATE_DEMAND = 0
RELEASE_RATE = '20'
DRIVE_RATE = '1/50'
DENSITY_CYCLES = 8
RETENTION = '1/4'
SURVIVAL = ('49/50',)*6
REFILL_ERRORS = ('-1/200','-1/200')
SMALL_COPY_SCALE = 300
SMALL_CYCLES = 3
SMALL_BATH_CAPACITY = 3000
SMALL_EVENT_BUDGET = 100000
RANDOM_SEED = 760023
REFERENCE_MOLAR = '1/1000'
BATH_REFERENCE_MOLAR = '1/1000'
TIME_UNIT_SECONDS = '60'
MANUSCRIPT_SHA256 = '5627f554fdb4e90cf31a8f38761fd154b3d9c13f3e8d712819e81c462c9c9a5f'

import argparse,csv,hashlib,json,platform
from pathlib import Path
from fractions import Fraction as Q
from math import ceil
import numpy as np
import scipy
from mpmath import mp,iv
from chemistry import Intervention,dot,A,B,I,Y
from finite_bath import Bath,FiniteBathReactor,DensityMission,CountMission,HistoryController,FoodMeter
from certificates import MissionCertificate,design,loaded_force_capacity,bath_energy,neighbour_balance,KAPPA,exp_lower
import paper_checks


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)
    def dump(name,value):(out/name).write_text(json.dumps(value,indent=2,default=str)+'\n')
    def table(name,rows):
        with (out/name).open('w',newline='') as f:w=csv.DictWriter(f,fieldnames=list(rows[0]));w.writeheader();w.writerows(rows)
    if not 0<Q(BATH_TOLERANCE)<1:raise ValueError('Bath tolerance must lie in (0,1).')
    reactor=FiniteBathReactor(RELEASE_RATE,DRIVE_RATE);R=max(1,ceil(Q(MISSION_CYCLES*(COPY_SCALE//10))/Q(BATH_TOLERANCE)));bath=Bath(R,R,0)
    cert=MissionCertificate(COPY_SCALE,reactor,bath,sharp=True)
    dump('configured_mission.json',dict(certificate=cert.evaluate(MISSION_CYCLES),inventory=cert.inventory(MISSION_CYCLES),all_free_seed_inventory=cert.inventory(MISSION_CYCLES,COPY_SCALE),bath_capacity=R,scope='Configured certificate requires the pure-total sharp-service class. No count simulation at this scale.'))
    requested=design(MISSION_CYCLES,FAILURE_TARGET,BATH_TOLERANCE,TEMPLATE_DEMAND,FREE_TEMPLATE_DEMAND);dump('inverse_design.json',requested)
    comparisons=[]
    for tag,V,refined,sharp in [('original',200000000000,False,False),('refined_reduced',96000000000,True,True),('refined_same_output',200000000000,True,True)]:
        c=MissionCertificate(V,FiniteBathReactor(),Bath(1,1,0),refined,sharp);inventory=c.inventory(100,V)
        comparisons.append(dict(case=tag,certificate=c.evaluate(100),inventory=inventory,reservoir=10*inventory['gross_service_allowance']))
    dump('canonical_comparison.json',comparisons)
    fixed=design(100,FAILURE_TARGET,BATH_TOLERANCE,357142857200,18518518600);dump('fixed_output_design.json',fixed)
    loaded=MissionCertificate(COPY_SCALE,reactor,Bath(R,R,R),sharp=False);loaded_inventory=loaded.inventory(MISSION_CYCLES)
    dump('loaded_bath_design.json',dict(certificate=loaded.evaluate(MISSION_CYCLES),inventory=loaded_inventory,capacity_for_ten_percent=10*loaded_inventory['gross_service_allowance'],capacity_for_force_01=loaded_force_capacity(loaded_inventory['gross_service_allowance'],'1/10'),scope='Loaded bath uses floor(V/5), not the pure bath floor(V/10); phase refinement still applies. Force capacity uses its actual gross allowance.'))
    dump('paper_arithmetic_replay.json',paper_checks.payload)
    pulse=Intervention(RETENTION,SURVIVAL,REFILL_ERRORS);c2=Q(1,28);initial=[Q(159,160)-2*c2,Q(159,160)-2*c2,0,0,c2,0]
    cases=[];summaries=[]
    for name,sigma,total,d in [('pure',1,1,DRIVE_RATE),('loaded',1,2,DRIVE_RATE),('small_capacity',.001,1,DRIVE_RATE),('zero_drive',1,1,'0')]:
        r=FiniteBathReactor(RELEASE_RATE,d);experiment=DensityMission(r,sigma,total);res=experiment.run(initial+[1,total-1],DENSITY_CYCLES,pulse)
        check=experiment.run(initial+[1,total-1],DENSITY_CYCLES,pulse,method='BDF')
        error=float(np.max(np.abs(np.array([h['endpoint'] for h in res['history']])-np.array([h['endpoint'] for h in check['history']]))))
        summaries.append(dict(case=name,R_per_V=sigma,total_activity=total,drive=d,min_collected_I=min(h['counters']['collected_I'] for h in res['history']),min_collected_X=min(h['counters']['collected_X'] for h in res['history']),forward=sum(h['counters']['forward'] for h in res['history']),reverse=sum(h['counters']['reverse'] for h in res['history']),final_fuel=res['history'][-1]['endpoint'][6],max_inventory_residual=max(abs(h['inventory_residual']) for h in res['history']),max_bath_residual=max(abs(h['bath_residual']) for h in res['history']),solver_difference=error))
        cases.append(dict(case=name,**res));print('Density',name,'minimum collected I:',summaries[-1]['min_collected_I'],flush=True)
    dump('density_histories.json',cases);table('density_summary.csv',summaries)
    table('density_trajectories.csv',[dict(case=c['case'],**row) for c in cases for row in c['trajectory']])
    small=CountMission(reactor);init=(0,0,SMALL_COPY_SCALE,0,0,0);smallbath=Bath(SMALL_BATH_CAPACITY,SMALL_BATH_CAPACITY,0)
    sample=small.run(init,SMALL_COPY_SCALE,smallbath,SMALL_CYCLES,HistoryController(),RANDOM_SEED,SMALL_EVENT_BUDGET,FoodMeter(5*SMALL_CYCLES*SMALL_COPY_SCALE+1,5*SMALL_CYCLES*SMALL_COPY_SCALE+1));dump('small_count_mission.json',sample)
    shortage=small.run(init,SMALL_COPY_SCALE,smallbath,SMALL_CYCLES,HistoryController(),RANDOM_SEED,SMALL_EVENT_BUDGET,FoodMeter(1,1));dump('metered_shutdown.json',shortage)
    checks=[]
    for f,p in [(0,100),(40,60),(100,0)]:
        b=Bath(100,f,p);N=(20,30,10,3,4,5);V=100
        assert reactor.exact_drift(N,V,b,A)==V-dot(A,N) and reactor.exact_drift(N,V,b,B)==V-dot(B,N)
        checks.append(dict(f=f,p=p,material_A=str(reactor.exact_drift(N,V,b,A)),material_B=str(reactor.exact_drift(N,V,b,B)),forward=str(reactor.propensities(N,V,b)[18]),reverse=str(reactor.propensities(N,V,b)[19])))
    dump('generator_and_boundary_checks.json',checks)
    thermo=[bath_energy(10,J) for J in [0,1,5,10]]+[bath_energy(10,J,True) for J in [-10,-1,0,1,10]]+[bath_energy(R,R//10)]
    dump('bath_endpoint_free_energy.json',thermo);dump('neighbour_detailed_balance.json',[neighbour_balance((2,3,4,0,0,0),100,Bath(10,10,0),reactor),neighbour_balance((2,3,4,0,0,0),100,Bath(10,4,6),reactor)])
    # Full channel-specific quadratic variation, as a diagnostic for the paper's
    # proposed next refinement. No new probability theorem is inferred.
    w=(0,0,Q(1,10),Q(1,20),Q(1,50),Q(1,8));N=(20,30,10,3,4,5);b=Bath(100,40,60)
    terms=[a*dot(w,c.jump)**2 for a,c in zip(reactor.propensities(N,100,b),reactor.internal.channels)]
    dump('channel_quadratic_variation.json',dict(state=N,V=100,weights=w,by_label=terms,total=sum(terms),scope='Exact local diagnostic, not a replacement for the proved phase envelope.'))
    NA=Q(602214076000000000000000);c=Q(REFERENCE_MOLAR);cb=Q(BATH_REFERENCE_MOLAR);tau=Q(TIME_UNIT_SECONDS)
    if min(c,cb,tau)<=0:raise ValueError('Positive physical reference choices required.')
    dump('dimensional_reading.json',dict(reactor_nL=float(Q(COPY_SCALE)/NA/c*10**9),bath_nL=float(Q(R)/NA/cb*10**9),source_minutes=float(4*MISSION_CYCLES*tau/60),scope='Illustrative reference choices, not chemical calibration; the separated-bath contact rate requires its own convention.'))
    plot(out,cases,summaries,comparisons)
    print('Source arithmetic checks:',paper_checks.payload['total'],'; count diagnostic:',sample['status'],sample.get('all_sharp_success'),flush=True)
    print('Configured certified joint lower:',cert.evaluate(MISSION_CYCLES)['joint_product_lower_rational'],flush=True)
    print('Fuel is carried forward and both directions are counted. Productive output is not a lower bound on fuel consumption. Imported process theorem; Lean not rerun.',flush=True)
    here=Path(__file__).parent;digest=lambda p:hashlib.sha256(p.read_bytes()).hexdigest()
    dump('run_metadata.json',dict(manuscript_sha256=MANUSCRIPT_SHA256,source_sha256=digest(Path(__file__)),python=platform.python_version(),numpy=np.__version__,scipy=scipy.__version__,module_sha256={p.name:digest(p) for p in sorted(here.glob('*.py')) if p.name not in ['example.py','test_example.py']},input_sha256={p.name:digest(p) for p in sorted(here.glob('*.json')) if p.name!='release.json'},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,cases,summaries,comparisons):
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
    fig,axs=plt.subplots(1,2,figsize=(11,4),layout='constrained')
    v=np.linspace(5e10,2.5e11,200);axs[0].plot(v/1e11,(np.log(10100)-float(KAPPA)*v)/np.log(10),label='Refined, valid in shown range');valid=v>=2e11;axs[0].plot(v[valid]/1e11,(np.log(10100)-v[valid]/1e10)/np.log(10),'--',label='Original, from its valid floor');axs[0].axhline(np.log10(21e-6),color='gray',ls=':');axs[0].set(xlabel='Copy scale / 100 billion',ylabel='log10 mission failure upper bound',title='Original and refined mission-failure bounds');axs[0].legend(fontsize=8)
    axs[1].bar(['Original','Reduced scale','Same output'],[r['reservoir']/1e12 for r in comparisons]);axs[1].set(ylabel='Sufficient pure fuel / trillion molecules',title='Sufficient fuel stocks under three mission\ndesigns');axs[1].text(1,30,'Reduced scale also\nreduces guaranteed output',ha='center',fontsize=9)
    for a in axs:a.grid(alpha=.2)
    fig.savefig(out/'mission_resource_design.png',dpi=180);fig.savefig(out/'mission_resource_design.svg');plt.close(fig)
    fig,axs=plt.subplots(1,2,figsize=(11,4),layout='constrained')
    for c in cases:
        axs[0].semilogy([r['time'] for r in c['trajectory']],[r['fuel_activity'] for r in c['trajectory']],label=c['case'].replace('_',' '))
        axs[1].plot([r['cycle'] for r in c['history']],[r['counters']['collected_I'] for r in c['history']],'o-',label=c['case'].replace('_',' '))
    axs[0].set(xlabel='Source-running time',ylabel='Fuel activity f/R',title='Fuel activity over successive operating\ncycles');axs[1].set(xlabel='Cycle',ylabel='Collected template equivalents / V',title='Collected output across bath and drive\nconditions')
    for a in axs:a.grid(alpha=.2);a.legend(fontsize=8)
    fig.savefig(out/'finite_bath_operation.png',dpi=180);fig.savefig(out/'finite_bath_operation.svg');plt.close(fig)


if __name__=='__main__':main()