#!/usr/bin/env python
#############################################################################
### mol2box.py - 2-atom box file defined by max extents of input mol
###    27 Feb 2006
#############################################################################
import os,sys,getopt
from openeye.oechem import *


#############################################################################
if __name__=='__main__':
  PROG=os.path.basename(sys.argv[0])
  usage='''
    %(PROG)s [options]
        options:
           --i=<infile>
           --o=<outfile> ... [.xyz]
           --plusx=<px>
           --minusx=<mx>
           --plusy=<py>
           --minusy=<my>
           --plusz=<pz>
           --minusz=<mz>
           --addbox=<d> ... extend all sides by d
           --v  ... verbose
           --h  ... help
  
  '''%{'PROG':PROG}
  ifile=""; ofile=".xyz"; verbose=0; addbox=None;
  plusx=None; minusx=None; plusy=None; minusy=None; plusz=None; minusz=None;
  opts,pargs = getopt.getopt(sys.argv[1:],'h',['h','v',
    'plusx=','minusx=','plusy=','minusy=','plusz=','minusz=',
    'addbox=','i=','o='])
  
  if not opts: OEThrow.Usage(usage)
  for opt,val in opts:
    if opt=='--h': OEThrow.Usage(usage)
    elif opt=='--i': ifile=val
    elif opt=='--o': ofile=val
    elif opt=='--addbox': addbox=float(val)
    elif opt=='--plusx': plusx=float(val)
    elif opt=='--plusy': plusy=float(val)
    elif opt=='--plusz': plusz=float(val)
    elif opt=='--minusx': minusx=float(val)
    elif opt=='--minusy': minusy=float(val)
    elif opt=='--minusz': minusz=float(val)
    elif opt=='--v': verbose=1
    else: OEThrow.Fatal('Illegal option: %s\n%s' % (opt,usage))
  
  ims=oemolistream()
  if not ifile: OEThrow.Fatal("Input file required.\n"+usage)
  if not ims.open(ifile): OEThrow.Fatal("Cannot open: %s"%ifile)
  oms=oemolostream()
  if not oms.open(ofile): OEThrow.Fatal("Cannot open: %s"%ofile)
  
  mol=OEGraphMol()
  OEReadMolecule(ims,mol)
  if mol.NumAtoms()==0:
    OEThrow.Fatal("mol has no atoms.")
  xmax,ymax,zmax=mol.GetCoords(mol.GetAtom(OEHasAtomIdx(0)))
  xmin,ymin,zmin=xmax,ymax,zmax
  for atom in mol.GetAtoms():
    x,y,z=mol.GetCoords(atom)
    if x>xmax: xmax=x
    if y>ymax: ymax=y
    if z>zmax: zmax=z
    if x<xmin: xmin=x
    if y<ymin: ymin=y
    if z<zmin: zmin=z


  if plusx: xmax+=plusx
  elif addbox: xmax+=addbox
  if plusy: ymax+=plusy
  elif addbox: ymax+=addbox
  if plusz: zmax+=plusz
  elif addbox: zmax+=addbox
  if minusx: xmin-=minusx
  elif addbox: xmin-=addbox
  if minusy: ymin-=minusy
  elif addbox: ymin-=addbox
  if minusz: zmin-=minusz
  elif addbox: zmin-=addbox

  boxmol=OEGraphMol()
  newatom=boxmol.NewAtom(6)
  boxmol.SetCoords(newatom,(xmax,ymax,zmax))
  newatom=boxmol.NewAtom(6)
  boxmol.SetCoords(newatom,(xmin,ymin,zmin))
  boxmol.SetTitle("%s"%(mol.GetTitle()))
  if addbox: boxmol.SetTitle("%s (addbox=%.2f)"%(boxmol.GetTitle(),addbox))

  OEWriteMolecule(oms,boxmol)