#  Copyright (C) 2014 Sereina Riniker
#
#  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.
#
""" Torsion Fingerprints (Deviation) (TFD)
    According to a paper from Schulz-Gasch et al., JCIM, 52, 1499-1512 (2012).

"""
from rdkit import rdBase
from rdkit import RDConfig
from rdkit import Geometry
from rdkit import Chem
from rdkit.Chem import rdchem
from rdkit.Chem import rdMolDescriptors
import math, os

def _doMatch(inv, atoms):
  """ Helper function to check if all atoms in the list are the same
      
      Arguments:
      - inv:    atom invariants (used to define equivalence of atoms)
      - atoms:  list of atoms to check

      Return: boolean
  """
  match = True
  for i in range(len(atoms)-1):
    for j in range(i+1, len(atoms)):
      if (inv[atoms[i].GetIdx()] != inv[atoms[j].GetIdx()]):
        match = False
        return match
  return match

def _doNotMatch(inv, atoms):
  """ Helper function to check if all atoms in the list are NOT the same
      
      Arguments:
      - inv:    atom invariants (used to define equivalence of atoms)
      - atoms:  list of atoms to check

      Return: boolean
  """
  match = True
  for i in range(len(atoms)-1):
    for j in range(i+1, len(atoms)):
      if (inv[atoms[i].GetIdx()] == inv[atoms[j].GetIdx()]):
        match = False
        return match
  return match

def _doMatchExcept1(inv, atoms):
  """ Helper function to check if two atoms in the list are the same, 
      and one not
      Note: Works only for three atoms
      
      Arguments:
      - inv:    atom invariants (used to define equivalence of atoms)
      - atoms:  list of atoms to check

      Return: atom that is different
  """
  if len(atoms) != 3:
    raise ValueError("Number of atoms must be three")
  a1 = atoms[0].GetIdx()
  a2 = atoms[1].GetIdx()
  a3 = atoms[2].GetIdx()
  if (inv[a1] == inv[a2] and inv[a1] != inv[a3] and inv[a2] != inv[a3]):
    return atoms[2]
  elif (inv[a1] != inv[a2] and inv[a1] == inv[a3] and inv[a2] != inv[a3]):
    return atoms[1]
  elif (inv[a1] != inv[a2] and inv[a1] != inv[a3] and inv[a2] == inv[a3]):
    return atoms[0]
  return None

def _getAtomInvariantsWithRadius(mol, radius):
  """ Helper function to calculate the atom invariants for each atom 
      with a given radius

      Arguments:
      - mol:    the molecule of interest
      - radius: the radius for the Morgan fingerprint

      Return: list of atom invariants
  """
  inv = []
  for i in range(mol.GetNumAtoms()):
    info = {}
    fp = rdMolDescriptors.GetMorganFingerprint(mol, radius, fromAtoms=[i], bitInfo=info)
    for k in info.keys():
      if info[k][0][1] == radius:
        inv.append(k)
  return inv

def _getHeavyAtomNeighbors(atom1, aid2=-1):
  """ Helper function to calculate the number of heavy atom neighbors.

      Arguments:
      - atom1:    the atom of interest
      - aid2:     atom index that should be excluded from neighbors (default: none)

      Return: a list of heavy atom neighbors of the given atom
  """
  if aid2 < 0:
    return [n for n in atom1.GetNeighbors() if n.GetSymbol()!='H']
  else:
    return [n for n in atom1.GetNeighbors() if (n.GetSymbol()!='H' and n.GetIdx()!=aid2)]

def _getIndexforTorsion(neighbors, inv):
  """ Helper function to calculate the index of the reference atom for 
      a given atom

      Arguments:
      - neighbors:  list of the neighbors of the atom
      - inv:        atom invariants

      Return: list of atom indices as reference for torsion
  """
  if len(neighbors) == 1: # atom has only one neighbor
    return [neighbors[0]]
  elif _doMatch(inv, neighbors): # atom has all symmetric neighbors
    return neighbors
  elif _doNotMatch(inv, neighbors): # atom has all different neighbors
    # simply use the first neighbor
    return [neighbors[0]]
  at = _doMatchExcept1(inv, neighbors) # two neighbors the same, one different