"""This script is used to visualize the behavior of the fft if the input is
purely real."""

import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt


def bit_reverse(val: int, num_bits: int) -> None:
    b = '{:0{num_bits}b}'.format(val, num_bits=num_bits)
    return int(b[::-1], 2)


def fft_reorder(x: np.array) -> np.array:
    result = np.zeros(x.size)
    for i in range(x.size):
        result[i] = x[bit_reverse(i, int(np.log2(x.size)))]
    return result


def fft(x: np.array) -> np.array:
    x = fft_reorder(x).astype('complex64')

    N0 = x.size

    # fig, axes = plt.subplots(2, int(np.log2(N0)))

    N = 2
    while N <= N0:
        first_subfft_index = 0
        while first_subfft_index < N0:
            temp = x[first_subfft_index + N // 2]
            x[first_subfft_index + N // 2] = x[first_subfft_index] - temp
            x[first_subfft_index] = x[first_subfft_index] + temp

            for k in range(1, N // 2):
                factor = np.exp(-2j * np.pi * k / N)

                x[first_subfft_index + k] \
                    = x[first_subfft_index + k] \
                      + factor * x[first_subfft_index + k + N // 2]

            for k in range(1, N // 2):
                x[first_subfft_index + N - k] \
                    = x[first_subfft_index + k].conj()
            #      \
            #         = x[first_subfft_index + k] + temp
            # for k in range(N // 2):
            #     factor = np.exp(- 2j * np.pi * k / N)
            #
            #     temp = factor * x[first_subfft_index + k + N // 2]
            #     x[first_subfft_index + k + N // 2] \
            #         = x[first_subfft_index + k] - temp
            #     x[first_subfft_index + k] \
            #         = x[first_subfft_index + k] + temp

            first_subfft_index += N
        # axes[0][int(np.log2(N)) - 1].plot(x.real[1:])
        # axes[1][int(np.log2(N)) - 1].plot(x.imag[1:])
        N *= 2

    plt.show()

    return x


def main():
    sns.set_theme()

    x = np.arange(16) + 1
    X = fft(x)
    X_reference = np.fft.fft(x)

    for X_i, X_ref_i in zip(X, X_reference):
        print(
            f"{X_i.real:06.2f} + j({X_i.imag:06.2f})\t\t{X_ref_i.real:06.2f} "
            f"+ j({X_ref_i.imag:06.2f})")

    print(np.allclose(X, X_reference, rtol=1e-05, atol=1e-08))


if __name__ == "__main__":
    main()