/*
 * retain.c
 *
 *  Created on: Aug 23, 2026
 *      Author: ozpv
 */

#include "retain.h"
#include "arm_math.h"

__attribute__((section(".audio_buffers"))) static float input[WINDOW_SIZE];
static size_t input_position;
static size_t samples_since_frame;

__attribute__((section(".RAM_D2"))) static float window_function[WINDOW_SIZE];
__attribute__((section(".RAM_D2"))) static float normalize[HOP_SIZE];

__attribute__((section(".audio_buffers"))) static float fft_input[WINDOW_SIZE];

__attribute__((section(".audio_buffers"))) static float spectral[WINDOW_SIZE];

__attribute__((section(".RAM_D2"))) static float output[WINDOW_SIZE];

__attribute__((section(".RAM_D2"))) static float output_history[OVERLAP_SIZE];
static size_t output_position;

__attribute__((section(".RAM_D2"))) static float retain_scratch[COMPLEX_BIN_COUNT];

static arm_rfft_fast_instance_f32 fft_instance;

static inline void create_normalize_buffer(float *out, const float *function, const size_t window_size, const size_t hop_size) {
    const size_t num_overlaps = window_size / hop_size;

    for (size_t i = 0; i < hop_size; ++i) {
        float sum = 0.0f;

        for (size_t j = 0; j < num_overlaps; ++j) {
            const float w = function[i + j * hop_size];

            sum += w * w;
        }

        out[i] = sum;
    }
}

void retain_init(void) {
    arm_blackman_harris_92db_f32(window_function, WINDOW_SIZE);

    create_normalize_buffer(normalize, window_function, WINDOW_SIZE, HOP_SIZE);

    arm_rfft_fast_init_f32(&fft_instance, WINDOW_SIZE);

    for (size_t i = 0; i < WINDOW_SIZE; ++i) {
        input[i] = 0.0f;
        fft_input[i] = 0.0f;
        output[i] = 0.0f;
        spectral[i] = 0.0f;
    }

    for (size_t i = 0; i < OVERLAP_SIZE; ++i) {
        output_history[i] = 0.0f;
    }

    input_position = 0;
    samples_since_frame = 0;
    output_position = 0;
}

static inline float select_nth_largest(float *x, const size_t len, const size_t n) {
    size_t left = 0;
    size_t right = len - 1;
    size_t target = n - 1;

    while (left < right) {
        size_t mid = left + ((right - left) >> 1);
        float a = x[left];
        float b = x[mid];
        float c = x[right];
        float pivot;

        if (a < b) {
            if (b < c) {
                pivot = b;
            } else if (a < c) {
                pivot = c;
            } else {
                pivot = a;
            }
        } else {
            if (a < c) {
                pivot = a;
            } else if (b < c) {
                pivot = c;
            } else {
                pivot = b;
            }
        }

        size_t i = left;
        size_t j = right;

        while (i <= j) {
            while (x[i] > pivot) {
                ++i;
            }

            while (x[j] < pivot) {
                --j;
            }

            if (i <= j) {
                float t = x[i];
                x[i] = x[j];
                x[j] = t;

                ++i;
                --j;
            }
        }

        if (target <= j) {
            right = j;
        } else if (target >= i) {
            left = i;
        } else {
            break;
        }
    }

    return x[target];
}

static inline void process(float *spectral, const size_t n) {
    if (n == 0) {
        for (size_t bin = 1; bin < WINDOW_SIZE / 2; ++bin) {
            const size_t i = 2 * bin;

            spectral[i] = 0.0f;
            spectral[i + 1] = 0.0f;
        }

        return;
    }

    if (n >= COMPLEX_BIN_COUNT) {
        return;
    }

    for (size_t bin = 1; bin < WINDOW_SIZE / 2; ++bin) {
        const size_t i = 2 * bin;

        const float re = spectral[i];
        const float im = spectral[i + 1];

        retain_scratch[bin - 1] = re * re + im * im;
    }

    const float threshold = select_nth_largest(retain_scratch, COMPLEX_BIN_COUNT, n);

    for (size_t bin = 1; bin < WINDOW_SIZE / 2; ++bin) {
        const size_t i = 2 * bin;

        const float re = spectral[i];
        const float im = spectral[i + 1];

        const float mag2 = re * re + im * im;

        if (mag2 < threshold) {
            spectral[i] = 0.0f;
            spectral[i + 1] = 0.0f;
        }
    }
}

static inline void process_frame(void) {
    /* build a contiguous FFT frame */
    size_t src = input_position;
    for (size_t i = 0; i < WINDOW_SIZE; ++i) {
        fft_input[i] = input[src];

        ++src;

        if (src == WINDOW_SIZE) {
            src = 0;
        }
    }

    /* apply window function */
    for (size_t i = 0; i < WINDOW_SIZE; ++i) {
        fft_input[i] *= window_function[i];
    }

    /* forward FFT */
    arm_rfft_fast_f32(&fft_instance, fft_input, spectral, 0);

    process(spectral, 5);

    /* inverse FFT */
    arm_rfft_fast_f32(&fft_instance, spectral, output, 1);

    /* synthesis window because of heavy processing */
    for (size_t i = 0; i < WINDOW_SIZE; ++i) {
        output[i] *= window_function[i];
    }

    /* overlap add */
    for (size_t i = 0; i < OVERLAP_SIZE; ++i) {
        output[i] += output_history[i];
    }

    /* save future overlap */
    for (size_t i = 0; i < OVERLAP_SIZE; ++i) {
        output_history[i] = output[i + HOP_SIZE];
    }

    /* normalize output samples */
    for (size_t i = 0; i < HOP_SIZE; ++i) {
        output[i] /= normalize[i];
    }

    output_position = 0;
    samples_since_frame = 0;
}

inline void retain(
    const int32_t *samples_in,
    int32_t *samples_out,
    const size_t range_min,
    const size_t range_max,
    const enum Channel channel
) {
    for (size_t i = range_min + channel; i < range_max; i += 2) {
        input[input_position] = int24_t_to_f32(samples_in[i]);

        ++input_position;

        if (input_position == WINDOW_SIZE) {
            input_position = 0;
        }

        ++samples_since_frame;

        if (samples_since_frame == HOP_SIZE) {
            process_frame();
        }

        samples_out[i] = f32_to_int24_t(output[output_position]);

        ++output_position;

        if (output_position == HOP_SIZE) {
            output_position = 0;
        }
    }
}