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

/*
 * export_numpy.cpp
 *
 *  Created on: Mar 23, 2019
 *      Author: heber
 */

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

//#include "CodePatterns/MemDebug.hpp"

#include<boost/bind.hpp>
#include<boost/function.hpp>
#include<boost/python.hpp>
#include<boost/python/numpy.hpp>

#include "CodePatterns/Assert.hpp"

#include "World.hpp"

namespace p = boost::python;
namespace np = boost::python::numpy;

unsigned int get_num_atoms()
{
  return World::getInstance().countSelectedAtoms();
}

np::ndarray allocate_ndarray(const unsigned int num_atoms)
{
  p::tuple shape = p::make_tuple(num_atoms, 3);
  np::dtype dtype = np::dtype::get_builtin<double>();
  np::ndarray array = np::zeros(shape, dtype);
  return array;
}

np::ndarray get_ndarray(boost::function<const Vector &(const atom &)> &_get_function)
{
  unsigned int num_atoms = get_num_atoms();
  //std::cout << num_atoms << std::endl;
  np::ndarray positions = allocate_ndarray(num_atoms);

  unsigned int ia=0;
  for (World::AtomSelectionConstIterator iter = World::getInstance().beginAtomSelection();
      iter != World::getInstance().endAtomSelection();
      ++iter) {
    const atom & current = *iter->second;
    for (unsigned int i=0;i<NDIM;++i)
      positions[ia][i] = _get_function(current)[i];
    ++ia;
    ASSERT(ia <= num_atoms, "get_ndarray() - more atoms selected than expected.");
  }

  return positions;
}

np::ndarray get_positions()
{
  static boost::function< const Vector&(const atom&) > get_vector =
      boost::bind(&AtomInfo::getPosition, _1);
  return get_ndarray(get_vector);
}

np::ndarray get_velocities()
{
  static boost::function< const Vector&(const atom&) > get_vector =
      boost::bind(&AtomInfo::getAtomicVelocity, _1);
  return get_ndarray(get_vector);
}

np::ndarray get_forces()
{
  static boost::function< const Vector&(const atom&) > get_vector =
      boost::bind(&AtomInfo::getAtomicForce, _1);
  return get_ndarray(get_vector);
}

np::ndarray get_masses()
{
  unsigned int num_atoms = get_num_atoms();
  //std::cout << num_atoms << std::endl;
  p::tuple shape = p::make_tuple(num_atoms);
  np::dtype dtype = np::dtype::get_builtin<double>();
  np::ndarray masses = np::zeros(shape, dtype);

  unsigned int ia=0;
  for (World::AtomSelectionConstIterator iter = World::getInstance().beginAtomSelection();
      iter != World::getInstance().endAtomSelection();
      ++iter) {
    const atom & current = *iter->second;
    masses[ia] = current.getMass();
    ++ia;
    ASSERT(ia <= num_atoms, "get_masses() - more atoms selected than expected.");
  }

  return masses;
}

void set_ndarray(
    const np::ndarray &_positions,
    boost::function<void (atom &, const Vector &)> &_set_function)
{
  unsigned int num_atoms = get_num_atoms();
  
  // check whether shape of array is correct 
  ASSERT( _positions.shape(0) == num_atoms,
    "pyMoleCuilder::set_ndarray() - numpy array has unexpected size.");

  np::ndarray new_positions = _positions.copy();
  unsigned int ia=0;
  Vector temp;
  for (World::AtomSelectionIterator iter = World::getInstance().beginAtomSelection();
      iter != World::getInstance().endAtomSelection();
      ++iter) {
    atom &current = *iter->second;
    for (unsigned int i=0;i<NDIM;++i) {
      //std::cout << p::extract<char const *>(p::str(new_positions[ia][i])) << std::endl;
      temp[i] = p::extract<double>(new_positions[ia][i]);
    }
    _set_function(current, temp);
    ++ia;
    ASSERT(ia <= num_atoms, "set_ndarray() - more atoms selected than expected.");
  }
}

void set_positions(const np::ndarray &new_positions)
{
  static boost::function< void (atom&, const Vector&) > set_vector =
      boost::bind(&AtomInfo::setPosition, _1, _2);
  set_ndarray(new_positions, set_vector);
}

void set_velocities(const np::ndarray &new_positions)
{
  static boost::function< void (atom&, const Vector&) > set_vector =
      boost::bind(&AtomInfo::setAtomicVelocity, _1, _2);
  set_ndarray(new_positions, set_vector);
}

void set_forces(const np::ndarray &new_positions)
{
  static boost::function< void (atom&, const Vector&) > set_vector =
      boost::bind(&AtomInfo::setAtomicForce, _1, _2);
  set_ndarray(new_positions, set_vector);
}

void export_numpy()
{
  p::def("get_positions", get_positions);
  p::def("get_velocities", get_velocities);
  p::def("get_forces", get_forces);
  p::def("get_masses", get_masses);
  p::def("set_positions", set_positions, p::args("position"));
  p::def("set_velocities", set_velocities, p::args("velocity"));
  p::def("set_forces", set_forces, p::args("force"));
}