#include "imgui.h"
#include "imgui_impl_glfw.h"
#include "imgui_impl_opengl2.h"
#include "GL_DRAWER.h"
#include <iostream>
#ifdef __APPLE__
#define GL_SILENCE_DEPRECATION
#endif
#include <GLFW/glfw3.h>
#include "../../ext/glm/glm.hpp"
#include "../../glm/gtc/type_ptr.hpp"
#include "../../glm/gtc/constants.hpp"
#include "../ext/argparse.hpp"

#include "LOCALSUPERHELIX.h"
#include "LSHDERIVS.h"
#include "LSHOPTIMIZER.h"
#include "HELIX_IO.h"
#include "OBJ_IO.h"

GLFWwindow* _window;
VEC3 _startLoc, _anchorStart, _lookAt, _anchorLookAt, _up, _targetTranslate, _endCOM;
ImVec4 _clearColor;

float _rx, _ry, _rz;
float _tx, _ty, _tz;

char _outfile[512];

int _samplesPerStrand;

std::vector<std::vector<VEC3>> _allStrands, _allBaseStrands;
std::vector<std::vector<superHelix>> _allLSHs, _editedLSHs;
std::vector<VEC3> _allp0s;
std::vector<VEC3> _allpEs;
std::vector<VEC3> _allTargetPoints;
std::vector<MATRIX3> _allF0s;

bool _hotStartLoaded = false;
bool _liveUpdates = false;
double _kSpring = 100.0;
VECTOR _hotStart;

void recalcTargetPoints()
{
    for(unsigned int i = 0; i < _allTargetPoints.size(); i++)
    {
        const VEC3 offsetFromCOM = _allpEs[i] - _endCOM;
        // std::cout << "offsetFromCOM: " << offsetFromCOM.transpose() << std::endl;
        const VEC3 transformedOffset = Eigen::AngleAxisd(_rx*M_PI/180.0, VEC3::UnitX()) * 
            Eigen::AngleAxisd(_ry*M_PI/180.0, VEC3::UnitY()) *
            Eigen::AngleAxisd(_rz*M_PI/180.0, VEC3::UnitZ()) * offsetFromCOM;
        // std::cout << "transformedOffset: " << transformedOffset.transpose() << std::endl;
        _allTargetPoints[i] = transformedOffset + _endCOM + _targetTranslate;
    }
}

static void glfw_error_callback(int error, const char* description)
{
    fprintf(stderr, "GLFW Error %d: %s\n", error, description);
}

void processMouseClickDrag(const ImGuiIO& io)
{
  _startLoc = _anchorStart + (_anchorLookAt - _anchorStart)*io.MouseWheel*0.01;
  if(ImGui::IsMouseDragging(1))
  {
    ImVec2 rightDrag = io.MouseDelta;
    VEC3 lookVec = _lookAt - _startLoc;
    VEC3 panLeft = _up.cross(lookVec).normalized();
    VEC3 panUp = lookVec.cross(panLeft).normalized();
    VEC3 transVec = rightDrag[0]*panLeft + rightDrag[1]*panUp;
    transVec *= 0.01;
    double goodOffset = (_anchorStart - _anchorLookAt).norm();
    _startLoc = (_anchorStart + transVec - _anchorLookAt).normalized()*goodOffset + _anchorLookAt;
  }
  if(ImGui::IsMouseDragging(2))
  {
    ImVec2 midDrag = io.MouseDelta;
    VEC3 lookVec = _lookAt - _startLoc;
    VEC3 panLeft = _up.cross(lookVec).normalized();
    VEC3 panUp = lookVec.cross(panLeft).normalized();
    VEC3 transVec = midDrag[0]*panLeft + midDrag[1]*panUp;
    transVec *= 0.01;
    _lookAt = _anchorLookAt + transVec;
    _startLoc = _anchorStart + transVec;
  }
  
  _anchorStart = _startLoc;
  _anchorLookAt = _lookAt;
}

void windowInit()
{
    glfwSetErrorCallback(glfw_error_callback);
    if (!glfwInit())
        exit(1);

    // Create window with graphics context
    float main_scale = 1.0;
    _window = glfwCreateWindow((int)(1280 * main_scale), (int)(800 * main_scale), "Local SHelix Viewer", nullptr, nullptr);
    if (_window == nullptr)
        exit(1);
    glfwMakeContextCurrent(_window);
    glfwSwapInterval(1); // Enable vsync

    IMGUI_CHECKVERSION();
    ImGui::CreateContext();

    // Setup Dear ImGui style
    ImGui::StyleColorsDark();

    // Setup scaling
    ImGuiStyle& style = ImGui::GetStyle();
    style.ScaleAllSizes(main_scale);        // Bake a fixed style scale. (until we have a solution for dynamic style scaling, changing this requires resetting Style + calling this again)

    // Setup Platform/Renderer backends
    ImGui_ImplGlfw_InitForOpenGL(_window, true);
    ImGui_ImplOpenGL2_Init();

    _startLoc = VEC3(1, 0, 0);
    _lookAt = VEC3(0, 0, 0);
    _up = VEC3(0,0,1);
    _anchorStart = _startLoc;
    _anchorLookAt = _lookAt;
}

void rendering()
{
    ImGui::Render();
    int display_w, display_h;
    glfwGetFramebufferSize(_window, &display_w, &display_h);
    glViewport(0, 0, display_w, display_h);
    glClearColor(_clearColor.x * _clearColor.w, _clearColor.y * _clearColor.w, _clearColor.z * _clearColor.w, _clearColor.w);
    glClear(GL_COLOR_BUFFER_BIT);

    glMatrixMode(GL_PROJECTION);
    glLoadIdentity();
    glm::mat4 projection = glm::perspective(glm::radians(45.0f), (float)display_w / (float)display_h, 0.1f, 100.0f);
    glLoadMatrixf(glm::value_ptr(projection));

    glMatrixMode(GL_MODELVIEW);
    glLoadIdentity();
    glm::mat4 look = glm::lookAt(glm::vec3(_startLoc[0], _startLoc[1], _startLoc[2]),
                                    glm::vec3(_lookAt[0], _lookAt[1], _lookAt[2]), 
                                    glm::vec3(_up[0], _up[1], _up[2]));
    glLoadMatrixf(glm::value_ptr(look));

    for(unsigned int i =0 ; i< _allStrands.size(); i++)
    {
        drawStrand(_allStrands[i]);
        drawStrand(_allBaseStrands[i], VEC3(1.0, 0.0, 1.0));
        drawPoint(_allTargetPoints[i], VEC3(1.0, 1.0, 1.0), 10.0);
    }
    drawAxes();

    ImGui_ImplOpenGL2_RenderDrawData(ImGui::GetDrawData());

    glfwMakeContextCurrent(_window);
    glfwSwapBuffers(_window);

}

void uiElements(const ImGuiIO& io)
{
    // little background/fps checker
    {
        ImGui::Begin("utils");                        
        ImGui::ColorEdit3("clear color", (float*)&_clearColor); // Edit 3 floats representing a color

        ImGui::Text("Application average %.3f ms/frame (%.1f FPS)", 1000.0f / io.Framerate, io.Framerate);
        ImGui::End();
    }

    // display settings
    {
        ImGui::Begin("display settings");                        
        if(ImGui::InputInt("_samplesPerStrand", &_samplesPerStrand))
        {
            for(unsigned int i = 0; i < _allStrands.size(); i++)
            {
                _allStrands[i] = sampleLocalSupHelixes(_editedLSHs[i], _allp0s[i], _allF0s[i], _samplesPerStrand);
                _allBaseStrands[i] = sampleLocalSupHelixes(_allLSHs[i], _allp0s[i], _allF0s[i], _samplesPerStrand);
            }
        }
        ImGui::InputDouble("eye0", &_anchorStart[0]);
        ImGui::InputDouble("eye1", &_anchorStart[1]);
        ImGui::InputDouble("eye2", &_anchorStart[2]);
        ImGui::InputDouble("lookat0", &_lookAt[0]);
        ImGui::InputDouble("lookat1", &_lookAt[1]);
        ImGui::InputDouble("lookat2", &_lookAt[2]);
        ImGui::InputDouble("up0", &_up[0]);
        ImGui::InputDouble("up1", &_up[1]);
        ImGui::InputDouble("up2", &_up[2]);
        ImGui::End();
    }

    // multiple display visibility test
    {
        ImGui::Begin("global affine transform of endpoints");
        bool Translated = false;
        if(ImGui::SliderFloat("targetTranslate0", &_tx, -20, 20))
        {
            _targetTranslate[0] = static_cast<double>(_tx);
            Translated = true;
        } 
        if(ImGui::SliderFloat("targetTranslate1", &_ty, -20, 20))
        {
            _targetTranslate[1] = static_cast<double>(_ty);
            Translated = true;
        } 
        if(ImGui::SliderFloat("targetTranslate2", &_tz, -20, 20))
        {
            _targetTranslate[2] = static_cast<double>(_tz);
            Translated = true;
        } 

        bool springChanged = false;
        if(ImGui::InputDouble("optimization spring constant: ", &_kSpring))
        {
            springChanged = true;
        }

        ImGui::Checkbox("enable live optimization updates", &_liveUpdates);

        if(Translated || springChanged) 
        {
            recalcTargetPoints();
            if(_liveUpdates)
            {
                adamParams ap;
                _editedLSHs = _allLSHs;
                for(unsigned int i = 0; i < _allLSHs.size(); i++)
                {
                    VECTOR optimizedDarbParams;
                    if(!_hotStartLoaded)
                        optimizedDarbParams = adamOptimizeLSHChainEnd(ap, _allLSHs[i], _allp0s[i], _allF0s[i], _allTargetPoints[i], false, _kSpring);
                    else{
                        optimizedDarbParams = adamOptimizeLSHChainEndHotStart(ap, _allLSHs[i], _allp0s[i], _allF0s[i], _allTargetPoints[i], _hotStart, false, _kSpring);
                    }
                    _hotStart = optimizedDarbParams;
                    _hotStartLoaded = true;
                    updateLSHsState(_editedLSHs[i], optimizedDarbParams);
                    _allStrands[i] = sampleLocalSupHelixes(_editedLSHs[i], _allp0s[i], _allF0s[i], _samplesPerStrand);
                }
            }
        }

        if(ImGui::Button("run optimization"))
        {
            recalcTargetPoints();
            adamParams ap;
            _editedLSHs = _allLSHs;
            for(unsigned int i = 0; i < _allLSHs.size(); i++)
            {
                    VECTOR optimizedDarbParams;
                    if(!_hotStartLoaded)
                        optimizedDarbParams = adamOptimizeLSHChainEnd(ap, _allLSHs[i], _allp0s[i], _allF0s[i], _allTargetPoints[i], false, _kSpring);
                    else{
                        optimizedDarbParams = adamOptimizeLSHChainEndHotStart(ap, _allLSHs[i], _allp0s[i], _allF0s[i], _allTargetPoints[i], _hotStart, false, _kSpring);
                    }
                    _hotStart = optimizedDarbParams;
                    _hotStartLoaded = true;
                    updateLSHsState(_editedLSHs[i], optimizedDarbParams);
                    _allStrands[i] = sampleLocalSupHelixes(_editedLSHs[i], _allp0s[i], _allF0s[i], _samplesPerStrand);
            }
        }

        ImGui::End();
    }

    // save settings
    {
        ImGui::Begin("save settings");

        ImGui::InputText("outfile", _outfile, 512);
        if(ImGui::Button("save OBJ"))
        {
            writeMultStrandOBJ(_allStrands, _outfile);
        }
        if(ImGui::Button("save local .shelix"))
        {
            writeMultLocalSupHelixStrand(_editedLSHs, _allp0s, _allF0s, _outfile);
        }

        ImGui::End();
    }

}

void mainLoop()
{
    ImGuiIO& io = ImGui::GetIO(); (void)io;                   // using ImGui for IO (including mouse drag)
    io.ConfigFlags |= ImGuiConfigFlags_NavEnableKeyboard;     // Enable Keyboard Controls

    while (!glfwWindowShouldClose(_window))
    {
        // Poll and handle events (inputs, window resize, etc.)
        // You can read the io.WantCaptureMouse, io.WantCaptureKeyboard flags to tell if dear imgui wants to use your inputs.
        // - When io.WantCaptureMouse is true, do not dispatch mouse input data to your main application, or clear/overwrite your copy of the mouse data.
        // - When io.WantCaptureKeyboard is true, do not dispatch keyboard input data to your main application, or clear/overwrite your copy of the keyboard data.
        // Generally you may always pass all inputs to dear imgui, and hide them from your application based on those two flags.
        glfwPollEvents();
        if (glfwGetWindowAttrib(_window, GLFW_ICONIFIED) != 0)
        {
            ImGui_ImplGlfw_Sleep(10);
            continue;
        }

        // Start the Dear ImGui frame
        ImGui_ImplOpenGL2_NewFrame();
        ImGui_ImplGlfw_NewFrame();
        ImGui::NewFrame();

        processMouseClickDrag(io);
        uiElements(io);

        // Rendering
        rendering();
    }
}

int main(int argc, char** argv)
{
    argparse::ArgumentParser parser("local superhelix viewer");
    parser.add_argument(".shelix input").help("The local superhelix .shelix file").metavar("SHELIX");

    try {
    parser.parse_args(argc, argv); 
    }
    catch (const std::exception& err) {
    std::cerr << err.what() << std::endl;
    std::cerr << parser;
    std::exit(1);
    }
    
    _samplesPerStrand = 200;
    readMultLocalSupHelixStrand(parser.get(".shelix input").c_str(), _allLSHs, _allp0s, _allF0s);
    _editedLSHs = _allLSHs;
    _allStrands.resize(_editedLSHs.size());
    _allBaseStrands.resize(_editedLSHs.size());
    _allTargetPoints.resize(_allStrands.size());
    _allpEs.resize(_allStrands.size());
    _endCOM = VEC3::Zero();
    for(unsigned int i = 0; i < _allStrands.size(); i++)
    {
        _allStrands[i] = sampleLocalSupHelixes(_editedLSHs[i], _allp0s[i], _allF0s[i], _samplesPerStrand);
        _allBaseStrands[i] = sampleLocalSupHelixes(_editedLSHs[i], _allp0s[i], _allF0s[i], _samplesPerStrand);
        _allpEs[i] = _allStrands[i][_allStrands[i].size() - 1];
        // std::cout << "_allpEs["<<i<<"]: " << _allpEs[i].transpose() << std::endl;
        _endCOM += _allpEs[i];
    }
    _endCOM = _endCOM / _allStrands.size();
    // std::cout << "_endCOM: " << _endCOM.transpose() << std::endl;
    _targetTranslate = VEC3::Zero();
    _rx = _ry = _rz = _tx = _ty = _tz = 0;
    recalcTargetPoints();

    windowInit();
    // Our state
    _clearColor = ImVec4(0.45f, 0.55f, 0.60f, 1.00f);

    // Main loop
    mainLoop();

    // Cleanup
    ImGui_ImplOpenGL2_Shutdown();
    ImGui_ImplGlfw_Shutdown();
    ImGui::DestroyContext();

    glfwDestroyWindow(_window);
    glfwTerminate();

    return 0;
}