import numpy as np
import matplotlib.pyplot as plt
from pybaselines.utils import gaussian
from pybaselines import Baseline


x = np.linspace(1, 1000, 500)
signal = (
    gaussian(x, 1, 180, 25)
    + gaussian(x, 8, 350, 10)
    + gaussian(x, 15, 400, 8)
    + gaussian(x, 13, 700, 12)
    + gaussian(x, 9, 800, 10)
)
real_baseline = 5 + 5 * np.exp(-x / 400)
noise = np.random.default_rng(1).normal(0, 0.2, x.size)
y = signal + real_baseline + noise

weights = np.ones(y.shape[0])
# a weight of 0 designates that it is a peak; ensures the small peak at ~180 is
# not classified as part of the baseline
weights[(x > 100) & (x < 250)] = 0

baseline_fitter = Baseline(x, check_finite=False)
unweighted_baseline = baseline_fitter.std_distribution(y, half_window=10, num_std=1.3)[0]
weighted_baseline = baseline_fitter.std_distribution(
    y, half_window=10, num_std=1.3, weights=weights
)[0]

fig, ax = plt.subplots(tight_layout={'pad': 0.2})
data_handle = ax.plot(y)
baseline_handle = ax.plot(unweighted_baseline, '-')
weighted_baseline_handle = ax.plot(weighted_baseline, '--')
ax.set_yticks([])
ax.set_xticks([])
ax.legend(
    (data_handle[0], baseline_handle[0], weighted_baseline_handle[0]),
    ('data', 'unweighted baseline', 'weighted baseline'), frameon=False
)
plt.show()