#include <polyfem/SplineBasis3d.hpp>

#include <polyfem/QuadraticBSpline3d.hpp>
#include <polyfem/HexQuadrature.hpp>

#include <polysolve/LinearSolver.hpp>
#include <polyfem/MeshNodes.hpp>
#include <polyfem/FEBasis3d.hpp>
#include <polyfem/Types.hpp>

#include <polyfem/Common.hpp>
#include <polyfem/auto_q_bases.hpp>

#include <Eigen/Sparse>

#include <cassert>
#include <iostream>
#include <vector>
#include <array>
#include <map>
#include <numeric>

//TODO carefull with simplices

namespace polyfem
{
    using namespace Eigen;
    using namespace polysolve;

    namespace
    {
        class SpaceMatrix
        {
        public:
            inline const int &operator()(const int i, const int j, const int k) const
            {
                return space_[k](i,j);
            }

            inline int &operator()(const int i, const int j, const int k)
            {
                return space_[k](i,j);
            }

            bool is_k_regular = false;
            int x,y,z;
            int edge_id;

            bool is_regular(const int xx, const int yy, const int zz) const
            {
                if(!is_k_regular) return true;

                return !((x == 1 && y == yy && z == zz) || (x == xx && y == 1 && z == zz) || (x == xx && y == yy && z == 1));
                // space(1, y, z).size() <= 1 && space(x, 1, z).size() <= 1 && space(x, y, 1).size() <= 1
            }

        private:
            std::array<Matrix<int, 3, 3> , 3> space_;
        };

        bool is_edge_singular(const Navigation3D::Index &index, const Mesh3D &mesh)
        {
            std::vector<int> ids;
            mesh.get_edge_elements_neighs(index.edge, ids);

            if(ids.size() == 4 || mesh.is_boundary_edge(index.edge))
                return false;

            for(auto idx : ids)
            {
                if(mesh.is_polytope(idx))
                    return false;
            }

            return true;
        }

        void print_local_space(const SpaceMatrix &space)
        {
            for(int k = 2; k >= 0; --k)
            {
                for(int j=2; j >=0; --j)
                {
                    for(int i=0; i < 3; ++i)
                    {
                        // if(space(i, j, k).size() > 0){
                            // for(std::size_t l = 0; l < space(i, j, k).size(); ++l)
                        std::cout<<space(i, j, k)<<"\t";
                        // }
                        // else
                            // std::cout<<"x\t";
                    }
                    std::cout<<std::endl;
                }

                std::cout<<"\n"<<std::endl;
            }
        }

        int node_id_from_face_index(const Mesh3D &mesh, MeshNodes &mesh_nodes, const Navigation3D::Index &index)
        {
            int el_id = mesh.switch_element(index).element;
            if(el_id >= 0 && mesh.is_cube(el_id))
            {
                return mesh_nodes.node_id_from_cell(el_id);
            }

            return mesh_nodes.node_id_from_face(index.face);
        }

        int node_id_from_edge_index(const Mesh3D &mesh, MeshNodes &mesh_nodes, const Navigation3D::Index &index)
        {
            Navigation3D::Index new_index = mesh.switch_element(index);
            int el_id = new_index.element;

            if(el_id < 0 || mesh.is_polytope(el_id))
            {
                new_index = mesh.switch_element(mesh.switch_face(index));
                el_id = new_index.element;

                if(el_id < 0 || mesh.is_polytope(el_id))
                {
                    return mesh_nodes.node_id_from_edge(index.edge);
                }

                return node_id_from_face_index(mesh, mesh_nodes, mesh.switch_face(new_index));
            }

            return node_id_from_face_index(mesh, mesh_nodes, mesh.switch_face(new_index));
        }

        int node_id_from_vertex_index_explore(const Mesh3D &mesh, const MeshNodes &mesh_nodes, const Navigation3D::Index &index, int &node_id)
        {
            Navigation3D::Index new_index = mesh.switch_element(index);

            int id = new_index.element;

            if(id < 0 || mesh.is_polytope(id))
            {
                // id = vertex_node_id(index.vertex);
                // node = node_from_vertex(index.vertex);
                node_id = mesh_nodes.primitive_from_vertex(index.vertex);
                return 3;
            }

            new_index = mesh.switch_element(mesh.switch_face(new_index));
            id = new_index.element;

            if(id < 0 || mesh.is_polytope(id))
            {
                // id = edge_node_id(switch_edge(new_index).edge);
                // node = node_from_edge(switch_edge(new_index).edge);
                node_id = mesh_nodes.primitive_from_edge(mesh.switch_edge(new_index).edge);
                return 2;
            }

            new_index = mesh.switch_element(mesh.switch_face(mesh.switch_edge(new_index)));
            id = new_index.element;

            if(id < 0 || mesh.is_polytope(id))
            {
                // id = face_node_id(new_index.face);
                // node = node_from_face(new_index.face);
                node_id = mesh_nodes.primitive_from_face(new_index.face);
                return 1;
            }

            // node = node_from_element(id);
            node_id = mesh_nodes.primitive_from_cell(id);
            return 0;
        }

        int node_id_from_vertex_index(const Mesh3D &mesh, MeshNodes &mesh_nodes, const Navigation3D::Index &index)
        {
            std::array<int, 6> path;
            std::array<int, 6> primitive_ids;

            path[0] = node_id_from_vertex_index_explore(mesh, mesh_nodes, index, primitive_ids[0]);
            path[1] = node_id_from_vertex_index_explore(mesh, mesh_nodes, mesh.switch_face(index), primitive_ids[1]);

            path[2] = node_id_from_vertex_index_explore(mesh, mesh_nodes, mesh.switch_edge(index), primitive_ids[2]);
            path[3] = node_id_from_vertex_index_explore(mesh, mesh_nodes, mesh.switch_face(mesh.switch_edge(index)), primitive_ids[3]);

            path[4] = node_id_from_vertex_index_explore(mesh, mesh_nodes, mesh.switch_edge(mesh.switch_face(index)), primitive_ids[4]);
            path[5] = node_id_from_vertex_index_explore(mesh, mesh_nodes, mesh.switch_face(mesh.switch_edge(mesh.switch_face(index))), primitive_ids[5]);

            const int min_path = *std::min_element(path.begin(), path.end());

            int primitive_id = 0;
            for(int i = 0 ; i < 6; ++i)
            {
                if(path[i]==min_path)
                {
                    primitive_id = primitive_ids[i];
                    break;
                }
            }

            return mesh_nodes.node_id_from_primitive(primitive_id);
        }

        void get_edge_elements_neighs(const Mesh3D &mesh, MeshNodes &mesh_nodes, const int element_id, const int edge_id, int dir, std::vector<int> &ids)
        {
            std::array<std::function<Navigation3D::Index(Navigation3D::Index)>, 12> to_edge;
            mesh.to_edge_functions(to_edge);

            Navigation3D::Index index;
            for(int i = 0; i < 12; ++i)
            {
                index = to_edge[i](mesh.get_index_from_element(element_id));

                if(index.edge == edge_id)
                    break;
            }

            assert(index.edge == edge_id);

            if(dir == 1)
            {
                int id;
                do
                {
                    ids.push_back(mesh_nodes.node_id_from_cell(index.element));
                    index = mesh.next_around_edge(index);
                }
                while(index.element != element_id);

                return;
            }

            if(dir == 0)
            {
                int id;
                do
                {
                    const Navigation3D::Index f_index = mesh.switch_face(mesh.switch_edge(index));
                    ids.push_back(node_id_from_face_index(mesh, mesh_nodes, f_index));

                    index = mesh.next_around_edge(index);
                }
                while(index.element != element_id);

                return;
            }

            if(dir == 2)
            {
                int id;
                do
                {
                    const Navigation3D::Index f_index = mesh.switch_face(mesh.switch_edge(mesh.switch_vertex(index)));
                    ids.push_back(node_id_from_face_index(mesh, mesh_nodes, f_index));

                    index = mesh.next_around_edge(index);
                }
                while(index.element != element_id);

                return;
            }

            assert(false);
        }

        void add_edge_id_for_poly(const Navigation3D::Index &index, const Mesh3D &mesh, MeshNodes &mesh_nodes, const int global_index, std::map<int, InterfaceData> &poly_face_to_data)
        {
           const int f1 = index.face;
           const int f2 = mesh.switch_face(index).face;

           const int id1 = mesh_nodes.primitive_from_face(f1);
           const int id2 = mesh_nodes.primitive_from_face(f2);

           if(mesh_nodes.is_primitive_interface(id1))
           {
            InterfaceData &data = poly_face_to_data[f1];
            data.local_indices.push_back(global_index);
        }

        if(mesh_nodes.is_primitive_interface(id2))
        {
            InterfaceData &data = poly_face_to_data[f2];
            data.local_indices.push_back(global_index);
        }
    }

    void add_vertex_id_for_poly(const Navigation3D::Index &index, const Mesh3D &mesh, MeshNodes &mesh_nodes, const int global_index, std::map<int, InterfaceData> &poly_face_to_data)
    {
       const int f1 = index.face;
       const int f2 = mesh.switch_face(index).face;
       const int f3 = mesh.switch_face(mesh.switch_edge(index)).face;

       const int id1 = mesh_nodes.primitive_from_face(f1);
       const int id2 = mesh_nodes.primitive_from_face(f2);
       const int id3 = mesh_nodes.primitive_from_face(f3);

       if(mesh_nodes.is_primitive_interface(id1))
       {
        InterfaceData &data = poly_face_to_data[f1];
        data.local_indices.push_back(global_index);
    }

    if(mesh_nodes.is_primitive_interface(id2))
    {
        InterfaceData &data = poly_face_to_data[f2];
        data.local_indices.push_back(global_index);
    }

    if(mesh_nodes.is_primitive_interface(id3))
    {
        InterfaceData &data = poly_face_to_data[f3];
        data.local_indices.push_back(global_index);
    }
}

void explore_edge(const Navigation3D::Index &index, const Mesh3D &mesh, MeshNodes &mesh_nodes, const int x, const int y, const int z, SpaceMatrix &space, LocalBoundary &local_boundary, std::map<int, InterfaceData> &poly_face_to_data)
{
    int node_id = node_id_from_edge_index(mesh, mesh_nodes, index);
    space(x, y, z) = node_id;
            // node(x, y, z) = mesh_nodes.node_position(node_id);

    if(is_edge_singular(index, mesh))
    {
        std::vector<int> ids;
        mesh.get_edge_elements_neighs(index.edge, ids);
                //irregular edge

        assert(!space.is_k_regular);
        space.is_k_regular = true;
        space.x = x; space.y = y; space.z = z;
        space.edge_id = index.edge;
    }

    // if(mesh_nodes.is_boundary(node_id))
        // bounday_nodes.push_back(node_id);

    add_edge_id_for_poly(index, mesh, mesh_nodes, 9*z + 3*y + x, poly_face_to_data);
}

void explore_vertex(const Navigation3D::Index &index, const Mesh3D &mesh, MeshNodes &mesh_nodes, const int x, const int y, const int z, SpaceMatrix &space, LocalBoundary &local_boundary, std::map<int, InterfaceData> &poly_face_to_data)
{
    int node_id = node_id_from_vertex_index(mesh, mesh_nodes, index);
    space(x, y, z) = node_id;
            // node(x, y, z) = mesh_nodes.node_position(node_id);

    // if(mesh_nodes.is_boundary(node_id))
        // bounday_nodes.push_back(node_id);

    add_vertex_id_for_poly(index, mesh, mesh_nodes, 9*z + 3*y + x, poly_face_to_data);
}

void explore_face(const Navigation3D::Index &index, const Mesh3D &mesh, MeshNodes &mesh_nodes, const int x, const int y, const int z,  SpaceMatrix &space, LocalBoundary &local_boundary, std::map<int, InterfaceData> &poly_face_to_data)
{
    int node_id = node_id_from_face_index(mesh, mesh_nodes, index);
    space(x, y, z) = node_id;
            // node(x, y, z) = mesh_nodes.node_position(node_id);

    if(mesh_nodes.is_boundary(node_id))
    {
        // local_boundary.add_boundary_primitive(index.face, FEBasis3d::quadr_hex_face_local_nodes(mesh, index)[8]-20);
        local_boundary.add_boundary_primitive(index.face, FEBasis3d::hex_face_local_nodes(false, 2, mesh, index)[8]-20);
        // bounday_nodes.push_back(node_id);
    }
    else if(mesh_nodes.is_interface(node_id))
    {
        InterfaceData &data = poly_face_to_data[index.face];
        data.local_indices.push_back(9*z + 3*y + x);
                // igl::viewer::Viewer &viewer = UIState::ui_state().viewer;
                // viewer.data.add_points(mesh_nodes.node_position(node_id), Eigen::MatrixXd::Constant(1, 3, 0));
    }
}

void build_local_space(const Mesh3D &mesh,  MeshNodes &mesh_nodes, const int el_index,  SpaceMatrix &space, std::vector<LocalBoundary> &local_boundary, std::map<int, InterfaceData> &poly_face_to_data)
{
    assert(mesh.is_volume());

    Navigation3D::Index start_index = mesh.get_index_from_element(el_index);
    Navigation3D::Index index;

    std::array<std::function<Navigation3D::Index(Navigation3D::Index)>, 6> to_face;
    mesh.to_face_functions(to_face);

    std::array<std::function<Navigation3D::Index(Navigation3D::Index)>, 12> to_edge;
    mesh.to_edge_functions(to_edge);

    std::array<std::function<Navigation3D::Index(Navigation3D::Index)>, 8> to_vertex;
    mesh.to_vertex_functions(to_vertex);

    const int node_id = mesh_nodes.node_id_from_cell(el_index);
    space(1, 1, 1) = node_id;
            // node(1, 1, 1) = mesh_nodes.node_position(node_id);

    LocalBoundary lb(el_index, BoundaryType::Quad);

            ///////////////////////
    index = to_face[1](start_index);
    explore_face(index, mesh, mesh_nodes, 1, 1, 0, space, lb, poly_face_to_data);

    index = to_face[0](start_index);
    explore_face(index, mesh, mesh_nodes, 1, 1, 2, space, lb, poly_face_to_data);

    index = to_face[3](start_index);
    explore_face(index, mesh, mesh_nodes, 0, 1, 1, space, lb, poly_face_to_data);

    index = to_face[2](start_index);
    explore_face(index, mesh, mesh_nodes, 2, 1, 1, space, lb, poly_face_to_data);

    index = to_face[5](start_index);
    explore_face(index, mesh, mesh_nodes, 1, 0, 1, space, lb, poly_face_to_data);

    index = to_face[4](start_index);
    explore_face(index, mesh, mesh_nodes, 1, 2, 1, space, lb, poly_face_to_data);

            ///////////////////////
    index = to_edge[0](start_index);
    explore_edge(index, mesh, mesh_nodes, 1, 0, 0, space, lb, poly_face_to_data);

    index = to_edge[1](start_index);
    explore_edge(index, mesh, mesh_nodes, 2, 1, 0, space, lb, poly_face_to_data);

    index = to_edge[2](start_index);
    explore_edge(index, mesh, mesh_nodes, 1, 2, 0, space, lb, poly_face_to_data);

    index = to_edge[3](start_index);
    explore_edge(index, mesh, mesh_nodes, 0, 1, 0, space, lb, poly_face_to_data);

    index = to_edge[4](start_index);
    explore_edge(index, mesh, mesh_nodes, 0, 0, 1, space, lb, poly_face_to_data);

    index = to_edge[5](start_index);
    explore_edge(index, mesh, mesh_nodes, 2, 0, 1, space, lb, poly_face_to_data);

    index = to_edge[6](start_index);
    explore_edge(index, mesh, mesh_nodes, 2, 2, 1, space, lb, poly_face_to_data);

    index = to_edge[7](start_index);
    explore_edge(index, mesh, mesh_nodes, 0, 2, 1, space, lb, poly_face_to_data);

    index = to_edge[8](start_index);
    explore_edge(index, mesh, mesh_nodes, 1, 0, 2, space, lb, poly_face_to_data);

    index = to_edge[9](start_index);
    explore_edge(index, mesh, mesh_nodes, 2, 1, 2, space, lb, poly_face_to_data);

    index = to_edge[10](start_index);
    explore_edge(index, mesh, mesh_nodes, 1, 2, 2, space, lb, poly_face_to_data);

    index = to_edge[11](start_index);
    explore_edge(index, mesh, mesh_nodes, 0, 1, 2, space, lb, poly_face_to_data);

            ////////////////////////////////////////////////////////////////////////
    index = to_vertex[0](start_index);
    explore_vertex(index, mesh, mesh_nodes, 0, 0, 0, space, lb, poly_face_to_data);

    index = to_vertex[1](start_index);
    explore_vertex(index, mesh, mesh_nodes, 2, 0, 0, space, lb, poly_face_to_data);

    index = to_vertex[2](start_index);
    explore_vertex(index, mesh, mesh_nodes, 2, 2, 0, space, lb, poly_face_to_data);

    index = to_vertex[3](start_index);
    explore_vertex(index, mesh, mesh_nodes, 0, 2, 0, space, lb, poly_face_to_data);

    index = to_vertex[4](start_index);
    explore_vertex(index, mesh, mesh_nodes, 0, 0, 2, space, lb, poly_face_to_data);

    index = to_vertex[5](start_index);
    explore_vertex(index, mesh, mesh_nodes, 2, 0, 2, space, lb, poly_face_to_data);

    index = to_vertex[6](start_index);
    explore_vertex(index, mesh, mesh_nodes, 2, 2, 2, space, lb, poly_face_to_data);

    index = to_vertex[7](start_index);
    explore_vertex(index, mesh, mesh_nodes, 0, 2, 2, space, lb, poly_face_to_data);

    if(!lb.empty())
        local_boundary.emplace_back(lb);
}

void setup_knots_vectors(MeshNodes &mesh_nodes, const SpaceMatrix &space, std::array<std::array<double, 4>, 3> &h_knots, std::array<std::array<double, 4>, 3> &v_knots, std::array<std::array<double, 4>, 3> &w_knots)
{
            //left and right neigh are absent
    if(mesh_nodes.is_boundary_or_interface(space(0, 1, 1)) && mesh_nodes.is_boundary_or_interface(space(2, 1, 1)))
    {
        h_knots[0] = {{0, 0, 0, 1}};
        h_knots[1] = {{0, 0, 1, 1}};
        h_knots[2] = {{0, 1, 1, 1}};
    }
             //left neigh is absent
    else if(mesh_nodes.is_boundary_or_interface(space(0, 1, 1)))
    {
        h_knots[0] = {{0, 0, 0, 1}};
        h_knots[1] = {{0, 0, 1, 2}};
        h_knots[2] = {{0, 1, 2, 3}};
    }
            //right neigh is absent
    else if(mesh_nodes.is_boundary_or_interface(space(2,1,1)))
    {
        h_knots[0] = {{-2, -1, 0, 1}};
        h_knots[1] = {{-1, 0, 1, 1}};
        h_knots[2] = {{0, 1, 1, 1}};
    }
    else
    {
        h_knots[0] = {{-2, -1, 0, 1}};
        h_knots[1] = {{-1, 0, 1, 2}};
        h_knots[2] = {{0, 1, 2, 3}};
    }

            //top and bottom neigh are absent
    if(mesh_nodes.is_boundary_or_interface(space(1,0,1)) && mesh_nodes.is_boundary_or_interface(space(1,2,1)))
    {
        v_knots[0] = {{0, 0, 0, 1}};
        v_knots[1] = {{0, 0, 1, 1}};
        v_knots[2] = {{0, 1, 1, 1}};
    }
            //bottom neigh is absent
    else if(mesh_nodes.is_boundary_or_interface(space(1,0,1)))
    {
        v_knots[0] = {{0, 0, 0, 1}};
        v_knots[1] = {{0, 0, 1, 2}};
        v_knots[2] = {{0, 1, 2, 3}};
    }
            //top neigh is absent
    else if(mesh_nodes.is_boundary_or_interface(space(1,2,1)))
    {
        v_knots[0] = {{-2, -1, 0, 1}};
        v_knots[1] = {{-1, 0, 1, 1}};
        v_knots[2] = {{0, 1, 1, 1}};
    }
    else
    {
        v_knots[0] = {{-2, -1, 0, 1}};
        v_knots[1] = {{-1, 0, 1, 2}};
        v_knots[2] = {{0, 1, 2, 3}};
    }

            //front and back neigh are absent
    if(mesh_nodes.is_boundary_or_interface(space(1,1,0)) && mesh_nodes.is_boundary_or_interface(space(1,1,2)))
    {
        w_knots[0] = {{0, 0, 0, 1}};
        w_knots[1] = {{0, 0, 1, 1}};
        w_knots[2] = {{0, 1, 1, 1}};
    }
            //back neigh is absent
    else if(mesh_nodes.is_boundary_or_interface(space(1,1,0)))
    {
        w_knots[0] = {{0, 0, 0, 1}};
        w_knots[1] = {{0, 0, 1, 2}};
        w_knots[2] = {{0, 1, 2, 3}};
    }
            //front neigh is absent
    else if(mesh_nodes.is_boundary_or_interface(space(1,1,2)))
    {
        w_knots[0] = {{-2, -1, 0, 1}};
        w_knots[1] = {{-1, 0, 1, 1}};
        w_knots[2] = {{0, 1, 1, 1}};
    }
    else
    {
        w_knots[0] = {{-2, -1, 0, 1}};
        w_knots[1] = {{-1, 0, 1, 2}};
        w_knots[2] = {{0, 1, 2, 3}};
    }
}

void basis_for_regular_hex(MeshNodes &mesh_nodes, const SpaceMatrix &space, const std::array<std::array<double, 4>, 3> &h_knots, const std::array<std::array<double, 4>, 3> &v_knots, const std::array<std::array<double, 4>, 3> &w_knots, ElementBases &b)
{
    for(int z = 0; z < 3; ++z)
    {
        for(int y = 0; y < 3; ++y)
        {
            for(int x = 0; x < 3; ++x)
            {
                        if(space.is_regular(x,y,z)) //space(1, y, z).size() <= 1 && space(x, 1, z).size() <= 1 && space(x, y, 1).size() <= 1)
                        {
                            const int local_index = 9*z + 3*y + x;
                            const int global_index = space(x, y, z);
                            const auto node = mesh_nodes.node_position(global_index);
                            // loc_nodes(x, y, z);

                            b.bases[local_index].init(2, global_index, local_index, node);

                            const QuadraticBSpline3d spline(h_knots[x], v_knots[y], w_knots[z]);

                            b.bases[local_index].set_basis([spline](const Eigen::MatrixXd &uv, Eigen::MatrixXd &val) { spline.interpolate(uv, val); });
                            b.bases[local_index].set_grad( [spline](const Eigen::MatrixXd &uv, Eigen::MatrixXd &val) { spline.derivative(uv, val); });
                        }
                    }
                }
            }
        }

        void basis_for_irregulard_hex(const int el_index, const Mesh3D &mesh, MeshNodes &mesh_nodes, const SpaceMatrix &space, const std::array<std::array<double, 4>, 3> &h_knots, const std::array<std::array<double, 4>, 3> &v_knots, const std::array<std::array<double, 4>, 3> &w_knots, ElementBases &b, std::map<int, InterfaceData> &poly_face_to_data)
        {
            for(int z = 0; z < 3; ++z)
            {
                for(int y = 0; y < 3; ++y)
                {
                    for(int x = 0; x < 3; ++x)
                    {
                        if(!space.is_regular(x,y,z)) //space(1, y, z).size() > 1 || space(x, 1, z).size() > 1 || space(x, y, 1).size() > 1)
                        {
                            const int local_index = 9*z + 3*y + x;

                            int mpx = -1;
                            int mpy = -1;
                            int mpz = -1;

                            int mmx = -1;
                            int mmy = -1;
                            int mmz = -1;

                            int xx = 1;
                            int yy = 1;
                            int zz = 1;

                            const int edge_id = space.edge_id;
                            int dir = -1;

                            if(space.x == x && space.y == y && space.z == 1)
                            {
                                mpx = 1;
                                mpy = y;
                                mpz = z;

                                mmx = x;
                                mmy = 1;
                                mmz = z;

                                zz = z;
                                dir = z;
                            }
                            else if(space.x == x && space.y == 1 && space.z == z)
                            {
                                mpx = 1;
                                mpy = y;
                                mpz = z;

                                mmx = x;
                                mmy = y;
                                mmz = 1;

                                yy = y;
                                dir = y;
                            }
                            else if(space.x == 1 && space.y == y && space.z == z)
                            {
                                mpx = x;
                                mpy = y;
                                mpz = 1;

                                mmx = x;
                                mmy = 1;
                                mmz = z;

                                xx = x;
                                dir = x;
                            }
                            else
                                assert(false);

                            const auto &center = b.bases[zz*9 + yy*3 + xx].global().front();

                            const auto &el1 = b.bases[mpz*9 + mpy*3 + mpx].global().front();
                            const auto &el2 = b.bases[mmz*9 + mmy*3 + mmx].global().front();

                            // std::cout<<"el_index "<<el_index<<std::endl;
                            // std::cout<<"center "<<center.index<<std::endl;
                            // std::cout<<"el1 "<<el1.index<<std::endl;
                            // std::cout<<"el2 "<<el2.index<<std::endl;

                            std::vector<int> ids;
                            get_edge_elements_neighs(mesh, mesh_nodes, el_index, edge_id, dir, ids);

                            if(ids.front() != center.index)
                            {
                                assert(dir != 1);
                                ids.clear();
                                get_edge_elements_neighs(mesh, mesh_nodes, el_index, edge_id, dir == 2 ? 0 : 2, ids);
                            }

                            assert(ids.front() == center.index);

                            std::vector<int> other_indices;
                            std::vector<Eigen::MatrixXd> other_nodes;
                            for(size_t i = 0; i < ids.size(); ++i)
                            {
                                const int node_id = ids[i];
                                if(node_id != center.index && node_id != el1.index && node_id != el2.index){
                                    other_indices.push_back(node_id);
                                }
                            }

                            // std::cout<<ids.size()<< " " << other_indices.size()<<std::endl;

                            auto &base = b.bases[local_index];

                            const int k = int(other_indices.size()) + 3;

                            // const bool is_interface = mesh_nodes.is_interface(center.index);
                            // const int face_id = is_interface ? mesh_nodes.face_from_node_id(center.index) : -1;

                            base.global().resize(k);

                            base.global()[0].index = center.index;
                            base.global()[0].val = (4. - k) / k;
                            base.global()[0].node = center.node;

                            base.global()[1].index = el1.index;
                            base.global()[1].val = (4. - k) / k;
                            base.global()[1].node = el1.node;

                            base.global()[2].index = el2.index;
                            base.global()[2].val = (4. - k) / k;
                            base.global()[2].node = el2.node;

                            // if(is_interface){
                                // poly_face_to_data[face_id].local_indices.push_back(local_index);
                            // }

                            for(std::size_t n = 0; n < other_indices.size(); ++n)
                            {
                                base.global()[3+n].index = other_indices[n];
                                base.global()[3+n].val = 4./k;
                                base.global()[3+n].node = mesh_nodes.node_position(other_indices[n]);
                            }

                            const QuadraticBSpline3d spline(h_knots[x], v_knots[y], w_knots[z]);

                            b.bases[local_index].set_basis([spline](const Eigen::MatrixXd &uv, Eigen::MatrixXd &val) { spline.interpolate(uv, val); });
                            b.bases[local_index].set_grad( [spline](const Eigen::MatrixXd &uv, Eigen::MatrixXd &val) { spline.derivative(uv, val); });
                        }
                    }
                }
            }
        }

        void create_q2_nodes(const Mesh3D &mesh, const int el_index, std::set<int> &vertex_id, std::set<int> &edge_id, std::set<int> &face_id, ElementBases &b, std::vector<LocalBoundary> &local_boundary, int &n_bases)
        {
            b.bases.resize(27);

            std::array<std::function<Navigation3D::Index(Navigation3D::Index)>, 6> to_face;
            mesh.to_face_functions(to_face);

            std::array<std::function<Navigation3D::Index(Navigation3D::Index)>, 12> to_edge;
            mesh.to_edge_functions(to_edge);

            std::array<std::function<Navigation3D::Index(Navigation3D::Index)>, 8> to_vertex;
            mesh.to_vertex_functions(to_vertex);

            LocalBoundary lb(el_index, BoundaryType::Quad);

            const Navigation3D::Index start_index = mesh.get_index_from_element(el_index);
            for (int j = 0; j < 8; ++j)
            {
                const Navigation3D::Index index = to_vertex[j](start_index);
                // const int loc_index = FEBasis3d::quadr_hex_face_local_nodes(mesh, index)[0];
                const int loc_index = FEBasis3d::hex_face_local_nodes(false, 2, mesh, index)[0];

                int current_vertex_node_id = -1;
                Eigen::MatrixXd current_vertex_node;

                //if the edge/vertex is boundary the it is a Q2 edge
                bool is_vertex_q2 = true;

                std::vector<int> vertex_neighs;
                mesh.get_vertex_elements_neighs(index.vertex, vertex_neighs);

                for(size_t i = 0; i < vertex_neighs.size(); ++i)
                {
                    if(mesh.is_spline_compatible(vertex_neighs[i]))
                    {
                        is_vertex_q2 = false;
                        break;
                    }
                }
                const bool is_vertex_boundary = mesh.is_boundary_vertex(index.vertex);

                if(is_vertex_q2)
                {
                    const bool is_new_vertex = vertex_id.insert(index.vertex).second;

                    if(is_new_vertex)
                    {
                        current_vertex_node_id = n_bases++;
                        current_vertex_node = mesh.point(index.vertex);

                        // if(is_vertex_boundary)//mesh.is_vertex_boundary(index.vertex))
                            // bounday_nodes.push_back(current_vertex_node_id);
                    }
                }

                //init new Q2 nodes
                if(current_vertex_node_id >= 0)
                    b.bases[loc_index].init(2, current_vertex_node_id, loc_index, current_vertex_node);

                b.bases[loc_index].set_basis([loc_index](const Eigen::MatrixXd &uv, Eigen::MatrixXd &val) { autogen::q_basis_value_3d     (2, loc_index, uv, val); });
                b.bases[loc_index].set_grad( [loc_index](const Eigen::MatrixXd &uv, Eigen::MatrixXd &val) { autogen::q_grad_basis_value_3d(2, loc_index, uv, val); });
            }

            for (int j = 0; j < 12; ++j)
            {
                Navigation3D::Index index = to_edge[j](start_index);

                int current_edge_node_id = -1;
                Eigen::Matrix<double, 1, 3> current_edge_node;
                // const int loc_index = FEBasis3d::quadr_hex_face_local_nodes(mesh, index)[1];
                const int loc_index = FEBasis3d::hex_face_local_nodes(false, 2, mesh, index)[4];

                bool is_edge_q2 = true;

                std::vector<int> edge_neighs;
                mesh.get_edge_elements_neighs(index.edge, edge_neighs);

                for(size_t i = 0; i < edge_neighs.size(); ++i)
                {
                    if(mesh.is_spline_compatible(edge_neighs[i]))
                    {
                        is_edge_q2 = false;
                        break;
                    }
                }
                const bool is_edge_boundary = mesh.is_boundary_edge(index.edge);

                if(is_edge_q2)
                {
                    const bool is_new_edge = edge_id.insert(index.edge).second;

                    if(is_new_edge)
                    {
                        current_edge_node_id = n_bases++;
                        current_edge_node = mesh.edge_barycenter(index.edge);

                        // if(is_edge_boundary)
                            // bounday_nodes.push_back(current_edge_node_id);
                    }
                }

                //init new Q2 nodes
                if(current_edge_node_id >= 0)
                    b.bases[loc_index].init(2, current_edge_node_id, loc_index, current_edge_node);

                b.bases[loc_index].set_basis([loc_index](const Eigen::MatrixXd &uv, Eigen::MatrixXd &val) { autogen::q_basis_value_3d     (2, loc_index, uv, val); });
                b.bases[loc_index].set_grad( [loc_index](const Eigen::MatrixXd &uv, Eigen::MatrixXd &val) { autogen::q_grad_basis_value_3d(2, loc_index, uv, val); });
            }

            for (int j = 0; j < 6; ++j)
            {
                Navigation3D::Index index = to_face[j](start_index);

                int current_face_node_id = -1;

                Eigen::Matrix<double, 1, 3> current_face_node;
                const int opposite_element = mesh.switch_element(index).element;
                const bool is_face_q2 = opposite_element < 0 || !mesh.is_spline_compatible(opposite_element);
                // const int loc_index = FEBasis3d::quadr_hex_face_local_nodes(mesh, index)[8];
                const int loc_index = FEBasis3d::hex_face_local_nodes(false, 2, mesh, index)[8];

                if (is_face_q2)
                {
                    const bool is_new_face = face_id.insert(index.face).second;

                    if(is_new_face)
                    {
                        current_face_node_id = n_bases++;
                        current_face_node = mesh.face_barycenter(index.face);

                        const int b_index = loc_index - 20;

                        if(opposite_element < 0) // && mesh.n_element_faces(opposite_element) == 6 && mesh.n_element_vertices(opposite_element) == 8)
                        {
                            // bounday_nodes.push_back(current_face_node_id);
                            lb.add_boundary_primitive(index.face, b_index);
                        }
                    }
                }

                //init new Q2 nodes
                if(current_face_node_id >= 0)
                    b.bases[loc_index].init(2, current_face_node_id, loc_index, current_face_node);

                b.bases[loc_index].set_basis([loc_index](const Eigen::MatrixXd &uv, Eigen::MatrixXd &val) { autogen::q_basis_value_3d     (2, loc_index, uv, val); });
                b.bases[loc_index].set_grad( [loc_index](const Eigen::MatrixXd &uv, Eigen::MatrixXd &val) { autogen::q_grad_basis_value_3d(2, loc_index, uv, val); });
            }

            // //central node always present
            b.bases[26].init(2, n_bases++, 26, mesh.cell_barycenter(el_index));
            b.bases[26].set_basis([](const Eigen::MatrixXd &uv, Eigen::MatrixXd &val) { autogen::q_basis_value_3d     (2, 26, uv, val); });
            b.bases[26].set_grad( [](const Eigen::MatrixXd &uv, Eigen::MatrixXd &val) { autogen::q_grad_basis_value_3d(2, 26, uv, val); });

            if(!lb.empty())
                local_boundary.emplace_back(lb);
        }

        void insert_into_global(const int el_index, const Local2Global &data, std::vector<Local2Global> &vec, const int size)
        {
            //ignore small weights
            if(fabs(data.val) <1e-10 )
                return;

            bool found = false;

            for(int i = 0; i < size; ++i)
            {
                if(vec[i].index == data.index)
                {
                    // if(fabs(vec[i].val - data.val) > 1e-10){
                    //     std::cout<<el_index <<" "<<vec[i].val <<" "<< data.val<<" "<<fabs(vec[i].val - data.val)<<std::endl;
                    //     // vec[i].val += data.val;
                    // }
                    // assert(fabs(vec[i].val - data.val) < 1e-10);
                    assert((vec[i].node - data.node).norm() < 1e-10);
                    found = true;
                    break;
                }
            }

            if(!found)
                vec.push_back(data);
        }

        void assign_q2_weights(const Mesh3D &mesh, const int el_index, std::vector< ElementBases > &bases)
        {
            // Eigen::MatrixXd eval_p;
            std::vector<AssemblyValues> eval_p;
            const Navigation3D::Index start_index = mesh.get_index_from_element(el_index);
            ElementBases &b = bases[el_index];

            std::array<std::function<Navigation3D::Index(Navigation3D::Index)>, 6> to_face;
            mesh.to_face_functions(to_face);
            for (int f = 0; f < 6; ++f)
            {
                const Navigation3D::Index index = to_face[f](start_index);
                const int opposite_element = mesh.switch_element(index).element;

                if(opposite_element < 0 || !mesh.is_cube(opposite_element))
                    continue;

                // const auto &param_p     = FEBasis3d::quadr_hex_face_local_nodes_coordinates(mesh, mesh.switch_element(index));
                Eigen::Matrix<double, 9, 3> param_p;
                {
                    Eigen::MatrixXd hex_loc_nodes; polyfem::autogen::q_nodes_3d(2, hex_loc_nodes);
                    const auto opposite_indices = FEBasis3d::hex_face_local_nodes(false, 2, mesh, mesh.switch_element(index));
                    for(int k = 0; k < 9; ++k)
                        param_p.row(k) = hex_loc_nodes.row(opposite_indices[k]);
                }
                const auto &other_bases = bases[opposite_element];

                // const auto &indices     = FEBasis3d::quadr_hex_face_local_nodes(mesh, index);
                const auto &indices     = FEBasis3d::hex_face_local_nodes(false, 2, mesh, index);

                std::array<int, 9> sizes;

                for(int l = 0; l < 9; ++l)
                    sizes[l] = b.bases[indices[l]].global().size();

                other_bases.evaluate_bases(param_p, eval_p);
                for(std::size_t i = 0; i < other_bases.bases.size(); ++i)
                {
                    const auto &other_b = other_bases.bases[i];

                    if(other_b.global().empty()) continue;

                    // other_b.basis(param_p, eval_p);
                    assert(eval_p[i].val.size() == 9);

                    //basis i of element opposite element is zero on this elements
                    if(eval_p[i].val.cwiseAbs().maxCoeff() <= 1e-10)
                        continue;

                    for(std::size_t k = 0; k < other_b.global().size(); ++k)
                    {
                        for(int l = 0; l < 9; ++l)
                        {
                            Local2Global glob = other_b.global()[k];
                            glob.val *= eval_p[i].val(l);

                            insert_into_global(el_index, glob, b.bases[indices[l]].global(), sizes[l]);
                        }
                    }
                }
            }
        }

        void setup_data_for_polygons(const Mesh3D &mesh, const int el_index, const ElementBases &b, std::map<int, InterfaceData> &poly_face_to_data)
        {
            const Navigation3D::Index start_index = mesh.get_index_from_element(el_index);
            std::array<std::function<Navigation3D::Index(Navigation3D::Index)>, 6> to_face;
            mesh.to_face_functions(to_face);
            for (int f = 0; f < 6; ++f)
            {
                const Navigation3D::Index index = to_face[f](start_index);

                const int opposite_element = mesh.switch_element(index).element;
                const bool is_neigh_poly = opposite_element >= 0 && mesh.is_polytope(opposite_element);

                if(is_neigh_poly)
                {
                    // auto e2l = FEBasis3d::quadr_hex_face_local_nodes(mesh, index);
                    auto e2l = FEBasis3d::hex_face_local_nodes(false, 2, mesh, index);

                    InterfaceData &data = poly_face_to_data[index.face];

                    for(int kk = 0; kk < e2l.size(); ++kk)
                    {
                        const auto idx = e2l(kk);
                        data.local_indices.push_back(idx);
                    }
                }
            }
        }
    }

    int SplineBasis3d::build_bases(const Mesh3D &mesh, const int quadrature_order, std::vector< ElementBases > &bases, std::vector< LocalBoundary > &local_boundary, std::map<int, InterfaceData> &poly_face_to_data)
    {
        using std::max;
        assert(mesh.is_volume());

        MeshNodes mesh_nodes(mesh, true, true, 1, 1, 1);

        const int n_els = mesh.n_elements();
        bases.resize(n_els);
        local_boundary.clear();

        // bounday_nodes.clear();

        // HexQuadrature hex_quadrature;

        std::array<std::array<double, 4>, 3> h_knots;
        std::array<std::array<double, 4>, 3> v_knots;
        std::array<std::array<double, 4>, 3> w_knots;

        for(int e = 0; e < n_els; ++e)
        {
            if(!mesh.is_spline_compatible(e))
                continue;

            SpaceMatrix space;

            build_local_space(mesh, mesh_nodes, e, space, local_boundary, poly_face_to_data);

            ElementBases &b=bases[e];
            b.set_quadrature([quadrature_order](Quadrature &quad){
                HexQuadrature hex_quadrature;
                hex_quadrature.get_quadrature(quadrature_order, quad);
            });
            // hex_quadrature.get_quadrature(quadrature_order, b.quadrature);
            b.bases.resize(27);

            b.set_local_node_from_primitive_func([e](const int primitive_id, const Mesh &mesh)
            {
                const auto &mesh3d = dynamic_cast<const Mesh3D &>(mesh);

                std::array<std::function<Navigation3D::Index(Navigation3D::Index)>, 6> to_face;
                mesh3d.to_face_functions(to_face);

                auto start_index = mesh3d.get_index_from_element(e);
                auto index = start_index;

                int lf;
                for(lf = 0; lf < mesh3d.n_cell_faces(e); ++lf)
                {
                    index = to_face[lf](start_index);
                    if(index.face == primitive_id)
                        break;
                }
                assert(index.face == primitive_id);

                static constexpr std::array<std::array<int, 9>, 6> face_to_index = {{
                    {{2*9+0*3+0, 2*9+1*3+0, 2*9+2*3+0, 2*9+0*3+1, 2*9+1*3+1, 2*9+2*3+1, 2*9+0*3+2, 2*9+1*3+2, 2*9+2*3+2}}, //0
                    {{0*9+0*3+0, 0*9+1*3+0, 0*9+2*3+0, 0*9+0*3+1, 0*9+1*3+1, 0*9+2*3+1, 0*9+0*3+2, 0*9+1*3+2, 0*9+2*3+2}}, //1

                    {{0*9+0*3+2, 0*9+1*3+2, 0*9+2*3+2, 1*9+0*3+2, 1*9+1*3+2, 1*9+2*3+2, 2*9+0*3+2, 2*9+1*3+2, 2*9+2*3+2}}, //2
                    {{0*9+0*3+0, 0*9+1*3+0, 0*9+2*3+0, 1*9+0*3+0, 1*9+1*3+0, 1*9+2*3+0, 2*9+0*3+0, 2*9+1*3+0, 2*9+2*3+0}}, //3

                    {{0*9+2*3+0, 0*9+2*3+1, 0*9+2*3+2, 1*9+2*3+0, 1*9+2*3+1, 1*9+2*3+2, 2*9+2*3+0, 2*9+2*3+1, 2*9+2*3+2}}, //4
                    {{0*9+0*3+0, 0*9+0*3+1, 0*9+0*3+2, 1*9+0*3+0, 1*9+0*3+1, 1*9+0*3+2, 2*9+0*3+0, 2*9+0*3+1, 2*9+0*3+2}}, //5
                }};

                Eigen::VectorXi res(9);

                for(int i = 0; i< 9; ++i)
                    res(i)=face_to_index[lf][i];

                return res;
            });

            setup_knots_vectors(mesh_nodes, space, h_knots, v_knots, w_knots);
            // print_local_space(space);

            basis_for_regular_hex(mesh_nodes, space, h_knots, v_knots, w_knots, b);
            basis_for_irregulard_hex(e, mesh, mesh_nodes, space, h_knots, v_knots, w_knots, b, poly_face_to_data);
        }

        int n_bases = mesh_nodes.n_nodes();

        std::set<int> face_id;
        std::set<int> edge_id;
        std::set<int> vertex_id;

        for(int e = 0; e < n_els; ++e)
        {
            if(mesh.is_polytope(e) || mesh.is_spline_compatible(e))
                continue;

            ElementBases &b=bases[e];
            // hex_quadrature.get_quadrature(quadrature_order, b.quadrature);
            b.set_quadrature([quadrature_order](Quadrature &quad){
                HexQuadrature hex_quadrature;
                hex_quadrature.get_quadrature(quadrature_order, quad);
            });

            b.set_local_node_from_primitive_func([e](const int primitive_id, const Mesh &mesh)
            {
                const auto &mesh3d = dynamic_cast<const Mesh3D &>(mesh);
                Navigation3D::Index index;

                for(int lf = 0; lf < 6; ++lf)
                {
                    index = mesh3d.get_index_from_element(e, lf, 0);
                    if(index.face == primitive_id)
                        break;
                }
                assert(index.face == primitive_id);

                // const auto indices = FEBasis3d::quadr_hex_face_local_nodes(mesh3d, index);
                const auto indices = FEBasis3d::hex_face_local_nodes(false, 2, mesh3d, index);
                Eigen::VectorXi res(indices.size());

                for(size_t i = 0; i< indices.size(); ++i)
                    res(i)=indices[i];

                return res;
            });

            create_q2_nodes(mesh, e, vertex_id, edge_id, face_id, b, local_boundary, n_bases);
        }

        bool missing_bases = false;
        do
        {
            missing_bases = false;
            for(int e = 0; e < n_els; ++e)
            {
                if(mesh.is_polytope(e) || mesh.is_spline_compatible(e))
                    continue;

                auto &b=bases[e];
                if(b.is_complete())
                    continue;

                assign_q2_weights(mesh, e, bases);

                missing_bases = missing_bases || b.is_complete();
            }
        }
        while(missing_bases);

        for(int e = 0; e < n_els; ++e)
        {
            if(mesh.is_polytope(e) || mesh.is_spline_compatible(e))
                continue;

            const ElementBases &b=bases[e];
            setup_data_for_polygons(mesh, e, b, poly_face_to_data);
        }

        // for(int e = 0; e < n_els; ++e)
        // {
        //     if(!mesh.is_polytope(e))
        //         continue;

        //     for (int lf = 0; lf < mesh.n_cell_faces(e); ++lf)
        //     {
        //         auto index = mesh.get_index_from_element(e, lf, 0);
        //         auto index2 = mesh.switch_element(index);
        //         if (index2.element >= 0) {
        //             auto &array = poly_face_to_data[index.face].local_indices;
        //             auto &b = bases[index2.element];
        //             array.resize(b.bases.size());
        //             std::iota(array.begin(), array.end(), 0);
        //         }
        //     }
        // }

        for(auto &k : poly_face_to_data)
        {
            auto &array = k.second.local_indices;
            std::sort(array.begin(), array.end());
            auto it = std::unique(array.begin(), array.end());
            array.resize(std::distance(array.begin(), it));
        }

        return n_bases;
    }

    void SplineBasis3d::fit_nodes(const Mesh3D &mesh, const int n_bases, std::vector< ElementBases > &gbases)
    {
        assert(false);
        // const int dim = 3;
        // const int n_constraints =  27;
        // const int n_elements = mesh.n_elements();

        // std::vector< Eigen::Triplet<double> > entries, entries_t;

        // MeshNodes nodes(mesh, true, true, 1, 1, 1);
        // // Eigen::MatrixXd tmp;
        // std::vector<AssemblyValues> tmp_val;

        // Eigen::MatrixXd node_rhs(n_constraints*n_elements, dim);
        // Eigen::MatrixXd samples(n_constraints, dim);

        // for(int i = 0; i < n_constraints; ++i)
        //     samples.row(i) = FEBasis3d::quadr_hex_local_node_coordinates(i);

        // for(int i = 0; i < n_elements; ++i)
        // {
        //     auto &base = gbases[i];

        //     if(!mesh.is_cube(i))
        //         continue;

        //     auto global_ids = FEBasis3d::quadr_hex_local_to_global(mesh, i);
        //     assert(global_ids.size() == n_constraints);

        //     for(int j = 0; j < n_constraints; ++j)
        //     {
        //         auto n_id = nodes.node_id_from_primitive(global_ids[j]);
        //         auto n = nodes.node_position(n_id);
        //         for(int d = 0; d < dim; ++d)
        //             node_rhs(n_constraints*i + j, d) = n(d);
        //     }

        //     base.evaluate_bases(samples, tmp_val);
        //     const auto &lbs = base.bases;

        //     const int n_local_bases = int(lbs.size());
        //     for(int j = 0; j < n_local_bases; ++j)
        //     {
        //         const Basis &b = lbs[j];
        //         const auto &tmp = tmp_val[j].val;

        //         for(std::size_t ii = 0; ii < b.global().size(); ++ii)
        //         {
        //             for (long k = 0; k < tmp.size(); ++k)
        //             {
        //                 entries.emplace_back(n_constraints*i + k, b.global()[ii].index, tmp(k)*b.global()[ii].val);
        //                 entries_t.emplace_back(b.global()[ii].index, n_constraints*i + k, tmp(k)*b.global()[ii].val);
        //             }
        //         }
        //     }
        // }

        // Eigen::MatrixXd new_nodes(n_bases, dim);
        // {
        //     StiffnessMatrix mat(n_constraints*n_elements, n_bases);
        //     StiffnessMatrix mat_t(n_bases, n_constraints*n_elements);

        //     mat.setFromTriplets(entries.begin(), entries.end());
        //     mat_t.setFromTriplets(entries_t.begin(), entries_t.end());

        //     StiffnessMatrix A = mat_t * mat;
        //     Eigen::MatrixXd b = mat_t * node_rhs;

        //     json params = {
        //         {"mtype", -2}, // matrix type for Pardiso (2 = SPD)
        //         // {"max_iter", 0}, // for iterative solvers
        //         // {"tolerance", 1e-9}, // for iterative solvers
        //     };
        //     auto solver = LinearSolver::create("", "");
        //     solver->setParameters(params);
        //     solver->analyzePattern(A);
        //     solver->factorize(A);

        //     for(int d = 0; d < dim; ++d)
        //         solver->solve(b.col(d), new_nodes.col(d));
        // }

        // for(int i = 0; i < n_elements; ++i)
        // {
        //     auto &base = gbases[i];

        //     if(!mesh.is_cube(i))
        //         continue;

        //     auto &lbs = base.bases;
        //     const int n_local_bases = int(lbs.size());
        //     for(int j = 0; j < n_local_bases; ++j)
        //     {
        //         Basis &b = lbs[j];

        //         for(std::size_t ii = 0; ii < b.global().size(); ++ii)
        //         {
        //             // if(nodes.is_primitive_boundary(b.global()[ii].index))
        //                 // continue;

        //             for(int d = 0; d < dim; ++d)
        //             {
        //                 b.global()[ii].node(d) = new_nodes(b.global()[ii].index, d);
        //             }
        //         }
        //     }
        // }
    }

}


#include "general_quantum_operator.hpp"

#include <Eigen/Dense>
#include <csim/stat_ops.hpp>
#include <csim/update_ops.hpp>
#include <csim/update_ops_dm.hpp>
#include <cstring>
#include <fstream>
#include <numeric>

#include "gate_factory.hpp"
#include "pauli_operator.hpp"
#include "state.hpp"
#include "type.hpp"
#include "utility.hpp"

#ifdef _USE_GPU
#include <gpusim/update_ops_cuda.h>
#endif

GeneralQuantumOperator::GeneralQuantumOperator(const UINT qubit_count)
    : _qubit_count(qubit_count), _is_hermitian(true) {}

GeneralQuantumOperator::~GeneralQuantumOperator() {
    for (auto& term : this->_operator_list) {
        delete term;
    }
}

void GeneralQuantumOperator::add_operator(const PauliOperator* mpt) {
    PauliOperator* _mpt = mpt->copy();
    if (!check_Pauli_operator(this, _mpt)) {
        std::stringstream error_message_stream;
        error_message_stream
            << "Error: GeneralQuantumOperator::add_operator(const "
               "PauliOperator*): pauli_operator applies target qubit of "
               "which the index is larger than qubit_count";
        throw std::invalid_argument(error_message_stream.str());
    }
    if (this->_is_hermitian && std::abs(_mpt->get_coef().imag()) > 0) {
        this->_is_hermitian = false;
    }
    this->_operator_list.push_back(_mpt);
}

void GeneralQuantumOperator::add_operator(
    CPPCTYPE coef, std::string pauli_string) {
    PauliOperator* _mpt = new PauliOperator(pauli_string, coef);
    if (!check_Pauli_operator(this, _mpt)) {
        std::stringstream error_message_stream;
        error_message_stream
            << "Error: "
               "GeneralQuantumOperator::add_operator(double,std::string):"
               " pauli_operator applies target qubit of which the index "
               "is larger than qubit_count";
        throw std::invalid_argument(error_message_stream.str());
    }
    if (this->_is_hermitian && std::abs(coef.imag()) > 0) {
        this->_is_hermitian = false;
    }
    this->add_operator(_mpt);
    delete _mpt;
}

CPPCTYPE GeneralQuantumOperator::get_expectation_value(
    const QuantumStateBase* state) const {
    if (this->_qubit_count > state->qubit_count) {
        std::stringstream error_message_stream;
        error_message_stream
            << "Error: GeneralQuantumOperator::get_expectation_value(const "
               "QuantumStateBase*): invalid qubit count";
        throw std::invalid_argument(error_message_stream.str());
    }
    auto sum = std::accumulate(this->_operator_list.cbegin(),
        this->_operator_list.cend(), (CPPCTYPE)0.0,
        [&](CPPCTYPE acc, PauliOperator* pauli) {
            return acc + pauli->get_expectation_value(state);
        });
    return sum;
}

CPPCTYPE GeneralQuantumOperator::get_expectation_value_single_thread(
    const QuantumStateBase* state) const {
    if (this->_qubit_count > state->qubit_count) {
        std::cerr
            << "Error: GeneralQuantumOperator::get_expectation_value(const "
               "QuantumStateBase*): invalid qubit count"
            << std::endl;
        return 0.;
    }
    auto sum = std::accumulate(this->_operator_list.cbegin(),
        this->_operator_list.cend(), (CPPCTYPE)0.0,
        [&](CPPCTYPE acc, PauliOperator* pauli) {
            return acc + pauli->get_expectation_value_single_thread(state);
        });
    return sum;
}

CPPCTYPE GeneralQuantumOperator::get_transition_amplitude(
    const QuantumStateBase* state_bra,
    const QuantumStateBase* state_ket) const {
    if (this->_qubit_count > state_bra->qubit_count ||
        state_bra->qubit_count != state_ket->qubit_count) {
        std::stringstream error_message_stream;
        error_message_stream
            << "Error: GeneralQuantumOperator::get_transition_amplitude(const "
               "QuantumStateBase*, const QuantumStateBase*): invalid qubit "
               "count";
        throw std::invalid_argument(error_message_stream.str());
    }

    auto sum = std::accumulate(this->_operator_list.cbegin(),
        this->_operator_list.cend(), (CPPCTYPE)0.0,
        [&](CPPCTYPE acc, PauliOperator* pauli) {
            return acc + pauli->get_transition_amplitude(state_bra, state_ket);
        });
    return sum;
}

void GeneralQuantumOperator::add_random_operator(const UINT operator_count) {
    const auto qubit_count = this->get_qubit_count();
    Random random;
    for (UINT operator_index = 0; operator_index < operator_count;
         operator_index++) {
        auto target_qubit_index_list = std::vector<UINT>(qubit_count, 0);
        auto target_qubit_pauli_list = std::vector<UINT>(qubit_count, 0);
        for (UINT qubit_index = 0; qubit_index < qubit_count; qubit_index++) {
            const UINT pauli_id = random.int32() % 4;
            target_qubit_index_list.at(qubit_index) = qubit_index;
            target_qubit_pauli_list.at(qubit_index) = pauli_id;
        }
        // -1.0 <= coef <= 1.0
        const CPPCTYPE coef = random.uniform() * 2 - 1.0;
        auto pauli_operator = PauliOperator(
            target_qubit_index_list, target_qubit_pauli_list, coef);
        this->add_operator(&pauli_operator);
    }
}

CPPCTYPE
GeneralQuantumOperator::solve_ground_state_eigenvalue_by_arnoldi_method(
    QuantumStateBase* state, const UINT iter_count, const CPPCTYPE mu) const {
    if (this->get_term_count() == 0) {
        std::stringstream error_message_stream;
        error_message_stream
            << "Error: "
               "GeneralQuantumOperator::solve_ground_state_eigenvalue_by_"
               "arnoldi_method("
               "QuantumStateBase * state, const UINT iter_count, const "
               "CPPCTYPE mu): At least one PauliOperator is required.";
        throw std::invalid_argument(error_message_stream.str());
    }

    // Implemented based on
    // https://files.transtutors.com/cdn/uploadassignments/472339_1_-numerical-linear-aljebra.pdf
    const auto qubit_count = this->get_qubit_count();
    auto present_state = QuantumState(qubit_count);
    auto tmp_state = QuantumState(qubit_count);
    auto multiplied_state = QuantumState(qubit_count);
    auto mu_timed_state = QuantumState(qubit_count);

    // Vectors composing Krylov subspace.
    std::vector<QuantumStateBase*> state_list;
    state_list.reserve(iter_count + 1);
    state->normalize(state->get_squared_norm());
    state_list.push_back(state->copy());

    CPPCTYPE mu_;
    if (mu == 0.0) {
        // mu is not changed from default value.
        mu_ = this->calculate_default_mu();
    } else {
        mu_ = mu;
    }

    ComplexMatrix hessenberg_matrix =
        ComplexMatrix::Zero(iter_count, iter_count);
    for (UINT i = 0; i < iter_count; i++) {
        mu_timed_state.load(state_list[i]);
        mu_timed_state.multiply_coef(-mu_);
        this->apply_to_state(&tmp_state, *state_list[i], &multiplied_state);
        multiplied_state.add_state(&mu_timed_state);

        for (UINT j = 0; j < i + 1; j++) {
            const auto coef = state::inner_product(
                static_cast<QuantumState*>(state_list[j]), &multiplied_state);
            hessenberg_matrix(j, i) = coef;
            tmp_state.load(state_list[j]);
            tmp_state.multiply_coef(-coef);
            multiplied_state.add_state(&tmp_state);
        }

        const auto norm = multiplied_state.get_squared_norm();
        if (i != iter_count - 1) {
            hessenberg_matrix(i + 1, i) = std::sqrt(norm);
        }
        multiplied_state.normalize(norm);
        state_list.push_back(multiplied_state.copy());
    }

    Eigen::ComplexEigenSolver<ComplexMatrix> eigen_solver(hessenberg_matrix);
    const auto eigenvalues = eigen_solver.eigenvalues();
    const auto eigenvectors = eigen_solver.eigenvectors();

    // Find ground state eigenvalue and eigenvector.
    UINT minimum_eigenvalue_index = 0;
    auto minimum_eigenvalue = eigenvalues[0];
    for (UINT i = 0; i < eigenvalues.size(); i++) {
        if (eigenvalues[i].real() < minimum_eigenvalue.real()) {
            minimum_eigenvalue_index = i;
            minimum_eigenvalue = eigenvalues[i];
        }
    }

    // Compose ground state vector and store it to `state`.
    present_state.multiply_coef(0.0);
    for (UINT i = 0; i < state_list.size() - 1; i++) {
        tmp_state.load(state_list[i]);
        tmp_state.multiply_coef(eigenvectors(i, minimum_eigenvalue_index));
        present_state.add_state(&tmp_state);
    }
    state->load(&present_state);

    // Free states allocated by `QuantumState::copy()`.
    for (auto used_state : state_list) {
        delete used_state;
    }
    return minimum_eigenvalue + mu_;
}

CPPCTYPE GeneralQuantumOperator::solve_ground_state_eigenvalue_by_power_method(
    QuantumStateBase* state, const UINT iter_count, const CPPCTYPE mu) const {
    if (this->get_term_count() == 0) {
        std::stringstream error_message_stream;
        error_message_stream
            << "Error: "
               "GeneralQuantumOperator::solve_ground_state_eigenvalue_by_"
               "power_method("
               "QuantumStateBase * state, const UINT iter_count, const "
               "CPPCTYPE mu): At least one PauliOperator is required.";
        throw std::invalid_argument(error_message_stream.str());
    }

    CPPCTYPE mu_;
    if (mu == 0.0) {
        // mu is not changed from default value.
        mu_ = this->calculate_default_mu();
    } else {
        mu_ = mu;
    }

    // Stores a result of A|a>
    auto multiplied_state = QuantumState(state->qubit_count);
    // Stores a result of -\mu|a>
    auto mu_timed_state = QuantumState(state->qubit_count);
    auto work_state = QuantumState(state->qubit_count);
    for (UINT i = 0; i < iter_count; i++) {
        mu_timed_state.load(state);
        mu_timed_state.multiply_coef(-mu_);

        multiplied_state.multiply_coef(0.0);
        this->apply_to_state(&work_state, *state, &multiplied_state);
        state->load(&multiplied_state);
        state->add_state(&mu_timed_state);
        state->normalize(state->get_squared_norm());
    }
    return this->get_expectation_value(state) + mu;
}

void GeneralQuantumOperator::apply_to_state(QuantumStateBase* work_state,
    const QuantumStateBase& state_to_be_multiplied,
    QuantumStateBase* dst_state) const {
    if (state_to_be_multiplied.qubit_count != dst_state->qubit_count) {
        throw std::invalid_argument(
            "Qubit count of state_to_be_multiplied and dst_state must be the "
            "same");
    }

    dst_state->multiply_coef(0.0);
    const auto term_count = this->get_term_count();
    for (UINT i = 0; i < term_count; i++) {
        work_state->load(&state_to_be_multiplied);
        const auto term = this->get_term(i);
        auto pauli_operator =
            gate::Pauli(term->get_index_list(), term->get_pauli_id_list());
        pauli_operator->update_quantum_state(work_state);
        work_state->multiply_coef(term->get_coef());
        dst_state->add_state(work_state);
        delete pauli_operator;
    }
}

void GeneralQuantumOperator::apply_to_state(
    QuantumStateBase* state, QuantumStateBase* dst_state) const {
    if (state->qubit_count != dst_state->qubit_count) {
        throw std::invalid_argument(
            "Qubit count of state_to_be_multiplied and dst_state must be the "
            "same");
    }

    dst_state->set_zero_norm_state();
    const auto term_count = this->get_term_count();
    for (UINT i = 0; i < term_count; i++) {
        const auto term = this->get_term(i);
        _apply_pauli_to_state(
            term->get_pauli_id_list(), term->get_index_list(), state);
        dst_state->add_state_with_coef(term->get_coef(), state);
        _apply_pauli_to_state(
            term->get_pauli_id_list(), term->get_index_list(), state);
    }
}

void GeneralQuantumOperator::apply_to_state_single_thread(
    QuantumStateBase* state, QuantumStateBase* dst_state) const {
    if (state->qubit_count != dst_state->qubit_count) {
        throw std::invalid_argument(
            "Qubit count of state_to_be_multiplied and dst_state must be the "
            "same");
    }

    dst_state->set_zero_norm_state();
    const auto term_count = this->get_term_count();
    for (UINT i = 0; i < term_count; i++) {
        const auto term = this->get_term(i);
        _apply_pauli_to_state_single_thread(
            term->get_pauli_id_list(), term->get_index_list(), state);
        dst_state->add_state_with_coef_single_thread(term->get_coef(), state);
        _apply_pauli_to_state_single_thread(
            term->get_pauli_id_list(), term->get_index_list(), state);
    }
}

void GeneralQuantumOperator::_apply_pauli_to_state(
    std::vector<UINT> pauli_id_list, std::vector<UINT> target_index_list,
    QuantumStateBase* state) const {
    // this function is same as the gate::Pauli update quantum state
    if (state->is_state_vector()) {
#ifdef _USE_GPU
        if (state->get_device_name() == "gpu") {
            multi_qubit_Pauli_gate_partial_list_host(target_index_list.data(),
                pauli_id_list.data(), (UINT)target_index_list.size(),
                state->data(), state->dim, state->get_cuda_stream(),
                state->device_number);
            // _update_func_gpu(this->_target_qubit_list[0].index(), _angle,
            // state->data(), state->dim);
            return;
        }
#endif
        multi_qubit_Pauli_gate_partial_list(target_index_list.data(),
            pauli_id_list.data(), (UINT)target_index_list.size(),
            state->data_c(), state->dim);
    } else {
        dm_multi_qubit_Pauli_gate_partial_list(target_index_list.data(),
            pauli_id_list.data(), (UINT)target_index_list.size(),
            state->data_c(), state->dim);
    }
}

void GeneralQuantumOperator::_apply_pauli_to_state_single_thread(
    std::vector<UINT> pauli_id_list, std::vector<UINT> target_index_list,
    QuantumStateBase* state) const {
    // this function is same as the gate::Pauli update quantum state
    if (state->is_state_vector()) {
#ifdef _USE_GPU
        if (state->get_device_name() == "gpu") {
            // TODO: make it single thread for this function
            multi_qubit_Pauli_gate_partial_list_host(target_index_list.data(),
                pauli_id_list.data(), (UINT)target_index_list.size(),
                state->data(), state->dim, state->get_cuda_stream(),
                state->device_number);
            // _update_func_gpu(this->_target_qubit_list[0].index(), _angle,
            // state->data(), state->dim);
            return;
        }
#endif
        multi_qubit_Pauli_gate_partial_list_single_thread(
            target_index_list.data(), pauli_id_list.data(),
            (UINT)target_index_list.size(), state->data_c(), state->dim);
    } else {
        // TODO: make it single thread for this function
        dm_multi_qubit_Pauli_gate_partial_list(target_index_list.data(),
            pauli_id_list.data(), (UINT)target_index_list.size(),
            state->data_c(), state->dim);
    }
}

CPPCTYPE GeneralQuantumOperator::calculate_default_mu() const {
    double mu = 0.0;
    const auto term_count = this->get_term_count();
    for (UINT i = 0; i < term_count; i++) {
        const auto term = this->get_term(i);
        mu += std::abs(term->get_coef().real());
    }
    return static_cast<CPPCTYPE>(mu);
}

GeneralQuantumOperator* GeneralQuantumOperator::copy() const {
    auto quantum_operator = new GeneralQuantumOperator(_qubit_count);
    for (auto pauli : this->_operator_list) {
        quantum_operator->add_operator(pauli->copy());
    }
    return quantum_operator;
}

GeneralQuantumOperator GeneralQuantumOperator::operator+(
    const GeneralQuantumOperator& target) const {
    auto res = this->copy();
    *res += target;
    return *res;
}

GeneralQuantumOperator* GeneralQuantumOperator::get_dagger() const {
    auto quantum_operator = new GeneralQuantumOperator(_qubit_count);
    for (auto pauli : this->_operator_list) {
        quantum_operator->add_operator(
            std::conj(pauli->get_coef()), pauli->get_pauli_string());
    }
    return quantum_operator;
}

GeneralQuantumOperator GeneralQuantumOperator::operator+(
    const PauliOperator& target) const {
    auto res = this->copy();
    *res += target;
    return *res;
}

GeneralQuantumOperator& GeneralQuantumOperator::operator+=(
    const GeneralQuantumOperator& target) {
    ITYPE i, j;
    auto terms = target.get_terms();
#pragma omp parallel for
    for (i = 0; i < _operator_list.size(); i++) {
        auto pauli_operator = _operator_list[i];
        for (j = 0; j < terms.size(); j++) {
            auto target_operator = terms[j];
            auto pauli_x = pauli_operator->get_x_bits();
            auto pauli_z = pauli_operator->get_z_bits();
            auto target_x = target_operator->get_x_bits();
            auto target_z = target_operator->get_z_bits();
            if (pauli_x.size() != target_x.size()) {
                UINT max_size = std::max(pauli_x.size(), target_x.size());
                pauli_x.resize(max_size);
                pauli_z.resize(max_size);
                target_x.resize(max_size);
                target_z.resize(max_size);
            }
            if (pauli_x == target_x && pauli_z == target_z) {
                _operator_list[i]->change_coef(_operator_list[i]->get_coef() +
                                               target_operator->get_coef());
            }
        }
    }
    for (j = 0; j < terms.size(); j++) {
        auto target_operator = terms[j];
        bool flag = true;
        for (i = 0; i < _operator_list.size(); i++) {
            auto pauli_operator = _operator_list[i];
            auto pauli_x = pauli_operator->get_x_bits();
            auto pauli_z = pauli_operator->get_z_bits();
            auto target_x = target_operator->get_x_bits();
            auto target_z = target_operator->get_z_bits();
            if (pauli_x.size() != target_x.size()) {
                UINT max_size = std::max(pauli_x.size(), target_x.size());
                pauli_x.resize(max_size);
                pauli_z.resize(max_size);
                target_x.resize(max_size);
                target_z.resize(max_size);
            }
            if (pauli_x == target_x && pauli_z == target_z) {
                flag = false;
            }
        }
        if (flag) {
            this->add_operator(target_operator->copy());
        }
    }
    return *this;
}

GeneralQuantumOperator& GeneralQuantumOperator::operator+=(
    const PauliOperator& target) {
    bool flag = true;
    ITYPE i;
#pragma omp parallel for
    for (i = 0; i < _operator_list.size(); i++) {
        auto pauli_operator = _operator_list[i];
        auto pauli_x = pauli_operator->get_x_bits();
        auto pauli_z = pauli_operator->get_z_bits();
        auto target_x = target.get_x_bits();
        auto target_z = target.get_z_bits();
        if (pauli_x.size() != target_x.size()) {
            UINT max_size = std::max(pauli_x.size(), target_x.size());
            pauli_x.resize(max_size);
            pauli_z.resize(max_size);
            target_x.resize(max_size);
            target_z.resize(max_size);
        }
        if (pauli_x == target_x && pauli_z == target_z) {
            _operator_list[i]->change_coef(
                _operator_list[i]->get_coef() + target.get_coef());
            flag = false;
        }
    }
    if (flag) {
        this->add_operator(target.copy());
    }
    return *this;
}

GeneralQuantumOperator GeneralQuantumOperator::operator-(
    const GeneralQuantumOperator& target) const {
    auto res = this->copy();
    *res -= target;
    return *res;
}

GeneralQuantumOperator GeneralQuantumOperator::operator-(
    const PauliOperator& target) const {
    auto res = this->copy();
    *res -= target;
    return *res;
}

GeneralQuantumOperator& GeneralQuantumOperator::operator-=(
    const GeneralQuantumOperator& target) {
    ITYPE i, j;
    auto terms = target.get_terms();
#pragma omp parallel for
    for (i = 0; i < _operator_list.size(); i++) {
        auto pauli_operator = _operator_list[i];
        for (j = 0; j < terms.size(); j++) {
            auto target_operator = terms[j];
            auto pauli_x = pauli_operator->get_x_bits();
            auto pauli_z = pauli_operator->get_z_bits();
            auto target_x = target_operator->get_x_bits();
            auto target_z = target_operator->get_z_bits();
            if (pauli_x.size() != target_x.size()) {
                UINT max_size = std::max(pauli_x.size(), target_x.size());
                pauli_x.resize(max_size);
                pauli_z.resize(max_size);
                target_x.resize(max_size);
                target_z.resize(max_size);
            }
            if (pauli_x == target_x && pauli_z == target_z) {
                _operator_list[i]->change_coef(_operator_list[i]->get_coef() -
                                               target_operator->get_coef());
            }
        }
    }
    for (j = 0; j < terms.size(); j++) {
        auto target_operator = terms[j];
        bool flag = true;
        for (i = 0; i < _operator_list.size(); i++) {
            auto pauli_operator = _operator_list[i];
            auto pauli_x = pauli_operator->get_x_bits();
            auto pauli_z = pauli_operator->get_z_bits();
            auto target_x = target_operator->get_x_bits();
            auto target_z = target_operator->get_z_bits();
            if (pauli_x.size() != target_x.size()) {
                UINT max_size = std::max(pauli_x.size(), target_x.size());
                pauli_x.resize(max_size);
                pauli_z.resize(max_size);
                target_x.resize(max_size);
                target_z.resize(max_size);
            }
            if (pauli_x == target_x && pauli_z == target_z) {
                flag = false;
            }
        }
        if (flag) {
            auto copy = target_operator->copy();
            copy->change_coef(-copy->get_coef());
            this->add_operator(copy);
        }
    }
    return *this;
}

GeneralQuantumOperator& GeneralQuantumOperator::operator-=(
    const PauliOperator& target) {
    bool flag = true;
    ITYPE i;
    for (i = 0; i < _operator_list.size(); i++) {
        auto pauli_operator = _operator_list[i];
        auto pauli_x = pauli_operator->get_x_bits();
        auto pauli_z = pauli_operator->get_z_bits();
        auto target_x = target.get_x_bits();
        auto target_z = target.get_z_bits();
        if (pauli_x.size() != target_x.size()) {
            UINT max_size = std::max(pauli_x.size(), target_x.size());
            pauli_x.resize(max_size);
            pauli_z.resize(max_size);
            target_x.resize(max_size);
            target_z.resize(max_size);
        }
        if (pauli_x == target_x && pauli_z == target_z) {
            _operator_list[i]->change_coef(
                _operator_list[i]->get_coef() - target.get_coef());
            flag = false;
        }
    }
    if (flag) {
        auto copy = target.copy();
        copy->change_coef(-copy->get_coef());
        this->add_operator(copy);
    }
    return *this;
}
GeneralQuantumOperator GeneralQuantumOperator::operator*(
    const GeneralQuantumOperator& target) const {
    auto res = this->copy();
    *res *= target;
    return *res;
}

GeneralQuantumOperator GeneralQuantumOperator::operator*(
    const PauliOperator& target) const {
    auto res = this->copy();
    *res *= target;
    return *res;
}
GeneralQuantumOperator GeneralQuantumOperator::operator*(
    CPPCTYPE target) const {
    auto res = this->copy();
    *res *= target;
    return *res;
}

GeneralQuantumOperator& GeneralQuantumOperator::operator*=(
    const GeneralQuantumOperator& target) {
    auto copy = this->copy();
    _operator_list.clear();
    auto terms = copy->get_terms();
    auto target_terms = target.get_terms();
    ITYPE i, j;
#pragma omp parallel for
    for (i = 0; i < terms.size(); i++) {
        auto pauli_operator = terms[i];
        for (j = 0; j < target_terms.size(); j++) {
            auto target_operator = target_terms[j];
            PauliOperator* product = new PauliOperator;
            *product = (*pauli_operator) * (*target_operator);
            *this += *product;
        }
    }
    return *this;
}

GeneralQuantumOperator& GeneralQuantumOperator::operator*=(
    const PauliOperator& target) {
    auto copy = this->copy();
    _operator_list.clear();
    ITYPE i;
    auto terms = copy->get_terms();
#pragma omp parallel for
    for (i = 0; i < terms.size(); i++) {
        auto pauli_operator = terms[i];
        PauliOperator* product = new PauliOperator;
        *product = (*pauli_operator) * (target);
        *this += *product;
    }
    return *this;
}

GeneralQuantumOperator& GeneralQuantumOperator::operator*=(CPPCTYPE target) {
    ITYPE i;
#pragma omp parallel for
    for (i = 0; i < _operator_list.size(); i++) {
        *_operator_list[i] *= target;
    }
    return *this;
}
namespace quantum_operator {
GeneralQuantumOperator* create_general_quantum_operator_from_openfermion_file(
    std::string file_path) {
    UINT qubit_count = 0;
    std::vector<CPPCTYPE> coefs;
    std::vector<std::string> ops;

    // loading lines and check qubit_count
    std::string str_buf;
    std::vector<std::string> index_list;

    std::ifstream ifs;
    std::string line;
    ifs.open(file_path);

    while (getline(ifs, line)) {
        std::tuple<double, double, std::string> parsed_items =
            parse_openfermion_line(line);
        const auto coef_real = std::get<0>(parsed_items);
        const auto coef_imag = std::get<1>(parsed_items);
        str_buf = std::get<2>(parsed_items);

        CPPCTYPE coef(coef_real, coef_imag);
        coefs.push_back(coef);
        ops.push_back(str_buf);
        index_list = split(str_buf, "IXYZ ");

        for (UINT i = 0; i < index_list.size(); ++i) {
            UINT n = std::stoi(index_list[i]) + 1;
            if (qubit_count < n) qubit_count = n;
        }
    }
    if (!ifs.eof()) {
        std::stringstream error_message_stream;
        error_message_stream << "ERROR: Invalid format";
        throw std::runtime_error(error_message_stream.str());
    }
    ifs.close();

    GeneralQuantumOperator* general_quantum_operator =
        new GeneralQuantumOperator(qubit_count);

    for (UINT i = 0; i < ops.size(); ++i) {
        general_quantum_operator->add_operator(
            new PauliOperator(ops[i].c_str(), coefs[i]));
    }

    return general_quantum_operator;
}

GeneralQuantumOperator* create_general_quantum_operator_from_openfermion_text(
    std::string text) {
    UINT qubit_count = 0;
    std::vector<CPPCTYPE> coefs;
    std::vector<std::string> ops;

    std::string str_buf;
    std::vector<std::string> index_list;

    std::vector<std::string> lines;
    lines = split(text, "\n");
    for (std::string line : lines) {
        std::tuple<double, double, std::string> parsed_items =
            parse_openfermion_line(line);
        const auto coef_real = std::get<0>(parsed_items);
        const auto coef_imag = std::get<1>(parsed_items);
        str_buf = std::get<2>(parsed_items);

        CPPCTYPE coef(coef_real, coef_imag);
        coefs.push_back(coef);
        ops.push_back(str_buf);
        index_list = split(str_buf, "IXYZ ");

        for (UINT i = 0; i < index_list.size(); ++i) {
            UINT n = std::stoi(index_list[i]) + 1;
            if (qubit_count < n) qubit_count = n;
        }
    }
    GeneralQuantumOperator* general_quantum_operator =
        new GeneralQuantumOperator(qubit_count);

    for (UINT i = 0; i < ops.size(); ++i) {
        general_quantum_operator->add_operator(
            new PauliOperator(ops[i].c_str(), coefs[i]));
    }

    return general_quantum_operator;
}

std::pair<GeneralQuantumOperator*, GeneralQuantumOperator*>
create_split_general_quantum_operator(std::string file_path) {
    UINT qubit_count = 0;
    std::vector<CPPCTYPE> coefs;
    std::vector<std::string> ops;

    std::ifstream ifs;
    ifs.open(file_path);

    if (!ifs) {
        std::stringstream error_message_stream;
        error_message_stream << "ERROR: Cannot open file";
        throw std::runtime_error(error_message_stream.str());
    }

    // loading lines and check qubit_count
    std::string str_buf;
    std::vector<std::string> index_list;

    std::string line;
    while (getline(ifs, line)) {
        std::tuple<double, double, std::string> parsed_items =
            parse_openfermion_line(line);
        const auto coef_real = std::get<0>(parsed_items);
        const auto coef_imag = std::get<1>(parsed_items);
        str_buf = std::get<2>(parsed_items);
        if (str_buf == (std::string)NULL) {
            continue;
        }
        CPPCTYPE coef(coef_real, coef_imag);
        coefs.push_back(coef);
        ops.push_back(str_buf);
        index_list = split(str_buf, "IXYZ ");

        for (UINT i = 0; i < index_list.size(); ++i) {
            UINT n = std::stoi(index_list[i]) + 1;
            if (qubit_count < n) qubit_count = n;
        }
    }
    if (!ifs.eof()) {
        std::stringstream error_message_stream;
        error_message_stream << "ERROR: Invalid format";
        throw std::runtime_error(error_message_stream.str());
    }
    ifs.close();

    GeneralQuantumOperator* general_quantum_operator_diag =
        new GeneralQuantumOperator(qubit_count);
    GeneralQuantumOperator* general_quantum_operator_non_diag =
        new GeneralQuantumOperator(qubit_count);

    for (UINT i = 0; i < ops.size(); ++i) {
        if (ops[i].find("X") != std::string::npos ||
            ops[i].find("Y") != std::string::npos) {
            general_quantum_operator_non_diag->add_operator(
                new PauliOperator(ops[i].c_str(), coefs[i]));
        } else {
            general_quantum_operator_diag->add_operator(
                new PauliOperator(ops[i].c_str(), coefs[i]));
        }
    }

    return std::make_pair(
        general_quantum_operator_diag, general_quantum_operator_non_diag);
}
}  // namespace quantum_operator

bool check_Pauli_operator(const GeneralQuantumOperator* quantum_operator,
    const PauliOperator* pauli_operator) {
    auto vec = pauli_operator->get_index_list();
    UINT val = 0;
    if (vec.size() > 0) {
        val = std::max(val, *std::max_element(vec.begin(), vec.end()));
    }
    return val < (quantum_operator->get_qubit_count());
}

std::string GeneralQuantumOperator::to_string() const {
    std::stringstream os;
    auto term_count = this->get_term_count();
    for (UINT index = 0; index < term_count; index++) {
        os << this->get_term(index)->get_coef() << " ";
        os << this->get_term(index)->get_pauli_string();
        if (index != term_count - 1) {
            os << " + ";
        }
    }
    return os.str();
}


/*! \file meshlist.cpp
    \brief アルゴンに対して、分子動力学シミュレーションを行うクラスの宣言

    Copyright ©  2017 @dc1394 All Rights Reserved.
    This software is released under the BSD 2-Clause License.
*/

#include "meshlist.h"
#include <algorithm>        // for std::fill
#include <boost/assert.hpp> // for BOOST_ASSERT

namespace moleculardynamics {
    MeshList::MeshList(double periodiclen) : periodiclen_(periodiclen)
    {
        auto const SL = SystemParam::RCUTOFF + SystemParam::MARGIN;

        m_ = static_cast<std::int32_t>(periodiclen / SL) - 1;
        mesh_size_ = static_cast<double>(periodiclen) / m_;

        BOOST_ASSERT(m_ > 2);
        BOOST_ASSERT(mesh_size_ > SL);

        number_of_mesh_ = m_ * m_ * m_;
        count_.resize(number_of_mesh_);
        indexes_.resize(number_of_mesh_);
    }

    void MeshList::make_pair(SystemParam::myatomvector & atoms, SystemParam::mypairvector & pairs)
    {
        pairs.clear();

        auto const pn = atoms.size();

        std::vector<std::int32_t> particle_position(pn, 0);
        std::vector<std::int32_t> pointer(number_of_mesh_, 0);

        std::fill(count_.begin(), count_.end(), 0);

        auto const im = 1.0 / mesh_size_;
        for (auto i = 0; i < pn; i++) {
            auto ix = static_cast<std::int32_t>(atoms[i].r[0] * im);
            auto iy = static_cast<std::int32_t>(atoms[i].r[1] * im);
            auto iz = static_cast<std::int32_t>(atoms[i].r[2] * im);

            if (ix < 0) {
                ix += m_;
            }
            else if (ix >= m_) {
                ix -= m_;
            }

            if (iy < 0) {
                iy += m_;
            }
            else if (iy >= m_) {
                iy -= m_;
            }
            if (iz < 0) {
                iz += m_;
            }
            else if (iz >= m_) {
                iz -= m_;
            }

            auto const index = ix + iy * m_ + iz * m_ * m_;

            BOOST_ASSERT(index >= 0);
            BOOST_ASSERT(index < number_of_mesh_);

            count_[index]++;
            particle_position[i] = index;
        }

        indexes_[0] = 0;
        auto sum = 0;

        for (auto i = 0; i < number_of_mesh_ - 1; i++) {
            sum += count_[i];
            indexes_[i + 1] = sum;
        }

        for (auto i = 0; i < pn; i++) {
            auto const pos = particle_position[i];
            auto const j = indexes_[pos] + pointer[pos];
            sorted_buffer[j] = i;
            ++pointer[pos];
        }

        for (auto i = 0; i < number_of_mesh_; i++) {
            search(i, atoms, pairs);
        }
    }

    void MeshList::search_other(std::int32_t id, std::int32_t ix, std::int32_t iy, std::int32_t iz, SystemParam::myatomvector & atoms, SystemParam::mypairvector & pairs)
    {
        if (ix < 0) {
            ix += m_;
        }
        else if (ix >= m_) {
            ix -= m_;
        }

        if (iy < 0) {
            iy += m_;
        }
        else if (iy >= m_) {
            iy -= m_;
        }

        if (iz < 0) {
            iz += m_;
        }
        else if (iz >= m_) {
            iz -= m_;
        }

        auto const id2 = ix + iy * m_ + iz * m_ * m_;

        for (auto k = indexes_[id]; k < indexes_[id] + count_[id]; k++) {
            for (auto m_ = indexes_[id2]; m_ < indexes_[id2] + count_[id2]; m_++) {
                auto const i = sorted_buffer[k];
                auto const j = sorted_buffer[m_];

                Eigen::Vector4d d = atoms[j].r - atoms[i].r;

                SystemParam::adjust_periodic(d, periodiclen_);

                if (d.squaredNorm() <= SystemParam::ML2) {
                    pairs.push_back(std::make_pair(i, j));
                }
            }
        }
    }

    void MeshList::search(std::int32_t id, SystemParam::myatomvector & atoms, SystemParam::mypairvector & pairs)
    {
        auto const ix = id % m_;
        auto const iy = (id / m_) % m_;
        auto const iz = (id / m_ / m_);

        search_other(id, ix + 1, iy, iz, atoms, pairs);
        search_other(id, ix - 1, iy + 1, iz, atoms, pairs);
        search_other(id, ix, iy + 1, iz, atoms, pairs);
        search_other(id, ix + 1, iy + 1, iz, atoms, pairs);

        search_other(id, ix - 1, iy, iz + 1, atoms, pairs);
        search_other(id, ix, iy, iz + 1, atoms, pairs);
        search_other(id, ix + 1, iy, iz + 1, atoms, pairs);

        search_other(id, ix - 1, iy - 1, iz + 1, atoms, pairs);
        search_other(id, ix, iy - 1, iz + 1, atoms, pairs);
        search_other(id, ix + 1, iy - 1, iz + 1, atoms, pairs);

        search_other(id, ix - 1, iy + 1, iz + 1, atoms, pairs);
        search_other(id, ix, iy + 1, iz + 1, atoms, pairs);
        search_other(id, ix + 1, iy + 1, iz + 1, atoms, pairs);

        // Registration of self box
        auto const si = indexes_[id];
        auto const n = count_[id];

        for (auto k = si; k < si + n - 1; k++) {
            for (auto m_ = k + 1; m_ < si + n; m_++) {
                auto const i = sorted_buffer[k];
                auto const j = sorted_buffer[m_];

                Eigen::Vector4d d = atoms[j].r - atoms[i].r;

                SystemParam::adjust_periodic(d, periodiclen_);

                if (d.squaredNorm() <= SystemParam::ML2) {
                    pairs.push_back(std::make_pair(i, j));
                }
            }
        }
    }
}