
class SPSModel:
    def __init__(self,files,axes,filters,redshifts):
        """
        All arguments (files,axes,filters,redshifts) are required.

        `files' is a list of filenames. Each filename points to a pickle'd
            set of SPS spectra.
        `axes' is a dictionary; the keys are each of the parameters of the
            SPS model. Each key points to a dictionary with two keys: `points'
            and `eval.' The `points' value is an array describing the point
            the axis is defined on, and `eval' is a spline model for the axis.
        `filters' is a string (or list of strings) containing the name(s) of
            the filter(s) to be evaulated.
        `redshifts' is the redshift (or list of redshifts) at which the SEDs
            should be evaluated.
        """
        from stellarpop import tools
        self.filters = {}
        if type(filters)==type([]):
            for filter in filters:
                self.filters[filter] = tools.filterfromfile(filter)
        else:
            self.filters[filter] = tools.filterfromfile(filter)
        self.filter_names = self.filters.keys()
        if type(redshifts)==type([]):
            self.redshifts = [r for r in redshifts]
        else:
            self.redshifts = [redshifts]
        self.corrections = self.luminosity_correction()

        self.axes_names = axes.keys()
        self.naxes = len(axes.keys())

        self.axes = axes
        self.files = files

        self.models = {}
        self.models = self.create_models()


    def luminosity_correction(self):
        """
        The galaxev templates are in L_sun/Angstrom, with L_sun = 3.826e26W
          so we change to the appropriate flux at each relevent redshift. We
          also multiply by 1e11, so that the flux is appropriate for a galaxy
          of mass 1e11 M_sun.
        
        Our SED code gives correct magnitudes if the SED is in units of
          ergs/s/cm/cm/Angstrom, and our distance class gives distances in
          Mpc/h, so we will need to use the conversion Mpc = 3.08568e24cm
          and assume a value for h.

        The conversion factor is conv = flux_density/luminosity_density and
          the templates should therefore be multiplied by conv to get the
          correct flux_density. The factor of (1+z) accounts for the fact that
          these are _densities_.

                    1e11   3.826e33
            conv =  ----- ----------
                    (1+z) 4 pi Dl**2
        """
        from stellarpop import distances
        from math import pi
        dist = distances.Distance()
        dist.h = 0.7
        dist.OMEGA_M = 0.3
        dist.OMEGA_L = 0.7
        cm_per_Mpc = 3.08568e24
        corr = []
        for z in self.redshifts:
            if z==0:
                dl = 1e-5*cm_per_Mpc
            else:
                dl = dist.Dl(z)*cm_per_Mpc
            conv = 1/(4*pi*dl**2)
            conv *= 3.826e33
            corr.append(conv)

        return corr


    def create_models(self):
        import scipy,cPickle
        from stellarpop import tools
        from stellarpop.ndinterp import ndInterp

        index = {}
        shape = []
        axes = {}
        axes_index = 0
        for key in self.axes_names:
            index[key] = {}
            shape.append(self.axes[key]['points'].size)
            axes[axes_index] = self.axes[key]['eval']
            axes_index += 1
            for i in range(self.axes[key]['points'].size):
                index[key][self.axes[key]['points'][i]] = i

        models = {}
        model = scipy.empty(shape)*scipy.nan
        for f in self.filter_names:
            models[f] = {}
            for z in self.redshifts:
                models[f][z] = model.copy()

        for file in self.files:
            f = open(file,'rb')
            data = cPickle.load(f)
            wave = cPickle.load(f)
            f.close()
            for key in data.keys():
                obj = data[key]
                jj = key
                spec = obj['sed']
                ind = []
                for key in self.axes_names:
                    try:
                        ind.append([index[key][obj[key]]])
                    except:
                        print key,index[key]
                        print obj
                        df
                for f in self.filter_names:
                    for i in range(len(self.redshifts)):
                        z = self.redshifts[i]
                        # correction is the units correction factor
                        correction = self.corrections[i]
                        sed = [wave,spec*correction]
                        mag = tools.ABFilterMagnitude(self.filters[f],sed,z)
                        if scipy.isnan(mag)==True:
                            df
                        models[f][z][ind] = mag

        for f in self.filter_names:
            for z in self.redshifts:
                model = models[f][z].copy()
                if scipy.isnan(model).any():
                    models[f][z] = None
                else:
                    models[f][z] = ndInterp(axes,model)
        return models

    def eval(self,points,filter,redshift):
        pnts = []
        npoints = len(points[self.axes_names[0]])
        for i in range(npoints):
            p = []
            for key in self.axes_names:
                p.append(points[key][i])
            pnts.append(p)

        return self.models[filter][redshift].eval(pnts)
