#!/usr/bin/env python
#
#   openmm2psf.py: use OpenMM to build a system from pdb + force field,
#                  then use that sytem to create a fake psf for analysis
#                  purposes
#   NOTE: The psf cannot be used to run a new simulation, because it doesn't
#         have angles, torsions, etc.
#
#   Dr. Alan Grossfield, University of Rochester Medical Center
#   2018
#

from __future__ import print_function

from simtk.openmm import app
import simtk.openmm as mm
from simtk import unit
import sys
import argparse

parser = argparse.ArgumentParser()
parser.add_argument('--pdb', help='PDB file used to build system')
parser.add_argument('--psf', help='Name of output fake PSF file')
parser.add_argument('--xml', nargs='+', help='One or more OpenMM XML force field files')
args = parser.parse_args()



# have OpenMM make a system
pdb = app.PDBFile(args.pdb)

# loop over the force field xml files
forcefield = app.ForceField(args.xml[0])
for ff in args.xml[1:]:
    forcefield.loadFile(ff)

# the other parameters here won't be reflected in the output
system = forcefield.createSystem(pdb.topology, nonbondedMethod=app.PME,
    nonbondedCutoff=1.0*unit.nanometers, constraints=app.HBonds, rigidWater=True,
    ewaldErrorTolerance=0.0005, )

# get the list of bonds
bond_forces = system.getForce(0)
num_bonds = bond_forces.getNumBonds()
bonds = []
for i in range(num_bonds):
    result = bond_forces.getBondParameters(i)
    bonds.append([result[0], result[1]])

# get the charges
non_bond_forces = system.getForce(3)
num_particles = non_bond_forces.getNumParticles()
charges = []
for i in range(num_particles):
    c = non_bond_forces.getParticleParameters(i)[0]
    charges.append(c.value_in_unit(unit.elementary_charge))

# get the masses
num_particles = system.getNumParticles()
masses = []
for i in range(num_particles):
    m = system.getParticleMass(i)
    masses.append(m.value_in_unit(unit.dalton))

names = []
residues = []
resnames = []
segments = []   # use the chain ID
# get residue names and numbers
for a in pdb.topology.atoms():
    names.append(a.name)
    resnames.append(a.residue.name)
    residues.append(a.residue.index)
    segments.append(a.residue.chain.id)

# write out the PSF
index = 1
with open(args.psf, "w") as psf:
    # header
    psf.write("PSF CMAP \n\n")
    psf.write("       4 !NTITLE\n")
    psf.write("REMARKS THIS IS NOT A REAL PSF, USE CAREFULLY\n")
    psf.write("REMARKS WRITTEN BY openmm2psf.py\n")
    psf.write("REMARKS OpenMM version: " + mm.__version__ + "\n")
    psf.write("REMARKS Generated from " + args.pdb + "\n\n")

    # atoms
    s = '{:8d}'.format(num_particles)
    psf.write(s + " !NATOM\n")
    for i in range(num_particles):
        s = '{:>8d} {:4s} {:<5d}{:<4s} {:<4s} {:<4s} {:10.6f}      {:>8.4f}           0\n'.format(i+1, segments[i], residues[i]+1, resnames[i], names[i], names[i], charges[i], masses[i])
        psf.write(s)

    psf.write("\n")

    # bonds
    s = '{:8d}'.format(num_bonds)
    psf.write(s + " !NBOND: bonds\n")
    bond_count = 0
    s = ""
    for i in range(num_bonds):
        b1 = bonds[i][0]
        b2 = bonds[i][1]
        if b1 > b2:
            tmp = b2
            b2 = b1
            b1 = tmp
        s += "{:>8d}{:>8d}".format(b1+1, b2+1)
        bond_count += 1
        if bond_count == 4:
            psf.write(s + "\n")
            s = ""
            bond_count = 0
    psf.write(s + "\n")

    psf.write("\n")
    psf.write("      0  !NTHETA \n\n")
    psf.write("      0  !NPHI \n\n")
    psf.write("      0  !NIMPHI \n\n")
    psf.write("      0  !NDON \n\n")
    psf.write("      0  !NACC \n\n")

    # Nonbond exclusion section
    psf.write("      0  !NNB \n\n")
    for i in range(0, num_particles, 8):
        psf.write("0       0       0       0       0       0       0       0\n")
    remainder = num_particles % 8
    s = ""
    for i in range(remainder):
        s += "0       "
    psf.write(s + "\n\n")
    psf.write("   1       0 !NGRP\n")
    psf.write("   0       0       0 \n\n")
    psf.write("      0  !NCRTERM \n\n")
    psf.close()