#ifndef GAIA_SAMPLING_FLEX_RESIDUE_HPP
#define GAIA_SAMPLING_FLEX_RESIDUE_HPP

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

namespace gaia {
namespace sampling {

struct Rotamer {
    std::vector<double> chi_angles;
    double probability;
};

struct Residue {
    int id;
    std::string type;
    std::vector<double> current_chi;
    int rotamer_index;
};

class FlexResidueSampler {
public:
    FlexResidueSampler() { build_rotamer_library(); }
    
    void add_flexible_residue(int id, const std::string& type) {
        Residue res;
        res.id = id;
        res.type = type;
        res.current_chi = std::vector<double>(4, 0.0);
        res.rotamer_index = 0;
        flexible_residues.push_back(res);
    }
    
    void sample_rotamer(Residue& res) {
        if (rotamer_library.find(res.type) == rotamer_library.end()) return;
        auto& rotamers = rotamer_library[res.type];
        std::random_device rd;
        std::mt19937 gen(rd());
        std::uniform_int_distribution<int> dist(0, (int)rotamers.size() - 1);
        int idx = dist(gen);
        res.rotamer_index = idx;
        res.current_chi = rotamers[idx].chi_angles;
    }
    
    void sample_all_rotamers() {
        for (auto& res : flexible_residues) {
            sample_rotamer(res);
        }
    }
    
    void print_residues() {
        std::cout << "Flexible residues:\n";
        for (const auto& res : flexible_residues) {
            std::cout << "  " << res.type << " " << res.id << ": χ = ";
            for (double chi : res.current_chi) {
                std::cout << chi << " ";
            }
            std::cout << "\n";
        }
    }

private:
    void build_rotamer_library() {
        rotamer_library["SER"] = {{{-60, 0, 0, 0}, 0.5}, {{60, 0, 0, 0}, 0.3}, {{180, 0, 0, 0}, 0.2}};
        rotamer_library["THR"] = {{{-60, 0, 0, 0}, 0.4}, {{60, 0, 0, 0}, 0.4}, {{180, 0, 0, 0}, 0.2}};
        rotamer_library["TYR"] = {{{-60, 90, 0, 0}, 0.3}, {{60, 90, 0, 0}, 0.3}, {{-60, -90, 0, 0}, 0.2}, {{60, -90, 0, 0}, 0.2}};
        rotamer_library["LYS"] = {{{-60, 60, 60, 0}, 0.3}, {{60, -60, -60, 0}, 0.3}, {{180, 180, 180, 0}, 0.2}};
    }
    
    std::map<std::string, std::vector<Rotamer>> rotamer_library;
    std::vector<Residue> flexible_residues;
};

} // namespace sampling
} // namespace gaia

#endif