-- FFT Implementation in PostgreSQL
-- This implementation follows the detailed plan to compute the FFT of a set of input points using SQL.
-- It creates the input table, reorders the data using bit reversal, and then applies the butterfly computations in a loop.
-- Finally, it stores the FFT output in the FinalResults table.
-- Note: This implementation assumes that the number of input points (N) is a power of 2.

-------------------------
-- 1. Data Preparation --
-------------------------
DROP TABLE IF EXISTS InputData;
CREATE TABLE InputData (
    id INTEGER PRIMARY KEY,
    real DOUBLE PRECISION,
    imaginary DOUBLE PRECISION
);

-- Insert sample data for 8 points (a power-of-2 data set)
INSERT INTO InputData (id, real, imaginary) VALUES
    (0, 1.0, 0.0),
    (1, 1.0, 0.0),
    (2, 1.0, 0.0),
    (3, 1.0, 0.0),
    (4, 1.0, 0.0),
    (5, 1.0, 0.0),
    (6, 1.0, 0.0),
    (7, 1.0, 0.0);

DROP TABLE IF EXISTS FinalResults;
CREATE TABLE FinalResults (
    id INTEGER PRIMARY KEY,
    real DOUBLE PRECISION,
    imaginary DOUBLE PRECISION
);

---------------------------------------
-- 2. Bit Reversal Implementation  --
---------------------------------------
-- The following UDF computes the bit-reversed index for a given integer 'x' over 'bits' bits.
CREATE OR REPLACE FUNCTION bit_reverse(x INTEGER, bits INTEGER) RETURNS INTEGER AS $$
DECLARE
    result INTEGER := 0;
    i INTEGER;
BEGIN
    FOR i IN 0..(bits - 1) LOOP
        result := (result << 1) | ((x >> i) & 1);
    END LOOP;
    RETURN result;
END;
$$ LANGUAGE plpgsql IMMUTABLE STRICT;

------------------------------------------------
-- 3. Complex Arithmetic Functions (UDFs)   --
------------------------------------------------
-- These functions implement complex addition and multiplication.
DROP FUNCTION IF EXISTS complex_add(DOUBLE PRECISION, DOUBLE PRECISION, DOUBLE PRECISION, DOUBLE PRECISION);
DROP FUNCTION IF EXISTS complex_multiply(DOUBLE PRECISION, DOUBLE PRECISION, DOUBLE PRECISION, DOUBLE PRECISION);

-- Complex addition: (r1 + i1*i) + (r2 + i2*i)
CREATE OR REPLACE FUNCTION complex_add(r1 DOUBLE PRECISION, i1 DOUBLE PRECISION, r2 DOUBLE PRECISION, i2 DOUBLE PRECISION)
RETURNS TABLE(comp_real DOUBLE PRECISION, comp_imaginary DOUBLE PRECISION) AS $$
BEGIN
    RETURN QUERY SELECT r1 + r2, i1 + i2;
END;
$$ LANGUAGE plpgsql;

-- Complex multiplication: (r1 + i1*i) * (r2 + i2*i)
CREATE OR REPLACE FUNCTION complex_multiply(r1 DOUBLE PRECISION, i1 DOUBLE PRECISION, r2 DOUBLE PRECISION, i2 DOUBLE PRECISION)
RETURNS TABLE(comp_real DOUBLE PRECISION, comp_imaginary DOUBLE PRECISION) AS $$
BEGIN
    RETURN QUERY SELECT r1 * r2 - i1 * i2, r1 * i2 + i1 * r2;
END;
$$ LANGUAGE plpgsql;

----------------------------------------------------
-- 4. Complete FFT Computation in a Stored Procedure
----------------------------------------------------
-- This procedure implements the full FFT algorithm.
-- It first reorders the input data in bit-reversed order, then performs iterative butterfly computations.
-- Finally, the computed FFT output is stored in the FinalResults table.
DROP PROCEDURE IF EXISTS perform_fft();
CREATE OR REPLACE PROCEDURE perform_fft() LANGUAGE plpgsql AS $$
DECLARE
    N INTEGER;
    total_stages INTEGER;
    stage INTEGER;
    m INTEGER;
    half_m INTEGER;
    b INTEGER;
    i INTEGER;
    idx1 INTEGER;
    idx2 INTEGER;
    r1 RECORD;
    r2 RECORD;
    twiddle_real DOUBLE PRECISION;
    twiddle_imag DOUBLE PRECISION;
    new1_real DOUBLE PRECISION;
    new1_imag DOUBLE PRECISION;
    new2_real DOUBLE PRECISION;
    new2_imag DOUBLE PRECISION;
BEGIN
    -- Determine number of input points
    SELECT count(*) INTO N FROM InputData;
    IF N = 0 THEN
       RAISE EXCEPTION 'No input data in InputData table';
    END IF;
    -- Calculate total stages: log2(N)
    total_stages := ceil(log(2, N))::integer;

    -- Create a temporary table to hold FFT intermediate results in bit-reversed order.
    DROP TABLE IF EXISTS fft_result;
    CREATE TEMP TABLE fft_result (
        pos INTEGER PRIMARY KEY,
        real DOUBLE PRECISION,
        imaginary DOUBLE PRECISION
    );

    -- Insert the input data reordered by bit-reversed index.
    INSERT INTO fft_result (pos, real, imaginary)
    SELECT bit_reverse(id, total_stages) AS pos, real, imaginary
    FROM InputData
    ORDER BY pos;

    -- Iteratively perform the butterfly computations for each stage.
    FOR stage IN 1..total_stages LOOP
        m := CAST(power(2, stage) AS INTEGER);
        half_m := m / 2;
        -- Process each block of size m.
        FOR b IN 0..(N - 1) BY m LOOP
            -- For each pair within the block
            FOR i IN 0..(half_m - 1) LOOP
                idx1 := b + i;
                idx2 := b + i + half_m;
                -- Retrieve the two elements to be combined
                SELECT real, imaginary INTO r1 FROM fft_result WHERE pos = idx1;
                SELECT real, imaginary INTO r2 FROM fft_result WHERE pos = idx2;
                -- Compute the twiddle factor: exp(-2*pi*i/m)
                twiddle_real := cos(2 * pi() * i / m);
                twiddle_imag := - sin(2 * pi() * i / m);
                -- Compute the butterfly outputs:
                -- new1 = r1 + twiddle * r2
                new1_real := r1.real + (twiddle_real * r2.real - twiddle_imag * r2.imaginary);
                new1_imag := r1.imaginary + (twiddle_real * r2.imaginary + twiddle_imag * r2.real);
                -- new2 = r1 - twiddle * r2
                new2_real := r1.real - (twiddle_real * r2.real - twiddle_imag * r2.imaginary);
                new2_imag := r1.imaginary - (twiddle_real * r2.imaginary + twiddle_imag * r2.real);
                -- Update the fft_result table with the computed values.
                UPDATE fft_result SET real = new1_real, imaginary = new1_imag WHERE pos = idx1;
                UPDATE fft_result SET real = new2_real, imaginary = new2_imag WHERE pos = idx2;
            END LOOP;
        END LOOP;
    END LOOP;

    -- Store the final FFT result into the FinalResults table.
    DELETE FROM FinalResults;
    INSERT INTO FinalResults (id, real, imaginary)
      SELECT pos, real, imaginary FROM fft_result ORDER BY pos;
END;
$$;

---------------------------
-- 5. Testing and Validation
---------------------------
-- Call the FFT procedure to compute the FFT on InputData.
CALL perform_fft();

-- Display the final FFT results.
SELECT * FROM FinalResults;

-- Additional tests for the helper functions:
-- Test Bit Reversal function.
SELECT bit_reverse(1, total_stages) AS test_bit_reverse_1,
       bit_reverse(3, total_stages) AS test_bit_reverse_3,
       bit_reverse(7, total_stages) AS test_bit_reverse_7
FROM (SELECT ceil(log(2, count(*)))::integer AS total_stages FROM InputData) t;

-- Test Complex Addition function with explicit casts.
SELECT * FROM complex_add(1.0::double precision, 2.0::double precision, 3.0::double precision, 4.0::double precision);

-- Test Complex Multiplication function with explicit casts.
SELECT * FROM complex_multiply(1.0::double precision, 2.0::double precision, 3.0::double precision, 4.0::double precision);