//  $Id$
//
//  Copyright (C) 2003,2004 Rational Discovery LLC
//   All Rights Reserved
//
#include "TemplEnum.h"
#include <Geometry/Transform3D.h>
#include <GraphMol/MolTransforms/MolTransforms.h>
#include <GraphMol/FileParsers/FileParsers.h>
#define FEQ(_a_,_b_) (fabs((_a_)-(_b_))<1e-4)

namespace TemplateEnum {
  using namespace RDKit;

  // ------------------------------------------------------------------
  //
  // transforms a sidechain so that it is oriented better for attachment
  // to a molecule
  //
  //  Arguments:
  //   mol: core molecule 
  //   sidechain: sidechain to attach.
  //   molConnectorIdx: the index of the attachment point atom in the
  //     molecule.
  //   sidechainConnectorIdx: the index of the attachment point atom
  //     in the sidechain.
  // 
  // ------------------------------------------------------------------
  void orientSidechain(RWMol *mol,RWMol *sidechain,
		       int molAttachIdx,int sidechainAttachIdx){
    PRECONDITION(mol,"bad molecule");
    PRECONDITION(sidechain,"bad molecule");

    // ---------
    // start by getting our 4 atoms
    // ---------
    Conformer &molConf=mol->getConformer();
    Conformer &sidechainConf=sidechain->getConformer();
    Atom *molConnAtom,*molAttachAtom;
    Atom *chainConnAtom,*chainAttachAtom;
    molAttachAtom = mol->getAtomWithIdx(molAttachIdx);
    chainAttachAtom = sidechain->getAtomWithIdx(sidechainAttachIdx);
    PRECONDITION(molAttachAtom->getDegree()==1,"attachment points must be degree 1");
    PRECONDITION(chainAttachAtom->getDegree()==1,"attachment points must be degree 1");
    RWMol::ADJ_ITER nbrIdx,endNbrs;
    boost::tie(nbrIdx,endNbrs) = mol->getAtomNeighbors(molAttachAtom);
    molConnAtom=mol->getAtomWithIdx(*nbrIdx);
    boost::tie(nbrIdx,endNbrs) = sidechain->getAtomNeighbors(chainAttachAtom);
    chainConnAtom=sidechain->getAtomWithIdx(*nbrIdx);

    //-----------------------------------------
    //  Notation:
    //    Pmc: molecule connection point (the atom that will be
    //     removed from the molecule).
    //    Pma: molecule attachment point (the atom to which we'll form
    //     the bond).
    //    Psc: sidechain connection point
    //    Psa: sidechain attachment point
    //    Vm: Pmc-Pma (molecular attachment vector)
    //    Vs: Psc-Psa (sidechain attachment vector)
    //
    //-----------------------------------------
    RDGeom::Transform3D sidechainTform,templateTform,tmpTform;

    RDGeom::Point3D Vm,Um,Pmc,Pma;
    RDGeom::Point3D Vs,Us,Psc,Psa;
    Pmc = molConf.getAtomPos(molConnAtom->getIdx());
    Pma = molConf.getAtomPos(molAttachAtom->getIdx());
    std::cerr << "p=array(["<<Pma.x<<","<<Pma.y<<","<<Pma.z<<"])" << std::endl;
    Psc = sidechainConf.getAtomPos(chainConnAtom->getIdx());
    Psa = sidechainConf.getAtomPos(chainAttachAtom->getIdx());

    templateTform.setToIdentity();
    Vm = Pmc - Pma;
    // note the opposite direction here:
    Vs = Psa - Psc;
    Um = Vm;
    Um.normalize();
    std::cerr << "Um=array(["<<Um.x<<","<<Um.y<<","<<Um.z<<"])" << std::endl;
    Us = Vs;
    Us.normalize();
    std::cerr << "Us=array(["<<Us.x<<","<<Us.y<<","<<Us.z<<"])" << std::endl;
    // translate Psc -> Pma
    //RDGeom::Point3D headTrans = Pma-Psc;
    
    templateTform.setToIdentity();

    tmpTform.setToIdentity();
    tmpTform.SetTranslation(Pma);
    templateTform *= tmpTform;

    double sinT,cosT;
    cosT = Us.dotProduct(Um);
    if(cosT>1.0) cosT = 1.0;
    if(fabs(cosT)<1.0){
      tmpTform.setToIdentity();
      sinT = sqrt(1.0-cosT*cosT);
      RDGeom::Point3D rotnAxis=Us.crossProduct(Um);
      rotnAxis.normalize();
      std::cerr << "ax=array(["<<rotnAxis.x<<","<<rotnAxis.y<<","<<rotnAxis.z<<"])" << std::endl;
      tmpTform.SetRotation(cosT,sinT,rotnAxis);
      templateTform *= tmpTform;
    } else if(cosT==1.0){
      RDGeom::Point3D normal(1,0,0);
      if(fabs(Us.dotProduct(normal))==1.0){
	normal = RDGeom::Point3D(0,1,0);
      }
      RDGeom::Point3D rotnAxis=Us.crossProduct(normal);
      templateTform.SetRotation(-1,0,rotnAxis);
    }
    
    tmpTform.setToIdentity();
    tmpTform.SetTranslation(Psc*-1.0);
    templateTform *= tmpTform;

    // ---------
    // transform the atomic positions in the sidechain:
    // ---------
    MolTransforms::transformMolsAtoms(sidechain,templateTform);


    // that's it!
  }
  // ------------------------------------------------------------------
  //
  // attaches a sidechain fragment to a molecule.
  //
  //  Arguments:
  //   mol: molecule to be modified
  //   sidechain: sidechain to attach.  The sidechain is copied in.
  //   molConnectorIdx: the index of the attachment point atom in the
  //     molecule.
  //   sidechainConnectorIdx: the index of the attachment point atom
  //     in the sidechain.
  //   bondType: type of the bond to form between the atoms
  //
  //  The connector atoms are *NOT* part of the final molecule, they
  //  merely serve to establish where things connect.
  // 
  // ------------------------------------------------------------------
  void molAddSidechain(RWMol *mol,RWMol *sidechain,
		       int molConnectorIdx,int sidechainConnectorIdx,
		       Bond::BondType bondType){
    PRECONDITION(mol,"bad molecule provided");
    PRECONDITION(sidechain,"bad sidechain provided");
    int origNumAtoms=mol->getNumAtoms();
    mol->insertMol(*sidechain);

    // get pointers to the two connectors (these are the atoms
    // we'll end up removing)
    Atom *molConnAtom,*sidechainConnAtom;
    molConnAtom = mol->getAtomWithIdx(molConnectorIdx);
    sidechainConnAtom = mol->getAtomWithIdx(sidechainConnectorIdx+origNumAtoms);

    // now use those pointers to get the atoms which will remain
    // (we're going to connect these) and remove the original
    // connection points from the molecule
    Atom *tmpAtom;
    RWMol::ADJ_ITER nbrIdx,endNbrs;
    boost::tie(nbrIdx,endNbrs) = mol->getAtomNeighbors(molConnAtom);

    // we are assuming that the first neighbor is the correct one.
    // Really we should be able to assume that there is only a single
    // attachment point.
    tmpAtom = molConnAtom;
    molConnAtom = mol->getAtomWithIdx(*nbrIdx);
    mol->removeAtom(tmpAtom);
    // repeat that process for the sidechain:
    boost::tie(nbrIdx,endNbrs) = mol->getAtomNeighbors(sidechainConnAtom);
    tmpAtom = sidechainConnAtom;
    sidechainConnAtom = mol->getAtomWithIdx(*nbrIdx);
    mol->removeAtom(tmpAtom);

    // finally connect the remaining atoms:
    mol->addBond(molConnAtom,sidechainConnAtom,bondType);
  }


  // used as a starting point for connection bookmarks
  const int CONNECT_BOOKMARK_START=0x23424223;

  // ------------------------------------------------------------------
  //
  // Loop through all the atoms and bookmark the attachment points
  //
  //  attachment points are assumed to have one of these properties:
  //   - common_properties::molFileAlias (set by the mol file parser)
  //   - common_properties::dummyLabel   (set by the SMILES parser)
  //  frontMarker is the "recognition character" to be used to pick
  //  valid labels.  e.g. if the frontMarker is 'X', a label beginning
  //  with 'Y' will not be marked.
  //
  //  In addition to bookmarking the attachment points, we also set
  //  the common_properties::maxAttachIdx property, which holds an integer with the
  //  maximum attachment bookmark index. This is used in the
  //  enumeration to prevent us from having to scan through all the
  //  molecule's bookmarks
  //
  //
  // ------------------------------------------------------------------
  void markAttachmentPoints(RWMOL_SPTR mol,char frontMarker){
    markAttachmentPoints(mol.get(),frontMarker);
  }
  void markAttachmentPoints(RWMol *mol,char frontMarker){
    PRECONDITION(mol,"bad molecule");

    RWMol::AtomIterator atomIt;
    int maxAttachIdx=0;

    // scan through the atoms and mark those that have aliases
    //  (these might be attachment points)
    for(atomIt=mol->beginAtoms();atomIt!=mol->endAtoms();atomIt++){
      // start by finding possible attachment point properties:
      std::string attachLabel="";
      if((*atomIt)->hasProp(common_properties::molFileAlias)){
	(*atomIt)->getProp(common_properties::molFileAlias,attachLabel);
      } else if((*atomIt)->hasProp(common_properties::dummyLabel)){
	(*atomIt)->getProp(common_properties::dummyLabel,attachLabel);
      }

      // if we got one and it starts with the appropriate front
      // marker, proceed:
      if(attachLabel!="" && attachLabel[0]==frontMarker ){
	// to avoid trouble later, we guarantee that the attachment
	// point has degree 1 (one bond to it).
	if((*atomIt)->getDegree()>1)
	  throw EnumException("More than one bond to an attachment point.");

	int offset = CONNECT_BOOKMARK_START;
	if(attachLabel.length()>1){
	  if(attachLabel[1]>='a' && attachLabel[1]<='z'){
	    offset += (int)attachLabel[1] - (int)'a';
	  }
	}
	mol->setAtomBookmark(*atomIt,offset);
	if(offset>maxAttachIdx) maxAttachIdx=offset;
      }
    }
    if(maxAttachIdx){
      mol->setProp(common_properties::maxAttachIdx,maxAttachIdx);
    }
  }

  // ------------------------------------------------------------------
  //
  // loops through the sidechain molecules and calls
  // _markAttachmentPoints()_ on each.
  //
  // ------------------------------------------------------------------
  void prepareSidechains(RWMOL_SPTR_VECT *sidechains,char frontMarker){
    PRECONDITION(sidechains,"bad sidechain list");
    RWMOL_SPTR_VECT::iterator mpvI;
    for(mpvI=sidechains->begin();mpvI!=sidechains->end();mpvI++){
      markAttachmentPoints(*mpvI,frontMarker);
    }
  }

  // ------------------------------------------------------------------
  //
  // Enumerates the library around a template and returns the result.
  //
  //
  // ------------------------------------------------------------------
  RWMOL_SPTR_VECT enumerateLibrary(RWMol *templateMol,
				   VECT_RWMOL_SPTR_VECT &sidechains,
				   bool orientSidechains){
    PRECONDITION(templateMol,"bad molecule");
    RWMOL_SPTR_VECT res,tmp;
    res.push_back(RWMOL_SPTR(new RWMol(*templateMol)));

    // if there's no attachment point on the molecule or no
    // sidechains, return now: 
    if(!templateMol->hasProp(common_properties::maxAttachIdx) || sidechains.size()==0 )
      return res;

    int maxIdx;
    templateMol->getProp(common_properties::maxAttachIdx,maxIdx);

    tmp.clear();
    // loop over the sidechains and attach them
    for(unsigned int i=0;i<sidechains.size();i++){
      int tgtMark=CONNECT_BOOKMARK_START+i;
      // here's another boundary condition
      if(tgtMark>maxIdx) break;

      /// loop over all atoms with the appropriate mark
      //  This means that if a mol has two attachment points with the
      //  same name (e.g. two Xa's) they'll always have the same
      //  sidechain attached to them.  This is a feature.
      RWMOL_SPTR_VECT::iterator sidechainIt;
      for(sidechainIt=sidechains[i].begin();
	  sidechainIt!=sidechains[i].end();
	  sidechainIt++){
	// we've got our sidechain, find the atom it attaches from
	if( (*sidechainIt)->hasAtomBookmark(CONNECT_BOOKMARK_START) ){
	  //
	  // NOTE: If there's more than one marked atom in the sidechain,
	  ///      we'll only use the first for the moment.
	  //
	  int sidechainAtomIdx = (*sidechainIt)->getAtomWithBookmark(CONNECT_BOOKMARK_START)->getIdx();
	
	  // now add the sidechain to each molecule
	  RWMOL_SPTR_VECT::iterator templMolIt;
	  // loop over all the mols we've generated to this point
	  for(templMolIt=res.begin();templMolIt!=res.end();templMolIt++){
	    RWMol *templ =  new RWMol(**templMolIt);
	    std::string name,tmpStr;
	    if(templ->hasProp(common_properties::_Name)){
	      templ->getProp(common_properties::_Name,tmpStr);
	      name = name + " " + tmpStr;
	    }
	    while(templ->hasAtomBookmark(tgtMark)){
	      // this is the atom we'll be replacing in the template
	      Atom *at = templ->getAtomWithBookmark(tgtMark);

	      // copy and transform the sidechain:
	      RWMol *sidechain;
	      if(orientSidechains){
		sidechain = new RWMol(*(sidechainIt->get()));
		orientSidechain(templ,sidechain,
				at->getIdx(),sidechainAtomIdx);
	      } else {
		sidechain = sidechainIt->get();
	      }
	      // FIX: need to use the actual bond order here:
	      molAddSidechain(templ,sidechain,
			      at->getIdx(),sidechainAtomIdx,
			      Bond::SINGLE);
	      if(sidechain->hasProp(common_properties::_Name)){
		sidechain->getProp(common_properties::_Name,tmpStr);
		name = name + " " + tmpStr;
	      }
	      templ->clearAtomBookmark(tgtMark,at);
	      if(orientSidechains){
		delete sidechain;
	      }
	    }
	    //std::cout << templ << "> " << MolToSmiles(*templ) << std::endl;
	    if(name != "") templ->setProp(common_properties::_Name,name);
	    tmp.push_back(RWMOL_SPTR(templ));
	  }
	}
      }

      //
      // if we just made any molecules, free up the memory used by the
      // existing result set and move the molecules we just generated
      // over
      if(tmp.size()){
#if 0      
	RWMOL_SPTR_VECT::iterator tmpMolIt;
	for(tmpMolIt=res.begin();tmpMolIt!=res.end();tmpMolIt++){
	  delete *tmpMolIt;
	}
#endif
	res = tmp;
	tmp.clear();
      }
    }
    return res;
  }

  // ------------------------------------------------------------------
  //
  //  Reads a template and library of sidechains from input files.
  //   the template file should be a mol file and the sidechain files
  //   SD files
  //
  // ------------------------------------------------------------------
  RWMOL_SPTR_VECT enumFromFiles(const char *templateName,
			    std::vector<const char *> &sidechainNames){
    PRECONDITION(templateName,"bad template file name passed in");

    // build and mark the template molecule
    RWMol *templ = MolFileToMol(templateName,false);
    if(!templ) throw EnumException("could not construct template molecule"); 
    markAttachmentPoints(templ,'X');

    // now build and mark each set of sidechains:
    RWMOL_SPTR_VECT sidechains;
    VECT_RWMOL_SPTR_VECT allSidechains;
    for(std::vector<const char*>::const_iterator i=sidechainNames.begin();
	i!=sidechainNames.end();i++){
      sidechains = SDFileToMols(*i,false);
      if(!sidechains.size()){
	std::string err="no sidechains read from file: ";
	err += *i;
	throw EnumException(err.c_str());
      }
      prepareSidechains(&sidechains,'X');
      allSidechains.push_back(sidechains);
    }

    // enumerate the library:
    RWMOL_SPTR_VECT library=enumerateLibrary(templ,allSidechains);


    //--------------------------
    //
    // Clean up the molecules and sidechains we constructed along the
    // way. 
    //
    //--------------------------
    delete templ;
#if 0
    VECT_RWMOL_SPTR_VECT::iterator vmpvI;
    for(vmpvI=allSidechains.begin();vmpvI!=allSidechains.end();vmpvI++){
      RWMOL_SPTR_VECT::iterator mpvI;
      for(mpvI=vmpvI->begin();mpvI!=vmpvI->end();mpvI++){
	delete *mpvI;
      }
      vmpvI->clear();
    }
#endif
    allSidechains.clear();

    return library;
  }  

} // end of TemplateEnum namespace

