/*
 * 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 <utility>
#include <vector>
#include <boost/assign.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::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
  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);
    ASSERT (targetiter != elementtargets.end(),
        "Extractors::_detail::realignElementTargets() - not enough chargeiters for element "
        +toString(element)+".");
    const chargeiters_t &chargeiters = targetiter->second;
    const chargeiter_t &chargeiter = chargeiters[count++];
    targets.push_back(chargeiter);
  }
  return targets;
}

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