#!/net/python/bin/python
from Function1D import *
from Header import *
import os

class Spectrum(Function1D):
    #note: redshift formulae assume flux is in units of
    #erg/cm^2/s/A or equivalent (ie, not per unit frequency)
    #  this is the case with the Vega SED provided by SALT
    #  (vega_bohlin2004.dat)
    def redshift(self,z):
        selfcopy = self.xshift(1.+z)
        selfcopy.y /= (1.+z)**4
        return selfcopy

    def unredshift(self,z):
        selfcopy = self.xshift( 1./(1.+z) )
        selfcopy.y *= (1.+z)**4

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

    def unblueshift(self,z):
        return self.redshift(z)

def SingleSpectrum(filename=NUGENT_SPECTRUM,time=0):
    """
    imports a single spectrum at given time from a spectral surface file.
    First column is time, second is wavelength, third is flux
    """
    spectrum = Spectrum()
    x = []
    y = []
    if not os.path.exists(filename):
        print 'SingleSpectrum: ' + filename + ' is not a valid file'
        return s
    found=False
    for line in open(filename,'r'):
        line=map(float,line.split())
                
        if len(line) == 2: #single spectrum, no time
            x.append(line[0])
            y.append(line[1])
                    
        elif len(line) == 3: #time-domain spectrum file
            if line[0]==time:          #if found correct time
                found=True
                x.append(line[1])
                y.append(line[2])
            elif found:             #if already past correct time
                break

    spectrum.x = num.array(x)
    spectrum.y = num.array(y)
    
    return spectrum

def GetRefSpectrum():
    """
    returns REF_SPECTRUM
    defines it if not already defined
    """
    #Note: REF_SPECTRUM is defined as None in Header.py
    # This allows us to create it only once, saving
    #  computation time
    global REF_SPECTRUM
    if REF_SPECTRUM == None:
        REF_SPECTRUM = Spectrum(VEGA_SED)
    return REF_SPECTRUM