"""
Created on Wed Aug  7 14:14:34 2024

@author: Ioanna Koutsouridou (University of Florence)
"""

import scipy.io
import numpy as np
from scipy.interpolate import interp1d

data = scipy.io.readsav('NLTE_Al_mean3D_NordlanderLind2017.sav')
#data = scipy.io.readsav('NLTE_Al_MARCS_NordlanderLind2017.sav')

data.keys()

variable_name = 'nltegrid'
if variable_name in data:
    variable_data = data[variable_name]
    #print(variable_data)
else:
    print(f"Variable '{variable_name}' not found in the .sav file.")
    
    
dtype=[(('nabund', 'NABUND'), '>i4'), (('nteff', 'NTEFF'), '>i4'), (('nlogg', 'NLOGG'), '>i4'), (('nfeh', 'NFEH'), '>i4'), (('abund', 'ABUND'), 'O'), (('teff', 'TEFF'), 'O'), (('logg', 'LOGG'), 'O'), (('feh', 'FEH'), 'O'), (('nlevels', 'NLEVELS'), '>i4'), (('nlines', 'NLINES'), '>i4'), (('conf', 'CONF'), 'O'), (('term', 'TERM'), 'O'), (('spec', 'SPEC'), 'O'), (('ec', 'EC'), 'O'), (('ev', 'EV'), 'O'), (('wavelength', 'WAVELENGTH'), 'O'), (('irad', 'IRAD'), 'O'), (('jrad', 'JRAD'), 'O'), (('gw', 'GW'), 'O'), (('gwrep', 'GWREP'), 'O'), (('gq', 'GQ'), 'O'), (('ga', 'GA'), 'O'), (('loggf', 'LOGGF'), 'O'), (('logew', 'LOGEW'), 'O'), (('converged', 'CONVERGED'), 'O'), (('interpolated', 'INTERPOLATED'), 'O')]


    



ltegrid = {
    'teff': data['ltegrid']['TEFF'],  
    'logg':  data['ltegrid']['logg'],      
    'feh': data['ltegrid']['feh'],   
    'abund': data['ltegrid']['abund'],     
    'logEW': data['ltegrid']['logew'],  
    'wave' : data['ltegrid']['wavelength']
}

nltegrid = {
    'teff': data['nltegrid']['TEFF'],  
    'logg':  data['nltegrid']['logg'],     
    'feh': data['nltegrid']['feh'],   
    'abund': data['nltegrid']['abund'],     
    'logEW': data['nltegrid']['logew'],  
    'wave' : data['nltegrid']['wavelength']
}





def find_ew(line, teff, logg, feh, xfe, grid, warnings=None):
    n = len(xfe)
    EWs = np.full(n, np.nan)
    
    for i in range(len(xfe)):
        ti, tfrac = find_index(grid['teff'][0], teff)
        gi, gfrac = find_index(grid['logg'][0], logg)
        fi, ffrac = find_index(grid['feh'][0], feh)
        
        cog1 = grid['logEW'][0][line[0], fi[0], gi[0], ti,  :] * gfrac + grid['logEW'][0][line[0], fi[0], gi[1], ti, :] * (1 - gfrac)
        cog1 = cog1[0,:] * tfrac + cog1[1,:] * (1 - tfrac)
        cog2 = grid['logEW'][0][line[0], fi[1], gi[0], ti,  :] * gfrac + grid['logEW'][0][line[0], fi[1], gi[1], ti, :] * (1 - gfrac)
        cog2 = cog2[0,:] * tfrac + cog2[1,:] * (1 - tfrac)
        cog = cog1 * ffrac + cog2 * (1 - ffrac)
    
 
        finite_cog = np.isfinite(cog)
        if np.sum(finite_cog) > 1:
            interpolator = interp1d(grid['abund'][0][finite_cog], cog[finite_cog], kind='linear', fill_value='extrapolate')

            EWs[i] = 10 ** float(interpolator(xfe[i]))
        else:
            if warnings is not None:
                warnings[i] = 1
            EWs[i] = np.nan

    if warnings is not None:
        warnings[0] = 0 if np.all(np.isfinite(EWs)) else 1

    return EWs

def find_index(array, x, log=False):
    if log:
        array = np.log10(array)
        x = np.log10(x)
    
    indices = np.argsort(np.abs(array - x))[:2]
    i0, i1 = np.sort(indices)
    
    if log:
        frac = (np.log10(array[i1]) - x) / (np.log10(array[i1]) - np.log10(array[i0]))
    else:
        frac = (array[i1] - x) / (array[i1] - array[i0])
    
    return np.array([i0, i1]), frac





warnings = np.zeros(1)  # Initialize warnings array

    

#####   EXAMPLE    ######   

teff = 4247
logg = 1.59
feh = -0.52    
line = [2,3,4,5,6,7,11,15,16,17,18,19,20,21]
LTE = np.array([6.266,6.214,6.217,6.207,6.241,6.245,6.186,6.349,6.421,6.290,6.281,6.266,6.344,6.344]) 



xfe = LTE - 6.43 - feh      #translate A(Al) to [Al/Fe]

for i in range(len(line)):
    LTE  = find_ew([line[i]], teff, logg, feh, [xfe[i]], grid=ltegrid,  warnings=warnings)
    NLTE = find_ew([line[i]], teff, logg, feh, [xfe[i]], grid=nltegrid,  warnings=warnings)

    
    nltes=np.zeros(len(nltegrid['abund'][0]))
    ltes=np.zeros(len(nltegrid['abund'][0]))
    for k in range(len(nltegrid['abund'][0])):
        nltes[k] = find_ew([line[i]], teff, logg, feh, [nltegrid['abund'][0][k]], grid=nltegrid,  warnings=warnings)[0]
        ltes[k] = find_ew([line[i]], teff, logg, feh, [nltegrid['abund'][0][k]], grid=ltegrid,  warnings=warnings)[0]
    
    interp1 = interp1d(nltes,nltegrid['abund'][0], kind='linear', fill_value='extrapolate')
    interp2 = interp1d(nltegrid['abund'][0], interp1(ltes) - nltegrid['abund'][0], kind='linear', fill_value='extrapolate')
    dNLTE = float(interp2(xfe[i]))

    print('wavelength=', nltegrid['wave'][0][line[i]],'NLTE=', xfe[i]+dNLTE, 'LTE=', xfe[i], 'dNLTE=',dNLTE,'EW_LTE=',10*float(LTE[0]))