/*
 * 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 "FunctionApproximation/TrainingData.hpp"
#include "FunctionApproximation/writeDistanceEnergyTable.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 getFirstGraphwithTimesSpecificElement(
    const HomologyContainer &homologies,
    const size_t _number,
    const size_t _times)
{
  for (HomologyContainer::container_t::const_iterator iter =
      homologies.begin(); iter != homologies.end(); ++iter) {
    if (iter->first.hasTimesAtomicNumber(_number,_times))
      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;
}

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 = getFirstGraphwithTimesSpecificElement(homologies,6,3);
    if (graph != HomologyGraph()) {
      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
      TrainingData AngleData(
          boost::bind(&Extractors::reorderArgumentsByIncreasingDistance,
              boost::bind(&Extractors::gatherAllSymmetricDistanceArguments,
                  boost::bind(&Extractors::gatherPositionOfTuples,
                      _1, Fragment::charges_t(3,6.)
                  ), _2 // gather carbon triples
              )
            )
          );
      AngleData(homologies.getHomologousGraphs(graph));
      LOG(1, "INFO: I gathered the following training data: " << AngleData);
      // 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(AngleData, model);
      if (model.isBoxConstraint() && approximator.checkParameterDerivatives())
        approximator(FunctionApproximation::ParameterDerivative);
      else {
        ELOG(0, "We require parameter derivatives for a box constraint minimization.");
        return 1;
      }
      params = model.getParameters();

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

  /******************** MORSE TRAINING ********************/
  {
    // then we ought to pick the right HomologyGraph ...
    const HomologyGraph graph = getFirstGraphwithTimesSpecificElement(homologies,6,2);
    if (graph != HomologyGraph()) {
      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::gatherAllSymmetricDistanceArguments,
              boost::bind(&Extractors::gatherPositionOfTuples,
                  _1, Fragment::charges_t(2,6.)
              ), _2 // 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;
      LOG(0, "INFO: Initial parameters are " << params << ".");
      morse.setParameters(params);
      FunctionModel &model = morse;
      FunctionApproximation approximator(MorseData, model); // we only give CC distance, hence 1 input dim
      if (model.isBoxConstraint() && approximator.checkParameterDerivatives())
        approximator(FunctionApproximation::ParameterDerivative);
      else {
        ELOG(0, "We require parameter derivatives for a box constraint minimization.");
        return 1;
      }
      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 = getFirstGraphwithTimesSpecificElement(homologies,6,1);
    if (graph != HomologyGraph()) {
      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
      Fragment::charges_t SaturatedCarbon;
      SaturatedCarbon += 6.,1.;
      TrainingData TersoffData(
//          boost::bind(&Extractors::gatherAllSymmetricDistanceArguments,
//              boost::bind(&Extractors::gatherPositionOfTuples,
//                  _1, boost::cref(SaturatedCarbon)
//              ), _2
//            )
//          );
          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
      params[SaturationPotential::all_energy_offset] = 1e+1*rand()/(double)RAND_MAX;//3.487900e+00;
      params[SaturationPotential::morse_spring_constant] = 1e+1*rand()/(double)RAND_MAX;//1.393600e+03;
      params[SaturationPotential::morse_equilibrium_distance] = 1e+1*rand()/(double)RAND_MAX;//3.467000e+02;
      params[SaturationPotential::morse_dissociation_energy] = 1e+1*rand()/(double)RAND_MAX;//3.487900e+00;
      params[SaturationPotential::angle_spring_constant] = 1e+1*rand()/(double)RAND_MAX;//3.487900e+00;
      params[SaturationPotential::angle_equilibrium_distance] = 2e+0*rand()/(double)RAND_MAX - 1.;//3.487900e+00;
      LOG(0, "INFO: Initial parameters are " << params << ".");

      SaturationPotential saturation(2.5, triplefunction);
      saturation.setParameters(params);
      FunctionModel &model = saturation;
      FunctionApproximation approximator(TersoffData, model); // CH4 has 5 atoms, hence 5*4/2 distances
      if (model.isBoxConstraint() && approximator.checkParameterDerivatives())
        approximator(FunctionApproximation::ParameterDerivative);
      else {
        ELOG(0, "We require parameter derivatives for a box constraint minimization.");
        return 1;
      }

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