#!/usr/bin/env python 
# -*- coding: utf-8 -*-
########################################################
#
# this script converts a snapshot from a namd trajectory
# to a Maestro input file.
#
# (C) 2010 Markus Dittrich, NRBSC, PSC, CMU
#
# call with:
#
# vmd -dispdev text -python -e convertNAMDtoMaestro.py \
#      -args -p ubq_wb.psf \
#      -c ubq_wb_eq.restart.coor -v ubq_wb_eq.restart.vel \
#      -x ubq_wb_eq.xsc -o outfile -s -S "protein"
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
#
########################################################

import sys
import optparse
from atomsel import *
from AtomSel import AtomSel
from Molecule import Molecule
from VMD import evaltcl


def parse_cmdline(cmdlineArgs):
    """
    This function initializes the command line parser.
    """

    parser = optparse.OptionParser("Usage: vmdt -python -e "
            "readVelocities.py -args [options]")

    parser.add_option("-p", "--psffile", action="store", dest="psfFile")
    parser.add_option("-c", "--coorfile", action="store", dest="coorFile")
    parser.add_option("-v", "--velfile", action="store", dest="velFile")
    parser.add_option("-x", "--xscfile", action="store", dest="xscFile")
    parser.add_option("-o", "--outputfile", action="store", dest="outFile")
    parser.add_option("-s", "--centerSystem", action="store_true", dest="doCenter")
    parser.add_option("-S", "--centerSelection", action="store", dest="centerSel")

    parser.set_defaults(doCenter = False, centerSel = "all")

    opts, args = parser.parse_args(cmdlineArgs)
    psfFile   = opts.psfFile
    coorFile  = opts.coorFile
    velFile   = opts.velFile
    xscFile   = opts.xscFile
    outFile   = opts.outFile
    doCenter  = opts.doCenter
    centerSel = opts.centerSel

    # all filenames are required
    if (psfFile == None) or (coorFile == None) or (velFile == None) \
       or (xscFile == None) or (outFile == None):

                parser.print_help()
                exit()


    return psfFile, coorFile, velFile, xscFile, outFile, doCenter, \
           centerSel



def load_velocities(psfFile, velFile):
    """
    Load the binary velocity file and extract velocities.
    """

    mol = Molecule()
    mol.load(psfFile)
    mol.load(velFile, "namdbin")

    allVelocities = atomsel('all')
    xVel = allVelocities.get('x')
    yVel = allVelocities.get('y')
    zVel = allVelocities.get('z')

    # conversion from binvel units to A/ps
    convFactor = 20.4582651391 
    xVel = [v * convFactor for v in xVel]
    yVel = [v * convFactor for v in yVel]
    zVel = [v * convFactor for v in zVel]

    mol.delete()
    return xVel, yVel, zVel



def load_system(psfFile, coorFile):
    """
    Load the main system.
    """

    mol = Molecule()
    mol.load(psfFile)
    mol.load(coorFile, "namdbin")
    return mol



def set_velocities(mol, xVel, yVel, zVel):
    """
    Add the molecule velocities to the system.
    """

    allAtoms = atomsel("all")
    allAtoms.set("vx", xVel)
    allAtoms.set("vy", yVel)
    allAtoms.set("vz", zVel)



def save_mol_as_maestro(mol, fileName):
    """
    Save the current molecule as maestro file.
    """

    mol.save(fileName + ".mae")



def set_pbc(xscFile):
    """
    Sets the systems periodic boundaries."
    """
    xscFile = open(xscFile,"r")
        
    for line in xscFile:
       continue

    items = line.split()
    xDim  = items[1]
    yDim  = items[5]
    zDim  = items[9]

    #set pbd
    pbcCommand = ("package require pbctools; pbc set { %s %s %s }" 
                % (xDim, yDim, zDim))
    evaltcl(pbcCommand)

    xscFile.close()
        


def center_system(selection):
    """
    Center the system around the selection.
    """

    centerSel = atomsel(selection)
    center = centerSel.center()
    negCenter = [-1.0 * item for item in center]

    moveSel = atomsel("all")
    moveSel.moveby(negCenter)



def remove_tip3p_hh_bond():
    """
    This removes the bond between hydrogen atoms in
    TIP3P water if present since viparr will introduce
    the proper constraint.
    """

    # it looks like atomsel doesn't support set/getbonds
    # so we have to use the deprecated AtomSel for now
    oh2Sel = AtomSel("resname TIP3 and name OH2", 1)
    h1Sel  = AtomSel("resname TIP3 and name H1", 1)
  
    oh2Indices = oh2Sel.get("index")
    bondlist = []
    for i in oh2Indices:
        bondlist.append([i])
    h1Sel.setbonds(bondlist) 



#####################################################
# main routine
#####################################################
if __name__ == "__main__":

    # parse the command line
    psfFile, coorFile, velFile, xscFile, outfile, doCenter, \
        centerSel = parse_cmdline(sys.argv[1:])

    # transform NAMD to Maestro
    vx, vy, vz = load_velocities(psfFile, velFile)
    mol = load_system(psfFile, coorFile)

    if doCenter:
        center_system(centerSel)

    set_velocities(mol, vx, vy, vz)
    set_pbc(xscFile)
    remove_tip3p_hh_bond()
    save_mol_as_maestro(mol, outfile)  
    
    exit()