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

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


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

#include "CodePatterns/MemDebug.hpp"

#include "FragmentationResults.hpp"

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

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

#include "Fragmentation/Converter/DataConverter.hpp"
#include "Fragmentation/KeySetsContainer.hpp"
#include "Fragmentation/Automation/createMatrixNrLookup.hpp"
#include "Fragmentation/Automation/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"


FragmentationResults::FragmentationResults(
    const std::map<JobId_t,MPQCData> &fragmentData,
    std::map<JobId_t,VMGData> &longrangeData,
    const std::vector<VMGData> &fullsolutionData,
    const std::string &KeySetFilename,
    const size_t NoAtoms,
    const std::vector<SamplingGrid> &full_sample)
{
  // create lookup from job nr to fragment number
  size_t MPQCFragmentCounter = 0;
  const std::vector<JobId_t> mpqcjobids = extractJobIds<MPQCData>(fragmentData);
  const std::map< JobId_t, size_t > MPQCMatrixNrLookup =
      createMatrixNrLookup(mpqcjobids, MPQCFragmentCounter);

  size_t VMGFragmentCounter = 0;
  const std::vector<JobId_t> vmgjobids = extractJobIds<VMGData>(longrangeData);
  const std::map< JobId_t, size_t > VMGMatrixNrLookup =
      createMatrixNrLookup(vmgjobids, VMGFragmentCounter);

  // initialise keysets
  KeySetsContainer KeySet;
  KeySetsContainer ForceKeySet;
  {
    // else needs keysets without hydrogens
    std::stringstream filename;
    filename << FRAGMENTPREFIX << KEYSETFILE;
#ifndef NDEBUG
    bool status =
#endif
        KeySet.ParseKeySets(KeySetFilename, filename.str(), MPQCFragmentCounter);
    ASSERT( status,
        "FragmentationResults::FragmentationResults() - ParseKeySets() of "
        +toString(KEYSETFILE)+" failed.");
  }

  {
    // forces need keysets including hydrogens
    std::stringstream filename;
    filename << FRAGMENTPREFIX << FORCESFILE;
#ifndef NDEBUG
    bool status =
#endif
        ForceKeySet.ParseKeySets(KeySetFilename, filename.str(), MPQCFragmentCounter);
    ASSERT( status,
        "FragmentationResults::FragmentationResults() - ParseKeySets() of "
        +toString(FORCESFILE)+" failed.");
  }

  /// prepare for OrthogonalSummation

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

  MaxLevel = subsetmap->getMaximumSetLevel();
  LOG(1, "INFO: Summing up results till level " << MaxLevel << ".");
  /// convert all MPQCData to MPQCDataMap_t
  {
    ASSERT( ForceKeySet.KeySets.size() == fragmentData.size(),
        "FragmentationResults::FragmentationResults() - ForceKeySet's KeySets and fragmentData differ in size.");

    typedef boost::mpl::remove<MPQCDataEnergyVector_t, MPQCDataFused::energy_eigenvalues>::type MPQCDataEnergyVector_noeigenvalues_t;
    Result_Energy_fused =
        OrthogonalSumUpPerLevel<MPQCDataEnergyMap_t, MPQCData, MPQCDataEnergyVector_t>(
            fragmentData, MPQCMatrixNrLookup, container, subsetmap);
    Result_Grid_fused =
      OrthogonalSumUpPerLevel<MPQCDataGridMap_t, MPQCData, MPQCDataGridVector_t>(
          fragmentData, MPQCMatrixNrLookup, container, subsetmap);
    Result_Time_fused =
      SumUpPerLevel<MPQCDataTimeMap_t, MPQCData, MPQCDataTimeVector_t>(
          fragmentData, MPQCMatrixNrLookup, container, subsetmap);
    Result_Fragment_fused =
      OrthogonalSumUpPerLevel<MPQCDataFragmentMap_t, MPQCData, MPQCDataFragmentVector_t>(
          fragmentData, MPQCMatrixNrLookup, container, subsetmap);

    // 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
    Result_LongRange_fused =
      OrthogonalSumUpPerLevel<VMGDataMap_t, VMGData, VMGDataVector_t>(
          longrangeData, VMGMatrixNrLookup, container, subsetmap);

    // force has extra data converter
    std::map<JobId_t, MPQCDataForceMap_t> MPQCData_Force_fused;
    convertMPQCDatatoForceMap(fragmentData, ForceKeySet, MPQCData_Force_fused);
    Result_Force_fused.resize(MaxLevel); // we need the results of correct size already
    AllLevelOrthogonalSummator<MPQCDataForceMap_t> forceSummer(
                subsetmap,
                MPQCData_Force_fused,
                container->getContainer(),
                MPQCMatrixNrLookup,
                Result_Force_fused);
    boost::mpl::for_each<MPQCDataForceVector_t>(boost::ref(forceSummer));

    Result_LongRangeIntegrated_fused.reserve(MaxLevel);
    {
      // NOTE: potential for level 1 is not calculated as saturation hydrogen
      // are not removed on this level yet
      VMGDataLongRangeMap_t instance;
      boost::fusion::at_key<VMGDataFused::electron_longrange>(instance) = 0.;
      boost::fusion::at_key<VMGDataFused::electron_shortrange>(instance) = 0.;
      boost::fusion::at_key<VMGDataFused::nuclei_longrange>(instance) = 0.;
      boost::fusion::at_key<VMGDataFused::nuclei_shortrange>(instance) = 0.;
      boost::fusion::at_key<VMGDataFused::total_longrange>(instance) = 0.;
      boost::fusion::at_key<VMGDataFused::total_shortrange>(instance) = 0.;
      Result_LongRangeIntegrated_fused.push_back(instance);
    }
    for (size_t level = 2; level <= MaxLevel; ++level) {
      // we have to fill in the remainder values in the LongRangeMap by hand
      // weight times correct charge density of the same level
      const SamplingGrid &charge_weight =
          boost::fusion::at_key<MPQCDataFused::sampled_grid>(Result_Grid_fused[level-1]);
      SamplingGrid full_sample_solution = fullsolutionData[level-2].sampled_potential;
      full_sample_solution *= charge_weight;
      const SamplingGrid &short_range_correction =
          boost::fusion::at_key<VMGDataFused::sampled_potential>(Result_LongRange_fused[level-1]);
      full_sample_solution -= short_range_correction;
      double full_solution_energy = fullsolutionData[level-2].e_long;
      const double short_range_energy =
          boost::fusion::at_key<VMGDataFused::energy_long>(Result_LongRange_fused[level-1]);
      full_solution_energy -= short_range_energy;

      // multiply element-wise with charge distribution
      VMGDataLongRangeMap_t instance;
      boost::fusion::at_key<VMGDataFused::electron_longrange>(instance) = .5*full_sample_solution.integral();
//      LOG(0, "Remaining long-range potential integral of level " << level << " is "
//          << full_sample_solution.integral() << ".");
      boost::fusion::at_key<VMGDataFused::electron_shortrange>(instance) = .5*short_range_correction.integral();
//      LOG(0, "Short-range correction potential integral of level " << level << " is "
//          << short_range_correction.integral() << ".");
      boost::fusion::at_key<VMGDataFused::nuclei_longrange>(instance) = full_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) = 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) - full_solution_energy;
      boost::fusion::at_key<VMGDataFused::total_shortrange>(instance) =
          boost::fusion::at_key<VMGDataFused::electron_shortrange>(instance) - short_range_energy;
      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!

  }
}