#
# Copyright (C) 2015 Greg Landrum
#
#   @@ All Rights Reserved @@
#  This file is part of the RDKit.
#  The contents are covered by the terms of the BSD license
#  which is included in the file license.txt, found at the root
#  of the RDKit source tree.

from rdkit import Chem
from rdkit.Chem import AllChem
import numpy as np
from numpy import linalg


def GetBestFitPlane(pts, weights=None):
  if weights is None:
    wSum = len(pts)
    origin = np.sum(pts, 0)
  origin /= wSum

  sumXX = 0
  sumXY = 0
  sumXZ = 0
  sumYY = 0
  sumYZ = 0
  sumZZ = 0
  sums = np.zeros((3, 3), np.double)
  for pt in pts:
    dp = pt - origin
    for i in range(3):
      sums[i, i] += dp[i] * dp[i]
      for j in range(i + 1, 3):
        sums[i, j] += dp[i] * dp[j]
        sums[j, i] += dp[i] * dp[j]
  sums /= wSum
  vals, vects = linalg.eigh(sums)
  order = np.argsort(vals)
  normal = vects[:, order[0]]
  plane = np.zeros((4, ), np.double)
  plane[:3] = normal
  plane[3] = -1 * normal.dot(origin)
  return plane


def PBFRD(mol, confId=-1):
  conf = mol.GetConformer(confId)
  if not conf.Is3D():
    return 0

  pts = np.array([list(conf.GetAtomPosition(x)) for x in range(mol.GetNumAtoms())])
  plane = GetBestFitPlane(pts)
  denom = np.dot(plane[:3], plane[:3])
  denom = denom**0.5
  # add up the distance from the plane for each point:
  res = 0.0
  for pt in pts:
    res += np.abs(pt.dot(plane[:3]) + plane[3])
  res /= denom
  res /= len(pts)
  return res


if __name__ == '__main__':
  suppl = Chem.SDMolSupplier('./testData/egfr.sdf', removeHs=False)
  expected = open('./testData/egfr.out', 'r')
  for m in suppl:
    res = PBFRD(m)
    inl = next(expected).strip().split()
    expect = float(inl[1])
    assert abs(res - expect) < 1e-4