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

def lift(m,h,t,eta,area):
    if not all(math.isfinite(x) for x in [m,h,t,eta,area]) or m<0 or h<0 or t<=0 or not 0<eta<=1 or area<=0:raise ValueError('Outside model domain')
    force=m*G;energy=force*h
    return dict(mass_kg=m,height_m=h,duration_s=t,efficiency=eta,area_m2=area,force_N=force,useful_energy_J=energy,average_useful_power_W=energy/t,input_energy_J=energy/eta,pressure_difference_Pa=force/area)

def save(name,rows):
    with (P/name).open('w',newline='',encoding='utf-8') as f:
        w=csv.DictWriter(f,fieldnames=list(rows[0]));w.writeheader();w.writerows(rows)
rows=[lift(m,k/20,t,e,a) for m,k,t,e,a in itertools.product([50,100,200],range(21),[10,20,40],[.5,.75,1],[.001,.005,.01])]
pressure=[dict(mass_kg=m,area_m2=k/100000,pressure_difference_Pa=m*G/(k/100000)) for m in [50,100,200] for k in range(100,1001,10)]
geometry=[dict(left_height_m=.6,right_height_m=k/100,separation_m=1.2,slope_angle_deg=math.degrees(math.atan2(k/100-.6,1.2))) for k in range(40,121)]
save('lift-grid.csv',rows);save('pressure-area.csv',pressure);save('support-geometry.csv',geometry)
for r in rows:
    assert math.isclose(r['pressure_difference_Pa']*r['area_m2'],r['force_N'],rel_tol=1e-14)
    assert math.isclose(r['force_N']*r['height_m'],r['useful_energy_J'],rel_tol=1e-14,abs_tol=1e-14)
    assert math.isclose(r['average_useful_power_W']*r['duration_s'],r['useful_energy_J'],rel_tol=1e-14,abs_tol=1e-14)
    assert math.isclose(r['input_energy_J']*r['efficiency'],r['useful_energy_J'],rel_tol=1e-14,abs_tol=1e-14)
    if r['height_m']==0:assert r['useful_energy_J']==0 and r['force_N']>0
for r in geometry:assert math.isclose(1.2*math.tan(math.radians(r['slope_angle_deg'])),r['right_height_m']-.6,abs_tol=1e-14)
assert lift(0,1,10,.5,.001)['force_N']==0
assert all(lift(0,1,10,.5,.001)[k]==0 for k in ['useful_energy_J','average_useful_power_W','input_energy_J','pressure_difference_Pa'])
a=lift(100,1,10,.5,.001);b=lift(100,1,20,.5,.001)
assert a['useful_energy_J']==b['useful_energy_J'] and a['average_useful_power_W']==2*b['average_useful_power_W']
invalid=[(-1,1,10,1,.001),(1,-1,10,1,.001),(1,1,0,1,.001),(1,1,10,0,.001),(1,1,10,1.1,.001),(1,1,10,1,0),(float('nan'),1,10,1,.001)]
for args in invalid:
    try:lift(*args)
    except ValueError:pass
    else:raise AssertionError(args)
plt.rcParams.update({'font.size':10,'axes.spines.top':False,'axes.spines.right':False,'savefig.dpi':180})
colors=['#3769a2','#7957a8','#278c89']
fig,ax=plt.subplots(1,2,figsize=(10.5,4.5),layout='constrained')
h=np.linspace(0,1,101);t=np.linspace(10,40,121)
for m,c in zip([50,100,200],colors):
    ax[0].plot(h,m*G*h,color=c,label=f'{m} kg')
    ax[1].plot(t,m*G/t,color=c,label=f'{m} kg')
ax[0].set(xlabel='Vertical rise (m)',ylabel='Useful lifting energy (J)',title='A  Energy increases with height')
ax[1].set(xlabel='Lift duration (s)',ylabel='Average useful power (W)',title='B  Longer duration, same 1 m rise')
for a in ax:a.legend();a.grid(alpha=.2)
fig.suptitle('Ideal slow direct lift: chosen masses, no counterweight or friction')
for ext in ['png','svg','pdf']:fig.savefig(P/f'lift-energy-power.{ext}')
plt.close(fig)
fig,ax=plt.subplots(figsize=(7,4.8),layout='constrained')
for m,c in zip([50,100,200],colors):
    rr=[r for r in pressure if r['mass_kg']==m]
    ax.plot([r['area_m2']*1e4 for r in rr],[r['pressure_difference_Pa']/1e3 for r in rr],color=c,label=f'{m} kg')
ax.set(xlabel='Effective piston area (cm²)',ylabel='Ideal pressure difference (kPa)',title='A larger area reduces pressure for the same force\nSimple direct piston, uniform pressure difference')
ax.legend();ax.grid(alpha=.2)
for ext in ['png','svg','pdf']:fig.savefig(P/f'pressure-area.{ext}')
plt.close(fig)
(P/'validation.json').write_text(json.dumps(dict(all_pass=True,lift_rows=1701,pressure_rows=273,geometry_rows=81,dimensional_identities=True,zero_height_retains_force=True,zero_mass_returns_zero=True,time_energy_and_power_scaling=True,geometry_inverse_check=True,invalid_inputs_rejected=len(invalid)),indent=2)+'\n')
(P/'study.json').write_text(json.dumps(dict(title='Ideal lifting and adjustable support geometry',synthetic=True,g_m_s2=G,masses_kg=[50,100,200],height_m=dict(start=0,end=1,step=.05),durations_s=[10,20,40],efficiencies=[.5,.75,1],areas_m2=[.001,.005,.01],pressure_figure_grid=dict(start_m2=.001,end_m2=.01,step_m2=.0001),geometry=dict(left_height_m=.6,right_height_start_m=.4,right_height_end_m=1.2,right_height_step_m=.01,horizontal_separation_m=1.2),csv_rows=2055,limitations=['No counterweight, friction, acceleration, structural strength, stability or safety assessment','Pressure difference assumes a simple direct piston and appropriate effective area; no rod-side pressure model','Efficiency is hypothetical and affects input energy only','Support angle is geometry only; not load-bearing or comfort evidence']),indent=2)+'\n')
print('PASS: 1701 lift + 273 pressure + 81 geometry rows')