// Copyright (c) 2013-2020, SIB - Swiss Institute of Bioinformatics and
//                          Biozentrum - University of Basel
// 
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
// 
//   http://www.apache.org/licenses/LICENSE-2.0
// 
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.


#ifndef PROMOD3_GRAPH_MINIMIZER_HH
#define PROMOD3_GRAPH_MINIMIZER_HH

#include <vector>
#include <limits.h>

#include <boost/shared_ptr.hpp>

#include <ost/base.hh>
#include <ost/stdint.hh>

#include <promod3/core/eigen_types.hh>
#include <promod3/core/graph.hh>


namespace promod3 { namespace core {

class GMEdge;
class GMNode;
class GraphMinimizer;

typedef boost::shared_ptr<GraphMinimizer> GraphMinimizerPtr;

class GMEdge {

public:
  GMEdge(uint index, GMNode* node_a, GMNode* node_b, const EMatXX& emat);

  void Reset();

  void Deactivate() {active_ = false; };

  bool IsActive() const { return active_; }

  uint GetIndex() const { return index_; }

  bool Decompose(Real thresh);

  Real GetPairwiseEnergy(uint i, uint j) const { return emat_(i, j); }

  GMNode* GetNode1() const { return node1_; }

  GMNode* GetNode2() const { return node2_; }

  GMNode* GetOtherNode(GMNode* node);

  // resolves the expression  min_x(EPair(idx1,x)-EPair(idx2,x))
  // in the DEE formalism, where x only considers the active
  // solutions from the OTHER node than node_ptr
  Real EMinDiff(const GMNode* node_ptr, uint idx1, uint idx2) const;

  void FillActiveEnergies(EMatXX& emat) const;

private:
  uint index_;
  bool active_;
  GMNode* node1_;
  GMNode* node2_;
  EMatXX emat_;
};

class GMNode {

public:
  GMNode(uint index, const std::vector<Real>& self_energies);

  void Reset();

  void Deactivate(int index = -1);

  bool IsActive() const { return active_; }

  uint GetIndex() const { return index_; }

  void AddEdge(GMEdge* edge);

  GMEdge* GetEdge(uint idx) const { return edges_.at(idx); }

  void RemoveFromActiveEdges(GMEdge* edge);

  void AddValue(uint node_idx, Real val);

  bool DEE(Real e_cut = 0.0);

  uint GetOverallIndex(uint active_index) const { 
                               return active_solutions_.at(active_index); }

  size_t GetNumActiveSolutions() const { return active_solutions_.size(); }

  size_t GetNumActiveEdges() const { return active_edges_.size(); }

  size_t GetNumSolutions() const { return self_energies_.size(); }

  size_t GetNumEdges() const { return edges_.size(); }

  Real GetSelfEnergy(uint index) const { return self_energies_[index]; }

  void FillActiveSelfEnergies(std::vector<Real>& self_energies) const;

  typedef std::vector<uint>::const_iterator const_iterator;
  typedef std::vector<uint>::iterator iterator;
  
  iterator active_solutions_begin() { return active_solutions_.begin(); }
  iterator active_solutions_end() { return active_solutions_.end(); }
  iterator active_edges_begin() { return active_edges_.begin(); }
  iterator active_edges_end() { return active_edges_.end(); }

  const_iterator active_solutions_begin() const { 
                                            return active_solutions_.begin(); }
  const_iterator active_solutions_end() const { 
                                            return active_solutions_.end(); }
  const_iterator active_edges_begin() const { return active_edges_.begin(); }
  const_iterator active_edges_end() const { return active_edges_.end(); }

private:

  void FillSortedActiveSolutions();

  uint index_;
  bool active_;
  std::vector<Real> self_energies_;
  std::vector<Real> original_self_energies_;
  std::vector<uint> active_solutions_;
  std::vector<uint> sorted_active_solutions_;
  std::vector<uint> active_edges_;
  std::vector<GMEdge*> edges_;
};


class GraphMinimizer {

public:

  GraphMinimizer() { }

  virtual ~GraphMinimizer();

  int AddNode(const std::vector<Real>& self_energies);

  int AddEdge(uint node_idx_one, uint node_idx_two, const EMatXX& emat);

  bool ApplyDEE(uint node_idx, Real e_thresh = 0.0);

  bool ApplyEdgeDecomposition(uint edge_idx, Real e_tresh);

  void Prune(Real epsilon, Real e_cut = 0.0, bool consider_all_nodes = true);

  void Reset();

  promod3::core::Graph ToRawGraph() const;

  bool HasActiveNodes() const;

  std::pair<std::vector<int>, Real> 
  TreeSolve(uint64_t max_complexity = std::numeric_limits<uint64_t>::max(),
            Real intial_epsilon = 0.02);

  std::pair<std::vector<std::vector<int> >, std::vector<Real> > 
  AStarSolve(Real e_thresh, uint max_n = 100, 
             uint max_visited_nodes = 100000000);

  std::pair<std::vector<std::vector<int> >, std::vector<Real> > 
  MCSolve(int n = 10, int mc_steps = 1000, Real start_temperature = 100.0, 
          int change_frequency = 100, Real cooling_factor = 0.9, 
          int seed = 0);

  // Don't complain when it takes ages...
  std::pair<std::vector<int>, Real> NaiveSolve();

private:

  std::vector<GMNode*> nodes_;
  std::vector<GMEdge*> edges_;
};

}} // ns

#endif
