#file SGfilter.py

import numpy


def Smooth(data, halfWinLen=10, order=4, deriv=0, decimate=1):
    """
    Smooth (Savitzky-Golay filter)

    Useage:
        smoothed = Smooth(data, halfWinLen=10, order=4, deriv=0, decimate=1)

    Parameters:
        - data => numpy array of data points
        - halfWinLen => int > 1. Half length of the smoothing window
        - order => int > 0. order of the polynomal.
        - deriv => 0,1,2.  0 for data, 1 for 1st derivative etc
        - decimate => 1 or more. 1=every point, 2=every other point etc
    """
    try:
        ParameterChecks(len(data), halfWinLen, order, deriv, decimate)
    except ValueError, msg:
        raise ValueError("\nError in parameters supplied:\n%s\n" % msg)

    ordList = range(order+1)
    winList = range(-halfWinLen, halfWinLen+1)

    b = numpy.mat([[k**i for i in ordList] for k in winList])
    win = numpy.linalg.pinv(b).A[deriv]

    return numpy.convolve(win, data, mode='same')[::decimate]


def ParameterChecks(dataLen, halfWinLen, order, deriv, decimate):

    try:
        kernel = abs(int(halfWinLen))
        order = abs(int(order))
    except ValueError, msg:
        raise ValueError("halfWinLen and order have to be of type int.")

    if halfWinLen < order//2 + 1:
        raise ValueError("halfWinLen is too small for the polynomal order\n"
                         "halfWinLen should be > order//2 + 1")

    if deriv not in [0,1,2]:
        raise ValueError("deriv must be 0, 1 or 2")

    if decimate < 1:
        raise ValueError("decimate must be 1 or more")


def TestSmooth():
    t = numpy.arange(0.0, 20.0, 0.002)
    x = numpy.sin(t) + 0.1*numpy.random.normal(size=len(t))
    dec = 10
    ts = t[::dec]
    xs = Smooth(x, halfWinLen=80, order=3, decimate=dec)
    return t, x, ts, xs
