#include <iostream>
#include <cstdlib>
#include <cstring>
#include <thread>
#include <atomic>
#include <mutex>
#include <vector>
#include <complex>
#include <chrono>

#include <GLFW/glfw3.h>

#include <imgui.h>
#include <imgui_impl_glfw.h>
#include <imgui_impl_opengl3.h>

#include <alsa/asoundlib.h>

#include <pocketfft_hdronly.h>


static std::atomic<bool> terminate_dsp;

static snd_pcm_t* pcm;

static std::mutex waterfall_data_mutex;
static int waterfall_width = 1600;
static int waterfall_height = 700;
static unsigned int waterfall_data[1600*700];

static std::atomic<int> fps;

static void dsp_thread() {
    using namespace std::literals::chrono_literals;

    size_t len = 6400;
    std::vector<short int> frames(len);

    std::vector<float> frames_f(len);
    pocketfft::shape_t shape{len};
    pocketfft::stride_t stride(shape.size());
    size_t tmp=sizeof(float);
    for (int i=shape.size()-1; i>=0; --i) {
        stride[i] = tmp;
        tmp *= shape[i];
    };
    pocketfft::shape_t axes;
    for(size_t i=0; i<shape.size(); ++i)
        axes.push_back(i);
    std::vector<std::complex<float>> frames_ft(len);
    int iterations = 0;
    auto last_second = std::chrono::steady_clock::now();
    while(!terminate_dsp) {
        int read_count = snd_pcm_readi(pcm, frames.data(), len);
        if(read_count < 0) {
            std::cerr << "read from audio interface failed: " << snd_strerror(read_count) << std::endl;
            break;
        }
        for(size_t i=0;i<len;i++)
            frames_f[i] = frames[i];
        pocketfft::r2c(shape, stride, stride, axes, true, frames_f.data(), frames_ft.data(), 1.0f);
        {
            std::lock_guard<std::mutex> guard(waterfall_data_mutex);
            std::memmove(&waterfall_data[waterfall_width], &waterfall_data[0], sizeof(int)*waterfall_width*(waterfall_height - 1));
            for(int i=0;i<waterfall_width;i++) {
                waterfall_data[i] = 0xff000000 | 0x010101*std::min(int(abs(frames_ft[i*len/(4*waterfall_width)])/10000.), 255);
            }
        }
        iterations++;
        if((std::chrono::steady_clock::now() - last_second) >= 1s) {
            fps = iterations;
            iterations = 0;
            last_second += 1s;
        }
    };
}

int main(int argc, char* argv[])
{
    if (!glfwInit()) {
        std::cerr << "failed to initialize GLFW" << std::endl;
        return 1;
    }
    std::atexit(glfwTerminate);
    glfwWindowHint(GLFW_CONTEXT_VERSION_MAJOR, 3);
    glfwWindowHint(GLFW_CONTEXT_VERSION_MINOR, 0);
    static GLFWwindow* window = glfwCreateWindow(1900, 720, "fastsa", nullptr, nullptr);
    if (window == nullptr) {
        std::cerr << "failed to create GLFW window" << std::endl;
        return 1;
    }
    static auto DestroyWindow = []() {
        glfwDestroyWindow(window);
    };
    std::atexit(DestroyWindow);
    glfwMakeContextCurrent(window);
    glfwSwapInterval(1);

    IMGUI_CHECKVERSION();
    ImGui::CreateContext();
    static auto ImGuiDestroyContext = []() {
        ImGui::DestroyContext();
    };
    std::atexit(ImGuiDestroyContext);
    ImGuiIO& io = ImGui::GetIO();
    io.ConfigFlags |= ImGuiConfigFlags_NavEnableKeyboard;
    ImGui::StyleColorsDark();
    ImGui_ImplGlfw_InitForOpenGL(window, true);
    std::atexit(ImGui_ImplGlfw_Shutdown);
    ImGui_ImplOpenGL3_Init("#version 130");
    std::atexit(ImGui_ImplOpenGL3_Shutdown);

    GLuint waterfall;
    glGenTextures(1, &waterfall);
    glBindTexture(GL_TEXTURE_2D, waterfall);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE);

    int err;
    if ((err = snd_pcm_open(&pcm, "default", SND_PCM_STREAM_CAPTURE, 0)) < 0) {
        std::cerr << "cannot open ALSA device: " << snd_strerror(err) << std::endl;
        return 1;
    }
    static auto PCMClose = []() {
        snd_pcm_close(pcm);
    };
    std::atexit(PCMClose);
    snd_pcm_hw_params_t* pcm_hw_params;
    if((err = snd_pcm_hw_params_malloc(&pcm_hw_params)) < 0) {
        std::cerr << "cannot allocate ALSA hardware parameter structure: " << snd_strerror(err) << std::endl;
        return 1;
    }
    if((err = snd_pcm_hw_params_any(pcm, pcm_hw_params)) < 0) {
        std::cerr << "cannot initialize ALSA hardware parameter structure: " << snd_strerror (err) << std::endl;
        snd_pcm_hw_params_free(pcm_hw_params);
        return 1;
    }
    if((err = snd_pcm_hw_params_set_access(pcm, pcm_hw_params, SND_PCM_ACCESS_RW_INTERLEAVED)) < 0) {
        std::cerr << "cannot set ALSA access type: " << snd_strerror (err) << std::endl;
        snd_pcm_hw_params_free(pcm_hw_params);
        return 1;
    }
    if((err = snd_pcm_hw_params_set_format(pcm, pcm_hw_params, SND_PCM_FORMAT_S16_LE)) < 0) {
        std::cerr << "cannot set sample format: " << snd_strerror (err) << std::endl;
        snd_pcm_hw_params_free(pcm_hw_params);
        return 1;
    }
    if((err = snd_pcm_hw_params_set_rate(pcm, pcm_hw_params, 192000, 0)) < 0) {
        std::cerr << "cannot set ALSA sample rate: " << snd_strerror (err) << std::endl;
        snd_pcm_hw_params_free(pcm_hw_params);
        return 1;
    }
    if((err = snd_pcm_hw_params_set_channels(pcm, pcm_hw_params, 1)) < 0) {
        std::cerr << "cannot set ALSA channel count: " << snd_strerror (err) << std::endl;
        snd_pcm_hw_params_free(pcm_hw_params);
        return 1;
    }
    if((err = snd_pcm_hw_params(pcm, pcm_hw_params)) < 0) {
        std::cerr << "cannot set ALSA parameters: " << snd_strerror (err) << std::endl;
        snd_pcm_hw_params_free(pcm_hw_params);
        return 1;
    }
    snd_pcm_hw_params_free(pcm_hw_params);
    if((err = snd_pcm_prepare(pcm)) < 0) {
        std::cerr << "cannot prepare ALSA interface: " << snd_strerror (err) << std::endl;
        return 1;
    }

    terminate_dsp = false;
    static std::thread dsp_thread_h = std::thread(dsp_thread);
    static auto TerminateDSP = []() {
        terminate_dsp = true;
        dsp_thread_h.join();
    };
    std::atexit(TerminateDSP);

    bool exit = false;
    while(!exit && !glfwWindowShouldClose(window)) {
        glfwPollEvents();
        glClearColor(0.0f, 0.0f, 0.0f, 1.0f);
        glClear(GL_COLOR_BUFFER_BIT);

        glBindTexture(GL_TEXTURE_2D, waterfall);
        #if defined(GL_UNPACK_ROW_LENGTH) && !defined(__EMSCRIPTEN__)
        glPixelStorei(GL_UNPACK_ROW_LENGTH, 0);
        #endif
        {
            std::lock_guard<std::mutex> guard(waterfall_data_mutex);
            glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA, waterfall_width, waterfall_height, 0, GL_RGBA, GL_UNSIGNED_BYTE, waterfall_data);
        }

        ImGui_ImplOpenGL3_NewFrame();
        ImGui_ImplGlfw_NewFrame();
        ImGui::NewFrame();

        ImGui::SetNextWindowPos(ImVec2(0.0f, 0.0f));
        ImGui::SetNextWindowSize(io.DisplaySize);
        ImGui::PushStyleVar(ImGuiStyleVar_WindowRounding, 0.0f);
        ImGui::Begin("fastsa", nullptr, ImGuiWindowFlags_NoDecoration | ImGuiWindowFlags_NoResize);

        ImGui::BeginTable("fastsa", 2, ImGuiTableFlags_SizingStretchSame | ImGuiTableFlags_Resizable);
        ImGui::TableSetupColumn("", 0, 280.0);
        ImGui::TableSetupColumn("", 0, 1600.0);
        ImGui::TableNextColumn();
        if(ImGui::Button("Exit"))
            exit = true;
        ImGui::Text("FPS: %d", (int)fps);
        ImGui::TableNextColumn();
        ImGui::Image((void*)(intptr_t)waterfall, ImVec2(waterfall_width, waterfall_height));
        ImGui::EndTable();

        ImGui::End();
        ImGui::PopStyleVar(1);

        ImGui::Render();
        ImGui_ImplOpenGL3_RenderDrawData(ImGui::GetDrawData());
        int display_w, display_h;
        glfwGetFramebufferSize(window, &display_w, &display_h);
        glViewport(0, 0, display_w, display_h);
        glfwSwapBuffers(window);
    }

    return 0;
}