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

#ifndef TOP3D_MESH_H
#define TOP3D_MESH_H

#include "Util.h"
#include "unsupported/Eigen/KroneckerProduct"
namespace top {
    class Mesh {
    public:
        Mesh(int len_x, int len_y, int len_z) : lx_(len_x), ly_(len_y), lz_(len_z),
                                                num_node_((lx_ + 1) * (ly_ + 1) * (lz_ + 1)),
                                                num_ele_(lx_ * ly_ * lz_) {
            ten_node_coord2node_id = Tensor3i(num_node_, 1, 1);
            for (int i = 0; i < num_node_; ++i) {
                ten_node_coord2node_id(i, 0, 0) = i;
            }
            ten_node_coord2node_id=ten_node_coord2node_id.reshape(Eigen::array<int, 3>{lx_ + 1, ly_ + 1, lz_ + 1});

            ten_ele_coord2ele_id = Tensor3i(num_ele_, 1, 1);
            for (int i = 0; i < num_ele_; ++i) {
                ten_ele_coord2ele_id(i, 0, 0) = i;
            }
            ten_ele_coord2ele_id=ten_ele_coord2ele_id.reshape(Eigen::array<int, 3>{lx_, ly_, lz_});

            mat_ele_id2dofs_.resize(num_ele_, NUM_NODES_EACH_ELE*3);
            static const Eigen::MatrixXi delta_coord = (Eigen::MatrixXi(8, 3) <<
                    0, 0, 0,
                    1, 0, 0,
                    1, 1, 0,
                    0, 1, 0,
                    0, 0, 1,
                    1, 0, 1,
                    1, 1, 1,
                    0, 1, 1
            ).finished();
            for (int k = 0; k < lz_; ++k) {
                for (int j = 0; j < ly_; ++j) {
                    for (int i = 0; i < lx_; ++i) {
                        int cur_ele_id = ten_ele_coord2ele_id(i, j, k);
                        Eigen::MatrixXi world_node_coords=delta_coord.rowwise()+Eigen::RowVector3i(i,j,k);
                        assert(world_node_coords.rows()==8 && world_node_coords.cols()==3);
                        Eigen::Vector<int,8> node_ids=ten_node_coord2node_id(world_node_coords);
                        for(int nodi=0;nodi<8;++nodi){
                            mat_ele_id2dofs_(cur_ele_id, {3*nodi,3*nodi+1,3*nodi+2}) = Eigen::Vector3i(3*node_ids(nodi),3*node_ids(nodi)+1,3*node_ids(nodi)+2);
                        }
                    }
                }
            }



        }

        int GetLx()const{
            return lx_;
        }
        int GetLy()const{
            return ly_;
        }
        int GetLz()const{
            return lz_;
        }

        int GetNumDofs() const{
            return num_node_*3;
        }
        int GetNumEles()const{
            return num_ele_;
        }
        int GetNumNodes()const{
            return num_node_;
        }
        Eigen::MatrixXi GetEleId2DofsMap()const{
            return mat_ele_id2dofs_;
        }

        Eigen::VectorXi MapNodeCoord2Id(const Eigen::MatrixXi & node_coords)const{
            return ten_node_coord2node_id(node_coords);
        }
        Eigen::VectorXi MapEleCoord2Id(const Eigen::MatrixXi& ele_coords)const{
            return ten_ele_coord2ele_id(ele_coords);
        }
        Eigen::MatrixXi MapEleId2Dofs(const Eigen::VectorXi &ele_ids)const{
            return mat_ele_id2dofs_(ele_ids, all);
        }
        Eigen::VectorXi MapEleId2Dofs(const int ele_id)const{
            return mat_ele_id2dofs_(ele_id,all);
        }

    private:
        static const int NUM_NODES_EACH_ELE = 8;
        int lx_;
        int ly_;
        int lz_;
        int num_ele_;
        int num_node_;
        Tensor3i ten_node_coord2node_id;
        Tensor3i ten_ele_coord2ele_id;
        Eigen::MatrixXi mat_ele_id2dofs_;// num_ele_x8

    };

} // top

#endif //TOP3D_MESH_H