/*
 * 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/>. 
 */

/*
 * SaturationPotential.cpp
 *
 *  Created on: Oct 11, 2012
 *      Author: heber
 */


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

#include "CodePatterns/MemDebug.hpp"

#include "SaturationPotential.hpp"

#include <boost/assign.hpp>
#include <boost/assign/list_of.hpp> // for 'map_list_of()'
#include <iostream>
#include <string>

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

#include "FunctionApproximation/Extractors.hpp"
#include "FunctionApproximation/TrainingData.hpp"
#include "Potentials/helpers.hpp"
#include "Potentials/ParticleTypeCheckers.hpp"

class Fragment;

using namespace boost::assign;

// static definitions
const SaturationPotential::ParameterNames_t
SaturationPotential::ParameterNames =
      boost::assign::list_of<std::string>
      ("all_energy_offset")
      ("")
      ("")
      ("")
      ("")
      ("")
    ;
const std::string SaturationPotential::potential_token("saturation");

SaturationPotential::SaturationPotential(
    const ParticleTypes_t &_ParticleTypes) :
  SerializablePotential(_ParticleTypes),
  morse(_ParticleTypes),
  angle(symmetrizeTypes(_ParticleTypes)),
  energy_offset(0.)
{
  // have some decent defaults for parameter_derivative checking
  // Morse and Angle have their own defaults, offset is set
  ASSERT( _ParticleTypes.size() == (size_t)2,
      "SaturationPotential::SaturationPotential() - exactly two types must be given.");
  ASSERT( _ParticleTypes[1] == 1,
      "SaturationPotential::SaturationPotential() - second type must be hydrogen.");
}

SaturationPotential::SaturationPotential(
    const ParticleTypes_t &_ParticleTypes,
    const double _all_energy_offset,
    const double _morse_spring_constant,
    const double _morse_equilibrium_distance,
    const double _morse_dissociation_energy,
    const double _angle_spring_constant,
    const double _angle_equilibrium_distance) :
  SerializablePotential(_ParticleTypes),
  morse(_ParticleTypes),
  angle(symmetrizeTypes(_ParticleTypes)),
  energy_offset(_all_energy_offset)
{
  ASSERT( _ParticleTypes.size() == (size_t)2,
      "SaturationPotential::SaturationPotential() - exactly two types must be given.");
  ASSERT( _ParticleTypes[1] == 1,
      "SaturationPotential::SaturationPotential() - second type must be hydrogen.");
  parameters_t morse_params(morse.getParameterDimension());
  morse_params[PairPotential_Morse::spring_constant] = _morse_spring_constant;
  morse_params[PairPotential_Morse::equilibrium_distance] = _morse_equilibrium_distance;
  morse_params[PairPotential_Morse::dissociation_energy] = _morse_dissociation_energy;
  morse_params[PairPotential_Morse::energy_offset] = 0.;
  morse.setParameters(morse_params);
  parameters_t angle_params(angle.getParameterDimension());
  angle_params[PairPotential_Angle::spring_constant] = _angle_spring_constant;
  angle_params[PairPotential_Angle::equilibrium_distance] = _angle_equilibrium_distance;
  angle_params[PairPotential_Angle::energy_offset] = 0.;
  angle.setParameters(angle_params);
}

void SaturationPotential::setParameters(const parameters_t &_params)
{
  const size_t paramsDim = _params.size();
  ASSERT( paramsDim <= getParameterDimension(),
      "SaturationPotential::setParameters() - we need not more than "
      +toString(getParameterDimension())+" parameters.");
//    LOG(1, "INFO: Setting new SaturationPotential params: " << _params);


  // offsets
  if (paramsDim > all_energy_offset)
    energy_offset = _params[all_energy_offset];

  // Morse
  {
    parameters_t morse_params(morse.getParameters());
    if (paramsDim > morse_spring_constant)
      morse_params[PairPotential_Morse::spring_constant] = _params[morse_spring_constant];
    if (paramsDim > morse_equilibrium_distance)
      morse_params[PairPotential_Morse::equilibrium_distance] = _params[morse_equilibrium_distance];
    if (paramsDim > morse_dissociation_energy)
      morse_params[PairPotential_Morse::dissociation_energy] = _params[morse_dissociation_energy];
    morse_params[PairPotential_Morse::energy_offset] = 0.;
    morse.setParameters(morse_params);
  }

  // Angle
  {
    parameters_t angle_params(angle.getParameters());
    if (paramsDim > angle_spring_constant)
      angle_params[PairPotential_Angle::spring_constant] = _params[angle_spring_constant];
    if (paramsDim > angle_equilibrium_distance)
      angle_params[PairPotential_Angle::equilibrium_distance] = _params[angle_equilibrium_distance];
    angle_params[PairPotential_Angle::energy_offset] = 0.;
    angle.setParameters(angle_params);
  }
#ifndef NDEBUG
  parameters_t check_params(getParameters());
  check_params.resize(paramsDim); // truncate to same size
  ASSERT( check_params == _params,
      "SaturationPotential::setParameters() - failed, mismatch in to be set "
      +toString(_params)+" and set "+toString(check_params)+" params.");
#endif
}

SaturationPotential::parameters_t SaturationPotential::getParameters() const
{
  parameters_t params(getParameterDimension());
  const parameters_t morse_params = morse.getParameters();
  const parameters_t angle_params = angle.getParameters();

  params[all_energy_offset] = energy_offset;

  params[morse_spring_constant] = morse_params[PairPotential_Morse::spring_constant];
  params[morse_equilibrium_distance] = morse_params[PairPotential_Morse::equilibrium_distance];
  params[morse_dissociation_energy] = morse_params[PairPotential_Morse::dissociation_energy];

  params[angle_spring_constant] = angle_params[PairPotential_Angle::spring_constant];
  params[angle_equilibrium_distance] = angle_params[PairPotential_Angle::equilibrium_distance];
  return params;
}

std::vector<FunctionModel::arguments_t>
triplefunction(
    const argument_t &argument,
    const FunctionModel::arguments_t& args)
{
  const size_t firstindex = argument.indices.first;
  const size_t secondindex = argument.indices.second;
//  LOG(2, "DEBUG: first index is " << firstindex << ", second index is " << secondindex << ".");

  // place all arguments that share either index into a lookup map
  typedef std::map< size_t, FunctionModel::arguments_t::const_iterator > IndexLookup_t;
  IndexLookup_t LookuptoFirst;
  IndexLookup_t LookuptoSecond;
  for (FunctionModel::arguments_t::const_iterator iter = args.begin();
      iter != args.end();
      ++iter) {
    if (((*iter).indices.first == argument.indices.first)
        && ((*iter).indices.second == argument.indices.second))
      continue;
    if (firstindex == (*iter).indices.first) {
      LookuptoFirst.insert( std::make_pair( (*iter).indices.second, iter) );
    }
    else if (firstindex == (*iter).indices.second) {
      LookuptoFirst.insert( std::make_pair( (*iter).indices.first, iter) );
    }
    if (secondindex == (*iter).indices.first) {
      LookuptoSecond.insert( std::make_pair( (*iter).indices.second, iter) );
    }
    else if (secondindex == (*iter).indices.second) {
      LookuptoSecond.insert( std::make_pair((*iter).indices.first, iter) );
    }
  }
//  {
//    std::stringstream lookupstream;
//    for (IndexLookup_t::const_iterator iter = LookuptoFirst.begin();
//        iter != LookuptoFirst.end();
//        ++iter) {
//      lookupstream << "(" << iter->first << "," << *(iter->second) << ") ";
//    }
//    LOG(2, "DEBUG: LookupToFirst is " << lookupstream.str() << ".");
//  }
//  {
//    std::stringstream lookupstream;
//    for (IndexLookup_t::const_iterator iter = LookuptoSecond.begin();
//        iter != LookuptoSecond.end();
//        ++iter) {
//      lookupstream << "(" << iter->first << "," << *(iter->second) << ") ";
//    }
//    LOG(2, "DEBUG: LookuptoSecond is " << lookupstream.str() << ".");
//  }

  // now go through the first lookup as the second argument and pick the
  // corresponding third argument by the matching index
  std::vector<FunctionModel::arguments_t> results;
  for (IndexLookup_t::const_iterator iter = LookuptoFirst.begin();
      iter != LookuptoFirst.end();
      ++iter) {
    IndexLookup_t::const_iterator otheriter = LookuptoSecond.find(iter->first);
    ASSERT( otheriter != LookuptoSecond.end(),
        "triplefunction() - cannot find index "+toString(iter->first)
        +" in LookupToSecond");
    FunctionModel::arguments_t result(1, argument);
    result.reserve(3);
    result.push_back(*(iter->second));
    result.push_back(*(otheriter->second));
    results.push_back(result);
  }

  return results;
}

SaturationPotential::results_t
SaturationPotential::operator()(
    const arguments_t &arguments
    ) const
{
  double result = 0.;
  const ParticleTypes_t &morse_types = morse.getParticleTypes();
  for(arguments_t::const_iterator argiter = arguments.begin();
      argiter != arguments.end();
      ++argiter) {
    const argument_t &r_ij = *argiter;
    if (((r_ij.types.first == morse_types[0]) && (r_ij.types.second == morse_types[1]))
        || ((r_ij.types.first == morse_types[1]) && (r_ij.types.second == morse_types[0]))) {
      arguments_t args(1, r_ij);

      // Morse contribution
      const double tmp = morse(args)[0];
//      LOG(3, "DEBUG: Morse yields " << tmp << " for << " << r_ij << ".");
      result += tmp;
      if (result != result)
        ELOG(1, "result is NAN.");

      // Angle contribution
      {
        typedef std::vector<FunctionModel::arguments_t> tripleargs_t;
        tripleargs_t tripleargs =
            triplefunction(r_ij, arguments);
        for (tripleargs_t::const_iterator iter = tripleargs.begin();
            iter != tripleargs.end();
            ++iter) {
          FunctionModel::arguments_t tempargs =
              Extractors::reorderArgumentsByParticleTypes(*iter, angle.getParticleTypes());
          // We get both angles, e.g. 0-4-1 and 1-4-0, hence multiply with 0.5
          const double tmp = .5*angle(tempargs)[0];  // as we have all distances we get both jk and kj
//          LOG(3, "DEBUG: angle yields " << tmp << " for << " << tempargs << ".");
          result += tmp;
          if (result != result)
            ELOG(1, "result is NAN.");
        }
      }
    }
  }
  return std::vector<result_t>(1, energy_offset + result);
}

SaturationPotential::derivative_components_t
SaturationPotential::derivative(
    const arguments_t &arguments
    ) const
{
  ASSERT( 0,
      "SaturationPotential::operator() - not implemented.");
  derivative_components_t result;
  return result;
}

SaturationPotential::results_t
SaturationPotential::parameter_derivative(
    const arguments_t &arguments,
    const size_t index
    ) const
{
  double result = 0.;
  if (index == all_energy_offset) {
    result = 1.;
  } else {
    const ParticleTypes_t &morse_types = morse.getParticleTypes();
    for(arguments_t::const_iterator argiter = arguments.begin();
        argiter != arguments.end();
        ++argiter) {
      const argument_t &r_ij = *argiter;
      if (((r_ij.types.first == morse_types[0]) && (r_ij.types.second == morse_types[1]))
          || ((r_ij.types.first == morse_types[1]) && (r_ij.types.second == morse_types[0]))) {
        arguments_t args(1, r_ij);

        switch (index) {
          case morse_spring_constant:
            result += morse.parameter_derivative(args, PairPotential_Morse::spring_constant)[0];
            break;
          case morse_equilibrium_distance:
            result += morse.parameter_derivative(args, PairPotential_Morse::equilibrium_distance)[0];
            break;
          case morse_dissociation_energy:
            result += morse.parameter_derivative(args, PairPotential_Morse::dissociation_energy)[0];
            break;
          case angle_spring_constant:
          case angle_equilibrium_distance:
          {
            typedef std::vector<FunctionModel::arguments_t> tripleargs_t;
            tripleargs_t tripleargs =
                triplefunction(r_ij, arguments);
            for (tripleargs_t::const_iterator iter = tripleargs.begin();
                iter != tripleargs.end();
                ++iter) {
              FunctionModel::arguments_t tempargs =
                  Extractors::reorderArgumentsByParticleTypes(*iter, angle.getParticleTypes());
              // We get both angles, e.g. 0-4-1 and 1-4-0, hence multiply with 0.5
              if (index == angle_spring_constant)
                result += .5*angle.parameter_derivative(tempargs, PairPotential_Angle::spring_constant)[0];
              else if (index == angle_equilibrium_distance)
                result += .5*angle.parameter_derivative(tempargs, PairPotential_Angle::equilibrium_distance)[0];
//              LOG(2, "DEBUG: angle yields " << tmp << " for << " << arguments << ".");
//              result += tmp;
              if (result != result)
                ELOG(1, "result is NAN.");
            }
          }
            break;
          default:
            ASSERT( 0, "SaturationPotential::parameter_derivative() - impossible to get here.");
            break;
        }
      }
    }
  }
  return SaturationPotential::results_t(1, result);
}

const SaturationPotential::ParticleTypes_t
SaturationPotential::symmetrizeTypes(const ParticleTypes_t &_ParticleTypes)
{
  ASSERT( _ParticleTypes.size() == (size_t)2,
      "SaturationPotential::symmetrizeTypes() - require initial _ParticleTypes with two elements.");
//  // insert before couple
//  ParticleTypes_t types(1, _ParticleTypes[1]);
//  types.insert(types.end(), _ParticleTypes.begin(), _ParticleTypes.end());
  // insert after the couple
  ParticleTypes_t types(_ParticleTypes);
  types.push_back( _ParticleTypes.back() );
  ASSERT( types.size() == (size_t)3,
      "SaturationPotential::symmetrizeTypes() - failed to generate three types for angle.");
  return types;
}

std::ostream& operator<<(std::ostream &ost, const SaturationPotential &potential)
{
  ost << potential.morse;
  ost << potential.angle;
  return ost;
}

std::istream& operator>>(std::istream &ist, SaturationPotential &potential)
{
  ist >> potential.morse;
  ist >> potential.angle;
  return ist;
}

FunctionModel::extractor_t
SaturationPotential::getFragmentSpecificExtractor(
    const charges_t &charges) const
{
  ASSERT(charges.size() == (size_t)2,
      "SaturationPotential::getFragmentSpecificExtractor() - requires 2 charges.");
  FunctionModel::extractor_t returnfunction;
  if (charges[0] == charges[1]) {
    // In case both types are equal there is only a single pair of possible
    // type combinations.
     returnfunction =
        boost::bind(&Extractors::gatherAllDistancesFromFragment,
            boost::bind(&Fragment::getPositions, _1),
            boost::bind(&Fragment::getCharges, _1),
            boost::cref(charges),
            _2);
  } else {
    // we have to chain here a rather complex "tree" of functions
    // as we only have a couple of ParticleTypes but need to get
    // all possible three pairs of the set of the two types.
    // Finally, we also need to arrange them in correct order
    // (for PairPotentiale_Angle).
    //charges_t firstpair(2, boost::cref(charges[0]));
    // only that saturation potential never has its middle element twice!
    // hence, we skip the firstpair but keep the code for later generalization
    charges_t secondpair(2, boost::cref(charges[1]));
    const charges_t &thirdpair = charges;
    returnfunction =
//        boost::bind(&Extractors::reorderArgumentsByParticleTypes,
          boost::bind(&Extractors::combineArguments,
//            boost::bind(&Extractors::combineArguments,
//              boost::bind(&Extractors::gatherAllDistancesFromFragment,
//                  boost::bind(&Fragment::getPositions, _1),
//                  boost::bind(&Fragment::getCharges, _1),
//                  firstpair,  // no crefs here as are temporaries!
//                  _2),
              boost::bind(&Extractors::gatherAllDistancesFromFragment,
                  boost::bind(&Fragment::getPositions, _1),
                  boost::bind(&Fragment::getCharges, _1),
                  secondpair,  // no crefs here as are temporaries!
                  _2),
//            ),
            boost::bind(&Extractors::gatherAllDistancesFromFragment,
                boost::bind(&Fragment::getPositions, _1),
                boost::bind(&Fragment::getCharges, _1),
                boost::cref(thirdpair), // only the last one is no temporary
                _2)
          );
//          boost::cref(angle.getParticleTypes())
//        );
}
  return returnfunction;
}

void
SaturationPotential::setParametersToRandomInitialValues(
    const TrainingData &data)
{
  energy_offset = data.getTrainingOutputAverage()[0];
  morse.setParametersToRandomInitialValues(data);
  angle.setParametersToRandomInitialValues(data);
}