#!/net/python/bin/python
from Function1D import *
import numpy as num
from AnalyticFit import AnalyticFit1D
from mathfuncs import splrep,splev,smooth

class Data1D(Function1D):
    """
    A class to hold a 1D set of data with error bars
    """
    def __init__(self,filename=''):
        self.getfromfile(filename)

    def has_dy(self):
        if len(self.dy) > 0:
            return True
        else:
            return False

    def getfromfile(self,filename):
        """
        imports a function from filename
        filename is assumed to be three columns,
        with x in 1, y in 2, dy in 3
        unused lines should start with #
        """
        self.x = []
        self.y = []
        self.dy = []
        if filename == '':
            pass
        
        elif not os.path.exists(filename):
            print "Function1D:getfromfile: cannot open path %s" % filename
            
        else:
            for line in open(filename,'r'):
                if line.startswith(('#','\n','@')):
                    continue
                line = map(float,line.split())
                self.x.append(line[0])
                self.y.append(line[1])
                self.dy.append(line[2])

        self.x = num.array(self.x)
        self.y = num.array(self.y)
        self.dy = num.array(self.dy)
        return self
        

    def copy(self):
        selfcopy = self.__class__()
        selfcopy.x = self.x.copy()
        selfcopy.y = self.y.copy()
        selfcopy.dy = self.dy.copy()

        return selfcopy

    def fourier_fit(self,N=20):
        fitfunc = AnalyticFit1D(self.x,self.y,self.dy,N)
        
        F = Function1D()
        xmin = self.x[0]
        xmax = self.x[-1]
        dx = 0.01*(xmax-xmin)
        F.x = num.arange(xmin,xmax+dx,dx)
        F.y = fitfunc(F.x)

        return F

    def smooth(self,s=0.1):
        F = Function1D()
        xmin = self.x[0]
        xmax = self.x[-1]
        dx = 0.01*(xmax-xmin)
        F.x = num.arange(xmin,xmax+dx,dx)
        
        y2 = smooth(self.x,self.y,self.dy,s=s)
        sp = splrep(self.x.tolist(),y2)
        F.y = num.array(splev(F.x,sp))
        return F

    def multiply(self,other):
        self.y *= other
        if len(self.dy) > 0:
            self.dy *= other

    def upperbound(self):
        selfcopy = self.copy()
        selfcopy.dy = num.array([])
        if self.has_dy():
            selfcopy.y += self.dy
        return selfcopy

    def lowerbound(self):
        selfcopy = self.copy()
        selfcopy.dy = num.array([])
        if self.has_dy():
            selfcopy.y -= self.dy
        return selfcopy

    def pointgraph(self,color='k'):
        if self.has_dy():
            p.errorbar(self.x,self.y,self.dy,fmt='.'+color)
        else:
            p.plot(self.x,self.y,'.'+color)

    def linegraph(self,color='k'):
        self.graph(color+'-')
        self.upperbound().graph(color + ':')
        self.lowerbound().graph(color + ':')