/*
 * 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)
{
  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)
{
  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>(
        fragmentData, 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>(
        longrangeData, VMGMatrixNrLookup, container, subsetmap,
        Result_LongRange_fused, Result_perIndexSet_LongRange);

    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 =
          boost::fusion::at_key<MPQCDataFused::sampled_grid>(Result_Grid_fused[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_LongRange_fused[level-1]);
      double electron_short_range_energy = short_range_correction.integral();
      full_sample_solution -= short_range_correction;
      electron_solution_energy -= electron_short_range_energy;
      ASSERT( fabs(electron_solution_energy - full_sample_solution.integral()) < 1e-7,
          "FragmentationLongRangeResults::operator() - integral and energy are not exchangeable.");

      // 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);
    }
//    {
//      //    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_LongRange_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!

  }
}