//
// Created by cflin on 4/7/23.
//

#ifndef STATICSIMGUI_H
#define STATICSIMGUI_H

#include <memory> // shared_ptr

#include <spdlog/spdlog.h>

#include <igl/Timer.h>
#include <igl/opengl/glfw/Viewer.h>
#include <igl/opengl/glfw/imgui/ImGuiPlugin.h>
#include <igl/opengl/glfw/imgui/ImGuiMenu.h>
#include <igl/opengl/glfw/imgui/ImGuiHelpers.h>
#include <iostream>
#include <igl/jet.h>
#include <igl/colormap.h>
#include "../static_sim/StaticSim.h"
#include "StaticMenu.h"

namespace ssim {

    struct GUICtrl {
        bool is_initialized = false;
        bool is_loaded_json = false;
        bool is_modified = false;
        bool is_solved = false;
        bool is_visible_BC = false;
        bool mesh_visible = true;
        bool single_color = true;
        ssim::SimTargetOption::Target target_to_evaluate = ssim::SimTargetOption::U_NORM;

        void set_target(ssim::SimTargetOption::Target target_to_set) {
            target_to_evaluate = target_to_set;
            is_modified = true;
            single_color = false;
        }

    };

    class StaticSimGUI {
    public:
        StaticSimGUI() : gui_ctrl_(), sp_StaticSim_(nullptr) {
            plugin_.widgets.push_back(&menu_);
            viewer_.plugins.push_back(&plugin_);
            init_screen();
        }

        void viewer_launch();

    protected:

        void init_screen() {
            viewer_.core().background_color << 0.9f, 0.9f, 0.9f, 0.4f;
        }

        void init_mesh() {
            // erase mesh to 1
            while (viewer_.erase_mesh(viewer_.selected_data_index)) { ;
            }
            // clear last mesh
            viewer_.data(0).clear();
            // add mesh to visible boundary
            for (int i = 0; i < sp_StaticSim_->DirichletBCs.size() + sp_StaticSim_->NeumannBCs.size(); ++i) {
                viewer_.append_mesh();
                if (i >= 1)
                    viewer_.data(i).show_lines = false;
            }
        }

        void menu_draw_viewer_menu() {
            // Draw parent menu content
//            menu_.draw_viewer_menu();

            auto io_menu = [&]() {
                if (ImGui::Button("加载文件##IO", ImVec2(-1, 0))) {
#ifndef DEBUG_SSIM
                    std::string fname = igl::file_dialog_open();
                    if (fname != "") {
                        sp_StaticSim_ = std::make_shared<ssim::StaticSim>(ssim::SimTargetOption(), fname);
                        init_screen();
                        // reset gui_ctrl_
                        gui_ctrl_ = GUICtrl();
                        gui_ctrl_.is_initialized = true;
                        gui_ctrl_.is_loaded_json = true;
                    }
#else
                    load(CMAKE_SOURCE_DIR "/sim-test/rigid-test/rocker-arm/config.json");
#endif
                }
            };
            auto model_menu = [&]() {
                if (ImGui::Button("开始计算")) {
                    sp_StaticSim_->simulation();
                    gui_ctrl_.is_modified = true;
                    gui_ctrl_.is_solved = true;
                    gui_ctrl_.single_color = false;
                    gui_ctrl_.is_visible_BC = false;
                }
                if (ImGui::Button(("原始模型"))) {
                    gui_ctrl_.single_color = true;
                    gui_ctrl_.is_modified = true;
                }

                if (ImGui::Checkbox("网格可视化", &gui_ctrl_.mesh_visible)) {
                    gui_ctrl_.is_modified = true;
                }

                // --------------------------------------------------------------------
                if (ImGui::CollapsingHeader("边界条件", ImGuiTreeNodeFlags_DefaultOpen)) {

                    if (ImGui::Checkbox("边界可视化", &gui_ctrl_.is_visible_BC)) {
                        gui_ctrl_.is_modified = true;
                    }

                    if (ImGui::CollapsingHeader(
                            "Dirichlet边界", ImGuiTreeNodeFlags_DefaultOpen)) {
                        // --------------------------------------------------------------------

                        for (int i = 0; i < sp_StaticSim_->DirichletBCs.size(); ++i) {
                            // each boundary
                            // ImGui::Checkbox("受力", &m_state.m_solve_collisions);
                            auto &i_DBC = sp_StaticSim_->DirichletBCs[i];
                            std::vector<float> v_min_point = {(float) i_DBC.relMinBBox.x(),
                                                              (float) i_DBC.relMinBBox.y(),
                                                              (float) i_DBC.relMinBBox.z()};
                            std::vector<float> v_max_point = {(float) i_DBC.relMaxBBox.x(),
                                                              (float) i_DBC.relMaxBBox.y(),
                                                              (float) i_DBC.relMaxBBox.z()};
                            if (ImGui::InputFloat3(("Dirichlet最小值 " + std::to_string(i + 1)).c_str(),
                                                   v_min_point.data()) ||
                                ImGui::InputFloat3(("Dirichlet最大值 " + std::to_string(i + 1)).c_str(),
                                                   v_max_point.data())) {
                                gui_ctrl_.is_modified = true;
                                i_DBC.relMinBBox = {v_min_point[0], v_min_point[1], v_min_point[2]};
                                i_DBC.relMaxBBox = {v_max_point[0], v_max_point[1], v_max_point[2]};
                            }
                        }
                        if (gui_ctrl_.is_modified) {
                            sp_StaticSim_->updateBC();
                        }
                    }

                    if (ImGui::CollapsingHeader(
                            "Neumann边界", ImGuiTreeNodeFlags_DefaultOpen)) {
                        // --------------------------------------------------------------------

                        for (int i = 0; i < sp_StaticSim_->NeumannBCs.size(); ++i) {
                            auto &i_NBC = sp_StaticSim_->NeumannBCs[i];
                            std::vector<float> v_min_point = {(float) i_NBC.relMinBBox.x(),
                                                              (float) i_NBC.relMinBBox.y(),
                                                              (float) i_NBC.relMinBBox.z()};
                            std::vector<float> v_max_point = {(float) i_NBC.relMaxBBox.x(),
                                                              (float) i_NBC.relMaxBBox.y(),
                                                              (float) i_NBC.relMaxBBox.z()};
                            std::vector<float> v_force = {(float) i_NBC.force.x(), (float) i_NBC.force.y(),
                                                          (float) i_NBC.force.z()};
                            if (
                                    ImGui::InputFloat3(("Neumann最小值 " + std::to_string(i + 1)).c_str(),
                                                       v_min_point.data()) ||
                                    ImGui::InputFloat3(("Neumann最大值 " + std::to_string(i + 1)).c_str(),
                                                       v_max_point.data()) ||
                                    ImGui::InputFloat3(("Neumann力 " + std::to_string(i + 1)).c_str(),
                                                       v_force.data())) {
                                gui_ctrl_.is_modified = true;
                                i_NBC.relMinBBox = {v_min_point[0], v_min_point[1], v_min_point[2]};
                                i_NBC.relMaxBBox = {v_max_point[0], v_max_point[1], v_max_point[2]};
                                i_NBC.force = {v_force[0], v_force[1], v_force[2]};

                            }
                        }
                        if (gui_ctrl_.is_modified) {
                            sp_StaticSim_->updateBC();
                        }

                    }

                }

                // ImGui::HelpMarker("yes - stop playing if step had a collision.");

                // --------------------------------------------------------------------
                if (ImGui::CollapsingHeader("材料设置", ImGuiTreeNodeFlags_DefaultOpen)) {

                    ImGui::InputFloat("杨氏模量", &sp_StaticSim_->get_material_property().Youngs_Modulus, 1.0f, 1e10f,
                                      "%.3f");
                    ImGui::InputFloat("泊松比", &sp_StaticSim_->get_material_property().Poisson_ratio, 0.01f, 1.0f,
                                      "%.3f");
                    ImGui::InputFloat("密度", &sp_StaticSim_->get_material_property().density, 1.0f, 1000.0f, "%.3f");
                }
            };
            auto result_menu = [&]() {
                ImGui::TextWrapped("柔顺度 (compliance): %.3e",
                                   sp_StaticSim_->EvaluateTarget(ssim::SimTargetOption::COMPLIANCE)(0, 0));

                if (ImGui::Button("位移范数", ImVec2(-1, 0))) {
                    gui_ctrl_.set_target(ssim::SimTargetOption::Target::U_NORM);
                }
                if (ImGui::Button("位移-X", ImVec2(-1, 0))) {
                    gui_ctrl_.set_target(ssim::SimTargetOption::Target::UX);
                }
                if (ImGui::Button("位移-Y", ImVec2(-1, 0))) {
                    gui_ctrl_.set_target(ssim::SimTargetOption::Target::UY);
                }
                if (ImGui::Button("位移-Z", ImVec2(-1, 0))) {
                    gui_ctrl_.set_target(ssim::SimTargetOption::Target::UZ);
                }

                if (ImGui::Button("应力范数", ImVec2(-1, 0))) {
                    gui_ctrl_.set_target(ssim::SimTargetOption::Target::S_NORM);
                }
                if (ImGui::Button("冯氏应力", ImVec2(-1, 0))) {
                    gui_ctrl_.set_target(ssim::SimTargetOption::Target::S_VON_Mises);
                }

                if (ImGui::Button("应力-X", ImVec2(-1, 0))) {
                    gui_ctrl_.set_target(ssim::SimTargetOption::Target::SX);
                }
                if (ImGui::Button("应力-Y", ImVec2(-1, 0))) {
                    gui_ctrl_.set_target(ssim::SimTargetOption::Target::SY);
                }
                if (ImGui::Button("应力-Z", ImVec2(-1, 0))) {
                    gui_ctrl_.set_target(ssim::SimTargetOption::Target::SZ);
                }
            };


            if (ImGui::CollapsingHeader("载入模型", ImGuiTreeNodeFlags_DefaultOpen)) {
                io_menu();
            }
            if (!gui_ctrl_.is_initialized && gui_ctrl_.is_loaded_json) {
                if (ImGui::CollapsingHeader("模型设置", ImGuiTreeNodeFlags_DefaultOpen)) {
                    model_menu();
                }

                if (gui_ctrl_.is_solved) {
                    if (ImGui::CollapsingHeader("仿真结果", ImGuiTreeNodeFlags_DefaultOpen)) {
                        result_menu();
                    }
                }
            }

        }

    private:
        bool pre_draw_loop() {
            auto load_first_buffer = [&]() {
                // get V, F
                ssim::Model model = sp_StaticSim_->get_mesh();
                const Eigen::MatrixXd &V = model.V;
                const Eigen::MatrixXi &F = model.F;

                // plot
                viewer_.data(0).clear();
                viewer_.data(0).is_visible = true;
                viewer_.data(0).set_mesh(V, F);

                std::cerr << "load first buffer?" << std::endl;
                int dim = V.cols();
                viewer_.core().trackball_angle.setIdentity();
                viewer_.core().set_rotation_type(
                        dim == 2 ? igl::opengl::ViewerCore::ROTATION_TYPE_NO_ROTATION
                                 : igl::opengl::ViewerCore::
                        ROTATION_TYPE_TWO_AXIS_VALUATOR_FIXED_UP);
                viewer_.core().align_camera_center(
                        V, F);
            };
            auto redraw_scene = [&]() {
                // get V, F
                ssim::Model model = sp_StaticSim_->get_mesh();
                const Eigen::MatrixXd &V = model.V;
                const Eigen::MatrixXi &F = model.F;
                Eigen::VectorXi DBC_vertex_idx = sp_StaticSim_->DBC_vertexIdx_;
                Eigen::VectorXi NBC_vertex_idx = sp_StaticSim_->NBC_vertexIdx_;


                // plot
                viewer_.data(0).set_mesh(V, F);


                viewer_.data(0).show_lines = gui_ctrl_.mesh_visible;


                if (!gui_ctrl_.single_color) {
                    //  colormap
                    Eigen::MatrixXd target = sp_StaticSim_->EvaluateTarget(gui_ctrl_.target_to_evaluate);
                    Eigen::MatrixXd C;
                    igl::jet(target, true, C);
                    if (0 && gui_ctrl_.is_visible_BC) {
                        for (int i_ = 0; i_ < DBC_vertex_idx.size(); ++i_) {
                            C.row(DBC_vertex_idx(i_)) = Eigen::RowVector3d(0, 0, 1);
                        }
                        for (int i_ = 0; i_ < NBC_vertex_idx.size(); ++i_) {
                            C.row(NBC_vertex_idx(i_)) = Eigen::RowVector3d(1, 0, 0);
                        }
                    }
                    viewer_.data(0).set_colors(C);
                } else {
                    Eigen::MatrixXd color(V.rows(), 3);
                    for (int i_ = 0; i_ < V.rows(); ++i_) {
                        color.row(i_) = Eigen::RowVector3d(1, 1, 0);
                    }
                    if (gui_ctrl_.is_visible_BC) {
                        for (int i_ = 0; i_ < DBC_vertex_idx.size(); ++i_) {
                            color.row(DBC_vertex_idx(i_)) = Eigen::RowVector3d(0, 0, 1);
                        }
                        for (int i_ = 0; i_ < NBC_vertex_idx.size(); ++i_) {
                            color.row(NBC_vertex_idx(i_)) = Eigen::RowVector3d(1, 0, 0);
                        }
                    }
                    // single color
                    viewer_.data(0).set_colors(color);
//                    colorbar_plugin_.draw_colorbar(igl::ColorMapType::COLOR_MAP_TYPE_JET,0,1,/*background_color=*/Eigen::Vector4f(1,1,1,0));
                }
            };
            auto redraw_BC_box = [&]() {
                auto get_V_F_from_bd_box = [](const Eigen::Vector3d &v_min,
                                              const Eigen::Vector3d &v_max) -> std::tuple<Eigen::MatrixXd, Eigen::MatrixXi> {
                    Eigen::MatrixXd V(8, 3);
                    V << -1, -1, -1,
                            1, -1, -1,
                            1, 1, -1,
                            -1, 1, -1,
                            -1, -1, 1,
                            1, -1, 1,
                            1, 1, 1,
                            -1, 1, 1;
                    // Define the faces of the cube
                    static Eigen::MatrixXi F(12, 3);
                    F << 0, 2, 1,
                            0, 3, 2,
                            4, 5, 6,
                            4, 6, 7,
                            3, 6, 2,
                            3, 7, 6,
                            0, 1, 5,
                            0, 5, 4,
                            5, 1, 2,
                            5, 2, 6,
                            4, 7, 3,
                            4, 3, 0;
                    Eigen::Vector3d transform = (v_max + v_min) / 2.0;
                    Eigen::Vector3d scaled_factor = (v_max - v_min) / 2.0;
                    V = V.array().rowwise() * scaled_factor.transpose().array();
                    V = V.array().rowwise() + transform.transpose().array();
                    return {V, F};

                };
                for (int i = 0; i < sp_StaticSim_->DirichletBCs.size(); ++i) {
                    auto [V, F] = get_V_F_from_bd_box(sp_StaticSim_->DirichletBCs[i].absMinBBox,
                                                      sp_StaticSim_->DirichletBCs[i].absMaxBBox);
                    viewer_.data(i + 1).clear();
                    viewer_.data(i + 1).show_lines = false;
                    viewer_.data(i + 1).set_mesh(V, F);
                    viewer_.data(i + 1).set_colors(Eigen::RowVector4d(0, 0, 0, 0.1));
                    viewer_.data(i + 1).is_visible = gui_ctrl_.is_visible_BC;

                }
                for (int i = 0; i < sp_StaticSim_->NeumannBCs.size(); ++i) {
                    auto [V, F] = get_V_F_from_bd_box(sp_StaticSim_->NeumannBCs[i].absMinBBox,
                                                      sp_StaticSim_->NeumannBCs[i].absMaxBBox);
                    viewer_.data(i + 1 + sp_StaticSim_->DirichletBCs.size()).clear();
                    viewer_.data(i + 1 + sp_StaticSim_->DirichletBCs.size()).show_lines = false;
                    viewer_.data(i + 1 + sp_StaticSim_->DirichletBCs.size()).set_mesh(V, F);
                    viewer_.data(i + 1 + sp_StaticSim_->DirichletBCs.size()).set_colors(
                            Eigen::RowVector4d(0, 0, 0, 0.1));
                    viewer_.data(i + 1 + sp_StaticSim_->DirichletBCs.size()).is_visible = gui_ctrl_.is_visible_BC;

                }
            };

            if (gui_ctrl_.is_initialized && gui_ctrl_.is_loaded_json) {
                init_mesh();
                load_first_buffer();
                gui_ctrl_.is_initialized = false;
            }

            if (gui_ctrl_.is_modified) {
                redraw_BC_box();
                redraw_scene();
                gui_ctrl_.is_modified = false;
            }
            return false;
        }

        bool post_draw_loop() {
            return false;
        }

    private:
        igl::opengl::glfw::Viewer viewer_;
        igl::opengl::glfw::imgui::ImGuiPlugin plugin_;
        StaticMenu menu_;
        GUICtrl gui_ctrl_;
        std::shared_ptr<ssim::StaticSim> sp_StaticSim_;

    };

} // ssim


#endif //STATICSIMGUI_H