system = prmtop.createSystem(
    nonbondedMethod=PME,
    nonbondedCutoff=1*nanometer,
    constraints=HBonds
)

# --- ADD HEAVY ATOM RESTRAINTS ---
# Changed 'periodicdist' to 'periodicdistance'
restraint_force = CustomExternalForce('k*periodicdistance(x, y, z, x0, y0, z0)^2')
restraint_force.addGlobalParameter('k', 10.0*kilocalories_per_mole/angstrom**2)
restraint_force.addPerParticleParameter('x0')
restraint_force.addPerParticleParameter('y0')
restraint_force.addPerParticleParameter('z0')

positions = inpcrd.getPositions()
for atom in prmtop.topology.atoms():
    if atom.element.symbol != 'H': # Restrain everything except Hydrogen
        restraint_force.addParticle(atom.index, positions[atom.index])

system.addForce(restraint_force)

# Standard Integrator
integrator = LangevinIntegrator(300*kelvin, 1/picosecond, 0.002*picoseconds)

rial_steps = 25000 
report_freq = 500 # Save data every 1 ps (50 total entries in log)
a
# --- STAGE 1:(Charges=0, Strong Restraints) ---
# Resolves the 0.414 Å clashes without the "infinite" pull of charges [cite: 64, 73]
nb_force = [f for f in system.getForces() if isinstance(f, NonbondedForce)][0]
original_params = [nb_force.getParticleParameters(i) for i in range(nb_force.getNumParticles())]

for i in range(nb_force.getNumParticles()):
    charge, sigma, epsilon = nb_force.getParticleParameters(i)
    nb_force.setParticleParameters(i, 0.0, sigma, epsilon)
nb_force.updateParametersInContext(simulation.context)

print("Stage 1: Removing severe overlaps (No-Elec)...")
simulation.minimizeEnergy(maxIterations=1000)

# --- STAGE 2: Structural Relaxation (Restore Charges, Keep Restraints) ---
# Restores the +8 total charge and original parameters [cite: 63]
for i in range(nb_force.getNumParticles()):
    nb_force.setParticleParameters(i, *original_params[i])
nb_force.updateParametersInContext(simulation.context)

print("Stage 2: Relaxation with full physics...")
simulation.minimizeEnergy(maxIterations=1000)

# --- STAGE 3: Final Clean-up (No Restraints) ---
# Release the "k" parameter defined in Cell 2 to settle the system
simulation.context.setParameter('k', 0.0) 
print("Stage 3: Final unconstrained minimization...")
simulation.minimizeEnergy()

simulation.reporters.append(DCDReporter(traj_file, report_freq))
simulation.reporters.append(
    StateDataReporter(
        log_file,
        report_freq,
        step=True,
        time=True,
        potentialEnergy=True, 
        temperature=True,    
        density=True,
        speed=True,
        totalSteps=trial_steps,
        separator="\t"
    )
)


print(f"Heating to 300K and running 50ps Trial ({trial_steps} steps)...")
simulation.context.setVelocitiesToTemperature(300*kelvin)
simulation.step(trial_steps)

print("Simulation finished.")
print(f"Check {log_file} to verify energy convergence.")