"""test what the scaling factore of numpy's fft is"""
import numpy as np
from customfunc.app import PlotApp


def setup_axes(n: int, lims=(-10, 10)) -> tuple[np.ndarray, float, np.ndarray, float]:
    x, dx = np.linspace(*lims, n, retstep=True)
    f = np.fft.fftfreq(n, dx)
    return x, dx, f, f[1] - f[0]


def abs2(z):
    return z.real ** 2 + z.imag ** 2


def g(x, w):
    return np.exp(-w * x ** 2)


def main():

    with PlotApp(exp_param=np.linspace(1, 50), n=np.arange(32, 1024)) as app:

        app.set_antialiasing(True)
        fax = app["Frequency"]

        @app.update
        def draw(exp_param, n):
            x, dx, f, df = setup_axes(n)
            ind = np.argsort(f)
            time_amp = 1

            freq_amp = np.sqrt(np.pi / exp_param) * time_amp / dx
            freq_w = np.pi ** 2 / exp_param

            signal = g(x, exp_param) * time_amp

            transformed = np.fft.fft(np.fft.fftshift(signal))
            fax.set_line_data("transformed", f[ind], transformed[ind].real)

            modeled = g(f, freq_w) * freq_amp
            fax.set_line_data("expected", f[ind], modeled[ind])


if __name__ == "__main__":
    main()