## Automatically adapted for numpy.oldnumeric Jun 27, 2008 by -c

# $Id$
#
# Copyright (C) 2007 by Greg Landrum 
#  All rights reserved
#
from __future__ import print_function
from rdkit import Geometry
from rdkit.Chem.Subshape import SubshapeObjects
import math
import numpy

#-----------------------------------------------------------------------------
def ComputeGridIndices(shapeGrid,winRad):
  if getattr(shapeGrid,'_indicesInSphere',None):
    return shapeGrid._indicesInSphere
  gridSpacing = shapeGrid.GetSpacing()
  dX = shapeGrid.GetNumX()
  dY = shapeGrid.GetNumY()
  dZ = shapeGrid.GetNumZ()
  radInGrid = int(winRad/gridSpacing)
  indicesInSphere=[]
  for i in range(-radInGrid,radInGrid+1):
    for j in range(-radInGrid,radInGrid+1):
      for k in range(-radInGrid,radInGrid+1):
        d=int(math.sqrt(i*i + j*j + k*k ))
        if d<=radInGrid:
          idx = (i*dY+j)*dX+k
          indicesInSphere.append(idx)
  shapeGrid._indicesInSphere = indicesInSphere
  return indicesInSphere

#-----------------------------------------------------------------------------
def ComputeShapeGridCentroid(pt,shapeGrid,winRad):
  count=0
  centroid = Geometry.Point3D(0,0,0)
  idxI = shapeGrid.GetGridPointIndex(pt)
  shapeGridVect = shapeGrid.GetOccupancyVect()
  
  indicesInSphere = ComputeGridIndices(shapeGrid,winRad)
  
  nGridPts = len(shapeGridVect)
  for idxJ in indicesInSphere:
    idx = idxI+idxJ;
    if idx>=0 and idx<nGridPts:
      wt = shapeGridVect[idx]
      tPt = shapeGrid.GetGridPointLoc(idx)
      centroid += tPt*wt
      count+=wt
  if not count:
    raise ValueError('found no weight in sphere')
  centroid /= count
  #print 'csgc:','(%2f,%2f,%2f)'%tuple(pt),'(%2f,%2f,%2f)'%tuple(centroid),count
  return count,centroid
      


#-----------------------------------------------------------------------------
def FindTerminalPtsFromShape(shape,winRad,fraction,maxGridVal=3):
  pts = Geometry.FindGridTerminalPoints(shape.grid,winRad,fraction)
  termPts = [SubshapeObjects.SkeletonPoint(location=x) for x in pts]
  return termPts

#-----------------------------------------------------------------------------
def FindTerminalPtsFromConformer(conf,winRad,nbrCount):
  mol = conf.GetOwningMol()
  nAts = conf.GetNumAtoms()
  nbrLists=[[] for x in range(nAts)]
  for i in range(nAts):
    if(mol.GetAtomWithIdx(i).GetAtomicNum()<=1): continue
    pi = conf.GetAtomPosition(i)
    nbrLists[i].append((i,pi))
    for j in range(i+1,nAts):
      if(mol.GetAtomWithIdx(j).GetAtomicNum()<=1): continue
      pj = conf.GetAtomPosition(j)
      dist = pi.Distance(conf.GetAtomPosition(j))
      if dist<winRad:
        nbrLists[i].append((j,pj))
        nbrLists[j].append((i,pi))
  termPts=[]
  #for i in range(nAts):
  #  if not len(nbrLists[i]): continue
  #  if len(nbrLists[i])>10:
  #    print i+1,len(nbrLists[i])
  #  else:
  #    print i+1,len(nbrLists[i]),[x[0]+1 for x in nbrLists[i]]
        
  while 1:
    for i in range(nAts):
      if not nbrLists[i]: continue
      pos = Geometry.Point3D(0,0,0)
      totWt=0.0
      if len(nbrLists[i])<nbrCount:
        nbrList = nbrLists[i]
        for j in range(0,len(nbrList)):
          nbrJ,posJ=nbrList[j]
          weight = 1.*len(nbrLists[i])/len(nbrLists[nbrJ])
          pos += posJ*weight
          totWt+=weight
        pos /= totWt
        termPts.append(SubshapeObjects.SkeletonPoint(location=pos))
    if not len(termPts):
      nbrCount += 1
    else:
      break
  return termPts

#-----------------------------------------------------------------------------
def FindGridPointBetweenPoints(pt1,pt2,shapeGrid,winRad):
  center = pt1+pt2
  center /= 2.0
  d=1e8
  while d>shapeGrid.GetSpacing():
    count,centroid=Geometry.ComputeGridCentroid(shapeGrid,center,winRad)
    d = center.Distance(centroid)
    center = centroid
  return center

#-----------------------------------------------------------------------------
def ClusterTerminalPts(pts,winRad,scale):
  res = []
  tagged = [(y,x) for x,y in enumerate(pts)]
  while tagged:
    head,headIdx = tagged.pop(0)
    currSet = [head]
    i=0
    while i<len(tagged):
      nbr,nbrIdx=tagged[i]
      if head.location.Distance(nbr.location)<scale*winRad:
        currSet.append(nbr)
        del tagged[i]
      else:
        i+=1
    pt = Geometry.Point3D(0,0,0)
    for o in currSet:
      pt += o.location
    pt /= len(currSet)
    res.append(SubshapeObjects.SkeletonPoint(location=pt))
  return res

def GetMoreTerminalPoints(shape,pts,winRad,maxGridVal,targetNumber=5):
  """ adds a set of new terminal points using a max-min algorithm
  """
  shapeGrid=shape.grid
  shapeVect = shapeGrid.GetOccupancyVect()
  nGridPts = len(shapeVect)
  # loop, taking the grid point with the maximum minimum distance, until
  # we have enough points
  while len(pts)<targetNumber:
    maxMin=-1
    for i in range(nGridPts):
      if shapeVect[i]<maxGridVal:
        continue
      minVal=1e8
      posI = shapeGrid.GetGridPointLoc(i)
      for currPt in pts:
        dst = posI.Distance(currPt.location)
        if dst<minVal:
          minVal=dst
      if minVal>maxMin:
        maxMin=minVal
        bestPt=posI
    count,centroid=Geometry.ComputeGridCentroid(shapeGrid,bestPt,winRad)
    pts.append(SubshapeObjects.SkeletonPoint(location=centroid))


def FindFarthestGridPoint(shape,loc,winRad,maxGridVal):
  """ find the grid point with max occupancy that is furthest from a 
    given location
  """
  shapeGrid=shape.grid
  shapeVect = shapeGrid.GetOccupancyVect()
  nGridPts = len(shapeVect)
  dMax=-1;
  for i in range(nGridPts):
    if shapeVect[i]<maxGridVal:
      continue
    posI = shapeGrid.GetGridPointLoc(i)
    dst = posI.Distance(loc)
    if dst>dMax:
      dMax=dst
      res=posI

  count,centroid=Geometry.ComputeGridCentroid(shapeGrid,res,winRad)
  res=centroid
  return res

def ExpandTerminalPts(shape,pts,winRad,maxGridVal=3.0,targetNumPts=5):
  """ find additional terminal points until a target number is reached
  """
  if len(pts)==1:
    # if there's only one point, find the grid point with max value that is
    # *farthest* from this one and use it:
    pt2=FindFarthestGridPoint(shape,pts[0].location,winRad,maxGridVal)
    pts.append(SubshapeObjects.SkeletonPoint(location=pt2))
  if len(pts)==2:
    # add a point roughly in the middle:
    shapeGrid=shape.grid
    pt1 = pts[0].location
    pt2 = pts[1].location
    center = FindGridPointBetweenPoints(pt1,pt2,shapeGrid,winRad)
    pts.append(SubshapeObjects.SkeletonPoint(location=center))
  if len(pts)<targetNumPts:
    GetMoreTerminalPoints(shape,pts,winRad,maxGridVal,targetNumPts)

   
#-----------------------------------------------------------------------------
def AppendSkeletonPoints(shapeGrid,termPts,winRad,stepDist,maxGridVal=3,
                         maxDistC=15.0,distTol=1.5,symFactor=1.5):
  nTermPts = len(termPts)
  skelPts=[]
  shapeVect = shapeGrid.GetOccupancyVect()
  nGridPts = len(shapeVect)
  # find all possible skeleton points
  print('generate all possible')
  for i in range(nGridPts):
    if shapeVect[i]<maxGridVal:
      continue
    posI = shapeGrid.GetGridPointLoc(i)
    ok=True
    for pt in termPts:
      dst = posI.Distance(pt.location)
      if dst<stepDist:
        ok=False
        break
    if ok:
      skelPts.append(SubshapeObjects.SkeletonPoint(location=posI))
  # now start removing them
  print('Compute centroids:',len(skelPts))
  gridBoxVolume=shapeGrid.GetSpacing()**3
  maxVol = 4.0*math.pi/3.0 * winRad**3 * maxGridVal / gridBoxVolume
  i=0
  while i<len(skelPts):
    pt = skelPts[i]
    count,centroid=Geometry.ComputeGridCentroid(shapeGrid,pt.location,winRad)
    #count,centroid=ComputeShapeGridCentroid(pt.location,shapeGrid,winRad)
    centroidPtDist=centroid.Distance(pt.location)
    if centroidPtDist>maxDistC:
      del skelPts[i]
    else:
      pt.fracVol = float(count)/maxVol
      pt.location.x = centroid.x
      pt.location.y = centroid.y
      pt.location.z = centroid.z
      i+=1

  print('remove points:',len(skelPts))
  res = termPts+skelPts
  i=0
  while i<len(res):
    p=-1
    mFrac=0.0
    ptI = res[i]

    startJ = max(i+1,nTermPts)
    for j in range(startJ,len(res)):
      ptJ=res[j]
      distC = ptI.location.Distance(ptJ.location)
      if distC<symFactor*stepDist:
        if ptJ.fracVol>mFrac:
          p=j
          mFrac=ptJ.fracVol
    #print i,len(res),p,mFrac
    if p>-1:
      ptP = res.pop(p)
      j = startJ
      while j < len(res):
        ptJ=res[j]
        distC = ptI.location.Distance(ptJ.location)
        if distC<symFactor*stepDist:
          del res[j]
        else:
          j+=1
      res.append(ptP)
      #print '% 3d'%i,'% 5.2f % 5.2f % 5.2f'%tuple(list(ptI.location)),' - ','% 5.2f % 5.2f % 5.2f'%tuple(list(ptJ.location))
    i+=1
  return res

#-----------------------------------------------------------------------------
def CalculateDirectionsAtPoint(pt,shapeGrid,winRad):
  shapeGridVect = shapeGrid.GetOccupancyVect()
  nGridPts = len(shapeGridVect)
  tmp = winRad/shapeGrid.GetSpacing()
  radInGrid=int(tmp)
  radInGrid2=int(tmp*tmp)
  covMat = numpy.zeros((3,3),numpy.float64)

  dX = shapeGrid.GetNumX()
  dY = shapeGrid.GetNumY()
  dZ = shapeGrid.GetNumZ()
  idx = shapeGrid.GetGridPointIndex(pt.location)
  idxZ = idx//(dX*dY)
  rem = idx%(dX*dY)
  idxY = rem//dX
  idxX = rem%dX
  totWt=0.0
  for i in range(-radInGrid,radInGrid):
    xi = idxX+i
    for j in range(-radInGrid,radInGrid):
      xj = idxY+j
      for k in range(-radInGrid,radInGrid):
        xk = idxZ+k
        d2 = i*i+j*j+k*k
        if d2>radInGrid2 and int(math.sqrt(d2))>radInGrid:
          continue
        gridIdx = (xk*dY+xj)*dX+xi
        if gridIdx>=0 and gridIdx<nGridPts:
          wtHere = shapeGridVect[gridIdx]
          totWt += wtHere
          ptInSphere = shapeGrid.GetGridPointLoc(gridIdx)
          ptInSphere -= pt.location
          covMat[0][0]+= wtHere*(ptInSphere.x*ptInSphere.x)
          covMat[0][1]+= wtHere*(ptInSphere.x*ptInSphere.y)
          covMat[0][2]+= wtHere*(ptInSphere.x*ptInSphere.z)
          covMat[1][1]+= wtHere*(ptInSphere.y*ptInSphere.y)
          covMat[1][2]+= wtHere*(ptInSphere.y*ptInSphere.z)
          covMat[2][2]+= wtHere*(ptInSphere.z*ptInSphere.z)
  covMat /= totWt
  covMat[1][0] = covMat[0][1]
  covMat[2][0] = covMat[0][2]
  covMat[2][1] = covMat[1][2]

  eVals,eVects = numpy.linalg.eigh(covMat)
  sv = zip(eVals,numpy.transpose(eVects))
  sv.sort(reverse=True)
  eVals,eVects=zip(*sv)
  pt.shapeMoments=tuple(eVals)
  pt.shapeDirs = tuple([Geometry.Point3D(p[0],p[1],p[2]) for p in eVects])

  #print '-------------'
  #print pt.location.x,pt.location.y,pt.location.z
  #for v in covMat:
  # print '  ',v
  #print '---'
  #print eVals
  #for v in eVects:
  #  print '  ',v
  #print '-------------'
  

#-----------------------------------------------------------------------------
def AssignMolFeatsToPoints(pts,mol,featFactory,winRad):
  feats = featFactory.GetFeaturesForMol(mol)
  for i,pt in enumerate(pts):
    for feat in feats:
      if feat.GetPos().Distance(pt.location)<winRad:
        typ = feat.GetFamily()
        if typ not in pt.molFeatures:
          pt.molFeatures.append(typ)
    print(i,pt.molFeatures)

