/*
 * 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/assign/list_of.hpp>
#include <boost/bind.hpp>
#include <boost/filesystem.hpp>
#include <boost/foreach.hpp>
#include <boost/function.hpp>
#include <boost/lambda/lambda.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/PotentialFactory.hpp"
#include "Potentials/PotentialRegistry.hpp"
#include "Potentials/Specifics/PairPotential_Morse.hpp"
#include "Potentials/Specifics/PairPotential_Angle.hpp"
#include "Potentials/Specifics/SaturationPotential.hpp"
#include "types.hpp"

namespace po = boost::program_options;

using namespace boost::assign;

HomologyGraph getFirstGraphwithSpecifiedElements(
    const HomologyContainer &homologies,
    const SerializablePotential::ParticleTypes_t &types)
{
  ASSERT( !types.empty(),
      "getFirstGraphwithSpecifiedElements() - charges is empty?");
  // create charges
  Fragment::charges_t charges;
  charges.resize(types.size());
  std::transform(types.begin(), types.end(),
      charges.begin(), boost::lambda::_1);
  // convert into count map
  Extractors::elementcounts_t counts_per_charge =
      Extractors::_detail::getElementCounts(charges);
  ASSERT( !counts_per_charge.empty(),
      "getFirstGraphwithSpecifiedElements() - charge counts are empty?");
  LOG(2, "DEBUG: counts_per_charge is " << counts_per_charge << ".");
  // we want to check each (unique) key only once
  HomologyContainer::const_key_iterator olditer = homologies.key_end();
  for (HomologyContainer::const_key_iterator iter =
      homologies.key_begin(); iter != homologies.key_end(); olditer = iter++) {
    // if it's the same as the old one, skip it
    if (*olditer == *iter)
      continue;
    // if it's a new key, check if every element has the right number of counts
    Extractors::elementcounts_t::const_iterator countiter = counts_per_charge.begin();
    for (; countiter != counts_per_charge.end(); ++countiter)
      if (!(*iter).hasTimesAtomicNumber(countiter->first,countiter->second))
        break;
    if( countiter == counts_per_charge.end())
      return *iter;
  }
  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;

//  setVerbosity(4);

  // 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")
      ("fit-potential", po::value< std::string >(), "potential type to fit")
      ("charges", po::value< SerializablePotential::ParticleTypes_t >()->multitoken(), "charges specifying the potential")
      ("fragment", po::value< SerializablePotential::ParticleTypes_t >()->multitoken(), "all charges in the fragment")
  ;

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

  // homology-file
  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 {
    ELOG(0, "homology file (homology-file) was not set.");
    return 1;
  }

  // type of potential to fit
  std::string potentialtype;
  if (vm.count("fit-potential")) {
    potentialtype = vm["fit-potential"].as<std::string>();
  } else {
    ELOG(0, "potential type to fit (fit-potential) was not set.");
    return 1;
  }

  // charges
  SerializablePotential::ParticleTypes_t charges;
  if (vm.count("charges")) {
    charges = vm["charges"].as< SerializablePotential::ParticleTypes_t >();
  } else {
    ELOG(0, "Vector of charges specifying the potential (charges) was not set.");
    return 1;
  }

  // fragment
  SerializablePotential::ParticleTypes_t fragment;
  if (vm.count("fragment")) {
    fragment = vm["fragment"].as< SerializablePotential::ParticleTypes_t >();
  } else {
    ELOG(0, "Vector of charges specifying the fragment (charges) was not set.");
    return 1;
  }

  // parse homologies into container
  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(0, "Failed to parse from " << homology_file.string() << ".");
      return 1;
    }
    returnstream.close();
  } else {
    ELOG(0, homology_file << " does not exist.");
    return 1;
  }

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

  LOG(0, "STATUS: I'm training now a " << potentialtype << " potential on charges "
      << charges << ".");

  /******************** TRAINING ********************/
  // fit potential
  FunctionModel *model =
      PotentialFactory::getInstance().createInstance(
          potentialtype,
          charges);
  ASSERT( model != NULL,
      "main() - model returned from PotentialFactory is NULL.");
  FunctionModel::parameters_t params(model->getParameterDimension(), 0.);
  {
    // then we ought to pick the right HomologyGraph ...
    const HomologyGraph graph = getFirstGraphwithSpecifiedElements(homologies,fragment);
    if (graph != HomologyGraph()) {
      LOG(1, "First representative graph containing fragment "
          << fragment << " is " << graph << ".");

      // Afterwards we go through all of this type and gather the distance and the energy value
      TrainingData data(model->getFragmentSpecificExtractor());
      data(homologies.getHomologousGraphs(graph));
      if (!data.getTrainingInputs().empty()) {
        // print which distance is which
        size_t counter=1;
        const FunctionModel::arguments_t &inputs = data.getTrainingInputs()[0];
        for (FunctionModel::arguments_t::const_iterator iter = inputs.begin();
            iter != inputs.end(); ++iter) {
          const argument_t &arg = *iter;
          LOG(1, "INFO: distance " << counter++ << " is between (#"
              << arg.indices.first << "c" << arg.types.first << ","
              << arg.indices.second << "c" << arg.types.second << ").");
        }

        // print table
        LOG(1, "INFO: I gathered the following training data:\n" <<
            _detail::writeDistanceEnergyTable(data.getDistanceEnergyTable()));
      }
      // NOTICE that distance are in bohrradi as they come from MPQC!

      // now perform the function approximation by optimizing the model function
      FunctionApproximation approximator(data, *model);
      if (model->isBoxConstraint() && approximator.checkParameterDerivatives()) {
        // we set parameters here because we want to test with default ones
        srand((unsigned)time(0)); // seed with current time
        model->setParametersToRandomInitialValues(data);
        LOG(0, "INFO: Initial parameters are " << model->getParameters() << ".");
        approximator(FunctionApproximation::ParameterDerivative);
      } else {
        ELOG(0, "We require parameter derivatives for a box constraint minimization.");
        return 1;
      }

      // create a map of each fragment with error.
      typedef std::multimap< double, size_t > WorseFragmentMap_t;
      WorseFragmentMap_t WorseFragmentMap;
      HomologyContainer::range_t fragmentrange = homologies.getHomologousGraphs(graph);
      // fragments make it into the container in reversed order, hence count from top down
      size_t index= std::distance(fragmentrange.first, fragmentrange.second)-1;
      for (HomologyContainer::const_iterator iter = fragmentrange.first;
          iter != fragmentrange.second;
          ++iter) {
        const Fragment& fragment = iter->second.first;
        const double &energy = iter->second.second;

        // create arguments from the fragment
        FunctionModel::extractor_t extractor = model->getFragmentSpecificExtractor();
        FunctionModel::arguments_t args = extractor(fragment, 1);

        // calculate value from potential
        const double fitvalue = (*model)(args)[0];

        // insert difference into map
        const double error = fabs(energy - fitvalue);
        WorseFragmentMap.insert( std::make_pair( error, index-- ) );

        {
          // give only the distances in the debugging text
          std::stringstream streamargs;
          BOOST_FOREACH (argument_t arg, args) {
            streamargs << " " << arg.distance*AtomicLengthToAngstroem;
          }
          LOG(2, "DEBUG: frag.#" << index+1 << "'s error is |" << energy << " - " << fitvalue
              << "| = " << error << " for args " << streamargs.str() << ".");
        }
      }
      LOG(0, "RESULT: WorstFragmentMap " << WorseFragmentMap << ".");

      params = model->getParameters();

      SerializablePotential *potential = dynamic_cast<SerializablePotential *>(model);
      if (potential != NULL) {
        LOG(1, "STATUS: Resulting parameters are " << std::endl << *potential << ".");
      } else {
        LOG(1, "INFO: FunctionModel is no serializable potential.");
      }
    }
  }
  delete model;
  // remove static instances
  PotentialFactory::purgeInstance();

  return 0;
}