/*
 * Project: MoleCuilder
 * Description: creates and alters molecular systems
 * Copyright (C)  2013 Frederik Heber. All rights reserved.
 * 
 *
 *   This file is part of MoleCuilder.
 *
 *    MoleCuilder 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, either version 2 of the License, or
 *    (at your option) any later version.
 *
 *    MoleCuilder 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.
 *
 *    You should have received a copy of the GNU General Public License
 *    along with MoleCuilder.  If not, see <http://www.gnu.org/licenses/>.
 */

/*
 * Interfragmenter.cpp
 *
 *  Created on: Jul 5, 2013
 *      Author: heber
 */

// include config.h
#ifdef HAVE_CONFIG_H
#include <config.h>
#endif

#include "CodePatterns/MemDebug.hpp"

#include "Interfragmenter.hpp"

#include <algorithm>

#include "CodePatterns/Assert.hpp"
#include "CodePatterns/Log.hpp"

#include "LinearAlgebra/Vector.hpp"

#include "Descriptors/AtomDescriptor.hpp"
#include "Element/element.hpp"
#include "Fragmentation/Graph.hpp"
#include "Fragmentation/KeySet.hpp"
#include "LinkedCell/LinkedCell_View.hpp"
#include "LinkedCell/types.hpp"
#include "World.hpp"

static Vector getAtomIdSetCenter(
    const AtomIdSet &_atoms)
{
  const molecule * const _mol = (*_atoms.begin())->getMolecule();
  const size_t atoms_size = _atoms.getAtomIds().size();
  Vector center;
  for (AtomIdSet::const_iterator iter = _atoms.begin();
      iter != _atoms.end(); ++iter) {
    center += (*iter)->getPosition();
    ASSERT ( _mol == (*iter)->getMolecule(),
        "getAtomIdSetCenter() - ids in same keyset belong to different molecule.");
  }
  center *= 1./(double)atoms_size;

  return center;
}

Interfragmenter::candidates_t Interfragmenter::getNeighborsOutsideMolecule(
    const AtomIdSet &_atoms,
    double _Rcut,
    const enum HydrogenTreatment _treatment) const
{
  /// go through linked cell and get all neighboring atoms up to Rcut
  const LinkedCell::LinkedCell_View view = World::getInstance().getLinkedCell(_Rcut);
  const Vector center = getAtomIdSetCenter(_atoms);
  const LinkedCell::LinkedList neighbors = view.getAllNeighbors(_Rcut, center);
  LOG(4, "DEBUG: Obtained " << neighbors.size() << " neighbors in distance of "
      << _Rcut << " around " << center << ".");

  /// remove all atoms that belong to same molecule as does one of the
  /// fragment's atoms
  const molecule * const _mol = (*_atoms.begin())->getMolecule();
  candidates_t candidates;
  candidates.reserve(neighbors.size());
  for (LinkedCell::LinkedList::const_iterator iter = neighbors.begin();
      iter != neighbors.end(); ++iter) {
    const atom * const _atom = static_cast<const atom * const >(*iter);
    ASSERT( _atom != NULL,
        "Interfragmenter::getNeighborsOutsideMolecule() - a neighbor is not actually an atom?");
    if ((_atom->getMolecule() != _mol)
        && (_atom->getPosition().DistanceSquared(center) < _Rcut*_Rcut)
        && ((_treatment == IncludeHydrogen) || (_atom->getType()->getAtomicNumber() != 1))) {
      candidates.push_back(_atom);
    }
  }
  LOG(3, "DEBUG: There remain " << candidates.size() << " candidates.");

  return candidates;
}

void Interfragmenter::combineFragments(
    const size_t _MaxOrder,
    const candidates_t &_candidates,
    const atomkeyset_t &_fragmentmap,
    const KeySet &_keyset,
    Graph &_InterFragments,
    int &_counter)
{
  for (candidates_t::const_iterator candidateiter = _candidates.begin();
      candidateiter != _candidates.end(); ++candidateiter) {
    const atom *_atom = *candidateiter;
    LOG(3, "DEBUG: Current candidate is " << *_atom << ".");
    atomkeyset_t::const_iterator finditer = _fragmentmap.find(_atom->getId());
    ASSERT( finditer != _fragmentmap.end(),
        "Interfragmenter::combineFragments() - could not find atom "
        +toString(_atom->getId())+" in fragment specific lookup.");
    // copy set to allow erase
    keysets_t othersets(finditer->second);
    ASSERT( !othersets.empty(),
        "Interfragmenter::combineFragments() - keysets to "+toString(_atom->getId())+
        "is empty.");
    keysets_t::iterator otheriter = othersets.begin();
    while (otheriter != othersets.end()) {
      const KeySet &otherset = *otheriter;
      LOG(3, "DEBUG: Current keyset is " << otherset << ".");
      // only add them one way round and not the other
      if (otherset < _keyset) {
        ++otheriter;
        continue;
      }
      // only add if combined they don't exceed the desired maxorder
      if (otherset.size() + _keyset.size() > _MaxOrder) {
        LOG(3, "INFO: Rejecting " << otherset << " as in sum their orders exceed " << _MaxOrder);
        ++otheriter;
        continue;
      }
      KeySet newset(otherset);
      newset.insert(_keyset.begin(), _keyset.end());
      LOG(3, "DEBUG: Inserting new combined set " << newset << ".");
      _InterFragments.insert( std::make_pair(newset, std::make_pair(++_counter, 1.)));
      // finally, remove the set such that no other combination exists
      otheriter = othersets.erase(otheriter);
    }
  }
}

void Interfragmenter::operator()(
    Graph &TotalGraph,
    const size_t MaxOrder,
    const double Rcut,
    const enum HydrogenTreatment treatment)
{
  AtomFragmentsMap& atomfragments_container = AtomFragmentsMap::getInstance();
  const atomkeyset_t &atomkeyset = atomfragments_container.getMap();

  Graph InterFragments;
  int counter = atomkeyset.size();

  /// go through all fragments up to MaxOrder
  LOG(1, "INFO: Creating inter-fragments.");
  for (atomkeyset_t::const_iterator atomiter = atomkeyset.begin();
      atomiter != atomkeyset.end(); ++atomiter) {
    const atomId_t &atomid = atomiter->first;
    LOG(2, "DEBUG: Current atomid is " << atomid);

    const AtomFragmentsMap::keysets_t &keysets = atomiter->second;
    for (AtomFragmentsMap::keysets_t::const_iterator keyiter = keysets.begin();
        keyiter != keysets.end(); ++keyiter) {
      const KeySet &keyset = *keyiter;
      const AtomIdSet atoms(keyset);
      const size_t atoms_size = atoms.getAtomIds().size();
      if ((atoms_size > MaxOrder) || (atoms_size == 0))
        continue;

      // get neighboring atoms outside the current molecule
      candidates_t candidates = getNeighborsOutsideMolecule(atoms, Rcut, treatment);

      // create a lookup specific to this fragment
      std::vector<atomId_t> atomids(candidates.size());
      std::transform(
          candidates.begin(), candidates.end(),
          atomids.begin(),
          boost::bind(&atom::getId, _1));
      atomkeyset_t fragmentmap = atomfragments_container.getMap(atomids, MaxOrder);

      /// combine each remaining fragment with current fragment to a new fragment
      /// if keyset is less (to prevent addding same inter-fragment twice)
      combineFragments(MaxOrder, candidates, fragmentmap, keyset, InterFragments, counter);
    }
  }

  /// eventually, add all new fragments to the Graph
  counter = atomkeyset.size();
  TotalGraph.InsertGraph(InterFragments, counter);
}

double Interfragmenter::findLargestCutoff(
    const size_t _MaxOrder,
    const double _upperbound,
    const enum HydrogenTreatment _treatment) const
{
  double Rcut = _upperbound*_upperbound;
  std::pair<atomId_t, atomId_t> ClosestPair;
  // place all atoms into LC grid with some upper bound
  const LinkedCell::LinkedCell_View view = World::getInstance().getLinkedCell(_upperbound);

  // go through each atom and find closest atom not in the same keyset
  AtomFragmentsMap& atomfragments_container = AtomFragmentsMap::getInstance();
  const atomkeyset_t &atomkeyset = atomfragments_container.getMap();
  for (atomkeyset_t::const_iterator atomiter = atomkeyset.begin();
      atomiter != atomkeyset.end(); ++atomiter) {
    const atomId_t &atomid = atomiter->first;
    LOG(2, "DEBUG: Current atomid is " << atomid);

    const AtomFragmentsMap::keysets_t &keysets = atomiter->second;
    for (AtomFragmentsMap::keysets_t::const_iterator keyiter = keysets.begin();
        keyiter != keysets.end(); ++keyiter) {
      const KeySet &keyset = *keyiter;
      const AtomIdSet atoms(keyset);

      // get neighboring atoms outside the current molecule
      const candidates_t candidates = getNeighborsOutsideMolecule(atoms, _upperbound, _treatment);
      const Vector center = getAtomIdSetCenter(atoms);

      for (candidates_t::const_iterator candidateiter = candidates.begin();
          candidateiter != candidates.end(); ++candidateiter) {
        atom const * const Walker = *candidateiter;
        // go through each atom in set and pick minimal distance
        for (AtomIdSet::const_iterator setiter = atoms.begin(); setiter != atoms.end(); ++setiter) {
          const double distanceSquared = Walker->getPosition().DistanceSquared(center);
          // pick the smallest compared to current Rcut if smaller
          if (distanceSquared < Rcut) {
            Rcut = distanceSquared;
            ClosestPair.first = (*setiter)->getId();
            ClosestPair.second = Walker->getId();
            LOG(2, "DEBUG: Found new pair " << ClosestPair << " with distance " << sqrt(Rcut));
          }
        }
      }
    }
  }
  const double largest_distance = sqrt(Rcut);
  LOG(1, "INFO: Largest inter-fragment distance to cause no additional fragments: "
      << largest_distance);

  return largest_distance;
}