#include <pybind11/pybind11.h>
#include "simulation.h"
#include "model.h"
#include <string>
#include <memory>
#include <variant>

namespace py = pybind11;


// using variant because my model comes in 1D, 2D, and 3D versions.
using ModelVariant = std::variant<
    std::unique_ptr<model<1>>, 
    std::unique_ptr<model<2>>, 
    std::unique_ptr<model<3>>
>;

// same for simulation
using SimuVariant = std::variant<
    std::unique_ptr<simulation<1>>, 
    std::unique_ptr<simulation<2>>, 
    std::unique_ptr<simulation<3>>
>;


// Wrapper class exposed to Python.
class PythonWrapper {

    public:
        
        PythonWrapper(int dim, /* other args*/) 
        {
            if (dim == 1){
                model_variant = std::make_unique<model<1>>(/*args*/);
                simulation_variant=  std::make_unique<simulation<1>>( *std::get<std::unique_ptr<model<1>>>(model_variant), /* other args */);
            }
            else if (dim == 2) {
                model_variant = std::make_unique<model<2>>(/*args*/);
                simulation_variant =  std::make_unique<simulation<2>>( *std::get<std::unique_ptr<model<2>>>(model_variant), /* other args */);
            }
            else if (dim == 3) {
                model_variant = std::make_unique<model<3>>(/*args*/);
                simulation_variant =  std::make_unique<simulation<3>>( *std::get<std::unique_ptr<model<3>>>(model_variant), /* other args */);
            }
            else throw std::runtime_error("Unsupported dimension. Only 1D, 2D, and 3D are allowed.");         
        }

        void run() {
            std::visit([](auto& simulation) { simulation->run(); }, simulation_variant);    // This will be much slower when run from Python for some reason.
        }

    private:

        ModelVariant model_variant;   
        SimuVariant simulation_variant;

};


// make wrapper class known to Python

PYBIND11_MODULE(my_module, m) {
    
    // Bind wrapper class.
    py::class_<PythonWrapper>(m, "Simulation")
        .def(py::init</*args*/>())
        .def("run", &PythonWrapper::run)
        
}