#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""OpenMM 隐式溶剂 MD：最小化 → 平衡 → 生产
用法：python md_run.py <pdb> <outdir> <ns>
"""
import os
import sys
import time

from openmm import LangevinMiddleIntegrator, Platform, unit
from openmm.app import (DCDReporter, HBonds, PDBFile, Simulation, StateDataReporter,
                        CutoffNonPeriodic)
import openmm.app as app
import pdbfixer

pdb_in, outdir, ns = sys.argv[1], sys.argv[2], float(sys.argv[3])
os.makedirs(outdir, exist_ok=True)

print('[1/5] 修复结构', flush=True)
fixer = pdbfixer.PDBFixer(pdb_in)
fixer.findMissingResidues()
fixer.findNonstandardResidues()
fixer.replaceNonstandardResidues()
fixer.removeHeterogens(keepWater=False)
fixer.findMissingAtoms()
fixer.addMissingAtoms()
fixer.addMissingHydrogens(7.0)

print('[2/5] 力场 + 系统', flush=True)
ff = app.ForceField('amber14-all.xml', 'implicit/gbn2.xml')
modeller = app.Modeller(fixer.topology, fixer.positions)
system = ff.createSystem(modeller.topology, nonbondedMethod=CutoffNonPeriodic,
                         nonbondedCutoff=2.0 * unit.nanometer, constraints=HBonds,
                         soluteDielectric=1.0, solventDielectric=78.5)

print('[3/5] 选择平台', flush=True)
platform = None
for name in ('OpenCL', 'CPU'):
    try:
        platform = Platform.getPlatformByName(name)
        print('  使用', name, flush=True)
        break
    except Exception:
        pass
integrator = LangevinMiddleIntegrator(300 * unit.kelvin, 1 / unit.picosecond, 0.002 * unit.picoseconds)
sim = Simulation(modeller.topology, system, integrator, platform)

print('[4/5] 最小化 + 平衡', flush=True)
sim.context.setPositions(modeller.positions)
sim.minimizeEnergy(maxIterations=500)
sim.context.setVelocitiesToTemperature(300 * unit.kelvin)
sim.step(5000)  # 10 ps 平衡

print('[5/5] 生产 %.2f ns' % ns, flush=True)
sim.reporters.append(DCDReporter(os.path.join(outdir, 'traj.dcd'), 2500))  # 每 5 ps 一帧
sim.reporters.append(StateDataReporter(os.path.join(outdir, 'log.csv'), 2500,
                                       step=True, time=True, potentialEnergy=True,
                                       temperature=True, speed=True))
steps = int(ns * 500000)  # 2 fs/step
t0 = time.time()
sim.step(steps)
el = time.time() - t0
print('完成：%.2f ns，用时 %.1f s，速度 %.2f ns/day' % (ns, el, ns / el * 86400), flush=True)

with open(os.path.join(outdir, 'final.pdb'), 'w') as f:
    app.PDBFile.writeFile(sim.topology, sim.context.getState(getPositions=True).getPositions(), f)
print('输出:', outdir)