/*
 * Project: MoleCuilder
 * Description: creates and alters molecular systems
 * Copyright (C)  2013 University of Bonn. 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/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.hasTimesAtomicNumber(
          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 parameters_t &model_params = model->getParameters();
    const size_t model_dim = model_params.size();
    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);
    }
  }
}

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();
    std::copy(params.begin(), params.end(), iter);
    ASSERT( iter != parameters.end(),
        "CompoundPotential::getParameters() - iter already 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))
  );
}

std::vector<CompoundPotential::arguments_t> CompoundPotential::splitUpArgumentsByModels(
    const arguments_t &arguments) const
{
  std::vector<arguments_t> partial_args(models.size());
  std::vector<arguments_t>::iterator partialiter = partial_args.begin();
  arguments_t::const_iterator argiter = arguments.begin();
  BOOST_FOREACH( const SerializablePotential::ParticleTypes_t &types, particletypes_per_model) {
    // we always expect N(N-1)/2 distances for N particle types
    arguments_t::const_iterator enditer = argiter+(types.size()*(types.size()-1)/2);
    ASSERT( argiter != arguments.end(),
        "CompoundPotential::splitUpArgumentsByModels() - incorrect number of arguments.");
    std::copy(argiter, enditer, std::back_inserter(*partialiter++));
    argiter = enditer;
  }
  ASSERT( argiter == arguments.end(),
      "CompoundPotential::splitUpArgumentsByModels() - incorrect number of arguments.");
  return partial_args;
}

CompoundPotential::results_t CompoundPotential::operator()(const arguments_t &arguments) const
{
  // first, we have to split up the given arguments
  std::vector<arguments_t> partial_args =
      splitUpArgumentsByModels(arguments);
  // then, with each bunch of arguments, we call the specific model
  results_t results(1,0.);
  std::vector<results_t> partial_results(models.size());
  std::transform(
      models.begin(), models.end(),
      partial_args.begin(),
      partial_results.begin(),
      boost::bind(&FunctionModel::operator(), _1, _2)
  );
  std::for_each(partial_results.begin(), partial_results.end(),
      std::cout << (boost::lambda::_1)[0] << "\t");
  std::for_each(partial_results.begin(), partial_results.end(),
      results[0] += (boost::lambda::_1)[0]);
  return results;
}

CompoundPotential::results_t CompoundPotential::parameter_derivative(const arguments_t &arguments, const size_t index) const
{
  // first, we have to split up the given arguments
  std::vector<arguments_t> partial_args =
      splitUpArgumentsByModels(arguments);
  // 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 value not less than index
  std::vector<size_t>::const_iterator iter =
      std::lower_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) );
  std::vector<arguments_t>::const_iterator argiter = partial_args.begin();
  std::advance(argiter,
      std::distance(const_cast<const std::vector<size_t> &>(dimensions).begin(), iter) );

  // evaluate with correct relative index and return
  const size_t indexbase = (iter == dimensions.begin()) ? 0 : *(--iter);
  CompoundPotential::results_t results =
      (*modeliter)->parameter_derivative(*argiter, index-indexbase);
  return results;
}

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();
    std::copy(params.begin(), params.end(), iter);
    ASSERT( iter != constraints.end(),
        "CompoundPotential::getParameters() - iter already 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();
    std::copy(params.begin(), params.end(), iter);
    ASSERT( iter != constraints.end(),
        "CompoundPotential::getParameters() - iter already at end.");
  }
  return constraints;
}

FunctionModel::extractor_t CompoundPotential::getFragmentSpecificExtractor() const
{
  // create initial returnfunction
  particletypes_per_model_t::const_iterator typesiter = particletypes_per_model.begin();
  Fragment::charges_t charges;
  {
    charges.resize((*typesiter).size());
    std::transform((*typesiter).begin(), (*typesiter).end(),
        charges.begin(), boost::lambda::_1);
  }
  FunctionModel::extractor_t returnfunction =
    boost::bind(&Extractors::gatherAllDistancesFromFragment,
        boost::bind(&Fragment::getPositions, _1),
        boost::bind(&Fragment::getCharges, _1),
        charges,  // no crefs here as are temporaries!
        _2);

  // every following fragments combines its arguments with the initial function
  for (; typesiter != particletypes_per_model.end(); ++typesiter) {
    Fragment::charges_t charges;
    {
      charges.resize((*typesiter).size());
      std::transform((*typesiter).begin(), (*typesiter).end(),
          charges.begin(), boost::lambda::_1);
    }

    returnfunction =
          boost::bind(&Extractors::concatenateArguments,
              boost::bind(returnfunction, _1, _2),
              boost::bind(&Extractors::gatherAllDistancesFromFragment,
                  boost::bind(&Fragment::getPositions, _1),
                  boost::bind(&Fragment::getCharges, _1),
                  charges,  // no crefs here as are temporaries!
                  _2)
        );
  }
  return returnfunction;
}