#!/net/python/bin/python

from Header import *
from FilterSet import *
from Data1D import *

class LightCurve:
    #make fluxes the native coordinate
    #write function SetZeroPoint(self,ZP)
    """
    A class to hold light curve data.  lc is a dictionary
    that holds Data1D objects, each containing a flux
    light curve
    """
    def __init__(self,Spec2D=None,filterset='SDSS',z=0.0,DistMod = 0.0,EBV=0.0):
        self.filterset = GetFilterSet(filterset)

        self.ZP = {}

        self.data = {}

        for f in self.filterset:
            self[f] = Data1D()
            self[f].x=[]
            self[f].y=[]
            self[f].dy = []
            self.ZP[f] = 30

        if not Spec2D:
            return

        i=0
        for spectrum in Spec2D:
            xval = Spec2D.x[i] * (1.+z) #time dialation
            for f in self.filterset.extincted(EBV):
                flux = self.filterset[f].CCDflux(spectrum.redshift(z),self.ZP[f])
                self[f].x.append(xval)
                self[f].y.append(flux)
            i+=1

        for f in self.filterset:
            self[f].x = num.array(self[f].x)
            self[f].y = num.array(self[f].y)
            self[f].dy = num.array(self[f].dy)

    def offset(self,time):
        """
        offsets tmax by time
        """
        for f in self:
            self[f].x += time
    
    def add_distance(self,DistMod):
        """
        adds DistMod to magnitudes
        """
        for f in self:
            self[f] *= 10**(-0.4*DistMod)

    def __imul__(self,other):
        assert type(other) in NumTypes #defined in Header.py
        for f in self:
            self[f].multiply(other)
        return self

    def SetZeroPoints(self,ZP):
        """
        changes the Zero Point of the given band
        ZP is a dictionary {'u':ZP_u,'g':ZP_g...}
        """
        for f in ZP:
            self[f].multiply( 10**(0.4*(ZP[f] - self.ZP[f]) ) )
            self.ZP[f] = ZP[f]

    def __getitem__(self,key):
        return self.data[key]

    def __setitem__(self,key,val):
        self.data[key] = val

    def __iter__(self):
        return iter(self.filterset)

    def copy(self):
        selfcopy = self.__class__()

        selfcopy.filterset = self.filterset
        selfcopy.data = {}
        selfcopy.ZP = {}
        for f in self.filterset:
            selfcopy[f] = self[f].copy()
            selfcopy.ZP[f] = self.ZP[f]

        return selfcopy

    def pointgraph(self, filters=None):
        if filters == None:
            filters = self.filterset.keys
        for filter in self:
            if filter not in filters:
                continue
            self[filter].pointgraph(color = Colors[filter])

    def linegraph(self, filters=None):
        if filters == None:
            filters = self.filterset.keys
        for filter in self.data:
            if filter not in filters:
                continue
            self[filter].linegraph(color = Colors[filter])
        

def ImportLightCurve(path,filterset='SDSS'):
    #normalize will convert telescope flux (ADU counts) to physical flux, based
    # on vega magnitudes
    for line in open('%s/lightfile' % path, 'r'):
        if line.startswith('Redshift'):
            z=float(line.split()[1])
    
    LC = LightCurve(filterset=filterset,z=z)

    #because we may need to delete some of these objects as we
    # iterate over them, make a copy of the keys:
    filters = LC.filterset.keys[:]
    for filter in filters:
        filepath = '%s/lc2fit_%s.dat' % (path,filter)
        if not os.path.exists(filepath):
            del LC.filterset[filter]
            continue
        x = []
        y = []
        dy = []
        
        for line in open(filepath,'r'):
            if line[0] in ('@','#','\n'):
                continue
            line = map(float,line.split())
            # mag = -2.5 log10(flux) + ZP
            # dmag = (2.5/log(10))dflux/flux
            ZP = line[3]
            x.append( line[0] )
            y.append( line[1] )
            dy.append( line[2] )
            ZP = line[3]

        LC[filter].x = num.array(x)
        LC[filter].y = num.array(y)
        LC[filter].dy = num.array(dy)
        LC.ZP[filter] = ZP
    
    return LC

def LC_with_error(LC_mean,LC_upper,LC_lower = None):
    """
    takes a mean Lightcurve, an upper bound, and optional
    lower bound, and returns a single light curve with
    appropriate error bars
    In the case of uneven error bars, dy is set to the maximum
    """
    assert AreEquivalent(LC_mean,LC_upper)
    LC = LC_mean.copy()
    if LC_lower is not None:
        assert AreEquivalent(LC_mean,LC_lower)
        for f in LC_mean:
            dy_upper = abs(LC_mean[f].y - LC_upper[f].y)
            dy_lower = abs(LC_mean[f].y - LC_lower[f].y)
            LC[f].dy = num.maximum(dy_upper,dy_lower)
    else:
        for f in LC_mean:
            LC[f].dy = abs(LC_mean[f].y - LC_upper[f].y)
            
    return LC

def AreEquivalent(LC1,LC2):
    """
    returns False if LC1 and LC2 have different filtersets
     or different time values,
    returns True otherwise
    """
    if LC1.filterset.keys != LC2.filterset.keys:
        return False
    for key in LC1:
        if not all(LC1[key].x == LC2[key].x):
            return False

    return True

if __name__ == '__main__':
    LC = ImportLightCurve('/astro/users/vanderplas/research/SDSSdata/10028')
    LC.pointgraph_flux()
    p.show()