#include <catch2/catch_test_macros.hpp>
#include <catch2/matchers/catch_matchers_floating_point.hpp>

#include <data/Molecule.h>
#include <data/Body.h>
#include <grid/Grid.h>
#include <grid/detail/GridMember.h>
#include <constants/Constants.h>
#include <utility/Utility.h>
#include <fitter/LinearFitter.h>
#include <settings/All.h>
#include <fitter/SmartFitter.h>
#include <hydrate/generation/RadialHydration.h>
#include <hist/histogram_manager/HistogramManager.h>
#include <hist/intensity_calculator/CompositeDistanceHistogram.h>
#include <hist/intensity_calculator/ExactDebyeCalculator.h>
#include <io/pdb/PDBStructure.h>
#include <rigidbody/BodySplitter.h>

#include <vector>
#include <string>
#include <iostream>
#include <fstream>

using namespace ausaxs;
using namespace data;
using std::cout, std::endl;

struct fixture {
    fixture() {
        settings::molecule::center = false;
        settings::molecule::implicit_hydrogens = false;
    }

    AtomFF a1 = AtomFF({-1, -1, -1}, form_factor::form_factor_t::C);
    AtomFF a2 = AtomFF({-1,  1, -1}, form_factor::form_factor_t::C);
    AtomFF a3 = AtomFF({ 1, -1, -1}, form_factor::form_factor_t::C);
    AtomFF a4 = AtomFF({ 1,  1, -1}, form_factor::form_factor_t::C);
    AtomFF a5 = AtomFF({-1, -1,  1}, form_factor::form_factor_t::C);
    AtomFF a6 = AtomFF({-1,  1,  1}, form_factor::form_factor_t::C);
    AtomFF a7 = AtomFF({ 1, -1,  1}, form_factor::form_factor_t::C);
    AtomFF a8 = AtomFF({ 1,  1,  1}, form_factor::form_factor_t::C);
    
    Water w1 = Water({-1, -1, -1});
    Water w2 = Water({-1,  1, -1});

    std::vector<AtomFF> b1 = {a1, a2};
    std::vector<AtomFF> b2 = {a3, a4};
    std::vector<AtomFF> b3 = {a5, a6};
    std::vector<AtomFF> b4 = {a7, a8};
    std::vector<Body> bodies = {Body(b1, std::vector{w1, w2}), Body(b2), Body(b3), Body(b4)};
};

/**
 * @brief Compare two histograms. 
 *        Only indices [0, p1.size()] are checked.
 */
bool compare_hist(Vector<double> p1, Vector<double> p2) {
    for (unsigned int i = 0; i < p1.size(); i++) {
        if (!utility::approx(p1[i], p2[i])) {
            std::cout << "Failed on index " << i << ". Values: " << p1[i] << ", " << p2[i] << std::endl;
            return false;
        }
    }
    return true;
}

TEST_CASE_METHOD(fixture, "Molecule::Molecule") {
    settings::general::verbose = false;

    SECTION("vector<Body>&&") {
        Molecule protein(std::move(bodies));
        REQUIRE(protein.size_body() == 4);
        CHECK(protein.get_body(0).size_atom() == 2);
        CHECK(protein.get_body(1).size_atom() == 2);
        CHECK(protein.get_body(2).size_atom() == 2);
        CHECK(protein.get_body(3).size_atom() == 2);
    }

    SECTION("vector<string>&") {
        std::vector<std::string> files = {"tests/files/2epe.pdb", "tests/files/2epe.pdb"};
        Molecule protein(files);
        Body body("tests/files/2epe.pdb"); // compare with body constructor

        REQUIRE(protein.size_body() == 2);
        CHECK(protein.size_atom() == 2002);
        CHECK(protein.size_water() == 96);
        CHECK(protein.get_body(0).equals_content(protein.get_body(1)));
        CHECK(protein.get_body(0).equals_content(body));
    }

    SECTION("ExistingFile&") {
        SECTION("fake data") {
            io::File path("temp/io/temp.pdb");
            path.create();

            std::ofstream pdb_file(path);
            pdb_file << "ATOM      1  CB  ARG A 129         2.1     3.2     4.3  0.50 42.04           C " << std::endl;
            pdb_file << "ATOM      2  CB  ARG A 129         3.2     4.3     5.4  0.50 42.04           C " << std::endl;
            pdb_file << "TER       3      ARG A 129                                                     " << std::endl;
            pdb_file << "HETATM    4  O   HOH A 130      30.117  29.049  34.879  0.94 34.19           O " << std::endl;
            pdb_file << "HETATM    5  O   HOH A 131      31.117  30.049  35.879  0.94 34.19           O " << std::endl;
            pdb_file.close();

            // check PDB io
            Molecule protein(path);

            REQUIRE(protein.size_atom() == 2);
            auto atoms = protein.get_atoms();
            CHECK(atoms[0].coordinates().x() == 2.1);
            CHECK(atoms[0].coordinates().y() == 3.2);
            CHECK(atoms[0].coordinates().z() == 4.3);
            CHECK(atoms[0].form_factor_type() == form_factor::form_factor_t::C);
            CHECK(atoms[0].weight() == 0.5*constants::charge::get_ff_charge(atoms[0].form_factor_type()));

            REQUIRE(protein.size_water() == 2);
            auto waters = protein.get_waters();
            CHECK(waters[0].coordinates().x() == 30.117);
            CHECK(waters[0].coordinates().y() == 29.049);
            CHECK(waters[0].coordinates().z() == 34.879);
            CHECK(waters[0].weight() == constants::charge::get_ff_charge(waters[0].form_factor_type()));
        }

        SECTION("real data") {
            io::ExistingFile file("tests/files/2epe.pdb");
            Molecule protein(file);
            Body body("tests/files/2epe.pdb"); // compare with body constructor

            REQUIRE(protein.size_body() == 1);
            CHECK(protein.size_atom() == 1001);
            CHECK(protein.size_water() == 48);
            CHECK(protein.get_body(0).equals_content(body));
        }
    }

    SECTION("vector<Body>&") {
        Molecule protein(bodies);
        REQUIRE(protein.size_body() == 4);
        CHECK(protein.get_body(0).size_atom() == 2);
        CHECK(protein.get_body(1).size_atom() == 2);
        CHECK(protein.get_body(2).size_atom() == 2);
        CHECK(protein.get_body(3).size_atom() == 2);
        REQUIRE(protein.size_water() == 2);
        auto waters = protein.get_waters();
        CHECK(waters[0] == w1);
        CHECK(waters[1] == w2);
    }

    SECTION("vector<Body>&") {
        Molecule protein(bodies);
        REQUIRE(protein.size_body() == 4);
        CHECK(protein.get_body(0).size_atom() == 2);
        CHECK(protein.get_body(1).size_atom() == 2);
        CHECK(protein.get_body(2).size_atom() == 2);
        CHECK(protein.get_body(3).size_atom() == 2);
    }
}

TEST_CASE("Molecule::get_Rg", "[files]") {
    // tests without effective charge are compared against the electron Rg from CRYSOL
    settings::general::verbose = false;
    settings::molecule::implicit_hydrogens = false;

    SECTION("2epe") {
        Molecule protein("tests/files/2epe.pdb");
        REQUIRE_THAT(protein.get_Rg(), Catch::Matchers::WithinAbs(13.89, 0.01));
    }

    SECTION("6lyz") {
        Molecule protein("tests/files/6lyz.pdb");
        REQUIRE_THAT(protein.get_Rg(), Catch::Matchers::WithinAbs(13.99, 0.01));
    }

    SECTION("LAR1-2") {
        Molecule protein("tests/files/LAR1-2.pdb");
        REQUIRE_THAT(protein.get_Rg(), Catch::Matchers::WithinAbs(28.91, 0.02));
    }

    settings::molecule::allow_unknown_atoms = true;
    SECTION("SASDJQ4") {
        Molecule protein("tests/files/SASDJQ4.pdb");
        REQUIRE_THAT(protein.get_Rg(), Catch::Matchers::WithinAbs(28.08, 0.02));
    }
    settings::molecule::allow_unknown_atoms = false;
}

TEST_CASE("Molecule::simulate_dataset", "[files]") {
    settings::axes::qmax = 0.4;
    settings::general::verbose = false;
    settings::em::sample_frequency = 2;
    Molecule protein("tests/files/2epe.pdb");

    SimpleDataset data = protein.simulate_dataset();
    fitter::LinearFitter fitter(data, protein.get_histogram());
    auto res = fitter.fit();
    REQUIRE_THAT(res->fval/res->dof, Catch::Matchers::WithinAbs(1., 0.5));
    // plots::PlotIntensityFit plot1(res);
    // plot1.save("figures/tests/protein/check_chi2_1.png");
}

TEST_CASE_METHOD(fixture, "Molecule::get_cm") {
    Molecule protein(bodies);
    protein.translate({1, 2, 3});
    protein.center();
    Vector3<double> cm = protein.get_cm();
    REQUIRE(cm == Vector3<double>{0, 0, 0});
}

// TEST_CASE_METHOD(fixture, "Molecule::get_volume", "[broken]") {
//     // broken since it was supposed to use the old protein.get_volume_acids() method
//     // since the protein does not consist of a complete amino acid, the volume is not correct
//     // TODO: create a protein containing a full amino acid and check if the volume is roughly correct
//     Molecule protein(bodies);
//     REQUIRE_THAT(protein.get_volume_grid(), Catch::Matchers::WithinRel(4*constants::volume::amino_acids.get("LYS")));
// }

TEST_CASE_METHOD(fixture, "Molecule::get_histogram", "[files]") {
    settings::general::verbose = false;

    SECTION("delegated to HistogramManager") {
        Molecule protein(bodies);
        REQUIRE(compare_hist(protein.get_histogram()->get_weighted_counts(), protein.get_histogram_manager()->calculate()->get_weighted_counts()));
    }
 
    SECTION("compare_debye") {
        std::vector<AtomFF> atoms = {
            AtomFF({-1, -1, -1}, form_factor::form_factor_t::C), AtomFF({-1, 1, -1}, form_factor::form_factor_t::C),
            AtomFF({ 1, -1, -1}, form_factor::form_factor_t::C), AtomFF({ 1, 1, -1}, form_factor::form_factor_t::C),
            AtomFF({-1, -1,  1}, form_factor::form_factor_t::C), AtomFF({-1, 1,  1}, form_factor::form_factor_t::C),
            AtomFF({ 1, -1,  1}, form_factor::form_factor_t::C), AtomFF({ 1, 1,  1}, form_factor::form_factor_t::C)
        };
        Molecule protein({Body{atoms}});

        std::vector<double> I_dumb = hist::exact_debye_transform(protein, constants::axes::q_axis.as_vector());
        std::vector<double> I_smart = protein.get_histogram()->debye_transform().get_counts();

        for (int i = 0; i < 8; i++) {
            if (!utility::approx(I_dumb[i], I_smart[i], 1e-1)) {
                cout << "Failed on index " << i << ". Values: " << I_dumb[i] << ", " << I_smart[i] << endl;
                REQUIRE(false);
            }
        }
        SUCCEED();
    }

    SECTION("compare_debye_real") {
        Molecule protein("tests/files/2epe.pdb");
        protein.clear_hydration();

        std::vector<double> I_dumb = hist::exact_debye_transform(protein, constants::axes::q_axis.as_vector());
        std::vector<double> I_smart = protein.get_histogram()->debye_transform().get_counts();

        for (int i = 0; i < 8; i++) {
            if (!utility::approx(I_dumb[i], I_smart[i], 1e-3, 0.05)) {
                cout << "Failed on index " << i << ". Values: " << I_dumb[i] << ", " << I_smart[i] << endl;
                REQUIRE(false);
            }
        }
        SUCCEED();
    }
}

TEST_CASE_METHOD(fixture, "Molecule::get_total_histogram") {
    Molecule protein(bodies);
    REQUIRE(compare_hist(protein.get_histogram()->get_weighted_counts(), protein.get_histogram_manager()->calculate_all()->get_weighted_counts()));
    // REQUIRE(protein.get_histogram() == protein.get_histogram_manager()->calculate_all());
}

TEST_CASE("Molecule::save", "[files]") {
    settings::general::verbose = false;
    settings::molecule::implicit_hydrogens = false;

    Molecule protein("tests/files/2epe.pdb");
    protein.save("temp/tests/protein_save_2epe.pdb");
    Molecule protein2("temp/tests/protein_save_2epe.pdb");
    auto atoms1 = protein.get_atoms();
    auto atoms2 = protein2.get_atoms();
    REQUIRE(atoms1.size() == atoms2.size());

    // we have to manually compare the data since saving & loading will necessarily round the doubles
    for (size_t i = 0; i < atoms1.size(); ++i) {
        bool ok = true;
        auto& a1 = atoms1[i];
        auto& a2 = atoms2[i];
        if (1e-3 < abs(a1.coordinates().x() - a2.coordinates().x()) + abs(a1.coordinates().y() - a2.coordinates().y()) + abs(a1.coordinates().z() - a2.coordinates().z())) {ok = false;}
        if (a1.form_factor_type() != a2.form_factor_type()) {ok = false;}
        REQUIRE(ok);
    }
}

TEST_CASE_METHOD(fixture, "Molecule::generate_new_hydration", "[files]") {
    settings::general::verbose = false;
    hydrate::RadialHydration::set_noise_generator([] () {return Vector3<double>{0, 0, 0};});

    SECTION("generates new waters") {
        // the molecule is really small, so we have to make sure there's enough space for the waters
        settings::grid::scaling = 5;
        Molecule protein(bodies);
        protein.generate_new_hydration();
        REQUIRE(protein.size_water() != 0);
        settings::grid::scaling = 0.25;
    }

    // we want to check that the hydration shells are consistent for fitting purposes
    SECTION("consistent hydration generation") {
        Molecule protein("tests/files/2epe.pdb");
        fitter::LinearFitter fitter(SimpleDataset{"tests/files/2epe.dat"}, protein.get_histogram());

        protein.generate_new_hydration();
        double chi2 = fitter.fit()->fval;

        for (int i = 0; i < 10; i++) {
            protein.generate_new_hydration();
            double _chi2 = fitter.fit()->fval;
            REQUIRE_THAT(chi2, Catch::Matchers::WithinRel(_chi2));
        }
    }
}

TEST_CASE("Molecule::get_volume_grid", "[files]") {
    settings::general::verbose = false;
    Molecule protein("tests/files/2epe.pdb");
    REQUIRE(protein.get_volume_grid() == protein.get_grid()->get_volume());
}

// TEST_CASE("Molecule::get_volume_calpha") {    
//     CHECK(false);
// }

TEST_CASE("Molecule::get_molar_mass", "[files]") {
    settings::general::verbose = false;
    Molecule protein("tests/files/2epe.pdb");
    REQUIRE(protein.get_molar_mass() == protein.get_absolute_mass()*constants::Avogadro);
}

TEST_CASE("Molecule::get_absolute_mass", "[files]") {
    settings::general::verbose = false;
    Molecule protein("tests/files/2epe.pdb");
    double sum = 0;
    for (auto& atom : protein.get_atoms()) {
        sum += constants::mass::get_mass(atom.form_factor_type());
    }
    for (auto& water : protein.get_waters()) {
        sum += constants::mass::get_mass(water.form_factor_type());
    }
    REQUIRE(protein.get_absolute_mass() == sum);
}

TEST_CASE("Molecule::get_total_atomic_charge", "[files]") {
    settings::general::verbose = false;
    Molecule protein("tests/files/2epe.pdb");
    double sum = 0;
    for (auto& atom : protein.get_atoms()) {
        sum += atom.weight();
    }
    REQUIRE(protein.get_total_atomic_charge() == sum);
}

TEST_CASE("Molecule::get_relative_charge_density", "[files]") {
    settings::general::verbose = false;
    Molecule protein("tests/files/2epe.pdb");
    REQUIRE_THAT(
        protein.get_relative_charge_density(), 
        Catch::Matchers::WithinAbs(
            (protein.get_total_atomic_charge() - constants::charge::density::water*protein.get_volume_grid())/protein.get_volume_grid(), 
            1e-6
        )
    );
}

TEST_CASE("Molecule::get_relative_mass_density", "[files]") {
    settings::general::verbose = false;
    Molecule protein("tests/files/2epe.pdb");
    Catch::Matchers::WithinAbs(
        protein.get_relative_mass_density() - (protein.get_absolute_mass() - constants::mass::density::water*protein.get_volume_grid())/protein.get_volume_grid(),
        1e-6
    );
}

TEST_CASE("Molecule::get_relative_charge", "[files]") {
    settings::general::verbose = false;
    Molecule protein("tests/files/2epe.pdb");
    Catch::Matchers::WithinAbs(
        protein.get_relative_charge() - (protein.get_total_atomic_charge() - protein.get_volume_grid()*constants::charge::density::water),
        1e-6    
    );
}

TEST_CASE_METHOD(fixture, "Molecule::get_grid") {
    Molecule protein(bodies);
    // we just want to test that the grid is created by default
    REQUIRE(protein.get_grid() != nullptr);
}

TEST_CASE_METHOD(fixture, "Molecule::set_grid") {
    Molecule protein(bodies);
    grid::Grid grid(Limit3D(0, 1, 0, 1, 0, 1));
    auto grid_dup = grid;
    protein.set_grid(std::move(grid_dup));
    REQUIRE(*protein.get_grid() == grid);
}

TEST_CASE_METHOD(fixture, "Molecule::clear_hydration") {
    Molecule protein2(bodies);
    protein2.get_body(0).set_hydration(hydrate::Hydration::create(std::vector{w1, w2}));
    REQUIRE(protein2.size_water() != 0);
    protein2.clear_hydration();
    REQUIRE(protein2.size_water() == 0);
}

TEST_CASE_METHOD(fixture, "Molecule::center") {
    Molecule protein(bodies);
    protein.translate({1, 2, 3});
    protein.center();
    REQUIRE(protein.get_cm() == Vector3<double>{0, 0, 0});

    protein.translate(Vector3<double>{1, 1, 1});
    REQUIRE(protein.get_cm() == Vector3<double>{1, 1, 1});
    
    protein.center();
    REQUIRE(protein.get_cm() == Vector3<double>{0, 0, 0});
}

TEST_CASE_METHOD(fixture, "Molecule::get_body") {
    Molecule protein(bodies);
    REQUIRE(protein.get_body(0) == protein.get_bodies()[0]);
    REQUIRE(protein.get_body(1) == protein.get_bodies()[1]);
    REQUIRE(protein.get_body(2) == protein.get_bodies()[2]);
    REQUIRE(protein.get_body(3) == protein.get_bodies()[3]);
}

TEST_CASE_METHOD(fixture, "Molecule::get_bodies") {
    Molecule protein(bodies);
    REQUIRE(protein.get_bodies() == bodies);
}

TEST_CASE_METHOD(fixture, "Molecule::get_atoms") {
    Molecule protein(bodies);
    REQUIRE(protein.get_atoms() == std::vector<AtomFF>{a1, a2, a3, a4, a5, a6, a7, a8});
}

TEST_CASE_METHOD(fixture, "Molecule::get_waters") {
    Molecule protein2(bodies);
    REQUIRE(protein2.get_waters() == std::vector<Water>{w1, w2});
}

TEST_CASE_METHOD(fixture, "Molecule::get_water") {
    Molecule protein2(bodies);
    auto waters = protein2.get_waters();
    REQUIRE(waters[0] == w1);
    REQUIRE(waters[1] == w2);
}

TEST_CASE_METHOD(fixture, "Molecule::create_grid") {
    Molecule protein(bodies);
    protein.clear_grid();
    auto grid = protein.get_grid();
    protein.create_grid();
    REQUIRE(protein.get_grid() != grid);
}

TEST_CASE_METHOD(fixture, "Molecule::size_body") {
    Molecule protein(bodies);
    CHECK(protein.size_body() == 4);
}

TEST_CASE_METHOD(fixture, "Molecule::size_atom") {
    Molecule protein(bodies);
    CHECK(protein.size_atom() == 8);
}

TEST_CASE_METHOD(fixture, "Molecule::size_water") {
    Molecule protein(bodies);
    protein.clear_hydration();
    CHECK(protein.size_water() == 0);
    Molecule protein2(bodies);
    CHECK(protein2.size_water() == 2);
}

TEST_CASE_METHOD(fixture, "Molecule::get_histogram_manager") {
    Molecule protein(bodies);
    CHECK(protein.get_histogram_manager() != nullptr);
}

// TEST_CASE_METHOD(fixture, "Molecule::set_histogram_manager") {
//     Molecule protein = Molecule(bodies, {});
//     auto hm = protein.get_histogram_manager();
// }

TEST_CASE("Molecule::translate", "[files]") {
    Molecule protein("tests/files/2epe.pdb");
    Vector3<double> cm = protein.get_cm();
    protein.translate(Vector3<double>{1, 1, 1});
    REQUIRE(protein.get_cm() == cm + Vector3<double>{1, 1, 1});
}

TEST_CASE("Molecule::histogram", "[files]") {
    SECTION("multiple bodies, simple") {
        // make the protein
        std::vector<AtomFF> b1 = {AtomFF({-1, -1, -1}, form_factor::form_factor_t::C), AtomFF({-1, 1, -1}, form_factor::form_factor_t::C)};
        std::vector<AtomFF> b2 = {AtomFF({ 1, -1, -1}, form_factor::form_factor_t::C), AtomFF({ 1, 1, -1}, form_factor::form_factor_t::C)};
        std::vector<AtomFF> b3 = {AtomFF({-1, -1,  1}, form_factor::form_factor_t::C), AtomFF({-1, 1,  1}, form_factor::form_factor_t::C)};
        std::vector<AtomFF> b4 = {AtomFF({ 1, -1,  1}, form_factor::form_factor_t::C), AtomFF({ 1, 1,  1}, form_factor::form_factor_t::C)};
        std::vector<Body> ap = {Body(b1), Body(b2), Body(b3), Body(b4)};
        Molecule many(ap);

        // make the body
        std::vector<AtomFF> ab = {
            AtomFF({-1, -1, -1}, form_factor::form_factor_t::C), AtomFF({-1, 1, -1}, form_factor::form_factor_t::C),
            AtomFF({ 1, -1, -1}, form_factor::form_factor_t::C), AtomFF({ 1, 1, -1}, form_factor::form_factor_t::C),
            AtomFF({-1, -1,  1}, form_factor::form_factor_t::C), AtomFF({-1, 1,  1}, form_factor::form_factor_t::C),
            AtomFF({ 1, -1,  1}, form_factor::form_factor_t::C), AtomFF({ 1, 1,  1}, form_factor::form_factor_t::C)
        };
        Molecule one({Body{ab}});

        // create some water molecules
        std::vector<Water> ws(10);
        for (size_t i = 0; i < ws.size(); i++) {
            ws[i] = Water(Vector3<double>(i, i, i));
        }

        many.get_waters() = ws;
        one.get_waters() = ws;

        // we now have a protein consisting of three bodies with the exact same contents as a single body.
        // the idea is now to compare the ScatteringHistogram output from their distance calculations, since it
        // is far easier to do for the single body. 
        auto d_m = many.get_histogram();
        auto d_o = one.get_histogram();

        // direct access to the histogram data (only p is defined)
        const std::vector<double>& p_m = d_m->get_weighted_counts();
        const std::vector<double>& p_o = d_o->get_weighted_counts();

        // compare each entry
        for (size_t i = 0; i < p_o.size(); i++) {
            if (!utility::approx(p_o[i], p_m[i])) {
                cout << "Failed on index " << i << ". Values: " << p_m[i] << ", " << p_o[i] << endl;
                REQUIRE(false);
            }
        }
        REQUIRE(true);
    }

    SECTION("multiple bodies, real input") {
        Body body("tests/files/2epe.pdb");
        
        // We iterate through the protein data from the body, and split it into multiple pieces of size 100.  
        std::vector<Body> patoms;           // vector containing the pieces we split it into
        std::vector<AtomFF> p_current(100); // vector containing the current piece
        unsigned int index = 0;             // current index in p_current
        for (unsigned int i = 0; i < body.get_atoms().size(); i++) {
            p_current[index] = body.get_atom(i);
            index++;
            if (index == 100) { // if index is 100, reset to 0
                patoms.emplace_back(p_current);
                index = 0;
            }
        }

        // add the final few atoms to our list
        if (index != 0) {
            p_current.resize(index);
            patoms.emplace_back(p_current);
        }

        // create the atom, and perform a sanity check on our extracted list
        Molecule protein(patoms);
        std::vector<AtomFF> protein_atoms = protein.get_atoms();
        std::vector<AtomFF> body_atoms = body.get_atoms();

        // sizes must be equal. this also serves as a separate consistency check on the body generation. 
        if (protein_atoms.size() != body_atoms.size()) {
            cout << "Sizes " << protein_atoms.size() << " and " << body_atoms.size() << " should be equal. " << endl;
            REQUIRE(false);
        }

        // stronger consistency check - we check that all atoms are equal, and appear in the exact same order
        for (unsigned int i = 0; i < protein_atoms.size(); i++) {
            if (protein_atoms[i] != body_atoms[i]) {
                cout << "Comparison failed on index " << i << endl;
                cout << protein_atoms[i].coordinates() << endl;
                cout << body_atoms[i].coordinates() << endl;
                REQUIRE(false);
            }
        }

        // generate a hydration layer for the protein, and copy it over to the body
        protein.generate_new_hydration();

        // generate the distance histograms
        auto d_p = protein.get_histogram();
        auto d_b = hist::HistogramManager<false, false>(&protein).calculate_all();

        // direct access to the histogram data (only p is defined)
        const std::vector<double>& p = d_p->get_weighted_counts();
        const std::vector<double>& b_tot = d_b->get_weighted_counts();

        // compare each entry
        for (unsigned int i = 0; i < b_tot.size(); i++) {
            if (!utility::approx(p[i], b_tot[i])) {
                cout << "Failed on index " << i << ". Values: " << p[i] << ", " << b_tot[i] << endl;
                REQUIRE(false);
            }
        }
        REQUIRE(true);
    }

    SECTION("equivalent to old approach") {
        std::vector<AtomFF> atoms = {
            AtomFF({-1, -1, -1}, form_factor::form_factor_t::C), AtomFF({-1, 1, -1}, form_factor::form_factor_t::C),
            AtomFF({ 1, -1, -1}, form_factor::form_factor_t::C), AtomFF({ 1, 1, -1}, form_factor::form_factor_t::C),
            AtomFF({-1, -1,  1}, form_factor::form_factor_t::C), AtomFF({-1, 1,  1}, form_factor::form_factor_t::C),
            AtomFF({ 1, -1,  1}, form_factor::form_factor_t::C), AtomFF({ 1, 1,  1}, form_factor::form_factor_t::C)
        };

        // new auto-scaling approach
        Molecule protein1({Body{atoms}});
        protein1.set_grid(grid::Grid(atoms));

        // old approach
        Molecule protein2({Body{atoms}});
        {
            grid::Grid grid2({-2, 2, -2, 2, -2, 2}); 
            grid2.add(Body{atoms});
            protein2.set_grid(std::move(grid2));
        }

        // generate the distance histograms
        auto h1 = protein1.get_histogram();
        auto h2 = protein2.get_histogram();

        // direct access to the histogram data (only p is defined)
        const std::vector<double>& p1 = h1->get_weighted_counts();
        const std::vector<double>& p2 = h2->get_weighted_counts();

        // compare each entry
        for (size_t i = 0; i < p1.size(); i++) {
            if (!utility::approx(p1[i], p2[i])) {
                cout << "Failed on index " << i << ". Values: " << p1[i] << ", " << p2[i] << endl;
                REQUIRE(false);
            }
        }
        REQUIRE(true);
    }
}

#include <data/state/StateManager.h>
#include <data/state/BoundSignaller.h>
#include <hist/histogram_manager/IPartialHistogramManager.h>
TEST_CASE_METHOD(fixture, "Molecule::bind_body_signallers") {
    settings::general::verbose = false;

    Molecule protein(bodies);
    protein.set_histogram_manager(settings::hist::HistogramManagerChoice::PartialHistogramManager);

    SECTION("at construction") {
        auto& bodies = protein.get_bodies();
        REQUIRE(bodies.size() == 4);
        auto manager = static_cast<hist::IPartialHistogramManager*>(protein.get_histogram_manager())->get_state_manager();
        for (unsigned int i = 0; i < bodies.size(); ++i) {
            CHECK(std::dynamic_pointer_cast<signaller::BoundSignaller>(bodies[i].get_signaller()) != nullptr);
            CHECK(manager->get_probe(i) == bodies[i].get_signaller());
        }

        manager->reset_to_false();
        for (unsigned int i = 0; i < bodies.size(); ++i) {
            bodies[i].get_signaller()->modified_external();
            CHECK(manager->is_externally_modified(i));
        }
    }

    SECTION("after construction") {
        auto& bodies = protein.get_bodies();
        REQUIRE(bodies.size() == 4);
        protein.set_histogram_manager(settings::hist::HistogramManagerChoice::PartialHistogramManager);
        auto manager = static_cast<hist::IPartialHistogramManager*>(protein.get_histogram_manager())->get_state_manager();

        for (unsigned int i = 0; i < bodies.size(); ++i) {
            CHECK(std::dynamic_pointer_cast<signaller::BoundSignaller>(bodies[i].get_signaller()) != nullptr);
            CHECK(manager->get_probe(i) == bodies[i].get_signaller());
        }
    }
}

TEST_CASE_METHOD(fixture, "Molecule::signal_modified_hydration_layer") {
    Molecule protein(bodies);
    protein.set_histogram_manager(settings::hist::HistogramManagerChoice::PartialHistogramManager);
    auto manager = static_cast<hist::IPartialHistogramManager*>(protein.get_histogram_manager())->get_state_manager();
    manager->reset_to_false();
    REQUIRE(manager->is_modified_hydration() == false);

    protein.signal_modified_hydration_layer();
    REQUIRE(manager->is_modified_hydration() == true);
}

TEST_CASE("Molecule: implicit hydrogens") {
    auto generate_molecule = [] () {
        std::vector<io::pdb::PDBAtom> a = {
            io::pdb::PDBAtom(1, "N",  "", "LYS", 'A', 1, "", Vector3<double>(0, 0, 0), 1, 0, constants::atom_t::N, "0"),
            io::pdb::PDBAtom(2, "CA", "", "LYS", 'A', 1, "", Vector3<double>(0, 0, 0), 1, 0, constants::atom_t::C, "0"),
            io::pdb::PDBAtom(3, "C",  "", "LYS", 'A', 1, "", Vector3<double>(0, 0, 0), 1, 0, constants::atom_t::C, "0"),
            io::pdb::PDBAtom(4, "O",  "", "LYS", 'A', 1, "", Vector3<double>(0, 0, 0), 1, 0, constants::atom_t::O, "0"),
            io::pdb::PDBAtom(5, "CB", "", "LYS", 'A', 1, "", Vector3<double>(0, 0, 0), 1, 0, constants::atom_t::C, "0"),
            io::pdb::PDBAtom(6, "CG", "", "LYS", 'A', 1, "", Vector3<double>(0, 0, 0), 1, 0, constants::atom_t::C, "0"),
            io::pdb::PDBAtom(7, "CD", "", "LYS", 'A', 1, "", Vector3<double>(0, 0, 0), 1, 0, constants::atom_t::C, "0"),
            io::pdb::PDBAtom(8, "CE", "", "LYS", 'A', 1, "", Vector3<double>(0, 0, 0), 1, 0, constants::atom_t::C, "0"),
            io::pdb::PDBAtom(9, "NZ", "", "LYS", 'A', 1, "", Vector3<double>(0, 0, 0), 1, 0, constants::atom_t::N, "0"),
        };
        return io::pdb::PDBStructure({a, {}});
    };

    SECTION("enabled") {
        settings::molecule::implicit_hydrogens = true;
        auto file = generate_molecule();
        file.add_implicit_hydrogens();
        auto res = file.reduced_representation();
        Molecule protein({Body{res.atoms, res.waters}});
        auto atoms = protein.get_atoms();

        // Weight is I0(grouped form factor) + hydrogen count
        CHECK(atoms[0].weight() == constants::charge::get_ff_charge(atoms[0].form_factor_type()) + 1);
        CHECK(atoms[0].form_factor_type() == form_factor::form_factor_t::NH);

        CHECK(atoms[1].weight() == constants::charge::get_ff_charge(atoms[1].form_factor_type()) + 1);
        CHECK(atoms[1].form_factor_type() == form_factor::form_factor_t::CH);

        CHECK(atoms[2].weight() == constants::charge::get_ff_charge(atoms[2].form_factor_type()) + 0);
        CHECK(atoms[2].form_factor_type() == form_factor::form_factor_t::C);

        CHECK(atoms[3].weight() == constants::charge::get_ff_charge(atoms[3].form_factor_type()) + 0);
        CHECK(atoms[3].form_factor_type() == form_factor::form_factor_t::O);

        CHECK(atoms[4].weight() == constants::charge::get_ff_charge(atoms[4].form_factor_type()) + 2);
        CHECK(atoms[4].form_factor_type() == form_factor::form_factor_t::CH2);

        CHECK(atoms[5].weight() == constants::charge::get_ff_charge(atoms[5].form_factor_type()) + 2);
        CHECK(atoms[5].form_factor_type() == form_factor::form_factor_t::CH2);

        CHECK(atoms[6].weight() == constants::charge::get_ff_charge(atoms[6].form_factor_type()) + 2);
        CHECK(atoms[6].form_factor_type() == form_factor::form_factor_t::CH2);

        CHECK(atoms[7].weight() == constants::charge::get_ff_charge(atoms[7].form_factor_type()) + 2);
        CHECK(atoms[7].form_factor_type() == form_factor::form_factor_t::CH2);

        CHECK(atoms[8].weight() == constants::charge::get_ff_charge(atoms[8].form_factor_type()) + 3);
        CHECK(atoms[8].form_factor_type() == form_factor::form_factor_t::NH3);
    }

    SECTION("disabled") {
        settings::molecule::implicit_hydrogens = false;
        auto file = generate_molecule();
        auto res = file.reduced_representation();
        Molecule protein({Body{res.atoms, res.waters}});

        for (auto a : protein.get_atoms()) {
            CHECK(a.weight() == constants::charge::get_ff_charge(a.form_factor_type()));
        }

        auto atoms = protein.get_atoms();
        CHECK(atoms[0].form_factor_type() == form_factor::form_factor_t::N);
        CHECK(atoms[1].form_factor_type() == form_factor::form_factor_t::C);
        CHECK(atoms[2].form_factor_type() == form_factor::form_factor_t::C);
        CHECK(atoms[3].form_factor_type() == form_factor::form_factor_t::O);
        CHECK(atoms[4].form_factor_type() == form_factor::form_factor_t::C);
        CHECK(atoms[5].form_factor_type() == form_factor::form_factor_t::C);
        CHECK(atoms[6].form_factor_type() == form_factor::form_factor_t::C);
        CHECK(atoms[7].form_factor_type() == form_factor::form_factor_t::C);
        CHECK(atoms[8].form_factor_type() == form_factor::form_factor_t::N);
    }
}

TEST_CASE("Molecule::iterate_atoms") {
    SECTION("empty") {
        Molecule molecule;
        size_t count = 0;
        for (const auto& atom : molecule.iterate_atoms()) {
            (void)atom; // to silence unused variable warning
            ++count;
        }
        REQUIRE(count == 0);
    }

    SECTION("simple") {
        std::vector<AtomFF> atoms = {
            AtomFF({-1, -1, -1}, form_factor::form_factor_t::C), AtomFF({-1, 1, -1}, form_factor::form_factor_t::NH2),
            AtomFF({ 1, -1, -1}, form_factor::form_factor_t::H), AtomFF({ 1, 1, -1}, form_factor::form_factor_t::SH),
            AtomFF({-1, -1,  1}, form_factor::form_factor_t::CH3), AtomFF({-1, 1,  1}, form_factor::form_factor_t::CH2),
            AtomFF({ 1, -1,  1}, form_factor::form_factor_t::N), AtomFF({ 1, 1,  1}, form_factor::form_factor_t::O)
        };
        Molecule molecule({Body{atoms}});

        size_t count = 0;
        for (const auto& atom : molecule.iterate_atoms()) {
            CHECK(atom == atoms[count++]);
        }
        REQUIRE(count == atoms.size());
    }

    SECTION("real system") {
        Molecule molecule("tests/files/2epe.pdb");
        auto atoms = molecule.get_atoms();
        size_t count = 0;
        for (const auto& atom : molecule.iterate_atoms()) {
            CHECK(atom == atoms[count++]);
        }
        REQUIRE(count == atoms.size());
    }

    SECTION("multiple bodies") {
        auto molecule = rigidbody::BodySplitter::split("tests/files/LAR1-2.pdb", {1, 15, 35, 55, 75, 95});
        auto atoms = molecule.get_atoms();
        size_t count = 0;
        for (const auto& atom : molecule.iterate_atoms()) {
            CHECK(atom == atoms[count++]);
        }
        REQUIRE(count == atoms.size());
    }
}

TEST_CASE("Molecule::iterate_waters") {
    SECTION("empty") {
        Molecule molecule;
        size_t count = 0;
        for (const auto& water : molecule.iterate_waters()) {
            (void)water; // to silence unused variable warning
            ++count;
        }
        REQUIRE(count == 0);
    }

    SECTION("simple") {
        std::vector<Water> waters = {
            Water(Vector3<double>(-1, -1, -1)), Water(Vector3<double>(-1, 1, -1)),
            Water(Vector3<double>(1, -1, -1)), Water(Vector3<double>(1, 1, -1)),
            Water(Vector3<double>(-1, -1, 1)), Water(Vector3<double>(-1, 1, 1)),
            Water(Vector3<double>(1, -1, 1)), Water(Vector3<double>(1, 1, 1))
        };
        Molecule molecule(std::vector{Body{{}, waters}});

        size_t count = 0;
        for (const auto& water : molecule.iterate_waters()) {
            CHECK(water == waters[count++]);
        }
        REQUIRE(count == waters.size());
    }

    SECTION("real system") {
        Molecule molecule("tests/files/2epe.pdb");
        molecule.generate_new_hydration();
        auto waters = molecule.get_waters();
        size_t count = 0;
        for (const auto& water : molecule.iterate_waters()) {
            CHECK(water == waters[count++]);
        }
        REQUIRE(count == waters.size());
    }

    SECTION("multiple bodies") {
        auto molecule = rigidbody::BodySplitter::split("tests/files/LAR1-2.pdb", {1, 15, 35, 55, 75, 95});
        molecule.generate_new_hydration();
        auto waters = molecule.get_waters();

        // assign to individual bodies
        int stride = waters.size() / molecule.size_body();
        for (size_t i = 0; i < molecule.size_body()-1; i++) {
            molecule.get_body(i).set_hydration(hydrate::Hydration::create(std::vector<Water>(waters.begin() + i*stride, waters.begin() + (i+1)*stride)));
        }
        // assign remainder to the last body
        molecule.get_bodies().back().set_hydration(hydrate::Hydration::create(std::vector<Water>(waters.begin() + (molecule.size_body()-1)*stride, waters.end())));

        size_t count = 0;
        for (const auto& water : molecule.iterate_waters()) {
            CHECK(water == waters[count++]);
        }
        REQUIRE(count == waters.size());
    }
}