import numpy as np
import math
import mdtraj
import multiprocessing.pool
from openmm import unit

def run_STed_AFED_simulation(Nrep, T_r, sim, lambda_pp, lambda_pw, pdb, nsteps=150000000, exch_freq=50):
    """
    Handles the MD loop, data logging, and REST2 replica exchange logic.
    """
    kb = 0.008314 * unit.kilojoules_per_mole * (1 / unit.kelvin)
    beta_r = {i: 1/(kb * T_r[i]) for i in range(Nrep)}

    # Setup tracking variables
    dum_ex = {i: exch_freq + (i * exch_freq) for i in range(Nrep)}
    Replica_CV = {i: open(f'Replica_CV{i}.file', 'w') for i in range(Nrep)}
    CV_file = open('CV.file', 'w')
    count = 0
    cluster_coord = []
    step_range = int(nsteps / 10)

    print(f"Starting simulation with step_range: {step_range}")

    for step in range(1, step_range):
        # Multiprocessing for replica steps
        gpu_ids = list(range(Nrep))
        with multiprocessing.pool.ThreadPool(processes=Nrep) as pool:
            pool.map(lambda x: sim[x].step(10), gpu_ids)

        # Monitor first replica (reference state)
        current_lambdas = [sim[i].context.getParameter("lambda00") for i in range(Nrep)]
        
        if step % 1 == 0:
            for i in range(Nrep):
                if current_lambdas[i] == lambda_pp[0]:
                    cv = sim[i].context.driving_force.getCollectiveVariableValues(sim[i].context)
                    print(step, *cv, file=CV_file)

        # Trajectory saving
        if step % 1000 == 0:
            for i in range(Nrep):
                if current_lambdas[i] == lambda_pp[0]:
                    state = sim[i].context.getState(getPositions=True)
                    pos = state.getPositions(asNumpy=True).reshape((-1, 3))
                    cluster_coord.append(pos)
                    
                    # Periodic DCD save (example uses 1912 atoms; update if topology changes)
                    traj_array = np.vstack(cluster_coord).reshape(-1, pdb.topology.getNumAtoms(), 3)
                    mdtraj.Trajectory(traj_array, pdb.topology).save_dcd("simulation_cluster.dcd")

        # Replica Exchange (REST2)
        for i in range(Nrep - 1):
            if step % exch_freq == 0 and step == dum_ex[i] + ((Nrep - 1) * exch_freq):
                dum_ex[i] = step
                
                # Identify which simulation objects currently hold the target lambdas
                idx_i = current_lambdas.index(lambda_pp[i])
                idx_j = current_lambdas.index(lambda_pp[i+1])

                # Get Current Energies
                U_i_Ri = sim[idx_i].context.getState(getEnergy=True).getPotentialEnergy()
                U_j_Rj = sim[idx_j].context.getState(getEnergy=True).getPotentialEnergy()

                # Swap Parameters to calculate cross-energies
                sim[idx_i].context.setParameter("lambda00", lambda_pp[i+1])
                sim[idx_i].context.setParameter("lambda01", lambda_pw[i+1])
                sim[idx_j].context.setParameter("lambda00", lambda_pp[i])
                sim[idx_j].context.setParameter("lambda01", lambda_pw[i])

                U_j_Ri = sim[idx_i].context.getState(getEnergy=True).getPotentialEnergy()
                U_i_Rj = sim[idx_j].context.getState(getEnergy=True).getPotentialEnergy()

                # Metropolis Criterion
                delta_mn = beta_r[0] * ((U_i_Rj - U_i_Ri) + (U_j_Ri - U_j_Rj))
                
                if np.random.random() < math.exp(-max(0, delta_mn)):
                    count += 1 # Accept: Keep swapped parameters
                else:
                    # Reject: Revert parameters
                    sim[idx_i].context.setParameter("lambda00", lambda_pp[i])
                    sim[idx_i].context.setParameter("lambda01", lambda_pw[i])
                    sim[idx_j].context.setParameter("lambda00", lambda_pp[i+1])
                    sim[idx_j].context.setParameter("lambda01", lambda_pw[i+1])

        if step % 1000 == 0:
            print(f"Step {step} | Acceptance %: {(count / (1000/exch_freq)) * 100}")
            count = 0
            for r in range(Nrep):
                sim[r].saveCheckpoint(f'res_checkpoint{r}.chk')
