#include <catch2/catch_test_macros.hpp>

#include <rigidbody/selection/RandomBodySelect.h>
#include <rigidbody/selection/SequentialBodySelect.h>
#include <rigidbody/selection/RandomConstraintSelect.h>
#include <rigidbody/selection/SequentialConstraintSelect.h>
#include <rigidbody/selection/ManualSelect.h>
#include <rigidbody/Rigidbody.h>
#include <rigidbody/constraints/ConstraintManager.h>
#include <rigidbody/constraints/DistanceConstraintCM.h>
#include <data/Molecule.h>
#include <data/Body.h>
#include <settings/All.h>

using namespace ausaxs;
using namespace ausaxs::data;
using namespace ausaxs::rigidbody;
using namespace ausaxs::rigidbody::selection;

struct SelectionStrategiesFixture {
    SelectionStrategiesFixture() {
        settings::general::verbose = false;
        settings::molecule::implicit_hydrogens = false;
        settings::rigidbody::constraint_generation_strategy = settings::rigidbody::ConstraintGenerationStrategyChoice::None;
        
        // Create test bodies
        AtomFF a1({0, 0, 0}, form_factor::form_factor_t::C);
        AtomFF a2({5, 0, 0}, form_factor::form_factor_t::C);
        AtomFF a3({10, 0, 0}, form_factor::form_factor_t::C);
        AtomFF a4({15, 0, 0}, form_factor::form_factor_t::C);
        
        Body b1(std::vector{a1});
        Body b2(std::vector{a2});
        Body b3(std::vector{a3});
        Body b4(std::vector{a4});
        
        rb = std::make_unique<Rigidbody>(Molecule{std::vector<Body>{b1, b2, b3, b4}});
        
        // Add some constraints for constraint-based selection
        rb->constraints->add_constraint(
            std::make_unique<constraints::DistanceConstraintCM>(&rb->molecule, 0, 1)
        );
        rb->constraints->add_constraint(
            std::make_unique<constraints::DistanceConstraintCM>(&rb->molecule, 1, 2)
        );
        rb->constraints->add_constraint(
            std::make_unique<constraints::DistanceConstraintCM>(&rb->molecule, 2, 3)
        );
    }
    
    std::unique_ptr<Rigidbody> rb;
};

TEST_CASE_METHOD(SelectionStrategiesFixture, "SelectionStrategies::RandomBodySelect") {
    RandomBodySelect selector(rb.get());
    
    SECTION("next returns valid body indices") {
        for (int i = 0; i < 10; ++i) {
            auto [ibody, iconstraint] = selector.next();
            CHECK(ibody < rb->molecule.get_bodies().size());
        }
    }
}

TEST_CASE_METHOD(SelectionStrategiesFixture, "SelectionStrategies::SequentialBodySelect") {
    SequentialBodySelect selector(rb.get());
    
    SECTION("next cycles through bodies in order") {
        unsigned int num_bodies = rb->molecule.get_bodies().size();
        
        for (unsigned int i = 0; i < num_bodies * 2; ++i) {
            auto [ibody, iconstraint] = selector.next();
            CHECK(ibody == i % num_bodies);
        }
    }
}
