#ifndef GAIA_REST2_ADAPTIVE_REST2_HPP
#define GAIA_REST2_ADAPTIVE_REST2_HPP

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

#ifdef _OPENMP
#include <omp.h>
#endif

namespace gaia {
namespace rest2 {

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

class AdaptiveReplica {
public:
    AdaptiveReplica(double temp, int seed) : temp(temp), seed(seed) {
        kT = 0.001987204258 * temp;
        energy = 0.0;
        positions.resize(100);
        std::mt19937 gen(seed);
        std::uniform_real_distribution<double> dist(-1.0, 1.0);
        for (auto& p : positions) {
            p = Vec3(dist(gen), dist(gen), dist(gen));
        }
    }
    
    void step(int n_steps) {
        std::mt19937 gen(seed + 1);
        std::uniform_real_distribution<double> dist(-0.01, 0.01);
        for (int s = 0; s < n_steps; s++) {
            for (auto& p : positions) {
                p.x += dist(gen);
                p.y += dist(gen);
                p.z += dist(gen);
            }
        }
        energy = 0.0;
        for (const auto& p : positions) {
            energy += p.x*p.x + p.y*p.y + p.z*p.z;
        }
        energy *= 0.5 * kT;
    }
    
    double get_energy() const { return energy; }
    double get_kT() const { return kT; }
    double get_temperature() const { return temp; }
    void set_temperature(double t) { temp = t; kT = 0.001987204258 * t; }
    
    void swap(AdaptiveReplica& other) {
        std::swap(positions, other.positions);
        std::swap(energy, other.energy);
        std::swap(temp, other.temp);
        std::swap(kT, other.kT);
    }

private:
    double temp;
    double kT;
    int seed;
    double energy;
    std::vector<Vec3> positions;
};

class AdaptiveREST2 {
public:
    AdaptiveREST2(int n_replicas = 8, double T_min = 300, double T_max = 500, int seed = 42)
        : n_replicas(n_replicas), T_min(T_min), T_max(T_max) {
        
        // Adaptive temperature spacing: geometric (better overlap)
        std::vector<double> temps;
        for (int i = 0; i < n_replicas; i++) {
            double t = T_min * pow(T_max / T_min, (double)i / (n_replicas - 1));
            temps.push_back(t);
        }
        
        replicas.reserve(n_replicas);
        for (int i = 0; i < n_replicas; i++) {
            replicas.emplace_back(temps[i], seed + i);
        }
        
        // Track acceptance for adaptive adjustment
        exchange_history.resize(n_replicas - 1, 0.0);
    }
    
    void run(int n_steps, int exchange_interval = 10) {
        int exchange_count = 0;
        int accepted_count = 0;
        std::vector<int> accepted_per_pair(n_replicas - 1, 0);
        std::vector<int> attempts_per_pair(n_replicas - 1, 0);
        
        for (int step = 0; step < n_steps; step++) {
            #ifdef _OPENMP
            #pragma omp parallel for
            #endif
            for (int i = 0; i < n_replicas; i++) {
                replicas[i].step(1);
            }
            
            if (step % exchange_interval == 0 && step > 0) {
                // Try exchanges at multiple pairs
                std::random_device rd;
                std::mt19937 gen(rd());
                std::uniform_int_distribution<int> dist(0, n_replicas - 2);
                
                // Try n_replicas/2 random exchanges
                for (int e = 0; e < n_replicas / 2; e++) {
                    int i = dist(gen);
                    int j = i + 1;
                    
                    double beta_i = 1.0 / replicas[i].get_kT();
                    double beta_j = 1.0 / replicas[j].get_kT();
                    
                    double delta = (beta_i - beta_j) *
                                   (replicas[j].get_energy() - replicas[i].get_energy());
                    
                    attempts_per_pair[i]++;
                    if (delta < 0 || std::exp(-delta) > uniform_random()) {
                        replicas[i].swap(replicas[j]);
                        accepted_count++;
                        accepted_per_pair[i]++;
                    }
                    exchange_count++;
                }
            }
        }
        
        double acceptance = (double)accepted_count / exchange_count;
        std::cout << "REST2 Exchange acceptance: " << acceptance * 100 << "%\n";
        
        // Per-pair acceptance
        std::cout << "Per-pair acceptance:\n";
        for (int i = 0; i < n_replicas - 1; i++) {
            double pair_accept = (double)accepted_per_pair[i] / (attempts_per_pair[i] + 1);
            std::cout << "  Pair " << i << "-" << i+1 << ": " << pair_accept * 100 << "%\n";
            exchange_history[i] = pair_accept;
        }
        
        // Adaptive temperature adjustment
        if (acceptance < 0.10) {
            std::cout << "⚠️  Acceptance <10% - Adjusting temperatures...\n";
            adjust_temperatures(acceptance);
        }
    }
    
    void adjust_temperatures(double acceptance) {
        // If acceptance is too low, reduce temperature range
        if (acceptance < 0.10) {
            double T_max_new = T_min + (T_max - T_min) * 0.6;
            for (int i = 0; i < n_replicas; i++) {
                double t = T_min * pow(T_max_new / T_min, (double)i / (n_replicas - 1));
                replicas[i].set_temperature(t);
            }
            std::cout << "  Reduced T_max to " << T_max_new << " K\n";
        }
        // If acceptance is too high, increase temperature range
        else if (acceptance > 0.50) {
            double T_max_new = T_min + (T_max - T_min) * 1.4;
            for (int i = 0; i < n_replicas; i++) {
                double t = T_min * pow(T_max_new / T_min, (double)i / (n_replicas - 1));
                replicas[i].set_temperature(t);
            }
            std::cout << "  Increased T_max to " << T_max_new << " K\n";
        }
    }
    
    void print_temperatures() {
        std::cout << "Replica temperatures:\n";
        for (int i = 0; i < n_replicas; i++) {
            std::cout << "  " << i << ": " << replicas[i].get_temperature() << " K\n";
        }
    }

private:
    double uniform_random() {
        static std::random_device rd;
        static std::mt19937 gen(rd());
        static std::uniform_real_distribution<double> dist(0.0, 1.0);
        return dist(gen);
    }
    
    int n_replicas;
    double T_min, T_max;
    std::vector<AdaptiveReplica> replicas;
    std::vector<double> exchange_history;
};

} // namespace rest2
} // namespace gaia

#endif