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

#include "Util.h"
#include <cassert>
#include <iostream>
#include <memory>
#include <igl/readOBJ.h>
#include <ratio>
#include "Boundary.h"
#include "Mesh/Mesh.h"

namespace da::sha {
    namespace top {



        void WriteTriVTK(const std::string &path, const Eigen::MatrixXd &V, const Eigen::MatrixXi &T,
                         const std::vector<double> &cell_data, const std::vector<double> &v_data) {
            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() << " double" << std::endl;
            for (int i = 0; i < V.rows(); ++i) {
                out << std::setprecision(18) << V.row(i).x() << " " << V.row(i).y() << " " << V.row(i).z() << std::endl;
            }
            out << "CELLS " << T.rows() << " " << T.rows() * (3 + 1) << std::endl;
            for (int i = 0; i < T.rows(); ++i) {
                out << "3 " << T.row(i).x() << " " << T.row(i).y() << " " << T.row(i).z() << std::endl;
            }
            out << "CELL_TYPES " << T.rows() << std::endl;
            for (int i = 0; i < T.rows(); ++i) {
                out << 5 << std::endl;
            }

            if(!cell_data.empty()){
                out << "CELL_DATA " << cell_data.size() << "\n" <<
                    "SCALARS cell_scalars double 1\n" <<
                    "LOOKUP_TABLE default" << std::endl;
                for (auto &d: cell_data) {
                    out << d << std::endl;
                }
            }

            if(!v_data.empty()){
                out << "POINT_DATA " << v_data.size() << "\n" <<
                    "SCALARS point_scalars double 1\n" <<
                    "LOOKUP_TABLE default" << std::endl;
                for (auto &d: v_data) {
                    out << d << std::endl;
                }
            }
        }

        void WriteHexahedralMatmeshToVtk(const fs_path &file_path, const HexahedralMatMesh &matmesh,
                                         const std::vector<double> &point_data, const std::vector<double> &cell_data) {
            WriteToVtk(file_path, matmesh.mat_coordinates, matmesh.mat_hexahedrons, point_data, cell_data,
                       12);
        }

        void
        WriteToVtk(const fs_path &file_path, const Eigen::MatrixXd &mat_coordinates, const Eigen::MatrixXi &mat_cells,
                   const std::vector<double> &point_data, const std::vector<double> &cell_data, int cell_type) {
            std::ofstream vtk_out_stream(file_path.string());
            if (not vtk_out_stream.is_open()) {
                spdlog::critical("File " + file_path.string() + " can not be opened.");
                exit(-7);
            }
            vtk_out_stream << "# vtk DataFile Version 3.0\n"
                              "Volume Mesh\n"
                              "ASCII\n"
                              "DATASET UNSTRUCTURED_GRID"
                           << std::endl;
            vtk_out_stream << "POINTS " << mat_coordinates.rows() << " float" << std::endl;
            for (index_t vertex_idx = 0; vertex_idx < mat_coordinates.rows(); ++vertex_idx) {
                vtk_out_stream << mat_coordinates.row(vertex_idx).x() << " "
                               << mat_coordinates.row(vertex_idx).y() << " "
                               << mat_coordinates.row(vertex_idx).z() << std::endl;
            }
            vtk_out_stream << "CELLS " << mat_cells.rows() << " " << mat_cells.rows() * (mat_cells.cols() + 1)
                           << std::endl;
            for (index_t face_idx = 0; face_idx < mat_cells.rows(); ++face_idx) {
                vtk_out_stream << mat_cells.cols();
                for (index_t col_idx = 0; col_idx < mat_cells.cols(); ++col_idx) {
                    vtk_out_stream << " " << mat_cells(face_idx, col_idx);
                }
                vtk_out_stream << std::endl;
            }
            vtk_out_stream << "CELL_TYPES " << mat_cells.rows() << std::endl;
            for (index_t face_idx = 0; face_idx < mat_cells.rows(); ++face_idx) {
                vtk_out_stream << cell_type << std::endl;
            }

            if (!point_data.empty()) {
                vtk_out_stream << "POINT_DATA " << point_data.size() << "\n"
                               << "SCALARS point_scalars double 1\n"
                               << "LOOKUP_TABLE default" << std::endl;
                for (auto &data: point_data) {
                    vtk_out_stream << data << std::endl;
                }
            }

            if (!cell_data.empty()) {
                vtk_out_stream << "CELL_DATA " << cell_data.size() << "\n"
                               << "SCALARS cell_scalars double 1\n"
                               << "LOOKUP_TABLE default" << std::endl;

                for (auto &data: cell_data) {
                    vtk_out_stream << data << std::endl;
                }
            }
        }

        void
        WriteTensorToVtk(const boost::filesystem::path &file_path, const Tensor3d &t3, std::shared_ptr<Mesh> sp_mesh,double threshold) {
            HexahedralMatMesh matmesh;
            Eigen::Vector3d origin = sp_mesh->GetOrigin();
            double pixel_len = sp_mesh->GetPixelLen();
            matmesh.mat_coordinates =
                    (sp_mesh->GetNodeId2CoordMap().cast<double>() * pixel_len).rowwise() + origin.transpose();
            matmesh.mat_hexahedrons = sp_mesh->GetEleId2NodeIdsMap();

            std::vector<double> v_ele_data;
            Eigen::VectorXi pixel_idx = sp_mesh->GetGlobalIdxOfEleInUse();
            for (int i = 0; i < pixel_idx.size(); ++i) {
                v_ele_data.push_back(*(t3.data() + pixel_idx[i]));
            }
            assert(matmesh.IsHexahedral());
            if(threshold!=-1){
                if(threshold<0 || threshold>1){
                    throw std::runtime_error("threshold should be in [0,1]");
                }
                for(auto &d:v_ele_data){
                    if(d>=threshold){
                        d=1;
                    }else{
                        d=0;
                    }
                }
            }
            WriteHexahedralMatmeshToVtk(file_path, matmesh, std::vector<double>(), v_ele_data);
        }

        void WriteUToVtk(const boost::filesystem::path &file_path,const Eigen::VectorXd &U, std::shared_ptr<Mesh> sp_mesh) {
            assert(sp_mesh->GetNumNodes()==U.size());
            HexahedralMatMesh matmesh;
            Eigen::Vector3d origin = sp_mesh->GetOrigin();
            double pixel_len = sp_mesh->GetPixelLen();
            matmesh.mat_coordinates =
                    (sp_mesh->GetNodeId2CoordMap().cast<double>() * pixel_len).rowwise() + origin.transpose();
            matmesh.mat_hexahedrons = sp_mesh->GetEleId2NodeIdsMap();
            assert(matmesh.IsHexahedral());
            WriteHexahedralMatmeshToVtk(file_path, matmesh, std::vector<double>(U.data(),U.data()+U.size()),std::vector<double>());
        }

        std::pair<MatMesh3, Eigen::VectorXd> GetMeshVertexPropty(const boost::filesystem::path &fs_obj,
                                                                 const Tensor3d &ten_rho_field,
                                                                 Boundary &r_bdr, bool offset_flg) {
            MatMesh3 obj_mesh;
            igl::readOBJ(fs_obj.string(), obj_mesh.mat_coordinates, obj_mesh.mat_faces);
            if (offset_flg) {
                // for box
                for (int i = 0; i < obj_mesh.mat_coordinates.rows(); ++i) {
                    obj_mesh.mat_coordinates(i, 0) += ten_rho_field.dimension(0) / 2.0;
                    obj_mesh.mat_coordinates(i, 1) += ten_rho_field.dimension(1) / 2.0;
                    obj_mesh.mat_coordinates(i, 2) += ten_rho_field.dimension(2) / 2.0;
                }
            }
            // trilinear interploation
            Eigen::MatrixXd top_space_coords = r_bdr.MapWorldCoords2TopologyCoords(obj_mesh.mat_coordinates);
            Eigen::MatrixXi center_coords = (top_space_coords.array() + 0.5).cast<int>();
            static Eigen::MatrixXi delta_coords =
                    (Eigen::MatrixXi(8, 3) << -1, -1, -1, 0, -1, -1, 0, 0, -1, -1, 0, -1,

                            -1, -1, 0, 0, -1, 0, 0, 0, 0, -1, 0, 0)
                            .finished();

            Eigen::VectorXd rho_each_vertex(top_space_coords.rows());
            for (int i = 0; i < center_coords.rows(); ++i) {
                // each coord
                Eigen::Vector3i i_center_coord = center_coords.row(i);
                Eigen::MatrixXi around_coords = i_center_coord.transpose().replicate(8, 1) + delta_coords;
                around_coords = r_bdr.ClampTopEleCoords(around_coords);
                Eigen::VectorXd rho_around = ten_rho_field(around_coords);  // 8x1
                double rho_max = rho_around.maxCoeff();
                for (int ri = 0; ri < rho_around.size(); ++ri) {
                    if (rho_around(ri) < 1e-3) {
                        rho_around(ri) = rho_max;
                    }
                }

                Eigen::Vector3d d_center_coord = top_space_coords.row(i);
                Eigen::Vector3d ratio = (d_center_coord - i_center_coord.cast<double>()).array() + 0.5;  // 3x1

                static auto LinearInterp = [](const Eigen::Vector2d &props, double ratio) {
                    return props(1) * ratio + props(0) * (1.0 - ratio);
                };
                static auto BilinearInterp = [](const Eigen::Vector4d &props, const Eigen::Vector2d &ratios) {
                    Eigen::Vector2d y_props;
                    y_props(0) = LinearInterp({props(0), props(1)}, ratios.x());
                    y_props(1) = LinearInterp({props(3), props(2)}, ratios.x());
                    return LinearInterp(y_props, ratios.y());
                };
                static auto TrilinearInterp = [](const Eigen::Vector<double, 8> &props,
                                                 const Eigen::Vector3d &ratios) {
                    Eigen::Vector2d xy_props;
                    xy_props(0) = BilinearInterp(props.topRows(4), ratios.topRows(2));
                    xy_props(1) = BilinearInterp(props.bottomRows(4), ratios.topRows(2));
                    return LinearInterp(xy_props, ratios.z());
                };
                rho_each_vertex(i) = TrilinearInterp(rho_around, ratio);
                // assert(rho_each_vertex(i)>=0 && rho_each_vertex(i)<=1);
            }

            return {obj_mesh, rho_each_vertex};
        }

    }  // namespace top
}  // namespace da::sha