import numpy as np
import hyperspy.api as hs

domain = 32 #size of the square domain
hfactor = 600
cent = (domain//2, domain//2)
y,x = np.ogrid[-cent[0]:domain-cent[0], -cent[1]:domain-cent[1]]

def gaussian2d(x, y, A=1, x0=0, y0=0, sigmax=20, sigmay=10):
    return A * np.exp(-((x-x0)**2 / 2 / sigmax ** 2 + (y-y0)**2 / 2 / sigmay ** 2))

center_narrow = 50 + 10 * np.sin(3 * np.pi * x / domain) * np.cos(4 * np.pi * y / domain)
center_wide = 50 + 10 * (-0.1 * np.sin(3 * np.pi * x / domain) * np.cos(4 * np.pi * y / domain))

r = np.sqrt(x**2 + y**2)
h_narrow = .5 * (.5 + np.sin(r)**2) * gaussian2d(x, y) *  hfactor
h_wide = (.5 + np.cos(r)**2) * gaussian2d(x, y) *  hfactor

s = hs.signals.Signal1D(np.ones((domain,domain, 1024)))
s.metadata.General.title = 'Two gaussians'
s.axes_manager[0].name = "x"
s.axes_manager[0].units = "nm"
s.axes_manager[1].name = "y"
s.axes_manager[1].units = "nm"

s.axes_manager[2].name = "Energy"
s.axes_manager[2].name = "eV"
s.axes_manager[2].scale = 0.1
m0 = s.create_model()

gs01 = hs.model.components1D.GaussianHF()
gs01.name = "wide"
m0.append(gs01)
gs01.fwhm.value = 60
gs01.centre.map['values'][:] = center_wide
gs01.centre.map['is_set'][:] = True
gs01.height.map['values'][:] = h_wide
gs01.height.map['is_set'][:] = True

gs02 = hs.model.components1D.GaussianHF()
gs02.name = "narrow"
m0.append(gs02)
gs02.fwhm.value = 6
gs02.centre.map['values'][:] = center_narrow
gs02.centre.map['is_set'][:] = True
gs02.height.map['values'][:] = h_narrow
gs02.height.map['is_set'][:] = True
s.data = m0.as_signal().data
s.add_poissonian_noise(random_state=0)
m0.store("ground truth")
s.save("two_peaks.hspy", overwrite=True)

import numpy as np
import hyperspy.api as hs

k = 1
alpha = 15
amp = 3

gaus_position = 15
gaus_width = 5
gaus_A = 300

gradient = 0.6
offset= 3

sin_component = hs.model.components1D.Expression('A * sin(k*x + alpha)', name='sin', k=k,
                                               alpha=alpha, A=amp)
gaus = hs.model.components1D.Gaussian(A=gaus_A, sigma=gaus_width, centre=gaus_position)

poly = hs.model.components1D.Polynomial(1)

poly.coefficients.value = (gradient, offset)

axis = np.linspace(0, 30, 3000, dtype='double')

result = sin_component.function(axis)+ gaus.function(axis) + poly.function(axis)
s = hs.signals.Signal1D(result)
s.axes_manager[0].name = 'x'
s.axes_manager[0].scale = 0.1
s.axes_manager[0].offset = 0

s.metadata.General.author = 'Tomas Ostasevicius'
s.metadata.General.title = 'Sin + poly(2) + Gaussian'
s.save('wobbly_peak', overwrite=True)