"""
This creates FOUR spectral plots used used in the appendix as Fig. 10-13.
This will only work for the four pre-selected spectra listed in __main__.

Dependencies
----------
github repository : https://github.com/bradlyke/utilities
Note: To get these plots, LaTeX must be installed and
      matplotlib must be able to compile LaTeX commands

Input file
----------
The most recent version of the DR16Q quasar-only catalog.

Parameters
----------

p : [10,11,12,13] The figure number from the paper

-h : Output help text to terminal (spec_plot.py -h)
-e : Plot the error spectrum in red
-k : Plot the sky spectrum in green
-x : Scale the sky flux to the quasar flux (so it's readable)
-s : Save the figure (does not display)

Output
----------
If selected, an EPS file of the plot.

"""

#Import the tools necessary to work with spectra files.
import cat_tools as ct #This is a tool I wrote, so you'll need the script in your path.
import numpy as np
import matplotlib.pyplot as plt
import matplotlib
import matplotlib.ticker as ticker
import sys
matplotlib.rc('text',usetex=True)

#A spectrum is a class. The object has 3 attributes right now: self, boxcar,
#paper_plot_small.
class spectrum:
    #initialize the object by loading the file. Needs a file name for infile.
    def __init__(self,infile):
        self.infile = infile #Load the filename as a thing for later use in plot saving.
        self.data = ct.file_load(infile) #Load the file as a numpy structured array.
        self.loglam = self.data['loglam'] #SDSS stores wavelengths as log10 values.
        self.lam = 10**self.loglam #Convert this to a decimal wavelength in Angstroms.
        self.flux = self.data['flux'] #Load the flux. Absolute, units are in plot.
        self.ivar = self.data['ivar'] #Load the inverse variance. Needed for plotting and boxcar.
        self.ferr = self.data['ivar']**(-0.5)

    '''
    We want to be able to smooth it fairly easily. Smooth_pct = 10 is good for
    emphasizing broad emission lines through noise. Smooth_pct = 1000 will
    remove (almost) all features and noise, leaving a useful continuum

    This uses a square window of width smooth_pct to perform a modified
    moving average on a spectrum using a boxcar (or "square" or "tophat")
    window. The box calculates a new flux for a point using the original
    signal flux, NOT the updated (smoothed) flux. This should only be used
    to make plots more readable, not for scientific analysis.
    '''
    def boxcar(self,flux_arr,flux_var,smooth_pct,weighted=False):
        spc = smooth_pct
        num_dpoints = len(flux_arr)
        #Initialize the new flux and error vectors.
        box_flux = np.zeros(num_dpoints,dtype='f8')
        box_err = np.zeros(num_dpoints,dtype='f8')

        #The window should be symmetrical, but won't be near the ends of the
        #spectrum, so it needs to be truncated properly, as padding won't
        #allow the edges to be properly averaged.
        for i in range(num_dpoints):
            #Define range of values to smooth
            #We have to check if the window can be symmetrical. If it is closer
            #to the edge of spectrum than half of smooth_pct it can't be
            #symmetrical, so truncate the denominator when averaging.
            if i < int(spc/2):
                lower = 0
            else:
                lower = -int(spc/2) + i
            if ((int(spc/2) + i) > num_dpoints):
                upper = num_dpoints
            else:
                upper = int(spc/2) + i
            #Smooth the values depending on smoothing type
            #Weighting use the flux error to do a weighted average, but if the
            #spectrum flux is smoothed, then so must the error spectrum, or else
            #it will be the wrong scale.
            #Note that this always calls a box of the ORIGINAL signal, not taking
            #previously smoothing pixels into account.
            if weighted==True:
                noise_temp = np.sqrt(flux_var[lower:upper])
                signal_temp = flux_arr[lower:upper]
                flux_temp = np.average(signal_temp,weights=noise_temp)
                ivar_temp = (np.average((signal_temp-flux_temp)**2, weights=noise_temp))**(-1.0)
            else:
                flux_temp = np.median(flux_arr[lower:upper])
                ivar_temp = flux_var[i]
            #If the smoothing is not weighted by error, the error isn't smoothed
            #and thus will no longer be properly scaled.

            #These hold the smoothed flux and (possibly) smoothed error.
            box_flux[i] = flux_temp
            box_err[i] = ivar_temp

        return box_flux,box_err

    #A "small paper plot" is just one of the attributes taken from the
    #specClass in the utilities repository. The emission line labels had
    #to be added by hand for each spectrum, so the original class wouldn't work
    #here.
    def paper_plot_small(self,z_in,spec_num,smooth=False,err=False,sky=False,scale_sky=False,save=False,rest=True):
        #These look at the data domain and range (wavelength and flux) and
        #find limits for the plot window.
        wobs = np.where((self.lam>=3700)&(self.lam<=10000))[0]
        flux_range = np.amax(self.flux[wobs]) - np.amin(self.flux[wobs])
        flux_pad = float(flux_range)/10
        y_lower = np.amin(self.flux[wobs]) - flux_pad
        y_upper = np.amax(self.flux[wobs]) + flux_pad
        x_lower = np.amin(self.lam)
        x_upper = np.amax(self.lam)

        '''
        For the four following spectra, the rest wavelength of the emission lines
        are taken from Vanden Berk et al. 2001 paper. That paper lists rest
        wavelengths in the lab and observed in spectra. The "official" SDSS list
        uses a blend of these two sources, so we chose lab or observed based on
        whichever SDSS chose.

        The following if/elif/else chain just builds the line_table array for
        the user-selected spectrum.
        '''

        #This is the FeLoBAL spectrum, Fig. 10.
        if spec_num==10:
            name_str = r'\textbf{SDSS J235134.38+031757.6, z = 2.230}'
            line_table = np.zeros(6,dtype=[('NAME','U26'),('LAM_REST','f8'),('LAM_OBS','f8'),('X_DISP','f8'),('Y_DISP','f8')])
            line_table['NAME'][0],line_table['LAM_REST'][0] = r'\textbf{Ly}$\alpha$',1216.25
            line_table['NAME'][1],line_table['LAM_REST'][1] = r'\textbf{N V}',1239.85
            line_table['NAME'][2],line_table['LAM_REST'][2] = r'\textbf{Si IV+O IV]}',1398.33
            line_table['NAME'][3],line_table['LAM_REST'][3] = r'\textbf{C IV}',1546.15
            line_table['NAME'][4],line_table['LAM_REST'][4] = r'\textbf{C III]}',1905.97
            line_table['NAME'][5],line_table['LAM_REST'][5] = r'\textbf{Mg II}',2800.26
            for i in range(6):
                #From the defined rest wavelength and given redshift, find the
                #observed wavelength.
                line_table['LAM_OBS'][i] = line_table['LAM_REST'][i] * (1+ z_in)
                #The exact calculated observed wavelength likely won't be in the data
                #so we need to find the closest point that is included.
                idx = np.argmin(np.absolute(line_table['LAM_OBS'][i] - self.lam)) #
                line_table['X_DISP'][i] = self.lam[idx] #The selected wavelength for the line
                line_table['Y_DISP'][i] = self.flux[idx] #The flux at that wavelength (hopefully the peak)

        #This is the BAL example spectrum, Fig. 11
        elif spec_num==11:
            name_str = r'\textbf{SDSS J003713.64+241121.5, z = 3.495}'
            line_table = np.zeros(6,dtype=[('NAME','U26'),('LAM_REST','f8'),('LAM_OBS','f8'),('X_DISP','f8'),('Y_DISP','f8')])
            line_table['NAME'][0],line_table['LAM_REST'][0] = r'\textbf{Ly}$\beta$',1033.03
            line_table['NAME'][1],line_table['LAM_REST'][1] = r'\textbf{Ly}$\alpha$',1216.25
            line_table['NAME'][2],line_table['LAM_REST'][2] = r'\textbf{O I}',1305.42
            line_table['NAME'][3],line_table['LAM_REST'][3] = r'\textbf{Si IV+O IV]}',1398.33
            line_table['NAME'][4],line_table['LAM_REST'][4] = r'\textbf{C IV}',1546.15
            line_table['NAME'][5],line_table['LAM_REST'][5] = r'\textbf{C III]}',1905.97
            for i in range(6):
                line_table['LAM_OBS'][i] = line_table['LAM_REST'][i] * (1+ z_in)
                idx = np.argmin(np.absolute(line_table['LAM_OBS'][i] - self.lam))
                line_table['X_DISP'][i] = self.lam[idx]
                line_table['Y_DISP'][i] = self.flux[idx]

        #This is the Mg II BAL spectrum, Fig. 12
        elif spec_num==12:
            name_str = r'\textbf{SDSS J212627.22+012321.0, z = 0.944}'
            line_table = np.zeros(8,dtype=[('NAME','U26'),('LAM_REST','f8'),('LAM_OBS','f8'),('X_DISP','f8'),('Y_DISP','f8')])
            line_table['NAME'][0],line_table['LAM_REST'][0] = r'\textbf{Fe II}',2626.92
            line_table['NAME'][1],line_table['LAM_REST'][1] = r'\textbf{Mg II}',2800.26
            line_table['NAME'][2],line_table['LAM_REST'][2] = r'\textbf{[O II]}',3729.66
            line_table['NAME'][3],line_table['LAM_REST'][3] = r'\textbf{H}$\delta$',4102.73
            line_table['NAME'][4],line_table['LAM_REST'][4] = r'\textbf{H}$\gamma$',4346.42
            line_table['NAME'][5],line_table['LAM_REST'][5] = r'\textbf{H}$\beta$',4853.13
            line_table['NAME'][6],line_table['LAM_REST'][6] = r'\textbf{[O III]}',4960.36
            line_table['NAME'][7],line_table['LAM_REST'][7] = r'\textbf{[O III]}',5008.22
            for i in range(8):
                line_table['LAM_OBS'][i] = line_table['LAM_REST'][i] * (1+ z_in)
                idx = np.argmin(np.absolute(line_table['LAM_OBS'][i] - self.lam))
                line_table['X_DISP'][i] = self.lam[idx]
                line_table['Y_DISP'][i] = self.flux[idx]

        #This is the LyA forest spectrum, Fig. 13
        elif spec_num==13:
            name_str = r'\textbf{SDSS J161016.71+411753.7, z = 5.005}'
            line_table = np.zeros(6,dtype=[('NAME','U26'),('LAM_REST','f8'),('LAM_OBS','f8'),('X_DISP','f8'),('Y_DISP','f8')])
            line_table['NAME'][0],line_table['LAM_REST'][0] = r'\textbf{Ly Limit}',912
            line_table['NAME'][1],line_table['LAM_REST'][1] = r'\textbf{Ly}$\beta$',1033.03
            line_table['NAME'][2],line_table['LAM_REST'][2] = r'\textbf{Ly}$\alpha$',1216.25
            line_table['NAME'][3],line_table['LAM_REST'][3] = r'\textbf{O I}',1305.42
            line_table['NAME'][4],line_table['LAM_REST'][4] = r'\textbf{Si IV+O IV]}',1398.33
            line_table['NAME'][5],line_table['LAM_REST'][5] = r'\textbf{C IV}',1546.15
            for i in range(6):
                line_table['LAM_OBS'][i] = line_table['LAM_REST'][i] * (1+ z_in)
                idx = np.argmin(np.absolute(line_table['LAM_OBS'][i] - self.lam))
                line_table['X_DISP'][i] = self.lam[idx]
                line_table['Y_DISP'][i] = self.flux[idx]

        #Do the "make it pretty" stuff
        matplotlib.rc('font',size=11)
        matplotlib.rcParams['text.latex.preamble'] = [r'\boldmath']
        fig1,ax1 = plt.subplots(figsize=(10,4))
        #We want the top and bottom x-axis to be wavelength, but different
        #reference frames (rest and observed respectively). Twiny() does this.
        ax1T = ax1.twiny() #This will be the rest wavelength x-axis object
        if smooth==True: #If smoothed, plot the smoothed on top of the raw flux.
            self.bflux,self.berr = self.boxcar(self.flux,self.ivar,10,weighted=True)
            ax1.plot(self.lam,self.flux,color='0.70',linewidth=0.8)
            ax1.plot(self.lam,self.bflux,color='black',linewidth=0.6)
        else:
            ax1.plot(self.lam,self.flux,color='black') #Or just raw flux if unsmoothed.
        if sky==True: #We can also plot the sky spectrum in green, if selected.
            if scale_sky==True:
                #Sky flux is usually much greater than the source flux, so we
                #can choose to scale it. The sky flux is less important than the
                #shape of the sky spectrum (to find patterns), so scaling is useful.
                self.sky = self.data['sky'] / 10.0
            else:
                self.sky = self.data['sky']
            ax1.plot(self.lam,self.sky,color='green',linewidth=0.7,alpha=0.5)
        if err==True: #And we can also plot the error spectrum in red if selected
            ax1.plot(self.lam,self.ferr,color='red',linewidth=0.6)
        ax1.set_xlim((x_lower,x_upper)) #Observerd wavelength x-axis
        ax1T.set_xlim((x_lower/(1+z_in),x_upper/(1+z_in))) #Rest wavelength x-axis
        ax1.set_ylim((y_lower,y_upper)) #The y-axis doesn't change.
        ax1.set_xlabel(r'\textbf{Observed Frame Wavelength (\AA)}')
        ax1T.set_xlabel(r'\textbf{Rest Frame Wavelength (\AA)}')
        ax1.set_ylabel(r'$f_{\lambda}$ ($10^{-17}$ \textbf{ergs s}$^{-1}$ \textbf{cm}$^{-2}$\,\textbf{\AA}$^{-1}$)')
        #The emission lines defined in the tables above need to be plotted with
        #marker lines and text. Because emission line redshift and quasar
        #redshift are always the same due to physical processes, these lines
        #needed to be hand-modified until they point to the right thing.
        #EXCEPT in Fig. 10. These lines are plotted to demonstrate how FeLoBAL
        #features do NOT align with common emission lines.
        if spec_num==10:
            ax1.annotate(line_table['NAME'][0],xy=(line_table['X_DISP'][0],line_table['Y_DISP'][0]),
                         xytext=(-10,30),textcoords='offset points',
                         bbox=dict(boxstyle='square,pad=0.2',fc='magenta',alpha=0.0),
                         arrowprops=dict(arrowstyle='-',connectionstyle='arc3,rad=0',color='blue'))
            #N V (line_table[1]) is usually blended with LyA, so the labels overlapped
            #ax1.annotate(line_table['NAME'][1],xy=(line_table['X_DISP'][1],line_table['Y_DISP'][1]),
                         #xytext=(-31,60),textcoords='offset points',
                         #bbox=dict(boxstyle='square,pad=0.2',fc='magenta',alpha=0.0),
                         #arrowprops=dict(arrowstyle='-',connectionstyle='arc3,rad=0',color='blue'))
            ax1.annotate(line_table['NAME'][2],xy=(line_table['X_DISP'][2],line_table['Y_DISP'][2]),
                         xytext=(-31,60),textcoords='offset points',
                         bbox=dict(boxstyle='square,pad=0.2',fc='magenta',alpha=0.0),
                         arrowprops=dict(arrowstyle='-',connectionstyle='arc3,rad=0',color='blue'))
            ax1.annotate(line_table['NAME'][3],xy=(line_table['X_DISP'][3],line_table['Y_DISP'][3]),
                         xytext=(-12,25),textcoords='offset points',
                         bbox=dict(boxstyle='square,pad=0.2',fc='magenta',alpha=0.0),
                         arrowprops=dict(arrowstyle='-',connectionstyle='arc3,rad=0',color='blue'))
            ax1.annotate(line_table['NAME'][4],xy=(line_table['X_DISP'][4]+10,line_table['Y_DISP'][4]+1),
                         xytext=(-13,30),textcoords='offset points',
                         bbox=dict(boxstyle='square,pad=0.2',fc='magenta',alpha=0.0),
                         arrowprops=dict(arrowstyle='-',connectionstyle='arc3,rad=0',color='blue'))
            ax1.annotate(line_table['NAME'][5],xy=(line_table['X_DISP'][5],line_table['Y_DISP'][5]+0.5),
                         xytext=(-14,50),textcoords='offset points',
                         bbox=dict(boxstyle='square,pad=0.2',fc='magenta',alpha=0.0),
                         arrowprops=dict(arrowstyle='-',connectionstyle='arc3,rad=0',color='blue'))
            #annotate calls mark the lines. The text call is what puts the
            #quasar name ('SDSS JHH:MM:SS.SS+DD:MM:SS.S) on the spectrum
            ax1.text(0.01,0.95,name_str,transform=ax1.transAxes, verticalalignment='top',
                    bbox=dict(boxstyle='square,pad=0.2',fc='magenta',alpha=0.0))
        elif spec_num==11:
            ax1.annotate(line_table['NAME'][0],xy=(line_table['X_DISP'][0],line_table['Y_DISP'][0]),
                         xytext=(-10,60),textcoords='offset points',
                         bbox=dict(boxstyle='square,pad=0.2',fc='magenta',alpha=0.0),
                         arrowprops=dict(arrowstyle='-',connectionstyle='arc3,rad=0',color='blue'))
            #The +35,+2 is the hand-modifying part, so the label matches the observed
            #position of the line in the plot. LyA and C IV are commonly shifted
            #a small amount from the value calculated using the "host" redshift
            ax1.annotate(line_table['NAME'][1],xy=(line_table['X_DISP'][1]+35,line_table['Y_DISP'][1]+2),
                         xytext=(-10,18),textcoords='offset points',
                         bbox=dict(boxstyle='square,pad=0.2',fc='magenta',alpha=0.0),
                         arrowprops=dict(arrowstyle='-',connectionstyle='arc3,rad=0',color='blue'))
            ax1.annotate(line_table['NAME'][2],xy=(line_table['X_DISP'][2],line_table['Y_DISP'][2]),
                         xytext=(-8,20),textcoords='offset points',
                         bbox=dict(boxstyle='square,pad=0.2',fc='magenta',alpha=0.0),
                         arrowprops=dict(arrowstyle='-',connectionstyle='arc3,rad=0',color='blue'))
            ax1.annotate(line_table['NAME'][3],xy=(line_table['X_DISP'][3],line_table['Y_DISP'][3]),
                         xytext=(-31,60),textcoords='offset points',
                         bbox=dict(boxstyle='square,pad=0.2',fc='magenta',alpha=0.0),
                         arrowprops=dict(arrowstyle='-',connectionstyle='arc3,rad=0',color='blue'))
            ax1.annotate(line_table['NAME'][4],xy=(line_table['X_DISP'][4],line_table['Y_DISP'][4]-2),
                         xytext=(-12,25),textcoords='offset points',
                         bbox=dict(boxstyle='square,pad=0.2',fc='magenta',alpha=0.0),
                         arrowprops=dict(arrowstyle='-',connectionstyle='arc3,rad=0',color='blue'))
            ax1.annotate(line_table['NAME'][5],xy=(line_table['X_DISP'][5]+10,line_table['Y_DISP'][5]),
                         xytext=(-13,20),textcoords='offset points',
                         bbox=dict(boxstyle='square,pad=0.2',fc='magenta',alpha=0.0),
                         arrowprops=dict(arrowstyle='-',connectionstyle='arc3,rad=0',color='blue'))
            ax1.text(0.6,0.95,name_str,transform=ax1.transAxes, verticalalignment='top',
                    bbox=dict(boxstyle='square,pad=0.2',fc='magenta',alpha=0.0))
        elif spec_num==12:
            ax1.annotate(line_table['NAME'][0],xy=(line_table['X_DISP'][0],line_table['Y_DISP'][0]),
                         xytext=(-11,18),textcoords='offset points',
                         bbox=dict(boxstyle='square,pad=0.2',fc='magenta',alpha=0.0),
                         arrowprops=dict(arrowstyle='-',connectionstyle='arc3,rad=0',color='blue'))
            ax1.annotate(line_table['NAME'][1],xy=(line_table['X_DISP'][1]+10,line_table['Y_DISP'][1]+6),
                         xytext=(-14,20),textcoords='offset points',
                         bbox=dict(boxstyle='square,pad=0.2',fc='magenta',alpha=0.0),
                         arrowprops=dict(arrowstyle='-',connectionstyle='arc3,rad=0',color='blue'))
            ax1.annotate(line_table['NAME'][2],xy=(line_table['X_DISP'][2],line_table['Y_DISP'][2]),
                         xytext=(-14,20),textcoords='offset points',
                         bbox=dict(boxstyle='square,pad=0.2',fc='magenta',alpha=0.0),
                         arrowprops=dict(arrowstyle='-',connectionstyle='arc3,rad=0',color='blue'))
            #H delta (line_table[3]) should be visible in the wavelength range,
            #but the flux was too small to see (looks like noise)
            #ax1.annotate(line_table['NAME'][3],xy=(line_table['X_DISP'][3],line_table['Y_DISP'][3]),
                         #xytext=(-31,60),textcoords='offset points',
                         #bbox=dict(boxstyle='square,pad=0.2',fc='magenta',alpha=0.0),
                         #arrowprops=dict(arrowstyle='-',connectionstyle='arc3,rad=0',color='blue'))
            ax1.annotate(line_table['NAME'][4],xy=(line_table['X_DISP'][4],line_table['Y_DISP'][4]),
                         xytext=(-7,25),textcoords='offset points',
                         bbox=dict(boxstyle='square,pad=0.2',fc='magenta',alpha=0.0),
                         arrowprops=dict(arrowstyle='-',connectionstyle='arc3,rad=0',color='blue'))
            ax1.annotate(line_table['NAME'][5],xy=(line_table['X_DISP'][5]+22,line_table['Y_DISP'][5]+2),
                         xytext=(-8,30),textcoords='offset points',
                         bbox=dict(boxstyle='square,pad=0.2',fc='magenta',alpha=0.0),
                         arrowprops=dict(arrowstyle='-',connectionstyle='arc3,rad=0',color='blue'))
            ax1.annotate('',xy=(line_table['X_DISP'][6],line_table['Y_DISP'][6]),
                         xytext=(0,89),textcoords='offset points',
                         bbox=dict(boxstyle='square,pad=0.2',fc='magenta',alpha=0.0),
                         arrowprops=dict(arrowstyle='-',connectionstyle='arc3,rad=0',color='blue'))
            ax1.annotate(line_table['NAME'][7],xy=(line_table['X_DISP'][7],line_table['Y_DISP'][7]-2.5),
                         xytext=(-16,25),textcoords='offset points',
                         bbox=dict(boxstyle='square,pad=0.2',fc='magenta',alpha=0.0),
                         arrowprops=dict(arrowstyle='-',connectionstyle='arc3,rad=0',color='blue'))
            ax1.text(0.01,0.95,name_str,transform=ax1.transAxes, verticalalignment='top',
                    bbox=dict(boxstyle='square,pad=0.2',fc='magenta',alpha=0.0))
        elif spec_num==13:
            ax1.annotate(line_table['NAME'][0],xy=(line_table['X_DISP'][0],line_table['Y_DISP'][0]),
                         xytext=(-22,30),textcoords='offset points',
                         bbox=dict(boxstyle='square,pad=0.2',fc='magenta',alpha=0.0),
                         arrowprops=dict(arrowstyle='-',connectionstyle='arc3,rad=0',color='blue'))
            ax1.annotate(line_table['NAME'][1],xy=(line_table['X_DISP'][1]+35,line_table['Y_DISP'][1]+2),
                         xytext=(-10,18),textcoords='offset points',
                         bbox=dict(boxstyle='square,pad=0.2',fc='magenta',alpha=0.0),
                         arrowprops=dict(arrowstyle='-',connectionstyle='arc3,rad=0',color='blue'))
            ax1.annotate(line_table['NAME'][2],xy=(line_table['X_DISP'][2]-10,line_table['Y_DISP'][2]+2),
                         xytext=(-10,20),textcoords='offset points',
                         bbox=dict(boxstyle='square,pad=0.2',fc='magenta',alpha=0.0),
                         arrowprops=dict(arrowstyle='-',connectionstyle='arc3,rad=0',color='blue'))
            ax1.annotate(line_table['NAME'][3],xy=(line_table['X_DISP'][3],line_table['Y_DISP'][3]),
                         xytext=(-8,30),textcoords='offset points',
                         bbox=dict(boxstyle='square,pad=0.2',fc='magenta',alpha=0.0),
                         arrowprops=dict(arrowstyle='-',connectionstyle='arc3,rad=0',color='blue'))
            ax1.annotate(line_table['NAME'][4],xy=(line_table['X_DISP'][4]-30,line_table['Y_DISP'][4]+0.5),
                         xytext=(-31,45),textcoords='offset points',
                         bbox=dict(boxstyle='square,pad=0.2',fc='magenta',alpha=0.0),
                         arrowprops=dict(arrowstyle='-',connectionstyle='arc3,rad=0',color='blue'))
            ax1.annotate(line_table['NAME'][5],xy=(line_table['X_DISP'][5],line_table['Y_DISP'][5]+0.5),
                         xytext=(-12,30),textcoords='offset points',
                         bbox=dict(boxstyle='square,pad=0.2',fc='magenta',alpha=0.0),
                         arrowprops=dict(arrowstyle='-',connectionstyle='arc3,rad=0',color='blue'))
            ax1.text(0.01,0.95,name_str,transform=ax1.transAxes, verticalalignment='top',
                    bbox=dict(boxstyle='square,pad=0.2',fc='magenta',alpha=0.0))
        #Ticks point inward in physics and astronomy
        #And minor ticks should be displayed
        ax1.tick_params(axis='both',direction='in')
        ax1.tick_params(axis='both',which='minor',direction='in')
        ax1.tick_params(right=True)
        ax1.tick_params(which='minor',right=True)
        ax1T.tick_params(direction='in')
        ax1T.tick_params(which='minor',direction='in')
        ax1.xaxis.set_minor_locator(ticker.MultipleLocator(100))
        ax1T.xaxis.set_minor_locator(ticker.MultipleLocator(50))
        ax1.yaxis.set_major_locator(ticker.MultipleLocator(5))
        ax1.yaxis.set_minor_locator(ticker.MultipleLocator(1))

        #Show or save?
        if save==True:
            fname_out = self.infile.replace('.fits','.eps')
            fname_out = fname_out.replace('../data/','../plots/')
            fig1.savefig(fname_out,bbox_inches='tight',pad_inches=0.03,format='eps')
            plt.close()
        else:
            plt.tight_layout()
            plt.show()

# There are a lot of options here. Use python spec_plot.py -h to see them.
# This automatically generates the name of the eps plot using the spectrum file
# name.
if __name__=='__main__':
    import argparse
    parser = argparse.ArgumentParser(description='Plot one of the spectra in Lyke et al. 2020')
    parser.add_argument('p', type=int, choices=[10,11,12,13], help='Paper figure number')
    parser.add_argument('-e', '--error', action='store_true',
                        help='Include the error spectrum in red')
    parser.add_argument('-k', '--sky', action='store_true',
                        help='Include the sky spectrum in green')
    parser.add_argument('-x', '--scale_sky', action='store_true',
                        help='Scale the sky flux to the quasar flux')
    parser.add_argument('-s', '--save', action='store_true',
                        help='Save the plot, without display')

    args = parser.parse_args()

    #Set up the plate-mjd-fiberid combinations with redshifts.
    spec_to_plot = np.zeros(4,dtype=[('FIG_NUM','i2'),('PMF','U16'),('Z','f8')])
    spec_to_plot['FIG_NUM'][0],spec_to_plot['PMF'][0],spec_to_plot['Z'][0] = 10,'11278-58395-0576',2.23
    spec_to_plot['FIG_NUM'][1],spec_to_plot['PMF'][1],spec_to_plot['Z'][1] = 11,'7672-57339-0394',3.495
    spec_to_plot['FIG_NUM'][2],spec_to_plot['PMF'][2],spec_to_plot['Z'][2] = 12,'9162-58040-0354',0.944
    spec_to_plot['FIG_NUM'][3],spec_to_plot['PMF'][3],spec_to_plot['Z'][3] = 13,'8528-57896-0104',5.005
    '''
    NOTE: The spectrum for 8528-57896-0104 is not the PRIMARY record in DR16Q. That's listed under: 6044-56090-0418
          The spectra, however, are pretty identical.
    NOTE: The spectrum for 11278-58395-0576 is not the PRIMARY record in DR16Q. That's listed under: 8741-57390-0450
          The spectra is identical, but that record has a Z_VI of 2.43. We are going to trust Vivek at Z = 2.23
    '''
    #This chooses which of the four spectra to use, then grabs the redshift.
    w = np.where(spec_to_plot['FIG_NUM']==args.p)[0]
    red = spec_to_plot['Z'][w]
    #Make the spectrum filename.
    spec_name = '../data/spec-{}.fits'.format(spec_to_plot['PMF'][w[0]])
    spec = spectrum(spec_name) #Load the spectrum using the class functionality.
    #Use the class attribute for spectra to make a 10x4 plot.
    spec.paper_plot_small(red,args.p,smooth=True,err=args.error,sky=args.sky,
                        scale_sky=args.scale_sky,save=args.save)


"""Contains classes to represent non-equilibrium ionization simulations."""

__all__ = ["NEI", "NEIError", "SimulationResults"]

from typing import Callable, Dict, List, Optional, Union

import astropy.units as u
import numpy as np
from scipy import interpolate, optimize

from plasmapy_nei.eigen import EigenData, eigen_data_dict

try:
    from plasmapy.atomic import IonizationStates, atomic_number
except ImportError:
    from plasmapy.particles import IonizationStates, atomic_number

import warnings

# TODO: Allow this to keep track of velocity and position too, and
# eventually to have density and temperature be able to be functions of
# position.  (and more complicated expressions for density and
# temperature too)

# TODO: Expand Simulation docstring

# TODO: Include the methods in the original Visualize class which is a
#       subclass of NEI in the NEI-modeling/NEI repo. These were deleted
#       temporarily to make it possible to get the NEI class itself
#       adapted into this package.

# TODO: In this file and test_nei.py, there are a few places with
#       initial.ionic_fractions.keys(), where initial is an instance
#       of IonizationStates.  This workaround exists because I forgot
#       to put in an `elements` attribute in IonizationStates, and
#       should be corrected.

class NEIError(Exception):
    """For when there are errors in setting up or performing NEI simulations."""

    pass

class SimulationResults:
    """
    Results from a non-equilibrium ionization simulation.

    Parameters
    ----------
    initial: plasmapy.atomic.IonizationStates
        The ``IonizationStates`` instance representing the ionization
        states of different elements and plasma properties as the
        initial conditions.

    n_init: astropy.units.Quantity
        The initial number density scaling factor.

    T_e_init: astropy.units.Quantity
        The initial electron temperature.

    max_steps: int
        The maximum number of time steps that the simulation can take
        before stopping.

    time_start: astropy.units.Quantity
        The time at the start of the simulation.

    """

    def __init__(
        self,
        initial: IonizationStates,
        n_init: u.Quantity,
        T_e_init: u.Quantity,
        max_steps: int,
        time_start: u.Quantity,
    ):

        self._elements = list(initial.ionic_fractions.keys())
        self._abundances = initial.abundances
        self._max_steps = max_steps

        self._nstates = {elem: atomic_number(elem) + 1 for elem in self.elements}

        self._ionic_fractions = {
            elem: np.full((max_steps + 1, self.nstates[elem]), np.nan, dtype=np.float64)
            for elem in self.elements
        }

        self._number_densities = {
            elem: np.full((max_steps + 1, self.nstates[elem]), np.nan, dtype=np.float64)
            * u.cm ** -3
            for elem in self.elements
        }

        self._n_elem = {
            elem: np.full(max_steps + 1, np.nan) * u.cm ** -3 for elem in self.elements
        }

        self._n_e = np.full(max_steps + 1, np.nan) * u.cm ** -3
        self._T_e = np.full(max_steps + 1, np.nan) * u.K
        self._time = np.full(max_steps + 1, np.nan) * u.s

        self._index = 0

        self._assign(
            new_time=time_start,
            new_ionfracs=initial.ionic_fractions,
            new_n=n_init,
            new_T_e=T_e_init,
        )

    def _assign(
        self,
        new_time: u.Quantity,
        new_ionfracs: Dict[str, np.ndarray],
        new_n: u.Quantity,
        new_T_e: u.Quantity,
    ):
        """
        Store results from a time step of a non-equilibrium ionization
        time advance in the `~plasmapy_nei.classes.NEI` class.

        Parameters
        ----------
        new_time
            The time associated with this time step.

        new_ionfracs: dict
            The new ionization fractions for this time step.  The keys
            of this `dict` are the atomic symbols of the elements being
            tracked, and with the corresponding value being an
            ``numpy.ndarray`` representing the ionic fractions.  Each
            element's array must have a length of the atomic number plus
            one, and be normalized to one with all values between zero
            and one.

        new_n
            The new number density scaling factor for this time step.
            The number densities of each ionic species will be the
            product of this scaling factor, the element's abundance, and
            the ionic fraction given in ``new_ionfracs``.

        new_T_e
            The new electron temperature.

        """

        try:
            index = self._index
            elements = self.elements
            self._time[index] = new_time
            self._T_e[index] = new_T_e

            for elem in elements:
                self._ionic_fractions[elem][index, :] = new_ionfracs[elem][:]

            # Calculate elemental and ionic number densities
            n_elem = {elem: new_n * self.abundances[elem] for elem in elements}
            number_densities = {
                elem: n_elem[elem] * new_ionfracs[elem] for elem in elements
            }

            # Calculate the electron number density
            n_e = 0.0 * u.cm ** -3
            for elem in elements:
                integer_charges = np.linspace(
                    0, self.nstates[elem] - 1, self.nstates[elem]
                )
                n_e += np.sum(number_densities[elem] * integer_charges)

            # Assign densities
            self._n_e[index] = n_e
            for elem in elements:
                self._n_elem[elem][index] = n_elem[elem]
                self._number_densities[elem][index, :] = number_densities[elem]

        except Exception as exc:
            raise NEIError(
                f"Unable to assign parameters to Simulation instance "
                f"for index {index} at time = {new_time}.  The "
                f"parameters are new_n = {new_n}, new_T_e = {new_T_e}, "
                f"and new_ionic_fractions = {new_ionfracs}."
            ) from exc
        finally:
            self._index += 1

    def _cleanup(self):
        """
        Clean up this class after the simulation is complete.

        This method removes the excess elements from each array that
        did not end up getting used for a time step in the simulation
        and sets the ``last_step`` attribute.

        """
        nsteps = self._index

        self._n_e = self._n_e[0:nsteps]
        self._T_e = self._T_e[0:nsteps]
        self._time = self._time[0:nsteps]

        for element in self.elements:
            self._ionic_fractions[element] = self._ionic_fractions[element][0:nsteps, :]
            self._number_densities[element] = self._number_densities[element][
                0:nsteps, :
            ]

        self._last_step = nsteps - 1

        self._index = None

    @property
    def max_steps(self) -> int:
        """
        The maximum number of time steps allowed for this simulation.
        """
        return self._max_steps

    @property
    def last_step(self) -> int:
        """The time index of the last step."""
        return self._last_step

    @property
    def nstates(self) -> Dict[str, int]:
        """
        Return the dictionary containing atomic symbols as keys and the
        number of ionic species for the corresponding element as the
        value.
        """
        return self._nstates

    @property
    def elements(self) -> List[str]:
        """The elements modeled by this simulation."""
        return self._elements

    @property
    def abundances(self) -> Dict[str, float]:
        """
        The relative elemental abundances of the elements modeled in
        this simulation.

        The keys are the atomic symbols and the values are a `float`
        representing that element's elemental abundance.
        """
        return self._abundances

    @property
    def ionic_fractions(self) -> Dict[str, np.ndarray]:
        """
        Return the ionic fractions over the course of the simulation.

        The keys of this dictionary are atomic symbols.  The values are
        2D arrays where the first index refers to the time step and the
        second index refers to the integer charge.
        """
        return self._ionic_fractions

    @property
    def number_densities(self) -> Dict[str, u.Quantity]:
        """
        Return the number densities over the course of the simulation.

        The keys of ``number_densities`` are atomic symbols.  The values
        are 2D arrays with units of number density where the first index
        refers to the time step and the second index is the integer
        charge.

        """
        return self._number_densities

    @property
    def n_elem(self) -> Dict[str, u.Quantity]:
        """
        The number densities of each element over the course of the
        simulation.

        The keys of ``n_elem`` are atomic symbols.  The values are 1D
        arrays with units of number density where the index refers to
        the time step.

        """
        return self._n_elem

    @property
    def n_e(self) -> u.Quantity:
        """
        The electron number density over the course of the simulation in
        units of number density.

        The index of this array corresponds to the time step.
        """
        return self._n_e

    @property
    def T_e(self) -> u.Quantity:
        """
        The electron temperature over the course of the simulation in
        kelvin.

        The index of this array corresponds to the time step.
        """
        return self._T_e

    @property
    def time(self) -> u.Quantity:
        """
        The time for each time step over the course of the simulation
        in units of seconds.
        """
        return self._time

class NEI:
    r"""
    Perform and analyze a non-equilibrium ionization simulation.

    Parameters
    ----------
    inputs

    T_e: astropy.units.Quantity or callable
        The electron temperature, which may be a constant, an array of
        temperatures corresponding to the times in `time_input`, or a
        function that yields the temperature as a function of time.

    n: astropy.units.Quantity or callable
        The number density multiplicative factor.  The number density of
        each element will be ``n`` times the abundance given in
        ``abundances``.  For example, if ``abundance['H'] = 1``, then this
        will correspond to the number density of hydrogen (including
        neutral hydrogen and protons).  This factor may be a constant,
        an array of number densities over time, or a function that
        yields a number density as a function of time.

    time_input: astropy.units.Quantity, optional
        An array containing the times associated with ``n`` and ``T_e`` in
        units of time.

    time_start: astropy.units.Quantity, optional
        The start time for the simulation.  If density and/or
        temperature are given by arrays, then this argument must be
        greater than ``time_input[0]``.  If this argument is not supplied,
        then ``time_start`` defaults to ``time_input[0]`` (if given) and
        zero seconds otherwise.

    time_max: astropy.units.Quantity
        The maximum time for the simulation.  If density and/or
        temperature are given by arrays, then this argument must be less
        than ``time_input[-1]``.

    max_steps: `int`
        The maximum number of time steps to be taken during a
        simulation.

    dt: astropy.units.Quantity
        The time step.  If ``adapt_dt`` is `False`, then ``dt`` is the
        time step for the whole simulation.

    dt_max: astropy.units.Quantity
        The maximum time step to be used with an adaptive time step.

    dt_min: astropy.units.Quantity
        The minimum time step to be used with an adaptive time step.

    adapt_dt: `bool`
        If `True`, change the time step based on the characteristic
        ionization and recombination time scales and change in
        temperature.  Not yet implemented.

    safety_factor: `float` or `int`
        A multiplicative factor to multiply by the time step when
        ``adapt_dt`` is `True`.  Lower values improve accuracy, whereas
        higher values reduce computational time.  Not yet implemented.

    tol: float
        The absolute tolerance to be used in comparing ionic fractions.

    verbose: bool, optional
        A flag stating whether or not to print out information for every
        time step. Setting ``verbose`` to `True` is useful for testing.
        Defaults to `False`.

    abundances: dict

    Examples
    --------

    >>> import numpy as np
    >>> import astropy.units as u

    >>> inputs = {'H': [0.9, 0.1], 'He': [0.9, 0.099, 0.001]}
    >>> abund = {'H': 1, 'He': 0.085}
    >>> n = u.Quantity([1e9, 1e8], u.cm**-3)
    >>> T_e = np.array([10000, 40000]) * u.K
    >>> time = np.array([0, 300]) * u.s
    >>> dt = 0.25 * u.s

    The initial conditions can be accessed using the initial attribute.

    >>> sim = NEI(inputs=inputs, abundances=abund, n=n, T_e=T_e, time_input=time, adapt_dt=False, dt=dt)

    After having inputted all of the necessary information, we can run
    the simulation.

    >>> results = sim.simulate()

    The initial results are stored in the ``initial`` attribute.

    >>> sim.initial.ionic_fractions['H']
    array([0.9, 0.1])

    The final results can be access with the ``final`` attribute.

    >>> sim.final.ionic_fractions['H']
    array([0.16665179, 0.83334821])
    >>> sim.final.ionic_fractions['He']
    array([0.88685261, 0.11218358, 0.00096381])
    >>> sim.final.T_e
    <Quantity 40000. K>

    Both ``initial`` and ``final`` are instances of the ``IonizationStates``
    class.

    Notes
    -----
    The ionization and recombination rates are from Chianti version
    8.7.  These rates include radiative and dielectronic recombination.
    Photoionization is not included.
    """

    def __init__(
        self,
        inputs,
        abundances: Union[Dict, str] = None,
        T_e: Union[Callable, u.Quantity] = None,
        n: Union[Callable, u.Quantity] = None,
        time_input: u.Quantity = None,
        time_start: u.Quantity = None,
        time_max: u.Quantity = None,
        max_steps: Union[int, np.integer] = 10000,
        tol: Union[int, float] = 1e-15,
        dt: u.Quantity = None,
        dt_max: u.Quantity = np.inf * u.s,
        dt_min: u.Quantity = 0 * u.s,
        adapt_dt: bool = None,
        safety_factor: Union[int, float] = 1,
        verbose: bool = False,
    ):

        try:

            self.time_input = time_input
            self.time_start = time_start
            self.time_max = time_max
            self.T_e_input = T_e
            self.n_input = n
            self.max_steps = max_steps
            self.dt_input = dt

            if self.dt_input is None:
                self._dt = self.time_max / max_steps
            else:
                self._dt = self.dt_input

            self.dt_min = dt_min
            self.dt_max = dt_max
            self.adapt_dt = adapt_dt
            self.safety_factor = safety_factor
            self.verbose = verbose

            T_e_init = self.electron_temperature(self.time_start)
            n_init = self.hydrogen_number_density(self.time_start)

            self.initial = IonizationStates(
                inputs=inputs,
                abundances=abundances,
                T_e=T_e_init,
                n=n_init,
                tol=tol,
            )

            self.tol = tol

            # TODO: Update IonizationStates in PlasmaPy to have elements attribute

            self.elements = list(self.initial.ionic_fractions.keys())

            if "H" not in self.elements:
                raise NEIError("Must have H in elements")

            self.abundances = self.initial.abundances

            self._eigen_data_dict = eigen_data_dict

            if self.T_e_input is not None and not isinstance(inputs, dict):
                for element in self.initial.ionic_fractions.keys():
                    self.initial.ionic_fractions[element] = self.eigen_data_dict[
                        element
                    ].equilibrium_state(T_e_init.value)

            self._temperature_grid = self._eigen_data_dict[
                self.elements[0]
            ].temperature_grid

            self._get_temperature_index = self._eigen_data_dict[
                self.elements[0]
            ]._get_temperature_index

            self._results = None

        except Exception as e:
            raise NEIError(
                f"Unable to create NEI object for:\n"
                f"     inputs = {inputs}\n"
                f" abundances = {abundances}\n"
                f"        T_e = {T_e}\n"
                f"          n = {n}\n"
                f" time_input = {time_input}\n"
                f" time_start = {time_start}\n"
                f"   time_max = {time_max}\n"
                f"  max_steps = {max_steps}\n"
            ) from e

    def equil_ionic_fractions(
        self,
        T_e: u.Quantity = None,
        time: u.Quantity = None,
    ) -> Dict[str, np.ndarray]:
        """
        Return the equilibrium ionic fractions for a temperature or at
        a given time.

        Parameters
        ----------
        T_e: astropy.units.Quantity, optional
            The electron temperature in units that can be converted to
            kelvin.

        time: astropy.units.Quantity, optional
            The time in units that can be converted to seconds.

        Returns
        -------
        equil_ionfracs: `dict`
            The equilibrium ionic fractions for the elements contained
            within this class

        Notes
        -----
        Only one of ``T_e`` and ``time`` may be included as an argument.
        If neither ``T_e`` or ``time`` is provided and the temperature
        for the simulation is given by a constant, the this method will
        assume that ``T_e`` is the temperature of the simulation.
        """

        if T_e is not None and time is not None:
            raise NEIError("Only one of T_e and time may be an argument.")

        if T_e is None and time is None:
            if self.T_e_input.isscalar:
                T_e = self.T_e_input
            else:
                raise NEIError

        try:
            T_e = T_e.to(u.K) if T_e is not None else None
            time = time.to(u.s) if time is not None else None
        except Exception as exc:
            raise NEIError("Invalid input to equilibrium_ionic_fractions.") from exc

        if time is not None:
            T_e = self.electron_temperature(time)

        if not T_e.isscalar:
            raise NEIError("Need scalar input for equil_ionic_fractions.")

        equil_ionfracs = {}
        for element in self.elements:
            equil_ionfracs[element] = self.eigen_data_dict[element].equilibrium_state(
                T_e.value
            )

        return equil_ionfracs

    @property
    def elements(self) -> List[str]:
        """A `list` of the elements."""
        return self._elements

    @elements.setter
    def elements(self, elements):
        # TODO: Update this
        self._elements = elements

    @property
    def abundances(self) -> Dict[str, Union[float, int]]:
        """Return the abundances."""
        return self._abundances

    @abundances.setter
    def abundances(self, abund: Dict[Union[str, int], Union[float, int]]):

        # TODO: Update initial, etc. when abundances is updated. The
        # checks within IonizationStates will also be checks for

        # TODO: Update initial and other attributes when abundances is
        # updated.

        self._abundances = abund

    @property
    def tol(self) -> float:
        """
        The tolerance for comparisons between different ionization
        states.
        """
        return self._tol

    @tol.setter
    def tol(self, value: Union[float, int]):
        try:
            value = float(value)
        except Exception as exc:
            raise TypeError(f"Invalid tolerance: {value}") from exc
        if not 0 <= value < 1:
            raise ValueError("Need 0 <= tol < 1.")
        self._tol = value

    @property
    def time_input(self) -> u.s:
        return self._time_input

    @time_input.setter
    def time_input(self, times: u.s):
        if times is None:
            self._time_input = None
        elif isinstance(times, u.Quantity):
            if times.isscalar:
                raise ValueError("time_input must be an array.")
            try:
                times = times.to(u.s)
            except u.UnitConversionError:
                raise u.UnitsError("time_input must have units of seconds.") from None
            if not np.all(times[1:] > times[:-1]):
                raise ValueError("time_input must monotonically increase.")
            self._time_input = times
        else:
            raise TypeError("Invalid time_input.")

    @property
    def time_start(self) -> u.s:
        """The start time of the simulation."""
        return self._time_start

    @time_start.setter
    def time_start(self, time: u.s):
        if time is None:
            self._time_start = 0.0 * u.s
        elif isinstance(time, u.Quantity):
            if not time.isscalar:
                raise ValueError("time_start must be a scalar")
            try:
                time = time.to(u.s)
            except u.UnitConversionError:
                raise u.UnitsError("time_start must have units of seconds") from None
            if (
                hasattr(self, "_time_max")
                and self._time_max is not None
                and self._time_max <= time
            ):
                raise ValueError("Need time_start < time_max.")
            if self.time_input is not None and self.time_input.min() > time:
                raise ValueError("time_start must be less than min(time_input)")
            self._time_start = time
        else:
            raise TypeError("Invalid time_start.") from None

    @property
    def time_max(self) -> u.s:
        """The maximum time allowed for the simulation."""
        return self._time_max

    @time_max.setter
    def time_max(self, time: u.s):
        if time is None:
            self._time_max = (
                self.time_input[-1] if self.time_input is not None else np.inf * u.s
            )
        elif isinstance(time, u.Quantity):
            if not time.isscalar:
                raise ValueError("time_max must be a scalar")
            try:
                time = time.to(u.s)
            except u.UnitConversionError:
                raise u.UnitsError("time_max must have units of seconds") from None
            if (
                hasattr(self, "_time_start")
                and self._time_start is not None
                and self._time_start >= time
            ):
                raise ValueError("time_max must be greater than time_start")
            self._time_max = time
        else:
            raise TypeError("Invalid time_max.") from None

    @property
    def adapt_dt(self) -> Optional[bool]:
        """
        Return `True` if the time step is set to be adaptive, `False`
        if the time step is set to not be adapted, and `None` if this
        attribute was not set.
        """
        return self._adapt_dt

    @adapt_dt.setter
    def adapt_dt(self, choice: Optional[bool]):
        if choice is None:
            self._adapt_dt = True if self.dt_input is None else False
        elif choice is True or choice is False:
            self._adapt_dt = choice
        else:
            raise TypeError("Invalid value for adapt_dt")

    @property
    def dt_input(self) -> u.s:
        """Return the inputted time step."""
        return self._dt_input

    @dt_input.setter
    def dt_input(self, dt: u.s):
        if dt is None:
            self._dt_input = None
        elif isinstance(dt, u.Quantity):
            try:
                dt = dt.to(u.s)
                if dt > 0 * u.s:
                    self._dt_input = dt
            except (AttributeError, u.UnitConversionError):
                raise NEIError("Invalid dt.")

    @property
    def dt_min(self) -> u.s:
        """The minimum time step."""
        return self._dt_min

    @dt_min.setter
    def dt_min(self, value: u.s):
        if not isinstance(value, u.Quantity):
            raise TypeError("dt_min must be a Quantity.")
        try:
            value = value.to(u.s)
        except u.UnitConversionError as exc:
            raise u.UnitConversionError("Invalid units for dt_min.") from exc
        if (
            hasattr(self, "_dt_input")
            and self.dt_input is not None
            and self.dt_input < value
        ):
            raise ValueError("dt_min cannot exceed the inputted time step.")
        if hasattr(self, "_dt_max") and self.dt_max < value:
            raise ValueError("dt_min cannot exceed dt_max.")
        self._dt_min = value

    @property
    def dt_max(self) -> u.s:
        return self._dt_max

    @dt_max.setter
    def dt_max(self, value: u.s):
        if not isinstance(value, u.Quantity):
            raise TypeError("dt_max must be a Quantity.")
        try:
            value = value.to(u.s)
        except u.UnitConversionError as exc:
            raise u.UnitConversionError("Invalid units for dt_max.") from exc
        if (
            hasattr(self, "_dt_input")
            and self.dt_input is not None
            and self.dt_input > value
        ):
            raise ValueError("dt_max cannot be less the inputted time step.")
        if hasattr(self, "_dt_min") and self.dt_min > value:
            raise ValueError("dt_min cannot exceed dt_max.")
        self._dt_max = value

    @property
    def safety_factor(self):
        """
        The multiplicative factor that the time step is to be multiplied
        by when using an adaptive time step.
        """
        return self._safety_factor

    @safety_factor.setter
    def safety_factor(self, value):
        if not isinstance(value, (float, np.float64, np.integer, int)):
            raise TypeError
        if 1e-3 <= value <= 1e3:
            self._safety_factor = value
        else:
            raise NEIError("Invalid safety factor.")

    @property
    def verbose(self) -> bool:
        """
        Return `True` if verbose output during a simulation is
        requested, and `False` otherwise.
        """
        return self._verbose

    @verbose.setter
    def verbose(self, choice: bool):
        if choice is True or choice is False:
            self._verbose = choice
        else:
            raise TypeError("Invalid choice for verbose.")

    @u.quantity_input
    def in_time_interval(self, time: u.s, buffer: u.s = 1e-9 * u.s):
        """
        Return `True` if the ``time`` is between ``time_start - buffer``
        and ``time_max + buffer`` , and `False` otherwise.

        Raises
        ------
        TypeError
            If ``time`` or ``buffer`` is not a ``astropy.units.Quantity``

        astropy.units.UnitsError
            If ``time`` or ``buffer`` is not in units of time.

        """
        return self.time_start - buffer <= time <= self.time_max + buffer

    @property
    def max_steps(self) -> int:
        """
        The maximum number of steps that a simulation will be allowed
        to take.
        """
        return self._max_steps

    @max_steps.setter
    def max_steps(self, n: int):
        if isinstance(n, (int, np.integer)) and 0 < n <= 1000000:
            self._max_steps = n
        else:
            raise TypeError(
                "max_steps must be an integer with 0 < max_steps <= 1000000"
            )

    @property
    def T_e_input(self) -> Union[u.Quantity, Callable]:
        """
        The temperature input.
        """
        return self._T_e_input

    @T_e_input.setter
    def T_e_input(self, T_e: Optional[Union[Callable, u.Quantity]]):
        """Set the input electron temperature."""
        if isinstance(T_e, u.Quantity):
            try:
                T_e = T_e.to(u.K, equivalencies=u.temperature_energy())
            except u.UnitConversionError:
                raise u.UnitsError("Invalid electron temperature.") from None
            if T_e.isscalar:
                self._T_e_input = T_e
                self._electron_temperature = lambda time: T_e
            else:
                if self._time_input is None:
                    raise TypeError("Must define time_input prior to T_e for an array.")
                time_input = self.time_input
                if len(time_input) != len(T_e):
                    raise ValueError("len(T_e) not equal to len(time_input).")
                f = interpolate.interp1d(
                    time_input.value,
                    T_e.value,
                    bounds_error=False,
                    fill_value="extrapolate",
                )
                self._electron_temperature = lambda time: f(time.value) * u.K
                self._T_e_input = T_e
        elif callable(T_e):
            if self.time_start is not None:
                try:
                    T_e(self.time_start).to(u.K)
                    T_e(self.time_max).to(u.K)
                except Exception:
                    raise ValueError("Invalid electron temperature function.")
            self._T_e_input = T_e
            self._electron_temperature = T_e
        elif T_e is None:
            self._electron_temperature = lambda: None
        else:
            raise TypeError("Invalid T_e")

    def electron_temperature(self, time: u.Quantity) -> u.Quantity:
        try:
            if not self.in_time_interval(time):
                warnings.warn(
                    f"{time} is not in the simulation time interval:"
                    f"[{self.time_start}, {self.time_max}]. "
                    f"May be extrapolating temperature."
                )
            T_e = self._electron_temperature(time.to(u.s))
            if np.isnan(T_e) or np.isinf(T_e) or T_e < 0 * u.K:
                raise NEIError(f"T_e = {T_e} at time = {time}.")
            return T_e
        except Exception as exc:
            raise NEIError(
                f"Unable to calculate a valid electron temperature " f"for time {time}"
            ) from exc

    @property
    def n_input(self) -> u.Quantity:
        """The number density factor input."""
        if "H" in self.elements:
            return self._n_input
        else:
            raise ValueError

    @n_input.setter
    def n_input(self, n: u.Quantity):
        if isinstance(n, u.Quantity):
            try:
                n = n.to(u.cm ** -3)
            except u.UnitConversionError:
                raise u.UnitsError("Invalid hydrogen density.")
            if n.isscalar:
                self._n_input = n
                self.hydrogen_number_density = lambda time: n
            else:
                if self._time_input is None:
                    raise TypeError("Must define time_input prior to n for an array.")
                time_input = self.time_input
                if len(time_input) != len(n):
                    raise ValueError("len(n) is not equal to len(time_input).")
                f = interpolate.interp1d(
                    time_input.value,
                    n.value,
                    bounds_error=False,
                    fill_value="extrapolate",
                )
                self._hydrogen_number_density = lambda time: f(time.value) * u.cm ** -3
                self._n_input = n
        elif callable(n):
            if self.time_start is not None:
                try:
                    n(self.time_start).to(u.cm ** -3)
                    n(self.time_max).to(u.cm ** -3)
                except Exception:
                    raise ValueError("Invalid number density function.")
            self._n_input = n
            self._hydrogen_number_density = n
        elif n is None:
            self._hydrogen_number_density = lambda: None
        else:
            raise TypeError("Invalid n.")

    def hydrogen_number_density(self, time: u.Quantity) -> u.Quantity:
        try:
            time = time.to(u.s)
        except (AttributeError, u.UnitsError):
            raise NEIError("Invalid time in hydrogen_density")
        return self._hydrogen_number_density(time)

    @property
    def eigen_data_dict(self) -> Dict[str, EigenData]:
        """
        Return a `dict` containing `~plasmapy_nei.eigen.EigenData` instances
        for each element.
        """
        return self._eigen_data_dict

    @property
    def initial(self) -> IonizationStates:
        """
        Return the ionization states of the plasma at the beginning of
        the simulation.
        """
        return self._initial

    @initial.setter
    def initial(self, initial_states: IonizationStates):
        if isinstance(initial_states, IonizationStates):
            self._initial = initial_states
            self._elements = (
                initial_states.ionic_fractions.keys()
            )  # TODO IonizationStates
        elif initial_states is None:
            self._ionstates = None
        else:
            raise TypeError("Expecting an IonizationStates instance.")

    @property
    def results(self) -> SimulationResults:
        """
        Return the `~plasmapy_nei.nei.SimulationResults` class instance that
        corresponds to the simulation results.

        """
        if self._results is not None:
            return self._results
        else:
            raise AttributeError("The simulation has not yet been performed.")

    @property
    def final(self) -> IonizationStates:
        """
        Return the ionization states of the plasma at the end of the
        simulation.
        """
        try:
            return self._final
        except AttributeError:
            raise NEIError("The simulation has not yet been performed.") from None

    def _initialize_simulation(self):

        self._results = SimulationResults(
            initial=self.initial,
            n_init=self.hydrogen_number_density(self.time_start),
            T_e_init=self.electron_temperature(self.time_start),
            max_steps=self.max_steps,
            time_start=self.time_start,
        )
        self._old_time = self.time_start.to(u.s)
        self._new_time = self.time_start.to(u.s)

    def simulate(self) -> SimulationResults:
        """
        Perform a non-equilibrium ionization simulation.

        Returns
        -------
        results: `~plasmapy_nei.classes.Simulation`
            The results from the simulation (which are also stored in
            the ``results`` attribute of the `~plasmapy_nei.nei.NEI`
            instance this method was called from.

        """

        self._initialize_simulation()

        for step in range(self.max_steps):
            try:
                self.set_timestep()
                self.time_advance()
            except StopIteration:
                break
            except Exception as exc:
                raise NEIError(f"Unable to complete simulation.") from exc

        self._finalize_simulation()

        # Is there a way to use the inspect package or something similar
        # to only return self.results if it is in an expression where

        return self.results

    def _finalize_simulation(self):
        self._results._cleanup()

        final_ionfracs = {
            element: self.results.ionic_fractions[element][-1, :]
            for element in self.elements
        }

        self._final = IonizationStates(
            inputs=final_ionfracs,
            abundances=self.abundances,
            n=np.sum(self.results.number_densities["H"][-1, :]),  # modify this later?,
            T_e=self.results.T_e[-1],
            tol=1e-6,
        )

        if not np.isclose(self.time_max / u.s, self.results.time[-1] / u.s):
            warnings.warn(
                f"The simulation ended at {self.results.time[-1]}, "
                f"which is prior to time_max = {self.time_max}."
            )

    def _set_adaptive_timestep(self):
        """Adapt the time step."""

        t = self._new_time if hasattr(self, "_new_time") else self.t_start

        # We need to guess the timestep in order to narrow down what the
        # timestep should be.  If we are in the middle of a simulation,
        # we can use the old timestep as a reasonable guess.  If we are
        # simulation, then we can either use the inputted timestep or
        # estimate it from other inputs.

        dt_guess = (
            self._dt
            if self._dt
            else self._dt_input
            if self._dt_input
            else self.time_max / self.max_steps
        )

        # Make sure that dt_guess does not lead to a time that is out
        # of the domain.

        dt_guess = dt_guess if t + dt_guess <= self.time_max - t else self.time_max - t

        # The temperature may start out exactly at the boundary of a
        # bin, so we check what bin it is in just slightly after to
        # figure out which temperature bin the plasma is entering.

        T = self.electron_temperature(t + 1e-9 * dt_guess)

        # Find the boundaries to the temperature bin.

        index = self._get_temperature_index(T.to(u.K).value)
        T_nearby = np.array(self._temperature_grid[index - 1 : index + 2]) * u.K
        T_boundary = (T_nearby[0:-1] + T_nearby[1:]) / 2

        # In order to use Brent's method, we must bound the root's
        # location.  Functions may change sharply or slowly, so we test
        # different times that are logarithmically spaced to find the
        # first one that is outside of the boundary.

        dt_spread = (
            np.geomspace(1e-9 * dt_guess.value, (self.time_max - t).value, num=100)
            * u.s
        )
        time_spread = t + dt_spread
        T_spread = [self.electron_temperature(time) for time in time_spread]
        in_range = [T_boundary[0] <= temp <= T_boundary[1] for temp in T_spread]

        # If all of the remaining temperatures are in the same bin, then
        # the temperature will be roughly constant for the rest of the
        # simulation.  Take one final long time step, unless it exceeds
        # dt_max.

        if all(in_range):
            new_dt = self.time_max - t
            self._dt = new_dt if new_dt <= self.dt_max else self.dt_max
            return

        # Otherwise, we need to find the first index in the spread that
        # corresponds to a temperature outside of the temperature bin
        # for this time step.

        first_false_index = in_range.index(False)

        # We need to figure out if the temperature is dropping so that
        # it crosses the low temperature boundary of the bin, or if it
        # is rising so that it crosses the high temperature of the bin.

        T_first_outside = self.electron_temperature(time_spread[first_false_index])

        if T_first_outside >= T_boundary[1]:
            boundary_index = 1
        elif T_first_outside <= T_boundary[0]:
            boundary_index = 0

        # Select the values for the time step in the spread just before
        # and after the temperature leaves the temperature bin as bounds
        # for the root finding method.

        dt_bounds = (dt_spread[first_false_index - 1 : first_false_index + 1]).value

        # Define a function for the difference between the temperature
        # and the temperature boundary as a function of the value of the
        # time step.

        T_val = lambda dtval: (
            self.electron_temperature(t + dtval * u.s) - T_boundary[boundary_index]
        ).value

        # Next we find the root.  This method should succeed as long as
        # the root is bracketed by dt_bounds.  Because astropy.units is
        # not fully compatible with SciPy, we temporarily drop units and
        # then reattach them.

        try:
            new_dt = (
                optimize.brentq(
                    T_val,
                    *dt_bounds,
                    xtol=1e-14,
                    maxiter=1000,
                    disp=True,
                )
                * u.s
            )
        except Exception as exc:
            raise NEIError(f"Unable to find new dt at t = {t}") from exc
        else:
            if np.isnan(new_dt.value):
                raise NEIError(f"new_dt = {new_dt}")

        # Enforce that the time step is in the interval [dt_min, dt_max].

        if new_dt < self.dt_min:
            new_dt = self.dt_min
        elif new_dt > self.dt_max:
            new_dt = self.dt_max

        # Store the time step as a private attribute so that it can be
        # used in the time advance.

        self._dt = new_dt.to(u.s)

    def set_timestep(self, dt: u.Quantity = None):
        """
        Set the time step for the next non-equilibrium ionization time
        advance.

        Parameters
        ----------
        dt: astropy.units.Quantity, optional
            The time step to be used for the next time advance.

        Notes
        -----
        If ``dt`` is not `None`, then the time step will be set to ``dt``.

        If ``dt`` is not set and the ``adapt_dt`` attribute of an
        `~plasmapy_nei.nei.NEI` instance is `True`, then this method will
        calculate the time step corresponding to how long it will be
        until the temperature rises or drops into the next temperature
        bin.  If this time step is between ``dtmin`` and ``dtmax``, then

        If ``dt`` is not set and the ``adapt_dt`` attribute is `False`,
        then this method will set the time step as what was inputted to
        the `~plasmapy_nei.nei.NEI` class upon instantiation in the
        ``dt`` argument or through the `~plasmapy_nei.nei.NEI` class's
        ``dt_input`` attribute.

        Raises
        ------
        ~plasmapy_nei.nei.NEIError
            If the time step cannot be set, for example if the ``dt``
            argument is invalid or the time step cannot be adapted.
        """

        if dt is not None:
            # Allow the time step to set as an argument to this method.
            try:
                dt = dt.to(u.s)
            except Exception as exc:
                raise NEIError(f"{dt} is not a valid time step.") from exc
            finally:
                self._dt = dt
        elif self.adapt_dt:
            try:
                self._set_adaptive_timestep()
            except Exception as exc:
                raise NEIError("Unable to adapt the time step.") from exc
        elif self.dt_input is not None:
            self._dt = self.dt_input
        else:
            raise NEIError("Unable to set the time step.")

        self._old_time = self._new_time
        self._new_time = self._old_time + self._dt

        if self._new_time > self.time_max:
            self._new_time = self.time_max
            self._dt = self._new_time - self._old_time

    def time_advance(self):
        """Advance the simulation by one time step."""
        # TODO: Expand docstring and include equations!

        # TODO: Fully implement units into this.

        step = self.results._index
        T_e = self.results.T_e[step - 1].value
        n_e = self.results.n_e[step - 1].value  # set average
        dt = self._dt.value

        if self.verbose:
            print(f"step={step}  T_e={T_e}  n_e={n_e}  dt={dt}")

        new_ionic_fractions = {}

        try:
            for elem in self.elements:
                nstates = self.results.nstates[elem]
                f0 = self.results._ionic_fractions[elem][self.results._index - 1, :]

                evals = self.eigen_data_dict[elem].eigenvalues(T_e=T_e)
                evect = self.eigen_data_dict[elem].eigenvectors(T_e=T_e)
                evect_inverse = self.eigen_data_dict[elem].eigenvector_inverses(T_e=T_e)

                diagonal_evals = np.zeros((nstates, nstates), dtype=np.float64)
                for ii in range(0, nstates):
                    diagonal_evals[ii, ii] = np.exp(evals[ii] * dt * n_e)

                matrix_1 = np.dot(diagonal_evals, evect)
                matrix_2 = np.dot(evect_inverse, matrix_1)

                ft = np.dot(f0, matrix_2)

                # Due to truncation errors in the solutions in the
                # eigenvalues and eigenvectors, there is a chance that
                # very slightly negative ionic fractions will arise.
                # These are not natural and will make the code grumpy.
                # For these reasons, the ionic fractions will be very
                # slightly unnormalized.  We set negative ionic
                # fractions to zero and renormalize.

                ft[np.where(ft < 0.0)] = 0.0
                new_ionic_fractions[elem] = ft / np.sum(ft)

        except Exception as exc:
            raise NEIError(f"Unable to do time advance for {elem}") from exc
        else:

            new_time = self.results.time[self.results._index - 1] + self._dt
            self.results._assign(
                new_time=new_time,
                new_ionfracs=new_ionic_fractions,
                new_T_e=self.electron_temperature(new_time),
                new_n=self.hydrogen_number_density(new_time),
            )

        if new_time >= self.time_max or np.isclose(new_time.value, self.time_max.value):
            raise StopIteration

    def save(self, filename: str = "nei.h5"):
        """
        Save the `~plasmapy_nei.nei.NEI` instance to an HDF5 file.  Not
        implemented.
        """
        raise NotImplementedError

    def index_to_time(self, index):
        """
        Returns the time value or array given the index/indices

        Parameters
        ------
        index: array-like
               A value or array of values representing the index of
               the time array created by the simulation

        Returns
        ------
        get_time: astropy.units.Quantity
                  The time value associated with index input(s)
        """

        return self.results.time[index]

    def time_to_index(self, time):
        """
        Returns the closest index value or array for the given time(s)

        Parameters
        ------
        time: array-like,
               A value or array of values representing the values of
               the time array created by the simulation

        Returns
        ------
        index: int or array-like,
                  The index value associated with the time input(s)
        """
        index = (np.abs(self.results.time.value - time)).argmin()

        return index


"""
Module for petitCODE atmospheric model spectra.
"""

import os
import tarfile
import urllib.request
import warnings
import zipfile

from typing import Optional, Tuple

import h5py
import numpy as np
import spectres

from typeguard import typechecked

from species.core import constants
from species.util import data_util, read_util

@typechecked
def add_petitcode_cool_clear(input_path: str,
                             database: h5py._hl.files.File,
                             wavel_range: Optional[Tuple[float, float]] = None,
                             teff_range: Optional[Tuple[float, float]] = None,
                             spec_res: Optional[float] = 1000.) -> None:
    """
    Function for adding the petitCODE cool clear atmospheric models to the database.

    Parameters
    ----------
    input_path : str
        Folder where the data is located.
    database : h5py._hl.files.File
        Database.
    wavel_range : tuple(float, float), None
        Wavelength range (um). The original wavelength points are used if set to None.
    teff_range : tuple(float, float), None
        Effective temperature range (K). All temperatures are selected if set to None.
    spec_res : float, None
        Spectral resolution. Not used if ``wavel_range`` is set to None.

    Returns
    -------
    NoneType
        None
    """

    if not os.path.exists(input_path):
        os.makedirs(input_path)

    data_folder = os.path.join(input_path, 'linder_molliere_grid/clear/specs/')

    url = 'http://mpia.de/~molliere/online_data/linder_molliere_grid.zip'

    data_file = os.path.join(input_path, 'linder_molliere_grid.zip')

    if not os.path.isfile(data_file):
        print('Downloading petitCODE cool model spectra (3.7 GB)...', end='', flush=True)
        urllib.request.urlretrieve(url, data_file)
        print(' [DONE]')

    print('Unpacking petitCODE cool model spectra (3.7 GB)...', end='', flush=True)

    with zipfile.ZipFile(data_file, 'r') as zip_ref:
        zip_ref.extractall(input_path)

    print(' [DONE]')

    teff = []
    logg = []
    feh = []
    flux = []

    if wavel_range is not None:
        wavelength = read_util.create_wavelengths(wavel_range, spec_res)
    else:
        wavelength = None

    for _, _, files in os.walk(data_folder):
        for filename in files:
            file_split = filename.split('_')

            teff_val = float(file_split[2])
            logg_val = float(file_split[4])
            feh_val = float(file_split[6])

            if teff_range is not None:
                if teff_val < teff_range[0] or teff_val > teff_range[1]:
                    continue

            print_message = f'Adding petitCODE cool clear model spectra... {filename}'
            print(f'\r{print_message:<87}', end='')

            data = np.loadtxt(os.path.join(data_folder, filename))

            teff.append(teff_val)
            logg.append(logg_val)
            feh.append(feh_val)

            if wavel_range is None:
                if wavelength is None:
                    # (cm) -> (um)
                    wavelength = data[:, 0]*1e4

                if np.all(np.diff(wavelength) < 0):
                    raise ValueError('The wavelengths are not all sorted by increasing value.')

                # (erg s-1 cm-2 Hz-1) -> (W m-2 um-1)
                flux.append(data[:, 1]*1e-9*constants.LIGHT/(wavelength*1e-6)**2)

            else:
                # (cm) -> (um)
                data_wavel = data[:, 0]*1e4

                # (erg s-1 cm-2 Hz-1) -> (W m-2 um-1)
                data_flux = data[:, 1]*1e-9*constants.LIGHT/(data_wavel*1e-6)**2

                try:
                    flux.append(spectres.spectres(wavelength, data_wavel, data_flux))
                except ValueError:
                    flux.append(np.zeros(wavelength.shape[0]))

                    warnings.warn('The wavelength range should fall within the range of the '
                                  'original wavelength sampling. Storing zeros instead.')

    print_message = 'Adding petitCODE cool clear model spectra... [DONE]'
    print(f'\r{print_message:<87}')

    data_sorted = data_util.sort_data(np.asarray(teff),
                                      np.asarray(logg),
                                      np.asarray(feh),
                                      None,
                                      None,
                                      wavelength,
                                      np.asarray(flux))

    data_util.write_data('petitcode-cool-clear',
                         ['teff', 'logg', 'feh'],
                         database,
                         data_sorted)

@typechecked
def add_petitcode_cool_cloudy(input_path: str,
                              database: h5py._hl.files.File,
                              wavel_range: Optional[Tuple[float, float]] = None,
                              teff_range: Optional[Tuple[float, float]] = None,
                              spec_res: Optional[float] = 1000.) -> None:
    """
    Function for adding the petitCODE cool cloudy atmospheric models to the database.

    Parameters
    ----------
    input_path : str
        Folder where the data is located.
    database : h5py._hl.files.File
        Database.
    wavel_range : tuple(float, float), None
        Wavelength range (um). The original wavelength points are used if set to None.
    teff_range : tuple(float, float), None
        Effective temperature range (K). All temperatures are selected if set to None.
    spec_res : float, None
        Spectral resolution. Not used if ``wavel_range`` is set to None.

    Returns
    -------
    NoneType
        None
    """

    if not os.path.exists(input_path):
        os.makedirs(input_path)

    data_folder = os.path.join(input_path, 'linder_molliere_grid/cloudy/specs/')

    url = 'http://mpia.de/~molliere/online_data/linder_molliere_grid.zip'

    data_file = os.path.join(input_path, 'linder_molliere_grid.zip')

    if not os.path.isfile(data_file):
        print('Downloading petitCODE cool model spectra (3.7 GB)...', end='', flush=True)
        urllib.request.urlretrieve(url, data_file)
        print(' [DONE]')

    print('Unpacking petitCODE cool model spectra (3.7 GB)...', end='', flush=True)

    with zipfile.ZipFile(data_file, 'r') as zip_ref:
        zip_ref.extractall(input_path)

    print(' [DONE]')

    teff = []
    logg = []
    feh = []
    fsed = []
    flux = []

    if wavel_range is not None:
        wavelength = read_util.create_wavelengths(wavel_range, spec_res)
    else:
        wavelength = None

    for _, _, files in os.walk(data_folder):
        for filename in files:
            file_split = filename.split('_')

            teff_val = float(file_split[2])
            logg_val = float(file_split[4])
            feh_val = float(file_split[6])
            fsed_val = float(file_split[8])

            if teff_range is not None:
                if teff_val < teff_range[0] or teff_val > teff_range[1]:
                    continue

            print_message = f'Adding petitCODE cool cloudy model spectra... {filename}'
            print(f'\r{print_message:<106}', end='')

            data = np.loadtxt(os.path.join(data_folder, filename))

            teff.append(teff_val)
            logg.append(logg_val)
            feh.append(feh_val)
            fsed.append(fsed_val)

            if wavel_range is None:
                if wavelength is None:
                    # (cm) -> (um)
                    wavelength = data[:, 0]*1e4

                if np.all(np.diff(wavelength) < 0):
                    raise ValueError('The wavelengths are not all sorted by increasing value.')

                # (erg s-1 cm-2 Hz-1) -> (W m-2 um-1)
                flux.append(data[:, 1]*1e-9*constants.LIGHT/(wavelength*1e-6)**2)

            else:
                # (cm) -> (um)
                data_wavel = data[:, 0]*1e4

                # (erg s-1 cm-2 Hz-1) -> (W m-2 um-1)
                data_flux = data[:, 1]*1e-9*constants.LIGHT/(data_wavel*1e-6)**2

                try:
                    flux.append(spectres.spectres(wavelength, data_wavel, data_flux))
                except ValueError:
                    flux.append(np.zeros(wavelength.shape[0]))

                    warnings.warn('The wavelength range should fall within the range of the '
                                  'original wavelength sampling. Storing zeros instead.')

    print_message = 'Adding petitCODE cool cloudy model spectra... [DONE]'
    print(f'\r{print_message:<106}')

    data_sorted = data_util.sort_data(np.asarray(teff),
                                      np.asarray(logg),
                                      np.asarray(feh),
                                      None,
                                      np.asarray(fsed),
                                      wavelength,
                                      np.asarray(flux))

    data_util.write_data('petitcode-cool-cloudy',
                         ['teff', 'logg', 'feh', 'fsed'],
                         database,
                         data_sorted)

@typechecked
def add_petitcode_hot_clear(input_path: str,
                            database: h5py._hl.files.File,
                            wavel_range: Optional[Tuple[float, float]] = None,
                            teff_range: Optional[Tuple[float, float]] = None,
                            spec_res: Optional[float] = 1000.) -> None:
    """
    Function for adding the petitCODE hot clear atmospheric models to the database.

    Parameters
    ----------
    input_path : str
        Folder where the data is located.
    database : h5py._hl.files.File
        Database.
    wavel_range : tuple(float, float), None
        Wavelength range (um). The original wavelength points are used if set to None.
    teff_range : tuple(float, float), None
        Effective temperature range (K). All temperatures are selected if set to None.
    spec_res : float, None
        Spectral resolution. Not used if ``wavel_range`` is set to None.

    Returns
    -------
    NoneType
        None
    """

    if not os.path.exists(input_path):
        os.makedirs(input_path)

    data_folder = os.path.join(input_path, 'petitcode-hot-clear/')

    url = 'https://people.phys.ethz.ch/~ipa/tstolker/petitcode-hot-clear.tgz'

    data_file = os.path.join(input_path, 'petitcode-hot-clear.tgz')

    if not os.path.isfile(data_file):
        print('Downloading petitCODE hot clear model spectra (93 MB)...', end='', flush=True)
        urllib.request.urlretrieve(url, data_file)
        print(' [DONE]')

    print('Unpacking petitCODE hot clear model spectra (93 MB)...', end='', flush=True)
    tar = tarfile.open(data_file)
    tar.extractall(data_folder)
    tar.close()
    print(' [DONE]')

    teff = []
    logg = []
    feh = []
    co_ratio = []
    flux = []

    if wavel_range is not None:
        wavelength = read_util.create_wavelengths(wavel_range, spec_res)
    else:
        wavelength = None

    for _, _, files in os.walk(data_folder):
        for filename in files:
            file_split = filename.split('_')

            teff_val = float(file_split[2])
            logg_val = float(file_split[4])
            feh_val = float(file_split[6])
            co_ratio_val = float(file_split[8])

            if teff_range is not None:
                if teff_val < teff_range[0] or teff_val > teff_range[1]:
                    continue

            print_message = f'Adding petitCODE hot clear model spectra... {filename}'
            print(f'\r{print_message:<99}', end='')

            data = np.loadtxt(os.path.join(data_folder, filename))

            teff.append(teff_val)
            logg.append(logg_val)
            feh.append(feh_val)
            co_ratio.append(co_ratio_val)

            if wavel_range is None:
                if wavelength is None:
                    # (cm) -> (um)
                    wavelength = data[:, 0]*1e4

                if np.all(np.diff(wavelength) < 0):
                    raise ValueError('The wavelengths are not all sorted by increasing value.')

                # (erg s-1 cm-2 Hz-1) -> (W m-2 um-1)
                flux.append(data[:, 1]*1e-9*constants.LIGHT/(wavelength*1e-6)**2)

            else:
                # (cm) -> (um)
                data_wavel = data[:, 0]*1e4

                # (erg s-1 cm-2 Hz-1) -> (W m-2 um-1)
                data_flux = data[:, 1]*1e-9*constants.LIGHT/(data_wavel*1e-6)**2

                try:
                    flux.append(spectres.spectres(wavelength, data_wavel, data_flux))
                except ValueError:
                    flux.append(np.zeros(wavelength.shape[0]))

                    warnings.warn('The wavelength range should fall within the range of the '
                                  'original wavelength sampling. Storing zeros instead.')

    print_message = 'Adding petitCODE hot clear model spectra... [DONE]'
    print(f'\r{print_message:<99}')

    data_sorted = data_util.sort_data(np.asarray(teff),
                                      np.asarray(logg),
                                      np.asarray(feh),
                                      np.asarray(co_ratio),
                                      None,
                                      wavelength,
                                      np.asarray(flux))

    data_util.write_data('petitcode-hot-clear',
                         ['teff', 'logg', 'feh', 'co'],
                         database,
                         data_sorted)

@typechecked
def add_petitcode_hot_cloudy(input_path: str,
                             database: h5py._hl.files.File,
                             wavel_range: Optional[Tuple[float, float]] = None,
                             teff_range: Optional[Tuple[float, float]] = None,
                             spec_res: Optional[float] = 1000.) -> None:
    """
    Function for adding the petitCODE hot cloudy atmospheric models to the database.

    Parameters
    ----------
    input_path : str
        Folder where the data is located.
    database : h5py._hl.files.File
        Database.
    wavel_range : tuple(float, float), None
        Wavelength range (um). The original wavelength points are used if set to None.
    teff_range : tuple(float, float), None
        Effective temperature range (K). All temperatures are selected if set to None.
    spec_res : float, None
        Spectral resolution. Not used if ``wavel_range`` is set to None.

    Returns
    -------
    NoneType
        None
    """

    if not os.path.exists(input_path):
        os.makedirs(input_path)

    data_folder = os.path.join(input_path, 'petitcode-hot-cloudy/')

    url = 'https://people.phys.ethz.ch/~ipa/tstolker/petitcode-hot-cloudy.tgz'

    data_file = os.path.join(input_path, 'petitcode-hot-cloudy.tgz')

    if not os.path.isfile(data_file):
        print('Downloading petitCODE hot cloudy model spectra (276 MB)...', end='', flush=True)
        urllib.request.urlretrieve(url, data_file)
        print(' [DONE]')

    print('Unpacking petitCODE hot cloudy model spectra (276 MB)...', end='', flush=True)
    tar = tarfile.open(data_file)
    tar.extractall(data_folder)
    tar.close()
    print(' [DONE]')

    teff = []
    logg = []
    feh = []
    co_ratio = []
    fsed = []
    flux = []

    if wavel_range is not None:
        wavelength = read_util.create_wavelengths(wavel_range, spec_res)
    else:
        wavelength = None

    for _, _, files in os.walk(data_folder):
        for filename in files:
            file_split = filename.split('_')

            teff_val = float(file_split[2])
            logg_val = float(file_split[4])
            feh_val = float(file_split[6])
            co_ratio_val = float(file_split[8])
            fsed_val = float(file_split[10])

            if teff_range is not None:
                if teff_val < teff_range[0] or teff_val > teff_range[1]:
                    continue

            print_message = f'Adding petitCODE hot cloudy model spectra... {filename}'
            print(f'\r{print_message:<111}', end='')

            data = np.loadtxt(os.path.join(data_folder, filename))

            teff.append(teff_val)
            logg.append(logg_val)
            feh.append(feh_val)
            co_ratio.append(co_ratio_val)
            fsed.append(fsed_val)

            if wavel_range is None:
                if wavelength is None:
                    # (cm) -> (um)
                    wavelength = data[:, 0]*1e4

                if np.all(np.diff(wavelength) < 0):
                    raise ValueError('The wavelengths are not all sorted by increasing value.')

                # (erg s-1 cm-2 Hz-1) -> (W m-2 um-1)
                flux.append(data[:, 1]*1e-9*constants.LIGHT/(wavelength*1e-6)**2)

            else:
                # (cm) -> (um)
                data_wavel = data[:, 0]*1e4

                # (erg s-1 cm-2 Hz-1) -> (W m-2 um-1)
                data_flux = data[:, 1]*1e-9*constants.LIGHT/(data_wavel*1e-6)**2

                try:
                    flux.append(spectres.spectres(wavelength, data_wavel, data_flux))
                except ValueError:
                    flux.append(np.zeros(wavelength.shape[0]))

                    warnings.warn('The wavelength range should fall within the range of the '
                                  'original wavelength sampling. Storing zeros instead.')

    print_message = 'Adding petitCODE hot cloudy model spectra... [DONE]'
    print(f'\r{print_message:<111}')

    data_sorted = data_util.sort_data(np.asarray(teff),
                                      np.asarray(logg),
                                      np.asarray(feh),
                                      np.asarray(co_ratio),
                                      np.asarray(fsed),
                                      wavelength,
                                      np.asarray(flux))

    data_util.write_data('petitcode-hot-cloudy',
                         ['teff', 'logg', 'feh', 'co', 'fsed'],
                         database,
                         data_sorted)