#!   /usr/bin/env   python
#    coding: utf8

# Copyright CERN, 2013
# Author: Matthieu Cattin <matthieu.cattin@cern.ch>
# Licence: GPL v2 or later.
# Website: http://www.ohwr.org


# Import system modules
import sys
import time
import os

# Import specific modules
from numpy import *
from pylab import *
from ctypes import *
import scipy.optimize as optimize
import scipy.fftpack as fftpack


"""
Sine wave fit test
"""


def gen_sine(frequency=1, amplitude=1, noise=0.0, cycles=1, rate=10, phase=0.0, offset=0.0):
    # Sample a noisy sine wave and return a list of sample amplitudes

    # Calculated derived values
    dt          = 1/float(frequency*rate)
    period      = 1/float(frequency)
    steps       = arange(0, cycles*period, dt)

    samples     = []    # create an empty list

    # calculate amplitude at each sample point
    # for given number of cycles
    for t in steps:
        angle   = 2*pi*t/period + phase     # in radians
        a       = amplitude*sin(angle) + noise*random() + offset
        samples.append(a)

    return samples


def my_sine(x, o, a, f, p):
    return o + a * sin(f*x + p)


def fit_sine(x, y, f):
    # Guessed amplitude is the max value of input array
    guess_ampl = max(y)
    # Guessed offset is the mean value of input array
    guess_offset = mean(y)
    # Guessed frequency is the highest peak in the input array fft
    yhat = fftpack.rfft(y)
    idx = (yhat**2).argmax()
    freqs = fftpack.rfftfreq(len(y), d = (x[1]-x[0])/(2*pi))
    guess_freq = freqs[idx]
    # Guessed phase is zero (doesn't really matters), could make a fisrt rising_edge detection
    #tmp = diff(sign(y))
    #rising_edges = where(select([tmp>0],[tmp]))[0]
    guess_phase = 0.0

    guess_params = [guess_offset, guess_ampl, guess_freq, guess_phase]
    print("Guessed params    => Offset: %f, Amplitude: %f, Frequency: %f, Phase: %f"%(guess_offset, guess_ampl, guess_freq, guess_phase))

    # Fit f to x,y input data
    (fit_offset, fit_ampl, fit_freq, fit_phase), pcov = optimize.curve_fit(f, x, y, guess_params)
    print("Fit params        => Offset: %f, Amplitude: %f, Frequency: %f, Phase: %f"%(fit_offset, fit_ampl, fit_freq, fit_phase))
    print("pcov:")
    print pcov

    return fit_offset, fit_ampl, fit_freq, fit_phase


def check_bounds(data, upper, lower):
    out_bounds = 0
    for i in range(len(data)):
        if data[i] > upper[i] or data[i] < lower[i]:
            out_bounds += 1
    return out_bounds


def plot_result(data, fit, sin_upper, sin_lower, ylimit, x_data):
    clf()
    plot(x_data, data, 'k', label='Data')
    plot(x_data, fit, 'r', label='Fit')
    plot(x_data, sin_upper, 'b', label='Upper limit')
    plot(x_data, sin_lower, 'g', label='Lower limit')
    ylim(-ylimit-(ylimit/10.0), ylimit+(ylimit/10.0))
    grid(which='both')
    legend()
    draw()
    show()
    return 0


def main (default_directory = '.'):

    # Creates fake input data (noisy sine wave)
    ampl = 10*random()
    freq = 1
    rate = 500
    phase = 2*pi*random()
    noise = 0.5
    cycles = 2
    offset = 0
    y = gen_sine(freq, ampl, noise, cycles, rate, phase, offset)
    x = arange(len(y))
    print("Input sine params => Offset: %f, Amplitude: %f, Frequency: %f, Phase: %f"%(offset, ampl, freq, phase))

    fit_offset, fit_ampl, fit_freq, fit_phase = fit_sine(x, y, my_sine)

    # Create a sine wave from the fit params
    fit = my_sine(x, fit_offset, fit_ampl, fit_freq, fit_phase)

    # Create upper and lower bound sine waves
    margin = 0.05
    upper_margin = fit_offset + abs(fit_ampl) * margin
    lower_margin = fit_offset - abs(fit_ampl) * margin
    upper = my_sine(x, upper_margin, fit_ampl, fit_freq, fit_phase)
    lower = my_sine(x, lower_margin, fit_ampl, fit_freq, fit_phase)

    # Check that input data are inside bounds
    out_bounds = check_bounds(y, upper, lower)
    if out_bounds != 0:
        print("######## Input waveform outside bounds! ############ %d"%out_bounds)

    # Plot all waveforms
    plot_result(y, fit, upper, lower, ampl, x)


if __name__ == '__main__' :
    main()
