/*
 * Project: MoleCuilder
 * Description: creates and alters molecular systems
 * Copyright (C)  2012 University of Bonn. All rights reserved.
 * Copyright (C)  2013 Frederik Heber. All rights reserved.
 * Please see the COPYING file or "Copyright notice" in builder.cpp for details.
 *
 *
 *   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/>.
 */

/*
 * Extractors.cpp
 *
 *  Created on: 15.10.2012
 *      Author: heber
 */

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

#include "CodePatterns/MemDebug.hpp"

#include <sstream>
#include <utility>
#include <vector>
#include <boost/assign.hpp>
#include <boost/bind.hpp>
#include <boost/foreach.hpp>

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

#include "LinearAlgebra/Vector.hpp"

#include "FunctionApproximation/Extractors.hpp"
#include "FunctionApproximation/FunctionArgument.hpp"

#include "Fragmentation/Homology/HomologyGraph.hpp"

using namespace boost::assign;

FunctionModel::arguments_t
Extractors::gatherAllSymmetricDistanceArguments(
    const Fragment::positions_t& positions,
    const Fragment::atomicnumbers_t& atomicnumbers,
    const size_t globalid)
{
  FunctionModel::arguments_t result;

  // go through current configuration and gather all other distances
  Fragment::positions_t::const_iterator firstpositer = positions.begin();
  for (;firstpositer != positions.end(); ++firstpositer) {
    Fragment::positions_t::const_iterator secondpositer = firstpositer;
    for (; secondpositer != positions.end(); ++secondpositer) {
      if (firstpositer == secondpositer)
        continue;
      argument_t arg;
      const Vector firsttemp((*firstpositer)[0],(*firstpositer)[1],(*firstpositer)[2]);
      const Vector secondtemp((*secondpositer)[0],(*secondpositer)[1],(*secondpositer)[2]);
      arg.distance = firsttemp.distance(secondtemp);
      arg.types = std::make_pair(
          (int)atomicnumbers[ std::distance(positions.begin(), firstpositer) ],
          (int)atomicnumbers[ std::distance(positions.begin(), secondpositer) ]
          );
      arg.indices = std::make_pair(
          std::distance(
              positions.begin(), firstpositer),
          std::distance(
              positions.begin(), secondpositer)
          );
      arg.globalid = globalid;
      LOG(3, "DEBUG: Created argument " << arg << ".");
      result.push_back(arg);
    }
  }

  return result;
}

Extractors::elementcounts_t 
Extractors::_detail::getElementCounts(
    const Fragment::atomicnumbers_t elements
  )
{
  elementcounts_t elementcounts;
  for (Fragment::atomicnumbers_t::const_iterator elementiter = elements.begin();
      elementiter != elements.end(); ++elementiter) {
    // insert new element
    std::pair< elementcounts_t::iterator, bool> inserter =
        elementcounts.insert( std::make_pair( *elementiter, 1) );
    // if already present, just increase its count
    if (!inserter.second)
      ++(inserter.first->second);
  }
  return elementcounts;
}

struct ParticleTypesComparator {
  bool operator()(const argument_t::types_t &a, const argument_t::types_t &b)
  {
    if (a.first < a.second) {
      if (b.first < b.second) {
        if (a.first < b.first)
          return true;
        else if (a.first > b.first)
          return false;
        else
          return (a.second < b.second);
      } else {
        if (a.first < b.second)
          return true;
        else if (a.first > b.second)
          return false;
        else
          return (a.second < b.first);
      }
    } else {
      if (b.first < b.second) {
        if (a.second < b.first)
          return true;
        else if (a.second > b.first)
          return false;
        else
          return (a.first < b.second);
      } else {
        if (a.second < b.second)
          return true;
        else if (a.second > b.second)
          return false;
        else
          return (a.first < b.first);
      }
    }
  }
};

std::ostream& operator<<(std::ostream &out, const argument_t::types_t &a)
{
  out << "[" << a.first << "," << a.second << "]";
  return out;
}

FunctionModel::list_of_arguments_t Extractors::reorderArgumentsByParticleTypes(
    const FunctionModel::list_of_arguments_t &listargs,
    const HomologyGraph &_graph,
    const ParticleTypes_t &_types,
    const HomologyGraph &_bindingmodel
    )
{
  FunctionModel::list_of_arguments_t returnargs;
  for (FunctionModel::list_of_arguments_t::const_iterator iter = listargs.begin();
      iter != listargs.end(); ++iter) {
    const FunctionModel::arguments_t &args = *iter;
    /// We  place all arguments into multimap according to particle type pair.
    // here, we need a special comparator such that types in key pair are always
    // properly ordered.
    typedef std::multimap<
        argument_t::types_t,
        argument_t,
        ParticleTypesComparator> TypePair_Argument_Map_t;
    TypePair_Argument_Map_t argument_map;
    for(FunctionModel::arguments_t::const_iterator iter = args.begin();
        iter != args.end(); ++iter) {
      argument_map.insert( std::make_pair(iter->types, *iter) );
    }
    LOG(4, "DEBUG: particle_type map is " << argument_map << ".");

    /// Then, we create the desired unique keys
    typedef std::vector<argument_t::types_t> UniqueTypes_t;
    UniqueTypes_t UniqueTypes;
    for (ParticleTypes_t::const_iterator firstiter = _types.begin();
        firstiter != _types.end();
        ++firstiter) {
      for (ParticleTypes_t::const_iterator seconditer = firstiter;
          seconditer != _types.end();
          ++seconditer) {
        if (seconditer == firstiter)
          continue;
        UniqueTypes.push_back( std::make_pair(*firstiter, *seconditer) );
      }
    }
    LOG(4, "DEBUG: Created unique types as keys " << UniqueTypes << ".");

    /// Finally, we use the unique key list to pick corresponding arguments from the map
    FunctionModel::arguments_t sortedargs;
    sortedargs.reserve(args.size());
    while (!argument_map.empty()) {
      // note that particle_types_t may be flipped, i.e. 1,8 is equal to 8,1, but we
      // must maintain the correct order in indices in accordance with the order
      // in _types, i.e. 1,8,1 must match with e.g. ids 1,0,2 where 1 has type 1,
      // 0 has type 8, and 2 has type 2.
      // In other words: We do not want to flip/modify arguments such that they match
      // with the specific type pair we seek but then this comes at the price that we
      // have flip indices when the types in a pair are flipped.

      typedef std::vector<size_t> indices_t;
      //!> here, we gather the indices as we discover them
      indices_t indices;
      indices.resize(_types.size(), (size_t)-1);

      // these are two iterators that create index pairs in the same way as we have
      // created type pairs. If a -1 is still present in indices, then the index is
      // still arbitrary but is then set by the next found index
      indices_t::iterator firstindex = indices.begin();
      indices_t::iterator secondindex = firstindex+1;

      //!> here, we gather the current bunch of arguments as we find them
      FunctionModel::arguments_t argumentbunch;
      argumentbunch.reserve(UniqueTypes.size());

      for (UniqueTypes_t::const_iterator typeiter = UniqueTypes.begin();
          typeiter != UniqueTypes.end(); ++typeiter) {
        // have all arguments to same type pair as list within the found range
        std::pair<
            TypePair_Argument_Map_t::iterator,
            TypePair_Argument_Map_t::iterator> range_t =
                argument_map.equal_range(*typeiter);
        LOG(4, "DEBUG: Set of arguments to current key [" << typeiter->first << ","
            << typeiter->second << "] is " << std::list<argument_t>(
                MapValueIterator<TypePair_Argument_Map_t::iterator>(range_t.first),
                MapValueIterator<TypePair_Argument_Map_t::iterator>(range_t.second)
                ) << ".");
        // the first key is always easy and is pivot which the rest has to be associated to
        if (typeiter == UniqueTypes.begin()) {
          const argument_t & arg = range_t.first->second;
          if ((typeiter->first == arg.types.first) && (typeiter->second == arg.types.second)) {
            // store in correct order
            *firstindex = arg.indices.first;
            *secondindex = arg.indices.second;
          } else {
            // store in flipped order
            *firstindex = arg.indices.second;
            *secondindex = arg.indices.first;
          }
          argumentbunch.push_back(arg);
          argument_map.erase(range_t.first);
          LOG(4, "DEBUG: Gathered first argument " << arg << ".");
        } else {
          // go through the range and pick the first argument matching the index constraints
          for (TypePair_Argument_Map_t::iterator argiter = range_t.first;
              argiter != range_t.second; ++argiter) {
            // seconditer may be -1 still
            const argument_t &arg = argiter->second;
            if (arg.indices.first == *firstindex) {
              if ((arg.indices.second == *secondindex) || (*secondindex == (size_t)-1)) {
                if (*secondindex == (size_t)-1)
                  *secondindex = arg.indices.second;
                argumentbunch.push_back(arg);
                argument_map.erase(argiter);
                LOG(4, "DEBUG: Gathered another argument " << arg << ".");
                break;
              }
            } else if ((arg.indices.first == *secondindex) || (*secondindex == (size_t)-1)) {
              if (arg.indices.second == *firstindex) {
                if (*secondindex == (size_t)-1)
                  *secondindex = arg.indices.first;
                argumentbunch.push_back(arg);
                argument_map.erase(argiter);
                LOG(4, "DEBUG: Gathered another (flipped) argument " << arg << ".");
                break;
              }
            }
          }
        }
        // move along in indices and check bounds
        ++secondindex;
        if (secondindex == indices.end()) {
          ++firstindex;
          if (firstindex != indices.end()-1)
            secondindex = firstindex+1;
        }
      }
      ASSERT( (firstindex == indices.end()-1) && (secondindex == indices.end()),
          "Extractors::reorderArgumentsByParticleTypes() - we have not gathered enough arguments.");
      ASSERT( argumentbunch.size() == UniqueTypes.size(),
          "Extractors::reorderArgumentsByParticleTypes() - we have not gathered enough arguments.");
      // place bunch of arguments in return args
      LOG(3, "DEBUG: Given types " << _types << " and found indices " << indices << ".");
      LOG(3, "DEBUG: Final bunch of arguments is " << argumentbunch << ".");
      sortedargs.insert(sortedargs.end(), argumentbunch.begin(), argumentbunch.end());
    }
    returnargs.push_back(sortedargs);
  }

  return returnargs;
}

FunctionModel::list_of_arguments_t Extractors::filterArgumentsByParticleTypes(
    const FunctionModel::arguments_t &args,
    const HomologyGraph &_graph,
    const ParticleTypes_t &_types,
    const HomologyGraph &_bindingmodel
    )
{
  typedef std::list< argument_t > ListArguments_t;
  ListArguments_t availableList(args.begin(), args.end());
  LOG(2, "DEBUG: Initial list of args is " << args << ".");


  // TODO: fill a lookup map such that we don't have O(M^3) scaling, if M is number
  // of types (and we always must have M(M-1)/2 args) but O(M^2 log(M)). However, as
  // M is very small (<=3), this is not necessary fruitful now.
//  typedef ParticleTypes_t firsttype;
//  typedef ParticleTypes_t secondtype;
//  typedef std::map< firsttype, std::map< secondtype, boost::ref(args) > > ArgsLookup_t;
//  ArgsLookup_t ArgsLookup;

  // basically, we have two choose any two pairs out of types but only those
  // where the first is less than the latter. Hence, we start the second
  // iterator at the current position of the first one and skip the equal case.
  FunctionModel::arguments_t allargs;
  allargs.reserve(args.size());
  for (ParticleTypes_t::const_iterator firstiter = _types.begin();
      firstiter != _types.end();
      ++firstiter) {
    for (ParticleTypes_t::const_iterator seconditer = firstiter;
        seconditer != _types.end();
        ++seconditer) {
      if (seconditer == firstiter)
        continue;
      LOG(3, "DEBUG: Looking for (" << *firstiter << "," << *seconditer << ") in all args.");

      // search the right one in _args (we might allow switching places of
      // firstiter and seconditer, as distance is symmetric).
      ListArguments_t::iterator iter = availableList.begin();
      while (iter != availableList.end()) {
        LOG(4, "DEBUG: Current args is " << *iter << ".");
        if ((iter->types.first == *firstiter)
              && (iter->types.second == *seconditer)) {
          allargs.push_back( *iter );
          iter = availableList.erase(iter);
          LOG(4, "DEBUG: Accepted argument.");
        } else if ((iter->types.first == *seconditer)
              && (iter->types.second == *firstiter)) {
          allargs.push_back( *iter );
          iter = availableList.erase(iter);
          LOG(4, "DEBUG: Accepted (flipped) argument.");
        } else {
          ++iter;
          LOG(4, "DEBUG: Rejected argument.");
        }
      }
    }
  }
  LOG(2, "DEBUG: Final list of args is " << allargs << ".");

  // first, we bring together tuples of distances that belong together
  FunctionModel::list_of_arguments_t singlelist_allargs;
  singlelist_allargs.push_back(allargs);
  FunctionModel::list_of_arguments_t sortedargs =
      reorderArgumentsByParticleTypes(singlelist_allargs, _graph, _types, _bindingmodel);
  ASSERT( sortedargs.size() == (size_t)1,
      "Extractors::filterArgumentsByParticleTypes() - reordering did not generate a single list.");
  // then we split up the tuples of arguments and place each into single list
  FunctionModel::list_of_arguments_t returnargs;
  FunctionModel::arguments_t::const_iterator argiter = sortedargs.begin()->begin();
  const size_t num_types = _types.size();
  const size_t args_per_tuple = num_types * (num_types-1) / 2;
  while (argiter != sortedargs.begin()->end()) {
    FunctionModel::arguments_t currenttuple(args_per_tuple);
    const FunctionModel::arguments_t::const_iterator startiter = argiter;
    std::advance(argiter, args_per_tuple);
#ifndef NDEBUG
    FunctionModel::arguments_t::const_iterator endoutiter =
#endif
        std::copy(startiter, argiter, currenttuple.begin());
    ASSERT( endoutiter == currenttuple.end(),
        "Extractors::filterArgumentsByParticleTypes() - currenttuple not initialized to right size.");
    returnargs.push_back(currenttuple);
  }

  LOG(2, "DEBUG: We have generated " << returnargs.size() << " tuples of distances.");

  return returnargs;
}


FunctionModel::arguments_t Extractors::combineArguments(
    const FunctionModel::arguments_t &firstargs,
    const FunctionModel::arguments_t &secondargs)
{
  FunctionModel::arguments_t args = concatenateArguments(firstargs, secondargs);
  std::sort(args.begin(), args.end(),
      boost::bind(&argument_t::operator<, _1, _2));
  FunctionModel::arguments_t::iterator iter =
      std::unique(args.begin(), args.end(), 
          boost::bind(&argument_t::operator==, _1, _2));
  args.erase(iter, args.end());
  return args;
}

FunctionModel::arguments_t Extractors::concatenateArguments(
    const FunctionModel::arguments_t &firstargs,
    const FunctionModel::arguments_t &secondargs)
{
  FunctionModel::arguments_t args(firstargs);
  args.insert(args.end(), secondargs.begin(), secondargs.end());
  return args;
}

FunctionModel::list_of_arguments_t Extractors::concatenateListOfArguments(
    const FunctionModel::list_of_arguments_t &firstlistargs,
    const FunctionModel::list_of_arguments_t &secondlistargs)
{
  FunctionModel::list_of_arguments_t listargs(firstlistargs);
  listargs.insert(listargs.end(), secondlistargs.begin(), secondlistargs.end());
  return listargs;
}