/*
 * Project: MoleCuilder
 * Description: creates and alters molecular systems
 * Copyright (C)  2012 University of Bonn. All rights reserved.
 * 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/>.
 */

/*
 * FragmentationLongRangeResults.cpp
 *
 *  Created on: Aug 31, 2012
 *      Author: heber
 */


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

#include "CodePatterns/MemDebug.hpp"

#include "FragmentationLongRangeResults.hpp"

#include <boost/mpl/for_each.hpp>
#include <boost/mpl/remove.hpp>
#include <sstream>

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

#include "Fragmentation/KeySetsContainer.hpp"
#include "Fragmentation/parseKeySetFile.hpp"
#include "Fragmentation/Summation/Converter/DataConverter.hpp"
#include "Fragmentation/Summation/Containers/createMatrixNrLookup.hpp"
#include "Fragmentation/Summation/Containers/extractJobIds.hpp"
#include "Fragmentation/Summation/AllLevelOrthogonalSummator.hpp"
#include "Fragmentation/Summation/IndexSetContainer.hpp"
#include "Fragmentation/Summation/OrthogonalSumUpPerLevel.hpp"
#include "Fragmentation/Summation/SubsetMap.hpp"
#include "Fragmentation/Summation/SumUpPerLevel.hpp"

#include "Helpers/defs.hpp"

FragmentationLongRangeResults::FragmentationLongRangeResults(
    const std::map<JobId_t,MPQCData> &fragmentData,
    std::map<JobId_t,VMGData> &longrangeData,
    const KeySetsContainer& _KeySet,
    const KeySetsContainer& _ForceKeySet) :
    KeySet(_KeySet),
    ForceKeySet(_ForceKeySet),
    hasForces((!longrangeData.empty()) && (longrangeData.begin()->second.hasForces))
{
  initLookups(fragmentData, longrangeData);

  // convert KeySetContainer to IndexSetContainer
  container.reset(new IndexSetContainer(KeySet));
  // create the map of all keysets
  subsetmap.reset(new SubsetMap(*container));
}

void FragmentationLongRangeResults::initLookups(
    const std::map<JobId_t,MPQCData> &fragmentData,
    std::map<JobId_t,VMGData> &longrangeData
    )
{
  // create lookup from job nr to fragment number
  {
    size_t MPQCFragmentCounter = 0;
    std::vector<bool> ValueMask(fragmentData.size(), true);
    const std::vector<JobId_t> mpqcjobids = extractJobIds<MPQCData>(fragmentData);
    MPQCMatrixNrLookup =
        createMatrixNrLookup(mpqcjobids, MPQCFragmentCounter, ValueMask);
  }

  {
    size_t VMGFragmentCounter = 0;
    std::vector<bool> ValueMask(longrangeData.size(), true);
    const std::vector<JobId_t> vmgjobids = extractJobIds<VMGData>(longrangeData);
    VMGMatrixNrLookup =
        createMatrixNrLookup(vmgjobids, VMGFragmentCounter, ValueMask);
  }
}

void FragmentationLongRangeResults::operator()(
    const std::map<JobId_t,MPQCData> &fragmentData,
    std::map<JobId_t,VMGData> &longrangeData,
    const std::vector<VMGData> &fullsolutionData,
    const std::vector<SamplingGrid> &full_sample,
    const SamplingGrid &_zero_globalgrid,
    const IndexedVectors::indices_t &_implicit_charges_indices)
{
  MaxLevel = subsetmap->getMaximumSetLevel();
  LOG(1, "INFO: Summing up results till level " << MaxLevel << ".");
  /// convert all MPQCData to MPQCDataMap_t
  {
    ASSERT( ForceKeySet.KeySets.size() == fragmentData.size(),
        "FragmentationLongRangeResults::FragmentationLongRangeResults() - ForceKeySet's KeySets and fragmentData differ in size.");

    OrthogonalSumUpPerLevel<MPQCDataGridMap_t, MPQCData, MPQCDataGridVector_t>
        chargesummation(fragmentData);
    // charges and potential use same grid, hence same zero instance
    chargesummation.setZeroInstance<MPQCDataFused::sampled_grid>(_zero_globalgrid);
    chargesummation(MPQCMatrixNrLookup, container, subsetmap, Result_Grid_fused, 
        Result_perIndexSet_Grid);

    // multiply each short-range potential with the respective charge
    std::map<JobId_t,MPQCData>::const_iterator mpqciter = fragmentData.begin();
    std::map<JobId_t,VMGData>::iterator vmgiter = longrangeData.begin();
    for (; vmgiter != longrangeData.end(); ++mpqciter, ++vmgiter) {
      vmgiter->second.sampled_potential *= mpqciter->second.sampled_grid;
    }
    // then sum up
    OrthogonalSumUpPerLevel<VMGDataMap_t, VMGData, VMGDataVector_t> vmgsummation(
        longrangeData);
    vmgsummation(VMGMatrixNrLookup, container, subsetmap,
        Result_LongRange_fused, Result_perIndexSet_LongRange);

    IndexedVectors::indices_t fullindices;
    if (hasLongRangeForces()) {
      // initialize zero instance map
      VMGDataForceMap_t ZeroInstances;
      ZeroInstanceInitializer<VMGDataForceMap_t> initZeroInstance(ZeroInstances);
      boost::mpl::for_each<VMGDataForceVector_t>(boost::ref(initZeroInstance));

      // force has extra data converter (this is similar to MPQCData's forces
      std::map<JobId_t, VMGDataForceMap_t> VMGData_Force_fused;
      convertDatatoForceMap<VMGData, VMGDataForceMap_t, VMGDataFused>(
          longrangeData, ForceKeySet, VMGData_Force_fused);
      Result_ForceLongRange_fused.resize(MaxLevel); // we need the results of correct size already
      AllLevelOrthogonalSummator<VMGDataForceMap_t> forceSummer(
                  subsetmap,
                  VMGData_Force_fused,
                  container->getContainer(),
                  VMGMatrixNrLookup,
                  Result_ForceLongRange_fused,
                  Result_perIndexSet_LongRange_Force,
                  ZeroInstances);
      boost::mpl::for_each<VMGDataForceVector_t>(boost::ref(forceSummer));
      // build full force index set
      KeySetsContainer::ArrayOfIntVectors::const_iterator arrayiter = ForceKeySet.KeySets.begin();
      std::set<IndexedVectors::index_t> sorted_indices;
      for (;arrayiter != ForceKeySet.KeySets.end(); ++arrayiter) {
        sorted_indices.insert(arrayiter->begin(), arrayiter->end());
      }
      // add additionally those from implicit charges which are not associated to
      // any fragment and hence are unknown so far.
      sorted_indices.insert(_implicit_charges_indices.begin(), _implicit_charges_indices.end());
      sorted_indices.erase(-1);
      fullindices.insert(fullindices.begin(), sorted_indices.begin(), sorted_indices.end());
    }

    // then sum up
    OrthogonalSumUpPerLevel<VMGDataGridMap_t, VMGData, VMGDataGridVector_t> vmggridsummation(
        longrangeData);
    vmggridsummation.setZeroInstance<VMGDataFused::sampled_potential>(_zero_globalgrid);
    vmggridsummation.setZeroInstance<VMGDataFused::both_sampled_potential>(_zero_globalgrid);
    vmggridsummation(VMGMatrixNrLookup, container, subsetmap,
        Result_GridLongRange_fused, Result_perIndexSet_LongRange_Grid);

    Result_LongRangeIntegrated_fused.reserve(MaxLevel);
    // NOTE: potential for level 1 is wrongly calculated within a molecule
    // as saturation hydrogen are not removed on this level yet
    for (size_t level = 1; level <= MaxLevel; ++level) {
      // We have calculated three different contributions: e-e, e-n+n-n, and n-n.
      // And we want to have e-e+e-n, n-n+n-e (where e-n = n-e).
      // For each of these three contributions we have a full solution and summed
      // up short range solutions.

      // first we obtain the full e-e energy as potential times charge on the
      // respective level.
      const SamplingGrid &charge_weight = full_sample[level-1];
      SamplingGrid full_sample_solution = fullsolutionData[level-1].sampled_potential;
      full_sample_solution *= charge_weight;
      double electron_solution_energy = full_sample_solution.integral();

      // then we subtract the summed-up short-range e-e interaction energy from
      // the full solution.
      const SamplingGrid &short_range_correction =
          boost::fusion::at_key<VMGDataFused::sampled_potential>(Result_GridLongRange_fused[level-1]);
      double electron_short_range_energy = short_range_correction.integral();
      electron_solution_energy -= electron_short_range_energy;
#ifndef NDEBUG
      {
        static const double round_offset(
            (std::numeric_limits<size_t>::round_style == std::round_toward_zero) ?
                0.5 : 0.); // need offset to get to round_toward_nearest behavior
        // we can only check equivalence if both have same level, otherwise
        // short_range_correction.integral() has higher precision because of finger grid
        const int surplus_level = full_sample_solution.getSurplusLevel(short_range_correction)+round_offset;
        if (full_sample_solution.level == short_range_correction.level+surplus_level) {
          SamplingGrid check_difference_full_solution = full_sample_solution;
          check_difference_full_solution -= short_range_correction;
          ASSERT( fabs(electron_solution_energy - check_difference_full_solution.integral()) < 1e-7,
              "FragmentationLongRangeResults::operator() - integral and energy are not exchangeable.");
        }
      }
#endif

      // then, we obtain the e-n+n-n full solution in the same way
      double nuclei_solution_energy = fullsolutionData[level-1].nuclei_long;
      double nuclei_short_range_energy =
          boost::fusion::at_key<VMGDataFused::nuclei_long>(Result_LongRange_fused[level-1]);
      nuclei_solution_energy -= nuclei_short_range_energy;

      // and also the e-n full solution
      double both_solution_energy = fullsolutionData[level-1].electron_long;
      double both_short_range_energy =
          boost::fusion::at_key<VMGDataFused::electron_long>(Result_LongRange_fused[level-1]);
      both_solution_energy -= both_short_range_energy;

      // energies from interpolation at nuclei position has factor of 1/2 already
      electron_solution_energy *= .5;
      electron_short_range_energy *= .5;

      // At last, we subtract e-n from n-n+e-n for full solution and short-range
      // correction.
      nuclei_solution_energy -= both_solution_energy;
      nuclei_short_range_energy -= both_short_range_energy;

      VMGDataLongRangeMap_t instance;
      boost::fusion::at_key<VMGDataFused::electron_longrange>(instance) = electron_solution_energy;
//      LOG(0, "Remaining long-range potential integral of level " << level << " is "
//          << full_sample_solution.integral() << ".");
      boost::fusion::at_key<VMGDataFused::electron_shortrange>(instance) = electron_short_range_energy;
//      LOG(0, "Short-range correction potential integral of level " << level << " is "
//          << short_range_correction.integral() << ".");
      boost::fusion::at_key<VMGDataFused::mixed_longrange>(instance) = both_solution_energy;
//      LOG(0, "Remaining long-range energy from potential integral of level " << level << " is "
//          << full_solution_energy << ".");
      boost::fusion::at_key<VMGDataFused::mixed_shortrange>(instance) = both_short_range_energy;
//      LOG(0, "Short-range correction energy from potential integral of level " << level << " is "
//          << short_range_energy << ".");
      boost::fusion::at_key<VMGDataFused::nuclei_longrange>(instance) = nuclei_solution_energy;
//      LOG(0, "Remaining long-range energy from potential integral of level " << level << " is "
//          << full_solution_energy << ".");
      boost::fusion::at_key<VMGDataFused::nuclei_shortrange>(instance) = nuclei_short_range_energy;
//      LOG(0, "Short-range correction energy from potential integral of level " << level << " is "
//          << short_range_energy << ".");
      boost::fusion::at_key<VMGDataFused::total_longrange>(instance) =
          boost::fusion::at_key<VMGDataFused::electron_longrange>(instance)
          + 2.*boost::fusion::at_key<VMGDataFused::mixed_longrange>(instance)
          + boost::fusion::at_key<VMGDataFused::nuclei_longrange>(instance);
      boost::fusion::at_key<VMGDataFused::total_shortrange>(instance) =
          boost::fusion::at_key<VMGDataFused::electron_shortrange>(instance)
          + 2.*boost::fusion::at_key<VMGDataFused::mixed_shortrange>(instance)
          + boost::fusion::at_key<VMGDataFused::nuclei_shortrange>(instance);
      Result_LongRangeIntegrated_fused.push_back(instance);

      if (hasLongRangeForces()) {
        VMGDataLongRangeForceMap_t forceinstance;
        IndexedVectors fullforces( fullindices, fullsolutionData[level-1].forces);
        IndexedVectors longrangeforces =
            boost::fusion::at_key<VMGDataFused::forces>(Result_ForceLongRange_fused[level-1]);
        boost::fusion::at_key<VMGDataFused::forces_shortrange>(forceinstance) =
            fullforces;
        fullforces -= longrangeforces;
        boost::fusion::at_key<VMGDataFused::forces_longrange>(forceinstance) =
            fullforces;
        Result_ForcesLongRangeIntegrated_fused.push_back(forceinstance);
      }
    }
//    {
//      //    LOG(0, "Remaining long-range energy from energy_potential is " << full_sample_solution.integral()-epotentialSummer.getFullContribution() << ".");
//      SamplingGrid full_sample_solution = fullsolutionData.back().sampled_potential;
//      const SamplingGrid &short_range_correction =
//          boost::fusion::at_key<VMGDataFused::sampled_potential>(Result_GridLongRange_fused.back()).getFullContribution();
//      full_sample_solution -= short_range_correction;
//      // multiply element-wise with charge distribution
//      LOG(0, "Remaining long-range potential integral is " << full_sample_solution.integral() << ".");
//      LOG(0, "Short-range correction potential integral of level is " << short_range_correction.integral() << ".");
//      LOG(0, "Remaining long-range energy from potential integral is "
//          << full_sample_solution.integral(full_sample.back()) << ".");
//      LOG(0, "Short-range correction energy from potential integral is "
//          << short_range_correction.integral(full_sample.back()) << ".");
//
//      double e_long = fullsolutionData.back().e_long;
//      e_long -= boost::fusion::at_key<VMGDataFused::energy_long>(Result_LongRange_fused.back()).getFullContribution();
//      LOG(0, "Remaining long-range energy is " << e_long << ".");
//    }

    // TODO: Extract long-range corrections to forces
    // NOTE: potential is in atomic length units, NOT IN ANGSTROEM!

  }
}