/*
 * 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 "Potentials/helpers.hpp"
#include "Potentials/ParticleTypeCheckers.hpp"

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,
    const double _saturation_cutoff,
    boost::function< std::vector<arguments_t>(const argument_t &, const double)> &_triplefunction) :
  SerializablePotential(_ParticleTypes),
  morse(_ParticleTypes),
  angle(_ParticleTypes),
  energy_offset(0.),
  triplefunction(_triplefunction),
  saturation_cutoff(_saturation_cutoff)
{}

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,
    const double _saturation_cutoff,
    boost::function< std::vector<arguments_t>(const argument_t &, const double)> &_triplefunction) :
  SerializablePotential(_ParticleTypes),
  morse(_ParticleTypes),
  angle(symmetrizeTypes(_ParticleTypes)),
  energy_offset(_all_energy_offset),
  triplefunction(_triplefunction),
  saturation_cutoff(_saturation_cutoff)
{
  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;
}

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(2, "DEBUG: Morse yields " << tmp << " for << " << r_ij << ".");
      result += tmp;
      if (result != result)
        ELOG(1, "result is NAN.");
    }
  }
  {
    // Angle contribution
    const double tmp = angle(arguments)[0];  // as we have all distances we get both jk and kj
//    LOG(2, "DEBUG: angle yields " << tmp << " for << " << arguments << ".");
    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.;
  switch (index) {
    case all_energy_offset:
    {
      result = 1.;
      break;
    }
    case morse_spring_constant:
    case morse_equilibrium_distance:
    case morse_dissociation_energy:
    {
      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;
            default:
              ASSERT( 0, "SaturationPotential::parameter_derivative() - impossible to get here.");
              break;
          }
        }
      }
      break;
    }
    case angle_spring_constant:
    {
      result = angle.parameter_derivative(arguments, PairPotential_Angle::spring_constant)[0];
      break;
    }
    case angle_equilibrium_distance:
    {
      result = angle.parameter_derivative(arguments, PairPotential_Angle::equilibrium_distance)[0];
      break;
    }
    default:
      ELOG(1, "SaturationPotential::parameter_derivative() - index " << index << " invalid.");
      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;
}