import matplotlib.pyplot as plt
from spectractor.tools import plot_spectrum_simple
from spectractor.extractor.spectrum import Spectrum, detect_lines
from spectractor.extractor.spectroscopy import *
lambdas = np.arange(300,1000,1)
spectrum = 1e4*np.exp(-((lambdas-600)/200)**2)
spectrum += HALPHA.gaussian_model(lambdas, A=5000, sigma=3)
spectrum += HBETA.gaussian_model(lambdas, A=3000, sigma=2)
spectrum += O2.gaussian_model(lambdas, A=-3000, sigma=7)
spectrum_err = np.sqrt(spectrum)
spec = Spectrum()
spec.lambdas = lambdas
spec.data = spectrum
spec.err = spectrum_err
fwhm_func = interp1d(lambdas, 0.01 * lambdas)
lines = Lines([HALPHA, HBETA, O2], hydrogen_only=True, atmospheric_lines=True, redshift=0, emission_spectrum=True)
global_chisq, res, ws = detect_lines(lines, lambdas, spectrum, spectrum_err, fwhm_func=fwhm_func)
spec.lines = lines
fig = plt.figure()
plot_spectrum_simple(plt.gca(), lambdas, spec.data, data_err=spec.err)
lines.plot_detected_lines(plt.gca())
plt.show()