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

#ifndef RIGIDIPC_STATICSIM_H
#define RIGIDIPC_STATICSIM_H

#include <vector>
#include <set>
#include <memory>
#include <fstream>
#include <iomanip>
#include <Eigen/Eigen>
#include <igl/readOBJ.h>
#include <spdlog/spdlog.h>
#include "SimTargetOption.h"

#ifdef LINSYSSOLVER_USE_CHOLMOD

#include "CHOLMODSolver.hpp"

#else
#include "EigenLibSolver.hpp"
#endif

#include "BoundaryConditions.hpp"
#include "Config.hpp"
#include "Utils.hpp"

#include <igl/read_triangle_mesh.h>
#include <igl/write_triangle_mesh.h>

namespace ssim {
    class StaticSimGUI;
}
namespace ssim {

    struct Model {
        Eigen::MatrixX3d V;
        Eigen::MatrixX3i F;

        Model() = default;

        Model(const Eigen::MatrixX3d &V, const Eigen::MatrixX3i &F) : V(V), F(F) {}

        int NumVertex() const {
            return V.rows();
        };
    };

    struct MaterialProperty {
        float Youngs_Modulus = 1.0f;
        float Poisson_ratio = 0.3f;
        float density = 1.0f;
    };


    class StaticSim {
        using MeshModel = Model;
        using Mesh = Model;

        friend class StaticSimGUI;

    public:

        StaticSim(const SimTargetOption &option, const std::string &jsonPath);

        ~StaticSim() {}

        void simulation();

        /**
         * update BC's absBBox
         */
        void updateBC();

        /**
         * Given query points Q, compute their displacements and stress
         * @param Q (nQ, 3), query points
         * @param QU return value, (nQ), norm of displacements
         * @param Qstress return value, (nQ, 6), stress size is (6) on each query point
         */
        void postprocess(Eigen::MatrixXd &Q,
                         Eigen::VectorXd &QU,
                         Eigen::MatrixXd &Qstress);
        /// Extract data to save directory
        /// \param save_dir  absolute address, like "/path/to/save_dir"
        /// \param v_is_extract if (v_is_extract[i]){ extract the target i; }
        /// \param extract_coord_flg
        void ExtractTarget(std::string save_dir,std::vector<bool> v_is_extract,bool extract_coord_flg){
            if(extract_coord_flg) {
                Eigen::MatrixXd V_surf(SVI.size(), 3);
                for (int svI = 0; svI < SVI.size(); ++svI) {
                    int vI = SVI[svI];
                    Eigen::Vector3d u = U.segment<3>(vI * 3);
                    V_surf.row(svI) = TV.row(vI) + u.transpose();
                }

                igl::write_triangle_mesh(save_dir + "/deformed-surf.obj", V_surf, F_surf);
                spdlog::info("save deformed surface mesh to 'deformed-surf.obj'");
            }
            for(int i=0;i<v_is_extract.size();++i){
                assert(i<SimTargetOption::Target::ENUM_SIZE);
                if(v_is_extract[i]){
                    Utils::writeMatrix(save_dir + "/" + SimTargetOption::GetTargetName(i) + ".txt",EvaluateTarget(static_cast<ssim::SimTargetOption::Target>(i)) /*map_target_to_evaluated_[i]*/);
                    spdlog::info("extract target {}",i);
                }
            }
        }

        Eigen::MatrixXd EvaluateTarget(SimTargetOption::Target target) {
            if (!option_.is_option_set(target) || map_target_to_evaluated_[target].size()==0) {
                // If no cache, update option_ and map_
                MapAppendTarget(target);
                option_.set_option(target);
            }
            // Return cache
            return map_target_to_evaluated_[target];
        }

        // return surf tri mesh of tet mesh
        Model get_mesh();

    private:

        void computeFeatures();

        void computeK();

        void solve();

        void setBC();

        void prepare_surf_result();

        void MapAppendTarget(int target);

        Eigen::VectorXi getSurfTriForBox(const Eigen::Vector3d &minBox,
                                         const Eigen::Vector3d &maxBox);

        MaterialProperty &get_material_property() {
            return material_property_;
        }


        // Return: #mesh.V.rows() x 1
        Eigen::MatrixXd EvaluateUNorm() const;

        // Return: #mesh.V.rows() x 1
        Eigen::MatrixXd EvaluateUX() const;

        // Return: #mesh.V.rows() x 1
        Eigen::MatrixXd EvaluateUY() const;

        // Return: #mesh.V.rows() x 1
        Eigen::MatrixXd EvaluateUZ() const;

        // Return: #mesh.V.rows() x 1
        Eigen::MatrixXd EvaluateSNorm() const;

        // Return: #mesh.V.rows() x 1
        Eigen::MatrixXd EvaluateSVonMises() const;

        // Return: #mesh.V.rows() x 1
        Eigen::MatrixXd EvaluateSX() const;

        // Return: #mesh.V.rows() x 1
        Eigen::MatrixXd EvaluateSY() const;

        // Return: #mesh.V.rows() x 1
        Eigen::MatrixXd EvaluateSZ() const;

        // Return: 1 x 1
        Eigen::MatrixXd EvaluateCompliance() const;

    private:
        MeshModel mesh_;
        MaterialProperty material_property_;
        SimTargetOption option_;
        std::vector<Eigen::MatrixXd> map_target_to_evaluated_;


    private:
        std::vector<DirichletBC> DirichletBCs;
        std::vector<NeumannBC> NeumannBCs;
        Eigen::Matrix<double, 6, 6> D; // constitutive matrix

        // owned data
        int nN, nEle, nDof;
        Eigen::MatrixXd TV; // vertices coordinates
        Eigen::MatrixXd TV1; // deformed vertices coordinates
        Eigen::MatrixXi TT; // vertice index of each tetrahedron
        Eigen::MatrixXi SF;

        int eleNodeNum;
        int eleDofNum;
        std::vector<Eigen::VectorXi> eDof;

        // owned features
        Eigen::VectorXd load; // load of each dof
        Eigen::VectorXd U; // dofs' displacement to be computed

        Eigen::VectorXi DBC_nI; // vertex in DBC
        Eigen::VectorXi isDBC; // 0: not in DBC, 1: in DBC

        Eigen::VectorXi SVI; // vertice indices of surface nodes
        Eigen::MatrixXi F_surf; // boundary vertice indices in surface triangles mesh
        std::unordered_map<int, int> vI2SVI;

        // indices for fast access
        std::vector<std::set<int>> vNeighbor; // records all vertices' indices adjacent to each vertice
        std::vector<std::set<std::pair<int, int>>> vFLoc;
        std::shared_ptr<LinSysSolver<Eigen::VectorXi, Eigen::VectorXd>> linSysSolver;

        // resulted data to evaluate
        double compliance_;
        Eigen::MatrixXd surf_U_;
        Eigen::MatrixXd surf_stress_;
        Eigen::VectorXd surf_vonstress_;

        Eigen::VectorXi DBC_faceIdx_;
        Eigen::VectorXi DBC_vertexIdx_;
        Eigen::VectorXi NBC_faceIdx_;
        Eigen::VectorXi NBC_vertexIdx_;

    };

} // ssim

#endif //RIGIDIPC_STATICSIM_H