import matplotlib.pyplot as plt
import numpy as np

import scgenerator as sc


def main():
    nf = 513
    nseg = 240
    ntot = (nf - 1) * (nseg - 1)

    fs = 44100
    nf_full = 45611
    f = np.linspace(0, fs // 2, nf_full)

    spec = 1 / (f + 1)
    spec += np.exp(-(((f - (0.1 * fs)) / (0.5 * fs)) ** 2))
    spec += np.exp(-(((f - (0.45 * fs)) / (0.02 * fs)) ** 2))
    # spec = np.ones_like(f) * 1e-5
    # spec += np.exp(-(((f - 15000) / 300) ** 2))

    noise = sc.noise.NoiseMeasurement(f, spec)
    t, y = noise.time_series(nf=nf, nseg=nseg)
    tf, yf = noise.time_series(nf=ntot // 2 + 1)
    print(y.std(), yf.std())

    plt.plot(t, y, ls="", marker=".")
    plt.plot(tf, yf, ls="", marker=".")
    plt.show()

    newnoise = noise.from_time_series(y, t[1] - t[0], nf=nf)
    newnoise_full = noise.from_time_series(yf, tf[1] - tf[0], nperseg=(nf - 1) * 2)

    print(newnoise.psd[1:-1].var())
    print(newnoise_full.psd[1:-1].var())

    fig, ax = plt.subplots()
    ax.plot(*noise.plottable())
    ax.plot(*newnoise.plottable())
    ax.plot(*newnoise_full.plottable(discard_nyquist=True))
    ax.set_xscale("log")

    axin = plt.gca().inset_axes(
        [0.1, 0.1, 0.3, 0.3],
        xlim=(1.5e4, fs / 1.9),
        ylim=(-3.1, 2.4),
        yticklabels=[],
    )
    # axin.set_xscale("log")
    axin.plot(*newnoise.plottable())
    axin.plot(*newnoise_full.plottable())
    axin.plot(*noise.plottable())
    ax.indicate_inset_zoom(axin)
    axin.tick_params(labelbottom=False)

    plt.show()


if __name__ == "__main__":
    main()