/*
 * 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 "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"


Interfragmenter::atomkeyset_t Interfragmenter::getAtomKeySetMap(
    size_t _MaxOrder
    ) const
{
  /// create a map of atom to keyset (below equal MaxOrder)
  atomkeyset_t atomkeyset;
  LOG(1, "INFO: Placing all atoms and their keysets into a map.");
  for (Graph::const_iterator keysetiter = TotalGraph.begin();
      keysetiter != TotalGraph.end(); ++keysetiter) {
    const KeySet &keyset = keysetiter->first;
    LOG(2, "DEBUG: Current keyset is " << keyset);
    const AtomIdSet atoms(keyset);
    const size_t atoms_size = atoms.getAtomIds().size();
    if ((atoms_size > _MaxOrder) || (atoms_size == 0))
      continue;
    for (AtomIdSet::const_iterator atomiter = atoms.begin();
        atomiter != atoms.end(); ++atomiter) {
      // either create new list ...
      std::pair<atomkeyset_t::iterator, bool> inserter =
          atomkeyset.insert( std::make_pair(*atomiter, keysets_t(1, &keyset) ));
      // ... or push to end
      if (inserter.second) {
        LOG(3, "DEBUG: Created new entry in map.");
      } else {
        LOG(3, "DEBUG: Added keyset to present entry.");
        inserter.first->second.push_back(&keyset);
      }
    }
  }
  LOG(2, "DEBUG: There are " << atomkeyset.size() << " entries in lookup.");

  return atomkeyset;
}

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;
}

Interfragmenter::atomkeyset_t Interfragmenter::getCandidatesSpecificKeySetMap(
    const candidates_t &_candidates,
    const atomkeyset_t &_atomkeyset) const
{
  atomkeyset_t fragmentmap;
  for (candidates_t::const_iterator candidateiter = _candidates.begin();
      candidateiter != _candidates.end(); ++candidateiter) {
    const atom * _atom = *candidateiter;
    atomkeyset_t::const_iterator iter = _atomkeyset.find(_atom);
    ASSERT( iter != _atomkeyset.end(),
        "Interfragmenter::getAtomSpecificKeySetMap() - could not find atom "
        +toString(_atom->getId())+" in lookup.");
    fragmentmap.insert( std::make_pair( _atom, iter->second) );
  }
  LOG(4, "DEBUG: Copied part of lookup map contains " << fragmentmap.size() << " keys.");

  return fragmentmap;
}

void Interfragmenter::combineFragments(
    const size_t _MaxOrder,
    const candidates_t &_candidates,
    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::iterator finditer = _fragmentmap.find(_atom);
    ASSERT( finditer != _fragmentmap.end(),
        "Interfragmenter::combineFragments() - could not find atom "
        +toString(_atom->getId())+" in fragment specific lookup.");
    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()(
    const size_t MaxOrder,
    const double Rcut,
    const enum HydrogenTreatment treatment)
{
  atomkeyset_t atomkeyset = getAtomKeySetMap(MaxOrder);

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

  /// go through all fragments up to MaxOrder
  LOG(1, "INFO: Creating inter-fragments.");
  for (Graph::const_iterator keysetiter = TotalGraph.begin();
      keysetiter != TotalGraph.end(); ++keysetiter) {
    const KeySet &keyset = keysetiter->first;
    LOG(2, "DEBUG: Current keyset is " << keyset);
    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
    atomkeyset_t fragmentmap = getCandidatesSpecificKeySetMap(candidates, atomkeyset);

    /// 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 = TotalGraph.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);

  atomkeyset_t atomkeyset = getAtomKeySetMap(_MaxOrder);

  // go through each atom and find closest atom not in the same keyset
  for (Graph::const_iterator keysetiter = TotalGraph.begin();
      keysetiter != TotalGraph.end(); ++keysetiter) {
    const KeySet &keyset = keysetiter->first;
    LOG(2, "DEBUG: Current keyset is " << keyset);
    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;
}