from pathlib import Path
import math,csv,json,sys,hashlib
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
P=Path(__file__).resolve().parent

def transmission(tau):
 if not math.isfinite(tau)or not 0<=tau<=8:raise ValueError('Optical depth must be finite in [0,8]')
 return math.exp(-tau)
def layers(a,b):
 transmission(a);transmission(b)
 if a+b>8:raise ValueError('Total optical depth exceeds study range')
 return math.exp(-(a+b))
rows=[dict(optical_depth=i/100,uncollided_fraction=transmission(i/100),removed_from_uncollided_fraction=0.0 if i==0 else -math.expm1(-i/100))for i in range(801)]
lr=[dict(first_depth=i/10,second_depth=j/10,total_depth=(i+j)/10,first_transmission=transmission(i/10),second_transmission=transmission(j/10),combined_transmission=layers(i/10,j/10))for i in range(41)for j in range(41)]
maxerr=max(abs(r['first_transmission']*r['second_transmission']-r['combined_transmission'])for r in lr)
assert maxerr<1e-14 and rows[0]['uncollided_fraction']==1
assert all(abs(r['uncollided_fraction']+r['removed_from_uncollided_fraction']-1)<1e-14 for r in rows)
assert all(a['uncollided_fraction']>b['uncollided_fraction']for a,b in zip(rows,rows[1:]))
# Independent midpoint integration of dT/dtau=-T at optical depth4.
errors=[]
for n in (100,200,400):
 h=4/n;y=1
 for _ in range(n):y+=h*(-(y-h*y/2))
 errors.append(abs(y-transmission(4)))
assert 3.8<errors[1]/errors[2]<4.2
assert math.isclose(math.exp(-math.log(2)),.5)
guards=0
for f,args in [(transmission,(-1,)),(transmission,(9,)),(transmission,(math.nan,)),(transmission,(math.inf,)),(layers,(5,4))]:
 try:f(*args)
 except ValueError:guards+=1
 else:raise AssertionError(args)
for name,data in [('attenuation.csv',rows),('layers.csv',lr)]:
 with(P/name).open('w',newline='',encoding='utf8')as f:w=csv.DictWriter(f,fieldnames=list(data[0]));w.writeheader();w.writerows(data)
plt.rcParams.update({'font.size':11,'axes.spines.top':False,'axes.spines.right':False,'savefig.dpi':180})
fig,axes=plt.subplots(1,2,figsize=(10,4.6),layout='constrained')
for ax in axes:ax.plot([r['optical_depth']for r in rows],[r['uncollided_fraction']for r in rows],color='#2462a0');ax.set(xlabel='Optical depth τ (dimensionless)',ylabel='Uncollided transmitted fraction T');ax.grid(alpha=.2)
axes[0].set_title('A  Linear scale');axes[1].set_yscale('log');axes[1].set_title('B  Logarithmic scale');fig.suptitle('Exponential removal from the uncollided beam\nAbstract model; not total radiation dose')
for e in ('png','svg','pdf'):fig.savefig(P/f'attenuation-curve.{e}')
plt.close(fig)
fig,ax=plt.subplots(figsize=(8,5),layout='constrained')
for a,ls in zip((0,1,2,4),(':','--','-.','-')):
 b=np.linspace(0,4,81);ax.plot(b,[layers(a,x)for x in b],ls=ls,label=f'First layer depth {a}')
ax.set(xlabel='Second layer optical depth (dimensionless)',ylabel='Combined uncollided transmitted fraction',title='Independent layer factors multiply\nTotal depth adds; physical layer properties are unspecified');ax.legend();ax.grid(alpha=.2)
for e in ('png','svg','pdf'):fig.savefig(P/f'layer-composition.{e}')
plt.close(fig)
(P/'validation.json').write_text(json.dumps(dict(all_pass=True,attenuation_rows=801,layer_rows=1681,layer_product_max_error=maxerr,ode_midpoint_steps=[100,200,400],ode_errors=errors,guards=guards,conservation_and_monotonicity=True),indent=2)+'\n')
(P/'study.json').write_text(json.dumps(dict(runtime=dict(python=sys.version,numpy=np.__version__,matplotlib=matplotlib.__version__),generator_sha256=hashlib.sha256(Path(__file__).read_bytes()).hexdigest(),dimensionless_only=True,attenuation_depth=[0,8,.01],layer_depths_each=[0,4,.1],total_dose_not_modeled=True,empirical_data=False),indent=2)+'\n')
print('PASS2482rows')