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

/*
 * StabilityEvaluator.cpp
 *
 *  Created on: Apr 18, 2021
 *      Author: heber
 */


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

//#include "CodePatterns/MemDebug.hpp"

#include "CodePatterns/Log.hpp"

#include "StabilityEvaluator.hpp"

#include <iostream>

#include "Atom/atom.hpp"
#include "Bond/bond.hpp"
#include "Fragmentation/Homology/HomologyContainer.hpp"
#include "Fragmentation/Homology/HomologyGraph.hpp"
#include "Fragmentation/KeySet.hpp"
#include "molecule.hpp"
#include "World.hpp"


StabilityEvaluator::StabilityEvaluator(const molecule * _mol):
  mol(_mol),
  container(World::getInstance().getHomologies())
{

}


/**
 * Graph discovery to find the keyset if we exclude a bond from a connected
 * graph.
 */
KeySet discoverKeySetsFromAtomOnExcludedBond(
    const atom *_startatom,
    const bond::ptr &_bond) {
  KeySet result;

  std::stack<const atom *> atoms_to_visit;
  atoms_to_visit.push(_startatom);

  result.insert(_startatom->getId());
  while (!atoms_to_visit.empty()) {
    const atom * Walker = atoms_to_visit.top();
    atoms_to_visit.pop();
    const BondList& ListOfBonds = Walker->getListOfBonds();
    for (BondList::const_iterator bonditer = ListOfBonds.begin();
            bonditer != ListOfBonds.end(); ++bonditer) {
      // exclude the given bond
      if (*bonditer != _bond) {
        const atom * OtherAtom = (*bonditer)->GetOtherAtom(Walker);
        // have we seen the atom already_
        if (result.count(OtherAtom->getId()) == 0) {
          result.insert(OtherAtom->getId());
          atoms_to_visit.push(OtherAtom);
        }
      }
    }
  }

  return result;
}


template <class N, class G>
void insertInto(const N &item, G &_list) {
  std::pair<typename G::iterator, bool> inserter = _list.insert(std::make_pair(item, (size_t)1));
  if (!inserter.second) {
    ++(inserter.first->second);
  }
}

std::string getFormulaFrom(const HomologyGraph &graph) {
  Formula formula;
  for (HomologyGraph::nodes_t::const_iterator iter = graph.getNodes().begin();
      iter != graph.getNodes().end(); ++iter)
    for (size_t i=0;i<(*iter).second;++i)
      formula += (*iter).first.getAtomicNumber();
  return formula.toString();
}

HomologyGraph createHydrogenMoleculeHomologyGraph() {
  HomologyGraph::nodes_t nodes;
  // two hydrogen atoms
  const FragmentNode node = FragmentNode((atomicNumber_t)1, 1);
  insertInto(node, nodes);
  insertInto(node, nodes);
  // one hydrogen bond/edge
  HomologyGraph::edges_t edges;
  const FragmentEdge edge = FragmentEdge((atomicNumber_t)1, (atomicNumber_t)1);
  insertInto(edge, edges);
  // and return
  return HomologyGraph(nodes, edges);

}

StabilityEvaluator::stabilities_t StabilityEvaluator::operator()() const {
  stabilities_t stabilities;

  // gather all non-hydrogen bonds of the molecule
  typedef std::set<bond::ptr> bondset_t;
  bondset_t non_hydrogen_bonds;
  for (molecule::const_iterator iter = mol->begin(); iter != mol->end(); ++iter) {
    const BondList& ListOfBonds = (*iter)->getListOfBonds();
    for (BondList::const_iterator bonditer = ListOfBonds.begin();
        bonditer != ListOfBonds.end(); ++bonditer) {
      if ((*bonditer)->HydrogenBond == 0)
        non_hydrogen_bonds.insert(*bonditer );
    }
  }
  stabilities.reserve(non_hydrogen_bonds.size());

  // create homology graph for full molecule
  const HomologyGraph fullgraph(mol->getAtomIds());
  const std::string fullgraph_formula = mol->getFormula().toString();
  const HomologyGraph hydrogengraph = createHydrogenMoleculeHomologyGraph();

  // get the energy of the educts
  const double energy_fullgraph = getLowestEnergyFromHomologyContainer(fullgraph);
  const double energy_hydrogen_molecule = getLowestEnergyFromHomologyContainer(hydrogengraph);

  // go through every discovered non-hydrogen bond
  for (bondset_t::iterator bonditer = non_hydrogen_bonds.begin();
      bonditer != non_hydrogen_bonds.end(); ++bonditer) {
    const bond::ptr &bondptr = (*bonditer);

    // create stability criterion structure to fill
    StabilityCriterion criterion;
    criterion.energy_educts = energy_fullgraph + bondptr->getDegree() * energy_hydrogen_molecule;
    criterion.formula_educt1 = fullgraph_formula;
    criterion.formula_educt2 = "H"+toString(bondptr->getDegree()*2);

    // create graphs for each set
    const HomologyGraph leftgraph(
        discoverKeySetsFromAtomOnExcludedBond(bondptr->leftatom, *bonditer)
        );
    const HomologyGraph rightgraph(
        discoverKeySetsFromAtomOnExcludedBond(bondptr->rightatom, *bonditer)
        );

    // look up energy for each graph in the homology container
    criterion.energy_products =
        getLowestEnergyFromHomologyContainer(leftgraph)
        + getLowestEnergyFromHomologyContainer(rightgraph);
    criterion.formula_product1 = getFormulaFrom(leftgraph);
    criterion.formula_product2 = getFormulaFrom(rightgraph);
    criterion.isStable = criterion.energy_educts < criterion.energy_products;
    stabilities.push_back(criterion);
  }

  return stabilities;
}

double StabilityEvaluator::getLowestEnergyFromHomologyContainer(const HomologyGraph &nodes_graph) const {
  HomologyContainer::range_t range = container.getHomologousGraphs(nodes_graph);

  if (range.first == range.second) {
    // range is empty, the fragment is unknown
    ELOG(1, "Cannot find fragment graph " << nodes_graph << " .");
    return 0.;
  } else {
    // list lowest energy
    const HomologyContainer::const_iterator lowest_contribution_graph =
        std::min_element(range.first, range.second, HomologyContainer::compareEnergy);
    const HomologyContainer::const_iterator highest_contribution_graph =
        std::max_element(range.first, range.second, HomologyContainer::compareEnergy);
    LOG(2, "INFO: Fragment graph " << nodes_graph << " has energy contributions from "
        << lowest_contribution_graph->second.fragmentenergy << " Ht till "
        << highest_contribution_graph->second.fragmentenergy << " Ht, picking lowest.");
    return lowest_contribution_graph->second.fragmentenergy;
  }
}

void createHomologyGraphFromMolecule() {
  HomologyGraph::nodes_t graph_nodes;
  HomologyGraph::edges_t graph_edges;
}

std::ostream& operator<<(std::ostream &ost, const StabilityEvaluator::StabilityCriterion &_stability) {
  ost << _stability.formula_educt1 << " + " << _stability.formula_educt2
      << " (" << _stability.energy_educts << " Ht) < "
     << _stability.formula_product1 << " + " << _stability.formula_product2
      << " (" << _stability.energy_products << " Ht) ? "
      << _stability.isStable;
  return ost;
}