"""Reproduce idealized thermal teaching figures. Python 3 + Matplotlib.
Writes beside this script. Calculations are not experimental observations.
"""
import csv
import json
import math
from pathlib import Path
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

OUT=Path(__file__).resolve().parent
GAMMA=1.4
THETA=5500*5/9
AREA=0.1
H=10.0
L2=0.002
K2=0.2
CAPACITY=2000.0
BOUNDARY=80.0
INITIAL=20.0

def gamma(T):
    q=THETA/T
    e=math.exp(-q)
    return 1+(GAMMA-1)/(1+(GAMMA-1)*q*q*e/(1-e)**2)

def mach_squared(T,Tt):
    return 2*Tt/T/gamma(T)*((GAMMA/(GAMMA-1))*(1-T/Tt)+(THETA/Tt)*(1/math.expm1(THETA/Tt)-1/math.expm1(THETA/T)))

def stagnation(T,M):
    if M==0:return T
    lo=T;hi=max(2*T,1000)
    while mach_squared(T,hi)<M*M:hi*=2
    for _ in range(100):
        mid=(lo+hi)/2
        if mach_squared(T,mid)<M*M:lo=mid
        else:hi=mid
    return (lo+hi)/2

def resistance(thickness,k):
    return 1/(H*AREA)+thickness/(k*AREA)+L2/(K2*AREA)

def payload(t,R):
    T=BOUNDARY+(INITIAL-BOUNDARY)*math.exp(-t/(R*CAPACITY))
    heat=(BOUNDARY-T)/R
    # Boundary -> convection -> layer 1 -> layer 2 -> payload
    return T,heat,CAPACITY*(T-INITIAL),BOUNDARY-heat/(H*AREA),T+heat*L2/(K2*AREA)

def save_csv(name,rows):
    with (OUT/name).open('w',newline='',encoding='utf-8') as f:
        w=csv.DictWriter(f,fieldnames=list(rows[0]));w.writeheader();w.writerows(rows)

def save_figure(fig,name):
    fig.savefig(OUT/(name+'.svg'))
    fig.savefig(OUT/(name+'.pdf'))
    fig.savefig(OUT/(name+'.png'),dpi=180)
    plt.close(fig)

def main():
    plt.rcParams.update({'font.family':'DejaVu Sans','font.size':10,'axes.spines.top':False,'axes.spines.right':False,'svg.fonttype':'none','axes.grid':True,'grid.alpha':0.2})
    flow=[]
    for T in (220.,300.):
        for i in range(101):
            M=i/20;p=T*(1+(GAMMA-1)*M*M/2);v=stagnation(T,M)
            flow.append(dict(static_temperature_K=T,mach=M,fixed_heat_capacity_K=p,vibrational_heat_capacity_K=v,fixed_minus_vibrational_K=p-v))
    save_csv('flow.csv',flow)
    fig,ax=plt.subplots(1,2,figsize=(10,4.6),layout='constrained')
    for T,col in ((220.,'#12658a'),(300.,'#aa267f')):
        r=[x for x in flow if x['static_temperature_K']==T]
        ax[0].plot([x['mach'] for x in r],[x['fixed_heat_capacity_K'] for x in r],color=col,ls='--',label=f'{T:g} K upstream: fixed heat capacity')
        ax[0].plot([x['mach'] for x in r],[x['vibrational_heat_capacity_K'] for x in r],color=col,label=f'{T:g} K upstream: vibrational model')
        ax[1].plot([x['mach'] for x in r],[x['fixed_minus_vibrational_K'] for x in r],color=col,label=f'{T:g} K upstream')
    ax[0].set(xlabel='Upstream Mach number',ylabel='Stagnation temperature (K)',title='a  Slowing the same upstream air');ax[0].legend(fontsize=8)
    ax[1].set(xlabel='Upstream Mach number',ylabel='Fixed minus vibrational prediction (K)',title='b  Model difference');ax[1].legend(fontsize=8)
    fig.suptitle('Idealized air energy balance — chemistry and wall heating excluded',fontsize=12)
    save_figure(fig,'flow-study')
    thermal=[];times=[]
    for k in (.04,.2,1.):
        for mm in range(1,51):
            R=resistance(mm/1000,k);T,q,E,surface,interface=payload(3600,R)
            thermal.append(dict(layer1_conductivity_W_per_mK=k,layer1_thickness_mm=mm,total_resistance_K_per_W=R,time_constant_s=R*CAPACITY,payload_after_1h_C=T,heat_rate_after_1h_W=q,energy_after_1h_J=E,outer_layer_surface_after_1h_C=surface,layer_interface_after_1h_C=interface))
        R=resistance(.01,k)
        for minutes in range(181):
            T,q,E,surface,interface=payload(minutes*60,R)
            times.append(dict(layer1_conductivity_W_per_mK=k,layer1_thickness_mm=10,elapsed_minutes=minutes,payload_C=T,heat_rate_W=q,accumulated_energy_J=E,outer_layer_surface_C=surface,layer_interface_C=interface))
    save_csv('barrier-sweep.csv',thermal);save_csv('payload-time.csv',times)
    fig,ax=plt.subplots(1,2,figsize=(10,4.6),layout='constrained')
    for k,col in ((.04,'#12658a'),(.2,'#aa267f'),(1.,'#866c00')):
        r=[x for x in times if x['layer1_conductivity_W_per_mK']==k]
        ax[0].plot([x['elapsed_minutes'] for x in r],[x['payload_C'] for x in r],color=col,label=f'k = {k:g} W/(m K)')
        ax[1].plot([x['elapsed_minutes'] for x in r],[x['accumulated_energy_J']/1000 for x in r],color=col,label=f'k = {k:g} W/(m K)')
    ax[0].axhline(BOUNDARY,color='#555',ls='--',label='Environment: 80 °C');ax[0].legend(fontsize=8)
    ax[0].set(xlabel='Elapsed time (min)',ylabel='Uniform payload temperature (°C)',title='a  Insulation delays temperature change')
    ax[1].set(xlabel='Elapsed time (min)',ylabel='Energy added to payload (kJ)',title='b  Energy follows the same balance');ax[1].legend(fontsize=8)
    fig.suptitle('Two massless layers + uniform inert payload — layer 1 thickness 10 mm',fontsize=12)
    save_figure(fig,'payload-study')
    fig,ax=plt.subplots(figsize=(7,4.6),layout='constrained')
    for k,col in ((.04,'#12658a'),(.2,'#aa267f'),(1.,'#866c00')):
        r=[x for x in thermal if x['layer1_conductivity_W_per_mK']==k]
        ax.plot([x['layer1_thickness_mm'] for x in r],[x['payload_after_1h_C'] for x in r],color=col,label=f'k = {k:g} W/(m K)')
    ax.set(xlabel='Layer 1 thickness (mm)',ylabel='Payload temperature after 1 hour (°C)',title='Same boundary and payload, different barriers');ax.legend(fontsize=9)
    save_figure(fig,'barrier-study')
    residual=max(abs(mach_squared(x['static_temperature_K'],x['vibrational_heat_capacity_K'])-x['mach']**2) for x in flow)
    assert residual<1e-10
    assert all(x['vibrational_heat_capacity_K']>=x['static_temperature_K'] for x in flow)
    for R in (1.,10.,100.):
        assert payload(0,R)[0]==INITIAL
        assert abs(payload(R*CAPACITY,R)[0]-(INITIAL+(BOUNDARY-INITIAL)*(1-math.exp(-1))))<1e-12
    result=dict(title='Following the heat: deterministic teaching studies',status='Calculated values, no measurements or fitted confidence intervals',counts=dict(flow=len(flow),barrier_sweep=len(thermal),payload_time=len(times)),constants=dict(gamma_fixed=GAMMA,vibrational_theta_K=THETA,area_m2=AREA,external_convection_W_per_m2K=H,layer2_thickness_m=L2,layer2_conductivity_W_per_mK=K2,payload_heat_capacity_J_per_K=CAPACITY,environment_C=BOUNDARY,initial_payload_C=INITIAL),formulas=dict(fixed='Tt=T*(1+(gamma-1)*M^2/2)',resistance='R=1/(h*A)+L1/(k1*A)+L2/(k2*A)',payload='Tp=Tb+(T0-Tb)*exp(-t/(R*C))',heat_rate='q=(Tb-Tp)/R',energy='E=C*(Tp-T0)'),checks=dict(max_mach_squared_residual=residual,zero_time_and_one_time_constant='passed'),limitations=['Gas stagnation temperature is not wall temperature or heat flux.','Vibrational model excludes dissociation, ionization, chemical reactions and nonequilibrium. Higher-temperature results are simplified-model comparisons, not validated real-air predictions.','Layers have negligible heat capacity and quasi-steady one-dimensional conduction.','Payload is uniform, inert, constant heat capacity, no phase change; parameters are illustrative.','No spatial cooling, freezing, biological preservation, human exposure, structural qualification or food safety demonstrated.'],sources=['https://www.grc.nasa.gov/WWW/BGH/stagtmp.html','https://web.mit.edu/16.unified/www/FALL/thermodynamics/notes/node118.html','https://web.mit.edu/16.unified/www/FALL/thermodynamics/notes/node124.html'])
    (OUT/'study.json').write_text(json.dumps(result,indent=2)+'\n',encoding='utf-8')
    print(json.dumps({'counts':result['counts'],'checks':result['checks']}))
if __name__=='__main__':main()
