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

#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"

#include "dsp_lib.hpp"
#include "fifo.hpp"
#include "kirdy.hpp"

#define SND_BITS 24
#define SND_PCHAN 2
#define SND_RCHAN 2
#define SND_RATE 192000
#define SND_BUFLEN 4096

static std::atomic<bool> shutdown_threads;

static std::atomic<double> frequency[SND_PCHAN];
static std::atomic<double> amplitude[SND_PCHAN];

#define HIST_DEPTH 512

static std::atomic<double> peak;
static std::atomic<bool> clipped;
static std::atomic<int> in_wave_trigger;
static std::mutex in_wave_mutex;
static float in_wave[HIST_DEPTH];

static std::atomic<int> multiplier[SND_PCHAN];
static std::atomic<double> lpf_bandwidth[SND_PCHAN];
static std::atomic<double> li_mag[SND_PCHAN];
static std::atomic<double> li_phase[SND_PCHAN];
static std::mutex li_hist_mutex;
static std::atomic<bool> li_hist_hold[SND_PCHAN];
static float li_hist_mag[SND_PCHAN][HIST_DEPTH];
static float li_hist_phase[SND_PCHAN][HIST_DEPTH];

enum class InWaveState { delay, wait_trigger, capturing };

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\n";
        return;
    }
    sio_initpar(&par);
    par.sig = 1;
    par.bits = SND_BITS;
    par.pchan = SND_PCHAN;
    par.rchan = SND_RCHAN;
    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\n";
        sio_close(hdl);
        return;
    }
    if(!sio_getpar(hdl, &par)) {
        std::cerr << "failed to get back sound device parameters\n";
        sio_close(hdl);
        return;
    }
    if(par.sig != 1
    || par.bits != SND_BITS
    || par.pchan != SND_PCHAN
    || par.rchan != SND_RCHAN
    || par.rate != SND_RATE
    || par.le != SIO_LE_NATIVE) {
        std::cerr << "sound device parameter mismatch\n";
        sio_close(hdl);
        return;
    }

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

    DDS dds[SND_PCHAN];
    int32_t buf_out[SND_BUFLEN*SND_PCHAN];
    size_t buf_out_offset = sizeof(buf_out);

    int32_t buf_in[SND_BUFLEN*SND_RCHAN];
    size_t buf_in_offset = sizeof(buf_in);

    FIFO<phase_t, 32> ftw_fifo[SND_PCHAN];

    int clipped_count = 0;
    double peak_running = 0.0;
    int peak_count = 0;
    phase_t last_trig_phase = 0;
    InWaveState in_wave_state = InWaveState::delay;
    int in_wave_count = 0;
    float in_wave_buf[HIST_DEPTH];

    Lockin<4> lockin[SND_PCHAN];
    for(int i=0;i<SND_PCHAN;i++)
        // input channels are averaged together to reduce uncorrelated noise
        lockin[i].set_scale(pow(0.5, SND_BITS-1)/SND_RCHAN);
    int li_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);

        for(int i=0;i<SND_PCHAN;i++) {
            lockin[i].set_multiplier(multiplier[i]);
            lockin[i].set_bandwidth(lpf_bandwidth[i]/SND_RATE);
        }

        if(revents & POLLOUT) {
            if(buf_out_offset == sizeof(buf_out)) {
                for(int i=0;i<SND_PCHAN;i++) {
                    phase_t ftw = frequency_to_ftw(frequency[i]/SND_RATE);
                    dds[i].ftw = ftw;
                    if(!ftw_fifo[i].push(ftw))
                        std::cerr << "FTW FIFO overflow\n";
                }
                for(int i=0;i<SND_PCHAN;i++) {
                    double scale = amplitude[i]*(pow(2.0, SND_BITS-1) - 1.0);
                    for(int j=0;j<SND_BUFLEN;j++)
                        buf_out[SND_PCHAN*j+i] = scale*dds[i].get();
                }
                buf_out_offset = 0;
            }
            size_t written = sio_write(hdl, (const char *)buf_out + buf_out_offset, sizeof(buf_out) - buf_out_offset);
            buf_out_offset += written;
        }

        if(revents & POLLIN) {
            size_t read = sio_read(hdl, buf_in, sizeof(buf_in));
            buf_in_offset += read;
            if(buf_in_offset >= sizeof(buf_in)) {
                for(int i=0;i<SND_PCHAN;i++) {
                    std::optional<phase_t> ftw = ftw_fifo[i].pull();
                    if(ftw.has_value())
                        lockin[i].ftw = ftw.value();
                    else
                        std::cerr << "FTW FIFO underflow\n";
                }
                buf_in_offset -= sizeof(buf_in);
            }
            for(int i=0;i<read/(SND_RCHAN*sizeof(buf_in[0]));i++) {
                double sample = 0.0;
                for(int j=0;j<SND_RCHAN;j++) {
                    double sample_d = (double)buf_in[SND_RCHAN*i+j];
                    sample += sample_d;
                    if(std::abs(sample_d) > 0.99*pow(2.0, SND_BITS-1))
                        // display the clipped indicator for about one second
                        clipped_count = SND_RATE;
                }
                if(clipped_count > 0) {
                    clipped_count--;
                    clipped = true;
                } else
                    clipped = false;

                double sample_abs = std::abs(sample);
                if(sample_abs > peak_running)
                    peak_running = sample_abs;
                peak_count++;
                if(peak_count == SND_RATE/10) {
                    peak_count = 0;
                    peak = peak_running/(pow(2.0, SND_BITS-1)*SND_RCHAN);
                    peak_running = 0.0;
                }

                phase_t trig_phase = lockin[in_wave_trigger].get_phase();
                bool trigger = trig_phase > last_trig_phase;
                last_trig_phase = trig_phase;
                switch(in_wave_state) {
                    case InWaveState::delay:
                        in_wave_count++;
                        if(in_wave_count == SND_RATE/10)
                            in_wave_state = InWaveState::wait_trigger;
                        break;
                    case InWaveState::wait_trigger:
                        if(trigger) {
                            in_wave_count = 0;
                            in_wave_state = InWaveState::capturing;
                        }
                        break;
                    case InWaveState::capturing:
                        in_wave_buf[in_wave_count++] = sample/(pow(2.0, SND_BITS-1)*SND_RCHAN);
                        if(in_wave_count == HIST_DEPTH) {
                            {
                                std::lock_guard<std::mutex> guard(in_wave_mutex);
                                std::memcpy(in_wave, in_wave_buf, sizeof(in_wave));
                            }
                            in_wave_count = 0;
                            in_wave_state = InWaveState::delay;
                        }
                        break;
                }

                for(int j=0;j<SND_PCHAN;j++) {
                    std::complex<double> lockin_out = lockin[j].update(sample);
                    double mag = std::abs(lockin_out);
                    double phase = std::arg(lockin_out);

                    li_mag[j] = mag;
                    li_phase[j] = phase;
                    if(!li_hist_hold[j]) {
                        li_count[j]++;
                        if(li_count[j] == 200) {
                            li_count[j] = 0;
                            std::lock_guard<std::mutex> guard(li_hist_mutex);
                            std::memmove(&li_hist_mag[j][0], &li_hist_mag[j][1], (HIST_DEPTH-1)*sizeof(float));
                            li_hist_mag[j][HIST_DEPTH-1] = mag;
                            std::memmove(&li_hist_phase[j][0], &li_hist_phase[j][1], (HIST_DEPTH-1)*sizeof(float));
                            li_hist_phase[j][HIST_DEPTH-1] = phase;
                        }
                    }
                }
            }
        }

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

    sio_stop(hdl);
    sio_close(hdl);
}

static const char *kirdies[SND_PCHAN][2] = {
    {"192.168.1.128", "1337"},
    {"192.168.1.126", "1337"},
};

static std::atomic<bool> servo_enable[SND_PCHAN];
static std::atomic<const char *> servo_state[SND_PCHAN];
static std::atomic<float> laser_temp[SND_PCHAN];
static std::atomic<float> init_current_cooling[SND_PCHAN];
static std::atomic<float> init_current_heating[SND_PCHAN];
static std::atomic<float> init_temp[SND_PCHAN];
static std::atomic<float> leadin_current[SND_PCHAN];
static std::atomic<float> leadin_mag[SND_PCHAN];

static void servo_thread(int channel)
{
    Clocker clocker = Clocker(std::chrono::milliseconds(100));
    Kirdy kirdy = Kirdy(kirdies[channel][0], kirdies[channel][1]);
    float temp;
    
    while(true) {
        servo_state[channel] = "DISABLED";
        while(!servo_enable[channel]) {
            clocker.tick();
            if(shutdown_threads)
                return;
            laser_temp[channel] = temp = kirdy.get_laser_temp();
        }

        servo_state[channel] = "INIT";
        if(temp > init_temp[channel]) {
            while(servo_enable[channel] && temp > init_temp[channel]) {
                kirdy.set_tec_current(init_current_cooling[channel]);
                clocker.tick();
                laser_temp[channel] = temp = kirdy.get_laser_temp();
            }
        } else {
            while(servo_enable[channel] && temp < init_temp[channel]) {
                kirdy.set_tec_current(init_current_heating[channel]);
                clocker.tick();
                laser_temp[channel] = temp = kirdy.get_laser_temp();
            }
        }
        if(!servo_enable[channel])
            continue;

        servo_state[channel] = "LEAD-IN";
        while(servo_enable[channel] && li_mag[channel] < leadin_mag[channel]) {
            kirdy.set_tec_current(leadin_current[channel]);
            clocker.tick();
            laser_temp[channel] = temp = kirdy.get_laser_temp();
        }

        servo_state[channel] = "LOCKING";
        while(servo_enable[channel]) {
            clocker.tick();
            laser_temp[channel] = temp = kirdy.get_laser_temp();
        }
    }
}

int main(int argc, char* argv[])
{
    if(!glfwInit()) {
        std::cerr << "failed to initialize GLFW\n";
        return 1;
    }
    std::atexit(glfwTerminate);
    glfwWindowHint(GLFW_CONTEXT_VERSION_MAJOR, 3);
    glfwWindowHint(GLFW_CONTEXT_VERSION_MINOR, 0);
    static GLFWwindow *window = glfwCreateWindow(1024, 1370, "Soundlocker", nullptr, nullptr);
    if(window == nullptr) {
        std::cerr << "failed to create GLFW window\n";
        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.0 + 202.0*i;
        amplitude[i] = 1.0;
        multiplier[i] = 1;
        lpf_bandwidth[i] = 10.0;
        servo_state[i] = "DISABLED";
        init_current_cooling[i] = 270.0f;
        init_current_heating[i] = -30.0f;
        init_temp[i] = 16.0f;
        leadin_current[i] = 230.0f;
        leadin_mag[i] = 0.007f;
    }

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

    static std::thread servo_thread_h[SND_PCHAN];
    for(int i=0;i<SND_PCHAN;i++)
        servo_thread_h[i] = std::thread(servo_thread, i);
    static auto JoinServo = []() {
        for(int i=0;i<SND_PCHAN;i++)
            servo_thread_h[i].join();
    };
    std::atexit(JoinServo);

    shutdown_threads = false;
    static auto SetShutdown = []() {
        for(int i=0;i<SND_PCHAN;i++)
            servo_enable[i] = false;
        shutdown_threads = true;
    };
    std::atexit(SetShutdown);

    float plot_scale[SND_PCHAN];
    for(int i=0;i<SND_PCHAN;i++)
        plot_scale[i] = 0.05f;
    while(!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("Soundlocker", nullptr, ImGuiWindowFlags_NoDecoration | ImGuiWindowFlags_NoResize);

        if(ImGui::CollapsingHeader("Modulation")) {
            for(int i=0;i<SND_PCHAN;i++) {
                char str[64];
                sprintf(str, "Channel %d", i);
                ImGui::SeparatorText(str);
                sprintf(str, "frequency##%d", i);
                float frequency_l = frequency[i];
                ImGui::SliderFloat(str, &frequency_l, 50.0f, 8000.0f);
                frequency[i] = frequency_l;
                sprintf(str, "amplitude##%d", i);
                float amplitude_l = amplitude[i];
                ImGui::SliderFloat(str, &amplitude_l, 0.0f, 1.0f);
                amplitude[i] = amplitude_l;
            }
        }
        if(ImGui::CollapsingHeader("Input monitor")) {
            ImGui::AlignTextToFramePadding();
            ImGui::Text("level:");
            ImGui::SameLine();
            ImGui::ProgressBar((float)peak, ImVec2(30.0f*ImGui::GetFontSize(), 0.0f));
            if(clipped) {
              ImGui::SameLine();
              ImGui::TextColored(ImVec4(1.0f, 0.0f, 0.0f, 1.0f), "clipped!");
            }
            ImGui::AlignTextToFramePadding();
            ImGui::Text("trigger:");
            ImGui::SameLine();
            int in_wave_trigger_l = in_wave_trigger;
            char str[64];
            for(int i=0;i<SND_PCHAN;i++) {
                sprintf(str, "ch%d", i);
                ImGui::RadioButton(str, &in_wave_trigger_l, i);
                if(i < (SND_PCHAN-1)) ImGui::SameLine();
            }
            in_wave_trigger = in_wave_trigger_l;
            {
                std::lock_guard<std::mutex> guard(in_wave_mutex);
                ImGui::PlotLines("waveform", in_wave, HIST_DEPTH, 0, 0, -1.0f, 1.0f, ImVec2(0.0f, 200.0f));
            }
        }
        if(ImGui::CollapsingHeader("Demodulation", ImGuiTreeNodeFlags_DefaultOpen)) {
            for(int i=0;i<SND_PCHAN;i++) {
                char str[64];

                sprintf(str, "Channel %d", i);
                ImGui::SeparatorText(str);
                ImGui::AlignTextToFramePadding();
                ImGui::Text("multiplier:");
                int multiplier_l = multiplier[i];
                ImGui::SameLine();
                sprintf(str, "f##%d", i);
                ImGui::RadioButton(str, &multiplier_l, 1);
                ImGui::SameLine();
                sprintf(str, "2f##%d", i);
                ImGui::RadioButton(str, &multiplier_l, 2);
                ImGui::SameLine();
                sprintf(str, "3f##%d", i);
                ImGui::RadioButton(str, &multiplier_l, 3);
                multiplier[i] = multiplier_l;
                ImGui::SameLine();
                ImGui::Text("/");
                bool hold_l = li_hist_hold[i];
                ImGui::SameLine();
                sprintf(str, "hold plot##%d", i);
                ImGui::Checkbox(str, &hold_l);
                li_hist_hold[i] = hold_l;
                sprintf(str, "LPF BW##%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, "mag scale##%d", i);
                ImGui::SliderFloat(str, &plot_scale[i], 0.01f, 0.1f);
                ImGui::PushItemWidth(ImGui::GetContentRegionAvail().x);
                {
                    std::lock_guard<std::mutex> guard(li_hist_mutex);
                    sprintf(str, "##magnitude%d", i);
                    ImGui::PlotLines(str, li_hist_mag[i], HIST_DEPTH, 0, 0, -0.0f, 0.1f*plot_scale[i], ImVec2(0.0f, 200.0f));
                    sprintf(str, "##phase%d", i);
                    ImGui::PlotLines(str, li_hist_phase[i], HIST_DEPTH, 0, 0, -M_PI, M_PI, ImVec2(0.0f, 200.0f));
                }
                ImGui::PopItemWidth();
                ImGui::Text("values: %8.5f %8.5f rad", (double)li_mag[i], (double)li_phase[i]);
            }
        }

        if(ImGui::CollapsingHeader("Laser servo", ImGuiTreeNodeFlags_DefaultOpen)) {
            for(int i=0;i<SND_PCHAN;i++) {
                char str[64];
                sprintf(str, "Channel %d", i);
                ImGui::SeparatorText(str);
                bool servo_enable_l = servo_enable[i];
                sprintf(str, "enable##%d", i);
                ImGui::Checkbox(str, &servo_enable_l);
                servo_enable[i] = servo_enable_l;
                ImGui::Text("servo state: %s", (const char *)servo_state[i]);
                ImGui::Text("laser temperature: %.4f °C", (float)laser_temp[i]);

                ImGui::AlignTextToFramePadding();
                ImGui::Text("init:   ");
                ImGui::PushItemWidth(ImGui::GetContentRegionAvail().x / 6.0f);
                ImGui::SameLine();
                sprintf(str, "cooling mA##%d", i);
                float init_current_cooling_l = init_current_cooling[i];
                ImGui::InputFloat(str, &init_current_cooling_l, 1.0f, 500.0f, "%.3f");
                init_current_cooling[i] = init_current_cooling_l;
                ImGui::SameLine();
                sprintf(str, "heating mA##%d", i);
                float init_current_heating_l = init_current_heating[i];
                ImGui::InputFloat(str, &init_current_heating_l, 1.0f, 500.0f, "%.3f");
                init_current_heating[i] = init_current_heating_l;
                ImGui::SameLine();
                sprintf(str, "temp °C##%d", i);
                float init_temp_l = init_temp[i];
                ImGui::InputFloat(str, &init_temp_l, 1.0f, 500.0f, "%.4f");
                init_temp[i] = init_temp_l;
                ImGui::PopItemWidth();
                
                ImGui::AlignTextToFramePadding();
                ImGui::Text("lead-in:");
                ImGui::PushItemWidth(ImGui::GetContentRegionAvail().x / 6.0f);
                ImGui::SameLine();
                sprintf(str, "current mA##%d", i);
                float leadin_current_l = leadin_current[i];
                ImGui::InputFloat(str, &leadin_current_l, 1.0f, 500.0f, "%.3f");
                leadin_current[i] = leadin_current_l;
                ImGui::SameLine();
                sprintf(str, "magnitude##%d", i);
                float leadin_mag_l = leadin_mag[i];
                ImGui::InputFloat(str, &leadin_mag_l, 0.0f, 0.001f, "%.6f");
                leadin_mag[i] = leadin_mag_l;
                ImGui::PopItemWidth();
            }
        }

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