#ifndef GAIA_REST2_PARALLEL_HPP
#define GAIA_REST2_PARALLEL_HPP

#include <vector>
#include <random>
#include <cmath>
#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 Replica {
public:
    Replica(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 swap(Replica& 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 ParallelREST2 {
public:
    ParallelREST2(int n_replicas, double T_min, double T_max, int seed)
        : n_replicas(n_replicas), T_min(T_min), T_max(T_max) {
        
        replicas.reserve(n_replicas);
        for (int i = 0; i < n_replicas; i++) {
            double temp = T_min + (T_max - T_min) * i / (n_replicas - 1);
            replicas.emplace_back(temp, seed + i);
        }
    }
    
    void run(int n_steps, int exchange_interval = 10) {
        int exchange_count = 0;
        int accepted_count = 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) {
                std::random_device rd;
                std::mt19937 gen(rd());
                std::uniform_int_distribution<int> dist(0, n_replicas - 2);
                
                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());
                
                if (delta < 0 || std::exp(-delta) > uniform_random()) {
                    replicas[i].swap(replicas[j]);
                    accepted_count++;
                }
                exchange_count++;
            }
        }
        
        double acceptance = (double)accepted_count / exchange_count;
        std::cout << "REST2 Exchange acceptance: " << acceptance * 100 << "%" << std::endl;
    }
    
    void print_temperatures() {
        std::cout << "Replica temperatures:" << std::endl;
        for (int i = 0; i < n_replicas; i++) {
            std::cout << "  " << i << ": " << replicas[i].get_temperature() << " K" << std::endl;
        }
    }

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<Replica> replicas;
};

} // namespace rest2
} // namespace gaia

#endif