/*
providing tools for structure manipulation
*/

#include "atomsciflow/utils.h"

#include "atomsciflow/base/crystal.h"
#include "atomsciflow/base/atom.h"

#include <algorithm>
#include <iostream>
#include <cstdlib>
#include <cmath>
#include <armadillo>

namespace atomsciflow {

std::string version() {
    return "0.0.0";
}

int move_along(atomsciflow::Crystal* structure, std::vector<int> atoms_to_move, std::vector<int> direc, double disp) {
    /*
    :param structure: an instance of atomsciflow::Crystal()
    :params atoms_to_move: a list of atoms to move, counting starts with 0
    :param direc: three int number to indicate the direction to move along
        namely the crystal orientation index
    :param disp: displacement of the atoms in unit of Angstrom
    */
    // do some checking
    if (*std::max_element(atoms_to_move.begin(), atoms_to_move.end()) > (structure->natom() - 1)) {
        std::cout << "===================================================================" << std::endl;
        std::cout << "                              WARNING \n";
        std::cout << "------------------------------------------------------------------------------\n";
        std::cout << "the atom you are trying to move is beyond the number of atoms in the structure\n";
        std::exit(1);
    }

    arma::vec direc_cartesian(3);
    arma::mat latcell(3, 3);
    std::vector<std::vector<double>> cell = structure->get_cell();
    latcell.row(0) = arma::conv_to<arma::rowvec>::from(cell[0]);
    latcell.row(1) = arma::conv_to<arma::rowvec>::from(cell[1]);
    latcell.row(2) = arma::conv_to<arma::rowvec>::from(cell[2]);
    direc_cartesian = latcell.row(0) * direc[0] + latcell.row(1) * direc[1] + latcell.row(2) * direc[2];

    // normalize
    double length = arma::norm(direc_cartesian);
    // length = np.sqrt(direc_cartesian[0]**2+direc_cartesian[1]**2+direc_cartesian[2]**2)
    direc_cartesian = direc_cartesian / length;

    double deltax = direc_cartesian[0] * disp;
    double deltay = direc_cartesian[1] * disp;
    double deltaz = direc_cartesian[2] * disp;

    for (auto i : atoms_to_move) {
        structure->atoms[i].x += deltax;
        structure->atoms[i].y += deltay;
        structure->atoms[i].z += deltaz;
    }
    return 0;
}

int remove_atoms(atomsciflow::Crystal* structure, std::vector<int> atoms_to_remove) {
    /*
    :param structure: an instance of atomsciflow::Crystal()
    :params atoms_to_remove: a list of atoms to remove, counting starts with 0
    */
    // do some checking
    if (*std::max_element(atoms_to_remove.begin(), atoms_to_remove.end()) > (structure->natom() - 1)) {
        std::cout << "===================================================================" << std::endl;
        std::cout << "                              WARNING \n";
        std::cout << "------------------------------------------------------------------------------\n";
        std::cout << "the atom you are trying to remove is beyond the number of atoms in the structure\n";
        std::exit(1);
    }

    structure->remove_atoms(atoms_to_remove);
    //end
    return 0;
}

int vacuum_layer(atomsciflow::Crystal* structure, int plane, double thickness) {
    /*
    :param structure: an instance of atomsciflow::Crystal()
    :param plane: the plane chosen to add vacuum layer, can be:
        1 -> ab plane
        2 -> ac plane
        3 -> bc plane
    :param thickness: thickness of the vacuum layer
    */
    //
    //

    arma::mat latcell(3, 3);
    latcell.row(0) = arma::conv_to<arma::rowvec>::from(structure->cell[0]);
    latcell.row(1) = arma::conv_to<arma::rowvec>::from(structure->cell[1]);
    latcell.row(2) = arma::conv_to<arma::rowvec>::from(structure->cell[2]);

    if (plane == 1) {
        std::cout << "the normal_of_ab" << std::endl;
        arma::vec normal_of_ab = arma::cross(arma::conv_to<arma::vec>::from(latcell.row(0)), arma::conv_to<arma::vec>::from(latcell.row(1)));
        std::cout << "got the normal_of_ab" << std::endl;
        // get the cosine of the angle between c and outer product of ab
        double cosangle = arma::dot(latcell.row(2), normal_of_ab) / arma::norm(latcell.row(2)) / arma::norm(normal_of_ab);
        double proj_c_on_ab_normal = arma::norm(latcell.row(2)) * std::abs(cosangle);
        std::vector<double> z_all;
        for (auto& atom : structure->atoms) {
            z_all.push_back(atom.z);
        }
        std::cout << "got all z for atoms" << std::endl;

        double factor = (*std::max_element(z_all.begin(), z_all.end()) - *std::min_element(z_all.begin(), z_all.end()) + thickness) / proj_c_on_ab_normal * 1.0;
        structure->cell[2] = arma::conv_to<std::vector<double>>::from(latcell.row(2) * factor);
    } else if (plane == 2) {
        arma::vec normal_of_ac = arma::cross(arma::conv_to<arma::vec>::from(latcell.row(0)), arma::conv_to<arma::vec>::from(latcell.row(2)));
        // get the cosine of the angle between b and outer product of ac
        double cosangle = arma::dot(latcell.row(1), normal_of_ac) / arma::norm(latcell.row(1)) / arma::norm(normal_of_ac);
        double proj_b_on_ac_normal = arma::norm(latcell.row(1)) * std::abs(cosangle);
        std::vector<double> y_all;
        for (auto& atom : structure->atoms) {
            y_all.push_back(atom.y);
        }
        double factor = (*std::max_element(y_all.begin(), y_all.end()) - *std::min_element(y_all.begin(), y_all.end()) + thickness) / proj_b_on_ac_normal * 1.0;
        structure->cell[1] = arma::conv_to<std::vector<double>>::from(latcell.row(1) * factor);
    } else if (plane == 3) {
        arma::vec normal_of_bc = arma::cross(arma::conv_to<arma::vec>::from(latcell.row(1)), arma::conv_to<arma::vec>::from(latcell.row(2)));
        // get the cosine of the angle between b and outer product of ac
        double cosangle = arma::dot(latcell.row(0), normal_of_bc) / arma::norm(latcell.row(0)) / arma::norm(normal_of_bc);
        double proj_a_on_bc_normal = arma::norm(latcell.row(0)) * std::abs(cosangle);
        std::vector<double> x_all;
        for (auto& atom : structure->atoms) {
            x_all.push_back(atom.x);
        }
        double factor = (*std::max_element(x_all.begin(), x_all.end()) - *std::min_element(x_all.begin(), x_all.end()) + thickness) / proj_a_on_bc_normal * 1.0;
        structure->cell[0] = arma::conv_to<std::vector<double>>::from(latcell.row(0) * factor);
    }

    // end
    return 0;
}

int set_frac_min_to_zero(atomsciflow::Crystal* structure) {
    /*
    :return an object of atomsciflow::Crystal()
    Note:
        set the fractional coordinate minimum to zero, this is a way of standardize the cif file
    */

    int i = 0;

    // now calc the fractional coordinates
    std::vector<atomsciflow::Atom> atoms_frac{structure->atoms};

    arma::mat latcell(3, 3);
    latcell.row(0) = arma::conv_to<arma::rowvec>::from(structure->cell[0]);
    latcell.row(1) = arma::conv_to<arma::rowvec>::from(structure->cell[1]);
    latcell.row(2) = arma::conv_to<arma::rowvec>::from(structure->cell[2]);

    arma::mat convmat = arma::inv(arma::trans(latcell));

    std::vector<double> cart_tmp(3);
    arma::vec frac_tmp(3);

    for (i = 0; i < structure->natom(); i++) {
        //
        cart_tmp.push_back(structure->atoms[i].x);
        cart_tmp.push_back(structure->atoms[i].y);
        cart_tmp.push_back(structure->atoms[i].z);
        //frac_tmp = arma::dot(convmat, arma::vec(cart_tmp));
        frac_tmp = convmat * arma::vec(cart_tmp);
        atoms_frac[i].x = frac_tmp[0];
        atoms_frac[i].y = frac_tmp[1];
        atoms_frac[i].z = frac_tmp[2];
    }

    // set the minimum of fractional coord to to 0
    std::vector<double> values;
    for (i = 0; i < structure->natom(); i++) {
        values.push_back(atoms_frac[i].x);
    }
    double min_frac_x;
    min_frac_x = *std::min_element(values.begin(), values.end());

    values.clear();
    for (i = 0; i < structure->natom(); i++) {
        values.push_back(atoms_frac[i].y);
    }
    double min_frac_y = *std::min_element(values.begin(), values.end());

    values.clear();
    for (i = 0; i < structure->natom(); i++) {
        values.push_back(atoms_frac[i].z);
    }
    double min_frac_z = *std::min_element(values.begin(), values.end());

    for (i = 0; i < structure->natom(); i++) {
        atoms_frac[i].x -= min_frac_x;
        atoms_frac[i].y -= min_frac_y;
        atoms_frac[i].z -= min_frac_z;
    }

    // now convert coord of atom in atoms_frac_within_new_cell to cartesian

    // update convmat for frac to cart conversion
    convmat = arma::trans(latcell);

    arma::vec cart_final;
    for (i = 0; i < structure->natom(); i++) {
        //
        frac_tmp.at(0) = atoms_frac[i].x;
        frac_tmp.at(1) = atoms_frac[i].y;
        frac_tmp.at(2) = atoms_frac[i].z;
        //cart_final = arma::dot(convmat, frac_tmp);
        cart_final = convmat * frac_tmp;
        structure->atoms[i].set_x(cart_final[0]);
        structure->atoms[i].set_y(cart_final[1]);
        structure->atoms[i].set_z(cart_final[2]);
    }

    return 0;
    //
}

int set_frac_within_zero_and_one(atomsciflow::Crystal* structure) {
    /*
    :return an object of atomsciflow::Crystal()
    Note:
        set the fractional coordinate within the range of 0 and 1, this is a way of standardize the cif file
    */

    int i = 0; // for iteration
    int j = 0; // for iteration

    // now conv to the fractional coordinates
    arma::mat latcell(3, 3);
    latcell.row(0) = arma::conv_to<arma::rowvec>::from(structure->cell[0]);
    latcell.row(1) = arma::conv_to<arma::rowvec>::from(structure->cell[1]);
    latcell.row(2) = arma::conv_to<arma::rowvec>::from(structure->cell[2]);

    arma::mat convmat = arma::inv(arma::trans(latcell));
    arma::vec tmp_arma_vec(3);
    //std::vector<double> tmp_cpp_vec(3);
    for (i = 0; i < structure->natom(); i ++) {
        tmp_arma_vec(0) = structure->atoms[i].x;
        tmp_arma_vec(1) = structure->atoms[i].y;
        tmp_arma_vec(2) = structure->atoms[i].z;
        //tmp_arma_vec = arma::dot(convmat, tmp_arma_vec);
        tmp_arma_vec = convmat * tmp_arma_vec;
        structure->atoms[i].x = tmp_arma_vec(0);
        structure->atoms[i].y = tmp_arma_vec(1);
        structure->atoms[i].z = tmp_arma_vec(2);
    }

    // set the fractional coordinates within 0 and 1
    for (i = 0; i < structure->atoms.size(); i++) {
        while (structure->atoms[i].x >= 1) {
            structure->atoms[i].x -= 1;
        }
        while (structure->atoms[i].x < 0) {
            structure->atoms[i].x += 1;
        }
        while (structure->atoms[i].y >= 1) {
            structure->atoms[i].y -= 1;
        }
        while (structure->atoms[i].y < 0) {
            structure->atoms[i].y += 1;
        }
        while (structure->atoms[i].z >= 1) {
            structure->atoms[i].z -= 1;
        }
        while (structure->atoms[i].z < 0) {
            structure->atoms[i].z += 1;
        }
    }

    // now convert coord of atom in atoms_frac_within_new_cell to cartesian
    // update convmat for convertion from frac to cart
    convmat = arma::trans(latcell);

    for (i = 0; i < structure->atoms.size(); i++) {
        tmp_arma_vec(0) = structure->atoms[i].x;
        tmp_arma_vec(1) = structure->atoms[i].y;
        tmp_arma_vec(2) = structure->atoms[i].z;
        //tmp_arma_vec = arma::dot(convmat, tmp_arma_vec);
        tmp_arma_vec = convmat * tmp_arma_vec;
        structure->atoms[i].x = tmp_arma_vec(0);
        structure->atoms[i].y = tmp_arma_vec(1);
        structure->atoms[i].z = tmp_arma_vec(2);
    }

    return 0;
}

int inverse_geo_center(atomsciflow::Crystal* structure) {
    /*
    calc the geometric center of the system and make an inversion against that center
    :param structure: an instance of atomsciflow::Crystal()
    */

    int i = 0; // for iteration

    // calc the geometric center
    double x = 0;
    double y = 0;
    double z = 0;

    for (auto atom : structure->atoms) {
        x += atom.x;
        y += atom.y;
        z += atom.z;
    }

    x /= structure->atoms.size();
    y /= structure->atoms.size();
    z /= structure->atoms.size();

    // now get the symmetry image against the geometric center
    for (i = 0; i < structure->natom(); i++) {
        structure->atoms[i].x = x * 2 - structure->atoms[i].x;
        structure->atoms[i].y = y * 2 - structure->atoms[i].y;
        structure->atoms[i].z = z * 2 - structure->atoms[i].z;
    }

    // end
    return 0;
}

int inverse_point(atomsciflow::Crystal* structure, std::vector<double> point) {
    /*
    calc the geometric center of the system and make an inversion against that center
    :param structure: an instance of pymatflow.structure.crystal.crystal()
    :param point: the inverse center point, like [0.0, 0.0, 0.0]
    */

    // now get the symmetry image against the inverse center
    for (auto& atom : structure->atoms) {
        atom.x = point[0] * 2 - atom.x;
        atom.y = point[1] * 2 - atom.y;
        atom.z = point[2] * 2 - atom.z;
    }

    // end
    return 0;
}

int inverse_cell_center(atomsciflow::Crystal* structure) {
    /*
    make an inversion against the cell center
    :param structure: an instance of pymatflow.structure.crystal.crystal()
    */

    int i = 0; // for iteration

    //first transfer to fractional coordinate and inverse against [0.5, 0.5, 0.5]
    arma::mat latcell(3, 3);
    latcell.row(0) = arma::conv_to<arma::rowvec>::from(structure->cell[0]);
    latcell.row(1) = arma::conv_to<arma::rowvec>::from(structure->cell[1]);
    latcell.row(2) = arma::conv_to<arma::rowvec>::from(structure->cell[2]);

    arma::mat convmat = arma::inv(arma::trans(latcell));

    arma::vec arma_vec(3);

    for (i = 0; i < structure->natom(); i++) {
        arma_vec[0] = structure->atoms[i].x;
        arma_vec[1] = structure->atoms[i].y;
        arma_vec[2] = structure->atoms[i].z;
        //arma_vec = arma::dot(convmat, arma_vec);
        arma_vec = convmat * arma_vec;
        structure->atoms[i].x = arma_vec[0];
        structure->atoms[i].y = arma_vec[1];
        structure->atoms[i].z = arma_vec[2];
    }

    for (i = 0; i < structure->natom(); i++) {
        structure->atoms[i].x = 0.5 * 2 - structure->atoms[i].x;
        structure->atoms[i].y = 0.5 * 2 - structure->atoms[i].y;
        structure->atoms[i].z = 0.5 * 2 - structure->atoms[i].z;
    }

    // convert frac to cartesian again
    // update convmat for convertion from frac to cart
    convmat = arma::trans(latcell);

    for (i = 0; i < structure->natom(); i++) {
        arma_vec[0] = structure->atoms[i].x;
        arma_vec[1] = structure->atoms[i].y;
        arma_vec[2] = structure->atoms[i].z;
        //arma_vec = arma::dot(convmat, arma_vec);
        arma_vec = convmat * arma_vec;
        structure->atoms[i].x = arma_vec[0];
        structure->atoms[i].y = arma_vec[1];
        structure->atoms[i].z = arma_vec[2];
    }

    //
    return 0;
}

int rotate_along_axis(atomsciflow::Crystal* structure, std::vector<int> rotate_atoms, std::vector<int> axis) {
    /*
    rotate the specified atoms along the specified axis
    :param structure: an instance of pymatflow.structure.crystal.crystal()
    */
    return 0;
}

int enlarge_atoms(atomsciflow::Crystal* structure) {
    /*
    :return out:
        atoms: [
                ["C", 0.00000, 0.000000, 0.0000],
                ["O", 0.00000, 0.500000, 0.0000],
                ...
            ]
    Note: will enlarge the atoms in the unit cell along both a, b, c and -a, -b, -c direction.
        The goal is to make sure when the cell rotate in the 3D space, it will always be filled
        with atoms.
    */

    /*
    from pymatflow.base.atom import Atom
    #
    cell = copy.deepcopy(structure.cell)
    a = np.linalg.norm(cell[0])
    b = np.linalg.norm(cell[1])
    c = np.linalg.norm(cell[2])

    n1 = np.ceil(np.max([a, b, c]) / a ) * 2 # maybe times 2 is not needed
    n2 = np.ceil(np.max([a, b, c]) / b ) * 2
    n3 = np.ceil(np.max([a, b, c]) / c ) * 2
    n = [int(n1), int(n2), int(n3)]
    print(n)

    atoms = copy.deepcopy(structure.atoms)
    # build supercell: replica in three vector one by one
    for i in range(3):
        natom_now = len(atoms)
        for j in range(n[i] - 1):
            for atom in atoms[:natom_now]:
                x = atom.x + float(j + 1) * structure.cell[i][0]
                y = atom.y + float(j + 1) * structure.cell[i][1]
                z = atom.z + float(j + 1) * structure.cell[i][2]
                atoms.append(Atom(atom.name, x, y, z))
        # replicate in the negative direction of structure.cell[i]
        for atom in atoms[:natom_now*n[i]]:
            x = atom.x - float(n[i]) * structure.cell[i][0]
            y = atom.y - float(n[i]) * structure.cell[i][1]
            z = atom.z - float(n[i]) * structure.cell[i][2]
            atoms.append(Atom(atom.name, x, y, z))
    return [[atom.name, atom.x, atom.y, atom.z] for atom in atoms]

    */
    return 0;
}

std::vector<atomsciflow::Atom> enlarge_atoms_new_cell(atomsciflow::Crystal* structure, std::vector<std::vector<double> > new_cell) {
    /*
    Note: will enlarge the atoms in the unit cell along both a, b, c and -a, -b, -c direction of the new_cell !!!
        but the cell is not redefined, the returned atoms is not used to form crystal, but to be
        tailored by redefine_lattice function to get atoms for the redfined lattice.
        The goal is to make sure when the new cell rotate in the 3D space, it will always be filled
        with atoms.
    */

    int i = 0;
    int j = 0;
    int k = 0;
    // std::cout << "testing..." << std::endl;
    //
    arma::mat latcell(3, 3);
    latcell.row(0) = arma::conv_to<arma::rowvec>::from(structure->cell[0]);
    latcell.row(1) = arma::conv_to<arma::rowvec>::from(structure->cell[1]);
    latcell.row(2) = arma::conv_to<arma::rowvec>::from(structure->cell[2]);

    double a = arma::norm(latcell.row(0));
    double b = arma::norm(latcell.row(1));
    double c = arma::norm(latcell.row(2));

    arma::mat latcell_new(3, 3);
    latcell_new.row(0) = arma::conv_to<arma::rowvec>::from(new_cell[0]);
    latcell_new.row(1) = arma::conv_to<arma::rowvec>::from(new_cell[1]);
    latcell_new.row(2) = arma::conv_to<arma::rowvec>::from(new_cell[2]);

    double new_a = arma::norm(latcell_new.row(0));
    double new_b = arma::norm(latcell_new.row(1));
    double new_c = arma::norm(latcell_new.row(2));

    std::cout << "new_a: " << new_a << "; new_b: " << new_b << "; new_c: " << new_c << std::endl;

    std::vector<double> cpp_vec(3);
    cpp_vec[0] = new_a;
    cpp_vec[1] = new_b;
    cpp_vec[2] = new_c;
    //std::cout << "testing..." << std::endl;
    int n1 = std::ceil(*std::max_element(cpp_vec.begin(), cpp_vec.end()) / a) * 2; // maye times 2 is not needed
    int n2 = std::ceil(*std::max_element(cpp_vec.begin(), cpp_vec.end()) / b) * 2;
    int n3 = std::ceil(*std::max_element(cpp_vec.begin(), cpp_vec.end()) / c) * 2;

    std::vector<int> n(3);
    n[0] = n1;
    n[1] = n2;
    n[2] = n3;

    std::vector<atomsciflow::Atom> atoms = structure->atoms;

    std::cout << "n1: " << n[0] << std::endl;
    std::cout << "n2: " << n[1] << std::endl;
    std::cout << "n3: " << n[2] << std::endl;

    //std::cout << atoms[0].name << std::endl;
    // build supercell: replica in three vector one by one
    std::cout << "atoms.size(): " << atoms.size() << std::endl;
    int natom_now = 0;
    double x, y, z;
    for (i = 0; i < 3; i++) {
        //
        natom_now = atoms.size();
        for (j = 0; j < (n[i] - 1); j++) {
            for (k = 0; k < natom_now; k++) {
                //std::cout << "testing" << std::endl;
                x = atoms[k].x + (j + 1) * structure->cell[i][0];
                y = atoms[k].y + (j + 1) * structure->cell[i][1];
                z = atoms[k].z + (j + 1) * structure->cell[i][2];
                //std::cout << "testing." << std::endl;
                atomsciflow::Atom atm;
                //std::cout << "testing.." << std::endl;
                atm.name = atoms[k].name;
                atm.x = x;
                atm.y = y;
                atm.z = z;
                //std::cout << "testing..." << std::endl;
                atoms.push_back(atm);
                //std::cout << "testing." << std::endl;
            }
        }

        //std::cout << "natom_now-> " << natom_now << std::endl;
        // std::cout << "n[" << i << "]-> " << n[i] << std::endl;
        // replicate in the negative direction of sturcture.cell[i]
        for (k = 0; k < (natom_now * n[i]); k++) {
            x = atoms[k].x - n[i] * structure->cell[i][0];
            y = atoms[k].y - n[i] * structure->cell[i][1];
            z = atoms[k].z - n[i] * structure->cell[i][2];
            atomsciflow::Atom atm;
            atm.name = atoms[k].name;
            atm.x = x;
            atm.y = y;
            atm.z = z;
            //std::cout << "testing.." << std::endl;
            atoms.push_back(atm);
        }

    }

    return atoms;
}

int redefine_lattice(atomsciflow::Crystal* structure, std::vector<int> a, std::vector<int> b, std::vector<int> c, double precision=1.0e-8) {
    /*
    :param a, b, c: new lattice vectors in terms of old.
        new_a = a[0] * old_a + a[1] * old_b + a[2] * old_c
        like a=[1, 0, 0], b=[0, 1, 0], c=[0, 0, 1] actually defines the
        same lattice as old.
    :param precision, a value that is less than 1 and infinitely close to 1
        used to judge whether one atom is in another periodic of the redefined cell

    Method:
        first make a large enough supercell, which guarantee that all the atoms in the new lattice are inside
        the supercell.
        then redfine the cell, and calc the fractional coord of all atoms with regarding the new cell
        finally remove those atoms who's fractional coord is not within range [0, 1), and we can convert fractional
        coords to cartesian.
    Note:
        relationship among convertion of coords. the most important point is that all coords actually have one common
        reference system, namely the General XYZ coordinate system. all the cell are defined with XYZ as reference,
        and the convmat build from the cell(with XYZ as reference) can be applied only to atoms also with XYZ as ference,
        finally we convert frac to cartesian using convmat also defined using cell with XYZ as reference, so we get
        the cartesian with general XYZ as reference. In the last all the coord of atoms and cell have the general
        XYZ  system as reference. So it works!
    */

    int i = 0; // for iteration
    int j = 0; // for iteration

    std::cout << "a: " << a[0] << " " << a[1] << " " << a[2] << std::endl;
    std::cout << "b: " << b[0] << " " << b[1] << " " << b[2] << std::endl;
    std::cout << "c: " << c[0] << " " << c[1] << " " << c[2] << std::endl;

    std::cout << "structure->cell[0]: " << structure->cell[0][0] << " " << structure->cell[0][1] << " " << structure->cell[0][2] << std::endl;
    std::cout << "structure->cell[1]: " << structure->cell[1][0] << " " << structure->cell[1][1] << " " << structure->cell[1][2] << std::endl;
    std::cout << "structure->cell[0]: " << structure->cell[2][0] << " " << structure->cell[2][1] << " " << structure->cell[2][2] << std::endl;

    arma::mat latcell_old(3, 3);
    latcell_old.row(0) = arma::conv_to<arma::rowvec>::from(structure->cell[0]);
    latcell_old.row(1) = arma::conv_to<arma::rowvec>::from(structure->cell[1]);
    latcell_old.row(2) = arma::conv_to<arma::rowvec>::from(structure->cell[2]);

    std::cout << "latcell_old[0]: " << latcell_old(0, 0) << " " << latcell_old(0, 1) << " " << latcell_old(0, 2) << std::endl;
    std::cout << "latcell_old[1]: " << latcell_old(1, 0) << " " << latcell_old(1, 1) << " " << latcell_old(1, 2) << std::endl;
    std::cout << "latcell_old[2]: " << latcell_old(2, 0) << " " << latcell_old(2, 1) << " " << latcell_old(2, 2) << std::endl;

    arma::mat latcell_new(3, 3);

    latcell_new.row(0) =  a[0] * latcell_old.row(0) + a[1] * latcell_old.row(1) + a[2] * latcell_old.row(2);
    latcell_new.row(1) =  b[0] * latcell_old.row(0) + b[1] * latcell_old.row(1) + b[2] * latcell_old.row(2);
    latcell_new.row(2) =  c[0] * latcell_old.row(0) + c[1] * latcell_old.row(1) + c[2] * latcell_old.row(2);

    std::cout << "latcell_new[0]: " << latcell_new(0, 0) << " " << latcell_new(0, 1) << " " << latcell_new(0, 2) << std::endl;
    std::cout << "latcell_new[1]: " << latcell_new(1, 0) << " " << latcell_new(1, 1) << " " << latcell_new(1, 2) << std::endl;
    std::cout << "latcell_new[2]: " << latcell_new(2, 0) << " " << latcell_new(2, 1) << " " << latcell_new(2, 2) << std::endl;

    // enlarge the system
    std::vector<std::vector<double> > new_cell;
    new_cell.push_back(arma::conv_to<std::vector<double> >::from(latcell_new.row(0)));
    new_cell.push_back(arma::conv_to<std::vector<double> >::from(latcell_new.row(1)));
    new_cell.push_back(arma::conv_to<std::vector<double> >::from(latcell_new.row(2)));

    std::cout << "new_cell[0]: " << new_cell[0][0] << " " << new_cell[0][1] << " " << new_cell[0][2] << std::endl;
    std::cout << "new_cell[1]: " << new_cell[1][0] << " " << new_cell[1][1] << " " << new_cell[1][2] << std::endl;
    std::cout << "new_cell[2]: " << new_cell[2][0] << " " << new_cell[2][1] << " " << new_cell[2][2] << std::endl;

    //std::cout << "testing..." << std::endl;

    std::vector<atomsciflow::Atom> atoms_container = enlarge_atoms_new_cell(structure, new_cell);

    std::cout << "redefine_lattice: enlarge_atoms_new_cell finished!" << std::endl;

    // now calc the fractional coordinates of all atoms in atoms_container with new_cell as reference
    arma::mat convmat(3, 3);
    convmat = arma::inv(arma::trans(latcell_new));
    arma::vec arma_vec(3);

    for (i = 0; i < atoms_container.size(); i++) {
        arma_vec(0) = atoms_container[i].x;
        arma_vec(1) = atoms_container[i].y;
        arma_vec(2) = atoms_container[i].z;
        arma_vec = convmat * arma_vec;
        atoms_container[i].x = arma_vec(0);
        atoms_container[i].y = arma_vec(1);
        atoms_container[i].z = arma_vec(2);
    }

    std::cout << "redefine_lattice: atoms_container convert to frac coord finished!" << std::endl;
    std::cout << "atoms_container.size(): " << atoms_container.size() << std::endl;

    std::vector<atomsciflow::Atom> atoms_frac_within_new_cell;

    for (i = 0; i < atoms_container.size(); i++) {

        if ((0 <= atoms_container[i].x) && (atoms_container[i].x < (1 - precision))) {
            if ((0 <= atoms_container[i].y) && (atoms_container[i].y < (1 - precision))) {
                if ((0 <= atoms_container[i].z) && (atoms_container[i].z < (1 - precision))) {
                    atoms_frac_within_new_cell.push_back(atoms_container[i]);
                }
            }
        }

    }

    std::cout << "atoms_frac_within_new_cell->size(): " << atoms_frac_within_new_cell.size() << std::endl;

    // now convert coord of atom in atoms_frac_within_new_cell to cartesian

    convmat = arma::trans(latcell_new);
    structure->atoms.clear();

    for (i = 0; i < atoms_frac_within_new_cell.size(); i++) {
        arma_vec[0] = atoms_frac_within_new_cell[i].x;
        arma_vec[1] = atoms_frac_within_new_cell[i].y;
        arma_vec[2] = atoms_frac_within_new_cell[i].z;
        arma_vec = convmat * arma_vec;
        atomsciflow::Atom atm;
        atm.name = atoms_frac_within_new_cell[i].name;
        atm.x = arma_vec[0];
        atm.y = arma_vec[1];
        atm.z = arma_vec[2];
        structure->atoms.push_back(atm);
    }

    structure->cell = new_cell;

    //
    return 0;
}

int cleave_surface(atomsciflow::Crystal* structure, std::vector<int> direction, double thickness=10, double precision=1.0e-8) {
    /*
    :param structure: an instance of crystal()
    :param direction: direction of the surface plane, like [0, 0, 1], the reference of it is three lattice
        vector a, b, c
    :param precision, a value that is less than 1 and infinitely close to 1
        used to judge whether one atom is in another periodic of the redefined cell used in cleave surface

    :return an object of crystal()

    Note: make use of redefine_lattice() to cleave surface.
        we try to find new_a and new_b which form the surface plane (the direction is normal to the plane)
    */

    arma::mat old_cell(3, 3);
    old_cell.row(0) = arma::conv_to<arma::rowvec>::from(structure->cell[0]);
    old_cell.row(1) = arma::conv_to<arma::rowvec>::from(structure->cell[1]);
    old_cell.row(2) = arma::conv_to<arma::rowvec>::from(structure->cell[2]);

    arma::vec direction_xyz_as_ref(3);
    direction_xyz_as_ref = arma::conv_to<arma::vec>::from(direction[0] * old_cell.row(0) + direction[1] * old_cell.row(1) + direction[2] * old_cell.row(2));
    std::cout << "direction_xyz_as_ref.size(): "  << direction_xyz_as_ref.size() << std::endl;

    std::vector<int> a_from_old;
    std::vector<int> b_from_old;
    std::vector<int> c_from_old;

    std::vector<int> iter_ijk;

    int i = 0;
    int j = 0;
    int k = 0;

    iter_ijk.push_back(0);

    for (i = 0; i < 16; i++) {
        iter_ijk.push_back(i);
        iter_ijk.push_back(-i);
    }

    arma::vec new_vec(3);
    arma::vec new_a(3);

    std::cout << "begin search for new a and b" << std::endl;

    for (auto iter_i : iter_ijk) {
        i = iter_i;
        if ((! a_from_old.empty()) && (! b_from_old.empty())) {
            break;
        }

        for (auto iter_j : iter_ijk) {
            j = iter_j;
            if ((! a_from_old.empty()) && (! b_from_old.empty())) {
                break;
            }
            for (auto iter_k : iter_ijk) {
                k = iter_k;
                if ((! a_from_old.empty()) && (! b_from_old.empty())) {
                    break;
                }
                if ((i == 0) && (j == 0) && (k == 0 )) {
                    continue;
                }
                new_vec = arma::conv_to<arma::vec>::from(i * old_cell.row(0) + j * old_cell.row(1) + k * old_cell.row(2));
                //std::cout << "get new vec..." << std::endl;
                //std::cout << "new_vec.size(): "  << new_vec.size() << std::endl;
                //std::cout << "direction_xyz_as_ref.size(): "  << direction_xyz_as_ref.size() << std::endl;
                if (arma::dot(new_vec, direction_xyz_as_ref) == 0) {
                    if (a_from_old.empty()) {
                        a_from_old.push_back(i);
                        a_from_old.push_back(j);
                        a_from_old.push_back(k);
                        new_a = arma::conv_to<arma::vec>::from(a_from_old[0] * old_cell.row(0) + a_from_old[1] * old_cell.row(1) + a_from_old[2] * old_cell.row(2));
                        std::cout << "get new a..." << std::endl;
                        std::cout << "get a_from_old" << std::endl;
                        continue;
                    } else if (b_from_old.empty()) {
                        double cosangle = arma::dot(new_vec, new_a) / arma::norm(new_vec) / arma::norm(new_a);
                        double sinangle = std::sin(std::acos(cosangle));
                        std::cout << "sin: " << sinangle << std::endl;
                        std::cout << "cos: " << cosangle << std::endl;
                        if (std::abs(sinangle - 0) < 1.0e-3) {
                            continue;
                        } else if (1.0e-5 < cosangle && cosangle < 1) {
                            // ignore angle between [0, 90] by default
                            // like we ignore angle  = 60 degree for a and b here, we seek for 120 instead
                            continue;
                        } else {
                            std::cout << "get b_from_old" << std::endl;
                            b_from_old.push_back(i);
                            b_from_old.push_back(j);
                            b_from_old.push_back(k);
                            continue;
                        }
                    }
                }
            }
        }
    }

    std::cout << "begin search for new c" << std::endl;

    arma::vec new_c(3);
    std::vector<std::vector<int>> new_c_ijk;
    std::vector<double> new_c_sin;

    for (auto i_iter : iter_ijk) {
        i = i_iter;
        for (auto j_iter : iter_ijk) {
            j = j_iter;
            for (auto k_iter : iter_ijk) {
                k = k_iter;
                new_c = arma::conv_to<arma::vec>::from(i * old_cell.row(0) + j * old_cell.row(1) + k * old_cell.row(2));
                double cosangle = arma::dot(new_c, direction_xyz_as_ref) / arma::norm(new_c) / arma::norm(direction_xyz_as_ref);
                double sinangle = std::sin(std::acos(cosangle));
                std::vector<int> ijk;
                ijk.push_back(i);
                ijk.push_back(j);
                ijk.push_back(k);
                new_c_ijk.push_back(ijk);
                new_c_sin.push_back(sinangle);
            }
        }
    }

    std::vector<double> new_c_sin_remove_nan;
    for (i = 0; i < new_c_sin.size(); i++) {
        if (std::isnan(new_c_sin[i])) {
            continue;
        } else {
            new_c_sin_remove_nan.push_back(new_c_sin[i]);
        }
    }

    double min_sin = *std::min_element(new_c_sin_remove_nan.begin(), new_c_sin_remove_nan.end());

    std::cout << "min sin is: " << min_sin << std::endl;

    for (i = 0; i < new_c_sin.size(); i++) {
        if (new_c_sin[i] == min_sin) {
            c_from_old = new_c_ijk[i];
            break;
        }
    }

    std::cout << "a-> " << a_from_old[0] << " " << a_from_old[1] << " " << a_from_old[2] << std::endl;
    std::cout << "b-> " << b_from_old[0] << " " << b_from_old[1] << " " << b_from_old[2] << std::endl;
    std::cout << "c-> " << c_from_old[0] << " " << c_from_old[1] << " " << c_from_old[2] << std::endl;

    // redefine lattice
    std::cout << "redefine the lattice" << std::endl;
    redefine_lattice(structure, a_from_old, b_from_old, c_from_old, precision);

    std::cout << "add vacuum layer" << std::endl;
    vacuum_layer(structure, 1, thickness);

    return 0;
}

atomsciflow::Crystal merge_layers(atomsciflow::Crystal* structure1, atomsciflow::Crystal* structure2, int use_cell=0, double distance=3.4, double thickness=10) {
    /*
    :param structure1: an instance of crystal()
    :param structure2: an instance of crystal()
    :param use_cell: use cell parameter of structure 1 or 2 or 0(average) to set the new a b cell parameter
        attention: c vector is not handled this way

    :param distance: the distance between layers
    :param thickness: the vaccum layer thickness of the combined system

    :return an object of crystal()
    Note:
        only merge layers with ab plane as the surface plane
    */

    int i = 0; // for iteration
    int j = 0; // for iteration
    int k = 0; // for iteration

    set_frac_within_zero_and_one(structure1);
    set_frac_within_zero_and_one(structure2);

    std::cout << "frac of structure 1 and 2 already set to be within zero and one" << std::endl;

    arma::mat old_cell_1(3, 3);
    old_cell_1.row(0) = arma::conv_to<arma::rowvec>::from(structure1->cell[0]);
    old_cell_1.row(1) = arma::conv_to<arma::rowvec>::from(structure1->cell[1]);
    old_cell_1.row(2) = arma::conv_to<arma::rowvec>::from(structure1->cell[2]);

    arma::mat old_cell_2(3, 3);
    old_cell_2.row(0) = arma::conv_to<arma::rowvec>::from(structure2->cell[0]);
    old_cell_2.row(1) = arma::conv_to<arma::rowvec>::from(structure2->cell[1]);
    old_cell_2.row(2) = arma::conv_to<arma::rowvec>::from(structure2->cell[2]);

    // first transfer to fractional coordinate
    std::cout << "first transfer to fractional coordinates" << std::endl;

    arma::mat latcell_old_1(3, 3);
    latcell_old_1.row(0) = arma::conv_to<arma::rowvec>::from(structure1->cell[0]);
    latcell_old_1.row(1) = arma::conv_to<arma::rowvec>::from(structure1->cell[1]);
    latcell_old_1.row(2) = arma::conv_to<arma::rowvec>::from(structure1->cell[2]);

    arma::mat latcell_old_2(3, 3);
    latcell_old_2.row(0) = arma::conv_to<arma::rowvec>::from(structure2->cell[0]);
    latcell_old_2.row(1) = arma::conv_to<arma::rowvec>::from(structure2->cell[1]);
    latcell_old_2.row(2) = arma::conv_to<arma::rowvec>::from(structure2->cell[2]);

    arma::mat convmat(3, 3);
    convmat = arma::inv(arma::trans(latcell_old_1));

    arma::vec arma_vec(3);

    std::vector<atomsciflow::Atom> frac_1;
    std::vector<atomsciflow::Atom> frac_2;

    for (i = 0; i < structure1->atoms.size(); i++) {
        arma_vec(0) = structure1->atoms[i].x;
        arma_vec(1) = structure1->atoms[i].y;
        arma_vec(2) = structure1->atoms[i].z;
        //arma_vec = arma::dot(convmat, arma_vec);
        arma_vec = convmat * arma_vec;
        atomsciflow::Atom atm;
        atm.name = structure1->atoms[i].name;
        atm.x = arma_vec(0);
        atm.y = arma_vec(1);
        atm.z = arma_vec(2);
        frac_1.push_back(atm);
    }

    convmat = arma::inv(arma::trans(latcell_old_2));

    for (i = 0; i < structure2->atoms.size(); i++) {
        arma_vec(0) = structure2->atoms[i].x;
        arma_vec(1) = structure2->atoms[i].y;
        arma_vec(2) = structure2->atoms[i].z;
        //arma_vec = arma::dot(convmat, arma_vec);
        arma_vec = convmat * arma_vec;
        atomsciflow::Atom atm;
        atm.name = structure2->atoms[i].name;
        atm.x = arma_vec(0);
        atm.y = arma_vec(1);
        atm.z = arma_vec(2);
        frac_2.push_back(atm);
    }

    //
    //
    arma::mat average_cell = (latcell_old_1 + latcell_old_2) / 2;

    atomsciflow::Crystal out;

    arma::mat latcell_frac_to_cart_1(3, 3);
    arma::mat latcell_frac_to_cart_2(3, 3);

    if (use_cell == 1) {
        latcell_frac_to_cart_1 = old_cell_1;
        latcell_frac_to_cart_2.row(0) = old_cell_1.row(0);
        latcell_frac_to_cart_2.row(1) = old_cell_1.row(1);
        latcell_frac_to_cart_2.row(2) = old_cell_2.row(2);
    } else if (use_cell == 2) {
        latcell_frac_to_cart_2 = old_cell_2;
        latcell_frac_to_cart_1.row(0) = old_cell_2.row(0);
        latcell_frac_to_cart_1.row(1) = old_cell_2.row(1);
        latcell_frac_to_cart_1.row(2) = old_cell_1.row(2);
    } else {
        //average_ab
        latcell_frac_to_cart_1.row(0) = (old_cell_1.row(0) + old_cell_2.row(0)) / 2;
        latcell_frac_to_cart_1.row(1) = (old_cell_1.row(1) + old_cell_2.row(1)) / 2;
        latcell_frac_to_cart_1.row(2) = old_cell_1.row(2);
        latcell_frac_to_cart_2.row(0) = (old_cell_1.row(0) + old_cell_2.row(0)) / 2;
        latcell_frac_to_cart_2.row(1) = (old_cell_1.row(1) + old_cell_2.row(1)) / 2;
        latcell_frac_to_cart_2.row(2) = old_cell_2.row(2);
    }

    // convert frac to cartesian again
    std::cout << "then transfer back to cartesian coordinates" << std::endl;
    arma::mat convmat_1 = arma::trans(latcell_frac_to_cart_1);
    arma::mat convmat_2 = arma::trans(latcell_frac_to_cart_2);

    std::vector<atomsciflow::Atom> cart_1;
    for (i = 0; i < frac_1.size(); i++) {
        arma_vec(0) = frac_1[i].x;
        arma_vec(1) = frac_1[i].y;
        arma_vec(2) = frac_1[i].z;
        //arma_vec = arma::dot(convmat_1, arma_vec);
        arma_vec = convmat_1 * arma_vec;
        atomsciflow::Atom atm;
        atm.name = frac_1[i].name;
        atm.x = arma_vec(0);
        atm.y = arma_vec(1);
        atm.z = arma_vec(2);
        cart_1.push_back(atm);
    }

    std::vector<atomsciflow::Atom> cart_2;
    for (i = 0; i < frac_2.size(); i++) {
        arma_vec(0) = frac_2[i].x;
        arma_vec(1) = frac_2[i].y;
        arma_vec(2) = frac_2[i].z;
        //arma_vec = arma::dot(convmat_2, arma_vec);
        arma_vec = convmat_2 * arma_vec;
        atomsciflow::Atom atm;
        atm.name = frac_2[i].name;
        atm.x = arma_vec(0);
        atm.y = arma_vec(1);
        atm.z = arma_vec(2);
        cart_2.push_back(atm);
    }

    // make distance gap between cart_1 and cart_2 is the value of distance
    std::cout << "make distance gap between cart_1 and cart_2 is the value of distance" << std::endl;
    std::vector<double> z_1;
    for (i = 0; i < cart_1.size(); i++) {
        z_1.push_back(cart_1[i].z);
    }

    std::vector<double> z_2;
    for (i = 0; i < cart_2.size(); i++) {
        z_2.push_back(cart_2[i].z);
    }

    double max_z_1 = *std::max_element(z_1.begin(), z_1.end());
    double min_z_2 = *std::min_element(z_2.begin(), z_2.end());

    std::cout << "max_z_1: " << max_z_1 << std::endl;
    std::cout << "min_z_1: " << min_z_2 << std::endl;

    for (i = 0; i < cart_2.size(); i++) {
        cart_2[i].z += distance - (min_z_2 - max_z_1);
    }

    //out.atoms = cart_1 + cart_2;
    std::cout << "merge atoms" << std::endl;
    out.atoms = cart_1;
    out.atoms.insert(out.atoms.end(), cart_2.begin(), cart_2.end());

    std::cout << "update cell information" << std::endl;

    double factor;

    if (use_cell == 1) {
        out.cell.clear();
        out.cell.push_back(arma::conv_to<std::vector<double>>::from(old_cell_1.row(0)));
        out.cell.push_back(arma::conv_to<std::vector<double>>::from(old_cell_1.row(1)));
        factor = (arma::norm(old_cell_1.row(2)) + arma::norm(old_cell_2.row(2))) / arma::norm(old_cell_1.row(2));
        out.cell.push_back(arma::conv_to<std::vector<double>>::from(old_cell_1.row(2) * factor));
    }  else if (use_cell == 2) {
        //
        out.cell.clear();
        out.cell.push_back(arma::conv_to<std::vector<double>>::from(old_cell_2.row(0)));
        out.cell.push_back(arma::conv_to<std::vector<double>>::from(old_cell_2.row(1)));
        factor = (arma::norm(old_cell_1.row(2)) + arma::norm(old_cell_2.row(2))) / arma::norm(old_cell_2.row(2));
        out.cell.push_back(arma::conv_to<std::vector<double>>::from(old_cell_2.row(2) * factor));
    } else {
        // average cell
        out.cell.clear();
        out.cell.push_back(arma::conv_to<std::vector<double>>::from((old_cell_1.row(0) + old_cell_2.row(0)) / 2));
        out.cell.push_back(arma::conv_to<std::vector<double>>::from((old_cell_1.row(1) + old_cell_2.row(1)) / 2));
        factor = (arma::norm(old_cell_1.row(2)) + arma::norm(old_cell_2.row(2))) / arma::norm((old_cell_1.row(2) + old_cell_2.row(2)) / 2);
        out.cell.push_back(arma::conv_to<std::vector<double>>::from((old_cell_1.row(2) + old_cell_2.row(2)) / 2 * factor));
    }

    std::cout << "add vacuum layer: " << thickness << " (Angstrom)" << std::endl;
    vacuum_layer(&out, 1, thickness);
    std::cout << "end add vacuum layer" << std::endl;

    return out;
}

atomsciflow::Crystal build_nanotube_ab(atomsciflow::Crystal* structure, std::string axis="b") {
    /*
    :param structure: an instance of crystal()
    :param axis: a or b
    :return an object of crystal()
    Note:
        build nanotube along an axis parallel to axis a or b
        only apply to film structure with ab as plane and a must be vertical to b.
        ab plane must be periodical plane
    */

    int i = 0; // for iteration
    int j = 0; // for iteration
    int k = 0; // for iteration

    set_frac_within_zero_and_one(structure);

    atomsciflow::Crystal out;
    out.atoms = structure->atoms;
    out.cell = structure->cell;

    double a = arma::norm(arma::conv_to<arma::vec>::from(out.cell[0]));
    double b = arma::norm(arma::conv_to<arma::vec>::from(out.cell[1]));
    double c = arma::norm(arma::conv_to<arma::vec>::from(out.cell[2]));

    std::vector<double> all_x;
    std::vector<double> all_y;
    std::vector<double> all_z;

    for (i = 0; i < out.atoms.size(); i++) {
        all_x.push_back(out.atoms[i].x);
        all_y.push_back(out.atoms[i].y);
        all_z.push_back(out.atoms[i].z);
    }

    double radius;

    if (axis == "b") {
        radius = a / (2 * arma::datum::pi);
    } else if (axis == "a") {
        radius = b / (2 * arma::datum::pi);
    }

    double middle_z = *std::min_element(all_z.begin(), all_z.end()) + (*std::max_element(all_z.begin(), all_z.end()) - *std::min_element(all_z.begin(), all_z.end())) / 2;
    double center_x;
    double center_y;
    double center_z;
    if (axis == "b") {
        center_x = *std::min_element(all_x.begin(), all_x.end()) = a / 2;
        center_z = middle_z + radius;
    } else if (axis == "a") {
        center_y = *std::min_element(all_y.begin(), all_y.end()) + b / 2;
        center_z = middle_z + radius;
    }

    double x, y, z;

    for (i = 0; i < out.atoms.size(); i++) {
        x = out.atoms[i].x;
        y = out.atoms[i].y;
        z = out.atoms[i].z;
        if (axis == "b") {
            //
            double new_x, new_z;
            if (x >= center_x) {
                double arc_len = x - center_x;
                // arc_len = radius * arc
                double arc = arc_len / radius;
                new_x = center_x + radius * std::sin(arc);
                new_z = center_z - radius * std::cos(arc);
                new_x += (middle_z - z) * std::cos(arc - arma::datum::pi / 2);
                new_z += (middle_z - z) * std::cos(arma::datum::pi - arc);
            }
            if (x < center_x) {
                double arc_len = center_x - x;
                // arc_len = radius * arc
                double arc = arc_len / radius;
                new_x = center_x - radius * std::sin(arc);
                new_z = center_z - radius * std::cos(arc);
                new_x += (z - middle_z) * std::cos(arc - arma::datum::pi / 2);
                new_z -= (z - middle_z) * std::cos(arma::datum::pi - arc);
            }
            out.atoms[i].x = new_x;
            out.atoms[i].z = new_z;
        } else if (axis == "a") {
            //
            double new_y, new_z;
            if (y >= center_y) {
                double arc_len = y - center_y;
                // arc_len = radius * arc
                double arc = arc_len / radius;
                new_y = center_y + radius * std::sin(arc);
                new_z = center_z - radius * std::cos(arc);
                new_y += (middle_z - z) * std::cos(arc - arma::datum::pi / 2);
                new_z += (middle_z - z) * std::cos(arma::datum::pi - arc);
            }
            if (y < center_y) {
                double arc_len = center_y - y;
                // arc_len = radius * arc
                double arc = arc_len / radius;
                new_y = center_y - radius * std::sin(arc);
                new_z = center_z - radius * std::cos(arc);
                new_y += (z - middle_z) * std::cos(arc - arma::datum::pi / 2);
                new_z -= (z - middle_z) * std::cos(arma::datum::pi - arc);
            }
            out.atoms[i].y = new_y;
            out.atoms[i].z = new_z;

        }
    }

    double c_factor = 2 * radius / arma::norm(arma::conv_to<arma::vec>::from(out.cell[2])) * 2.5;
    out.cell[2] = arma::conv_to<std::vector<double>>::from(arma::conv_to<arma::vec>::from(out.cell[2]) * c_factor);

    return out;
}

atomsciflow::Crystal build_nanotube_ac(atomsciflow::Crystal* structure, std::string axis="a") {
    /*
    :param structure: an instance of crystal()
    :param axis: a or c
    :return an object of crystal()
    Note:
        build nanotube along an axis parallel to axis a or c
        only apply to film structure with ac as plane and a must be vertical to c.
        ac plane must be periodical plane
    */

    int i = 0; // for iteration
    int j = 0; // for iteration

    atomsciflow::Crystal out;
    out.atoms = structure->atoms;
    out.cell = structure->cell;

    double a = arma::norm(arma::conv_to<arma::vec>::from(out.cell[0]));
    double b = arma::norm(arma::conv_to<arma::vec>::from(out.cell[1]));
    double c = arma::norm(arma::conv_to<arma::vec>::from(out.cell[2]));

    std::vector<double> all_x;
    std::vector<double> all_y;
    std::vector<double> all_z;

    for (i = 0; i < out.atoms.size(); i++) {
        all_x.push_back(out.atoms[i].x);
        all_y.push_back(out.atoms[i].y);
        all_z.push_back(out.atoms[i].z);
    }

    double radius;

    if (axis == "c") {
        radius = a / (2 * arma::datum::pi);
    } else if (axis == "a") {
        radius = c / (2 * arma::datum::pi);
    }

    double middle_y = *std::min_element(all_y.begin(), all_y.end()) + (*std::max_element(all_y.begin(), all_y.end()) - *std::min_element(all_y.begin(), all_y.end())) / 2;

    double center_x, center_y, center_z;

    if (axis == "c") {
        center_x = *std::min_element(all_x.begin(), all_x.end()) + a / 2;
        center_y = middle_y + radius;
    } else if (axis == "a") {
        center_z = *std::min_element(all_z.begin(), all_z.end()) + c / 2;
        center_y = middle_y + radius;
    }

    double x, y, z;

    for (i = 0; i < out.atoms.size(); i++) {
        x = out.atoms[i].x;
        y = out.atoms[i].y;
        z = out.atoms[i].z;
        if (axis == "c") {
            double new_x, new_y;
            if (x >= center_x) {
                double arc_len = x - center_x;
                // arc_len = radius * arc
                double arc = arc_len / radius;
                new_x = center_x + radius * std::sin(arc);
                new_y = center_y - radius * std::cos(arc);
                new_x += (middle_y - y) * std::cos(arc - arma::datum::pi / 2);
                new_y += (middle_y - y) * std::cos(arma::datum::pi - arc);
            }
            if (x < center_x) {
                double arc_len = center_x - x;
                // arc_len = radius * arc
                double arc = arc_len / radius;
                new_x = center_x - radius * std::sin(arc);
                new_y = center_y - radius * std::cos(arc);
                new_x += (y - middle_y) * std::cos(arc - arma::datum::pi / 2);
                new_y -= (y - middle_y) * std::cos(arma::datum::pi - arc);
            }
            out.atoms[i].x = new_x;
            out.atoms[i].y = new_y;
        } else if (axis == "a") {
            double new_y, new_z;
            if (z >= center_z) {
                double arc_len = z - center_z;
                // arc_len = radius * arc
                double arc = arc_len / radius;
                new_z = center_z + radius * std::sin(arc);
                new_y = center_y - radius * std::cos(arc);
                new_z += (middle_y - y) * std::cos(arc - arma::datum::pi / 2);
                new_y += (middle_y - y) * std::cos(arma::datum::pi - arc);
            }
            if (z < center_z) {
                double arc_len = center_z - z;
                // arc_len = radius * arc
                double arc = arc_len / radius;
                new_z = center_z - radius * std::sin(arc);
                new_y = center_y - radius * std::cos(arc);
                new_z += (y - middle_y) * std::cos(arc - arma::datum::pi / 2);
                new_y -= (y - middle_y) * std::cos(arma::datum::pi - arc);
            }
            out.atoms[i].y = new_y;
            out.atoms[i].z = new_z;
        }
    }

    double b_factor = 2 * radius / arma::norm(arma::conv_to<arma::vec>::from(out.cell[1])) * 2.5;
    out.cell[1] = arma::conv_to<std::vector<double>>::from(arma::conv_to<arma::vec>::from(out.cell[1]) * b_factor);

    return out;
}

atomsciflow::Crystal build_nanotube_bc(atomsciflow::Crystal* structure, std::string axis="c") {
    /*
    :param structure: an instance of crystal()
    :param axis: b or c
    :return an object of crystal()
    Note:
        build nanotube along an axis parallel to axis b or c
        only apply to film structure with bc as plane and b must be vertical to c.
        bc plane must be periodical plane
    */

    int i = 0; // for iteration
    int j = 0; // for iteartion

    atomsciflow::Crystal out;

    out.atoms = structure->atoms;
    out.cell = structure->cell;

    double a = arma::norm(arma::conv_to<arma::vec>::from(out.cell[0]));
    double b = arma::norm(arma::conv_to<arma::vec>::from(out.cell[1]));
    double c = arma::norm(arma::conv_to<arma::vec>::from(out.cell[2]));

    std::vector<double> all_x;
    std::vector<double> all_y;
    std::vector<double> all_z;

    for (i = 0; i < out.atoms.size(); i++) {
        all_x.push_back(out.atoms[i].x);
        all_y.push_back(out.atoms[i].y);
        all_z.push_back(out.atoms[i].z);
    }

    double radius;
    if (axis == "b") {
        radius = c / (2 * arma::datum::pi);
    } else if (axis == "c") {
        radius = b / (2 * arma::datum::pi);
    }

    double middle_x = *std::min_element(all_x.begin(), all_x.end()) + (*std::max_element(all_x.begin(), all_x.end()) - *std::min_element(all_x.begin(), all_x.end())) / 2;

    double center_x, center_y, center_z;
    if (axis == "b") {
        center_z = *std::min_element(all_z.begin(), all_z.end()) + c / 2;
        center_x = middle_x + radius;
    } else if (axis == "c") {
        center_y = *std::min_element(all_y.begin(), all_y.end()) + b / 2;
        center_x = middle_x  + radius;
    }

    double x, y, z;

    for (i = 0; i < out.atoms.size(); i++) {
        x = out.atoms[i].x;
        y = out.atoms[i].y;
        z = out.atoms[i].z;
        if (axis == "b") {
            double new_x, new_z;
            if (z >= center_z) {
                double arc_len = z - center_z;
                // arc_len = radius * arc
                double arc = arc_len / radius;
                new_z = center_z + radius * std::sin(arc);
                new_x = center_x - radius * std::cos(arc);
                new_z += (middle_x - x) * std::cos(arc - arma::datum::pi / 2);
                new_x += (middle_x - x) * std::cos(arma::datum::pi - arc);
            }
            if (z < center_z) {
                double arc_len = center_z - z;
                // arc_len = radius * arc
                double arc = arc_len / radius;
                new_z = center_z - radius * std::sin(arc);
                new_x = center_x - radius * std::cos(arc);
                new_z += (x - middle_x) * std::cos(arc - arma::datum::pi / 2);
                new_x -= (x - middle_x) * std::cos(arma::datum::pi - arc);
            }
            out.atoms[i].x = new_x;
            out.atoms[i].z = new_z;
        } else if (axis == "c") {
            double new_x, new_y;
            if (y >= center_y) {
                double arc_len  = y - center_y;
                // arc_len = radius * arc
                double arc = arc_len / radius;
                new_y = center_y + radius * std::sin(arc);
                new_x = center_x = radius * std::cos(arc);
                new_y += (middle_x - x) * std::cos(arc - arma::datum::pi / 2);
                new_x += (middle_x - x) * std::cos(arma::datum::pi - arc);
            }
            if (y < center_y) {
                double arc_len = center_y - y;
                // arc_len = radius * arc
                double arc = arc_len / radius;
                new_y = center_y - radius * std::sin(arc);
                new_x = center_x - radius * std::cos(arc);
                new_y += (x - middle_x) * std::cos(arc - arma::datum::pi / 2);
                new_x -= (x - middle_x) * std::cos(arma::datum::pi - arc);
            }
            out.atoms[i].y = new_y;
            out.atoms[i].x = new_x;
        }
    }

    double a_factor = 2 * radius / arma::norm(arma::conv_to<arma::vec>::from(out.cell[0])) * 2.5;
    out.cell[0] = arma::conv_to<std::vector<double>>::from(arma::conv_to<arma::vec>::from(out.cell[0]) * a_factor);

    return out;
}

} //namespace atomsciflow


#include <boost/geometry/formulas/vincenty_direct.hpp>


#pragma once

#include <stdio.h>
#include <Eigen/Dense>
#include <iostream>

#include <PnC/RobotSystem/RobotSystem.hpp>
#include <Utils/IO/IOUtilities.hpp>
#include <Utils/Math/MathUtilities.hpp>

//! Contact Frame's z-axis should correspond normal vector to the ground

class ContactSpec {
   public:
    ContactSpec(RobotSystem* _robot, const int& _dim) {
        robot_ = _robot;
        dim_contact_ = _dim;
        b_set_contact_ = false;
        idx_Fz_ = dim_contact_ - 1;
        Jc_ = Eigen::MatrixXd::Zero(dim_contact_, robot_->getNumDofs());
        JcDotQdot_ = Eigen::VectorXd::Zero(dim_contact_);
    }

    virtual ~ContactSpec() {}

    void getContactJacobian(Eigen::MatrixXd& Jc) { Jc = Jc_; }
    void getJcDotQdot(Eigen::VectorXd& JcDotQdot) { JcDotQdot = JcDotQdot_; }
    int getDim() { return dim_contact_; }
    void unsetContact() { b_set_contact_ = false; }

    bool updateContactSpec() {
        _UpdateJc();
        _UpdateJcDotQdot();
        _UpdateUf();
        _UpdateInequalityVector();
        b_set_contact_ = true;
        return true;
    }

    virtual int getDimRFConstratint() { return Uf_.rows(); }
    void getRFConstraintMtx(Eigen::MatrixXd& Uf) { Uf = Uf_; }
    void getRFConstraintVec(Eigen::VectorXd& ieq_vec) { ieq_vec = ieq_vec_; }

    int getFzIndex() { return idx_Fz_; }

   protected:
    virtual bool _UpdateJc() = 0;
    virtual bool _UpdateJcDotQdot() = 0;
    virtual bool _UpdateUf() = 0;
    virtual bool _UpdateInequalityVector() = 0;

    RobotSystem* robot_;
    Eigen::MatrixXd Jc_;
    Eigen::VectorXd JcDotQdot_;
    int dim_contact_;
    int idx_Fz_;
    bool b_set_contact_;

    Eigen::MatrixXd Uf_;
    Eigen::VectorXd ieq_vec_;
};


/* ----------------------------------------------------------------------------

 * GTSAM Copyright 2010, Georgia Tech Research Corporation,
 * Atlanta, Georgia 30332-0415
 * All Rights Reserved
 * Authors: Frank Dellaert, et al. (see THANKS for the full author list)

 * See LICENSE for the license information

 * -------------------------------------------------------------------------- */

/**
 * @file testOrdering
 * @author Alex Cunningham
 * @author Andrew Melim
 */

#include <gtsam/inference/Symbol.h>
#include <gtsam/symbolic/SymbolicFactorGraph.h>
#include <gtsam/inference/Ordering.h>
#include <gtsam/inference/MetisIndex.h>
#include <gtsam/base/TestableAssertions.h>
#include <CppUnitLite/TestHarness.h>

#include <boost/assign/std.hpp>

using namespace std;
using namespace gtsam;
using namespace boost::assign;

namespace example {
SymbolicFactorGraph symbolicChain() {
  SymbolicFactorGraph symbolicGraph;
  symbolicGraph.push_factor(0, 1);
  symbolicGraph.push_factor(1, 2);
  symbolicGraph.push_factor(2, 3);
  symbolicGraph.push_factor(3, 4);
  symbolicGraph.push_factor(4, 5);
  return symbolicGraph;
}
}
/* ************************************************************************* */
TEST(Ordering, constrained_ordering) {
  // create graph with wanted variable set = 2, 4
  SymbolicFactorGraph symbolicGraph = example::symbolicChain();

  // unconstrained version
  {
  Ordering actual = Ordering::Colamd(symbolicGraph);
  Ordering expected = Ordering(list_of(0)(1)(2)(3)(4)(5));
  EXPECT(assert_equal(expected, actual));
  }

  // constrained version - push one set to the end
  {
    Ordering actual = Ordering::ColamdConstrainedLast(symbolicGraph, list_of(2)(4));
    Ordering expected = Ordering(list_of(0)(1)(5)(3)(4)(2));
    EXPECT(assert_equal(expected, actual));
  }

  // constrained version - push one set to the start
  {
    Ordering actual = Ordering::ColamdConstrainedFirst(symbolicGraph, list_of(2)(4));
    Ordering expected = Ordering(list_of(2)(4)(0)(1)(3)(5));
    EXPECT(assert_equal(expected, actual));
  }

  // Make sure giving empty constraints does not break the code
  {
    Ordering actual = Ordering::ColamdConstrainedLast(symbolicGraph, {});
    Ordering expected = Ordering(list_of(0)(1)(2)(3)(4)(5));
    EXPECT(assert_equal(expected, actual));
  }
  {
    Ordering actual = Ordering::ColamdConstrainedFirst(symbolicGraph, {});
    Ordering expected = Ordering(list_of(0)(1)(2)(3)(4)(5));
    EXPECT(assert_equal(expected, actual));
  }

  // Make sure giving empty graph does not break the code
  SymbolicFactorGraph emptyGraph;
  Ordering empty;
  {
    Ordering actual = Ordering::ColamdConstrainedLast(emptyGraph, list_of(2)(4));
    EXPECT(assert_equal(empty, actual));
  }
  {
    Ordering actual = Ordering::ColamdConstrainedFirst(emptyGraph, list_of(2)(4));
    EXPECT(assert_equal(empty, actual));
  }
}

/* ************************************************************************* */
TEST(Ordering, grouped_constrained_ordering) {

  // create graph with constrained groups:
  // 1: 2, 4
  // 2: 5
  SymbolicFactorGraph symbolicGraph = example::symbolicChain();

  // constrained version - push one set to the end
  FastMap<Key, int> constraints;
  constraints[2] = 1;
  constraints[4] = 1;
  constraints[5] = 2;

  Ordering actual = Ordering::ColamdConstrained(symbolicGraph, constraints);
  Ordering expected = list_of(0)(1)(3)(2)(4)(5);
  EXPECT(assert_equal(expected, actual));
}

/* ************************************************************************* */
TEST(Ordering, csr_format) {
  // Example in METIS manual
  SymbolicFactorGraph symbolicGraph;
  symbolicGraph.push_factor(0, 1);
  symbolicGraph.push_factor(1, 2);
  symbolicGraph.push_factor(2, 3);
  symbolicGraph.push_factor(3, 4);
  symbolicGraph.push_factor(5, 6);
  symbolicGraph.push_factor(6, 7);
  symbolicGraph.push_factor(7, 8);
  symbolicGraph.push_factor(8, 9);
  symbolicGraph.push_factor(10, 11);
  symbolicGraph.push_factor(11, 12);
  symbolicGraph.push_factor(12, 13);
  symbolicGraph.push_factor(13, 14);

  symbolicGraph.push_factor(0, 5);
  symbolicGraph.push_factor(5, 10);
  symbolicGraph.push_factor(1, 6);
  symbolicGraph.push_factor(6, 11);
  symbolicGraph.push_factor(2, 7);
  symbolicGraph.push_factor(7, 12);
  symbolicGraph.push_factor(3, 8);
  symbolicGraph.push_factor(8, 13);
  symbolicGraph.push_factor(4, 9);
  symbolicGraph.push_factor(9, 14);

  MetisIndex mi(symbolicGraph);

  vector<int> xadjExpected, adjExpected;
  xadjExpected += 0, 2, 5, 8, 11, 13, 16, 20, 24, 28, 31, 33, 36, 39, 42, 44;
  adjExpected += 1, 5, 0, 2, 6, 1, 3, 7, 2, 4, 8, 3, 9, 0, 6, 10, 1, 5, 7, 11, 2, 6, 8, 12, 3, 7, 9, 13, 4, 8, 14, 5, 11, 6, 10, 12, 7, 11, 13, 8, 12, 14, 9, 13;

  EXPECT(xadjExpected == mi.xadj());
  EXPECT(adjExpected.size() == mi.adj().size());
  EXPECT(adjExpected == mi.adj());
}

/* ************************************************************************* */
TEST(Ordering, csr_format_2) {
  SymbolicFactorGraph symbolicGraph;

  symbolicGraph.push_factor(0);
  symbolicGraph.push_factor(0, 1);
  symbolicGraph.push_factor(1, 2);
  symbolicGraph.push_factor(2, 3);
  symbolicGraph.push_factor(3, 4);
  symbolicGraph.push_factor(4, 1);

  MetisIndex mi(symbolicGraph);

  vector<int> xadjExpected, adjExpected;
  xadjExpected += 0, 1, 4, 6, 8, 10;
  adjExpected += 1, 0, 2, 4, 1, 3, 2, 4, 1, 3;

  EXPECT(xadjExpected == mi.xadj());
  EXPECT(adjExpected.size() == mi.adj().size());
  EXPECT(adjExpected == mi.adj());
}

/* ************************************************************************* */
TEST(Ordering, csr_format_3) {
  SymbolicFactorGraph symbolicGraph;

  symbolicGraph.push_factor(100);
  symbolicGraph.push_factor(100, 101);
  symbolicGraph.push_factor(101, 102);
  symbolicGraph.push_factor(102, 103);
  symbolicGraph.push_factor(103, 104);
  symbolicGraph.push_factor(104, 101);

  MetisIndex mi(symbolicGraph);

  vector<int> xadjExpected, adjExpected;
  xadjExpected += 0, 1, 4, 6, 8, 10;
  adjExpected += 1, 0, 2, 4, 1, 3, 2, 4, 1, 3;
  //size_t minKey = mi.minKey();

  vector<int> adjAcutal = mi.adj();

  // Normalize, subtract the smallest key
  //std::transform(adjAcutal.begin(), adjAcutal.end(), adjAcutal.begin(),
  //    std::bind2nd(std::minus<size_t>(), minKey));

  EXPECT(xadjExpected == mi.xadj());
  EXPECT(adjExpected.size() == mi.adj().size());
  EXPECT(adjExpected == adjAcutal);
}

/* ************************************************************************* */
#ifdef GTSAM_SUPPORT_NESTED_DISSECTION
TEST(Ordering, csr_format_4) {
  SymbolicFactorGraph symbolicGraph;

  symbolicGraph.push_factor(Symbol('x', 1));
  symbolicGraph.push_factor(Symbol('x', 1), Symbol('x', 2));
  symbolicGraph.push_factor(Symbol('x', 2), Symbol('x', 3));
  symbolicGraph.push_factor(Symbol('x', 3), Symbol('x', 4));
  symbolicGraph.push_factor(Symbol('x', 4), Symbol('x', 5));
  symbolicGraph.push_factor(Symbol('x', 5), Symbol('x', 6));

  MetisIndex mi(symbolicGraph);

  vector<int> xadjExpected, adjExpected;
  xadjExpected += 0, 1, 3, 5, 7, 9, 10;
  adjExpected += 1, 0, 2, 1, 3, 2, 4, 3, 5, 4;

  vector<int> adjAcutal = mi.adj();
  vector<int> xadjActual = mi.xadj();

  EXPECT(xadjExpected == mi.xadj());
  EXPECT(adjExpected.size() == mi.adj().size());
  EXPECT(adjExpected == adjAcutal);

  Ordering metOrder = Ordering::Metis(symbolicGraph);

  // Test different symbol types
  symbolicGraph.push_factor(Symbol('l', 1));
  symbolicGraph.push_factor(Symbol('x', 1), Symbol('l', 1));
  symbolicGraph.push_factor(Symbol('x', 2), Symbol('l', 1));
  symbolicGraph.push_factor(Symbol('x', 3), Symbol('l', 1));
  symbolicGraph.push_factor(Symbol('x', 4), Symbol('l', 1));

  Ordering metOrder2 = Ordering::Metis(symbolicGraph);
}
#endif
/* ************************************************************************* */
#ifdef GTSAM_SUPPORT_NESTED_DISSECTION
TEST(Ordering, metis) {

  SymbolicFactorGraph symbolicGraph;

  symbolicGraph.push_factor(0);
  symbolicGraph.push_factor(0, 1);
  symbolicGraph.push_factor(1, 2);

  MetisIndex mi(symbolicGraph);

  vector<int> xadjExpected, adjExpected;
  xadjExpected += 0, 1, 3, 4;
  adjExpected += 1, 0, 2, 1;

  EXPECT(xadjExpected == mi.xadj());
  EXPECT(adjExpected.size() == mi.adj().size());
  EXPECT(adjExpected == mi.adj());

  Ordering metis = Ordering::Metis(symbolicGraph);
}
#endif
/* ************************************************************************* */
#ifdef GTSAM_SUPPORT_NESTED_DISSECTION
TEST(Ordering, MetisLoop) {

  // create linear graph
  SymbolicFactorGraph symbolicGraph = example::symbolicChain();

  // add loop closure
  symbolicGraph.push_factor(0, 5);

  // METIS
#if !defined(__APPLE__)
  {
    Ordering actual = Ordering::Create(Ordering::METIS, symbolicGraph);
    //  - P( 0 4 1)
    //  | - P( 2 | 4 1)
    //  | | - P( 3 | 4 2)
    //  | - P( 5 | 0 1)
    Ordering expected = Ordering(list_of(3)(2)(5)(0)(4)(1));
    EXPECT(assert_equal(expected, actual));
  }
#else
  {
    Ordering actual = Ordering::Create(Ordering::METIS, symbolicGraph);
    //  - P( 1 0 3)
    //  | - P( 4 | 0 3)
    //  | | - P( 5 | 0 4)
    //  | - P( 2 | 1 3)
    Ordering expected = Ordering(list_of(5)(4)(2)(1)(0)(3));
    EXPECT(assert_equal(expected, actual));
  }
#endif
}
#endif
/* ************************************************************************* */
#ifdef GTSAM_SUPPORT_NESTED_DISSECTION
TEST(Ordering, MetisEmptyGraph) {
  SymbolicFactorGraph symbolicGraph;

  Ordering actual = Ordering::Create(Ordering::METIS, symbolicGraph);
  Ordering expected;
  EXPECT(assert_equal(expected, actual));
}
#endif
/* ************************************************************************* */
#ifdef GTSAM_SUPPORT_NESTED_DISSECTION
TEST(Ordering, MetisSingleNode) {
  // create graph with a single node
  SymbolicFactorGraph symbolicGraph;
  symbolicGraph.push_factor(7);

  Ordering actual = Ordering::Create(Ordering::METIS, symbolicGraph);
  Ordering expected = Ordering(list_of(7));
  EXPECT(assert_equal(expected, actual));
}
#endif
/* ************************************************************************* */
TEST(Ordering, Create) {

  // create chain graph
  SymbolicFactorGraph symbolicGraph = example::symbolicChain();

  // COLAMD
  {
    //- P( 4 5)
    //| - P( 3 | 4)
    //| | - P( 2 | 3)
    //| | | - P( 1 | 2)
    //| | | | - P( 0 | 1)
    Ordering actual = Ordering::Create(Ordering::COLAMD, symbolicGraph);
    Ordering expected = Ordering(list_of(0)(1)(2)(3)(4)(5));
    EXPECT(assert_equal(expected, actual));
  }

#ifdef GTSAM_SUPPORT_NESTED_DISSECTION
  // METIS
  {
    Ordering actual = Ordering::Create(Ordering::METIS, symbolicGraph);
    //- P( 1 0 2)
    //| - P( 3 4 | 2)
    //| | - P( 5 | 4)
    Ordering expected = Ordering(list_of(5)(3)(4)(1)(0)(2));
    EXPECT(assert_equal(expected, actual));
  }
#endif

  // CUSTOM
  CHECK_EXCEPTION(Ordering::Create(Ordering::CUSTOM, symbolicGraph), runtime_error);
}

/* ************************************************************************* */
int main() {
  TestResult tr;
  return TestRegistry::runAllTests(tr);
}
/* ************************************************************************* */


#ifndef BNI_SAMPLER_HPP
#define BNI_SAMPLER_HPP

#include <fstream>
#include <sstream>
#include <functional>
#include <vector>
#include <numeric>
#include <unordered_map>
#include <boost/optional.hpp>
#include <boost/algorithm/string.hpp>
#include <bayesian/graph.hpp>
#include <bayesian/hash.hpp>

namespace bn {

class sampler {
public:
    sampler()
        : filename_(), table_(), sampling_size_(0)
    {
    }

    sampler(std::string const& filename)
        : filename_(filename), table_(), sampling_size_(0)
    {
    }

    bool load_sample(std::unordered_map<condition_t, std::size_t> const& table)
    {
        table_ = table;

        sampling_size_ = 0;
        for(auto const& p : table) sampling_size_ += p.second;

        return true;
    }

    // 事前に設定したファイルからCSVデータを読み込み，一般化したテーブルを作る
    // 引数: node_list: ファイルに記述されている順番を指定するvector
    // 戻り値: 成功した場合のみ true
    bool load_sample(std::vector<vertex_type> const& node_list)
    {
        // ファイルオープン
        std::ifstream ifs(filename_);
        if(!ifs.is_open()) return false;

        // 読み込み先
        std::size_t sampling_size = 0;
        std::unordered_map<condition_t, std::size_t> table;

        // 1行ずつサンプリング
        std::string line_str;
        condition_t sample;
        while(std::getline(ifs, line_str))
        {
            std::vector<std::string> line;
            boost::algorithm::split(line, line_str, boost::is_space(), boost::algorithm::token_compress_on);

            auto select_it = line.cbegin() + 1;
            for(auto it = node_list.cbegin(); it != node_list.cend();)
                sample[*it++] = std::stoi(*select_it++);

            auto const sample_num = std::stoi(line[0]);

            auto const it = table.find(sample);
            if(it == table.cend()) table[sample] = sample_num;
            else (it->second) += sample_num;

            sampling_size += sample_num;
        }

        sampling_size_ = sampling_size;
        table_ = std::move(table);
        return true;
    }

	// load_sampleによって生成したテーブルから，CPTを計算し，graphに格納する
	// 引数: graph: 読み込みしたCPTを保存するグラフ構造
	// 戻り値: 成功した場合のみ true
    bool make_cpt(graph_t const& graph) const
    {
        // sampleが読み込まれているかどうか
        if(sampling_size() == 0) return false;

        // CPTの初期化
        for(auto const& node : graph.vertex_list())
        {
            auto const in_vertex = graph.in_vertexes(node);
            node->cpt.assign(in_vertex, node);
        }

        // カウンタ
        std::unordered_map<
            std::pair<vertex_type, condition_t>,
            std::vector<std::size_t>,
            bn::hash<std::pair<vertex_type, condition_t>>
        > counter;

        // カウント
        for(auto const& sample : table())
        {
            for(auto const& node : graph.vertex_list())
            {
                auto const parents = graph.in_vertexes(node);

                // 条件を畳み込む
                condition_t conditional;
                for(auto const& parent : parents) conditional[parent] = sample.first.at(parent);

                // counterのカウントアップ
                auto container_first = std::make_pair(node, conditional);
                auto const it = counter.find(container_first);
                if(it == counter.end())
                {
                    // 作ってcounterに要素を追加
                    std::vector<std::size_t> local_counter(node->selectable_num, 0);
                    local_counter[sample.first.at(node)] = sample.second;
                    counter.insert(std::make_pair(std::move(container_first), std::move(local_counter)));
                }
                else
                {
                    // 足しあわせ
                    it->second[sample.first.at(node)] += sample.second;
                }
            }
        }

        // カウント値を用い，CPTを算出
        for(auto const& node : graph.vertex_list())
        {
            for(auto const& conditional : node->cpt.pattern())
            {
                // カウンタからデータを引きずりだし，確率1にすべき母数を算出
                auto const key = std::make_pair(node, conditional);
                auto const it = counter.find(key);
                auto const select_counter =
                    (it != counter.end()) ? counter.at(std::make_pair(node, conditional))
                                          : std::vector<std::size_t>(node->selectable_num, 0);

                double const parameter = static_cast<double>(std::accumulate(select_counter.begin(), select_counter.end(), static_cast<std::size_t>(0)));

                // 本題のCPT作成
                std::vector<double> cpt(node->selectable_num);
                if(parameter == 0)
                {
                    // 同じ値を与える
                    for(std::size_t i = 0; i < node->selectable_num; ++i)
                        cpt[i] = /*0.0;*/static_cast<double>(1.0) / node->selectable_num;
                }
                else
                {
                    // CPTを正規化する
                    for(std::size_t i = 0; i < node->selectable_num; ++i)
                        cpt[i] = static_cast<double>(select_counter.at(i)) / parameter;
                }

                node->cpt[conditional].second = std::move(cpt);
            }
        }

        return true;
    }

    // 読み込み対象のファイルを指定するgetter
    std::string filename() const
    {
        return filename_;
    }

    // 読み込み対象のファイルを指定するsetter
    void set_filename(std::string const& filename)
    {
        filename_ = filename;
        sampling_size_ = 0;
        table_.clear();
    }

	// テーブルに関するgetter
	std::unordered_map<condition_t, std::size_t> table() const
    {
        return table_;
    }

	// サンプル数に関するgetter
    std::size_t sampling_size() const
    {
        return sampling_size_;
    }

private:
    std::string filename_;
    std::unordered_map<condition_t, std::size_t> table_;
    std::size_t sampling_size_;

    // 与えられた確率変数全ての組み合わせに対し，functionを実行するというインターフェースを提供する
    // TODO: class bpの中と重複しているので，後で整理する
    void all_combination_pattern(
        std::vector<vertex_type> const& combination,
        std::function<void(condition_t const&)> const& function
        )
    {
        typedef std::vector<vertex_type>::const_iterator iterator_type;
        std::function<void(iterator_type const, iterator_type const&)> recursive;

        condition_t condition;
        recursive = [&](iterator_type const it, iterator_type const& end)
        {
            if(it == end)
            {
                function(condition);
            }
            else
            {
                for(int i = 0; i < (*it)->selectable_num; ++i)
                {
                    condition[*it] = i;
                    recursive(it + 1, end);
                }
            }
        };

        recursive(combination.cbegin(), combination.cend());
    }
};

} // namespace bn

#endif // #ifndef BNI_SAMPLER_HPP


/*
Lightweight Automated Planning Toolkit
Copyright (C) 2012
Miquel Ramirez <miquel.ramirez@rmit.edu.au>
Nir Lipovetzky <nirlipo@gmail.com>

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program.  If not, see <http://www.gnu.org/licenses/>.
*/

#ifndef __H_2__
#define __H_2__

#include <aptk/search_prob.hxx>
#include <aptk/heuristic.hxx>
#include <aptk/bit_set.hxx>
#include <strips_state.hxx>
#include <strips_prob.hxx>
#include <boost/circular_buffer.hpp>
#include <vector>
#include <list>
#include <iosfwd>

namespace aptk {

namespace agnostic {

namespace H2_Helper {

	inline int	pair_index( unsigned p, unsigned q ) {
		return ( p >= q ? p*(p+1)/2 + q : q*(q+1)/2 + p);
	}

}

enum class H2_Cost_Function { Zero_Costs, Unit_Costs, Use_Costs };

template <typename Search_Model, H2_Cost_Function cost_opt = H2_Cost_Function::Use_Costs >
class H2_Heuristic : public Heuristic<State> {

public:
	H2_Heuristic( const Search_Model& prob )
	: Heuristic<State>( prob ), m_strips_model( prob.task() ) {
		unsigned F = m_strips_model.num_fluents();
		m_values.resize( (F*F + F)/2 );
		m_op_values.resize( m_strips_model.num_actions() );
		m_interfering_ops.resize( F );
		for ( unsigned p = 0; p < m_interfering_ops.size(); p++ ) {
			m_interfering_ops[p] = new Bit_Set( m_strips_model.num_actions() );
			for ( unsigned op = 0; op < m_strips_model.num_actions(); op++ ) 	{
				const Action* op_ptr = m_strips_model.actions()[op];
				if ( op_ptr->add_set().isset( p ) || op_ptr->del_set().isset( p ) )
					m_interfering_ops[p]->set(op);
			}
		}
		// NIR: Set up the relevant actions once here so we don't need
		//      to iterate through all of them when evaluating, similar to h1
		m_already_updated.resize( (F*F + F)/2 );
		m_updated.resize( (F*F + F)/2 );
		m_relevant_actions.resize( (F*F + F)/2 );

		for ( unsigned i = 0; i < m_strips_model.num_actions(); i++ ) {

			const Action& a = *(m_strips_model.actions()[i]);

			// Relevant if the fluent is in the precondition
			for ( unsigned p = 0; p < a.prec_vec().size(); ++p ) {
				for ( unsigned q = p; q < a.prec_vec().size(); ++q ) {

					m_relevant_actions[ H2_Helper::pair_index( a.prec_vec()[p], a.prec_vec()[q] ) ].insert(i);
				}
			}

			// Relevant if the fluent is in the head of a conditional effect
			for ( unsigned j = 0; j < a.ceff_vec().size(); ++j ) {

				const Conditional_Effect& ceff = *(a.ceff_vec()[j]);
				for ( unsigned p = 0; p < ceff.prec_vec().size(); ++p) {
					for ( unsigned q = p; q < ceff.prec_vec().size(); ++q) {
						m_relevant_actions[ H2_Helper::pair_index( ceff.prec_vec()[p], ceff.prec_vec()[q] ) ].insert(i);
					}
				}
			}
		}

	}

	virtual ~H2_Heuristic() {
	}

	template <typename Search_Node>
        void eval( const Search_Node* n, float& h_val, std::vector<Action_Idx>& pref_ops) {
		eval(n->state(), h_val, pref_ops);
	}

        template <typename Search_Node>
        void eval( const Search_Node* n, float& h_val ) {

		eval(n->state(),h_val);
	}

	virtual	void	eval( const State& s, float& h_val ) {
		m_already_updated.reset();
		m_updated.clear();
		initialize( s );
		compute();
		h_val = eval( m_strips_model.goal() );
	}

	virtual void eval( const State& s, float& h_val,  std::vector<Action_Idx>& pref_ops ) {
		eval( s, h_val );
	}

	float&	op_value( unsigned a ) 		{ return m_op_values.at(a); }
	float   op_value( unsigned a ) const 	{ return m_op_values.at(a); }

	float& value( unsigned p, unsigned q ) 	{
		assert( H2_Helper::pair_index(p,q) < (int)m_values.size() );
		return m_values[H2_Helper::pair_index(p,q)];
	}

	float value( unsigned p, unsigned q ) const 	{
		assert( H2_Helper::pair_index(p,q) < (int)m_values.size() );
		return m_values[H2_Helper::pair_index(p,q)];
	}

	float& value( unsigned p ) 	{
		assert( H2_Helper::pair_index(p,p) < (int)m_values.size() );
		return m_values[H2_Helper::pair_index(p,p)];
	}

	float value( unsigned p ) const 	{
		assert( H2_Helper::pair_index(p,p) < (int)m_values.size() );
		return m_values[H2_Helper::pair_index(p,p)];
	}

	float eval( const Fluent_Vec& s ) const {
		float v = 0;
		for ( unsigned i = 0; i < s.size(); i++ )
			for ( unsigned j = i; j < s.size(); j++ ){
				v = std::max( v, value( s[i], s[j] ) );
				if(v == infty ) return infty;
			}

		return v;
	}

	bool  is_mutex( const Fluent_Vec& s ) const {
		return eval(s) == infty;
	}

	bool	is_mutex( unsigned p, unsigned q ) const {
		return value( p, q ) == infty;
	}

	float eval( const Fluent_Vec& s, unsigned p ) const {
		float v = 0;
		for ( unsigned k = 0; k < s.size(); k++ )
			v = std::max( v, value( s[k], p ) );
		float v2 = 0;
		for ( unsigned i = 0; i < s.size(); i++ )
			for ( unsigned j = i; j < s.size(); j++ )
				v2 = std::max( v2, value( s[i], s[j] ) );
		return std::max( v, v2 );
	}

	bool interferes( unsigned a, unsigned p ) const {
		return m_interfering_ops[p]->isset(a);
	}

	void print_values( std::ostream& os ) const {
		for ( unsigned p = 0; p < m_strips_model.fluents().size(); p++ )
			for ( unsigned q = p; q < m_strips_model.fluents().size(); q++ ) {
				os << "h²({ ";
				os << m_strips_model.fluents()[p]->signature();
				os << ", ";
				os << m_strips_model.fluents()[q]->signature();
				os << "}) = " << value(p,q) << std::endl;
			}
	}

	void compute_edeletes( STRIPS_Problem& prob ){
		m_already_updated.reset();
		m_updated.clear();
		initialize( prob.init() );
		compute_mutexes_only();
		extract_edeletes( prob );

#ifdef DEBUG
		print_values(std::cout);
		prob.print_actions(std::cout);
#endif
	}

	void compute_edeletes_aij( STRIPS_Problem& prob ){
		compute_mutexes_only_aij();
		extract_edeletes( prob );

#ifdef DEBUG
		print_values(std::cout);
		prob.print_actions(std::cout);
#endif
	}

	void 	goal_mutex_closure( STRIPS_Problem& prob ){
		unsigned fsize = prob.num_fluents();
		Fluent_Vec new_goal;
		for( auto g : prob.goal()  )
			new_goal.push_back(g);

		for( unsigned f_idx = 0; f_idx != fsize; f_idx++ ){
			Fluent* p =  prob.fluents()[f_idx];

			if( ! prob.is_in_init( p->index() ) ) continue;

			for( auto g : prob.goal()  ){
				if( is_mutex(p->index(),g) ){
					unsigned notp_idx = STRIPS_Problem::add_fluent( prob,"not-"+p->signature());
					new_goal.push_back(notp_idx);

					for( auto const_a : prob.actions_adding( p->index() ) ){
						Action* a = prob.actions()[const_a->index()];
						a->del_vec().push_back(notp_idx);
						a->del_set().set( notp_idx );
						//prob.actions_deleting( notp_idx ).push_back( a );
					}
					for( auto const_a : prob.actions_deleting( p->index() ) ){
						Action* a = prob.actions()[const_a->index()];
						a->add_vec().push_back(notp_idx);
						a->add_set().set( notp_idx );
						//prob.actions_adding( notp_idx ).push_back( a );
					}

					break;
				}
			}
		}
		//STRIPS_Problem::set_goal( prob, new_goal );
		prob.make_action_tables(false);

	}

	void
	compute_mutexes_only_aij() {
		for ( unsigned k = 0; k < m_values.size(); k++ )
			m_values[k] = 0.0f;
		typedef std::pair< unsigned, unsigned > Fluent_Pair;
		typedef std::list< Fluent_Pair > Pair_List;
		Pair_List M;

		// M_A set: pairs p,q s.t. there is at least one a where Add(a) = p and Del(a) = q

		for ( auto a : m_strips_model.actions() )
			for ( auto p : a->add_vec() )
				for ( auto q : a->del_vec() ) {
					if ( value(p,q) == infty ) continue;
					M.push_back( Fluent_Pair(p,q) );
					value(p,q) = infty;
				}

		Pair_List M_B;
		// M_B set:

		for ( auto P : M ) {
			for ( auto a : m_strips_model.actions() ) {
				unsigned p = P.first;
				unsigned q = P.second;
				if ( !a->asserts(p) ) continue;
				for ( auto r : a->prec_vec() ) {
					if ( value(r,q) == infty ) continue;
					M_B.push_back( Fluent_Pair(r,q) );
					value(r, q) = infty;
				}
			}
		}

		for ( auto P : M_B )
			M.push_back( P );

		bool changed;

		do {
			changed = false;

			Pair_List M2;

			for ( auto P : M ) {
				unsigned p = P.first;
				unsigned q = P.second;
				if ( m_strips_model.is_in_init( p ) && m_strips_model.is_in_init(q) ) {
					value(p,q) = 0.0f;
					changed = true;
					continue;
				}

				bool needs_to_be_removed = false;
				for ( auto adding : m_strips_model.actions_adding( p ) ) {
					bool is_deleting = adding->retracts(q);
					bool is_not_adding = !adding->asserts(q);
					if ( !is_deleting && !is_not_adding ) {
						needs_to_be_removed = true;
						break;
					}
					bool r_q_in_M = false;
					for ( auto r : adding->prec_vec() ) {
						if ( value(r,q) == infty ) {
							r_q_in_M = true;
							break;
						}
					}
					if ( !r_q_in_M ) {
						needs_to_be_removed = true;
						break;
					}
				}
				if ( needs_to_be_removed ) {
					value(p,q) = 0.0f;
					changed = true;
				}
				else {
					M2.push_back( P );
				}
			}
			M = M2;

		} while (changed);

	}

protected:

	void extract_edeletes( STRIPS_Problem& prob ){
		for ( unsigned p = 0 ; p < m_strips_model.num_fluents(); p++ ){
			for ( unsigned a = 0; a < m_strips_model.num_actions(); a++ ){
				bool is_edelete = false;
				Action& action = *(prob.actions()[a]);

				for ( unsigned i = 0; i < action.add_vec().size(); i++ ){
					unsigned q = action.add_vec()[i];
					if ( value(p,q) == infty ){
						is_edelete = true;
						prob.actions()[a]->edel_vec().push_back( p );
						prob.actions()[a]->edel_set().set( p );
						prob.actions_edeleting( p ).push_back( (const Action*) &action );
						break;
					}
				}

				if ( is_edelete ) continue;

				for ( unsigned i = 0; i < action.prec_vec().size(); i++ ){
					unsigned r = action.prec_vec()[i];
					if ( !action.add_set().isset(p) && value( p, r ) == infty ){
						prob.actions()[a]->edel_vec().push_back( p );
						prob.actions()[a]->edel_set().set( p );
						prob.actions_edeleting( p ).push_back( (const Action*) &action );
						break;
					}
				}

				if ( !action.edel_set().isset(p) && action.del_set().isset(p) ){
					prob.actions()[a]->edel_vec().push_back( p );
					prob.actions()[a]->edel_set().set( p );
					prob.actions_edeleting( p ).push_back( (const Action*) &action );
				}

			}
		}

	}

	void initialize_ceffs_and_emtpy_precs(){
		//conditional effects and empty precs
		for ( unsigned k = 0; k < m_strips_model.empty_prec_actions().size(); k++ ) {
			const Action& a = *(m_strips_model.empty_prec_actions()[k]);
			float v =  ( cost_opt == H2_Cost_Function::Unit_Costs ? 1.0f :
					( cost_opt == H2_Cost_Function::Use_Costs ? (float)a.cost()  : 1.0f + (float)a.cost()) );

			for ( unsigned i = 0; i< a.add_vec().size(); i++ ){
				for ( unsigned j = i; i < a.add_vec().size(); j++ ){
					unsigned p = a.add_vec()[i];
					unsigned q = a.add_vec()[j];
					value(p,q) = v;

					int curr_idx = H2_Helper::pair_index(p,q);
					if ( !m_already_updated.isset( curr_idx ) ) {
						m_updated.push_back( curr_idx );
						m_already_updated.set( curr_idx );
					}
				}
			}

			// Conditional effects
			for ( unsigned j = 0; j < a.ceff_vec().size(); j++ )
			{
				const Conditional_Effect& ceff = *(a.ceff_vec()[j]);
				if ( !ceff.prec_vec().empty() ) continue;
				float v_eff = v;

				for ( unsigned i = 0; i< ceff.add_vec().size(); i++ ){
					for ( unsigned j = i; i < ceff.add_vec().size(); j++ ){
						unsigned p = ceff.add_vec()[i];
						unsigned q = ceff.add_vec()[j];
						value(p,q) = v_eff;

						int curr_idx = H2_Helper::pair_index(p,q);
						if ( !m_already_updated.isset( curr_idx ) ) {
							m_updated.push_back( curr_idx );
							m_already_updated.set( curr_idx );
						}
					}
				}

			}
		}
	}
	void initialize( const State& s ) {
		for ( unsigned k = 0; k < m_values.size(); k++ )
			m_values[k] = infty;
		for ( unsigned k = 0; k < m_op_values.size(); k++ )
			m_op_values[k] = infty;

		initialize_ceffs_and_emtpy_precs();

		for ( unsigned i = 0; i < s.fluent_vec().size(); i++ )
		{
			unsigned p = s.fluent_vec()[i];
			value(p,p) = 0.0f;

			int curr_idx = H2_Helper::pair_index(p,p);
			if ( !m_already_updated.isset( curr_idx ) ) {
				m_updated.push_back( curr_idx );
				m_already_updated.set( curr_idx );
			}

			for ( unsigned j = i+1; j < s.fluent_vec().size(); j++ )
			{
				unsigned q = s.fluent_vec()[j];
				value(p,q) = 0.0f;

				int curr_idx = H2_Helper::pair_index(p,q);
				if ( !m_already_updated.isset( curr_idx ) ) {
					m_updated.push_back( curr_idx );
					m_already_updated.set( curr_idx );
				}
			}
		}
	}

	void initialize( const Fluent_Vec& f ) {
		for ( unsigned k = 0; k < m_values.size(); k++ )
			m_values[k] = infty;
		for ( unsigned k = 0; k < m_op_values.size(); k++ )
			m_op_values[k] = infty;

		initialize_ceffs_and_emtpy_precs();

		for ( unsigned i = 0; i < f.size(); i++ )
		{
			unsigned p = f[i];
			value(p,p) = 0.0f;

			int curr_idx = H2_Helper::pair_index(p,p);
			if ( !m_already_updated.isset( curr_idx ) ) {
				m_updated.push_back( curr_idx );
				m_already_updated.set( curr_idx );
			}

			for ( unsigned j = i+1; j < f.size(); j++ )
			{
				unsigned q = f[j];
				value(p,q) = 0.0f;

				int curr_idx = H2_Helper::pair_index(p,q);
				if ( !m_already_updated.isset( curr_idx ) ) {
					m_updated.push_back( curr_idx );
					m_already_updated.set( curr_idx );
				}
			}
		}
	}

	void compute() {
		while ( !m_updated.empty() ) {

			unsigned p = m_updated.front();

			m_updated.pop_front();
			m_already_updated.unset(p);

			for ( std::set<unsigned>::iterator action_it = m_relevant_actions[p].begin(); action_it != m_relevant_actions[p].end(); ++action_it) {

				const Action& action = *(m_strips_model.actions()[*action_it]);
				unsigned a = action.index();

				op_value(a) = eval( action.prec_vec() );
				if ( op_value(a) == infty ) continue;

				for ( unsigned i = 0; i < action.add_vec().size(); i++ ) {
					unsigned p = action.add_vec()[i];
					for ( unsigned j = i; j < action.add_vec().size(); j++ ) {
						unsigned q = action.add_vec()[j];
						float curr_value = value(p,q);
						if ( curr_value == 0.0f ) continue;
						float v = op_value(a);
						if ( cost_opt == H2_Cost_Function::Unit_Costs ) v += 1.0f;
						if ( cost_opt == H2_Cost_Function::Use_Costs ) v += action.cost();
						if ( v < curr_value ) {
							value(p,q) = v;
							int curr_idx = H2_Helper::pair_index(p,q);
							if ( !m_already_updated.isset( curr_idx ) ) {
								m_updated.push_back( curr_idx );
								m_already_updated.set( curr_idx );
							}
							curr_idx = H2_Helper::pair_index(p,p);
							if ( !m_already_updated.isset( curr_idx ) ) {
								m_updated.push_back( curr_idx );
								m_already_updated.set( curr_idx );
							}
							curr_idx = H2_Helper::pair_index(q,q);
							if ( !m_already_updated.isset( curr_idx ) ) {
								m_updated.push_back( curr_idx );
								m_already_updated.set( curr_idx );
							}
						}
					}

					for ( unsigned r = 0; r < m_strips_model.num_fluents(); r++ ) {
						if ( interferes( a, r ) || value( p, r ) == 0.0f ) continue;
						float h2_pre_noop = std::max( op_value(a), value(r,r) );
						if ( h2_pre_noop == infty ) continue;
						for ( unsigned j = 0; j < action.prec_vec().size(); j++ ) {
							unsigned s = action.prec_vec()[j];
							h2_pre_noop = std::max( h2_pre_noop, value(r,s) );
						}
						float v = h2_pre_noop;
						if ( cost_opt == H2_Cost_Function::Unit_Costs ) v += 1.0f;
						if ( cost_opt == H2_Cost_Function::Use_Costs ) v += action.cost();
						if ( v < value(p,r) )
						{
							value(p,r) = v;
							int curr_idx = H2_Helper::pair_index(p,r);
							if ( !m_already_updated.isset( curr_idx ) ) {
								m_updated.push_back( curr_idx );
								m_already_updated.set( curr_idx );
							}
							curr_idx = H2_Helper::pair_index(r,r);
							if ( !m_already_updated.isset( curr_idx ) ) {
								m_updated.push_back( curr_idx );
								m_already_updated.set( curr_idx );
							}
						}
					}

				}
			}

		}

	}

	void compute_mutexes_only() {
		while ( !m_updated.empty() ) {

			unsigned p = m_updated.front();

			m_updated.pop_front();
			m_already_updated.unset(p);

			for ( std::set<unsigned>::iterator action_it = m_relevant_actions[p].begin(); action_it != m_relevant_actions[p].end(); ++action_it) {

				const Action& action = *(m_strips_model.actions()[*action_it]);
				unsigned a = action.index();

				op_value(a) = eval( action.prec_vec() );
				if ( op_value(a) == infty ) continue;

				for ( unsigned i = 0; i < action.add_vec().size(); i++ ) {
					unsigned p = action.add_vec()[i];
					for ( unsigned j = i; j < action.add_vec().size(); j++ ) {
						unsigned q = action.add_vec()[j];
						float curr_value = value(p,q);
						if ( curr_value == 0.0f ) continue;
						value(p,q) = 0.0f;
						int curr_idx = H2_Helper::pair_index(p,q);
						if ( !m_already_updated.isset( curr_idx ) ) {
							m_updated.push_back( curr_idx );
							m_already_updated.set( curr_idx );
						}
						curr_idx = H2_Helper::pair_index(p,p);
						if ( !m_already_updated.isset( curr_idx ) ) {
							m_updated.push_back( curr_idx );
							m_already_updated.set( curr_idx );
						}
						curr_idx = H2_Helper::pair_index(q,q);
						if ( !m_already_updated.isset( curr_idx ) ) {
							m_updated.push_back( curr_idx );
							m_already_updated.set( curr_idx );
						}

					}

					for ( unsigned r = 0; r < m_strips_model.num_fluents(); r++ ) {
						if ( interferes( a, r ) || value( p, r ) == 0.0f ) continue;
						float h2_pre_noop = std::max( op_value(a), value(r,r) );
						if ( h2_pre_noop == infty ) continue;

						for ( unsigned j = 0; j < action.prec_vec().size(); j++ ) {
							unsigned s = action.prec_vec()[j];
							h2_pre_noop = std::max( h2_pre_noop, value(r,s) );
							if ( h2_pre_noop == infty ) break;
						}

						if ( h2_pre_noop == infty ) continue;
						value(p,r) = 0.0f;
						int curr_idx = H2_Helper::pair_index(p,r);
						if ( !m_already_updated.isset( curr_idx ) ) {
							m_updated.push_back( curr_idx );
							m_already_updated.set( curr_idx );
						}
						curr_idx = H2_Helper::pair_index(r,r);
						if ( !m_already_updated.isset( curr_idx ) ) {
							m_updated.push_back( curr_idx );
							m_already_updated.set( curr_idx );
						}
					}

				}
			}

		}

	}

protected:

	const STRIPS_Problem&			m_strips_model;
	std::vector<float>			m_values;
	std::vector<float>			m_op_values;
	std::vector< Bit_Set* >			m_interfering_ops;

	std::vector< std::set<unsigned> >       m_relevant_actions;
	boost::circular_buffer<int>		m_updated;
	Bit_Set					m_already_updated;
};

}

}
#endif // h_2.hxx


#include <iostream>
#include <fstream>
#include <sstream>
#include <string>
#include <vector>
#include <iterator>
#include <algorithm>
#include <cmath>
#include <boost/algorithm/string.hpp>

//<deque>, <list>, <map>, <set>, <queue>, <stack>, <vector>
//<vector>: .size(), .push_back(), .pop_back()
//int a[7] vs int* p = new int[7]

using namespace std;

vector<double> convertStringVectortoDoubleVector(const vector<string>& stringVector) {
	vector<double> doubleVector(stringVector.size());
	transform(stringVector.begin(), stringVector.end(), doubleVector.begin(), [](const string& val)
		{
			return stod(val);
		});
	return doubleVector;
}

string double_array_to_string(double double_array[], int size_of_array) {
	string s = "";
	for (int i = 0; i < size_of_array; i++) {
		string temp;
		ostringstream convert;
		convert << double_array[i];
		s += convert.str() + ", ";
	}
	s.erase(s.end()-2, s.end());
	return s;
}

void print_double_vector(vector<double>& v) {
	copy(v.begin(), v.end(), std::ostream_iterator<double>(cout, " "));
	cout << "" << endl;
}

void print_int_vector(vector<int>& v) {
	copy(v.begin(), v.end(), std::ostream_iterator<int>(cout, " "));
	cout << "" << endl;
}

void print_array (int arg[], int length) {
	for (int n=0; n<length; ++n)
		cout << arg[n] << ' ';
	cout << '\n';
}

void print_array (double arg[], int length) {
	for (int n=0; n<length; ++n)
		cout << arg[n] << ' ';
	cout << '\n';
}

//Ignore threshold_adjust and pair_check for now
string query_distance(const char* libfile, const char* values, int threshold_adjust_counts=100, float threshold=0.4) {

	bool threshold_adjust = threshold_adjust_counts == -1 ? false:true;

	//Splitting template distance values into vector
	string dist_str(values);
	vector<string> dist_split_str;
	boost::split(dist_split_str, dist_str, boost::algorithm::is_any_of(","));
	vector<double> dists = convertStringVectortoDoubleVector(dist_split_str);

	int num_dists = dists.size();

	//RMSD values and AA identities for every IP in library file
	vector<int> aas;
	vector<double> rmsds;

	//Keep track of lowest 100 RMSDs
	vector<double> thresholds;

	string line;
	ifstream infile;
	infile.open(libfile);
	while (!infile.eof()) {
		getline(infile, line);
		vector<std::string> split1;
		boost::split(split1, line, boost::algorithm::is_any_of(";"));
		if (split1.size() > 1) {
			string lib_dists_str(split1.at(0));
			vector<std::string> split2;
			boost::split(split2, lib_dists_str, boost::algorithm::is_any_of(","));
			vector<double> lib_dists = convertStringVectortoDoubleVector(split2);

			float rmsd = 0;
			for (int i = 0; i < num_dists; i++) {
				rmsd += (lib_dists.at(i) - dists.at(i)) * (lib_dists.at(i) - dists.at(i));
			}
			rmsd = sqrt(rmsd/(num_dists-1));
			rmsds.push_back(rmsd);

			if (threshold_adjust) {
				if (thresholds.size() < threshold_adjust_counts) {
					thresholds.push_back(rmsd);
					push_heap(thresholds.begin(), thresholds.end());

				} else if (rmsd < thresholds.front()) {
					pop_heap (thresholds.begin(), thresholds.end());
					thresholds.pop_back();

					thresholds.push_back(rmsd);
					push_heap(thresholds.begin(), thresholds.end());
				}
			}

			int aa = stoi(split1.at(1));
			aas.push_back(aa);
		}
		cout << "\n";
	}

	const float ip_probs[20] = {0.064, 0.018, 0.038, 0.043, 0.065, 0.052, 0.025, 0.084, 0.04, 0.127, 0.024, 0.033, 0.038, 0.029, 0.052, 0.044, 0.05, 0.095, 0.025, 0.055};
	infile.close();

	threshold = threshold_adjust ? thresholds.front(): threshold;
	cout << "Using threshold " << threshold << endl;

	int match_counts[20] = {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0};
	double match_probs[20] = {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0};

	//Tally up matches within RMSD threshold
	for (unsigned int i = 0; i < aas.size(); i++)
		if (rmsds[i] < threshold) match_counts[aas[i]] += 1;

	int match_total = 0;
	for (int i = 0; i < 20; i++) {
		if (match_counts[i] == 0) match_counts[i] += 1;
		match_total += match_counts[i];
	}

	for (int i = 0; i < 20; i++)
		match_probs[i] = match_counts[i]/match_total;

	double energies[20];
	for (int i = 0; i < 20; i++) {
		if (abs(match_probs[i] - ip_probs[i]) < 0.0001 || match_probs[i] == 0.0) {
			energies[i] = 0.0;
		} else {
			energies[i] = -1.0*log10(match_probs[i] / ip_probs[i]);
		}
	}

	return double_array_to_string(energies, 20);
}

#include <boost/python/module.hpp>
#include <boost/python/def.hpp>
using namespace boost::python;

BOOST_PYTHON_MODULE(statium_cpp)
{
	def("query_distance", query_distance);
}