#ifndef GAIA_METADYNAMICS_FUNNEL_HPP
#define GAIA_METADYNAMICS_FUNNEL_HPP

#include <vector>
#include <cmath>
#include <random>
#include <iostream>

namespace gaia {
namespace metadynamics {

struct Vec3 {
    double x, y, z;
    Vec3(double x=0, double y=0, double z=0) : x(x), y(y), z(z) {}
};

struct Hill {
    double cv1, cv2;
    double height;
    double width;
};

class Metadynamics {
public:
    Metadynamics(double height = 0.5, double width = 0.3, int freq = 1000)
        : hill_height(height), hill_width(width), hill_freq(freq), step(0) {}
    
    double get_bias(const Vec3& position, const std::vector<double>& cv) {
        double bias = 0.0;
        for (const auto& hill : hills) {
            double d1 = cv[0] - hill.cv1;
            double d2 = cv[1] - hill.cv2;
            double r2 = d1*d1 + d2*d2;
            bias += hill.height * std::exp(-r2 / (2 * hill.width * hill.width));
        }
        return bias;
    }
    
    void add_hill_if_due(const Vec3& position, const std::vector<double>& cv) {
        step++;
        if (step % hill_freq == 0) {
            hills.push_back({cv[0], cv[1], hill_height, hill_width});
            std::cout << "Hill added at CV: (" << cv[0] << ", " << cv[1] << ")" << std::endl;
        }
    }
    
    std::vector<Hill> get_hills() const { return hills; }

private:
    double hill_height;
    double hill_width;
    int hill_freq;
    int step;
    std::vector<Hill> hills;
};

class FunnelMetadynamics : public Metadynamics {
public:
    FunnelMetadynamics(int n_funnels = 8, double height = 0.5, double width = 0.3, int freq = 1000)
        : Metadynamics(height, width, freq), n_funnels(n_funnels) {}
    
    void add_funnel_hill(const Vec3& position, const std::vector<double>& cv) {
        if (cv[0] < 15.0) {
            Metadynamics::add_hill_if_due(position, cv);
        }
    }
    
private:
    int n_funnels;
};

} // namespace metadynamics
} // namespace gaia

#endif