from scipy import optimize
from numpy import *
import math
import matplotlib.pyplot as plt

class Parameter:
  def __init__(self, value):
    self.value = value

  def set(self, value):
    self.value = value

  def __call__(self):
    return self.value

def fit(function, parameters, y, x = None):
  def f(params):
    i = 0
    for p in parameters:
      p.set(params[i])
      i += 1
    return y - function(x)

  if x is None: x = arange(y.shape[0])
  p = [param() for param in parameters]
  optimize.leastsq(f, p)

#   Now fitting becomes really easy, for example fitting to a gaussian:
data = loadtxt('refData.res')
# Take the column (0 is first)
x = data[:,0]
y = data[:,3]

# giving initial parameters
height = Parameter(max(y))
height2 = Parameter(max(y))
mu = Parameter(0.05)
sigma = Parameter(0.005)
mu2 = Parameter(0.35)
sigma2 = Parameter(0.005)

# define your function:
def g(x): return 1 + height() * exp(-((x-mu())/sigma())**2) + height2() * exp(-((x-mu2())/sigma2())**2)

# fit! (given that data is an array with the data to fit)
fit(g, [mu, sigma, height, mu2, sigma2, height2], y,x)

plt.plot(x,y,'o', label='original')
plt.plot(x,g(x), label='gaussian fit')
refFile = open('./refParams.lua', 'w')
refFile.write( '-- Fit to double Gaussian resulted in \n')
refFile.write( 'GPampl = %0.15f\n' % height())
refFile.write( 'GPmu   = %0.15f\n' % mu())
refFile.write( 'GPsigma  = %0.15f\n' % sigma())
refFile.write( 'GPmu2    = %0.15f\n' % mu2())
refFile.write( 'GPsigma2 = %0.15f\n' % sigma2())
refFile.write( '-- Parameters found for '+ '\n')
refFile.write( 'function refFunction( x ) '+ '\n')
refFile.write( '   return 1 + GPampl * math.exp(-((x-GPmu)/GPsigma)^2) + GPampl* math.exp(-((x-GPmu2)/GPsigma2)^2) \n')
refFile.write( 'end\n')
refFile.close
plt.show()
