// $Id$
//
//  Copyright (C) 2009-2012 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.
//
//  Includes contributions from Dave Cosgrove (davidacosgroveaz@gmail.com)
//
#ifndef _RD_MOLDRAWING_H_
#define _RD_MOLDRAWING_H_

#include <vector>
#include <boost/foreach.hpp>
#include <boost/lexical_cast.hpp>
#include <GraphMol/RDKitBase.h>
#include <GraphMol/Depictor/RDDepictor.h>
#include <Geometry/point.h>

/***********
  Return Format: vector of ints

  RESOLUTION dots_per_angstrom
  BOUNDS x1 y1 x2 y2
  LINE width dashed atom1_atnum atom2_atnum x1 y1 x2 y2
  WEDGE dashed atom1_atnum atom2_atnum x1 y1 x2 y2 x3 y3 
  ATOM idx atnum x y num_chars char1-charx orient



*************/

namespace RDKit {
  namespace Drawing {
    typedef int ElementType;

    typedef enum {
      LINE=1,
      WEDGE,
      ATOM,
      BOUNDS,
      RESOLUTION
    } PrimType;
    typedef enum {
      C=0,
      N,
      E,
      S,
      W
    } OrientType;

    namespace detail {
      // **************************************************************************
      void drawLine( std::vector<ElementType> &res ,
                     int atnum1 , int atnum2 , int lineWidth , int dashed ,
                     double x1 , double y1 ,
                     double x2 , double y2 ) {

        res.push_back( LINE );
        res.push_back( static_cast<ElementType>(lineWidth) );
        res.push_back(dashed);
        res.push_back( static_cast<ElementType>(atnum1) );
        res.push_back( static_cast<ElementType>(atnum2) );
        res.push_back( static_cast<ElementType>(x1) );
        res.push_back( static_cast<ElementType>(y1) );
        res.push_back( static_cast<ElementType>(x2) );
        res.push_back( static_cast<ElementType>(y2) );

      }
      std::pair<std::string,OrientType> getAtomSymbolAndOrientation(const Atom &atom,RDGeom::Point2D nbrSum){
        std::string symbol="";
        OrientType orient=C;
        int isotope = atom.getIsotope();
        if(atom.getAtomicNum()!=6 ||
           atom.getFormalCharge()!=0 ||
           isotope ||
           atom.getNumRadicalElectrons()!=0 ||
           atom.hasProp(common_properties::molAtomMapNumber) ||
           atom.getDegree()==0 ){
          symbol=atom.getSymbol();
          bool leftToRight=true;
          if(atom.getDegree()==1 && nbrSum.x>0){
            leftToRight=false;
          }
          if(isotope){
            symbol = boost::lexical_cast<std::string>(isotope)+symbol;
          }
          if(atom.hasProp(common_properties::molAtomMapNumber)){
            std::string mapNum;
            atom.getProp(common_properties::molAtomMapNumber,mapNum);
            symbol += ":" + mapNum;
          }
          int nHs=atom.getTotalNumHs();
          if(nHs>0){
            std::string h="H";
            if(nHs>1) {
              h += boost::lexical_cast<std::string>(nHs);
            }
            if(leftToRight) symbol += h;
            else symbol = h+symbol;
          }
          if( atom.getFormalCharge()!=0 ){
            int chg=atom.getFormalCharge();
            std::string sgn="+";
            if(chg<0){
              sgn="-";
            }
            chg=abs(chg);
            if(chg>1){
              sgn += boost::lexical_cast<std::string>(chg);
            } 
            if(leftToRight) symbol+=sgn;
            else symbol = sgn+symbol;
          }

          if(atom.getDegree()==1){
            double islope=0;
            if(fabs(nbrSum.y)>1){
              islope=nbrSum.x/fabs(nbrSum.y);
            } else {
              islope=nbrSum.x;
            }
            if(fabs(islope)>.85){
              if(islope>0){
                orient=W;
              } else {
                orient=E;
              }
            } else {
              if(nbrSum.y>0){
                orient=N;
              } else {
                orient=S;
              }
            }
          }
        }
        return std::make_pair(symbol,orient);
      }
    } // end of detail namespace
    // **************************************************************************
    std::vector<ElementType> DrawMol(const ROMol &mol,int confId=-1,
                                     const std::vector<int> *highlightAtoms=0,
                                     bool includeAtomCircles=false,
                                     unsigned int dotsPerAngstrom=100,
                                     double dblBondOffset=0.3,
                                     double dblBondLengthFrac=0.8,
                                     double angstromsPerChar=0.20
                                     ){
      if(!mol.getRingInfo()->isInitialized()){
        MolOps::findSSSR(mol);
      }
      std::vector<ElementType> res;
      res.push_back(RESOLUTION);
      res.push_back(static_cast<ElementType>(dotsPerAngstrom));
      
      const Conformer &conf=mol.getConformer(confId);
      const RDGeom::POINT3D_VECT &locs=conf.getPositions();

      // get atom symbols and orientations
      // (we need them for the bounding box calculation)
      std::vector< std::pair<std::string,OrientType> > atomSymbols;
      ROMol::VERTEX_ITER bAts,eAts;
      boost::tie(bAts,eAts)=mol.getVertices();
      while(bAts!=eAts){
        ROMol::OEDGE_ITER nbr,endNbrs;
        RDGeom::Point2D nbrSum(0,0);
        boost::tie(nbr,endNbrs) = mol.getAtomBonds(mol[*bAts].get());
        RDGeom::Point2D a1(locs[mol[*bAts]->getIdx()].x,
                           locs[mol[*bAts]->getIdx()].y);
        while(nbr!=endNbrs){
          const BOND_SPTR bond=mol[*nbr];
          ++nbr;
          int a2Idx=bond->getOtherAtomIdx(mol[*bAts]->getIdx());
          RDGeom::Point2D a2(locs[a2Idx].x,locs[a2Idx].y);
          nbrSum+=a2-a1;
        }
        atomSymbols.push_back(detail::getAtomSymbolAndOrientation(*mol[*bAts],nbrSum));
        ++bAts;
      }


      //------------
      // do the bounding box
      //------------
      double minx=1e6,miny=1e6,maxx=-1e6,maxy=-1e6;
      for(unsigned int i=0;i<mol.getNumAtoms();++i){
        RDGeom::Point3D pt=locs[i];
        std::string symbol;
        OrientType orient;
        boost::tie(symbol,orient)=atomSymbols[i];
        if(symbol!=""){
          // account for a possible expansion of the bounding box by the symbol
          if(pt.x<=minx){
            switch(orient){
            case C:
            case N:
            case S:
            case E:
              minx = pt.x-symbol.size()/2*angstromsPerChar;
              break;
            case W:
              minx = pt.x-symbol.size()*angstromsPerChar;
              break;
            }
          }
          if(pt.x>=maxx){
            switch(orient){
            case C:
            case N:
            case S:
            case W:
              maxx = pt.x+symbol.size()/2*angstromsPerChar;
              break;
            case E:
              maxx = pt.x+symbol.size()*angstromsPerChar;
              break;
            }
          }

          if(pt.y<=miny){
            miny = pt.y-1.5*angstromsPerChar;
          }
          if(pt.y>=maxy){
            maxy = pt.y+angstromsPerChar;
          }
        } else {
          minx=std::min(pt.x,minx);
          miny=std::min(pt.y,miny);
          maxx=std::max(pt.x,maxx);
          maxy=std::max(pt.y,maxy);
        }
      }
      double dimx=(maxx-minx),dimy=(maxy-miny);
      res.push_back(BOUNDS);
      res.push_back(static_cast<ElementType>(dotsPerAngstrom*0));
      res.push_back(static_cast<ElementType>(dotsPerAngstrom*0));
      res.push_back(static_cast<ElementType>(dotsPerAngstrom*dimx));
      res.push_back(static_cast<ElementType>(dotsPerAngstrom*dimy));

      // loop over atoms:
      boost::tie(bAts,eAts)=mol.getVertices();
      while(bAts!=eAts){
        int a1Idx=mol[*bAts]->getIdx();
        RDGeom::Point2D a1(locs[a1Idx].x-minx,locs[a1Idx].y-miny);
        ROMol::OEDGE_ITER nbr,endNbrs;
        RDGeom::Point2D nbrSum(0,0);
        boost::tie(nbr,endNbrs) = mol.getAtomBonds(mol[*bAts].get());
        while(nbr!=endNbrs){
          const BOND_SPTR bond=mol[*nbr];
          ++nbr;
          int a2Idx=bond->getOtherAtomIdx(a1Idx);
          int lineWidth=1;
          if(highlightAtoms
             &&
             std::find(highlightAtoms->begin(),highlightAtoms->end(),a1Idx)
              != highlightAtoms->end()
             &&
             std::find(highlightAtoms->begin(),highlightAtoms->end(),a2Idx)
              != highlightAtoms->end() ){
            lineWidth=3;
          }
          RDGeom::Point2D a2(locs[a2Idx].x-minx,locs[a2Idx].y-miny);
          nbrSum+=a2-a1;
          if(a2Idx<a1Idx) continue;

          // draw bond from a1 to a2.
          int atnum1 = mol[*bAts]->getAtomicNum();
          int atnum2 = mol.getAtomWithIdx(a2Idx)->getAtomicNum();

          if( !mol.getRingInfo()->numBondRings(bond->getIdx()) && bond->getBondType()!=Bond::AROMATIC ) {
            // acyclic bonds
            RDGeom::Point2D obv=a2-a1;
            RDGeom::Point2D perp=obv;
            perp.rotate90();
            perp.normalize();

            if( bond->getBondType()==Bond::DOUBLE || bond->getBondType()==Bond::TRIPLE)  {
              RDGeom::Point2D startP=a1,endP=a2;              
              if( bond->getBondType()==Bond::TRIPLE){
                perp *= dblBondOffset;
                startP+=(obv*(1.-dblBondLengthFrac)/2);
              endP-=(obv*(1.-dblBondLengthFrac)/2);
              } else {
                perp *= 0.5 * dblBondOffset;
              }
              detail::drawLine( res , atnum1 , atnum2 , lineWidth , 0 ,
                                dotsPerAngstrom*(startP.x+perp.x) ,
                                dotsPerAngstrom*(startP.y+perp.y) ,
                                dotsPerAngstrom*(endP.x+perp.x) ,
                                dotsPerAngstrom*(endP.y+perp.y) );
              if(bond->getBondType() != Bond::AROMATIC){
                detail::drawLine( res , atnum1 , atnum2 , lineWidth , 0 ,
                                  dotsPerAngstrom*(startP.x-perp.x) ,
                                  dotsPerAngstrom*(startP.y-perp.y) ,
                                  dotsPerAngstrom*(endP.x-perp.x) ,
                                  dotsPerAngstrom*(endP.y-perp.y) );
              } else {
                detail::drawLine( res , atnum1 , atnum2 , lineWidth , 1 ,
                                  dotsPerAngstrom*(startP.x-perp.x) ,
                                  dotsPerAngstrom*(startP.y-perp.y) ,
                                  dotsPerAngstrom*(endP.x-perp.x) ,
                                  dotsPerAngstrom*(endP.y-perp.y) );

              }
            }
            if( bond->getBondType()==Bond::SINGLE || bond->getBondType()==Bond::TRIPLE ) {
              detail::drawLine( res , atnum1 , atnum2 , lineWidth , 0 ,
                                dotsPerAngstrom*(a1.x) ,
                                dotsPerAngstrom*(a1.y) ,
                                dotsPerAngstrom*(a2.x) ,
                                dotsPerAngstrom*(a2.y) );
            } else if( bond->getBondType()!=Bond::DOUBLE ) {
              detail::drawLine( res , atnum1 , atnum2 , lineWidth , 2 ,
                                dotsPerAngstrom*(a1.x) ,
                                dotsPerAngstrom*(a1.y) ,
                                dotsPerAngstrom*(a2.x) ,
                                dotsPerAngstrom*(a2.y) );
            }
          } else {
            // cyclic bonds
              detail::drawLine( res , atnum1 , atnum2 , lineWidth , 0 ,
                                dotsPerAngstrom*a1.x ,
                                dotsPerAngstrom*a1.y ,
                                dotsPerAngstrom*a2.x ,
                                dotsPerAngstrom*a2.y );

            if(bond->getBondType()==Bond::DOUBLE ||
               bond->getBondType()==Bond::AROMATIC ||
               bond->getBondType()==Bond::TRIPLE ){
              RDGeom::Point2D obv=a2-a1;
              RDGeom::Point2D perp=obv;
              perp.rotate90();
              perp.normalize();

              if( (bond->getBondType()==Bond::DOUBLE ||
                   bond->getBondType()==Bond::AROMATIC) &&
                  mol.getRingInfo()->numBondRings(bond->getIdx())){
                // we're in a ring... we might need to flip sides:
                ROMol::OEDGE_ITER nbr2,endNbrs2;
                boost::tie(nbr2,endNbrs2) = mol.getAtomBonds(mol[*bAts].get());
                while(nbr2!=endNbrs2){
                  const BOND_SPTR bond2=mol[*nbr2];
                  ++nbr2;
                  if(bond2->getIdx()==bond->getIdx() ||
                     !mol.getRingInfo()->numBondRings(bond2->getIdx())) continue;
                  bool sharedRing=false;
                  BOOST_FOREACH(const INT_VECT &ring,mol.getRingInfo()->bondRings()){
                    if(std::find(ring.begin(),ring.end(),bond->getIdx())!=ring.end() &&
                       std::find(ring.begin(),ring.end(),bond2->getIdx())!=ring.end()){
                      sharedRing=true;
                      break;
                    }
                  }
                  if(sharedRing){
                    // these two bonds share a ring.
                    int a3Idx=bond2->getOtherAtomIdx(a1Idx);
                    if(a3Idx!=a2Idx){
                      RDGeom::Point2D a3(locs[a3Idx].x-minx,locs[a3Idx].y-miny);
                      RDGeom::Point2D obv2=a3-a1;
                      if(obv2.dotProduct(perp)<0){
                        perp*=-1;
                      }
                    }
                  }
                }
              }
              perp *= dblBondOffset;

              RDGeom::Point2D offsetStart=a1 + obv*(.5*(1.-dblBondLengthFrac));

              obv *= dblBondLengthFrac;

              detail::drawLine( res , atnum1 , atnum2 , lineWidth , (bond->getBondType()==Bond::AROMATIC),
                                dotsPerAngstrom*(offsetStart.x+perp.x) ,
                                dotsPerAngstrom*(offsetStart.y+perp.y) ,
                                dotsPerAngstrom*(offsetStart.x+obv.x+perp.x) ,
                                dotsPerAngstrom*(offsetStart.y+obv.y+perp.y) );
            }
          }
        }
        std::string symbol;
        OrientType orient;
        boost::tie(symbol,orient)=atomSymbols[a1Idx];
        if(symbol!="" || includeAtomCircles ){
          res.push_back(ATOM);
          res.push_back(mol[*bAts]->getAtomicNum());
          res.push_back(static_cast<ElementType>(dotsPerAngstrom*a1.x));
          res.push_back(static_cast<ElementType>(dotsPerAngstrom*a1.y));
          res.push_back(static_cast<ElementType>(symbol.length()));
          if(symbol.length()){
            BOOST_FOREACH(char c, symbol){
              res.push_back(static_cast<ElementType>(c));
            }
          }
          res.push_back(static_cast<ElementType>(orient));
        }

        ++bAts;
      }

      return res;
    }

    std::vector<int> MolToDrawing(const RDKit::ROMol &mol,const std::vector<int> *highlightAtoms=0,
                                  bool kekulize=true,bool includeAtomCircles=false){
      RDKit::RWMol *cp = new RDKit::RWMol(mol);
      if(kekulize){
        try{
          RDKit::MolOps::Kekulize(*cp);
        } catch (...) {
          delete cp;
          cp = new RDKit::RWMol(mol);
        }
      }
      if(!mol.getNumConformers()) {
        RDDepict::compute2DCoords(*cp);
      }
      std::vector<int> drawing=DrawMol(*cp,-1,highlightAtoms,includeAtomCircles);
      delete cp;
      return drawing;
    }

  }   // end of namespace Drawing
} // end of namespace RDKit

#endif
