import warnings
import numpy as np
import matplotlib.pyplot as plt
import astropy.units as u
#
from specutils.spectra.spectrum1d import Spectrum1D
from specutils.fitting.continuum import fit_continuum
#
np.random.seed(0)
x = np.linspace(0., 10., 200)
y = 3 * np.exp(-0.5 * (x - 6.3) ** 2 / 0.1 ** 2)
y += np.random.normal(0., 0.2, x.shape)
y += 3.2 * np.exp(-0.5 * (x - 5.6) ** 2 / 4.8 ** 2)
#
spectrum = Spectrum1D(flux=y * u.Jy, spectral_axis=x * u.um)
region = [(1 * u.um, 5 * u.um), (7 * u.um, 10 * u.um)]
with warnings.catch_warnings():  # Ignore warnings
    warnings.simplefilter('ignore')
    fitted_continuum = fit_continuum(spectrum, window=region)
y_fit = fitted_continuum(x*u.um)
#
f, ax = plt.subplots()  # doctest: +IGNORE_OUTPUT
ax.plot(x, y)  # doctest: +IGNORE_OUTPUT
ax.plot(x, y_fit)  # doctest: +IGNORE_OUTPUT
ax.set_title("Continuum Fitting")  # doctest: +IGNORE_OUTPUT
plt.grid(True)  # doctest: +IGNORE_OUTPUT
