#include <iostream>
#include <cstdlib>
#include <cstring>
#include <thread>
#include <atomic>
#include <mutex>
#include <vector>
#include <complex>
#include <chrono>
#include <format>
#include <experimental/net>

#include <GLFW/glfw3.h>

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

#include <libbladeRF.h>

#define POCKETFFT_CACHE_SIZE (16*1024*1024)
#define POCKETFFT_NO_MULTITHREADING  // MT doesn't seem to help at small sizes
#include <pocketfft_hdronly.h>

namespace net = std::experimental::net;

static std::atomic<bool> terminate_dsp;

struct bladerf* bladerf_dev;

static std::mutex waterfall_data_mutex;
static int waterfall_width = 1000;
static int waterfall_height = 2000;
static unsigned int waterfall_data[1000*2000];

static std::atomic<float> freq_setpoint;
static std::atomic<float> freq_peak;
static std::atomic<float> tec_bias;
static std::atomic<float> tec_p;
static std::atomic<float> tec_current;

static std::atomic<int> fps;

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

    net::io_context io_context;
    net::ip::tcp::socket tec_socket(io_context);
    tec_socket.connect(net::ip::tcp::endpoint(net::ip::make_address("192.168.1.27"), 23));

    size_t len = 16384;
    std::vector<std::complex<int16_t>> frames(len);
    std::vector<std::complex<float>> frames_f(len);
    pocketfft::shape_t shape{len};
    pocketfft::stride_t stride(1);
    stride[0] = sizeof(std::complex<float>);
    pocketfft::shape_t axes;
    axes.push_back(0);
    std::vector<std::complex<float>> frames_ft_rev(len);
    std::vector<float> frames_mag(len);

    int iterations = 0;
    auto last_second = std::chrono::steady_clock::now();
    auto last_tec = std::chrono::steady_clock::now();
    while(!terminate_dsp) {
        int status;
        if((status = bladerf_sync_rx(bladerf_dev, frames.data(), len, NULL, 0)) != 0) {
            std::cerr << "failed to receive samples from bladeRF: " << bladerf_strerror(status) << std::endl;
            break;
        }

        // FFT
        for(size_t i=0;i<len;i++)
            frames_f[i] = frames[i];
        pocketfft::c2c(shape, stride, stride, axes, pocketfft::FORWARD, frames_f.data(), frames_ft_rev.data(), 1.0f);
        // negative frequencies first
        for(size_t i=0;i<len/2;i++)
            frames_mag[i] = abs(frames_ft_rev[i + len/2]);
        for(size_t i=0;i<len/2;i++)
            frames_mag[i + len/2] = abs(frames_ft_rev[i]);

        // stabilize laser
        bool tick = false;
        if((std::chrono::steady_clock::now() - last_tec) >= 100ms) {
            float freq_peak_local;
            freq_peak_local = 40.0f*float(distance(frames_mag.begin(), max_element(frames_mag.begin(), frames_mag.end())))/float(len);
            freq_peak = freq_peak_local;
            float freq_error = freq_peak_local - freq_setpoint;
            float tec_current_local = std::max(tec_bias+tec_p*freq_error, 0.0f);
            tec_current = tec_current_local;
            // FIXME: net::write seems unimplemented as of libstdc++ 13
            tec_socket.write_some(net::buffer(std::format("pwm 0 i_set {:.6f}\n", tec_current_local)));
            std::string reply;
            net::read(tec_socket, net::dynamic_buffer(reply),
                [&reply](auto ec, auto n) -> std::size_t
                {
                    if(ec || (reply.size() > 0 && reply.compare(reply.size()-1, 1, "\n") == 0))
                        return 0;
                    else
                        return 1;
                });
            last_tec += 100ms;
            tick = true;
        }

        // update waterfall
        {
            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++) {
                int j = i*len/waterfall_width;
                waterfall_data[i] = 0xff000000 | 0x010101*std::min(int(frames_mag[j]/900.0f), 255);
            }
            waterfall_data[int(freq_setpoint*waterfall_width)/40] = 0xff0000ff;
            if(tick)
                for(int i=0;i<100;i++)
                    waterfall_data[i] = 0xffff0000;
        }

        // FPS counter
        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(1280, 2200, "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 status;
    struct bladerf_devinfo dev_info;
    bladerf_init_devinfo(&dev_info);
    if((status = bladerf_open_with_devinfo(&bladerf_dev, &dev_info)) != 0) {
        std::cerr << "cannot open bladeRF device: " << bladerf_strerror(status) << std::endl;
        return 1;
    }
    static auto BladeRFClose = []() {
        bladerf_close(bladerf_dev);
    };
    std::atexit(BladeRFClose);
    if((status = bladerf_set_frequency(bladerf_dev, BLADERF_CHANNEL_RX(0), 1'655'000'000)) != 0) {
        std::cerr << "failed to set bladeRF frequency: " << bladerf_strerror(status) << std::endl;
        return 1;
    }
    if((status = bladerf_set_sample_rate(bladerf_dev, BLADERF_CHANNEL_RX(0), 40'000'000, NULL)) != 0) {
        std::cerr << "failed to set bladeRF sample rate: " << bladerf_strerror(status) << std::endl;
        return 1;
    }
    if((status = bladerf_set_bandwidth(bladerf_dev, BLADERF_CHANNEL_RX(0), 35'000'000, NULL)) != 0) {
        std::cerr << "failed to set bladeRF bandwidth: " << bladerf_strerror(status) << std::endl;
        return 1;
    }
    if((status = bladerf_set_gain(bladerf_dev, BLADERF_CHANNEL_RX(0), 20)) != 0) {
        std::cerr << "failed to set bladeRF gain: " << bladerf_strerror(status) << std::endl;
        return 1;
    }
    if((status = bladerf_sync_config(bladerf_dev, BLADERF_RX_X1, BLADERF_FORMAT_SC16_Q11,
            16,   /* num_buffers */
            8192, /* buffer_size */
            8,    /* num_transfers */
            3500  /* timeout_ms */)) != 0) {
        std::cerr << "failed to set bladeRF sync settings: " << bladerf_strerror(status) << std::endl;
        return 1;
    }
    if((status = bladerf_enable_module(bladerf_dev, BLADERF_RX, true) != 0)) {
        std::cerr << "failed to enable bladeRF RX: " << bladerf_strerror(status) << std::endl;
        return 1;
    }
    static auto BladeRFDisableRX = []() {
        bladerf_enable_module(bladerf_dev, BLADERF_RX, false);
    };
    std::atexit(BladeRFDisableRX);

    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;
    bool update = true;
    float freq_setpoint_local = 0.0f;
    float tec_bias_local = 0.02f;
    float tec_p_local = 0.001f;
    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
        if(update) {
            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.0f);
        ImGui::TableSetupColumn("", 0, 1000.0f);
        ImGui::TableNextColumn();
        ImGui::Checkbox("Update waterfall", &update);
        ImGui::Text("Baseband peak: %.3f MHz", (float)freq_peak);
        ImGui::Text("TEC current: %.6f mA", (float)tec_current);
        ImGui::SliderFloat("Setpoint", &freq_setpoint_local, 0.0f, 40.0f);
        ImGui::SliderFloat("TEC bias", &tec_bias_local, 0.0f, 0.5f);
        ImGui::SliderFloat("TEC P", &tec_p_local, -1.0f, 1.0f);
        freq_setpoint = freq_setpoint_local;
        tec_bias = tec_bias_local;
        tec_p = 0.05f*tec_p_local;
        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;
}