/**
 * @file tests/qdafn_test.cpp
 * @author Ryan Curtin
 *
 * Test the QDAFN functionality.
 *
 * mlpack is free software; you may redistribute it and/or modify it under the
 * terms of the 3-clause BSD license.  You should have received a copy of the
 * 3-clause BSD license along with mlpack.  If not, see
 * http://www.opensource.org/licenses/BSD-3-Clause for more information.
 */
#include <boost/test/unit_test.hpp>
#include "test_tools.hpp"
#include "serialization.hpp"

#include <mlpack/core.hpp>
#include <mlpack/methods/approx_kfn/qdafn.hpp>
#include <mlpack/methods/neighbor_search/neighbor_search.hpp>

using namespace std;
using namespace arma;
using namespace mlpack;
using namespace mlpack::neighbor;

BOOST_AUTO_TEST_SUITE(QDAFNTest);

/**
 * With one reference point, make sure that is the one that is returned.
 */
BOOST_AUTO_TEST_CASE(QDAFNTrivialTest)
{
  arma::mat refSet(5, 1);
  refSet.randu();

  // 5 tables, 1 point.
  QDAFN<> qdafn(refSet, 5, 1);

  arma::mat querySet(5, 5);
  querySet.randu();

  arma::Mat<size_t> neighbors;
  arma::mat distances;
  qdafn.Search(querySet, 1, neighbors, distances);

  // Check sizes.
  BOOST_REQUIRE_EQUAL(neighbors.n_rows, 1);
  BOOST_REQUIRE_EQUAL(neighbors.n_cols, 5);
  BOOST_REQUIRE_EQUAL(distances.n_rows, 1);
  BOOST_REQUIRE_EQUAL(distances.n_cols, 5);

  for (size_t i = 0; i < 5; ++i)
  {
    BOOST_REQUIRE_EQUAL(neighbors[i], 0);
    const double dist = metric::EuclideanDistance::Evaluate(querySet.col(i),
        refSet.col(0));
    BOOST_REQUIRE_CLOSE(distances[i], dist, 1e-5);
  }
}

/**
 * Given a random uniform reference set, ensure that we get a neighbor and
 * distance within 10% of the actual true furthest neighbor distance at least
 * 70% of the time.
 */
BOOST_AUTO_TEST_CASE(QDAFNUniformSet)
{
  arma::mat uniformSet = arma::randu<arma::mat>(25, 1000);

  QDAFN<> qdafn(uniformSet, 10, 30);

  // Get the actual neighbors.
  KFN kfn(uniformSet);
  arma::Mat<size_t> trueNeighbors;
  arma::mat trueDistances;

  kfn.Search(999, trueNeighbors, trueDistances);

  arma::Mat<size_t> qdafnNeighbors;
  arma::mat qdafnDistances;

  qdafn.Search(uniformSet, 1, qdafnNeighbors, qdafnDistances);

  BOOST_REQUIRE_EQUAL(qdafnNeighbors.n_rows, 1);
  BOOST_REQUIRE_EQUAL(qdafnNeighbors.n_cols, 1000);
  BOOST_REQUIRE_EQUAL(qdafnDistances.n_rows, 1);
  BOOST_REQUIRE_EQUAL(qdafnDistances.n_cols, 1000);

  size_t successes = 0;
  for (size_t i = 0; i < 999; ++i)
  {
    // Find the true neighbor.
    size_t trueIndex = 999;
    for (size_t j = 0; j < 999; ++j)
    {
      if (trueNeighbors(j, i) == qdafnNeighbors(0, i))
      {
        trueIndex = j;
        break;
      }
    }

    BOOST_REQUIRE_NE(trueIndex, 999);
    if (0.9 * trueDistances(0, i) <= qdafnDistances(0, i))
      ++successes;
  }

  BOOST_REQUIRE_GE(successes, 695);
}

/**
 * Test re-training method.
 */
BOOST_AUTO_TEST_CASE(RetrainTest)
{
  arma::mat dataset = arma::randu<arma::mat>(25, 500);
  arma::mat newDataset = arma::randu<arma::mat>(15, 600);

  QDAFN<> qdafn(dataset, 20, 60);

  qdafn.Train(newDataset, 10, 50);

  BOOST_REQUIRE_EQUAL(qdafn.NumProjections(), 10);
  for (size_t i = 0; i < 10; ++i)
  {
    BOOST_REQUIRE_EQUAL(qdafn.CandidateSet(i).n_rows, 15);
    BOOST_REQUIRE_EQUAL(qdafn.CandidateSet(i).n_cols, 50);
  }
}

/**
 * Test serialization of QDAFN.
 */
BOOST_AUTO_TEST_CASE(SerializationTest)
{
  // Use a random dataset.
  arma::mat dataset = arma::randu<arma::mat>(15, 300);

  QDAFN<> qdafn(dataset, 10, 50);

  arma::mat fakeDataset1 = arma::randu<arma::mat>(10, 200);
  arma::mat fakeDataset2 = arma::randu<arma::mat>(50, 500);
  QDAFN<> qdafnXml(fakeDataset1, 5, 10);
  QDAFN<> qdafnText(6, 50);
  QDAFN<> qdafnBinary(7, 15);
  qdafnBinary.Train(fakeDataset2);

  // Serialize the objects.
  SerializeObjectAll(qdafn, qdafnXml, qdafnText, qdafnBinary);

  // Check that the tables are all the same.
  BOOST_REQUIRE_EQUAL(qdafnXml.NumProjections(), qdafn.NumProjections());
  BOOST_REQUIRE_EQUAL(qdafnText.NumProjections(), qdafn.NumProjections());
  BOOST_REQUIRE_EQUAL(qdafnBinary.NumProjections(), qdafn.NumProjections());

  for (size_t i = 0; i < qdafn.NumProjections(); ++i)
  {
    BOOST_REQUIRE_EQUAL(qdafnXml.CandidateSet(i).n_rows,
        qdafn.CandidateSet(i).n_rows);
    BOOST_REQUIRE_EQUAL(qdafnText.CandidateSet(i).n_rows,
        qdafn.CandidateSet(i).n_rows);
    BOOST_REQUIRE_EQUAL(qdafnBinary.CandidateSet(i).n_rows,
        qdafn.CandidateSet(i).n_rows);

    BOOST_REQUIRE_EQUAL(qdafnXml.CandidateSet(i).n_cols,
        qdafn.CandidateSet(i).n_cols);
    BOOST_REQUIRE_EQUAL(qdafnText.CandidateSet(i).n_cols,
        qdafn.CandidateSet(i).n_cols);
    BOOST_REQUIRE_EQUAL(qdafnBinary.CandidateSet(i).n_cols,
        qdafn.CandidateSet(i).n_cols);

    for (size_t j = 0; j < qdafn.CandidateSet(i).n_elem; ++j)
    {
      if (std::abs(qdafn.CandidateSet(i)[j]) < 1e-5)
      {
        BOOST_REQUIRE_SMALL(qdafnXml.CandidateSet(i)[j], 1e-5);
        BOOST_REQUIRE_SMALL(qdafnText.CandidateSet(i)[j], 1e-5);
        BOOST_REQUIRE_SMALL(qdafnBinary.CandidateSet(i)[j], 1e-5);
      }
      else
      {
        const double value = qdafn.CandidateSet(i)[j];
        BOOST_REQUIRE_CLOSE(qdafnXml.CandidateSet(i)[j], value, 1e-5);
        BOOST_REQUIRE_CLOSE(qdafnText.CandidateSet(i)[j], value, 1e-5);
        BOOST_REQUIRE_CLOSE(qdafnBinary.CandidateSet(i)[j], value, 1e-5);
      }
    }
  }
}

// Make sure QDAFN works with sparse data.
BOOST_AUTO_TEST_CASE(SparseTest)
{
  arma::sp_mat dataset;
  dataset.sprandu(200, 1000, 0.3);

  // Create a sparse version.
  QDAFN<arma::sp_mat> sparse(dataset, 15, 50);

  // Make sure the results are of the right shape.  It's hard to test anything
  // more than that because we don't have easy-to-check performance guarantees.
  arma::Mat<size_t> neighbors;
  arma::mat distances;
  sparse.Search(dataset, 3, neighbors, distances);

  BOOST_REQUIRE_EQUAL(neighbors.n_rows, 3);
  BOOST_REQUIRE_EQUAL(neighbors.n_cols, 1000);
  BOOST_REQUIRE_EQUAL(distances.n_rows, 3);
  BOOST_REQUIRE_EQUAL(distances.n_cols, 1000);
}

BOOST_AUTO_TEST_SUITE_END();


#ifndef GRAPH_HPP
#define GRAPH_HPP

#include <algorithm>           // std::all_of
#include <Eigen/Dense>
#include <limits>

#include "readGraphOBJ.hpp"
#include "writeGraphOBJ.hpp"
#include "graphOptions.hpp"

/* TODO: check adjacency_list
 *       add linear triconnectivity implementation
 *
 * 		- merge 	Graph(Eigen::MatrixXd nodes, Eigen::MatrixXi adjacency_matrix)
 * 					Graph(Eigen::MatrixXd nodes, Eigen::Matrix<int, Eigen::Dynamic, 2> edges)
 */
namespace libgraphcpp
{

	class Graph
	{
	private:
		graphOptions opts_;                                                      // Store all visualization infos

		// basic structure for edges and nodes
		Eigen::MatrixXd nodes_;                                             // should be a N by 3 matrix of doubles
		Eigen::MatrixXi edges_;                                             // should be a M by 2 matrix of integers
		int num_nodes_;                                                     // set to N
		int num_edges_;                                                     // set to M
		double scale_;                                                      // used to trim the tree in simplify_tree
		double cycle_ratio_ = 2;
        double triangle_ratio_ = 0.9;

		// structures used for fast circulation through data
		std::vector< std::vector<int> > adjacency_list_;                    // contains for each nodes, its nodes neighbors
		std::vector< std::vector<int> > adjacency_edge_list_;               // contains for each nodes, its edges neighbors
		Eigen::VectorXd edges_length_;                                      // used only for Dijkstra
		Eigen::MatrixXi adjacency_matrix_;

		// connectivity properties
		int is_connected_ = -1;                                             // 0 no, 1 yes, -1 undefined
		int is_biconnected_ = -1;                                           // 0 no, 1 yes, -1 undefined
		int is_triconnected_ = -1;                                          // 0 no, 1 yes, -1 undefined
		int has_bridges_ = -1;                                              // 0 no, 1 yes, -1 undefined
		int has_cycle_ = -1;                                                // 0 no, 1 yes, -1 undefined

		// storing of the cut sets
		std::vector< int > one_cut_vertices_;                               // set of articulation points : vector of nodes ids
		std::vector< std::pair<int, int> > two_cut_vertices_;               // set of two-cut vertices    : vector of nodes ids pair
		std::vector<int> bridges_;                                          // set of briges              : vector of edges ids
		std::vector< std::vector<int> > cycles_;                            // set of cycles

		// internal functions: tools (defined at the bottom of the file)
		inline void removeRow(Eigen::MatrixXd& matrix, unsigned int rowToRemove);
		inline void removeRow(Eigen::MatrixXi& matrix, unsigned int rowToRemove);
		inline void removeDuplicates(std::vector<std::pair<int, int>>& v);
		inline bool is_element_in_vector(int a, std::vector<int> & A);

		// internal functions: iterative functions (defined at the bottom of the file)
		inline void DFSUtil(int u, std::vector< std::vector<int> > adj, std::vector<bool> &visited);
		inline void APUtil(int u, std::vector<bool> & visited, int disc[], int low[], std::vector<int> & parent, std::vector<bool> & ap);
		inline void bridgeUtil(int u, std::vector<bool> & visited, int disc[], int low[], std::vector<int> & parent, std::vector<int> & bridges);

	public:
		// creators
		inline Graph(std::string file_name);
		inline Graph(std::string file_name, graphOptions opts);
		inline Graph(Eigen::MatrixXd nodes, Eigen::MatrixXi edges);
		inline Graph(Eigen::MatrixXd nodes, Eigen::MatrixXi edges, graphOptions opts);

		// destructor
		inline ~Graph(){};

		// initialisation of the private variables
		inline void init();
		inline void set_adjacency_lists();

		// save graph as OBJ file
		inline void save(std::string output_file);
		inline void print_isolated_vertices();

		// accessors
		int num_nodes();
		int num_edges();
		Eigen::MatrixXd get_nodes();
		Eigen::Vector3d get_node(int i);
		Eigen::MatrixXi get_edges();
		Eigen::Vector2i get_edge(int i);
		std::vector <int> get_adjacency_list(int i);
		int get_adjacency_list(int i, int j);
        int find_edge_from_nodes(int node_1, int node_2);

		// modifiers for nodes
		inline void add_node(Eigen::Vector3d node, std::vector<int> neighbours);
		inline void remove_node(int nodeToRemove);
		inline void merge_nodes(std::vector<int> nodes);
		inline void update_node(int node_id, Eigen::Vector3d new_node);
		inline void update_nodes(Eigen::MatrixXd new_nodes);

		// modifiers for edges
		inline void add_edge(Eigen::Vector2i edge);
		inline void remove_edge(int edgeToRemove);
		inline void collapse_edge(int edge_id);

		/* CONNECTIVITY TESTS */
		inline void connectivity_tests();
		inline bool is_connected();
		inline bool is_biconnected(std::vector<int>& one_cut_vertices);
		inline bool is_biconnected();
		inline bool is_triconnected(std::vector< std::pair<int, int> >& two_cut_vertices);
		inline bool is_triconnected();
		inline bool has_bridges(std::vector<int>& bridges);
		inline bool has_bridges();
		inline bool has_cycles(std::vector <std::vector<int>>& cycle_basis, std::vector<double>& cycle_lengths);
		inline bool has_cycles();

		// graph manipulation
		inline std::vector<std::vector <int> > make_tree(Eigen::MatrixXi& deleted_edges);
		inline std::vector<std::vector <int> > make_tree();
		inline void simplify_tree();
		inline void symplify_graph();
        inline void remove_flat_triangles();
		inline void transform (double scale, Eigen::Vector3d move);
		inline double dijkstra(int source, int target, std::vector<int>& node_path);
		inline double dijkstra(int source, int target);

		/* TO BE REMOVED? This apply only for directional graph */
		inline int edge_source(int i);
		inline int edge_target(int i);
		inline void swap_edge(int i);

	};

	inline Graph::Graph(std::string file_name)
	{
		readGraphOBJ(file_name, nodes_, edges_);

		init();
	};

	// overload with options
	inline Graph::Graph(std::string file_name, graphOptions opts) : Graph(file_name)
	{
		opts_ = opts;
	};

	inline Graph::Graph(Eigen::MatrixXd nodes, Eigen::MatrixXi edges)
	{
		// test if edges is m by 2 (explicit edges) or n by n (adjacency matrix)
		if ( edges.cols() == 2 )
		{
			nodes_ = nodes;
			edges_ = edges;
		}
		else
		{
			Eigen::MatrixXi adjacency_matrix;
			adjacency_matrix = edges;

			if (adjacency_matrix.rows() != adjacency_matrix.cols() || adjacency_matrix.rows() != nodes.rows()) {
				std::cout << "\nLibGraphCpp error: wrong input size in the class definition\n";
				std::cout << "Error thrown while assessing: adjacency_matrix.rows() != adjacency_matrix.cols() || adjacency_matrix.rows() != nodes.rows()\n ";
				std::cout << "size of the nodes matrix: " << nodes.rows() << ", " << nodes.cols() << "\n";
				std::cout << "size of the edges matrix: " << edges.rows() << ", " << edges.cols() << "\n";
				std::exit(EXIT_FAILURE);
			}
			if (adjacency_matrix.transpose() != adjacency_matrix)
			{
				std::cout << "\nLibGraphCpp error: the adjacency_matrix should be symmetric\n ";
				std::exit(EXIT_FAILURE);
			}

			nodes_ = nodes;
			Eigen::MatrixXi edges((adjacency_matrix.array() != 0).count(), 2);

			// explore the upper triangle of the adjacency_matrix (without the diagonal)
			int num_edges = 0;
			for (int i=0; i<adjacency_matrix.rows(); i++)
				for (int j=i+1; j<adjacency_matrix.cols(); j++) {
					if (adjacency_matrix(i,j) != 0)
					{
						edges(num_edges, 0) = i;
						edges(num_edges, 1) = j;
						num_edges ++;
					}
				}
			edges.conservativeResize(num_edges, 2);

			edges_ = edges;
			adjacency_matrix_  = adjacency_matrix;
		}

		init();
	};

	// overload with options
	inline Graph::Graph(Eigen::MatrixXd nodes, Eigen::MatrixXi edges, graphOptions opts) : Graph(nodes, edges)
	{
		opts_ = opts;
	};

	// initialisation of the private variables
	inline void Graph::init()
	{
		if (nodes_.cols()!=3 || edges_.cols()!=2) {
			std::cout << "\nLibGraphCpp error: wrong graph dimensions in the class initialization" << std::endl;
			std::cout << "nodes size: " << nodes_.rows() << " * " << nodes_.cols()<< std::endl;
			std::cout << "edges size: " << edges_.rows() << " * " << edges_.cols()<< std::endl;
			std::exit(EXIT_FAILURE);
		}

		one_cut_vertices_.clear();
		two_cut_vertices_.clear();
		bridges_.clear();

		// set up properties
		num_nodes_ = nodes_.rows();
		num_edges_ = edges_.rows();

		// get scale
		Eigen::MatrixXd max_point = nodes_.colwise().maxCoeff();
		Eigen::MatrixXd min_point = nodes_.colwise().minCoeff();
		scale_ = (max_point - min_point).norm();

		// set up edges_length
		edges_length_ = Eigen::VectorXd::Zero(num_edges_);
		for (int i=0; i<num_edges_; i++) {
			if ( (edges_(i, 0) >= num_nodes_) || (edges_(i, 1) >= num_nodes_) ||
					(edges_(i, 0) < 0) || (edges_(i, 1) < 0) ) {
				std::cout << "\nLibGraphCpp error: wrong edge given in the class initialization" << std::endl;
				std::cout << "the edge: (" << edges_(i, 0) << ", " << edges_(i, 1) << ") does not works for " << num_nodes_ << " nodes." << std::endl;
				std::exit(EXIT_FAILURE);
			}
			edges_length_(i) = (nodes_.row(edges_(i,0)) - nodes_.row(edges_(i,1)) ).norm();
		}

		// set up the adjacency list
		set_adjacency_lists();
	};

	inline void Graph::set_adjacency_lists()
	{
		// TODO: - the if statement should be removed for undirected graphs as they are not needed
		//       - is the adjacency_edge_list_ needed?

		// make sure the lists are empty to start with
		adjacency_list_.clear();
		adjacency_edge_list_.clear();

		adjacency_list_.resize(num_nodes_);
		adjacency_edge_list_.resize(num_nodes_);

		for (int i=0; i<num_edges_; i++)
		{
			// the if statements check if the node has already been inserted
			if (std::find (adjacency_list_[edges_(i, 0)].begin(), adjacency_list_[edges_(i, 0)].end(), edges_(i, 1))==adjacency_list_[edges_(i, 0)].end())
				adjacency_list_[edges_(i, 0)].push_back(edges_(i, 1));
			if (std::find (adjacency_list_[edges_(i, 1)].begin(), adjacency_list_[edges_(i, 1)].end(), edges_(i, 0))==adjacency_list_[edges_(i, 1)].end())
				adjacency_list_[edges_(i, 1)].push_back(edges_(i, 0));
			adjacency_edge_list_[edges_(i, 0)].push_back(i);
			adjacency_edge_list_[edges_(i, 1)].push_back(i);
		}
	};

	// save graph as OBJ file
	inline void Graph::save(std::string output_file)
	{
		writeGraphOBJ(nodes_, edges_, output_file);
	};

	inline void Graph::print_isolated_vertices()
	{
		for (int i=0; i<num_nodes_; i++)
			if (adjacency_list_[i].size() == 0)
				std::cout << "Isolated vertices at: " << i << std::endl;
	};

	// accessors
	inline int Graph::num_nodes() { return num_nodes_; };
	inline int Graph::num_edges() { return num_edges_; };
	inline Eigen::MatrixXd Graph::get_nodes() { return nodes_; };
	inline Eigen::Vector3d Graph::get_node(int i) { return nodes_.row(i); };
	inline Eigen::MatrixXi Graph::get_edges() { return edges_; };
	inline Eigen::Vector2i Graph::get_edge(int i) { return edges_.row(i); };
	inline std::vector <int> Graph::get_adjacency_list(int i) { return adjacency_list_[i]; };
	inline int Graph::get_adjacency_list(int i, int j) { return adjacency_list_[i][j]; };

	inline int Graph::find_edge_from_nodes(int node_1, int node_2)
	{
		for (int i=0; i<num_edges_; i++) {
			if (edges_(i, 0) == node_1 && edges_(i, 1) == node_2)
				return i;
			if (edges_(i, 0) == node_2 && edges_(i, 1) == node_1)
				return i;
		}
		return -1;
	};

	// modifiers for nodes
	inline void Graph::add_node(Eigen::Vector3d node, std::vector<int> neighbours)
	{
		// add node
		Eigen::MatrixXd temp_nodes = nodes_;
		nodes_.resize(nodes_.rows()+1, 3);
		nodes_ << temp_nodes, node.transpose();

		// add related edges
		Eigen::MatrixXi temp_edges = edges_;
		Eigen::MatrixXi edges_to_add(neighbours.size(), 2);
		for (int edge_iterator=0; edge_iterator<neighbours.size(); edge_iterator++)
			edges_to_add.row(edge_iterator) << nodes_.rows()-1, neighbours[edge_iterator];

		edges_.resize(temp_edges.rows()+edges_to_add.rows(), 2);
		edges_ << temp_edges, edges_to_add;

		init();
	};

	inline void Graph::remove_node(int nodeToRemove)
	{
		// remove node
		removeRow(nodes_, nodeToRemove);

		// remove related edges
		std::vector<int> edges_to_remove = adjacency_edge_list_[nodeToRemove];
		std::sort(edges_to_remove.begin(), edges_to_remove.end(), std::greater<int>());
		for (int edge_iterator=0; edge_iterator<edges_to_remove.size(); edge_iterator++)
			removeRow(edges_, edges_to_remove[edge_iterator]);

		// update the edge nodes index (we have one less node now)
		for (int i=0; i<edges_.rows(); i++)
			for (int j=0; j<edges_.cols(); j++)
				if (edges_(i,j) >= nodeToRemove)
					edges_(i,j)--;

		init();
	};

	inline void Graph::merge_nodes(std::vector<int> nodes)
	{
		// make sure there is no duplicates
		sort( nodes.begin(), nodes.end() );
		nodes.erase( unique( nodes.begin(), nodes.end() ), nodes.end() );

		// store connected elements first
		std::vector<int> neighbours;
		for (int node : nodes)
			for (int neighour : adjacency_list_[node])
				neighbours.push_back(neighour);

		std::sort(neighbours.begin(), neighbours.end());
		neighbours.erase( std::unique( neighbours.begin(), neighbours.end() ), neighbours.end() );

		// find node position
		Eigen::Vector3d merged_node;
		merged_node << 0,0,0;
		for (int node : nodes)
			merged_node += nodes_.row(node);
		merged_node /= nodes.size();

		// add merged node (it is stacked at the end so it has to be done before deleting the nodes)
		add_node(merged_node, neighbours);

		// delete other nodes
		std::sort(nodes.begin(), nodes.end(), std::greater<int> ());
		for (int node : nodes)
			remove_node(node);
	};

	inline void Graph::update_node(int node_id, Eigen::Vector3d new_node)
	{
		nodes_.row(node_id) = new_node;
	}

	inline void Graph::update_nodes(Eigen::MatrixXd new_nodes)
	{
		if (nodes_.rows() == new_nodes.rows() && nodes_.cols() == new_nodes.cols()) {
			nodes_ = new_nodes;
		} else {
			std::cout << "Error: wrong dimension:\n";
			std::cout << "nodes_.rows() == new_nodes.rows() && nodes_.cols() == new_nodes.cols() returned False \n ";
			std::exit(0);
		}
	};

	// modifiers for edges
	inline void Graph::add_edge(Eigen::Vector2i edge)
	{
		// add node
		Eigen::MatrixXi temp_edges = edges_;
		edges_.resize(edges_.rows()+1, 3);
		edges_ << temp_edges, edge.transpose();

		init();
	};

	inline void Graph::remove_edge(int edgeToRemove)
	{
		// remove edge
		removeRow(edges_, edgeToRemove);

		init();
	};

	// replace an edge and its two connected nodes by a single node
	inline void Graph::collapse_edge(int edge_id)
	{
		// get nodes_id:
		int node_1 = edges_(edge_id, 0);
		int node_2 = edges_(edge_id, 1);

		// merge the two previous ones
		Eigen::Vector3d fused_node = (nodes_.row(node_1) + nodes_.row(node_2) ) / 2;

		// get the connected neighbours
		std::vector<int> neighbours;
		neighbours = adjacency_list_[node_1];
		neighbours.insert( neighbours.end(), adjacency_list_[node_2].begin(), adjacency_list_[node_2].end() );

		// remove old nodes
		neighbours.erase(std::remove(neighbours.begin(), neighbours.end(), node_1), neighbours.end());
		neighbours.erase(std::remove(neighbours.begin(), neighbours.end(), node_2), neighbours.end());

		sort( neighbours.begin(), neighbours.end() );
		neighbours.erase( unique( neighbours.begin(), neighbours.end() ), neighbours.end() );

		// add node
		add_node(fused_node, neighbours);

		// remove the previous nodes (the order matter)
		if (node_2<node_1) {
			remove_node(node_1);
			remove_node(node_2);
		} else {
			remove_node(node_2);
			remove_node(node_1);
		}
	};

	/* CONNECTIVITY TESTS */
	inline void Graph::connectivity_tests()
	{
		is_connected();
		is_biconnected();
		is_triconnected();
		has_bridges();
	};

	inline bool Graph::is_connected()
	{
		if (is_connected_ == -1) {
			// check for graph connectivity using DFS:
			std::vector<bool> visited(num_nodes_, false);
			DFSUtil(0, adjacency_list_, visited);
			is_connected_ = std::all_of(visited.begin(), visited.end(), [](bool v) { return v; });

			if (opts_.verbose)
				std::cout << "graph is connected: " << is_connected_ << std::endl;

			// if not connected, update the next ones
			if (is_connected_ == 0) {
				is_biconnected_ = 0;
				is_triconnected_ = 0;
			}
		}
		return is_connected_;
	};

	// return the set of one cut vertices
	inline bool Graph::is_biconnected(std::vector<int>& one_cut_vertices)
	{
		if (is_biconnected_ == -1) {
			// check for 2 node connectivity:
			std::vector<bool> visited(num_nodes_, false);
			std::vector<int> parent(num_nodes_, -1);
			std::vector<bool> ap(num_nodes_, false);
			int *disc = new int[num_nodes_];
			int *low = new int[num_nodes_];
			// Call the recursive helper function to find articulation points
			// in DFS tree rooted with vertex '0'
			for (int i = 0; i < num_nodes_; i++)
				if (visited[i] == false)
					APUtil(i, visited, disc, low, parent, ap);

			// Now ap[] contains articulation points, print them and store them into one_cut_vertices_
			one_cut_vertices.clear();
			for (int i = 0; i < num_nodes_; i++)
				if (ap[i] == true) {
					one_cut_vertices.push_back(i);
					if (opts_.verbose)
						std::cout << "Articulation point " << i << " at: "  << nodes_.row(i) << std::endl;
				}
			is_biconnected_ = std::none_of(ap.begin(), ap.end(), [](bool v) { return v; });

			one_cut_vertices_ = one_cut_vertices;
		} else {
			one_cut_vertices = one_cut_vertices_;
		}

		if (opts_.verbose)
			std::cout << "graph is 2-connected: " << is_biconnected_ << std::endl;

		return is_biconnected_;
	};

	// overload is_biconnected
	inline bool Graph::is_biconnected()
	{
		std::vector< int > one_cut_vertices;
		return is_biconnected(one_cut_vertices);
	};

	// return the set of two cut vertices
	inline bool Graph::is_triconnected(std::vector< std::pair<int, int> >& two_cut_vertices)
	{
		/* this function run in quadratic time, this is not the most efficient way to do it
		* for alternative, see these papers:
		* - Finding the Triconnected Components of a Graph (1972)
		* - A Linear Time Implementation of SPQR-Trees (2000)
		* and these implementations:
		* - http://www.ogdf.net/doku.php
		* - https://github.com/adrianN/Triconnectivity
		*/
		graphOptions opts = opts_;
		opts.verbose = false;
		if (is_triconnected_ == -1) {
			is_triconnected_ = true;
			for (int i=0; i<num_nodes_; i++) {
				std::vector< int > one_cut_vertices;
				Graph reduced_graph(nodes_, edges_, opts);
				reduced_graph.init();
				reduced_graph.remove_node(i);
				bool reduced_graph_is_biconnected = reduced_graph.is_biconnected(one_cut_vertices);

				if (not reduced_graph_is_biconnected)
					for (int j=0; j<one_cut_vertices.size(); j++) {
						int node_offset = (one_cut_vertices[j]>=i); // offset needed because a node has been deleted
						two_cut_vertices.push_back(std::make_pair(i, one_cut_vertices[j]+node_offset));
					}
				is_triconnected_ &= reduced_graph_is_biconnected;
			}

			removeDuplicates(two_cut_vertices);
			two_cut_vertices_ = two_cut_vertices;
		} else {
			two_cut_vertices = two_cut_vertices_;
		}

		if (opts_.verbose) {
			std::cout << "graph is 3-connected: " << is_triconnected_ << std::endl;
			for (int i = 0; i<two_cut_vertices.size(); i++)
				std::cout << "two_cut_vertices between node " << two_cut_vertices[i].first <<" and node " << two_cut_vertices[i].second << std::endl;
		}

		return is_triconnected_;
	};

	// overload is_triconnected
	inline bool Graph::is_triconnected()
	{
		std::vector< std::pair<int, int> > two_cut_vertices;
		return is_triconnected(two_cut_vertices);
	};

	// return the set of bridges
	inline bool Graph::has_bridges(std::vector<int>& bridges)
	{
		if (has_bridges_ == -1) {
			// check for bridges:

			std::vector<bool> visited(num_nodes_, false);
			std::vector<int> parent(num_nodes_, -1);
			int *disc = new int[num_nodes_];
			int *low = new int[num_nodes_];

			// Call the recursive helper function to find Bridges
			// in DFS tree rooted with vertex 'i'
			for (int i = 0; i < num_nodes_; i++)
				if (visited[i] == false)
					bridgeUtil(i, visited, disc, low, parent, bridges);

			delete[] disc;
			delete[] low;

			if (bridges.size() == 0)
				has_bridges_ = 0;
			else
				has_bridges_ = 1;

			bridges_ = bridges;
		} else {
			bridges = bridges_;
		}

		if (opts_.verbose)
			std::cout << "graph has bridges: " << has_bridges_ << std::endl;

		return has_bridges_;
	};

	// overload has_bridges
	inline bool Graph::has_bridges()
	{
		std::vector<int> bridges;
		return has_bridges(bridges);
	};

	// return a fundamental cycles basis (see: An Algorithm for Finding a Fundamental Set of Cycles of a Graph)
	// should use DFS or BFS for generating the tree
	inline bool Graph::has_cycles(std::vector <std::vector<int>>& cycle_basis, std::vector<double>& cycle_lengths)
	{
		graphOptions opts = opts_;
		opts.verbose = false;
		Graph spanning_tree(nodes_, edges_, opts);
		Eigen::MatrixXi deleted_edges;
		spanning_tree.make_tree(deleted_edges);

		for (int i=0; i<deleted_edges.rows(); i++) {
			std::vector<int> cycle;
			double cycle_length;
			cycle_length = spanning_tree.dijkstra(deleted_edges(i, 0), deleted_edges(i, 1), cycle);

			double deleted_edge_length = (spanning_tree.get_node(deleted_edges(i, 0)) - spanning_tree.get_node(deleted_edges(i, 1))).norm();
			cycle_length = cycle_length + deleted_edge_length;
			cycle_basis.push_back(cycle);
			cycle_lengths.push_back(cycle_length);
		}

		if (cycle_basis.size() != 0)
			return true;
		else
			return false;
	};

	// overload has_cycle
	inline bool Graph::has_cycles()
	{
		std::vector <std::vector <int>> cycles;
		std::vector<double> cycle_lengths;
		return has_cycles(cycles, cycle_lengths);
	};

	inline std::vector<std::vector <int> > Graph::make_tree(Eigen::MatrixXi& deleted_edges)
	{
		std::vector<std::vector <int> > nodes_references(num_nodes_);
		for (int i=0; i<num_nodes_; i++)
			nodes_references[i].push_back(i);

		bool verbose_temp = opts_.verbose;
		opts_.verbose = false;

		// check for bridges
		has_bridges_ = -1;
		has_bridges();

		deleted_edges.resize(num_edges_ - num_nodes_ + 1, 2); // cycle rank * 2
		int counter = 0;

		while (bridges_.size() != num_edges_)
		{

			// list all non bridges
			std::vector <int> edges_to_edit;
			std::vector <double> edges_length;
			for (int i=0; i<num_edges_; i++)
				if (find (bridges_.begin(), bridges_.end(), i) == bridges_.end()) {
					edges_to_edit.push_back(i);
					edges_length.push_back(edges_length_(i));
				}

			// sort edges by decreasing length
			std::vector<std::pair<double, int>> sorting_container;
			sorting_container.reserve(edges_to_edit.size());
			std::transform(edges_length.begin(), edges_length.end(), edges_to_edit.begin(), std::back_inserter(sorting_container),
				[](double a, int b) { return std::make_pair(a, b); });

			std::sort(sorting_container.begin(), sorting_container.end());
			std::reverse(sorting_container.begin(), sorting_container.end());

			// store edges
			deleted_edges.row(counter) << edges_.row(sorting_container[0].second);
			counter ++;

			// actually store the edges
			//collapse_edge(edges_to_edit[0]);
			remove_edge(sorting_container[0].second);

			// check for bridges
			has_bridges_ = -1;
			has_bridges();
		}

		opts_.verbose = verbose_temp;
		return nodes_references;
	};

	// overload make_tree if deleted edges are not needed
	inline std::vector<std::vector <int> > Graph::make_tree()
	{
		Eigen::MatrixXi deleted_edges;
		return make_tree(deleted_edges);
	};

	inline void Graph::simplify_tree()
	{
		/*
			* for each junctions:
			* if the distance between two junction is "small"
			* then merge all the nodes in the path
			*/
		std::vector<int> junction_list;
		start_junction_checking:
		junction_list.clear();
		for (int i=0; i<adjacency_list_.size(); i++)
			if (adjacency_list_[i].size()>2)
				junction_list.push_back(i);

		for (int i = 0; i<junction_list.size(); i++) {
			for (int j = i+1; j<junction_list.size(); j++) {
				std::vector<int> path;
				double junctions_distance = dijkstra(junction_list[i], junction_list[j], path);

				if (junctions_distance < scale_/10)
				{
					merge_nodes(path);
					goto start_junction_checking; // went throwing up after writing this
				}
			}
		}

		/*
			* for each leaf:
			* if the distance to the closest junction is "small"
			* then delete all nodes until this node is reached
			*/
		std::vector<int> leaves_list;
		junction_list.clear();
		for (int i=0; i<adjacency_list_.size(); i++){
			if (adjacency_list_[i].size()==1)
				leaves_list.push_back(i);

			if (adjacency_list_[i].size()>2)
				junction_list.push_back(i);
		}

		// for each leaves check the distance to each junction
		std::vector<int> nodes_to_delete;
		for (int leaf: leaves_list) {
			for (int junction: junction_list) {
				std::vector<int> path;
				double leaf_to_junction_distance = dijkstra(leaf, junction, path);

				if (leaf_to_junction_distance < scale_/10)
				{
					for (int node_id : path)
						if (node_id != junction)
							nodes_to_delete.push_back(node_id);
				}
			}
		}

		if (nodes_to_delete.size() != 0) {
			std::sort(nodes_to_delete.begin(), nodes_to_delete.end(), std::greater<int>());
			nodes_to_delete.erase( std::unique( nodes_to_delete.begin(), nodes_to_delete.end() ), nodes_to_delete.end() );
			for (int node_id : nodes_to_delete)
				remove_node(node_id);
		}

	};

	inline void Graph::symplify_graph()
	{
		// list all cycles
		std::vector< std::vector< int > > cycle_basis;
		std::vector< double > cycle_lengths;
		has_cycles(cycle_basis, cycle_lengths);

		// remove large cycles
		for(int i=0; i<cycle_lengths.size();) {
			if(cycle_lengths[i] > scale_/cycle_ratio_) {
				cycle_lengths.erase(cycle_lengths.begin()+i);
				cycle_basis.erase(cycle_basis.begin() + i);
			} else {
				++i;
			}
		}

		// find set of clustered nodes
		std::vector< std::vector< int > > clusters;
		bool clusters_found = false;
		while (!clusters_found) {
			std::vector <int> cluster_temp;
			cluster_temp = cycle_basis[0];
			cycle_basis.erase(cycle_basis.begin());

			for(int i=0; i<cycle_basis.size(); )
			{
				// check for intersection
				std::vector <int> v1, v2;
				v1 = cluster_temp;
				v2 = cycle_basis[i];

				sort(v1.begin(), v1.end());
				sort(v2.begin(), v2.end());

				std::vector<int> v(v1.size() + v2.size());
				std::vector<int>::iterator it, st;

				it = set_intersection(v1.begin(),
										v1.end(),
										v2.begin(),
										v2.end(),
										v.begin());

				bool each_cycle_have_common_elements = false;
				for (st = v.begin(); st != it; ++st)
					each_cycle_have_common_elements = true;

				if(each_cycle_have_common_elements) {
					cluster_temp.insert( cluster_temp.end(), cycle_basis[i].begin(), cycle_basis[i].end() );
					cycle_basis.erase(cycle_basis.begin()+i);
					i = 0;
				} else {
					++i;
				}
			}

			// add sort + unique
			sort( cluster_temp.begin(), cluster_temp.end() );
			cluster_temp.erase( unique( cluster_temp.begin(), cluster_temp.end() ), cluster_temp.end() );
			clusters.push_back(cluster_temp);

			if (cycle_basis.size() == 0)
				clusters_found = true;
		}

		// merge nodes
		for (int i=0; i<clusters.size(); i++) {
			std::vector<int> empty_set;
			merge_nodes(clusters[i]);
			// each time a node is removed, the other lists should be updated to avoid deleting the wrong nodes
			for (int j=i+1; j<clusters.size(); j++)
				for (int k=0; k<clusters[j].size(); k++) {
					int offset = 0;
					for (int merged_node : clusters[i]) {
						// first check if the point is shared (it would then be the last node)
						if (clusters[j][k] == merged_node)
						{
							offset = 0;
							clusters[j][k] = num_nodes_-1;
							break;
						}
						// otherwise check if the merged point is lower, then the node has to be shifted of one position
						if (clusters[j][k] > merged_node)
							offset ++;
					}
					clusters[j][k] -= offset;
				}
		}
	};

	// the listing of the triangles should be performed differently
	// -> listing the fundamental cycles basis can be inacurate
	// -> or replace by the minimal cycle basis
	inline void Graph::remove_flat_triangles()
	{
		// list all cycles
		std::vector< std::vector< int > > cycle_basis;
		std::vector< double > cycle_lengths;
		has_cycles(cycle_basis, cycle_lengths);

		// go through each cycles
		for (int i=0; i<cycle_basis.size(); i++) {
			if (cycle_basis[i].size() == 3) {
				int edge_1 = find_edge_from_nodes(cycle_basis[i][0], cycle_basis[i][1]);
				int edge_2 = find_edge_from_nodes(cycle_basis[i][1], cycle_basis[i][2]);
				int edge_3 = find_edge_from_nodes(cycle_basis[i][2], cycle_basis[i][0]);

				std::vector<std::pair<double, int>> triangle_distance(3);
				triangle_distance[0] = std::make_pair( edges_length_(edge_1), edge_1);
				triangle_distance[1] = std::make_pair( edges_length_(edge_2), edge_2);
				triangle_distance[2] = std::make_pair( edges_length_(edge_3), edge_3);

				sort(triangle_distance.begin(), triangle_distance.end());

				if (triangle_distance[2].first > (triangle_distance[0].first + triangle_distance[1].first)*triangle_ratio_)
					remove_edge(triangle_distance[2].second);

			}
		}
	};

	inline void Graph::transform (double scale, Eigen::Vector3d move)
	{
		nodes_ /= scale;
		nodes_ += move.transpose();
	};

	inline double Graph::dijkstra(int source, int target, std::vector<int>& node_path)
	{
		// initialization
		double distance_source_to_target;
		std::vector<double> min_distance(num_nodes_, std::numeric_limits<double>::infinity());
		std::vector<int> previous_node(num_nodes_, -1);

		// set the set of visited nodes:
		std::vector<int> visited;
		std::vector<int> to_visit;

		// initialize the node to start from
		int u = source;
		min_distance.at(source) = 0;

		// start searching
		bool target_found;
		while (!target_found) {
			// check all neighbours of "u"
			for (int i = 0; i < adjacency_list_.at(u).size(); ++i) {
				int neighbour_node = adjacency_list_.at(u).at(i);
				double edge_length = edges_length_(adjacency_edge_list_.at(u).at(i));
				if (not is_element_in_vector(neighbour_node, to_visit))
					if (not is_element_in_vector(neighbour_node, visited))
						to_visit.push_back(neighbour_node);

				if ( min_distance.at(u) + edge_length < min_distance.at( neighbour_node ) ) {
					min_distance.at( neighbour_node ) = min_distance.at(u) +  edge_length;
					previous_node.at( neighbour_node ) = u;
				}
			}

			visited.push_back(u);

			// check if the visited point is in sub_V
			target_found = is_element_in_vector(target, visited);

			// set next u
			int index_of_next_point = 0;
			for (int i = 0; i < to_visit.size(); ++i)
				if (min_distance.at( to_visit.at(i) ) < min_distance.at( to_visit.at( index_of_next_point) ) )
					index_of_next_point = i;

			// check if all vertices have been visited:
			if (to_visit.size() == 0)
				break;

			u = to_visit.at(index_of_next_point);
			to_visit.erase(to_visit.begin() + index_of_next_point);
		}

		// backtracking to generate path
		if (min_distance.at(target) != std::numeric_limits<double>::infinity()) {
			int backtracking_node = target;
			node_path.insert(node_path.begin(), backtracking_node);

			while (backtracking_node != source) {
				node_path.insert(node_path.begin(), previous_node[backtracking_node]);
				backtracking_node = previous_node[backtracking_node];
			}
		}

		return min_distance.at(target);
	};

	inline double Graph::dijkstra(int source, int target)
	{
		std::vector<int> node_path;
		return  dijkstra(source, target, node_path);
	};

	/* TO BE REMOVED? This apply just for directional graph */

	inline int Graph::edge_source(int i)
	{
		return edges_(i, 0);
	};

	inline int Graph::edge_target(int i)
	{
		return edges_(i, 1);
	};

	inline void Graph::swap_edge(int i)
	{
		int temp = edges_(i,0);
		edges_(i,0) = edges_(i,1);
		edges_(i,1) = temp;
	};

	/*
	 * PRIVATE FUNCTIONS:
	 */

	//https://stackoverflow.com/a/46303314/2562693
	void Graph::removeRow(Eigen::MatrixXd& matrix, unsigned int rowToRemove)
	{
		unsigned int numRows = matrix.rows()-1;
		unsigned int numCols = matrix.cols();

		if( rowToRemove < numRows )
			matrix.block(rowToRemove,0,numRows-rowToRemove,numCols) = matrix.bottomRows(numRows-rowToRemove);

		matrix.conservativeResize(numRows,numCols);
	};

	void Graph::removeRow(Eigen::MatrixXi& matrix, unsigned int rowToRemove)
	{
		unsigned int numRows = matrix.rows()-1;
		unsigned int numCols = matrix.cols();

		if( rowToRemove < numRows )
			matrix.block(rowToRemove,0,numRows-rowToRemove,numCols) = matrix.bottomRows(numRows-rowToRemove);

		matrix.conservativeResize(numRows,numCols);
	};

	//https://stackoverflow.com/a/32842128/2562693
	void Graph::removeDuplicates(std::vector<std::pair<int, int>>& v)
	{

		//Normalize == sort the pair members
		for(auto& p : v){
			int x = std::max(p.first, p.second), y = std::min(p.first, p.second);
			p.first = x; p.second = y;
		}

		//Sort the pairs
		std::sort(v.begin(), v.end());

		//Unique the vector
		auto last = unique(v.begin(), v.end() );
		v.erase(last, v.end());
	};

	bool Graph::is_element_in_vector(int a, std::vector<int> & A)
	{
		auto it = std::find(A.begin(), A.end(), a);
		return it != A.end();
	}

	// used for checking connectivity:
	//
	// for reference, see Kosaraju's algorithm (https://en.wikipedia.org/wiki/Kosaraju%27s_algorithm)
	//
	// from: https://www.geeksforgeeks.org/graph-implementation-using-stl-for-competitive-programming-set-1-dfs-of-unweighted-and-undirected/
	// A utility function to do DFS of graph
	// recursively from a given vertex u.
	void Graph::DFSUtil(int u, std::vector< std::vector<int> > adj, std::vector<bool> &visited)
	{
		visited[u] = true;
		for (int i=0; i<adj[u].size(); i++)
			if (visited[adj[u][i]] == false)
				DFSUtil(adj[u][i], adj, visited);
	};

	// used for finding articulation points (1-node-connectivity):
	//
	// for reference, see Tarjan’s algorithm
	//
	// from: https://www.geeksforgeeks.org/articulation-points-or-cut-vertices-in-a-graph/
	// A recursive function that find articulation points using DFS traversal
	// u --> The vertex to be visited next
	// visited[] --> keeps tract of visited vertices
	// disc[] --> Stores discovery times of visited vertices
	// parent[] --> Stores parent vertices in DFS tree
	// ap[] --> Store articulation points
	void Graph::APUtil(int u, std::vector<bool> & visited, int disc[],
										int low[], std::vector<int> & parent, std::vector<bool> & ap)
	{
		// A static variable is used for simplicity, we can avoid use of static
		// variable by passing a pointer.
		static int time = 0;

		// Count of children in DFS Tree
		int children = 0;

		// Mark the current node as visited
		visited[u] = true;

		// Initialize discovery time and low value
		disc[u] = low[u] = ++time;

		for (int i=0; i<adjacency_list_[u].size(); i++) {
			int v = adjacency_list_[u][i];  // v is current adjacent of u

			// If v is not visited yet, then make it a child of u
			// in DFS tree and recur for it
			if (!visited[v]) {
				children++;
				parent[v] = u;
				APUtil(v, visited, disc, low, parent, ap);

				// Check if the subtree rooted with v has a connection to
				// one of the ancestors of u
				low[u]  = std::min(low[u], low[v]);

				// u is an articulation point in following cases

				// (1) u is root of DFS tree and has two or more chilren.
				if (parent[u] == -1 && children > 1)
				ap[u] = true;

				// (2) If u is not root and low value of one of its child is more
				// than discovery value of u.
				if (parent[u] != -1 && low[v] >= disc[u])
				ap[u] = true;
			}

			// Update low value of u for parent function calls.
			else if (v != parent[u])
				low[u]  = std::min(low[u], disc[v]);
		}
	};

	// used for finding articulation points (1-node-connectivity):
	//
	// for reference, see Tarjan’s algorithm
	//
	// from: https://www.geeksforgeeks.org/bridge-in-a-graph/
	// A recursive function that finds and prints bridges using
	// DFS traversal
	// u --> The vertex to be visited next
	// visited[] --> keeps tract of visited vertices
	// disc[] --> Stores discovery times of visited vertices
	// parent[] --> Stores parent vertices in DFS tree
	void Graph::bridgeUtil(int u,
						std::vector<bool> & visited,
						int disc[],
						int low[],
						std::vector<int> & parent,
						std::vector<int> & bridges)
	{
		// A static variable is used for simplicity, we can
		// avoid use of static variable by passing a pointer.
		static int time = 0;

		// Mark the current node as visited
		visited[u] = true;

		// Initialize discovery time and low value
		disc[u] = low[u] = ++time;

		// Go through all vertices aadjacent to this
		for (int i=0; i<adjacency_list_[u].size(); i++) {

			int v = adjacency_list_[u][i];

			// If v is not visited yet, then recur for it
			if (!visited[v]) {
				parent[v] = u;

				bridgeUtil(v, visited, disc, low, parent, bridges);

				// Check if the subtree rooted with v has a
				// connection to one of the ancestors of u
				low[u]  = std::min(low[u], low[v]);

				if (low[v] > disc[u]) {
					bridges.push_back(adjacency_edge_list_[u][i]);
				}

				// If the lowest vertex reachable from subtree
				// under v is  below u in DFS tree, then u-v
				// is a bridge
				if (low[v] > disc[u] && opts_.verbose)
				std::cout << "Bridge at: " << u << " " << v << std::endl;
			}

			// Update low value of u for parent function calls.
			else if (v != parent[u])
				low[u]  = std::min(low[u], disc[v]);
		}
	};

};

#endif


#ifndef NNQS_OPTIMIZERS_ADADELTA_HPP
#define NNQS_OPTIMIZERS_ADADELTA_HPP
#include <Eigen/Dense>
#include <type_traits>
#include <complex>
#include "Utilities/type_traits.hpp"
#include "Optimizers/Optimizer.hpp"

namespace yannq
{

template<typename T>
class AdaDelta
	: public OptimizerGeometry<T>
{
public:
	using typename OptimizerGeometry<T>::RealT;
	using typename OptimizerGeometry<T>::Vector;
	using typename OptimizerGeometry<T>::RealVector;

private:
	const double rho_;

	int t;
	RealVector v_;
	RealVector x_;

public:
	static constexpr double DEFAULT_PARAMS[1] = {0.95};

	static nlohmann::json defaultParams()
	{
		return nlohmann::json
		{
			{"name", "AdaDelta"},
			{"rho", DEFAULT_PARAMS[0]},
		};
	}

	nlohmann::json params() const override
	{
		return nlohmann::json
		{
			{"name", "AdaDelta"},
			{"rho", rho_},
		};
	}

	AdaDelta(const nlohmann::json& params)
		: rho_(params.value("rho", DEFAULT_PARAMS[0])),
		t{}
	{
	}

	AdaDelta(double rho = DEFAULT_PARAMS[0])
		: rho_(rho)
	{
	}

	Vector getUpdate(const Vector& grad, const Vector& oloc) override
	{
		const double eps = 1e-6;
		auto toRMS = [eps](RealT x){ return sqrt(x)+eps; };
		auto norm = [](T x) -> RealT { return std::norm(x); };

		if(t == 0)
		{
			v_ = RealVector::Zero(grad.rows());
			x_ = RealVector::Zero(grad.rows());
		}
		++t;

		RealVector g2 = oloc.unaryExpr(norm);
		v_ *= (1-rho_);
		v_ += rho_*g2;

		RealVector num = x_.unaryExpr(toRMS);
		RealVector denom = v_.unaryExpr(toRMS);
		Vector res = -grad.cwiseProduct(num).cwiseQuotient(denom);

		x_ *= (1-rho_);
		x_ += rho_*res.unaryExpr(norm);

		return res;
	}

};
template<typename T>
constexpr double AdaDelta<T>::DEFAULT_PARAMS[];
} //namespace yannq
#endif//NNQS_OPTIMIZERS_ADADELTA_HPP


/**
* \file auxiliaryFunctions.cpp
* \brief functions for the module auxiliaryFunctions
* \author Guillaume St-Onge
* \version 1.0
* \date 08/11/2017
*/

#include "auxiliaryFunctions.hpp"
#include <iostream>
#include <fstream>
#include <sstream>
#include <cmath>
#include <boost/random/uniform_int_distribution.hpp>
#include <boost/random/discrete_distribution.hpp>

using namespace std;

namespace DynNet
{//start of namespace DynNet

typedef boost::random::uniform_int_distribution<> uniform_int;

/**
 * \fn void order(edge p_edge)
 * \brief Assure the edge is in ascending order
 * \param[in] p_edge
 */
void order(edge& p_edge)
{
	if (p_edge.first > p_edge.second)
	{
		swap(p_edge.first,p_edge.second);
	}
}

/**
 * \fn vector<edge> input_edgeList(string p_path)
 * \brief Input the edge list from file
 * \param[in] p_path path name to the file
 */
vector<edge> input_edgeList(string p_path)
{
	ifstream inStream;
	inStream.open(p_path, ios::in);
	string line;

	vector<edge> edgeList;

	while(getline(inStream,line))
	{
		edge e;
		stringstream line_stream(line);
		line_stream >> e.first >> e.second;
		//swap edge to assure ascending order
		order(e);
		edgeList.push_back(e);
	}
	inStream.close();

	return edgeList;
}

/**
 * \fn pair<double,double> update_estimator(DynamicNetwork& p_net, edge& p_chosenEdge,
	pair<double,double>& p_est_param)ble,double>& p_model_param,
	edge& p_chosenEdge, double& p_degreeNormalization)
 * \brief Get likelihood for the new edge and update normalisation.
 * \param[in] p_net Structure of a network in reconstruction.
 * \param[in] p_chosenEdge new edge
 * \param[in] p_est_param Estimated parameter
 */
pair<double,double> update_estimator(DynamicNetwork& p_net, edge& p_chosenEdge,
	pair<double,double>& p_est_param)
{
	double n = p_net.get_currentDegreeMap().size();
	double m = p_net.get_currentEdgeList().size();
	double kavg = 2*m/n;
	double ksum = kavg*n;
	double k2avg = p_est_param.first + kavg*kavg;
	double k2sum = n*k2avg;

	unordered_map<node, unsigned int> currentDegreeMap = p_net.get_currentDegreeMap();

	//update parameter according to new node
	m += 1;
	ksum += 2;
	if (currentDegreeMap.find(p_chosenEdge.first) !=
		currentDegreeMap.end())
	{
		if (currentDegreeMap.find(p_chosenEdge.second) !=
		currentDegreeMap.end())
		{
			//2 old nodes, n stays the same
			//check if self-loop
			if (p_chosenEdge.first == p_chosenEdge.second)
			{
				k2sum -= (currentDegreeMap[p_chosenEdge.first]) * (currentDegreeMap[p_chosenEdge.first]);
				k2sum += (currentDegreeMap[p_chosenEdge.first] + 2) * (currentDegreeMap[p_chosenEdge.first] + 2);
				// k2sum -= pow(currentDegreeMap[p_chosenEdge.first],2.);
				// k2sum += pow(currentDegreeMap[p_chosenEdge.first] + 2,2.);
			}
			else
			{
			// 	k2sum -= pow(currentDegreeMap[p_chosenEdge.first],2.);
			// 	k2sum += pow(currentDegreeMap[p_chosenEdge.first] + 1,2.);
			// 	k2sum -= pow(currentDegreeMap[p_chosenEdge.second],2.);
			// 	k2sum += pow(currentDegreeMap[p_chosenEdge.second] + 1,2.);
				k2sum -= (currentDegreeMap[p_chosenEdge.first]) * (currentDegreeMap[p_chosenEdge.first]);
				k2sum += (currentDegreeMap[p_chosenEdge.first] + 1) * (currentDegreeMap[p_chosenEdge.first] + 1);
				k2sum -= (currentDegreeMap[p_chosenEdge.second]) * (currentDegreeMap[p_chosenEdge.second]);
				k2sum += (currentDegreeMap[p_chosenEdge.second] + 1) * (currentDegreeMap[p_chosenEdge.second] + 1);
			}
		}
		else
		{
			//first is old, second is new
			n += 1;
			// k2sum -= pow(currentDegreeMap[p_chosenEdge.first],2.);
			// k2sum += pow(currentDegreeMap[p_chosenEdge.first] + 1,2.);
			k2sum -= (currentDegreeMap[p_chosenEdge.first]) *  (currentDegreeMap[p_chosenEdge.first]);
			k2sum += (currentDegreeMap[p_chosenEdge.first] + 1) * (currentDegreeMap[p_chosenEdge.first] + 1);
			k2sum += 1;
		}
	}
	else
	{
		//first is new, second is old
		n += 1;
		// k2sum -= pow(currentDegreeMap[p_chosenEdge.second],2.);
		// k2sum += pow(currentDegreeMap[p_chosenEdge.second] + 1,2.);
		k2sum -= (currentDegreeMap[p_chosenEdge.second]) * (currentDegreeMap[p_chosenEdge.second]);
		k2sum += (currentDegreeMap[p_chosenEdge.second] + 1) * (currentDegreeMap[p_chosenEdge.second] + 1);
		k2sum += 1;
	}

	// pair<double,double> new_est_param(k2sum/n - pow(ksum/n, 2.),(n-2)/(m-1));
	pair<double,double> new_est_param(k2sum/n - (ksum/n) * (ksum/n), (n-2)/(m-1));

	return new_est_param;
}

/**
 * \fn double update_modelLikelihood(DynamicNetwork& p_net, pair<double,
 	double>& p_model_param,	edge& p_chosenEdge, double& p_degreeNormalization)
 * \brief Get likelihood for the new edge and update normalisation.
 * \param[in] p_net Structure of a network in reconstruction.
 * \param[in] p_model_param Model parameters
 * \param[in] p_chosenEdge new edge
 * \param[in] p_degreeNormalization Normalization factor for probability
 */
double update_modelLikelihood(DynamicNetwork& p_net,
	pair<double,double>& p_model_param, edge& p_chosenEdge,
	double& p_degreeNormalization)
{
	double model_likelihood = 1.;
	if (p_net.get_currentDegreeMap().find(p_chosenEdge.first) !=
		p_net.get_currentDegreeMap().end())
	{
		if (p_net.get_currentDegreeMap().find(p_chosenEdge.second) !=
		p_net.get_currentDegreeMap().end())
		{
			//2 old nodes
			double k1 = p_net.get_currentDegreeMap().at(p_chosenEdge.first);
			double k2 = p_net.get_currentDegreeMap().at(p_chosenEdge.second);
			model_likelihood *= (1 - p_model_param.second)*pow(k1*k2,
				p_model_param.first)/pow(p_degreeNormalization,2.);
			//check if self-loop
			if (p_chosenEdge.first == p_chosenEdge.second)
			{
				p_degreeNormalization -= pow(k1,p_model_param.first);
				p_degreeNormalization += pow(k1+2,p_model_param.first);
			}
			else
			{
				p_degreeNormalization -= pow(k1,p_model_param.first);
				p_degreeNormalization -= pow(k2,p_model_param.first);
				p_degreeNormalization += pow(k1+1,p_model_param.first);
				p_degreeNormalization += pow(k2+1,p_model_param.first);
			}
		}
		else
		{
			//first is old, second is new
			double k1 = p_net.get_currentDegreeMap().at(p_chosenEdge.first);
			model_likelihood *= p_model_param.second*pow(k1,
				p_model_param.first)/p_degreeNormalization;
			p_degreeNormalization -= pow(k1,p_model_param.first);
			p_degreeNormalization += pow(k1+1,p_model_param.first)+1;
		}
	}
	else
	{
		//first is new, second is old
		double k2 = p_net.get_currentDegreeMap().at(p_chosenEdge.second);
		model_likelihood *= p_model_param.second*pow(k2,
			p_model_param.first)/p_degreeNormalization;
		p_degreeNormalization -= pow(k2,p_model_param.first);
		p_degreeNormalization += pow(k2+1,p_model_param.first)+1;
	}

	return model_likelihood;
}

/**
 * \fn void add_edge(DynamicNetwork& p_net, pair<double,double>& p_model_param,
	pair<double,double>& p_target_param, pair<double,double>& p_est_param,
	pair<double,double>& p_bias_param, RNGType& gen)
 * \brief Add a new edge to the network following to the growth model.
 * \param[in] p_net Structure of a network in reconstruction.
 * \param[in] p_model_param Model parameters
 * \param[in] p_target_param Target parameter for bias
 * \param[in] p_est_param Estimated parameter for bias
 * \param[in] p_bias_param Tunable parameter for bias
 * \param[in] p_degreeNormalization Normalization factor for probability
 * \param[in] gen Random number generator.
 */
void add_edge(DynamicNetwork& p_net, pair<double,double>& p_model_param,
	pair<double,double>& p_target_param, pair<double,double>& p_est_param,
	pair<double,double>& p_bias_param, double& p_degreeNormalization,
	RNGType& gen)
{
	//this add_edge method is used with important_sampling

	edgeSet reachableEdgeSet = p_net.get_reachableEdgeSet();
	unordered_map<unsigned int, edge> indexMap;
	vector<double> weightList(reachableEdgeSet.size(), 0);
	// vector<double> weightList;
	double weightNorm = 0.;
	bool oldSeed = false;
	bool oldTarget = false;

	//get weights
	unsigned int index = 0;
	for (auto iter = reachableEdgeSet.begin();
		iter != reachableEdgeSet.end(); ++iter)
	{
		indexMap[index] = *iter;

		double weight = p_net.get_remainingEdgeMap().at(*iter); //bias for multiple edges
		edge possibleEdge = *iter;
		pair<double,double> new_est_param = update_estimator(p_net, possibleEdge,
			p_est_param);
		// weight *= exp(-p_bias_param.first*(new_est_param.second
		// 	- p_target_param.second));
		// weight *= exp(-p_bias_param.second*(new_est_param.first
		// 	- p_target_param.first));
		weight *= exp(-p_bias_param.first*(new_est_param.second- p_target_param.second) * (new_est_param.second- p_target_param.second));  // density
		weight *= exp(-p_bias_param.second*(new_est_param.first- p_target_param.first) * (new_est_param.first- p_target_param.first));  // variance

		// weightList.push_back(weight);

		weightList[index] = weight;
		weightNorm += weight;

		index += 1;
	}

	//Get a new random edge
	boost::random::discrete_distribution<int> edgeDist(weightList.begin(), weightList.end());
	unsigned int newEdgeIndex = edgeDist(gen);
	edge chosenEdge = indexMap[newEdgeIndex];

	//determine the model likelihood
	double model_likelihood = update_modelLikelihood(p_net, p_model_param, chosenEdge,
		p_degreeNormalization);

	//update the estimator
	p_est_param = update_estimator(p_net, chosenEdge, p_est_param);

	// cout << p_est_param.first << " " << p_est_param.second << "\n";

	//add the edge with according probability ratio (model_likelihood/bias_prob)
	double ratio = model_likelihood*weightNorm/(weightList[newEdgeIndex]);
	//account for multiple edge indistinguishability
	ratio *= p_net.get_remainingEdgeMap().at(chosenEdge);

	p_net.add(chosenEdge, ratio);
}

/**
 * \fn void add_edge(DynamicNetwork& p_net, pair<double,double>& p_param,
	double& p_degreeNormalization, RNGType& gen)
 * \brief Add a new edge to the network following to the growth model.
 * \param[in] p_net Structure of a network in reconstruction.
 * \param[in] p_param Model parameters
 * \param[in] p_degreeNormalization Normalization factor for probability
 * \param[in] gen Random number generator.
 */
void add_edge(DynamicNetwork& p_net, pair<double,double>& p_param,
	double& p_degreeNormalization, RNGType& gen)
{
	edgeSet reachableEdgeSet = p_net.get_reachableEdgeSet();
	uniform_int indexDist(0, p_net.get_reachableEdge()-1);

	// cout << reachableEdgeSet.size() << " " << p_net.get_reachableEdge() << endl;
	int index = indexDist(gen);

	auto iter = reachableEdgeSet.begin();
	//count multiple edge with remaining map
	while (index >= 0)
	{
		unsigned int remaining = p_net.get_remainingEdgeMap().at(*iter);
		index -= remaining;
		if (index >= 0)
		{
			iter++ ;
		}
	}

	bool oldSeed = false;
	bool oldTarget = false;
	if (p_net.get_currentDegreeMap().find(iter->first) !=
		p_net.get_currentDegreeMap().end())
	{
		oldSeed = true;
	}
	else
	{
		oldSeed = false;
	}
	if (p_net.get_currentDegreeMap().find(iter->second) !=
		p_net.get_currentDegreeMap().end())
	{
		oldTarget = true;
	}
	else
	{
		oldTarget = false;
	}

	double weight = p_net.get_reachableEdge();
	if (oldSeed)
	{
		if (oldTarget)
		{
			double k1 = p_net.get_currentDegreeMap().at(iter->first);
			double k2 = p_net.get_currentDegreeMap().at(iter->second);
			weight *= (1- p_param.second)/pow(p_degreeNormalization,2.);
			weight *= pow(k1, p_param.first);
			weight *= pow(k2, p_param.first);
			if (iter->second == iter->first)
			{
				p_degreeNormalization -= pow(k1, p_param.first);
				p_degreeNormalization += pow(k1+2, p_param.first);
			}
			else
			{
				p_degreeNormalization -= pow(k1, p_param.first);
				p_degreeNormalization -= pow(k2, p_param.first);
				p_degreeNormalization += pow(k1+1., p_param.first);
				p_degreeNormalization += pow(k2+1., p_param.first);
			}
		}
		else
		{
			double k1 = p_net.get_currentDegreeMap().at(iter->first);
			weight *= p_param.second/p_degreeNormalization;
			weight *= pow(k1, p_param.first);
			p_degreeNormalization -= pow(k1, p_param.first);
			p_degreeNormalization += pow(k1+1., p_param.first) + 1;
		}
	}
	else
	{
		if (oldTarget)
		{
			double k2 = p_net.get_currentDegreeMap().at(iter->second);
			weight *= p_param.second/p_degreeNormalization;
			weight *= pow(k2, p_param.first);
			// weight *= p_net.get_remainingEdgeMap().at(*iter);//bias to account
			p_degreeNormalization -= pow(k2, p_param.first);
			p_degreeNormalization += pow(k2+1., p_param.first) + 1;
		}
		else
		{
			//impossible to have 2 new nodes
			cout << "ERROR : two new nodes" << endl;
		}
	}

	p_net.add(*iter, weight);
}

/**
 * \fn void update_meanMarginal(DynamicNetwork& p_net, unordered_map< edge,
 	vector<double>, boost::hash<edge> > p_meanMarginalMap,
 	double meanLogweight, double& p_effectiveWeightSum)
 * \brief Update mean marginal evaluation from a network sample.
 * \param[in] p_net Structure of a network in reconstruction.
 * \param[in] p_meanMarginalMap Map of the mean marginal.
 * \param[in] p_meanLogweight Mean value of the logweight.
 * \param[in] p_effectiveWeightSum Normalization value to update.
 */
void update_meanMarginal(DynamicNetwork& p_net, unordered_map< edge,
	vector<double>, boost::hash<edge> >& p_meanMarginalMap,
	double meanLogweight, double& p_effectiveWeightSum)
{
	edgeIntMap edgeCountMap;
	double effectiveWeight = exp(p_net.get_logweight()-meanLogweight);
	double t = 0.;
	for (auto iter = p_net.get_currentEdgeList().begin();
		iter != p_net.get_currentEdgeList().end(); ++iter)
	{
		p_meanMarginalMap[*iter][edgeCountMap[*iter]] += t*effectiveWeight;
		edgeCountMap[*iter] += 1;
		t += 1.;
	}
	p_effectiveWeightSum += effectiveWeight;
}

/**
 * \fn bool update_meanMarginal(DynamicNetwork& p_net, unordered_map< edge,
 	vector<double>, boost::hash<edge> > p_meanMarginalMap,
 	double meanLogweight, double& p_effectiveWeightSum, double p_rtol)
 * \brief Update mean marginal evaluation from a network sample.
 * \param[in] p_net Structure of a network in reconstruction.
 * \param[in] p_meanMarginalMap Map of the mean marginal.
 * \param[in] p_meanLogweight Mean value of the logweight.
 * \param[in] p_effectiveWeightSum Normalization value to update.
 * \param[in] p_rtol Tolerance for convergence.
 */
bool update_meanMarginal(DynamicNetwork& p_net, unordered_map< edge,
	vector<double>, boost::hash<edge> >& p_meanMarginalMap,
	double meanLogweight, double& p_effectiveWeightSum, double p_rtol,
	unsigned int sampleSize)
{
	edgeIntMap edgeCountMap;
	bool converged = true;
	double effectiveWeight = exp(p_net.get_logweight()-meanLogweight);
	p_effectiveWeightSum += effectiveWeight;
	double t = 0.;
	for (auto iter = p_net.get_currentEdgeList().begin();
		iter != p_net.get_currentEdgeList().end(); ++iter)
	{
		p_meanMarginalMap[*iter][edgeCountMap[*iter]] += t*effectiveWeight;
		double avg_t = (p_meanMarginalMap[*iter][edgeCountMap[*iter]]
			/p_effectiveWeightSum);
		if (abs(t-avg_t)/sqrt(sampleSize) > p_rtol) //comportement moyen
		{
			converged = false;
		}
		edgeCountMap[*iter] += 1;
		t += 1.;
	}

	return converged;
}

}//end of namespace DynNet


#include <iostream>
#include <fstream>

#include <boost/numeric/ublas/matrix.hpp>
#include <boost/numeric/ublas/vector.hpp>
#include <boost/numeric/ublas/io.hpp>

#include <vector>
#include <cmath>

#include "NRGclasses.hpp"
#include "NRGfunctions.hpp"
#include "TwoChQSz.hpp"

#ifndef pi
#define pi 3.141592653589793238462643383279502884197169
#endif

bool AeqB(CNRGbasisarray *pAeigCut, int ibl1,int ibl2){

  return (ibl1==ibl2);

}

double AtimesB(CNRGbasisarray *pAeig, CNRGbasisarray *pSingleSite,
	       int ibl1,int ibl2){

  return((double)ibl1*ibl2);

}

int main (){

  CNRGarray Aeig(2);

  CNRGbasisarray AeigCut(2);

  CNRGbasisarray Abasis(2);

  CNRGbasisarray SingleSite(2);

  // STL vector

  CNRGmatrix Qm1fNQ[2];

  CNRGmatrix MQQp1;

  double U,ed,Gamma1,Gamma2;
  vector<double> Params;
  double Lambda;
  double HalfLambdaFactor;
  double Dband=1.0;
  int calcdens;

  double DN=0.0;
  double TM=0.0;
  double betabar=0.727;
  double Temp=0.0;
  double Sus=0.0;
  //vector<double> SuscepChain;
  char arqSus[32],arqname[32];

  double chi_m1,chi_N[2];
  double daux[4];

  int Nsites,Nsitesmax=2;

  int Ncutoff=700;
  int UpdateBefCut=0;
  int auxIn;

  // outstream
  ofstream OutFile;
  // instream
  ifstream InFile;

  int ii,jj,i1,i2;

  // STL iterator:

  vector<double>::iterator diter;

  // Test
  CNRGmatrix HN;

  HN.CheckForMatEl=AeqB;
  HN.CalcMatEl=AtimesB;
  if (HN.CheckForMatEl(&Abasis,2,2)) cout << "TRUE!" << endl;
  else cout << "FALSE!" << endl;
  cout << HN.CalcMatEl(&Abasis,&SingleSite,3,4) << endl;
  exit(0);

  ///            ///
  /// Begin code ///
  ///            ///
  U=0.5;
  ed=-0.5*U;
  Gamma1=0.0282691;
  Gamma2=0.0;
  Lambda=2.5;

  // Steps: Input parameters
  InFile.open("nrg_input_TwoChAnderson.dat");
  if (InFile.is_open())
    {
      InFile >> Nsitesmax;
      InFile >> Ncutoff;
      InFile >> U;
      InFile >> Gamma1;
      InFile >> Gamma2;
      InFile >> ed;
      InFile >> Lambda;
      InFile >> Dband;
      InFile >> auxIn;
      InFile >> UpdateBefCut;
      InFile >> calcdens;
    }
  else cout << "can't open nrg_input_TwoChAnderson.dat" << endl;

  InFile.close();

  ///////////////////////////
  ///////////////////////////

  if (calcdens==2)
    {
      //SuscepChain.clear();
      strcpy(arqSus,"SuscepImp_25_726.dat");

      strcpy(arqname,"SuscepChain.dat");
      InFile.open(arqname);
      if (InFile.fail())
	{
	  cout << "Can't find " << arqname << endl;
	  strcpy(arqSus,"SuscepChain.dat");
	  calcdens=3;
	}
///////////////////////
//       else
// 	{
// 	  while (!InFile.eof())
// 	    {
// 	      InFile >> Temp >> daux[0] >> daux[0] >> Sus;
// 	      SuscepChain.push_back(Sus);
// 	    }
// 	  SuscepChain.pop_back();
// 	  if (SuscepChain.size()<Nsitesmax-1)
// 	    {
// 	      cout << "Nsuscep = " << SuscepChain.size() << " Nsitesmax = " << Nsitesmax << endl;
// 	      cout << " Longer SuscepChain needed! Doing it all over..." << endl;
// 	      SuscepChain.clear();
// 	      strcpy(arqSus,"SuscepChain.dat");
// 	      calcdens=3;
// 	    }
// 	}
//       InFile.close();
//////////////////////////
     }
  if (calcdens==3)
    {
      strcpy(arqSus,"SuscepChain.dat");
      U=0.0;ed=0.0;Gamma1=0.0;Gamma2=0.0;
    }

  ///////////////////////////
  ///////////////////////////

  HalfLambdaFactor=0.5*(1.0+(1.0/Lambda));
  chi_m1=sqrt(2.0*Gamma1/pi)/(sqrt(Lambda)*HalfLambdaFactor);

  double U_tilde=0.5*U/HalfLambdaFactor;
  double ed_tilde=ed/HalfLambdaFactor;
  double Gamma1_tilde=Gamma1*(2.0/pi)/(HalfLambdaFactor*HalfLambdaFactor);
  double Gamma2_tilde=Gamma2*(2.0/pi)/(HalfLambdaFactor*HalfLambdaFactor);

  OutFile.open("NRG_in.txt");
  OutFile << "Begin NRG 2-ch calculation" << endl;
  OutFile << " Nsitesmax    = " << Nsitesmax-1 << endl;
  OutFile << " Ncutoff      = " << Ncutoff << endl;
  OutFile << " U            = " << U << endl;
  OutFile << " Gamma1        = " << Gamma1 << endl;
  OutFile << " Gamma2        = " << Gamma2 << endl;
  OutFile << " ed           = " << ed << endl;
  OutFile << " Lambda       = " << Lambda << endl;
  OutFile << " Dband        = " << Dband << endl;
  OutFile << " UpdateBefCut = " << UpdateBefCut << endl;
  OutFile << " calcdens     = " << calcdens << endl;
  OutFile << "=================================" << endl;
  OutFile << "U~ = " << U_tilde << endl;
  OutFile << "ed~ = " <<  ed_tilde << endl;
  OutFile << "Gamma1~ = " <<  Gamma1_tilde<< endl;
  OutFile << "Gamma2~ = " <<  Gamma2_tilde<< endl;
  OutFile << "=================================" << endl;
  OutFile.close();

  // Define H0 (impurity + 1s site)
  // Output:
  //       Aeig,
  //       fd_{1sigma} and fd_{2sigma} matrix elements
  //

  // Set single site (use pointers in the subroutines!!)

  TwoChQSz_SetSingleSite(&SingleSite);

  Params.push_back(U_tilde/sqrt(Lambda));
  Params.push_back(ed_tilde/sqrt(Lambda));
  Params.push_back(sqrt(Gamma1_tilde/Lambda));
  Params.push_back(sqrt(Gamma2_tilde/Lambda));

  Nsites=0;
  DN=HalfLambdaFactor*pow(Lambda,(-(Nsites-1)/2.0) );
  TM=DN/betabar;
  cout << "DN = " << DN << "TM = " << TM << endl;

  //TwoChQSz_SetH0Anderson(Params,&SingleSite,&Aeig,Qm1fNQ);
  TwoChQSz_SetH0Anderson(Params,&SingleSite,&Aeig,&Abasis);

  Aeig.PrintEn();

  // Calculate Susceptibility

  if ( (calcdens==2)||(calcdens==3) )
    {
      TM=DN/betabar;
      Params.clear();
      Params.push_back(betabar);
      Sus=CalcSuscep(Params,&Aeig,1,false);
      double Sus0=0.0;
      int nlines=0;
      if (calcdens==2)
	{
	  InFile.open("SuscepChain.dat");
	  InFile.clear();
	  InFile.seekg(0, ios::beg); // rewind
	  if (InFile.fail())
	    {
	      cout << "Can't open  SuscepChain.dat"<< endl;
	    }
	  else
	    {
	      while ( (!InFile.eof())&&(nlines<=Nsites) )
		{
		  InFile >> daux[0] >> daux[0] >> daux[0] >> Sus0;
		  nlines++;
		}
	    }
	  InFile.close();
	}
      //if (calcdens==2) Sus0=SuscepChain[Nsites];
      else Sus-=1.0/8.0; // exclude dot site in the chain.
      if (Nsites==0) OutFile.open(arqSus);
      else OutFile.open(arqSus,ofstream::app);
      OutFile.precision(20);
      OutFile << scientific << TM << " " << Sus << " " << Sus0 << " " << Sus-Sus0 << endl;
      OutFile.close();

    }

  TwoChQSz_UpdateQm1fQ(&SingleSite,&Aeig,&Abasis,Qm1fNQ);

  // Check Matrix
  for (int ich=1;ich<=2;ich++)
    {
      cout << " Printing Qm1fQ channel: " << ich << endl;
      for (int ibl=0; ibl<Qm1fNQ[ich-1].NumMatBlocks();ibl++)
	{
	  Qm1fNQ[ich-1].PrintMatBlock(ibl);
	}
    }

  // Loop on Nsites: start from Nsites=0 (imp+1)
  //

  Nsites=1;
  while (Nsites<=Nsitesmax)
    {

      DN=HalfLambdaFactor*pow(Lambda,(-(Nsites-1)/2.0) );
      TM=DN/betabar;
      cout << "DN = " << DN << "TM = " << TM << endl;

      // 0 - Update chi_N, eps_N

      daux[0]=(double)( 1.0-pow(Lambda,(-Nsites)) );
      daux[1]=(double)sqrt( 1.0-pow(Lambda,-(2*Nsites-1)) );
      daux[2]=(double)sqrt( 1.0-pow(Lambda,-(2*Nsites+1)) );
      daux[3]=0.5*(1.0+(1.0/Lambda))*(double)sqrt(Lambda);

      chi_N[0]=daux[0]/(daux[1]*daux[2]);
      chi_N[1]=daux[0]/(daux[1]*daux[2]);

      cout << "chi_N = " << chi_N[0] << "  " << chi_N[1] << endl;

      cout << "Nsites = " << Nsites << endl;
      cout << "BEG Eig Nshell = " << Aeig.Nshell << endl;

      // 1 - Eliminate states and Build Abasis

      cout << "Cutting states..." << endl;
      AeigCut.CNRGbasisarray::ClearAll();
      AeigCut=CutStates(&Aeig, Ncutoff);

      cout << "... done cutting states." << endl;

      // Calculate new matrix elements using the CUT basis:
      // update Qm1fNQ, Qm1cdQ, etc.

      if ( (UpdateBefCut==0)&&(Nsites>1) )
	{
	  cout << "Updating matrices after cutting... " << endl;
	  TwoChQSz_UpdateMatrixAfterCutting(&SingleSite,
					    &AeigCut, &Abasis, Qm1fNQ, &MQQp1);
	  cout << "... done updating matrices. " << endl;
	}

      //Q1Q2Sz_BuildBasis(&AeigCut,&Abasis,&SingleSite,UpdateBefCut);
      // doesnt work
      QSz_BuildBasis(&AeigCut,&Abasis,&SingleSite,UpdateBefCut);

      cout << "No blocks = " << Abasis.NumBlocks() << endl;

      cout << "No states = " << Abasis.Nstates() << endl;

      //Abasis.PrintAll();

      // 2 - Build and diagonalize H_N+1

      Params.clear();
      Params.push_back(chi_N[0]);
      Params.push_back(chi_N[1]);
      Params.push_back(Lambda);

      cout << "Diagonalizing HN... " << endl;

      TwoChQSz_DiagHN(Params,&Abasis,&SingleSite,Qm1fNQ,&Aeig);

      Aeig.PrintEn();

      cout << "..done diagonalizing HN. " << endl;

      // 3 - Update Qm1f1NQ, Qm1f2Q, Qm1cdQ, etc.

      if ( (UpdateBefCut==1)&&(Nsites<Nsitesmax) )
	{
	  cout << "Updating matrices before cutting... " << endl;
	  TwoChQSz_UpdateQm1fQ(&SingleSite,&Aeig,&Abasis,Qm1fNQ);
	  cout << "... done updating matrices. " << endl;

	}

      // Calculate Susceptibility

      if ( (calcdens==2)||(calcdens==3) )
	{
	  TM=DN/betabar;
	  Params.clear();
	  Params.push_back(betabar);
	  Sus=CalcSuscep(Params,&Aeig,1,false);
	  double Sus0=0.0;
	  int nlines=0;
	  // Gets SuscepChain iteration by iteration
	  if (calcdens==2)
	    {
	      InFile.open("SuscepChain.dat");
	      InFile.clear();
	      InFile.seekg(0, ios::beg); // rewind
	      if (InFile.fail())
		{
		  cout << "Can't find " << arqname << endl;
		}
	      else
		{
		  while ( (!InFile.eof())&&(nlines<=Nsites) )
		    {
		      InFile >> daux[0] >> daux[0] >> daux[0] >> Sus0;
		      nlines++;
		    }
		}
	      InFile.close();
	    }
	  //if (calcdens==2) Sus0=SuscepChain[Nsites];
	  else Sus-=1.0/8.0; // exclude dot site in the chain.

	  if (Nsites==0) OutFile.open(arqSus);
	  else OutFile.open(arqSus,ofstream::app);
	  OutFile.precision(20);
	  OutFile << scientific << TM << " " << Sus << " " << Sus0 << " " << Sus-Sus0 << endl;
	  OutFile.close();
	}

      // 4 - Update Nsites

      Nsites++;

    }

  cout << "=== Calculation Finished! ==== "<< endl;
  OutFile.open("NRG_end.txt");
  OutFile << "END NRG calculation" << endl;
  OutFile.close();

  cout << "Calling destructors " << endl;

}
//END code


// Filename: CirculationModel.cpp
// Created on 20 Aug 2007 by Boyce Griffith

// Modified 2019, Alexander D. Kaiser

#include "CirculationModel.h"

/////////////////////////////// INCLUDES /////////////////////////////////////

#ifndef included_IBAMR_config
#include <IBAMR_config.h>
#define included_IBAMR_config
#endif

#ifndef included_SAMRAI_config
#include <SAMRAI_config.h>
#define included_SAMRAI_config
#endif

// SAMRAI INCLUDES
#include <CartesianGridGeometry.h>
#include <CartesianPatchGeometry.h>
#include <PatchLevel.h>
#include <SideData.h>
#include <tbox/RestartManager.h>
#include <tbox/SAMRAI_MPI.h>
#include <tbox/Utilities.h>

// C++ STDLIB INCLUDES
#include <cassert>
#include <cmath>

#include <Eigen/Dense>
using namespace Eigen;

/////////////////////////////// NAMESPACE ////////////////////////////////////

/////////////////////////////// STATIC ///////////////////////////////////////

#define MIN_PER_L_T0_SEC_PER_ML 60.0e-3   // 60/1000
#define MMHG_TO_CGS 1333.22368

#define USE_WINDKESSEL
// If not defined then this only computes fluxes
// Pressure is set to zero
// Currently hard coded for upper boundary

namespace
{
    // Name of output file.
    static const string DATA_FILE_NAME = "bc_data.m";

    // constants
    static const double C_PA         =  4.12; // Pulmonary artery compliance, ml / mmHg
    static const double C_LA_relaxed =  3*1.6;  // Left atrial compliance ml / mmHg
    static const double C_PV         = 10.0 - C_LA_relaxed;  // Pulmonary vein compliance, ml / mmHg

    static const double R_P  = (9.0/5.6) * MIN_PER_L_T0_SEC_PER_ML; // Pulmonary resistance, mmHg / (ml/s)

    static const double beat_time = 0.8;
    static const double T_on = .53;   // Pulmonary valve open
    static const double T_off = .75;  // Pulmonary valve closes
    static const double T_peak = T_on + 0.4 * (T_off - T_on); // Peak pulmonary valve flow
    static const double stroke_volume = 75.0; // ml
    static const double h = 2.0 * stroke_volume / (T_off - T_on); // Peak flow to get given stroke volume

    static const bool   atrial_kick_on = true;
    static const double atrial_kick_center = .44;
    static const double atrial_kick_time_radius = .09;
    static const double atrial_kick_time_width = 2.0 * atrial_kick_time_radius;

    inline double compute_Q_R(double t){
        // Triangle wave flux

        double t_reduced = t - beat_time * floor(t/beat_time);

        if (t_reduced <= T_on)
            return 0.0;
        else if (t_reduced <= T_peak)
            return ( h/(T_peak - T_on) )*t_reduced - (h/(T_peak - T_on)  )*T_on;
        else if (t_reduced <= T_off)
            return (-h/(T_off - T_peak))*t_reduced + (h/(T_off -  T_peak))*T_off;
        else if (t_reduced <= beat_time)
            return 0.0;

        TBOX_ERROR("Valid time for flux not found.");
        return 0.0;
    }

    inline double compute_C_LA(double t){

        if(atrial_kick_on){
            double t_reduced = t - beat_time * floor(t/beat_time);

            if (abs(t_reduced - atrial_kick_center) < atrial_kick_time_radius)
                return C_LA_relaxed * (1.0 - pow(cos(M_PI * (t_reduced - atrial_kick_center) / atrial_kick_time_width),2));
            else
                return C_LA_relaxed;
        }

        return C_LA_relaxed;
    }

    // Backward Euler update for windkessel model.
    inline void
    windkessel_be_update(double& Q_R, double& P_PA, double& Q_P, double& P_LA, const double Q_mi, const double t, const double dt)
    {

        Q_R = compute_Q_R(t);

        double C_LA_current = compute_C_LA(t);
        double C_LA_next    = compute_C_LA(t + dt);

        double a = C_PA/dt + 1/R_P;
        double b = -1/R_P;
        double c = -1/R_P;
        double d = (C_PV + C_LA_next)/dt + 1/R_P;

        double rhs[2];
        rhs[0] = (C_PA/dt)*P_PA + Q_R;
        rhs[1] = ((C_PV + C_LA_current)/dt)*P_LA - Q_mi;

        double det = a*d - b*c;

        // Closed form linear system solution
        P_PA = (1/det) * ( d*rhs[0] + -b*rhs[1]);
        P_LA = (1/det) * (-c*rhs[0] +  a*rhs[1]);

        Q_P = (1/R_P) * (P_PA - P_LA);

        return;
    } // windkessel_be_update

}

/////////////////////////////// PUBLIC ///////////////////////////////////////

CirculationModel::CirculationModel(const string& object_name, double P_PA_0, double P_LA_0, double t, bool register_for_restart)
    : d_object_name(object_name),
      d_registered_for_restart(register_for_restart),
      d_time(t),
      d_nsrc(1),           // number of sets of variables
      d_psrc(d_nsrc, 0.0), // pressure
      d_qsrc(d_nsrc, 0.0), // flux
      d_srcname(d_nsrc),
      d_P_PA(P_PA_0),
      d_P_LA(P_LA_0),
      d_Q_R(0.0),
      d_Q_P(0.0),
      d_Q_mi(0.0),
      d_bdry_interface_level_number(numeric_limits<int>::max())
{
#if !defined(NDEBUG)
    assert(!object_name.empty());
#endif
    if (d_registered_for_restart)
    {
        RestartManager::getManager()->registerRestartItem(d_object_name, this);
    }

    // Initialize object with data read from the input and restart databases.
    const bool from_restart = RestartManager::getManager()->isFromRestart();
    if (from_restart)
    {
        getFromRestart();
    }
    else
    {
        //   nsrcs = the number of sources in the valve tester:
        //           (1) left atrium
        d_srcname[0] = "left atrium       ";
    }
    return;
} // CirculationModel

CirculationModel::~CirculationModel()
{
    return;
} // ~CirculationModel

void
CirculationModel::advanceTimeDependentData(const double dt,
                                           const double Q_mi)
{
    /*
    // Compute the mean flow rates in the vicinity of the inflow and outflow
    // boundaries.
    std::fill(d_qsrc.begin(), d_qsrc.end(), 0.0);
    for (int ln = 0; ln <= hierarchy->getFinestLevelNumber(); ++ln)
    {
        Pointer<PatchLevel<NDIM> > level = hierarchy->getPatchLevel(ln);
        for (PatchLevel<NDIM>::Iterator p(level); p; p++)
        {
            Pointer<Patch<NDIM> > patch = level->getPatch(p());
            Pointer<CartesianPatchGeometry<NDIM> > pgeom = patch->getPatchGeometry();
            if (pgeom->getTouchesRegularBoundary())
            {
                Pointer<SideData<NDIM, double> > U_data = patch->getPatchData(U_idx);
                Pointer<SideData<NDIM, double> > wgt_sc_data = patch->getPatchData(wgt_sc_idx);
                const Box<NDIM>& patch_box = patch->getBox();
                // const double* const x_lower = pgeom->getXLower();
                const double* const dx = pgeom->getDx();
                double dV = 1.0;
                for (int d = 0; d < NDIM; ++d)
                {
                    dV *= dx[d];
                }

                static const int axis = 2;  // Always z axis here
                const int side = 1;         // Compute flux at the top only

                const bool is_lower = side == 0;
                if (pgeom->getTouchesRegularBoundary(axis, side))
                {

                    Vector n;
                    for (int d = 0; d < NDIM; ++d)
                    {
                        n[d] = axis == d ? (is_lower ? -1.0 : +1.0) : 0.0;
                    }
                    Box<NDIM> side_box = patch_box;
                    if (is_lower)
                    {
                        side_box.lower(axis) = patch_box.lower(axis);
                        side_box.upper(axis) = patch_box.lower(axis);
                    }
                    else
                    {
                        side_box.lower(axis) = patch_box.upper(axis) + 1;
                        side_box.upper(axis) = patch_box.upper(axis) + 1;
                    }
                    for (Box<NDIM>::Iterator b(side_box); b; b++)
                    {
                        const Index<NDIM>& i = b();

                        // no conditional here, just add the flux in
                        const SideIndex<NDIM> i_s(i, axis, SideIndex<NDIM>::Lower);
                        if ((*wgt_sc_data)(i_s) > std::numeric_limits<double>::epsilon())
                        {
                            double dA = n[axis] * dV / dx[axis];
                            d_qsrc[0] += (*U_data)(i_s)*dA;

                            // pout << "adding " << (*U_data)(i_s)*dA << "to the flux\n";

                        }
                    }
                }
            }
        }
    }
    SAMRAI_MPI::sumReduction(&d_qsrc[0], d_nsrc);

    // pout << "computed flux = " << d_qsrc[0] << "\n";
    */

    // The downstream (Atrial) pressure is determined by a zero-d model
    const double t = d_time;

    double& Q_R  = d_Q_R;
    double& P_PA = d_P_PA;
    double& Q_P  = d_Q_P;
    double& P_LA = d_P_LA;

    // Mitral flux passed in
    d_Q_mi = Q_mi;

    windkessel_be_update(Q_R, P_PA, Q_P, P_LA, d_Q_mi, t, dt);

    // model in mmHg, body force converts
    d_psrc[0] = d_P_LA;

    // Update the current time.
    d_time += dt;

    // Output the updated values.
    const long precision = plog.precision();
    plog.unsetf(ios_base::showpos);
    plog.unsetf(ios_base::scientific);

    plog.precision(12);

    plog << "============================================================================\n"
         << "Circulation model variables at time " << d_time << ":\n";

    plog << "P_PA (mmHg)\t P_LA (mmHg)\t Q_R (ml/s)\t Q_P (ml/s)\t Q_mi (ml/s)\n";
    plog.setf(ios_base::showpos);
    plog.setf(ios_base::scientific);

    plog << d_P_PA << ",\t " << d_P_LA << ",\t " << d_Q_R << ",\t " << d_Q_P << ",\t " << d_Q_mi << "\n";
    plog << "============================================================================\n";

    plog.unsetf(ios_base::showpos);
    plog.unsetf(ios_base::scientific);
    plog.precision(precision);

    // Write the current state to disk.
    writeDataFile();
    return;
} // advanceTimeDependentData

void
CirculationModel::putToDatabase(Pointer<Database> db)
{
    db->putDouble("d_time", d_time);
    db->putInteger("d_nsrc", d_nsrc);
    db->putDoubleArray("d_qsrc", &d_qsrc[0], d_nsrc);
    db->putDoubleArray("d_psrc", &d_psrc[0], d_nsrc);
    db->putStringArray("d_srcname", &d_srcname[0], d_nsrc);
    db->putDouble("d_P_PA", d_P_PA);
    db->putDouble("d_P_LA", d_P_LA);
    db->putDouble("d_Q_R", d_Q_R);
    db->putDouble("d_Q_P", d_Q_P);
    db->putDouble("d_Q_mi", d_Q_mi);
    db->putInteger("d_bdry_interface_level_number", d_bdry_interface_level_number);
    return;
} // putToDatabase

void CirculationModel::write_plot_code()
{
    static const int mpi_root = 0;
    if (SAMRAI_MPI::getRank() == mpi_root)
    {
        ofstream fout(DATA_FILE_NAME.c_str(), ios::app);
        fout.setf(ios_base::scientific);
        fout.setf(ios_base::showpos);
        fout.precision(10);
        fout << "];\n";
        fout << "fig = figure;\n";
        fout << "subplot(3,2,1)\n";
        fout << "plot(bc_vals(:,1), bc_vals(:,2))\n";
        fout << "title('P_{PA}')\n";
        fout << "subplot(3,2,2)\n";
        fout << "plot(bc_vals(:,1), bc_vals(:,3))\n";
        fout << "title('P_{LA}')\n";
        fout << "subplot(3,2,3)\n";
        fout << "plot(bc_vals(:,1), bc_vals(:,4))\n";
        fout << "title('Q_{R}')\n";
        fout << "subplot(3,2,4)\n";
        fout << "plot(bc_vals(:,1), bc_vals(:,5))\n";
        fout << "title('Q_{P}')\n";
        fout << "subplot(3,2,5)\n";
        fout << "plot(bc_vals(:,1), bc_vals(:,6))\n";
        fout << "title('Q_{mi}')\n";
        fout << "dt = bc_vals(2,1) - bc_vals(1,1);\n";
        fout << "net_flux = dt*cumsum(bc_vals(:,6));\n";
        fout << "subplot(3,2,6)\n";
        fout << "plot(bc_vals(:,1), net_flux)\n";
        fout << "title('net Q')\n";
        fout << "printfig(fig, 'bc_model_variables')\n";
    }
    return;
}

/////////////////////////////// PROTECTED ////////////////////////////////////

/////////////////////////////// PRIVATE //////////////////////////////////////

void
CirculationModel::writeDataFile() const
{
    static const int mpi_root = 0;
    if (SAMRAI_MPI::getRank() == mpi_root)
    {
        static bool file_initialized = false;
        const bool from_restart = RestartManager::getManager()->isFromRestart();
        if (!from_restart && !file_initialized)
        {
            ofstream fout(DATA_FILE_NAME.c_str(), ios::out);
            fout << "% time \t P_PA (mmHg)\t d_P_LA (mmHg)\t Q_R (ml/s)\t Q_P (ml/s)\t Q_mi (ml/s)"
                 << "\n"
                 << "bc_vals = [";
            file_initialized = true;
        }

        ofstream fout(DATA_FILE_NAME.c_str(), ios::app);
        for (int n = 0; n < d_nsrc; ++n)
        {
            fout << d_time;
            fout.setf(ios_base::scientific);
            fout.setf(ios_base::showpos);
            fout.precision(10);
            fout << " " << d_P_PA << " " << d_P_LA << " " << d_Q_R << " " << d_Q_P << " " << d_Q_mi << "; \n";
        }
    }
    return;
} // writeDataFile

void
CirculationModel::getFromRestart()
{
    Pointer<Database> restart_db = RestartManager::getManager()->getRootDatabase();
    Pointer<Database> db;
    if (restart_db->isDatabase(d_object_name))
    {
        db = restart_db->getDatabase(d_object_name);
    }
    else
    {
        TBOX_ERROR("Restart database corresponding to " << d_object_name << " not found in restart file.");
    }

    d_time = db->getDouble("d_time");
    d_nsrc = db->getInteger("d_nsrc");
    d_qsrc.resize(d_nsrc);
    d_psrc.resize(d_nsrc);
    d_srcname.resize(d_nsrc);
    db->getDoubleArray("d_qsrc", &d_qsrc[0], d_nsrc);
    db->getDoubleArray("d_psrc", &d_psrc[0], d_nsrc);
    db->getStringArray("d_srcname", &d_srcname[0], d_nsrc);
    d_P_PA = db->getDouble("d_P_PA");
    d_P_LA = db->getDouble("d_P_LA");
    d_Q_R  = db->getDouble("d_Q_R");
    d_Q_P  = db->getDouble("d_Q_P");
    d_Q_mi = db->getDouble("d_Q_mi");
    d_bdry_interface_level_number = db->getInteger("d_bdry_interface_level_number");
    return;
} // getFromRestart

/////////////////////////////// NAMESPACE ////////////////////////////////////

/////////////////////////////// TEMPLATE INSTANTIATION ///////////////////////

//////////////////////////////////////////////////////////////////////////////


/**
 *
 * Copyright (C) 2021 Mohammad Javad Dousti, Qing Xie, Mahdi Nazemi,
 * and Massoud Pedram. All rights reserved.
 *
 * Please refer to the LICENSE file for terms of use.
 *
 */

#pragma once

#include "general.hpp"
#include <vector>
#include <Eigen/Core>
#include <Eigen/Sparse>

typedef Eigen::Matrix<VALUE, Eigen::Dynamic, 1> state_type;
typedef Eigen::SparseMatrix<VALUE> matrix_type;

#if USE_GPU == 1

#include <cuda_runtime.h>
#include <cusolverDn.h>
#include <cusolverSp.h>
#include <cusparse.h>

#else
#define EIGEN_SUPERLU_SUPPORT
#define EIGEN_USE_MKL_ALL

// A quick and dirty fix to support MKL 11.2
// http://eigen.tuxfamily.org/bz/show_bug.cgi?id=874
#ifndef MKL_BLAS
#define MKL_BLAS MKL_DOMAIN_BLAS
#endif

#include <boost/numeric/odeint/external/eigen/eigen.hpp>

#include <Eigen/Dense>
#include <Eigen/IterativeLinearSolvers>
#include <boost/numeric/odeint.hpp>
#include <boost/numeric/odeint/stepper/runge_kutta_dopri5.hpp>
// #include <boost/numeric/odeint/external/eigen/eigen_algebra.hpp>
// #include <boost/numeric/odeint/algebra/vector_space_algebra.hpp>

// #include <boost/numeric/odeint/algebra/algebra_dispatcher.hpp>
// #include <boost/numeric/odeint/algebra/vector_space_algebra.hpp>

#include <unsupported/Eigen/IterativeSolvers>

// using namespace Eigen;
// using namespace boost::numeric::odeint;

#endif

namespace solver {

#if USE_GPU == 1
void gpuSolveDenseSteady(const matrix_type& matrix, const state_type& p_vector,
                         state_type t_vector, bool use_cholesky = true);

void gpuSolveSparseSteady(const matrix_type& matrix, const state_type& p_vector,
                                  state_type& t_vector);

void gpuSolveTransient();
#else
void cpuSolveSteady(const matrix_type &matrix,
                    const state_type &p_vector, state_type &t_vector);
void cpuSolveTransient(
    const matrix_type &a_matrix,
    const Eigen::DiagonalMatrix<VALUE, Eigen::Dynamic, Eigen::Dynamic> &invC,
    const state_type &p_vector, state_type &t_vector, VALUE time_start,
    VALUE time_end);

#endif

void writeTemp(const state_type &t_vector_, const VALUE t);
}; // namespace solver


/**
 * @file BtreeTests.cpp
 *
 * @breif Experimentation with a Neural Network querying position
 *
 * @date 1/04/2018
 * @author Ben Caine
 */

#define BOOST_TEST_DYN_LINK
#define BOOST_TEST_MODULE NetworkIndexTests

#include <boost/test/unit_test.hpp>
#include <chrono>
#include <random>
#include <iomanip>
#include <fstream>
#include <unordered_set>
#include "../src/utils/DataGenerators.h"
#include "../external/nn_cpp/nn/Net.h"

std::unordered_set<size_t> getRandomSubset(int numValues, int maxValue) {
    std::unordered_set<size_t> values;
    auto seed = std::chrono::high_resolution_clock::now().time_since_epoch().count();
    std::mt19937 rng(seed);
    std::uniform_int_distribution<size_t> gen(0, maxValue);

    while (values.size() < numValues) {
        auto val = gen(rng);
        values.insert(val);
    }
    return values;
}

BOOST_AUTO_TEST_CASE(basic_linear_net) {

    /**
     * This is a test of a simple linear model predicting the position in a sorted list of integer
     * lognormals generated from a lognormal distribution of mean = 0, std = 2.0, where
     * we scale all the values such that the max value is equal to maxValue.
     *
     * This gives us a test of convergence, hyperparam sensitivity, ext for predicting a general location
     * of a value in a sorted list, and acts as a precursor to the second level of the
     * Recursive Model Index
     */
    // Hyperparams
    int batchSize = 64;
    float learningRate = 0.01;
    int numEpochs = 10000;
    const size_t datasetSize = 1000;
    float maxValue = 100;

    auto values = getIntegerLognormals<size_t, datasetSize>(maxValue);

    // File to write loss data to
    std::ofstream outputFile("loss.csv");

    bool useBias = true;
    // Simple linear model
    nn::Net<float> net;
    net.add(new nn::Dense<float, 2>(batchSize, 1, 1, useBias, nn::InitializationScheme::GlorotNormal));

    // NOTE: MSE blows up the gradients. Smooth L1 Loss behaves much nicer here
    nn::HuberLoss<float, 2> lossFunction;

    // NOTE: SGD simply does not work. Adam works in Pytorch, with a lot of sensitivity to learning rate
    net.registerOptimizer(new nn::Adam<float>(learningRate));

    Eigen::Tensor<float, 2> input(batchSize, 1);
    Eigen::Tensor<float, 2> positions(batchSize, 1);

    auto startTime = std::chrono::system_clock::now();
    for (unsigned int currentEpoch = 0; currentEpoch < numEpochs; ++currentEpoch) {
        auto newBatch = getRandomSubset(batchSize, datasetSize);
        unsigned int ii = 0;
        for (auto idx : newBatch) {
            input(ii, 0) = static_cast<float>(values[idx]);
            positions(ii, 0) = static_cast<float>(idx);
            ii ++;
        }

        auto result = net.forward<2, 2>(input);
        result = result * result.constant(datasetSize);

        auto loss = lossFunction.loss(result, positions);
        std::cout << "Epoch: " << currentEpoch << " Loss: " << loss << std::endl;
        outputFile << currentEpoch << ", " << loss << "\n";

        auto lossBack = lossFunction.backward(result, positions);
        lossBack = lossBack / lossBack.constant(datasetSize);
        net.backward<2>(lossBack);
        net.step();
    }
    auto endTime = std::chrono::system_clock::now();
    std::chrono::duration<float> duration = endTime - startTime;
    std::cout << "Total training of " << numEpochs << " iters took: " << duration.count() << "s" << std::endl;

    outputFile.close();

    // Test
    auto newBatch = getRandomSubset(batchSize, datasetSize);
    unsigned int ii = 0;
    for (auto idx : newBatch) {
        input(ii, 0) = static_cast<float>(values[idx]);
        positions(ii, 0) = static_cast<float>(idx);
        ii ++;
    }
    auto result = net.forward<2, 2>(input);

    for (unsigned int row = 0; row < batchSize; ++row) {
        std::cout << std::fixed << std::setprecision(0) << input(row, 0) << ", " << positions(row, 0) << ", " << result(row, 0) * datasetSize << std::endl;
    }
}

BOOST_AUTO_TEST_CASE(basic_net) {

    /**
     * This is a test of a simple network predicting the position in a sorted list of integer
     * lognormals generated from a lognormal distribution of mean = 0, std = 2.0, where
     * we scale all the values such that the max value is equal to maxValue.
     *
     * This gives us a test of convergence, hyperparam sensitivity, ext for predicting a general location
     * of a value in a sorted list, and acts as a precursor to the first level of the
     * Recursive Model Index
     */
    // Hyperparams
    int batchSize = 256;
    int numNeurons = 8;
    float learningRate = 0.01;
    int numEpochs = 25000;
    const size_t datasetSize = 100000;
    float maxValue = 1e5;

    auto values = getIntegerLognormals<size_t, datasetSize>(maxValue);

    // File to write loss data to
    std::ofstream outputFile("loss.csv");

    bool useBias = true;
    // Simple linear model
    nn::Net<float> net;
    net.add(new nn::Dense<float, 2>(batchSize, 1, numNeurons, useBias, nn::InitializationScheme::GlorotNormal));
    net.add(new nn::Relu<float, 2>());
    net.add(new nn::Dense<float, 2>(batchSize, numNeurons, 1, useBias, nn::InitializationScheme::GlorotNormal));

    // NOTE: MSE blows up the gradients. Smooth L1 Loss behaves much nicer here
    nn::HuberLoss<float, 2> lossFunction;

    // NOTE: SGD simply does not work. Adam works in Pytorch, with a lot of sensitivity to learning rate
    net.registerOptimizer(new nn::Adam<float>(learningRate));

    Eigen::Tensor<float, 2> input(batchSize, 1);
    Eigen::Tensor<float, 2> positions(batchSize, 1);

    auto startTime = std::chrono::system_clock::now();
    for (unsigned int currentEpoch = 0; currentEpoch < numEpochs; ++currentEpoch) {
        auto newBatch = getRandomSubset(batchSize, datasetSize);
        unsigned int ii = 0;
        for (auto idx : newBatch) {
            input(ii, 0) = static_cast<float>(values[idx]);
            positions(ii, 0) = static_cast<float>(idx);
            ii ++;
        }

        auto result = net.forward<2, 2>(input);
        result = result * result.constant(datasetSize);

        auto loss = lossFunction.loss(result, positions);
        std::cout << "Epoch: " << currentEpoch << " Loss: " << loss << std::endl;
        outputFile << currentEpoch << ", " << loss << "\n";

        auto lossBack = lossFunction.backward(result, positions);
        lossBack = lossBack / lossBack.constant(datasetSize);
        net.backward<2>(lossBack);
        net.step();
    }
    auto endTime = std::chrono::system_clock::now();
    std::chrono::duration<float> duration = endTime - startTime;
    std::cout << "Total training of " << numEpochs << " iters took: " << duration.count() << "s" << std::endl;

    outputFile.close();

    // Test
    auto newBatch = getRandomSubset(batchSize, datasetSize);
    unsigned int ii = 0;
    for (auto idx : newBatch) {
        input(ii, 0) = static_cast<float>(values[idx]);
        positions(ii, 0) = static_cast<float>(idx);
        ii ++;
    }
    auto result = net.forward<2, 2>(input);

    for (unsigned int row = 0; row < batchSize; ++row) {
        std::cout << std::fixed << std::setprecision(0) << input(row, 0) << ", " << positions(row, 0) << ", " << result(row, 0) * datasetSize << std::endl;
    }
}