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

def wall(K,A,H,L):
 if not all(math.isfinite(v)for v in [K,A,H,L])or K<0 or A<0 or L<=0:raise ValueError('Invalid wall inputs')
 q=K*A*H/L;return q,q*3600000

def bucket(I,days=30):
 if not math.isfinite(I)or not 0<=I<=20:raise ValueError('Invalid input')
 s=20.;rr=[]
 for day in range(1,days+1):
  before=s;available=s+I;drain=max(available-40,0);withdraw=min(5,available-drain);unmet=5-withdraw;s=available-drain-withdraw
  rr.append(dict(infiltration_mm_day=I,day=day,storage_before_mm=before,incoming_mm=I,drainage_mm=drain,withdrawal_mm=withdraw,unmet_mm=unmet,storage_after_mm=s))
 return rr
wr=[dict(K_m_s=K,area_m2=.1,head_difference_m=k/50,thickness_m=.005,flow_m3_s=wall(K,.1,k/50,.005)[0],flow_L_h=wall(K,.1,k/50,.005)[1])for K in [1e-9,1e-8,1e-7]for k in range(-100,101)]
br=[r for I in range(21)for r in bucket(I)]
summ=[]
for I in range(21):
 rr=bucket(I);summ.append(dict(infiltration_mm_day=I,final_storage_mm=rr[-1]['storage_after_mm'],total_incoming_mm=30*I,total_drainage_mm=sum(r['drainage_mm']for r in rr),total_withdrawal_mm=sum(r['withdrawal_mm']for r in rr),total_unmet_mm=sum(r['unmet_mm']for r in rr)))
 for r in rr:assert abs(r['storage_before_mm']+I-r['drainage_mm']-r['withdrawal_mm']-r['storage_after_mm'])<1e-12 and 0<=r['storage_after_mm']<=40
 assert abs(20+30*I-summ[-1]['total_drainage_mm']-summ[-1]['total_withdrawal_mm']-summ[-1]['final_storage_mm'])<1e-12
assert wall(1e-8,.1,0,.005)==(0,0) and wall(0,.1,2,.005)==(0,0)
assert wall(1e-8,.1,-1,.005)[0]==-wall(1e-8,.1,1,.005)[0]
assert math.isclose(wall(1e-8,.1,1,.005)[1],.72)
assert bucket(0)[3]['storage_after_mm']==0 and bucket(0)[4]['unmet_mm']==5
assert bucket(5)[-1]['storage_after_mm']==20
# Event ordering: incoming20 into prior35 produces15 drainage before5 withdrawal.
assert bucket(20)[1]['drainage_mm']==15 and bucket(20)[1]['storage_after_mm']==35
# Independent transcription check against Siyal & Skaggs (2009), Table 1,
# printed p. 900 (USDA-hosted PDF p. 8). ME is defined there as simulated minus observed.
published=[
 dict(pressure_head_cm=25,mean_error_cm3_cm3=-.001,rmse_cm3_cm3=.011,reported_r_squared=.974),
 dict(pressure_head_cm=50,mean_error_cm3_cm3=-.004,rmse_cm3_cm3=.013,reported_r_squared=.970),
 dict(pressure_head_cm=100,mean_error_cm3_cm3=-.004,rmse_cm3_cm3=.012,reported_r_squared=.984),
 dict(pressure_head_cm=200,mean_error_cm3_cm3=-.006,rmse_cm3_cm3=.014,reported_r_squared=.983),
]
assert [r['mean_error_cm3_cm3'] for r in published]==[-.001,-.004,-.004,-.006]
assert [r['rmse_cm3_cm3'] for r in published]==[.011,.013,.012,.014]
assert [r['reported_r_squared'] for r in published]==[.974,.970,.984,.983]
bad=0
for fn,args in [(wall,(-1,.1,1,.005)),(wall,(1,-1,1,.005)),(wall,(1,1,1,0)),(wall,(1,1,math.nan,1)),(bucket,(-1,)),(bucket,(21,))]:
 try:fn(*args)
 except ValueError:bad+=1
 else:raise AssertionError(args)
def save(n,r):
 with(P/n).open('w',newline='',encoding='utf8')as f:w=csv.DictWriter(f,fieldnames=list(r[0]));w.writeheader();w.writerows(r)
for n,r in [('wall-flow.csv',wr),('bucket-days.csv',br),('bucket-summary.csv',summ),('published-fit-study.csv',published)]:save(n,r)
plt.rcParams.update({'font.size':10,'axes.spines.top':False,'axes.spines.right':False,'savefig.dpi':180})
fig,ax=plt.subplots(figsize=(8,4.8),layout='constrained')
for K,dash in zip([1e-9,1e-8,1e-7],[':','--','-']):
 rr=[r for r in wr if r['K_m_s']==K];ax.plot([r['head_difference_m']for r in rr],[r['flow_L_h']for r in rr],ls=dash,label=f'Assigned K = {K:.0e} m/s')
ax.axhline(0,color='gray',lw=.7);ax.set(xlabel='Head on first side minus second side (m)',ylabel='Signed flow from first to second side (L/h)',title='Flat saturated wall: area 0.1 m², thickness 0.005 m\nPrescribed heads; soil and pipe curvature omitted');ax.legend();ax.grid(alpha=.2)
for e in ['png','svg','pdf']:fig.savefig(P/f'wall-flow.{e}')
plt.close(fig)
for size,suffix in [((8,4.5),''),((5.5,4.5),'-mobile')]:
 fig,ax=plt.subplots(figsize=size,layout='constrained')
 x=np.arange(len(published)); width=.34
 ax.bar(x-width/2,[r['mean_error_cm3_cm3'] for r in published],width,label='Mean error (simulated − observed)',color='#6b4c9a')
 ax.bar(x+width/2,[r['rmse_cm3_cm3'] for r in published],width,label='Root mean square error',color='#007c91')
 ax.axhline(0,color='#24334a',lw=.8);ax.set_xticks(x,[str(r['pressure_head_cm']) for r in published]);ax.set(xlabel='Pipe pressure head (cm)',ylabel='Water-content error (cm³/cm³)',title='Published HYDRUS comparison after five days\nSiyal & Skaggs (2009), Table 1; one loam field site');ax.legend(fontsize=8);ax.grid(axis='y',alpha=.2)
 for e in ['png','svg','pdf']:fig.savefig(P/f'published-error-study{suffix}.{e}')
 plt.close(fig)
 svg_path=P/f'published-error-study{suffix}.svg'; svg=svg_path.read_text(encoding='utf8')
 desc='Four source-reported mean errors (simulated minus observed) are negative: minus 0.001, 0.004, 0.004 and 0.006 cubic centimetres per cubic centimetre at 25, 50, 100 and 200 centimetres. RMSE values are 0.011, 0.013, 0.012 and 0.014. One five-day loam-field study; not water savings.'
 svg=re.sub(r'<svg\b([^>]*)>',r'<svg\1 role="img" aria-labelledby="study-error-title study-error-desc"><title id="study-error-title">Published soil water model error statistics</title><desc id="study-error-desc">'+desc+'</desc>',svg,count=1)
 svg_path.write_text(svg,encoding='utf8')
fig,ax=plt.subplots(1,2,figsize=(11,4.8),layout='constrained')
for I,dash in zip([0,3,5,10,20],[':', '--','-','-.',(0,(3,1,1,1))]):
 rr=bucket(I);ax[0].plot([0]+[r['day']for r in rr],[20]+[r['storage_after_mm']for r in rr],ls=dash,label=f'I = {I} mm/day')
ax[0].set(xlabel='Day',ylabel='Stored water depth (mm)',title='A  End-of-day storage');ax[0].legend(fontsize=8)
for key,label,dash in [('total_drainage_mm','Drainage','-'),('total_unmet_mm','Unmet assigned demand','--')]:ax[1].plot([r['infiltration_mm_day']for r in summ],[r[key]for r in summ],ls=dash,label=label)
ax[1].set(xlabel='Assigned infiltrated irrigation (mm/day)',ylabel='30-day cumulative depth (mm)',title='B  Different destinations in the ledger');ax[1].legend(fontsize=8);fig.suptitle('Synthetic daily bucket: 20 mm initially; 40 mm hard cap; 5 mm/day demand\nDrainage occurs before demand; this is not a field-capacity model')
for e in ['png','svg','pdf']:fig.savefig(P/f'root-zone-ledger.{e}')
plt.close(fig)
(P/'validation.json').write_text(json.dumps(dict(all_pass=True,wall_rows=603,bucket_daily_rows=630,bucket_summaries=21,per_day_and_telescoping_conservation=True,signed_flow_and_units=True,event_order=True,zero_and_balanced_inputs=True,invalid_guards=bad,published_table1_rows=len(published),published_table1_values_match_source_transcription=True,published_table1_source_locator='Siyal & Skaggs (2009), Table 1, printed p. 900 / USDA-hosted PDF p. 8'),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(),wall=dict(equation='Q=K*A*deltaH/L',area_m2=.1,thickness_m=.005,K_m_s=[1e-9,1e-8,1e-7],head_range_m=[-2,2],limitations=['steady','saturated homogeneous flat wall','prescribed heads','no pipe curvature','no soil feedback','no unsaturated flow','no network losses']),bucket=dict(initial_mm=20,hard_capacity_mm=40,demand_mm_day=5,days=30,order='Add infiltrated irrigation; drain above cap; withdraw up to demand',other_fluxes='rain capillary rise runoff lateral exchange all zero',limitations=['accounting rule, not field capacity','no crop response','no ET physics','not coupled to wall model']),published_field_study=dict(citation='Siyal & Skaggs (2009), Agricultural Water Management 96, 893-904, doi:10.1016/j.agwat.2008.11.013',source_url='https://www.ars.usda.gov/arsuserfiles/20360500/pdf_pubs/P2266.pdf',source_locator='USDA-hosted PDF p. 8; printed p. 900; Table 1; ME definition Eq. 5 on PDF p. 8',conditions='one loam field site; four pipe pressure-head treatments; after five days',mean_error_definition='mean(simulated minus observed); negative indicates underprediction',metrics=published,csv_sha256=hashlib.sha256((P/'published-fit-study.csv').read_bytes()).hexdigest(),figure_generator='published-error-study.svg/.png/.pdf generated in this script',interpretation_limits=['transcribed reported statistics, not raw observations','not water saving','not crop yield','not universal performance']),independent_models=True),indent=2)+'\n');print('PASS1254rows; PASS4 published Table 1 rows')