#!/net/python/bin/python
from Function2D import *
from Spectrum import *
import os

class Spectrum2D(Function2D):
    #note: redshift formulae assume flux is in units of
    #erg/cm^2/s/A or equivalent (i.e., not per unit frequency)
    #  this is the case with the Vega SED provided by SALT
    #  (vega_bohlin2004.dat)

    #self.x is the time coordinate
    #self.y is the wavelength coordinate
    #self.z is the flux value, assumed to be in erg/cm^2/A
    def redshift(self,z):
        selfcopy = self.copy()
        
        selfcopy.x *= (1.+z)
        selfcopy.y *= (1.+z)
        selfcopy.z /= (1.+z)**4

        return selfcopy

    def blueshift(self,z):
        selfcopy = self.copy()
        
        selfcopy.x /= (1.+z)
        selfcopy.y /= (1.+z)
        selfcopy.z *= (1.+z)**4

        return selfcopy

    def unredshift(self,z):
        return self.blueshift(z)

    def unblueshift(self,z):
        return self.redshift(z)
    
    def getslice_t(self,tval):
        return self.getslice_x(tval)

    def getslice_lam(self, lval):
        return self.getslice_y(lval)

    def __iter__(self):
        """
        allows iterator constructions like
          for spec1D in spec2D:
              foo
        """
        self.index = 0
        return self

    def next(self):
        """
        Companion function to __iter__
        """
        if self.index >= len(self.x):
            raise StopIteration
        S = Spectrum()
        S.x = num.array(self.y)
        S.y = num.array(self.z[self.index,:])
        self.index += 1
        return S

def GetTemplate():
    """
    returns SN_TEMPLATE
    defines it if not already defined
    """
    #note: SN_TEMPLATE is defined as None in Header.py
    # this variable will store the template so it has to
    # be created only once
    global SN_TEMPLATE
    if SN_TEMPLATE is None:
        SN_TEMPLATE = [None]*len(TEMPLATE_FILES)
        for i in range(len(TEMPLATE_FILES)):
            SN_TEMPLATE[i] = Spectrum2D('%s/%s'%\
                               (TEMPLATE_DIR,TEMPLATE_FILES[i]) )
    return SN_TEMPLATE
    
if __name__ == '__main__':
    S=Spectrum2D()
    S.plotsurface(xmin=5,ymin=2500,ymax=6500)
    p.show()