"""module for performing calculations on Spectral Energy Distributions (SEDs)

(c) 2007-2013 Matt Hilton

U{http://astlib.sourceforge.net}

This module provides classes for manipulating SEDs, in particular the Bruzual &
Charlot 2003, Maraston 2005, and Percival et al 2009 stellar population
synthesis models are currently supported. Functions are provided for
calculating the evolution of colours and magnitudes in these models with
redshift etc., and for fitting broadband photometry using these models.

@var VEGA: The SED of Vega, used for calculation of magnitudes on the Vega system.
@type VEGA: L{SED} object
@var AB: Flat spectrum SED, used for calculation of magnitudes on the AB system.
@type AB: L{SED} object
@var SOL: The SED of the Sun.
@type SOL: L{SED} object

"""

#-----------------------------------------------------------------------------
import numpy
import math
import operator
try:
    from scipy import interpolate
    from scipy import ndimage
except ImportError:
    print("WARNING: astSED: failed to import scipy modules - some functions "
          "will not work.")
import astLib
from astLib import astCalc
import os
try:
    import matplotlib
    from matplotlib import pylab
    matplotlib.interactive(False)
except ImportError:
    print("WARNING: astSED: failed to import matplotlib - some functions will "
          "not work.")
import glob

#-----------------------------------------------------------------------------
class Passband:
    """This class describes a filter transmission curve. Passband objects are
    created by loading data from from text files containing wavelength in
    angstroms in the first column, relative transmission efficiency in the
    second column (whitespace delimited). For example, to create a Passband
    object for the 2MASS J filter:

    passband=astSED.Passband("J_2MASS.res")

    where "J_2MASS.res" is a file in the current working directory that
    describes the filter.

    Wavelength units can be specified as 'angstroms', 'nanometres' or
    'microns'; if either of the latter, they will be converted to angstroms.

    """

    def __init__(self, fileName, normalise=True, inputUnits='angstroms'):

        inFile = open(fileName, "r")
        lines = inFile.readlines()

        wavelength = []
        transmission = []
        for line in lines:

            if line[0] != "#" and len(line) > 3:

                bits = line.split()
                transmission.append(float(bits[1]))
                wavelength.append(float(bits[0]))

        self.wavelength = numpy.array(wavelength)
        self.transmission = numpy.array(transmission)

        if inputUnits == 'angstroms':
            pass
        elif inputUnits == 'nanometres':
            self.wavelength = self.wavelength * 10.0
        elif inputUnits == 'microns':
            self.wavelength = self.wavelength * 10000.0
        elif inputUnits == 'mm':
            self.wavelength = self.wavelength * 1e7
        elif inputUnits == 'GHz':
            self.wavelength = 3e8 / (self.wavelength * 1e9)
            self.wavelength = self.wavelength * 1e10
        else:
            raise Exception("didn't understand passband input units")

        # Sort into ascending order of wavelength otherwise normalisation will be wrong
        merged = numpy.array([self.wavelength, self.transmission]).transpose()
        sortedMerged = numpy.array(sorted(merged, key=operator.itemgetter(0)))
        self.wavelength = sortedMerged[:, 0]
        self.transmission = sortedMerged[:, 1]

        if normalise:
            self.transmission = self.transmission / numpy.trapz(
                self.transmission, self.wavelength)

        # Store a ready-to-go interpolation object to speed calculation of fluxes up
        self.interpolator = interpolate.interp1d(self.wavelength,
                                                 self.transmission,
                                                 kind='linear')

    def asList(self):
        """Returns a two dimensional list of [wavelength, transmission],
        suitable for plotting by gnuplot.

        @rtype: list
        @return: list in format [wavelength, transmission]

        """

        listData = []
        for l, f in zip(self.wavelength, self.transmission):
            listData.append([l, f])

        return listData

    def rescale(self, maxTransmission):
        """Rescales the passband so that maximum value of the transmission is
        equal to maxTransmission. Useful for plotting.

        @type maxTransmission: float
        @param maxTransmission: maximum value of rescaled transmission curve

        """

        self.transmission = self.transmission * (maxTransmission /
                                                 self.transmission.max())

    def plot(self, xmin='min', xmax='max', maxTransmission=None):
        """Plots the passband, rescaling the maximum of the tranmission curve
        to maxTransmission if required.

        @type xmin: float or 'min'
        @param xmin: minimum of the wavelength range of the plot
        @type xmax: float or 'max'
        @param xmax: maximum of the wavelength range of the plot
        @type maxTransmission: float
        @param maxTransmission: maximum value of rescaled transmission curve

        """

        if maxTransmission is not None:
            self.rescale(maxTransmission)

        pylab.matplotlib.interactive(True)
        pylab.plot(self.wavelength, self.transmission)

        if xmin == 'min':
            xmin = self.wavelength.min()
        if xmax == 'max':
            xmax = self.wavelength.max()

        pylab.xlim(xmin, xmax)
        pylab.xlabel("Wavelength")
        pylab.ylabel("Relative Flux")

    def effectiveWavelength(self):
        """Calculates effective wavelength for the passband. This is the same
        as equation (3) of Carter et al. 2009.

        @rtype: float
        @return: effective wavelength of the passband, in Angstroms

        """

        a = numpy.trapz(self.transmission * self.wavelength, self.wavelength)
        b = numpy.trapz(self.transmission / self.wavelength, self.wavelength)
        effWavelength = numpy.sqrt(a / b)

        return effWavelength

#-----------------------------------------------------------------------------
class TopHatPassband(Passband):
    """This class generates a passband with a top hat response between the
    given wavelengths.

    """

    def __init__(self, wavelengthMin, wavelengthMax, normalise=True):
        """Generates a passband object with top hat response between
        wavelengthMin, wavelengthMax. Units are assumed to be Angstroms.

        @type wavelengthMin: float
        @param wavelengthMin: minimum of the wavelength range of the passband
        @type wavelengthMax: float
        @param wavelengthMax: maximum of the wavelength range of the passband
        @type normalise: bool
        @param normalise: if True, scale such that total area under the
        passband over the wavelength
        range is 1.

        """

        self.wavelength = numpy.arange(
            wavelengthMin, wavelengthMax + 10,
            10, dtype=float)
        self.transmission = numpy.ones(self.wavelength.shape, dtype=float)

        if normalise:
            self.transmission = self.transmission / numpy.trapz(
                self.transmission, self.wavelength)

        # Store a ready-to-go interpolation object to speed calculation of fluxes up
        self.interpolator = interpolate.interp1d(self.wavelength,
                                                 self.transmission,
                                                 kind='linear')

#-----------------------------------------------------------------------------
class SED:
    """This class describes a Spectral Energy Distribution (SED).

    To create a SED object, lists (or numpy arrays) of wavelength and relative
    flux must be provided. The SED can optionally be redshifted. The wavelength
    units of SEDs are assumed to be Angstroms - flux calculations using
    Passband and SED objects specified with different wavelength units will be
    incorrect.

    The L{StellarPopulation} class (and derivatives) can be used to extract
    SEDs for specified ages from e.g. the Bruzual & Charlot 2003 or Maraston
    2005 models.

    """

    def __init__(self,
                 wavelength=[],
                 flux=[],
                 z=0.0,
                 ageGyr=None,
                 normalise=False,
                 label=None):

        # We keep a copy of the wavelength, flux at z = 0, as it's more robust
        # to copy that to self.wavelength, flux and redshift it, rather than
        # repeatedly redshifting the same arrays back and forth
        self.z0wavelength = numpy.array(wavelength)
        self.z0flux = numpy.array(flux)
        self.wavelength = numpy.array(wavelength)
        self.flux = numpy.array(flux)
        self.z = z
        self.label = label  # plain text label, handy for using in photo-z codes

        # Store the intrinsic (i.e. unextincted) flux in case we change
        # extinction
        self.EBMinusV = 0.0
        self.intrinsic_z0flux = numpy.array(flux)

        if normalise:
            self.normalise()

        if z != 0.0:
            self.redshift(z)

        self.ageGyr = ageGyr

    def copy(self):
        """Copies the SED, returning a new SED object

        @rtype: L{SED} object
        @return: SED

        """

        newSED = SED(wavelength=self.z0wavelength,
                     flux=self.z0flux,
                     z=self.z,
                     ageGyr=self.ageGyr,
                     normalise=False,
                     label=self.label)

        return newSED

    def loadFromFile(self, fileName):
        """Loads SED from a white space delimited file in the format
        wavelength, flux. Lines beginning with # are ignored.

        @type fileName: string
        @param fileName: path to file containing wavelength, flux data

        """

        inFile = open(fileName, "r")
        lines = inFile.readlines()
        inFile.close()
        wavelength = []
        flux = []
        for line in lines:
            if line[0] != "#" and len(line) > 3:
                bits = line.split()
                wavelength.append(float(bits[0]))
                flux.append(float(bits[1]))

        # Sort SED so wavelength is in ascending order
        if wavelength[0] > wavelength[-1]:
            wavelength.reverse()
            flux.reverse()

        self.z0wavelength = numpy.array(wavelength)
        self.z0flux = numpy.array(flux)
        self.wavelength = numpy.array(wavelength)
        self.flux = numpy.array(flux)

    def writeToFile(self, fileName):
        """Writes SED to a white space delimited file in the format wavelength,
        flux.

        @type fileName: string
        @param fileName: path to file

        """

        outFile = open(fileName, "w")
        for l, f in zip(self.wavelength, self.flux):
            outFile.write(str(l) + " " + str(f) + "\n")
        outFile.close()

    def asList(self):
        """Returns a two dimensional list of [wavelength, flux], suitable for
        plotting by gnuplot.

        @rtype: list
        @return: list in format [wavelength, flux]

        """

        listData = []
        for l, f in zip(self.wavelength, self.flux):
            listData.append([l, f])

        return listData

    def plot(self, xmin='min', xmax='max'):
        """Produces a simple (wavelength, flux) plot of the SED.

        @type xmin: float or 'min'
        @param xmin: minimum of the wavelength range of the plot
        @type xmax: float or 'max'
        @param xmax: maximum of the wavelength range of the plot

        """

        pylab.matplotlib.interactive(True)
        pylab.plot(self.wavelength, self.flux)

        if xmin == 'min':
            xmin = self.wavelength.min()
        if xmax == 'max':
            xmax = self.wavelength.max()

        # Sensible y scale
        plotMask = numpy.logical_and(
            numpy.greater(self.wavelength, xmin), numpy.less(self.wavelength,
                                                             xmax))
        plotMax = self.flux[plotMask].max()
        pylab.ylim(0, plotMax * 1.1)
        pylab.xlim(xmin, xmax)
        pylab.xlabel("Wavelength")
        pylab.ylabel("Relative Flux")

    def integrate(self, wavelengthMin='min', wavelengthMax='max'):
        """Calculates flux in SED within given wavelength range.

        @type wavelengthMin: float or 'min'
        @param wavelengthMin: minimum of the wavelength range
        @type wavelengthMax: float or 'max'
        @param wavelengthMax: maximum of the wavelength range
        @rtype: float
        @return: relative flux

        """

        if wavelengthMin == 'min':
            wavelengthMin = self.wavelength.min()
        if wavelengthMax == 'max':
            wavelengthMax = self.wavelength.max()

        mask = numpy.logical_and(numpy.greater(self.wavelength, wavelengthMin),
                                 numpy.less(self.wavelength, wavelengthMax))
        flux = numpy.trapz(self.flux[mask], self.wavelength[mask])

        return flux

    def smooth(self, smoothPix):
        """Smooths SED.flux with a uniform (boxcar) filter of width smoothPix.
        Cannot be undone.

        @type smoothPix: int
        @param smoothPix: size of uniform filter applied to SED, in pixels

        """
        smoothed = ndimage.uniform_filter1d(self.flux, smoothPix)
        self.flux = smoothed

    def redshift(self, z):
        """Redshifts the SED to redshift z.

        @type z: float
        @param z: redshift

        """

        # We have to conserve energy so the area under the redshifted SED has
        # to be equal to the area under the unredshifted SED, otherwise
        # magnitude calculations will be wrong when comparing SEDs at different
        # zs
        self.wavelength = numpy.zeros(self.z0wavelength.shape[0])
        self.flux = numpy.zeros(self.z0flux.shape[0])
        self.wavelength = self.wavelength + self.z0wavelength
        self.flux = self.flux + self.z0flux

        z0TotalFlux = numpy.trapz(self.z0wavelength, self.z0flux)
        self.wavelength = self.wavelength * (1.0 + z)
        zTotalFlux = numpy.trapz(self.wavelength, self.flux)
        self.flux = self.flux * (z0TotalFlux / zTotalFlux)
        self.z = z

    def normalise(self, minWavelength='min', maxWavelength='max'):
        """Normalises the SED such that the area under the specified wavelength
        range is equal to 1.

        @type minWavelength: float or 'min'
        @param minWavelength: minimum wavelength of range over which to
        normalise SED
        @type maxWavelength: float or 'max'
        @param maxWavelength: maximum wavelength of range over which to
        normalise SED

        """
        if minWavelength == 'min':
            minWavelength = self.wavelength.min()
        if maxWavelength == 'max':
            maxWavelength = self.wavelength.max()

        lowCut = numpy.greater(self.wavelength, minWavelength)
        highCut = numpy.less(self.wavelength, maxWavelength)
        totalCut = numpy.logical_and(lowCut, highCut)
        sedFluxSlice = self.flux[totalCut]
        sedWavelengthSlice = self.wavelength[totalCut]

        self.flux = self.flux / numpy.trapz(
            abs(sedFluxSlice), sedWavelengthSlice)  # self.wavelength)

    def normaliseToMag(self, ABMag, passband):
        """Normalises the SED to match the flux equivalent to the given AB
        magnitude in the given passband.

        @type ABMag: float
        @param ABMag: AB magnitude to which the SED is to be normalised at the
        given passband
        @type passband: an L{Passband} object
        @param passband: passband at which normalisation to AB magnitude is
        calculated

        """

        magFlux = mag2Flux(ABMag, 0.0, passband)
        sedFlux = self.calcFlux(passband)
        norm = magFlux[0] / sedFlux
        self.flux = self.flux * norm
        self.z0flux = self.z0flux * norm

    def matchFlux(self, matchSED, minWavelength, maxWavelength):
        """Matches the flux in the wavelength range given by minWavelength,
        maxWavelength to the flux in the same region in matchSED. Useful for
        plotting purposes.

        @type matchSED: L{SED} object
        @param matchSED: SED to match flux to
        @type minWavelength: float
        @param minWavelength: minimum of range in which to match flux of
        current SED to matchSED
        @type maxWavelength: float
        @param maxWavelength: maximum of range in which to match flux of
        current SED to matchSED

        """

        interpMatch = interpolate.interp1d(matchSED.wavelength,
                                           matchSED.flux,
                                           kind='linear')
        interpSelf = interpolate.interp1d(self.wavelength,
                                          self.flux,
                                          kind='linear')

        wavelengthRange = numpy.arange(minWavelength, maxWavelength, 5.0)

        matchFlux = numpy.trapz(interpMatch(wavelengthRange), wavelengthRange)
        selfFlux = numpy.trapz(interpSelf(wavelengthRange), wavelengthRange)

        self.flux = self.flux * (matchFlux / selfFlux)

    def calcFlux(self, passband):
        """Calculates flux in the given passband.

        @type passband: L{Passband} object
        @param passband: filter passband through which to calculate the flux
        from the SED
        @rtype: float
        @return: flux

        """
        lowCut = numpy.greater(self.wavelength, passband.wavelength.min())
        highCut = numpy.less(self.wavelength, passband.wavelength.max())
        totalCut = numpy.logical_and(lowCut, highCut)
        sedFluxSlice = self.flux[totalCut]
        sedWavelengthSlice = self.wavelength[totalCut]

        # Use linear interpolation to rebin the passband to the same dimensions as the
        # part of the SED we're interested in
        sedInBand = passband.interpolator(sedWavelengthSlice) * sedFluxSlice
        totalFlux = numpy.trapz(sedInBand * sedWavelengthSlice,
                                sedWavelengthSlice)
        totalFlux = totalFlux /\
                    numpy.trapz(passband.interpolator(sedWavelengthSlice) *
                    sedWavelengthSlice, sedWavelengthSlice)

        return totalFlux

    def calcMag(self, passband, addDistanceModulus=True, magType="Vega"):
        """Calculates magnitude in the given passband. If addDistanceModulus ==
        True, then the distance modulus (5.0*log10*(dl*1e5), where dl is the
        luminosity distance in Mpc at the redshift of the L{SED}) is added.

        @type passband: L{Passband} object
        @param passband: filter passband through which to calculate the
        magnitude from the SED
        @type addDistanceModulus: bool
        @param addDistanceModulus: if True, adds 5.0*log10*(dl*1e5) to the mag
        returned, where dl is the luminosity distance (Mpc) corresponding to
        the SED z
        @type magType: string
        @param magType: either "Vega" or "AB"
        @rtype: float
        @return: magnitude through the given passband on the specified
        magnitude system

        """
        f1 = self.calcFlux(passband)
        if magType == "Vega":
            f2 = VEGA.calcFlux(passband)
        elif magType == "AB":
            f2 = AB.calcFlux(passband)

        mag = -2.5 * math.log10(f1 / f2)
        if magType == "Vega":
            # Add 0.026 because Vega has V=0.026 (e.g. Bohlin & Gilliland 2004)
            mag += 0.026

        if self.z > 0.0 and addDistanceModulus:
            appMag = 5.0 * math.log10(astCalc.dl(self.z) * 1e5) + mag
        else:
            appMag = mag

        return appMag

    def calcColour(self, passband1, passband2, magType="Vega"):
        """Calculates the colour passband1-passband2.

        @type passband1: L{Passband} object
        @param passband1: filter passband through which to calculate the first
        magnitude
        @type passband2: L{Passband} object
        @param passband1: filter passband through which to calculate the second
        magnitude
        @type magType: string
        @param magType: either "Vega" or "AB"
        @rtype: float
        @return: colour defined by passband1 - passband2 on the specified
        magnitude system

        """
        mag1 = self.calcMag(passband1,
                            magType=magType,
                            addDistanceModulus=True)
        mag2 = self.calcMag(passband2,
                            magType=magType,
                            addDistanceModulus=True)

        colour = mag1 - mag2
        return colour

    def getSEDDict(self, passbands):
        """This is a convenience function for pulling out fluxes from a SED for
        a given set of passbands
        in the same format as made by L{mags2SEDDict} - designed to make
        fitting code simpler.

        @type passbands: list of L{Passband} objects
        @param passbands: list of passbands through which fluxes will be
        calculated

        """

        flux = []
        wavelength = []
        for p in passbands:
            flux.append(self.calcFlux(p))
            wavelength.append(p.effectiveWavelength())

        SEDDict = {}
        SEDDict['flux'] = numpy.array(flux)
        SEDDict['wavelength'] = numpy.array(wavelength)

        return SEDDict

    def extinctionCalzetti(self, EBMinusV):
        """Applies the Calzetti et al. 2000 (ApJ, 533, 682) extinction law to
        the SED with the given E(B-V) amount of extinction. R_v' = 4.05 is
        assumed (see equation (5) of Calzetti et al.).

        @type EBMinusV: float
        @param EBMinusV: extinction E(B-V), in magnitudes

        """

        self.EBMinusV = EBMinusV

        # All done in rest frame
        self.z0flux = self.intrinsic_z0flux

        # Allow us to set EBMinusV == 0 to turn extinction off
        if EBMinusV > 0:
            # Note that EBMinusV is assumed to be Es as in equations (2) - (5)
            # Note here wavelength units have to be microns for constants to
            # make sense
            RvPrime = 4.05  # equation (5) of Calzetti et al. 2000
            shortWavelengthMask =\
                numpy.logical_and(numpy.greater_equal(self.z0wavelength, 1200),
                              numpy.less(self.z0wavelength, 6300))
            longWavelengthMask =\
                numpy.logical_and(numpy.greater_equal(self.z0wavelength, 6300),
                              numpy.less_equal(self.z0wavelength, 22000))
            wavelengthMicrons = numpy.array(self.z0wavelength / 10000.0,
                                            dtype=numpy.float64)
            kPrime = numpy.zeros(self.z0wavelength.shape[0],
                                 dtype=numpy.float64)
            kPrimeLong = (2.659 * (-1.857 + 1.040 / wavelengthMicrons)) +\
                        RvPrime
            kPrimeShort = (2.659 * (-2.156 + 1.509 / wavelengthMicrons -
                                    0.198 / wavelengthMicrons**2 + 0.011 /
                                    wavelengthMicrons**3)) + RvPrime
            kPrime[longWavelengthMask] = kPrimeLong[longWavelengthMask]
            kPrime[shortWavelengthMask] = kPrimeShort[shortWavelengthMask]

            # Here we extrapolate kPrime in similar way to what HYPERZ does
            # Short wavelengths
            try:
                interpolator = interpolate.interp1d(self.z0wavelength,
                                                    kPrimeShort,
                                                    kind='linear')
                slope = (interpolator(1100.0) - interpolator(1200.0)) / (
                    1100.0 - 1200.0)
                intercept = interpolator(1200.0) - (slope * 1200.0)
                mask = numpy.less(self.z0wavelength, 1200.0)
                kPrime[mask] = slope * self.z0wavelength[mask] + intercept

                # Long wavelengths
                interpolator = interpolate.interp1d(self.z0wavelength,
                                                    kPrimeLong,
                                                    kind='linear')
                slope = (interpolator(21900.0) - interpolator(22000.0)) / (
                    21900.0 - 22000.0)
                intercept = interpolator(21900.0) - (slope * 21900.0)
                mask = numpy.greater(self.z0wavelength, 22000.0)
                kPrime[mask] = slope * self.z0wavelength[mask] + intercept
            except:
                raise Exception("This SED has a wavelength range that doesn't "
                                "cover ~1200-22000 Angstroms")

            # Never let go negative
            kPrime[numpy.less_equal(kPrime, 0.0)] = 1e-6

            reddening = numpy.power(10, 0.4 * EBMinusV * kPrime)
            self.z0flux = self.z0flux / reddening

        self.redshift(self.z)

#-----------------------------------------------------------------------------
class VegaSED(SED):
    """This class stores the SED of Vega, used for calculation of magnitudes on the Vega system.

    The Vega SED used is taken from Bohlin 2007
    (http://adsabs.harvard.edu/abs/2007ASPC..364..315B), and is available from
    the STScI CALSPEC library
    (http://www.stsci.edu/hst/observatory/cdbs/calspec.html).

    """

    def __init__(self, normalise=False):

        VEGA_SED_PATH = astLib.__path__[
            0] + os.path.sep + "data" + os.path.sep + "bohlin2006_Vega.sed"  # from HST CALSPEC

        inFile = open(VEGA_SED_PATH, "r")
        lines = inFile.readlines()

        wavelength = []
        flux = []
        for line in lines:

            if line[0] != "#" and len(line) > 3:

                bits = line.split()
                flux.append(float(bits[1]))
                wavelength.append(float(bits[0]))

        self.wavelength = numpy.array(wavelength)
        self.flux = numpy.array(flux, dtype=numpy.float64)

        # We may want to redshift reference SEDs to calculate rest-frame colors
        # from SEDs at different zs
        self.z0wavelength = numpy.array(wavelength)
        self.z0flux = numpy.array(flux, dtype=numpy.float64)
        self.z = 0.0

        #if normalise == True:
        #self.flux=self.flux/numpy.trapz(self.flux, self.wavelength)
        #self.z0flux=self.z0flux/numpy.trapz(self.z0flux, self.z0wavelength)

#-----------------------------------------------------------------------------
class StellarPopulation:
    """This class describes a stellar population model, either a Simple Stellar
    Population (SSP) or a Composite Stellar Population (CSP), such as the
    models of Bruzual & Charlot 2003 or Maraston 2005.

    The constructor for this class can be used for generic SSPs or CSPs stored
    in white space delimited text files, containing columns for age,
    wavelength, and flux. Columns are counted from 0 ... n. Lines starting
    with # are ignored.

    The classes L{M05Model} (for Maraston 2005 models), L{BC03Model} (for
    Bruzual & Charlot 2003 models), and L{P09Model} (for Percival et al. 2009
    models) are derived from this class. The only difference between them is
    the code used to load in the model data.

    """

    def __init__(self,
                 fileName,
                 ageColumn=0,
                 wavelengthColumn=1,
                 fluxColumn=2):

        inFile = open(fileName, "r")
        lines = inFile.readlines()
        inFile.close()

        self.fileName = fileName

        # Extract a list of model ages and valid wavelengths from the file
        self.ages = []
        self.wavelengths = []
        for line in lines:
            if line[0] != "#" and len(line) > 3:
                bits = line.split()
                age = float(bits[ageColumn])
                wavelength = float(bits[wavelengthColumn])
                if age not in self.ages:
                    self.ages.append(age)
                if wavelength not in self.wavelengths:
                    self.wavelengths.append(wavelength)

        # Construct a grid of flux - rows correspond to each wavelength, columns to age
        self.fluxGrid = numpy.zeros([len(self.ages), len(self.wavelengths)])
        for line in lines:
            if line[0] != "#" and len(line) > 3:
                bits = line.split()
                sedAge = float(bits[ageColumn])
                sedWavelength = float(bits[wavelengthColumn])
                sedFlux = float(bits[fluxColumn])

                row = self.ages.index(sedAge)
                column = self.wavelengths.index(sedWavelength)

                self.fluxGrid[row][column] = sedFlux

    def getSED(self, ageGyr, z=0.0, normalise=False, label=None):
        """Extract a SED for given age. Do linear interpolation between models
        if necessary.

        @type ageGyr: float
        @param ageGyr: age of the SED in Gyr
        @type z: float
        @param z: redshift the SED from z = 0 to z = z
        @type normalise: bool
        @param normalise: normalise the SED to have area 1
        @rtype: L{SED} object
        @return: SED

        """

        if ageGyr in self.ages:

            flux = self.fluxGrid[self.ages.index(ageGyr)]
            sed = SED(self.wavelengths,
                      flux,
                      z=z,
                      normalise=normalise,
                      label=label)
            return sed

        else:

            # Use interpolation, iterating over each wavelength column
            flux = []
            for i in range(len(self.wavelengths)):
                interpolator = interpolate.interp1d(self.ages,
                                                    self.fluxGrid[:, i],
                                                    kind='linear')
                sedFlux = interpolator(ageGyr)
                flux.append(sedFlux)
            sed = SED(self.wavelengths,
                      flux,
                      z=z,
                      normalise=normalise,
                      label=label)
            return sed

    def getColourEvolution(self,
                           passband1,
                           passband2,
                           zFormation,
                           zStepSize=0.05,
                           magType="Vega"):
        """Calculates the evolution of the colour observed through passband1 -
        passband2 for the StellarPopulation with redshift, from z = 0 to z =
        zFormation.

        @type passband1: L{Passband} object
        @param passband1: filter passband through which to calculate the first
        magnitude
        @type passband2: L{Passband} object
        @param passband2: filter passband through which to calculate the second
        magnitude
        @type zFormation: float
        @param zFormation: formation redshift of the StellarPopulation
        @type zStepSize: float
        @param zStepSize: size of interval in z at which to calculate model
        colours
        @type magType: string
        @param magType: either "Vega" or "AB"
        @rtype: dictionary
        @return: dictionary of numpy.arrays in format {'z', 'colour'}

        """

        zSteps = int(math.ceil(zFormation / zStepSize))
        zData = []
        colourData = []
        for i in range(1, zSteps):
            zc = i * zStepSize
            age = astCalc.tl(zFormation) - astCalc.tl(zc)
            sed = self.getSED(age, z=zc)
            colour = sed.calcColour(passband1, passband2, magType=magType)
            zData.append(zc)
            colourData.append(colour)

        zData = numpy.array(zData)
        colourData = numpy.array(colourData)

        return {'z': zData, 'colour': colourData}

    def getMagEvolution(self,
                        passband,
                        magNormalisation,
                        zNormalisation,
                        zFormation,
                        zStepSize=0.05,
                        onePlusZSteps=False,
                        magType="Vega"):
        """Calculates the evolution with redshift (from z = 0 to z =
        zFormation) of apparent magnitude in the observed frame through the
        passband for the StellarPopulation, normalised to magNormalisation
        (apparent) at z = zNormalisation.

        @type passband: L{Passband} object
        @param passband: filter passband through which to calculate the
        magnitude
        @type magNormalisation: float
        @param magNormalisation: sets the apparent magnitude of the SED at
        zNormalisation
        @type zNormalisation: float
        @param zNormalisation: the redshift at which the magnitude
        normalisation is carried out
        @type zFormation: float
        @param zFormation: formation redshift of the StellarPopulation
        @type zStepSize: float
        @param zStepSize: size of interval in z at which to calculate model
        magnitudes
        @type onePlusZSteps: bool
        @param onePlusZSteps: if True, zSteps are (1+z)*zStepSize, otherwise
        zSteps are linear
        @type magType: string
        @param magType: either "Vega" or "AB"
        @rtype: dictionary
        @return: dictionary of numpy.arrays in format {'z', 'mag'}

        """

        # Count upwards in z steps as interpolation doesn't work if array ordered z decreasing
        zSteps = int(math.ceil(zFormation / zStepSize))
        zData = []
        magData = []
        absMagData = []
        zc0 = 0.0
        for i in range(1, zSteps):
            if not onePlusZSteps:
                zc = i * zStepSize
            else:
                zc = zc0 + (1 + zc0) * zStepSize
                zc0 = zc
                if zc >= zFormation:
                    break
            age = astCalc.tl(zFormation) - astCalc.tl(zc)
            sed = self.getSED(age, z=zc)
            mag = sed.calcMag(passband,
                              magType=magType,
                              addDistanceModulus=True)
            zData.append(zc)
            magData.append(mag)
            absMagData.append(sed.calcMag(passband, addDistanceModulus=False))

        zData = numpy.array(zData)
        magData = numpy.array(magData)

        # Do the normalisation
        interpolator = interpolate.interp1d(zData, magData, kind='linear')
        modelNormMag = interpolator(zNormalisation)
        normConstant = magNormalisation - modelNormMag
        magData = magData + normConstant

        return {'z': zData, 'mag': magData}

    def calcEvolutionCorrection(self,
                                zFrom,
                                zTo,
                                zFormation,
                                passband,
                                magType="Vega"):
        """Calculates the evolution correction in magnitudes in the rest frame
        through the passband from redshift zFrom to redshift zTo, where the
        stellarPopulation is assumed to be formed at redshift zFormation.

        @type zFrom: float
        @param zFormation: redshift to evolution correct from
        @type zTo: float
        @param zTo: redshift to evolution correct to
        @type zFormation: float
        @param zFormation: formation redshift of the StellarPopulation
        @type passband: L{Passband} object
        @param passband: filter passband through which to calculate magnitude
        @type magType: string
        @param magType: either "Vega" or "AB"
        @rtype: float
        @return: evolution correction in magnitudes in the rest frame

        """

        ageFrom = astCalc.tl(zFormation) - astCalc.tl(zFrom)
        ageTo = astCalc.tl(zFormation) - astCalc.tl(zTo)

        fromSED = self.getSED(ageFrom)
        toSED = self.getSED(ageTo)

        fromMag = fromSED.calcMag(passband,
                                  magType=magType,
                                  addDistanceModulus=False)
        toMag = toSED.calcMag(passband,
                              magType=magType,
                              addDistanceModulus=False)

        return fromMag - toMag

#-----------------------------------------------------------------------------
class M05Model(StellarPopulation):
    """This class describes a Maraston 2005 stellar population model. To load a
    composite stellar population model (CSP) for a tau = 0.1 Gyr burst of star
    formation, solar metallicity, Salpeter IMF:

    m05csp = astSED.M05Model(M05_DIR+"/csp_e_0.10_z02_salp.sed_agb")

    where M05_DIR is set to point to the directory where the Maraston 2005
    models are stored on your system.

    The file format of the Maraston 2005 simple stellar poulation (SSP) models
    is different to the file format used for the CSPs, and this needs to be
    specified using the fileType parameter. To load a SSP with solar
    metallicity, red horizontal branch morphology:

    m05ssp = astSED.M05Model(M05_DIR+"/sed.ssz002.rhb", fileType = "ssp")

    The wavelength units of SEDs from M05 models are Angstroms, with flux in
    units of erg/s/Angstrom.

    """

    def __init__(self, fileName, fileType="csp"):

        self.modelFamily = "M05"

        inFile = open(fileName, "r")
        lines = inFile.readlines()
        inFile.close()

        self.fileName = fileName

        if fileType == "csp":
            ageColumn = 0
            wavelengthColumn = 1
            fluxColumn = 2
        elif fileType == "ssp":
            ageColumn = 0
            wavelengthColumn = 2
            fluxColumn = 3
        else:
            raise Exception("fileType must be 'ssp' or 'csp'")

        # Extract a list of model ages and valid wavelengths from the file
        self.ages = []
        self.wavelengths = []
        for line in lines:
            if line[0] != "#" and len(line) > 3:
                bits = line.split()
                age = float(bits[ageColumn])
                wavelength = float(bits[wavelengthColumn])
                if age not in self.ages:
                    self.ages.append(age)
                if wavelength not in self.wavelengths:
                    self.wavelengths.append(wavelength)

        # Construct a grid of flux - rows correspond to each wavelength, columns to age
        self.fluxGrid = numpy.zeros([len(self.ages), len(self.wavelengths)])
        for line in lines:
            if line[0] != "#" and len(line) > 3:
                bits = line.split()
                sedAge = float(bits[ageColumn])
                sedWavelength = float(bits[wavelengthColumn])
                sedFlux = float(bits[fluxColumn])

                row = self.ages.index(sedAge)
                column = self.wavelengths.index(sedWavelength)

                self.fluxGrid[row][column] = sedFlux

#-----------------------------------------------------------------------------
class BC03Model(StellarPopulation):
    """This class describes a Bruzual & Charlot 2003 stellar population model,
    extracted from a GALAXEV .ised file using the galaxevpl program that is
    included in GALAXEV. The file format is white space delimited, with
    wavelength in the first column. Subsequent columns contain the model fluxes
    for SEDs of different ages, as specified when running galaxevpl. The age
    corresponding to each flux column is taken from the comment line beginning
    "# Age (yr)", and is converted to Gyr.

    For example, to load a tau = 0.1 Gyr burst of star formation,  solar
    metallicity, Salpeter IMF model stored in a file (created by galaxevpl)
    called "csp_lr_solar_0p1Gyr.136":

    bc03model = BC03Model("csp_lr_solar_0p1Gyr.136")

    The wavelength units of SEDs from BC03 models are Angstroms. Flux is
    converted into units of erg/s/Angstrom (the units in the files output by
    galaxevpl are LSun/Angstrom).

    """

    def __init__(self, fileName):

        self.modelFamily = "BC03"
        self.fileName = fileName

        inFile = open(fileName, "r")
        lines = inFile.readlines()
        inFile.close()

        # Extract a list of model ages - BC03 ages are in years, so convert to Gyr
        self.ages = []
        for line in lines:
            if line.find("# Age (yr)") != -1:
                rawAges = line[line.find("# Age (yr)") + 10:].split()
                for age in rawAges:
                    self.ages.append(float(age) / 1e9)

        # Extract a list of valid wavelengths from the file
        # If we have many ages in the file, this is more complicated...
        lambdaLinesCount = 0
        startFluxDataLine = None
        for i in range(len(lines)):
            line = lines[i]
            if "# Lambda(A)" in line:
                lambdaLinesCount = lambdaLinesCount + 1
            if line[0] != "#" and len(line) > 3 and startFluxDataLine is None:
                startFluxDataLine = i
        self.wavelengths = []
        for i in range(startFluxDataLine, len(lines), lambdaLinesCount):
            line = lines[i]
            bits = line.split()
            self.wavelengths.append(float(bits[0]))

        # Construct a grid of flux - rows correspond to each wavelength, columns to age
        self.fluxGrid = numpy.zeros([len(self.ages), len(self.wavelengths)])
        for i in range(startFluxDataLine, len(lines), lambdaLinesCount):
            line = lines[i]
            bits = []
            for k in range(i, i + lambdaLinesCount):
                bits = bits + lines[k].split()
            ageFluxes = bits[1:]
            sedWavelength = float(bits[0])
            column = self.wavelengths.index(sedWavelength)
            for row in range(len(ageFluxes)):
                sedFlux = float(ageFluxes[row])
                self.fluxGrid[row][column] = sedFlux

        # Convert flux into erg/s/Angstrom - native units of galaxevpl files are LSun/Angstrom
        self.fluxGrid = self.fluxGrid * 3.826e33

#-----------------------------------------------------------------------------
class P09Model(StellarPopulation):
    """This class describes a Percival et al 2009 (BaSTI;
    http://albione.oa-teramo.inaf.it) stellar population model. We assume that
    the synthetic spectra for each model are unpacked under the directory
    pointed to by fileName.

    The wavelength units of SEDs from P09 models are converted to Angstroms.
    Flux is converted into units of erg/s/Angstrom (the units in the BaSTI
    low-res spectra are 4.3607e-33 erg/s/m).

    """

    def __init__(self, fileName):

        self.modelFamily = "P09"

        files = glob.glob(fileName + os.path.sep + "*.t??????")

        self.fileName = fileName

        # Map end of filenames to ages in Gyr
        extensionAgeMap = {}
        self.ages = []
        for f in files:
            ext = f.split(".")[-1]
            ageGyr = float(f[-5:]) / 1e3
            self.ages.append(ageGyr)
            extensionAgeMap[ext] = ageGyr
        self.ages.sort()

        # Construct a grid of flux - rows correspond to each wavelength, columns to age
        self.wavelengths = None
        self.fluxGrid = None
        for i in range(len(self.ages)):
            for e in extensionAgeMap.keys():
                if extensionAgeMap[e] == self.ages[i]:
                    inFileName = glob.glob(fileName + os.path.sep + "*." + e)[
                        0]
                    inFile = open(inFileName, "r")
                    lines = inFile.readlines()
                    inFile.close()
                    wavelength = []
                    flux = []
                    for line in lines:
                        bits = line.split()
                        wavelength.append(
                            float(bits[0]) *
                            10.0)  # units in file are nm, not angstroms
                        flux.append(float(bits[1]))
                    if self.wavelengths is None:
                        self.wavelengths = wavelength
                    if self.fluxGrid is None:
                        self.fluxGrid = numpy.zeros(
                            [len(self.ages), len(self.wavelengths)])
                    self.fluxGrid[i] = flux

        # Convert flux into erg/s/Angstrom - native units in BaSTI files are
        # 4.3607e-33 erg/s/m
        self.fluxGrid = self.fluxGrid / 4.3607e-33 / 1e10

#-----------------------------------------------------------------------------
def makeModelSEDDictList(modelList,
                         z,
                         passbandsList,
                         labelsList=[],
                         EBMinusVList=[0.0],
                         forceYoungerThanUniverse=True):
    """This routine makes a list of SEDDict dictionaries (see L{mags2SEDDict})
    for fitting using L{fitSEDDict}. This speeds up the fitting as this allows
    us to calculate model SED magnitudes only once, if all objects to be fitted
    are at the same redshift. We add some meta data to the modelSEDDicts (e.g.
    the model file names).

    The effect of extinction by dust (assuming the Calzetti et al. 2000 law)
    can be included by giving a list of E(B-V) values.

    If forceYoungerThanUniverse == True, ages which are older than the universe
    at the given z will not be included.

    @type modelList: list of L{StellarPopulation} model objects
    @param modelList: list of StellarPopulation models to include
    @type z: float
    @param z: redshift to apply to all stellar population models in modelList
    @type EBMinusVList: list
    @param EBMinusVList: list of E(B-V) extinction values to apply to all
    models, in magnitudes
    @type labelsList: list
    @param labelsList: optional list used for labelling passbands in output
    SEDDicts
    @type forceYoungerThanUniverse: bool
    @param forceYoungerThanUniverse: if True, do not allow models that exceed
    the age of the universe at z
    @rtype: list
    @return: list of dictionaries containing model fluxes, to be used as input
    to L{fitSEDDict}.

    """

    # Otherwise if this is the case we won't actually make any model SEDDicts ...
    if EBMinusVList == []:
        EBMinusVList = [0.0]

    modelSEDDictList = []
    for m in range(len(modelList)):
        testAges = numpy.array(modelList[m].ages)
        if forceYoungerThanUniverse:
            testAges = testAges[numpy.logical_and(
                numpy.less(testAges, astCalc.tz(z)), numpy.greater(testAges,
                                                                   0))]
        for t in testAges:
            s = modelList[m].getSED(
                t,
                z=z,
                label=modelList[m].fileName + " - age=" + str(t) + " Gyr")
            for EBMinusV in EBMinusVList:
                try:
                    s.extinctionCalzetti(EBMinusV)
                except:
                    raise Exception(
                        "Model %s has a wavelength range that doesn't cover ~1200-22000 Angstroms"
                        % (modelList[m].fileName))
                modelSEDDict = s.getSEDDict(passbandsList)
                modelSEDDict['labels'] = labelsList
                modelSEDDict['E(B-V)'] = EBMinusV
                modelSEDDict['ageGyr'] = t
                modelSEDDict['z'] = z
                modelSEDDict['fileName'] = modelList[m].fileName
                modelSEDDict['modelListIndex'] = m
                modelSEDDictList.append(modelSEDDict)

    return modelSEDDictList

#-----------------------------------------------------------------------------
def fitSEDDict(SEDDict, modelSEDDictList):
    """Fits the given SED dictionary (made using L{mags2SEDDict}) with the
    given list of model SED dictionaries. The latter should be made using
    L{makeModelSEDDictList}, and entries for fluxes should correspond directly
    between the model and SEDDict.

    Returns a dictionary with best fit values.

    @type SEDDict: dictionary, in format of L{mags2SEDDict}
    @param SEDDict: dictionary of observed fluxes and uncertainties, in format
    of L{mags2SEDDict}
    @type modelSEDDictList: list of dictionaries, in format of
    L{makeModelSEDDictList}
    @param modelSEDDictList: list of dictionaries containing fluxes of models
    to be fitted to the observed fluxes listed in the SEDDict. This should be
    made using L{makeModelSEDDictList}.
    @rtype: dictionary
    @return: results of the fitting - keys:
             - 'minChiSq': minimum chi squared value of best fit
             - 'chiSqContrib': corresponding contribution at each passband to
             the minimum chi squared value
             - 'ageGyr': the age in Gyr of the best fitting model
             - 'modelFileName': the file name of the stellar population model
             corresponding to the best fit
             - 'modelListIndex': the index of the best fitting model in the
             input modelSEDDictList
             - 'norm': the normalisation that the best fit model should be
             multiplied by to match the SEDDict
             - 'z': the redshift of the best fit model
             - 'E(B-V)': the extinction, E(B-V), in magnitudes, of the best fit
             model

    """

    modelFlux = []
    for modelSEDDict in modelSEDDictList:
        modelFlux.append(modelSEDDict['flux'])
    modelFlux = numpy.array(modelFlux)
    sedFlux = numpy.array([SEDDict['flux']] * len(modelSEDDictList))
    sedFluxErr = numpy.array([SEDDict['fluxErr']] * len(modelSEDDictList))

    # Analytic expression below is for normalisation at minimum chi squared (see note book)
    norm = numpy.sum(
        (modelFlux * sedFlux) /
        (sedFluxErr**2), axis=1) / numpy.sum(modelFlux**2 / sedFluxErr**2,
                                             axis=1)
    norms = numpy.array([norm] * modelFlux.shape[1]).transpose()
    chiSq = numpy.sum(((sedFlux - norms * modelFlux)**2) / sedFluxErr**2,
                      axis=1)
    chiSq[numpy.isnan(
        chiSq)] = 1e6  # throw these out, should check this out and handle more gracefully
    minChiSq = chiSq.min()
    bestMatchIndex = numpy.equal(chiSq, minChiSq).nonzero()[0][0]
    bestNorm = norm[bestMatchIndex]
    bestChiSq = minChiSq
    bestChiSqContrib = ((sedFlux[bestMatchIndex] - norms[bestMatchIndex] *
                         modelFlux[bestMatchIndex])**2) /\
                            sedFluxErr[bestMatchIndex]**2

    resultsDict = {'minChiSq': bestChiSq,
                   'chiSqContrib': bestChiSqContrib,
                   'allChiSqValues': chiSq,
                   'ageGyr': modelSEDDictList[bestMatchIndex]['ageGyr'],
                   'modelFileName':
                   modelSEDDictList[bestMatchIndex]['fileName'],
                   'modelListIndex':
                   modelSEDDictList[bestMatchIndex]['modelListIndex'],
                   'norm': bestNorm,
                   'z': modelSEDDictList[bestMatchIndex]['z'],
                   'E(B-V)': modelSEDDictList[bestMatchIndex]['E(B-V)']}

    return resultsDict

#-----------------------------------------------------------------------------
def mags2SEDDict(ABMags, ABMagErrs, passbands):
    """Takes a set of corresponding AB magnitudes, uncertainties, and
    passbands, and returns a dictionary with keys 'flux', 'fluxErr'
    'wavelength'. Fluxes are in units of erg/s/cm^2/Angstrom, wavelength in
    Angstroms. These dictionaries are the staple diet of the L{fitSEDDict}
    routine.

    @type ABMags: list or numpy array
    @param ABMags: AB magnitudes, specified in corresponding order to passbands
    and ABMagErrs
    @type ABMagErrs: list or numpy array
    @param ABMagErrs: AB magnitude errors, specified in corresponding order to
    passbands and ABMags
    @type passbands: list of L{Passband} objects
    @param passbands: passband objects, specified in corresponding order to
    ABMags and ABMagErrs
    @rtype: dictionary
    @return: dictionary with keys {'flux', 'fluxErr', 'wavelength'}, suitable
    for input to L{fitSEDDict}

    """

    flux = []
    fluxErr = []
    wavelength = []
    for m, e, p in zip(ABMags, ABMagErrs, passbands):
        f, err = mag2Flux(m, e, p)
        flux.append(f)
        fluxErr.append(err)
        wavelength.append(p.effectiveWavelength())

    SEDDict = {}
    SEDDict['flux'] = numpy.array(flux)
    SEDDict['fluxErr'] = numpy.array(fluxErr)
    SEDDict['wavelength'] = numpy.array(wavelength)

    return SEDDict

#-----------------------------------------------------------------------------
def mag2Flux(ABMag, ABMagErr, passband):
    """Converts given AB magnitude and uncertainty into flux, in
    erg/s/cm^2/Angstrom.

    @type ABMag: float
    @param ABMag: magnitude on AB system in passband
    @type ABMagErr: float
    @param ABMagErr: uncertainty in AB magnitude in passband
    @type passband: L{Passband} object
    @param passband: L{Passband} object at which ABMag was measured
    @rtype: list
    @return: [flux, fluxError], in units of erg/s/cm^2/Angstrom

    """

    fluxJy = (10**23.0) * 10**(-(ABMag + 48.6) / 2.5)  # AB mag
    aLambda = 3e-13  # for conversion to erg s-1 cm-2 angstrom-1 with lambda in microns
    effLMicron = passband.effectiveWavelength() * (1e-10 / 1e-6)
    fluxWLUnits = aLambda * fluxJy / effLMicron**2

    fluxJyErr = (10**23.0) * 10**(-(ABMag - ABMagErr + 48.6) / 2.5)  # AB mag
    fluxWLUnitsErr = aLambda * fluxJyErr / effLMicron**2
    fluxWLUnitsErr = fluxWLUnitsErr - fluxWLUnits

    return [fluxWLUnits, fluxWLUnitsErr]

#-----------------------------------------------------------------------------
def flux2Mag(flux, fluxErr, passband):
    """Converts given flux and uncertainty in erg/s/cm^2/Angstrom into AB
    magnitudes.

    @type flux: float
    @param flux: flux in erg/s/cm^2/Angstrom in passband
    @type fluxErr: float
    @param fluxErr: uncertainty in flux in passband, in erg/s/cm^2/Angstrom
    @type passband: L{Passband} object
    @param passband: L{Passband} object at which ABMag was measured
    @rtype: list
    @return: [ABMag, ABMagError], in AB magnitudes

    """

    # aLambda = 3x10-5 for effective wavelength in angstroms
    aLambda = 3e-13  # for conversion to erg s-1 cm-2 angstrom-1 with lambda in microns
    effLMicron = passband.effectiveWavelength() * (1e-10 / 1e-6)

    fluxJy = (flux * effLMicron**2) / aLambda
    mag = -2.5 * numpy.log10(fluxJy / 10**23) - 48.6

    fluxErrJy = (fluxErr * effLMicron**2) / aLambda
    magErr = mag - (-2.5 * numpy.log10((fluxJy + fluxErrJy) / 10**23) - 48.6)

    return [mag, magErr]

#-----------------------------------------------------------------------------
def mag2Jy(ABMag):
    """Converts an AB magnitude into flux density in Jy

    @type ABMag: float
    @param ABMag: AB magnitude
    @rtype: float
    @return: flux density in Jy

    """

    fluxJy = ((10**23) * 10**(-(float(ABMag) + 48.6) / 2.5))

    return fluxJy

#-----------------------------------------------------------------------------
def Jy2Mag(fluxJy):
    """Converts flux density in Jy into AB magnitude

    @type fluxJy: float
    @param fluxJy: flux density in Jy
    @rtype: float
    @return: AB magnitude

    """

    ABMag = -2.5 * (numpy.log10(fluxJy) - 23.0) - 48.6

    return ABMag

#-----------------------------------------------------------------------------
# Data
VEGA = VegaSED()

# AB SED has constant flux density 3631 Jy
AB = SED(wavelength=numpy.logspace(1, 8, 1e5), flux=numpy.ones(int(1e6)))
AB.flux = (3e-5 * 3631) / (AB.wavelength**2)
AB.z0flux = AB.flux[:]

# Solar SED from HST CALSPEC (http://www.stsci.edu/hst/observatory/cdbs/calspec.html)
SOL = SED()
SOL.loadFromFile(astLib.__path__[0] + os.path.sep + "data" + os.path.sep +
                 "sun_reference_stis_001.ascii")


# -*- coding: utf-8 -*-
"""
Created on Fri Apr 22 11:42:36 2016

@author: Will Kew
will.kew@gmail.com

    Copyright Will Kew, 2016

    This file is part of FTMS Visualisation (also known as i-van Krevelen).

    FTMS Visualisation is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    FTMS Visualisation is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with FTMS Visualisation.  If not, see <http://www.gnu.org/licenses/>.

This script will read in an assigned peaklist (example input file included) and calculate the heteroatomic class distribution.
The output is a vbar plot of heteroamtic class versus count. You can also have the calculated numbers output in a format for replotting.
This tool uses Seaborn - http://seaborn.pydata.org/

A number of (partially tested) other functions to plot output are included, though commented out.

This tool was used in our recent paper on Scotch Whisky - https://link.springer.com/article/10.1007/s13361-016-1513-y
The prompt for the user about whisky samples is thus borne from this - it also serves as an example of how to customise which classes to include.
"""

from __future__ import print_function    # Python 2 compatibility
from __future__ import absolute_import   # Python 2 compatibility

import os, sys
import pandas as pd
from collections import Counter
import matplotlib.pyplot as plt
import seaborn as sns

"""
# We import also the FTMSVizProcessingModule which contains a few useful functions.
# here we define where the scripts are stored.
# Make sure to change this to where you have saved these scripts.
"""
try: #test if running in ipython
    __IPYTHON__
except NameError: #if not running in ipython....
    import FTMSVizProcessingModule as FTPM
    path  = os.getcwd()+"data\\" #example data location
else: #if running in ipython
    scriptlocation = "/LOCAL/FTMSVis/FTMSVisualization-master/"
    sys.path.append(scriptlocation)
    import FTMSVizProcessingModule as FTPM
    path = "/LOCAL/FTMSVis/data/"

whisky = input("Are these Whisky samples - Y or N?" )
if whisky.upper() == "Y":
    whisky = True
else:
    whisky = False

inputpath = path +"OutputCSV/"
outputpath = path + "Images/Classes/"
FTPM.make_sure_path_exists(outputpath) #this function checks the output directory exists; if it doesnt, it creates it.
print("Looking for CSVs in " + inputpath)
filesA = os.listdir(inputpath)
filesB = []
for y in filesA:
    if y[-8:] =="hits.csv" and y[-10:] != "nohits.csv" and y[-11:] !="isohits.csv":
        filesB.append(y)
nfiles = len(filesB)

samplenames=[]
for x in filesB:
    samplenames.append(x[:-9])

heteroclasses=[]
for z in filesB:
    df1 = pd.read_csv(inputpath+z,index_col=0)
    hetclas = df1["HeteroClass"]
    hetclaslist = hetclas.tolist()
    heteroclasses.append(hetclaslist)

heteroclasses = [item for sublist in heteroclasses for item in sublist]
hetclasset = list(set(heteroclasses))

indexlist = []
for i in samplenames:
    for n in range(len(hetclasset)):
        indexlist.append(i)

###This section is relevant to my whisky samples
if whisky == True:
    columnnames = ["Sample","Class","WoodType","Region","Age","Peated","HeteroClass","HeteroClassCount"]
    df4 = pd.read_csv(path+"SampleInfo-Dict.csv",index_col=0)
    df4 = df4.T
    dict4 = df4.to_dict()

    outputdata = pd.DataFrame(index = range(len(indexlist)), columns=columnnames)
    a = 0
    for y in filesB:
        df2 = pd.read_csv(inputpath+y,index_col=0)
        counter = Counter(df2["HeteroClass"])
        for x in counter:
            outputdata.iloc[a][0] = y[:-9]
            outputdata.iloc[a][1] = dict4[y[:-9]]["Class"]
            outputdata.iloc[a][2] = dict4[y[:-9]]["Total Wood"]
            outputdata.iloc[a][3] = dict4[y[:-9]]["Region"]
            outputdata.iloc[a][4] = dict4[y[:-9]]["Age"]
            outputdata.iloc[a][5] = dict4[y[:-9]]["Peated"]
            outputdata.iloc[a][6] = x
            outputdata.iloc[a][7] = counter[x]
            a = a+1
    outputdata = outputdata.dropna(how="all",axis=0)

else:
    columnnames = ["Sample","Class","HeteroClass","HeteroClassCount"]
    outputdata = pd.DataFrame(index = range(len(indexlist)), columns=columnnames)
    a = 0
    for y in filesB:
        df2 = pd.read_csv(inputpath+y,index_col=0)
        counter = Counter(df2["HeteroClass"])
        for x in counter:
            outputdata.iloc[a][0] = y[:-9]
            outputdata.iloc[a][1] = y[:-9] #this is the Class variable, and should be defined as approrpriate for what you're plotting. In the case of single samples, it can be the sample name.
            outputdata.iloc[a][2] = x
            outputdata.iloc[a][3] = counter[x]
            a = a+1
    outputdata = outputdata.dropna(how="all",axis=0)

pd.to_numeric(outputdata["HeteroClassCount"],errors="raise")

saveoutputdata = input("Do you want to save the output data in a text file for later re-processing - Y or N? ")
if saveoutputdata.upper() == "Y":
    outputdata.to_excel(inputpath+"HetClassByClass-longform.xlsx") #this saves the info out in a longform for plotting.

#outputdata = pd.read_excel(inputpath+"HetClassByClass-longform.xlsx") #this reads that data back in. Only necessary for manually re-running bits of script.

# This section creates a unique, naturally sorted list of heteroatom classes for plotting. Only really works for CHO formula.
# If you have exotic heteroatoms, will need to refigure this yourself, or just hardcode the order you want. easy to do in Excel.

order = outputdata["HeteroClass"].tolist()
order= list(set(order))
order.sort(key=FTPM.natural_sort_key) # this natural sort function ensures a logical order to your barplot.

if whisky == True:
    CHOorder = ["O2","O3","O4","O5","O6","O7","O8","O9","O10","O11","O12","O13","O14","O15","O16","O17","O18","O19"]
    Fullorder = ["O2","O3","O4","O5","O6","O7","O8","O9","O10","O11","O12","O13","O14","O15","O16","O17","O18",
    "O19","O1S1","O2S1","O3S1","O4S1","O5S1","O6S1","O7S1","O8S1","O9S1","O10S1","O11S1","O12S1"]
    CHOSorder =["O1S1","O2S1","O3S1","O4S1","O5S1","O6S1","O7S1","O8S1","O9S1","O10S1","O11S1","O12S1"]
    CHOSorderNew = ["O2","O3","O4","O5","O6","O7","O8","O9","O10","O11","O12","O13","O14","O15","O16","O17","O18","O19","OnS"]
    labels = ["O2","O3","O4","O5","O6","O7","O8","O9","O10","O11","O12","O13","O14","O15","O16","O17","O18","O19",r'O$\mathregular {_n}$S']

else:
    df = outputdata
    #colours = ["#a6cee3","#1f78b4","#b2df8a"] #colorblind and print friendly colours picked from http://colorbrewer2.org/
    colours = ["#1b9e77","#d95f02","#7570b3"] #as above, but brighter

def barplot():
    sns.set_style("white")
    sns.set_context("paper",font_scale=2)
    ax = sns.barplot(x="HeteroClass",y="HeteroClassCount",hue="Class",
                     data=outputdata,order=order,palette=sns.color_palette(colours))
    ax.set(xlabel='Heteroatomic Class', ylabel='Count')
    handles, labels = ax.get_legend_handles_labels()
    if len(labels) == 1:
        ax.legend_.remove()
    sns.despine()
    fig = ax.get_figure()
    plt.xticks(rotation=90)
    fig.set_size_inches(8, 6, forward=True)
    fig.savefig(outputpath+"Barplot.png",dpi=600,bbox_inches="tight")
    fig.savefig(outputpath+"Barplot.eps",dpi=600,bbox_inches="tight")

barplot() #plots a barplot.

"""
# Here are some further examples of the Seaborn Plotting library applied to this problem.
# Most of these rely on having many samples across a small number of classes you wish to compare
def violinplot():
    sns.set_style("white")
    sns.set_context("paper",font_scale=2)
    ax = sns.violinplot(x="HeteroClass",y="HeteroClassCount",hue="Class",data=outputdata,
                        order=order,
                        palette=sns.color_palette("bright"),
                        split=False,bw="silverman",scale_hue=True,scale="width",
                        cut=2,linewidth=1.5,inner="quartiles",saturation=1)
    ax.set(xlabel='Heteroatomic Class', ylabel='Count')
    sns.despine()
    fig = ax.get_figure()
    locs, labels = plt.xticks()
    plt.xticks(locs, labels, rotation=90)
    cur_ylim = ax.get_ylim()
    ax.set_ylim(0,cur_ylim[1])
    fig.set_size_inches((POPM.mm2inch(171,80)), forward=True)
    fig.savefig(outputpath+"violinplot-scalewidth.png",dpi=600,bbox_inches="tight")
    fig.savefig(outputpath+"violinplot-scalewidth.eps",dpi=600,bbox_inches="tight")

def boxplot():
    sns.set_style("white")
    sns.set_context("paper",font_scale=2)
    ax = sns.boxplot(x="HeteroClass",y="HeteroClassCount",hue="Class",data=outputdata,order=order,palette=sns.color_palette("bright"))
    ax.set(xlabel='Heteroatomic Class', ylabel='Count')
    sns.despine()
    fig = ax.get_figure()
    plt.xticks(rotation=90)
    fig.set_size_inches(8, 6, forward=True)
    fig.savefig(outputpath+"Boxplot-comparison-CHO-only.png",dpi=300,bbox_inches="tight")

def swarmplot():
    sns.set_style("white")
    sns.set_context("paper",font_scale=2)
    ax = sns.swarmplot(x="HeteroClass",y="HeteroClassCount",hue="Class",data=outputdata,order=order,palette=sns.color_palette("bright"))
    ax.set(xlabel='Heteroatomic Class', ylabel='Average Count')
    sns.despine()
    fig = ax.get_figure()
    plt.xticks(rotation=90)
    fig.set_size_inches(8, 6, forward=True)
    fig.savefig(outputpath+"swarmplot-comparison-CHO-only.png",dpi=300,bbox_inches="tight")

def stripplot():
    sns.set_style("white")
    sns.set_context("paper",font_scale=2)
    ax = sns.stripplot(x="HeteroClass",y="HeteroClassCount",hue="Class",data=outputdata,order=order,palette=sns.color_palette("bright"),jitter=False,split=True)
    ax.set(xlabel='Heteroatomic Class', ylabel='Average Count')
    sns.despine()
    fig = ax.get_figure()
    plt.xticks(rotation=90)
    fig.set_size_inches(8, 6, forward=True)
    fig.savefig(outputpath+"striplot-comparison-CHO-only.png",dpi=300,bbox_inches="tight")

"""
#EOF


# Author: Brian M. Clapper, G Varoquaux
# License: BSD

import numpy as np

# XXX we should be testing the public API here
from sklearn.utils.linear_assignment_ import _hungarian

def test_hungarian():
    matrices = [
        # Square
        ([[400, 150, 400],
          [400, 450, 600],
          [300, 225, 300]],
         850  # expected cost
         ),

        # Rectangular variant
        ([[400, 150, 400, 1],
          [400, 450, 600, 2],
          [300, 225, 300, 3]],
         452  # expected cost
         ),

        # Square
        ([[10, 10,  8],
          [9,  8,  1],
          [9,  7,  4]],
         18
         ),

        # Rectangular variant
        ([[10, 10,  8, 11],
          [9, 8, 1, 1],
          [9, 7, 4, 10]],
         15
         ),

        # n == 2, m == 0 matrix
        ([[], []],
         0
         ),
    ]

    for cost_matrix, expected_total in matrices:
        cost_matrix = np.array(cost_matrix)
        indexes = _hungarian(cost_matrix)
        total_cost = 0
        for r, c in indexes:
            x = cost_matrix[r, c]
            total_cost += x
        assert expected_total == total_cost

        indexes = _hungarian(cost_matrix.T)
        total_cost = 0
        for c, r in indexes:
            x = cost_matrix[r, c]
            total_cost += x
        assert expected_total == total_cost


from pymongo import MongoClient
from matplotlib import pyplot as plt
import os
from datetime import datetime, date, time, timedelta

client = MongoClient()

# using wowtest.auctiondata
db = client.wowtest
posts = db.auctiondata

auctions = posts.find().limit(10)
#time.time() into datetime --->
#datetime.datetime.fromtimestamp('xxxx').strftime('%c')

def dt_to_timestamp(dt):
    #timestamp = (dt - datetime(1970, 1, 1).total_seconds())
    return (int(dt.strftime('%s')))

def getdata(num, quantum):
    valid = []
    today = datetime.combine(date.today(), time())
    for i in range(num+1):
        day = today - i*quantum
        gte = dt_to_timestamp(day)
        lt = dt_to_timestamp(day+quantum)
        time_query = {'$gte':gte, '$lt':lt}
        valid.insert(0, posts.find({'viewtime':time_query}).count())
    return valid

def format_date(x, n):
    today = datetime.combine(date.today(), time())
    day = today - timedelta(hours=n-x-1)
    return day.strftime('%m%d%H')

def plotbar(data, color):
    plt.bar(range(len(data)), data, align='center', color=color)

# run
n = 48
val = getdata(n, timedelta(hours=1))
plotbar(val, '#4788d2')

plt.xticks(range(n), [format_date(i, n) for i in range(n)], size='small', rotation=90)
plt.grid(axis='y')

plt.show()


"""Nearest Neighbor Regression"""

# Authors: Jake Vanderplas <vanderplas@astro.washington.edu>
#          Fabian Pedregosa <fabian.pedregosa@inria.fr>
#          Alexandre Gramfort <alexandre.gramfort@inria.fr>
#          Sparseness support by Lars Buitinck
#          Multi-output support by Arnaud Joly <a.joly@ulg.ac.be>
#
# License: BSD 3 clause (C) INRIA, University of Amsterdam

import numpy as np

from .base import _get_weights, _check_weights, NeighborsBase, KNeighborsMixin
from .base import RadiusNeighborsMixin, SupervisedFloatMixin
from ..base import RegressorMixin
from ..utils import check_array

class KNeighborsRegressor(NeighborsBase, KNeighborsMixin,
                          SupervisedFloatMixin,
                          RegressorMixin):
    """Regression based on k-nearest neighbors.

    The target is predicted by local interpolation of the targets
    associated of the nearest neighbors in the training set.

    Read more in the :ref:`User Guide <regression>`.

    Parameters
    ----------
    n_neighbors : int, optional (default = 5)
        Number of neighbors to use by default for :meth:`kneighbors` queries.

    weights : str or callable
        weight function used in prediction.  Possible values:

        - 'uniform' : uniform weights.  All points in each neighborhood
          are weighted equally.
        - 'distance' : weight points by the inverse of their distance.
          in this case, closer neighbors of a query point will have a
          greater influence than neighbors which are further away.
        - [callable] : a user-defined function which accepts an
          array of distances, and returns an array of the same shape
          containing the weights.

        Uniform weights are used by default.

    algorithm : {'auto', 'ball_tree', 'kd_tree', 'brute'}, optional
        Algorithm used to compute the nearest neighbors:

        - 'ball_tree' will use :class:`BallTree`
        - 'kd_tree' will use :class:`KDtree`
        - 'brute' will use a brute-force search.
        - 'auto' will attempt to decide the most appropriate algorithm
          based on the values passed to :meth:`fit` method.

        Note: fitting on sparse input will override the setting of
        this parameter, using brute force.

    leaf_size : int, optional (default = 30)
        Leaf size passed to BallTree or KDTree.  This can affect the
        speed of the construction and query, as well as the memory
        required to store the tree.  The optimal value depends on the
        nature of the problem.

    metric : string or DistanceMetric object (default='minkowski')
        the distance metric to use for the tree.  The default metric is
        minkowski, and with p=2 is equivalent to the standard Euclidean
        metric. See the documentation of the DistanceMetric class for a
        list of available metrics.

    p : integer, optional (default = 2)
        Power parameter for the Minkowski metric. When p = 1, this is
        equivalent to using manhattan_distance (l1), and euclidean_distance
        (l2) for p = 2. For arbitrary p, minkowski_distance (l_p) is used.

    metric_params : dict, optional (default = None)
        Additional keyword arguments for the metric function.

    n_jobs : int, optional (default = 1)
        The number of parallel jobs to run for neighbors search.
        If ``-1``, then the number of jobs is set to the number of CPU cores.
        Doesn't affect :meth:`fit` method.

    Examples
    --------
    >>> X = [[0], [1], [2], [3]]
    >>> y = [0, 0, 1, 1]
    >>> from sklearn.neighbors import KNeighborsRegressor
    >>> neigh = KNeighborsRegressor(n_neighbors=2)
    >>> neigh.fit(X, y) # doctest: +ELLIPSIS
    KNeighborsRegressor(...)
    >>> print(neigh.predict([[1.5]]))
    [ 0.5]

    See also
    --------
    NearestNeighbors
    RadiusNeighborsRegressor
    KNeighborsClassifier
    RadiusNeighborsClassifier

    Notes
    -----
    See :ref:`Nearest Neighbors <neighbors>` in the online documentation
    for a discussion of the choice of ``algorithm`` and ``leaf_size``.

    .. warning::

       Regarding the Nearest Neighbors algorithms, if it is found that two
       neighbors, neighbor `k+1` and `k`, have identical distances but
       but different labels, the results will depend on the ordering of the
       training data.

    https://en.wikipedia.org/wiki/K-nearest_neighbor_algorithm
    """

    def __init__(self, n_neighbors=5, weights='uniform',
                 algorithm='auto', leaf_size=30,
                 p=2, metric='minkowski', metric_params=None, n_jobs=1,
                 **kwargs):
        self._init_params(n_neighbors=n_neighbors,
                          algorithm=algorithm,
                          leaf_size=leaf_size, metric=metric, p=p,
                          metric_params=metric_params, n_jobs=n_jobs, **kwargs)
        self.weights = _check_weights(weights)

    def predict(self, X):
        """Predict the target for the provided data

        Parameters
        ----------
        X : array-like, shape (n_query, n_features), \
                or (n_query, n_indexed) if metric == 'precomputed'
            Test samples.

        Returns
        -------
        y : array of int, shape = [n_samples] or [n_samples, n_outputs]
            Target values
        """
        X = check_array(X, accept_sparse='csr')

        neigh_dist, neigh_ind = self.kneighbors(X)

        weights = _get_weights(neigh_dist, self.weights)

        _y = self._y
        if _y.ndim == 1:
            _y = _y.reshape((-1, 1))

        if weights is None:
            y_pred = np.mean(_y[neigh_ind], axis=1)
        else:
            y_pred = np.empty((X.shape[0], _y.shape[1]), dtype=np.float64)
            denom = np.sum(weights, axis=1)

            for j in range(_y.shape[1]):
                num = np.sum(_y[neigh_ind, j] * weights, axis=1)
                y_pred[:, j] = num / denom

        if self._y.ndim == 1:
            y_pred = y_pred.ravel()

        return y_pred

class RadiusNeighborsRegressor(NeighborsBase, RadiusNeighborsMixin,
                               SupervisedFloatMixin,
                               RegressorMixin):
    """Regression based on neighbors within a fixed radius.

    The target is predicted by local interpolation of the targets
    associated of the nearest neighbors in the training set.

    Read more in the :ref:`User Guide <regression>`.

    Parameters
    ----------
    radius : float, optional (default = 1.0)
        Range of parameter space to use by default for :meth:`radius_neighbors`
        queries.

    weights : str or callable
        weight function used in prediction.  Possible values:

        - 'uniform' : uniform weights.  All points in each neighborhood
          are weighted equally.
        - 'distance' : weight points by the inverse of their distance.
          in this case, closer neighbors of a query point will have a
          greater influence than neighbors which are further away.
        - [callable] : a user-defined function which accepts an
          array of distances, and returns an array of the same shape
          containing the weights.

        Uniform weights are used by default.

    algorithm : {'auto', 'ball_tree', 'kd_tree', 'brute'}, optional
        Algorithm used to compute the nearest neighbors:

        - 'ball_tree' will use :class:`BallTree`
        - 'kd_tree' will use :class:`KDtree`
        - 'brute' will use a brute-force search.
        - 'auto' will attempt to decide the most appropriate algorithm
          based on the values passed to :meth:`fit` method.

        Note: fitting on sparse input will override the setting of
        this parameter, using brute force.

    leaf_size : int, optional (default = 30)
        Leaf size passed to BallTree or KDTree.  This can affect the
        speed of the construction and query, as well as the memory
        required to store the tree.  The optimal value depends on the
        nature of the problem.

    metric : string or DistanceMetric object (default='minkowski')
        the distance metric to use for the tree.  The default metric is
        minkowski, and with p=2 is equivalent to the standard Euclidean
        metric. See the documentation of the DistanceMetric class for a
        list of available metrics.

    p : integer, optional (default = 2)
        Power parameter for the Minkowski metric. When p = 1, this is
        equivalent to using manhattan_distance (l1), and euclidean_distance
        (l2) for p = 2. For arbitrary p, minkowski_distance (l_p) is used.

    metric_params : dict, optional (default = None)
        Additional keyword arguments for the metric function.

    Examples
    --------
    >>> X = [[0], [1], [2], [3]]
    >>> y = [0, 0, 1, 1]
    >>> from sklearn.neighbors import RadiusNeighborsRegressor
    >>> neigh = RadiusNeighborsRegressor(radius=1.0)
    >>> neigh.fit(X, y) # doctest: +ELLIPSIS
    RadiusNeighborsRegressor(...)
    >>> print(neigh.predict([[1.5]]))
    [ 0.5]

    See also
    --------
    NearestNeighbors
    KNeighborsRegressor
    KNeighborsClassifier
    RadiusNeighborsClassifier

    Notes
    -----
    See :ref:`Nearest Neighbors <neighbors>` in the online documentation
    for a discussion of the choice of ``algorithm`` and ``leaf_size``.

    https://en.wikipedia.org/wiki/K-nearest_neighbor_algorithm
    """

    def __init__(self, radius=1.0, weights='uniform',
                 algorithm='auto', leaf_size=30,
                 p=2, metric='minkowski', metric_params=None, **kwargs):
        self._init_params(radius=radius,
                          algorithm=algorithm,
                          leaf_size=leaf_size,
                          p=p, metric=metric, metric_params=metric_params,
                          **kwargs)
        self.weights = _check_weights(weights)

    def predict(self, X):
        """Predict the target for the provided data

        Parameters
        ----------
        X : array-like, shape (n_query, n_features), \
                or (n_query, n_indexed) if metric == 'precomputed'
            Test samples.

        Returns
        -------
        y : array of int, shape = [n_samples] or [n_samples, n_outputs]
            Target values
        """
        X = check_array(X, accept_sparse='csr')

        neigh_dist, neigh_ind = self.radius_neighbors(X)

        weights = _get_weights(neigh_dist, self.weights)

        _y = self._y
        if _y.ndim == 1:
            _y = _y.reshape((-1, 1))

        if weights is None:
            y_pred = np.array([np.mean(_y[ind, :], axis=0)
                               for ind in neigh_ind])
        else:
            y_pred = np.array([(np.average(_y[ind, :], axis=0,
                                           weights=weights[i]))
                               for (i, ind) in enumerate(neigh_ind)])

        if self._y.ndim == 1:
            y_pred = y_pred.ravel()

        return y_pred


try:
    import cPickle as pickle
except:
    import pickle

import os
import time
import signal
import sys

import lockfile

from HPOlibConfigSpace import configuration_space

from autosklearn.data.data_manager import DataManager
import autosklearn.models.holdout_evaluator
from autosklearn.models.paramsklearn import get_class

def store_and_or_load_data(outputdir, dataset, data_dir):
    save_path = os.path.join(outputdir, dataset + "_Manager.pkl")
    if not os.path.exists(save_path):
        lock = lockfile.LockFile(save_path)
        while not lock.i_am_locking():
            try:
                lock.acquire(timeout=60)    # wait up to 60 seconds
            except lockfile.LockTimeout:
                lock.break_lock()
                lock.acquire()
        print "I locked", lock.path
        #It is not yet sure, whether the file already exists
        try:
            if not os.path.exists(save_path):
                D = DataManager(dataset, data_dir, verbose=True)
                fh = open(save_path, 'w')
                pickle.dump(D, fh, -1)
                fh.close()
            else:
                D = pickle.load(open(save_path, 'r'))
        except:
            raise
        finally:
            lock.release()
    else:
        D = pickle.load(open(save_path, 'r'))
        print "Loaded data"
    return D

# signal handler seem to work only if they are globally defined
# to give it access to the evaluator class, the evaluator name has to
# be a global name. It's not the cleanest solution, but works for now.
evaluator = None
def signal_handler(signum, frame):
    print "Aborting Training!"
    global evaluator
    evaluator.finish_up()
    exit(0)

signal.signal(15, signal_handler)

def main(basename, input_dir, params):

    output_dir = os.getcwd()
    D = store_and_or_load_data(data_dir=input_dir, dataset=basename,
                               outputdir=output_dir)

    cs = get_class(D.info).get_hyperparameter_search_space()
    configuration = configuration_space.Configuration(cs, **params)

    global evaluator
    evaluator = autosklearn.models.holdout_evaluator.HoldoutEvaluator(
        Datamanager=D, configuration=configuration, with_predictions=True,
        all_scoring_functions=True, output_dir=output_dir)
    evaluator.fit()
    evaluator.finish_up()

if __name__ == "__main__":
    outer_starttime = time.time()

    instance_name = sys.argv[1]
    instance_specific_information = sys.argv[2]  # = 0
    cutoff_time = float(sys.argv[3])  # = inf
    cutoff_length = int(float(sys.argv[4]))  # = 2147483647
    seed = int(float(sys.argv[5]))

    params = dict()
    for i in range(6, len(sys.argv), 2):
        p_name = str(sys.argv[i])
        if p_name[0].startswith("-"):
            p_name = p_name[1:]
        params[p_name] = sys.argv[i + 1].strip()

    for key in params:
        try:
            params[key] = float(params[key])
        except:
            pass

    dataset = os.path.basename(instance_name)
    data_dir = os.path.dirname(instance_name)
    main(basename=dataset, input_dir=data_dir, params=params)

    sys.exit(0)