#include <catch2/catch_test_macros.hpp>

#include <rigidbody/controller/SimpleController.h>
#include <rigidbody/controller/MetropolisController.h>
#include <rigidbody/controller/ControllerFactory.h>
#include <rigidbody/Rigidbody.h>
#include <rigidbody/BodySplitter.h>
#include <rigidbody/detail/MoleculeTransformParametersAbsolute.h>
#include <data/Molecule.h>
#include <data/Body.h>
#include <settings/All.h>
#include <io/ExistingFile.h>

using namespace ausaxs;
using namespace ausaxs::data;
using namespace ausaxs::rigidbody;
using namespace ausaxs::rigidbody::controller;

struct ControllerFixture {
    ControllerFixture() {
        settings::general::verbose = false;
        settings::molecule::implicit_hydrogens = false;
        settings::grid::min_bins = 250;
        settings::rigidbody::constraint_generation_strategy = settings::rigidbody::ConstraintGenerationStrategyChoice::Linear;
        
        // Create a rigidbody for testing
        auto bodies = BodySplitter::split("tests/files/SASDJG5.pdb");
        rb = std::make_unique<Rigidbody>(std::move(bodies));
    }
    
    std::unique_ptr<Rigidbody> rb;
};

TEST_CASE_METHOD(ControllerFixture, "Controllers::SimpleController basic functionality") {
    SimpleController ctrl(rb.get());
    
    SECTION("Setup initializes controller") {
        REQUIRE_NOTHROW(ctrl.setup(io::ExistingFile("tests/files/SASDJG5.dat")));
        CHECK(ctrl.get_fitter() != nullptr);
        CHECK(ctrl.get_current_best_config() != nullptr);
    }
    
    SECTION("Prepare and finish step") {
        ctrl.setup(io::ExistingFile("tests/files/SASDJG5.dat"));
        
        // Run a few optimization steps
        for (int i = 0; i < 5; ++i) {
            bool accepted = ctrl.prepare_step();
            ctrl.finish_step();
            
            // Step acceptance is deterministic for SimpleController (better chi2 = accepted)
            // We just check it doesn't crash
            REQUIRE((accepted == true || accepted == false));
        }
    }
    
    SECTION("Best configuration is updated on improvement") {
        ctrl.setup(io::ExistingFile("tests/files/SASDJG5.dat"));
        
        double initial_chi2 = ctrl.get_current_best_config()->chi2;
        
        // Run several steps
        bool improvement_found = false;
        for (int i = 0; i < 100; ++i) {
            bool accepted = ctrl.prepare_step();
            ctrl.finish_step();
            
            if (accepted) {
                improvement_found = true;
                double new_chi2 = ctrl.get_current_best_config()->chi2;
                CHECK(new_chi2 <= initial_chi2);
                initial_chi2 = new_chi2;
                break;
            }
        }
        
        // With 100 steps, at least one should be accepted
        CHECK(improvement_found);
    }
}

TEST_CASE("Controllers::ControllerFactory") {
    settings::general::verbose = false;
    settings::molecule::implicit_hydrogens = false;
    settings::rigidbody::constraint_generation_strategy = settings::rigidbody::ConstraintGenerationStrategyChoice::None;
    
    AtomFF a1({0, 0, 0}, form_factor::form_factor_t::C);
    AtomFF a2({5, 0, 0}, form_factor::form_factor_t::C);
    Rigidbody rb(Molecule{std::vector<Body>{Body(std::vector{a1}), Body(std::vector{a2})}});
    
    SECTION("Create SimpleController") {
        auto ctrl = factory::create_controller(&rb, settings::rigidbody::ControllerChoice::Classic);
        REQUIRE(ctrl != nullptr);
        CHECK(dynamic_cast<SimpleController*>(ctrl.get()) != nullptr);
    }
}
