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

/*
 * StretchBondAction.cpp
 *
 *  Created on: Sep 26, 2012
 *      Author: heber
 */

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

#include <boost/graph/adjacency_list.hpp>
#include <boost/graph/breadth_first_search.hpp>
#include <boost/graph/subgraph.hpp>

//#include "CodePatterns/MemDebug.hpp"

#include "Actions/MoleculeAction/StretchBondAction.hpp"

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

#include "LinearAlgebra/Plane.hpp"

#include "Atom/atom.hpp"
#include "Bond/bond.hpp"
#include "Descriptors/AtomIdDescriptor.hpp"
#include "molecule.hpp"
#include "World.hpp"

using namespace MoleCuilder;

// and construct the stuff
#include "StretchBondAction.def"
#include "Action_impl_pre.hpp"


/**
 * I have no idea why this is so complicated with BGL ...
 *
 * This is taken from the book "The Boost Graph Library: User Guide and Reference Manual, Portable Documents",
 * chapter "Basic Graph Algorithms", example on calculating the bacon number.
 */
template <typename DistanceMap>
class distance_recorder : public boost::default_bfs_visitor
{
public:
  distance_recorder(DistanceMap dist) : d(dist) {}

  template <typename Edge, typename Graph>
  void tree_edge(Edge e, const Graph &g) const {
    typename boost::graph_traits<Graph>::vertex_descriptor u = source(e,g), v = target(e,g);
    d[v] = d[u] + 1;
  }

private:
  DistanceMap d;
};

template <typename DistanceMap>
distance_recorder<DistanceMap> record_distance(DistanceMap d)
{
  return distance_recorder<DistanceMap>(d);
}

/** =========== define the function ====================== */
ActionState::ptr MoleculeStretchBondAction::performCall()
{
  // check preconditions
  const std::vector< atom *> atoms = World::getInstance().getSelectedAtoms();
  if (atoms.size() != 2) {
    STATUS("Exactly two atoms must be selected.");
    return Action::failure;
  }
  std::vector<atomId_t> atomids(2);
  atomids[0] = atoms[0]->getId();
  atomids[1] = atoms[1]->getId();
  std::sort(atomids.begin(), atomids.end());
  LOG(1, "DEBUG: Selected nodes are " << atomids);
  molecule *mol = World::getInstance().
      getMolecule(MoleculeById(atoms[0]->getMolecule()->getId()));
  if (mol != atoms[1]->getMolecule()) {
    STATUS("The two selected atoms must belong to the same molecule.");
    return Action::failure;
  }
  // gather undo information
  const double olddistance = atoms[0]->getPosition().distance(atoms[1]->getPosition());
  const double newdistance = params.bonddistance.get();
  LOG(1, "INFO: Old bond distance is " << olddistance << ", stretching to " << newdistance << ".");

  // Assume the selected bond splits the molecule into two parts, each one on
  // either side of the bond. We need to perform a BFS from each bond partner
  // not using the selected bond. Therefrom, we obtain two sets of atoms/nodes.
  // If both are disjoint, the bond is not contained in a cycle and we simply
  // shift either set as desired. If not, then we simply shift each atom,
  // leaving the other positions untouched.
  typedef boost::adjacency_list < boost::vecS, boost::vecS, boost::undirectedS,
      boost::no_property, boost::no_property > UndirectedGraph;

  // convert BondGraph into boost::graph
  UndirectedGraph molgraph(mol->getAtomCount());
  size_t no_edges = 0;
  for(molecule::const_iterator iter = const_cast<const molecule * const>(mol)->begin();
      iter != const_cast<const molecule * const>(mol)->end(); ++iter) {
    LOG(2, "DEBUG: Looking at node " << (*iter)->getId());
    const BondList& ListOfBonds = (*iter)->getListOfBonds();
    for(BondList::const_iterator bonditer = ListOfBonds.begin();
        bonditer != ListOfBonds.end(); ++bonditer) {
      const unsigned int &leftid = (*bonditer)->leftatom->getId();
      const unsigned int &rightid = (*bonditer)->rightatom->getId();
      // only add each edge once and do not add selected edge
      if ((leftid == (*iter)->getId())
          && ((leftid != atomids[0]) || (rightid != atomids[1]))) {
        LOG(1, "DEBUG: ADDING edge " << leftid << " <-> " << rightid);
        boost::add_edge(leftid, rightid, molgraph);
        ++no_edges;
      } else {
        LOG(1, "DEBUG: Discarding edge " << leftid << " <-> " << rightid);
      }
    }
  }
  typedef boost::property_map < boost::adjacency_list <>, boost::vertex_index_t >::type index_map_t;
  index_map_t index_map = boost::get(boost::vertex_index, molgraph);
  const size_t num_vertices = boost::num_vertices(molgraph);
  LOG(2, "DEBUG: We have " << num_vertices << " nodes and " << no_edges
      << " edges in the molecule graph.");

  std::vector< std::vector<size_t> > distances;
  for (size_t i=0;i<2;++i) {
    distances.push_back(std::vector<size_t>(num_vertices, num_vertices+1)); // set distance to num+1
    distances[i][atomids[i]] = 0;
    boost::breadth_first_search(
        molgraph,
        boost::vertex(atomids[i], molgraph),
        boost::visitor(record_distance(&(distances[i][0]))));
    LOG(3, "DEBUG: From atom #" << atomids[i]
        << " BFS discovered the following distances " << distances[i]);
  }

  const Vector NormalVector = (atoms[0]->getPosition() - atoms[1]->getPosition())* (1./olddistance);
  const double shift = 0.5*(newdistance - olddistance);
  std::vector<Vector> Shift(2);
  Shift[0] = shift * NormalVector;
  Shift[1] = -shift * NormalVector;
  Box &domain = World::getInstance().getDomain();
  std::vector< std::vector<size_t> > bondside_sets(2);

  // Check whether there are common nodes in each set of distances
  for (size_t i=0;i<num_vertices;++i)
    if ((distances[0][i] != (num_vertices+1)) && (distances[1][i] != (num_vertices+1))) {
      ELOG(2, "Node #" << i << " is reachable from either side of bond, hence must by cyclic."
          << " Shifting only bond partners.");
      for(size_t j=0;j<2;++j) {
        bondside_sets[j].push_back(atoms[j]->getId());
        const Vector &position = atoms[j]->getPosition();
        atoms[j]->setPosition( domain.enforceBoundaryConditions(position+Shift[j]) );
      }
      break;
    }

  // go through the molecule and stretch each atom in either set of nodes
  if (bondside_sets[0].empty()) {
    for (molecule::iterator iter = mol->begin(); iter != mol->end(); ++iter) {
      const Vector &position = (*iter)->getPosition();
      // for each atom determine in which set of nodes it is and shift accordingly
      size_t i=0;
      const size_t nodeindex = boost::get(index_map, boost::vertex((*iter)->getId(), molgraph));
      for (;i<2;++i) {
        if (distances[i][nodeindex] != (num_vertices+1)) {
          (*iter)->setPosition( domain.enforceBoundaryConditions(position+Shift[i]) );
          bondside_sets[i].push_back((*iter)->getId());
          break;
        }
      }
      if (i==2) {
        ELOG(1, "Atom " << *iter << " is not contained on either side of bond? Undoing done shifts");
        // Have to undo shifts
        for (i=0;i<2;++i) {
          for (std::vector<size_t>::const_iterator iter = bondside_sets[i].begin();
              iter != bondside_sets[i].end(); ++iter) {
            atom &walker = *World::getInstance().getAtom(AtomById(*iter));
            const Vector &position = walker.getPosition();
            walker.setPosition( domain.enforceBoundaryConditions(position-Shift[i]) );
          }
        }
        return Action::failure;
      }
    }
  }

  MoleculeStretchBondState *UndoState = new MoleculeStretchBondState(Shift, bondside_sets, mol, params);
  return ActionState::ptr(UndoState);
}

ActionState::ptr MoleculeStretchBondAction::performUndo(ActionState::ptr _state) {
  MoleculeStretchBondState *state = assert_cast<MoleculeStretchBondState*>(_state.get());

  // use given plane to undo
  Box &domain = World::getInstance().getDomain();
  for (size_t i=0;i<2;++i) {
    for (std::vector<size_t>::const_iterator iter = state->bondside_sets[i].begin();
        iter != state->bondside_sets[i].end(); ++iter) {
      atom &walker = *World::getInstance().getAtom(AtomById(*iter));
      const Vector &position = walker.getPosition();
      walker.setPosition( domain.enforceBoundaryConditions(position-state->Shift[i]) );
    }
  }

  return ActionState::ptr(_state);
}

ActionState::ptr MoleculeStretchBondAction::performRedo(ActionState::ptr _state){
  MoleculeStretchBondState *state = assert_cast<MoleculeStretchBondState*>(_state.get());

  Box &domain = World::getInstance().getDomain();
  for (size_t i=0;i<2;++i) {
    for (std::vector<size_t>::const_iterator iter = state->bondside_sets[i].begin();
        iter != state->bondside_sets[i].end(); ++iter) {
      atom &walker = *World::getInstance().getAtom(AtomById(*iter));
      const Vector &position = walker.getPosition();
      walker.setPosition( domain.enforceBoundaryConditions(position+state->Shift[i]) );
    }
  }

  return ActionState::ptr(_state);
}

bool MoleculeStretchBondAction::canUndo() {
  return true;
}

bool MoleculeStretchBondAction::shouldUndo() {
  return true;
}
/** =========== end of function ====================== */