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

Copyright (C) 2016 by NextMove Software

This program is free software; you can redistribute it and/or modify

it under the terms of the GNU General Public License as published by

the Free Software Foundation version 2 of the License.

This program is distributed in the hope that it will be useful,

but WITHOUT ANY WARRANTY; without even the implied warranty of

MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the

GNU General Public License for more details.

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

#include <openbabel/babelconfig.h>

#include <openbabel/oberror.h>

#include <openbabel/mol.h>

#include <openbabel/atom.h>

#include <openbabel/bond.h>

#include <openbabel/obiter.h>

#include <openbabel/fingerprint.h>

#include <openbabel/obiter.h>

#include <openbabel/elements.h>

#include <vector>

#include <algorithm>

using namespace std;

namespace OpenBabel

{

/// \brief Fingerprint based on ECFP

class fingerprintECFP : public OBFingerprint

{

public:

  fingerprintECFP(const char* ID, bool IsDefault=false,

                  unsigned int radius=4, bool keepdups = true)

    : OBFingerprint(ID, IsDefault),

      _radius(radius), _keepdups(keepdups), _flags(0){};

  const char* Description() override

  {

          // Important! The second line is used by some output formats (e.g. FPS)

    // to determine the default size

    return "Extended-Connectivity Fingerprints (ECFPs)\n"

                 "4096 bits.\n"

                 "Circular topological fingerprints of specified radius\n";

  }

  //Calculates the fingerprint

  bool GetFingerprint(OBBase* pOb, vector<unsigned int>&fp, int nbits=0) override;

  /// \returns fragment info unless SetFlags(OBFingerprint::FPT_NOINFO) has been called before GetFingerprint() called.

  std::string DescribeBits(const std::vector<unsigned int> fp, bool bSet=true) override

  { return _ss.str(); }

  unsigned int Flags() override { return _flags; }

  void SetFlags(unsigned int f) override { _flags=f; }

private:

  std::vector<unsigned int> v;

  stringstream _ss;

  unsigned int _radius;

  bool _keepdups;

  unsigned int _flags;

};

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

//Make the global instances

  fingerprintECFP theECFP0("ECFP0",false, 0, true);

  fingerprintECFP theECFP2("ECFP2",false, 1, true);

  fingerprintECFP theECFP4("ECFP4",false, 2, true);

  fingerprintECFP theECFP6("ECFP6",false, 3, true);

  fingerprintECFP theECFP8("ECFP8",false, 4, true);

  fingerprintECFP theECFP10("ECFP10",false, 5, true);

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

#define mix32(a,b,c) \

{ \

  a -= b; a -= c; a ^= (c>>13); \

  b -= c; b -= a; b ^= (a<<8); \

  c -= a; c -= b; c ^= (b>>13); \

  a -= b; a -= c; a ^= (c>>12);  \

  b -= c; b -= a; b ^= (a<<16); \

  c -= a; c -= b; c ^= (b>>5); \

  a -= b; a -= c; a ^= (c>>3);  \

  b -= c; b -= a; b ^= (a<<10); \

  c -= a; c -= b; c ^= (b>>15); \

}

static unsigned int ECFPHash(unsigned char *ptr, unsigned int length)

{

  unsigned int a = 0;

  unsigned int b = 0;

  unsigned int c = 0;

  unsigned int len = length;

  while (len >= 12) {

    a += ptr[0] + ((unsigned int)ptr[1]<<8)

                + ((unsigned int)ptr[2]<<16)

                + ((unsigned int)ptr[3]<<24);

    b += ptr[4] + ((unsigned int)ptr[5]<<8)

                + ((unsigned int)ptr[6]<<16)

                + ((unsigned int)ptr[7]<<24);

    c += ptr[8] + ((unsigned int)ptr[9]<<8)

                + ((unsigned int)ptr[10]<<16)

                + ((unsigned int)ptr[11]<<24);

    mix32(a,b,c);

    ptr += 12;

    len -= 12;

  }

  c += length;

  switch (len) {

  case 11: c += ((unsigned int)ptr[10]<<24);

  case 10: c += ((unsigned int)ptr[9]<<16);

  case  9: c += ((unsigned int)ptr[8]<<8);

  case  8: b += ((unsigned int)ptr[7]<<24);

  case  7: b += ((unsigned int)ptr[6]<<16);

  case  6: b += ((unsigned int)ptr[5]<<8);

  case  5: b += ptr[4];

  case  4: a += ((unsigned int)ptr[3]<<24);

  case  3: a += ((unsigned int)ptr[2]<<16);

  case  2: a += ((unsigned int)ptr[1]<<8);

  case  1: a += ptr[0];

  }

  mix32(a,b,c);

  return c;

}

static unsigned int ECFPHash(std::vector<unsigned int> &v)

{

  unsigned int len = (unsigned int)v.size();

  unsigned int a = 0;

  unsigned int b = 0;

  unsigned int c = 0;

  unsigned int i = 0;

  while (i+2 < len) {

    a += v[i];

    b += v[i+1];

    c += v[i+2];

    mix32(a,b,c);

    i += 3;

  }

  c += len;

  switch (len - i) {

  case 2:  b += v[i+1];

  case 1:  a += v[i];

  }

  mix32(a,b,c);

  return c;

}

struct AtomInfo {

  unsigned int e[5]; // hashcodes for ecfp0, 2, 4... This will segfault at ECFP12

  std::vector<unsigned int> b[4]; // for duplicate removal as described in paper?

};

struct NborInfo {

  unsigned int order;

  unsigned int idx;

  NborInfo(unsigned int o, unsigned int i) : order(o), idx(i) {}

  bool operator < (const NborInfo &x) const

  {

    if (order != x.order)

      return order < x.order;

    return idx < x.idx;

  }

};

static void ECFPPass(OpenBabel::OBMol &mol,

                     AtomInfo *ainfo, unsigned int pass)

{

  FOR_ATOMS_OF_MOL(atom, mol) {

    if (atom->GetAtomicNum() == OBElements::Hydrogen)

      continue;

    OpenBabel::OBAtom* aptr = &(*atom);

    unsigned int idx = aptr->GetIdx();

    AtomInfo *ptr = &ainfo[idx];

    std::vector<NborInfo> nbrs;

    FOR_BONDS_OF_ATOM(bptr, aptr) {

      OpenBabel::OBAtom* nptr = bptr->GetNbrAtom(aptr);

      if (nptr->GetAtomicNum() == OBElements::Hydrogen)

        continue;

      unsigned int order;

      if (!bptr->IsAromatic()) {

        switch (bptr->GetBondOrder()) {

        case 2:  order = 2;  break;

        case 3:  order = 3;  break;

        default: order = 1;

        }

      } else order = 4;

      unsigned int nidx = nptr->GetIdx();

      nbrs.push_back(NborInfo(order,ainfo[nidx].e[pass-1]));

      // for duplicate removal as described in paper (?)

      if (pass == 1)

        ptr->b[pass-1].push_back(bptr->GetIdx());

    }

    std::sort(nbrs.begin(),nbrs.end());

    std::vector<unsigned int> vint;

    vint.push_back(pass);

    vint.push_back(ptr->e[pass-1]);

    std::vector<NborInfo>::const_iterator ni;

    for (ni=nbrs.begin(); ni!=nbrs.end(); ++ni) {

      vint.push_back(ni->order);

      vint.push_back(ni->idx);

    }

    ptr->e[pass] = ECFPHash(vint);

  }

}

static void ECFPFirstPass(OpenBabel::OBMol &mol,

                          AtomInfo *ainfo)

{

  unsigned char buffer[8];

  /* First Pass: ECFP_0 */

  FOR_ATOMS_OF_MOL(atom, mol) {

    if (atom->GetAtomicNum() == OBElements::Hydrogen)

      continue;

    OpenBabel::OBAtom* aptr = &(*atom);

    unsigned int idx = aptr->GetIdx();

    buffer[0] = aptr->GetHvyDegree(); // degree of heavy atom connections

    buffer[1] = aptr->GetExplicitValence() - aptr->ExplicitHydrogenCount(); // valence of heavy atom connections

    buffer[2] = aptr->GetAtomicNum();

    buffer[3] = (unsigned char)aptr->GetIsotope();

    buffer[4] = (unsigned char)aptr->GetFormalCharge();

    buffer[5] = (unsigned char)(aptr->ExplicitHydrogenCount() + aptr->GetImplicitHCount());

    buffer[6] = aptr->IsInRing() ? 1 : 0;

    buffer[7] = 0;  // aptr->IsAromatic() ? 1 : 0;

    ainfo[idx].e[0] = ECFPHash(buffer,8);

  }

}

bool fingerprintECFP::GetFingerprint(OBBase* pOb, vector<unsigned int>&fp, int nbits)

{

  OBMol* pmol = dynamic_cast<OBMol*>(pOb);

  if(!pmol) return false;

  // default fingeprint size

  if (nbits <= 0)

    nbits = 4096;

  fp.resize(0); // clear without deallocating memory

  fp.resize(nbits/Getbitsperint());

  _ss.str("");

  unsigned int pass;

  unsigned int count = pmol->NumAtoms();

  if (count == 0) return true;

  // Access this using the Atom::Idx()

  AtomInfo *ainfo = new AtomInfo[count+1];

  ECFPFirstPass(*pmol,ainfo);

  for (pass=1; pass<= _radius; pass++)

    ECFPPass(*pmol,ainfo,pass);

  // Duplicate removal - this is a simplified version of what's in the paper

  FOR_ATOMS_OF_MOL(atom, pmol) {

    if (atom->GetAtomicNum() == OBElements::Hydrogen)

      continue;    

    unsigned int idx = atom->GetIdx();

    for (pass=0; pass <= _radius; pass++) {

      unsigned int bit = (ainfo[idx].e[pass] % nbits) & 0x7fffffff; 

      SetBit(fp, bit);

    }

  }

  delete[] ainfo;

  return true;

}

} //namespace OpenBabel

//! \file fingerecfp.cpp

//! \brief ECFP fingerprint definition and implementation