/*
 * 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/filesystem.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/FunctionApproximation.hpp"
#include "FunctionApproximation/FunctionModel.hpp"
#include "Helpers/defs.hpp"
#include "Potentials/Specifics/PairPotential_Morse.hpp"
#include "Potentials/Specifics/SaturationPotential.hpp"

namespace po = boost::program_options;

HomologyGraph getFirstGraphWithTwoCarbons(const HomologyContainer &homologies)
{
  FragmentNode SaturatedCarbon(6,4); // 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,3); // 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();
}

FunctionModel::arguments_t
gatherAllDistanceArguments(
    const Fragment::charges_t &charges,
    const Fragment::positions_t &positions,
    const size_t globalid)
{
  FunctionModel::arguments_t result;

  // go through current configuration and gather all other distances
  Fragment::charges_t::const_iterator firstchargeiter = charges.begin();
  Fragment::positions_t::const_iterator firstpositer = positions.begin();
  for (;firstchargeiter != charges.end();
      ++firstchargeiter, ++firstpositer) {
    Fragment::charges_t::const_iterator secondchargeiter = charges.begin();//firstchargeiter;
    Fragment::positions_t::const_iterator secondpositer = positions.begin();//firstpositer;
    for (;
        secondchargeiter != charges.end();
        ++secondchargeiter, ++secondpositer) {
      if (firstchargeiter == secondchargeiter)
        continue;
      argument_t arg;
      const Vector firsttemp((*firstpositer)[0],(*firstpositer)[1],(*firstpositer)[2]);
      const Vector secondtemp((*secondpositer)[0],(*secondpositer)[1],(*secondpositer)[2]);
      arg.distance = firsttemp.distance(secondtemp);
      arg.indices = std::make_pair(
          std::distance(
              charges.begin(), firstchargeiter),
          std::distance(
              charges.begin(), secondchargeiter)
          );
      arg.globalid = globalid;
      result.push_back(arg);
    }
    ASSERT( secondpositer == positions.end(),
        "gatherAllDistanceArguments() - there are not as many positions as charges.");
  }
  ASSERT( firstpositer == positions.end(),
      "gatherAllDistanceArguments() - there are not as many positions as charges.");

  return result;
}

/** 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;
}


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 << ".");
  }

  /******************** 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
    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 iter = range.first; iter != range.second; ++iter) {
      // get distance out of Fragment
      const double &energy = iter->second.second;
      const Fragment &fragment = iter->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)2) {
        argument_t arg;
        arg.indices.first = DistanceVectors[0].second;
        arg.indices.second = DistanceVectors[1].second;
        arg.distance = DistanceVectors[0].first.distance(DistanceVectors[1].first);
        arg.globalid = DistanceEnergyVector.first.size();
        DistanceEnergyVector.first.push_back( FunctionModel::arguments_t(1,arg) );
        DistanceEnergyVector.second.push_back( FunctionModel::results_t(1,energy) );
      } else {
        ELOG(2, "main() - found not exactly two 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) {
        LOG(1, "INFO: (" << (*initer)[0].indices.first << "," << (*initer)[0].indices.second
            << ") " << (*initer)[0].distance << " with energy " << *outiter);
      }
    }
    // 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(1, 1, 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 << ".");
  }

  /******************* 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
    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);
    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 double &energy = iter->second.second;
      const Fragment &fragment = iter->second.first;
      const Fragment::charges_t charges = fragment.getCharges();
      const Fragment::positions_t positions = fragment.getPositions();
      FunctionModel::arguments_t args =
          gatherAllDistanceArguments(charges, positions, DistanceEnergyVector.first.size());
      DistanceEnergyVector.first.push_back( args );
      DistanceEnergyVector.second.push_back( FunctionModel::results_t(1,energy-EnergySum) );
    }
    // 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;
        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;
        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
    boost::function< std::vector<FunctionModel::arguments_t>(const argument_t &, const double)> triplefunction =
        boost::bind(&getTripleFromArgument, DistanceEnergyVector.first, _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(
        DistanceEnergyVector.first.begin()->size(), 
        DistanceEnergyVector.second.begin()->size(), 
        model); // CH4 has 5 atoms, hence 5*4/2 distances
    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 << ".");

//    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
    double L2sum = 0.;
    double Lmax = 0.;
    size_t maxindex = -1;
    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) {
      const FunctionModel::results_t result = model((*initer));
      const double temp = fabs((*outiter)[0] - result[0]);
      LOG(2, "DEBUG: L2 contribution = " << (*outiter)[0] << "-" << result[0] << "=" << temp);
      if (temp > Lmax) {
        Lmax = temp;
        maxindex = std::distance(const_cast<const FunctionApproximation::inputs_t &>(DistanceEnergyVector.first).begin(), initer);
      }
      L2sum += temp*temp;
    }
    LOG(1, "INFO: L2sum = " << L2sum << ", LMax = " << Lmax << " from " << maxindex);
  }

  return 0;
}