"""Mass/energy-conserving thermal-pressure prototype; SI units, dry ideal gas.

Uniform control volume + finite wall + compressible check-valve ports. An optional
downstream idealized expander takes an imposed pressure drop. All outputs are
predictions conditional on input parameters, never physical measurements.
"""
import csv,json,math
from pathlib import Path
import numpy as np
R=Path(__file__).resolve().parents[1]
KEYS=['m','U','Tw','Tc','Min','Mout','Hin','Hexhaust','Qhot','Qcold','Qambient','Wshaft','Bhot','Bcold','Bexhaust','Sexhaust','Qhot_over_Th','Qcold_over_Tc','Qambient_over_Ta']
I={k:i for i,k in enumerate(KEYS)}

def config():return json.loads((R/'design/config.json').read_text())

def properties(c):
    cv=c['gas_constant_JkgK']/(c['gamma']-1);return cv,cv+c['gas_constant_JkgK']

def port(pu,tu,pd,c):
    if pu<=pd:return 0.
    g=c['gamma'];r=max(pd/pu,1e-9);critical=(2/(g+1))**(g/(g-1))
    if r<=critical:f=math.sqrt(g)*(2/(g+1))**((g+1)/(2*(g-1)))
    else:f=math.sqrt(max(0,2*g/(g-1)*(r**(2/g)-r**((g+1)/g))))
    a=math.pi*c['port_diameter_m']**2/4
    smooth=math.sqrt((pu-pd)/(pu-pd+c['low_pressure_regularization_Pa']))
    return c['discharge_coefficient']*a*pu/math.sqrt(c['gas_constant_JkgK']*tu)*f*smooth

def state(y,c):
    cv,cp=properties(c);m,U,Tw,Tc=y[:4];T=U/(m*cv);p=(c['gamma']-1)*U/c['volume_each_m3']
    return T,p

def stored(y,c):
    cv,cp=properties(c);m,U,Tw,Tc=y[:4];T,p=state(y,c);ta=c['ambient_temperature_K'];pa=c['ambient_pressure_Pa']
    S=m*(cp*math.log(T/ta)-c['gas_constant_JkgK']*math.log(p/pa))+c['wall_heat_capacity_JK']*math.log(Tw/ta)
    B=U+pa*c['volume_each_m3']-m*cp*ta-ta*m*(cp*math.log(T/ta)-c['gas_constant_JkgK']*math.log(p/pa))
    B+=c['wall_heat_capacity_JK']*(Tw-ta-ta*math.log(Tw/ta))
    return U+c['wall_heat_capacity_JK']*Tw,S,B

def rates(y,c,cold,mode,stage=0):
    cv,cp=properties(c);m,U,Tw,Tc=y[:4];T,p=state(y,c);ta=c['ambient_temperature_K'];pa=c['ambient_pressure_Pa'];th=c['hot_temperature_K']
    timed=c.get('timed_valves',False) and mode not in ['sealed','hold','null']
    inlet_allowed=mode!='sealed' and (not timed or stage==1)
    outlet_allowed=mode not in ['sealed','hold','null'] and (not timed or stage==3)
    mi=mo=work=0.;ti=ta;te=T
    if inlet_allowed and pa>p:
        pb=pa-c.get('head_fraction',0)*max(0,pa-p) if c.get('load_method')=='fractional_head' else pa
        ti=ta*(1-c['expander_efficiency_assumed']*(1-(pb/pa)**((c['gamma']-1)/c['gamma'])))
        mi=port(pb,ti,p,c);work+=mi*cp*(ta-ti)
    if outlet_allowed and p>pa:
        pb=p-c.get('head_fraction',0)*max(0,p-pa) if c.get('load_method')=='fractional_head' else p-c['load_pressure_Pa']
        if pb>pa:
            te=T*(1-c['expander_efficiency_assumed']*(1-(pb/p)**((c['gamma']-1)/c['gamma'])))
            mo=port(pb,te,pa,c);work+=mo*cp*(T-te)
    qh=0 if cold else c['hot_wall_conductance_WK']*max(0,th-Tw)
    qc=c['cold_wall_conductance_WK']*max(0,Tw-Tc) if cold else 0
    qa=c['ambient_wall_conductance_WK']*(ta-Tw)
    qg=c['gas_wall_conductance_WK']*(Tw-T)
    a=np.zeros(len(KEYS));a[0]=mi-mo;a[1]=mi*cp*ti-mo*cp*T+qg
    a[2]=(qh-qc+qa-qg)/c['wall_heat_capacity_JK'];a[3]=qc/c['finite_cold_heat_capacity_JK'] if mode=='finite' else 0
    sex=cp*math.log(te/ta)
    for k,v in {'Min':mi,'Mout':mo,'Hin':mi*cp*ta,'Hexhaust':mo*cp*te,'Qhot':qh,'Qcold':qc,'Qambient':qa,'Wshaft':work,'Bhot':qh*(1-ta/th),'Bcold':qc*(ta/Tc-1),'Bexhaust':mo*cp*(te-ta-ta*math.log(te/ta)),'Sexhaust':mo*sex,'Qhot_over_Th':qh/th,'Qcold_over_Tc':qc/Tc,'Qambient_over_Ta':qa/ta}.items():a[I[k]]=v
    return a,{'T_K':T,'p_Pa':p,'inlet_kgs':mi,'outlet_kgs':mo,'shaft_W':work,'cold_W':qc,'hot_W':qh}

def simulate(c,mode='cycle',phase=0.,duration=None,record=True):
    cv,cp=properties(c);ta=c['ambient_temperature_K'];pa=c['ambient_pressure_Pa'];dt=c['step_s'];period=c['period_s']
    end=duration if duration is not None else c['cycles']*period
    n=round(end/dt);dt=end/n
    y=np.zeros(len(KEYS));y[0]=pa*c['volume_each_m3']/(c['gas_constant_JkgK']*ta);y[1]=y[0]*cv*ta;y[2]=ta;y[3]=c['cold_temperature_K']
    initial=y.copy();history=[];cycle_start=None
    stride=max(1,round(.1/dt))
    for j in range(n+1):
        t=j*dt
        if j==round((end-period)/dt):cycle_start=y.copy()
        cold= mode in ['hold','sealed','null'] or ((t+phase)%period<period/2)
        if record and (j%stride==0 or j==n):
            stage=int(((t+phase)%period)/(period/4))
            _,r=rates(y,c,cold,mode,stage);history.append(dict(time_s=t,wall_K=y[2],cold_reservoir_K=y[3],**r))
        if j==n:break
        # Thermal switches align exactly with steps; all RK stages use one state.
        cold= mode in ['hold','sealed','null'] or (((t+dt/2)+phase)%period<period/2)
        stage=int((((t+dt/2)+phase)%period)/(period/4))
        k1=rates(y,c,cold,mode,stage)[0];k2=rates(y+dt*k1/2,c,cold,mode,stage)[0]
        k3=rates(y+dt*k2/2,c,cold,mode,stage)[0];k4=rates(y+dt*k3,c,cold,mode,stage)[0]
        y+=dt*(k1+2*k2+2*k3+k4)/6
        if y[0]<=0 or state(y,c)[0]<100 or not np.isfinite(y).all():raise ValueError('Invalid gas state or integration')
    e0,s0,b0=stored(initial,c);e1,s1,b1=stored(y,c);ta=c['ambient_temperature_K']
    z=dict(zip(KEYS,y.tolist()));z.update(mode=mode,duration_s=end,final_T_K=state(y,c)[0],final_p_Pa=state(y,c)[1],
       mass_residual_kg=y[0]-initial[0]-z['Min']+z['Mout'],
       energy_residual_J=e1-e0-(z['Hin']-z['Hexhaust']+z['Qhot']-z['Qcold']+z['Qambient']-z['Wshaft']),
       entropy_generation_JK=s1-s0+z['Sexhaust']-z['Qhot_over_Th']+z['Qcold_over_Tc']-z['Qambient_over_Ta'])
    z['exergy_residual_J']=b1-b0-(z['Bhot']+z['Bcold']-z['Wshaft']-z['Bexhaust']-ta*z['entropy_generation_JK'])
    z['minimum_refrigeration_reset_work_J']=z['Bcold']
    if cycle_start is not None:
        delta=y-cycle_start
        z['last_cycle']={k:delta[I[k]] for k in KEYS[4:]}
        z['last_cycle'].update(mean_inlet_gs=1000*delta[I['Min']]/period,mean_shaft_W=delta[I['Wshaft']]/period,
            minimum_cold_reset_power_W=delta[I['Bcold']]/period,
            gas_and_wall_endpoint_change_J=stored(y,c)[0]-stored(cycle_start,c)[0],
            gas_mass_endpoint_change_kg=y[0]-cycle_start[0])
    return z,history

def write_csv(path,records):
    with path.open('w',newline='') as f:
        w=csv.DictWriter(f,fieldnames=list(records[0]));w.writeheader();w.writerows(records)

if __name__=='__main__':
    c=config();z,h=simulate(c);write_csv(R/'output/single-chamber.csv',h);print(json.dumps(z,indent=2))