# Generated by pandoc-plot 1.8.0
import matplotlib.pyplot as plt
import numpy as np

np.random.seed(2019)

from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import pywt
from numpy.random import default_rng

def discrete_colors(it):
    """Returns a list of discrete colors to plot, for example, various time-traces."""
    cmap = plt.get_cmap("inferno")
    mi, ma = 0.11, 0.75

    it = list(it)

    if len(it) == 1:
        yield from zip([cmap(mi)], it)
        return
    elif len(it) == 2:
        yield from zip([cmap(mi), cmap(ma)], it)
        return

    step = (ma - mi) / (len(it) - 1)
    yield from zip([cmap(mi + i * step) for i in range(len(it))], it)

random = default_rng(seed=2022)

data = np.loadtxt(
    Path("files") / "wavelet-filter" / "diffraction.csv", delimiter=",", skiprows=1
)
t = data[:, 0]
signal = data[:, 1]

noise = random.normal(size=t.shape[0], loc=-0.5, scale=1)

banks = pywt.wavedec(signal + noise, pywt.Wavelet('sym17'), level=3, mode='zero')
assert len(banks) == 4

fig, axes = plt.subplots(len(banks), 1, figsize=(8,9))
for color, (ax, bank, label) in discrete_colors(zip(axes, banks, ["Approx. level=3", "Detail level=3", "Detail level=2", "Detail level=1"])):
    ax.plot(bank, color=color)
    ax.xaxis.set_visible(False)
    ax.set_title(label)
    ax.set_ylabel("Amplitude [a.u.]")
    ax.set_xlim([0, len(bank)])