from __future__ import print_function, division

"""
.. note::
         These are the plotting functions for the SPLAT code
"""

# imports: internal
import copy
import glob
import os

# imports: external
import matplotlib.cm as cm
import matplotlib.colors as colmap
import matplotlib.patches as patches
import matplotlib.pyplot as plt
from matplotlib.backends.backend_pdf import PdfPages
import numpy
from scipy.interpolate import interp1d
from scipy import ndimage

# imports: splat
from .initialize import *
from .utilities import *
from .core import Spectrum, classifyByStandard

def plotSpectrum(*args, **kwargs):
    '''
    :Purpose: ``Primary plotting program for Spectrum objects.``

    :Input
    Spectrum objects, either sequentially, in list, or in list of lists
            - Spec1, Spec2, ...: plot multiple spectra together, or separately if multiplot = True
            - [Spec1, Spec2, ...]: plot multiple spectra together, or separately if multiplot = True
            - [[Spec1, Spec2], [Spec3, Spec4], ..]: plot multiple sets of spectra (multiplot forced to be True)

    :Parameters:
    title = ''
        string giving plot title  - NOTE: THIS IS NOT WORKING
    xrange = [0.85,2.42]:
        plot range for wavelength axis
    yrange = [-0.02,1.2]*fluxMax:
        plot range for wavelength axis
    xlabel:
        wavelength axis label; by default set by wlabel and wunit keywords in first spectrum object
    ylabel:
        flux axis label; by default set by fscale, flabel and funit keywords in first spectrum object

    features:
        a list of strings indicating chemical features to label on the spectra
        options include H2O, CH4, CO, TiO, VO, FeH, H2, HI, KI, NaI, SB (for spectral binary)
    mdwarf, ldwarf, tdwarf, young, binary = False:
        add in features characteristic of these classes
    telluric = False:
        mark telluric absorption features
    legend, legends, label or labels:
        list of strings providing legend-style labels for each spectrum plotted
    legendLocation or labelLocation = 'upper right':
        place of legend; options are 'upper left', 'center middle', 'lower right' (variations thereof) and 'outside'
    legendfontscale = 1:
        sets the scale factor for the legend fontsize (defaults to fontscale)
    grid = False:
        add a grid

    stack = 0:
        set to a numerical offset to stack spectra on top of each other
    zeropoint = [0,...]:
        list of offsets for each spectrum, giving finer control than stack
    showZero = True:
        plot the zeropoint(s) of the spectra

    comparison:
        a comparison Spectrum to compare in each plot, useful for common reference standard
    noise, showNoise or uncertainty = False:
        plot the uncertainty for each spectrum
    residual = False:
        plots the residual between two spectra

    color or colors:
        color of plot lines; by default all black
    colorUnc or colorsUnc:
        color of uncertainty lines; by default same as line color but reduced opacity
    colorScheme or colorMap:
        color map to apply based on matplotlib colormaps;
        see http://matplotlib.org/api/pyplot_summary.html?highlight=colormaps#matplotlib.pyplot.colormaps
    linestyle:
        line style of plot lines; by default all solid
    fontscale = 1:
        sets a scale factor for the fontsize

    inset = False:
        place an inset panel showing a close up region of the spectral data
    inset_xrange = False:
        wavelength range for inset panel
    inset_position = [0.65,0.60,0.20,0.20]
        position of inset planet in normalized units, in order left, bottom, width, height
    inset_features = False
        list of features to label in inset plot

    file or filename or output:
        filename or filename base for output
    filetype = 'pdf':
        output filetype, generally determined from filename
    multiplot = False:
        creates multiple plots, depending on format of input (optional)
    multipage = False:
        spreads plots across multiple pages; output file format must be PDF
        if not set and plots span multiple pages, these pages are output sequentially as separate files
    layout or multilayout = [1,1]:
        defines how multiple plots are laid out on a page
    figsize:
        set the figure size; set to default size if not indicated
    interactive = False:
        if plotting to window, set this to make window interactive

    :Example 1: A simple view of a random spectrum
       >>> import splat
       >>> spc = splat.getSpectrum(spt = 'T5', lucky=True)[0]
       >>> spc.plot()                       # this automatically generates a "quicklook" plot
       >>> splat.plotSpectrum(spc)          # does the same thing
       >>> splat.plotSpectrum(spc,uncertainty=True,tdwarf=True)     # show the spectrum uncertainty and T dwarf absorption features

    :Example 2: Viewing a set of spectra for a given object
        In this case we'll look at all of the spectra of TWA 30B in the library, sorted by year and compared to the first epoch data
        This is an example of using multiplot and multipage

       >>> splist = splat.getSpectrum(name = 'TWA 30B')         # get all spectra of TWA 30B
       >>> junk = [sp.normalize() for sp in splist]             # normalize the spectra
       >>> dates = [sp.date for sp in splist]                   # observation dates
       >>> spsort = [s for (s,d) in sorted(zip(dates,splis))]   # sort spectra by dates
       >>> dates.sort()                                         # don't forget to sort dates!
       >>> splat.plotSpectrum(spsort,multiplot=True,layout=[2,2],multipage=True,\   # here's our plot statement
           comparison=spsort[0],uncertainty=True,mdwarf=True,telluric=True,legends=dates,\
           legendLocation='lower left',output='TWA30B.pdf')

    :Example 3: Display the spectra sequence of L dwarfs
        This example uses the list of standard files contained in SPLAT, and illustrates the stack feature

       >>> spt = [splat.typeToNum(i+20) for i in range(10)] # generate list of L spectral types
       >>> splat.initiateStandards()                        # initiate standards
       >>> splist = [splat.SPEX_STDS[s] for s in spt]       # extact just L dwarfs
       >>> junk = [sp.normalize() for sp in splist]         # normalize the spectra
       >>> labels = [sp.shortname for sp in splist]         # set labels to be names
       >>> splat.plotSpectrum(splist,figsize=[10,20],labels=labels,stack=0.5,\  # here's our plot statement
           colorScheme='copper',legendLocation='outside',telluric=True,output='lstandards.pdf')

    '''

# keyword parameters
#    from .splat import Spectrum
    nsamples = kwargs.get('nsamples',1000)
    multiplot = kwargs.get('multiplot',False)           # create multiple plots
    multipage = kwargs.get('multipage',False)           # create a multiple page sequence of plots
    multilayout = kwargs.get('multilayout',[1,1])       # layout of multiple plots, [# horizontal, # vertical]
    multilayout = kwargs.get('layout',multilayout)
    stack = kwargs.get('stack',0)                   # stack spectra on top of each other
    grid = kwargs.get('grid',False)                 # plot internal grid lines
    filename = kwargs.get('filename','')            # output filename
    filename = kwargs.get('file',filename)
    filename = kwargs.get('output',filename)
    if filename == False:
        filename = ''
    title = kwargs.get('title','')
    fontscale = kwargs.get('fontscale',1)
    legendfontscale = kwargs.get('legendfontscale',0.8*fontscale)
    fb = filename.split('.')[:-1]               # filebase for multiple files
    if filename != '' and len(fb) > 1:
        filebase = fb[0]
        for x in fb[1:]:
            filebase+='.'+x
    elif filename != '' and len(fb) == 1:
        filebase = fb[0]
    else:
        filebase = filename
    filetype = kwargs.get('format',filename.split('.')[-1])
    filetype.lower()
    if filetype == '':
        filetype = 'pdf'
    comparison = kwargs.get('comparison',False)
    if not isinstance(comparison,Spectrum):
        comparison = False
    residual = kwargs.get('residual',False)
    inset = kwargs.get('inset',False)
    inset_xrange = kwargs.get('inset_xrange',False)
    inset_position = kwargs.get('inset_position',[.65, .6, .2, .2])
#    inset_color = kwargs.get('inset_color','k')
    inset_features = kwargs.get('inset_features',False)

#    mask = kwargs.get('mask',False)                # not yet implemented

# features to label on spectra
    feature_labels = { \
        'h2o': {'label': r'H$_2$O', 'type': 'band', 'wavelengths': [[0.92,0.95],[1.08,1.20],[1.325,1.550],[1.72,2.14]]}, \
        'ch4': {'label': r'CH$_4$', 'type': 'band', 'wavelengths': [[1.1,1.24],[1.28,1.44],[1.6,1.76],[2.2,2.35]]}, \
        'co': {'label': r'CO', 'type': 'band', 'wavelengths': [[2.29,2.39]]}, \
        'tio': {'label': r'TiO', 'type': 'band', 'wavelengths': [[0.76,0.80],[0.825,0.831]]}, \
        'vo': {'label': r'VO', 'type': 'band', 'wavelengths': [[1.04,1.08]]}, \
#        'feh': {'label': r'FeH', 'type': 'band', 'wavelengths': [[0.86,0.90],[0.98,1.03],[1.19,1.25],[1.57,1.64]]}, \
        'feh': {'label': r'FeH', 'type': 'band', 'wavelengths': [[0.98,1.03],[1.19,1.25],[1.57,1.64]]}, \
        'h2': {'label': r'H$_2$', 'type': 'band', 'wavelengths': [[1.5,2.4]]}, \
        'sb': {'label': r'*', 'type': 'band', 'wavelengths': [[1.6,1.64]]}, \
        'h': {'label': r'H I', 'type': 'line', 'wavelengths': [[1.004,1.005],[1.093,1.094],[1.281,1.282],[1.944,1.945],[2.166,2.166]]},\
        'hi': {'label': r'H I', 'type': 'line', 'wavelengths': [[1.004,1.005],[1.093,1.094],[1.281,1.282],[1.944,1.945],[2.166,2.166]]},\
        'h1': {'label': r'H I', 'type': 'line', 'wavelengths': [[1.004,1.005],[1.093,1.094],[1.281,1.282],[1.944,1.945],[2.166,2.166]]},\
        'na': {'label': r'Na I', 'type': 'line', 'wavelengths': [[0.8186,0.8195],[1.136,1.137],[2.206,2.209]]}, \
        'nai': {'label': r'Na I', 'type': 'line', 'wavelengths': [[0.8186,0.8195],[1.136,1.137],[2.206,2.209]]}, \
        'na1': {'label': r'Na I', 'type': 'line', 'wavelengths': [[0.8186,0.8195],[1.136,1.137],[2.206,2.209]]}, \
        'mg': {'label': r'Mg I', 'type': 'line', 'wavelengths': [[1.7113336,1.7113336],[1.5745017,1.5770150],[1.4881595,1.4881847,1.5029098,1.5044356],[1.1831422,1.2086969],]}, \
        'mgi': {'label': r'Mg I', 'type': 'line', 'wavelengths': [[1.7113336,1.7113336],[1.5745017,1.5770150],[1.4881595,1.4881847,1.5029098,1.5044356],[1.1831422,1.2086969],]}, \
        'mg1': {'label': r'Mg I', 'type': 'line', 'wavelengths': [[1.7113336,1.7113336],[1.5745017,1.5770150],[1.4881595,1.4881847,1.5029098,1.5044356],[1.1831422,1.2086969],]}, \
        'ca': {'label': r'Ca I', 'type': 'line', 'wavelengths': [[2.263110,2.265741],[1.978219,1.985852,1.986764],[1.931447,1.945830,1.951105]]}, \
        'cai': {'label': r'Ca I', 'type': 'line', 'wavelengths': [[2.263110,2.265741],[1.978219,1.985852,1.986764],[1.931447,1.945830,1.951105]]}, \
        'ca1': {'label': r'Ca I', 'type': 'line', 'wavelengths': [[2.263110,2.265741],[1.978219,1.985852,1.986764],[1.931447,1.945830,1.951105]]}, \
        'caii': {'label': r'Ca II', 'type': 'line', 'wavelengths': [[1.184224,1.195301],[0.985746,0.993409]]}, \
        'ca2': {'label': r'Ca II', 'type': 'line', 'wavelengths': [[1.184224,1.195301],[0.985746,0.993409]]}, \
        'al': {'label': r'Al I', 'type': 'line', 'wavelengths': [[1.672351,1.675511],[1.3127006,1.3154345]]}, \
        'ali': {'label': r'Al I', 'type': 'line', 'wavelengths': [[1.672351,1.675511],[1.3127006,1.3154345]]}, \
        'al1': {'label': r'Al I', 'type': 'line', 'wavelengths': [[1.672351,1.675511],[1.3127006,1.3154345]]}, \
        'fe': {'label': r'Fe I', 'type': 'line', 'wavelengths': [[1.5081407,1.5494570],[1.25604314,1.28832892],[1.14254467,1.15967616,1.16107501,1.16414462,1.16931726,1.18860965,1.18873357,1.19763233]]}, \
        'fei': {'label': r'Fe I', 'type': 'line', 'wavelengths': [[1.5081407,1.5494570],[1.25604314,1.28832892],[1.14254467,1.15967616,1.16107501,1.16414462,1.16931726,1.18860965,1.18873357,1.19763233]]}, \
        'fe1': {'label': r'Fe I', 'type': 'line', 'wavelengths': [[1.5081407,1.5494570],[1.25604314,1.28832892],[1.14254467,1.15967616,1.16107501,1.16414462,1.16931726,1.18860965,1.18873357,1.19763233]]}, \
        'k': {'label': r'K I', 'type': 'line', 'wavelengths': [[0.7699,0.7665],[1.169,1.177],[1.244,1.252]]}, \
        'ki': {'label': r'K I', 'type': 'line', 'wavelengths': [[0.7699,0.7665],[1.169,1.177],[1.244,1.252]]}, \
        'k1': {'label': r'K I', 'type': 'line', 'wavelengths': [[0.7699,0.7665],[1.169,1.177],[1.244,1.252]]}}

    features = kwargs.get('features',[])
    if not isinstance(features,list):
        features = [features]
    if (kwargs.get('ldwarf',False) or kwargs.get('mdwarf',False)):
        features.extend(['k','na','feh','tio','co','h2o','h2'])
    if (kwargs.get('tdwarf',False)):
        features.extend(['k','ch4','h2o','h2'])
    if (kwargs.get('young',False)):
        features.extend(['vo'])
    if (kwargs.get('binary',False)):
        features.extend(['sb'])
# clean repeats while maintaining order - set does not do this
    fea = []
    for i in features:
        if i not in fea:
            fea.append(i)
    features = fea

# error check - make sure you're plotting something
    if (len(args) < 1):
        print('plotSpectrum needs at least one Spectrum object to plot')
        return

# if a list is passed, use this list
    elif (len(args) == 1 and isinstance(args[0],list)):
        splist = args[0]

# if a set of objects is passed, turn into a list
    else:
        splist = []
        for a in args:
            if isinstance(a,Spectrum):      # a spectrum object
                splist.append(a)
            elif isinstance(a,list):
                splist.append(a)
            else:
                print('\nplotSpectrum: Ignoring input object {} as it is neither a Spectrum object nor a list\n\n'.format(a))

# set up for multiplot
    if (len(splist) == 1):
        multiplot = False

# array of lists => force multiplot
    elif (len(splist) > 1 and isinstance(splist[0],list)):
        multiplot = True

# reformat array of spectra of multiplot is used (i.e., user forgot to set)
    if multiplot == True and isinstance(splist[0],Spectrum):
        splist = [[s] for s in splist]

    elif multiplot == False and isinstance(splist[0],Spectrum):
        splist = [splist]

# flatten array if multiplot is not set
    elif multiplot == False and isinstance(splist[0],list) and len(splist) > 1:
        splist = [[item for sublist in splist for item in sublist]]       # flatten

    tot_sp = len([item for sublist in splist for item in sublist])    # Total number of spectra

# prep legend
    legend = kwargs.get('legend',[str() for x in numpy.arange(tot_sp)])
    legend = kwargs.get('legends',legend)
    legend = kwargs.get('label',legend)
    legend = kwargs.get('labels',legend)
    if not isinstance(legend,list):
        legend = [legend]
    if(len(legend) < tot_sp):
        legend.extend([str() for x in numpy.arange(tot_sp-len(legend))])
    legendLocation = kwargs.get('legendLocation','upper right')       # sets legend location
    legendLocation = kwargs.get('labelLocation',legendLocation)       # sets legend location

# now run a loop through the input subarrays
    plt.close('all')

# set up here for multiple file output
    nplot = 1
    if multipage == True or multiplot == True:
        nplot = multilayout[0]*multilayout[1]
        numpages = int(len(splist) / nplot) + 1
        if (len(splist) % nplot == 0):
               numpages -= 1
        fig = []

    if multipage == True and filetype == 'pdf':
        pdf_pages = PdfPages(filename)

    if multipage == False:
        if len(splist) > 1:
            files = [filebase+'{}.'.format(i+1)+filetype for i in numpy.arange(len(splist))]
        else:
            files = [filebase+'.'+filetype]

    pg_n = 0        # page counter
    plt_n = 0       # plot per page counter
    lg_n = 0        # legend per plot counter
    for plts,sp in enumerate(splist):
# set specific plot parameters
#        print(sp[0],Spectrum,isinstance(sp[0],Spectrum))
        if not isinstance(sp[0],Spectrum):
            raise ValueError('\nInput to plotSpectrum has wrong format:\n\n{}\n\n'.format(sp[0]))
        zeropoint = kwargs.get('zeropoint',[0. for x in numpy.arange(len(sp))])

# settings that work if the spectrum was read in as legitmate Spectrum object
        try:
            xlabel = kwargs.get('xlabel','{} ({})'.format(sp[0].wlabel,sp[0].wunit))
            ylabel = kwargs.get('ylabel','{} {} ({})'.format(sp[0].fscale,sp[0].flabel,sp[0].funit))
        except:
            xlabel = kwargs.get('xlabel','Wavelength (unknown units)')
            ylabel = kwargs.get('ylabel','Flux (unknown units)')
        xrange = kwargs.get('xrange',[0.85,2.42])
        bound = xrange
        ymax = [s.fluxMax().value for s in sp]
        yrng = kwargs.get('yrange',map(lambda x: x*numpy.nanmax(ymax)+numpy.nanmax(zeropoint),[-0.02,1.2]))
        bound.extend(yrng)
        linestyle = kwargs.get('linestyle',['steps' for x in numpy.arange(len(sp))])
        linestyle = kwargs.get('linestyles',linestyle)
        if (len(linestyle) < len(sp)):
            linestyle.extend(['steps' for x in numpy.arange(len(sp)-len(linestyle))])

# colors
# by default all black lines
        colors = kwargs.get('colors',['k' for x in numpy.arange(len(sp))])
        colors = kwargs.get('color',colors)
        if not isinstance(colors,list):
            colors = [colors]
        if (len(colors) < len(sp)):
            while len(colors) < len(sp):
                colors.extend(colors[-1])
        colorScheme = kwargs.get('colorScheme',None)
        colorScheme = kwargs.get('colorMap',colorScheme)
        if (colorScheme != None):
            values = numpy.arange(len(sp))
            color_map = plt.get_cmap(colorScheme)
            norm  = colmap.Normalize(vmin=0, vmax=1.0*values[-1])
            scalarMap = cm.ScalarMappable(norm=norm, cmap=color_map)
            for i in numpy.arange(len(sp)):
                colors[i] = scalarMap.to_rgba(values[i])
        colorsUnc = kwargs.get('colorsUnc',colors)
        colorsUnc = kwargs.get('colorUnc',colorsUnc)
        if (len(colorsUnc) < len(sp)):
            while len(colorsUnc) < len(sp):
                colorsUnc.extend(colors[-1])

# show uncertainties
        showNoise = kwargs.get('showNoise',[False for x in numpy.arange(len(sp))])
        showNoise = kwargs.get('noise',showNoise)
        showNoise = kwargs.get('uncertainty',showNoise)
        if not isinstance(showNoise, list):
            showNoise = [showNoise]
        if (len(showNoise) < len(sp)):
            showNoise.extend([True for x in numpy.arange(len(sp)-len(showNoise))])

# zero points - by default true
        showZero = kwargs.get('showZero',[True for x in numpy.arange(len(sp))])
        if not isinstance(showZero, list):
            showZero = [showZero]
        if (len(showZero) < len(sp)):
            while len(showZero) < len(sp):
                showZero.extend(showZero[-1])

# GENERATE PLOTS
        if (multiplot == True or multipage == True):
            plt_n = plts % nplot
            if (plt_n == 0):# and plts != len(splist)):
#                ax = range(nplot)
#                t = tuple([tuple([i+b*multilayout[1] for i in range(multilayout[1])]) for b in range(multilayout[0])])
#                fig[pg_n], ax = plt.subplots(multilayout[0], multilayout[1], sharex = True, sharey = True)
#
# NOTE THE FOLLOWING LINE IS HAVING PROBLEMS IN PYTHON3
#
#
                fig.append(plt.figure())
                pg_n += 1
            ax = fig[pg_n-1].add_subplot(multilayout[0], multilayout[1], plt_n+1)

# plotting a single plot with all spectra
        else:
            plt.close('all')
#            ax = range(1)
            plt_n = 0
            fig = []
            if (kwargs.get('figsize') != None):
                fig.append(plt.figure(figsize = kwargs.get('figsize')))
            else:
                fig.append(plt.figure())
            ax = fig[0].add_subplot(111)

        for ii, a in enumerate(sp):
            flx = [i+zeropoint[ii] for i in a.flux.value]
#stack
            if stack > 0:
                flx = [f + (len(sp)-ii-1)*stack for f in flx]
#                if kwargs.get('yrange') == None:
#                    bound[3] = bound[3] + stack

            ax.plot(a.wave.value,flx,color=colors[ii],linestyle=linestyle[ii], zorder = 10, label = legend[lg_n])

# add comparison
            if comparison != False:
                colorComparison = kwargs.get('colorComparison',colors[0])
                linestyleComparison = kwargs.get('linestyleComparison',linestyle[0])
                cflx = [i+zeropoint[ii] for i in comparison.flux.value]

                if stack > 0:
                    cflx = [f + (len(sp)-ii-1)*stack for f in cflx]

                ax.plot(comparison.wave.value,cflx,color=colorComparison,linestyle=linestyleComparison,alpha=0.5, zorder = 10)

# add residual
            if residual == True and len(sp) == 2:
                # Save flux values from first spectrum
                if ii == 0:
                    flx0 = [f - (len(sp)-ii-1)*stack for f in flx]

                # Subtract fluxes and plot
                elif ii == 1:
                    res = [flx0[f_n] - f for f_n, f in enumerate(flx)]
                    ax.plot(a.wave.value, res, alpha = 0.3, color = 'g')

                    # Fix bound[2] if res goes below 0
                    if min(res) < 0:
                        b0 = numpy.argmax(a.wave.value > bound[0])
                        b1 = numpy.argmin(a.wave.value < bound[1])
                        bound[2] = bound[2] + min(res[b0:b1])

# noise
            if (showNoise[ii]):
                ns = [i+zeropoint[ii] for i in a.noise.value]
                ax.plot(a.wave.value,ns,color=colorsUnc[ii],linestyle=linestyle[ii],alpha=0.3, zorder = 10)

# zeropoint
            if (showZero[ii]):
                ze = numpy.ones(len(a.flux))*zeropoint[ii]
                ax.plot(a.wave.value,ze,color=colors[ii],linestyle=':',alpha=0.3, zorder = 10)

# determine maximum flux for all spectra
            f = interp1d(a.wave,flx,bounds_error=False,fill_value=0.)
            if (ii == 0):
                wvmax = numpy.arange(bound[0],bound[1],0.001)
                flxmax = f(wvmax)
            else:
                flxmax = numpy.maximum(flxmax,f(wvmax))

# legend counter
            lg_n = lg_n + 1 # Increment lg_n

# label features
# THIS NEEDS TO BE FIXED WITH GRETEL'S STUFF
        yoff = 0.02*(bound[3]-bound[2])
        fontsize = 10-numpy.min([(multilayout[0]*multilayout[1]-1),6])
        for ftr in features:
            ftr = ftr.lower()
            if ftr in feature_labels:
                for ii,waveRng in enumerate(feature_labels[ftr]['wavelengths']):
                    if (numpy.min(waveRng) > bound[0] and numpy.max(waveRng) < bound[1]):
                        x = (numpy.arange(0,nsamples+1.0)/nsamples)* \
                            (numpy.nanmax(waveRng)-numpy.nanmin(waveRng)+0.04)+numpy.nanmin(waveRng)-0.02
                        f = interp1d(wvmax,flxmax,bounds_error=False,fill_value=0.)
                        y = numpy.nanmax(f(x))+0.5*yoff

                        if feature_labels[ftr]['type'] == 'band':
                            ax.plot(waveRng,[y+yoff]*2,color='k',linestyle='-')
                            ax.plot([waveRng[0]]*2,[y,y+yoff],color='k',linestyle='-')
                            ax.text(numpy.mean(waveRng),y+1.5*yoff,feature_labels[ftr]['label'],horizontalalignment='center',fontsize=fontsize)
                        else:
                            for w in waveRng:
                                ax.plot([w]*2,[y,y+yoff],color='k',linestyle='-')
                            ax.text(numpy.mean(waveRng),y+1.5*yoff,feature_labels[ftr]['label'],horizontalalignment='center',fontsize=fontsize)
                            waveRng = [waveRng[0]-0.02,waveRng[1]+0.02]   # for overlap

# update offset
                        foff = [y+3*yoff if (w >= waveRng[0] and w <= waveRng[1]) else 0 for w in wvmax]
                        flxmax = [numpy.max([xx,yy]) for xx, yy in zip(flxmax, foff)]
        bound[3] = numpy.max([numpy.max(flxmax)+1.*yoff,bound[3]])
        ax.axis(bound)

# grid
        if (grid):
            ax.grid()

# axis labels
        fontsize = (13-numpy.min([(multilayout[0]*multilayout[1]-1),8])) * fontscale        # Added in fontscale
        legendfontsize = (13-numpy.min([(multilayout[0]*multilayout[1]-1),8])) * legendfontscale        # Added in fontscale
        ax.set_xlabel(xlabel, fontsize = fontsize)
        ax.set_ylabel(ylabel, fontsize = fontsize)
        ax.tick_params(axis='x', labelsize=fontsize)
        ax.tick_params(axis='y', labelsize=fontsize)

# place legend
        if len(legend) > 0:
            if legendLocation == 'outside':
                box = ax.get_position()
                ax.set_position([box.x0, box.y0 + box.height * 0.15, box.width * 0.7, box.height * 0.7])
                ax.legend(loc='center left', bbox_to_anchor=(1, 0.5), prop={'size':legendfontsize})
            else:
                ax.legend(loc=legendLocation, prop={'size':legendfontsize})
                bound[3] = bound[3]+0.1*(bound[3]-bound[2])     # extend axis for in-plot legends
            ax.axis(bound)

# overplot telluric absorption
        if (kwargs.get('telluric',False) == True):
            twv = [[1.1,1.2],[1.3,1.5],[1.75,2.0]]
            for waveRng in twv:
                rect = patches.Rectangle((waveRng[0],bound[2]),waveRng[1]-waveRng[0],bound[3]-bound[2],facecolor='0.95',alpha=0.2,color='0.95')
                ax.add_patch(rect)
                ax.text(numpy.mean(waveRng),bound[2]+3*yoff,r'$\oplus$',horizontalalignment='center',fontsize=fontsize)

# place inset - RIGHT NOW ONLY SETTING LIMITS WITH FIRST SPECTRUM IN LIST
        if inset == True and inset_xrange != False:
            ax_inset = fig[pg_n-1].add_axes(inset_position) #, axisbg='white')
            bound2 = inset_xrange
            b0 = numpy.argmax(sp[0].wave.value > bound2[0])
            b1 = numpy.argmin(sp[0].wave.value < bound2[1])
            bound2.extend([min(sp[0].flux.value[b0:b1]),max(sp[0].flux.value[b0:b1])])
            db = (bound2[3]-bound2[2])
            bound2[2] = bound2[2]-0.05*db
            bound2[3] = bound2[3]+0.05*db
            ax_inset.axis(bound2)
            inset_fontsize = fontsize*0.7

            for ii,a in enumerate(sp):
                flx = [i+zeropoint[ii] for i in a.flux.value]
                ax_inset.plot(a.wave.value,flx,color=colors[ii],linestyle=linestyle[ii])
                ax_inset.set_xlabel('')
                ax_inset.set_ylabel('')
                ax_inset.tick_params(axis='x', labelsize=inset_fontsize)
                ax_inset.tick_params(axis='y', labelsize=inset_fontsize)
#                ax_inset.legend()

# inset feature labels
            if inset_features != False:
                f = interp1d(sp[0].wave,flx,bounds_error=False,fill_value=0.)
                wvmax = numpy.arange(bound2[0],bound2[1],0.001)
                flxmax = f(wvmax)
                yoff = 0.05*(bound2[3]-bound2[2])
                for ftr in inset_features:
                    ftr = ftr.lower()
                    if ftr in feature_labels:
                        for ii,waveRng in enumerate(feature_labels[ftr]['wavelengths']):
                            if (numpy.min(waveRng) > bound2[0] and numpy.max(waveRng) < bound2[1]):
                                x = (numpy.arange(0,nsamples+1.0)/nsamples)* \
                                    (numpy.nanmax(waveRng)-numpy.nanmin(waveRng)+0.04)+numpy.nanmin(waveRng)-0.02
                                f = interp1d(wvmax,flxmax,bounds_error=False,fill_value=0.)
                                y = numpy.nanmax(f(x))+0.5*yoff

                                if feature_labels[ftr]['type'] == 'band':
                                    ax_inset.plot(waveRng,[y+yoff]*2,color='k',linestyle='-')
                                    ax_inset.plot([waveRng[0]]*2,[y,y+yoff],color='k',linestyle='-')
                                    ax_inset.text(numpy.mean(waveRng),y+2*yoff,feature_labels[ftr]['label'],horizontalalignment='center',fontsize=inset_fontsize)
                                else:
                                    for w in waveRng:
                                        ax_inset.plot([w]*2,[y,y+yoff],color='k',linestyle='-')
                                    ax_inset.text(numpy.mean(waveRng),y+2*yoff,feature_labels[ftr]['label'],horizontalalignment='center',fontsize=inset_fontsize)
                                    waveRng = [waveRng[0]-0.02,waveRng[1]+0.02]   # for overlap

# update offset
                                foff = [y+3*yoff if (w >= waveRng[0] and w <= waveRng[1]) else 0 for w in wvmax]
                                flxmax = [numpy.max([xx,yy]) for xx, yy in zip(flxmax, foff)]
                bound2[3] = numpy.max([bound2[3],numpy.max(flxmax)+5.*yoff])
                ax_inset.axis(bound2)

# save to file or display
        if multipage == False:
            if filebase != '' and (plts % nplot == 3 or plts == len(splist)-1):
                plt.savefig(files[plts], format=filetype)
    if filename == '' and not kwargs.get('web',False):
        plt.show()
        if (kwargs.get('interactive',False) != False):
            plt.ion()        # make window interactive
        else:
            plt.ioff()

# save figures in multipage format and write off pdf file
    if (multipage == True):
        for pg_n in numpy.arange(numpages):
#            fig[pg_n].text(0.5, 0.04, xlabel, ha = 'center', va = 'center')
#            fig[pg_n].text(0.06, 0.5, ylabel, ha = 'center', va = 'center', rotation = 'vertical')
            fig[pg_n].tight_layout
            fig[pg_n].suptitle(title, fontsize = 14, fontweight = 'bold')
            pdf_pages.savefig(fig[pg_n])
        if filetype == 'pdf':
            pdf_pages.close()

    return fig

def plotBatch(*args, **kwargs):
    '''
    :Purpose: Plots a batch of spectra into a 2x2 set of PDF files, with options of overplotting comparison spectra, including best-match spectral standards.

    :param input: A single or list of Spectrum objects or filenames, or the glob search string for a set of files (e.g., '/Data/myspectra/*.fits').
    :type input: required

    :param output: Filename for PDF file output; full path should be include if not saving to current directory
    :type output: optional, default = 'spectra_plot.pdf'
    :param comparisons: list of Spectrum objects or filenames for comparison spectra. If comparisons list is shorter than source list, then last comparison source will be repeated. If the comparisons list is longer, the list will be truncated.
    :type comparisons: optional, default = None
    :param classify: Set to True to classify sources based on comparison to MLT spectral standards following the method of `Kirkpatrick et al. (2010) <http://adsabs.harvard.edu/abs/2010ApJS..190..100K>`_. This option normalizes the spectra by default
    :type classify: optional, default = False
    :param normalize: Set to True to normalize source and (if passed) comparison spectra.
    :type normalize: optional, default = False
    :param legend: Set to list of legends for plots. The number of legends should equal the number of sources and comparisons (if passed) in an alternating sequence. T
    :type legend: optional, default = displays file name for source and object name for comparison (if passed)

    Relevant parameters for plotSpectrum may also be passed

    :Example:
       >>> import glob, splat
       >>> files = glob.glob('/home/mydata/*.fits')
       >>> sp = splat.plotBatch(files,classify=True,output='comparison.pdf')
       >>> sp = splat.plotBatch('/home/mydata/*.fits',classify=True,output='comparison.pdf')
       >>> sp = splat.plotBatch([splat.Spectrum(file=f) for f in files],classify=True,output='comparison.pdf')
       All three of these commands produce the same result
    '''

# keyword check
    parameters = ['output','comparisons','classify','normalize','legend']
    checkKeys(kwargs,parameters,forcekey=False)
    kwargs['layout']=kwargs.get('layout',[2,2])
    kwargs['fontscale']=kwargs.get('fontscale',0.7)

# input check
    if len(args) == 0:
        print('\nNeed to provide a list of spectra or filenames, or a file search string, to use plotBatch')
        return
    kwargs['output'] = kwargs.get('output','spectra_plot.pdf')
#    from .splat import Spectrum
# break out a list
    if isinstance(args[0],list):
        inputlist = args[0]
    else:
        inputlist = args

# if filenames, read in each file to a spectrum object
    if isinstance(inputlist[0],str):
# try a glob search string
        files = glob.glob(inputlist[0])
        if len(files) > 1 or (len(files) == 1 and inputlist[0].find('*') != -1):
            inputlist = files
# try reading in files into Spectrum object
        try:
            splist = [Spectrum(file = f) for f in inputlist]
        except:
            raise ValueError('\nCould not read in list of files - make sure the full path is specified and the files are correctly formatted')

# if filenames, read in each file to a spectrum object
    elif isinstance(inputlist[0],Spectrum):
        splist = copy.deepcopy(inputlist)
    else:
        print('\nInput should be list of Spectra objects or filenames')
        return

# comparison files are present
    compflag = False
    complist = []
    if kwargs.get('comparisons',False) != False:
        comp = kwargs.get('comparisons')
        if not isinstance(comp,list):
            comp = [comp]
        if len(comp) < len(splist):
            while len(comp) < len(splist):
                comp.append(comp[-1])
        if isinstance(comp[0],str):
            try:
                complist = [Spectrum(file = f) for f in comp]
                compflag = True
            except:
                print('\nCould not read in comparison files: ignoring comparisons')
        if isinstance(comp[0],Spectrum):
            complist = comp
            compflag = True

# set comparison files to be standards for spectral classification
    if kwargs.get('classify',False):
        complist = []
        for sp in splist:
#            from .splat import classifyByStandard
            spt = classifyByStandard(sp,method='kirkpatrick')
            complist.append(STDS_DWARF_SPEX[spt[0]])
        compflag = True
        kwargs['normalize'] = kwargs.get('normalize',True)

# normalize if desired
    if kwargs.get('normalize',False):
        tmp = [sp.normalize() for sp in splist]
        if compflag == True:
            tmp = [sp.normalize() for sp in complist]

# prep for plotting
    plotlist = []
    clist = []
    for i,sp in enumerate(splist):
        if compflag == True:
            plotlist.append([sp,complist[i]])
            clist.extend(['k','r'])
        else:
            plotlist.append([sp])
            clist.extend(['k'])

# manage legends
    if kwargs.get('legend',False) != False:
        legends = kwargs.get('legend')
        if not isinstance(legends,list):
            legends = [legends]
        if len(legends) < (len(splist)+len(complist)):
# guess: just left out the comparison legends
            if len(complist) > 0 and len(legends) == len(splist):
                legtmp = []
                for i,l in enumerate(legends):
                    legtmp.extend([l,'{}'.format(complist[i].name)])
                legends = legtmp
            else:
# otherwise: pad the remaining legends with the last legend (pairs)
                while len(legends) < (len(splist)+len(complist)):
                    if compflag:
                        legends.extend([legends[-2],legends[-1]])
                    else:
                        legends.extend([legends[-1]])
        if len(legends) > (len(splist)+len(complist)):
            legends = legends[0:(len(splist)+len(complist))]
    else:
        legends = []
        for i,sp in enumerate(splist):
            if compflag == True:
                legends.extend(['{}'.format(os.path.basename(sp.filename)),'{}'.format(complist[i].name)])
            else:
                legends.extend(['{}'.format(os.path.basename(sp.filename))])

# generate pdf
    plotSpectrum(plotlist,multiplot=True,multipage=True,legends=legends,colors=clist,**kwargs)

    return

def plotIndices(*args, **kwargs):
    '''
    :Purpose: ``Plot index-index plots.``

    Not currently implemented
    '''
    pass
    return

def plotSED(*args, **kwargs):
    '''
    :Purpose: ``Plot SED photometry with SpeX spectrum.``

    Not currently implemented
    '''
    pass
    return

def plotSequence(*args, **kwargs):
    '''
    :Purpose: Compares a spectrum to a sequence of standards a batch of spectra into a 2x2 set of PDF files, with options of overplotting comparison spectra, including best-match spectral standards.

    :param input: A single or list of Spectrum objects or filenames, or the glob search string for a set of files (e.g., '/Data/myspectra/*.fits').
    :type input: required

    :param type_range: Number of subtypes to consider above and below best-fit spectral type
    :type type_range: optional, default = 2
    :param spt: Default spectral type for source; this input skips classifyByStandard
    :type spt: optional, default = None
    :param output: Filename for output; full path should be include if not saving to current directory. If blank, plot is shown on screen
    :type output: optional, default = None (screen display)

    Relevant parameters for plotSpectrum may also be passed

    :Example:
       >>> import glob, splat
       >>> files = glob.glob('/home/mydata/*.fits')
       >>> sp = splat.plotBatch(files,classify=True,output='comparison.pdf')
       >>> sp = splat.plotBatch('/home/mydata/*.fits',classify=True,output='comparison.pdf')
       >>> sp = splat.plotBatch([splat.Spectrum(file=f) for f in files],classify=True,output='comparison.pdf')
       All three of these commands produce the same result
    '''

# check inputs
    if len(args) == 0:
        raise ValueError('\nNeed to provide a spectrum object or filename for plotStandardSequence')

#    from .splat import classifyByStandard, Spectrum
    parameters = ['type_range']
    checkKeys(kwargs,parameters,forcekey=False)
    type_range =kwargs.get('type_range',2)
    kwargs['stack']=kwargs.get('stack',0.5)
#    kwargs['legendLocation']=kwargs.get('legendLocation','outside')
    kwargs['telluric']=kwargs.get('telluric',True)
    kwargs['method']=kwargs.get('method','kirkpatrick')
#    kwargs['color']=kwargs.get('color','r')
    kwargs['colorComparison']=kwargs.get('colorComparison','k')
    kwargs['colorScheme']=kwargs.get('colorScheme','winter')
    kwargs['figsize']=kwargs.get('figsize',[10,10*type_range*kwargs['stack']])
    kwargs['fontscale']=kwargs.get('fontscale',1.5)

# process input
    if isinstance(args[0],str):
        if len(glob.glob(args[0])) == 0:
            raise ValueError('\nCannot find input file {} - make sure full path is included'.format(args[0]))
        try:
            sp = Spectrum(file = args[0])
        except:
            raise ValueError('\nCould not read in file {} - make sure the file is correctly formatted'.format(args[0]))
    elif isinstance(args[0],Spectrum):
        sp = copy.deepcopy(args[0])
    else:
        raise ValueError('\nInput should be a Spectrum object or filename')
    sp.normalize()

# classify by comparison to standards
    spt = kwargs.get('spt',classifyByStandard(sp,**kwargs)[0])
    if not isinstance(spt,str):
        spt = typeToNum(spt)

# produce range of standards for plot
    stdnum = numpy.arange(numpy.floor(typeToNum(spt)-type_range),numpy.ceil(typeToNum(spt)+type_range)+1)
    if numpy.max(stdnum) > 39:
        stdnum-=(numpy.max(stdnum)-39)
    if numpy.min(stdnum) < 10:
        stdnum+=(10-numpy.min(stdnum))
    stdspt = [typeToNum(i) for i in stdnum]
    stds = [STDS_DWARF_SPEX[s] for s in stdspt]
    stdlabels = ['{} Standard'.format(s) for s in stdspt]
    plotlist = []
    labels = []
    colors = []
    for i,stdsp in enumerate(stds):
        plotlist.append([stdsp,sp])
        labels.extend([])
    fig = plotSpectrum(stds,comparison=sp,labels=stdlabels,**kwargs)

    return fig

# main testing of program
if __name__ == '__main__':
    out_folder = '/Users/adam/projects/splat/code/testing/'
    def test_plotBatch():
        data_folder = '/Users/adam/projects/splat/adddata/done/daniella/'
        files = glob.glob(data_folder+'*.fits')
        plotBatch(files,classify=True,output=out_folder+'plotBatch_test1.pdf',telluric=True)
        plotBatch(data_folder+'*.fits',classify=True,output=out_folder+'plotBatch_test2.pdf',noise=True)
        splist = [Spectrum(file=f) for f in files]
        plotBatch(splist,classify=True,output=out_folder+'plotBatch_test3.pdf',features=['h2o','feh','co'],legend=[s.name for s in splist])
        return

    def test_plotSequence():
        sp = Spectrum(10001)
        plotSequence(sp,output=out_folder+'plotSequence_test1.pdf')
        plotSequence(sp,type_range=3,output=out_folder+'plotSequence_test2.png')
        data_folder = '/Users/adam/projects/splat/adddata/done/daniella/'
        files = glob.glob(data_folder+'*.fits')
        plotSequence(files[0],telluric=True,stack=0.7,spt='M0',output=out_folder+'plotSequence_test3.eps')
        sp = getSpectrum(shortname='0415-0935')[0]
        plotSequence(sp,telluric=True,stack=0.3,output=out_folder+'plotSequence_test4.eps')
        plotSequence(sp,telluric=True,stack=0.3)
        return

    test_plotBatch()
#    test_plotSequence()


import os
import logging
import math
import numpy as np
import time
import typing

from tqdm import tqdm

from smac.callbacks import IncorporateRunResultCallback
from smac.configspace import Configuration
from smac.epm.rf_with_instances import RandomForestWithInstances
from smac.initial_design.initial_design import InitialDesign
from smac.intensification.abstract_racer import AbstractRacer, RunInfoIntent
from smac.optimizer import pSMAC
from smac.optimizer.acquisition import AbstractAcquisitionFunction
from smac.optimizer.random_configuration_chooser import ChooserNoCoolDown, RandomConfigurationChooser
from smac.optimizer.ei_optimization import AcquisitionFunctionMaximizer
from smac.optimizer.epm_configuration_chooser import EPMChooser
from smac.runhistory.runhistory import RunHistory, RunInfo, RunValue
from smac.runhistory.runhistory2epm import AbstractRunHistory2EPM
from smac.scenario.scenario import Scenario
from smac.stats.stats import Stats
from smac.utils.constants import MAXINT
from smac.tae import (
    FirstRunCrashedException,
    StatusType,
    TAEAbortException,
)
from smac.tae.base import BaseRunner
from smac.utils.io.traj_logging import TrajLogger
from smac.utils.validate import Validator
from smac.utils.neptune import configspace_to_dict
import smac.utils.neptune as smac_neptune

__author__ = "Aaron Klein, Marius Lindauer, Matthias Feurer"
__copyright__ = "Copyright 2015, ML4AAD"
__license__ = "3-clause BSD"

class SMBO(object):
    """Interface that contains the main Bayesian optimization loop

    Attributes
    ----------
    logger
    incumbent
    scenario
    config_space
    stats
    initial_design
    runhistory
    intensifier
    num_run
    rng
    initial_design_configs
    epm_chooser
    tae_runner
    """

    def __init__(self,
                 scenario: Scenario,
                 stats: Stats,
                 initial_design: InitialDesign,
                 runhistory: RunHistory,
                 runhistory2epm: AbstractRunHistory2EPM,
                 intensifier: AbstractRacer,
                 num_run: int,
                 model: RandomForestWithInstances,
                 acq_optimizer: AcquisitionFunctionMaximizer,
                 acquisition_func: AbstractAcquisitionFunction,
                 rng: np.random.RandomState,
                 tae_runner: BaseRunner,
                 restore_incumbent: Configuration = None,
                 random_configuration_chooser: typing.Union[RandomConfigurationChooser] = ChooserNoCoolDown(2.0),
                 predict_x_best: bool = True,
                 min_samples_model: int = 1,
                 neptune_run = None):
        """
        Interface that contains the main Bayesian optimization loop

        Parameters
        ----------
        scenario: smac.scenario.scenario.Scenario
            Scenario object
        stats: Stats
            statistics object with configuration budgets
        initial_design: InitialDesign
            initial sampling design
        runhistory: RunHistory
            runhistory with all runs so far
        runhistory2epm : AbstractRunHistory2EPM
            Object that implements the AbstractRunHistory2EPM to convert runhistory
            data into EPM data
        intensifier: Intensifier
            intensification of new challengers against incumbent configuration
            (probably with some kind of racing on the instances)
        num_run: int
            id of this run (used for pSMAC)
        model: RandomForestWithInstances
            empirical performance model (right now, we support only RandomForestWithInstances)
        acq_optimizer: AcquisitionFunctionMaximizer
            Optimizer of acquisition function.
        acquisition_func : AcquisitionFunction
            Object that implements the AbstractAcquisitionFunction (i.e., infill criterion for acq_optimizer)
        restore_incumbent: Configuration
            incumbent to be used from the start. ONLY used to restore states.
        rng: np.random.RandomState
            Random number generator
        tae_runner : smac.tae.base.BaseRunner Object
            target algorithm run executor
        random_configuration_chooser
            Chooser for random configuration -- one of
            * ChooserNoCoolDown(modulus)
            * ChooserLinearCoolDown(start_modulus, modulus_increment, end_modulus)
        predict_x_best: bool
            Choose x_best for computing the acquisition function via the model instead of via the observations.
        min_samples_model: int
            Minimum number of samples to build a model.
        """

        self.logger = logging.getLogger(
            self.__module__ + "." + self.__class__.__name__)
        self.incumbent = restore_incumbent

        self.scenario = scenario
        self.config_space = scenario.cs  # type: ignore[attr-defined] # noqa F821
        self.stats = stats
        self.initial_design = initial_design
        self.runhistory = runhistory
        self.intensifier = intensifier
        self.num_run = num_run
        self.rng = rng
        self._min_time = 10 ** -5
        self.tae_runner = tae_runner
        if neptune_run is None:
            neptune_run = smac_neptune.get_run()
        self.neptune_run = neptune_run

        self.initial_design_configs = []  # type: typing.List[Configuration]

        # initialize the chooser to get configurations from the EPM
        self.epm_chooser = EPMChooser(scenario=scenario,
                                      stats=stats,
                                      runhistory=runhistory,
                                      runhistory2epm=runhistory2epm,
                                      model=model,
                                      acq_optimizer=acq_optimizer,
                                      acquisition_func=acquisition_func,
                                      rng=rng,
                                      restore_incumbent=restore_incumbent,
                                      random_configuration_chooser=random_configuration_chooser,
                                      predict_x_best=predict_x_best,
                                      min_samples_model=min_samples_model)

        # Internal variable - if this is set to True it will gracefully stop SMAC
        self._stop = False

        # Callbacks. All known callbacks have a key. If something does not have a key here, there is
        # no callback available.
        self._callbacks = {
            '_incorporate_run_results': list()
        }  # type: typing.Dict[str, typing.List[typing.Callable]]
        self._callback_to_key = {
            IncorporateRunResultCallback: '_incorporate_run_results',
        }  # type: typing.Dict[typing.Type, str]

    def start(self) -> None:
        """Starts the Bayesian Optimization loop.
        Detects whether the optimization is restored from a previous state.
        """
        self.stats.start_timing()

        # Initialization, depends on input
        if self.stats.submitted_ta_runs == 0 and self.incumbent is None:
            self.logger.info('Running initial design')
            # Intensifier initialization
            self.initial_design_configs = self.initial_design.select_configurations()

            # to be on the safe side, never return an empty list of initial configs
            if not self.initial_design_configs:
                self.initial_design_configs = [self.config_space.get_default_configuration()]

        elif self.stats.submitted_ta_runs > 0 and self.incumbent is None:
            raise ValueError("According to stats there have been runs started, "
                             "but the optimizer cannot detect an incumbent. Did "
                             "you set the incumbent (e.g. after restoring state)?")
        elif self.stats.submitted_ta_runs == 0 and self.incumbent is not None:
            raise ValueError("An incumbent is specified, but there are no runs "
                             "recorded as started in the Stats-object. If you're "
                             "restoring a state, please provide the Stats-object.")
        else:
            # Restoring state!
            self.logger.info("State Restored! Starting optimization with "
                             "incumbent %s", self.incumbent)
            self.logger.info("State restored with following budget:")
            self.stats.print_stats()

    def run(self) -> Configuration:
        """Runs the Bayesian optimization loop

        Returns
        ----------
        incumbent: np.array(1, H)
            The best found configuration
        """
        self.neptune_run['config_space'] = configspace_to_dict(self.config_space)
        self.neptune_run['num_run'] = self.num_run

        self.start()

        # Main BO loop
        bo_iterations = 0
        t = tqdm(desc='Optimization')
        while True:
            self.neptune_run['bo_iterations'].log(bo_iterations)
            self.neptune_run['initial_design_configs_count'].log(len(self.initial_design_configs))
            for k in ['continue_challenger', 'current_challenger', 'elapsed_time', 'iteration_done', 'n_iters',
                      'num_chall_run', 'num_run']:
                v = getattr(self.intensifier, k)
                if k == 'current_challenger':
                    if v is not None:
                        for config_k, config_v in v.get_dictionary().items():
                            self.neptune_run[f'intensifier/{k}/{config_k}'].log(config_v, bo_iterations)
                    continue
                self.neptune_run[f'intensifier/{k}'].log(v)
            self.neptune_run['intensifier/to_run_count'].log(len(self.intensifier.to_run))
            if self.scenario.shared_model:  # type: ignore[attr-defined] # noqa F821
                pSMAC.read(run_history=self.runhistory,
                           output_dirs=self.scenario.input_psmac_dirs,  # type: ignore[attr-defined] # noqa F821
                           configuration_space=self.config_space,
                           logger=self.logger)

            start_time = time.time()

            # sample next configuration for intensification
            # Initial design runs are also included in the BO loop now.
            intent, run_info = self.intensifier.get_next_run(
                challengers=self.initial_design_configs,
                incumbent=self.incumbent,
                chooser=self.epm_chooser,
                run_history=self.runhistory,
                repeat_configs=self.intensifier.repeat_configs,
                num_workers=self.tae_runner.num_workers(),
            )

            # remove config from initial design challengers to not repeat it again
            self.initial_design_configs = [c for c in self.initial_design_configs if c != run_info.config]

            # update timebound only if a 'new' configuration is sampled as the challenger
            if self.intensifier.num_run == 0:
                time_spent = time.time() - start_time
                time_left = self._get_timebound_for_intensification(time_spent, update=False)
                self.logger.debug('New intensification time bound: %f', time_left)
            else:
                old_time_left = time_left
                time_spent = time_spent + (time.time() - start_time)
                time_left = self._get_timebound_for_intensification(time_spent, update=True)
                self.logger.debug('Updated intensification time bound from %f to %f', old_time_left, time_left)
            self.neptune_run['intensification_time_bound'].log(time_left)

            # Skip starting new runs if the budget is now exhausted
            if self.stats.is_budget_exhausted():
                intent = RunInfoIntent.SKIP

            # Skip the run if there was a request to do so.
            # For example, during intensifier intensification, we
            # don't want to rerun a config that was previously ran
            if intent == RunInfoIntent.RUN:
                # Track the fact that a run was launched in the run
                # history. It's status is tagged as RUNNING, and once
                # completed and processed, it will be updated accordingly
                self.runhistory.add(
                    config=run_info.config,
                    cost=float(MAXINT),
                    time=0.0,
                    status=StatusType.RUNNING,
                    instance_id=run_info.instance,
                    seed=run_info.seed,
                    budget=run_info.budget,
                )

                run_info.config.config_id = self.runhistory.config_ids[run_info.config]

                self.tae_runner.submit_run(run_info=run_info)

                # There are 2 criteria that the stats object uses to know
                # if the budged was exhausted.
                # The budget time, which can only be known when the run finishes,
                # And the number of ta executions. Because we submit the job at this point,
                # we count this submission as a run. This prevent for using more
                # runner runs than what the scenario allows
                self.stats.submitted_ta_runs += 1

            elif intent == RunInfoIntent.SKIP:
                # No launch is required
                # This marks a transition request from the intensifier
                # To a new iteration
                pass
            elif intent == RunInfoIntent.WAIT:
                # In any other case, we wait for resources
                # This likely indicates that no further decision
                # can be taken by the intensifier until more data is
                # available
                self.tae_runner.wait()
            else:
                raise NotImplementedError("No other RunInfoIntent has been coded!")
            for k, v in run_info._asdict().items():
                if k == 'config':
                    if v is not None:
                        for config_k, config_v in v.get_dictionary().items():
                            self.neptune_run[f'run_info/{k}/{config_k}'].log(config_v, bo_iterations)
                    continue
                self.neptune_run[f'run_info/{k}'].log(v)

            # Check if there is any result, or else continue
            for run_info, result in self.tae_runner.get_finished_runs():

                # Add the results of the run to the run history
                # Additionally check for new incumbent
                self._incorporate_run_results(run_info, result, time_left)

            stats = {}
            for name, config in {'inc': self.incumbent, 'chal': self.intensifier.current_challenger}.items():
                perf = self.runhistory.get_cost(config)
                config_stats = {
                    'perf': perf,
                    'perf_normalized': self.scenario.normalize_cost(perf),
                    'runs': len(self.runhistory.get_runs_for_config(config, only_max_observed_budget=True)),
                    'config_id': config.config_id if config is not None else None,
                    'origin': config.origin if config is not None else None
                }
                stats.update({f'{name}/{k}': v for k, v in config_stats.items()})
            stats.update({
                'n_configs': self.stats.n_configs,
                'inc_changed': self.stats.inc_changed,
                'intensifier/N': self.intensifier.N,
                'intensifier/to_run': len(self.intensifier.to_run),
                'intensifier/stage': self.intensifier.stage.name,
                'runhistory_size': len(self.runhistory.data),
                'intent': intent.name,
                'ta_runs/submitted': self.stats.submitted_ta_runs,
                'ta_runs/finished': self.stats.finished_ta_runs
            })

            t.set_postfix({k: stats[k] for k in ['inc/perf', 'inc/runs', 'chal/perf', 'chal/runs', 'inc_changed',
                                                 'intensifier/N', 'intensifier/to_run', 'intensifier/stage', 'intent']})

            for k, v in stats.items():
                if k in ['inc/perf', 'inc/perf_normalized', 'chal/perf', 'chal/perf_normalized'] and not math.isfinite(
                        v):
                    v = -1
                if k in ['inc/config_id', 'chal/config_id'] and v is None:
                    v = -1
                self.neptune_run[k].log(v)

            if self.scenario.shared_model:  # type: ignore[attr-defined] # noqa F821
                assert self.scenario.output_dir_for_this_run is not None  # please mypy
                pSMAC.write(run_history=self.runhistory,
                            output_directory=self.scenario.output_dir_for_this_run,  # type: ignore[attr-defined] # noqa F821
                            logger=self.logger)

            self.logger.debug("Remaining budget: %f (wallclock), %f (ta costs), %f (target runs)" % (
                self.stats.get_remaing_time_budget(),
                self.stats.get_remaining_ta_budget(),
                self.stats.get_remaining_ta_runs()))

            self.stats.log(self.neptune_run)

            if self.stats.is_budget_exhausted() or self._stop:
                if self.stats.is_budget_exhausted():
                    self.logger.debug("Exhausted configuration budget")
                else:
                    self.logger.debug("Shutting down because a configuration or callback returned status STOP")

                # The budget can be exhausted  for 2 reasons: number of ta runs or
                # time. If the number of ta runs is reached, but there is still budget,
                # wait for the runs to finish
                while self.tae_runner.pending_runs():

                    self.tae_runner.wait()

                    for run_info, result in self.tae_runner.get_finished_runs():
                        # Add the results of the run to the run history
                        # Additionally check for new incumbent
                        self._incorporate_run_results(run_info, result, time_left)

                # Break from the intensification loop,
                # as there are no more resources
                break

            # print stats at the end of each intensification iteration
            if self.intensifier.iteration_done:
                self.stats.print_stats(debug_out=True)

            bo_iterations += 1
            t.update()
        t.close()

        self.neptune_run['incumbent'] = self.incumbent
        return self.incumbent

    def validate(self,
                 config_mode: typing.Union[str, typing.List[Configuration]] = 'inc',
                 instance_mode: typing.Union[str, typing.List[str]] = 'train+test',
                 repetitions: int = 1,
                 use_epm: bool = False,
                 n_jobs: int = -1,
                 backend: str = 'threading') -> RunHistory:
        """Create validator-object and run validation, using
        scenario-information, runhistory from smbo and tae_runner from intensify

        Parameters
        ----------
        config_mode: str or list<Configuration>
            string or directly a list of Configuration
            str from [def, inc, def+inc, wallclock_time, cpu_time, all]
            time evaluates at cpu- or wallclock-timesteps of:
            [max_time/2^0, max_time/2^1, max_time/2^3, ..., default]
            with max_time being the highest recorded time
        instance_mode: string
            what instances to use for validation, from [train, test, train+test]
        repetitions: int
            number of repetitions in nondeterministic algorithms (in
            deterministic will be fixed to 1)
        use_epm: bool
            whether to use an EPM instead of evaluating all runs with the TAE
        n_jobs: int
            number of parallel processes used by joblib

        Returns
        -------
        runhistory: RunHistory
            runhistory containing all specified runs
        """
        if isinstance(config_mode, str):
            assert self.scenario.output_dir_for_this_run is not None  # Please mypy
            traj_fn = os.path.join(self.scenario.output_dir_for_this_run, "traj_aclib2.json")
            trajectory = (
                TrajLogger.read_traj_aclib_format(fn=traj_fn, cs=self.config_space)
            )  # type: typing.Optional[typing.List[typing.Dict[str, typing.Union[float, int, Configuration]]]]
        else:
            trajectory = None
        if self.scenario.output_dir_for_this_run:
            new_rh_path = os.path.join(self.scenario.output_dir_for_this_run, "validated_runhistory.json")  # type: typing.Optional[str] # noqa E501
        else:
            new_rh_path = None

        validator = Validator(self.scenario, trajectory, self.rng)
        if use_epm:
            new_rh = validator.validate_epm(config_mode=config_mode,
                                            instance_mode=instance_mode,
                                            repetitions=repetitions,
                                            runhistory=self.runhistory,
                                            output_fn=new_rh_path)
        else:
            new_rh = validator.validate(config_mode, instance_mode, repetitions,
                                        n_jobs, backend, self.runhistory,
                                        self.tae_runner,
                                        output_fn=new_rh_path)
        return new_rh

    def _get_timebound_for_intensification(self, time_spent: float, update: bool) -> float:
        """Calculate time left for intensify from the time spent on
        choosing challengers using the fraction of time intended for
        intensification (which is specified in
        scenario.intensification_percentage).

        Parameters
        ----------
        time_spent : float

        update : bool
            Only used to check in the unit tests how this function was called

        Returns
        -------
        time_left : float
        """
        frac_intensify = self.scenario.intensification_percentage  # type: ignore[attr-defined] # noqa F821
        if frac_intensify <= 0 or frac_intensify >= 1:
            raise ValueError("The value for intensification_percentage-"
                             "option must lie in (0,1), instead: %.2f" %
                             frac_intensify)
        total_time = time_spent / (1 - frac_intensify)
        time_left = frac_intensify * total_time
        self.logger.debug("Total time: %.4f, time spent on choosing next "
                          "configurations: %.4f (%.2f), time left for "
                          "intensification: %.4f (%.2f)" %
                          (total_time, time_spent, (1 - frac_intensify), time_left, frac_intensify))
        return time_left

    def _incorporate_run_results(self, run_info: RunInfo, result: RunValue,
                                 time_left: float) -> None:
        """
        The SMBO submits a config-run-request via a RunInfo object.
        When that config run is completed, a RunValue, which contains
        all the relevant information obtained after running a job, is returned.
        This method incorporates the status of that run into
        the stats/runhistory objects so that other consumers
        can advance with their task.

        Additionally, it checks for a new incumbent via the intensifier process results,
        which also has the side effect of moving the intensifier to a new state

        Parameters
        ----------
        run_info: RunInfo
            Describes the run (config) from which to process the results
        result: RunValue
            Contains relevant information regarding the execution of a config
        time_left: float
            time in [sec] available to perform intensify
        """

        # update SMAC stats
        self.stats.ta_time_used += float(result.time)
        self.stats.finished_ta_runs += 1

        self.logger.debug(
            "Return: Status: %r, cost: %f, time: %f, additional: %s" % (
                result.status, result.cost, result.time, str(result.additional_info)
            )
        )

        self.runhistory.add(
            config=run_info.config,
            cost=result.cost,
            time=result.time,
            status=result.status,
            instance_id=run_info.instance,
            seed=run_info.seed,
            budget=run_info.budget,
            starttime=result.starttime,
            endtime=result.endtime,
            force_update=True,
            additional_info=result.additional_info,
        )
        self.stats.n_configs = len(self.runhistory.config_ids)

        if result.status == StatusType.ABORT:
            raise TAEAbortException("Target algorithm status ABORT - SMAC will "
                                    "exit. The last incumbent can be found "
                                    "in the trajectory-file.")
        elif result.status == StatusType.STOP:
            self._stop = True
            return

        if self.scenario.abort_on_first_run_crash:  # type: ignore[attr-defined] # noqa F821
            if self.stats.finished_ta_runs == 1 and result.status == StatusType.CRASHED:
                raise FirstRunCrashedException(
                    "First run crashed, abort. Please check your setup -- we assume that your default "
                    "configuration does not crash. (To deactivate this exception, use the SMAC scenario option "
                    "'abort_on_first_run_crash'). Additional run info: %s" % result.additional_info
                )

        # Update the intensifier with the result of the runs
        self.incumbent, inc_perf = self.intensifier.process_results(
            run_info=run_info,
            incumbent=self.incumbent,
            run_history=self.runhistory,
            time_bound=max(self._min_time, time_left),
            result=result,
        )

        for callback in self._callbacks['_incorporate_run_results']:
            response = callback(smbo=self, run_info=run_info, result=result, time_left=time_left)
            # If a callback returns False, the optimization loop should be interrupted
            # the other callbacks are still being called
            if response is False:
                self.logger.debug("An IncorporateRunResultCallback returned False, requesting abort.")
                self._stop = True

        if self.scenario.save_results_instantly:  # type: ignore[attr-defined] # noqa F821
            self.save()

        return

    def save(self) -> None:
        """
        Saves the current stats and runhistory.
        """
        self.stats.save()

        output_dir = self.scenario.output_dir_for_this_run
        if output_dir is not None:
            self.runhistory.save_json(
                fn=os.path.join(output_dir, "runhistory.json")
            )
            self.neptune_run['runhistory'].upload(os.path.join(output_dir, "runhistory.json"))


#Code to run for testing on new protein sequences

import matplotlib as mpl
import matplotlib.pyplot as plt
from xgboost import XGBClassifier
import xgboost as xgb
import numpy as np
import pandas as pd
np.random.seed(1337)
import io
import pickle
import csv
import argparse
import os
from Bio import SeqIO
import shap

parser = argparse.ArgumentParser()
parser.add_argument('file', type=argparse.FileType('r'))
parser.add_argument('file1', type=argparse.FileType('r'))
parser.add_argument('directory', help='the destination')
parser.add_argument("--encoding", default="utf_8")
args = parser.parse_args()

csv_reader = csv.reader(
    args.file,
    delimiter=",",
    quotechar='"'
)

Seq_id = []
with args.file1 as file:
        fasta_sequences = SeqIO.parse(file,'fasta')
        for i,fasta in enumerate(fasta_sequences, 1):
            if (i <= 5):
                name, sequence = fasta.id, str(fasta.seq)
                Seq_id.append(name)
            else:
                break

#Get and Manipulate Sequences
x = []
for row in csv_reader:
    x.append(row)

X_test1 = []
for i in range(1,len(x)):
    X_test1.append(x[i][1:])

df = pd.read_csv('Feature_Names.csv')
X1 = []
for i in range(len(df.values)):
    X1.append(df.values[i][0])

X_test = []
X_test.append(X1)
for i in range(len(X_test1)):
    X_test.append(np.array(X_test1[i]).tolist())

#Add feature name information on test set
df = pd.DataFrame(X_test)
df.to_csv("Metadata_Test.csv", header=False)
df = pd.read_csv('Metadata_Test.csv')
feature_names = []
for i in range(1, len(df.columns)):
    feature_names.append(df.columns[i])
X_test2 = df[feature_names]

X_test1 = X_test2
X_test2 = np.array(X_test2)
dtest = xgb.DMatrix(X_test2)

#Run the model on new test set
bs = xgb.Booster({'nthread': 4})  # init model
bs.load_model('xgb_cross_best_score_15_less90.model')  # load data
yprob = bs.predict(dtest)

scores = []
for i in range(len(yprob)):
    scores.append(yprob[i])

out_directory = 'mkdir '+args.directory
if os.path.exists(args.directory) is False:
	os.system(out_directory)

#Output the results in <output-dir>/prediction.csv file
with open(args.directory+'/prediction.csv', 'w') as f:
    write = csv.writer(f)
    write.writerows(zip(Seq_id, yprob))
f.close()

#Plot the top 10 features driving crystallization propensity for each protein in test file
colors = ["#7fffd4","#ff0000"]
for i in range(0, (len(Seq_id))):
    data_for_prediction = X_test1.iloc[[i]]
    data_for_prediction.columns = feature_names
    explainer = shap.TreeExplainer(bs)
    shap_values = explainer.shap_values(data_for_prediction.values)
    data_for_prediction = np.round(data_for_prediction, 3)
    shap.initjs()
    fig = shap.force_plot(explainer.expected_value, shap_values, data_for_prediction, plot_cmap=colors, matplotlib=True, show=False)
    plt.figure(figsize=(11.5, 8), dpi=300)

    x2 = np.argsort(np.abs(shap_values[0]))
    x2 = x2[::-1]
    yaxis = []
    for j in range(10):
        yaxis.append(shap_values[0][x2[j]])

    xax2 = []
    for j in range(10):
        xax2.append(feature_names[x2[j]] + "=" + str(np.round(X_test2[i][x2[j]], 3)))
    yaxis = yaxis[::-1]
    xaxis = xax2[::-1]
    yaxis_pos = []
    yaxis_neg = []
    xaxis_pos = []
    xaxis_neg = []
    colors = []
    labels = []
    for j in range(len(yaxis)):
        if(yaxis[j] >=0):
            colors.append("#7fffd4")
            yaxis_pos.append(yaxis[j])
            xaxis_pos.append(xaxis[j])
            labels.append("-> Crys")
        else:
            colors.append("#ff0000")
            yaxis_neg.append(yaxis[j])
            xaxis_neg.append(xaxis[j])
            labels.append("-> Non-Crys")
    if len(yaxis_pos) > 0 & len(yaxis_neg) > 0  :
        plt.barh(xaxis,yaxis, color=colors, label=labels)
        plt.legend(["-> Crys","-> Non-Crys"],loc=4)
    if len(yaxis_neg) > 0 & len(yaxis_pos)==0:
        plt.barh(xaxis_neg,yaxis_neg, color="#ff0000",label="-> Non-Crys")
        plt.legend(loc=4)
    if len(yaxis_pos)> 0 & len(yaxis_neg)==0:
        plt.barh(xaxis_pos,yaxis_pos, color="#7fffd4",label="-> Crys")
        plt.legend(loc=4)
    plt.xlim(-1.3, 1)
    plt.xlabel("SHAP Values")
    plt.title('Top Features for: '+Seq_id[i]+', score = '+str(yprob[i]))
    plt.savefig(args.directory+'/bar_plot_'+str(i+1)+'.png', bbox_inches='tight')

os.system('rm Metadata_Test.csv')
print("DONE")


import sys

sys.path.append(".")

import torch
import argparse
import torch.backends.cudnn as cudnn

import os
import numpy as np
from torchvision import datasets, transforms

os.environ['CUDA_VISIBLE_DEVICES'] = '-1'  # use only cpu for tf-lite evaluation
import tensorflow as tf

tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)

# Training settings
parser = argparse.ArgumentParser()
parser.add_argument('--tflite_path', default=None, help='path to tflite file for evaluation')

# dataset args.
parser.add_argument('--dataset', default='imagenet', type=str)
parser.add_argument('--val-dir', default='/dataset/imagenet/val',
                    help='path to validation data')
parser.add_argument('--batch-size', type=int, default=256,
                    help='input batch size for training')
parser.add_argument('-j', '--workers', default=16, type=int, metavar='N',
                    help='number of data loading workers')

args = parser.parse_args()

cudnn.benchmark = True

def get_val_dataset(resolution):
    # NOTE: we do not use normalization for tf-lite evaluation; the input is normalized to 0-1
    kwargs = {'num_workers': args.workers, 'pin_memory': False}
    if args.dataset == 'imagenet':
        val_dataset = \
            datasets.ImageFolder(args.val_dir,
                                 transform=transforms.Compose([
                                     transforms.Resize(int(resolution * 256 / 224)),
                                     transforms.CenterCrop(resolution),
                                     transforms.ToTensor(),
                                 ]))
    elif args.dataset == 'vww':
        raise NotImplementedError
        from lib.vww import VWWDataset
        val_transform = transforms.Compose([
            transforms.Resize([resolution, resolution]),
            transforms.ToTensor(),
        ])
        val_dataset = VWWDataset(split='minival', transform=val_transform)
    else:
        raise NotImplementedError
    val_loader = torch.utils.data.DataLoader(
        val_dataset, batch_size=args.batch_size,
        shuffle=False, **kwargs)
    return val_loader

def eval_image(data):
    image, target = data
    if len(image.shape) == 3:
        image = image.unsqueeze(0)
    image = image.permute(0, 2, 3, 1)
    image_np = image.cpu().numpy()
    image_np = (image_np * 255 - 128).astype(np.int8)
    interpreter.set_tensor(
        input_details[0]['index'], image_np.reshape(*input_shape))
    interpreter.invoke()
    output_data = interpreter.get_tensor(
        output_details[0]['index'])
    output = torch.from_numpy(output_data).view(1, -1)
    from utils import accuracy
    acc1, acc5 = accuracy(output, target.view(1), topk=(1, 5))

    return acc1.item(), acc5.item()

if __name__ == '__main__':
    interpreter = tf.lite.Interpreter(model_path=args.tflite_path)
    interpreter.allocate_tensors()

    # get input & output tensors
    input_details = interpreter.get_input_details()
    output_details = interpreter.get_output_details()

    input_shape = input_details[0]['shape']
    resolution = input_shape[1]

    # we first cache the whole test set into memory for faster data loading
    print(' * start caching the test set...', end='')
    val_loader = get_val_dataset(resolution)  # range [0, 1]
    val_loader_cache = [v for v in val_loader]
    images = torch.cat([v[0] for v in val_loader_cache], dim=0)
    targets = torch.cat([v[1] for v in val_loader_cache], dim=0)

    val_loader_cache = [[x, y] for x, y in zip(images, targets)]
    print('done.')
    print(' * dataset size:', len(val_loader_cache))

    # use multi-processing for faster evaluation
    n_thread = 32
    from multiprocessing import Pool
    from tqdm import tqdm

    p = Pool(n_thread)
    correctness1 = []
    correctness5 = []

    pbar = tqdm(p.imap_unordered(eval_image, val_loader_cache), total=len(val_loader_cache),
                desc='Evaluating...')
    for idx, (acc1, acc5) in enumerate(pbar):
        correctness1.append(acc1)
        correctness5.append(acc5)
        pbar.set_postfix({
            'top1': sum(correctness1) / len(correctness1),
            'top5': sum(correctness5) / len(correctness5),
        })
    print('* top1: {:.2f}%, top5: {:.2f}%'.format(
        sum(correctness1) / len(correctness1),
        sum(correctness5) / len(correctness5)
    ))