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


#include <boost/python.hpp>
#include <boost/python/register_ptr_to_python.hpp>
#include <promod3/core/graph_minimizer.hh>
#include <promod3/core/eigen_types.hh>
#include <promod3/core/export_helper.hh>
#include <promod3/core/message.hh>

using namespace boost::python;

namespace {


int WrapAddNode(promod3::core::GraphMinimizerPtr graph, 
                const boost::python::list& self_energies) {
  
  std::vector<Real> v_self_energies;
  promod3::core::ConvertListToVector(self_energies, v_self_energies);
  return graph->AddNode(v_self_energies);
}


int WrapAddEdge(promod3::core::GraphMinimizerPtr graph, 
                int node_idx_one, int node_idx_two,
                const boost::python::list& pairwise_energies) {

  int num_rows = boost::python::len(pairwise_energies);

  if(num_rows == 0) {
    throw promod3::Error("Pairwise energies must not be empty!");
  }

  int num_cols = len(extract<list>(pairwise_energies[0]));

  promod3::core::EMatXX emat = promod3::core::EMatXX::Zero(num_rows, num_cols);

  for(int i = 0; i < num_rows; ++i) {
    list current_list = extract<list>(pairwise_energies[i]);
    if(len(current_list) != num_cols) {
      throw promod3::Error("All lists must have exactly the same length!");
    }
    for(int j = 0; j < num_cols; ++j) {
      emat(i,j) = extract<Real>(current_list[j]);
    }
  }

  return graph->AddEdge(node_idx_one, node_idx_two, emat);
}


boost::python::tuple WrapTreeSolve(promod3::core::GraphMinimizerPtr graph, 
                                   uint64_t max_complexity,
                                   Real initial_epsilon) {

  std::pair<std::vector<int>, Real> 
  full_solution = graph->TreeSolve(max_complexity, initial_epsilon);
  std::vector<int> solution = full_solution.first;
  Real energy = full_solution.second;
  boost::python::list return_list;
  promod3::core::AppendVectorToList(solution, return_list);
  return boost::python::make_tuple(return_list, energy);
}


boost::python::tuple WrapAStarSolve(promod3::core::GraphMinimizerPtr graph,
                                    Real e_thresh, uint max_n, 
                                    uint max_visited_nodes) {

  std::pair<std::vector<std::vector<int> >, std::vector<Real> > 
  full_solution = graph->AStarSolve(e_thresh, max_n, max_visited_nodes);

  std::vector<std::vector<int> > solutions = full_solution.first;
  std::vector<Real> energies = full_solution.second;

  boost::python::list solution_list;
  boost::python::list energy_list;

  for(uint i = 0; i < solutions.size(); ++i){
    boost::python::list actual_solution;
    promod3::core::AppendVectorToList(solutions[i], actual_solution);
    solution_list.append(actual_solution);
  }
  promod3::core::AppendVectorToList(energies, energy_list);

  return boost::python::make_tuple(solution_list, energy_list);
}


boost::python::tuple WrapMCSolve(promod3::core::GraphMinimizerPtr graph, int n,
                                 int mc_steps, Real start_temperature,
                                 int change_frequency, Real cooling_factor,
                                 int seed) {

  std::pair<std::vector<std::vector<int> >, std::vector<Real> > 
  full_solution = graph->MCSolve(n, mc_steps, start_temperature, 
                                 change_frequency, cooling_factor, seed);

  std::vector<std::vector<int> > solutions = full_solution.first;
  std::vector<Real> energies = full_solution.second;

  boost::python::list solution_list;
  boost::python::list energy_list;

  for(uint i = 0; i < solutions.size(); ++i){
    boost::python::list actual_solution;
    promod3::core::AppendVectorToList(solutions[i], actual_solution);
    solution_list.append(actual_solution);
  }
  promod3::core::AppendVectorToList(energies, energy_list);

  return boost::python::make_tuple(solution_list, energy_list);
}


boost::python::tuple WrapNaiveSolve(promod3::core::GraphMinimizerPtr graph) {

  std::pair<std::vector<int>, Real> 
  full_solution = graph->NaiveSolve();
  std::vector<int> solution = full_solution.first;
  Real energy = full_solution.second;
  boost::python::list return_list;
  promod3::core::AppendVectorToList(solution, return_list);
  return boost::python::make_tuple(return_list, energy);
}

}


void export_graph_minimizer()
{

  class_<promod3::core::GraphMinimizer, 
         boost::noncopyable>("GraphMinimizer", init<>())

    .def("AddNode", &WrapAddNode, (arg("self_energies")))
    .def("AddEdge", &WrapAddEdge, (arg("node_idx_one"), arg("node_idx_two"), 
                                   arg("pairwise_energies")))
    .def("ApplyDEE", &promod3::core::GraphMinimizer::ApplyDEE,
                     (arg("node_idx"), arg("e_thresh")=0.0))
    .def("ApplyEdgeDecomposition", 
         &promod3::core::GraphMinimizer::ApplyEdgeDecomposition,
         (arg("edge_idx"), arg("e_thresh")=0.02))
    .def("Prune", &promod3::core::GraphMinimizer::Prune, 
         (arg("epsilon")=0.0, arg("e_cut")=0.0, arg("consider_all_nodes")=false))
    .def("Reset", &promod3::core::GraphMinimizer::Reset)
    .def("TreeSolve", &WrapTreeSolve,
         (arg("max_complexity")=std::numeric_limits<uint64_t>::max(),
          arg("initial_epsilon")=0.02))
    .def("AStarSolve", &WrapAStarSolve,
         (arg("e_thresh")=1.0, arg("max_n")=100, 
          arg("max_visited_nodes")=100000000))
    .def("MCSolve", &WrapMCSolve,
         (arg("n")=100, arg("mc_steps")=100000, arg("start_temperature")=1000.0,
          arg("change_frequency")=1000, arg("cooling_factor")=0.9, arg("seed")=0))
    .def("NaiveSolve", &WrapNaiveSolve)
  ;  

  register_ptr_to_python<promod3::core::GraphMinimizerPtr>();
}
