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

/*
 * LevMartester.cpp
 *
 *  Created on: Sep 27, 2012
 *      Author: heber
 */


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

#include <boost/archive/text_iarchive.hpp>

#include "CodePatterns/MemDebug.hpp"

#include <boost/assign.hpp>
#include <boost/bind.hpp>
#include <boost/filesystem.hpp>
#include <boost/function.hpp>
#include <boost/program_options.hpp>

#include <cstdlib>
#include <ctime>
#include <fstream>
#include <iostream>
#include <iterator>
#include <list>
#include <vector>

#include <levmar.h>

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

#include "LinearAlgebra/Vector.hpp"

#include "Fragmentation/Homology/HomologyContainer.hpp"
#include "Fragmentation/SetValues/Fragment.hpp"
#include "FunctionApproximation/Extractors.hpp"
#include "FunctionApproximation/FunctionApproximation.hpp"
#include "FunctionApproximation/FunctionModel.hpp"
#include "Helpers/defs.hpp"
#include "Potentials/Specifics/PairPotential_Morse.hpp"
#include "Potentials/Specifics/PairPotential_Angle.hpp"
#include "Potentials/Specifics/SaturationPotential.hpp"

namespace po = boost::program_options;

using namespace boost::assign;

HomologyGraph getFirstGraphWithThreeCarbons(const HomologyContainer &homologies)
{
  FragmentNode SaturatedCarbon(6,4); // carbon has atomic number 6 and should have 4 bonds for C3H8
  FragmentNode DanglingCarbon(6,3); // carbon has atomic number 6 and should have 3 pure bonds for C3H8
  for (HomologyContainer::container_t::const_iterator iter =
      homologies.begin(); iter != homologies.end(); ++iter) {
    if ((iter->first.hasNode(SaturatedCarbon,2)) && (iter->first.hasNode(DanglingCarbon,1)))
      return iter->first;
  }
  return HomologyGraph();
}

HomologyGraph getFirstGraphWithTwoCarbons(const HomologyContainer &homologies)
{
  FragmentNode SaturatedCarbon(6,3); // carbon has atomic number 6 and should have 4 bonds for C2H6
  for (HomologyContainer::container_t::const_iterator iter =
      homologies.begin(); iter != homologies.end(); ++iter) {
    if (iter->first.hasNode(SaturatedCarbon,2))
      return iter->first;
  }
  return HomologyGraph();
}

HomologyGraph getFirstGraphWithOneCarbon(const HomologyContainer &homologies)
{
  FragmentNode SaturatedCarbon(6,2); // carbon has atomic number 6 and has 3 bonds (to other Hs)
  for (HomologyContainer::container_t::const_iterator iter =
      homologies.begin(); iter != homologies.end(); ++iter) {
    if (iter->first.hasNode(SaturatedCarbon,1))
      return iter->first;
  }
  return HomologyGraph();
}


/** This function returns the elements of the sum over index "k" for an
 * argument containing indices "i" and "j"
 * @param inputs vector of all configuration (containing each a vector of all arguments)
 * @param arg argument containing indices "i" and "j"
 * @param cutoff cutoff criterion for sum over k
 * @return vector of argument pairs (a vector) of ik and jk for at least all k
 *        within distance of \a cutoff to i
 */
std::vector<FunctionModel::arguments_t>
getTripleFromArgument(const FunctionApproximation::inputs_t &inputs, const argument_t &arg, const double cutoff)
{
  typedef std::list<argument_t> arg_list_t;
  typedef std::map<size_t, arg_list_t > k_args_map_t;
  k_args_map_t tempresult;
  ASSERT( inputs.size() > arg.globalid,
      "getTripleFromArgument() - globalid "+toString(arg.globalid)
      +" is greater than all inputs "+toString(inputs.size())+".");
  const FunctionModel::arguments_t &listofargs = inputs[arg.globalid];
  for (FunctionModel::arguments_t::const_iterator argiter = listofargs.begin();
      argiter != listofargs.end();
      ++argiter) {
    // first index must be either i or j but second index not
    if (((argiter->indices.first == arg.indices.first)
        || (argiter->indices.first == arg.indices.second))
      && ((argiter->indices.second != arg.indices.first)
          && (argiter->indices.second != arg.indices.second))) {
      // we need arguments ik and jk
      std::pair< k_args_map_t::iterator, bool> inserter =
          tempresult.insert( std::make_pair( argiter->indices.second, arg_list_t(1,*argiter)));
      if (!inserter.second) {
        // is present one ik or jk, if ik insert jk at back
        if (inserter.first->second.begin()->indices.first == arg.indices.first)
          inserter.first->second.push_back(*argiter);
        else // if jk, insert ik at front
          inserter.first->second.push_front(*argiter);
      }
    }
//    // or second index must be either i or j but first index not
//    else if (((argiter->indices.first != arg.indices.first)
//              && (argiter->indices.first != arg.indices.second))
//            && ((argiter->indices.second == arg.indices.first)
//                || (argiter->indices.second == arg.indices.second))) {
//      // we need arguments ki and kj
//      std::pair< k_args_map_t::iterator, bool> inserter =
//          tempresult.insert( std::make_pair( argiter->indices.first, arg_list_t(1,*argiter)));
//      if (!inserter.second) {
//        // is present one ki or kj, if ki insert kj at back
//        if (inserter.first->second.begin()->indices.second == arg.indices.first)
//          inserter.first->second.push_back(*argiter);
//        else // if kj, insert ki at front
//          inserter.first->second.push_front(*argiter);
//      }
//    }
  }
  // check that i,j are NOT contained
  ASSERT( tempresult.count(arg.indices.first) == 0,
      "getTripleFromArgument() - first index of argument present in k_args_map?");
  ASSERT( tempresult.count(arg.indices.second) == 0,
      "getTripleFromArgument() - first index of argument present in k_args_map?");

  // convert
  std::vector<FunctionModel::arguments_t> result;
  for (k_args_map_t::const_iterator iter = tempresult.begin();
      iter != tempresult.end();
      ++iter) {
    ASSERT( iter->second.size() == 2,
        "getTripleFromArgument() - for index "+toString(iter->first)+" we did not find both ik and jk.");
    result.push_back( FunctionModel::arguments_t(iter->second.begin(), iter->second.end()) );
  }
  return result;
}

double
function_angle(
    const double &r_ij,
    const double &r_ik,
    const double &r_jk
  )
{
//  Info info(__func__);
  const double angle = pow(r_ij,2.) + pow(r_ik,2.) - pow(r_jk,2.);
  const double divisor = 2.* r_ij * r_ik;

//  LOG(2, "DEBUG: cos(theta)= " << angle/divisor);
  if (divisor == 0.)
    return 0.;
  else
    return angle/divisor;
}

/** This class encapsulates the training data for a given potential function
 * to learn.
 *
 * The data is added piece-wise by calling the operator() with a specific
 * Fragment.
 */
class TrainingData
{
public:
  //!> typedef for a range within the HomologyContainer at which fragments to look at
  typedef std::pair<
      HomologyContainer::const_iterator,
      HomologyContainer::const_iterator> range_t;
  //!> Training tuple input vector pair
  typedef FunctionApproximation::inputs_t InputVector_t;
  //!> Training tuple output vector pair
  typedef FunctionApproximation::outputs_t OutputVector_t;
  //!> Typedef for a function containing how to extract required information from a Fragment.
  typedef boost::function< FunctionModel::arguments_t (const Fragment &, const size_t)> extractor_t;

public:
  /** Constructor for class TrainingData.
   *
   */
  explicit TrainingData(const extractor_t &_extractor) :
      extractor(extractor)
  {}
  /** Destructor for class TrainingData.
   *
   */
  ~TrainingData()
  {}

  /** We go through the given \a range of homologous fragments and call
   * TrainingData::extractor on them in order to gather the distance and
   * the energy value, stored internally.
   *
   * \param range given range within a HomologyContainer of homologous fragments
   */
  void operator()(const range_t &range) {
    double EnergySum = 0.; //std::numeric_limits<double>::max();
    size_t counter = 0.;
    for (HomologyContainer::const_iterator iter = range.first; iter != range.second; ++iter) {
      const double &energy = iter->second.second;
//      if (energy <= EnergySum)
//        EnergySum = energy;
      EnergySum += energy;
      ++counter;
    }
    EnergySum *= 1./(double)counter;
    for (HomologyContainer::const_iterator iter = range.first; iter != range.second; ++iter) {
      // get distance out of Fragment
      const Fragment &fragment = iter->second.first;
      FunctionModel::arguments_t args = extractor(
            fragment,
            DistanceVector.size()
          );
      DistanceVector.push_back( args );
      const double &energy = iter->second.second;
      EnergyVector.push_back( FunctionModel::results_t(1, energy-EnergySum) );
    }
  }

  /** Getter for const access to internal training data inputs.
   *
   * \return const ref to training tuple of input vector
   */
  const InputVector_t& getTrainingInputs() const {
    return DistanceVector;
  }

  /** Getter for const access to internal training data outputs.
   *
   * \return const ref to training tuple of output vector
   */
  const OutputVector_t& getTrainingOutputs() const {
    return EnergyVector;
  }

  /** Calculate the L2 error of a given \a model against the stored training data.
   *
   * \param model model whose L2 error to calculate
   * \return sum of squared differences at training tuples
   */
  const double getL2Error(const FunctionModel &model) const
  {
    double L2sum = 0.;

    FunctionApproximation::inputs_t::const_iterator initer = DistanceVector.begin();
    FunctionApproximation::outputs_t::const_iterator outiter = EnergyVector.begin();
    for (; initer != DistanceVector.end(); ++initer, ++outiter) {
      const FunctionModel::results_t result = model((*initer));
      const double temp = fabs((*outiter)[0] - result[0]);
      L2sum += temp*temp;
    }
    return L2sum;
  }

  /** Calculate the Lmax error of a given \a model against the stored training data.
   *
   * \param model model whose Lmax error to calculate
   * \return maximum difference over all training tuples
   */
  const double getLMaxError(const FunctionModel &model) const
  {
    double Lmax = 0.;
    size_t maxindex = -1;
    FunctionApproximation::inputs_t::const_iterator initer = DistanceVector.begin();
    FunctionApproximation::outputs_t::const_iterator outiter = EnergyVector.begin();
    for (; initer != DistanceVector.end(); ++initer, ++outiter) {
      const FunctionModel::results_t result = model((*initer));
      const double temp = fabs((*outiter)[0] - result[0]);
      if (temp > Lmax) {
        Lmax = temp;
        maxindex = std::distance(
            const_cast<const FunctionApproximation::inputs_t &>(DistanceVector).begin(),
            initer
            );
      }
    }
    return Lmax;
  }

private:
  // prohibit use of default constructor, as we always require extraction functor.
  TrainingData();

private:
  //!> private training data vector
  InputVector_t DistanceVector;
  OutputVector_t EnergyVector;
  //!> function to be used for training input data extraction from a fragment
  const extractor_t extractor;
};

// print training data for debugging
std::ostream &operator<<(std::ostream &out, const TrainingData &data)
{
  const TrainingData::InputVector_t &DistanceVector = data.getTrainingInputs();
  const TrainingData::OutputVector_t &EnergyVector = data.getTrainingOutputs();
  out << "(" << DistanceVector.size()
      << "," << EnergyVector.size() << ") data pairs: ";
  FunctionApproximation::inputs_t::const_iterator initer = DistanceVector.begin();
  FunctionApproximation::outputs_t::const_iterator outiter = EnergyVector.begin();
  for (; initer != DistanceVector.end(); ++initer, ++outiter) {
    for (size_t index = 0; index < (*initer).size(); ++index)
       out << "(" << (*initer)[index].indices.first << "," << (*initer)[index].indices.second
          << ") " << (*initer)[index].distance;
    out << " with energy " << *outiter;
  }
  return out;
}

int main(int argc, char **argv)
{
  std::cout << "Hello to the World from LevMar!" << std::endl;

  // load homology file
  po::options_description desc("Allowed options");
  desc.add_options()
      ("help", "produce help message")
      ("homology-file", po::value< boost::filesystem::path >(), "homology file to parse")
  ;

  po::variables_map vm;
  po::store(po::parse_command_line(argc, argv, desc), vm);
  po::notify(vm);

  if (vm.count("help")) {
      std::cout << desc << "\n";
      return 1;
  }

  boost::filesystem::path homology_file;
  if (vm.count("homology-file")) {
    homology_file = vm["homology-file"].as<boost::filesystem::path>();
    LOG(1, "INFO: Parsing " << homology_file.string() << ".");
  } else {
    LOG(0, "homology-file level was not set.");
  }
  HomologyContainer homologies;
  if (boost::filesystem::exists(homology_file)) {
    std::ifstream returnstream(homology_file.string().c_str());
    if (returnstream.good()) {
      boost::archive::text_iarchive ia(returnstream);
      ia >> homologies;
    } else {
      ELOG(2, "Failed to parse from " << homology_file.string() << ".");
    }
    returnstream.close();
  } else {
    ELOG(0, homology_file << " does not exist.");
  }

  // first we try to look into the HomologyContainer
  LOG(1, "INFO: Listing all present homologies ...");
  for (HomologyContainer::container_t::const_iterator iter =
      homologies.begin(); iter != homologies.end(); ++iter) {
    LOG(1, "INFO: graph " << iter->first << " has Fragment "
        << iter->second.first << " and associated energy " << iter->second.second << ".");
  }

  /******************** Angle TRAINING ********************/
  {
    // then we ought to pick the right HomologyGraph ...
    const HomologyGraph graph = getFirstGraphWithThreeCarbons(homologies);
    LOG(1, "First representative graph containing three saturated carbons is " << graph << ".");

    // Afterwards we go through all of this type and gather the distance and the energy value
    typedef std::pair<
        FunctionApproximation::inputs_t,
        FunctionApproximation::outputs_t> InputOutputVector_t;
    InputOutputVector_t DistanceEnergyVector;
    std::pair<HomologyContainer::const_iterator, HomologyContainer::const_iterator> range =
        homologies.getHomologousGraphs(graph);
    for (HomologyContainer::const_iterator fragiter = range.first; fragiter != range.second; ++fragiter) {
      // get distance out of Fragment
      const double &energy = fragiter->second.second;
      const Fragment &fragment = fragiter->second.first;
      const Fragment::charges_t charges = fragment.getCharges();
      const Fragment::positions_t positions = fragment.getPositions();
      std::vector< std::pair<Vector, size_t> > DistanceVectors;
      for (Fragment::charges_t::const_iterator chargeiter = charges.begin();
          chargeiter != charges.end(); ++chargeiter) {
        if (*chargeiter == 6) {
          Fragment::positions_t::const_iterator positer = positions.begin();
          const size_t steps = std::distance(charges.begin(), chargeiter);
          std::advance(positer, steps);
          DistanceVectors.push_back(
              std::make_pair(Vector((*positer)[0], (*positer)[1], (*positer)[2]),
                  steps));
        }
      }
      if (DistanceVectors.size() == (size_t)3) {
        FunctionModel::arguments_t args(3);
        // we require specific ordering of the carbons: ij, ik, jk
        typedef std::vector< std::pair<size_t, size_t> > indices_t;
        indices_t indices;
        indices += std::make_pair(0,1), std::make_pair(0,2), std::make_pair(1,2);
        // create the three arguments
        for (indices_t::const_iterator iter = indices.begin(); iter != indices.end(); ++iter) {
          const size_t &firstindex = iter->first;
          const size_t &secondindex = iter->second;
          argument_t &arg = args[(size_t)std::distance(const_cast<const indices_t&>(indices).begin(), iter)];
          arg.indices.first = DistanceVectors[firstindex].second;
          arg.indices.second = DistanceVectors[secondindex].second;
          arg.distance = DistanceVectors[firstindex].first.distance(DistanceVectors[secondindex].first);
          arg.globalid = DistanceEnergyVector.first.size();
        }
        // make largest distance last to create correct angle 
        // (this would normally depend on the order of the nodes in the subgraph)
        std::list<argument_t> sorted_args;
        double greatestdistance = 0.;
        for(FunctionModel::arguments_t::const_iterator iter = args.begin(); iter != args.end(); ++iter)
          greatestdistance = std::max(greatestdistance, iter->distance);
        for(FunctionModel::arguments_t::const_iterator iter = args.begin(); iter != args.end(); ++iter)
          if (iter->distance == greatestdistance)
            sorted_args.push_back(*iter);
          else
            sorted_args.push_front(*iter);
        // and add the training pair
        DistanceEnergyVector.first.push_back( FunctionModel::arguments_t(sorted_args.begin(), sorted_args.end()) );
        DistanceEnergyVector.second.push_back( FunctionModel::results_t(1,energy) );
      } else {
        ELOG(2, "main() - found not exactly three carbon atoms in fragment "
            << fragment << ".");
      }
    }
    // print training data for debugging
    {
      LOG(1, "INFO: I gathered the following (" << DistanceEnergyVector.first.size()
          << "," << DistanceEnergyVector.second.size() << ") data pairs: ");
      FunctionApproximation::inputs_t::const_iterator initer = DistanceEnergyVector.first.begin();
      FunctionApproximation::outputs_t::const_iterator outiter = DistanceEnergyVector.second.begin();
      for (; initer != DistanceEnergyVector.first.end(); ++initer, ++outiter) {
        std::stringstream stream;
        const double cos_angle = function_angle((*initer)[0].distance,(*initer)[1].distance,(*initer)[2].distance);
        for (size_t index = 0; index < (*initer).size(); ++index)
           stream << " (" << (*initer)[index].indices.first << "," << (*initer)[index].indices.second
              << ") " << (*initer)[index].distance;
        stream << " with energy " << *outiter << " and cos(angle) " << cos_angle;
        LOG(1, "INFO:" << stream.str());
      }
    }
    // NOTICE that distance are in bohrradi as they come from MPQC!

    // now perform the function approximation by optimizing the model function
    FunctionModel::parameters_t params(PairPotential_Angle::MAXPARAMS, 0.);
    params[PairPotential_Angle::energy_offset] =  -1.;
    params[PairPotential_Angle::spring_constant] =  1.;
    params[PairPotential_Angle::equilibrium_distance] =  0.2;
    PairPotential_Angle angle;
    LOG(0, "INFO: Initial parameters are " << params << ".");
    angle.setParameters(params);
    FunctionModel &model = angle;
    FunctionApproximation approximator(
        DistanceEnergyVector.first.begin()->size(), 
        DistanceEnergyVector.second.begin()->size(), 
        model);
    approximator.setTrainingData(DistanceEnergyVector.first,DistanceEnergyVector.second);
    if (model.isBoxConstraint() && approximator.checkParameterDerivatives())
      approximator(FunctionApproximation::ParameterDerivative);
    else
      ELOG(0, "We require parameter derivatives for a box constraint minimization.");
    params = model.getParameters();

    LOG(0, "RESULT: Best parameters are " << params << ".");
  }

  /******************** MORSE TRAINING ********************/
  {
    // then we ought to pick the right HomologyGraph ...
    const HomologyGraph graph = getFirstGraphWithTwoCarbons(homologies);
    LOG(1, "First representative graph containing two saturated carbons is " << graph << ".");

    // Afterwards we go through all of this type and gather the distance and the energy value
    TrainingData MorseData(
        boost::bind(&Extractors::gatherFirstDistance, _1, _2, 6, 6) // gather first carbon pair
        );
    MorseData(homologies.getHomologousGraphs(graph));
    LOG(1, "INFO: I gathered the following training data: " << MorseData);
    // NOTICE that distance are in bohrradi as they come from MPQC!

    // now perform the function approximation by optimizing the model function
    FunctionModel::parameters_t params(PairPotential_Morse::MAXPARAMS, 0.);
    params[PairPotential_Morse::dissociation_energy] =  0.5;
    params[PairPotential_Morse::energy_offset] =  -1.;
    params[PairPotential_Morse::spring_constant] =  1.;
    params[PairPotential_Morse::equilibrium_distance] =  2.9;
    PairPotential_Morse morse;
    morse.setParameters(params);
    FunctionModel &model = morse;
    FunctionApproximation approximator(
        MorseData.getTrainingInputs().begin()->size(), 
        MorseData.getTrainingOutputs().begin()->size(), 
        model);
    approximator.setTrainingData(MorseData.getTrainingInputs(),MorseData.getTrainingOutputs());
    if (model.isBoxConstraint() && approximator.checkParameterDerivatives())
      approximator(FunctionApproximation::ParameterDerivative);
    else
      ELOG(0, "We require parameter derivatives for a box constraint minimization.");
    params = model.getParameters();

    LOG(0, "RESULT: Best parameters are " << params << ".");
  }

  /******************* SATURATION TRAINING *******************/
  FunctionModel::parameters_t params(SaturationPotential::MAXPARAMS, 0.);
  {
    // then we ought to pick the right HomologyGraph ...
    const HomologyGraph graph = getFirstGraphWithOneCarbon(homologies);
    LOG(1, "First representative graph containing one saturated carbon is " << graph << ".");

    // Afterwards we go through all of this type and gather the distance and the energy value
    TrainingData TersoffData(
        TrainingData::extractor_t(&Extractors::gatherAllDistances) // gather first carbon pair
        );
    TersoffData( homologies.getHomologousGraphs(graph) );
    LOG(1, "INFO: I gathered the following training data: " << TersoffData);
    // NOTICE that distance are in bohrradi as they come from MPQC!

    // now perform the function approximation by optimizing the model function
    boost::function< std::vector<FunctionModel::arguments_t>(const argument_t &, const double)> triplefunction =
        boost::bind(&getTripleFromArgument, boost::cref(TersoffData.getTrainingInputs()), _1, _2);
    srand((unsigned)time(0)); // seed with current time
    LOG(0, "INFO: Initial parameters are " << params << ".");

    SaturationPotential saturation(triplefunction);
    saturation.setParameters(params);
    FunctionModel &model = saturation;
    FunctionApproximation approximator(
        TersoffData.getTrainingInputs().begin()->size(),
        TersoffData.getTrainingOutputs().begin()->size(),
        model);
    approximator.setTrainingData(TersoffData.getTrainingInputs(),TersoffData.getTrainingOutputs());
    if (model.isBoxConstraint() && approximator.checkParameterDerivatives())
      approximator(FunctionApproximation::ParameterDerivative);
    else
      ELOG(0, "We require parameter derivatives for a box constraint minimization.");

    params = model.getParameters();

    LOG(0, "RESULT: Best parameters are " << params << ".");

//    std::cout << "\tsaturationparticle:";
//    std::cout << "\tparticle_type=C,";
//    std::cout << "\tA=" << params[SaturationPotential::A] << ",";
//    std::cout << "\tB=" << params[SaturationPotential::B] << ",";
//    std::cout << "\tlambda=" << params[SaturationPotential::lambda] << ",";
//    std::cout << "\tmu=" << params[SaturationPotential::mu] << ",";
//    std::cout << "\tbeta=" << params[SaturationPotential::beta] << ",";
//    std::cout << "\tn=" << params[SaturationPotential::n] << ",";
//    std::cout << "\tc=" << params[SaturationPotential::c] << ",";
//    std::cout << "\td=" << params[SaturationPotential::d] << ",";
//    std::cout << "\th=" << params[SaturationPotential::h] << ",";
////    std::cout << "\toffset=" << params[SaturationPotential::offset] << ",";
//    std::cout << "\tR=" << saturation.R << ",";
//    std::cout << "\tS=" << saturation.S << ";";
//    std::cout << std::endl;

    // check L2 and Lmax error against training set
    LOG(1, "INFO: L2sum = " << TersoffData.getL2Error(model)
        << ", LMax = " << TersoffData.getLMaxError(model) << ".");
  }

  return 0;
}