# -*- coding: utf-8 -*-
"""
Created on Sun Jul 17 14:31:57 2022

@author: frede
"""

import main_program as mp
from astropy.table import Table
import four_levels_model as four_lvl
import matplotlib.pyplot as plt
import numpy as np

'''================================ initialization of parameters ================================'''

filename = './converted/kp00_30000.txt'
star_radius = 5
t_gas = 500
n_e = 1.6
n_c = 54
parsec = 6.5e-8
fc_pah = 0.1
ISRF = False

# test3 = mp.HeatingGas(filename, star_radius, t_gas, n_e, n_c, parsec, fc_pah, ISRF)
# test3.parameters()
# test4 = mp.HeatingGas(filename, star_radius, t_gas, n_e, n_c, parsec, fc_pah, ISRF)
# test4.parameters()

''' parameter '''
distance_list = np.array([parsec])
for j in range(2,1002,1):
    d = distance_list[0]
    distance_list = np.append(distance_list, j * d) #in parsec

test = four_lvl.HeatingGas(filename, star_radius, t_gas, n_e, n_c, parsec, fc_pah, ISRF)
test.parameters()

# calculus=0
# for i in distance_list:
#     test = four_lvl.HeatingGas(filename, star_radius, t_gas, n_e, n_c, i, fc_pah, ISRF)
#     test.parameters()
#     calculus += test.total_gas_heating*1e-7/(test.g_0 * test.heating_efficiency * 0.1)

# print('value = ',calculus/len(distance_list))    

'''================================ plots ================================'''  

# ''' intensity per wavelength '''
# plt.figure()
# plt.xlim([0,1500])
# plt.plot(test.wavelength,test.wavelength_intensity*1e-3,'k')
# plt.ylabel('intensity spectrum (W m$^{-2}$ sr$^{-1}$ nm$^{-1}$)')
# plt.xlabel('wavelength (nm)')

''' average of 3 sizes categories of pah of photoabsorption sigma 
and photoionization sigma_ion cross sections per C atom of pahs
as functions of photon energy in eV'''

''' conversion from mb/C atom to cm²/C atom '''
test.pah_cross_a = test.pah_cross_a*1e-18
test.pah_cross_n = test.pah_cross_n*1e-18
test.pah_cross_c = test.pah_cross_c*1e-18
test.pah_cross_dc = test.pah_cross_dc*1e-18

''' small n_c '''
test_small_n_c_1 = four_lvl.HeatingGas(filename, star_radius, t_gas, n_e, 32, parsec, fc_pah, ISRF)
test_small_n_c_2 = four_lvl.HeatingGas(filename, star_radius, t_gas, n_e, 36, parsec, fc_pah, ISRF)
test_small_n_c_3 = four_lvl.HeatingGas(filename, star_radius, t_gas, n_e, 38, parsec, fc_pah, ISRF)

test_small_n_c_1.parameters()
test_small_n_c_2.parameters()
test_small_n_c_3.parameters()

''' medium n_c '''
test_medium_n_c_1 = four_lvl.HeatingGas(filename, star_radius, t_gas, n_e, 40, parsec, fc_pah, ISRF)
test_medium_n_c_2 = four_lvl.HeatingGas(filename, star_radius, t_gas, n_e, 42, parsec, fc_pah, ISRF)

test_medium_n_c_1.parameters()
test_medium_n_c_2.parameters()

''' large n_c '''
test_large_n_c_1 = four_lvl.HeatingGas(filename, star_radius, t_gas, n_e, 48, parsec, fc_pah, ISRF)
test_large_n_c_2 = four_lvl.HeatingGas(filename, star_radius, t_gas, n_e, 54, parsec, fc_pah, ISRF)
test_large_n_c_3 = four_lvl.HeatingGas(filename, star_radius, t_gas, n_e, 66, parsec, fc_pah, ISRF)

test_large_n_c_1.parameters()
test_large_n_c_2.parameters()
test_large_n_c_3.parameters()

''' average of first and second photoionization yield '''
yield_of_first_photoionization = (test_small_n_c_1.yield_of_first_photoionization +\
                                  test_small_n_c_2.yield_of_first_photoionization +\
                                  test_small_n_c_3.yield_of_first_photoionization +\
                                  test_medium_n_c_1.yield_of_first_photoionization +\
                                2*test_medium_n_c_2.yield_of_first_photoionization +\
                                  test_large_n_c_1.yield_of_first_photoionization +\
                                  test_large_n_c_2.yield_of_first_photoionization +\
                                  test_large_n_c_3.yield_of_first_photoionization )/9
yield_of_second_photoionization = (test_small_n_c_1.yield_of_second_photoionization +\
                                    test_small_n_c_2.yield_of_second_photoionization +\
                                    test_small_n_c_3.yield_of_second_photoionization +\
                                    test_medium_n_c_1.yield_of_second_photoionization +\
                                  2*test_medium_n_c_2.yield_of_second_photoionization +\
                                    test_large_n_c_1.yield_of_second_photoionization +\
                                    test_large_n_c_2.yield_of_second_photoionization +\
                                    test_large_n_c_3.yield_of_second_photoionization )/9

''' average of ionization cross-section for the cases of neutral molecules, and 
cations '''
ionization_cross_a = test.pah_cross_a * test.detachment_yield
ionization_cross_n = test.pah_cross_n * yield_of_first_photoionization
ionization_cross_c = test.pah_cross_c * yield_of_second_photoionization
#ionization_cross_dc is too small to be traced

''' average of energy per cross-section for the cases of neutral molecules, 
cations and dications '''
energy = (test.energy_negative_charged + test.energy_neutral + test.energy_charged + test.energy_double_charged)/4

''' plot sigma and sigma_ion in cm²/(C atom) as a function of photon energy in eV
for the cases of neutral, charged and double charged molecules '''
plt.figure(figsize=(3.57,9.0))

plt.subplot(4,1,1)
plt.plot(test.energy_negative_charged,test.pah_cross_a,'k', label='$\sigma(E,Z)$')
plt.plot(energy,ionization_cross_a,'k',ls='--', label='$\sigma_{ion}(E,Z)$')
plt.title('Z=-1', fontsize = 10)
plt.legend(fontsize = 8)

plt.subplot(4,1,2)
plt.plot(test.energy_neutral,test.pah_cross_n,'k')
plt.plot(energy,ionization_cross_n,'k',ls='--')
plt.title('Z=0', fontsize = 10)

plt.subplot(4,1,3)
plt.plot(test.energy_neutral,test.pah_cross_c,'k')
plt.plot(energy,ionization_cross_c,'k',ls='--') 
plt.title('Z=1', fontsize = 10)
plt.ylabel('Cross sections (cm$^2$/C)')

plt.subplot(4,1,4)
plt.plot(test.energy_neutral,test.pah_cross_dc,'k') 
plt.title('Z=2', fontsize = 10)
plt.xlabel('Photon energy (eV)')

plt.subplots_adjust(hspace = 0.35)
plt.savefig('cross_ioni_section_categories.pdf', bbox_inches='tight')
if ISRF == False:

    ''' evolution of the different populations of neutral pah molecules, 
    # cations and dications as a function of the distance to the star considered '''
    
    ''' parameter '''
    # population_frac3 = np.zeros([1, len(distance_list)])
    # population_frac4 = np.zeros([4, len(distance_list)])
    
    ''' indentation '''
    # j=0
    # gamma_list3 = []
    # gamma_list4 = []
    # for i in distance_list:
    #     population_plot_test3 = mp.HeatingGas(filename, star_radius, t_gas, n_e, n_c, i, fc_pah, ISRF)
    #     population_plot_test3.parameters()
    #     population_frac3[0,j] =(population_plot_test3.frac_charged +\
    #                             population_plot_test3.frac_double_charged)/\
    #                           (population_plot_test3.frac_neutral +\
    #                             population_plot_test3.frac_charged +\
    #                             population_plot_test3.frac_double_charged)
    #     # population_frac3[0,j] = population_plot_test3.frac_neutral
    #     # population_frac3[1,j] = population_plot_test3.frac_charged
    #     # population_frac3[2,j] = population_plot_test3.frac_double_charged
    #     gamma_list3.append(population_plot_test3.gamma)         
                         
        # population_plot_test4 = four_lvl.HeatingGas(filename, star_radius, t_gas, n_e, n_c, i, fc_pah, ISRF)
        # population_plot_test4.parameters()
        # population_frac4[0,para] =(population_plot_test4.frac_anion +\
        #                         population_plot_test4.frac_charged +\
        #                         population_plot_test4.frac_double_charged)/\
        #                       (population_plot_test4.frac_anion +\
        #                         population_plot_test4.frac_neutral +\
        #                         population_plot_test4.frac_charged +\
        #                         population_plot_test4.frac_double_charged)
        # population_frac4[0,j] = population_plot_test4.frac_anion
        # population_frac4[1,j] = population_plot_test4.frac_neutral
        # population_frac4[2,j] = population_plot_test4.frac_charged
        # population_frac4[3,j] = population_plot_test4.frac_double_charged
        # gamma_list4.append(population_plot_test4.gamma)
        # j=j+1
    
    ''' population fractions as function of the distance ''' 
    # plt.figure(figsize=(5.90,3.93))
    # plt.plot(gamma_list3,population_frac3[0,:],c='k',label='3 levels model')
    # plt.plot(gamma_list4,population_frac4[0,:],ls='--', c='k',label='4 levels model')
    # plt.xlim([1e1, 1e7])
    # plt.plot(gamma_list3, population_frac3[0,:], color = 'green', label='Z=0, 3 levels model')
    # plt.plot(gamma_list3, population_frac3[1,:], color = 'blue', label='Z=1, 3 levels model')
    # plt.plot(gamma_list3, population_frac3[2,:], color = 'red', label='Z=2, 3 levels model')
    # plt.plot(gamma_list4, population_frac4[0,:], ls='--', color = 'k', label='Z=-1, 4 levels model')
    # plt.plot(gamma_list4, population_frac4[1,:], ls='--', color = 'green', label='Z=0, 4 levels model')
    # plt.plot(gamma_list4, population_frac4[2,:], ls='--', color = 'k', label='Z=1, 4 levels model')
    # plt.plot(gamma_list4, population_frac4[3,:], ls='--', color = 'red', label='Z=2, 4 levels model')
    # plt.plot(gamma_list3, population_frac3[0,:]+population_frac3[1,:]+population_frac3[2,:],'blue')
    # plt.plot(gamma_list4, population_frac4[0,:]+population_frac4[1,:]+population_frac4[2,:]+population_frac4[3,:],'darkblue')
    # plt.xlabel('$\gamma$ ($G_0\sqrt{T}/n_e$)')
    # plt.ylabel('population fraction R$_i$')
    # plt.legend()
    
    ''' photoelectric efficiencies with the pah heating model ''' 
    
    # gamma_list_T_one3 = []
    # heating_efficiency_list_T_one3 = []
    
    # gamma_list_T_two3 = []
    # heating_efficiency_list_T_two3 = []
    
    # gamma_list_T_one4 = []
    # heating_efficiency_list_T_one4 = []
      
    # gamma_list_T_two4 = []
    # heating_efficiency_list_T_two4 = []
    
    # ratio1_1 = []
    # ratio2_1 = []
    # gamma_mean1_1 = []
    # gamma_mean2_1 = []
    # gamma_list_1 = []
    # j = 0
    # for i in distance_list:
    #     j = j+1
    #     if j%100 == 0:
    #         print('Progression à {}'.format(j/10),'%')
    #         print('')
        # test_for_efficiencies_T_one3 = mp.HeatingGas(filename, star_radius, 100, n_e, 54, i, fc_pah, ISRF)
        # test_for_efficiencies_T_one3.parameters()
        # gamma_list_T_one3.append(test_for_efficiencies_T_one3.gamma)
        # heating_efficiency_list_T_one3.append(test_for_efficiencies_T_one3.heating_efficiency)
    
        # test_for_efficiencies_T_two3 = mp.HeatingGas(filename, star_radius, 1000, n_e, 54, i, fc_pah, ISRF)
        # test_for_efficiencies_T_two3.parameters()
        # gamma_list_T_two3.append(test_for_efficiencies_T_two3.gamma)
        # heating_efficiency_list_T_two3.append(test_for_efficiencies_T_two3.heating_efficiency)
    
        # test_for_efficiencies_T_one4 = four_lvl.HeatingGas(filename, star_radius, 100, n_e, n_c, i, fc_pah, ISRF)
        # test_for_efficiencies_T_one4.parameters()
        # gamma_list_T_one4.append(test_for_efficiencies_T_one4.gamma)
        # heating_efficiency_list_T_one4.append(test_for_efficiencies_T_one4.heating_efficiency)   
        
        # test_for_efficiencies_T_two4 = four_lvl.HeatingGas(filename, star_radius, 1000, n_e, n_c, i, fc_pah, ISRF)
        # test_for_efficiencies_T_two4.parameters()
        # gamma_list_T_two4.append(test_for_efficiencies_T_two4.gamma)
        # heating_efficiency_list_T_two4.append(test_for_efficiencies_T_two4.heating_efficiency)
 
    #         ratio1_1.append( (test_for_efficiencies_T_one3.heating_efficiency/test_for_efficiencies_T_one4.heating_efficiency)*100)
    #         ratio2_1.append( (test_for_efficiencies_T_two3.heating_efficiency/test_for_efficiencies_T_two4.heating_efficiency)*100)
    #         gamma_mean1_1.append( (test_for_efficiencies_T_one3.gamma + test_for_efficiencies_T_one4.gamma)/2 )
    #         gamma_mean2_1.append( (test_for_efficiencies_T_two3.gamma + test_for_efficiencies_T_two4.gamma)/2 )
    
    # plt.figure()
    # plt.xlim([1e1, 1e6])
    # plt.loglog(gamma_mean1_1, ratio1_1, 'blue', label='ratio of 3 and 4 levels model $\epsilon_{PAH}$ (T = 100K) ')
    # plt.loglog(gamma_mean2_1, ratio2_1, 'red', ls='--' , label='ratio of 3 and 4 levels model $\epsilon_{PAH}$ (T = 1000K)')
    # plt.loglog(gamma_list_T_one3, heating_efficiency_list_T_one3, 'blue', label='$\epsilon_{PAH}$ (T = 100K) 3 levels model')
    # plt.loglog(gamma_list_T_two3, heating_efficiency_list_T_two3, 'red', ls='--' , label='$\epsilon_{PAH}$ (T = 1000K) 3 levels model')
    # plt.loglog(gamma_list_T_one4, heating_efficiency_list_T_one4, 'darkblue', label='$\epsilon_{PAH}$ (T = 100K) 4 levels model')
    # plt.loglog(gamma_list_T_two4, heating_efficiency_list_T_two4, 'darkred', ls='--' , label='$\epsilon_{PAH}$ (T = 1000K) 4 levels model')
    # plt.xlabel('$\gamma(G_0\sqrt{T}/n_e)$')
    # plt.ylabel('$\epsilon_{PAH}$')
    # plt.ylabel('$\epsilon_{PAH, 3lvlmod}$ / $\epsilon_{PAH, 4lvlmod}$ (%)')
    # plt.legend()
    
    # ''' Total photoelectric heating rates of the gas '''
    
    # gamma_list_fc_one3 = []
    # total_gas_heating_per_rf_list_fc_one3 = []
    
    # gamma_list_fc_two3 = []
    # total_gas_heating_per_rf_list_fc_two3 = []
    
    # gamma_list_fc_one4 = []
    # total_gas_heating_per_rf_list_fc_one4 = []
    
    # gamma_list_fc_two4 = []
    # total_gas_heating_per_rf_list_fc_two4 = []
    
    # ratio1_2 = []
    # ratio2_2 = []
    # gamma_mean1_2 = []
    # gamma_mean2_2 = []
    # Table_list_g_fc5 = []
    # Table_list_Gpg0_fc5 = []
    
    # Table_list_g_fc10 = []
    # Table_list_Gpg0_fc10 = []
    
    # for i in distance_list:
    #     test_for_total_heating_rates_fc_one3 = mp.HeatingGas(filename, star_radius, t_gas, n_e, n_c, i, 0.05, ISRF)
    #     test_for_total_heating_rates_fc_one3.parameters()
    #     gamma_list_fc_one3.append(test_for_total_heating_rates_fc_one3.gamma)
    #     total_gas_heating_per_rf_list_fc_one3.append( (test_for_total_heating_rates_fc_one3.total_gas_heating*1e-7) / test_for_total_heating_rates_fc_one3.g_0)
        
    #     test_for_total_heating_rates_fc_two3 = mp.HeatingGas(filename, star_radius, t_gas, n_e, n_c, i, 0.1, ISRF)
    #     test_for_total_heating_rates_fc_two3.parameters()
    #     gamma_list_fc_two3.append(test_for_total_heating_rates_fc_two3.gamma)
    #     total_gas_heating_per_rf_list_fc_two3.append( (test_for_total_heating_rates_fc_two3.total_gas_heating*1e-7) / test_for_total_heating_rates_fc_two3.g_0)
        
        # test_for_total_heating_rates_fc_one4 = four_lvl.HeatingGas(filename, star_radius, 500, n_e, 54, i, 0.05, ISRF)
        # test_for_total_heating_rates_fc_one4.parameters()
        # gamma_list_fc_one4.append(test_for_total_heating_rates_fc_one4.gamma)
        # total_gas_heating_per_rf_list_fc_one4.append( (test_for_total_heating_rates_fc_one4.total_gas_heating*1e-7) / test_for_total_heating_rates_fc_one4.g_0)
        
        # test_for_total_heating_rates_fc_two4 = four_lvl.HeatingGas(filename, star_radius, 500, n_e, 54, i, 0.1, ISRF)
        # test_for_total_heating_rates_fc_two4.parameters()
        # gamma_list_fc_two4.append(test_for_total_heating_rates_fc_two4.gamma)
        # total_gas_heating_per_rf_list_fc_two4.append( (test_for_total_heating_rates_fc_two4.total_gas_heating*1e-7) / test_for_total_heating_rates_fc_two4.g_0)

    #     ratio1_2.append( ( ( (test_for_total_heating_rates_fc_one3.total_gas_heating*1e-7) /\
    #                           test_for_total_heating_rates_fc_one3.g_0) /\
    #                        ( (test_for_total_heating_rates_fc_one4.total_gas_heating*1e-7) /\
    #                           test_for_total_heating_rates_fc_one4.g_0) )*100 )
    #     ratio2_2.append( ( ( (test_for_total_heating_rates_fc_two3.total_gas_heating*1e-7) /\
    #                           test_for_total_heating_rates_fc_two3.g_0) /\
    #                        ( (test_for_total_heating_rates_fc_two4.total_gas_heating*1e-7) /\
    #                           test_for_total_heating_rates_fc_two4.g_0) )*100 )
    #     gamma_mean1_2.append( (test_for_total_heating_rates_fc_one3.gamma+test_for_total_heating_rates_fc_one4.gamma)/2 )
    #     gamma_mean2_2.append( (test_for_total_heating_rates_fc_two3.gamma+test_for_total_heating_rates_fc_two4.gamma)/2 )
        
    # plt.figure()
    # plt.xlim([1e1, 1e7])
    # plt.loglog(gamma_mean1_2,ratio1_2, 'red', label='ratio of 3 and 4 levels model $\Gamma_{tot}$ (5% of C in PAHs)')
    # plt.loglog(gamma_mean2_2,ratio2_2, 'blue', label='ratio of 3 and 4 levels model $\Gamma_{tot}$ (10% of C in PAHs)')
    # plt.loglog(gamma_list_fc_one3, total_gas_heating_per_rf_list_fc_one3, c='red', label='This work (5% of C in PAHs) 3 levels model')
    # plt.loglog(gamma_list_fc_two3, total_gas_heating_per_rf_list_fc_two3,ls='--', c='blue', label='This work (10% of C in PAHs) 3 levels model')
    # plt.loglog(gamma_list_fc_one4, total_gas_heating_per_rf_list_fc_one4, c='darkred', label='This work (5% of C in PAHs) 4 levels model')
    # plt.loglog(gamma_list_fc_two4, total_gas_heating_per_rf_list_fc_two4,ls='--', c='darkblue', label='This work (10% of C in PAHs) 4 levels model')
    # plt.xlabel('$\gamma(G_0\sqrt{T}/n_e)$')
    # # plt.ylabel('$\Gamma_{tot}$ (W H$^{-1}$ G$_0^{-1}$)')
    # plt.ylabel('$\Gamma_{tot, 3lvlmod}$ / $\Gamma_{tot, 4lvlmod}$ (%)')
    # plt.legend()
    
    # Table_values.append(heating_efficiency_list_T_one4)
    # Table_list_g_fc5 = gamma_list_fc_one4[::-1]
    # Table_list_Gpg0_fc5 = total_gas_heating_per_rf_list_fc_one4[::-1]
    
    # Table_list_g_fc10 = gamma_list_fc_two4[::-1]
    # Table_list_Gpg0_fc10 = total_gas_heating_per_rf_list_fc_two4[::-1]