import math
import openmm
from openmm import app, unit
import ufedmm
import nonbondedslicing as nbs

def initialize_STed_AFED_system(pdb_file, forcefield_files, Nrep, Pro_atoms, T_min, T_max):
    """
    Sets up the REST2 systems, scales nonbonded forces, and prepares UFED variables.
    """
    pdb = app.PDBFile(pdb_file)
    forcefield = app.ForceField(*forcefield_files)

    # Temperature and Lambda Scaling
    temp_list = [T_min + (T_max - T_min) * (math.exp(float(i) / float(Nrep-1)) - 1.0) / (math.e - 1.0) for i in range(Nrep)]
    T_r = {i: temp * unit.kelvin for i, temp in enumerate(temp_list)}
    lambda_pp = {i: 300 * unit.kelvin / T_r[i] for i in range(Nrep)}
    lambda_pw = {i: math.sqrt(300 * unit.kelvin / T_r[i]) for i in range(Nrep)}

    sys_Dict = {}
    
    for i in range(Nrep):
        system = forcefield.createSystem(
            pdb.topology, 
            nonbondedMethod=app.PME, 
            nonbondedCutoff=1.0 * unit.nanometers, 
            constraints=None, 
            rigidWater=False,
            removeCMMotion=False
        )
        
        # Apply REST2 SlicedNonbondedForce
        forces = system.getForces()
        for c, f in enumerate(forces):
            if isinstance(f, openmm.NonbondedForce):
                sliced_force = nbs.SlicedNonbondedForce(f, 2)
                for j in range(f.getNumParticles()):
                    # Subset 0 for Solute, Subset 1 for Solvent
                    sliced_force.setParticleSubset(j, 0 if j < Pro_atoms else 1)
                
                sliced_force.addGlobalParameter("lambda00", lambda_pp[i])
                sliced_force.addGlobalParameter("lambda01", lambda_pw[i])
                sliced_force.addGlobalParameter("lambda11", 1.0)
                
                sliced_force.addScalingParameter("lambda00", 0, 0, True, True)
                sliced_force.addScalingParameter("lambda01", 0, 1, True, True)
                sliced_force.addScalingParameter("lambda11", 1, 1, True, True)
                
                system.addForce(sliced_force)
                system.removeForce(c)
        
        sys_Dict[i] = system

    return sys_Dict, T_r, lambda_pp, lambda_pw, pdb
