/*
 * 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/>.
 */

/*
 * InterfaceVMGJob.cpp
 *
 *  Created on: 10.06.2012
 *      Author: Frederik Heber
 */

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

#ifdef HAVE_MPI
#include "mpi.h"
#endif

#include "base/vector.hpp"
#include "base/math.hpp"
#include "comm/comm.hpp"
#include "grid/grid.hpp"
#include "grid/multigrid.hpp"
#include "units/particle/comm_mpi_particle.hpp"
#include "units/particle/interpolation.hpp"
#include "units/particle/linked_cell_list.hpp"
#include "mg.hpp"

#include "InterfaceVMGJob.hpp"

#include "CodePatterns/MemDebug.hpp"

#include <cmath>
#include <iostream>
#include <limits>

#include "CodePatterns/Log.hpp"

#include "Fragmentation/Summation/SetValues/FragmentForces.hpp"
#include "Jobs/WindowGrid_converter.hpp"
#include "Jobs/ChargeSmearer.hpp"

using namespace VMG;
using VMGInterfaces::InterfaceVMGJob;

InterfaceVMGJob::InterfaceVMGJob(const SamplingGrid &_sampled_input,
    VMGData &_returndata,
    const std::vector< std::vector<double> > &_particle_positions,
    const std::vector< double > &_particle_charges,
    VMG::Boundary boundary,
    int levelMin,
    int levelMax,
    const VMG::Vector &_box_begin,
    vmg_float _box_end,
    const int& near_field_cells,
    const ImportParticles_t _ImportParticles,
    const bool _DoPrintDebug,
    const bool _DoSmearCharges,
    int coarseningSteps,
    double alpha) :
  VMG::Interface(boundary, levelMin, levelMax,
      _box_begin, _box_end, coarseningSteps, alpha),
  nfc(near_field_cells),
  meshwidth(Extent(MaxLevel()).MeshWidth().Max()),
  spl(nfc, meshwidth),
  sampled_input(_sampled_input),
  returndata(_returndata),
  level(levelMax),
  ImportParticles(_ImportParticles),
  DoPrintDebug(_DoPrintDebug),
  OpenBoundaryCondition(boundary[0] == VMG::Open),
  DoSmearCharges(_DoSmearCharges)
{
  for (size_t i=0;i<3;++i) {
    box_begin[i] = _box_begin[i];
    box_end[i] = _box_end;
  }
  std::vector< std::vector<double> >::const_iterator positer = _particle_positions.begin();
  std::vector<double>::const_iterator chargeiter = _particle_charges.begin();
  double pos[3];
  for (; positer != _particle_positions.end(); ++positer, ++chargeiter) {
    ASSERT( (*positer).size() == 3,
        "InterfaceVMGJob::InterfaceVMGJob() - particle "
        +toString(distance(_particle_positions.begin(), positer))+" has not exactly 3 coordinates.");
    for (size_t i=0;i<3;++i)
      pos[i] = (*positer)[i];
    particles.push_back(Particle::Particle(pos, *chargeiter));
  }
}

void InterfaceVMGJob::ImportRightHandSide(Multigrid& multigrid)
{
  Index i;
  Vector pos;
  //  VMG::TempGrid *temp_grid = new VMG::TempGrid(129, 0, 0., 1.);

  Grid& grid = multigrid(multigrid.MaxLevel());
  grid.Clear();
  //grid.ClearBoundary(); // we don't have a boundary under periodic boundary conditions

  // print debugging info on grid size
  LOG(1, "INFO: Mesh has extent " << grid.Extent().MeshWidth() << ".");
  const int gridpoints = pow(2, level);
  LOG(1, "INFO: gridpoints on finest level are " << gridpoints << ".");
  LOG(1, "INFO: "
      << "X in [" << grid.Local().Begin().X() << "," << grid.Local().End().X() << "],"
      << "Y in [" << grid.Local().Begin().Y() << "," << grid.Local().End().Y() << "],"
      << "Z in [" << grid.Local().Begin().Z() << "," << grid.Local().End().Z() << "].");

  /// 1. assign nuclei as smeared-out charges to the grid

  /*
   * Charge assignment on the grid
   */
  Particle::CommMPI& comm = *dynamic_cast<Particle::CommMPI*>(MG::GetComm());
  Grid& particle_grid = comm.GetParticleGrid();
  particle_grid.Clear();

  // distribute particles
  particles.clear();
  comm.CommParticles(grid, particles);

  assert(particle_grid.Global().LocalSize().IsComponentwiseGreater(
      VMG::MG::GetFactory().GetObjectStorageVal<int>("PARTICLE_NEAR_FIELD_CELLS")));

  if (ImportParticles == DoImportParticles) {
    // create smeared-out particle charges on particle_grid via splines
    LOG(1, "INFO: Creating particle grid for " << particles.size() << " particles.");
    for (std::list<Particle::Particle>::iterator iter = particles.begin();
        iter != particles.end(); ++iter) {
      LOG(2, "DEBUG: Current particle is at " << (*iter).Pos()
          << " with charge " << (*iter).Charge() << ".");
      spl.SetSpline(particle_grid, *iter);
    }
  }

  // Communicate charges over halo
  comm.CommFromGhosts(particle_grid);

  if (DoPrintDebug) {
    // print nuclei grid to vtk
    comm.PrintGrid(particle_grid, "Sampled Nuclei Density");
  }

  // add sampled electron charge density onto grid
  if (DoSmearCharges) {
    ChargeSmearer &smearer = ChargeSmearer::getInstance();
    smearer.initializeSplineArray(spl, nfc, meshwidth);
  }
  WindowGrid_converter::addWindowOntoGrid(
      grid,
      sampled_input,
      1.,
      OpenBoundaryCondition,
      DoSmearCharges);

  if (DoPrintDebug) {
    // print electron grid to vtk
    comm.PrintGrid(grid, "Sampled Electron Density");
  }

  // add particle_grid onto grid
  for (int i=0; i<grid.Local().Size().X(); ++i)
    for (int j=0; j<grid.Local().Size().Y(); ++j)
      for (int k=0; k<grid.Local().Size().Z(); ++k)
  grid(grid.Local().Begin().X() + i,
       grid.Local().Begin().Y() + j,
       grid.Local().Begin().Z() + k) = 4.0 * VMG::Math::pi * (
           grid(grid.Local().Begin().X() + i,
                  grid.Local().Begin().Y() + j,
                  grid.Local().Begin().Z() + k) +
    particle_grid.GetVal(particle_grid.Local().Begin().X() + i,
             particle_grid.Local().Begin().Y() + j,
             particle_grid.Local().Begin().Z() + k));

  // calculate sum over grid times h^3 as check, should be roughly zero
  const double element_volume = grid.Extent().MeshWidth().Product();
  double charge_sum = 0.0;
  for (Grid::iterator grid_iter = grid.Iterators().Local().Begin();
      grid_iter != grid.Iterators().Local().End();
      ++grid_iter)
    charge_sum += grid.GetVal(*grid_iter);
  charge_sum = element_volume * comm.GlobalSum(charge_sum);
  comm.PrintOnce(Debug, "Grid charge integral: %e", charge_sum/(4.0 * VMG::Math::pi));

  if (DoPrintDebug) {
    // print total grid to vtk
    comm.PrintGrid(grid, "Total Charge Density");
  }

//  delete temp_grid;
}

void InterfaceVMGJob::ExportSolution(Grid& grid)
{
  /// sample the obtained potential to evaluate with the electron charge density

  // grid now contains the sough-for potential
  //Comm& comm = *MG::GetComm();
  Particle::CommMPI& comm = *dynamic_cast<Particle::CommMPI*>(MG::GetComm());


  if (DoPrintDebug) {
    // print output grid to vtk
    comm.PrintGrid(grid, "Potential Solution");
  }

  // obtain sampled potential from grid
  returndata.sampled_potential.setWindow(
      box_begin,
      box_end
      );
  WindowGrid_converter::addGridOntoWindow(
      grid,
      returndata.sampled_potential,
      +1.,
      OpenBoundaryCondition
      );

  // calculate integral over potential as long-range energy contribution
  const double element_volume =
      grid.Extent().MeshWidth().X() * grid.Extent().MeshWidth().Y() * grid.Extent().MeshWidth().Z();
  Grid::iterator grid_iter;
  double potential_sum = 0.0;
  for (grid_iter=grid.Iterators().Local().Begin(); grid_iter!=grid.Iterators().Local().End(); ++grid_iter)
    potential_sum += grid.GetVal(*grid_iter);
  potential_sum = element_volume * comm.GlobalSum(potential_sum);
  comm.PrintOnce(Debug, "Grid potential sum: %e", potential_sum);

  {
    Grid::iterator grid_iter = grid.Iterators().Local().Begin();
    comm.PrintOnce(Debug, "Grid potential at (0,0,0): %e", grid.GetVal(*grid_iter));
  }

  //Particle::CommMPI& comm = *dynamic_cast<Particle::CommMPI*>(MG::GetComm());  returndata.e_long = potential_sum;

  /// Calculate potential energy of nuclei

  vmg_float e = 0.0;
  vmg_float e_long = 0.0;
  vmg_float e_self = 0.0;
  vmg_float e_short_peak = 0.0;
  vmg_float e_short_spline = 0.0;

  Factory& factory = MG::GetFactory();

  /*
   * Get parameters and arrays
   */
  const vmg_int& near_field_cells = factory.GetObjectStorageVal<int>("PARTICLE_NEAR_FIELD_CELLS");
  const vmg_int& interpolation_degree = factory.GetObjectStorageVal<int>("PARTICLE_INTERPOLATION_DEGREE");

  Particle::Interpolation ip(interpolation_degree);

  const vmg_float r_cut = near_field_cells * grid.Extent().MeshWidth().Max();

  /*
   * Copy potential values to a grid with sufficiently large halo size.
   * This may be optimized in future.
   * The parameters of this grid have been set in the import step.
   */
  Grid& particle_grid = comm.GetParticleGrid();

  {
    Index i;
    for (i.X()=0; i.X()<grid.Local().Size().X(); ++i.X())
      for (i.Y()=0; i.Y()<grid.Local().Size().Y(); ++i.Y())
        for (i.Z()=0; i.Z()<grid.Local().Size().Z(); ++i.Z())
          particle_grid(i + particle_grid.Local().Begin()) = grid.GetVal(i + grid.Local().Begin());
    comm.CommToGhosts(particle_grid);
  }

  /*
   * Compute potentials
   */
  Particle::LinkedCellList lc(particles, near_field_cells, grid);
  Particle::LinkedCellList::iterator p1, p2;
  Grid::iterator iter;

  comm.CommLCListToGhosts(lc);

  for (int i=lc.Local().Begin().X(); i<lc.Local().End().X(); ++i)
    for (int j=lc.Local().Begin().Y(); j<lc.Local().End().Y(); ++j)
      for (int k=lc.Local().Begin().Z(); k<lc.Local().End().Z(); ++k) {

  if (lc(i,j,k).size() > 0)
    ip.ComputeCoefficients(particle_grid, Index(i,j,k) - lc.Local().Begin() + particle_grid.Local().Begin());

  for (p1=lc(i,j,k).begin(); p1!=lc(i,j,k).end(); ++p1) {

    // Interpolate long-range part of potential and electric field
    ip.Evaluate(**p1);

    // Subtract self-induced potential
    (*p1)->Pot() -= (*p1)->Charge() * spl.GetAntiDerivativeAtZero();

    e_long += 0.5 * (*p1)->Charge() * ip.EvaluatePotentialLR(**p1);
    e_self += 0.5 * (*p1)->Charge() * (*p1)->Charge() * spl.GetAntiDerivativeAtZero();

    for (int dx=-1*near_field_cells; dx<=near_field_cells; ++dx)
      for (int dy=-1*near_field_cells; dy<=near_field_cells; ++dy)
        for (int dz=-1*near_field_cells; dz<=near_field_cells; ++dz) {

    for (p2=lc(i+dx,j+dy,k+dz).begin(); p2!=lc(i+dx,j+dy,k+dz).end(); ++p2)

      if (*p1 != *p2) {

        const Vector dir = (*p1)->Pos() - (*p2)->Pos();
        const vmg_float length = dir.Length();

        if ((length < r_cut) && (length > std::numeric_limits<double>::epsilon())) {

          (*p1)->Pot() += (*p2)->Charge() / length * (1.0 + spl.EvaluatePotential(length));
          (*p1)->Field() += (*p2)->Charge() * dir * spl.EvaluateField(length);

          e_short_peak += 0.5 * (*p1)->Charge() * (*p2)->Charge() / length;
          e_short_spline += 0.5 * (*p1)->Charge() * (*p2)->Charge() / length * spl.EvaluatePotential(length);
        }
      }
        }
  }
      }

  const vmg_int& num_particles_local = factory.GetObjectStorageVal<vmg_int>("PARTICLE_NUM_LOCAL");

  /* Remove average force term */
//  if (!particles.empty()) {
//    Vector average_force = 0.0;
//    for (std::list<Particle::Particle>::const_iterator iter=particles.begin(); iter!=particles.end(); ++iter)
//      average_force += iter->Charge() * iter->Field();
//    const vmg_int num_particles_global = comm.GlobalSum(num_particles_local);
//    average_force /= (double)num_particles_global;
//    comm.GlobalSumArray(average_force.vec(), 3);
//    for (std::list<Particle::Particle>::iterator iter=particles.begin(); iter!=particles.end(); ++iter)
//      iter->Field() -= average_force / iter->Charge();
//    comm.PrintOnce(Debug, "Average force term is:         %e %e %e", average_force[0], average_force[1], average_force[2]);
//  }

  comm.CommParticlesBack(particles);

  vmg_float* q = factory.GetObjectStorageArray<vmg_float>("PARTICLE_CHARGE_ARRAY");
  const vmg_float* p = factory.GetObjectStorageArray<vmg_float>("PARTICLE_POTENTIAL_ARRAY");
  const vmg_float* f = factory.GetObjectStorageArray<vmg_float>("PARTICLE_FIELD_ARRAY");

  // extract forces
  if (!particles.empty()) {
    size_t index = 0;
    returndata.forces.resize(
        num_particles_local, FragmentForces::force_t(3, 0.) );
    for (FragmentForces::forces_t::iterator iter = returndata.forces.begin();
        iter != returndata.forces.end(); ++iter) {
      comm.PrintOnce(Debug, "%d force vector:        %e %e %e", (index/3)+1, f[index+0], f[index+1], f[index+2]);
      for (size_t i=0;i<3;++i)
          (*iter)[i] = f[index++];
    }
    returndata.hasForces = true;
  }

  e_long = comm.GlobalSumRoot(e_long);
  e_short_peak = comm.GlobalSumRoot(e_short_peak);
  e_short_spline = comm.GlobalSumRoot(e_short_spline);
  e_self = comm.GlobalSumRoot(e_self);

  for (int j=0; j<num_particles_local; ++j)
    e += 0.5 * p[j] * q[j];
  e = comm.GlobalSumRoot(e);

  comm.PrintOnce(Debug, "E_long:         %e", e_long);
  comm.PrintOnce(Debug, "E_short_peak:   %e", e_short_peak);
  comm.PrintOnce(Debug, "E_short_spline: %e", e_short_spline);
  comm.PrintOnce(Debug, "E_self:         %e", e_self);
  comm.PrintOnce(Debug, "E_total:        %e", e);
  comm.PrintOnce(Debug, "E_total*:       %e", e_long + e_short_peak + e_short_spline - e_self);

  returndata.nuclei_long = e_long;
  returndata.electron_long = e_long;
}