#!/usr/bin/env python3

from openmm import *
from openmm.app import *
from openmm.unit import *
import MDAnalysis
import MDAnalysis.transformations as trans
import sys, argparse

parser = argparse.ArgumentParser(description='Wrap waters in a trajectory to center the protein')
parser.add_argument("topology",help="Topology PDB file from OpenMM")
parser.add_argument("trajectory",help="Trajectory file")
parser.add_argument("output",default="wrapped",help="Output prefix")
parser.add_argument("--step",default=1,type=int,help="Step")

args = parser.parse_args()

# The PDB file must be read by OpenMM, due to differences in how
# implementations deal with overflowing fixed-width fields
U = MDAnalysis.Universe(PDBFile(args.topology),args.trajectory)
protein = U.select_atoms('protein')
system = U.select_atoms('all') 

transforms = [trans.unwrap(protein),
              trans.center_in_box(protein),
              trans.wrap(U.select_atoms('not protein'),compound='residues')]
              
U.trajectory.add_transformations(*transforms)

with MDAnalysis.Writer(args.output+'.dcd', system.n_atoms) as W:
    system.write(args.output+'.pdb')  # write first frame as wrapped PDB topology
    for ts in U.trajectory[0::args.step]: # remaining frames go in trajectory
        W.write(system)