import smart
import numpy as np
import sys, os, os.path, time
from astropy.table import Table
from numpy.linalg import inv, det
from ..utils.interpolations import trilinear_interpolation

##############################################################################################################

def InterpModel(teff, logg=4, metal=0, alpha=0, modelset='marcs-apogee-dr15', instrument='nirspec', order=33):

    FULL_PATH  = os.path.realpath(__file__)
    BASE, NAME = os.path.split(FULL_PATH)

    # Check the model set and instrument
    if instrument == 'nirspec':
        if modelset.lower() == 'btsettl08':
            path     = BASE + '/../libraries/btsettl08/NIRSPEC-O%s-RAW/'%order
            Gridfile = BASE + '/../libraries/btsettl08/btsettl08_gridparams.csv'
        elif modelset.lower() == 'phoenix-aces-agss-cond-2011':
            path     = BASE + '/../libraries/PHOENIX_ACES_AGSS_COND_2011/NIRSPEC-O%s-RAW/'%order
            Gridfile = BASE + '/../libraries/PHOENIX_ACES_AGSS_COND_2011/PHOENIX_ACES_AGSS_COND_2011_gridparams.csv'
        elif modelset.lower() == 'sonora-2018':
            path     = BASE + '/../libraries/SONORA_2018/NIRSPEC-O%s-RAW/'%order
            Gridfile = BASE + '/../libraries/SONORA_2018/SONORA_2018_gridparams.csv'

    elif instrument == 'apogee':
        if modelset.lower() == 'btsettl08':
            path     = BASE + '/../libraries/btsettl08/APOGEE-RAW/'
            Gridfile = BASE + '/../libraries/btsettl08/btsettl08_gridparams_apogee.csv'
        elif modelset.lower() == 'phoenix-btsettl-cifist2011-2015':
            path     = BASE + '/../libraries/PHOENIX_BTSETTL_CIFIST2011_2015/APOGEE-RAW/'
            Gridfile = BASE + '/../libraries/PHOENIX_BTSETTL_CIFIST2011_2015/PHOENIX_BTSETTL_CIFIST2011_2015_gridparams_apogee.csv'
        elif modelset.lower() == 'phoenix-aces-agss-cond-2011' :
            path     = BASE + '/../libraries/PHOENIX_ACES_AGSS_COND_2011/APOGEE-RAW/'
            Gridfile = BASE + '/../libraries/PHOENIX_ACES_AGSS_COND_2011/PHOENIX_ACES_AGSS_COND_2011_gridparams_apogee.csv'
        elif modelset.lower() == 'marcs-apogee-dr15' :
            path     = BASE + '/../libraries/MARCS_APOGEE_DR15/APOGEE-RAW/'
            Gridfile = BASE + '/../libraries/MARCS_APOGEE_DR15/MARCS_APOGEE_DR15_gridparams_apogee.csv'

    # Read the grid file
    T1 = Table.read(Gridfile)

    ###################################################################################

    def GetModel(temp, wave=False, **kwargs):

        logg       = kwargs.get('logg', 4.5)
        metal      = kwargs.get('metal', 0)
        alpha      = kwargs.get('alpha', 0)
        gridfile   = kwargs.get('gridfile', None)
        instrument = kwargs.get('instrument', 'nirspec')
        order      = kwargs.get('order', None)

        if gridfile is None:
            raise ValueError('Model gridfile must be provided.')

        if instrument == 'nirspec':
            if modelset.lower() == 'btsettl08':
                filename = 'btsettl08_t'+ str(int(temp.data[0])) + '_g' + '{0:.2f}'.format(float(logg)) + '_z-' + '{0:.2f}'.format(float(metal)) + '_en' + '{0:.2f}'.format(float(alpha)) + '_NIRSPEC-O' + str(order) + '-RAW.txt'
            elif modelset.lower() == 'phoenix-aces-agss-cond-2011':
                filename = 'PHOENIX_ACES_AGSS_COND_2011_t{0:03d}'.format(int(temp.data[0])) + '_g{0:.2f}'.format(float(logg)) + '_z{0:.2f}'.format(float(metal)) + '_alpha{0:.2f}'.format(float(alpha)) + '_NIRSPEC-O' + str(order) + '-RAW.txt'
            elif modelset.lower() == 'sonora-2018':
                filename = 'SONORA_2018_t{0:03d}'.format(int(temp.data[0])) + '_g{0:.2f}'.format(float(logg)) + '_FeH{0:.2f}'.format(0) + '_Y{0:.2f}'.format(0.28) + '_CO{0:.2f}'.format(1.00) + '_NIRSPEC-O' + str(order) + '-RAW.txt'

        if instrument == 'apogee':
            filename = gridfile['File'][np.where( (gridfile['Temp']==temp) & (gridfile['Logg']==logg) & (gridfile['Metal']==metal) & (gridfile['Alpha']==alpha) )].data[0]

        Tab = Table.read(path+filename, format='ascii.tab', names=['wave', 'flux'])

        if wave:
            return Tab['wave']
        else:
            return Tab['flux']

    ###################################################################################

    # Check if the model already exists (grid point)
    if modelset.lower() == 'sonora-2018':
        if (teff, logg) in zip(T1['Temp'], T1['Logg']):
            metal, ys = 0, 0.28
            index0 = np.where( (T1['Temp'] == teff) & (T1['Logg'] == logg) & (T1['FeH'] == metal) & (T1['Y'] == ys) )
            #flux2  = GetModel(T1['Temp'][index0], T1['Logg'][index0], T1['Metal'][index0], modelset=modelset )
            #waves2 = GetModel(T1['Temp'][index0], T1['Logg'][index0], T1['Metal'][index0], modelset=modelset, wave=True)
            flux2  = GetModel(T1['Temp'][index0], logg=T1['Logg'][index0], metal=T1['FeH'][index0], alpha=T1['Y'][index0], instrument=instrument, order=order, gridfile=T1)
            waves2 = GetModel(T1['Temp'][index0], logg=T1['Logg'][index0], metal=T1['FeH'][index0], alpha=T1['Y'][index0], instrument=instrument, order=order, gridfile=T1, wave=True)
            return waves2, flux2
    else:
        if (teff, logg, metal, alpha) in zip(T1['Temp'], T1['Logg'], T1['Metal'], T1['Alpha']):
            index0 = np.where( (T1['Temp'] == teff) & (T1['Logg'] == logg) & (T1['Metal'] == metal) & (T1['Alpha'] == alpha) )
            #flux2  = GetModel(T1['Temp'][index0], T1['Logg'][index0], T1['Metal'][index0], modelset=modelset )
            #waves2 = GetModel(T1['Temp'][index0], T1['Logg'][index0], T1['Metal'][index0], modelset=modelset, wave=True)
            flux2  = GetModel(T1['Temp'][index0], logg=T1['Logg'][index0], metal=T1['Metal'][index0], alpha=T1['Alpha'][index0], instrument=instrument, order=order, gridfile=T1)
            waves2 = GetModel(T1['Temp'][index0], logg=T1['Logg'][index0], metal=T1['Metal'][index0], alpha=T1['Alpha'][index0], instrument=instrument, order=order, gridfile=T1, wave=True)
            return waves2, flux2

    try:
        if modelset.lower() == 'sonora-2018':
            metal, alpha = 0, 0.28
            # Get the nearest models to the gridpoint (Temp)
            x0 = np.max(T1['Temp'][np.where(T1['Temp'] <= teff)])
            x1 = np.min(T1['Temp'][np.where(T1['Temp'] >= teff)])
            #print(x0, x1)

            # Get the nearest grid point to Logg
            y0 = np.max(list(set(T1['Logg'][np.where( (T1['Temp'] == x0) & (T1['Logg'] <= logg) )]) &
                             set(T1['Logg'][np.where( (T1['Temp'] == x1) & (T1['Logg'] <= logg) )])))
            y1 = np.min(list(set(T1['Logg'][np.where( (T1['Temp'] == x0) & (T1['Logg'] >= logg) )]) &
                             set(T1['Logg'][np.where( (T1['Temp'] == x1) & (T1['Logg'] >= logg) )])))
            #print(y0, y1)

            # Get the nearest grid point to [M/H]
            z0 = np.max(list(set(T1['FeH'][np.where( (T1['Temp'] == x0) & (T1['Logg'] == y0) & (T1['FeH'] <= metal) )]) &
                             set(T1['FeH'][np.where( (T1['Temp'] == x1) & (T1['Logg'] == y1) & (T1['FeH'] <= metal) )])))
            z1 = np.min(list(set(T1['FeH'][np.where( (T1['Temp'] == x0) & (T1['Logg'] == y0) & (T1['FeH'] >= metal) )]) &
                             set(T1['FeH'][np.where( (T1['Temp'] == x1) & (T1['Logg'] == y1) & (T1['FeH'] >= metal) )])))
            #print(z0, z1)

            # Get the nearest grid point to Alpha
            t0 = np.max(list(set(T1['Y'][np.where( (T1['Temp'] == x0) & (T1['Logg'] == y0) & (T1['FeH'] == z0) & (T1['Y'] <= alpha) )]) &
                             set(T1['Y'][np.where( (T1['Temp'] == x1) & (T1['Logg'] == y1) & (T1['FeH'] == z1) & (T1['Y'] <= alpha) )])))
            t1 = np.min(list(set(T1['Y'][np.where( (T1['Temp'] == x0) & (T1['Logg'] == y0) & (T1['FeH'] == z0) & (T1['Y'] >= alpha) )]) &
                             set(T1['Y'][np.where( (T1['Temp'] == x1) & (T1['Logg'] == y1) & (T1['FeH'] == z1) & (T1['Y'] >= alpha) )])))
            #print(t0, t1)

        else:
            # Get the nearest models to the gridpoint (Temp)
            x0 = np.max(T1['Temp'][np.where(T1['Temp'] <= teff)])
            x1 = np.min(T1['Temp'][np.where(T1['Temp'] >= teff)])
            #print('teff:', x0, teff, x1)
            # Get the nearest grid point to Logg
            y0 = np.max(list(set(T1['Logg'][np.where( (T1['Temp'] == x0) & (T1['Logg'] <= logg) )]) &
                             set(T1['Logg'][np.where( (T1['Temp'] == x1) & (T1['Logg'] <= logg) )])))
            y1 = np.min(list(set(T1['Logg'][np.where( (T1['Temp'] == x0) & (T1['Logg'] >= logg) )]) &
                             set(T1['Logg'][np.where( (T1['Temp'] == x1) & (T1['Logg'] >= logg) )])))
            #print('logg:', y0, logg, y1)
            # Get the nearest grid point to [M/H]
            #print(metal)
            #print(list(set(T1['Metal'][np.where( (T1['Temp'] == x0) & (T1['Logg'] == y0) )])))
            #print(list(set(T1['Metal'][np.where( (T1['Temp'] == x1) & (T1['Logg'] == y1) )])))
            #print(list(set(T1['Metal'][np.where( (T1['Temp'] == x0) & (T1['Logg'] == y0) & (T1['Metal'] <= metal))])))
            #print(list(set(T1['Metal'][np.where( (T1['Temp'] == x1) & (T1['Logg'] == y1) & (T1['Metal'] <= metal))])))
            #print(list(set(T1['Metal'][np.where( (T1['Temp'] == x0) & (T1['Logg'] == y0) & (T1['Metal'] >= metal))])))
            #print(list(set(T1['Metal'][np.where( (T1['Temp'] == x1) & (T1['Logg'] == y1) & (T1['Metal'] >= metal))])))
            z0 = np.max(list(set(T1['Metal'][np.where( (T1['Temp'] == x0) & (T1['Logg'] == y0) & (T1['Metal'] <= metal) )]) &
                             set(T1['Metal'][np.where( (T1['Temp'] == x1) & (T1['Logg'] == y1) & (T1['Metal'] <= metal) )])))
            z1 = np.min(list(set(T1['Metal'][np.where( (T1['Temp'] == x0) & (T1['Logg'] == y0) & (T1['Metal'] >= metal) )]) &
                             set(T1['Metal'][np.where( (T1['Temp'] == x1) & (T1['Logg'] == y1) & (T1['Metal'] >= metal) )])))
            #print('metal:', z0, metal, z1)
            # Get the nearest grid point to Alpha
            #print(list(set(T1['Alpha'][np.where( (T1['Temp'] == x0) & (T1['Logg'] == y0) & (T1['Metal'] == z0) )])))
            #print(list(set(T1['Alpha'][np.where( (T1['Temp'] == x1) & (T1['Logg'] == y1) & (T1['Metal'] == z1) )])))
            #print(list(set(T1['Alpha'][np.where( (T1['Temp'] == x0) & (T1['Logg'] == y0) & (T1['Metal'] == z0) & (T1['Alpha'] <= alpha) )])))
            #print(list(set(T1['Alpha'][np.where( (T1['Temp'] == x1) & (T1['Logg'] == y1) & (T1['Metal'] == z1) & (T1['Alpha'] <= alpha) )])))
            #print(list(set(T1['Alpha'][np.where( (T1['Temp'] == x0) & (T1['Logg'] == y0) & (T1['Metal'] == z0) & (T1['Alpha'] >= alpha) )])))
            #print(list(set(T1['Alpha'][np.where( (T1['Temp'] == x1) & (T1['Logg'] == y1) & (T1['Metal'] == z1) & (T1['Alpha'] >= alpha) )])))
            t0 = np.max(list(set(T1['Alpha'][np.where( (T1['Temp'] == x0) & (T1['Logg'] == y0) & (T1['Metal'] == z0) & (T1['Alpha'] <= alpha) )]) &
                             set(T1['Alpha'][np.where( (T1['Temp'] == x1) & (T1['Logg'] == y1) & (T1['Metal'] == z1) & (T1['Alpha'] <= alpha) )])))
            t1 = np.min(list(set(T1['Alpha'][np.where( (T1['Temp'] == x0) & (T1['Logg'] == y0) & (T1['Metal'] == z0) & (T1['Alpha'] >= alpha) )]) &
                             set(T1['Alpha'][np.where( (T1['Temp'] == x1) & (T1['Logg'] == y1) & (T1['Metal'] == z1) & (T1['Alpha'] >= alpha) )])))
            #print('alpha:', z0, alpha, z1)
    except:
        raise ValueError('Model Parameters Teff: %0.3f, Logg: %0.3f, [M/H]: %0.3f, Alpha: %0.3f are outside the model grid.'%(teff, logg, metal, alpha))

    if modelset.lower() == 'sonora-2018':
        # Get the 16 points
        ind0000 = np.where( (T1['Temp'] == x0) & (T1['Logg'] == y0) & (T1['FeH'] == z0) & (T1['Y'] == t0) ) # 0000
        ind1000 = np.where( (T1['Temp'] == x1) & (T1['Logg'] == y0) & (T1['FeH'] == z0) & (T1['Y'] == t0) ) # 1000
        ind0100 = np.where( (T1['Temp'] == x0) & (T1['Logg'] == y1) & (T1['FeH'] == z0) & (T1['Y'] == t0) ) # 0100
        ind0010 = np.where( (T1['Temp'] == x0) & (T1['Logg'] == y0) & (T1['FeH'] == z1) & (T1['Y'] == t0) ) # 0010
        ind0001 = np.where( (T1['Temp'] == x0) & (T1['Logg'] == y0) & (T1['FeH'] == z0) & (T1['Y'] == t1) ) # 0001
        ind1001 = np.where( (T1['Temp'] == x1) & (T1['Logg'] == y0) & (T1['FeH'] == z0) & (T1['Y'] == t1) ) # 1001
        ind0101 = np.where( (T1['Temp'] == x0) & (T1['Logg'] == y1) & (T1['FeH'] == z0) & (T1['Y'] == t1) ) # 0101
        ind0011 = np.where( (T1['Temp'] == x0) & (T1['Logg'] == y0) & (T1['FeH'] == z1) & (T1['Y'] == t1) ) # 0011
        ind1011 = np.where( (T1['Temp'] == x1) & (T1['Logg'] == y0) & (T1['FeH'] == z1) & (T1['Y'] == t1) ) # 1011
        ind0111 = np.where( (T1['Temp'] == x0) & (T1['Logg'] == y1) & (T1['FeH'] == z1) & (T1['Y'] == t1) ) # 0111
        ind1111 = np.where( (T1['Temp'] == x1) & (T1['Logg'] == y1) & (T1['FeH'] == z1) & (T1['Y'] == t1) ) # 1111
        ind0110 = np.where( (T1['Temp'] == x0) & (T1['Logg'] == y1) & (T1['FeH'] == z1) & (T1['Y'] == t0) ) # 0110
        ind1010 = np.where( (T1['Temp'] == x1) & (T1['Logg'] == y0) & (T1['FeH'] == z1) & (T1['Y'] == t0) ) # 1010
        ind1100 = np.where( (T1['Temp'] == x1) & (T1['Logg'] == y1) & (T1['FeH'] == z0) & (T1['Y'] == t0) ) # 1100
        ind1101 = np.where( (T1['Temp'] == x1) & (T1['Logg'] == y1) & (T1['FeH'] == z0) & (T1['Y'] == t1) ) # 1101
        ind1110 = np.where( (T1['Temp'] == x1) & (T1['Logg'] == y1) & (T1['FeH'] == z1) & (T1['Y'] == t0) ) # 1110
        Points =  [ [np.log10(T1['Temp'][ind0000]), T1['Logg'][ind0000], T1['FeH'][ind0000], T1['Y'][ind0000],
                     np.log10(GetModel(T1['Temp'][ind0000], logg=T1['Logg'][ind0000], metal=T1['FeH'][ind0000], alpha=T1['Y'][ind0000], instrument=instrument, order=order, gridfile=T1))],
                    [np.log10(T1['Temp'][ind1000]), T1['Logg'][ind1000], T1['FeH'][ind1000], T1['Y'][ind1000],
                     np.log10(GetModel(T1['Temp'][ind1000], logg=T1['Logg'][ind1000], metal=T1['FeH'][ind1000], alpha=T1['Y'][ind1000], instrument=instrument, order=order, gridfile=T1))],
                    [np.log10(T1['Temp'][ind0100]), T1['Logg'][ind0100], T1['FeH'][ind0100], T1['Y'][ind0100],
                     np.log10(GetModel(T1['Temp'][ind0100], logg=T1['Logg'][ind0100], metal=T1['FeH'][ind0100], alpha=T1['Y'][ind0100], instrument=instrument, order=order, gridfile=T1))],
                    [np.log10(T1['Temp'][ind0010]), T1['Logg'][ind0010], T1['FeH'][ind0010], T1['Y'][ind0010],
                     np.log10(GetModel(T1['Temp'][ind0010], logg=T1['Logg'][ind0010], metal=T1['FeH'][ind0010], alpha=T1['Y'][ind0010], instrument=instrument, order=order, gridfile=T1))],
                    [np.log10(T1['Temp'][ind0001]), T1['Logg'][ind0001], T1['FeH'][ind0001], T1['Y'][ind0001],
                     np.log10(GetModel(T1['Temp'][ind0001], logg=T1['Logg'][ind0001], metal=T1['FeH'][ind0001], alpha=T1['Y'][ind0001], instrument=instrument, order=order, gridfile=T1))],
                    [np.log10(T1['Temp'][ind1001]), T1['Logg'][ind1001], T1['FeH'][ind1001], T1['Y'][ind1001],
                     np.log10(GetModel(T1['Temp'][ind1001], logg=T1['Logg'][ind1001], metal=T1['FeH'][ind1001], alpha=T1['Y'][ind1001], instrument=instrument, order=order, gridfile=T1))],
                    [np.log10(T1['Temp'][ind0101]), T1['Logg'][ind0101], T1['FeH'][ind0101], T1['Y'][ind0101],
                     np.log10(GetModel(T1['Temp'][ind0101], logg=T1['Logg'][ind0101], metal=T1['FeH'][ind0101], alpha=T1['Y'][ind0101], instrument=instrument, order=order, gridfile=T1))],
                    [np.log10(T1['Temp'][ind0011]), T1['Logg'][ind0011], T1['FeH'][ind0011], T1['Y'][ind0011],
                     np.log10(GetModel(T1['Temp'][ind0011], logg=T1['Logg'][ind0011], metal=T1['FeH'][ind0011], alpha=T1['Y'][ind0011], instrument=instrument, order=order, gridfile=T1))],
                    [np.log10(T1['Temp'][ind1011]), T1['Logg'][ind1011], T1['FeH'][ind1011], T1['Y'][ind1011],
                     np.log10(GetModel(T1['Temp'][ind1011], logg=T1['Logg'][ind1011], metal=T1['FeH'][ind1011], alpha=T1['Y'][ind1011], instrument=instrument, order=order, gridfile=T1))],
                    [np.log10(T1['Temp'][ind0111]), T1['Logg'][ind0111], T1['FeH'][ind0111], T1['Y'][ind0111],
                     np.log10(GetModel(T1['Temp'][ind0111], logg=T1['Logg'][ind0111], metal=T1['FeH'][ind0111], alpha=T1['Y'][ind0111], instrument=instrument, order=order, gridfile=T1))],
                    [np.log10(T1['Temp'][ind1111]), T1['Logg'][ind1111], T1['FeH'][ind1111], T1['Y'][ind1111],
                     np.log10(GetModel(T1['Temp'][ind1111], logg=T1['Logg'][ind1111], metal=T1['FeH'][ind1111], alpha=T1['Y'][ind1111], instrument=instrument, order=order, gridfile=T1))],
                    [np.log10(T1['Temp'][ind0110]), T1['Logg'][ind0110], T1['FeH'][ind0110], T1['Y'][ind0110],
                     np.log10(GetModel(T1['Temp'][ind0110], logg=T1['Logg'][ind0110], metal=T1['FeH'][ind0110], alpha=T1['Y'][ind0110], instrument=instrument, order=order, gridfile=T1))],
                    [np.log10(T1['Temp'][ind1010]), T1['Logg'][ind1010], T1['FeH'][ind1010], T1['Y'][ind1010],
                     np.log10(GetModel(T1['Temp'][ind1010], logg=T1['Logg'][ind1010], metal=T1['FeH'][ind1010], alpha=T1['Y'][ind1010], instrument=instrument, order=order, gridfile=T1))],
                    [np.log10(T1['Temp'][ind1100]), T1['Logg'][ind1100], T1['FeH'][ind1100], T1['Y'][ind1100],
                     np.log10(GetModel(T1['Temp'][ind1100], logg=T1['Logg'][ind1100], metal=T1['FeH'][ind1100], alpha=T1['Y'][ind1100], instrument=instrument, order=order, gridfile=T1))],
                    [np.log10(T1['Temp'][ind1101]), T1['Logg'][ind1101], T1['FeH'][ind1101], T1['Y'][ind1101],
                     np.log10(GetModel(T1['Temp'][ind1101], logg=T1['Logg'][ind1101], metal=T1['FeH'][ind1101], alpha=T1['Y'][ind1101], instrument=instrument, order=order, gridfile=T1))],
                    [np.log10(T1['Temp'][ind1110]), T1['Logg'][ind1110], T1['FeH'][ind1110], T1['Y'][ind1110],
                     np.log10(GetModel(T1['Temp'][ind1110], logg=T1['Logg'][ind1110], metal=T1['FeH'][ind1110], alpha=T1['Y'][ind1110], instrument=instrument, order=order, gridfile=T1))],
                  ]
        #print(Points)
        waves2 = GetModel(T1['Temp'][ind1111], logg=T1['Logg'][ind1111], metal=T1['FeH'][ind1111], alpha=T1['Y'][ind1111], instrument=instrument, order=order, gridfile=T1, wave=True)
    else:
        # Get the 16 points
        ind0000 = np.where( (T1['Temp'] == x0) & (T1['Logg'] == y0) & (T1['Metal'] == z0) & (T1['Alpha'] == t0) ) # 0000
        ind1000 = np.where( (T1['Temp'] == x1) & (T1['Logg'] == y0) & (T1['Metal'] == z0) & (T1['Alpha'] == t0) ) # 1000
        ind0100 = np.where( (T1['Temp'] == x0) & (T1['Logg'] == y1) & (T1['Metal'] == z0) & (T1['Alpha'] == t0) ) # 0100
        ind0010 = np.where( (T1['Temp'] == x0) & (T1['Logg'] == y0) & (T1['Metal'] == z1) & (T1['Alpha'] == t0) ) # 0010
        ind0001 = np.where( (T1['Temp'] == x0) & (T1['Logg'] == y0) & (T1['Metal'] == z0) & (T1['Alpha'] == t1) ) # 0001
        ind1001 = np.where( (T1['Temp'] == x1) & (T1['Logg'] == y0) & (T1['Metal'] == z0) & (T1['Alpha'] == t1) ) # 1001
        ind0101 = np.where( (T1['Temp'] == x0) & (T1['Logg'] == y1) & (T1['Metal'] == z0) & (T1['Alpha'] == t1) ) # 0101
        ind0011 = np.where( (T1['Temp'] == x0) & (T1['Logg'] == y0) & (T1['Metal'] == z1) & (T1['Alpha'] == t1) ) # 0011
        ind1011 = np.where( (T1['Temp'] == x1) & (T1['Logg'] == y0) & (T1['Metal'] == z1) & (T1['Alpha'] == t1) ) # 1011
        ind0111 = np.where( (T1['Temp'] == x0) & (T1['Logg'] == y1) & (T1['Metal'] == z1) & (T1['Alpha'] == t1) ) # 0111
        ind1111 = np.where( (T1['Temp'] == x1) & (T1['Logg'] == y1) & (T1['Metal'] == z1) & (T1['Alpha'] == t1) ) # 1111
        ind0110 = np.where( (T1['Temp'] == x0) & (T1['Logg'] == y1) & (T1['Metal'] == z1) & (T1['Alpha'] == t0) ) # 0110
        ind1010 = np.where( (T1['Temp'] == x1) & (T1['Logg'] == y0) & (T1['Metal'] == z1) & (T1['Alpha'] == t0) ) # 1010
        ind1100 = np.where( (T1['Temp'] == x1) & (T1['Logg'] == y1) & (T1['Metal'] == z0) & (T1['Alpha'] == t0) ) # 1100
        ind1101 = np.where( (T1['Temp'] == x1) & (T1['Logg'] == y1) & (T1['Metal'] == z0) & (T1['Alpha'] == t1) ) # 1101
        ind1110 = np.where( (T1['Temp'] == x1) & (T1['Logg'] == y1) & (T1['Metal'] == z1) & (T1['Alpha'] == t0) ) # 1110
        Points =  [ [np.log10(T1['Temp'][ind0000]), T1['Logg'][ind0000], T1['Metal'][ind0000], T1['Alpha'][ind0000],
                     np.log10(GetModel(T1['Temp'][ind0000], logg=T1['Logg'][ind0000], metal=T1['Metal'][ind0000], alpha=T1['Alpha'][ind0000], instrument=instrument, order=order, gridfile=T1))],
                    [np.log10(T1['Temp'][ind1000]), T1['Logg'][ind1000], T1['Metal'][ind1000], T1['Alpha'][ind1000],
                     np.log10(GetModel(T1['Temp'][ind1000], logg=T1['Logg'][ind1000], metal=T1['Metal'][ind1000], alpha=T1['Alpha'][ind1000], instrument=instrument, order=order, gridfile=T1))],
                    [np.log10(T1['Temp'][ind0100]), T1['Logg'][ind0100], T1['Metal'][ind0100], T1['Alpha'][ind0100],
                     np.log10(GetModel(T1['Temp'][ind0100], logg=T1['Logg'][ind0100], metal=T1['Metal'][ind0100], alpha=T1['Alpha'][ind0100], instrument=instrument, order=order, gridfile=T1))],
                    [np.log10(T1['Temp'][ind0010]), T1['Logg'][ind0010], T1['Metal'][ind0010], T1['Alpha'][ind0010],
                     np.log10(GetModel(T1['Temp'][ind0010], logg=T1['Logg'][ind0010], metal=T1['Metal'][ind0010], alpha=T1['Alpha'][ind0010], instrument=instrument, order=order, gridfile=T1))],
                    [np.log10(T1['Temp'][ind0001]), T1['Logg'][ind0001], T1['Metal'][ind0001], T1['Alpha'][ind0001],
                     np.log10(GetModel(T1['Temp'][ind0001], logg=T1['Logg'][ind0001], metal=T1['Metal'][ind0001], alpha=T1['Alpha'][ind0001], instrument=instrument, order=order, gridfile=T1))],
                    [np.log10(T1['Temp'][ind1001]), T1['Logg'][ind1001], T1['Metal'][ind1001], T1['Alpha'][ind1001],
                     np.log10(GetModel(T1['Temp'][ind1001], logg=T1['Logg'][ind1001], metal=T1['Metal'][ind1001], alpha=T1['Alpha'][ind1001], instrument=instrument, order=order, gridfile=T1))],
                    [np.log10(T1['Temp'][ind0101]), T1['Logg'][ind0101], T1['Metal'][ind0101], T1['Alpha'][ind0101],
                     np.log10(GetModel(T1['Temp'][ind0101], logg=T1['Logg'][ind0101], metal=T1['Metal'][ind0101], alpha=T1['Alpha'][ind0101], instrument=instrument, order=order, gridfile=T1))],
                    [np.log10(T1['Temp'][ind0011]), T1['Logg'][ind0011], T1['Metal'][ind0011], T1['Alpha'][ind0011],
                     np.log10(GetModel(T1['Temp'][ind0011], logg=T1['Logg'][ind0011], metal=T1['Metal'][ind0011], alpha=T1['Alpha'][ind0011], instrument=instrument, order=order, gridfile=T1))],
                    [np.log10(T1['Temp'][ind1011]), T1['Logg'][ind1011], T1['Metal'][ind1011], T1['Alpha'][ind1011],
                     np.log10(GetModel(T1['Temp'][ind1011], logg=T1['Logg'][ind1011], metal=T1['Metal'][ind1011], alpha=T1['Alpha'][ind1011], instrument=instrument, order=order, gridfile=T1))],
                    [np.log10(T1['Temp'][ind0111]), T1['Logg'][ind0111], T1['Metal'][ind0111], T1['Alpha'][ind0111],
                     np.log10(GetModel(T1['Temp'][ind0111], logg=T1['Logg'][ind0111], metal=T1['Metal'][ind0111], alpha=T1['Alpha'][ind0111], instrument=instrument, order=order, gridfile=T1))],
                    [np.log10(T1['Temp'][ind1111]), T1['Logg'][ind1111], T1['Metal'][ind1111], T1['Alpha'][ind1111],
                     np.log10(GetModel(T1['Temp'][ind1111], logg=T1['Logg'][ind1111], metal=T1['Metal'][ind1111], alpha=T1['Alpha'][ind1111], instrument=instrument, order=order, gridfile=T1))],
                    [np.log10(T1['Temp'][ind0110]), T1['Logg'][ind0110], T1['Metal'][ind0110], T1['Alpha'][ind0110],
                     np.log10(GetModel(T1['Temp'][ind0110], logg=T1['Logg'][ind0110], metal=T1['Metal'][ind0110], alpha=T1['Alpha'][ind0110], instrument=instrument, order=order, gridfile=T1))],
                    [np.log10(T1['Temp'][ind1010]), T1['Logg'][ind1010], T1['Metal'][ind1010], T1['Alpha'][ind1010],
                     np.log10(GetModel(T1['Temp'][ind1010], logg=T1['Logg'][ind1010], metal=T1['Metal'][ind1010], alpha=T1['Alpha'][ind1010], instrument=instrument, order=order, gridfile=T1))],
                    [np.log10(T1['Temp'][ind1100]), T1['Logg'][ind1100], T1['Metal'][ind1100], T1['Alpha'][ind1100],
                     np.log10(GetModel(T1['Temp'][ind1100], logg=T1['Logg'][ind1100], metal=T1['Metal'][ind1100], alpha=T1['Alpha'][ind1100], instrument=instrument, order=order, gridfile=T1))],
                    [np.log10(T1['Temp'][ind1101]), T1['Logg'][ind1101], T1['Metal'][ind1101], T1['Alpha'][ind1101],
                     np.log10(GetModel(T1['Temp'][ind1101], logg=T1['Logg'][ind1101], metal=T1['Metal'][ind1101], alpha=T1['Alpha'][ind1101], instrument=instrument, order=order, gridfile=T1))],
                    [np.log10(T1['Temp'][ind1110]), T1['Logg'][ind1110], T1['Metal'][ind1110], T1['Alpha'][ind1110],
                     np.log10(GetModel(T1['Temp'][ind1110], logg=T1['Logg'][ind1110], metal=T1['Metal'][ind1110], alpha=T1['Alpha'][ind1110], instrument=instrument, order=order, gridfile=T1))],
                  ]
        #print(Points)
        waves2 = GetModel(T1['Temp'][ind1111], logg=T1['Logg'][ind1111], metal=T1['Metal'][ind1111], alpha=T1['Alpha'][ind1111], instrument=instrument, order=order, gridfile=T1, wave=True)

    return waves2, smart.utils.interpolations.quadlinear_interpolation(np.log10(teff), logg, metal, alpha, Points)

################################################################################################################################################################################################

def InterpModel_3D(Teff, Logg, Metal, modelset='marcs-apogee-dr15', instrument='nirspec', order=33):
    Alpha = 0.0
    FULL_PATH  = os.path.realpath(__file__)
    BASE, NAME = os.path.split(FULL_PATH)

    # Check the model set
    if instrument == 'nirspec':

        if modelset.lower() == 'btsettl08':
            path = BASE + '/../libraries/btsettl08/NIRSPEC-O%s-RAW/'%order

        elif modelset.lower() == 'phoenixaces' :
            path = BASE + '/../libraries/phoenixaces/NIRSPEC-O%s-RAW/'%order

    elif instrument == 'apogee':

        if modelset.lower() == 'btsettl08':
            path = BASE + '/../libraries/btsettl08/APOGEE-RAW/'

        elif modelset.lower() == 'phoenix-btsettl-cifist2011-2015':
            path = BASE + '/../libraries/PHOENIX_BTSETTL_CIFIST2011_2015/APOGEE-RAW/'

        elif modelset.lower() == 'phoenix-aces-agss-cond-2011' :
            path = BASE + '/../libraries/PHOENIX_ACES_AGSS_COND_2011/APOGEE-RAW/'

        elif modelset.lower() == 'marcs-apogee-dr15' :
            path = BASE + '/../libraries/MARCS_APOGEE_DR15/APOGEE-RAW_3D/'

    def GetModel(temp, logg, metal, modelset='marcs-apogee-dr15', wave=False):
        en = 0.00
        if instrument == 'nirspec':

            if modelset.lower() == 'btsettl08':
                # alpha enhancement correction
                if metal == -0.5:
                    en = 0.2
                elif (metal == -1.0) or (metal == -1.5) or (metal == -2.0) or (metal == -2.5):
                    en = 0.4

                if metal == 0.0:
                    filename = 'btsettl08_t'+ str(int(temp.data[0])) + '_g' + '{0:.2f}'.format(float(logg)) + '_z-' + '{0:.2f}'.format(float(metal)) + '_en' + '{0:.2f}'.format(float(en)) + '_NIRSPEC-O' + str(order) + '-RAW.txt'
                elif metal == 0.5:
                    filename = 'btsettl08_t'+ str(int(temp.data[0])) + '_g' + '{0:.2f}'.format(float(logg)) + '_z' + '{0:.2f}'.format(float(metal)) + '_en' + '{0:.2f}'.format(float(en)) + '_NIRSPEC-O' + str(order) + '-RAW.txt'
                else:
                    filename = 'btsettl08_t'+ str(int(temp.data[0])) + '_g' + '{0:.2f}'.format(float(logg)) + '_z' + '{0:.2f}'.format(float(metal)) + '_en' + '{0:.2f}'.format(float(en)) + '_NIRSPEC-O' + str(order) + '-RAW.txt'

            elif modelset.lower() == 'phoenixaces':
                filename = 'phoenixaces_t{0:03d}'.format(int(temp.data[0])) + '_g{0:.2f}'.format(float(logg)) + '_z-{0:.2f}'.format(float(metal)) + '_en{0:.2f}'.format(float(en)) + '_NIRSPEC-O' + str(order) + '-RAW.txt'

        elif instrument == 'apogee':
            if modelset.lower() == 'btsettl08':
                # alpha enhancement correction
                if metal == -0.5:
                    en = 0.2
                elif (metal == -1.0) or (metal == -1.5) or (metal == -2.0) or (metal == -2.5):
                    en = 0.4

                if metal == 0.0:
                    filename = 'btsettl08_t'+ str(int(temp.data[0])) + '_g' + '{0:.2f}'.format(float(logg)) + '_z-' + '{0:.2f}'.format(float(metal)) + '_en' + '{0:.2f}'.format(float(en)) + '_APOGEE-RAW.txt'
                elif metal > 0.0:
                    filename = 'btsettl08_t'+ str(int(temp.data[0])) + '_g' + '{0:.2f}'.format(float(logg)) + '_z' + '{0:.2f}'.format(float(metal)) + '_en' + '{0:.2f}'.format(float(en)) + '_APOGEE-RAW.txt'
                else:
                    filename = 'btsettl08_t'+ str(int(temp.data[0])) + '_g' + '{0:.2f}'.format(float(logg)) + '_z' + '{0:.2f}'.format(float(metal)) + '_en' + '{0:.2f}'.format(float(en)) + '_APOGEE-RAW.txt'

            elif modelset.lower() == 'phoenix-btsettl-cifist2011-2015':
                filename = 'PHOENIX_BTSETTL_CIFIST2011_2015_t{0:03d}'.format(int(temp.data[0])) + '_g{0:.2f}'.format(float(logg)) + '_z{0:.2f}'.format(float(metal)) + '_en{0:.2f}'.format(float(en)) + '_APOGEE-RAW.txt'

            elif modelset.lower() == 'phoenix-aces-agss-cond-2011':
                filename = 'PHOENIX_ACES_AGSS_COND_2011_t{0:03d}'.format(int(temp.data[0])) + '_g{0:.2f}'.format(float(logg)) + '_z{0:.2f}'.format(float(metal)) + '_en{0:.2f}'.format(float(en)) + '_APOGEE-RAW.txt'

            elif modelset.lower() == 'marcs-apogee-dr15':
                filename = 'MARCS_APOGEE_DR15_t{0:03d}'.format(int(temp.data[0])) + '_g{0:.2f}'.format(float(logg)) + '_z{0:.2f}'.format(float(metal)) + '_en{0:.2f}'.format(float(en)) + '_APOGEE-RAW.txt'

        Tab = Table.read(path+filename, format='ascii.tab', names=['wave', 'flux'])

        if wave:
            return Tab['wave']
        else:
            return Tab['flux']

    if instrument == 'nirspec':

        if modelset.lower() == 'btsettl08':
            Gridfile = BASE + '/../libraries/btsettl08/btsettl08_gridparams.csv'

        elif modelset.lower() == 'phoenixaces':
            Gridfile = BASE + '/../libraries/phoenixaces/phoenixaces_gridparams.csv'

    elif instrument == 'apogee':

        if modelset.lower() == 'btsettl08':
            Gridfile = BASE + '/../libraries/btsettl08/btsettl08_gridparams_apogee.csv'

        elif modelset.lower() == 'phoenix-btsettl-cifist2011-2015':
            Gridfile = BASE + '/../libraries/PHOENIX_BTSETTL_CIFIST2011_2015/PHOENIX_BTSETTL_CIFIST2011_2015_gridparams_apogee.csv'

        elif modelset.lower() == 'phoenix-aces-agss-cond-2011':
            Gridfile = BASE + '/../libraries/PHOENIX_ACES_AGSS_COND_2011/PHOENIX_ACES_AGSS_COND_2011_gridparams_apogee.csv'

        elif modelset.lower() == 'marcs-apogee-dr15':
            Gridfile = BASE + '/../libraries/MARCS_APOGEE_DR15/MARCS_APOGEE_DR15_3D_gridparams_apogee.csv'

    T1 = Table.read(Gridfile)
    # Check if the model already exists (grid point)
    if (Teff, Logg, Metal) in zip(T1['Temp'], T1['Logg'], T1['Metal']):
        if modelset.lower() == 'btsettl08':
            if Metal == -0.5:
                Alpha = 0.2
            elif (Metal == -1.0) or (Metal == -1.5) or (Metal == -2.0) or (Metal == -2.5):
                Alpha = 0.4
        index0 = np.where( (T1['Temp'] == Teff) & (T1['Logg'] == Logg) & (T1['Metal'] == Metal) & (T1['Alpha'] == Alpha) )
        flux2  = GetModel(T1['Temp'][index0], T1['Logg'][index0], T1['Metal'][index0], modelset=modelset )
        waves2 = GetModel(T1['Temp'][index0], T1['Logg'][index0], T1['Metal'][index0], modelset=modelset, wave=True)
        return waves2, flux2

    # Get the nearest models to the gridpoint (Temp)
    x0 = np.max(T1['Temp'][np.where(T1['Temp'] <= Teff)])
    x1 = np.min(T1['Temp'][np.where(T1['Temp'] >= Teff)])
    #print(x0, Teff, x1)
    #y0 = T1['Logg'][np.where( ( (T1['Temp'] == x0) | (T1['Temp'] == x1) ) & (T1['Logg'] <= Logg) )][-1]
    #y1 = T1['Logg'][np.where( ( (T1['Temp'] == x0) | (T1['Temp'] == x1) ) & (T1['Logg'] >= Logg) )][0]
    #print(x0, list(set(T1['Logg'][np.where( ( (T1['Temp'] == x0) & (T1['Logg'] <= Logg) ) )])))
    #print(x1, list(set(T1['Logg'][np.where( ( (T1['Temp'] == x1) & (T1['Logg'] <= Logg) ) )])))
    #print(x0, list(set(T1['Logg'][np.where( ( (T1['Temp'] == x0) & (T1['Logg'] >= Logg) ) )])))
    #print(x1, list(set(T1['Logg'][np.where( ( (T1['Temp'] == x1) & (T1['Logg'] >= Logg) ) )])))
    y0 = np.max(list(set(T1['Logg'][np.where( ( (T1['Temp'] == x0) & (T1['Logg'] <= Logg) ) )]) & set(T1['Logg'][np.where( ( (T1['Temp'] == x1) & (T1['Logg'] <= Logg) ) )])))
    y1 = np.min(list(set(T1['Logg'][np.where( ( (T1['Temp'] == x0) & (T1['Logg'] >= Logg) ) )]) & set(T1['Logg'][np.where( ( (T1['Temp'] == x1) & (T1['Logg'] >= Logg) ) )])))
    #print(y0, Logg, y1)
    #z0 = T1['pgs'][np.where( ( (T1['Temp'] == x0) | (T1['Temp'] == x1) ) & ( (T1['Logg'] == y0) | (T1['Logg'] == y1) ) & (T1['pgs'] <= PGS) )][-1]
    #z1 = T1['pgs'][np.where( ( (T1['Temp'] == x0) | (T1['Temp'] == x1) ) & ( (T1['Logg'] == y0) | (T1['Logg'] == y1) ) & (T1['pgs'] >= PGS) )][0]
    #print(x0, y0, list(set(T1['pgs'][np.where( ( (T1['Temp'] == x0) & (T1['Logg'] == y0) ) & (T1['pgs'] <= PGS))])))
    #print(x1, y1, list(set(T1['pgs'][np.where( ( (T1['Temp'] == x1) & (T1['Logg'] == y1) ) & (T1['pgs'] <= PGS))])))
    #print(x0, y0, list(set(T1['pgs'][np.where( ( (T1['Temp'] == x0) & (T1['Logg'] == y0) ) & (T1['pgs'] >= PGS))])))
    #print(x1, y1, list(set(T1['pgs'][np.where( ( (T1['Temp'] == x1) & (T1['Logg'] == y1) ) & (T1['pgs'] >= PGS))])))
    z0 = np.max(list(set(T1['Metal'][np.where( ( (T1['Temp'] == x0) & (T1['Logg'] == y0) ) & (T1['Metal'] <= Metal) )]) & set(T1['Metal'][np.where( ( (T1['Temp'] == x1) & (T1['Logg'] == y1) & (T1['Metal'] <= Metal)) )])))
    z1 = np.min(list(set(T1['Metal'][np.where( ( (T1['Temp'] == x0) & (T1['Logg'] == y0) ) & (T1['Metal'] >= Metal) )]) & set(T1['Metal'][np.where( ( (T1['Temp'] == x1) & (T1['Logg'] == y1) & (T1['Metal'] >= Metal)) )])))
    #print(z0, PGS, z1)

    # Check if the gridpoint exists within the model ranges
    for x in [x0, x1]:
        for y in [y0, y1]:
            for z in [z0, z1]:
                if (x, y, z) not in zip(T1['Temp'], T1['Logg'], T1['Metal']):
                    print('No Model', x, y, z)
                    return 1
    '''
    print(np.where( (T1['Temp'] == x1) & (T1['Logg'] == y1)))
    print(np.where( (T1['Temp'] == x1) & (T1['Logg'] == y2)))
    print(np.where( (T1['Temp'] == x2) & (T1['Logg'] == y1)))
    print(np.where( (T1['Temp'] == x2) & (T1['Logg'] == y2)))
    print(np.log10(T1['Temp'][np.where( (T1['Temp'] == x1) & (T1['Logg'] == y1))]), np.log10(T1['Logg'][np.where( (T1['Temp'] == x1) & (T1['Logg'] == y1))]))
    print(np.log10(T1['Temp'][np.where( (T1['Temp'] == x1) & (T1['Logg'] == y2))]), np.log10(T1['Logg'][np.where( (T1['Temp'] == x1) & (T1['Logg'] == y2))]))
    print(np.log10(T1['Temp'][np.where( (T1['Temp'] == x2) & (T1['Logg'] == y1))]), np.log10(T1['Logg'][np.where( (T1['Temp'] == x2) & (T1['Logg'] == y1))]))
    print(np.log10(T1['Temp'][np.where( (T1['Temp'] == x2) & (T1['Logg'] == y2))]), np.log10(T1['Logg'][np.where( (T1['Temp'] == x2) & (T1['Logg'] == y2))]))
    '''
    # Get the 16 points
    if modelset.lower() == 'btsettl08':
        if z0 == -0.5:
            Alpha0 = 0.2
        elif (z0 == -1.0) or (z0 == -1.5) or (z0 == -2.0) or (z0 == -2.5):
            Alpha0 = 0.4
        else:
            Alpha0 = 0.0
        if z1 == -0.5:
            Alpha1 = 0.2
        elif (z1 == -1.0) or (z1 == -1.5) or (z1 == -2.0) or (z1 == -2.5):
            Alpha1 = 0.4
        else:
            Alpha1 = 0.0
        ind000 = np.where( (T1['Temp'] == x0) & (T1['Logg'] == y0) & (T1['Metal'] == z0) & (T1['Alpha'] == Alpha0) ) # 000
        ind100 = np.where( (T1['Temp'] == x1) & (T1['Logg'] == y0) & (T1['Metal'] == z0) & (T1['Alpha'] == Alpha0) ) # 100
        ind010 = np.where( (T1['Temp'] == x0) & (T1['Logg'] == y1) & (T1['Metal'] == z0) & (T1['Alpha'] == Alpha0) ) # 010
        ind110 = np.where( (T1['Temp'] == x1) & (T1['Logg'] == y1) & (T1['Metal'] == z0) & (T1['Alpha'] == Alpha0) ) # 110
        ind001 = np.where( (T1['Temp'] == x0) & (T1['Logg'] == y0) & (T1['Metal'] == z1) & (T1['Alpha'] == Alpha1) ) # 001
        ind101 = np.where( (T1['Temp'] == x1) & (T1['Logg'] == y0) & (T1['Metal'] == z1) & (T1['Alpha'] == Alpha1) ) # 101
        ind011 = np.where( (T1['Temp'] == x0) & (T1['Logg'] == y1) & (T1['Metal'] == z1) & (T1['Alpha'] == Alpha1) ) # 011
        ind111 = np.where( (T1['Temp'] == x1) & (T1['Logg'] == y1) & (T1['Metal'] == z1) & (T1['Alpha'] == Alpha1) ) # 111

    else:
        ind000 = np.where( (T1['Temp'] == x0) & (T1['Logg'] == y0) & (T1['Metal'] == z0) & (T1['Alpha'] == Alpha) ) # 000
        ind100 = np.where( (T1['Temp'] == x1) & (T1['Logg'] == y0) & (T1['Metal'] == z0) & (T1['Alpha'] == Alpha) ) # 100
        ind010 = np.where( (T1['Temp'] == x0) & (T1['Logg'] == y1) & (T1['Metal'] == z0) & (T1['Alpha'] == Alpha) ) # 010
        ind110 = np.where( (T1['Temp'] == x1) & (T1['Logg'] == y1) & (T1['Metal'] == z0) & (T1['Alpha'] == Alpha) ) # 110
        ind001 = np.where( (T1['Temp'] == x0) & (T1['Logg'] == y0) & (T1['Metal'] == z1) & (T1['Alpha'] == Alpha) ) # 001
        ind101 = np.where( (T1['Temp'] == x1) & (T1['Logg'] == y0) & (T1['Metal'] == z1) & (T1['Alpha'] == Alpha) ) # 101
        ind011 = np.where( (T1['Temp'] == x0) & (T1['Logg'] == y1) & (T1['Metal'] == z1) & (T1['Alpha'] == Alpha) ) # 011
        ind111 = np.where( (T1['Temp'] == x1) & (T1['Logg'] == y1) & (T1['Metal'] == z1) & (T1['Alpha'] == Alpha) ) # 111

    Points =  [ [np.log10(T1['Temp'][ind000]), T1['Logg'][ind000], T1['Metal'][ind000],
                 np.log10(GetModel(T1['Temp'][ind000], T1['Logg'][ind000], T1['Metal'][ind000], modelset=modelset))],
                [np.log10(T1['Temp'][ind100]), T1['Logg'][ind100], T1['Metal'][ind100],
                 np.log10(GetModel(T1['Temp'][ind100], T1['Logg'][ind100], T1['Metal'][ind100], modelset=modelset))],
                [np.log10(T1['Temp'][ind010]), T1['Logg'][ind010], T1['Metal'][ind010],
                 np.log10(GetModel(T1['Temp'][ind010], T1['Logg'][ind010], T1['Metal'][ind010], modelset=modelset))],
                [np.log10(T1['Temp'][ind110]), T1['Logg'][ind110], T1['Metal'][ind110],
                 np.log10(GetModel(T1['Temp'][ind110], T1['Logg'][ind110], T1['Metal'][ind110], modelset=modelset))],
                [np.log10(T1['Temp'][ind001]), T1['Logg'][ind001], T1['Metal'][ind001],
                 np.log10(GetModel(T1['Temp'][ind001], T1['Logg'][ind001], T1['Metal'][ind001], modelset=modelset))],
                [np.log10(T1['Temp'][ind101]), T1['Logg'][ind101], T1['Metal'][ind101],
                 np.log10(GetModel(T1['Temp'][ind101], T1['Logg'][ind101], T1['Metal'][ind101], modelset=modelset))],
                [np.log10(T1['Temp'][ind011]), T1['Logg'][ind011], T1['Metal'][ind011],
                 np.log10(GetModel(T1['Temp'][ind011], T1['Logg'][ind011], T1['Metal'][ind011], modelset=modelset))],
                [np.log10(T1['Temp'][ind111]), T1['Logg'][ind111], T1['Metal'][ind111],
                 np.log10(GetModel(T1['Temp'][ind111], T1['Logg'][ind111], T1['Metal'][ind111], modelset=modelset))],
              ]
    #print(Points)
    waves2 = GetModel(T1['Temp'][ind000], T1['Logg'][ind000], T1['Metal'][ind000], wave=True, modelset=modelset)

    return waves2, trilinear_interpolation(np.log10(Teff), Logg, Metal, Points)

    ################################################################################################


import sys
import pylidc as pl
import numpy as np
import keras
import tensorflow as tf
from tqdm import tqdm
import keras
from keras import Model
from keras.backend.tensorflow_backend import set_session
from preprocessing import scan_index_split
from keras.models import load_model
import pickle
import itertools

sys.path.append('../')
from keras_retinanet import models
from keras_retinanet.utils.eval import _get_detections as get_fast_detection
from keras_retinanet.preprocessing.generator import Generator

sys.path.append('/home/toosyou/projects/LungTumor')
import data_util

FAST_DETECTION_PARM = {
    'path': 'working_models/augmented_smallanchor64.h5',
    'backbone': 'p3d',
    'convert_model': True
}

FPR_MODEL_PATH = 'fp_reduction/working_models/resnet3d_amsgrad_097.h5'
# FPR_MODEL_PATH = 'fp_reduction/model_checkpoints/95.h5'

config = tf.ConfigProto()
config.gpu_options.allow_growth = True
set_session(tf.Session(config=config))

def false_positive_reduction(volume, fast_detection, fpr_model):
    def extract_patch(volume, x, y, z, patch_size=(64, 64, 16)):
        xs = np.arange(-patch_size[0]//2, patch_size[0]//2, dtype=np.int) + x
        ys = np.arange(-patch_size[1]//2, patch_size[1]//2, dtype=np.int) + y
        zs = np.arange(-patch_size[2]//2, patch_size[2]//2, dtype=np.int) + z

        patch = volume.take( ys, mode='wrap', axis=0).take(
                                xs, mode='wrap', axis=1).take(
                                zs, mode='warp', axis=2)
        return patch

    # remove predictions outside the lung
    lung_mask = data_util.lung_mask(volume, times_dilation=20, times_erosion=15, verbose=True)
    if lung_mask.any():
        for z in np.arange(len(fast_detection)) + 8:
            tmp_detection = list()
            for d in fast_detection[z-8][0]:
                if lung_mask[int((d[1]+d[3])/2), int((d[0]+d[2])/2), z]:
                    tmp_detection.append(d)
            fast_detection[z-8][0] = tmp_detection

    # generate batches for fpr
    result = list()
    batches = list()
    for z in tqdm(np.arange(len(fast_detection)) + 8, desc='generation fpr batches'):
        for d in fast_detection[z-8][0]:
            batches.append(extract_patch(volume, int((d[0]+d[2])/2), int((d[1]+d[3])/2), int(z)))
            result.append([d[0], d[1], d[2] ,d[3], z, d[4]])

    batches = np.array(batches)[..., np.newaxis]
    batches = (batches - 418.) / 414. # normalize

    # do the prediction
    print('making fpr predictions...')
    fpr_score = fpr_model.predict(batches, 128)
    for i, fs in enumerate(fpr_score[:, 0]):
        result[i].append(fs)

    # sort by the score
    result = np.array(result)
    result = result[result[:, -1].argsort()[::-1]]

    return result

# Malisiewicz et al.
def nms_2d(dets, thresh):
    x1 = dets[:, 0]
    y1 = dets[:, 1]
    x2 = dets[:, 2]
    y2 = dets[:, 3]
    scores = dets[:, -1]

    #每一个检测框的面积
    areas = (x2 - x1 + 1) * (y2 - y1 + 1)
    #按照score置信度降序排序
    order = scores.argsort()[::-1]

    keep = [] #保留的结果框集合
    while order.size > 0:
        i = order[0]
        keep.append(i) #保留该类剩余box中得分最高的一个
        #得到相交区域,左上及右下
        xx1 = np.maximum(x1[i], x1[order[1:]])
        yy1 = np.maximum(y1[i], y1[order[1:]])
        xx2 = np.minimum(x2[i], x2[order[1:]])
        yy2 = np.minimum(y2[i], y2[order[1:]])

        #计算相交的面积,不重叠时面积为0
        w = np.maximum(0.0, xx2 - xx1 + 1)
        h = np.maximum(0.0, yy2 - yy1 + 1)
        inter = w * h
        #计算IoU：重叠面积 /（面积1+面积2-重叠面积）
        ovr = inter / (areas[i] + areas[order[1:]] - inter)
        #保留IoU小于阈值的box
        inds = np.where(ovr <= thresh)[0]
        order = order[inds + 1] #因为ovr数组的长度比order数组少一个,所以这里要将所有下标后移一位

    return dets[keep]

def group_nodule(result, nms2d_threshold=0.1, group3d_threshold=0.1):
    # group by z axis
    layer_result = np.array( [result[result[:, 4] == i] for i in range(8, int(result[:, 4].max()+1))] )

    nodules = list()
    last_boxes, last_box_nodule_index = None, None
    for i in range(layer_result.shape[0]):
        nmsed_boxes = nms_2d(layer_result[i], nms2d_threshold)
        if i == 0: # treat every predict as a nodule
            for b in nmsed_boxes:
                nodules.append([b])
            box_nodule_index = list(range(nmsed_boxes.shape[0]))
        else:
            box_nodule_index = [-1] * nmsed_boxes.shape[0]
            last_areas = (last_boxes[:, 2] - last_boxes[:, 0] + 1) * (last_boxes[:, 3] - last_boxes[:, 1] + 1)# (x2 - x1 + 1) * (y2 - y1 + 1)
            for index_box, b in enumerate(nmsed_boxes):
                # calcualte iou between this and last layer
                x1, y1, x2, y2 = b[0], b[1], b[2], b[3]
                areas = (x2 - x1 + 1) * (y2 - y1 + 1)

                xx1 = np.maximum(x1, last_boxes[:, 0])
                yy1 = np.maximum(y1, last_boxes[:, 1])
                xx2 = np.minimum(x2, last_boxes[:, 2])
                yy2 = np.minimum(y2, last_boxes[:, 3])

                #计算相交的面积,不重叠时面积为0
                w = np.maximum(0.0, xx2 - xx1 + 1)
                h = np.maximum(0.0, yy2 - yy1 + 1)
                inter = w * h
                #计算IoU：重叠面积 /（面积1+面积2-重叠面积）
                ovr = inter / (areas + last_areas- inter)

                if ovr.max() > group3d_threshold: # merge to a existing nodule
                    index_max_overlap = ovr.argmax()
                    index_nodule = last_box_nodule_index[index_max_overlap]
                    nodules[index_nodule].append(b)
                    box_nodule_index[index_box] = index_nodule
                else: # new nodule
                    box_nodule_index[index_box] = len(nodules)
                    nodules.append([b])

        last_box_nodule_index = box_nodule_index
        last_boxes = nmsed_boxes

    # calculate the nodule score
    nodule_scores = np.zeros((len(nodules), 3))
    for index_nodule, n in enumerate(nodules):
        boxes = np.array(n).reshape((-1, 7))
        nodule_scores[index_nodule][0] = boxes[:, -2].mean()
        nodule_scores[index_nodule][1] = boxes[:, -1].mean()
        nodule_scores[index_nodule][2] = np.unique(boxes[:, 4]).shape[0] # number of layers crossed

    # sort nodules by score
    sorted_index = nodule_scores[:, 1].argsort()[::-1]
    nodules = np.array(nodules)[sorted_index]
    nodule_scores = nodule_scores[sorted_index]

    return nodules, nodule_scores

class ScanGenerator:
    def __init__(self, volume):
        def preprocess_image(image):
            """ Preprocess image and its annotations.
            """
            MEAN, STD = 174., 825.
            # image = (image - image.mean()) / image.std()
            image = (image - MEAN) / STD
            return image

        self.volume = volume
        self.z_indices = list(range(16//2, self.volume.shape[2]-16//2))
        self.preprocess_image = preprocess_image

    def num_classes(self):
        """ Number of classes in the dataset.
        """
        return 1

    def name_to_label(self, name):
        """ Map name to label.
        """
        return 0

    def label_to_name(self, label):
        """ Map label to name.
        """
        return 'nodule'

    def size(self):
        return len(self.z_indices)

    def load_image(self, image_index, repeat=False):
        """ Load an image at the image_index.
        """
        z = self.z_indices[image_index]
        return self.volume[:,:,z-16//2: z+16//2].reshape(512, 512, 16, 1)

    def resize_image(self, image):
        """ Resize an image using image_min_side and image_max_side.
        """
        return image, 1.

def predict(set, index):
    index_scan = scan_index_split(1018)[{'train': 0, 'valid': 1, 'test': 2}[set]][index]
    scans = pl.query(pl.Scan).filter()

    volume = scans[index_scan].to_volume()

    # load the models
    print('Loading model, this may take a second...')
    fast_detection_model = models.load_model(
                            FAST_DETECTION_PARM['path'],
                            backbone_name=FAST_DETECTION_PARM['backbone'],
                            nms=True,
                            convert=FAST_DETECTION_PARM['convert_model'])
    # fast_detection_model.summary()

    fpr_model = load_model(FPR_MODEL_PATH)
    fpr_model = keras.utils.multi_gpu_model(fpr_model)

    fast_detection_generator = ScanGenerator(volume)
    fast_detection = get_fast_detection(
                            generator=fast_detection_generator,
                            model=fast_detection_model,
                            score_threshold=0.05,
                            max_detections=30,
                            verbose=True,
                            save_path=None,
                            do_draw_annotations=False)

    # false positive reduction
    result = false_positive_reduction(volume, fast_detection, fpr_model)

    with open('result.pl', 'wb') as f:
        pickle.dump(result, f)

    return index_scan, volume, result

if __name__ == '__main__':
    result = pickle.load(open('result.pl', 'rb'))
    nodules, nodule_scores = group_nodule(result)
    print(nodule_scores)
    # predict('test', 85)


import sys
import os
import gc
import math
import numpy as np
import cv2
import glob
import imageio
import tensorflow as tf
from PIL import Image
from random import choice
from tensorflow.keras.utils import Sequence
import matplotlib.pyplot as plt
from tensorflow.keras import backend as K
from tensorflow.keras.layers import Lambda, Input
from tensorflow.keras.models import Model
from subprocess import Popen, PIPE
from timeit import default_timer as timer
from losses import psnr2 as psnr

class DataLoader(Sequence):
    def __init__(self, datapath, batch_size, height_hr, width_hr,
         scale, crops_per_image, media_type,channels=3,colorspace='RGB',stage=None,qf=30):
        """
        :param string datapath: filepath to training images
        :param int height_hr: Height of high-resolution images
        :param int width_hr: Width of high-resolution images
        :param int height_hr: Height of low-resolution images
        :param int width_hr: Width of low-resolution images
        :param int scale: Upscaling factor
        """

        # Store the datapath
        self.datapath = datapath
        self.batch_size = batch_size
        self.height_hr = height_hr
        self.height_lr = int(height_hr / scale)
        self.width_hr = width_hr
        self.width_lr = int(width_hr / scale)
        self.channels = channels
        self.colorspace = colorspace
        self.scale = scale
        self.crops_per_image = crops_per_image
        self.media_type  = media_type
        self.time_step=1
        self.total_imgs = None
        self.stage = stage
        self.qf=qf

        # Options for resizing
        self.options = [Image.NEAREST, Image.BILINEAR, Image.BICUBIC, Image.LANCZOS]

        # Check data source
        self.img_paths = []

        if os.path.isdir(self.datapath):
            self.get_paths()

    def get_paths(self):
        for dirpath, _, filenames in os.walk(self.datapath):
            for filename in [f for f in filenames if any(filetype in f.lower() for filetype in ['jpeg', 'png', 'jpg','mp4','264','webm','wma'])]:
                self.img_paths.append(os.path.join(dirpath, filename))
        self.total_imgs = len(self.img_paths)
        print(">> Found {} images in dataset".format(self.total_imgs))

    def random_crop(self, img, random_crop_size):
        # Note: image_data_format is 'channel_last'
        assert img.shape[2] == 3
        height, width = img.shape[0], img.shape[1]
        dy, dx = random_crop_size
        x = np.random.randint(0, width - dx + 1)
        y = np.random.randint(0, height - dy + 1)
        return img[y:(y+dy), x:(x+dx), :]

    def fix_crop(self, img, dy, dx, y, x):
        return img[y:(y+dy), x:(x+dx), :]

    @staticmethod
    def scale_lr_imgs(imgs):
        """Scale low-res images prior to passing to SRGAN"""
        return imgs / 255.

    @staticmethod
    def unscale_lr_imgs(imgs):
        """Un-Scale low-res images"""
        imgs = imgs * 255
        imgs = np.clip(imgs, 0., 255.)
        return imgs.astype('uint8')

    @staticmethod
    def scale_hr_imgs(imgs):
        """Scale high-res images prior to passing to SRGAN"""
        return imgs / 255.

    @staticmethod
    def unscale_hr_imgs(imgs):
        """Un-Scale high-res images"""
        imgs = imgs * 255
        imgs = np.clip(imgs, 0., 255.)
        return imgs.astype('uint8')

    def count_frames_manual(self,cap):
        count=0
        while(cap.isOpened()):
            ret, frame = cap.read()
            if(ret):
                count +=1
            else:
                break
        return count

    def count_frames(self,cap):
        '''Count total frames in video'''
        try:
            total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        except:
            total = self.count_frames_manual(cap)
        return total

    def get_random_frames(self,n_fms=1,cap=None,time_step=1):
        """Get random number of video frames"""
        np.random.seed(int(1000000*(timer()%1)))
        t_frames = self.count_frames(cap)
        if(time_step==1):
            choiced_frames = np.random.randint(t_frames, size=n_fms)
        else:
            choiced_frames = np.random.randint(t_frames-time_step, size=n_fms)
        return choiced_frames

    @staticmethod
    def load_img(path,colorspace='YCbCr'):
        img = Image.open(path)
        if (colorspace ==  'YCbCr' and img.mode != 'YCbCr'):
            img = img.convert('YCbCr')
        if (colorspace ==  'RGB' and img.mode != 'RGB'):
            img = img.convert('RGB')
        return np.array(img)

    def load_frame(self,videopath,time_step=1,colorspace='YCbCr'):
        """Get n_imgs random frames from the video"""
        cap = cv2.VideoCapture(videopath)
        if(not cap.isOpened()):
            print("Error to open video: ",videopath)
            return -1
        choiced_frame=self.get_random_frames(1,cap,time_step)
        cap.set(1,choiced_frame)
        frames = []
        for i in list(range(time_step)):
            ret, frame = cap.read()
            if not ret:
                print(">> Erro to access frames.")
                return -1
            if colorspace=='YCbCr':
                frame = cv2.cvtColor(frame, cv2.COLOR_BGR2YCrCb)
                frames.append(frame)
            else:
                frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
        cap.release()
        return np.array(frames)

    def resize(self,shape,scale,image):

        # 3 channel images of arbitrary shape
        inp = Input(shape=shape)
        try:
            out = Lambda(lambda image: K.tf.image.resize_bicubic(image, (shape[0]//scale, shape[1]//scale)))(inp)
        except :
            # if you have older version of tensorflow
            out = Lambda(lambda image: K.tf.image.resize_bicubic(image, shape[0]//scale, shape[1]//scale))(inp)
        model = Model(inputs=inp, outputs=out)
        out = model.predict(image[np.newaxis, ...])
        return np.squeeze(out, axis=0).astype(np.uint8)

    def __len__(self):
        return int(self.total_imgs / float(self.batch_size))

    def __getitem__(self, idx):
        return self.load_batch(idx=idx)

    def compression_artefacts(self,img=None,qf=30):
        encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), qf]
        _, encimg = cv2.imencode('.jpg', img, encode_param)
        decimg = cv2.imdecode(encimg,cv2.IMREAD_COLOR)
        return decimg

    def load_batch(self,idx=0, img_paths=None, training=True, bicubic=False):
        """ Loads a batch of images or video depending on the media type """
        if(self.media_type=='i'):
            #print("1. Media type image")
            imgs_lr, imgs_hr = self.load_batch_image(idx, img_paths=img_paths, training=training,bicubic=bicubic)
        elif(self.media_type=='v' and os.path.isdir(self.datapath)):
            #print("2. Media type video folder")
            imgs_lr, imgs_hr = self.load_batch_video(idx, img_paths=None, training=training, bicubic=bicubic)
        elif(self.media_type=='v' and not os.path.isdir(self.datapath)):
            #print("3. Media type single video")
            imgs_lr, imgs_hr = self.load_batch_video(idx, img_paths=self.datapath, training=training, bicubic=bicubic)
        else:
            print("Error: type of media fail")
        return imgs_lr, imgs_hr

    def load_batch_video(self, idx=0, img_paths=None, training=True, bicubic=True):
        """Loads a batch of frames from videos folder"""
        # Starting index to look in
        cur_idx = 0
        if not img_paths:
            cur_idx = idx*self.batch_size
        #print('cur_idx:',cur_idx)

        # Scale and pre-process images
        imgs_hr, imgs_lr = [], []
        while True:

            # Check if done with batch
            if img_paths is None:
                if cur_idx >= self.total_imgs:
                    cur_idx = 0
                if len(imgs_hr) >= self.batch_size:
                    break
            if img_paths is not None and len(imgs_hr) == len(img_paths):
                break

            try:
                # Load image
                img_hr = None
                if img_paths:
                    img_hr = self.load_frame(img_paths[cur_idx])
                else:
                    img_hr = self.load_frame(self.img_paths[cur_idx])

                # Create HR images to go through
                img_crops = []
                if training:
                    for i in range(self.crops_per_image):
                        img_crops.append(self.random_crop(img_hr, (self.height_hr, self.width_hr)))
                else:
                    img_crops = [img_hr]

                # Downscale the HR images and save
                for img_hr in img_crops:

                    # TODO: Refactor this so it does not occur multiple times
                    if img_paths is None:
                        if cur_idx >= self.total_imgs:
                            cur_idx = 0
                        if len(imgs_hr) >= self.batch_size:
                            break
                    if img_paths is not None and len(imgs_hr) == len(img_paths):
                        break

                    # For LR, do bicubic downsampling
                    method = Image.BICUBIC if bicubic else choice(self.options)
                    lr_shape = (int(img_hr.shape[1]/self.scale), int(img_hr.shape[0]/self.scale))
                    img_lr = Image.fromarray(img_hr.astype(np.uint8))
                    img_lr = np.array(img_lr.resize(lr_shape, method))

                    # Scale color values
                    img_hr = self.scale_hr_imgs(img_hr)
                    img_lr = self.scale_lr_imgs(img_lr)

                    # Store images
                    imgs_hr.append(img_hr)
                    imgs_lr.append(img_lr)

            except Exception as e:
                print(e)
                pass
            finally:
                cur_idx += 1

        # Convert to numpy arrays when we are training
        # Note: all are cropped to same size, which is not the case when not training
        if training:
            imgs_hr = np.array(imgs_hr)
            imgs_lr = np.array(imgs_lr)

        # Return image batch
        return imgs_lr, imgs_hr

    def load_batch_image(self, idx=0, img_paths=None, training=True, bicubic=False):
        """Loads a batch of images from datapath folder"""

        # Starting index to look in
        cur_idx = 0

        if not img_paths:
            cur_idx = idx*self.batch_size
        # Scale and pre-process images
        imgs_hr, imgs_lr = [], []
        while True:
            # Check if done with batch
            if img_paths is None:
                if cur_idx >= self.total_imgs:
                    cur_idx = 0
                if len(imgs_hr) >= self.batch_size:
                    break
            if img_paths is not None and len(imgs_hr) == len(img_paths):
                break
            try:
                # Load image
                img_hr = None
                if img_paths:
                    img_hr = self.load_img(img_paths[cur_idx],self.colorspace)
                else:
                    img_hr = self.load_img(self.img_paths[cur_idx],self.colorspace)
                # Create HR images to go through
                img_crops = []
                if training:
                    for i in range(self.crops_per_image):
                        #print(idx, cur_idx, "Loading crop: ", i)
                        img_crops.append(self.random_crop(img_hr, (self.height_hr, self.width_hr)))
                else:
                    img_crops = [img_hr]
                # Downscale the HR images and save
                for img_hr in img_crops:

                    # TODO: Refactor this so it does not occur multiple times
                    if img_paths is None:
                        if cur_idx >= self.total_imgs:
                            cur_idx = 0
                        if len(imgs_hr) >= self.batch_size:
                            break
                    if img_paths is not None and len(imgs_hr) == len(img_paths):
                        break

                    #img_lr = cv2.resize(cv2.GaussianBlur(img_hr,(5,5),cv2.BORDER_DEFAULT),lr_shape, interpolation = cv2.INTER_CUBIC)
                    #img_lr = cv2.resize(cv2.medianBlur(img_hr,5),lr_shape, interpolation = cv2.INTER_CUBIC)

                    lr_shape = (int(img_hr.shape[1]/self.scale), int(img_hr.shape[0]/self.scale))
                    if(self.stage=='dbcnn'):
                        img_hr = cv2.resize(img_hr,lr_shape, interpolation = cv2.INTER_CUBIC)
                        #img_lr = cv2.resize(cv2.GaussianBlur(img_hr,(3,3),cv2.BORDER_DEFAULT),lr_shape, interpolation = cv2.INTER_CUBIC)
                        img_lr = cv2.resize(img_hr,lr_shape, interpolation = cv2.INTER_CUBIC)
                        img_lr = self.compression_artefacts(img_lr,qf=self.qf)
                    else:
                        #img_lr = cv2.resize(cv2.GaussianBlur(img_hr,(3,3),cv2.BORDER_DEFAULT),lr_shape, interpolation = cv2.INTER_CUBIC)
                        img_lr = cv2.resize(img_hr,lr_shape, interpolation = cv2.INTER_CUBIC)
                        img_lr = self.compression_artefacts(img_lr,qf=self.qf)
                        #img_lr = cv2.resize(self.compression_artefacts(img_hr,qf=self.qf),lr_shape, interpolation = cv2.INTER_CUBIC)

                    # Scale color values
                    img_hr = self.scale_hr_imgs(img_hr)
                    img_lr = self.scale_lr_imgs(img_lr)

                    # Store images
                    #print(img_hr[:,:,:self.channels].shape,img_lr[:,:,:self.channels].shape)
                    #imgs_hr.append(img_hr[6:-6,6:-6,:self.channels])
                    imgs_hr.append(img_hr[:,:,:self.channels])
                    imgs_lr.append(img_lr[:,:,:self.channels])

            except Exception as e:
                print(e)
                pass
            finally:
                cur_idx += 1

        # Convert to numpy arrays when we are training
        # Note: all are cropped to same size, which is not the case when not training
        if training:
            imgs_hr = np.array(imgs_hr)
            imgs_lr = np.array(imgs_lr)

        # Return image batch
        return imgs_lr, imgs_hr

def plot_test_images(model, loader, datapath_test, test_output, epoch, name='SRGAN', channels = 3,colorspace='RGB'):

    try:
        # Get the location of test images
        test_images = [os.path.join(datapath_test, f) for f in os.listdir(datapath_test) if any(filetype in f.lower() for filetype in ['jpeg','mp4','264', 'png', 'jpg'])]

        # Load the images to perform test on images
        imgs_lr, imgs_hr = loader.load_batch(img_paths=test_images, training=False, bicubic=True)
        # Create super resolution and bicubic interpolation images
        imgs_sr = []
        imgs_bi = []
        srgan_psnr = []
        bi_psnr = []
        for i in range(len(test_images)):

            pre=np.squeeze(
                    model.predict(
                        np.expand_dims(imgs_lr[i], 0),
                        batch_size=1
                    ),
                    axis=0
                )
            imgs_sr.append(pre)

        # Unscale colors values
        if channels == 1:
            imgs_lr = [loader.unscale_lr_imgs(img[:,:,0]).astype(np.uint8) for img in imgs_lr]
            imgs_hr = [loader.unscale_hr_imgs(img[:,:,0]).astype(np.uint8) for img in imgs_hr]
            imgs_sr = [loader.unscale_hr_imgs(img[:,:,0]).astype(np.uint8) for img in imgs_sr]
        else:
            if(colorspace == 'YCbCr'):
                imgs_lr = [cv2.cvtColor(loader.unscale_lr_imgs(img).astype(np.uint8), cv2.COLOR_YCrCb2BGR) for img in imgs_lr]
                imgs_hr = [cv2.cvtColor(loader.unscale_hr_imgs(img).astype(np.uint8), cv2.COLOR_YCrCb2BGR) for img in imgs_hr]
                imgs_sr = [cv2.cvtColor(loader.unscale_hr_imgs(img).astype(np.uint8), cv2.COLOR_YCrCb2BGR) for img in imgs_sr]

            else:
                imgs_lr = [loader.unscale_lr_imgs(img).astype(np.uint8) for img in imgs_lr]
                imgs_hr = [loader.unscale_hr_imgs(img).astype(np.uint8) for img in imgs_hr]
                imgs_sr = [loader.unscale_hr_imgs(img).astype(np.uint8) for img in imgs_sr]

        # Loop through images
        for img_hr, img_lr, img_sr, img_path in zip(imgs_hr, imgs_lr, imgs_sr, test_images):
            # Get the filename
            filename = os.path.basename(img_path).split(".")[0]

            # Bicubic upscale
            hr_shape = (int(img_hr.shape[1]), int(img_hr.shape[0]))
            img_bi = cv2.resize(img_lr,hr_shape, interpolation = cv2.INTER_CUBIC)

            # Images and titles
            images = {
                'Low Resoluiton': [img_lr, img_hr],
                'Bicubic': [img_bi, img_hr],
                name: [img_sr, img_hr],
                'Original': [img_hr,img_hr]
            }
            srgan_psnr.append(psnr(img_sr,img_hr,255.))
            bi_psnr.append(psnr(img_bi,img_hr,255.))
            # Plot the images. Note: rescaling and using squeeze since we are getting batches of size 1
            fig, axes = plt.subplots(1, 4, figsize=(40, 10))
            for i, (title, img) in enumerate(images.items()):
                axes[i].imshow(img[0])
                axes[i].set_title("{} - {} {}".format(title, img[0].shape, ("- psnr: "+str(round(psnr(img[0],img[1],255.),2)) if (title == name or title == 'Bicubic' ) else " ")))
                #axes[i].set_title("{} - {}".format(title, img.shape))
                axes[i].axis('off')
            plt.suptitle('{} - Epoch: {}'.format(filename, epoch))

            # Save directory
            savefile = os.path.join(test_output, "{}-Epoch{}.png".format(filename, epoch))
            fig.savefig(savefile)
            plt.close()
            gc.collect()
        print('test {} psnr: {} - test bi psnr: {}'.format(name,np.mean(srgan_psnr),np.mean(bi_psnr)))
    except Exception as e:
        print(e)


from collections import OrderedDict

import numpy as np

from env.ant.ant import AntEnv
from env.transform_utils import up_vector_from_quat, forward_vector_from_quat, \
    l2_dist, cos_dist, sample_quat

class AntPushEnv(AntEnv):
    def __init__(self, **kwargs):
        self.name = 'ant_push'
        super().__init__('ant_push.xml', **kwargs)

        # Env info
        #self.ob_shape = OrderedDict([("ant_1_shared_pos", 19), ("ant_1_lower_body", 24), ("ant_2_shared_pos", 19), ("ant_2_lower_body", 24), ("box_pos", 3)])
        self.ob_shape = OrderedDict([("ant_1", 41), ("ant_2", 41),
                                     ("box_1", 6), ("box_2", 6),
                                     ("goal_1", 3), ("goal_2", 3)])
        self.action_space.decompose(OrderedDict([("ant_1", 8), ("ant_2", 8)]))

        # Env config
        self._env_config.update({
            'random_ant_pos': 0.01,
            'random_goal_pos': 0.01,
            #'random_goal_pos': 0.5,
            'dist_threshold': 0.5,
            'quat_threshold': 0.9,
            'ant_dist_reward': 50,
            'goal_dist_reward': 200,
            'goal1_dist_reward': 2,
            'goal2_dist_reward': 2,
            'quat_reward': 10, # quat_dist usually between 0.95 ~ 1
            'height_reward': 0.5,
            'upright_reward': 0.5,
            'alive_reward': 0.,
            'success_reward': 500,
            'die_penalty': 10,
            'sparse_reward': 0,
            'init_randomness': 0.01,
            #'max_episode_steps': 400,
            'max_episode_steps': 200,
        })
        self._env_config.update({ k:v for k,v in kwargs.items() if k in self._env_config })

        self._ant1_push = False
        self._ant2_push = False
        self._success_count = 0

    def _step(self, a):
        ant1_pos_before = self._get_pos('ant_1_torso_geom')
        ant2_pos_before = self._get_pos('ant_2_torso_geom')
        box1_pos_before = self._get_pos('box_geom1')
        box2_pos_before = self._get_pos('box_geom2')
        box_pos_before = self._get_pos('box')
        box_quat_before = self._get_quat('box')

        self.do_simulation(a)

        ant1_pos = self._get_pos('ant_1_torso_geom')
        ant2_pos = self._get_pos('ant_2_torso_geom')
        box1_pos = self._get_pos('box_geom1')
        box2_pos = self._get_pos('box_geom2')
        box_pos = self._get_pos('box')
        box_quat = self._get_quat('box')
        goal1_pos = self._get_pos('goal_geom1')
        goal2_pos = self._get_pos('goal_geom2')
        goal_pos = self._get_pos('goal')
        goal_quat = self._get_quat('goal')

        ob = self._get_obs()
        done = False
        ctrl_reward = self._ctrl_reward(a)

        goal_forward = forward_vector_from_quat(goal_quat)
        box_forward = forward_vector_from_quat(box_quat)
        box_forward_before = forward_vector_from_quat(box_quat_before)

        goal1_dist = l2_dist(goal1_pos, box1_pos)
        goal1_dist_before = l2_dist(goal1_pos, box1_pos_before)
        goal2_dist = l2_dist(goal2_pos, box2_pos)
        goal2_dist_before = l2_dist(goal2_pos, box2_pos_before)
        goal_dist = l2_dist(goal_pos, box_pos)
        goal_dist_before = l2_dist(goal_pos, box_pos_before)
        goal_quat = cos_dist(goal_forward, box_forward)
        goal_quat_before = cos_dist(goal_forward, box_forward_before)

        ant1_height = self.data.qpos[2]
        ant2_height = self.data.qpos[17]
        ant1_quat = self._get_quat('ant_1_torso')
        ant1_up = up_vector_from_quat(ant1_quat)
        ant1_upright = cos_dist(ant1_up, np.array([0, 0, 1]))
        ant2_quat = self._get_quat('ant_2_torso')
        ant2_up = up_vector_from_quat(ant2_quat)
        ant2_upright = cos_dist(ant2_up, np.array([0, 0, 1]))

        ant1_dist = l2_dist(ant1_pos_before[:2], box1_pos_before[:2]) - l2_dist(ant1_pos[:2], box1_pos[:2])
        ant2_dist = l2_dist(ant2_pos_before[:2], box2_pos_before[:2]) - l2_dist(ant2_pos[:2], box2_pos[:2])

        if l2_dist(ant1_pos[:2], box1_pos[:2]) < 2.0: self._ant1_push = True
        if l2_dist(ant2_pos[:2], box2_pos[:2]) < 2.0: self._ant2_push = True

        success_reward = 0
        ant1_dist_reward = self._env_config["ant_dist_reward"] * ant1_dist if not self._ant1_push else 0
        ant2_dist_reward = self._env_config["ant_dist_reward"] * ant2_dist if not self._ant2_push else 0
        #goal1_dist_reward = self._env_config["goal1_dist_reward"] * (goal1_dist_before - goal1_dist)
        #goal2_dist_reward = self._env_config["goal2_dist_reward"] * (goal2_dist_before - goal2_dist)
        goal1_dist_reward = self._env_config["goal1_dist_reward"] * 1 if goal1_dist < self._env_config["dist_threshold"] else 0
        goal2_dist_reward = self._env_config["goal2_dist_reward"] * 1 if goal2_dist < self._env_config["dist_threshold"] else 0
        goal_dist_reward = self._env_config["goal_dist_reward"] * (goal_dist_before - goal_dist)
        quat_reward = -self._env_config["quat_reward"] * (1 - goal_quat)
        #quat_reward = self._env_config["quat_reward"] * (goal_quat - goal_quat_before) if goal_pos[0] >= box_pos[0] else 0
        alive_reward = self._env_config["alive_reward"]
        ant1_height_reward = -self._env_config["height_reward"] * np.abs(ant1_height - 0.6)
        ant2_height_reward = -self._env_config["height_reward"] * np.abs(ant2_height - 0.6)
        ant1_upright_reward = self._env_config["upright_reward"] * ant1_upright
        ant2_upright_reward = self._env_config["upright_reward"] * ant2_upright

        if goal1_dist < self._env_config["dist_threshold"] and \
                goal2_dist < self._env_config["dist_threshold"] and \
                goal_quat > self._env_config["quat_threshold"]:
            self._success_count += 1
            success_reward = 10
            if self._success_count >= 1:
                self._success = True
                success_reward = self._env_config["success_reward"]

        # fail
        done = not np.isfinite(self.data.qpos).all() or \
            0.35 > ant1_height or ant1_height > 0.9 or \
            0.35 > ant2_height or ant2_height > 0.9
        #done = done or abs(ant1_pos[0]) > 7 or abs(ant2_pos[0]) > 7
        #done = done or abs(ant1_pos[1]) > 6 or abs(ant2_pos[1]) > 6
        die_penalty = -self._env_config["die_penalty"] if done else 0
        done = done or self._success

        if self._env_config['sparse_reward']:
            self._reward = reward = self._success == 1
        else:
            self._reward = reward = ant1_dist_reward + ant2_dist_reward + \
                ant1_height_reward + ant2_height_reward + ant1_upright_reward + ant2_upright_reward + \
                goal_dist_reward + goal1_dist_reward + goal2_dist_reward + quat_reward + \
                alive_reward + success_reward + ctrl_reward + die_penalty

        info = {"success": self._success,
                "reward_ant1_dist": ant1_dist_reward,
                "reward_ant2_dist": ant2_dist_reward,
                "reward_ant1_height": ant1_height_reward,
                "reward_ant2_height": ant2_height_reward,
                "reward_ant1_upright": ant1_upright_reward,
                "reward_ant2_upright": ant2_upright_reward,
                "reward_goal_dist": goal_dist_reward,
                "reward_goal1_dist": goal1_dist_reward,
                "reward_goal2_dist": goal2_dist_reward,
                "reward_quat": quat_reward,
                "reward_alive": alive_reward,
                "reward_ctrl": ctrl_reward,
                "die_penalty": die_penalty,
                "reward_success": success_reward,
                "ant1_pos": ant1_pos,
                "ant2_pos": ant2_pos,
                "ant1_up": ant1_up,
                "ant2_up": ant2_up,
                "goal_quat": goal_quat,
                "goal_dist": goal_dist,
                "goal1_dist": goal1_dist,
                "goal2_dist": goal2_dist,
                "box1_ob": ob['box_1'],
                "box2_ob": ob['box_2'],
                "goal1_ob": ob['goal_1'],
                "goal2_ob": ob['goal_2'],
                }

        return ob, reward, done, info

    def _get_obs(self):
        # ant
        qpos = self.data.qpos
        qvel = self.data.qvel
        qacc = self.data.qacc
        ant_pos1 = self._get_pos('ant_1_torso')
        ant_pos2 = self._get_pos('ant_2_torso')

        # box
        box_pos1 = self._get_pos('box_geom1')
        box_pos2 = self._get_pos('box_geom2')
        #box_quat1 = self._get_quat('box')
        #box_quat2 = self._get_quat('box')
        box_forward1 = self._get_forward_vector('box_geom1')
        box_forward2 = self._get_forward_vector('box_geom2')

        # goal
        goal_pos1 = self._get_pos('goal_geom1')
        goal_pos2 = self._get_pos('goal_geom2')

        obs = OrderedDict([
            #('ant_1_shared_pos', np.concatenate([qpos[:7], qvel[:6], qacc[:6]])),
            #('ant_1_lower_body', np.concatenate([qpos[7:15], qvel[6:14], qacc[6:14]])),
            #('ant_2_shared_pos', np.concatenate([qpos[15:22], qvel[14:22], qacc[14:22]])),
            #('ant_2_lower_body', np.concatenate([qpos[22:30], qvel[22:30], qacc[22:30]])),
            ('ant_1', np.concatenate([qpos[2:15], qvel[:14], qacc[:14]])),
            ('ant_2', np.concatenate([qpos[17:30], qvel[14:28], qacc[14:28]])),
            ('box_1', np.concatenate([box_pos1 - ant_pos1, box_forward1])),
            ('box_2', np.concatenate([box_pos2 - ant_pos2, box_forward2])),
            ('goal_1', goal_pos1 - box_pos1),
            ('goal_2', goal_pos2 - box_pos2),
        ])

        def ravel(x):
            obs[x] = obs[x].ravel()
        map(ravel, obs.keys())

        return obs

    @property
    def _init_qpos(self):
        # 3 for (x, y, z), 4 for (x, y, z, w), and 2 for each leg
        return np.array([0, 0., 0.58, 1., 0., 0., 0., 0., 1., 0., -1., 0., -1., 0., 1.,
                         0, 0., 0.58, 1., 0., 0., 0., 0., 1., 0., -1., 0., -1., 0., 1.,
                         0., 0., 0.8, 1., 0., 0., 0.])

    @property
    def _init_qvel(self):
        return np.zeros(self.model.nv)

    def _reset(self):
        qpos = self._init_qpos + self._init_random(self.model.nq)
        qvel = self._init_qvel + self._init_random(self.model.nv)

        qpos[2] = 0.58
        qpos[17] = 0.58

        self.set_state(qpos, qvel)

        self._reset_ant_box()

        self._ant1_push = False
        self._ant2_push = False
        self._success_count = 0

        return self._get_obs()

    def _reset_ant_box(self):
        qpos = self.data.qpos.ravel().copy()
        qvel = self.data.qvel.ravel().copy()

        # Initialize box
        init_box_pos = np.asarray([0, 0, 0.8])
        #init_box_quat = sample_quat(low=-np.pi/32, high=np.pi/32)
        init_box_quat = sample_quat(low=0, high=0)
        qpos[30:33] = init_box_pos
        qpos[33:37] = init_box_quat

        # Initialize ant
        #qpos[0:2] = [-4, 2] + np.random.uniform(-1, 1, size=(2,)) * self._env_config["random_ant_pos"]
        #qpos[15:17] = [-4, -2] + np.random.uniform(-1, 1, size=(2,)) * self._env_config["random_ant_pos"]
        qpos[0:2] = [-2, 2] + np.random.uniform(-1, 1, size=(2,)) * self._env_config["random_ant_pos"]
        qpos[15:17] = [-2, -2] + np.random.uniform(-1, 1, size=(2,)) * self._env_config["random_ant_pos"]
        #qpos[0:2] = [-2 + np.random.uniform(-1, 1) * self._env_config["random_ant_pos"], 2]
        #qpos[15:17] = [-2 + np.random.uniform(-1, 1) * self._env_config["random_ant_pos"], -2]

        # Initialize goal
        #x = 4.5 + np.random.uniform(-1, 1) * self._env_config["random_goal_pos"]
        x = 3 + np.random.uniform(-1, 1) * self._env_config["random_goal_pos"]
        y = 0 + np.random.uniform(-1, 1) * self._env_config["random_goal_pos"]
        z = 0.8
        goal_pos = np.asarray([x, y, z])
        #goal_quat = sample_quat()
        self._set_pos('goal', goal_pos)
        #self._set_quat('goal', goal_quat)

        self.set_state(qpos, qvel)

    def _render_callback(self):
        lookat = [0.5, 0, 0]
        cam_pos = lookat + np.array([-4.5, -12, 10])

        cam_id = self._camera_id
        self._set_camera_position(cam_id, cam_pos)
        self._set_camera_rotation(cam_id, lookat)

        self.sim.forward()


import serial
import time
import io
import re
import cv2
import tensorflow as tf
import numpy as np
#loading model
from tensorflow import keras
from keras.models import model_from_json

# input: [(x,y),(x,y),(x,y)]

COM_printer = "COM3"
COM_touch = "COM4"

#set this Path right !!!!
path= "C:\\Users\\Jan Lukas\\Desktop\\BioNx7\\model"
print("This path is used to locate the model: ", path)

def object_detection():
    model = keras.models.load_model(path)
    img_name ="image.png"
    model.summary()

    #choose right camera either 0 or 1
    cam = cv2.VideoCapture(1)
    img_counter = 0
    cv2.namedWindow("test")
    while True:
        ret, frame = cam.read()
        if not ret:
            print("failed to grab frame")
            break
        cv2.imshow("test", frame)

        k = cv2.waitKey(1)
        if k%256 == 27:
            # ESC pressed
            print("Escape hit, closing...")
            break
        elif k%256 == 32:
            # SPACE pressed
            image = frame
            img_counter = img_counter+1
            cv2.imwrite(img_name,image)
            break

    cam.release()

    from PIL import Image

    def crop_Images(path):
      # Opens a image in RGB mode

      import cv2
      import os
      import glob
      img = Image.open(img_name)
      width, height = img.size

      # Setting the points for cropped image
      left = 240
      top = 100
      right = width
      bottom =  height

      # Cropped image of above dimension
      # (It will not change orginal image)
      im1 = img.crop((left, top, right, bottom))
      im1.save(img_name)

    #testing

    image = cv2.imread(img_name)
    import matplotlib.pyplot as plt
    plt.imshow(image.astype("uint8"))
    plt.show()

    crop_Images(img_name)

    image = cv2.imread(img_name)

    import matplotlib.pyplot as plt
    plt.imshow(image.astype("uint8"))
    plt.show()

    image_size = (400, 380)

    img = keras.preprocessing.image.load_img(
          img_name, target_size=image_size
      )
    img_array = keras.preprocessing.image.img_to_array(img)
    percentages = model.predict(tf.expand_dims(img_array, 0))

    # 0-> 6er Wellplate | 1 -> 12er Wellplate | 2-> 96er Wellplate
    print(percentages)
    well_type= np.argmax(percentages)
    print(well_type)

    #TODO: weiterer Messpunkt obere Kante hinzufügen am Ende
    # Well_1 & Well_2 umrechnen von Relativ zu Absolut
    well_0 = [[-2.35,2.25],[-6.35,2.25],[-10.35,2.25],[-2.35,6.25],[-6.35,6.25],[-10.35,6.25]]
    well_1 = [[-2.35,1.55],[-2.7,0],[-2.7,0],[-2.7,0],[0,2.7],[2.7,0],[2.7,0],[2.7,0],[0,2.7],[-2.7,0],[-2.7,0],[-2.7,0]]
    well_2 = [[-1.4,1.1],[-0.9,0],[-0.9,0],[-0.9,0],[-0.9,0],[-0.9,0],[-0.9,0],[-0.9,0],[-0.9,0],[-0.9,0],[-0.9,0],[-0.9,0],[0,0.9],[0.9,0],[0.9,0],[0.9,0],[0.9,0],[0.9,0],[0.9,0],[0.9,0],[0.9,0],[0.9,0],[0.9,0],[0.9,0], [0,0.9],[-0.9,0],[-0.9,0],[-0.9,0],[-0.9,0],[-0.9,0],[-0.9,0],[-0.9,0],[-0.9,0],[-0.9,0],[-0.9,0],[-0.9,0],[0,0.9],[0.9,0],[0.9,0],[0.9,0],[0.9,0],[0.9,0],[0.9,0],[0.9,0],[0.9,0],[0.9,0],[0.9,0],[0.9,0], [0,0.9],[-0.9,0],[-0.9,0],[-0.9,0],[-0.9,0],[-0.9,0],[-0.9,0],[-0.9,0],[-0.9,0],[-0.9,0],[-0.9,0],[-0.9,0],[0,0.9],[0.9,0],[0.9,0],[0.9,0],[0.9,0],[0.9,0],[0.9,0],[0.9,0],[0.9,0],[0.9,0],[0.9,0],[0.9,0],[0,0.9],[-0.9,0],[-0.9,0],[-0.9,0],[-0.9,0],[-0.9,0],[-0.9,0],[-0.9,0],[-0.9,0],[-0.9,0],[-0.9,0],[-0.9,0]]
    array = []

    if(well_type==0):
        array = [[j*10 for j in i] for i in well_0]
    elif(well_type==1):
        array = [[j*10 for j in i] for i in well_1]
    else:
        array = [[j*10 for j in i] for i in well_2]

    print(array)
    return array

class PrinterControll():
    def __init__(self):

        #init Ports
        self.ser = serial.Serial(COM_printer, baudrate = 250000, timeout = 5)
        self.ser_touch = serial.Serial(COM_touch, baudrate = 9600, timeout = 5)
        self.homing()
        self.ser.write(b"G0 Y100 Z180\r\n")
        #object detection
        self.printer_script = object_detection()

    def homing(self):
        print("homing")
        time.sleep(2)
        self.ser_touch.write(b"up\r\n")
        self.ser.write(b"G28\r\n")
        time.sleep(30)
        pos = self.getposition()
        print(pos)
        self.ser.close()
        self.ser = serial.Serial(COM_printer, baudrate = 250000, timeout = 5)
        time.sleep(2)

    def readmove(self, value):
        return "G1 X" + str(value[0]) + " Y" + str(value[1]) + "\r\n"

    def start_run(self):
        print("start run")
        positions = []
        for pos in self.printer_script:
            self.ser.write(b"G1 Z40\r\n")
            gcode = ("G1 X" + str(155.0+pos[0]) + " Y" + str(38.0+pos[1]) + " Z40\r\n").encode("UTF-8")
            print(gcode)
            time.sleep(2)
            self.ser.write(gcode)
            pos = self.measure()
            positions.append(pos)
            print(pos)
        return positions
            #return

    def measure(self):
        print("measuring")
        self.ser_touch.write(b"down\r\n")
        self.ser.write(b"G1 Z0\r\n")
        while 1:
            data = self.ser_touch.readline().decode("ascii")
            if data == "high\r\n":
                self.ser_touch.write(b"up\r\n")
                self.ser_touch.write(b"stop\r\n")
                pos = self.getposition()
                self.ser_touch.write(b"release\r\n")
                #time.sleep(1) #ggf remove
                #self.ser.close()
                #self.ser = serial.Serial(COM_printer, baudrate = 250000, timeout = 5)
                #time.sleep(2)
                #gcode_actual_pos = ("G91 X" + str(pos[0]) + " Y" + str(pos[1]) + " Z" + str(pos[2])).encode("UTF-8")
                ##self.ser.write(gcode_actual_pos)
                gcode = ("G1 Z" + str(40-pos[2]) + "\r\n").encode("UTF-8")
                self.ser.write(gcode)
                self.homing()
                return pos

    def getposition(self):
        self.ser.write(b"M114\r\n")
        for x in range(20):
            res = self.ser.readline().decode("ascii")
            if res[0] == "X":
                pos = self.getpositionvalues(res)
                return pos
        return "mistake"

    def getpositionvalues(self, position_raw):
        reg = '[XYZE]'
        result = re.finditer(reg, position_raw)
        pos = []
        for match in result:
            pos.append(match.span())
            #print(match.group())
            #print(match.span())
        posX = pos[4][0]
        posY = pos[5][0]
        posZ = pos[6][0]
        #posE = pos[3][0]
        return [float(position_raw[posX+2:posY]),float(position_raw[posY+2:posZ]),float(position_raw[posZ+2:len(position_raw)-1])]

printerControll = PrinterControll()
print(printerControll.start_run())


# -*- coding: utf-8 -*-
"""
Created on Thu Jun 17 09:45:00 2021

@author: Narmin Ghaffari Laleh
"""

import random
import numpy as np
import pandas as pd
import os
import argparse
from Survival_Utils import CreateCleanedTable, GetTiles, Predictor
from tensorflow.keras.applications.resnet50 import ResNet50
from tensorflow.keras.layers import Dense, GlobalAveragePooling2D
from tensorflow.keras.models import Model
from sksurv.metrics import concordance_index_censored
import gc
from keras_preprocessing.image import ImageDataGenerator
from tqdm import tqdm

##############################################################################

parser = argparse.ArgumentParser(description = 'Survival Prediction')
parser.add_argument('--datadir_test', type = str, default = r'')
parser.add_argument('--slide_dir', type = str, default = r"")
parser.add_argument('--clini_dir', type = str, default = r"")

parser.add_argument('--outputPath', type = str, default = r"")

parser.add_argument('--modelPath', type = str, default = r"")

parser.add_argument('--num_classes', type = int, default = 1)
parser.add_argument('--batch_size', type = int, default ="64")
parser.add_argument('--maxBlockNum', type = int, default ="150")

parser.add_argument('--gpuNo', type = int, default = 1)

parser.add_argument('--evenCol', type = str, default = "")
parser.add_argument('--timeCol', type = str, default = "")
parser.add_argument('--timeInDays', type = bool, default = False)
parser.add_argument('--patientCol', type = str, default = "")

##############################################################################

if __name__ == '__main__':

    args = parser.parse_args()
    os.environ["CUDA_VISIBLE_DEVICES"]=str(args.gpuNo)

    random.seed(23)
    args.projectFolder = args.outputPath

    if not os.path.exists(args.projectFolder):
        os.mkdir(args.projectFolder)

    cleanedTable, _ = CreateCleanedTable(imagesPath = args.datadir_test, cliniTablePath = args.clini_dir, slideTablePath = args.slide_dir,
                                         targetNames = [args.evenCol, args.timeCol],outputPath = args.outputPath)

    cleanedTable.dropna(subset = [args.evenCol], inplace = True)
    cleanedTable.dropna(subset = [args.timeCol], inplace = True)
    cleanedTable = cleanedTable[~cleanedTable[args.timeCol].isin(['na', 0])]
    cleanedTable = cleanedTable[~cleanedTable[args.evenCol].isin(['na'])]
    patientID = list(cleanedTable[args.patientCol].unique())

    time = list(cleanedTable[args.timeCol])
    time = [float(i) for i in time]

    if args.timeInDays:
        time = [i / 30 for i in time]

    event = list(cleanedTable[args.evenCol])
    if 'yes' in event:
        event = [1 if i == 'yes' else 0 for i in event]
    else:
        event = [int(i) for i in event]

    patientID = np.asarray(cleanedTable[args.patientCol].unique())
    time = np.asarray(time)
    event = np.asarray(event)

    args.split_dir = os.path.join(args.projectFolder, 'SPLITS')
    if not os.path.exists(args.split_dir):
        os.mkdir(args.split_dir)

    args.results = os.path.join(args.projectFolder, 'RESULTS')
    if not os.path.exists(args.results):
        os.mkdir(args.results)

    test_data = GetTiles(patients = patientID, time = time, event = event, imgsList = args.datadir_test,
                         cleanedTable = cleanedTable, maxBlockNum = args.maxBlockNum, test = True)

    test_x = list(test_data['TILEPATH'])
    test_time = list(test_data['TIME'])
    test_event = list(test_data['EVENT'])
    test_pid = list(test_data['PATIENT'])

    df = pd.DataFrame(list(zip(test_x, test_time, test_event)), columns =['TILEPATH', 'TIME', 'EVENT'])
    df.to_csv(os.path.join(args.split_dir, 'SPLIT_TEST_TOTAL.csv'), index = False)

    base_model = ResNet50(weights='imagenet', include_top=False)

    x = base_model.output
    x = GlobalAveragePooling2D()(x)
    x = Dense(1024, activation='relu')(x)
    predictions = Dense(args.num_classes, activation='linear')(x)

    model = Model(inputs = base_model.input, outputs=predictions)

    temp_df = pd.DataFrame(test_x)

    test_datagen = ImageDataGenerator()

    test_generator=test_datagen.flow_from_dataframe(
                                dataframe=temp_df,
                                x_col = 0,
                                batch_size = args.batch_size,
                                seed = 42,
                                shuffle = False,
                                class_mode = None,
                                target_size = (224,224))

    test_time = np.asarray(test_time)
    test_event = np.asarray(test_event)

    predictor = Predictor(model, args.modelPath)
    sample_predictions = []

    #ds = tf.data.Dataset.from_tensor_slices(test_x).batch(64)
    #it = tf.compat.v1.data.make_one_shot_iterator(ds)
    batches = 0
    for tensor in tqdm(test_generator):
        sample_predictions.extend(predictor.predict(tensor))
        gc.collect()
        batches += 1
        if batches >= len(test_generator):
            # we need to break the loop by hand because
            # the generator loops indefinitely
            break
        #tensor = it.get_next()

    df = pd.DataFrame(list(zip(test_pid, test_x, list(sample_predictions), test_time, test_event)),columns =['PATIENT', 'TILEPATH', 'SCORE', 'TIME', 'EVENT'])
    df.to_csv(os.path.join(args.results, 'TEST_RESULTS_TILE_BASED_FULL.csv'), index = False)
    scores = []
    patients = []
    t = []
    e = []

    patientID_unique = set(test_pid)
    for pi in patientID_unique:
        patients.append(pi)
        data_temp = df.loc[df['PATIENT'] == pi]
        data_temp = data_temp.reset_index()
        score = np.mean(data_temp['SCORE'])[0]
        scores.append(score)
        t.append(data_temp['TIME'][0])
        e.append(data_temp['EVENT'][0])

    label = []
    hazard = []
    surv_time = []
    for i in range(len(scores)):
        if not np.isnan(scores[i]):
            label.append(e[i])
            hazard.append(scores[i])
            surv_time.append(t[i])

    new_label = np.asarray(label)
    new_hazard = np.asarray(hazard)
    new_surv = np.asarray(surv_time)

    cindex = concordance_index_censored(
      event_indicator=new_label == 1,
      event_time=new_surv,
      estimate=new_hazard)

    df = pd.DataFrame(list(zip(patients, t, e, scores)), columns = ['PATIENT', 'TIME', 'EVENT', 'SCORE'])
    df.to_csv(os.path.join(args.results, 'TEST_RESULTS_PATIENT_FULL.csv'), index = False)

    print('\n')
    print('#################################################################')
    print('Patient-Wise Concordance_Index For Test Data Set {}:'.format(cindex))
    print('#################################################################')