#include <iostream>
#include <thread>
#include <cstdint>
#include <complex>

#include <GLFW/glfw3.h>
#include <sndio.h>
#include <poll.h>
#include <math.h>

#include "imgui.h"
#include "imgui_impl_glfw.h"
#include "imgui_impl_opengl3.h"

#define SND_BITS 16
#define SND_PCHAN 2
#define SND_RATE 44100
#define SND_BUFLEN 4096

static std::atomic<bool> shutdown_threads;
static std::atomic<float> frequency[SND_PCHAN];
static std::atomic<float> lpf_bandwidth[SND_PCHAN];

static std::mutex lpf_hist_mutex;
static float lpf_hist[SND_PCHAN][512];

static void dsp_thread()
{
    struct sio_hdl *hdl;
    struct sio_par par;

    hdl = sio_open(SIO_DEVANY, SIO_PLAY|SIO_REC, 1);
    if(hdl == nullptr) {
        std::cerr << "failed to open sound device" << std::endl;
        return;
    }
    sio_initpar(&par);
    par.sig = 1;
    par.bits = SND_BITS;
    par.pchan = SND_PCHAN;
    par.rchan = 1;
    par.rate = SND_RATE;
    par.le = SIO_LE_NATIVE;
    par.xrun = SIO_ERROR;
    if(!sio_setpar(hdl, &par)) {
        std::cerr << "failed to set sound device parameters" << std::endl;
        sio_close(hdl);
        return;
    }
    if(!sio_getpar(hdl, &par)) {
        std::cerr << "failed to get back sound device parameters" << std::endl;
        sio_close(hdl);
        return;
    }

    if(!sio_start(hdl)) {
        std::cerr << "failed to start sound device" << std::endl;
        sio_close(hdl);
        return;
    }

    uint32_t phase_out[SND_PCHAN] = { 0 };
    int16_t buf_out[SND_BUFLEN*SND_PCHAN];
    size_t buf_out_idx = SND_BUFLEN*SND_PCHAN;

    uint32_t phase_in[SND_PCHAN] = { 0 };
    std::complex<float> lpf_y[SND_PCHAN];
    int lpf_count[SND_PCHAN] = { 0 };

    nfds_t nfds = sio_nfds(hdl);
    struct pollfd pfd[nfds];
    while(!shutdown_threads) {
        sio_pollfd(hdl, pfd, POLLOUT|POLLIN);
        poll(pfd, nfds, INFTIM);
        int revents = sio_revents(hdl, pfd);

        uint32_t ftw[SND_PCHAN];
        float lpf_k[SND_PCHAN];
        for(int i=0;i<SND_PCHAN;i++) {
            ftw[i] = frequency[i]*(float)UINT32_MAX/SND_RATE;
            lpf_k[i] = 2.0f*(float)M_PI*lpf_bandwidth[i]/SND_RATE;
        }

        if(revents & POLLOUT) {
            float scale = powf(2.0f, SND_BITS-1) - 1.0f;
            if(buf_out_idx == SND_BUFLEN*SND_PCHAN) {
                for(int i=0;i<SND_PCHAN;i++)
                    for(int j=0;j<SND_BUFLEN;j++) {
                        buf_out[2*j+i] = scale*sinf(phase_out[i]*2.0f*(float)M_PI/(float)UINT32_MAX);
                        phase_out[i] += ftw[i];  // wraps on overflow
                    }
                buf_out_idx = 0;
            }
            size_t written = sio_write(hdl, &buf_out[buf_out_idx], sizeof(buf_out) - buf_out_idx*sizeof(buf_out[0]));
            buf_out_idx += written/sizeof(buf_out[0]);
        }

        if(revents & POLLIN) {
            int16_t buf_in[SND_BUFLEN];
            size_t read = sio_read(hdl, buf_in, sizeof(buf_in));
            float scale = powf(0.5f, SND_BITS-1);
            for(int i=0;i<SND_PCHAN;i++)
                for(int j=0;j<read/sizeof(buf_in[0]);j++) {
                    std::complex<float> rotated;
                    rotated = (float)buf_in[j]*std::polar(scale, phase_in[i]*2.0f*(float)M_PI/(float)UINT32_MAX);
                    phase_in[i] += ftw[i];  // wraps on underflow

                    lpf_y[i] += (rotated - lpf_y[i])*lpf_k[i];
                    float mag = std::abs(lpf_y[i]);

                    lpf_count[i]++;
                    if(lpf_count[i] == 800) {
                        lpf_count[i] = 0;
                        std::lock_guard<std::mutex> guard(lpf_hist_mutex);
                        std::memmove(&lpf_hist[i][0], &lpf_hist[i][1], 511*sizeof(float));
                        lpf_hist[i][511] = mag;
                    }
                }
        }

        if(sio_eof(hdl)) {
            std::cerr << "sound I/O error" << std::endl;
            sio_close(hdl);
            return;
        }
    }

    sio_stop(hdl);
    sio_close(hdl);
}

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(1024, 768, "sndlock", 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);

    for(int i=0;i<SND_PCHAN;i++) {
        frequency[i] = 441.0f + 202.0f*i;
        lpf_bandwidth[i] = 10.0f;
    }

    static std::thread dsp_thread_h = std::thread(dsp_thread);
    static auto JoinDSP = []() {
        dsp_thread_h.join();
    };
    std::atexit(JoinDSP);

    shutdown_threads = false;
    static auto SetShutdown = []() {
        shutdown_threads = true;
    };
    std::atexit(SetShutdown);

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

        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("sndlock", nullptr, ImGuiWindowFlags_NoDecoration | ImGuiWindowFlags_NoResize);

        if(ImGui::CollapsingHeader("Modulation", ImGuiTreeNodeFlags_DefaultOpen)) {
            for(int i=0;i<SND_PCHAN;i++) {
                char str[64];
                sprintf(str, "Frequency ch%d", i);
                float frequency_l = frequency[i];
                ImGui::SliderFloat(str, &frequency_l, 50.0f, 8000.0f);
                frequency[i] = frequency_l;
            }
        }
        if(ImGui::CollapsingHeader("Demodulation", ImGuiTreeNodeFlags_DefaultOpen)) {
            for(int i=0;i<SND_PCHAN;i++) {
                char str[64];
                sprintf(str, "LPF bandwidth ch%d", i);
                float lpf_bandwidth_l = lpf_bandwidth[i];
                ImGui::SliderFloat(str, &lpf_bandwidth_l, 0.5f, 200.0f);
                lpf_bandwidth[i] = lpf_bandwidth_l;

                sprintf(str, "Output ch%d", i);
                {
                    std::lock_guard<std::mutex> guard(lpf_hist_mutex);
                    ImGui::PlotLines(str, lpf_hist[i], 512, 0, 0, -0.0f, 0.01f, ImVec2(0, 200.0f));
                }
            }
        }
        if(ImGui::Button("Exit"))
            exit = true;

        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;
}