#include <assert.h>
#include <complex.h>
#include <math.h>
#include <raylib.h>
#include <raymath.h>
#include <stddef.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#define STEP 1.06f
#define LOWF 1.0f
#define LERP 0.4f

typedef struct {
    float left;
    float right;
} Frame;

#define N (1<<13)
Frame g_frames[N];
size_t g_frames_count = 0;

float g_in[N];
float complex g_out1[N];
float complex g_out2[N];
float complex *g_out_front = g_out1;
float complex *g_out_back = g_out2;

void dft(float in[], float complex out[], size_t n)
{
    for (size_t f = 0; f < n; ++f) {
        out[f] = 0.0f;
        for (size_t i = 0; i < n; ++i) {
            float t = (float)i / n;
            out[f] += in[i] * cexpf(2*I*M_PI*f*t);
        }
    }
}

void fft(float in[], size_t stride, float complex out[], size_t n)
{
    assert(n > 0);

    if (n == 1) {
        out[0] = in[0];
        return;
    }

    fft(in, stride*2, out, n/2);
    fft(in + stride, stride*2, out + n/2, n/2);

    for (size_t k = 0; k < n/2; ++k) {
        float t = (float)k/n;
        float complex v = cexpf(-2*I*M_PI*t) * out[k + n/2];
        float complex e = out[k];
        out[k] = e + v;
        out[k + n/2] = e - v;
    }
}

void callback(void *buffer_data, unsigned int frames)
{
    if (frames <= N - g_frames_count) {
        memcpy(g_frames + g_frames_count, buffer_data, sizeof(Frame) * frames);
        g_frames_count += frames;
    } else if (frames <= N) {
        memmove(g_frames, g_frames + frames, sizeof(Frame) * (N - frames));
        memcpy(g_frames + (N - frames), buffer_data, sizeof(Frame) * frames);
    } else {
        memcpy(g_frames, buffer_data, sizeof(Frame) * N);
        g_frames_count = N;
    }
}

void render_wave(void)
{
    float w = GetScreenWidth();
    float h = GetScreenHeight();
    float cell_width = (float)w/g_frames_count;
    Color color = {50, 50, 50, 20};
#if 1
    for (size_t i = 0; i < N; ++i) {
        float sample = fabsf(g_frames[i].left);
        // if (sample > 0.0f) {
        //     DrawRectangle(i*cell_width, h/2 - h/2*sample, 1, h/2*sample, color);
        // } else {
        //     DrawRectangle(i*cell_width, h/2, 1, h/2*-sample, color);
        // }
        float height = h/2*sample/2;
        // color.a = 40 + 128*sample;
        color.r = 50*sample;
        color.g = 64 - color.r;
        color.b = 64 - color.r;
        DrawRectangle(i*cell_width, h/2 - height, 1, 2*height, color);
    }
#else
    size_t to = g_frames_count == 0 ? 1 : g_frames_count;
    for (size_t i = 0; i < to - 1; ++i) {
        float t1 = (float)i/g_frames_count;
        float t2 = (float)(i+1)/g_frames_count;
        // float hann1 = 0.5f - 0.5f * cosf(2*M_PI*t1);
        // float hann2 = 0.5f - 0.5f * cosf(2*M_PI*t2);
        float sample1 = g_frames[i].left/3/* * hann1*/;
        float sample2 = g_frames[i+1].left/3/* * hann2*/;
        DrawLine(i*cell_width, h/2 + h/2*sample1, (i+1)*cell_width, h/2 + h/2*sample2,
                 color);
    }
#endif
}

float amp(float complex z)
{
    float a = crealf(z);
    float b = cimagf(z);
    return logf(a*a + b*b);
}

float complex clerp(float complex start, float complex end, float amount)
{
    float complex result = start + amount*(end - start);
    return result;
}

float g_max_amp = 0.0f;
void render_fft(void)
{
    float w = GetScreenWidth();
    float h = GetScreenHeight();

    for (size_t i = 0; i < g_frames_count; ++i) {
        float t = (float)i/g_frames_count;
        float hann = 0.5f - 0.5f * cosf(2*M_PI*t);
        g_in[i] = g_frames[i].left * hann;
    }

    float complex *tmp = g_out_back;
    g_out_back = g_out_front;
    g_out_front = tmp;

    fft(g_in, 1, g_out_back, N);

    for (size_t i = 0; i < N; ++i) {
        float a = amp(g_out_back[i]);
        if (a > g_max_amp) {
            g_max_amp = a;
        }

        g_out_back[i] = clerp(g_out_front[i], g_out_back[i], LERP);
    }

    float step = STEP;
    float lowf = LOWF;
    size_t m = 0;
    for (float f = lowf; (size_t)f < N/2; f = ceilf(f*step)) {
        m += 1;
    }

    float cell_width = (float)w/m;
    float prev_t = 0.0f;
    m = 0;
    for (float f = lowf; (size_t)f < N/2; f = ceilf(f*step)) {
        float f1 = ceilf(f*step);
        float a = 0.0f;
        for (size_t q = (size_t)f; q < N/2 && q < (size_t)f1; ++q) {
            float b = amp(g_out_back[q]);
            if (b > a) {
                a = b;
            }
        }
        float t = a/g_max_amp;
        float tf = f / N/2;
        Color c = {255*t, 255*tf, 255*m, 255*t};
        // DrawRectangle(m*cell_width, h - h/2*t, cell_width, h/2*t, c);
        // RLAPI void DrawRectangle(int posX, int posY, int width, int height, Color color);
        // RLAPI void DrawLine(int startPosX, int startPosY, int endPosX, int endPosY, Color color);
        float x1 = m*cell_width;
        float y1 = h - h/2*prev_t;
        float x2 = x1 + cell_width;
        float y2 = h - h/2*t;
        // DrawLine(x1, y1, x2, y2, c);
        Vector2 v1 = {x1, h};
        Vector2 v2 = {x2, h};
        Vector2 v3 = {x1, y1};
        Vector2 v4 = {x2, y2};
        DrawTriangle(v1, v2, v4, c);
        DrawTriangle(v4, v3, v1, c);
        prev_t = t;
        m += 1;
    }
}

int main(int argc, char **argv)
{
    if (argc <= 1) {
        fprintf(stderr, "ERROR: Incorrect usage\n");
        return EXIT_FAILURE;
    }

    SetConfigFlags(FLAG_WINDOW_RESIZABLE);
    InitWindow(800, 600, "Fast Fourier Transform");
    SetTargetFPS(60);
    InitAudioDevice();

    Music music = LoadMusicStream(argv[1]);
    printf("music.frameCount = %u\n", music.frameCount);
    printf("music.stream.sampleRate = %u\n", music.stream.sampleRate);
    printf("music.stream.sampleSize = %u\n", music.stream.sampleSize);
    printf("music.stream.channels = %u\n", music.stream.channels);

    AttachAudioStreamProcessor(music.stream, callback);
    SetMusicVolume(music, 0.5);
    PlayMusicStream(music);

    bool toggle_wave = true;
    bool toggle_fft = true;
    float pitch = 1.0f;

    while (!WindowShouldClose()) {
        UpdateMusicStream(music);

        if (IsKeyPressed(KEY_SPACE)) {
            if (IsMusicStreamPlaying(music)) {
                PauseMusicStream(music);
            } else {
                ResumeMusicStream(music);
            }
        }
        if (IsKeyPressed(KEY_UP)) {
            pitch += 0.05f;
            SetMusicPitch(music, pitch);
        }
        if (IsKeyPressed(KEY_DOWN)) {
            pitch -= 0.05f;
            SetMusicPitch(music, pitch);
        }
        if (IsKeyPressed(KEY_LEFT)) {
            SeekMusicStream(music, GetMusicTimePlayed(music) - 10.0f);
        }
        if (IsKeyPressed(KEY_RIGHT)) {
            SeekMusicStream(music, GetMusicTimePlayed(music) + 10.0f);
        }
        if (IsKeyPressed(KEY_W)) {
            toggle_wave = !toggle_wave;
        }
        if (IsKeyPressed(KEY_F)) {
            toggle_fft = !toggle_fft;
        }

        BeginDrawing();
        ClearBackground(BLACK);
        if (toggle_wave) {
            render_wave();
        }
        if (toggle_fft) {
            render_fft();
        }
        EndDrawing();
    }

    CloseAudioDevice();
    CloseWindow();

    return EXIT_SUCCESS;
}