/*
 * Project: MoleCuilder
 * Description: creates and alters molecular systems
 * Copyright (C)  2013 Frederik Heber. All rights reserved.
 * Please see the LICENSE file or "Copyright notice" in builder.cpp for details.
 *
 *
 *   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/>.
 */

/*
 * PartialNucleiChargeFitter.cpp
 *
 *  Created on: 12.05.2013
 *      Author: heber
 */

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

#include "CodePatterns/MemDebug.hpp"

#include "PartialNucleiChargeFitter.hpp"

#include <cmath>
#include <fstream>
#include <limits>
#include <numeric>

#include "LinearAlgebra/MatrixContent.hpp"
#include "LinearAlgebra/VectorContent.hpp"

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

#include "Fragmentation/Summation/SetValues/SamplingGrid.hpp"

PartialNucleiChargeFitter::dimensions_t
PartialNucleiChargeFitter::getGridDimensions(const SamplingGrid &grid) const
{
  // convert sampled potential into a vector
  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
  dimensions_t total(3,0);
  for(size_t index=0;index<3;++index) {
    const double delta = grid.getDeltaPerAxis(index);
    // delta is conversion factor from box length to discrete length, i.e. number of points
    total[index] = (grid.end[index] - grid.begin[index])/delta+round_offset;
  }
  return total;
}

PartialNucleiChargeFitter::PartialNucleiChargeFitter(
    const SamplingGrid &grid,
    const positions_t &_positions,
    const double _threshold) :
      total(getGridDimensions(grid)),
      SampledPotential(std::accumulate(total.begin(), total.end(), 1, std::multiplies<double>())),
      grid_properties(static_cast<const SamplingGridProperties &>(grid)),
      positions(_positions),
      PotentialFromCharges(NULL),
      PartialCharges(NULL),
      threshold(_threshold)
{
  // we must take care of the "window", i.e. there may be less entries in sampled_grid
  // vector as we would expect from size of grid ((2^level)^3) as 0-entries have been
  // omitted.
  size_t pre_offset[3];
  size_t post_offset[3];
  size_t length[3];
  size_t total[3];
  grid.getDiscreteWindowIndices(
      grid.begin, grid.end,
      grid.begin_window, grid.end_window,
      pre_offset,
      post_offset,
      length,
      total
      );
  const size_t calculated_size = length[0]*length[1]*length[2];
  ASSERT( calculated_size == grid.sampled_grid.size(),
      "PartialNucleiChargeFitter::PartialNucleiChargeFitter() - grid does not match size indicated by its window.");

  const double potential_sum = std::accumulate(grid.sampled_grid.begin(), grid.sampled_grid.end(), 0.);
  if ( fabs(potential_sum) > std::numeric_limits<double>::epsilon()*1e4 ) {
    ELOG(2, "Potential sum is not less than "
        << std::numeric_limits<double>::epsilon()*1e4 << " but " 
        << potential_sum << ".");
  }

  SamplingGrid::sampledvalues_t::const_iterator griditer = grid.sampled_grid.begin();
  size_t index=0;
  size_t N[3];
  Vector grid_position; // position of grid point in real domain
  size_t masked_points = 0;
  // store step length per axis
  double delta[3];
  for (size_t i=0;i<3;++i)
    delta[i] = grid_properties.getDeltaPerAxis(i);
  /// convert sampled potential into a vector
  grid_position[0] = grid_properties.begin[0];
  for(N[0]=0; N[0] < pre_offset[0]; ++N[0]) {
    grid_position[1] = grid_properties.begin[1];
    for(N[1]=0; N[1] < total[1]; ++N[1]) {
      grid_position[2] = grid_properties.begin[2];
      for(N[2]=0; N[2] < total[2]; ++N[2]) {
        const_cast<VectorContent &>(SampledPotential)[index++] = 0.;
        grid_position[2] += delta[2];
      }
      grid_position[1] += delta[1];
    }
    grid_position[0] += delta[0];
  }
  for(N[0]=0; N[0] < length[0]; ++N[0]) {
    grid_position[1] = grid_properties.begin[1];
    for(N[1]=0; N[1] < pre_offset[1]; ++N[1]) {
      grid_position[2] = grid_properties.begin[2];
      for(N[2]=0; N[2] < total[2]; ++N[2]) {
        const_cast<VectorContent &>(SampledPotential)[index++] = 0.;
        grid_position[2] += delta[2];
      }
      grid_position[1] += delta[1];
    }
    for(N[1]=0; N[1] < length[1]; ++N[1]) {
      grid_position[2] = grid_properties.begin[2];
      for(N[2]=0; N[2] < pre_offset[2]; ++N[2]) {
        const_cast<VectorContent &>(SampledPotential)[index++] = 0.;
        grid_position[2] += delta[2];
      }
      for(N[2]=0; N[2] < length[2]; ++N[2]) {
        if (isGridPointSettable(positions, grid_position))
          const_cast<VectorContent &>(SampledPotential)[index++] = *griditer++;
        else {
          // skip point
          ++griditer;
          ++masked_points;
          const_cast<VectorContent &>(SampledPotential)[index++] = 0.;
        }
        grid_position[2] += delta[2];
      }
      for(N[2]=0; N[2] < post_offset[2]; ++N[2]) {
        const_cast<VectorContent &>(SampledPotential)[index++] = 0.;
        grid_position[2] += delta[2];
      }
      grid_position[1] += delta[1];
    }
    for(N[1]=0; N[1] < post_offset[1]; ++N[1]) {
      grid_position[2] = grid_properties.begin[2];
      for(N[2]=0; N[2] < total[2]; ++N[2]) {
        const_cast<VectorContent &>(SampledPotential)[index++] = 0.;
        grid_position[2] += delta[2];
      }
      grid_position[1] += delta[1];
    }
    grid_position[0] += delta[0];
  }
  for(N[0]=0; N[0] < post_offset[0]; ++N[0]) {
    grid_position[1] = grid_properties.begin[1];
    for(N[1]=0; N[1] < total[1]; ++N[1]) {
      grid_position[2] = grid_properties.begin[2];
      for(N[2]=0; N[2] < total[2]; ++N[2]) {
        const_cast<VectorContent &>(SampledPotential)[index++] = 0.;
        grid_position[2] += delta[2];
      }
      grid_position[1] += delta[1];
    }
    grid_position[0] += delta[0];
  }
  // set remainder of points to zero
  ASSERT( index == SampledPotential.getDimension(),
      "PartialNucleiChargeFitter::PartialNucleiChargeFitter() - not enough or more than calculated sample points.");

#ifndef NDEBUG
  // write vector as paraview csv file file
  {
    size_t N[3];
    size_t index = 0;
    std::ofstream paraview_output("solution.csv");
    paraview_output << "x coord,y coord,z coord,scalar" << std::endl;
    for(N[0]=0; N[0] < total[0]; ++N[0]) {
      for(N[1]=0; N[1] < total[1]; ++N[1]) {
        for(N[2]=0; N[2] < total[2]; ++N[2]) {
          paraview_output
              << (double)N[0]/(double)total[0] << ","
              << (double)N[1]/(double)total[1] << ","
              << (double)N[2]/(double)total[2] << ","
              << SampledPotential.at(index++) << std::endl;
        }
      }
    }
    paraview_output.close();
  }
#endif

  LOG(1, "INFO: I masked " << masked_points << " points in right-hand-side.");
//  LOG(4, "DEBUG: Right-hand side vector is " << SampledPotential << ".");
}

bool PartialNucleiChargeFitter::isGridPointSettable(
    const positions_t &_positions,
    const Vector &grid_position) const
{
  bool status = true;
  for (positions_t::const_iterator iter = _positions.begin();
      iter != _positions.end(); ++iter) {
    status &= grid_position.DistanceSquared(*iter) > threshold*threshold;
  }
  return status;
}

PartialNucleiChargeFitter::~PartialNucleiChargeFitter()
{
  if (PartialCharges != NULL)
    delete PartialCharges;

  if (PotentialFromCharges != NULL)
    delete PotentialFromCharges;
}


void PartialNucleiChargeFitter::constructMatrix()
{
  const size_t rows = SampledPotential.getDimension();
  const size_t cols = positions.size();

  // allocate memory for PotentialFromCharges
  if (PotentialFromCharges != NULL) {
    delete PotentialFromCharges;
    PotentialFromCharges = NULL;
  }
  PotentialFromCharges = new MatrixContent( rows, cols );
  // store step length per axis
  double delta[3];
  for (size_t i=0;i<3;++i)
    delta[i] = grid_properties.getDeltaPerAxis(i);
  // then for each charge ...
  size_t masked_points = 0;
  for (size_t nuclei_index = 0; nuclei_index < cols; ++nuclei_index) {
    // ... calculate potential at each grid position,
    // i.e. step through grid and calculate distance to charge position
    Vector grid_position; // position of grid point in real domain
    grid_position[0] = grid_properties.begin[0];
    size_t N[3];      // discrete grid position
    size_t index = 0; // component of column vector
    for(N[0]=0; N[0] < total[0]; ++N[0]) {
      grid_position[1] = grid_properties.begin[1];
      for(N[1]=0; N[1] < total[1]; ++N[1]) {
        grid_position[2] = grid_properties.begin[2];
        for(N[2]=0; N[2] < total[2]; ++N[2]) {
          if (isGridPointSettable(positions, grid_position)) {
            const double distance = positions[nuclei_index].distance(grid_position);
            ASSERT( distance >= 0,
                "PartialNucleiChargeFitter::constructMatrix() - distance is negative?");
            // Coulomb's constant is 1 in atomic units, see http://en.wikipedia.org/wiki/Atomic_units
            const double epsilon0_au = 1.; //4.*M_PI*0.007957747154594767;
            // ... with epsilon_0 in atom units from http://folk.uio.no/michalj/node72.html
            const double value = 1./(epsilon0_au*distance);
            PotentialFromCharges->at(index++, nuclei_index) = value;
          } else {
            ++masked_points;
            PotentialFromCharges->at(index++, nuclei_index) = 0.;
          }
          grid_position[2] += delta[2];
        }
        grid_position[1] += delta[1];
      }
      grid_position[0] += delta[0];
    }
    ASSERT( index == PotentialFromCharges->getRows(),
        "PartialNucleiChargeFitter::operator() - number of sampled positions "
        +toString(index)+" unequal to set rows "
        +toString(PotentialFromCharges->getRows())+".");
  }

  LOG(1, "INFO: I masked " << masked_points/cols << " points in matrix.");
}

double PartialNucleiChargeFitter::operator()()
{
  // prepare PartialCharges
  if (PartialCharges != NULL) {
    delete PartialCharges;
    PartialCharges = NULL;
  }
  PartialCharges = new VectorContent(positions.size());

  // set up over-determined system's problem matrix A for Ax=b
  // i.e. columns represent potential of a single charge at grid positions
  constructMatrix();

  // solve for x
  *PartialCharges =
      PotentialFromCharges->solveOverdeterminedLinearEquation(
          SampledPotential);

//  LOG(4, "DEBUG: Solution vector is " << (*PotentialFromCharges) * (*PartialCharges) << ".");

  // calculate residual vector
  VectorContent residuum = (*PotentialFromCharges) * (*PartialCharges) - SampledPotential;

#ifndef NDEBUG
  // write solution to file
  writeMatrix();

  // write vector as paraview csv file file
  {
    size_t N[3];
    size_t index = 0;
    std::ofstream paraview_output("residuum.csv");
    paraview_output << "x coord,y coord,z coord,scalar" << std::endl;
    for(N[0]=0; N[0] < total[0]; ++N[0]) {
      for(N[1]=0; N[1] < total[1]; ++N[1]) {
        for(N[2]=0; N[2] < total[2]; ++N[2]) {
          paraview_output
              << (double)N[0]/(double)total[0] << ","
              << (double)N[1]/(double)total[1] << ","
              << (double)N[2]/(double)total[2] << ","
              << residuum.at(index++) << std::endl;
        }
      }
    }
    paraview_output.close();
  }
#endif

  // calculate L1 and L2 errors
  double residuum_l1 = 0.;
  for (size_t i=0; i< residuum.getDimension(); ++i)
    if (residuum_l1 < residuum[i])
      residuum_l1 = residuum[i];
  LOG(1, "INFO: L2-Norm of residuum is " << residuum.Norm() << ".");
  LOG(1, "INFO: L1-Norm of residuum is " << residuum_l1 << ".");

  return residuum.Norm();
}

void PartialNucleiChargeFitter::writeMatrix()
{
  constructMatrix();

  // write matrix as paraview csv file file
    size_t N[3];
    size_t index=0;
    std::string filename = std::string("potential.csv");
    std::ofstream paraview_output(filename.c_str());
    paraview_output << "x coord,y coord,z coord,scalar" << std::endl;
    for(N[0]=0; N[0] < total[0]; ++N[0]) {
      for(N[1]=0; N[1] < total[1]; ++N[1]) {
        for(N[2]=0; N[2] < total[2]; ++N[2]) {
          double sum = 0.;
          for (size_t nuclei_index = 0; nuclei_index < positions.size(); ++nuclei_index) {
            sum+= PotentialFromCharges->at(index, nuclei_index)*PartialCharges->at(nuclei_index);
          }
          paraview_output
              << (double)N[0]/(double)total[0] << ","
              << (double)N[1]/(double)total[1] << ","
              << (double)N[2]/(double)total[2] << ","
              << sum << std::endl;
          index++;
        }
      }
    }
    paraview_output.close();
}

PartialNucleiChargeFitter::charges_t
PartialNucleiChargeFitter::getSolutionAsCharges_t() const
{
  ASSERT( PartialCharges != NULL,
      "PartialNucleiChargeFitter::getSolutionAsCharges_t() - PartialCharges requested prior to calculation.");
  charges_t return_charges(positions.size(), 0.);
  for (size_t i = 0; i < return_charges.size(); ++i)
    return_charges[i] = PartialCharges->at(i);
  return return_charges;
}