import sys, os
import copy as cp
import numpy as np
from pmx import *
from pmx import geometry
from pmx import ndx
from pmx import library
from pmx.options import *
#from pmx.forcefield2 import *
from AbsRestraints import *

def main(argv):

    desc=('Read ligand and protein structure and topologies.',
          'Add distance, angle and dihedral restraints.',
          '',)

# define input/output files

    files= [
        FileOption("-f", "r",["pdb","gro"],"system.pdb", "input system structure"),
        FileOption("-ilig", "r/o",["itp","top"],"", "input ligand topology"),
        FileOption("-iprot", "r/o",["itp","top"],"", "input protein topology"),
#        FileOption("-fprot", "r",["pdb","gro"],"protein.pdb", "input protein structure"),
        FileOption("-n", "r",["ndx"],"ndx_restr.ndx", "input index file"),
#        FileOption("-oitp", "w",["itp"],"out.itp", "output renumbered ligand.itp"),
        FileOption("-oii", "w",["itp"],"outii.itp", "intermolecular_interactions itp to be added to the end of .top"),
        FileOption("-dg", "w",["dat"],"out_dg.dat", "write contribution of the restraints to the free energy"),
        FileOption("-orestr", "w/o",["pdb"],"restrained_atoms.pdb", "optional output of restrained atoms only"),
        ]

# define options

    options=[
       Option( "-manLig", "bool", False, "manual selection of ligand atoms for restraints"),
       Option( "-manProt", "bool", False, "manual selection of protein atoms for restraints"),
       Option( "-ligCenter", "string", False, "manually provide only central ligand atom"),
       Option( "-ligMaxSpread", "bool", False, "for automated ligand atom detection spread atoms to maximize distance between them"),
       Option( "-T", "float", 300.0, "temperature"),
       Option( "-kBond", "float", 4184.0, "force constant for bond, kJ/mol"),
       Option( "-kAngle", "float", 41.84, "force constant for angles, kJ/mol"),
       Option( "-kDih", "float", 41.84, "force constant for dihedrals, kJ/mol"),
       Option( "-martini", "bool", False, "martini force field"),
        ]

    help_text = ('Read ligand+protein structure and topologies.',
          'Add distance, angle and dihedral restraints.',
          '',)

# pass options, files and the command line to pymacs

    cmdl = Commandline( argv, options = options, fileoptions = files, program_desc = help_text, check_for_existing_files = False, version = "0.0" )

#    system = Model(cmdl['-f'])
#    protModel = Model(cmdl['-fprot'])


    absRestr = AbsRestraints( strFile=cmdl['-f'], topLigFile=cmdl['-ilig'], topProtFile=cmdl['-iprot'], ndxFile=cmdl['-n'], outiiFile=cmdl['-oii'], bManLig=cmdl['-manLig'], bManProt=cmdl['-manProt'], T=cmdl['-T'], dgFile=cmdl['-dg'], ligCenterByName=cmdl['-ligCenter'],bLigMaxSpread=cmdl['-ligMaxSpread'], kBond=cmdl['-kBond'], kAngle=cmdl['-kAngle'], kDih=cmdl['-kDih'], bMartini=cmdl['-martini'], pdbRestrAtoms=cmdl['-orestr'] )    


main( sys.argv )

