/*
 * Project: MoleCuilder
 * Description: creates and alters molecular systems
 * Copyright (C)  2012 University of Bonn. 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/foreach.hpp>

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

#include "LinearAlgebra/Vector.hpp"

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

using namespace boost::assign;

FunctionModel::arguments_t
Extractors::gatherAllDistanceArguments(
    const Fragment::positions_t& positions,
    const Fragment::charges_t& charges,
    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 = positions.begin();//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(
          charges[ std::distance(positions.begin(), firstpositer) ],
          charges[ std::distance(positions.begin(), secondpositer) ]
          );
      arg.indices = std::make_pair(
          std::distance(
              positions.begin(), firstpositer),
          std::distance(
              positions.begin(), secondpositer)
          );
      arg.globalid = globalid;
      result.push_back(arg);
    }
  }

  return result;
}

FunctionModel::arguments_t
Extractors::gatherAllSymmetricDistanceArguments(
    const Fragment::positions_t& positions,
    const Fragment::charges_t& charges,
    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(
          charges[ std::distance(positions.begin(), firstpositer) ],
          charges[ std::distance(positions.begin(), secondpositer) ]
          );
      arg.indices = std::make_pair(
          std::distance(
              positions.begin(), firstpositer),
          std::distance(
              positions.begin(), secondpositer)
          );
      arg.globalid = globalid;
      result.push_back(arg);
    }
  }

  return result;
}

Fragment::positions_t Extractors::_detail::gatherPositionsFromTargets(
    const Fragment::positions_t& positions,
    const Fragment::charges_t& charges,
    const chargeiters_t &targets
    )
{
  Fragment::positions_t filtered_positions;
  for (chargeiters_t::const_iterator firstpairiter = targets.begin();
      firstpairiter != targets.end(); ++firstpairiter) {
    Fragment::positions_t::const_iterator positer = positions.begin();
    const size_t steps = std::distance(charges.begin(), *firstpairiter);
    std::advance(positer, steps);
    filtered_positions.push_back(*positer);
  }
  return filtered_positions;
}

FunctionModel::arguments_t Extractors::_detail::gatherDistancesFromTargets(
    const Fragment::positions_t& positions,
    const Fragment::charges_t& charges,
    const chargeiters_t &targets,
    const size_t globalid
    )
{
  Fragment::positions_t filtered_positions;
  Fragment::charges_t filtered_charges;
  for (chargeiters_t::const_iterator firstpairiter = targets.begin();
      firstpairiter != targets.end(); ++firstpairiter) {
    Fragment::positions_t::const_iterator positer = positions.begin();
    const size_t steps = std::distance(charges.begin(), *firstpairiter);
    std::advance(positer, steps);
    filtered_positions.push_back(*positer);
    filtered_charges.push_back(**firstpairiter);
  }
  return Extractors::gatherAllSymmetricDistanceArguments(
      filtered_positions,
      filtered_charges,
      globalid);
}

Extractors::elementcounts_t 
Extractors::_detail::getElementCounts(
    const Fragment::charges_t elements
  )
{
  elementcounts_t elementcounts;
  for (Fragment::charges_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;
}

Extractors::elementtargets_t 
Extractors::_detail::convertElementcountsToTargets(
    const Fragment::charges_t &charges,
    const elementcounts_t &elementcounts
    )
{
  elementtargets_t elementtargets;
  for (elementcounts_t::const_iterator countiter = elementcounts.begin();
      countiter != elementcounts.end();
      ++countiter) {
    chargeiter_t chargeiter = charges.begin();
    const element_t &element = countiter->first;
    const count_t &count = countiter->second;
    for (count_t i = 0; i < count; ++i) {
      chargeiter_t tempiter = std::find(chargeiter, charges.end(), element);
      if (tempiter != charges.end()) {
        // try to insert new list
        std::pair< elementtargets_t::iterator, bool> inserter =
          elementtargets.insert( std::make_pair( countiter->first, chargeiters_t(1, tempiter)) );
        // if already present, append to it
        if (!inserter.second) {
          inserter.first->second.push_back(tempiter);
        } else { // if created, increase vector's reserve to known size
          inserter.first->second.reserve(countiter->second);
        }
        // search from this element onwards then
        chargeiter = ++tempiter;
      } else {
        //ELOG(1, "Could not find desired number of elements " << count << " in fragment.");
        return Extractors::elementtargets_t();
      }
    }
  }
  return elementtargets;
}

Extractors::elementtargets_t
Extractors::_detail::convertChargesToTargetMap(
    const Fragment::charges_t& charges,
    Fragment::charges_t elements
    )
{
  // place each charge into a map
  elementtargets_t completeelementtargets;
  for (chargeiter_t chargeiter = charges.begin();
      chargeiter != charges.end();
      ++chargeiter) {
    std::pair< elementtargets_t::iterator, bool> inserter =
        completeelementtargets.insert( std::make_pair( *chargeiter, chargeiters_t(1, chargeiter)) );
    // if already present, append to it
    if (!inserter.second) {
      inserter.first->second.push_back(chargeiter);
    }
  }
  // pick out desired charges only
  std::sort(elements.begin(), elements.end());
  Fragment::charges_t::iterator eraseiter =
      std::unique(elements.begin(), elements.end());
  elements.erase(eraseiter, elements.end());
  elementtargets_t elementtargets;
  for(Fragment::charges_t::const_iterator iter = elements.begin();
      iter != elements.end();
      ++iter) {
    elementtargets_t::const_iterator finditer = completeelementtargets.find(*iter);
    ASSERT( finditer != completeelementtargets.end(),
        "Extractors::_detail::convertChargesToTargetMap() - no element "+toString(*iter)+" present?");
    std::pair< elementtargets_t::iterator, bool> inserter =
        elementtargets.insert( std::make_pair( finditer->first, finditer->second) );
    ASSERT( inserter.second,
        "Extractors::_detail::convertChargesToTargetMap() - key twice?");
  }
  return elementtargets;
}


Extractors::chargeiters_t 
Extractors::_detail::realignElementtargets(
    const elementtargets_t &elementtargets,
    const Fragment::charges_t elements,
    const elementcounts_t &elementcounts
    )
{
  chargeiters_t targets;
  elementcounts_t counts; // how many chargeiters of this element have been used
  if (!elements.empty()) {  // skip if no elements given
    targets.reserve(elements.size());
    for (Fragment::charges_t::const_iterator elementiter = elements.begin();
        elementiter != elements.end(); ++elementiter) {
      const element_t &element = *elementiter;
      count_t &count = counts[element]; // if not present, std::map creates instances with default of 0
#ifndef NDEBUG
      {
        elementcounts_t::const_iterator testiter = elementcounts.find(element);
        ASSERT( (testiter != elementcounts.end()) && (count < testiter->second),
            "Extractors::_detail::realignElementTargets() - we want to use more chargeiters for element "
            +toString(element)+" than we counted initially.");
      }
#endif
      elementtargets_t::const_iterator targetiter = elementtargets.find(element);
      if (targetiter != elementtargets.end()) {
        const chargeiters_t &chargeiters = targetiter->second;
        const chargeiter_t &chargeiter = chargeiters[count++];
        targets.push_back(chargeiter);
      }
    }
  }
  return targets;
}

FunctionModel::arguments_t
Extractors::gatherAllDistancesFromFragment(
    const Fragment::positions_t& positions,
    const Fragment::charges_t& charges,
    const Fragment::charges_t elements,
    const size_t globalid
    )
{
  /// The main problem here is that we have to know how many same
  /// elements (but different atoms!) we are required to find. Hence,
  /// we first have to count same elements, then get different targets
  /// for each and then associated them in correct order back again.

  // 1. we have to place each charge into a map as unique chargeiter, i.e. map
  elementtargets_t elementtargets =
      Extractors::_detail::convertChargesToTargetMap(
          charges,
          elements);

  // 2. now we have to combine elementcounts out of elementtargets per desired element
  //    in a combinatorial fashion
  targets_per_combination_t combinations =
      Extractors::_detail::CombineChargesAndTargets(
          elements,
          elementtargets);

  // 3. finally, convert chargeiters into argument list
  FunctionModel::arguments_t args =
      Extractors::_detail::convertTargetsToArguments(
          positions,
          charges,
          combinations,
          globalid);

  return args;
}

Extractors::targets_per_combination_t
Extractors::_detail::CombineChargesAndTargets(
    const Fragment::charges_t& elements,
    const elementtargets_t& elementtargets
    )
{
  // recursively create all correct combinations of targets
  targets_per_combination_t combinations;
  chargeiters_t currenttargets;
  boost::function<void (const chargeiters_t &currenttargets)> addFunction =
      boost::bind(&targets_per_combination_t::push_back,
          boost::ref(combinations),
          _1);
  pickLastElementAsTarget(elements, elementtargets, currenttargets, addFunction);

  return combinations;
}

FunctionModel::arguments_t
Extractors::_detail::convertTargetsToArguments(
    const Fragment::positions_t& positions,
    const Fragment::charges_t& charges,
    const targets_per_combination_t combinations,
    const size_t globalid
    )
{
  FunctionModel::arguments_t args;
  // create arguments from each combination
  for (targets_per_combination_t::const_iterator iter = combinations.begin();
      iter != combinations.end();
      ++iter) {
    Fragment::positions_t partpositions;
    Fragment::charges_t partcharges;
    partpositions.reserve(iter->size());
    partcharges.reserve(iter->size());
    for(chargeiters_t::const_iterator targetiter = iter->begin();
        targetiter != iter->end();
        ++targetiter) {
      Fragment::positions_t::const_iterator positer = positions.begin();
      std::advance(positer, std::distance(charges.begin(), *targetiter));
      partpositions.push_back(*positer);
      partcharges.push_back(**targetiter);
    }
    FunctionModel::arguments_t partargs =
        gatherAllSymmetricDistanceArguments(partpositions, partcharges, globalid);
    args.insert(args.end(), partargs.begin(), partargs.end());
  }

  return args;
}

void
Extractors::_detail::pickLastElementAsTarget(
    Fragment::charges_t elements,
    elementtargets_t elementtargets,
    chargeiters_t &currenttargets,
    boost::function<void (const chargeiters_t &currenttargets)> &addFunction
    )
{
  // get last element from charges
  const Fragment::charge_t charge = elements.back();
  elements.pop_back();
  elementtargets_t::iterator iter = elementtargets.find(charge);
  if (iter == elementtargets.end())
    return;
  bool NotEmpty = !iter->second.empty();
  while (NotEmpty) {
    // get last target from the vector of chargeiters
    chargeiter_t target = iter->second.back();
    iter->second.pop_back();
    // remove this key if empty
    if (iter->second.empty()) {
      elementtargets.erase(iter);
      NotEmpty = false;
    }
    currenttargets.push_back(target);
    if (elements.empty()) {
      // call add function
      {
        std::stringstream targetstream;
        BOOST_FOREACH( chargeiter_t target, currenttargets ) {
          targetstream << " " << *target;
        }
        LOG(3, "DEBUG: Adding set" << targetstream.str() << ".");
      }
      addFunction(currenttargets);
    } else {
      // if not, call us recursively
      pickLastElementAsTarget(elements, elementtargets, currenttargets, addFunction);
    }
    // pop last in currenset again
    currenttargets.pop_back();
  }
}

Extractors::chargeiters_t
Extractors::_detail::gatherTargetsFromFragment(
    const Fragment::charges_t& charges,
    const Fragment::charges_t elements
    )
{
  /// The main problem here is that we have to know how many same
  /// elements (but different atoms!) we are required to find. Hence,
  /// we first have to count same elements, then get different targets
  /// for each and then associated them in correct order back again.

  // 1. we have to make elements unique with counts, hence convert to map
  elementcounts_t elementcounts =
      Extractors::_detail::getElementCounts(elements);

  // 2. then for each element we need as many targets (chargeiters) as counts
  elementtargets_t elementtargets =
      Extractors::_detail::convertElementcountsToTargets(charges, elementcounts);

  // 3. we go again through elements and use one found target for each count
  // in that order
  chargeiters_t targets =
      Extractors::_detail::realignElementtargets(elementtargets, elements, elementcounts);

#ifndef NDEBUG
  // check all for debugging
  for (chargeiters_t::const_iterator chargeiter = targets.begin();
      chargeiter != targets.end();
      ++chargeiter)
    ASSERT( *chargeiter != charges.end(),
        "Extractors::gatherTargetsFromFragment() - we have not found enough targets?!");
#endif

  return targets;
}

Fragment::positions_t
Extractors::gatherPositionsFromFragment(
    const Fragment::positions_t positions,
    const Fragment::charges_t charges,
    const Fragment::charges_t& elements
    )
{
  // 1.-3. gather correct charge positions
  chargeiters_t targets =
      Extractors::_detail::gatherTargetsFromFragment(charges, elements);
  // 4. convert position_t to Vector
  return Extractors::_detail::gatherPositionsFromTargets(
          positions,
          charges,
          targets);
}

FunctionModel::arguments_t
Extractors::gatherDistancesFromFragment(
    const Fragment::positions_t positions,
    const Fragment::charges_t charges,
    const Fragment::charges_t& elements,
    const size_t globalid
    )
{
  // 1.-3. gather correct charge positions
  chargeiters_t targets =
      Extractors::_detail::gatherTargetsFromFragment(charges, elements);
  // 4. convert position_t to Vector
  return Extractors::_detail::gatherDistancesFromTargets(
          positions,
          charges,
          targets,
          globalid);
}

FunctionModel::arguments_t Extractors::reorderArgumentsByIncreasingDistance(
    const FunctionModel::arguments_t &args
    )
{
  FunctionModel::arguments_t returnargs(args);
  std::sort(returnargs.begin(), returnargs.end(), argument_t::DistanceComparator);
  return returnargs;
}

FunctionModel::arguments_t Extractors::reorderArgumentsByParticleTypes(
    const FunctionModel::arguments_t &args,
    const ParticleTypes_t &_types
    )
{
  typedef std::list< argument_t > ListArguments_t;
  ListArguments_t availableList(args.begin(), args.end());
  FunctionModel::arguments_t returnargs;
  returnargs.reserve(args.size());

  // 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 letter. Hence, we start the second
  // iterator at the current position of the first one and skip the equal case.
  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;

      // 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();
      for (;iter != availableList.end(); ++iter) {
        LOG(3, "DEBUG: Current args is " << *iter << ".");
        if ((iter->types.first == *firstiter)
              && (iter->types.second == *seconditer)) {
          returnargs.push_back( *iter );
          break;
        }
        else if ((iter->types.first == *seconditer)
              && (iter->types.second == *firstiter)) {
          argument_t flippedtypes(*iter);
          std::swap( flippedtypes.indices.first, flippedtypes.indices.second );
          std::swap( flippedtypes.types.first, flippedtypes.types.second );
          returnargs.push_back( flippedtypes );
          break;
        }
      }
      ASSERT( iter != availableList.end(),
          "Extractors::reorderArgumentsByParticleTypes() - could not find arguments to "
          +toString(*firstiter)+","+toString(*seconditer)+".");
      availableList.erase(iter);
    }
  }
//  LOG(2, "DEBUG: Final list of args is " << returnargs << ".");

  return returnargs;
}