from __future__ import division, absolute_import, print_function
from past.builtins import xrange

import numpy as np
import scipy.interpolate as interpolate
import scipy.integrate as integrate
import os
import sys
from pkg_resources import resource_filename

from .modtranGenerator import ModtranGenerator
from .fgcmAtmosphereTable import FgcmAtmosphereTable

from .sharedNumpyMemManager import SharedNumpyMemManager as snmm
from .fgcmLogger import FgcmLogger

class FgcmLUTMaker(object):
    """
    Class to make a look-up table.

    parameters
    ----------
    lutConfig: dict
       Dictionary with LUT config variables
    makeSeds: bool, default=False
       Make a SED-table in the look-up table (experimental)
    """

    def __init__(self,lutConfig,makeSeds=False):
        self._checkLUTConfig(lutConfig)

        self.magConstant = 2.5/np.log(10)
        self._setThroughput = False
        self.makeSeds = makeSeds

        try:
            self.stellarTemplateFile = resource_filename(__name__,'data/templates/stellar_templates_master.fits')
        except:
            raise IOError("Could not find stellar template file")

        if (not os.path.isfile(self.stellarTemplateFile)):
            raise IOError("Could not find stellar template file")

        if 'logger' in lutConfig:
            self.fgcmLog = lutConfig['logger']
        else:
            self.fgcmLog = FgcmLogger('dummy.log', 'INFO', printLogger=True)

    def _checkLUTConfig(self,lutConfig):
        """
        Internal method to check the lutConfig dictionary

        parameters
        ----------
        lutConfig: dict
        """

        self.runModtran = False

        requiredKeys = ['filterNames', 'stdFilterNames', 'nCCD']

        # first: check if there is a tableName here!
        if 'atmosphereTableName' in lutConfig:
            # Can we find this table?

            # load parameters from it and stuff into config dict
            self.atmosphereTable = FgcmAtmosphereTable.initWithTableName(lutConfig['atmosphereTableName'])

            # look for consistency between configs?
            # Note that we're assuming if somebody asked for a table they wanted
            # what was in the table
            for key in self.atmosphereTable.atmConfig:
                if key in lutConfig:
                    if 'Range' in key:
                        if (not np.isclose(lutConfig[key][0], self.atmosphereTable.atmConfig[key][0]) or
                            not np.isclose(lutConfig[key][1], self.atmosphereTable.atmConfig[key][1])):
                            print("Warning: input config %s is %.5f-%.5f but precomputed table is %.5f-%.5f" %
                                  (key, lutConfig[key][0], lutConfig[key][1],
                                   self.atmosphereTable.atmConfig[key][0],
                                   self.atmosphereTable.atmConfig[key][1]))
                    else:
                        if not np.isclose(lutConfig[key], self.atmosphereTable.atmConfig[key]):
                            print("Warning: input config %s is %.5f but precomputed table is %.5f" %
                                  (key, lutConfig[key], self.atmosphereTable.atmConfig[key]))

        else:
            # regular config with parameters
            self.runModtran = True
            self.atmosphereTable = FgcmAtmosphereTable(lutConfig)

        for key in requiredKeys:
            if (key not in lutConfig):
                raise ValueError("required %s not in lutConfig" % (key))
            if ('Range' in key):
                if (len(lutConfig[key]) != 2):
                    raise ValueError("%s must have 2 elements" % (key))

        self.lutConfig = lutConfig

        self.filterNames = self.lutConfig['filterNames']
        self.stdFilterNames = self.lutConfig['stdFilterNames']

        if len(self.filterNames) != len(self.stdFilterNames):
            raise ValueError("Length of filterNames must be same as stdFilterNames")

        for stdFilterName in self.stdFilterNames:
            if stdFilterName not in self.filterNames:
                raise ValueError("stdFilterName %s not in list of filterNames" % (stdFilterName))

        self.nCCD = self.lutConfig['nCCD']
        self.nCCDStep = self.nCCD+1

        # and record the standard values out of the config
        #  (these will also come out of the save file)

        self.pmbStd = self.atmosphereTable.atmConfig['pmbStd']
        self.pwvStd = self.atmosphereTable.atmConfig['pwvStd']
        self.o3Std = self.atmosphereTable.atmConfig['o3Std']
        self.tauStd = self.atmosphereTable.atmConfig['tauStd']
        self.lnTauStd = np.log(self.tauStd)
        self.alphaStd = self.atmosphereTable.atmConfig['alphaStd']
        self.secZenithStd = self.atmosphereTable.atmConfig['airmassStd']
        self.zenithStd = np.arccos(1./self.secZenithStd)*180./np.pi

        if ('lambdaRange' in self.atmosphereTable.atmConfig):
            self.lambdaRange = np.array(self.atmosphereTable.atmConfig['lambdaRange'])
        else:
            self.lambdaRange = np.array([3000.0,11000.0])

        if ('lambdaStep' in self.atmosphereTable.atmConfig):
            self.lambdaStep = self.atmosphereTable.atmConfig['lambdaStep']
        else:
            self.lambdaStep = 0.5

        if ('lambdaNorm' in self.atmosphereTable.atmConfig):
            self.lambdaNorm = self.atmosphereTable.atmConfig['lambdaNorm']
        else:
            self.lambdaNorm = 7750.0

    def setThroughputs(self, throughputDict):
        """
        Set the throughputs per CCD

        parameters
        ----------
        throughputDict: dict
           Dict with throughput information

        The throughput dict should have one entry for each filterName.
        Each of these elements should be a dictionary with the following keys:
           'LAMBDA': numpy float array with wavelength values
           ccd_index: numpy float array with throughput values for the ccd_index
        There should be one entry for each CCD.
        """

        self.inThroughputs=[]
        for filterName in self.filterNames:
            try:
                lam = throughputDict[filterName]['LAMBDA']
            except:
                raise ValueError("Wavelength LAMBDA not found for filter %s in throughputDict!" % (filterName))

            tput = np.zeros(lam.size, dtype=[('LAMBDA','f4'),
                                             ('THROUGHPUT_AVG','f4'),
                                             ('THROUGHPUT_CCD','f4',self.nCCD)])
            tput['LAMBDA'][:] = lam
            for ccdIndex in xrange(self.nCCD):
                try:
                    tput['THROUGHPUT_CCD'][:,ccdIndex] = throughputDict[filterName][ccdIndex]
                except:
                    raise ValueError("CCD Index %d not found for filter %s in throughputDict!" % (ccdIndex,filterName))

            # check if the average is there, if not compute it
            if ('AVG' in throughputDict[filterName]):
                tput['THROUGHPUT_AVG'][:] = throughputDict[filterName]['AVG']
            else:
                self.fgcmLog.info("Average throughput not found in throughputDict for filter %s.  Computing now..." % (filterName))
                for i in xrange(lam.size):
                    use,=np.where(tput['THROUGHPUT_CCD'][i,:] > 0.0)
                    if (use.size > 0):
                        tput['THROUGHPUT_AVG'][i] = np.mean(tput['THROUGHPUT_CCD'][i,use])

            self.inThroughputs.append(tput)

        self._setThroughput = True

    def makeLUT(self):
        """
        Make the look-up table.  This can either be saved with saveLUT or accessed via
         attributes.

        parameters
        ----------
        None

        output attributes
        -----------------
        pmb: float array
           Pressure (millibars)
        pmbFactor: float array
           Pressure factor
        pmbElevation: float
           Standard PMB at elevation
        pwv: float array
           Water vapor array
        o3: float array
           Ozone array
        tau: float array
           Aerosol optical index array
        lambdaNorm: float
           Aerosol normalization wavelength
        alpha: float array
           Aerosol slope array
        zenith: float array
           Zenith angle array
        nccd: int
           Number of CCDs in table
        pmbStd: float
           Standard PMB
        pwvStd: float
           Standard PWV
        o3Std: float
           Standard O3
        tauStd: float
           Standard tau
        alphaStd: float
           Standard alpha
        zenithStd: float
           Standard zenith angle (deg)
        lambdaRange: numpy float array
           Wavelength range (A)
        lambdaStep: float
           Wavelength step (A)
        lambdaStd: numpy float array
           Standard wavelength for each filterName (A) at each standard band
        lambdaStdRaw: numpy float array
           Standard wavelength for each filterName (A)
        I0Std: numpy float array
           Standard value for I0 for each filterName
        I1Std: numpy float array
           Standard value for I1 for each filterName
        I10Std: numpy float array
           Standard value for I10 for each filterName
        lambdaB: numpy float array
           Standard wavelength for each filterName (A) assuming no atmosphere
        atmLambda: numpy float array
           Standard atmosphere wavelength array (A)
        atmStdTrans: numpy float array
           Standard atmosphere transmission array
        lut: Look-up table recarray
           lut['I0']: I0 LUT (multi-dimensional)
           lut['I1']: I1 LUT (multi-dimensional)
        lutDeriv: Derivative recarray
           lutDeriv['D_PMB']: I0 PMB derivative
           lutDeriv['D_PWV']: I0 PWV derivative
           lutDeriv['D_O3']: I0 O3 derivative
           lutDeriv['D_LNTAU']: I0 log(tau) derivative
           lutDeriv['D_ALPHA']: I0 alpha derivative
           lutDeriv['D_SECZENITH']: I0 sec(zenith) derivative
           lutDeriv['D_PMB_I1']: I1 PMB derivative
           lutDeriv['D_PWV_I1']: I1 PWV derivative
           lutDeriv['D_O3_I1']: I1 O3 derivative
           lutDeriv['D_LNTAU_I1']: I1 log(tau) derivative
           lutDeriv['D_ALPHA_I1']: I1 alpha derivative
           lutDeriv['D_SECZENITH_I1']: I0 sec(zenith) derivative
        """

        if not self._setThroughput:
            raise ValueError("Must set the throughput before running makeLUT")

        if self.runModtran:
            # need to build the table
            self.atmosphereTable.generateTable()
        else:
            # load from data
            self.atmosphereTable.loadTable()

        # and grab from the table
        self.atmLambda = self.atmosphereTable.atmLambda
        self.atmStdTrans = self.atmosphereTable.atmStdTrans

        self.pmbElevation = self.atmosphereTable.pmbElevation

        self.pmb = self.atmosphereTable.pmb
        self.pmbDelta = self.atmosphereTable.pmbDelta
        pmbPlus = np.append(self.pmb, self.pmb[-1] + self.pmbDelta)

        self.pwv = self.atmosphereTable.pwv
        self.pwvDelta = self.atmosphereTable.pwvDelta
        pwvPlus = np.append(self.pwv, self.pwv[-1] + self.pwvDelta)

        self.o3 = self.atmosphereTable.o3
        self.o3Delta = self.atmosphereTable.o3Delta
        o3Plus = np.append(self.o3, self.o3[-1] + self.o3Delta)

        self.lnTau = self.atmosphereTable.lnTau
        self.lnTauDelta = self.atmosphereTable.lnTauDelta
        self.tau = np.exp(self.lnTau)
        lnTauPlus = np.append(self.lnTau, self.lnTau[-1] + self.lnTauDelta)
        tauPlus = np.exp(lnTauPlus)

        self.alpha = self.atmosphereTable.alpha
        self.alphaDelta = self.atmosphereTable.alphaDelta
        alphaPlus = np.append(self.alpha, self.alpha[-1] + self.alphaDelta)

        self.secZenith = self.atmosphereTable.secZenith
        self.secZenithDelta = self.atmosphereTable.secZenithDelta
        self.zenith = np.arccos(1./self.secZenith)*180./np.pi
        secZenithPlus = np.append(self.secZenith, self.secZenith[-1] + self.secZenithDelta)
        zenithPlus = np.arccos(1./secZenithPlus)*180./np.pi

        # and compute the proper airmass...
        self.airmass = self.secZenith - 0.0018167*(self.secZenith-1.0) - 0.002875*(self.secZenith-1.0)**2.0 - 0.0008083*(self.secZenith-1.0)**3.0
        airmassPlus = secZenithPlus - 0.0018167*(secZenithPlus-1.0) - 0.002875*(secZenithPlus-1.0)**2.0 - 0.0008083*(secZenithPlus-1.0)**3.0

        pwvAtmTable = self.atmosphereTable.pwvAtmTable
        o3AtmTable = self.atmosphereTable.o3AtmTable
        o2AtmTable = self.atmosphereTable.o2AtmTable
        rayleighAtmTable = self.atmosphereTable.rayleighAtmTable

        # get the filters over the same lambda ranges...
        self.fgcmLog.info("\nInterpolating filters...")
        self.throughputs = []
        for i in xrange(len(self.filterNames)):
            inLam = self.inThroughputs[i]['LAMBDA']

            tput = np.zeros(self.atmLambda.size, dtype=[('LAMBDA','f4'),
                                                        ('THROUGHPUT_AVG','f4'),
                                                        ('THROUGHPUT_CCD','f4',self.nCCD)])
            tput['LAMBDA'][:] = self.atmLambda

            for ccdIndex in xrange(self.nCCD):
                ifunc = interpolate.interp1d(inLam, self.inThroughputs[i]['THROUGHPUT_CCD'][:,ccdIndex])
                tput['THROUGHPUT_CCD'][:,ccdIndex] = np.clip(ifunc(self.atmLambda),
                                                             0.0,
                                                             1e100)
            ifunc = interpolate.interp1d(inLam, self.inThroughputs[i]['THROUGHPUT_AVG'])
            tput['THROUGHPUT_AVG'][:] = np.clip(ifunc(self.atmLambda), 0.0, 1e100)

            self.throughputs.append(tput)

        # and now we can get the standard atmosphere and lambda_b
        self.fgcmLog.info("Computing lambdaB")
        self.lambdaB = np.zeros(len(self.filterNames))
        for i in xrange(len(self.filterNames)):
            num = integrate.simps(self.atmLambda * self.throughputs[i]['THROUGHPUT_AVG'] / self.atmLambda, self.atmLambda)
            denom = integrate.simps(self.throughputs[i]['THROUGHPUT_AVG'] / self.atmLambda, self.atmLambda)
            self.lambdaB[i] = num / denom
            self.fgcmLog.info("Filter: %s, lambdaB = %.3f" % (self.filterNames[i], self.lambdaB[i]))

        self.fgcmLog.info("Computing lambdaStdFilter")
        self.lambdaStdFilter = np.zeros(len(self.filterNames))
        for i in xrange(len(self.filterNames)):
            num = integrate.simps(self.atmLambda * self.throughputs[i]['THROUGHPUT_AVG'] * self.atmStdTrans / self.atmLambda, self.atmLambda)
            denom = integrate.simps(self.throughputs[i]['THROUGHPUT_AVG'] * self.atmStdTrans / self.atmLambda, self.atmLambda)
            self.lambdaStdFilter[i] = num / denom
            self.fgcmLog.info("Filter: %s, lambdaStdFilter = %.3f" % (self.filterNames[i],self.lambdaStdFilter[i]))

        # now compute lambdaStd based on the desired standards...
        self.fgcmLog.info("Calculating lambdaStd")
        self.lambdaStd = np.zeros(len(self.filterNames))

        for i, filterName in enumerate(self.filterNames):
            ind = self.filterNames.index(self.stdFilterNames[i])
            #ind, = np.where(self.filterNames == self.stdFilterNames[i])
            #self.lambdaStd[i] = self.lambdaStdFilter[ind[0]]
            self.lambdaStd[i] = self.lambdaStdFilter[ind]
            self.fgcmLog.info("Filter: %s (from %s) lambdaStd = %.3f" %
                              (filterName, self.stdFilterNames[i], self.lambdaStd[i]))

        self.fgcmLog.info("Computing I0Std/I1Std")
        self.I0Std = np.zeros(len(self.filterNames))
        self.I1Std = np.zeros(len(self.filterNames))

        for i in xrange(len(self.filterNames)):
            self.I0Std[i] = integrate.simps(self.throughputs[i]['THROUGHPUT_AVG'] * self.atmStdTrans / self.atmLambda, self.atmLambda)
            self.I1Std[i] = integrate.simps(self.throughputs[i]['THROUGHPUT_AVG'] * self.atmStdTrans * (self.atmLambda - self.lambdaStd[i]) / self.atmLambda, self.atmLambda)

        self.I10Std = self.I1Std / self.I0Std

        #################################
        ## Make the I0/I1 LUT
        #################################

        self.fgcmLog.info("Building look-up table...")
        lutPlus = np.zeros((len(self.filterNames),
                            pwvPlus.size,
                            o3Plus.size,
                            tauPlus.size,
                            alphaPlus.size,
                            zenithPlus.size,
                            self.nCCDStep),
                           dtype=[('I0','f4'),
                                  ('I1','f4')])

        # pre-compute pmb factors
        pmbMolecularScattering = np.exp(-(pmbPlus - self.pmbElevation)/self.pmbElevation)
        pmbMolecularAbsorption = pmbMolecularScattering ** 0.6
        pmbFactorPlus = pmbMolecularScattering * pmbMolecularAbsorption
        self.pmbFactor = pmbFactorPlus[:-1]

        # this set of nexted for loops could probably be vectorized in some way
        for i in xrange(len(self.filterNames)):
            self.fgcmLog.info("Working on filter %s" % (self.filterNames[i]))
            for j in xrange(pwvPlus.size):
                self.fgcmLog.info("  and on pwv #%d" % (j))
                for k in xrange(o3Plus.size):
                    self.fgcmLog.info("   and on o3 #%d" % (k))
                    for m in xrange(tauPlus.size):
                        for n in xrange(alphaPlus.size):
                            for o in xrange(zenithPlus.size):
                                aerosolTauLambda = np.exp(-1.0*tauPlus[m]*airmassPlus[o]*(self.atmLambda/self.lambdaNorm)**(-alphaPlus[n]))
                                for p in xrange(self.nCCDStep):
                                    if (p == self.nCCD):
                                        Sb = self.throughputs[i]['THROUGHPUT_AVG'] * o2AtmTable[o,:] * rayleighAtmTable[o,:] * pwvAtmTable[j,o,:] * o3AtmTable[k,o,:] * aerosolTauLambda
                                    else:
                                        Sb = self.throughputs[i]['THROUGHPUT_CCD'][:,p] * o2AtmTable[o,:] * rayleighAtmTable[o,:] * pwvAtmTable[j,o,:] * o3AtmTable[k,o,:] * aerosolTauLambda

                                    lutPlus['I0'][i,j,k,m,n,o,p] = integrate.simps(Sb / self.atmLambda, self.atmLambda)
                                    lutPlus['I1'][i,j,k,m,n,o,p] = integrate.simps(Sb * (self.atmLambda - self.lambdaStd[i]) / self.atmLambda, self.atmLambda)

        # and create the LUT (not plus)
        self.lut = np.zeros((len(self.filterNames),
                             self.pwv.size,
                             self.o3.size,
                             self.tau.size,
                             self.alpha.size,
                             self.zenith.size,
                             self.nCCDStep),
                            dtype=lutPlus.dtype)

        temp = np.delete(lutPlus['I0'], self.pwv.size, axis=1)
        temp = np.delete(temp, self.o3.size, axis=2)
        temp = np.delete(temp, self.tau.size, axis=3)
        temp = np.delete(temp, self.alpha.size, axis=4)
        temp = np.delete(temp, self.zenith.size, axis=5)

        self.lut['I0'] = temp

        temp = np.delete(lutPlus['I1'], self.pwv.size, axis=1)
        temp = np.delete(temp, self.o3.size, axis=2)
        temp = np.delete(temp, self.tau.size, axis=3)
        temp = np.delete(temp, self.alpha.size, axis=4)
        temp = np.delete(temp, self.zenith.size, axis=5)

        self.lut['I1'] = temp

        #################################
        ## Make the I0/I1 derivative LUTs
        #################################

        # This is *not* done plus-size

        self.lutDeriv = np.zeros((len(self.filterNames),
                                  self.pwv.size,
                                  self.o3.size,
                                  self.tau.size,
                                  self.alpha.size,
                                  self.zenith.size,
                                  self.nCCDStep),
                                 dtype=[('D_PMB','f4'),
                                        ('D_PWV','f4'),
                                        ('D_O3','f4'),
                                        ('D_LNTAU','f4'),
                                        ('D_ALPHA','f4'),
                                        ('D_SECZENITH','f4'),
                                        ('D_PMB_I1','f4'),
                                        ('D_PWV_I1','f4'),
                                        ('D_O3_I1','f4'),
                                        ('D_LNTAU_I1','f4'),
                                        ('D_ALPHA_I1','f4'),
                                        ('D_SECZENITH_I1','f4')])

        self.fgcmLog.info("Computing derivatives...")

        ## FIXME: figure out PMB derivative?

        for i in xrange(len(self.filterNames)):
            self.fgcmLog.info("Working on filter %s" % (self.filterNames[i]))
            for j in xrange(self.pwv.size):
                for k in xrange(self.o3.size):
                    for m in xrange(self.tau.size):
                        for n in xrange(self.alpha.size):
                            for o in xrange(self.zenith.size):
                                for p in xrange(self.nCCDStep):
                                    self.lutDeriv['D_PWV'][i,j,k,m,n,o,p] = (
                                        ((lutPlus['I0'][i,j+1,k,m,n,o,p] -
                                          lutPlus['I0'][i,j,k,m,n,o,p]) /
                                         self.pwvDelta)
                                        )
                                    self.lutDeriv['D_PWV_I1'][i,j,k,m,n,o,p] = (
                                        ((lutPlus['I1'][i,j+1,k,m,n,o,p] -
                                          lutPlus['I1'][i,j,k,m,n,o,p]) /
                                         self.pwvDelta)
                                        )
                                    self.lutDeriv['D_O3'][i,j,k,m,n,o,p] = (
                                        ((lutPlus['I0'][i,j,k+1,m,n,o,p] -
                                          lutPlus['I0'][i,j,k,m,n,o,p]) /
                                         self.o3Delta)
                                        )
                                    self.lutDeriv['D_O3_I1'][i,j,k,m,n,o,p] = (
                                        ((lutPlus['I1'][i,j,k+1,m,n,o,p] -
                                          lutPlus['I1'][i,j,k,m,n,o,p]) /
                                         self.o3Delta)
                                        )
                                    self.lutDeriv['D_LNTAU'][i,j,k,m,n,o,p] = (
                                        ((lutPlus['I0'][i,j,k,m+1,n,o,p] -
                                          lutPlus['I0'][i,j,k,m,n,o,p]) /
                                         self.lnTauDelta)
                                        )
                                    self.lutDeriv['D_LNTAU_I1'][i,j,k,m,n,o,p] = (
                                        ((lutPlus['I1'][i,j,k,m+1,n,o,p] -
                                          lutPlus['I1'][i,j,k,m,n,o,p]) /
                                         self.lnTauDelta)
                                        )
                                    self.lutDeriv['D_ALPHA'][i,j,k,m,n,o,p] = (
                                        ((lutPlus['I0'][i,j,k,m,n+1,o,p] -
                                          lutPlus['I0'][i,j,k,m,n,o,p]) /
                                         self.alphaDelta)
                                        )
                                    self.lutDeriv['D_ALPHA_I1'][i,j,k,m,n,o,p] = (
                                        ((lutPlus['I1'][i,j,k,m,n+1,o,p] -
                                          lutPlus['I1'][i,j,k,m,n,o,p]) /
                                         self.alphaDelta)
                                        )
                                    self.lutDeriv['D_SECZENITH'][i,j,k,m,n,o,p] = (
                                        ((lutPlus['I0'][i,j,k,m,n,o+1,p] -
                                          lutPlus['I0'][i,j,k,m,n,o,p]) /
                                         self.secZenithDelta)
                                        )
                                    self.lutDeriv['D_SECZENITH_I1'][i,j,k,m,n,o,p] = (
                                        ((lutPlus['I1'][i,j,k,m,n,o+1,p] -
                                          lutPlus['I1'][i,j,k,m,n,o,p]) /
                                         self.secZenithDelta)
                                        )

        if (self.makeSeds):
            # and the SED LUT
            self.fgcmLog.info("Building SED LUT")

            # arbitrary.  Configure?  Fit?  Seems stable...
            delta = 600.0

            # blah on fits here...
            import fitsio

            # how many extensions?
            fits=fitsio.FITS(self.stellarTemplateFile)
            fits.update_hdu_list()
            extNames = []
            for hdu in fits.hdu_list:
                extName = hdu.get_extname()
                if ('TEMPLATE_' in extName):
                    extNames.append(extName)

            # set up SED look-up table
            nTemplates = len(extNames)

            self.sedLUT = np.zeros(nTemplates, dtype=[('TEMPLATE','i4'),
                                                      ('SYNTHMAG','f4',len(self.filterNames)),
                                                      ('FPRIME','f4',len(self.filterNames))])

            # now do it...looping is no problem since there aren't that many.

            for i in xrange(nTemplates):
                data = fits[extNames[i]].read()

                templateLambda = data['LAMBDA']
                templateFLambda = data['FLUX']
                templateFnu = templateFLambda * templateLambda * templateLambda

                parts=extNames[i].split('_')
                self.sedLUT['TEMPLATE'][i] = int(parts[1])

                # interpolate to atmLambda
                intFunc = interpolate.interp1d(templateLambda, templateFnu)
                fnu = np.zeros(self.atmLambda.size)
                good,=np.where((self.atmLambda >= templateLambda[0]) &
                               (self.atmLambda <= templateLambda[-1]))
                fnu[good] = intFunc(self.atmLambda[good])

                # out of range, let it hit the limit
                lo,=np.where(self.atmLambda < templateLambda[0])
                if (lo.size > 0):
                    fnu[lo] = intFunc(self.atmLambda[good[0]])
                hi,=np.where(self.atmLambda > templateLambda[-1])
                if (hi.size > 0):
                    fnu[hi] = intFunc(self.atmLambda[good[-1]])

                # compute synthetic mags
                for j in xrange(len(self.filterNames)):
                    num = integrate.simps(fnu * self.throughputs[j]['THROUGHPUT_AVG'][:] * self.atmStdTrans / self.atmLambda, self.atmLambda)
                    denom = integrate.simps(self.throughputs[j]['THROUGHPUT_AVG'][:] * self.atmStdTrans / self.atmLambda, self.atmLambda)

                    self.sedLUT['SYNTHMAG'][i,j] = -2.5*np.log10(num/denom)

                # and compute fprimes
                for j in xrange(len(self.filterNames)):
                    use,=np.where((templateLambda >= (self.lambdaStd[j]-delta)) &
                                  (templateLambda <= (self.lambdaStd[j]+delta)))

                    fit = np.polyfit(templateLambda[use] - self.lambdaStd[j],
                                     templateFnu[use],
                                     1)

                    self.sedLUT['FPRIME'][i,j] = fit[0] / fit[1]

            fits.close()

    def saveLUT(self,lutFile,clobber=False):
        """
        """

        import fitsio

        if (os.path.isfile(lutFile) and not clobber):
            self.fgcmLog.info("lutFile %s already exists, and clobber is False." % (lutFile))
            return

        self.fgcmLog.info("Saving LUT to %s" % (lutFile))

        # first, save the LUT itself
        fitsio.write(lutFile,self.lut.flatten(),extname='LUT',clobber=True)

        # and now save the indices
        maxFilterLen = len(max(self.filterNames, key=len))

        indexVals = np.zeros(1,dtype=[('FILTERNAMES', 'a%d' % (maxFilterLen), len(self.filterNames)),
                                      ('STDFILTERNAMES', 'a%d' % (maxFilterLen), len(self.stdFilterNames)),
                                      ('PMB','f8',self.pmb.size),
                                      ('PMBFACTOR','f8',self.pmb.size),
                                      ('PMBELEVATION','f8'),
                                      ('PWV','f8',self.pwv.size),
                                      ('O3','f8',self.o3.size),
                                      ('TAU','f8',self.tau.size),
                                      ('LAMBDANORM','f8'),
                                      ('ALPHA','f8',self.alpha.size),
                                      ('ZENITH','f8',self.zenith.size),
                                      ('NCCD','i4')])
        indexVals['FILTERNAMES'] = self.filterNames
        indexVals['STDFILTERNAMES'] = self.stdFilterNames
        indexVals['PMB'] = self.pmb
        indexVals['PMBFACTOR'] = self.pmbFactor
        indexVals['PMBELEVATION'] = self.pmbElevation
        indexVals['PWV'] = self.pwv
        indexVals['O3'] = self.o3
        indexVals['TAU'] = self.tau
        indexVals['LAMBDANORM'] = self.lambdaNorm
        indexVals['ALPHA'] = self.alpha
        indexVals['ZENITH'] = self.zenith
        indexVals['NCCD'] = self.nCCD

        fitsio.write(lutFile,indexVals,extname='INDEX')

        # and the standard values
        stdVals = np.zeros(1,dtype=[('PMBSTD','f8'),
                                    ('PWVSTD','f8'),
                                    ('O3STD','f8'),
                                    ('TAUSTD','f8'),
                                    ('ALPHASTD','f8'),
                                    ('ZENITHSTD','f8'),
                                    ('LAMBDARANGE','f8',2),
                                    ('LAMBDASTEP','f8'),
                                    ('LAMBDASTD','f8',len(self.filterNames)),
                                    ('LAMBDASTDFILTER','f8',len(self.filterNames)),
                                    ('LAMBDANORM','f8'),
                                    ('I0STD','f8',len(self.filterNames)),
                                    ('I1STD','f8',len(self.filterNames)),
                                    ('I10STD','f8',len(self.filterNames)),
                                    ('LAMBDAB','f8',len(self.filterNames)),
                                    ('ATMLAMBDA','f8',self.atmLambda.size),
                                    ('ATMSTDTRANS','f8',self.atmStdTrans.size)])
        stdVals['PMBSTD'] = self.pmbStd
        stdVals['PWVSTD'] = self.pwvStd
        stdVals['O3STD'] = self.o3Std
        stdVals['TAUSTD'] = self.tauStd
        stdVals['ALPHASTD'] = self.alphaStd
        stdVals['ZENITHSTD'] = self.zenithStd
        stdVals['LAMBDARANGE'] = self.lambdaRange
        stdVals['LAMBDASTEP'] = self.lambdaStep
        stdVals['LAMBDASTD'][:] = self.lambdaStd
        stdVals['LAMBDASTDFILTER'][:] = self.lambdaStdFilter
        stdVals['LAMBDANORM'][:] = self.lambdaNorm
        stdVals['I0STD'][:] = self.I0Std
        stdVals['I1STD'][:] = self.I1Std
        stdVals['I10STD'][:] = self.I10Std
        stdVals['LAMBDAB'][:] = self.lambdaB
        stdVals['ATMLAMBDA'][:] = self.atmLambda
        stdVals['ATMSTDTRANS'][:] = self.atmStdTrans

        fitsio.write(lutFile,stdVals,extname='STD')

        # and the derivatives
        self.fgcmLog.info("Writing Derivative LUT")
        fitsio.write(lutFile,self.lutDeriv.flatten(),extname='DERIV')

        # and the SED LUT
        if (self.makeSeds):
            self.fgcmLog.info("Writing SED LUT")
            fitsio.write(lutFile,self.sedLUT,extname='SED')

class FgcmLUT(object):
    """
    Class to hold the main throughput look-up table and apply it.  If loading from
     a fits table, initialize with initFromFits(lutFile).

    parameters
    ----------
    indexVals: numpy recarray
       With LUT index values
    lutFlat: numpy recarray
       Flattened I0/I1 arrays
    lutDerivFlat: numpy recarray
       Flattened I0/I1 derivative arrays
    stdVals: numpy recarray
       Standard atmosphere and associated values
    sedLUT: bool, default=False
       Use SED look-up table instead of colors (experimental).
    filterToBand: dict, optional
       Dictionary to map filterNames to bands if not unique
    """

    def __init__(self, indexVals, lutFlat, lutDerivFlat, stdVals, sedLUT=None, filterToBand=None):

        #self.filterNames = indexVals['FILTERNAMES'][0]
        #self.stdFilterNames = indexVals['STDFILTERNAMES'][0]
        self.filterNames = [n.decode('utf-8') for n in indexVals['FILTERNAMES'][0]]
        self.stdFilterNames = [n.decode('utf-8') for n in indexVals['STDFILTERNAMES'][0]]
        self.pmb = indexVals['PMB'][0]
        self.pmbFactor = indexVals['PMBFACTOR'][0]
        self.pmbDelta = self.pmb[1] - self.pmb[0]
        self.pmbElevation = indexVals['PMBELEVATION'][0]
        self.lambdaNorm = indexVals['LAMBDANORM'][0]

        self.pwv = indexVals['PWV'][0]
        self.pwvDelta = self.pwv[1] - self.pwv[0]
        self.o3 = indexVals['O3'][0]
        self.o3Delta = self.o3[1] - self.o3[0]
        self.tau = indexVals['TAU'][0]
        self.lnTau = np.log(self.tau)
        self.lnTauDelta = self.lnTau[1] - self.lnTau[0]
        self.alpha = indexVals['ALPHA'][0]
        self.alphaDelta = self.alpha[1] - self.alpha[0]
        self.zenith = indexVals['ZENITH'][0]
        self.secZenith = 1./np.cos(self.zenith*np.pi/180.)
        self.secZenithDelta = self.secZenith[1] - self.secZenith[0]
        self.nCCD = indexVals['NCCD'][0]
        self.nCCDStep = self.nCCD+1

        # make shared memory arrays for LUTs
        sizeTuple = (len(self.filterNames),self.pwv.size,self.o3.size,
                     self.tau.size,self.alpha.size,self.zenith.size,self.nCCDStep)

        self.lutI0Handle = snmm.createArray(sizeTuple,dtype='f4')
        snmm.getArray(self.lutI0Handle)[:,:,:,:,:,:,:] = lutFlat['I0'].reshape(sizeTuple)

        self.lutI1Handle = snmm.createArray(sizeTuple,dtype='f4')
        snmm.getArray(self.lutI1Handle)[:,:,:,:,:,:,:] = lutFlat['I1'].reshape(sizeTuple)

        # and read in the derivatives

        # create shared memory
        self.lutDPWVHandle = snmm.createArray(sizeTuple,dtype='f4')
        self.lutDO3Handle = snmm.createArray(sizeTuple,dtype='f4')
        self.lutDLnTauHandle = snmm.createArray(sizeTuple,dtype='f4')
        self.lutDAlphaHandle = snmm.createArray(sizeTuple,dtype='f4')
        self.lutDSecZenithHandle = snmm.createArray(sizeTuple,dtype='f4')

        self.lutDPWVI1Handle = snmm.createArray(sizeTuple,dtype='f4')
        self.lutDO3I1Handle = snmm.createArray(sizeTuple,dtype='f4')
        self.lutDLnTauI1Handle = snmm.createArray(sizeTuple,dtype='f4')
        self.lutDAlphaI1Handle = snmm.createArray(sizeTuple,dtype='f4')
        self.lutDSecZenithI1Handle = snmm.createArray(sizeTuple,dtype='f4')

        snmm.getArray(self.lutDPWVHandle)[:,:,:,:,:,:,:] = lutDerivFlat['D_PWV'].reshape(sizeTuple)
        snmm.getArray(self.lutDO3Handle)[:,:,:,:,:,:,:] = lutDerivFlat['D_O3'].reshape(sizeTuple)
        snmm.getArray(self.lutDLnTauHandle)[:,:,:,:,:,:,:] = lutDerivFlat['D_LNTAU'].reshape(sizeTuple)
        snmm.getArray(self.lutDAlphaHandle)[:,:,:,:,:,:,:] = lutDerivFlat['D_ALPHA'].reshape(sizeTuple)
        snmm.getArray(self.lutDSecZenithHandle)[:,:,:,:,:,:,:] = lutDerivFlat['D_SECZENITH'].reshape(sizeTuple)

        self.hasI1Derivatives = False
        try:
            snmm.getArray(self.lutDPWVI1Handle)[:,:,:,:,:,:,:] = lutDerivFlat['D_PWV_I1'].reshape(sizeTuple)
            snmm.getArray(self.lutDO3I1Handle)[:,:,:,:,:,:,:] = lutDerivFlat['D_O3_I1'].reshape(sizeTuple)
            snmm.getArray(self.lutDLnTauI1Handle)[:,:,:,:,:,:,:] = lutDerivFlat['D_LNTAU_I1'].reshape(sizeTuple)
            snmm.getArray(self.lutDAlphaI1Handle)[:,:,:,:,:,:,:] = lutDerivFlat['D_ALPHA_I1'].reshape(sizeTuple)
            snmm.getArray(self.lutDSecZenithI1Handle)[:,:,:,:,:,:,:] = lutDerivFlat['D_SECZENITH_I1'].reshape(sizeTuple)
            self.hasI1Derivatives = True
        except:
            # just fill with zeros
            pass
            #print("No I1 derivative information")

        # get the standard values

        self.pmbStd = stdVals['PMBSTD'][0]
        self.pwvStd = stdVals['PWVSTD'][0]
        self.o3Std = stdVals['O3STD'][0]
        self.tauStd = stdVals['TAUSTD'][0]
        self.lnTauStd = np.log(self.tauStd)
        self.alphaStd = stdVals['ALPHASTD'][0]
        self.zenithStd = stdVals['ZENITHSTD'][0]
        self.secZenithStd = 1./np.cos(np.radians(self.zenithStd))
        self.lambdaRange = stdVals['LAMBDARANGE'][0]
        self.lambdaStep = stdVals['LAMBDASTEP'][0]
        self.lambdaStd = stdVals['LAMBDASTD'][0]
        self.lambdaStdFilter = stdVals['LAMBDASTDFILTER'][0]
        self.I0Std = stdVals['I0STD'][0]
        self.I1Std = stdVals['I1STD'][0]
        self.I10Std = stdVals['I10STD'][0]
        self.lambdaB = stdVals['LAMBDAB'][0]
        self.atmLambda = stdVals['ATMLAMBDA'][0]
        self.atmStdTrans = stdVals['ATMSTDTRANS'][0]

        self.magConstant = 2.5/np.log(10)

        if (filterToBand is None):
            # just set up a 1-1 mapping
            self.filterToBand = {}
            for filterName in self.filterNames:
                self.filterToBand[filterName] = filterName
        else:
            self.filterToBand = filterToBand

        # finally, read in the sedLUT
        ## this is experimental
        self.hasSedLUT = False
        if (sedLUT is not None):
            self.sedLUT = sedLUT

            self.hasSedLUT = True

            ## FIXME: make general
            self.sedColor = self.sedLUT['SYNTHMAG'][:,0] - self.sedLUT['SYNTHMAG'][:,2]
            st = np.argsort(self.sedColor)

            self.sedColor = self.sedColor[st]
            self.sedLUT = self.sedLUT[st]

    @classmethod
    def initFromFits(cls, lutFile, filterToBand=None):
        """
        Initials FgcmLUT using fits file.

        parameters
        ----------
        lutFile: string
           Name of the LUT file
        """

        import fitsio

        lutFlat = fitsio.read(lutFile, ext='LUT')
        lutDerivFlat = fitsio.read(lutFile, ext='DERIV')
        indexVals = fitsio.read(lutFile, ext='INDEX')
        stdVals = fitsio.read(lutFile, ext='STD')

        try:
            sedLUT = fitsio.read(lutFile, ext='SED')
        except:
            sedLUT = None

        return cls(indexVals, lutFlat, lutDerivFlat, stdVals,
                   sedLUT=sedLUT, filterToBand=filterToBand)

    def getIndices(self, filterIndex, pwv, o3, lnTau, alpha, secZenith, ccdIndex, pmb):
        """
        Compute indices in the look-up table.  These are in regular (non-normalized) units.

        parameters
        ----------
        filterIndex: int array
           Array with values pointing to the filterName index
        pwv: float array
        o3: float array
        lnTau: float array
        alpha: float array
        secZenith: float array
        ccdIndex: int array
           Array with values point to the ccd index
        pmb: float array
        """

        return (filterIndex,
                np.clip(((pwv - self.pwv[0])/self.pwvDelta).astype(np.int32), 0,
                        self.pwv.size-1),
                np.clip(((o3 - self.o3[0])/self.o3Delta).astype(np.int32), 0,
                        self.o3.size-1),
                np.clip(((lnTau - self.lnTau[0])/self.lnTauDelta).astype(np.int32), 0,
                        self.lnTau.size-1),
                np.clip(((alpha - self.alpha[0])/self.alphaDelta).astype(np.int32), 0,
                        self.alpha.size-1),
                np.clip(((secZenith - self.secZenith[0])/self.secZenithDelta).astype(np.int32), 0,
                        self.secZenith.size-1),
                ccdIndex,
                (np.exp(-(pmb - self.pmbElevation)/self.pmbElevation)) ** 1.6)

    def computeI0(self, pwv, o3, lnTau, alpha, secZenith, pmb, indices):
        """
        Compute I0 from the look-up table.

        parameters
        ----------
        pwv: float array
        o3: float array
        lnTau: float array
        alpha: float array
        secZenith: float array
        pmb: float array
        indices: tuple, from getIndices()
        """

        # do a simple linear interpolation
        dPWV = pwv - (self.pwv[0] + indices[1] * self.pwvDelta)
        dO3 = o3 - (self.o3[0] + indices[2] * self.o3Delta)
        dlnTau = lnTau - (self.lnTau[0] + indices[3] * self.lnTauDelta)
        dAlpha = alpha - (self.alpha[0] + indices[4] * self.alphaDelta)
        dSecZenith = secZenith - (self.secZenith[0] + indices[5] * self.secZenithDelta)

        indicesSecZenithPlus = np.array(indices[:-1])
        indicesSecZenithPlus[5] += 1
        indicesPWVPlus = np.array(indices[:-1])
        indicesPWVPlus[1] = np.clip(indicesPWVPlus[1] + 1, 0, self.pwv.size-1)

        # also include cross-terms for tau and pwv
        # and a second-derivative term for pwv
        #  note that indices[-1] is the PMB vactor

        return indices[-1]*(snmm.getArray(self.lutI0Handle)[indices[:-1]] +
                            dPWV * snmm.getArray(self.lutDPWVHandle)[indices[:-1]] +
                            dO3 * snmm.getArray(self.lutDO3Handle)[indices[:-1]] +
                            dlnTau * snmm.getArray(self.lutDLnTauHandle)[indices[:-1]] +
                            dAlpha * snmm.getArray(self.lutDAlphaHandle)[indices[:-1]] +
                            dSecZenith * snmm.getArray(self.lutDSecZenithHandle)[indices[:-1]] +
                            dlnTau * dSecZenith * (snmm.getArray(self.lutDLnTauHandle)[tuple(indicesSecZenithPlus)] -
                                                   snmm.getArray(self.lutDLnTauHandle)[indices[:-1]])/self.secZenithDelta +
                            dPWV * dSecZenith * (snmm.getArray(self.lutDPWVHandle)[tuple(indicesSecZenithPlus)] -
                                                 snmm.getArray(self.lutDPWVHandle)[indices[:-1]])/self.secZenithDelta +
                            dPWV * (dPWV - self.pwvDelta) * (snmm.getArray(self.lutDPWVHandle)[tuple(indicesPWVPlus)] -
                                                             snmm.getArray(self.lutDPWVHandle)[indices[:-1]]))

    def computeI1(self, pwv, o3, lnTau, alpha, secZenith, pmb, indices):
        """
        Compute I1 from the look-up table.

        parameters
        ----------
        pwv: float array
        o3: float array
        lnTau: float array
        alpha: float array
        secZenith: float array
        pmb: float array
        indices: tuple, from getIndices()
        """

        # do a simple linear interpolation
        dPWV = pwv - (self.pwv[0] + indices[1] * self.pwvDelta)
        dO3 = o3 - (self.o3[0] + indices[2] * self.o3Delta)
        dlnTau = lnTau - (self.lnTau[0] + indices[3] * self.lnTauDelta)
        dAlpha = alpha - (self.alpha[0] + indices[4] * self.alphaDelta)
        dSecZenith = secZenith - (self.secZenith[0] + indices[5] * self.secZenithDelta)

        indicesSecZenithPlus = np.array(indices[:-1])
        indicesSecZenithPlus[5] += 1
        indicesPWVPlus = np.array(indices[:-1])
        indicesPWVPlus[1] = np.clip(indicesPWVPlus[1] + 1, 0, self.pwv.size-1)

        # also include a cross-term for tau
        #  note that indices[-1] is the PMB vactor

        return indices[-1]*(snmm.getArray(self.lutI1Handle)[indices[:-1]] +
                            dPWV * snmm.getArray(self.lutDPWVI1Handle)[indices[:-1]] +
                            dO3 * snmm.getArray(self.lutDO3I1Handle)[indices[:-1]] +
                            dlnTau * snmm.getArray(self.lutDLnTauI1Handle)[indices[:-1]] +
                            dAlpha * snmm.getArray(self.lutDAlphaI1Handle)[indices[:-1]] +
                            dSecZenith * snmm.getArray(self.lutDSecZenithI1Handle)[indices[:-1]] +
                            dlnTau * dSecZenith * (snmm.getArray(self.lutDLnTauI1Handle)[tuple(indicesSecZenithPlus)] -
                                                   snmm.getArray(self.lutDLnTauI1Handle)[indices[:-1]])/self.secZenithDelta +
                            dPWV * dSecZenith * (snmm.getArray(self.lutDPWVI1Handle)[tuple(indicesSecZenithPlus)] -
                                                 snmm.getArray(self.lutDPWVI1Handle)[indices[:-1]])/self.secZenithDelta +
                            dPWV * (dPWV - self.pwvDelta) * (snmm.getArray(self.lutDPWVI1Handle)[tuple(indicesPWVPlus)] -
                                                             snmm.getArray(self.lutDPWVI1Handle)[indices[:-1]]))

    def computeI1Old(self, indices):
        """
        Unused
        """
        return indices[-1] * snmm.getArray(self.lutI1Handle)[indices[:-1]]

    def computeLogDerivatives(self, indices, I0):
        """
        Compute log derivatives.  Used in FgcmChisq.

        parameters
        ----------
        indices: tuple, from getIndices()
        I0: float array, from computeI0()
        """

        # dL(i,j|p) = d/dp(2.5*log10(LUT(i,j|p)))
        #           = 1.086*(LUT'(i,j|p)/LUT(i,j|p))
        return (self.magConstant*snmm.getArray(self.lutDPWVHandle)[indices[:-1]] / I0,
                self.magConstant*snmm.getArray(self.lutDO3Handle)[indices[:-1]] / I0,
                self.magConstant*snmm.getArray(self.lutDLnTauHandle)[indices[:-1]] / (I0),
                self.magConstant*snmm.getArray(self.lutDAlphaHandle)[indices[:-1]] / I0)

    def computeLogDerivativesI1(self, indices, I0, I10, sedSlope):
        """
        Compute log derivatives for I1.  Used in FgcmChisq.

        parameters
        ----------
        indices: tuple, from getIndices()
        I0: float array, from computeI0()
        I10: float array, from computeI1()/computeI0()
        sedSlope: float array, fnuprime
        """

        # dL(i,j|p) += d/dp(2.5*log10((1+F'*I10^obs) / (1+F'*I10^std)))
        #  the std part cancels...
        #            = 1.086*(F'/(1+F'*I10)*((I0*LUT1' - I1*LUT0')/(I0^2))

        preFactor = (self.magConstant * (sedSlope / (1 + sedSlope*I10))) / I0**2.

        return (preFactor * (I0 * snmm.getArray(self.lutDPWVHandle)[indices[:-1]] -
                             I10 * I0 * snmm.getArray(self.lutDPWVI1Handle)[indices[:-1]]),
                preFactor * (I0 * snmm.getArray(self.lutDO3Handle)[indices[:-1]] -
                             I10 * I0 * snmm.getArray(self.lutDO3I1Handle)[indices[:-1]]),
                preFactor * (I0 * snmm.getArray(self.lutDLnTauHandle)[indices[:-1]] -
                             I10 * I0 * snmm.getArray(self.lutDLnTauI1Handle)[indices[:-1]]),
                preFactor * (I0 * snmm.getArray(self.lutDAlphaHandle)[indices[:-1]] -
                             I10 * I0 * snmm.getArray(self.lutDAlphaI1Handle)[indices[:-1]]))

    def computeSEDSlopes(self, objectSedColor):
        """
        Compute SED slopes using the SED look-up table.  Experimental.

        parameters
        ----------
        objectSedColor: float array
           Color used for SED look-up (typically g-i)
        """

        indices = np.clip(np.searchsorted(self.sedColor, objectSedColor),0,self.sedColor.size-2)
        # right now, a straight matching to the nearest sedColor (g-i)
        #  though I worry about this.
        #  in fact, maybe the noise will make this not work?  Or this is real?
        #  but the noise in g-r is going to cause things to bounce around.  Pout.

        return self.sedLUT['FPRIME'][indices,:]

    def computeStepUnits(self, stepUnitReference, stepGrain, meanNightDuration,
                         meanWashIntervalDuration, fitBands, bands, nCampaignNights):
        """
        Compute normalization factors for fit step units.  Note that this might need
         to be tweaked.

        parameters
        ----------
        stepUnitReference: float
           How much should a typical step move things?  0.001 mag is default.
        stepGrain: float
           Additional fudge factor to apply to all steps.
        meanNightDuration: float
           Mean duration of a night (days).
        meanWashIntervalDuration: float
           Mean duration between washes (days).
        fitBands: string array
           Which bands are used for the fit?
        bands: string array
           What are all the bands?
        nCampaignNights: int
           Total number of nights in observing campaign to be calibrated.
        """

        unitDict = {}

        # bigger unit, smaller step

        # compute tau units

        deltaMagLnTau = (2.5*np.log10(np.exp(-self.secZenithStd*np.exp(self.lnTauStd))) -
                         2.5*np.log10(np.exp(-self.secZenithStd*np.exp(self.lnTauStd+1.0))))

        unitDict['lnTauUnit'] = np.abs(deltaMagLnTau) / stepUnitReference / stepGrain
        unitDict['lnTauUnit'] /= 5.0

        # FIXME?
        unitDict['lnTauSlopeUnit'] = unitDict['lnTauUnit'] * meanNightDuration

        # look for first use of 'g' or 'r' band in filterToBand...
        #  this is the reference filter for tau/alpha

        alphaFilterIndex = -1
        for i,filterName in enumerate(self.filterNames):
            if (self.filterToBand[filterName] == 'g' or
                self.filterToBand[filterName] == 'r'):
                alphaFilterIndex = i
                break

        if alphaFilterIndex == -1:
            # We don't have anything here...
            # Just set this to 1.0, since it's not sensitive?
            unitDict['alphaUnit'] = 1.0 / stepUnitReference / stepGrain
        else:
            deltaMagAlpha = (2.5*np.log10(np.exp(-self.secZenithStd*self.tauStd*(self.lambdaStd[alphaFilterIndex]/self.lambdaNorm)**self.alphaStd)) -
                             2.5*np.log10(np.exp(-self.secZenithStd*self.tauStd*(self.lambdaStd[alphaFilterIndex]/self.lambdaNorm)**(self.alphaStd+1.0))))
            unitDict['alphaUnit'] = np.abs(deltaMagAlpha) / stepUnitReference / stepGrain

            # and scale these by fraction of bands affected...
            alphaNAffectedBands = 0
            for filterName in self.filterNames:
                if ((self.filterToBand[filterName] == 'u' and
                     'u' in fitBands) or
                    (self.filterToBand[filterName] == 'g' and
                     'g' in fitBands) or
                    (self.filterToBand[filterName] == 'r' and
                     'r' in fitBands)):
                    alphaNAffectedBands += 1

            unitDict['alphaUnit'] *= float(alphaNAffectedBands) / float(len(fitBands))

        # pwv units -- reference to z or y or Y
        pwvFilterIndex = -1
        for i,filterName in enumerate(self.filterNames):
            if (self.filterToBand[filterName] == 'z' or
                self.filterToBand[filterName] == 'y' or
                self.filterToBand[filterName] == 'Y'):
                pwvFilterIndex = i
                break

        if pwvFilterIndex == -1:
            unitDict['pwvUnit'] = 1.0 / stepUnitReference / stepGrain
        else:
            indicesStd = self.getIndices(pwvFilterIndex,self.pwvStd,self.o3Std,np.log(self.tauStd),self.alphaStd,self.secZenithStd,self.nCCD,self.pmbStd)
            i0Std = self.computeI0(self.pwvStd,self.o3Std,np.log(self.tauStd),self.alphaStd,self.secZenithStd,self.pmbStd,indicesStd)
            indicesPlus = self.getIndices(pwvFilterIndex,self.pwvStd+1.0,self.o3Std,np.log(self.tauStd),self.alphaStd,self.secZenithStd,self.nCCD,self.pmbStd)
            i0Plus = self.computeI0(self.pwvStd+1.0,self.o3Std,np.log(self.tauStd),self.alphaStd,self.secZenithStd,self.pmbStd,indicesPlus)
            deltaMagPWV = 2.5*np.log10(i0Std) - 2.5*np.log10(i0Plus)

            unitDict['pwvUnit'] = np.abs(deltaMagPWV) / stepUnitReference / stepGrain

            # scale by fraction of bands affected
            pwvNAffectedBands = 0
            for filterName in self.filterNames:
                if ((self.filterToBand[filterName] == 'z' and
                     'z' in fitBands) or
                    (self.filterToBand[filterName] == 'y' and
                     'y' in fitBands) or
                    (self.filterToBand[filterName] == 'Y' and
                     'Y' in fitBands)):
                    pwvNAffectedBands += 1
            unitDict['pwvUnit'] *= float(pwvNAffectedBands) / float(len(fitBands))

        # PWV percent slope units
        unitDict['pwvPerSlopeUnit'] = unitDict['pwvUnit'] * meanNightDuration * self.pwvStd

        # PWV Global step units
        unitDict['pwvGlobalUnit'] = unitDict['pwvUnit'] * nCampaignNights

        # O3 units -- reference to r
        o3FilterIndex = -1
        for i,filterName in enumerate(self.filterNames):
            if (self.filterToBand[filterName] == 'u' or
                self.filterToBand[filterName] == 'r'):
                o3FilterIndex = i
                break

        if o3FilterIndex == -1:
            unitDict['o3Unit'] = 1.0 / stepUnitReference / stepGrain
        else:
            indicesStd = self.getIndices(o3FilterIndex,self.pwvStd,self.o3Std,np.log(self.tauStd),self.alphaStd,self.secZenithStd,self.nCCD,self.pmbStd)
            i0Std = self.computeI0(self.pwvStd,self.o3Std,np.log(self.tauStd),self.alphaStd,self.secZenithStd,self.pmbStd,indicesStd)
            indicesPlus = self.getIndices(o3FilterIndex,self.pwvStd,self.o3Std+1.0,np.log(self.tauStd),self.alphaStd,self.secZenithStd,self.nCCD,self.pmbStd)
            i0Plus = self.computeI0(self.pwvStd,self.o3Std+1.0,np.log(self.tauStd),self.alphaStd,self.secZenithStd,self.pmbStd,indicesPlus)
            deltaMagO3 = 2.5*np.log10(i0Std) - 2.5*np.log10(i0Plus)

            unitDict['o3Unit'] = np.abs(deltaMagO3) / stepUnitReference / stepGrain

            # scale by fraction of bands that are affected
            o3NAffectedBands = 0
            for filterName in self.filterNames:
                if (self.filterToBand[filterName] == 'u' or
                    self.filterToBand[filterName] == 'r'):
                    o3NAffectedBands += 1
            unitDict['o3Unit'] *= float(o3NAffectedBands) / float(len(fitBands))

        # wash parameters units...
        unitDict['qeSysUnit'] = 1.0 / stepUnitReference / stepGrain
        unitDict['qeSysSlopeUnit'] = unitDict['qeSysUnit'] * meanWashIntervalDuration

        return unitDict


import pickle
import os
import time
import numpy as np
import sys
import shutil
from openmmlib import openmmlib
from openmmlib import polymerutils
from openmmlib.polymerutils import scanBlocks
from openmmlib.openmmlib import Simulation
from openmmlib.polymerutils import grow_rw
sys.path.insert(0, "/net/levsha/share/homes/aafke/Documents/PolymerCode")
import pyximport; pyximport.install()
from smcTranslocator_switch_zloop_biasedloading import smcTranslocatorDirectional
'''Use proper smcTranslocator for biased loading'''

# -------defining parameters----------
# -- basic loop extrusion parameters--
GPU = sys.argv[1]
LIFETIME = int(sys.argv[2])   # Processivity
SEPARATION = int(sys.argv[3]) # Separation LEFs in number of monomers
TADSizes=[400,100,200,400,100,200,400,100,200,400,100,200,400,100,200,400,100,200,400,100,200,400,100,200]
'lleft and lright are defined in SMCTranslocator!'
lleft=(np.cumsum(TADSizes,dtype=np.int64)-1) #Load in with left arm active
lright=(np.cumsum(TADSizes,dtype=np.int64)[0:-1]+1) #Load with right arm active
N = sum(TADSizes) # number of monomers
smcStepsPerBlock = 1 # I take something like 1/steppingprobability, because stepping is not determistic. I usually choose the probability of stepping to be max 0.1. Usually I take 20 for genes speeding up and 10 for genes slowing down
stiff = 0                 # Polymer siffness in unit of bead size
dens = 0.2 #float(sys.argv[4]) # 0.2 # density in beads / volume. The density can roughly be estimated by looking at the amount of chromatin in a nuclear volume.
box = (N / dens) ** 0.33  # Define size of the bounding box for Periodic Boundary Conditions
data = polymerutils.grow_rw(N, int(box) - 2)  # creates a compact conformation
block = 0  # starting block
kon=0#float(sys.argv[5]) # CTCF binding rate
koff=1#float(sys.argv[6]) #CTCF unbinding rate
CTCFsites=np.array([0])#CTCF binding sites (take np.cumsum(TADsizes) if you want to compare to CTCF sites that are stall sites).
CTCFCohesinEn=0# float(sys.argv[7]) # Cohesin-CTCF interaction energy in kT, without a neighbouring cohesin, the unbinding rate of CTCF is koff, if a cohesin is right next to a CTCF site, the CTCF unbinding rate is koff*exp(-CTCFCohesinEn)

kswitch=0 #float(sys.argv[4])
CTCFPause=1
zloops=1
'this variable is controlled in smcTranslocator'

steps = int(200*(smcStepsPerBlock)) # nr of 3D simulation blocks btw advancing LEFs. For deterministic stepping choose 200-250 steps per block, otherwise, rescale with stepping probability. When genes are sparse, smcStepsPerBLock*dt is approximately the number of smc steps per smc block.
ConvGenes = True #True: Convergent genes, False: isolated genes
IncreasedBinding =1#int(sys.argv[4])
'SET dirLoading=1 in smc Translocator if you want directional loading'
# float(sys.argv[4])# Factor of increased binding at CTCF site
stg = 0.8#float(sys.argv[7]) #Probability of stalling at TAD boundary

saveEveryBlocks = int(100/(smcStepsPerBlock))   # 10 save every 10 blocks (saving every block is now too much almost)
skipSavedBlocksBeginning = int(20/(smcStepsPerBlock))  # how many blocks (saved) to skip after you restart LEF positions
totalSavedBlocks = 200#4000  # how many blocks to save (number of blocks done is totalSavedBlocks * saveEveryBlocks)
restartMilkerEveryBlocks = int(100/(smcStepsPerBlock))
#Only one Hamiltonian can be loaded at a time to the simkt, but we want to update the bonds every time a LEF steps. Loading a new Hamiltonian costs a lot of time. Instead we precalculate bonds and load all positions at once as one big Hamiltonian and just change the prefactors.

# parameters for smc bonds
smcBondWiggleDist = 0.2
smcBondDist = 0.5

if len(sys.argv)!=4:
    sys.exit('Number of input arguments is not correct')

folder = "filepath"

FullFileName=os.path.join(folder,"loopextrusion_Lifetime={0}_separation={1}_density={2}_stg={3}_N={4}_totalSavedBlocks={5}_kon=koff={6}_CTCFEn={7}_kswitch={8}_IncreasedBinding={9}_TwoLoadingSites_OneSided_NoUnstall_CTCFPause={10}_zloops={11}_NEW".format(LIFETIME, SEPARATION, dens,stg,N,totalSavedBlocks,kon,CTCFCohesinEn,kswitch,IncreasedBinding,CTCFPause,zloops))

if os.path.isdir(FullFileName):
    var=input('Folder already exists, press 9 to stop or any other key to empty and recreate the folder and continue')
    if eval(var)==9:
        sys.exit('Stopped simulation')
    else:
        shutil.rmtree(FullFileName) # Remove folder and content

# assertions for easy managing code below

assert restartMilkerEveryBlocks % saveEveryBlocks == 0
assert (skipSavedBlocksBeginning * saveEveryBlocks) % restartMilkerEveryBlocks == 0
assert (totalSavedBlocks * saveEveryBlocks) % restartMilkerEveryBlocks == 0
assert smcStepsPerBlock<6 # max number of steps per smc block should not be too large to prevent 'jerky' polymer motion

savesPerMilker = restartMilkerEveryBlocks // saveEveryBlocks
milkerInitsSkip = saveEveryBlocks * skipSavedBlocksBeginning  // restartMilkerEveryBlocks
milkerInitsTotal  = (totalSavedBlocks + skipSavedBlocksBeginning) * saveEveryBlocks // restartMilkerEveryBlocks
print("Time step = 1, Milker will be initialized {0} times, first {1} will be skipped".format(milkerInitsTotal, milkerInitsSkip))

# create filenames for Ekin, Epot and time

Ekin_fname = os.path.join(FullFileName,'Ekin.txt')
Epot_fname = os.path.join(FullFileName,'Epot.txt')
time_fname = os.path.join(FullFileName,'time.txt')
Rg_fname   = os.path.join(FullFileName,'Rg.txt')
Par_fname  = os.path.join(FullFileName,'Pars.txt')

def save_Es_ts_Rg():
    with open(time_fname, "a+") as time_file:
        time_file.write('%f\n'%(a.state.getTime()/openmmlib.ps))
    with open(Ekin_fname, "a+") as Ekin_file:
        Ekin_file.write('%f\n'%((a.state.getKineticEnergy())/a.N/a.kT))
    with open(Epot_fname, "a+") as Epot_file:
        Epot_file.write('%f\n'%((a.state.getPotentialEnergy()) /a.N /a.kT))
    with open(Rg_fname, "a+") as Rg_file:
        Rg_file.write('%f\n'%(analysis_plot_lib.rg(a.getData())))

class smcTranslocatorMilker(object):

    def __init__(self, smcTransObject):
        """
        :param smcTransObject: smc translocator object to work with
        """
        self.smcObject = smcTransObject
        self.allBonds = []

    def setParams(self, activeParamDict, inactiveParamDict):
        """
        A method to set parameters for bonds.
        It is a separate method because you may want to have a Simulation object already existing

        :param activeParamDict: a dict (argument:value) of addBond arguments for active bonds
        :param inactiveParamDict:  a dict (argument:value) of addBond arguments for inactive bonds

        """
        self.activeParamDict = activeParamDict
        self.inactiveParamDict = inactiveParamDict

    def setup(self, bondForce,  blocks = 100, smcStepsPerBlock = 1):
        """
        A method that milks smcTranslocator object
        and creates a set of unique bonds, etc.

        :param bondForce: a bondforce object (new after simulation restart!)
        :param blocks: number of blocks to precalculate
        :param smcStepsPerBlock: number of smcTranslocator steps per block
        :return:
        """

        if len(self.allBonds) != 0:
            raise ValueError("Not all bonds were used; {0} sets left".format(len(self.allBonds)))

        self.bondForce = bondForce

        #precalculating all bonds
        allBonds = []
        for dummy in range(blocks):
            self.smcObject.steps(smcStepsPerBlock)
            left, right = self.smcObject.getSMCs()
            bonds = [(int(i), int(j)) for i,j in zip(left, right)]
            allBonds.append(bonds)

        self.allBonds = allBonds
        self.uniqueBonds = list(set(sum(allBonds, [])))

        #adding forces and getting bond indices
        self.bondInds = []
        self.curBonds = allBonds.pop(0)

        for bond in self.uniqueBonds:
            paramset = self.activeParamDict if (bond in self.curBonds) else self.inactiveParamDict
            ind = bondForce.addBond(bond[0], bond[1], **paramset)
            self.bondInds.append(ind)
        self.bondToInd = {i:j for i,j in zip(self.uniqueBonds, self.bondInds)}
        return self.curBonds,[]

    def step(self, context, verbose=False):
        """
        Update the bonds to the next step.
        It sets bonds for you automatically!
        :param context:  context
        :return: (current bonds, previous step bonds); just for reference
        """
        if len(self.allBonds) == 0:
            raise ValueError("No bonds left to run; you should restart simulation and run setup  again")

        pastBonds = self.curBonds
        self.curBonds = self.allBonds.pop(0)  # getting current bonds
        bondsRemove = [i for i in pastBonds if i not in self.curBonds]
        bondsAdd = [i for i in self.curBonds if i not in pastBonds]
        bondsStay = [i for i in pastBonds if i in self.curBonds]
        if verbose:
            print("{0} bonds stay, {1} new bonds, {2} bonds removed".format(len(bondsStay),
                                                                            len(bondsAdd), len(bondsRemove)))
        bondsToChange = bondsAdd + bondsRemove
        bondsIsAdd = [True] * len(bondsAdd) + [False] * len(bondsRemove)
        for bond, isAdd in zip(bondsToChange, bondsIsAdd):
            ind = self.bondToInd[bond]
            paramset = self.activeParamDict if isAdd else self.inactiveParamDict
            self.bondForce.setBondParameters(ind, bond[0], bond[1], **paramset)  # actually updating bonds
        self.bondForce.updateParametersInContext(context)  # now run this to update things in the context
        return self.curBonds, pastBonds

def initModel():

    birthArray = np.ones(N, dtype=np.double)*0.1
    deathArray = np.zeros(N, dtype=np.double) + 1. / (LIFETIME)
    stallArray = np.zeros(N, dtype=np.double)
    pauseArray=np.ones(N, dtype=np.double)
    stallDeathArray = np.zeros(N, dtype=np.double) + 1 / (LIFETIME)
    smcNum = N // SEPARATION
    konprob=kon #CTCF binding probability per timestep
    koffprob=koff
    curPos = 0

    for size in TADSizes: # setting positions & strength of boundary elements, looping over each TAD
        curPos += size
        if curPos < len(stallArray):
            stallArray[curPos] = stg
            birthArray[curPos+1] = birthArray[curPos]*IncreasedBinding
            birthArray[curPos-1] = birthArray[curPos]*IncreasedBinding

    SMCTran = smcTranslocatorDirectional(birthArray, deathArray, stallArray,pauseArray, stallDeathArray,smcNum,CTCFsites,konprob,koffprob,switch=kswitch)
    return SMCTran

SMCTran = initModel()  # defining actual smc translocator object

# now polymer simulation code starts

# ------------feed smcTran to the milker---
SMCTran.steps(1000000)#(1000000)  # first steps to "equilibrate" SMC dynamics. If desired of course.
milker = smcTranslocatorMilker(SMCTran)   # now feed this thing to milker (do it once!)
#--------- end new code ------------

for milkerCount in range(milkerInitsTotal):
    doSave = milkerCount >= milkerInitsSkip

    # simulation parameters are defined below
    a = Simulation(timestep=80,thermostat=0.01)#, thermostat=0.01)#Collision rate in inverse picoseconds, low collistion rate means ballistic like motion, default in openmmpolymer is 0.001. Motion polymer is not diffusive, this is ok for statistical average,
    #but not for dynamics of the polymer
    a.setup(platform="CUDA", PBC=True, PBCbox=[box, box, box], GPU=GPU, precision="mixed")  # set up GPU here, PBC=Periodic Boundary Conditions
    a.saveFolder(FullFileName)
    a.load(data)
    a.addHarmonicPolymerBonds(wiggleDist=0.1) # WiggleDist controls distance at which energy of bond equals kT
    if stiff > 0:
        a.addGrosbergStiffness(stiff) # Chain stiffness is introduced by an angular potential U(theta)=stiff(1-cos(theta-Pi))
    a.addPolynomialRepulsiveForce(trunc=1.5, radiusMult=1.05) #Polynomial repulsive potential between particles. Has value trunc=3.0 at zero, stays flat until 0.6-0.7 and then drops to zero. For attraction between a selective set of particles, use LeonardJones or addSelectiveSSWForce (see blocks.py or ask Johannes)
    a.step = block

    # ------------ initializing milker; adding bonds ---------
    # copied from addBond
    kbond = a.kbondScalingFactor / (smcBondWiggleDist ** 2)
    bondDist = smcBondDist * a.length_scale

    activeParams = {"length":bondDist,"k":kbond}
    inactiveParams = {"length":bondDist, "k":0}
    milker.setParams(activeParams, inactiveParams)

    # this step actually puts all bonds in and sets first bonds to be what they should be
    milker.setup(bondForce=a.forceDict["HarmonicBondForce"],
                blocks=restartMilkerEveryBlocks,   # default value; milk for 100 blocks
                 smcStepsPerBlock=smcStepsPerBlock)  #
    print("Restarting milker")

    a.doBlock(steps=steps, increment=False)  # do block for the first time with first set of bonds in
    for i in range(restartMilkerEveryBlocks - 1):
        curBonds, pastBonds = milker.step(a.context)  # this updates bonds. You can do something with bonds here
        if i % saveEveryBlocks == (saveEveryBlocks - 2):
            a.doBlock(steps=steps, increment = doSave)
            if doSave:
                a.save()
                pickle.dump(curBonds, open(os.path.join(a.folder, "SMC{0}.dat".format(a.step)),'wb'))
                save_Es_ts_Rg() # save energies and time
        else:
            a.integrator.step(steps)  # do steps without getting the positions from the GPU (faster)

    data = a.getData()  # save data and step, and delete the simulation
    block = a.step
    del a

    time.sleep(0.2)  # wait 200ms for sanity (to let garbage collector do its magic)

with open(Par_fname,"a+") as Parfile:
    Parfile.write(" tau="+str(LIFETIME)+"\n Separation="+str(SEPARATION)+"\n N="+str(N)+"\n smcStepsPerBlock="+str(smcStepsPerBlock)+"\n stiff="+str(stiff)+"\n dens="+str(dens)+"\n block="+str(block)+"\n SaveEveryBlocks="+str(saveEveryBlocks)+"\n skipSavedBlocksBeginning="+str(skipSavedBlocksBeginning)+"\n totalSavedBlocks="+str(totalSavedBlocks)+"\n restartMilkerEveryBlocks="+str(restartMilkerEveryBlocks)+"\n smcBondWiggleDist="+str(smcBondWiggleDist)+"\n smcBondDist="+str(smcBondDist)+"\n SmcTimestep=" + str(dt)+"\n Mean TADsize"+str(np.mean(TADSizes)))


#!/usr/bin/env python3
# -*- coding: utf-8 -*-

''' Class to keep track of the seeds,
    i.e. all the anions, electrons and avalanches.

    This class is used to modify and get the properties of the seeds
    as well as classify, move and multiply them.
'''

# General imports
import numpy as np
import logging

# Import from project files
from .streamer_head import SHList

# settings
logger = logging.getLogger(__name__)
logger.addHandler(logging.NullHandler())

# idea: refactor to a general class for particles
#       create a new class to manage the particles

class Seeds(object):
    def __init__(self, mu_e, mu_ion, efield_ion, efield_crit, time_step,
                 Q_crit, efield_dtype, micro_step_no
                 ):
        # initiate the seeds object without adding any seeds

        # define the logger
        self.logger = logging.getLogger(__name__ + '.Seeds')

        # set the dtype for electric field calculation
        if efield_dtype in ['sp', 'float32',]:
            self.efield_dtype = np.float32
        elif efield_dtype in ['dp', 'float64',]:
            self.efield_dtype = np.float64
        else:
            self.efield_dtype = np.float32
            msg = 'Can not use dtype "{}". Using "{}"" instead.'
            self.logger.warning(msg.format(efield_dtype, self.efield_dtype))

        # set the properties of the seeds
        self.mu_e         = mu_e        # electron mobility
        self.mu_ion       = mu_ion      # cation mobility
        self.efield_ion   = efield_ion  # electric field for ion detachment
        self.efield_crit  = efield_crit  # electric field for multiplication
        self.time_step    = time_step   # simulation max time step
        self.dt           = time_step   # simulation actual time step
        self.micro_step_no = micro_step_no   # micro time steps
        self.Q_crit       = Q_crit      # avalanche to streamer criterion

        # initialize the seed variables, start with no seeds
        self.pos = np.array([]).reshape(3, 0)   # [x, y, z] position
        self.Q = np.array([])                   # avalanche size
        self.clear()                            # initialize other variables

        # log initiation
        self.logger.debug('Initated Seeds')
        self.logger.log(5, 'Seeds.__dict__')
        for k, v in self.__dict__.items():
            self.logger.log(5, '  "{}": {}'.format(k, v))

    def clear(self):
        # Clear variables

        # rename for better readability below
        no = self.no
        nz = np.zeros  # return zeroes
        nf = lambda shape: np.zeros(shape, dtype=bool)  # return false shape

        # Should be explicitly reset
        self.no_added       = 0
        self.no_removed     = 0
        self.pos_to_append  = nz((3, 0))
        self.is_to_remove   = nz(no, dtype=bool)

        # Should be updated each iteration
        self.ds             = nz((3, no))   # last movement
        self.ds_abs         = nz(no)        # movement, absolute
        self.ds_max         = 0             # movement, maximum
        self.dQ             = nz(no)        # change in charge
        self.mu             = nz(no)        # mobility
        self.e_dir          = nz((3, no))   # electric field, unit vector
        self.e_vec          = nz((3, no))   # electric field, vector
        self.e_str          = nz(no)        # electric field, strength
        # bools
        self.is_in_streamer = nf(no)        # inside streamer
        self.is_critical    = nf(no)        # avalanche termination
        self.is_behind_roi  = nf(no)        # behind ROI
        self.is_avalanche   = nf(no)        # avalanche seed
        self.is_ion         = nf(no)        # electron seed
        self.is_electron    = nf(no)        # ion seed

    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    #       DERIVED ATTIBUTES                                         #
    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    # The user need to ensure that relevant variables are updated
    # before getting these properties

    Q_max = property(lambda self: self.Q.max())
    dQ_max = property(lambda self: self.dQ.max())
    charge = property(lambda self: np.exp(self.Q).sum())
    charge_gen = property(lambda self: (self.dQ * np.exp(self.Q)).sum())
    charge_rem = property(lambda self: np.exp(self.Q_to_remove).sum())

    no = property(lambda self: self.pos.shape[1])  # total number
    no_to_append = property(lambda self: self.pos_to_append.shape[1])

    no_ion = property(lambda self: self.is_ion.sum())
    no_electron = property(lambda self: self.is_electron.sum())
    no_avalanche = property(lambda self: self.is_avalanche.sum())
    no_critical = property(lambda self: self.is_critical.sum())
    no_behind_roi = property(lambda self: self.is_behind_roi.sum())
    no_in_streamer = property(lambda self: self.is_in_streamer.sum())
    no_to_remove = property(lambda self: self.is_to_remove.sum())

    pos_ion = property(lambda self: self.pos[:, self.is_ion])
    pos_electron = property(lambda self: self.pos[:, self.is_electron])
    pos_avalanche = property(lambda self: self.pos[:, self.is_avalanche])
    pos_critical = property(lambda self: self.pos[:, self.is_critical])
    pos_behind_roi = property(lambda self: self.pos[:, self.is_behind_roi])
    pos_in_streamer = property(lambda self: self.pos[:, self.is_in_streamer])
    pos_to_remove = property(lambda self: self.pos[:, self.is_to_remove])

    Q_ion = property(lambda self: self.Q[self.is_ion])
    Q_electron = property(lambda self: self.Q[self.is_electron])
    Q_avalanche = property(lambda self: self.Q[self.is_avalanche])
    Q_critical = property(lambda self: self.Q[self.is_critical])
    Q_behind_roi = property(lambda self: self.Q[self.is_behind_roi])
    Q_in_streamer = property(lambda self: self.Q[self.is_in_streamer])
    Q_to_remove = property(lambda self: self.Q[self.is_to_remove])

    dQ_ion = property(lambda self: self.dQ[self.is_ion])
    dQ_electron = property(lambda self: self.dQ[self.is_electron])
    dQ_avalanche = property(lambda self: self.dQ[self.is_avalanche])
    dQ_critical = property(lambda self: self.dQ[self.is_critical])
    dQ_behind_roi = property(lambda self: self.dQ[self.is_behind_roi])
    dQ_in_streamer = property(lambda self: self.dQ[self.is_in_streamer])
    dQ_to_remove = property(lambda self: self.dQ[self.is_to_remove])

    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    #       UPDATE                                                    #
    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    # Updating of bools should be called explicitly
    # since they may be expensive to calculate
    # or it has to be done at specific times
    # Maintain a comment to when they are updated!

    # `is_ion` is set during movement
    def update_is_ion(self, e_str=None):
        e_str = self.e_str if (e_str is None) else e_str
        self.is_ion = (e_str < self.efield_ion)
        return self.is_ion

    # ´is_electron´ is set during movement
    def update_is_electron(self, e_str=None):
        e_str = self.e_str if (e_str is None) else e_str
        self.is_electron = (e_str >= self.efield_ion)
        return self.is_electron

    # ´is_avalanche´ is set during movement
    def update_is_avalanche(self, e_str=None):
        e_str = self.e_str if (e_str is None) else e_str
        self.is_avalanche = (e_str >= self.efield_crit)
        return self.is_avalanche

    # Invoked in main simulation loop
    def update_is_critical(self, Q=None):
        Q = self.Q if (Q is None) else Q
        self.is_critical = (Q >= self.Q_crit)
        return self.is_critical

    # Invoked in main simulation loop
    def update_is_behind_roi(self, roi):
        self.is_behind_roi = roi.is_pos_behind(self.pos)
        return self.is_behind_roi

    # ´is_in_streamer´ is set during movement
    def update_is_in_streamer(self, heads):
        self.is_in_streamer = SHList(heads).is_inside(self.pos)
        return self.is_in_streamer

    def update_mu(self):
        self.mu = np.ones(self.no) * self.mu_ion
        self.mu[self.is_electron] = self.mu_e

    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    #       OTHER                                                     #
    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #

    def eprop(self, heads, pos=None):
        ''' Return `edir`, `estr`, `evec`, and `is_inside`.
            Convenience function to ensure correct dtype for seeds.
        '''
        if pos is None:
            pos = self.pos
        return SHList(heads).eprop(pos, dtype=self.efield_dtype)

    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    #       MOVE & MULTIPLY                                           #
    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #

    def _update_ds(self, e_str, e_vec):
        ''' Calculate movement. Update `ds`.
            Also update `is_electron`, `is_ion`, `mu`.
        '''
        self.is_electron = (e_str > self.efield_ion)
        self.is_ion = ~self.is_electron  # note: need to be somewhere
        self.mu = self.mu_ion * np.ones(self.no)
        self.mu[self.is_electron] = self.mu_e
        self.ds = e_vec * self.mu * self.dt

    def _update_ds_abs(self, thresh):
        ''' Calculate the max movement. Update `ds_abs` and `ds_max`.
            Give warning if any movement exceeds `thresh`.
        '''
        self.ds_abs = np.linalg.norm(self.ds, axis=0)
        self.ds_max = self.ds_abs.max()
        if self.ds_max > thresh:
            mask = (self.ds_abs > thresh)
            msg = 'Warning! {}x large movement!'.format(mask.sum())
            self.logger.warning(msg)
            self.logger.debug('Movements {}'.format(self.ds_abs[mask]))

    # note: previously used function
    def move_multiply_single_step(self, heads, calc_alpha):
        ''' Move all seeds and multiply all avalanches, single step. '''

        # treat seeds below the plane as if they were at the plane
        pos = self.pos.copy()
        mask = (self.pos[2, :] < 0)
        pos[2, mask] = 0

        # Find the field at electrons!
        self.e_dir, self.e_str, self.e_vec, self.is_in_streamer = (
            self.eprop(heads, pos))

        # Move seeds
        self._update_ds(self.e_str, self.e_vec)
        self._update_ds_abs(thresh=heads[0].rp)
        self.pos += self.ds

        # Update electron numbers
        self.dQ = calc_alpha(self.e_str) * self.ds_abs
        self.dQ[self.e_str < self.efield_crit] = 0
        self.dQ[self.pos[2, :] <= 0] = 0
        self.Q += self.dQ

    def move_multiply(self, heads, calc_alpha):
        ''' Move all seeds, multiply electron in avalanches.

        Parameters
        ----------
        heads :     SHList
                    list of streamer heads used for field calculation
        calc_alpha : CalcAlpha
                     class instance for calculating electron multiplication

        Notes
        -----
        Seeds are classified by the electric field strength.
        Avalanches are moved in an inner loop.
        Critical avalanches and collisions with the streamer
        are saved and these break the inner loop.
        Other seeds are moved according to the time spent in the inner loop.
        Long seed movements give warnings.
        Positions and charge numbers are updated.
        '''

        # treat seeds below the plane as if they were at the plane
        # this implies that seeds below plane will move straight upwards
        pos = self.pos.copy()
        mask = (self.pos[2, :] < 0)
        pos[2, mask] = 0

        # find the field
        self.e_dir, self.e_str, self.e_vec, self.is_in_streamer = (
            self.eprop(heads, pos))

        # classify the seeds
        self.update_is_electron()   # uses self.e_str
        self.update_is_ion()        # uses self.e_str
        self.update_is_avalanche()  # uses self.e_str

        # set all mobilities
        self.update_mu()            # uses self.is_electron

        # extract information about just the avalanches
        is_avalanche = self.is_avalanche
        pos_avalanche = self.pos[:, is_avalanche]
        Q_avalanche = self.Q[is_avalanche]
        e_str = self.e_str[is_avalanche]
        e_vec = self.e_vec[:, is_avalanche]
        e_dir = self.e_dir[:, is_avalanche]
        is_in_streamer = self.is_in_streamer

        time_spent = 0

        def _get_ds(e_vec, mu=None, time_step=None):
            # return the change in avalanche positions
            if mu is None:
                mu = self.mu
            if time_step is None:
                time_step = self.time_step
            ds = e_vec * mu * time_step
            ds_abs = np.linalg.norm(ds, axis=0)
            return ds, ds_abs

        def _warn_large_ds(ds_abs, threshold, what):
            ds_max = ds_abs.max()
            if (ds_max < threshold):
                return ds_max
            mask = (ds_abs > threshold)
            msg = f'Warning! {mask.sum()}x large {what} movement!'
            self.logger.warning(msg)
            self.logger.info(f'Max movement: {(ds_max * 1e6):.2g} um')
            self.logger.log(5, f'Movements {ds_abs[mask]}')
            return ds_max

        # skip this part when it is not needed
        if any(is_avalanche):

            # do one iteration with the field already calculated
            time_spent += self.time_step        # update the time
            ds, ds_abs = _get_ds(e_vec, mu=self.mu_e) # find the movements
            _warn_large_ds(ds_abs, threshold=heads[0].rp/10, what='avalanches')
            pos_avalanche += ds                 # update positions
            dQ = calc_alpha(e_str) * ds_abs     # get avalanche change
            dQ[pos_avalanche[2, :] < 0] = 0     # no multiplication for z<0
            Q_avalanche += dQ                   # update avalanche electrons

            # do the rest of the iterations
            for i in range(self.micro_step_no - 1):  # one is already done
                # end when critical avalanche is reached
                if any(Q_avalanche > self.Q_crit):
                    logger.debug('Critical avalanche!')
                    break
                # end on collision with streamer
                if any(is_in_streamer):
                    if i != 0:  # not needed (or allowed) the first loop
                        self.is_in_streamer[is_avalanche] = is_in_streamer
                    logger.debug('Collision with streamer!')
                    break

                # treat seeds below the plane as if they were at the plane
                pos_avalanche_tmp = pos_avalanche.copy()
                mask = (pos_avalanche[2, :] < 0)
                pos_avalanche_tmp[2, mask] = 0

                # find the field
                e_dir, e_str, e_vec, is_in_streamer = (
                    self.eprop(heads, pos_avalanche_tmp))

                time_spent += self.time_step        # update the time
                ds, ds_abs = _get_ds(e_vec, mu=self.mu_e) # find the movements
                _warn_large_ds(
                    ds_abs, threshold=heads[0].rp/10, what='avalanches')
                pos_avalanche += ds                 # update positions
                dQ = calc_alpha(e_str) * ds_abs     # get avalanche change
                dQ[pos_avalanche[2, :] < 0] = 0     # no multiplication for z<0
                Q_avalanche += dQ                   # update avalanche electrons

        else:  # no avalanches, should be a rare event
            time_spent = self.time_step * self.micro_step_no
            logger.debug('No avalanches!')

        # update avalanche numbers
        self.is_avalanche = is_avalanche  # should already be correct
        self.dQ[is_avalanche] = Q_avalanche - self.Q[is_avalanche]
        self.Q[is_avalanche] = Q_avalanche

        # get seed movements
        self.e_vec[:, is_avalanche] = 0  # do not move avalanches here
        ds, ds_abs = _get_ds(self.e_vec, time_step=time_spent) # find movements
        _warn_large_ds(ds_abs, threshold=heads[0].rp, what='seed')

        # update positions
        pos_old = self.pos.copy()   # save old position for ds calculation
        self.pos += ds              # add ds for all seeds
        self.pos[:, is_avalanche] = pos_avalanche  # set these explicitly

        # store actual movements
        self.ds = self.pos - pos_old
        self.ds_abs = np.linalg.norm(self.ds, axis=0)
        self.ds_max = self.ds_abs.max()

        # store the time spent here
        self.dt = time_spent

        # log results, as required
        logger.log(5, f'Number of avalanches, {is_avalanche.sum()}')
        logger.log(5, f'Time spent, {(time_spent * 1e9):.3f} ns')
        logger.log(5, f'Maximum movement, {(self.ds_max * 1e6):.3f} um')
        if any(is_avalanche):
            msg = f'Maximum avalanche growth, {(self.dQ.max()):.3f}'
        else:
            msg = f'Maximum avalanche growth, 0'
        logger.log(5, msg)

    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    #       APPEND/REMOVE                                             #
    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #

    def append(self, pos):
        ''' Append the given pos, and update vars as required.'''
        if not pos.size:
            self.logger.log(5, 'Tried to add zero seeds.')
            return
        assert pos.shape[0] == 3, 'pos.shape[0] == 3'
        Q = np.zeros(pos.shape[1])
        self.pos = np.hstack((self.pos, pos))
        self.Q = np.hstack((self.Q, Q))
        self.no_added += pos.shape[1]

    def remove(self, idx):
        ''' Remove pos and Q at the given indices.'''
        if idx.sum() == 0:
            self.logger.log(5, 'Tried to remove zero seeds.')
            return
        self.pos = self.pos[:, idx == 0]
        self.Q = self.Q[idx == 0]
        self.no_removed += idx.sum()

    def append_at_end(self, pos):
        ''' Appends the given pos to pos_to_append.'''
        assert pos.shape[0] == 3, 'pos.shape[0] == 3'
        pos = np.array(pos).reshape(3, -1)
        self.pos_to_append = np.hstack((self.pos_to_append, pos))

    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    #       CLEAN                                                     #
    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #

    def clean(self):

        # Remove seeds and reset var
        self.remove(self.is_to_remove)

        # Append seeds and reset var
        self.append(self.pos_to_append)

        # Clear other variables
        self.clear()

#