/*
 * Project: MoleCuilder
 * Description: creates and alters molecular systems
 * Copyright (C)  2013 University of Bonn. All rights reserved.
 * 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/>.
 */

/*
 * CompoundPotential.cpp
 *
 *  Created on: May 8, 2013
 *      Author: heber
 */

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

#include "CodePatterns/MemDebug.hpp"

#include "Potentials/CompoundPotential.hpp"

#include <algorithm>
#include <boost/bind.hpp>
#include <boost/foreach.hpp>
#include <boost/lambda/lambda.hpp>
#include <iterator>
#include <numeric>

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

#include "Element/element.hpp"
#include "Fragmentation/Homology/HomologyGraph.hpp"
#include "Fragmentation/Summation/SetValues/Fragment.hpp"
#include "FunctionApproximation/Extractors.hpp"
#include "Potentials/EmpiricalPotential.hpp"
#include "Potentials/helpers.hpp"
#include "Potentials/PotentialRegistry.hpp"


CompoundPotential::CompoundPotential(const HomologyGraph &graph)
{
  LOG(1, "INFO: Creating CompoundPotential for graph " << graph << ".");
  // look though graph and place all matching FunctionModel's in
  // PotentialRegistry in models
  PotentialRegistry::const_iterator potentialiter =
      PotentialRegistry::getInstance().getBeginIter();
  while (potentialiter != PotentialRegistry::getInstance().getEndIter()) {
    // get model and types
    EmpiricalPotential * const potential = potentialiter->second;
    const SerializablePotential::ParticleTypes_t &types =
        potential->getParticleTypes();

    // create charges
    Fragment::charges_t charges;
    charges.resize(types.size());
    std::transform(types.begin(), types.end(),
        charges.begin(), boost::lambda::_1);
    // convert into count map
    Extractors::elementcounts_t counts_per_charge =
        Extractors::_detail::getElementCounts(charges);
//    ASSERT( !counts_per_charge.empty(),
//        "getFirstGraphwithSpecifiedElements() - charge counts are empty?");
    LOG(2, "DEBUG: counts_per_charge is " << counts_per_charge << ".");

    // check whether graph contains suitable types
    Extractors::elementcounts_t::const_iterator countiter = counts_per_charge.begin();
    for (; countiter != counts_per_charge.end(); ++countiter)
      if (!graph.hasGreaterEqualTimesAtomicNumber(
          static_cast<size_t>(countiter->first),
          static_cast<size_t>(countiter->second))
          )
        break;
    // if we have a match for every count, store model
    if( countiter == counts_per_charge.end()) {
      LOG(1, "INFO: Potential " << potentialiter->first << " matches with fragment.");
      models.push_back(static_cast<FunctionModel*>(potential));
      particletypes_per_model.push_back(types);
    }
    ++potentialiter;
  }

  // check that models and particletypes_per_model match
  ASSERT( models.size() == particletypes_per_model.size(),
      "CompoundPotential::CompoundPotential() - particletypes not stored for all models?");
}

CompoundPotential::~CompoundPotential()
{
  // clear all models and internally stored particletypes
  models.clear();
  particletypes_per_model.clear();
}

void CompoundPotential::setParameters(const parameters_t &_params)
{
  size_t dim = _params.size();
  parameters_t::const_iterator iter = _params.begin();
  BOOST_FOREACH( FunctionModel* model, models) {
    const size_t model_dim = model->getParameterDimension();
    if (dim > 0) {
      parameters_t subparams;
      if (dim < model_dim) {
        std::copy(iter, iter+dim, std::back_inserter(subparams));
        iter += dim;
        dim = 0;
      } else {
        std::copy(iter, iter+model_dim, std::back_inserter(subparams));
        iter += model_dim;
        dim -= model_dim;
      }
      model->setParameters(subparams);
    }
  }

#ifndef NDEBUG
  parameters_t check_params(getParameters());
  check_params.resize(_params.size()); // truncate to same size
  ASSERT( check_params == _params,
      "CompoundPotential::setParameters() - failed, mismatch in to be set "
      +toString(_params)+" and set "+toString(check_params)+" params.");
#endif
}

CompoundPotential::parameters_t CompoundPotential::getParameters() const
{
  const size_t dimension = getParameterDimension();
  CompoundPotential::parameters_t parameters(dimension);
  CompoundPotential::parameters_t::iterator iter = parameters.begin();
  BOOST_FOREACH( const FunctionModel* model, models) {
    const CompoundPotential::parameters_t &params = model->getParameters();
    ASSERT( iter != parameters.end(),
        "CompoundPotential::getParameters() - iter already at end.");
    iter = std::copy(params.begin(), params.end(), iter);
  }
  ASSERT( iter == parameters.end(),
      "CompoundPotential::getParameters() - iter not at end.");
  return parameters;
}

void CompoundPotential::setParametersToRandomInitialValues(const TrainingData &data)
{
  std::for_each(models.begin(), models.end(),
      boost::bind(&FunctionModel::setParametersToRandomInitialValues, _1, boost::cref(data))
  );
}

size_t CompoundPotential::getParameterDimension() const
{
  std::vector<size_t> dimensions(models.size(), 0);
  std::transform(models.begin(), models.end(), dimensions.begin(),
      boost::bind(&FunctionModel::getParameterDimension, _1));
  return std::accumulate(dimensions.begin(), dimensions.end(), 0, std::plus<size_t>());
}

void CompoundPotential::setTriplefunction(triplefunction_t &_triplefunction)
{
  std::for_each(models.begin(), models.end(),
      boost::bind(&FunctionModel::setTriplefunction, _1, boost::ref(_triplefunction))
  );
}

bool CompoundPotential::areValidArguments(
    const SerializablePotential::ParticleTypes_t &_types,
    const arguments_t &args) const
{
  // /this function does much the same as Extractors::reorderArgumentsByParticleTypes()
  typedef std::list< argument_t > ListArguments_t;
  ListArguments_t availableList(args.begin(), args.end());

  /// 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 (SerializablePotential::ParticleTypes_t::const_iterator firstiter = _types.begin();
      firstiter != _types.end();
      ++firstiter) {
    for (SerializablePotential::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).
      // we remove the matching argument to make sure we don't pick it twice
      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)) {
          availableList.erase(iter);
          break;
        }
        else if ((iter->types.first == *seconditer)
              && (iter->types.second == *firstiter)) {
          availableList.erase(iter);
          break;
        }
      }
      if ( iter == availableList.end())
        return false;
    }
  }

  return true;
}

CompoundPotential::arguments_by_model_t CompoundPotential::splitUpArgumentsByModelsFilter(
    const arguments_t &arguments) const
{
  arguments_by_model_t partial_args;
  // go through each model and have it filter out its arguments, this already
  // returns a list of tuples associated with the specific model
  for(models_t::const_iterator modeliter = models.begin();
      modeliter != models.end(); ++modeliter) {
    FunctionModel::filter_t filterfunction = (*modeliter)->getSpecificFilter();
    list_of_arguments_t tempargs = filterfunction(arguments);
    // then split up all the bunches, too.
    for (list_of_arguments_t::const_iterator argiter = tempargs.begin();
        argiter != tempargs.end(); ++argiter) {
      const arguments_t &args = *argiter;
      partial_args.push_back(
          std::make_pair(
              *modeliter,
              args
            )
          );
    }
  }

  return partial_args;
}

CompoundPotential::arguments_by_model_t CompoundPotential::splitUpArgumentsByModels(
    const list_of_arguments_t &listarguments) const
{
  arguments_by_model_t partial_args;
  particletypes_per_model_t::const_iterator typesiter = particletypes_per_model.begin();
  models_t::const_iterator modeliter = models.begin();

  /// add constant model (which is always first model) with empty args if present
  if (typesiter->empty()) {
    partial_args.push_back(
        std::pair<FunctionModel *, arguments_t>(*modeliter, arguments_t())
        );
    ++modeliter;
    ++typesiter;
  }

  // then check other models
  /// we only have to check whether the current model still matches or whether
  /// have to use the next model.
  for (list_of_arguments_t::const_iterator argiter = listarguments.begin();
      argiter != listarguments.end(); ++argiter) {
    const arguments_t &arguments = *argiter;
    if (typesiter+1 != particletypes_per_model.end()) {
      // check whether next argument bunch is for same model or different one
      // we extract both partial_arguments, if the latter fits, we take the latter.
      const SerializablePotential::ParticleTypes_t &types = *typesiter;
      const SerializablePotential::ParticleTypes_t &nexttypes = *(typesiter+1);

      // we always expect N(N-1)/2 distances for N particle types
      // check first from sizes alone
      const size_t tuplesize = types.size()*(types.size()-1)/2;
      const size_t nexttuplesize = nexttypes.size()*(nexttypes.size()-1)/2;
      if ((tuplesize != nexttuplesize)) {
        if ((arguments.size() == tuplesize) &&  areValidArguments(types, arguments)) {
          // only former still matches, don't increment
          partial_args.push_back(
              std::make_pair(*modeliter, arguments)
              );
        } else if ((arguments.size() == nexttuplesize) &&  areValidArguments(nexttypes, arguments)) {
          // latter matches, increment
          ++typesiter;
          partial_args.push_back(
              std::make_pair(*(++modeliter), arguments)
              );
        } else {
          ASSERT(0,
              "CompoundPotential::splitUpArgumentsByModels() - neither this model nor next model match (size) with current tuple.");
        }
      } else { // same size, now we have to check the types individually
        size_t encodeValidity = 0;
        encodeValidity += 1*areValidArguments(types, arguments);
        encodeValidity += 2*areValidArguments(nexttypes, arguments);

        switch (encodeValidity) {
          case 1:
            // only former still matches, don't increment
            partial_args.push_back(
                std::make_pair(*modeliter, arguments)
                );
            break;
          case 2:
            ++typesiter;
            partial_args.push_back(
                std::make_pair(*(++modeliter), arguments)
                );
            break;
          case 0:
          case 3:
          default:
            ASSERT(0,
                "CompoundPotential::splitUpArgumentsByModels() - neither this model nor next model match (type) with current tuple.");
            break;
        }
      }
    } else {
      const SerializablePotential::ParticleTypes_t &types = *typesiter;
      if (areValidArguments(types, arguments)) {
        // only former matches, don't increment
        partial_args.push_back(
            std::make_pair(*modeliter, arguments)
            );
      } else {
        ASSERT(0,
            "CompoundPotential::splitUpArgumentsByModels() - last model does not match with current tuple.");
      }
    }
  }

  return partial_args;
}

CompoundPotential::results_t CompoundPotential::operator()(
    const list_of_arguments_t &listarguments) const
{
  /// first, we have to split up the given arguments
  arguments_by_model_t partial_args =
      splitUpArgumentsByModels(listarguments);
  // print split up argument list for debugging
  if (DoLog(4)) {
    LOG(4, "Arguments by model are: ");
    for(arguments_by_model_t::const_iterator iter = partial_args.begin();
        iter != partial_args.end(); ++iter) {
      LOG(4, "\tModel with " << iter->first->getParameterDimension()
          << " parameters " << iter->first->getParameters()
          << " and arguments: " << iter->second);
    }
  }

  /// then, with each bunch of arguments, we call the specific model
  results_t results(1,0.);
  std::vector<results_t> partial_results;
  for(arguments_by_model_t::const_iterator iter = partial_args.begin();
      iter != partial_args.end(); ++iter) {
    partial_results.push_back(
        (*iter->first)(
            FunctionModel::list_of_arguments_t(1, iter->second))
    );
  }
  // print partial results for debugging
  if (DoLog(4)) {
    std::stringstream output;
    output << "Partial results are: ";
    std::for_each(partial_results.begin(), partial_results.end(),
        output << (boost::lambda::_1)[0] << "\t");
    LOG(4, output.str());
  }

  /// Finally, sum up all results and return
  std::for_each(partial_results.begin(), partial_results.end(),
      results[0] += (boost::lambda::_1)[0]);
  return results;
}

CompoundPotential::results_t CompoundPotential::parameter_derivative(
    const list_of_arguments_t &listarguments,
    const size_t index) const
{
  // first, we have to split up the given arguments
  arguments_by_model_t partial_args =
      splitUpArgumentsByModels(listarguments);
  // then, with each bunch of arguments, we call the specific model
  // get parameter dimensions per model
  std::vector<size_t> dimensions(models.size(), 0);
  std::transform(models.begin(), models.end(), dimensions.begin(),
      boost::bind(&FunctionModel::getParameterDimension, _1));

  // convert to index end+1 per model
  std::partial_sum(dimensions.begin(), dimensions.end(), dimensions.begin());

  // look for first value greater than index
  std::vector<size_t>::const_iterator iter =
      std::upper_bound(dimensions.begin(), dimensions.end(), index);

  // step forward to same model
  models_t::const_iterator modeliter = models.begin();
  std::advance(modeliter,
      std::distance(const_cast<const std::vector<size_t> &>(dimensions).begin(), iter) );

  CompoundPotential::results_t returnresults;
  for(arguments_by_model_t::const_iterator argiter = partial_args.begin();
      argiter != partial_args.end(); ++argiter) {
    const FunctionModel *model = argiter->first;

    // for every matching model evaluate
    if (model == *modeliter) {
      // evaluate with correct relative index and return
      const size_t indexbase = (iter == dimensions.begin()) ? 0 : *(iter-1);
      CompoundPotential::results_t results =
          model->parameter_derivative(
              FunctionModel::list_of_arguments_t(1, argiter->second), index-indexbase);

      // either set results or add
      if (returnresults.empty())
        returnresults = results;
      else
        std::transform(
            results.begin(), results.end(),
            returnresults.begin(),
            returnresults.begin(),
            std::plus<FunctionModel::result_t>());
    }
  }

  return returnresults;
}

bool CompoundPotential::isBoxConstraint() const
{
  std::vector<bool> constraints(models.size(), 0);
  std::transform(models.begin(), models.end(), constraints.begin(),
      boost::bind(&FunctionModel::getParameterDimension, _1));
  return std::accumulate(constraints.begin(), constraints.end(), true,
      std::logical_and<bool>());
}

CompoundPotential::parameters_t CompoundPotential::getLowerBoxConstraints() const
{
  const size_t dimension = getParameterDimension();
  CompoundPotential::parameters_t constraints(dimension);
  CompoundPotential::parameters_t::iterator iter = constraints.begin();
  BOOST_FOREACH( FunctionModel* model, models) {
    const CompoundPotential::parameters_t params = model->getLowerBoxConstraints();
    ASSERT( iter != constraints.end(),
        "CompoundPotential::getLowerBoxConstraints() - iter already at end.");
    iter = std::copy(params.begin(), params.end(), iter);
  }
  ASSERT( iter == constraints.end(),
      "CompoundPotential::getLowerBoxConstraints() - iter not at end.");
  return constraints;
}

CompoundPotential::parameters_t CompoundPotential::getUpperBoxConstraints() const
{
  const size_t dimension = getParameterDimension();
  CompoundPotential::parameters_t constraints(dimension);
  CompoundPotential::parameters_t::iterator iter = constraints.begin();
  BOOST_FOREACH( FunctionModel* model, models) {
    const CompoundPotential::parameters_t params = model->getUpperBoxConstraints();
    ASSERT( iter != constraints.end(),
        "CompoundPotential::getUpperBoxConstraints() - iter already at end.");
    iter = std::copy(params.begin(), params.end(), iter);
  }
  ASSERT( iter == constraints.end(),
      "CompoundPotential::getUpperBoxConstraints() - iter not at end.");
  return constraints;
}

FunctionModel::filter_t CompoundPotential::getSpecificFilter() const
{
  // we must concatenate all filtered arguments here
  // create initial returnfunction
  FunctionModel::filter_t returnfunction =
      boost::bind(&Helpers::returnEmptyListArguments);

  // every following fragments combines its arguments with the initial function
  for (models_t::const_iterator modeliter = models.begin();
      modeliter != models.end(); ++modeliter) {
    returnfunction =
          boost::bind(&Extractors::concatenateListOfArguments,
              boost::bind(returnfunction, _1),
              boost::bind((*modeliter)->getSpecificFilter(), _1)
        );
  }
  return returnfunction;
}

size_t CompoundPotential::getSpecificArgumentCount() const
{
  std::vector<size_t> argument_counts(models.size(), 0);
  std::transform(models.begin(), models.end(), argument_counts.begin(),
      boost::bind(&FunctionModel::getSpecificArgumentCount, _1));
  return std::accumulate(argument_counts.begin(), argument_counts.end(), 0,
      std::plus<size_t>());
}