import numpy as np
import matplotlib.pyplot as plt
import interp_tools as model

# EXAMPLE 1: Reading synthetic spectra in the grid 
# Input:
path='./' #specify the directory where the fits files are
Teff=6000.0 # Teff_nodes=[4000, 4500, 5000, 5500, 6000, 6500]
logg= 4.5 # logg_nodes=[1.5, 2.0, 2.5, 3.0, 4.0, 4.5, 5.0]
met=0.0 # [Fe/H]_nodes=[-4.0, -3.0, -2.0, -1.0, 0.0, 0.5]
ab= 6.25 # absolute abundance A(Na)
# Output:
wavelength, fluxNLTE = model.read_fluxgrid(Teff, logg, met, ab, ilte='NLTE') #3D NLTE flux closest to the input parameters
wavelength, fluxLTE = model.read_fluxgrid(Teff, logg, met, ab, ilte='LTE') #3D LTE flux closest to the input parameters

###----------------------------------------------------------------------------------###

# EXAMPLE 2: get a line profile (normalized flux) at a given abundance 
# Input:
Teff = 6000.0 #Kelvin, effective temperature
logg = 4.5 #cgs, surface gravity
fe_h = 0.0 #dex, metallicity
abund= 6.22 # absolute sodium abundance (A(Na)=log(N_Na/N_H)+12)
line= 5682 #Na line in Ångstrom
wl, flux1 = model.pred_flux(Teff, logg, fe_h, abund, line)

# Plot flux at different abundances:
ab_list = [5.72, 5.97, 6.22, 6.47, 6.72] #dex, absolute Na abundance (A(Na)=log(N_Na/N_H)+12)
for iab, abund in enumerate(ab_list):
    wl, flux = model.pred_flux(Teff, logg, fe_h, abund, line)   
    plt.plot(wl, flux, lw=1.0, label=f'A(Na)={abund}')
plt.title('$T_\mathrm{eff}=$'+str(Teff)+' K, $\log g=$'+str(logg)+', [Fe/H]='+str(fe_h))
plt.xlabel('Wavelength (Å)')
plt.ylabel('Normalized Flux')
plt.legend()
plt.show()

###----------------------------------------------------------------------------------###

# EXAMPLE 3: get a 3D non-LTE abundance from a given equivalent width (EW) of a line, or the reduced equivalent width (REW=log10(EW/lambda))
# Input: stellar parameters and the EW or REW of a Na line
Teff = 6000.0 #Kelvin, effective temperature
logg = 4.5 #cgs, surface gravity
fe_h = 0.0 #dex, metallicity
ew = 0.100 #Ångstrom, equivalent width (EW)
line= 5682 #Ångstrom, the line for which the EW is measured
rew = np.log10(ew/line)
# Find the 3D NLTE Na abundance:
ab3dnlte = model.pred_abund(Teff, logg, fe_h, rew, line, model='3DNLTE')
print('A(Na)_3DNLTE= ', ab3dnlte,' dex')

###----------------------------------------------------------------------------------###

# EXAMPLE 4: get abundance correction (3D NLTE- 1DLTE)
# Input: stellar parameters and the 1D LTE abundance of the Na line
Teff = 6000.0 #Kelvin, effective temperature
logg = 4.5 #cgs, surface gravity
fe_h = 0.0 #dex, metallicity
vmic = 1.0 #km/s, microturbulence velocity for the 1D LTE model
ab1dlte = 6.22 #dex, the 1D LTE Na abundance 
line= 5682 #Ångstrom
# find NLTE correction from the 1D LTE abundance:
nltec = model.nlte_correction(Teff, logg, fe_h, vmic, ab1dlte, line)
print('Abundance correction (3D NLTE - 1D LTE)= ', nltec, ' dex')
# or find NLTE correction from the measured EW:
ew = 0.100 #Ångstrom
nltec_ew = model.nlte_correctionEW(Teff, logg, fe_h, vmic, ew, line)
print('Abundance correction (3D NLTE - 1D LTE) from EW= ', nltec_ew, ' dex')