//
// 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 <igl/read_triangle_mesh.h>

namespace ipc::rigid {
    class UIStaticSimState;
}
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 ipc::rigid::UIStaticSimState;

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

        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;

        void writePntVTK(const std::string &path, const Eigen::MatrixXd &V) {
            std::ofstream out(path);
            out << "# vtk DataFile Version 3.0\n"
                   "Volume Mesh\n"
                   "ASCII\n"
                   "DATASET UNSTRUCTURED_GRID" << std::endl;
            out << "POINTS " << V.rows() << " float" << std::endl;
            for (int i = 0; i < V.rows(); ++i) {
                out << std::setprecision(4) << V.row(i).x() << " " << V.row(i).y() << " " << V.row(i).z() << std::endl;
            }
            out << "CELLS " << V.rows() << " " << V.rows() * (1 + 1) << std::endl;
            for (int i = 0; i < V.rows(); ++i) {
                out << "1 " << i << std::endl;
            }
            out << "CELL_TYPES " << V.rows() << std::endl;
            for (int i = 0; i < V.rows(); ++i) {
                out << 1 << std::endl;
            }
        }

    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