#  Copyright (c) 2017 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.

import copy
import os
import sys
import unittest

from rdkit import Chem, Geometry
from rdkit.Chem import rdCoordGen, rdMolAlign


def compareConfs(c1, c2, match, tol=1e-2, alignIt=False):
  for i, j in enumerate(match):
    pi = c2.GetAtomPosition(i)
    pj = c1.GetAtomPosition(j)
    if (pj - pi).Length() >= tol:
      return False
  return True


class TestCase(unittest.TestCase):

  def test_basics(self):
    mol = Chem.MolFromSmiles('CCOC')
    self.assertEqual(mol.GetNumConformers(), 0)
    rdCoordGen.AddCoords(mol)
    self.assertEqual(mol.GetNumConformers(), 1)
    rdCoordGen.AddCoords(mol)
    self.assertEqual(mol.GetNumConformers(), 1)
    rwmol = Chem.RWMol(mol)
    rdCoordGen.AddCoords(rwmol)
    self.assertEqual(rwmol.GetNumConformers(), 1)

  def test_template1(self):
    template = Chem.MolFromSmiles('C1OCOCOCOCCCNCCC1')
    template.SetProp("_Name", 'template')
    mol = Chem.MolFromSmiles('C1OCOCOCOCCCNC(OC(=O)C2CC2)CC1')
    mol.SetProp("_Name", "mol")
    rdCoordGen.AddCoords(template)
    rdCoordGen.AddCoords(mol)
    self.assertFalse(
      compareConfs(mol.GetConformer(), template.GetConformer(), mol.GetSubstructMatch(template)))
    match = mol.GetSubstructMatch(template)
    mapd = dict()
    for i, aid in enumerate(match):
      p = template.GetConformer().GetAtomPosition(i)
      mapd[aid] = Geometry.Point2D(p.x, p.y)
    ps = rdCoordGen.CoordGenParams()
    ps.SetCoordMap(mapd)
    ps.dbg_useFixed = True
    rdCoordGen.AddCoords(mol, ps)
    self.assertTrue(
      compareConfs(mol.GetConformer(), template.GetConformer(), mol.GetSubstructMatch(template)))

  def test_template2(self):
    # the easier way...
    template = Chem.MolFromSmiles('C1OCOCOCOCCCNCCC1')
    template.SetProp("_Name", 'template')
    mol = Chem.MolFromSmiles('C1OCOCOCOCCCNC(OC(=O)C2CC2)CC1')
    mol.SetProp("_Name", "mol")
    rdCoordGen.AddCoords(template)
    ps = rdCoordGen.CoordGenParams()
    ps.SetTemplateMol(template)
    ps.dbg_useFixed = True
    rdCoordGen.AddCoords(mol, ps)
    self.assertTrue(
      compareConfs(mol.GetConformer(), template.GetConformer(), mol.GetSubstructMatch(template)))

  def test_template3(self):
    # the easier way, test lifetime...
    template = Chem.MolFromSmiles('C1OCOCOCOCCCNCCC1')
    mol = Chem.MolFromSmiles('C1OCOCOCOCCCNC(OC(=O)C2CC2)CC1')
    rdCoordGen.AddCoords(template)
    template2 = Chem.Mol(template)
    ps = rdCoordGen.CoordGenParams()
    ps.SetTemplateMol(template2)
    template2 = None
    ps.dbg_useFixed = True
    rdCoordGen.AddCoords(mol, ps)
    self.assertTrue(
      compareConfs(mol.GetConformer(), template.GetConformer(), mol.GetSubstructMatch(template)))


if __name__ == '__main__':
  unittest.main()
