#!/net/python/bin/python
import os
import numpy as num
from mathfuncs import splrep,splev
import pylab as p
import types
from Header import i_max,NumTypes

class Function1D:
    """
    A class to hold an equally sampled 1D function
    """
    def __init__(self,filename=''):
        #filename is the path to a file that has two columns of data:
        #  Column 1 is the x value, column 2 is the y value
        self.getfromfile(filename)
        
    def getfromfile(self,filename):
        """
        imports a function from filename
        filename is assumed to be two columns,
        with x in column 1, y in column 2
        unused lines should start with #
        """
        self.x = []
        self.y = []

        if filename == '':
            pass
        elif not os.path.exists(filename):
            print "Function1D:getfromfile: cannot open path %s" % filename
            return

        else:
            for line in open(filename,'r'):
                if line.startswith(('#','\n','@')):
                    continue

                line = map(float,line.split())
                if len(line) < 2:
                    continue

                self.x.append(line[0])
                self.y.append(line[1])

        self.x = num.array(self.x)
        self.y = num.array(self.y)
        return self
    
    def yval(self,xval):
        """
        use linear interpolation to find y(xval)
        """
        if not self.inrange(xval):
            print "Function1D:yval: xval out of range"
            return None
        
        i = self.x.searchsorted(xval)
        
        if xval == self.x[i]:
            return self.y[i]
        else:
            dydx = (self.y[i]-self.y[i-1])/(self.x[i]-self.x[i-1])
            return self.y[i-1] + dydx * (xval-self.x[i-1])

    def inrange(self,val):
        """
        return True if xmin <= val <= xmax
               False otherwise
        """
        if len(self.x) == 0:
            return 0
        else:
            return (val >= self.x[0]) and (val <= self.x[-1])

    def contract(self,xmin,xmax):
        selfcopy = self.__class__() #new object
        xtemp = self.x[num.where(self.x >= xmin)]
        ytemp = self.y[num.where(self.x >= xmin)]
        selfcopy.x = xtemp[num.where(xtemp <= xmax)]
        selfcopy.y = ytemp[num.where(xtemp <= xmax)]

        return selfcopy

    def splrep(self):
        """
        returns the spline representation of self
        checks to make sure there are no duplicate values
        in x, as this returns an error
        """
        x=[self.x[0]]
        y=[self.y[0]]
        for i in range(1,len(self.x)):
            if x[-1] != self.x[i]:
                x.append(self.x[i])
                y.append(self.y[i])

        return splrep(x,y)
            
    
    def resample(self,xrange=None):
        """
        Uses a cubic interpolation scheme to resample the function.
        If xrange extends beyond xmin and xmax, range is truncated
        """

        if xrange == None:
            xrange = num.arange( self.x[0], self.x[-1]+0.1,0.1 )
        
        xrange_s = [val for val in xrange if self.inrange(val)]
        
        #prepare the cubic spline
        sp = self.splrep()

        #run the interpolation
        new_y = splev(xrange_s,sp)

        #create the new object
        newFunction=self.__class__()
        newFunction.x = num.array(xrange_s)
        newFunction.y = num.array(new_y)

        return newFunction

    def maximum(self,a=None,b=None,dx=None,tol=1E-4):
        """
        A brute force algorithm to find the absolute maximum of sp
        between a and b
         - (a,b) are the endpoints of the interval
         - dx is the range in which the function does not change much
             note: dx too large may miss an absolute maximum
                   dx too small will take a long time
         - tol gives the maximum final error in xmax
        """
        if a == None:
            a = min(self.x)
        if b == None:
            b = max(self.x)
        if dx == None:
            dx = 0.01*(b-a)
            
        sp = self.splrep()
        while (b-a) > tol:
            xrange = num.arange(a,b+dx,dx)
            yrange = splev(xrange,sp)
            max_i = i_max(yrange)
        
            i_a = max(0,max_i-1)
            a = xrange[i_a]
            fa = yrange[i_a]
        
            i_b = min(len(xrange)-1,max_i+1)
            b = xrange[i_b]
            fb = xrange[i_b]
        
            dx *= 0.1

        if fb > fa:
            return b
        else:
            return a

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

        return selfcopy

    def integrate(self,xmin=None,xmax=None):
        if xmin == None:
            xmin = self.x[0]
        if xmax == None:
            xmax = self.x[-1]
        selfcopy = self.contract(xmin,xmax)
        
        dx  = num.array( [selfcopy.x[i]-selfcopy.x[i-1]\
                          for i in range(1,len(selfcopy))] )
        y = num.array( [0.5*(selfcopy.y[i-1] + selfcopy.y[i])\
                        for i in range(1,len(selfcopy))] )

        return sum(y*dx)

    def xshift(self,other):
        """
        returns a copy of self with each x value multiplied by other
        """
        if type(other) in NumTypes: #int,float,long,numpy64
            selfcopy = self.copy()
            selfcopy.x *= other
        elif type(other) in types.FunctionType:
            try: y = other(self.x[0])
            except:
                print "Function1D:shiftx:cannot support ",other
                return None
            selfcopy = self.copy()
            for i in range(len(self.x)):
                selfcopy.x[i] *= other(selfcopy.x[i])
            
        return selfcopy
    
    def graph(self,style = '-',fill=False,fcolor='#eeeeee',linewidth=0.5):
        if fill:
            p.fill(self.x,self.y,fcolor)
        p.plot(self.x,self.y,style,linewidth=linewidth)

    def ymean(self):
        """
        returns the mean y value
        """
        return self.integrate/(self.x[-1]-self.x[0])

    def xmean(self):
        """
        returns the mean x value
        """
        return sum(self.x*self.y)/sum(self.y)

    def __len__(self):
        return len(self.x)
            
    def __mul__(self, other):
        if type(other) in NumTypes: #int,float,long,numpy64
            selfcopy = self.copy()
            selfcopy.y *= other

        elif type(other) == types.FunctionType:
            try: y = other(self.x[0])
            except:
                print "Function1D:__mul__:cannot support ",other
                return None
            selfcopy = self.copy()
            for i in range(len(self)):
                selfcopy.y[i] *= other(selfcopy.x[i])

        elif type(other) == type(self):
            try: othercopy = other.resample(self.x)
            except:
                print "Function1D:__mul__:cannot support ",other
                return None
            selfcopy = self.contract(othercopy.x[0],othercopy.x[-1])
            selfcopy.y *= othercopy.y

        return selfcopy

    def __rmul__(self,other):
        return self.__mul__(other)

    def __imul__(self,other):
        if type(other) in NumTypes: #int,float,long,numpy64
            self.y *= other

        elif type(other) == types.FunctionType:
            try: y = other(self.x[0])
            except:
                print "Function1D:__imul__:cannot support ",other
                return None
            for i in range(len(self)):
                self.y[i] *= other(self.x[i])

        elif type(other) == type(self):
            try: othercopy = other.resample(self.x)
            except:
                print "Function1D:__imul__:cannot support ",other
                return None
            self = self.contract(othercopy.x[0],othercopy.x[-1])
            self.y *= othercopy.y

        return self
            
    def __div__(self, other):
        if type(other) in NumTypes: #int,float,long,numpy64
            selfcopy = self.copy()
            selfcopy.y /= other

        elif type(other) == types.FunctionType:
            try: y = other(self.x[0])
            except:
                print "Function1D:__div__:cannot support ",other
                return None
            selfcopy = self.copy()
            for i in range(len(self)):
                selfcopy.y[i] /= other(selfcopy.x[i])

        elif type(other) == type(self):
            try: othercopy = other.resample(self.x)
            except:
                print "Function1D:__div__:cannot support ",other
                return None
            selfcopy = self.contract(othercopy.x[0],othercopy.x[-1])
            selfcopy.y /= othercopy.y

        return selfcopy

    def __idiv__(self,other):
        if type(other) in NumTypes: #int,float,long,numpy64
            self.y /= other

        elif type(other) == types.FunctionType:
            try: y = other(self.x[0])
            except:
                print "Function1D:__idiv__:cannot support ",other
                return None
            for i in range(len(self)):
                self.y[i] /= other(self.x[i])

        elif type(other) == type(self):
            try: othercopy = other.resample(self.x)
            except:
                print "Function1D:__idiv__:cannot support ",other
                return None
            self = self.contract(othercopy.x[0],othercopy.x[-1])
            self.y /= othercopy.y

        return self

    def __add__(self,other):
        if type(other) in NumTypes: #int,float,long,numpy64
            selfcopy = self.copy()
            selfcopy.y += other

        elif type(other) == types.FunctionType:
            try: y = other(self.x[0])
            except:
                print "Function1D:__add__:cannot support ",other
                return None
            selfcopy = self.copy()
            for i in range(len(self)):
                selfcopy.y[i] += other(selfcopy.x[i])

        elif type(other) == type(self):
            try: othercopy = other.resample(self.x)
            except:
                print "Function1D:__add__:cannot support ",other
                return None
            selfcopy = self.contract(othercopy.x[0],othercopy.x[-1])
            selfcopy.y += othercopy.y

        return selfcopy

    def __radd__(self,other):
        return self.__add__(other)

    def __iadd__(self,other):
        if type(other) in NumTypes: #int,float,long,numpy64
            self.y += other

        elif type(other) == types.FunctionType:
            try: y = other(self.x[0])
            except:
                print "Function1D:__radd__:cannot support ",other
                return None
            for i in range(len(self)):
                self.y[i] += other(self.x[i])

        elif type(other) == type(self):
            try: othercopy = other.resample(self.x)
            except:
                print "Function1D:__radd__:cannot support ",other
                return None
            self = self.contract(othercopy.x[0],othercopy.x[-1])
            self.y += othercopy.y

        return self

    def __neg__(self):
        selfcopy = self.copy()
        selfcopy.y = -selfcopy.y
        return selfcopy

    def __sub__(self,other):
        if type(other) in NumTypes: #int,float,long,numpy64
            selfcopy = self.copy()
            selfcopy.y -= other

        elif type(other) == types.FunctionType:
            try: y = other(self.x[0])
            except:
                print "Function1D:__sub__:cannot support ",other
                return None
            selfcopy = self.copy()
            for i in range(len(self)):
                selfcopy.y[i] -= other(selfcopy.x[i])

        elif type(other) == type(self):
            try: othercopy = other.resample(self.x)
            except:
                print "Function1D:__sub__:cannot support ",other
                return None
            selfcopy = self.contract(othercopy.x[0],othercopy.x[-1])
            selfcopy.y -= othercopy.y

        return selfcopy

    def __rsub__(self,other):
        return -self.__sub__(other)

    def __isub__(self,other):
        if type(other) in NumTypes: #int,float,long,numpy64
            self.y -= other

        elif type(other) == types.FunctionType:
            try: y = other(self.x[0])
            except:
                print "Function1D:__rsub__:cannot support ",other
                return None
            for i in range(len(self)):
                self.y[i] -= other(self.x[i])

        elif type(other) == type(self):
            try: othercopy = other.resample(self.x)
            except:
                print "Function1D:__rsub__:cannot support ",other
                return None
            self = self.contract(othercopy.x[0],othercopy.x[-1])
            self.y -= othercopy.y

        return self

    def __pow__(self,other):
        if type(other) in NumTypes: #int,float,long,numpy64
            selfcopy = self.copy()
            selfcopy.y = selfcopy.y**other
            return selfcopy
        else:
            print "Function1D:__pow__:cannot support ",other
            return None

def FromAnalytic(func,xmin=0.0,xmax=1E4,xstep=100.0):
    newFunction = self.__class__()
    newFunction.x = num.arange(xmin,xmax+xstep,xstep)
    newFunction.y = num.array([func(x) for x in newFunction.x])
    
    return newFunction
        
            

if __name__ == '__main__':
    f = '../code/filters/SNLS/g.dat'
    v = '../code/Vega/Vega.sed'