import matplotlib.pyplot as plt

from pyathena import set_units
import astropy.constants as c
import astropy.units as u

plt.rcParams['figure.figsize'] = (10,6)
plt.rcParams['figure.dpi'] = 150

plt.rcParams['font.size'] = 18
plt.rcParams['font.weight'] = 300

label_units={'density':r'$\,[{\rm cm}^{-3}]$',
        'pressure':r'$\,[{\rm cm}^{-3}\,{\rm K}]$',
        'massflux':r'$\,[M_\odot\, {\rm pc^{-2}\, Myr^{-1}}]$',
        'velocity':r'$\,[{\rm km/s}]$',
        'magnetic_field':r'$\,[\mu{\rm G}]$',
        'specific_energy':r'$\,[({\rm km/s})^2]$',
        'surf':r'$\,[{\rm M_\odot pc^{-2}}]$',
        'sfr':r'$\,[{\rm M_\odot pc^{-2} Myr^{-1}}]$',
        'snr':r'$\,[{\rm pc^{-2} Myr^{-1}}]$',
        'scaleH':r'$\,[\rm pc]$',
       }

labels={'density':r'$n_H$',
        'number_density':r'$n_H$',
        'd':r'$n_H$',
        'pressure':r'$P/k_B$',
        'P':r'$P/k_B$',
        'Pth':r'$P/k_B$',
        'Ek3':r'$P_{\rm turb}/k_B$',
        'Pturb':r'$P_{\rm turb}/k_B$',
        'Pmag':r'$P_{\rm B}/k_B$',
        'Pmag_mean':r'$P_{\rm \overline{B}}/k_B$',
        'Pmag_turb':r'$P_{\rm \delta B}/k_B$',
        'Pimag':r'$\Pi_{\rm B}/k_B$',
        'Pimag_mean':r'$\Pi_{\rm \overline{B}}/k_B$',
        'Pimag_turb':r'$\Pi_{\rm \delta B}/k_B$',
        'massflux_wind':r'$\dot\Sigma_{\rm wind}$',
        'massflux_out':r'$\dot\Sigma_{\rm out}$',
        'massflux':r'$\dot\Sigma$',
        'metalflux':r'$\dot\Sigma_Z$',
        'M3':r'$\dot\Sigma$',
        'v3':r'$v_{\rm z}$',
        'sigma_z':r'$\sigma_{\rm z}$',
        'sigma_eff_2p':r'$\sigma_{\rm eff,2p}$',
        'v3_2p':r'$v_{\rm z,2p}$',
        'vA_2p':r'$v_{\rm A,2p}$',
        'cs_2p':r'$c_{\rm s,2p}$',
        'B2':r'$\overline{B}_{\rm y}$',
        'B1':r'$\overline{B}_{\rm x}$',
        'dB3':r'$\delta{B}_{\rm z}$',
        'dB2':r'$\delta{B}_{\rm y}$',
        'dB1':r'$\delta{B}_{\rm x}$',
        'ske':r'specific kinetic energy',
        'ste':r'specific enthalphy',
        'pressure_mid_2p':r'$P_{\rm mid,2p}/k_B$',
        'PDE':r'$P_{\rm DE,gas}(1+\chi)$',
        'PDE2':r'$P_{\rm DE}$',
        'W_2p':r'$\mathcal{W}_{\rm 2p}$',
        'W':r'$\mathcal{W}$',
        'Pmid_2p':r'$P_{\rm tot}$',
        'Pturb_mid_2p':r'$P_{\rm turb}$',
        'Pth_mid_2p':r'$P_{\rm th}$',
        'Pimag_mean_mid_2p':r'$\Pi_{\rm \overline{B}}$',
        'Pimag_turb_mid_2p':r'$\Pi_{\rm \delta{B}}$',
        'surf':r'$\Sigma$',
        'sfr':r'$\Sigma_{\rm SFR}$',
        'sfr1':r'$\Sigma_{\rm SFR,1}$',#(<1\,{\rm Myr})$',
        'sfr10':r'$\Sigma_{\rm SFR,10}$',#(<10\,{\rm Myr})$',
        'sfr40':r'$\Sigma_{\rm SFR,40}$',#(<40\,{\rm Myr})$',
        'sfr100':r'$\Sigma_{\rm SFR,100}$',#(<100\,{\rm Myr})$',
	'snr':r'$\dot\mathcal{N}_{\rm SN}$',
        'snr1':r'$\dot\mathcal{N}_{\rm SN,1}$',#(<1\,{\rm Myr})$',
        'snr10':r'$\dot\mathcal{N}_{\rm SN,10}$',#(<10\,{\rm Myr})$',
        'snr40':r'$\dot\mathcal{N}_{\rm SN,40}$',#(<40\,{\rm Myr})$',
        'snr100':r'$\dot\mathcal{N}_{\rm SN,100}$',#(<100\,{\rm Myr})$',
        'scaleH':r'$H$',
        'H_2p':r'$H_{\rm 2p}$',
        'H_c':r'$H_{\rm c}$',
        'H_u':r'$H_{\rm u}$',
        'H_w':r'$H_{\rm w}$',
        'massflux_bd_d':r'$\dot{\Sigma}({|z|=L_z/2})$',
        'massflux_bd_d_h':r'$\dot{\Sigma}_{\rm hot}(|z|=L_z/2$)',
        'massflux_bd_d_2p':r'$\dot{\Sigma}_{\rm 2p}(|z|=L_z/2$)',
        'massflux_bd_s1':r'$\dot{\Sigma}_{\rm Z}({|z|=L_z/2})$',
        'massflux_bd_s1_h':r'$\dot{\Sigma}_{\rm Z,hot}(|z|=L_z/2$)',
        'massflux_bd_s1_2p':r'$\dot{\Sigma}_{\rm Z,2p}(|z|=L_z/2$)',
        'massflux_bd_s2':r'$\dot{\Sigma}_{\rm Z,cl}({|z|=L_z/2})$',
        'massflux_bd_s2_h':r'$\dot{\Sigma}_{\rm Z,cl,hot}(|z|=L_z/2$)',
        'massflux_bd_s2_2p':r'$\dot{\Sigma}_{\rm Z,cl,2p}(|z|=L_z/2$)',
        'massflux_bd_s3':r'$\dot{\Sigma}_{\rm Z,run}({|z|=L_z/2})$',
        'massflux_bd_s3_h':r'$\dot{\Sigma}_{\rm Z,run,hot}(|z|=L_z/2$)',
        'massflux_bd_s3_2p':r'$\dot{\Sigma}_{\rm Z,run,2p}(|z|=L_z/2$)',
        'massflux_out_5_h':r'$\dot{\Sigma}_{\rm hot}(|z|=500{\rm \,pc})$',
        'massflux_out_10_h':r'$\dot{\Sigma}_{\rm hot}(|z|=1{\rm \,kpc})$',
       }

unit=set_units(muH=1.4271)
units={'massflux':(unit['density']*unit['velocity']).to('Msun*pc^(-2)*Myr^(-1)').value,
       'M3':(unit['density']*unit['velocity']).to('Msun*pc^(-2)*Myr^(-1)').value,
       'P':(unit['pressure']/c.k_B).to('K*cm^(-3)').value,
       'Pth':(unit['pressure']/c.k_B).to('K*cm^(-3)').value,
       'Ek3':(unit['pressure']/c.k_B).to('K*cm^(-3)').value,
       'Pturb':(unit['pressure']/c.k_B).to('K*cm^(-3)').value,
       'Pmag':(unit['pressure']/c.k_B).to('K*cm^(-3)').value,
       'Pmag_mean':(unit['pressure']/c.k_B).to('K*cm^(-3)').value,
       'Pmag_turb':(unit['pressure']/c.k_B).to('K*cm^(-3)').value,
       'Pimag':(unit['pressure']/c.k_B).to('K*cm^(-3)').value,
       'Pimag_mean':(unit['pressure']/c.k_B).to('K*cm^(-3)').value,
       'Pimag_turb':(unit['pressure']/c.k_B).to('K*cm^(-3)').value,
       'B1':unit['magnetic_field'].value,
       'B2':unit['magnetic_field'].value,
       'B3':unit['magnetic_field'].value,
       'Myr':unit['time'].to('Myr').value,
       'Msun':unit['mass'].to('Msun').value,
      }
