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

#include <GLFW/glfw3.h>

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

#include <libbladeRF.h>

namespace net = std::experimental::net;
using namespace std::literals::chrono_literals;


void fft_mag(std::complex<float>* in, float* out, size_t len);


static const size_t block_len = 16384;
static const int wf_width = 1000;
static const int wf_height = 1200;

static std::atomic<bool> shutdown_threads;

static std::atomic<float> tec_current;

static void tec_thread()
{
    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));

    while(!shutdown_threads) {
        // FIXME: net::write seems unimplemented as of libstdc++ 13
        tec_socket.write_some(net::buffer(std::format("pwm 0 i_set {:.6f}\n", (float)tec_current)));
        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;
            });
        std::this_thread::sleep_for(100ms);
    }
}

static std::atomic<bool> poh_accept_input;
static std::binary_semaphore poh_input_ready{0};
static std::complex<float> poh_input_block[block_len];
static std::mutex wf_poh_mutex;
static unsigned int wf_poh[wf_width*wf_height];

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 void poh_thread()
{
    std::vector<float> block_mag(block_len);

    while(!shutdown_threads) {
        poh_accept_input = true;
        poh_input_ready.acquire();

        fft_mag(poh_input_block, block_mag.data(), block_len);

        // stabilize laser
        float freq_peak_local;
        freq_peak_local = 40.0f*float(distance(block_mag.begin(), max_element(block_mag.begin(), block_mag.end())))/float(block_len);
        freq_peak = freq_peak_local;
        float freq_error = freq_peak_local - freq_setpoint;
        tec_current = std::max(tec_bias+tec_p*freq_error, 0.0f);

        // update waterfall
        {
            std::lock_guard<std::mutex> guard(wf_poh_mutex);
            std::memmove(&wf_poh[wf_width], &wf_poh[0], sizeof(int)*wf_width*(wf_height - 1));
            for(int i=0;i<wf_width;i++) {
                int j = i*block_len/wf_width;
                wf_poh[i] = 0xff000000 | 0x010101*std::min(int(block_mag[j]/900.0f), 255);
            }
            wf_poh[int(freq_setpoint*wf_width)/40] = 0xff0000ff;
        }
    }
}

static std::atomic<bool> ls_accept_input;
static std::binary_semaphore ls_input_ready{0};
static std::complex<float> ls_input_block[block_len];
static std::mutex wf_ls_mutex;
static unsigned int wf_ls[wf_width*wf_height];

static void ls_thread()
{
    std::vector<float> block_mag(block_len);

    while(!shutdown_threads) {
        ls_accept_input = true;
        ls_input_ready.acquire();
        fft_mag(ls_input_block, block_mag.data(), block_len);
        {
            std::lock_guard<std::mutex> guard(wf_ls_mutex);
            std::memmove(&wf_ls[wf_width], &wf_ls[0], sizeof(int)*wf_width*(wf_height - 1));
            for(int i=0;i<wf_width;i++) {
                int j = i*block_len/wf_width;
                wf_ls[i] = 0xff000000 | 0x010101*std::min(int(block_mag[j]/900.0f), 255);
            }
        }
    }
}

static std::atomic<int> frames_processed;
static std::atomic<int> frames_dropped;

static auto brf_last_second = std::chrono::time_point<std::chrono::steady_clock>::min();
static int brf_frames_processed_cnt = 0;
static int brf_frames_dropped_cnt = 0;

static void* brf_dispatch(struct bladerf* dev, struct bladerf_stream* stream,
    struct bladerf_metadata* meta, void* samples_v, size_t num_samples, void* user_data)
{
    if(poh_accept_input && ls_accept_input) {
        std::complex<int16_t>* samples = (std::complex<int16_t>*)samples_v;
        for(size_t i=0;i<block_len;i++) {
            poh_input_block[i] = samples[2*i];
            ls_input_block[i] = samples[2*i+1];
        }
        poh_accept_input = false;
        poh_input_ready.release();
        ls_accept_input = false;
        ls_input_ready.release();
        brf_frames_processed_cnt++;
    } else
        brf_frames_dropped_cnt++;

    if((std::chrono::steady_clock::now() - brf_last_second) >= 1s) {
        frames_processed = brf_frames_processed_cnt;
        frames_dropped = brf_frames_dropped_cnt;
        brf_frames_processed_cnt = 0;
        brf_frames_dropped_cnt = 0;
        brf_last_second += 1s;
    }

    if(shutdown_threads)
        return BLADERF_STREAM_SHUTDOWN;
    else
        return samples_v;
}

static void brf_thread()
{
    struct bladerf_devinfo brf_dev_info;
    struct bladerf* brf_dev;
    struct bladerf_stream* brf_stream;
    void** brf_buffers;
    int status;

    bladerf_init_devinfo(&brf_dev_info);
    if((status = bladerf_open_with_devinfo(&brf_dev, &brf_dev_info)) != 0) {
        std::cerr << "cannot open bladeRF device: " << bladerf_strerror(status) << std::endl;
        return;
    }
    for(int i=0;i<2;i++) {
        if((status = bladerf_set_frequency(brf_dev, BLADERF_CHANNEL_RX(i), 1'655'000'000)) != 0) {
            std::cerr << "failed to set bladeRF frequency: " << bladerf_strerror(status) << std::endl;
            bladerf_close(brf_dev);
            return;
        }
        if((status = bladerf_set_sample_rate(brf_dev, BLADERF_CHANNEL_RX(i), 40'000'000, NULL)) != 0) {
            std::cerr << "failed to set bladeRF sample rate: " << bladerf_strerror(status) << std::endl;
            bladerf_close(brf_dev);
            return;
        }
        if((status = bladerf_set_bandwidth(brf_dev, BLADERF_CHANNEL_RX(i), 35'000'000, NULL)) != 0) {
            std::cerr << "failed to set bladeRF bandwidth: " << bladerf_strerror(status) << std::endl;
            bladerf_close(brf_dev);
            return;
        }
        if((status = bladerf_set_gain(brf_dev, BLADERF_CHANNEL_RX(i), 20)) != 0) {
            std::cerr << "failed to set bladeRF gain: " << bladerf_strerror(status) << std::endl;
            bladerf_close(brf_dev);
            return;
        }
    }
    if((status = bladerf_init_stream(&brf_stream, brf_dev, brf_dispatch, &brf_buffers,
            16, BLADERF_FORMAT_SC16_Q11, 2*block_len, 8, NULL)) != 0) {
        std::cerr << "failed to init bladeRF stream: " << bladerf_strerror(status) << std::endl;
        return;
    }
    if((status = bladerf_enable_module(brf_dev, BLADERF_CHANNEL_RX(0), true) != 0)) {
        std::cerr << "failed to enable bladeRF RX: " << bladerf_strerror(status) << std::endl;
        bladerf_deinit_stream(brf_stream);
        bladerf_close(brf_dev);
        return;
    }
    if((status = bladerf_enable_module(brf_dev, BLADERF_CHANNEL_RX(1), true) != 0)) {
        std::cerr << "failed to enable bladeRF RX: " << bladerf_strerror(status) << std::endl;
        bladerf_enable_module(brf_dev, BLADERF_CHANNEL_RX(0), false);
        bladerf_deinit_stream(brf_stream);
        bladerf_close(brf_dev);
        return;
    }

    brf_last_second = std::chrono::steady_clock::now();
    if((status = bladerf_stream(brf_stream, BLADERF_RX_X2)) != 0) {
        std::cerr << "bladeRF stream failed: " << bladerf_strerror(status) << std::endl;
        bladerf_enable_module(brf_dev, BLADERF_CHANNEL_RX(1), false);
        bladerf_enable_module(brf_dev, BLADERF_CHANNEL_RX(0), false);
        bladerf_deinit_stream(brf_stream);
        bladerf_close(brf_dev);
        return;
    };

    bladerf_enable_module(brf_dev, BLADERF_CHANNEL_RX(1), false);
    bladerf_enable_module(brf_dev, BLADERF_CHANNEL_RX(0), false);
    bladerf_deinit_stream(brf_stream);
    bladerf_close(brf_dev);
}

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(1500, 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 wftex[2];
    for(size_t i=0;i<2;i++) {
        glGenTextures(1, &wftex[i]);
        glBindTexture(GL_TEXTURE_2D, wftex[i]);
        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);
    }

    static std::thread brf_thread_h = std::thread(brf_thread);
    static auto JoinBRF = []() {
        brf_thread_h.join();
    };
    std::atexit(JoinBRF);

    static std::thread poh_thread_h = std::thread(poh_thread);
    static auto JoinPOH = []() {
        poh_thread_h.join();
    };
    std::atexit(JoinPOH);

    static std::thread ls_thread_h = std::thread(ls_thread);
    static auto JoinLS = []() {
        ls_thread_h.join();
    };
    std::atexit(JoinLS);

    static std::thread tec_thread_h = std::thread(tec_thread);
    static auto JoinTEC = []() {
        tec_thread_h.join();
    };
    std::atexit(JoinTEC);

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

    bool exit = false;
    bool update_wf = 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);

        if(update_wf) {
            #if defined(GL_UNPACK_ROW_LENGTH) && !defined(__EMSCRIPTEN__)
            glPixelStorei(GL_UNPACK_ROW_LENGTH, 0);
            #endif
            glBindTexture(GL_TEXTURE_2D, wftex[0]);
            {
                std::lock_guard<std::mutex> guard(wf_poh_mutex);
                glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA, wf_width, wf_height, 0, GL_RGBA, GL_UNSIGNED_BYTE, wf_poh);
            }
            glBindTexture(GL_TEXTURE_2D, wftex[1]);
            {
                std::lock_guard<std::mutex> guard(wf_ls_mutex);
                glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA, wf_width, wf_height, 0, GL_RGBA, GL_UNSIGNED_BYTE, wf_ls);
            }
        }

        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", 3, ImGuiTableFlags_SizingStretchSame | ImGuiTableFlags_Resizable);
        ImGui::TableSetupColumn("SBS laser control", 0, 280.0f);
        ImGui::TableSetupColumn("Pump/output heterodyne", 0, 1000.0f);
        ImGui::TableSetupColumn("Lineshape", 0, 1000.0f);
        ImGui::TableHeadersRow();
        ImGui::TableNextColumn();
        ImGui::Checkbox("Update waterfall plots", &update_wf);
        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("Frames processed: %d", (int)frames_processed);
        ImGui::Text("Frames dropped: %d", (int)frames_dropped);
        ImGui::TableNextColumn();
        ImGui::Image((void*)(intptr_t)wftex[0], ImVec2(wf_width, wf_height));
        ImGui::TableNextColumn();
        ImGui::Image((void*)(intptr_t)wftex[1], ImVec2(wf_width, wf_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;
}