#ifndef GAIA_REST2_ADAPTIVE_REST2_FINAL_HPP
#define GAIA_REST2_ADAPTIVE_REST2_FINAL_HPP

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

#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 AdaptiveReplicaFinal {
public:
    AdaptiveReplicaFinal(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(AdaptiveReplicaFinal& 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 AdaptiveREST2Final {
public:
    AdaptiveREST2Final(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), seed(seed) {
        
        // Use geometric spacing
        update_temperatures_geometric();
        
        replicas.reserve(n_replicas);
        for (int i = 0; i < n_replicas; i++) {
            replicas.emplace_back(current_temps[i], seed + i);
        }
        
        exchange_history.resize(n_replicas - 1, 0.0);
        adjustment_count = 0;
    }
    
    void update_temperatures_geometric() {
        current_temps.resize(n_replicas);
        for (int i = 0; i < n_replicas; i++) {
            current_temps[i] = T_min * pow(T_max / T_min, (double)i / (n_replicas - 1));
        }
    }
    
    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) {
                for (int i = 0; i < n_replicas - 1; i++) {
                    double beta_i = 1.0 / replicas[i].get_kT();
                    double beta_j = 1.0 / replicas[i+1].get_kT();
                    
                    double delta = (beta_i - beta_j) *
                                   (replicas[i+1].get_energy() - replicas[i].get_energy());
                    
                    attempts_per_pair[i]++;
                    if (delta < 0 || std::exp(-delta) > uniform_random()) {
                        replicas[i].swap(replicas[i+1]);
                        accepted_count++;
                        accepted_per_pair[i]++;
                    }
                    exchange_count++;
                }
            }
        }
        
        double acceptance = (double)accepted_count / exchange_count;
        std::cout << "REST2 Exchange acceptance: " << acceptance * 100 << "%\n";
        
        std::cout << "Per-pair acceptance:\n";
        double avg_accept = 0.0;
        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;
            avg_accept += pair_accept;
        }
        avg_accept /= (n_replicas - 1);
        
        // Adaptive optimization
        if (acceptance < 0.20 || avg_accept < 0.20) {
            std::cout << "⚠️  Acceptance <20% - Optimizing per-pair temperatures...\n";
            optimize_temperatures_per_pair(exchange_history);
            apply_temperatures();
        }
    }
    
    void optimize_temperatures_per_pair(const std::vector<double>& pair_acceptance) {
        std::vector<double> new_temps(n_replicas);
        new_temps[0] = T_min;
        new_temps[n_replicas-1] = T_max;
        
        for (int i = 1; i < n_replicas - 1; i++) {
            double accept_left = pair_acceptance[i-1];
            double accept_right = pair_acceptance[i];
            
            double target = 0.25;
            double factor = 1.0;
            
            if (accept_left < 0.10 && accept_right < 0.10) {
                factor = 0.6;
            } else if (accept_left > 0.40 && accept_right > 0.40) {
                factor = 1.4;
            } else if (accept_left < 0.10) {
                factor = 0.7;
            } else if (accept_right < 0.10) {
                factor = 0.7;
            } else if (accept_left < 0.20) {
                factor = 0.85;
            } else if (accept_right < 0.20) {
                factor = 0.85;
            }
            
            // Per-pair adjustment
            double left_temp = new_temps[i-1];
            double right_temp = T_min + (T_max - T_min) * (double)(i+1) / (n_replicas - 1);
            
            // Only adjust if we're not at the boundary
            if (i > 1 && i < n_replicas - 2) {
                double base = left_temp + (right_temp - left_temp) * 0.5;
                new_temps[i] = base * factor + (1.0 - factor) * current_temps[i];
            } else {
                new_temps[i] = left_temp + (right_temp - left_temp) * factor * 0.5;
            }
            
            // Ensure bounds
            double min_gap = (T_max - T_min) / (n_replicas * 2);
            new_temps[i] = std::max(new_temps[i], left_temp + min_gap);
            new_temps[i] = std::min(new_temps[i], right_temp - min_gap);
        }
        
        current_temps = new_temps;
        adjustment_count++;
        
        std::cout << "  Optimization #" << adjustment_count << " complete\n";
    }
    
    void apply_temperatures() {
        for (int i = 0; i < n_replicas; i++) {
            replicas[i].set_temperature(current_temps[i]);
        }
        std::cout << "  New temperatures applied:\n";
        for (int i = 0; i < n_replicas; i++) {
            std::cout << "    " << i << ": " << current_temps[i] << " 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;
    int seed;
    int adjustment_count;
    std::vector<AdaptiveReplicaFinal> replicas;
    std::vector<double> exchange_history;
    std::vector<double> current_temps;
};

} // namespace rest2
} // namespace gaia

#endif