// Copyright (C) 2014 The Regents of the University of California (Regents).
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:
//
//     * Redistributions of source code must retain the above copyright
//       notice, this list of conditions and the following disclaimer.
//
//     * Redistributions in binary form must reproduce the above
//       copyright notice, this list of conditions and the following
//       disclaimer in the documentation and/or other materials provided
//       with the distribution.
//
//     * Neither the name of The Regents or University of California nor the
//       names of its contributors may be used to endorse or promote products
//       derived from this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDERS OR CONTRIBUTORS BE
// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
// POSSIBILITY OF SUCH DAMAGE.
//
// Please contact the author of this library if you have any questions.
// Author: Chris Sweeney (cmsweeney@cs.ucsb.edu)

// BSD 3-Clause License

// Copyright (c) 2021, Chenyu
// All rights reserved.

// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are met:

// 1. Redistributions of source code must retain the above copyright notice, this
//    list of conditions and the following disclaimer.

// 2. Redistributions in binary form must reproduce the above copyright notice,
//    this list of conditions and the following disclaimer in the documentation
//    and/or other materials provided with the distribution.

// 3. Neither the name of the copyright holder nor the names of its
//    contributors may be used to endorse or promote products derived from
//    this software without specific prior written permission.

// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
// FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
// CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
// OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

#include "translation_averaging/lud_position_estimator.h"

#include <unordered_set>

#include <ceres/rotation.h>
#include <Eigen/SparseCore>
#include <glog/logging.h>

#include "solver/constrained_l1_solver.h"
#include "util/map_util.h"

namespace gopt {
namespace {

Eigen::Vector3d GetRotatedTranslation(
    const Eigen::Vector3d& rotation_angle_axis,
    const Eigen::Vector3d& translation) {
  Eigen::Matrix3d rotation;
  ceres::AngleAxisToRotationMatrix(
      rotation_angle_axis.data(),
      ceres::ColumnMajorAdapter3x3(rotation.data()));
  return rotation.transpose() * translation;
}

}  // namespace

LUDPositionEstimator::LUDPositionEstimator(
    const LUDPositionEstimator::Options& options)
    : options_(options) {
  CHECK_GT(options_.max_num_iterations, 0);
  CHECK_GT(options_.max_num_reweighted_iterations, 0);
}

bool LUDPositionEstimator::EstimatePositions(
    const std::unordered_map<ImagePair, TwoViewGeometry>& view_pairs,
    const std::unordered_map<image_t, Eigen::Vector3d>& orientations,
    std::unordered_map<image_t, Eigen::Vector3d>* positions) {
  CHECK_NOTNULL(positions)->clear();

  InitializeIndexMapping(view_pairs, orientations);
  const size_t num_views = view_id_to_index_.size();
  const size_t num_view_pairs = view_id_pair_to_index_.size();

  // Set up the linear system.
  SetupConstraintMatrix(view_pairs, orientations);
  Eigen::VectorXd solution;
  solution.setZero(constraint_matrix_.cols());

  // Create the lower bound constraint enforcing that all scales are > 1.
  Eigen::SparseMatrix<double> geq_mat(num_view_pairs,
                                      constraint_matrix_.cols());
  for (size_t i = 0; i < num_view_pairs; i++) {
    geq_mat.insert(i, 3 * (num_views - 1) + i) = 1.0;
  }
  Eigen::VectorXd geq_vec(num_view_pairs);
  geq_vec.setConstant(1.0);

  Eigen::VectorXd b(constraint_matrix_.rows());
  b.setZero();

  // Solve for camera positions by solving a constrained L1 problem to enforce
  // all relative translations scales > 1.
  ConstrainedL1Solver::Options l1_options;
  ConstrainedL1Solver solver(
      l1_options, constraint_matrix_, b, geq_mat, geq_vec);
  solver.Solve(&solution);

  // Set the estimated positions.
  for (const auto& view_id_index : view_id_to_index_) {
    const int index = view_id_index.second;
    const image_t view_id = view_id_index.first;
    if (index == kConstantViewIndex) {
      (*positions)[view_id] = Eigen::Vector3d::Zero();
    } else {
      (*positions)[view_id] = solution.segment<3>(index);
    }
  }

  return true;
}

void LUDPositionEstimator::InitializeIndexMapping(
    const std::unordered_map<ImagePair, TwoViewGeometry>& view_pairs,
    const std::unordered_map<image_t, Eigen::Vector3d>& orientations) {
  std::unordered_set<image_t> views;
  for (const auto& view_pair : view_pairs) {
    if (ContainsKey(orientations, view_pair.first.first) &&
        ContainsKey(orientations, view_pair.first.second)) {
      views.insert(view_pair.first.first);
      views.insert(view_pair.first.second);
    }
  }

  // Create a mapping from the view id to the index of the linear system.
  int index = kConstantViewIndex;
  view_id_to_index_.reserve(views.size());
  for (const image_t view_id : views) {
    view_id_to_index_[view_id] = index;
    index += 3;
  }

  // Create a mapping from the view id pair to the index of the linear system.
  view_id_pair_to_index_.reserve(view_pairs.size());
  for (const auto& view_pair : view_pairs) {
    if (ContainsKey(view_id_to_index_, view_pair.first.first) &&
        ContainsKey(view_id_to_index_, view_pair.first.second)) {
      view_id_pair_to_index_[view_pair.first] = index;
      ++index;
    }
  }
}

void LUDPositionEstimator::SetupConstraintMatrix(
    const std::unordered_map<ImagePair, TwoViewGeometry>& view_pairs,
    const std::unordered_map<image_t, Eigen::Vector3d>& orientations) {
  constraint_matrix_.resize(
      3 * view_id_pair_to_index_.size(),
      3 * (view_id_to_index_.size() - 1) + view_pairs.size());

  // Add the camera to camera constraints.
  std::vector<Eigen::Triplet<double>> triplet_list;
  triplet_list.reserve(9 * view_pairs.size());
  int row = 0;
  for (const auto& view_pair : view_pairs) {
    const ImagePair view_id_pair = view_pair.first;
    if (!ContainsKey(view_id_to_index_, view_id_pair.first) ||
        !ContainsKey(view_id_to_index_, view_id_pair.second)) {
      continue;
    }

    const int view1_index = FindOrDie(view_id_to_index_, view_id_pair.first);
    const int view2_index = FindOrDie(view_id_to_index_, view_id_pair.second);
    const int scale_index =
        FindOrDieNoPrint(view_id_pair_to_index_, view_id_pair);

    // Rotate the relative translation so that it is aligned to the global
    // orientation frame.
    const Eigen::Vector3d translation_direction =
        GetRotatedTranslation(FindOrDie(orientations, view_id_pair.first),
                              view_pair.second.translation_2);

    // Add the constraint for view 1 in the minimization:
    //   position2 - position1 - scale_1_2 * translation_direction.
    if (view1_index != kConstantViewIndex) {
      triplet_list.emplace_back(row + 0, view1_index + 0, -1.0);
      triplet_list.emplace_back(row + 1, view1_index + 1, -1.0);
      triplet_list.emplace_back(row + 2, view1_index + 2, -1.0);
    }

    // Add the constraint for view 2 in the minimization:
    //   position2 - position1 - scale_1_2 * translation_direction.
    if (view2_index != kConstantViewIndex) {
      triplet_list.emplace_back(row + 0, view2_index + 0, 1.0);
      triplet_list.emplace_back(row + 1, view2_index + 1, 1.0);
      triplet_list.emplace_back(row + 2, view2_index + 2, 1.0);
    }

    // Add the constraint for scale in the minimization:
    //   position2 - position1 - scale_1_2 * translation_direction.
    triplet_list.emplace_back(row + 0, scale_index, -translation_direction[0]);
    triplet_list.emplace_back(row + 1, scale_index, -translation_direction[1]);
    triplet_list.emplace_back(row + 2, scale_index, -translation_direction[2]);

    row += 3;
  }

  constraint_matrix_.setFromTriplets(triplet_list.begin(), triplet_list.end());

  VLOG(2) << view_pairs.size() << " camera to camera constraints were added "
                                  "to the position estimation problem.";
}

}  // namespace gopt


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

#ifndef NETKET_MOMENTUM_HPP
#define NETKET_MOMENTUM_HPP

#include <Eigen/Core>
#include <Eigen/Dense>
#include <cassert>
#include <cmath>
#include <complex>
#include <iostream>
#include "abstract_stepper.hpp"

namespace netket {

class Momentum : public AbstractStepper {
  int npar_;

  double eta_;
  double beta_;

  Eigen::VectorXd mt_;

  int mynode_;

  const std::complex<double> I_;

 public:
  // Json constructor
  explicit Momentum(const json &pars)
      : eta_(FieldOrDefaultVal(pars["Learning"], "LearningRate", 0.001)),
        beta_(FieldOrDefaultVal(pars["Learning"], "Beta", 0.9)),
        I_(0, 1) {
    npar_ = -1;

    PrintParameters();
  }

  void PrintParameters() {
    MPI_Comm_rank(MPI_COMM_WORLD, &mynode_);
    if (mynode_ == 0) {
      std::cout << "# Momentum stepper initialized with these parameters : "
                << std::endl;
      std::cout << "# Learning rate = " << eta_ << std::endl;
      std::cout << "# Beta = " << beta_ << std::endl;
    }
  }

  void Init(const Eigen::VectorXd &pars) override {
    npar_ = pars.size();
    mt_.setZero(npar_);

  }

  void Init(const Eigen::VectorXcd &pars) override {
    npar_ = 2 * pars.size();
    mt_.setZero(npar_);

  }

  void Update(const Eigen::VectorXd &grad, Eigen::VectorXd &pars) override {
    assert(npar_ > 0);

    mt_=beta_*mt_+(1.-beta_)*grad;

    for(int i=0;i<npar_;i++){
      pars(i)-=eta_*mt_(i);
    }
  }

  void Update(const Eigen::VectorXcd &grad, Eigen::VectorXd &pars) override {
    Update(Eigen::VectorXd(grad.real()), pars);
  }

  void Update(const Eigen::VectorXcd &grad, Eigen::VectorXcd &pars) override {
    assert(npar_ == 2 * pars.size());

    for(int i=0;i<pars.size();i++){
      mt_(2*i)=beta_*mt_(2*i)+(1.-beta_)*grad(i).real();
      mt_(2*i+1)=beta_*mt_(2*i+1)+(1.-beta_)*grad(i).imag();
    }

    for(int i=0;i<pars.size();i++){
      pars(i)-=eta_*mt_(2*i);
      pars(i)-=eta_*I_*mt_(2*i+1);
    }
  }

  void Reset() override {
    mt_ = Eigen::VectorXd::Zero(npar_);
  }

};

}  // namespace netket

#endif


#include <boost/graph/transitive_reduction.hpp>


/*
Copyright 2008 Intel Corporation

Use, modification and distribution are subject to the Boost Software License,
Version 1.0. (See accompanying file LICENSE_1_0.txt or copy at
http://www.boost.org/LICENSE_1_0.txt).
*/
#include <boost/polygon/polygon.hpp>
#include <cassert>
#include <list>
namespace gtl = boost::polygon;
using namespace boost::polygon::operators;

//first lets turn our polygon usage code into a generic
//function parameterized by polygon type
template <typename Polygon>
void test_polygon() {
  //lets construct a 10x10 rectangle shaped polygon
  typedef typename gtl::polygon_traits<Polygon>::point_type Point;
  Point pts[] = {gtl::construct<Point>(0, 0),
		 gtl::construct<Point>(10, 0),
		 gtl::construct<Point>(10, 10),
		 gtl::construct<Point>(0, 10) };
  Polygon poly;
  gtl::set_points(poly, pts, pts+4);

  //now lets see what we can do with this polygon
  assert(gtl::area(poly) == 100.0f);
  assert(gtl::contains(poly, gtl::construct<Point>(5, 5)));
  assert(!gtl::contains(poly, gtl::construct<Point>(15, 5)));
  gtl::rectangle_data<int> rect;
  assert(gtl::extents(rect, poly)); //get bounding box of poly
  assert(gtl::equivalence(rect, poly)); //hey, that's slick
  assert(gtl::winding(poly) == gtl::COUNTERCLOCKWISE);
  assert(gtl::perimeter(poly) == 40.0f);

  //add 5 to all coords of poly
  gtl::convolve(poly, gtl::construct<Point>(5, 5));
  //multiply all coords of poly by 2
  gtl::scale_up(poly, 2);
  gtl::set_points(rect, gtl::point_data<int>(10, 10),
		  gtl::point_data<int>(30, 30));
  assert(gtl::equivalence(poly, rect));
}

//Now lets declare our own polygon class
//Oops, we need a point class to support our polygon, lets borrow
//the CPoint example
struct CPoint {
  int x;
  int y;
};

//we have to get CPoint working with boost polygon to make our polygon
//that uses CPoint working with boost polygon
namespace boost { namespace polygon {
    template <>
    struct geometry_concept<CPoint> { typedef point_concept type; };
    template <>
    struct point_traits<CPoint> {
      typedef int coordinate_type;

      static inline coordinate_type get(const CPoint& point,
					orientation_2d orient) {
	if(orient == HORIZONTAL)
	  return point.x;
	return point.y;
      }
    };

    template <>
    struct point_mutable_traits<CPoint> {
      typedef int coordinate_type;

      static inline void set(CPoint& point, orientation_2d orient, int value) {
	if(orient == HORIZONTAL)
	  point.x = value;
	else
	  point.y = value;
      }
      static inline CPoint construct(int x_value, int y_value) {
	CPoint retval;
	retval.x = x_value;
	retval.y = y_value;
	return retval;
      }
    };
  } }

//I'm lazy and use the stl everywhere to avoid writing my own classes
//my toy polygon is a std::list<CPoint>
typedef std::list<CPoint> CPolygon;

//we need to specialize our polygon concept mapping in boost polygon
namespace boost { namespace polygon {
    //first register CPolygon as a polygon_concept type
    template <>
    struct geometry_concept<CPolygon>{ typedef polygon_concept type; };

    template <>
    struct polygon_traits<CPolygon> {
      typedef int coordinate_type;
      typedef CPolygon::const_iterator iterator_type;
      typedef CPoint point_type;

      // Get the begin iterator
      static inline iterator_type begin_points(const CPolygon& t) {
	return t.begin();
      }

      // Get the end iterator
      static inline iterator_type end_points(const CPolygon& t) {
	return t.end();
      }

      // Get the number of sides of the polygon
      static inline std::size_t size(const CPolygon& t) {
	return t.size();
      }

      // Get the winding direction of the polygon
      static inline winding_direction winding(const CPolygon& t) {
	return unknown_winding;
      }
    };

    template <>
    struct polygon_mutable_traits<CPolygon> {
      //expects stl style iterators
      template <typename iT>
      static inline CPolygon& set_points(CPolygon& t,
					 iT input_begin, iT input_end) {
	t.clear();
	t.insert(t.end(), input_begin, input_end);
	return t;
      }

    };
  } }

//now there's nothing left to do but test that our polygon
//works with library interfaces
int main() {
  test_polygon<CPolygon>(); //woot!
  return 0;
}


#include <boost/math/bindings/mpfr.hpp>


///////////////////////////////////////////////////////////////////////////////////
// distribution::common::concept::concept::unary::log_unnormalized_pdf.hpp       //
//                                                                               //
//  (C) Copyright 2009 Erwann Rogard                                             //
//  Use, modification and distribution are subject to the                        //
//  Boost Software License, Version 1.0. (See accompanying file                  //
//  LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)             //
///////////////////////////////////////////////////////////////////////////////////
#ifndef BOOST_STATISTICS_DETAIL_DISTRIBUTION_COMMON_CONCEPT_UNARY_LOG_UNNORMALIZED_PDF_HPP_ER_2009
#define BOOST_STATISTICS_DETAIL_DISTRIBUTION_COMMON_CONCEPT_UNARY_LOG_UNNORMALIZED_PDF_HPP_ER_2009
#include <boost/concept_check.hpp>

namespace boost{
namespace statistics{
namespace detail{
namespace distribution{
namespace concept{

    template<typename B>
    class LogUnnormalizedPdf : public B{
        typedef typename B::distribution_type   d_;
        typedef typename B::variable_type       x_;
        typedef typename B::probability_type    p_;

        public:

        BOOST_CONCEPT_USAGE(LogUnnormalizedPdf){
        	p = log_unnormalized_pdf(d,x);
        }

        private:
        d_ d;
        p_ p;
        x_ x;

    };

}// concept
}// distribution
}// detail
}// statistics
}// boost

#endif


#ifndef STAN_MATH_TEST_FIXTURE_ODE_LORENZ_HPP
#define STAN_MATH_TEST_FIXTURE_ODE_LORENZ_HPP

#include <stan/math/rev.hpp>
#include <boost/numeric/odeint.hpp>
#include <gtest/gtest.h>
#include <test/unit/math/prim/functor/lorenz.hpp>
#include <test/unit/math/rev/functor/test_fixture_ode.hpp>
#include <test/unit/util.hpp>
#include <iostream>
#include <sstream>
#include <vector>
#include <limits>
#include <string>

struct lorenz_ode_base {
  struct lorenz_rhs {
    template <typename T0, typename T1, typename T2>
    inline Eigen::Matrix<stan::return_type_t<T1, T2>, -1, 1> operator()(
        const T0& t_in, const T1& y_in, std::ostream* msgs,
        const T2& theta) const {
      Eigen::Matrix<stan::return_type_t<T1, T2>, -1, 1> res(3);
      res << theta.at(0) * (y_in(1) - y_in(0)),
          theta.at(1) * y_in(0) - y_in(1) - y_in(0) * y_in(2),
          -theta.at(2) * y_in(2) + y_in(0) * y_in(1);
      return res;
    }
  };

  lorenz_rhs f;
  Eigen::VectorXd y0;
  std::vector<double> theta;
  double t0;
  std::vector<double> ts;
  double rtol;
  double atol;
  int max_num_step;

  lorenz_ode_base()
      : f(),
        y0(3),
        theta{10.0, 28.0, 8.0 / 3.0},
        t0(0.0),
        ts(100),
        rtol(1.e-8),
        atol(1.e-10),
        max_num_step(100000) {
    y0 << 10.0, 1.0, 1.0;
    for (int i = 0; i < 100; i++)
      ts[i] = t0 + (0.1 * (i + 1));
  }

  std::vector<double>& times() { return ts; }
  Eigen::VectorXd init() { return y0; }
  std::vector<double> param() { return theta; }
};

/**
 * Inheriting base type, various fixtures differs by the type of ODE
 * functor used in apply_solver calls, intended for
 * different kind of tests.
 *
 */
template <typename T>
struct lorenz_test : public lorenz_ode_base,
                     public ODETestFixture<lorenz_test<T>> {
  lorenz_test() : lorenz_ode_base() {}

  auto apply_solver() {
    std::tuple_element_t<0, T> sol;
    return sol(this->f, this->y0, this->t0, this->ts, nullptr, this->theta);
  }

  template <typename T1, typename T2>
  auto apply_solver(T1&& init, T2&& theta_in) {
    std::tuple_element_t<0, T> sol;
    return sol(this->f, init, this->t0, this->ts, nullptr, theta_in);
  }

  auto apply_solver_tol() {
    std::tuple_element_t<1, T> sol;
    return sol(this->f, this->y0, this->t0, this->ts, this->rtol, this->atol,
               this->max_num_step, nullptr, this->theta);
  }
};

#endif


/* Copyright 2020, 2021 Evandro Chagas Ribeiro da Rosa <evandro.crr@posgrad.ufsc.br>
 * Copyright 2020, 2021 Rafael de Santiago <r.santiago@ufsc.br>
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

//! \file ket_bitwise.hpp

#pragma once
#include <array>
#include <boost/config.hpp>
#include <boost/functional/hash.hpp>
#include <boost/unordered_map.hpp>
#include <complex>
#include <iostream>
#include <string>
#include <unordered_map>
#include <vector>

namespace ket {
    using complex = std::complex<double>; //!< Complex number.

    using ctrl_list = std::vector<size_t>; //!< List of qubits of control.

    /*! \brief Key for the quantum state hashmap.
     *  \class Index
     */
    class Index {
    public:

        //! Bitwise-or.
        Index operator|(const Index& other) const;

        //! Flip the bit in the idx possition.
        void flip(size_t idx);

        //! True if the bit in the idx possition is 1.
        bool is_one(size_t idx) const;

        //! True if the bit in the idx possition is 0.
        bool is_zero(size_t idx) const;

        //! Get a 64 bits slice: [64*idx : 64*(idx+1)] (const).
        uint64_t operator[](size_t idx) const;

        //! Get a 64 bits slice: [64*idx : 64*(idx+1)].
        uint64_t& operator[](size_t idx);

        //! Less than comparison for the hashmap indexation
        bool operator<(const Index& other) const;

        const static size_t size = 20; //!< Size of the Index: size*64.

    private:
        std::array<uint64_t, size> data;
    };

    //! Hash of an Index.
    size_t hash_value(const Index& idx);

    //! True if Index a and b are equal.
    bool operator==(const ket::Index& a, const ket::Index& b);

    //! Print the Index.
    std::ostream& operator<<(std::ostream &os, const ket::Index& idx);

    using map = boost::unordered_map<Index, complex>; //!< Quantum state hashmap.

    using dump_t = std::map<std::vector<uint64_t>, std::vector<complex>>; //! Quantum state dump.

    /*! \brief Store and manipulate the quantum state hashmap.
     *  \class Bitwise
     */
    class Bitwise {
    public:

        //! Initialize a hashmap in the state \f$\left|0\right>\f$.
        Bitwise();

        //! Tensor product between hashmaps.
        Bitwise(const Bitwise& a, const Bitwise& b);

        //! Apply a Pauli X gate with optional quibits of control.
        void x(size_t idx, const ctrl_list& ctrl = {});

        //! Apply a Pauli Y gate with optional quibits of control.
        void y(size_t idx, const ctrl_list& ctrl = {});

        //! Apply a Pauli Z gate with optional quibits of control.
        void z(size_t idx, const ctrl_list& ctrl = {});

        //! Apply a Hadamard gate with optional quibits of control.
        void h(size_t idx, const ctrl_list& ctrl = {});

        //! Apply a Phase gate with optional quibits of control.
        void s(size_t idx, const ctrl_list& ctrl = {});

        //! Apply a Phase Dagger gate with optional quibits of control.
        void sd(size_t idx, const ctrl_list& ctrl = {});

        //! Apply a T gate with optional quibits of control.
        void t(size_t idx, const ctrl_list& ctrl = {});

        //! Apply a T Dagger gate with optional quibits of control.
        void td(size_t idx, const ctrl_list& ctrl = {});

        //! Apply a CNOT gate with optional quibits of control.
        void cnot(size_t ctrl, size_t target, const ctrl_list& ctrl2 = {});

        //! Apply a Phase gate with optional quibits of control.
        void p(double lambda, size_t idx, const ctrl_list& ctrl={});

        //! Apply a U2 gate with optional quibits of control.
        void u2(double phi, double lambda, size_t idx, const ctrl_list& ctrl={});

        //! Apply a U3 gate with optional quibits of control.
        void u3(double theta, double phi, double lambda, size_t idx, const ctrl_list& ctrl={});

        //! Apply a RX gate with optional quibits of control.
        void rx(double theta, size_t idx, const ctrl_list& ctrl={});

        //! Apply a RY gate with optional quibits of control.
        void ry(double theta, size_t idx, const ctrl_list& ctrl={});

        //! Apply a RZ gate with optional quibits of control.
        void rz(double theta, size_t idx, const ctrl_list& ctrl={});

        //! Measure a qubit.
        int measure(size_t idx);

        void measure_zero(size_t idx);

        //! Apply a SWAP gate.
        void swap(size_t a, size_t b);

        //! Get the quantum state hashmap.
        map& get_map();

        //! Dump the quantum state of the first `size` qubits.
        dump_t dump(size_t size) const;

    private:
        map qbits; //<! Quantum state hashmap.

        friend std::ostream& operator<<(std::ostream &os, const ket::Bitwise& q);
    };

    //! Print the quantum state hashmap.
    std::ostream& operator<<(std::ostream &os, const ket::Bitwise& q);

    /*! \brief Abstract class for plugin interface.
     *  \class bitwise_api
     */
    class BOOST_SYMBOL_VISIBLE bitwise_api {
    public:

        /*! \brief Execute plugin.
         *
         *  The plugin must modify the qbits param.  We suggest to populate a
         *  new map and perfor a swap at the end.
         *
         *  The input qubits are ordend from size-1 to 0 in the hashmap.
         *  Operate with qubits >= size is undefined behavior.
         *
         *  The control qubits are ordend from size+ctrl-1 to size in the hashmap.
         *  Do not modify the control qubits.
         *
         *  \param qbits Quantum state hashmap
         *  \param size  Number of qubits affected.
         *  \param args  Argument string.
         *  \param adj   True when the plugin must execute its inverse.
         *  \param ctrl  Number of control qubits.
         */
        virtual void run(map& qbits, size_t size, std::string args, bool adj, size_t ctrl) const = 0 ;
    };

}

/*! \def bitwise_plugin(plugin_class)
 *  \brief Expose the class to be load as a plugin.
 */
#define bitwise_plugin(plugin_class) extern "C" BOOST_SYMBOL_EXPORT plugin_class plugin; plugin_class plugin


//  (C) Copyright 2014 Alvaro J. Genial (http://alva.ro)
//  Use, modification and distribution are subject to the Boost Software License, Version 1.0.
//  (See accompanying file LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt).

#ifndef AJG_SYNTH_DETAIL_IS_INTEGER_HPP_INCLUDED
#define AJG_SYNTH_DETAIL_IS_INTEGER_HPP_INCLUDED

#include <cmath>

#include <boost/utility/enable_if.hpp>
#include <boost/type_traits/is_integral.hpp>

namespace ajg {
namespace synth {
namespace detail {

//
// has_fraction:
//     Determines whether a floating-point number has a fractional part ("decimals").
////////////////////////////////////////////////////////////////////////////////////////////////////

template <class T>
inline typename boost::enable_if<boost::is_integral<T>, bool>::type has_fraction(T const) { return false; }

template <class T>
inline typename boost::disable_if<boost::is_integral<T>, bool>::type has_fraction(T const t) {
    T integral_part;
    return (std::modf)(t, &integral_part) != static_cast<T>(0.0);
}

}}} // namespace ajg::synth::detail

#endif // AJG_SYNTH_DETAIL_IS_INTEGER_HPP_INCLUDED


// Copyright Nick Thompson, 2017
// Use, modification and distribution are subject to the
// Boost Software License, Version 1.0.
// (See accompanying file LICENSE_1_0.txt
// or copy at http://www.boost.org/LICENSE_1_0.txt)

/*
 * This class performs sinh-sinh quadrature over the entire real line.
 *
 * References:
 *
 * 1) Tanaka, Ken'ichiro, et al. "Function classes for double exponential integration formulas." Numerische Mathematik 111.4 (2009): 631-655.
 */

#ifndef BOOST_MATH_QUADRATURE_SINH_SINH_HPP
#define BOOST_MATH_QUADRATURE_SINH_SINH_HPP

#include <cmath>
#include <limits>
#include <memory>
#include <boost/math/quadrature/detail/sinh_sinh_detail.hpp>

namespace boost{ namespace math{ namespace quadrature {

template<class Real, class Policy = boost::math::policies::policy<> >
class sinh_sinh
{
public:
    sinh_sinh(size_t max_refinements = 9)
        : m_imp(std::make_shared<detail::sinh_sinh_detail<Real, Policy> >(max_refinements)) {}

    template<class F>
    auto integrate(const F f, Real tol = boost::math::tools::root_epsilon<Real>(), Real* error = nullptr, Real* L1 = nullptr, std::size_t* levels = nullptr)->decltype(std::declval<F>()(std::declval<Real>())) const
    {
        return m_imp->integrate(f, tol, error, L1, levels);
    }

private:
    std::shared_ptr<detail::sinh_sinh_detail<Real, Policy>> m_imp;
};

}}}
#endif


#define BOOST_TEST_MODULE TEST_NODES

#include <dynet/functors.h>
#include <dynet/dynet.h>
#include <dynet/expr.h>
#include <dynet/grad-check.h>
#include <dynet/index-tensor.h>
#include <boost/test/unit_test.hpp>
#include "test.h"
#include <stdexcept>

using namespace dynet;
using namespace std;

struct TensorTest {
  TensorTest() {
    // initialize if necessary
    if (default_device == nullptr) {
      for (auto x : {"TensorTest", "--dynet-mem", "10"}) {
        av.push_back(strdup(x));
      }
      ADD_EXTRA_ARGUMENTS(av)
      char **argv = &av[0];
      int argc = av.size();
      dynet::initialize(argc, argv);
    }

    batch_vals = {1.f, 2.f, 1.f, 4.f, 5.f, 6.f};
  }
  ~TensorTest() {
    for (auto x : av) free(x);
  }

  template <class T>
  std::string print_vec(const std::vector<T> vec) {
    ostringstream oss;
    if (vec.size()) oss << vec[0];
    for (size_t i = 1; i < vec.size(); i++)
      oss << ' ' << vec[i];
    return oss.str();
  }

  std::vector<char*> av;
  std::vector<float> batch_vals;
  dynet::ParameterCollection mod;
};

// define the test suite
BOOST_FIXTURE_TEST_SUITE(tensor_test, TensorTest);

// test whether argmax is working properly
BOOST_AUTO_TEST_CASE( argmax ) {
  dynet::ComputationGraph cg;
  Expression x1 = input(cg, Dim({ 3 }, 2), batch_vals);
  IndexTensor idx_tens = TensorTools::argmax(x1.value());
  vector<Eigen::DenseIndex> idx_act = as_vector(idx_tens);
  vector<Eigen::DenseIndex> idx_exp = {1, 2};
  BOOST_CHECK_EQUAL_COLLECTIONS(idx_exp.begin(), idx_exp.end(),
                                idx_act.begin(), idx_act.end());
}

// test whether topk is working properly
BOOST_AUTO_TEST_CASE(topk) {
  dynet::ComputationGraph cg;
  Expression x1 = input(cg, Dim({3}, 2), batch_vals);
  std::pair<Tensor, IndexTensor> rets = TensorTools::topk(x1.value());
  // idx
  vector<Eigen::DenseIndex> idx_act = as_vector(rets.second);
  vector<Eigen::DenseIndex> idx_exp = {1, 2};
  BOOST_CHECK_EQUAL_COLLECTIONS(idx_exp.begin(), idx_exp.end(),
    idx_act.begin(), idx_act.end());
  // val
  vector<float> val_act = as_vector(rets.first);
  vector<float> val_exp = {2.f, 6.f};
  BOOST_CHECK_EQUAL_COLLECTIONS(val_act.begin(), val_act.end(),
    val_exp.begin(), val_exp.end());
}

// for now, just make sure that things don't die
BOOST_AUTO_TEST_CASE( categorical_sample_log_prob ) {
  dynet::ComputationGraph cg;
  Expression x1 = input(cg, Dim({ 3 }, 2), batch_vals);
  IndexTensor idx_tens = TensorTools::categorical_sample_log_prob(x1.value());
  vector<Eigen::DenseIndex> idx_act = as_vector(idx_tens);
  vector<Eigen::DenseIndex> idx_exp = {0, 2};
}

BOOST_AUTO_TEST_SUITE_END()


#include <iostream>
#include <utility>
#include <iterator>
#include <vector>
#include <algorithm>
#include <fstream>
#include <boost/graph/depth_first_search.hpp>
#include <boost/graph/adjacency_list.hpp>
#include <boost/graph/graph_traits.hpp>

using namespace std;

typedef boost::adjacency_list< boost::vecS, boost::vecS, boost::undirectedS> MyGraph;
typedef boost::graph_traits<MyGraph>::vertex_descriptor MyVertex;

//Used for start and end position
struct Position{
	int x_pos;
	int y_pos;
	int level;
};

//Each square in the maze, and it's properties.
struct Square{
	//1 = yellow, 0 = white
	char yellow_sq;
	//1 = wall, 0 = no wall
	char N;
	char S;
	char E;
	char W;
};

//Edge information, beg_vertex -> end_vertex
struct Edge{
	int beg_vertex;
	int end_vertex;
};

int main()
{
	ifstream inFile;
	inFile.open("input.txt", inFile.in);

	ofstream outFile;
	outFile.open("Bell.txt", outFile.out);

	int instances = 0, loop = 0;
	int levels = 0, width = 0, height = 0;
	int beg_level = 0, beg_x = 0, beg_y = 0;
	int end_level = 0, end_x = 0, end_y = 0;
	//char in_dir, out_dir;

	inFile >> instances;

	while (loop < instances){
		inFile >> levels >> width >> height;

		inFile >> beg_level >> beg_x >> beg_y;

		//Start position, x y z co-ords
		Position start;
		start.x_pos = beg_x;
		start.y_pos = beg_y;
		start.level = beg_level;

		inFile >> end_level >> end_x >> end_y;

		//Finish position, x y z co-ords
		Position finish;
		finish.x_pos = end_x;
		finish.y_pos = end_y;
		finish.level = end_level;

		int i = 0, j = 0;
		//Total height of the array, making 3D into 2D
		//Allowing double fors instead of triple
		//Increasing efficiency
		int total_height = height * levels;

		int exit = (finish.x_pos * finish.y_pos * finish.level);
		int entry = (start.x_pos * start.y_pos * start.level);

		//2D array, used to read into our graph
		Square** n = new Square*[total_height+1];
		for (i = 0; i < total_height+1; i++){
			n[i] = new Square[width+1];
		}

		// Read data into 2D Array of squares.
		// This will be used to look at the wall/ floor config. for each index.
		for (i = 1; i <= total_height; i++){
			for (j = 1; j <= width; j++){
				string(ch);
				inFile >> ch;
				n[i][j].yellow_sq = ch[0];
				n[i][j].N = ch[1];
				n[i][j].S = ch[2];
				n[i][j].E = ch[3];
				n[i][j].W = ch[4];
			}
		}

		//undirected graph to be passed to boost library
		MyGraph g(total_height*width+2);

		//These make it so the proper output is printed from the graph
		//example, 80 -> 81 or 20 -> 81
		boost::add_edge(entry, (total_height*width) + 1, g);

		//example, 1 -> 0, 21 -> 0
		boost::add_edge(finish.x_pos + ((finish.y_pos-1)*total_height), 0, g);

		//Iterates through Edge array
		for (i = 1; i <= total_height; i++){
			for (j = 1; j <= width; j++){
				//floor = 1, no floor = 0
				if (n[i][j].yellow_sq == '0'){
					boost::add_edge((i + ((j - 1)*total_height)), (i + ((j - 1)*total_height)) + height, g);
				}
				//wall = 1, no wall = 0
				if (n[i][j].N == '0'){
					if (i != 1){
						boost::add_edge(i + ((j - 1)*total_height), (i + ((j - 1)*total_height) - 1), g);
					}
				}
				//wall = 1, no wall = 0
				if (n[i][j].S == '0'){
					if (i != total_height){
						boost::add_edge(i + ((j - 1)*total_height), (i + ((j - 1)*total_height)) + 1, g);
					}
				}
				//wall = 1, no wall = 0
				if (n[i][j].W == '0'){
					if ((i + (j - 1)*total_height) >= total_height && j != 1){
						boost::add_edge(i + ((j - 1)*total_height), (i + ((j - 1)*total_height)) - total_height, g);
					}
				}

				//wall = 1, no wall = 0
				if (n[i][j].E == '0'){
					if (i < i + ((width - 1)*total_height) && j != width){
						boost::add_edge(i + ((j - 1)*total_height), (i + ((j - 1)*total_height)) + total_height, g);
					}
				}
			}
		}

		//Used for depth_first_search
		std::vector<MyVertex> predecessor;
		//Used to process output path, both numerical and directional
		std::vector<int> path;

		predecessor.resize(total_height*width+2);

		//Boost depth first search algorithm
		boost::depth_first_search(g, boost::visitor(boost::make_dfs_visitor(boost::record_predecessors(&predecessor[0], boost::on_tree_edge()))));

		int begin = (total_height*width) + 1;
		path.resize(predecessor.size());
		int next = 0;

		//Properly put the path into the path array
		while (begin != 0){
			path[next] = begin;
			begin = predecessor.at(begin);
			next++;
		}

		//Ignore edge going into graph
		i = 1;

		//Put number co-ordinates into directional co-ordinates.
		//Will stop once spider has reached location, for example
		//(1,1,1) will stop at (1,1,1)

		while (i != next){
			//Location reached, ignore edge going out
			if (path[i+1] == 0){
				break;
			}
			//West
			if (path[i] - total_height == path[i + 1]){
				outFile << "W ";
				i++;
			}
			//East
			else if (path[i] + total_height == path[i + 1]){
				outFile << "E ";
				i++;
			}
			//South
			else if (path[i] + 1 == path[i + 1]){
				outFile << "S ";
				i++;
			}
			//North
			else if (path[i] - 1 == path[i + 1]){
				outFile << "N ";
				i++;
			}
			//Down
			else if (path[i] + height == path[i + 1]){
				outFile << "D ";
				i++;
			}
			//Up
			else if (path[i] - height == path[i + 1]){
				outFile << "U ";
				i++;
			}
			else
				i++;
		}

		outFile << endl;

		//Free memory
		for (i = 0; i < total_height+1; i++){
			delete[] n[i];
		}
		delete[] n;

		loop++;
	}

	return 0;

}


// __BEGIN_LICENSE__
//  Copyright (c) 2006-2013, United States Government as represented by the
//  Administrator of the National Aeronautics and Space Administration. All
//  rights reserved.
//
//  The NASA Vision Workbench is licensed under the Apache License,
//  Version 2.0 (the "License"); you may not use this file except in
//  compliance with the License. You may obtain a copy of the License at
//  http://www.apache.org/licenses/LICENSE-2.0
//
//  Unless required by applicable law or agreed to in writing, software
//  distributed under the License is distributed on an "AS IS" BASIS,
//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//  See the License for the specific language governing permissions and
//  limitations under the License.
// __END_LICENSE__

#ifdef _MSC_VER
#pragma warning(disable:4244)
#pragma warning(disable:4267)
#pragma warning(disable:4996)
#endif

#ifdef NDEBUG
#undef NDEBUG
#endif

#include <vw/Core/System.h>
#include <vw/Core/Log.h>
#include <vw/Core/FundamentalTypes.h>
#include <vw/Math/EulerAngles.h>
#include <vw/Image/Algorithms.h>
#include <vw/Image/ImageViewRef.h>
#include <vw/Image/Filter.h>
#include <vw/Image/PixelMask.h>
#include <vw/Image/MaskViews.h>
#include <vw/Image/PerPixelAccessorViews.h>
#include <vw/FileIO/DiskImageView.h>
#include <vw/FileIO/DiskImageResource.h>
#include <vw/Cartography/GeoReference.h>
#include <vw/tools/Common.h>

#include <boost/program_options.hpp>
#include <boost/filesystem/path.hpp>

namespace po = boost::program_options;
namespace fs = boost::filesystem;

using namespace vw;

struct Options {
  Options() : nodata_value(std::numeric_limits<double>::quiet_NaN()), blur_sigma(std::numeric_limits<double>::quiet_NaN()) {}
  // Input
  std::string input_file_name;

  // Settings
  std::string output_file_name;
  double azimuth, elevation, scale;
  double nodata_value;
  double blur_sigma;
};

// ----------------------------------------------------------------------------

template <class PixelT>
void do_hillshade(Options& opt) {

  cartography::GeoReference georef;
  cartography::read_georeference(georef, opt.input_file_name);

  // Select the pixel scale.
  float u_scale, v_scale;
  if (opt.scale == 0) {
    if (georef.is_projected()) {
      u_scale = georef.transform()(0,0);
      v_scale = georef.transform()(1,1);
    } else {
      double meters_per_degree = 2*M_PI*georef.datum().semi_major_axis()/360.0;
      u_scale = georef.transform()(0,0) * meters_per_degree;
      v_scale = georef.transform()(1,1) * meters_per_degree;
    }
  } else {
    u_scale = opt.scale;
    v_scale = -opt.scale;
  }

  // Set the direction of the light source.
  Vector3f light_0(1,0,0);
  Vector3f light = vw::math::euler_to_rotation_matrix(opt.elevation*M_PI/180, opt.azimuth*M_PI/180, 0, "yzx") * light_0;

  // Compute the surface normals
  vw_out() << "Loading: " << opt.input_file_name << ".\n";
  DiskImageView<PixelT > disk_dem_file(opt.input_file_name);

  ImageViewRef<PixelMask<PixelT > > dem;
  boost::scoped_ptr<SrcImageResource> disk_dem_rsrc(DiskImageResource::open(opt.input_file_name));
  if ( !std::isnan(opt.nodata_value) ) {
    vw_out() << "\t--> Masking pixel value: " << opt.nodata_value << ".\n";
    dem = create_mask(disk_dem_file, opt.nodata_value);
  } else if ( disk_dem_rsrc->has_nodata_read() ) {
    opt.nodata_value = disk_dem_rsrc->nodata_read();
    vw_out() << "\t--> Extracted nodata value from file: "
             << opt.nodata_value << ".\n";
    dem = create_mask(disk_dem_file, opt.nodata_value);
  } else {
    dem = pixel_cast<PixelMask<PixelT > >(disk_dem_file);
  }

  if ( !std::isnan(opt.blur_sigma) ) {
    vw_out() << "\t--> Blurring pixel with gaussian kernal.  Sigma = "
             << opt.blur_sigma << "\n";
    dem = gaussian_filter(dem, opt.blur_sigma);
  }

  // The final result is the dot product of the light source with the normals
  ImageViewRef<PixelMask<PixelGray<uint8> > > shaded_image =
    channel_cast_rescale<uint8>(clamp(dot_prod(compute_normals(dem, u_scale, v_scale), light)));

  // Save the result
  vw_out() << "Writing shaded relief image: " << opt.output_file_name << "\n";

  boost::scoped_ptr<DiskImageResource> r(DiskImageResource::create(opt.output_file_name,
                                                                   shaded_image.format()));
  if ( r->has_block_write() )
    r->set_block_write_size( Vector2i( vw_settings().default_tile_size(),
                                       vw_settings().default_tile_size() ) );
  write_georeference( *r, georef );
  block_write_image( *r, shaded_image,
                     TerminalProgressCallback( "tools.hillshade", "Writing:") );
}

void handle_arguments( int argc, char *argv[], Options& opt ) {
  po::options_description desc("Description: Outputs image of a DEM lighted as specified\n\nUsage: hillshade [options] <input file> \n\nOptions");
  desc.add_options()
    ("input-file", po::value(&opt.input_file_name), "Explicitly specify the input file")
    ("output-file,o", po::value(&opt.output_file_name), "Specify the output file")
    ("azimuth,a", po::value(&opt.azimuth)->default_value(0), "Sets the direction tha the light source is coming from (in degrees).  Zero degrees is to the right, with positive degree counter-clockwise.")
    ("elevation,e", po::value(&opt.elevation)->default_value(45), "Set the elevation of the light source (in degrees).")
    ("scale,s", po::value(&opt.scale)->default_value(0), "Set the scale of a pixel (in the same units as the DTM height values.")
    ("nodata-value", po::value(&opt.nodata_value), "Remap the DEM default value to the min altitude value.")
    ("blur", po::value(&opt.blur_sigma), "Pre-blur the DEM with the specified sigma.")
    ("help,h", "Display this help message");

  po::positional_options_description p;
  p.add("input-file", 1);

  po::variables_map vm;
  try {
    po::store( po::command_line_parser( argc, argv ).options(desc).positional(p).run(), vm );
    po::notify( vm );
  } catch (const po::error& e) {
    vw_throw( ArgumentErr() << "Error parsing input:\n\t"
              << e.what() << desc );
  }

  if( vm.count("help") )
    vw_throw( ArgumentErr() << desc );

  if ( vm.count("input-file") != 1 )
    vw_throw( ArgumentErr() << "Error: Must specify exactly one input file!\n\n" << desc );

  if ( opt.output_file_name.empty() )
    opt.output_file_name =
      fs::path(opt.input_file_name).replace_extension().string() + "_HILLSHADE.tif";
}

int main( int argc, char *argv[] ) {

  Options opt;
  try {
    handle_arguments( argc, argv, opt );

    ImageFormat fmt = tools::image_format(opt.input_file_name);

    switch(fmt.pixel_format) {
    case VW_PIXEL_GRAY:
    case VW_PIXEL_GRAYA:
      switch (fmt.channel_type) {
      case VW_CHANNEL_UINT8:  do_hillshade<PixelGray<uint8>  >( opt ); break;
      case VW_CHANNEL_INT16:  do_hillshade<PixelGray<int16>  >( opt ); break;
      case VW_CHANNEL_UINT16: do_hillshade<PixelGray<uint16> >( opt ); break;
      default:                do_hillshade<PixelGray<float>  >( opt ); break;
      }
      break;
    default:
      vw_throw( ArgumentErr() << "Unsupported pixel format. The DEM image must have only one channel." );
    }
  } catch ( const ArgumentErr& e ) {
    vw_out() << e.what() << std::endl;
    return 1;
  } catch ( const Exception& e ) {
    std::cerr << "Error: " << e.what() << std::endl;
    return 1;
  }

  return 0;
}


//
// Modified by jungwon on 19. 1. 21.
// Add side constraint on 19. 5. 14.
// grid free on 19. 8. 29.
//

#ifndef SWARM_PLANNER_ENVIRONMENT_H
#define SWARM_PLANNER_ENVIRONMENT_H

#include <ecbs.hpp>
#include <boost/functional/hash.hpp>
#include <boost/program_options.hpp>
#include <math.h>

using namespace libMultiRobotPlanning;

namespace libMultiRobotPlanning {
    struct State {
        State(int time, int x, int y, int zz) : time(time), x(x), y(y), z(zz) {
            if (zz==4)
                z=0;
            if (zz==-1)
                z=3;
        }

        bool operator==(const State &s) const {
            return time == s.time && x == s.x && y == s.y && z == s.z;
        }

        State operator-(const State &s) const {
            return {time, x - s.x, y - s.y, z - s.z};
        }

        State operator+(const State &s) const {
            return {time, x + s.x, y + s.y, z + s.z};
        }

        bool equalExceptTime(const State &s) const { return x == s.x && y == s.y; }

        friend std::ostream &operator<<(std::ostream &os, const State &s) {
            return os << s.time << ": (" << s.x << "," << s.y << "," << s.z << ")";
            // return os << "(" << s.x << "," << s.y << "," << s.z <<")";
        }

        void set(int timer, int xr, int yr, int zr){
            x=xr;
            y=yr;
            time=timer;
            z=zr;
        }

        int time;
        int x;
        int y;
        int z;
    };

    struct Vector {
        Vector(double x, double y, double z) : x(x), y(y) {}
        Vector (const State& s): x((double)s.x), y((double)s.y) {}

        bool operator==(const Vector &v) const {
            return x == v.x && y == v.y;
        }

        Vector operator-(const Vector &v) const {
            return {x - v.x, y - v.y, 0};
        }

        double norm(){
            return sqrt(pow(x, 2) + pow(y, 2));
        }

        void normalize(){
            double vector_norm = norm();
            x /= vector_norm;
            y /= vector_norm;
        }

        double dot(const Vector& v){
            return x*v.x+y*v.y;
        }

        double min_dist_to_origin(Vector b){
            Vector a = *this;
            double min_dist = norm();

            if(!(a == b)) {
                double dist = b.norm();
                if (min_dist > dist) {
                    min_dist = dist;
                }

                Vector n = b - a;
                n.normalize();
                double a_dot_n = a.dot(n);
                n.x *= a_dot_n;
                n.y *= a_dot_n;
                Vector c = a - n;
                dist = c.norm();

                if ((c - a).dot(c - b) < 0 && min_dist > dist) {
                    min_dist = dist;
                }
            }
            return min_dist;
        }

        double x;
        double y;
    };
}

namespace std {
    template <>
    struct hash<State> {
        size_t operator()(const State& s) const {
            size_t seed = 0;
            boost::hash_combine(seed, s.time);
            boost::hash_combine(seed, s.x);
            boost::hash_combine(seed, s.y);
            boost::hash_combine(seed, s.z);
            return seed;
        }
    };
}  // namespace std

namespace libMultiRobotPlanning {
    enum class Action {
        Forward,
        Left,
        Right,
        Wait,
        ForwardLeft,
        ForwardRight
    };

    std::ostream &operator<<(std::ostream &os, const Action &a) {
        switch (a) {
            case Action::Forward:
                os << "Forward";
                break;
            case Action::Left:
                os << "Left";
                break;
            case Action::Right:
                os << "Right";
                break;
            case Action::Wait:
                os << "Wait";
                break;
            case Action::ForwardLeft:
                os << "ForwardLeft";
                break;
            case Action::ForwardRight:
                os << "ForwardRight";
                break;
        }
        return os;
    }

    struct Conflict {
        enum Type {
            Vertex,
            Edge,
        };

        int time;
        size_t agent1;
        size_t agent2;
        Type type;

        int x1;
        int y1;
        int z1;
        int x1_2;
        int y1_2;
        int z1_2;
        int x2;
        int y2;
        int z2;
        int x2_2;
        int y2_2;
        int z2_2;

        friend std::ostream &operator<<(std::ostream &os, const Conflict &c) {
            switch (c.type) {
                case Vertex:
                    return os << c.time << ": (" << c.agent1 << "," << c.agent2
                              << "): Vertex(" << c.x1 << "," << c.y1 << "," << c.z1 << ")";
                case Edge:
                    return os << c.time << ": (" << c.agent1 << "," << c.agent2
                              << "): Edge(" << c.x1 << "," << c.y1 << "," << c.z1 << ","
                              << c.x1_2 << "," << c.y1_2 << "," << c.z1_2 << ","
                              << c.x2 << "," << c.y2 << "," << c.z2 << ","
                              << c.x2_2 << "," << c.y2_2 << "," << c.z2_2 << ")";
            }
            return os;
        }
    };

    struct VertexConstraint {
        VertexConstraint(int time, int x, int y) : time(time), x(x), y(y) {}

        int time;
        int x;
        int y;

        bool operator<(const VertexConstraint &other) const {
            return std::tie(time, x, y) < std::tie(other.time, other.x, other.y);
        }

        bool operator==(const VertexConstraint &other) const {
            return std::tie(time, x, y) == std::tie(other.time, other.x, other.y);
        }

        friend std::ostream &operator<<(std::ostream &os, const VertexConstraint &c) {
            return os << "VC(" << c.time << "," << c.x << "," << c.y << ")";
        }
    };
}

namespace std {
    template <>
    struct hash<VertexConstraint> {
        size_t operator()(const VertexConstraint& s) const {
            size_t seed = 0;
            boost::hash_combine(seed, s.time);
            boost::hash_combine(seed, s.x);
            boost::hash_combine(seed, s.y);
            return seed;
        }
    };
}  // namespace std

namespace libMultiRobotPlanning {
    struct EdgeConstraint {
        EdgeConstraint(int time, int x1, int y1, int x2, int y2)
                : time(time), x1(x1), y1(y1), x2(x2), y2(y2) {}

        int time;
        int x1;
        int y1;
        int x2;
        int y2;

        bool operator<(const EdgeConstraint &other) const {
            return std::tie(time, x1, y1, x2, y2) <
                   std::tie(other.time, other.x1, other.y1, other.x2, other.y2);
        }

        bool operator==(const EdgeConstraint &other) const {
            return std::tie(time, x1, y1, x2, y2) ==
                   std::tie(other.time, other.x1, other.y1, other.x2, other.y2);
        }

        friend std::ostream &operator<<(std::ostream &os, const EdgeConstraint &c) {
            return os << "EC(" << c.time << "," << c.x1 << "," << c.y1 << "," << ","
                      << c.x2 << "," << c.y2 << ")";
        }
    };
}

namespace std {
    template <>
    struct hash<EdgeConstraint> {
        size_t operator()(const EdgeConstraint& s) const {
            size_t seed = 0;
            boost::hash_combine(seed, s.time);
            boost::hash_combine(seed, s.x1);
            boost::hash_combine(seed, s.y1);
            boost::hash_combine(seed, s.x2);
            boost::hash_combine(seed, s.y2);
            return seed;
        }
    };
}  // namespace std

namespace libMultiRobotPlanning {
    struct Constraints {
        std::unordered_set<VertexConstraint> vertexConstraints;
        std::unordered_set<EdgeConstraint> edgeConstraints;

        void add(const Constraints &other) {
            vertexConstraints.insert(other.vertexConstraints.begin(),
                                     other.vertexConstraints.end());
            edgeConstraints.insert(other.edgeConstraints.begin(),
                                   other.edgeConstraints.end());
        }

        bool overlap(const Constraints &other) {
            std::vector<VertexConstraint> vertexIntersection;
            std::vector<EdgeConstraint> edgeIntersection;
            std::set_intersection(vertexConstraints.begin(), vertexConstraints.end(),
                                  other.vertexConstraints.begin(),
                                  other.vertexConstraints.end(),
                                  std::back_inserter(vertexIntersection));
            std::set_intersection(edgeConstraints.begin(), edgeConstraints.end(),
                                  other.edgeConstraints.begin(),
                                  other.edgeConstraints.end(),
                                  std::back_inserter(edgeIntersection));
            return !vertexIntersection.empty() || !edgeIntersection.empty();
        }

        friend std::ostream &operator<<(std::ostream &os, const Constraints &c) {
            for (const auto &vc : c.vertexConstraints) {
                os << vc << std::endl;
            }
            for (const auto &ec : c.edgeConstraints) {
                os << ec << std::endl;
            }
            return os;
        }
    };

    struct Location {
        Location(int x, int y, int z) : x(x), y(y), z(z) {}

        int x;
        int y;
        int z;

        bool operator<(const Location &other) const {
            return std::tie(x, y) < std::tie(other.x, other.y);
        }

        bool operator==(const Location &other) const {
            return std::tie(x, y) == std::tie(other.x, other.y);
        }

        friend std::ostream &operator<<(std::ostream &os, const Location &c) {
            return os << "(" << c.x << "," << c.y << ")";
        }
    };
}
namespace std {
    template <>
    struct hash<Location> {
        size_t operator()(const Location& s) const {
            size_t seed = 0;
            boost::hash_combine(seed, s.x);
            boost::hash_combine(seed, s.y);
            return seed;
        }
    };
}  // namespace std

///
namespace libMultiRobotPlanning {
    class Environment {
    public:
        Environment(size_t dimx, size_t dimy, size_t dimz,
                    std::unordered_set<Location> obstacles,
                    std::vector<Location> goals,
                    std::vector<double> quad_size,
                    double grid_size)
                : m_dimx(dimx),
                  m_dimy(dimy),
                  m_dimz(dimz),
                  m_obstacles(std::move(obstacles)),
                  m_goals(std::move(goals)),
                  m_agentIdx(0),
                  m_constraints(nullptr),
                  m_lastGoalConstraint(-1),
                  m_highLevelExpanded(0),
                  m_lowLevelExpanded(0),
                  m_quad_size(std::move(quad_size)),
                  m_grid_size(grid_size) {}

        Environment(const Environment &) = delete;

        Environment &operator=(const Environment &) = delete;

        //find last goal constraint!
        void setLowLevelContext(size_t agentIdx, const Constraints *constraints) {
            assert(constraints);
            m_agentIdx = agentIdx;
            m_constraints = constraints;
            m_lastGoalConstraint = -1;
            for (const auto &vc : constraints->vertexConstraints) {
                if (vc.x == m_goals[m_agentIdx].x && vc.y == m_goals[m_agentIdx].y) {
                    m_lastGoalConstraint = std::max(m_lastGoalConstraint, vc.time);
                }
            }
        }

        int admissibleHeuristic(const State &s) {
            return std::abs(s.x - m_goals[m_agentIdx].x) +
                   std::abs(s.y - m_goals[m_agentIdx].y);
        }

        // low-level, get numConflict(equal state) from given solution
        int focalStateHeuristic(
                const State &s, int /*gScore*/,
                const std::vector<PlanResult<State, Action, int> > &solution) {
            int numConflicts = 0;
            for (size_t i = 0; i < solution.size(); ++i) {
                if (i != m_agentIdx && !solution[i].states.empty()) {
                    State state2 = getState(i, solution, s.time);
                    if (isVertexConflict(m_agentIdx, i, s, state2)) {
                        ++numConflicts;
                    }
                }
            }
            return numConflicts;
        }

        // low-level, get numConflict(s1a <-> s1b) from given solution
        int focalTransitionHeuristic(
                const State &s1a, const State &s1b, int /*gScoreS1a*/, int /*gScoreS1b*/,
                const std::vector<PlanResult<State, Action, int> > &solution) {
            int numConflicts = 0;
            for (size_t i = 0; i < solution.size(); ++i) {
                if (i != m_agentIdx && !solution[i].states.empty()) {
                    State s2a = getState(i, solution, s1a.time);
                    State s2b = getState(i, solution, s1b.time);
                    if (isEdgeConflict(m_agentIdx, i, s1a, s1b, s2a, s2b)) {
                        ++numConflicts;
                    }
                }
            }
            return numConflicts;
        }

        // Count all conflicts
        int focalHeuristic(
                const std::vector<PlanResult<State, Action, int> > &solution) {
            int numConflicts = 0;

            int max_t = 0;
            for (const auto &sol : solution) {
                max_t = std::max<int>(max_t, sol.states.size() - 1);
            }

            for (int t = 0; t < max_t; ++t) {
                // check drive-drive vertex collisions
                for (size_t i = 0; i < solution.size(); ++i) {
                    State state1 = getState(i, solution, t);
                    for (size_t j = i + 1; j < solution.size(); ++j) {
                        State state2 = getState(j, solution, t);
                        if (isVertexConflict(i, j, state1, state2)) {
                            ++numConflicts;
                        }
                    }
                }
                // drive-drive edge (swap)
                for (size_t i = 0; i < solution.size(); ++i) {
                    State state1a = getState(i, solution, t);
                    State state1b = getState(i, solution, t + 1);
                    for (size_t j = i + 1; j < solution.size(); ++j) {
                        State state2a = getState(j, solution, t);
                        State state2b = getState(j, solution, t + 1);
                        //TODO: fix it
                        if (isEdgeConflict(i, j, state1a, state1b, state2a, state2b)) {
                            ++numConflicts;
                        }
                    }
                }
            }
            return numConflicts;
        }

        bool isSolution(const State &s) {
            return s.x == m_goals[m_agentIdx].x && s.y == m_goals[m_agentIdx].y &&
                   s.time > m_lastGoalConstraint;
        }

        void getNeighbors(const State &s,
                          std::vector<Neighbor<State, Action, int> > &neighbors) {
            // std::cout << "#VC " << constraints.vertexConstraints.size() << std::endl;
            // for(const auto& vc : constraints.vertexConstraints) {
            //   std::cout << "  " << vc.time << "," << vc.x << "," << vc.y << "," << vc.z <<
            //   std::endl;
            // }
            neighbors.clear();
            {
                State n(s.time + 1, s.x, s.y, s.z);
                if (stateValid(n) && transitionValid(s, n)) {
                    neighbors.emplace_back(
                            Neighbor<State, Action, int>(n, Action::Wait, 1));
                }
            }
            {
                State n(s.time + 1, s.x, s.y, s.z + 1);
                if (stateValid(n) && transitionValid(s, n)) {
                    neighbors.emplace_back(
                            Neighbor<State, Action, int>(n, Action::Left, 2));
                }
            }
            {
                State n(s.time + 1, s.x, s.y, s.z - 1);
                if (stateValid(n) && transitionValid(s, n)) {
                    neighbors.emplace_back(
                            Neighbor<State, Action, int>(n, Action::Right, 2));
                }
            }
            {
                State n(0,0,0,0);
                switch (s.z){
                    case 0:
                        n.set(s.time + 1, s.x + 1, s.y, s.z);
                        break;
                    case 1:
                        n.set(s.time + 1, s.x, s.y + 1, s.z);
                        break;
                    case 2:
                        n.set(s.time + 1, s.x - 1, s.y, s.z);
                        break;
                    case 3:
                        n.set(s.time + 1, s.x, s.y - 1, s.z);
                        break;
                }
                if (stateValid(n) && transitionValid(s, n)) {
                    neighbors.emplace_back(
                            Neighbor<State, Action, int>(n, Action::Forward, 1));
                }
            }
            {
                State n(0,0,0,0);
                switch (s.z){
                    case 0:
                        n.set(s.time + 1, s.x + 1, s.y + 1, s.z+1);
                        break;
                    case 1:
                        n.set(s.time + 1, s.x - 1, s.y + 1, s.z+1);
                        break;
                    case 2:
                        n.set(s.time + 1, s.x - 1, s.y - 1, s.z+1);
                        break;
                    case 3:
                        n.set(s.time + 1, s.x + 1, s.y - 1, s.z-3);
                        break;
                }
                if (stateValid(n) && transitionValid(s, n)) {
                    neighbors.emplace_back(
                            Neighbor<State, Action, int>(n, Action::ForwardLeft, 2));
                }
            }
            {
                State n(0,0,0,0);
                switch (s.z){
                    case 0:
                        n.set(s.time + 1, s.x + 1, s.y - 1, s.z+3);
                        break;
                    case 1:
                        n.set(s.time + 1, s.x + 1, s.y + 1, s.z-1);
                        break;
                    case 2:
                        n.set(s.time + 1, s.x - 1, s.y + 1, s.z-1);
                        break;
                    case 3:
                        n.set(s.time + 1, s.x - 1, s.y - 1, s.z-1);
                        break;
                }
                if (stateValid(n) && transitionValid(s, n)) {
                    neighbors.emplace_back(
                            Neighbor<State, Action, int>(n, Action::ForwardRight, 2));
                }
            }
            /*
            {
                State n(s.time + 1, s.x, s.y, s.z + 1);
                if (stateValid(n) && transitionValid(s, n)) {
                    neighbors.emplace_back(
                            Neighbor<State, Action, int>(n, Action::Top, 1));
                }
            }
            {
                State n(s.time + 1, s.x, s.y, s.z - 1);
                if (stateValid(n) && transitionValid(s, n)) {
                    neighbors.emplace_back(
                            Neighbor<State, Action, int>(n, Action::Bottom, 1));
                }
            }
            */
        }

        bool getFirstConflict(
                const std::vector<PlanResult<State, Action, int> > &solution,
                Conflict &result) {
            int max_t = 0;
            for (const auto &sol : solution) {
                max_t = std::max<int>(max_t, sol.states.size() - 1);
            }

            for (int t = 0; t < max_t; ++t) {
                // check drive-drive vertex collisions
                for (size_t i = 0; i < solution.size(); ++i) {
                    State state1 = getState(i, solution, t);
                    for (size_t j = i + 1; j < solution.size(); ++j) {
                        State state2 = getState(j, solution, t);
                        if (isVertexConflict(i, j, state1, state2)) {
                            result.time = t;
                            result.agent1 = i;
                            result.agent2 = j;
                            result.type = Conflict::Vertex;
                            result.x1 = state1.x;
                            result.y1 = state1.y;
                            result.z1 = state1.z;
                            result.x2 = state2.x;
                            result.y2 = state2.y;
                            result.z2 = state2.z;
                            // std::cout << "VC " << t << "," << state1.x << "," << state1.y << "," << state1.z <<
                            // std::endl;
                            return true;
                        }
                    }
                }
                // drive-drive edge (swap)
                for (size_t i = 0; i < solution.size(); ++i) {
                    State state1a = getState(i, solution, t);
                    State state1b = getState(i, solution, t + 1);
                    for (size_t j = i + 1; j < solution.size(); ++j) {
                        State state2a = getState(j, solution, t);
                        State state2b = getState(j, solution, t + 1);

                        if (isEdgeConflict(i, j, state1a, state1b, state2a, state2b)) {
                            result.time = t;
                            result.agent1 = i;
                            result.agent2 = j;
                            result.type = Conflict::Edge;
                            result.x1 = state1a.x;
                            result.y1 = state1a.y;
                            result.z1 = state1a.z;
                            result.x1_2 = state1b.x;
                            result.y1_2 = state1b.y;
                            result.z1_2 = state1b.z;
                            result.x2 = state2a.x;
                            result.y2 = state2a.y;
                            result.z2 = state2a.z;
                            result.x2_2 = state2b.x;
                            result.y2_2 = state2b.y;
                            result.z2_2 = state2b.z;
                            return true;
                        }
                    }
                }
            }

            return false;
        }

        void createConstraintsFromConflict(
                const Conflict &conflict, std::map<size_t, Constraints> &constraints) {
            if (conflict.type == Conflict::Vertex) {
                Constraints c1, c2;
                c1.vertexConstraints.emplace(
                        VertexConstraint(conflict.time, conflict.x1, conflict.y1));
                c2.vertexConstraints.emplace(
                        VertexConstraint(conflict.time, conflict.x2, conflict.y2));
                constraints[conflict.agent1] = c1;
                constraints[conflict.agent2] = c2;
            } else if (conflict.type == Conflict::Edge) {
                Constraints c1, c2;
                c1.edgeConstraints.emplace(EdgeConstraint(
                        conflict.time, conflict.x1, conflict.y1, conflict.x1_2, conflict.y1_2));
                c2.edgeConstraints.emplace(EdgeConstraint(
                        conflict.time, conflict.x2, conflict.y2, conflict.x2_2, conflict.y2_2));
                constraints[conflict.agent1] = c1;
                constraints[conflict.agent2] = c2;
            }
        }

        void onExpandHighLevelNode(int /*cost*/) { m_highLevelExpanded++; }

        void onExpandLowLevelNode(const State & /*s*/, int /*fScore*/,
                                  int /*gScore*/) {
            m_lowLevelExpanded++;
        }

        int highLevelExpanded() { return m_highLevelExpanded; }

        int lowLevelExpanded() const { return m_lowLevelExpanded; }

    private:
        State getState(size_t agentIdx,
                       const std::vector<PlanResult<State, Action, int> > &solution,
                       size_t t) {
            assert(agentIdx < solution.size());
            if (t < solution[agentIdx].states.size()) {
                return solution[agentIdx].states[t].first;
            }
            assert(!solution[agentIdx].states.empty());
            return solution[agentIdx].states.back().first;
        }

        bool stateValid(const State &s) {
            assert(m_constraints);
            const auto &con = m_constraints->vertexConstraints;
            return s.x >= 0 && s.x < m_dimx && s.y >= 0 && s.y < m_dimy &&
                   m_obstacles.find(Location(s.x * 5, s.y * 5, s.z)) == m_obstacles.end() &&
                   con.find(VertexConstraint(s.time, s.x, s.y)) == con.end();
        }

        bool transitionValid(const State &s1, const State &s2) {
            assert(m_constraints);
            const auto &con = m_constraints->edgeConstraints;
            return m_obstacles.find(Location(s1.x * 2 + s2.x * 3, s1.y * 2 + s2.y * 3, s1.z)) == m_obstacles.end() &&
                   m_obstacles.find(Location(s1.x * 3 + s2.x * 2, s1.y * 3 + s2.y * 2, s1.z)) == m_obstacles.end() &&
                   m_obstacles.find(Location(s1.x * 1 + s2.x * 4, s1.y * 1 + s2.y * 4, s1.z)) == m_obstacles.end() &&
                   m_obstacles.find(Location(s1.x * 4 + s2.x * 1, s1.y * 4 + s2.y * 1, s1.z)) == m_obstacles.end() &&
                   con.find(EdgeConstraint(s1.time, s1.x, s1.y, s2.x, s2.y)) == con.end();
        }

        bool isParallel(const State &state1a, const State &state1b, const State &state2a, const State &state2b) {
            return (state1b.x - state1a.x) == (state2b.x - state2a.x) &&
                   (state1b.y - state1a.y) == (state2b.y - state2a.y);
        }

        bool isVertexConflict(int i, int j, const State &state1, const State &state2) {

            if (m_quad_size[i] + m_quad_size[j] < m_grid_size) {
                return state1.equalExceptTime(state2);
            }
            else{
                Vector v(state2 - state1);
                return v.norm() * m_grid_size < m_quad_size[i] + m_quad_size[j];
            }

            return false;
        }

        bool isEdgeConflict(int i, int j, const State &state1a, const State &state1b,
                                          const State &state2a, const State &state2b){

            if (m_quad_size[i] + m_quad_size[j] < m_grid_size * 0.5) {
                return (state1a.equalExceptTime(state2b) && state1b.equalExceptTime(state2a))
                ||((state1a+state1b).equalExceptTime(state2a+state2b));
            }
            else if(m_quad_size[i] + m_quad_size[j] < m_grid_size) {
                return (state1a.equalExceptTime(state2b) || state2a.equalExceptTime(state1b)) &&
                       !isParallel(state1a, state1b, state2a, state2b);
            }
            else{
                Vector a(state2a - state1a);
                Vector b(state2b - state1b);
                double min_dist = a.min_dist_to_origin(b);
                return min_dist * m_grid_size <= m_quad_size[i] + m_quad_size[j];
            }

            return false;

        }

    private:
        int m_dimx;
        int m_dimy;
        int m_dimz;
        std::unordered_set<Location> m_obstacles;
        std::vector<Location> m_goals;
        size_t m_agentIdx;
        const Constraints *m_constraints;
        int m_lastGoalConstraint;
        int m_highLevelExpanded;
        int m_lowLevelExpanded;
        std::vector<double> m_quad_size;
        double m_grid_size;
    };
}
#endif //SWARM_PLANNER_ENVIRONMENT_H


/* ----------------------------------------------------------------------------

* GTSAM Copyright 2010-2020, Georgia Tech Research Corporation,
* Atlanta, Georgia 30332-0415
* All Rights Reserved
* Authors: Frank Dellaert, et al. (see THANKS for the full author list)

* See LICENSE for the license information

* -------------------------------------------------------------------------- */

/*
 * @file testDiscreteBayesTree.cpp
 * @date sept 15, 2012
 * @author Frank Dellaert
 */

#include <gtsam/base/Vector.h>
#include <gtsam/discrete/DiscreteBayesNet.h>
#include <gtsam/discrete/DiscreteBayesTree.h>
#include <gtsam/discrete/DiscreteFactorGraph.h>
#include <gtsam/inference/BayesNet.h>

#include <boost/assign/std/vector.hpp>
using namespace boost::assign;

#include <CppUnitLite/TestHarness.h>

#include <vector>

using namespace std;
using namespace gtsam;

static bool debug = false;

/* ************************************************************************* */

TEST_UNSAFE(DiscreteBayesTree, ThinTree) {
  const int nrNodes = 15;
  const size_t nrStates = 2;

  // define variables
  vector<DiscreteKey> key;
  for (int i = 0; i < nrNodes; i++) {
    DiscreteKey key_i(i, nrStates);
    key.push_back(key_i);
  }

  // create a thin-tree Bayesnet, a la Jean-Guillaume
  DiscreteBayesNet bayesNet;
  bayesNet.add(key[14] % "1/3");

  bayesNet.add(key[13] | key[14] = "1/3 3/1");
  bayesNet.add(key[12] | key[14] = "3/1 3/1");

  bayesNet.add((key[11] | key[13], key[14]) = "1/4 2/3 3/2 4/1");
  bayesNet.add((key[10] | key[13], key[14]) = "1/4 3/2 2/3 4/1");
  bayesNet.add((key[9] | key[12], key[14]) = "4/1 2/3 F 1/4");
  bayesNet.add((key[8] | key[12], key[14]) = "T 1/4 3/2 4/1");

  bayesNet.add((key[7] | key[11], key[13]) = "1/4 2/3 3/2 4/1");
  bayesNet.add((key[6] | key[11], key[13]) = "1/4 3/2 2/3 4/1");
  bayesNet.add((key[5] | key[10], key[13]) = "4/1 2/3 3/2 1/4");
  bayesNet.add((key[4] | key[10], key[13]) = "2/3 1/4 3/2 4/1");

  bayesNet.add((key[3] | key[9], key[12]) = "1/4 2/3 3/2 4/1");
  bayesNet.add((key[2] | key[9], key[12]) = "1/4 8/2 2/3 4/1");
  bayesNet.add((key[1] | key[8], key[12]) = "4/1 2/3 3/2 1/4");
  bayesNet.add((key[0] | key[8], key[12]) = "2/3 1/4 3/2 4/1");

  if (debug) {
    GTSAM_PRINT(bayesNet);
    bayesNet.saveGraph("/tmp/discreteBayesNet.dot");
  }

  // create a BayesTree out of a Bayes net
  auto bayesTree = DiscreteFactorGraph(bayesNet).eliminateMultifrontal();
  if (debug) {
    GTSAM_PRINT(*bayesTree);
    bayesTree->saveGraph("/tmp/discreteBayesTree.dot");
  }

  // Check frontals and parents
  for (size_t i : {13, 14, 9, 3, 2, 8, 1, 0, 10, 5, 4}) {
    auto clique_i = (*bayesTree)[i];
    EXPECT_LONGS_EQUAL(i, *(clique_i->conditional_->beginFrontals()));
  }

  auto R = bayesTree->roots().front();

  // Check whether BN and BT give the same answer on all configurations
  vector<DiscreteValues> allPosbValues = cartesianProduct(
      key[0] & key[1] & key[2] & key[3] & key[4] & key[5] & key[6] & key[7] &
      key[8] & key[9] & key[10] & key[11] & key[12] & key[13] & key[14]);
  for (size_t i = 0; i < allPosbValues.size(); ++i) {
    DiscreteValues x = allPosbValues[i];
    double expected = bayesNet.evaluate(x);
    double actual = bayesTree->evaluate(x);
    DOUBLES_EQUAL(expected, actual, 1e-9);
  }

  // Calculate all some marginals for DiscreteValues==all1
  Vector marginals = Vector::Zero(15);
  double joint_12_14 = 0, joint_9_12_14 = 0, joint_8_12_14 = 0, joint_8_12 = 0,
         joint82 = 0, joint12 = 0, joint24 = 0, joint45 = 0, joint46 = 0,
         joint_4_11 = 0, joint_11_13 = 0, joint_11_13_14 = 0,
         joint_11_12_13_14 = 0, joint_9_11_12_13 = 0, joint_8_11_12_13 = 0;
  for (size_t i = 0; i < allPosbValues.size(); ++i) {
    DiscreteValues x = allPosbValues[i];
    double px = bayesTree->evaluate(x);
    for (size_t i = 0; i < 15; i++)
      if (x[i]) marginals[i] += px;
    if (x[12] && x[14]) {
      joint_12_14 += px;
      if (x[9]) joint_9_12_14 += px;
      if (x[8]) joint_8_12_14 += px;
    }
    if (x[8] && x[12]) joint_8_12 += px;
    if (x[2]) {
      if (x[8]) joint82 += px;
      if (x[1]) joint12 += px;
    }
    if (x[4]) {
      if (x[2]) joint24 += px;
      if (x[5]) joint45 += px;
      if (x[6]) joint46 += px;
      if (x[11]) joint_4_11 += px;
    }
    if (x[11] && x[13]) {
      joint_11_13 += px;
      if (x[8] && x[12]) joint_8_11_12_13 += px;
      if (x[9] && x[12]) joint_9_11_12_13 += px;
      if (x[14]) {
        joint_11_13_14 += px;
        if (x[12]) {
          joint_11_12_13_14 += px;
        }
      }
    }
  }
  DiscreteValues all1 = allPosbValues.back();

  // check separator marginal P(S0)
  auto clique = (*bayesTree)[0];
  DiscreteFactorGraph separatorMarginal0 =
      clique->separatorMarginal(EliminateDiscrete);
  DOUBLES_EQUAL(joint_8_12, separatorMarginal0(all1), 1e-9);

  // check separator marginal P(S9), should be P(14)
  clique = (*bayesTree)[9];
  DiscreteFactorGraph separatorMarginal9 =
      clique->separatorMarginal(EliminateDiscrete);
  DOUBLES_EQUAL(marginals[14], separatorMarginal9(all1), 1e-9);

  // check separator marginal of root, should be empty
  clique = (*bayesTree)[11];
  DiscreteFactorGraph separatorMarginal11 =
      clique->separatorMarginal(EliminateDiscrete);
  LONGS_EQUAL(0, separatorMarginal11.size());

  // check shortcut P(S9||R) to root
  clique = (*bayesTree)[9];
  DiscreteBayesNet shortcut = clique->shortcut(R, EliminateDiscrete);
  LONGS_EQUAL(1, shortcut.size());
  DOUBLES_EQUAL(joint_11_13_14 / joint_11_13, shortcut.evaluate(all1), 1e-9);

  // check shortcut P(S8||R) to root
  clique = (*bayesTree)[8];
  shortcut = clique->shortcut(R, EliminateDiscrete);
  DOUBLES_EQUAL(joint_11_12_13_14 / joint_11_13, shortcut.evaluate(all1), 1e-9);

  // check shortcut P(S2||R) to root
  clique = (*bayesTree)[2];
  shortcut = clique->shortcut(R, EliminateDiscrete);
  DOUBLES_EQUAL(joint_9_11_12_13 / joint_11_13, shortcut.evaluate(all1), 1e-9);

  // check shortcut P(S0||R) to root
  clique = (*bayesTree)[0];
  shortcut = clique->shortcut(R, EliminateDiscrete);
  DOUBLES_EQUAL(joint_8_11_12_13 / joint_11_13, shortcut.evaluate(all1), 1e-9);

  // calculate all shortcuts to root
  DiscreteBayesTree::Nodes cliques = bayesTree->nodes();
  for (auto clique : cliques) {
    DiscreteBayesNet shortcut = clique.second->shortcut(R, EliminateDiscrete);
    if (debug) {
      clique.second->conditional_->printSignature();
      shortcut.print("shortcut:");
    }
  }

  // Check all marginals
  DiscreteFactor::shared_ptr marginalFactor;
  for (size_t i = 0; i < 15; i++) {
    marginalFactor = bayesTree->marginalFactor(i, EliminateDiscrete);
    double actual = (*marginalFactor)(all1);
    DOUBLES_EQUAL(marginals[i], actual, 1e-9);
  }

  DiscreteBayesNet::shared_ptr actualJoint;

  // Check joint P(8, 2)
  actualJoint = bayesTree->jointBayesNet(8, 2, EliminateDiscrete);
  DOUBLES_EQUAL(joint82, actualJoint->evaluate(all1), 1e-9);

  // Check joint P(1, 2)
  actualJoint = bayesTree->jointBayesNet(1, 2, EliminateDiscrete);
  DOUBLES_EQUAL(joint12, actualJoint->evaluate(all1), 1e-9);

  // Check joint P(2, 4)
  actualJoint = bayesTree->jointBayesNet(2, 4, EliminateDiscrete);
  DOUBLES_EQUAL(joint24, actualJoint->evaluate(all1), 1e-9);

  // Check joint P(4, 5)
  actualJoint = bayesTree->jointBayesNet(4, 5, EliminateDiscrete);
  DOUBLES_EQUAL(joint45, actualJoint->evaluate(all1), 1e-9);

  // Check joint P(4, 6)
  actualJoint = bayesTree->jointBayesNet(4, 6, EliminateDiscrete);
  DOUBLES_EQUAL(joint46, actualJoint->evaluate(all1), 1e-9);

  // Check joint P(4, 11)
  actualJoint = bayesTree->jointBayesNet(4, 11, EliminateDiscrete);
  DOUBLES_EQUAL(joint_4_11, actualJoint->evaluate(all1), 1e-9);
}

/* ************************************************************************* */
int main() {
  TestResult tr;
  return TestRegistry::runAllTests(tr);
}
/* ************************************************************************* */