#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Tue Jun 30 09:54:20 2020
Library of the photoionization model

@author: sacha.foschino@hotmail.fr
"""
import matplotlib.pyplot as plt
import numpy as np
from astropy.table import Table

def preparation_Fe(d,Av,wave_range,quiet):
    '''Function adapting the radiation field (RF) to its use by the PE model
    d: float, distance between the region modeled and the UV source, in cm, e.g. 4.413e+17cm for the 0.143pc separating HD200775 and the PDR front,
    Av: float, Extinction in the V band, in mag,
    wave_range: list of 2 elements, limits of the wavelength interval considered in nm, e.g. [91.16,2000],
    quiet: bool, if False, plot the RF at each step of the conversion
    
    returns Fe_E_,E,Fe_l,wave ; i.e. respectively the RF in erg/sec/cm2/eV/sr, the photon energy vector, 
            the RF in erg/sec/cm2/nm/sr and the corresponding wavelength vector
    '''
    # =============================================================================
    # preparation of HD200775 radiation field: dillution + extinction
    '''radiation field on HD200775'''
#    wave_rad,spec_rad=np.loadtxt(rf_name,unpack=True,skiprows=4) #if another RF than that of HD200775 is used (e.g. Draine, Habing or Mathis)
    wave_rad, spec_rad=np.loadtxt('./radiation_fields/HD200775_RF.txt', unpack=True)
     
    if not(quiet):
        plt.figure(1)
        plt.loglog(wave_rad,spec_rad,label='d={:.0f}mpc, {:.0f}" '.format(d/3.086e15,d/(3.086e18*3.4268e-3)))
        plt.title('virgin spectrum')
        plt.xlabel('wavelength (nm)')
        plt.ylabel('Specific intensity (erg/sec/cm$^2$/nm/sr)')
        plt.subplots_adjust(top=0.986,bottom=0.111,left=0.126,right=0.991,hspace=0.2,wspace=0.2)
        plt.legend()
        plt.tight_layout()

    
    '''wavelengths ranges for the G0 energies''' 
    #selection of the wavelength range
    spec7023=spec_rad[np.where((wave_rad>wave_range[0]) & (wave_rad<wave_range[1]))[0]]
    wave=wave_rad[np.where((wave_rad>wave_range[0]) & (wave_rad<wave_range[1]))[0]]

    #intensity of HD200775
    if not(quiet):
        plt.figure(2)
        plt.loglog(wave,spec7023,label='d={:.0f}mpc, {:.0f}" '.format(d/3.086e15,d/(3.086e18*3.4268e-3)))
        plt.title('virgin spectrum cut')
        plt.xlabel('wavelength (nm)')
        plt.ylabel('Specific intensity (erg/sec/cm$^2$/nm/sr)')
        plt.subplots_adjust(top=0.986,bottom=0.111,left=0.126,right=0.991,hspace=0.2,wspace=0.2)
        plt.legend()
        plt.tight_layout()


    '''interpolation of the extinction coefficients with HD200775's RF wavelength vector'''
    wave_draine__, albe, cos, cext__ = np.loadtxt('./extinction_curves/draine_curve_55.txt', unpack=True, usecols=[0,1,2,3], skiprows=80) 
    Cext=np.interp(wave,wave_draine__[::-1]*1e3,cext__[::-1]) #/!\ reverse the order of the cext__ vector
    
    '''geometrical dillution'''
    R=10*70000000000 #radius of the star in cm, Pilleri et al. 2012
 #   d=0.143*3.086e18
    intensity_diluted=spec7023*(R/d)**2
    
    if not(quiet):
        plt.figure(3)
        plt.loglog(wave,intensity_diluted,label='d={:.0f}mpc, {:.0f}" '.format(d/3.086e15,d/(3.086e18*3.4268e-3)))
        plt.title('diluted spectrum')
        plt.xlabel('wavelength (nm)')
        plt.ylabel('Specific intensity (erg/sec/cm$^2$/nm/sr)')
        plt.subplots_adjust(top=0.986,bottom=0.111,left=0.126,right=0.991,hspace=0.2,wspace=0.2)
        plt.legend()
        plt.tight_layout()
    
    '''extinction'''
    NH=2e21 #conlumn density in cm-2, a column density of 2e21cm-2 implies an Av of 1mag
#    Av=0
    taunu=NH*Av*Cext #opacity    
    Fe_l=intensity_diluted*np.exp(-taunu) #intensity diluted and extincted
    if not(quiet):
        plt.figure(4)
        plt.loglog(wave,Fe_l,label='d={:.0f}mpc, {:.0f}" '.format(d/3.086e15,d/(3.086e18*3.4268e-3)))
        plt.title('ext spectrum')
        plt.xlabel('wavelength (nm)')
        plt.ylabel('Specific intensity (erg/sec/cm$^2$/nm/sr)')
        plt.legend()
        plt.tight_layout()
    
    

    '''conversion in erg/sec/cm2/sr/eV'''
    c=299792458000000000 #light speed in nm
    h=4.135667516e-15 #planck constant in eV*sec
    
    E=(h*c/wave)[::-1] #ev
    Fe_E_=Fe_l[::-1]*h*c/(E**2)
    
    if not(quiet):
        plt.figure(5)
        plt.loglog(E,Fe_E_,label='d={:.0f}mpc, {:.0f}" '.format(d/3.086e15,d/(3.086e18*3.4268e-3)))
#        plt.loglog(E,Fe_E_,label='d={:.0f}mpc, {:.0f}" '.format(d/3.086e15,d/(3.086e18*3.4268e-3)))
        plt.title('energy spectrum')
        plt.xlabel('photon energy (eV)')
        plt.ylabel('Specific intensity (erg/sec/cm$^2$/eV/sr)')
        plt.subplots_adjust(top=0.986,bottom=0.111,left=0.126,right=0.991,hspace=0.2,wspace=0.2)
        plt.legend()
        plt.tight_layout()

    return Fe_E_,E,Fe_l,wave


    
    
def yield_n_to_p(E,IP):
    from numpy import where, concatenate,zeros,full
    '''adapt the photoelectric yield Neutral -> Cation
    version from Joachims et al 1996
    E: 1xn array, photonenergy vector
    IP: float, first ionization potential of the molecule
    
    returns the yield of the first photoionization
    '''

    first_part=where(E<IP)[0]
    scdn_part=where((E>=IP) & (E<(IP+9.2)))[0]
    third_part=where((E>IP+9.2))[0]
    
    Y_1=zeros([len(first_part)])
    Y_2=(E[scdn_part]-IP)/9.2
    Y_3=full(len(third_part), 1)

    Yntop=concatenate((Y_1,Y_2,Y_3))
    return Yntop    


def yield_p_to_2p(E,IP2,Nc,alpha=0.3):
    from numpy import where, concatenate,zeros,full
    '''adapt the photoelectric yield Cation -> Dication to the molecule
    version from Wenzel et al 2020
    E: 1xn array, photonenergy vector
    IP2: float, second ionization potential of the molecule
    Nc: int or float, size of the PAH in number of carbon atoms 
    alpha: steepness coefficient, see Wenzel et al. 2020, this parameter should not change
    
    returns the yield of the second photoionization
    '''
    
    if (Nc>=32) & (Nc<50):
        beta=0.59+8.1e-3*Nc
    if Nc>=50:
        beta=1

    first_part=where(E<IP2)[0]
    scdn_part=where((E>=IP2) & (E<(11.3)))[0]
    third_part=where((E>=11.3) & (E<(12.9)))[0]
    fourth_part=where((E>=12.9) & (E<(13.6)))[0]
#    print(Nc)
    Y_1=zeros([len(first_part)])
    Y_2=alpha/(11.3-IP2)*(E[scdn_part]-IP2)
    Y_3=full(len(third_part), alpha)
    Y_4=(beta-alpha)/2.1*(E[fourth_part]-12.9)+alpha
    Yptopp=concatenate((Y_1,Y_2,Y_3,Y_4))
    
    return Yptopp
    
def IP_estimate(Nc,Z):
    '''gives the estimate of the IP from Z to Z+1 according to Eq.1 of Wenzel et al. 2020
    Nc: int of float, size of the PAH in carbon atoms
    Z: positive int, the charge state of the molecule 
    
    returns the IP of the molecule'''
    
    eps0=8.85e-21 #en F/nm
    e=1.6e-19 #en C
    
    IP=3.9+(1/(4*np.pi*eps0))*((Z+0.5)*(e**2)/((Nc/468)**(1/3))+(Z+2)*((e**2)/((Nc/468)**(1/3)))*0.03/((Nc/468)**(1/3)))*1/e #en eV
    
    return IP




def cross_secs(Nc,quiet):
    '''==========================|building of the cross section|======================='''
    ''' derives a mean photoabsorption cross section of the molecule considered, in 3 size ranges 
    Nc: int, size of the molecule in carbon atoms
    quiet: bool, if False plot the resulting cross sections'''
    #absorption cross section
    if (Nc>=32) & (Nc<40):
        #32 a 40
        eVN_,crossN_1_1=np.loadtxt('./neutrals/ovalene_neutral.txt',unpack=True) #C32
        eVC_,crossC_1_1=np.loadtxt('./cations/ovalene_cation.txt',unpack=True) #C32
        eVDC_,crossDC_1_1=np.loadtxt('./dications/ovalene_dication.txt',unpack=True) #C32
        
        eVN_,crossN_2_1=np.loadtxt('./neutrals/tetrabenzocoronene_neutral.txt',unpack=True) #C36
        eVC_,crossC_2_1=np.loadtxt('./cations/tetrabenzocoronene_cation.txt',unpack=True) #C36
        eVDC_,crossDC_2_1=np.loadtxt('./dications/tetrabenzocoronene_dication.txt',unpack=True) #C36
        
        eVN_,crossN_3_1=np.loadtxt('./neutrals/circumbiphenyl_neutral.txt',unpack=True) #C38
        eVC_,crossC_3_1=np.loadtxt('./cations/circumbiphenyl_cation.txt',unpack=True) #C38
        eVDC_,crossDC_3_1=np.loadtxt('./dications/circumbiphenyl_dication.txt',unpack=True) #C38
        
        cross_N_=(((crossN_1_1/32)+(crossN_2_1/36)+(crossN_3_1/38))/3)*Nc
        cross_C_=(((crossC_1_1/32)+(crossC_2_1/36)+(crossC_3_1/38))/3)*Nc
        cross_DC_=(((crossDC_1_1/32)+(crossDC_2_1/36)+(crossDC_3_1/38))/3)*Nc
        
        E_range=np.where(eVN_<13.6)[0]
        
        if not(quiet):
            plt.figure(121)
            plt.plot(eVN_[E_range],cross_N_[E_range],'b',label='Z=0')
            plt.plot(eVC_[E_range],cross_C_[E_range],'r',label='Z=1')
            plt.plot(eVDC_[E_range],cross_DC_[E_range],'orange',label='Z=2')
            plt.ylabel('$\sigma$ (Mb)')
            plt.xlabel('photon energy (eV)')
            plt.legend()
#        plt.plot(eVN_[E_range],crossN_1_1[E_range])
#        plt.plot(eVN_[E_range],crossN_2_1[E_range])
#        plt.plot(eVN_[E_range],crossN_3_1[E_range])
    elif (Nc>=40) & (Nc<48):
        #40 a 48

        eVN_,crossN_1_2=np.loadtxt('./neutrals/circumanthracene_neutral.txt',unpack=True) #C40
        eVC_,crossC_1_2=np.loadtxt('./cations/circumanthracene_cation.txt',unpack=True) #C40
        eVDC_,crossDC_1_2=np.loadtxt('./dications/circumanthracene_dication.txt',unpack=True) #C40
        
        eVN_,crossN_2_2=np.loadtxt('./neutrals/circumpyrene_neutral.txt',unpack=True) #C42
        eVC_,crossC_2_2=np.loadtxt('./cations/circumpyrene_cation.txt',unpack=True) #C42
        eVDC_,crossDC_2_2=np.loadtxt('./dications/circumpyrene_dication.txt',unpack=True) #C42
        
        eVN_,crossN_3_2=np.loadtxt('./neutrals/hexabenzocoronene_neutral.txt',unpack=True) #C42
        eVC_,crossC_3_2=np.loadtxt('./cations/hexabenzocoronene_cation.txt',unpack=True) #C42
        eVDC_,crossDC_3_2=np.loadtxt('./dications/hexabenzocoronene_dication.txt',unpack=True) #C42    

        E_range=np.where(eVN_<13.6)[0]
        
        cross_N_=(((crossN_1_2/40)+(crossN_2_2/42)+(crossN_3_2/42))/3)*Nc
        cross_C_=(((crossC_1_2/40)+(crossC_2_2/42)+(crossC_3_2/42))/3)*Nc
        cross_DC_=(((crossDC_1_2/40)+(crossDC_2_2/42)+(crossDC_3_2/42))/3)*Nc
        if not(quiet):
        
            plt.figure(121)
            plt.plot(eVN_[E_range],cross_N_[E_range],'b',label='Z=0')
            plt.plot(eVC_[E_range],cross_C_[E_range],'r',label='Z=1')
            plt.plot(eVDC_[E_range],cross_DC_[E_range],'orange',label='Z=2')
            plt.ylabel('$\sigma$ (Mb)')
            plt.xlabel('photon energy (eV)')
            plt.legend()
#        plt.plot(eVN_[E_range],crossN_1_1[E_range])
#        plt.plot(eVN_[E_range],crossN_2_1[E_range])
#        plt.plot(eVN_[E_range],crossN_3_1[E_range])
    elif (Nc>=48):
        #48 a 66
        eVN_,crossN_1_3=np.loadtxt('./neutrals/dicoronylene_neutral.txt',unpack=True) #C48
        eVC_,crossC_1_3=np.loadtxt('./cations/dicoronylene_cation.txt',unpack=True) #C48
        eVDC_,crossDC_1_3=np.loadtxt('./dications/dicoronylene_dication.txt',unpack=True) #C48
        
        eVN_,crossN_2_3=np.loadtxt('./neutrals/circumcoronene_neutral.txt',unpack=True) #C54
        eVC_,crossC_2_3=np.loadtxt('./cations/circumcoronene_cation.txt',unpack=True) #C54
        eVDC_,crossDC_2_3=np.loadtxt('./dications/circumcoronene_dication.txt',unpack=True) #C54
        
        eVN_,crossN_3_3=np.loadtxt('./neutrals/circumovalene_neutral.txt',unpack=True) #C66
        eVC_,crossC_3_3=np.loadtxt('./cations/circumovalene_cation.txt',unpack=True) #C66
        eVDC_,crossDC_3_3=np.loadtxt('./dications/circumovalene_dication.txt',unpack=True) #C66
        
        E_range=np.where(eVN_<13.6)[0]
        
        cross_N_=(((crossN_1_3/48)+(crossN_2_3/54)+(crossN_3_3/66))/3)*Nc
        cross_C_=(((crossC_1_3/48)+(crossC_2_3/54)+(crossC_3_3/66))/3)*Nc
        cross_DC_=(((crossDC_1_3/48)+(crossDC_2_3/54)+(crossDC_3_3/66))/3)*Nc
        if not(quiet):
            plt.figure(121)
            plt.plot(eVN_[E_range],cross_N_[E_range],'b',label='Z=0')
            plt.plot(eVC_[E_range],cross_C_[E_range],'r',label='Z=1')
            plt.plot(eVDC_[E_range],cross_DC_[E_range],'orange',label='Z=2')
            plt.ylabel('$\sigma$ (Mb)')
            plt.xlabel('photon energy (eV)')
            plt.legend()
    return eVN_[E_range],eVC_[E_range],eVDC_[E_range],cross_N_[E_range],cross_C_[E_range],cross_DC_[E_range]