# Purpose: Provide cross validation between legacy capabilities and current streamlined models
# Authors: Thomas A. Scott

import scipy
import numpy as np
from pathlib import Path
from importlib.resources import files as imp_files

import matplotlib.pyplot as plt

import hrap.core as core
import hrap.fluid as fluid
from hrap.tank    import *
from hrap.grain   import *
from hrap.chamber import *
from hrap.nozzle  import *
from hrap.units   import _in, _ft, _lbf, _atm


jax.config.update("jax_enable_x64", True)
# np.random.seed(42)
hrap_root = Path(imp_files('hrap'))

# Experimentally curve-fit saturated nitrous properties, applicable range is -90 to 36 Celsius
# Ported from HRAP, originally from ?
def get_sat_nos_props(T):
    #--------------------------------------------------------------------------
    # Critical Properties
    #--------------------------------------------------------------------------
    Pc   = 7251000
    Tc   = 309.57
    rhoc = 452
    R    = 188.91
    
    #--------------------------------------------------------------------------
    # Vapor Pressure
    #--------------------------------------------------------------------------
    # Applicable Range: -90 to 36 Celsius
    
    a1 = -6.71893
    a2 =  1.35966
    a3 = -1.3779
    a4 = -4.051
    Pv = Pc*jnp.exp((1/(T/Tc))*(a1*(1-T/Tc) + \
        a2*(1-(T/Tc))**(3/2) + a3*(1-(T/Tc))**(5/2) + a4*(1-(T/Tc))**5))
    
    #--------------------------------------------------------------------------
    # Saturated Nitrous Oxide Density
    #--------------------------------------------------------------------------
    # Applicable Range: -90 to 36 Celsius
    
    # Liquid Density
    b1 =  1.72328
    b2 = -0.83950
    b3 =  0.51060
    b4 = -0.10412
    rho_l = rhoc*jnp.exp(b1*(1-(T/Tc))**(1/3) + \
        b2*(1-(T/Tc))**(2/3) + b3*(1-(T/Tc)) + b4*(1-(T/Tc))**(4/3))
    
    # Vapor Density
    c1 = -1.00900
    c2 = -6.28792
    c3 =  7.50332
    c4 = -7.90463
    c5 =  0.629427
    rho_v = rhoc*jnp.exp(c1*((Tc/T)-1)**(1/3) + \
        c2*((Tc/T)-1)**(2/3) + c3*((Tc/T)-1) + c4*((Tc/T)-1)**(4/3) + \
        c5*((Tc/T)-1)**(5/3))
    
    #--------------------------------------------------------------------------
    # Specific Enthalpy
    #--------------------------------------------------------------------------
    # Applicable Range: -90 to 35 Celsius
    
    # Liquid Enthalpy
    d1 = -200
    d2 =  116.043
    d3 = -917.225
    d4 =  794.779
    d5 = -589.587
    
    # Vapor Enthalpy
    e1 = -200
    e2 =  440.055
    e3 = -459.701
    e4 =  434.081
    e5 = -485.338
    
    # Latent Heat of Vaporization
    Hv = (e1-d1) + (e2-d2)*(1-(T/Tc))**(1/3) + (e3-d3)* \
        (1-(T/Tc))**(2/3) + (e4-d4)*(1-(T/Tc)) + (e5-d5)*(1-(T/Tc))**(4/3)
    
    #--------------------------------------------------------------------------
    # Specific Heat Capacity of Saturated Liquid
    #--------------------------------------------------------------------------
    
    f1 =  2.49973
    f2 =  0.023454
    f3 = -3.80136
    f4 =  13.0945
    f5 = -14.5180
    
    # Specific_Heat_Capacity = 
    Cp = f1*(1 + f2*(1-(T/Tc))**(-1) + f3* \
        (1-(T/Tc)) + f4*(1-(T/Tc))**2 + f5*(1-(T/Tc))**3)
    
    #--------------------------------------------------------------------------
    # Saturated Vapor Compressibility Factor (Tabulated from empirical data)
    #--------------------------------------------------------------------------
    
    Z = (Pc*jnp.exp((1/(T/Tc))*(a1*(1-T/Tc) + \
        a2*(1-(T/Tc))**(3/2) + a3*(1-(T/Tc))**(5/2) + a4*(1-(T/Tc))**5)))/ \
        ((rhoc*jnp.exp(c1*((Tc/T)-1)**(1/3) + c2*((Tc/T)-1)**(2/3) + c3*((Tc/T)-1) \
        + c4*((Tc/T)-1)**(4/3) + c5*((Tc/T)-1)**(5/3)))*R*T)
    
    return { 'Pv': Pv, 'rho_l': rho_l, 'rho_v': rho_v, 'Hv': Hv, 'Cp': Cp, 'Z': Z }



print('Loading chemistry table...')
chem = scipy.io.loadmat(hrap_root/'resources'/'propellant_configs'/'ABS.mat')
chem = chem['s'][0][0]
chem_OF = chem['prop_OF'].ravel()
chem_Pc = chem['prop_Pc'].ravel()
chem_k, chem_M, chem_T = chem['prop_k'], chem['prop_M'], chem['prop_T']
if chem_k.size == 1: chem_k = np.full_like(chem_T, chem_k.item())



# Initialization
emp_tnk = make_sat_tank(
    get_sat_nos_props,
    V = (np.pi/4 * 4.75**2 * _in**2) * (7.0 * _ft),
    inj_CdA= 0.22 * (np.pi/4 * 0.5**2 * _in**2),
    m_ox=2,
    T = 294,
)
cp_tnk = make_sat_tank(
    fluid.bake_sat_coolprop('NitrousOxide', np.linspace(183.0, 309.0, 20)),
    V = (np.pi/4 * 4.75**2 * _in**2) * (7.0 * _ft),
    inj_CdA= 0.22 * (np.pi/4 * 0.5**2 * _in**2),
    m_ox=2,
    T = 294,
)

shape = make_circle_shape(
    ID = 2.5 * _in,
)
constOF_grn = make_constOF_grain(
    shape,
    OF = 3.5,
    OD = 4.5 * _in,
    L = 30.0 * _in,
    rho = 1117.0,
)
shiftOF_grn = make_shiftOF_grain(
    shape,
    Reg = jnp.array([0.198, 0.325, 0.0]), # ABS regression coefficient (mm/s), regression exponent, length exponent
    OD = 4.5 * _in,
    L = 30.0 * _in,
    rho = 1117.0,
)

prepost_ID = 4.25*_in                              # Inner diameter of pre and post combustion chambers (m)
prepost_V  = (3.5+1.7)*_in * np.pi/4*prepost_ID**2 # Empty volume of pre and post combustion chambers (m^3)
rings_V    = 3 * (1/8*_in) * np.pi*(2.5/2 * _in)**2  # Empty volume of phenolic rings (m^3)
fuel_V     = (30.0 * _in) * np.pi*(4.5/2 * _in)**2   # Empty volume of grain footprint (m^3)
cmbr = make_chamber(
    # V0 =  prepost_V + rings_V + fuel_V,            # Volume of chamber w/o grain (m^3)
    V0 = 0.0, # Sim can be a bit unstable with this and incompressible injetor
    cstar_eff = 1.0,#0.95
)

noz = make_cd_nozzle(
    thrt = 1.75 * _in, # Throat diameter
    ER = 4.99,         # Exit/throat area ratio
    eff = 0.97,
    C_d = 0.995,
)

from jax.scipy.interpolate import RegularGridInterpolator
chem_interp_k = RegularGridInterpolator((chem_OF, chem_Pc), chem_k, fill_value=1.4)
chem_interp_M = RegularGridInterpolator((chem_OF, chem_Pc), chem_M, fill_value=29.0)
chem_interp_T = RegularGridInterpolator((chem_OF, chem_Pc), chem_T, fill_value=293.0)



# Ensure results folder exists
results_path = Path('./results')
results_path.mkdir(parents=True, exist_ok=True)

file_prefix = 'nitrous_compare'

# Visualization
fig, axs = plt.subplots(nrows=2, ncols=1, figsize=(12,7))
axs = np.array(axs).ravel()

for grn_name, linestyle, grn in [('constOF', '-', constOF_grn), ('shiftOF', '--', shiftOF_grn)]:
    for tnk_name, color, tnk in [('empirical', 'black', emp_tnk), ('coolprop', 'blue', cp_tnk)]:
        print('Starting case', grn_name, tnk_name)
        s, x, method = core.make_engine(
            tnk, grn, cmbr, noz,
            chem_interp_k=chem_interp_k, chem_interp_M=chem_interp_M, chem_interp_T=chem_interp_T,
            Pa=101e3,
        )

        fire_engine = core.make_integrator(
            core.step_fe,
            method,
        )

        # Integrate the engine state
        T = 12.0
        # T = 10E-2
        print('Running...')
        import time
        t1 = time.time()
        t, _x, xstack = fire_engine(s, x, dt=1E-3, T=T)
        jax.block_until_ready(xstack)
        t2  = time.time()
        t, x, xstack = fire_engine(s, x, dt=1E-3, T=T)
        jax.block_until_ready(xstack)
        t3 = time.time()
        print('done, first run was {a:.2f}s, second run was {b:.2f}s'.format(a=t2-t1, b=t3-t2))

        # Unpack the dynamic engine state
        N_t = xstack.shape[0]
        _tnk, _grn, _cmbr, _noz = core.unpack_engine(s, xstack, method)
        
        # Plot thrust
        axs[0].plot(np.linspace(0.0, T, N_t), _noz['thrust'], color=color, linestyle=linestyle, label=grn_name+', '+tnk_name)
        axs[0].set_title('Thrust')
        axs[0].legend()
        
        axs[1].plot(np.linspace(0.0, T, N_t), _tnk['P'], color=color, linestyle=linestyle)
        axs[1].set_title('Tank Pressure')
# print('tnk', tnk.keys())
# print('grn', grn.keys())
# print('cmbr', cmbr.keys())
# print('noz', noz.keys())
# print()
# print('Post-run arrays:')
# for name, obj in (('tnk', tnk), ('grn', grn), ('cmbr', cmbr), ('noz', noz)):
#     print(name+':')
#     for key, val in obj.items():
#         print(key+':', val)
#     print()

# # Plot oxidizer flow rate
# axs[1].plot(np.linspace(0.0, T, N_t), tnk['mdot_ox'], label='mdot_ox')
# axs[1].plot(np.linspace(0.0, T, N_t), tnk['mdot_inj'], label='mdot_inj')
# axs[1].plot(np.linspace(0.0, T, N_t), tnk['mdot_vnt'], label='mdot_vnt')
# axs[1].plot(np.linspace(0.0, T, N_t), grn['mdot'], label='mdot_grn')
# axs[1].plot(np.linspace(0.0, T, N_t), noz['mdot'], label='mdot_noz')
# axs[1].plot(np.linspace(0.0, T, N_t), cmbr['mdot_g'], label='mdot_cmbr')
# axs[1].legend(loc='upper right')
# axs[1].set_title('mdot')

# axs[2].plot(np.linspace(0.0, T, N_t), cmbr['P'], label='chamber')
# axs[2].plot(np.linspace(0.0, T, N_t), tnk['P'], label='tank')
# axs[2].plot(np.linspace(0.0, T, N_t), noz['Pe'], label='noz exit')
# axs[2].legend(loc='upper right')
# axs[2].set_title('P')

# axs[3].plot(np.linspace(0.0, T, N_t), tnk['T'])
# axs[3].set_title('T tank')

# axs[4].plot(np.linspace(0.0, T, N_t), tnk['m_ox_liq'], label='ox liq')
# axs[4].plot(np.linspace(0.0, T, N_t), tnk['m_ox_vap'], label='ox vap')
# axs[4].plot(np.linspace(0.0, T, N_t), cmbr['m_g'], label='cmbr stored')
# axs[4].plot(np.linspace(0.0, T, N_t), grn['V']*grn['rho'], label='grain')
# axs[4].legend()
# axs[4].set_title('m')

# D = (4.5 - 2.5)*_in # TODO: get
# axs[5].plot([0.0, T], [D]*2, label='grain thickness')
# axs[5].plot(np.linspace(0.0, T, N_t), grn['d'], label='net regression')
# # axs[5].plot(np.linspace(0.0, T, N_t), grn['V'], label='grain volume')
# axs[5].legend()

# axs[6].plot(np.linspace(0.0, T, N_t), noz['Me'], label='Mach exit')
# axs[6].legend()

# # axs[7].plot(np.linspace(0.0, T, N_t), cmbr['V0'] - 0*grn['V'], label='cmbr V0')
# # axs[7].plot(np.linspace(0.0, T, N_t), grn['V'], label='grain V')
# axs[7].plot(np.linspace(0.0, T, N_t), cmbr['cstar'], label='cstar')
# axs[7].plot(np.linspace(0.0, T, N_t), cmbr['T'], label='cmbr T')
# # axs[7].plot(np.linspace(0.0, T, N_t), cmbr['V0'] - grn['V'], label='Empty cmbr V')
# # axs[5].plot(np.linspace(0.0, T, N_t), grn['V'], label='grain volume')
# axs[7].legend()

# # axs[8].plot(np.linspace(0.0, T, N_t), cmbr['k'], label='cmbr k')
# axs[8].plot(np.linspace(0.0, T, N_t), cmbr['V0'] - grn['V'], label='Empty cmbr V')
# # axs[8].plot(np.linspace(0.0, T, N_t), grn['V'], label='grain volume')
# # axs[8].plot(np.linspace(0.0, T, N_t), cmbr['Pdot'], label='Pc dot')
# # axs[8].plot(np.linspace(0.0, T, N_t), grn['Vdot'], label='grain V dot')
# # axs[8].plot(np.linspace(0.0, T, N_t), grn['Vdot']/(cmbr['V0'] - grn['V']), label='Pc dot, V comb')
# # Pc*(mdot_g/m_g - dV/V)
# axs[8].legend()



# Save plot
fig.tight_layout()
fig.savefig(results_path/(file_prefix+'_plots.pdf'), format='pdf', bbox_inches='tight', pad_inches=0.1)

plt.show()
