include/docking/ensemble.hpp

#ifndef GAIA_DOCKING_ENSEMBLE_HPP
#define GAIA_DOCKING_ENSEMBLE_HPP

#include <vector>
#include <string>
#include <cmath>
#include <algorithm>
#include <iostream>

namespace gaia {
namespace docking {

struct Conformation {
    std::string name;
    double energy;
    double weight;
};

class EnsembleDocking {
public:
    EnsembleDocking(double temperature = 300.0) 
        : kT(0.001987204258 * temperature) {}
    
    void add_conformation(const std::string& name, double energy = 0.0) {
        Conformation conf;
        conf.name = name;
        conf.energy = energy;
        conf.weight = 1.0;
        conformations.push_back(conf);
    }
    
    void compute_weights() {
        double min_energy = conformations[0].energy;
        for (const auto& c : conformations) {
            if (c.energy < min_energy) min_energy = c.energy;
        }
        double Z = 0.0;
        for (auto& c : conformations) {
            c.weight = exp(-(c.energy - min_energy) / kT);
            Z += c.weight;
        }
        for (auto& c : conformations) {
            c.weight /= Z;
        }
    }
    
    double compute_ensemble_score(const std::vector<double>& ligand_scores) {
        double ensemble_score = 0.0;
        for (size_t i = 0; i < conformations.size() && i < ligand_scores.size(); i++) {
            ensemble_score += conformations[i].weight * ligand_scores[i];
        }
        return ensemble_score;
    }
    
    void print_conformations() {
        std::cout << "Ensemble conformations:\n";
        for (const auto& c : conformations) {
            std::cout << "  " << c.name << ": weight = " << c.weight * 100 << "%\n";
        }
    }

private:
    double kT;
    std::vector<Conformation> conformations;
};

} // namespace docking
} // namespace gaia

#endif