import sys
import numpy

import openbabel


openbabel.obErrorLog.SetOutputLevel(-1)

def OBMolFromFilename(filename):
  obmol = openbabel.OBMol()
  obconv = openbabel.OBConversion()
  obconv.SetInFormat("pdb")
  obconv.ReadFile(obmol, filename)
  return obmol

def OBMolToFilename(filename, mol):
  obconv = openbabel.OBConversion()
  obconv.SetOutFormat("pdb")
  obconv.WriteFile(mol, filename)

def get_center_of_mass(mol):
  x = 0.0
  y = 0.0
  z = 0.0
  m = 0.0

  for residue in openbabel.OBResidueIter(mol):
    for atom in openbabel.OBResidueAtomIter(residue):
      ms = atom.GetAtomicMass()
      x += atom.GetX()*ms
      y += atom.GetY()*ms
      z += atom.GetZ()*ms
      m += ms

  x /= m
  y /= m
  z /= m
  return numpy.array([x,y,z])

def get_distances_to_cm(cm, mol):
  """This is per molecule (residue), not per atom
     and it is the minimum distance
  """
  Rs = []
  Is = []
  for residue in openbabel.OBResidueIter(mol):
    r = get_residue_distance_to_cm(cm, residue)
    Rs.append(r)
    Is.append(residue.GetNum())

  Rs = numpy.array(Rs)
  Is = numpy.array(Is)

  ind = numpy.argsort(Rs)
  Is  = Is[ind]
  Rs  = Rs[ind]

  return (Rs, Is)

def get_residue_distance_to_cm(cm, residue):
  r = 30.0e8
  for atom in openbabel.OBResidueAtomIter(residue):
    d = get_atom_distance_to_cm(cm, atom)
    if d < r: r = d
  return r

def get_min_distances_from(idx, mol):
  Rs = [0.0]
  Is = [idx]
  res = None
  for ri in openbabel.OBResidueIter(mol):
    rii = ri.GetNum()
    if idx == rii: res = ri

  #print idx, res

  #for residue in openbabel.OBResidueIter(mol):
  #  print residue.GetIdx(), idx
  #  if residue.GetIdx()+1 == idx: res = residue

  #print Rs, Is, res
  if res is None: exit("Error: something terrible went wrong. Aborting.")
  for residue in openbabel.OBResidueIter(mol):
    if residue.GetNum() == idx: continue
    r = get_residue_distance_to(res, residue)
    Rs.append(r)
    Is.append(residue.GetNum())

  Rs = numpy.array(Rs)
  Is = numpy.array(Is)
  ind = numpy.argsort(Rs)
  Is  = Is[ind]
  Rs  = Rs[ind]

  return (Rs, Is)

def get_residue_distance_to(r, residue):
  d = 30e8
  for atom_i in openbabel.OBResidueAtomIter(r):
    ic = numpy.array(get_atom_coords(atom_i))
    for atom_j in openbabel.OBResidueAtomIter(residue):
      r = get_atom_distance_to_cm(ic, atom_j)
      d = min(r,d)
  return d

def get_atom_distance_to_cm(cm, atom):
  x,y,z = get_atom_coords(atom)
  R = numpy.array([x,y,z])
  v = R - cm
  return numpy.sqrt(numpy.dot(v,v))

def get_atom_coords(atom):
  return atom.GetX(), atom.GetY(), atom.GetZ()

def get_residue_cm(idx, mol):
  for residue in openbabel.OBResidueIter(mol):
    if residue.GetNum() == idx:
      break
  xm,ym,zm,mm = 0.0, 0.0, 0.0, 0.0
  for atom in openbabel.OBResidueAtomIter(residue):
    m = atom.GetAtomicMass()
    x,y,z = get_atom_coords(atom)
    xm += x*m
    ym += y*m
    zm += z*m
    mm += m

  xm /= mm
  ym /= mm
  zm /= mm
  return numpy.array([xm,ym,zm])

args = sys.argv[1:]

if len(args) == 0:
  print "usage: cm.py [options] filename"
  print
  print "     : -d FLOAT print Res# for all atoms within d"
  print
  print "     : -c cuts out all residues with a distance larger than d. Requires d."
  print
  print "     : -m uses minimum distance instead of center of mass"
  print
  print "     : -o specify output filename"
  sys.exit()

dist = -1.0
cut = False
min_distance = False
verbose = False
ofile = None
while len(args) > 1:
  if args[0] == "-d":
    dist = float(args[1])
    args = args[2:]

  if args[0] == "-c":
    cut = True
    args = args[1:]

  if args[0] == "-m":
    min_distance = True
    args = args[1:]

  if args[0] == "-v":
    verbose = True
    args = args[1:]

  if args[0] == "-o":
    ofile = args[1]
    args = args[2:]

if dist < 0.0 and (cut or min_distance):
  print "Error: You forgot to specify distance using -d"
  sys.exit()

if verbose:
  print "Options:"
  print " dist: %6.2f" % dist
  if cut: print "will make cutout"
  if min_distance: print "using mininum interfragment distances"

filename = args[0]
mol = OBMolFromFilename(filename)

cm = get_center_of_mass(mol)
(R,I) = get_distances_to_cm(cm, mol)

if dist > 0.0:
  V = get_residue_cm(I[0], mol)
  if not min_distance:
    (Ri,Ii) = get_distances_to_cm(V, mol)
  else:
    (Ri,Ii) = get_min_distances_from(I[0], mol)

# print the indices of the
if dist > 0.0 and not cut:
  idx = numpy.where(Ri<=dist)
  for val in Ii[idx]:
    print "%i" % (val),

# if the user requested to cut the cluster out
# of some larger bulk
if dist > 0.0 and cut:
  idx = numpy.where(Ri>dist)
  atoms_to_delete = []
  for residue in openbabel.OBResidueIter(mol):
      if residue.GetNum() in Ii[idx]:
        for atom in openbabel.OBResidueAtomIter(residue):
          atoms_to_delete.append(atom)
  if len(atoms_to_delete) == 0:
    print "Error: No atoms to remove"
    sys.exit()

  for atom in atoms_to_delete:
    mol.DeleteAtom(atom)

  # attempt to fix that damn naming scheme
  ff = filename[:-4].split(".")
  idx = int(ff[1])
  fname = "{0:s}_{1:04d}.pdb".format(ff[0], idx)
  if ofile is not None: fname = ofile
  OBMolToFilename(fname, mol)