#include <Window.h>

#include <utility>

Window::Window(std::unique_ptr<Renderer> renderer, std::unique_ptr<Volume> volume,
               std::unique_ptr<Volume> importanceSDF, std::unique_ptr<Camera> camera,
               std::unique_ptr<TransferFunction> transfer_function, std::vector<CUDA::PointLight> lights,
               const CUDA::OptimizationParameters& optParams, int width, int height, float4 volumeProperties) :
    renderer(std::move(renderer)), volume(std::move(volume)), importanceSDF(std::move(importanceSDF)),
    camera(std::move(camera)), transferFunction(std::move(transfer_function)) {
    width_ = width;
    height_ = height;
    this->lights = std::move(lights);

    this->currItem = (int) std::distance(Renderer::get_registry().begin(),
                                         Renderer::get_registry().find(this->renderer->getType()));

    temp_density = volumeProperties.x;
    temp_absorption = volumeProperties.y;
    temp_scattering = volumeProperties.z;
    temp_g = volumeProperties.w;

    glfwInit();
    glfwWindowHint(GLFW_CONTEXT_VERSION_MAJOR, 4);
    glfwWindowHint(GLFW_CONTEXT_VERSION_MINOR, 3);
    glfwWindowHint(GLFW_OPENGL_PROFILE, GLFW_OPENGL_CORE_PROFILE);

    window = glfwCreateWindow(width_, height_, "Volume Rendering", nullptr, nullptr);
    if (window == nullptr) {
        throw std::runtime_error("Window couldn't be created");
    }
    glfwMakeContextCurrent(window);

    if (gladLoadGLLoader(reinterpret_cast<GLADloadproc>(glfwGetProcAddress)) == 0) {
        throw std::runtime_error("Failed to initialize GLAD");
    }

    glfwSetWindowUserPointer(window, this);
    glfwSetInputMode(window, GLFW_CURSOR, GLFW_CURSOR_NORMAL);
    glfwSwapInterval(1);
    setCallbacks();

    interop.init();
    interopDenoised.init();

    initShader();

    initImGui();
    initOpenGL();

    updateTexture();

    cudaEventCreate(&start);
    cudaEventCreate(&stop);
#ifdef ENABLE_OPTIX
    denoiser = std::make_unique<Denoiser>();
#endif
    applyOptParams(optParams);

    volumeDataTex = this->volume->getTexture();
    if (this->importanceSDF != nullptr) {
        importanceSDFTex = this->importanceSDF->getTexture();
    }
    transferData = transferFunction->getData();
    submit();
}

Window::~Window() {
    camera.reset();
    volume.reset();
    transferFunction.reset();
    renderer.reset();
#ifdef ENABLE_OPTIX
    denoiser.reset();
#endif
    glDeleteTextures(1, &texture_id);
    glDeleteTextures(1, &texture_id_denoised);
    glDeleteVertexArrays(1, &VAO);
    glDeleteBuffers(1, &VBO);
    glDeleteProgram(shaderID);
    interop.destroy();
    interopDenoised.destroy();
    cudaEventDestroy(start);
    cudaEventDestroy(stop);
    ImGui_ImplOpenGL3_Shutdown();
    ImGui_ImplGlfw_Shutdown();
    ImGui::DestroyContext();
    glfwTerminate();
}

void Window::initShader() {
    const char* vShaderCode = "#version 430 core\n"
                              "layout (location = 0) in vec2 aPos;\n"
                              "layout (location = 1) in vec2 aTexCoord;\n"
                              "out vec2 TexCoord;\n"
                              "void main() {\n"
                              "   gl_Position = vec4(aPos.x, aPos.y, 0.0, 1.0);\n"
                              "   TexCoord = aTexCoord;\n"
                              "}\0";
    const char* fShaderCode = "#version 430 core\n"
                              "in vec2 TexCoord;\n"
                              "out vec4 FragColor;\n"
                              "uniform sampler2D textureSampler;\n"
                              "void main() {\n"
                              "   vec3 color = texture(textureSampler, TexCoord).rgb;\n"
                              "   FragColor = vec4(color, 1);\n"
                              "}\0";

    int success;
    char infoLog[512];
    // vertex Shader
    const unsigned int vertex = glCreateShader(GL_VERTEX_SHADER);
    glShaderSource(vertex, 1, &vShaderCode, nullptr);
    glCompileShader(vertex);
    glGetShaderiv(vertex, GL_COMPILE_STATUS, &success);
    if (success == 0) {
        glGetShaderInfoLog(vertex, 512, nullptr, infoLog);
        std::cout << "ERROR::SHADER::VERTEX::COMPILATION_FAILED\n" << infoLog << std::endl;
    }
    // fragment shader
    const unsigned int fragment = glCreateShader(GL_FRAGMENT_SHADER);
    glShaderSource(fragment, 1, &fShaderCode, nullptr);
    glCompileShader(fragment);
    glGetShaderiv(fragment, GL_COMPILE_STATUS, &success);
    if (success == 0) {
        glGetShaderInfoLog(fragment, 512, nullptr, infoLog);
        std::cout << "ERROR::SHADER::FRAGMENT::COMPILATION_FAILED\n" << infoLog << std::endl;
    }

    shaderID = glCreateProgram();
    glAttachShader(shaderID, vertex);
    glAttachShader(shaderID, fragment);
    glLinkProgram(shaderID);
    glGetProgramiv(shaderID, GL_LINK_STATUS, &success);
    if (success == 0) {
        glGetProgramInfoLog(shaderID, 512, nullptr, infoLog);
        std::cout << "ERROR::SHADER::PROGRAM::LINKING_FAILED\n" << infoLog << std::endl;
    }

    glDeleteShader(vertex);
    glDeleteShader(fragment);
}

void Window::initImGui() const {
    IMGUI_CHECKVERSION();
    ImGui::CreateContext();
    ImGuiIO& io = ImGui::GetIO();
    io.ConfigFlags |= ImGuiConfigFlags_NavEnableKeyboard;
    io.ConfigFlags |= ImGuiConfigFlags_NavEnableGamepad;
    io.ConfigFlags |= ImGuiConfigFlags_DockingEnable;
    io.ConfigFlags |= ImGuiConfigFlags_ViewportsEnable;

    ImGui_ImplGlfw_InitForOpenGL(window, true);
    ImGui_ImplOpenGL3_Init();
}

void Window::initOpenGL() {
    glDisable(GL_DEPTH_TEST);
    glClearColor(0.0, 0.0, 0.0, 1.0);

    float vertices[12] = {-1.0f, 3.0f, 0.0f, 2.0f, -1.0f, -1.0f, 0.0f, 0.0f, 3.0f, -1.0f, 2.0f, 0.0f};

    // Vertex Buffer Object (VBO) and Vertex Array Object (VAO) setup
    glGenVertexArrays(1, &VAO);
    glGenBuffers(1, &VBO);

    glBindVertexArray(VAO);

    glBindBuffer(GL_ARRAY_BUFFER, VBO);
    glBufferData(GL_ARRAY_BUFFER, sizeof(vertices), vertices, GL_STATIC_DRAW);

    // Position attribute
    glVertexAttribPointer(0, 2, GL_FLOAT, GL_FALSE, 4 * sizeof(float), nullptr);
    glEnableVertexAttribArray(0);

    // Texture coordinate attribute
    glVertexAttribPointer(1, 2, GL_FLOAT, GL_FALSE, 4 * sizeof(float), reinterpret_cast<void*>(2 * sizeof(float)));
    glEnableVertexAttribArray(1);

    glGenTextures(1, &texture_id);
    glGenTextures(1, &texture_id_denoised);
}

void Window::updateTexture() {
    interop.unregisterImage();

    glViewport(0, 0, width_, height_);

    glBindTexture(GL_TEXTURE_2D, texture_id);

    glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA32F, width_, height_, 0, GL_RGBA, GL_FLOAT, nullptr);

    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST);
    glBindTexture(GL_TEXTURE_2D, 0);

    interop.registerImage(texture_id, GL_TEXTURE_2D);

    interopDenoised.unregisterImage();
    glViewport(0, 0, width_, height_);

    glBindTexture(GL_TEXTURE_2D, texture_id_denoised);

    glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA32F, width_, height_, 0, GL_RGBA, GL_FLOAT, nullptr);

    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST);
    glBindTexture(GL_TEXTURE_2D, 0);
    interopDenoised.registerImage(texture_id_denoised, GL_TEXTURE_2D);
}

void Window::setCallbacks() const {
    glfwSetFramebufferSizeCallback(window, [](GLFWwindow* w, const int new_width, const int new_height) {
        static_cast<Window*>(glfwGetWindowUserPointer(w))->framebufferSizeCallback(new_width, new_height);
    });
    glfwSetCursorPosCallback(window, [](GLFWwindow* w, const double xpos, const double ypos) {
        static_cast<Window*>(glfwGetWindowUserPointer(w))->mouseCursorCallback(xpos, ypos);
    });
    glfwSetMouseButtonCallback(window, [](GLFWwindow* w, const int button, const int action, const int mods) {
        static_cast<Window*>(glfwGetWindowUserPointer(w))->mouseButtonCallback(button, action, mods);
    });
    glfwSetScrollCallback(window, [](GLFWwindow* w, const double xoffset, const double yoffset) {
        static_cast<Window*>(glfwGetWindowUserPointer(w))->scrollCallback(xoffset, yoffset);
    });
    glfwSetWindowIconifyCallback(window, [](GLFWwindow* w, int iconified) {
        static_cast<Window*>(glfwGetWindowUserPointer(w))->minimizeCallback(iconified);
    });
    glfwSetKeyCallback(window, [](GLFWwindow* w, int key, int scancode, int action, int mods) {
        static_cast<Window*>(glfwGetWindowUserPointer(w))->keyCallback(w, key, scancode, action, mods);
    });
}

void Window::framebufferSizeCallback(const int width, const int height) {
    if (sleep || (width_ == width && height_ == height)) {
        return;
    }
    width_ = width;
    height_ = height;
    updateTexture();
    renderer->resetRendering(width_, height_);
}

void Window::mouseCursorCallback(const double xposIn, const double yposIn) {
    if (sleep || ImGui::GetIO().WantCaptureMouse) {
        return;
    }
    if (lbutton_down) {
        const auto xpos = static_cast<float>(xposIn);
        const auto ypos = static_cast<float>(yposIn);

        if (firstMouse) {
            lastX = xpos;
            lastY = ypos;
            firstMouse = false;
        }

        camera->processMouseMovement((xpos - lastX), (ypos - lastY));

        lastX = xpos;
        lastY = ypos;
    }
}

void Window::mouseButtonCallback(const int button, const int action, int mods) {
    if (sleep || ImGui::GetIO().WantCaptureMouse) {
        return;
    }
    if (GLFW_MOUSE_BUTTON_LEFT == button) {
        if (GLFW_PRESS == action) {
            lbutton_down = true;
        } else if (GLFW_RELEASE == action) {
            lbutton_down = false;
            firstMouse = true;
            renderer->resetRendering(width_, height_);
        }
    }
}

void Window::scrollCallback(double xoffset, const double yoffset) const {
    if (sleep || ImGui::GetIO().WantCaptureMouse) {
        return;
    }

    camera->processMouseScroll(static_cast<float>(yoffset) * speed);
    renderer->resetRendering(width_, height_);
}

void Window::minimizeCallback(int iconified) {
    sleep = GLFW_TRUE == iconified;
}

void Window::keyCallback(GLFWwindow* window, int key, int scancode, int action, int mods) {
    if (sleep || ImGui::GetIO().WantTextInput) {
        return;
    }
    if (key == GLFW_KEY_ESCAPE && action == GLFW_PRESS) {
        glfwSetWindowShouldClose(window, 1);
    }
    if (key == GLFW_KEY_C && action == GLFW_PRESS) {
        camera->print();
    }
    if (key == GLFW_KEY_F && action == GLFW_PRESS) {
        camera->switchCam();
        renderer->resetRendering(width_, height_);
    }
    if (key == GLFW_KEY_LEFT_CONTROL && action == GLFW_PRESS) {
        speed = 0.1f;
    }
    if (key == GLFW_KEY_LEFT_CONTROL && action == GLFW_RELEASE) {
        speed = 1.f;
    }
    // start movement
    if (key == GLFW_KEY_W && action == GLFW_PRESS) {
        movementDir = movementDir | (1 << 0);
    }
    if (key == GLFW_KEY_S && action == GLFW_PRESS) {
        movementDir = movementDir | (1 << 1);
    }
    if (key == GLFW_KEY_A && action == GLFW_PRESS) {
        movementDir = movementDir | (1 << 2);
    }
    if (key == GLFW_KEY_D && action == GLFW_PRESS) {
        movementDir = movementDir | (1 << 3);
    }
    if (key == GLFW_KEY_SPACE && action == GLFW_PRESS) {
        movementDir = movementDir | (1 << 4);
    }
    if (key == GLFW_KEY_LEFT_SHIFT && action == GLFW_PRESS) {
        movementDir = movementDir | (1 << 5);
    }
    // end movement
    if (key == GLFW_KEY_W && action == GLFW_RELEASE) {
        movementDir = movementDir & ~(1 << 0);
    }
    if (key == GLFW_KEY_S && action == GLFW_RELEASE) {
        movementDir = movementDir & ~(1 << 1);
    }
    if (key == GLFW_KEY_A && action == GLFW_RELEASE) {
        movementDir = movementDir & ~(1 << 2);
    }
    if (key == GLFW_KEY_D && action == GLFW_RELEASE) {
        movementDir = movementDir & ~(1 << 3);
    }
    if (key == GLFW_KEY_SPACE && action == GLFW_RELEASE) {
        movementDir = movementDir & ~(1 << 4);
    }
    if (key == GLFW_KEY_LEFT_SHIFT && action == GLFW_RELEASE) {
        movementDir = movementDir & ~(1 << 5);
    }
}
void Window::updateCamera() const {
    if ((movementDir & (1 << 0)) != 0) {
        if (camera->processKeyboard(FORWARD, speed * deltaTime)) {
            renderer->resetRendering(width_, height_);
        }
    }
    if ((movementDir & (1 << 1)) != 0) {
        if (camera->processKeyboard(BACKWARD, speed * deltaTime)) {
            renderer->resetRendering(width_, height_);
        }
    }
    if ((movementDir & (1 << 2)) != 0) {
        if (camera->processKeyboard(LEFT, speed * deltaTime)) {
            renderer->resetRendering(width_, height_);
        }
    }
    if ((movementDir & (1 << 3)) != 0) {
        if (camera->processKeyboard(RIGHT, speed * deltaTime)) {
            renderer->resetRendering(width_, height_);
        }
    }
    if ((movementDir & (1 << 4)) != 0) {
        if (camera->processKeyboard(UP, speed * deltaTime)) {
            renderer->resetRendering(width_, height_);
        }
    }
    if ((movementDir & (1 << 5)) != 0) {
        if (camera->processKeyboard(DOWN, speed * deltaTime)) {
            renderer->resetRendering(width_, height_);
        }
    }
}

void Window::WindowGUI() {
    ImGui::Begin("Window Parameter");
    if (ImGui::CollapsingHeader("Renderer##001", ImGuiTreeNodeFlags_DefaultOpen)) {
        std::vector<const char*> test;
        for (auto& iter: Renderer::get_registry()) {
            test.push_back(iter.first.c_str());
        }

        if (ImGui::Combo("Renderer", &currItem, test.data(), static_cast<int>(test.size()))) {
            renderer = Renderer::get_registry().at(test[currItem])(width_, height_, 5000, 1, lights);
            RenderParameters params = {};
            params.lightSamples = -1;
            renderer->setParameters(params);
            renderer->setVolumeSize(volume->getSize());
        }
#ifdef ENABLE_OPTIX
        ImGui::Checkbox("Denoise", &denoise);
#endif
    }
    if (ImGui::CollapsingHeader("Volume Parameters", ImGuiTreeNodeFlags_DefaultOpen)) {
        ImGui::InputFloat("Density", &temp_density, 1.f);
        ImGui::SliderFloat("Absorption", &temp_absorption, 0, 1);
        ImGui::SliderFloat("Scattering", &temp_scattering, 0, 1);
        ImGui::SliderFloat("G", &temp_g, -1, 1);
        ImGui::Checkbox("Trilinear Interpolation", &temp_trilinear);
    }
    if (ImGui::CollapsingHeader("Optimization Parameters", ImGuiTreeNodeFlags_DefaultOpen)) {
        ImGui::Checkbox("Optimize View", &temp_optimizeView);
        ImGui::Checkbox("Optimize Light", &temp_optimizeLight);
        ImGui::SliderFloat("p", &temp_p, 0.001, 10);
        ImGui::SliderFloat("q", &temp_q, 0, 100);
        ImGui::SliderFloat("r", &temp_r, 0, 100);
        ImGui::SliderFloat("lambda", &temp_lambda, 0, 10);
        ImGui::DragFloatRange2("Threshold", &temp_thresholdMin, &temp_thresholdMax, 0.001f, 0.0f, 1.0f, "Min: %.3f",
                               "Max: %.3f", ImGuiSliderFlags_AlwaysClamp);
        ImGui::Checkbox("Phong Shading", &temp_phong);
        ImGui::SliderInt("shiny", &temp_shiny, 0, 100);
        ImGui::SliderFloat("scale", &temp_scale, 0, 10);
        if (importanceSDF != nullptr) {
            ImGui::Checkbox("SDF", &temp_sdf);
        }
        ImGui::Checkbox("Position (Magnetic Reconection)", &temp_position);
    }
    bool reset = false;
    if (transferFunction) {
        ImGui::SeparatorText("Transfer Function");
        if (transferFunction->TransferFunctionWindow()) {
            transferData = transferFunction->getData();
            submit();
            renderer->resetRendering(width_, height_);
        }

    } else {
        if (ImGui::Button("Reset")) {
            reset = true;
        }
    }

    if (temp_optimizeView != optimizeView || temp_optimizeLight != optimizeLight || temp_p != p || temp_q != q ||
        temp_r != r || temp_lambda != lambda || temp_thresholdMin != thresholdMin ||
        temp_thresholdMax != thresholdMax || temp_phong != phong || temp_shiny != (int) shiny || temp_scale != scale ||
        temp_sdf != sdf || temp_position != position || temp_absorption != absorption ||
        temp_scattering != scattering || temp_density != density || temp_g != g || temp_trilinear != trilinear) {
        reset = true;
    }

    if (reset) {
        submit();
        renderer->resetRendering(width_, height_);
    }

    ImGui::End();
}
void Window::submit() {
    density = temp_density;
    absorption = temp_absorption;
    scattering = temp_scattering;
    g = temp_g;

    thresholdMin = temp_thresholdMin;
    thresholdMax = temp_thresholdMax;

    p = temp_p;
    q = temp_q;
    r = temp_r;
    lambda = temp_lambda;

    if (trilinear != temp_trilinear) {
        trilinear = temp_trilinear;
        if (trilinear) {
            volume->updateFiltering(cudaFilterModeLinear);
        } else {
            volume->updateFiltering(cudaFilterModePoint);
        }
        volumeDataTex = volume->getTexture();
    }

    float sigma_t = absorption + scattering;
    volumeData = {volumeDataTex,           importanceSDFTex,     transferData,
                  volume->getMinBB(),      volume->getMaxBB(),   density,
                  absorption / sigma_t,    scattering / sigma_t, g,
                  volume->getGridSpacing()};


    optimizeView = temp_optimizeView;
    optimizeLight = temp_optimizeLight;
    phong = temp_phong;
    scale = temp_scale;
    shiny = (float) temp_shiny;
    sdf = temp_sdf;
    position = temp_position;

    optParams = {p,     q,     r,     lambda, thresholdMin, thresholdMax, optimizeView, optimizeLight,
                 phong, shiny, scale, sdf,    position};

    renderer->setVolumeSize(volume->getSize());
}

void Window::render() {
    ImGuiIO& io = ImGui::GetIO();
    while (glfwWindowShouldClose(window) == 0) {
        glfwPollEvents();
        if (sleep) {
            std::this_thread::sleep_for(std::chrono::milliseconds(5));
            continue;
        }

        const auto currentFrame = static_cast<float>(glfwGetTime());
        deltaTime = currentFrame - lastFrame;
        lastFrame = currentFrame;

        // input
        updateCamera();
        float3 U;
        float3 V;
        float3 W;
        camera->UVWFrame(U, V, W, (float) width_ / (float) height_);
        CUDA::Camera cameraData = {camera->getEye(), U, V, W};

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

        WindowGUI();
        renderer->drawGui(deltaTime, frametime);
        ImGui::Render();

        unsigned int texture = texture_id;
        // rendering
        interop.map();

        cudaEventRecord(start);
        // Render Call
        if (lbutton_down) {
            renderer->renderSimple(interop.getImage(), cameraData, volumeData, optParams);
        } else {
            renderer->render(interop.getImage(), cameraData, volumeData, optParams);
        }
#ifdef ENABLE_OPTIX
        if (denoise && denoiser != nullptr) {
            interopDenoised.map();

            denoiser->denoiseTexture(interopDenoised.getImage(), renderer->getImageData(), width_, height_);

            interopDenoised.unmap();
            texture = texture_id_denoised;
        }
#endif
        cudaEventRecord(stop);
        cudaEventSynchronize(stop);

        cudaEventElapsedTime(&frametime, start, stop);

        interop.unmap();

        renderOpengl(texture);

        ImGui_ImplOpenGL3_RenderDrawData(ImGui::GetDrawData());
        if ((io.ConfigFlags & ImGuiConfigFlags_ViewportsEnable) != 0) {
            GLFWwindow* backup_current_context = glfwGetCurrentContext();
            ImGui::UpdatePlatformWindows();
            ImGui::RenderPlatformWindowsDefault();
            glfwMakeContextCurrent(backup_current_context);
        }
        glfwSwapBuffers(window);
    }
}

void Window::renderOpengl(unsigned int texture) const {
    glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);

    glUseProgram(shaderID);

    glBindTexture(GL_TEXTURE_2D, texture);

    glBindVertexArray(VAO);
    glDrawArrays(GL_TRIANGLES, 0, 3);

    glBindTexture(GL_TEXTURE_2D, 0);
}
