#!/usr/bin/env python
#------------------------------------------------------------
# Script demonstrates how to interpolate irregular data so that
# it can be used as input for gauest.py
# Vog, 7 feb 2012, improved at 03-08-2019
#------------------------------------------------------------

import numpy
from matplotlib.pyplot import figure, show, rc
from kapteyn.profiles import gauest

def my_model(p, x, ncomp, zerolev):
   #-----------------------------------------------------------------------
   # This describes the model and its parameters for which we want to find
   # the best fit. 'p' is a sequence of parameters (array/list/tuple).
   #-----------------------------------------------------------------------
   y = 0.0
   for i in range(ncomp):
      A, mu, sigma = p[i*3:(i+1)*3]
      print("A, mu, sig=", A, mu, sigma)
      y += A * numpy.exp(-(x-mu)*(x-mu)/(2.0*sigma*sigma))
   return y + zerolev


# Artificial data
est_zerolev = 10.0
N = 100
numpy.random.seed(0)
x = numpy.random.random(N)*15 - 5  # Random numbers between -5 and +10
truepars1 = [10.0, 5.0, 1.0, 3.0, -1.0, 1.5, 0.0]
y = my_model(truepars1, x, 2, est_zerolev) + 0.3*numpy.random.randn(len(x))

# Analysis of x. function gauest() requires data points yi sampled on 
# a regular grid (i.e. sorted and at equidistant x values)
# It sets a flag that indicates whether an array must be sorted and
# a flag is set when the sorted array x does not represent a regular
# grid. Then interpolation is required. 
# These flags needs to set each time a new array 'x' is detected.

sorted = numpy.all(numpy.diff(x) > 0)
print("Is x sorted?", sorted)
if sorted:
   xs = x
   ys = y
else:
   sortindx = numpy.argsort(x)
   xs = x[sortindx]
   ys = y[sortindx]   # Keep x and y in phase

step = xs[1] - xs[0]
equidis = numpy.all((numpy.diff(xs)-step)< 0.01*step)
print("Are x values (almost) on regular grid? ", equidis)

if not equidis:
  sizex = len(xs)
  delta = (xs[-1]-xs[0])/sizex
  xse = numpy.linspace(xs[0], xs[-1], sizex)
  yse = numpy.interp(xse, xs, ys)
else:
  xse = xs
  yse = ys


step = xse[1] - xse[0]
equidis = numpy.all((numpy.diff(xse)-step)< 0.01*step)
print("Test 2: Are x values (almost) on regular grid? ", equidis)

# thresholds for gauest filters
cutamp = 0.1*yse.max()
cutsig = 5.0
rms = 0.3

# Gauest returns a list with up to ncomp tuples 
# of which each tuple contains the amplitude, 
# the centre and the dispersion of the gaussian, in that order.

ncomps = 0
Q = 6
# Make all x positive and remove offset in y
est_zerolev = yse.min()    # Estimated base level from the data
xse2 = xse-xse[0]
yse2 = yse-est_zerolev

while ncomps != 2 and Q < 15:
   # comp is (amplitude, center, dispersion)
   comps = gauest(xse2, yse2, rms, cutamp, cutsig, q=Q, ncomp=2)
   ncomps = len(comps)
   print("Finding {} components for Q={}".format(ncomps, Q))
   Q += 1

if ncomps == 2:
   d = xse[1] - xse[0]
   p0 = []
   for comp in comps:          # Scale to real x range
      print(comp)
      p0.append(comp[0])
      p0.append(xse[0] + comp[1]), 
      p0.append(comp[2])

   print("Gauest with cutamp, cutsig, rms", cutamp, cutsig, rms)
   print("Number of components found:", ncomps)
   print("Value of Q for which 2 comps. were found:", Q-1)
   print("Found ampl, center, dispersion:", p0)
else:
   print("Could not find any components!")


fig = figure()
rc('legend', fontsize=8)
frame = fig.add_subplot(1,1,1)
frame.plot(x, y, 'or', alpha=0.5, label="Data")
frame.plot(xse, yse, 'g+', label="Interpolated data points")
frame.plot(xse, yse, 'g')
frame.plot(xse, my_model(p0,xse,ncomps,est_zerolev), 'b', lw=2, 
           label="Estimate with gauest()")
leg = frame.legend(loc=2)
show()