#!/usr/bin/env python3

import numpy as np

from rich.traceback import install
install()

# We want to solve the FFT and inverse FFT in an easily accessible way, not necessarily efficient.
#
# The coefficient representation:
# - Every polynomial of degree d is of the form P_d(x) = c_0 * x**0 + c_1 * x**1 + ... + c_d * x**d
# - There exists a bijective mapping between the set of (d + 1)-tuples of coefficients and the set of polynomials
# - If we store a list of coefficients where the index corresponds to the degree we have a very simple representation of a polynomial
# - Example: P_2(x) = x**2 + 5 has the coefficient representation [5, 0, 1]
#
# The value representation:
# - There exists a surjective mapping between the set of (d + 1)-tuples of points from the polynomial's image to the set of polynomials
# - Example: P_2(x) = x**2 could have the value representation [(-1, 1), (1, 1)] or [(-2, 4), (2, 4)]
#
# Roots of unity:
# - The n-th roots of unity are all complex numbers satisfying the equation z**n = 1.
# - They can be visualized as n equally spaced points on the complex unit circle
#   - In the real domain the unit circle is defined by the constraint x**2 + y**2 = 1, meaning every point on the real unit circle has the distance 1 from the origin
#     Different formulation: We have a point (x, y) on the real unit circle with the angle p: x = cos(p), y = sin(p) and we know that cos(p)**2 + sin(p)**2 = 1, which is the distance to the origin (pythagorean theorem)
#   - In the complex domain the pythagorean theorem works the same: We have a complex number c = (c_real, c_img) on the complex unit circle with the angle p: c_real = cos(p), c_img = sin(p)
#     If we take Eulers formula: e**(i * p) = cos(p) + i * sin(p), we can easily define the r-th root of the n-th roots of unity:
#     - The angle between the n-th roots is (2pi)/n, because all n roots are equally spaced
#     - The angle of the r-th root can be written as (2pi * r)/n
#     - The r-th root can thus be written as cos((2pi * r)/n) + i * sin((2pi * r)/n), which is equal to e**((i * 2pi * r)/n)
#     - The "distance from the origin = 1" constraint holds (obviously): e**((i * 2pi * r)/n) = (e**((i * pi * r)/n))**2 = cos((2pi * r)/n)**2 + i * sin((2pi * r)/n)**2 = 1
# - With this visualization it is also easy to define them in a closed form term: w**r = e**((i * 2pi * r)/n) (w**r is the r-th of the n-th roots of unity)
# - It is also clear that the roots always come in +/- pairs, because if they are equally spaced on the complex unit circle every root is connected to another root by a line through the origin
#
# Calculating the value representation (FFT):
# - The naive approach is to evaluate polynomial of degree d in (d + 1) distinct points,
#   but this would have a runtime of O(d**2), as every evaluation has a runtime of O(d)
# - To lower the runtime we can make use of a polynomial's symmetry: Even polynomials are symmetric regarding the imaginary axis, odd polynomials are symmetric regarding the origin
# - Every polynomial can be split into a sum of an even and an odd polynomial, so if we do this a single time, we have a runtime of O((d/2) * d), which is still quadratic
#   - We split e.g. P_3(x) = 3 * x**3 + x**2 + 2 * x + 5 into P_3_even(x) = x + 5 and P_3_odd(x) = 3 * x + 2,
#     so we get the alternative representation P_3(x) = P_3_even(x**2) + x * P_3_odd(x**2) = x**2 + 5 + x * (3 * x**2 + 2) = 3 * x**3 + x**2 + 2 * x + 5.
#     This way we can always make use of the symmetry of even polynomials
#   - We only need to evaluate the polynomial in d/2 points, because by using the symmetry every point can be transformed into its counterpart "on the other side":
#     If we take the polynomial P_2(x) = x**2, the value representation could be [(-1, 1), (1, 1)], where (-1, 1) was inferred from one evaluation at x = 1
# - In general, we can write the value representation for a polynomial of 2nd degree as [(-x, P_2(x)), (x, P_2(x))]
# - To lower the runtime further one idea is obvious: Recursively split the polynomials to reach a runtime of O(d * log(d))
#   - This raises a problem though: To infer all needed points from the polynomials image we always have to evaluate the polynomial in +/- pairs, e.g. x = 1 and x = -1.
#     If we try to evaluate the split polynomial (P_3_even(x**2) and P_3_odd(x**2)) we can no longer choose arbitrary values, as x**2 is always positive in the real domain
#   - The solution is to use the roots of unity as evaluation points :
#     - If we want to calculate the FFT of P_2(x) = x**2 with P_2_even(x**2) = x and P_2_odd(x**2) = 0 we choose the 4th roots of unity as evaluation points for P_2(x)
#       (4th because we need at least 3 points for the value representation of a degree 2 polynomial), and the 2nd roots of unity as evaluation points for P_2_even(x**2) and P_2_odd(x**2)
# - The last step is to combine the split value representations back into one:
#
# Calculating the coefficient representation (IFFT):


# Main Algorithms


def FFT_recursive(n: int, coefficient_repr: list):
    if n == 1:
        return coefficient_repr # In this case we already have the value representation as the function is constant

    # Split the function in even and odd parts of degree n/2
    even_coefficients = coefficient_repr[0::2]
    odd_coefficients = coefficient_repr[1::2]

    # This step is the important part: We recursively apply FFT to reach the runtime of O(n log n). The result are two value representations evaluated in the (n/2)ths roots of unity.
    even_values = FFT_recursive(int(n/2), even_coefficients)
    odd_values = FFT_recursive(int(n/2), odd_coefficients)

    w = np.exp(1j * 2 * np.pi / n) # w**r is our r-th root of the n-th roots of unity

    value_repr = [0] * n
    for i in range(int(n/2)):
        # Note that the odd values have to be multiplied by the evaluation point because we extracted one "x" from the term to get an even polynomial again
        value_repr[i] = even_values[i] + w**i * odd_values[i] # Evaluate the even part of our polynomial (symmetry regarding x-axis => +, only the evaluation point is negative)
        value_repr[i + int(n/2)] = even_values[i] - w**i * odd_values[i] # Evaluate the odd part of our polynomial (symmetry regarding origin => -, the evaluation point and evaluation value are negative)

    return value_repr

def IFFT_recursive(n: int, value_repr: list):
    if n == 1:
        return value_repr

    even_values = value_repr[0::2]
    odd_values = value_repr[1::2]

    even_coefficients = IFFT_recursive(int(n/2), even_values)
    odd_coefficients = IFFT_recursive(int(n/2), odd_values)

    w = np.exp(-1j * 2 * np.pi / n) # The DFT can be inverted by just conjugating the exponants and normalizing everything later by 1/n

    coefficient_repr = [0] * n
    for i in range(int(n/2)):
        coefficient_repr[i] = even_coefficients[i] + w**i * odd_coefficients[i]
        coefficient_repr[i + int(n/2)] = even_coefficients[i] - w**i * odd_coefficients[i]

    return coefficient_repr


# Helper functions


def power_of_two(n):
    "Find the smallest power of two that is larger than n."
    power = 0
    while 2**power < n:
        power += 1
    return power

def pad_coefficients(coefficient_repr):
    "Transform a coefficient representation to a representation of higher degree (FFT handles polynomials with degrees of powers of two)."
    power = power_of_two(len(coefficient_repr))
    for _ in range(2**power - len(coefficient_repr)):
        coefficient_repr += [(0+0j)]
    return coefficient_repr

# FFT basically converts from the frequency domain to the time/signal domain
# The coefficient representation
def FFT(coefficient_repr):
    "Convert a polynomial from its coefficient representation to its value representation."
    n = len(coefficient_repr)
    return FFT_recursive(n, coefficient_repr)

# FFT basically converts from the time/signal domain to the frequency domain
# The value representation is the signal: A lot of values at different points in time
def IFFT(value_repr):
    "Convert a polynomial from its coefficient representation to its value representation."
    n = len(value_repr)
    unnormalized = IFFT_recursive(n, value_repr)

    # The second part of the DFT matrix inversion: Normalize by 1/n
    normalized = [(0+0j)] * n
    for i in range(n):
        normalized[i] = unnormalized[i] / n

    return normalized


# FFT:
print("FFT:")
coefficient_repr = pad_coefficients([(0+0j), (0+0j), (1+0j), (0+0j)]) # P_2(x) = x**2
value_repr = FFT(coefficient_repr)
print("Coefficient Representation:", coefficient_repr)
print("Value Representation:", value_repr)

# IFFT:
print("IFFT:")
coefficient_repr = IFFT(value_repr)
print("Value Representation:", value_repr)
print("Coefficient Representation:", coefficient_repr)