/*
 * Project: MoleCuilder
 * Description: creates and alters molecular systems
 * Copyright (C)  2013 Frederik Heber. All rights reserved.
 * 
 *
 *   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/>.
 */

/*
 * FitParticleChargesAction.cpp
 *
 *  Created on: Jul 03, 2013
 *      Author: heber
 */

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

// needs to come before MemDebug due to placement new
#include <boost/archive/text_iarchive.hpp>

#include "CodePatterns/MemDebug.hpp"

#include "Atom/atom.hpp"
#include "CodePatterns/Log.hpp"
#include "Fragmentation/Exporters/ExportGraph_ToFiles.hpp"
#include "Fragmentation/Graph.hpp"
#include "World.hpp"

#include <boost/filesystem.hpp>
#include <boost/foreach.hpp>
#include <algorithm>
#include <functional>
#include <iostream>
#include <string>

#include "Actions/PotentialAction/FitParticleChargesAction.hpp"

#include "Potentials/PartialNucleiChargeFitter.hpp"

#include "Element/element.hpp"
#include "Fragmentation/Homology/HomologyContainer.hpp"
#include "Fragmentation/Homology/HomologyGraph.hpp"
#include "Fragmentation/Summation/SetValues/SamplingGrid.hpp"
#include "FunctionApproximation/Extractors.hpp"
#include "Potentials/PartialNucleiChargeFitter.hpp"
#include "Potentials/SerializablePotential.hpp"
#include "World.hpp"

using namespace MoleCuilder;

// and construct the stuff
#include "FitParticleChargesAction.def"
#include "Action_impl_pre.hpp"
/** =========== define the function ====================== */

inline
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(
          static_cast<size_t>(countiter->first),
          static_cast<size_t>(countiter->second))
          )
        break;
    if( countiter == counts_per_charge.end())
      return *iter;
  }
  return HomologyGraph();
}

ActionState::ptr PotentialFitParticleChargesAction::performCall() {

  // fragment specifies the homology fragment to use
  SerializablePotential::ParticleTypes_t fragmentnumbers;
  {
    const std::vector<const element *> &fragment = params.fragment.get();
    std::transform(fragment.begin(), fragment.end(), std::back_inserter(fragmentnumbers),
        boost::bind(&element::getAtomicNumber, _1));
  }

  // parse homologies into container
  HomologyContainer &homologies = World::getInstance().getHomologies();
  const HomologyGraph graph = getFirstGraphwithSpecifiedElements(homologies,fragmentnumbers);
  HomologyContainer::range_t range = homologies.getHomologousGraphs(graph);
  // for the moment just use the very first fragment
  if (range.first == range.second) {
    STATUS("HomologyContainer does not contain specified fragment.");
    return Action::failure;
  }

  // average partial charges over all fragments
  HomologyContainer::const_iterator iter = range.first;
  if (!iter->second.containsGrids) {
    STATUS("This HomologyGraph does not contain sampled grids.");
    return Action::failure;
  }
  PartialNucleiChargeFitter::charges_t averaged_charges;
  averaged_charges.resize(iter->second.fragment.getCharges().size(), 0.);
  size_t NoFragments = 0;
  for (;
      iter != range.second; ++iter, ++NoFragments) {
    if (!iter->second.containsGrids) {
      ELOG(2, "This HomologyGraph does not contain sampled grids,\ndid you forget to add '--store-grids 1' to AnalyseFragmentResults.");
      return Action::failure;
    }
    const Fragment &fragment = iter->second.fragment;
  //  const double &energy = iter->second.energy;
  //  const SamplingGrid &charge = iter->second.charge_distribution;
    const SamplingGrid &potential = iter->second.potential_distribution;
    if ((potential.level == 0)
        || ((potential.begin[0] == potential.end[0])
            && (potential.begin[1] == potential.end[1])
            && (potential.begin[2] == potential.end[2]))) {
      ELOG(1, "Sampled grid contains grid made of zero points.");
      return Action::failure;
    }

    // then we extract positions from fragment
    PartialNucleiChargeFitter::positions_t positions;
    Fragment::positions_t fragmentpositions = fragment.getPositions();
    positions.reserve(fragmentpositions.size());
    BOOST_FOREACH( Fragment::position_t pos, fragmentpositions) {
      positions.push_back( Vector(pos[0], pos[1], pos[2]) );
    }
    PartialNucleiChargeFitter fitter(potential, positions, params.radius.get());
    fitter();
    PartialNucleiChargeFitter::charges_t return_charges = fitter.getSolutionAsCharges_t();
    std::transform(
        return_charges.begin(), return_charges.end(),
        averaged_charges.begin(),
        averaged_charges.begin(),
        std::plus<PartialNucleiChargeFitter::charge_t>());
  }
  std::transform(averaged_charges.begin(),averaged_charges.end(),
      averaged_charges.begin(),
      std::bind1st(std::multiplies<PartialNucleiChargeFitter::charge_t>(),1./NoFragments)
  );


  // output fitted charges
  LOG(0, "STATUS: We have fitted the following charges " << averaged_charges
      << ", averaged over " << NoFragments << " fragments.");

  return Action::success;
}

ActionState::ptr PotentialFitParticleChargesAction::performUndo(ActionState::ptr _state) {
  return Action::success;
}

ActionState::ptr PotentialFitParticleChargesAction::performRedo(ActionState::ptr _state){
  return Action::success;
}

bool PotentialFitParticleChargesAction::canUndo() {
  return false;
}

bool PotentialFitParticleChargesAction::shouldUndo() {
  return false;
}
/** =========== end of function ====================== */