import hyperspy.api as hs
import numpy as np
import matplotlib.pyplot as plt

def Gauss2D(X, Y, A=1, cx=0, cy=0, sx=1, sy=1):
    return A*np.exp(-( ((X - cx)**2)/(2*sx**2) + ((Y - cy)**2)/(2*sy**2)))

def Gauss(x, A=1, c=0, s=1):
    return A*np.exp(-((x - c)**2)/(2*s**2))

# Create some moving gaussians
np.random.seed(0)
X,Y = np.meshgrid(np.arange(20), np.arange(20))
g2d = Gauss2D(X,Y, sx=3, sy=4, cx=7, cy=7)
g2d /= g2d.max()
g2d += 0.5
g2d *= 15

g2db = Gauss2D(X,Y, sx=3, sy=4, cx=15, cy=9)
g2db /= g2db.max()
g2db += 0.5
g2db *= 15

g2dc = Gauss2D(X,Y, sx=3.5, sy=5, cx=15, cy=9)
g2dc /= g2dc.max()
g2dc += 1
g2dc *= 40

x = np.linspace(0, 100, 1000)
data = np.zeros(g2d.shape + x.shape)

for index, val in np.ndenumerate(g2d):
    i,j = index
    c = g2d[i,j]
    a2 = g2db[i,j]
    c2 = g2dc[i,j]
    data[i,j] += Gauss(x, A=c, c=c, s=3)
    data[i,j] += Gauss(x, A=a2/2, c=c2, s=3)

 # add a lot of noise
noise = 5 * (np.random.random(data.shape) - 0.5)
data += noise

s = hs.signals.Signal1D(data)
s.axes_manager[-1].scale = np.diff(x)[0]

# Create and fit normal model
m = s.create_model()
g1 = hs.model.components1D.Gaussian(centre=8.)
g2 = hs.model.components1D.Gaussian(centre=40.)
g1.name = 'g1'
g2.name = 'g2'
m.extend([g1,g2])

m.multifit()
normal = g1.A.as_signal() + g2.A.as_signal()

# Create and fit zigzag model
m = s.create_model()
g1 = hs.model.components1D.Gaussian(centre=8.)
g2 = hs.model.components1D.Gaussian(centre=40.)
m.extend([g1,g2])

m.multifit(zigzag=True)
zigzag = g1.A.as_signal() + g2.A.as_signal()

hs.plot.plot_images([s.T.sum(), normal, zigzag], axes_decor=None, label=['Raw Signal', 'Normal', 'Zigzag'])