// Copyright Nick Thompson, 2019
// 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_MATH_INTERPOLATORS_CARDINAL_QUINTIC_B_SPLINE_DETAIL_HPP
#define BOOST_MATH_INTERPOLATORS_CARDINAL_QUINTIC_B_SPLINE_DETAIL_HPP
#include <cmath>
#include <vector>
#include <utility>
#include <boost/math/special_functions/cardinal_b_spline.hpp>

namespace boost{ namespace math{ namespace interpolators{ namespace detail{

template <class Real>
class cardinal_quintic_b_spline_detail
{
public:
    // If you don't know the value of the derivative at the endpoints, leave them as nans and the routine will estimate them.
    // y[0] = y(a), y[n -1] = y(b), step_size = (b - a)/(n -1).

    cardinal_quintic_b_spline_detail(const Real* const y,
                                     size_t n,
                                     Real t0 /* initial time, left endpoint */,
                                     Real h  /*spacing, stepsize*/,
                                     std::pair<Real, Real> left_endpoint_derivatives,
                                     std::pair<Real, Real> right_endpoint_derivatives)
    {
        static_assert(!std::is_integral<Real>::value, "The quintic B-spline interpolator only works with floating point types.");
        if (h <= 0) {
            throw std::logic_error("Spacing must be > 0.");
        }
        m_inv_h = 1/h;
        m_t0 = t0;

        if (n < 8) {
            throw std::logic_error("The quintic B-spline interpolator requires at least 8 points.");
        }

        using std::isnan;
        // This interpolator has error of order h^6, so the derivatives should be estimated with the same error.
        // See: https://en.wikipedia.org/wiki/Finite_difference_coefficient
        if (isnan(left_endpoint_derivatives.first)) {
            Real tmp = -49*y[0]/20 + 6*y[1] - 15*y[2]/2 + 20*y[3]/3 - 15*y[4]/4 + 6*y[5]/5 - y[6]/6;
            left_endpoint_derivatives.first = tmp/h;
        }
        if (isnan(right_endpoint_derivatives.first)) {
            Real tmp = 49*y[n-1]/20 - 6*y[n-2] + 15*y[n-3]/2 - 20*y[n-4]/3 + 15*y[n-5]/4 - 6*y[n-6]/5 + y[n-7]/6;
            right_endpoint_derivatives.first = tmp/h;
        }
        if(isnan(left_endpoint_derivatives.second)) {
            Real tmp = 469*y[0]/90 - 223*y[1]/10 + 879*y[2]/20 - 949*y[3]/18 + 41*y[4] - 201*y[5]/10 + 1019*y[6]/180 - 7*y[7]/10;
            left_endpoint_derivatives.second = tmp/(h*h);
        }
        if (isnan(right_endpoint_derivatives.second)) {
            Real tmp = 469*y[n-1]/90 - 223*y[n-2]/10 + 879*y[n-3]/20 - 949*y[n-4]/18 + 41*y[n-5] - 201*y[n-6]/10 + 1019*y[n-7]/180 - 7*y[n-8]/10;
            right_endpoint_derivatives.second = tmp/(h*h);
        }

        // This is really challenging my mental limits on by-hand row reduction.
        // I debated bringing in a dependency on a sparse linear solver, but given that that would cause much agony for users I decided against it.

        std::vector<Real> rhs(n+4);
        rhs[0] = 20*y[0] - 12*h*left_endpoint_derivatives.first +  2*h*h*left_endpoint_derivatives.second;
        rhs[1] = 60*y[0] - 12*h*left_endpoint_derivatives.first;
        for (size_t i = 2; i < n + 2; ++i) {
            rhs[i] = 120*y[i-2];
        }
        rhs[n+2] = 60*y[n-1] + 12*h*right_endpoint_derivatives.first;
        rhs[n+3] = 20*y[n-1] + 12*h*right_endpoint_derivatives.first +  2*h*h*right_endpoint_derivatives.second;

        std::vector<Real> diagonal(n+4, 66);
        diagonal[0] = 1;
        diagonal[1] = 18;
        diagonal[n+2] = 18;
        diagonal[n+3] = 1;

        std::vector<Real> first_superdiagonal(n+4, 26);
        first_superdiagonal[0] = 10;
        first_superdiagonal[1] = 33;
        first_superdiagonal[n+2] = 1;
        // There is one less superdiagonal than diagonal; make sure that if we read it, it shows up as a bug:
        first_superdiagonal[n+3] = std::numeric_limits<Real>::quiet_NaN();

        std::vector<Real> second_superdiagonal(n+4, 1);
        second_superdiagonal[0] = 9;
        second_superdiagonal[1] = 8;
        second_superdiagonal[n+2] = std::numeric_limits<Real>::quiet_NaN();
        second_superdiagonal[n+3] = std::numeric_limits<Real>::quiet_NaN();

        std::vector<Real> first_subdiagonal(n+4, 26);
        first_subdiagonal[0] = std::numeric_limits<Real>::quiet_NaN();
        first_subdiagonal[1] = 1;
        first_subdiagonal[n+2] = 33;
        first_subdiagonal[n+3] = 10;

        std::vector<Real> second_subdiagonal(n+4, 1);
        second_subdiagonal[0] = std::numeric_limits<Real>::quiet_NaN();
        second_subdiagonal[1] = std::numeric_limits<Real>::quiet_NaN();
        second_subdiagonal[n+2] = 8;
        second_subdiagonal[n+3] = 9;

        for (size_t i = 0; i < n+2; ++i) {
            Real di = diagonal[i];
            diagonal[i] = 1;
            first_superdiagonal[i] /= di;
            second_superdiagonal[i] /= di;
            rhs[i] /= di;

            // Eliminate first subdiagonal:
            Real nfsub = -first_subdiagonal[i+1];
            // Superfluous:
            first_subdiagonal[i+1] /= nfsub;
            // Not superfluous:
            diagonal[i+1] /= nfsub;
            first_superdiagonal[i+1] /= nfsub;
            second_superdiagonal[i+1] /= nfsub;
            rhs[i+1] /= nfsub;

            diagonal[i+1] += first_superdiagonal[i];
            first_superdiagonal[i+1] += second_superdiagonal[i];
            rhs[i+1] += rhs[i];
            // Superfluous, but clarifying:
            first_subdiagonal[i+1] = 0;

            // Eliminate second subdiagonal:
            Real nssub = -second_subdiagonal[i+2];
            first_subdiagonal[i+2] /= nssub;
            diagonal[i+2] /= nssub;
            first_superdiagonal[i+2] /= nssub;
            second_superdiagonal[i+2] /= nssub;
            rhs[i+2] /= nssub;

            first_subdiagonal[i+2] += first_superdiagonal[i];
            diagonal[i+2] += second_superdiagonal[i];
            rhs[i+2] += rhs[i];
            // Superfluous, but clarifying:
            second_subdiagonal[i+2] = 0;
        }

        // Eliminate last subdiagonal:
        Real dnp2 = diagonal[n+2];
        diagonal[n+2] = 1;
        first_superdiagonal[n+2] /= dnp2;
        rhs[n+2] /= dnp2;
        Real nfsubnp3 = -first_subdiagonal[n+3];
        diagonal[n+3] /= nfsubnp3;
        rhs[n+3] /= nfsubnp3;

        diagonal[n+3] += first_superdiagonal[n+2];
        rhs[n+3] += rhs[n+2];

        m_alpha.resize(n + 4, std::numeric_limits<Real>::quiet_NaN());

        m_alpha[n+3] = rhs[n+3]/diagonal[n+3];
        m_alpha[n+2] = rhs[n+2] - first_superdiagonal[n+2]*m_alpha[n+3];
        for (int64_t i = int64_t(n+1); i >= 0; --i) {
            m_alpha[i] = rhs[i] - first_superdiagonal[i]*m_alpha[i+1] - second_superdiagonal[i]*m_alpha[i+2];
        }

    }

    Real operator()(Real t) const {
        using std::ceil;
        using std::floor;
        using boost::math::cardinal_b_spline;
        // tf = t0 + (n-1)*h
        // alpha.size() = n+4
        if (t < m_t0 || t > m_t0 + (m_alpha.size()-5)/m_inv_h) {
            const char* err_msg = "Tried to evaluate the cardinal quintic b-spline outside the domain of of interpolation; extrapolation does not work.";
            throw std::domain_error(err_msg);
        }
        Real x = (t-m_t0)*m_inv_h;
        // Support of B_5 is [-3, 3]. So -3 < x - j + 2 < 3, so x-1 < j < x+5.
        // TODO: Zero pad m_alpha so that only the domain check is necessary.
        int64_t j_min = std::max(int64_t(0), int64_t(ceil(x-1)));
        int64_t j_max = std::min(int64_t(m_alpha.size() - 1), int64_t(floor(x+5)) );
        Real s = 0;
        for (int64_t j = j_min; j <= j_max; ++j) {
            // TODO: Use Cox 1972 to generate all integer translates of B5 simultaneously.
            s += m_alpha[j]*cardinal_b_spline<5, Real>(x - j + 2);
        }
        return s;
    }

    Real prime(Real t) const {
        using std::ceil;
        using std::floor;
        using boost::math::cardinal_b_spline_prime;
        if (t < m_t0 || t > m_t0 + (m_alpha.size()-5)/m_inv_h) {
            const char* err_msg = "Tried to evaluate the cardinal quintic b-spline outside the domain of of interpolation; extrapolation does not work.";
            throw std::domain_error(err_msg);
        }
        Real x = (t-m_t0)*m_inv_h;
        // Support of B_5 is [-3, 3]. So -3 < x - j + 2 < 3, so x-1 < j < x+5
        int64_t j_min = std::max(int64_t(0), int64_t(ceil(x-1)));
        int64_t j_max = std::min(int64_t(m_alpha.size() - 1), int64_t(floor(x+5)) );
        Real s = 0;
        for (int64_t j = j_min; j <= j_max; ++j) {
            s += m_alpha[j]*cardinal_b_spline_prime<5, Real>(x - j + 2);
        }
        return s*m_inv_h;

    }

    Real double_prime(Real t) const {
        using std::ceil;
        using std::floor;
        using boost::math::cardinal_b_spline_double_prime;
        if (t < m_t0 || t > m_t0 + (m_alpha.size()-5)/m_inv_h) {
            const char* err_msg = "Tried to evaluate the cardinal quintic b-spline outside the domain of of interpolation; extrapolation does not work.";
            throw std::domain_error(err_msg);
        }
        Real x = (t-m_t0)*m_inv_h;
        // Support of B_5 is [-3, 3]. So -3 < x - j + 2 < 3, so x-1 < j < x+5
        int64_t j_min = std::max(int64_t(0), int64_t(ceil(x-1)));
        int64_t j_max = std::min(int64_t(m_alpha.size() - 1), int64_t(floor(x+5)) );
        Real s = 0;
        for (int64_t j = j_min; j <= j_max; ++j) {
            s += m_alpha[j]*cardinal_b_spline_double_prime<5, Real>(x - j + 2);
        }
        return s*m_inv_h*m_inv_h;
    }

    Real t_max() const {
        return m_t0 + (m_alpha.size()-5)/m_inv_h;
    }

private:
    std::vector<Real> m_alpha;
    Real m_inv_h;
    Real m_t0;
};

}}}}
#endif


/**************************************************************************
 * Copyright (c) 2017-2019 by the mfmg authors                            *
 * All rights reserved.                                                   *
 *                                                                        *
 * This file is part of the mfmg library. mfmg is distributed under a BSD *
 * 3-clause license. For the licensing terms see the LICENSE file in the  *
 * top-level directory                                                    *
 *                                                                        *
 * SPDX-License-Identifier: BSD-3-Clause                                  *
 *************************************************************************/

#include <mfmg/common/exceptions.hpp>

#include <deal.II/base/timer.h>

#include <boost/program_options.hpp>

#include "test_hierarchy_helpers.hpp"

template <int dim, int fe_degree>
void matrix_free_two_grids(std::shared_ptr<boost::property_tree::ptree> params)
{
  // In case the operator was configured to act as a solver, i.e.
  // "is preconditioner" is false, we perform a set number of V-cycles and
  // report the convergence rate. Otherwise, we solve the Laplace problem using
  // the Hierarchy object as a solver and report the number of iterations
  // required.
  bool const test_preconditioner = params->get<bool>("is preconditioner");

  using DVector = dealii::LinearAlgebra::distributed::Vector<double>;

  MPI_Comm comm = MPI_COMM_WORLD;

  dealii::ConditionalOStream pcout(
      std::cout, dealii::Utilities::MPI::this_mpi_process(comm) == 0);

  // Print a table with the timings when the destructor is called
  auto timer = std::make_shared<dealii::TimerOutput>(
      comm, pcout, dealii::TimerOutput::summary,
      dealii::TimerOutput::wall_times);

  auto material_property =
      MaterialPropertyFactory<dim>::create_material_property(
          params->get<std::string>("material_property.type"));
  Source<dim> source;

  auto const &laplace_ptree = params->get_child("laplace");
  LaplaceMatrixFree<dim, fe_degree, double> mf_laplace(comm);
  mf_laplace.setup_system(laplace_ptree, *material_property);

  auto const &locally_owned_dofs = mf_laplace._locally_owned_dofs;
  DVector solution(locally_owned_dofs, comm);
  DVector rhs(mf_laplace._system_rhs);

  std::default_random_engine generator;
  std::uniform_real_distribution<typename DVector::value_type> distribution(0.,
                                                                            1.);
  for (auto const index : locally_owned_dofs)
  {
    // Make the solution satisfy the Dirichlet conditions because these should
    // be treated outside the preconditioner
    if (mf_laplace._constraints.is_constrained(index))
      solution[index] = 0.;
    else
      solution[index] = distribution(generator);
  }

  auto evaluator =
      std::make_shared<TestMFMeshEvaluator<dim, fe_degree, double>>(
          mf_laplace._dof_handler, mf_laplace._constraints,
          mf_laplace._laplace_operator, material_property);

  mfmg::Hierarchy<DVector> hierarchy(comm, evaluator, params, timer);

  if (!test_preconditioner)
  {
    // We want to do 20 V-cycle iterations. The rhs of is zero.
    // Use D(istributed)Vector because deal has its own Vector class
    DVector residual(rhs);
    unsigned int const n_cycles = 20;
    std::vector<double> res(n_cycles + 1);

    mf_laplace._laplace_operator.vmult(residual, solution);
    residual.sadd(-1., 1., rhs);
    auto const residual0_norm = residual.l2_norm();

    std::cout << std::scientific;
    res[0] = 1.0;
    for (unsigned int i = 0; i < n_cycles; ++i)
    {
      hierarchy.vmult(solution, rhs);

      mf_laplace._laplace_operator.vmult(residual, solution);
      residual.sadd(-1., 1., rhs);
      double rel_residual = residual.l2_norm() / residual0_norm;
      res[i + 1] = rel_residual;
    }

    double const conv_rate = res[n_cycles] / res[n_cycles - 1];
    pcout << "Convergence rate: " << std::fixed << std::setprecision(2)
          << conv_rate << std::endl;
  }
  else
  {
    auto const solver_tolerance =
        params->get<double>("solver.tolerance", 1.e-6);
    dealii::SolverControl solver_control(solution.size(), solver_tolerance);
    dealii::SolverCG<DVector> solver(solver_control);

    // We frequently like to compare to not applying a preconditioner at all.
    solver.solve(mf_laplace._laplace_operator, solution, rhs,
                 //                 dealii::PreconditionIdentity()
                 hierarchy);
    pcout << "Converging after " << solver_control.last_step()
          << " iterations.\n";
  }
}

template <int dim>
void matrix_based_two_grids(std::shared_ptr<boost::property_tree::ptree> params)
{
  bool const test_preconditioner = params->get<bool>("is preconditioner");

  using DVector = dealii::LinearAlgebra::distributed::Vector<double>;
  using MeshEvaluator = mfmg::DealIIMeshEvaluator<dim>;

  MPI_Comm comm = MPI_COMM_WORLD;

  dealii::ConditionalOStream pcout(
      std::cout, dealii::Utilities::MPI::this_mpi_process(comm) == 0);

  // Print a table with the timings when the destructor is called
  auto timer = std::make_shared<dealii::TimerOutput>(
      comm, pcout, dealii::TimerOutput::summary,
      dealii::TimerOutput::wall_times);

  auto material_property =
      MaterialPropertyFactory<dim>::create_material_property(
          params->get<std::string>("material_property.type"));
  Source<dim> source;

  auto laplace_ptree = params->get_child("laplace");
  auto fe_degree = laplace_ptree.get<unsigned>("fe_degree", 4);
  Laplace<dim, DVector> laplace(comm, fe_degree);
  laplace.setup_system(laplace_ptree);
  laplace.assemble_system(source, *material_property);

  auto const &a = laplace._system_matrix;
  auto const locally_owned_dofs = laplace._locally_owned_dofs;
  DVector solution(locally_owned_dofs, comm);
  DVector rhs(laplace._system_rhs);

  std::default_random_engine generator;
  std::uniform_real_distribution<typename DVector::value_type> distribution(0.,
                                                                            1.);
  for (auto const index : locally_owned_dofs)
  {
    // Make the solution satisfy the Dirichlet conditions because these should
    // be treated outside the preconditioner
    if (laplace._constraints.is_constrained(index))
      solution[index] = 0.;
    else
      solution[index] = distribution(generator);
  }

  std::shared_ptr<MeshEvaluator> evaluator(
      new TestMeshEvaluator<mfmg::DealIIMeshEvaluator<dim>>(
          laplace._dof_handler, laplace._constraints, fe_degree, a,
          material_property));
  mfmg::Hierarchy<DVector> hierarchy(comm, evaluator, params, timer);

  if (!test_preconditioner)
  {
    // We want to do 20 V-cycle iterations. The rhs of is zero.
    // Use D(istributed)Vector because deal has its own Vector class
    DVector residual(rhs);
    unsigned int const n_cycles = 20;
    std::vector<double> res(n_cycles + 1);

    a.vmult(residual, solution);
    residual.sadd(-1., 1., rhs);
    auto const residual0_norm = residual.l2_norm();

    std::cout << std::scientific;
    res[0] = 1.0;
    for (unsigned int i = 0; i < n_cycles; ++i)
    {
      hierarchy.vmult(solution, rhs);

      a.vmult(residual, solution);
      residual.sadd(-1., 1., rhs);
      double rel_residual = residual.l2_norm() / residual0_norm;
      res[i + 1] = rel_residual;
    }

    double const conv_rate = res[n_cycles] / res[n_cycles - 1];
    pcout << "Convergence rate: " << std::fixed << std::setprecision(2)
          << conv_rate << std::endl;
  }
  else
  {
    auto const solver_tolerance =
        params->get<double>("solver.tolerance", 1.e-6);
    dealii::SolverControl solver_control(solution.size(), solver_tolerance);
    dealii::SolverCG<DVector> solver(solver_control);

    // We frequently like to compare to not applying a preconditioner at all.
    solver.solve(a, solution, rhs,
                 //                 dealii::PreconditionIdentity()
                 hierarchy);
    pcout << "Converging after " << solver_control.last_step()
          << " iterations.\n";
  }
}

int main(int argc, char *argv[])
{
  namespace boost_po = boost::program_options;

  dealii::Utilities::MPI::MPI_InitFinalize mpi_init(argc, argv);

  boost_po::options_description cmd("Available options");
  cmd.add_options()("help,h", "produce help message");
  cmd.add_options()("filename,f", boost_po::value<std::string>()->multitoken(),
                    "file containing input parameters");
  cmd.add_options()("dim,d", boost_po::value<int>(), "dimension");
  cmd.add_options()("matrix_free,m", boost_po::value<bool>(),
                    "use matrix-free algorithm");
  cmd.add_options()("tolerance,t", boost_po::value<double>(),
                    "tolerance to use for the solver");

  boost_po::variables_map vm;
  boost_po::store(boost_po::parse_command_line(argc, argv, cmd), vm);
  boost_po::notify(vm);

  if (vm.count("help"))
  {
    std::cout << cmd << std::endl;

    return 0;
  }

  std::string filename = {"hierarchy_input.info"};
  if (vm.count("filename"))
    filename = vm["filename"].as<std::string>();

  int dim = 2;
  if (vm.count("dim"))
    dim = vm["dim"].as<int>();
  mfmg::ASSERT(dim == 2 || dim == 3, "Dimension must be 2 or 3");

  bool matrix_free = true;
  if (vm.count("matrix_free"))
    matrix_free = vm["matrix_free"].as<bool>();

  auto const params = std::make_shared<boost::property_tree::ptree>();
  boost::property_tree::info_parser::read_info(filename, *params);

  // "Uncover" params for running ML if instructed to do so
  if (!matrix_free && params->get("use_raw_ml", false))
  {
    auto const &ml_options = params->get_child("hidden");
    for (auto const &opt : ml_options)
    {
      params->put_child(opt.first, opt.second);
    }
  }

  int const fe_degree = params->get<unsigned int>("laplace.fe_degree", 4);
  params->put("fast_ap", true);
  params->put("eigensolver.type", "anasazi");

  double solver_tolerance = 1.e-6;
  if (vm.count("tolerance"))
  {
    solver_tolerance = vm["tolerance"].as<double>();
  }
  params->put("solver.tolerance", solver_tolerance);

  std::cout << "input file: " << filename << ", dimension: " << dim
            << ", matrix-free: " << matrix_free << ", fe_degree: " << fe_degree
            << ", solver_tolerance: " << solver_tolerance << std::endl;

  if (matrix_free)
  {
    params->put("smoother.type", "Chebyshev");
    if (dim == 2)
    {
      switch (fe_degree)
      {
      case 1:
      {
        matrix_free_two_grids<2, 1>(params);

        break;
      }
      case 2:
      {
        matrix_free_two_grids<2, 2>(params);

        break;
      }
      case 3:
      {
        matrix_free_two_grids<2, 3>(params);

        break;
      }
      case 4:
      {
        matrix_free_two_grids<2, 4>(params);

        break;
      }
      case 5:
      {
        matrix_free_two_grids<2, 5>(params);

        break;
      }
      case 6:
      {
        matrix_free_two_grids<2, 6>(params);

        break;
      }
      case 7:
      {
        matrix_free_two_grids<2, 7>(params);

        break;
      }
      case 8:
      {
        matrix_free_two_grids<2, 8>(params);

        break;
      }
      case 9:
      {
        matrix_free_two_grids<2, 9>(params);

        break;
      }
      case 10:
      {
        matrix_free_two_grids<2, 10>(params);

        break;
      }
      default:
      {
        mfmg::ASSERT(false, "The fe_degree should be between 1 and 10");
      }
      }
    }
    else
    {
      switch (fe_degree)
      {
      case 1:
      {
        matrix_free_two_grids<3, 1>(params);

        break;
      }
      case 2:
      {
        matrix_free_two_grids<3, 2>(params);

        break;
      }
      case 3:
      {
        matrix_free_two_grids<3, 3>(params);

        break;
      }
      case 4:
      {
        matrix_free_two_grids<3, 4>(params);

        break;
      }
      case 5:
      {
        matrix_free_two_grids<3, 5>(params);

        break;
      }
      case 6:
      {
        matrix_free_two_grids<3, 6>(params);

        break;
      }
      case 7:
      {
        matrix_free_two_grids<3, 7>(params);

        break;
      }
      case 8:
      {
        matrix_free_two_grids<3, 8>(params);

        break;
      }
      case 9:
      {
        matrix_free_two_grids<3, 9>(params);

        break;
      }
      case 10:
      {
        matrix_free_two_grids<3, 10>(params);

        break;
      }
      default:
      {
        mfmg::ASSERT(false, "The fe_degree should be between 1 and 10");
      }
      }
    }
  }
  else
  {
    dealii::MultithreadInfo::set_thread_limit(1);
    if (dim == 2)
      matrix_based_two_grids<2>(params);
    else
      matrix_based_two_grids<3>(params);
  }

  return 0;
}


#include <climits>
#include <cmath>
#include <ctime>
#include <fstream>
#include <iostream>
#include <iomanip>
#include <sstream>
#include <unistd.h>
#include <algorithm>
#include <vector>
#include <string>

#include "../include/domains/boxpushing/BoxAction.h"
#include "../include/domains/boxpushing/BoxState.h"
#include "../include/domains/boxpushing/BoxProblem.h"
#include "../include/domains/navigation/NavProblem.h"
#include "../include/domains/navigation/NavAction.h"
#include "../include/domains/navigation/NavState.h"

#include "../include/solvers/Solver.h"
#include "../include/solvers/Exploration.h"
#include "../include/solvers/MOLAOStarSolver.h"
#include "../include/solvers/LAOStarSolver.h"
#include "../include/util/general.h"
#include "../include/util/flags.h"
#include "../include/util/graph.h"

#include <boost/shared_ptr.hpp>
#include <boost/algorithm/string.hpp>
#include <boost/accumulators/accumulators.hpp>
#include <boost/accumulators/statistics/stats.hpp>
#include <boost/accumulators/statistics/mean.hpp>
#include <boost/accumulators/statistics/variance.hpp>
using namespace boost::accumulators;

using namespace mlsolvers;
using namespace mdplib_mobj_solvers;
using namespace mdplib;
using namespace mlcore;
using namespace mlmobj;
using namespace std;

int verbosity = 0;
MOProblem* problem = nullptr;
std::string testing_file = "";
std::string training_file = "";
std::string RL_training_file = "";
std::string processed_file = "";
std::string policy_file = "";
bool useLearnedValues = false;
int max_episodes=200;
Solver* solver = nullptr;
string problem_name  = "";
///////////////////////////////////////////////////////////////////////////////
//                              PROBLEM SETUP                                //
///////////////////////////////////////////////////////////////////////////////
void setupBoxPushing(){
    string boxFile = flag_value("box");
    testing_file = flag_value("box").substr(0, flag_value("box").find(".bp")) + "_Testing.txt";
    training_file = flag_value("box").substr(0, flag_value("box").find(".bp")) + "_Samples.txt";
    RL_training_file = flag_value("box").substr(0, flag_value("box").find(".bp")) + "_RLSamples.txt";
    processed_file = flag_value("box").substr(0, flag_value("box").find(".bp")) + "_Processed.csv";
    policy_file = flag_value("box").substr(0, flag_value("box").find(".bp")) + "_Policy.txt";

    int numObj = stoi(flag_value("numObj"));
    std::vector<double> slack(numObj, 0);

    if(flag_is_registered_with_value("slack")){
        std::string line = flag_value("slack");
        std::vector<std::string> vec;
        boost::algorithm::split(vec, line, boost::is_any_of(","));
        int slack_length = vec.size();
        int limit = std::min(numObj, slack_length);
        for(int i = 0; i < limit; i++){
            slack.at(i) = stod(vec.at(i));
        }
    }
    double gamma = 0.99;
    if(flag_is_registered_with_value("gamma"))
        gamma = stod(flag_value("gamma"));

    string costfunc = "optimal";
    if(flag_is_registered_with_value("cost"))
        costfunc = flag_value("cost");

    if(flag_is_registered_with_value("testing")){
        useLearnedValues = true;
        processed_file = flag_value("testing");
    }
    problem = new BoxProblem(boxFile.c_str(), numObj, costfunc.c_str(), useLearnedValues, processed_file);
    problem->slack(slack);
    problem->gamma(gamma);
    problem->ProblemName("boxPushing");
    problem->StateDimensions(3); // used for RL algorithm function approximation
    if(useLearnedValues)
         problem->use_learnedReward(useLearnedValues);
    if(flag_is_registered("impactFree")){
        BoxProblem* bp = static_cast<BoxProblem*> (problem);
        bp->Set_impactFreeTrans(true);
    }
}
void setupNavigation(){
    string navFile = flag_value("nav");
    testing_file = flag_value("nav").substr(0, flag_value("nav").find(".nav")) + "_Testing.txt";
    training_file = flag_value("nav").substr(0, flag_value("nav").find(".nav")) + "_Samples.txt";
    RL_training_file = flag_value("nav").substr(0, flag_value("nav").find(".nav")) + "_RLSamples.txt";
    processed_file = flag_value("nav").substr(0, flag_value("nav").find(".nav")) + "_Processed.csv";
    policy_file = flag_value("nav").substr(0, flag_value("nav").find(".nav")) + "_Policy.txt";

    int numObj = stoi(flag_value("numObj"));
    std::vector<double> slack(numObj, 0);

    if(flag_is_registered_with_value("slack")){
        std::string line = flag_value("slack");
        std::vector<std::string> vec;
        boost::algorithm::split(vec, line, boost::is_any_of(","));
        int slack_length = vec.size();
        int limit = std::min(numObj, slack_length);
        for(int i = 0; i < limit; i++){
            slack.at(i) = stod(vec.at(i));
        }
    }
    double gamma = 0.99;
    if(flag_is_registered_with_value("gamma"))
        gamma = stod(flag_value("gamma"));

    string costfunc = "optimal";
    if(flag_is_registered_with_value("cost"))
        costfunc = flag_value("cost");

    if(flag_is_registered_with_value("testing")){
        useLearnedValues = true;
        processed_file = flag_value("testing");
    }
    problem = new NavProblem(navFile.c_str(), numObj, costfunc.c_str(), useLearnedValues, processed_file);
    problem->slack(slack);
    problem->gamma(gamma);
    problem->ProblemName("navigation");
    problem->StateDimensions(3); // used for RL algorithm function approximation
    if(useLearnedValues)
         problem->use_learnedReward(useLearnedValues);
    if(flag_is_registered("impactFree")){
        NavProblem* np = static_cast<NavProblem*> (problem);
        np->Set_impactFreeTrans(true);
    }
}

void setupProblem()
{
    if (verbosity > 100)
        cout << "Setting up problem" << endl;
    if (flag_is_registered_with_value("box")) {
        problem_name = "boxpushing";
        setupBoxPushing();
    } else if(flag_is_registered_with_value("nav")){
        problem_name = "navigation";
        setupNavigation();
    }

    else {
        cerr << "Invalid problem." << endl;
        exit(-1);
    }
}
///////////////////////////////////////////////////////////////////////////////
//                               Exploration                                     //
///////////////////////////////////////////////////////////////////////////////
void setupExploration(string alg, Exploration*& learner){
    double epsilon = 0.1;
    if(flag_is_registered_with_value("epsilon"))
        epsilon = stod(flag_value("epsilon"));
    learner = new Exploration(problem,epsilon,RL_training_file);
    if(useLearnedValues)
        learner->generateSamples(false);
}

///////////////////////////////////////////////////////////////////////////////
//                              SOLVER SETUP                                 //
///////////////////////////////////////////////////////////////////////////////

void setupSolver(string alg, Solver*& solver){
    if (alg == "LLAO"){
        solver = new MOLAOStarSolver(problem);
    }

    else if(alg == "deviation"){
        problem->deviationScore("True");
        std::vector<double> weights;
        if(flag_is_registered_with_value("weights")){
            std::string line = flag_value("weights");
            std::vector<std::string> vec;
            boost::algorithm::split(vec, line, boost::is_any_of(","));
            int wt_length = vec.size();
            for(int i = 0; i < wt_length; i++)
                weights.push_back(stod(vec.at(i)));
        }
        else{
            weights.push_back(1); // weights for action reward wrt. o_1
            // corresponds to beta value in "Penalizing side effects using stepwise relative reachability"
            weights.push_back(0.1);
        }
        problem->weights(weights);
        solver = new MOLAOStarSolver(problem,false,1.0e-6,1000000,true);
    }
    else if(alg == "LAO"){
        double tol = 1.0e-3;
        solver = new LAOStarSolver(problem,tol, 100000000);
    }
    else{
        std::cout << " Algorithm not found" << std::endl;
        exit(1);
    }
}
/** Generates samples with true labels and creates a file for learning (with all s,a and without labels) **/
void generateSamples(int trials, Solver*& temp_solver){
    std::cout << " generating trajectories for querying " << std::endl;
    ofstream output_file;
    output_file.open(training_file);
    int trial_count = 0;
    mlcore::State* s = problem->initialState();
    int traj_length = 0;
    while(trial_count < trials){
        double prob = 0.0;
        mlcore::Action* a = s->bestAction();
        output_file << s << " " << (a)->hashValue() << std::endl;
        mlcore::State* succ = mlsolvers::randomSuccessor(problem,s,a, &prob);
        traj_length ++;
        s = succ;
        if(problem->goal(succ) || traj_length > 500){
            trial_count++;
            s = problem->initialState();
        }
    }
     output_file.close();
}
/* Generates all (s,a).. A part of this will be used for training, remaining will be used for testing*/
void generateSamples(){
    ofstream output_file;
     output_file.open(training_file);
     for(mlcore::State* s : problem->states()){
        mlmobj::MOState* mos = static_cast<mlmobj::MOState*>(s);
        std::list<mlcore::Action*> filteredActions = mos->filteredActions();
        for(mlcore::Action* a : filteredActions){
        // for(mlcore::Action* a : problem->actions()){
            if(problem->applicable(s,a)){
                 output_file << s << " " << a << " " << a->hashValue() << " " <<
                                problem->cost(s,a,1) << std::endl;
            }
        }
     }
     output_file.close();
}
void generateSamples(std::string filename){
    ofstream output_file;
     output_file.open(filename);
     for(mlcore::State* s : problem->states()){
        for(mlcore::Action* a : problem->actions()){
            if(problem->applicable(s,a)){
                 output_file << s << " " << a << " " << a->hashValue() << " " <<
                                problem->cost(s,a,1) << std::endl;
            }
        }
     }
     output_file.close();
}

void printPolicy(){
    ofstream output_file;
    output_file.open(policy_file);
    for(mlcore::State* s : problem->states()){
        if(s->bestAction() != nullptr){
            mlcore::Action* a = s->bestAction();
            output_file << s << " " << a << " " << a->hashValue() << " " << problem->cost(s,a,1) << std::endl;
        }
    }
    output_file.close();
}
///////////////////////////////////////////////////////////////////////////////
//                            SIMULATE                                       //
///////////////////////////////////////////////////////////////////////////////
//Simulates the computed policy
void simulate(int max_trials){
    int low_nse = 0, high_nse = 0; // tracks the number of low and high-impact nse (s,a) encountered.
    problem->use_learnedReward(false);

    if(problem->getProblemName() == "boxPushing"){
        BoxProblem* bp = static_cast<BoxProblem*>(problem);
        bp->ResetCostFunction(); //resets using HA values.
    }
    else if(problem->getProblemName() == "navigation"){
        NavProblem* np = static_cast<NavProblem*>(problem);
        np->ResetCostFunction(); //resets using HA values.
    }
    vector <double> o1_cost, o2_cost;
    vector <double> lowNse, highNse;
    double TrialCost_1 = 0, TrialCost_2 = 0;

    for (int i = 0; i < max_trials; i++) {
        TrialCost_1 = 0, TrialCost_2 = 0;
        low_nse = 0, high_nse = 0;
        mlcore::State* tmp = problem->initialState();
        while (!problem->goal(tmp)) {
                Action* a = greedyAction(problem, tmp);
                 if(problem->getProblemName() == "boxPushing"){
                        BoxProblem* bp = static_cast<BoxProblem*>(problem);
                        if(bp->isLow_NSE(tmp,a))
                            low_nse++;
                        if(bp->isHigh_NSE(tmp,a))
                            high_nse++;
                    }
                else if(problem->getProblemName() == "navigation"){
                      NavProblem* cp = static_cast<NavProblem*>(problem);
                        if(cp->isLow_NSE(tmp,a)){
                            low_nse++;
                          }
                        if(cp->isHigh_NSE(tmp,a))
                            high_nse++;
                    }

                TrialCost_1 += problem->cost(tmp,a,0) * problem->gamma();
                TrialCost_2 +=  problem->cost(tmp,a,1) * problem->gamma();
                double prob = 0.0;
                State* aux = randomSuccessor(problem, tmp, a, &prob);
                tmp = aux;
                if(TrialCost_1 >= 500)
                    break;

            }
            o1_cost.push_back(TrialCost_1);
            o2_cost.push_back(TrialCost_2);
            lowNse.push_back(low_nse);
            highNse.push_back(high_nse);
    }
    accumulator_set<double, stats<tag::variance> > acc_o1_cost;
    for (auto &u : o1_cost){acc_o1_cost(u);}
    accumulator_set<double, stats<tag::variance> > acc_o2_cost;
    for (auto &u : o2_cost){acc_o2_cost(u);}

    accumulator_set<double, stats<tag::variance> > acc_lowNse;
    for (auto &u : lowNse){acc_lowNse(u);}
    accumulator_set<double, stats<tag::variance> > acc_highNse;
    for (auto &u : highNse){acc_highNse(u);}

    cout << "Avg. Exec cost " << mean(acc_o1_cost) << "  " << mean(acc_o2_cost);
    cout << "  Std. Dev. " << sqrt(variance(acc_o1_cost)) << "  " << sqrt(variance(acc_o2_cost)) << endl;

    cout << "Avg. NSE " << mean(acc_lowNse) << "  " << mean(acc_highNse);
    cout << "  Std. Dev. " << sqrt(variance(acc_lowNse)) << "  " << sqrt(variance(acc_highNse)) << endl;
}

void Exploration_feedback(int max_trials,Exploration*& learner){
    mlmobj::MOState* mos = static_cast<mlmobj::MOState*> (problem->initialState());
    for (int trial = 0; trial < max_trials; trial++) {
        for (int episodes = 0; episodes < max_episodes; episodes++) {
            learner->gatherFeedback(problem, mos,1);
        }
     }
     if(!useLearnedValues)
        learner->writeSamples();
}
///////////////////////////////////////////////////////////////////////////////
//                            MAIN                                           //
///////////////////////////////////////////////////////////////////////////////
int main(int argc, char* args[]){
    register_flags(argc, args);
    setupProblem();
    problem->generateAll();
    assert(flag_is_registered_with_value("algorithm"));
    string algorithm = flag_value("algorithm");
    if (flag_is_registered_with_value("v"))
        verbosity = stoi(flag_value("v"));
    if(flag_is_registered_with_value("episodes"))
        max_episodes = stoi(flag_value("episodes"));

    if(algorithm == "Exploration"){
        Solver* temp_solver = new MOLAOStarSolver(problem,true);  /* solve for level 0 optimally */
        temp_solver->solve(problem->initialState());
        delete temp_solver;
        Exploration* learner = nullptr;
        setupExploration(algorithm, learner);
        Exploration_feedback(1,learner);
        delete learner;
    }
    else if (algorithm == "OF") // Oracle Feedback
    {
        if(useLearnedValues) // plan with learned values
        {
            Solver* solver = new MOLAOStarSolver(problem);
            double start_time = clock();
            solver->solve(problem->initialState());
            double end_time = clock();
            std::cout << " total time taken (s) = " << (end_time-start_time)/ CLOCKS_PER_SEC << std::endl;
            int max_trials = 100;
            simulate(max_trials);
        }
        else{
            Solver* temp_solver = new MOLAOStarSolver(problem,true);  /* solve for level 0 optimally */
            temp_solver->solve(problem->initialState());
            double trials = 10;
            if (flag_is_registered_with_value("samples"))
                trials = stoi(flag_value("samples"));

            double budget = 100000000;
            if (flag_is_registered_with_value("budget"))
                budget = stoi(flag_value("budget"));

            generateSamples();
            delete temp_solver;
        }

    }else if(algorithm == "LLAO" && problem->size() == 1){
         Solver* temp_solver = new MOLAOStarSolver(problem,true);  /* solve for level 0 optimally */
         temp_solver->solve(problem->initialState());
         double return_val = ((MOState *) problem->initialState())->mobjCost()[0];
         std::cout << return_val << std::endl;
        if(verbosity >= 100)
            simulate(100);
    }

    else{
        setupSolver(algorithm, solver);
        solver->solve(problem->initialState());
        if(algorithm == "LAO"){
            std::cout << problem->initialState()->cost() << std::endl;
        }
        else{
            std::cout << " Estimated cost "
             << ((MOState *) problem->initialState())->mobjCost()[0] << " "
             << ((MOState *) problem->initialState())->mobjCost()[1] << endl;
         }
        if(algorithm == "LLAO" && flag_is_registered("demo")){
            int trials = 10;
           if (flag_is_registered_with_value("trials"))
                trials = stoi(flag_value("trials"));
             generateSamples(trials,solver);
             generateSamples(testing_file);
         }
         if(algorithm == "LAO" && flag_is_registered("demo")){
            int trials = 10;
           if (flag_is_registered_with_value("trials"))
                trials = stoi(flag_value("trials"));
             generateSamples(trials,solver);
             generateSamples(testing_file);
         }

        if(algorithm == "LLAO" && flag_is_registered("corrections") ){
           int trials = 10;
           if (flag_is_registered_with_value("trials"))
                trials = stoi(flag_value("trials"));
             generateSamples(trials,solver);
             generateSamples(testing_file);
        }

        if(verbosity > 200)
            printPolicy();

        if(verbosity >= 100){
            simulate(100);
            printPolicy();
        }
        if(verbosity >= 10)
            generateSamples(testing_file);
    }
    if(algorithm == "NLLAO" && !flag_is_registered("demo")){
        int lownse_count = 0, highnse_count = 0, sa_count = 0;
        for(mlcore::State* s: problem->states()){
            for(mlcore::Action* a: problem->actions()){
                if(problem->applicable(s,a)){
                    sa_count++;
                    mlmobj::MOState* mos = static_cast<mlmobj::MOState*> (s);
                    if(problem->isLow_NSE(s,a))
                        lownse_count++;
                    if(problem->isHigh_NSE(s,a))
                        highnse_count ++;
                }
            }
        }

        std::cout << "Low Nse fraction= " << (double(lownse_count)/double(sa_count))
                    << "  High Nse fraction= " << (double(highnse_count)/double(sa_count)) << std::endl;
    }

    delete solver;
    delete problem;
}


/*
    This file is part of Leela Zero.
    Copyright (C) 2017-2019 Gian-Carlo Pascutto and contributors

    Leela Zero is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    Leela Zero is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with Leela Zero.  If not, see <http://www.gnu.org/licenses/>.

    Additional permission under GNU GPL version 3 section 7

    If you modify this Program, or any covered work, by linking or
    combining it with NVIDIA Corporation's libraries from the
    NVIDIA CUDA Toolkit and/or the NVIDIA CUDA Deep Neural
    Network library and/or the NVIDIA TensorRT inference library
    (or a modified version of those libraries), containing parts covered
    by the terms of the respective license agreement, the licensors of
    this Program grant you additional permission to convey the resulting
    work.
*/

#include "config.h"

#ifdef __APPLE__
#include <Accelerate/Accelerate.h>
#endif
#ifdef USE_MKL
#include <mkl.h>
#endif
#ifdef USE_OPENBLAS
#include <cblas.h>
#endif
#ifndef USE_BLAS
#include <Eigen/Dense>
#endif

#include "CPUPipe.h"
#include "Network.h"
#include "Im2Col.h"

#ifndef USE_BLAS
// Eigen helpers
template <typename T>
using EigenMatrixMap =
    Eigen::Map<Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic>>;
template <typename T>
using ConstEigenMatrixMap =
    Eigen::Map<const Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic>>;
#endif

void CPUPipe::initialize(int channels) {
    m_input_channels = channels;
}

void CPUPipe::winograd_transform_in(const std::vector<float>& in,
                                    std::vector<float>& V,
                                    const int C) {
    constexpr auto W = BOARD_SIZE;
    constexpr auto H = BOARD_SIZE;
    constexpr auto WTILES = WINOGRAD_WTILES;
    constexpr auto P = WINOGRAD_P;

    constexpr auto Wpad = 2 + WINOGRAD_M * WTILES;

    constexpr auto buffersize = 32;

    std::array<std::array<float, Wpad>, Wpad> in_pad{0.0f};

    std::array<float, buffersize * WINOGRAD_ALPHA * WINOGRAD_ALPHA> buffer;
    auto buffer_offset = 0;
    auto buffer_entries = 0;

    // multiple vector [i0..i5] by Bt and produce [o0..o5]
    // const auto Bt = std::array<float, WINOGRAD_TILE>
    //           {1.0f,  0.0f,     -5.0f/2.0f,  0.0f,      1.0f, 0.0f,
    //            0.0f, -SQ2,      -2.0f,       SQ2/2.0f,  1.0f, 0.0f,
    //            0.0f,  SQ2,      -2.0f,      -SQ2/2.0f,  1.0f, 0.0f,
    //            0.0f, -SQ2/2.0f, -1.0f/2.0f,  SQ2,       1.0f, 0.0f,
    //            0.0f,  SQ2/2.0f, -1.0f/2.0f, -SQ2,       1.0f, 0.0f,
    //            0.0f,  1.0f,      0.0f,      -5.0f/2.0f, 0.0f, 1.0f};
    auto multiply_bt = [](
        float & o0, float & o1, float & o2, float & o3, float & o4, float & o5,
        float i0, float i1, float i2, float i3, float i4, float i5
    ) {
        auto i3m1 = i1 * -SQ2 + i3 * (SQ2 / 2.0f);
        auto i4m2 = i2 * -2.0f + i4 * 1.0f;

        o0 = i0 + i2 * (-5.0f/2.0f) + i4;
        o1 = i3m1 + i4m2;
        o2 = -i3m1 + i4m2;

        auto i3m1_2 = i3 * (SQ2) + i1 * (-SQ2/2.0f);
        auto i4m2_2 = i2 * (-1.0f/2.0f) + i4;

        o3 = i3m1_2 + i4m2_2;
        o4 = -i3m1_2 + i4m2_2;

        o5 = i1 + i3 * (-5.0f/2.0f) + i5;
    };

    for (auto ch = 0; ch < C; ch++) {
        for (auto yin = 0; yin < H; yin++) {
            for (auto xin = 0; xin < W; xin++) {
                in_pad[yin + 1][xin + 1] = in[ch*(W*H) + yin*W + xin];
            }
        }
        for (auto block_y = 0; block_y < WTILES; block_y++) {
            // Tiles overlap by 2
            const auto yin = WINOGRAD_M * block_y;
            for (auto block_x = 0; block_x < WTILES; block_x++) {
                const auto xin = WINOGRAD_M * block_x;
#define DECL_T1(XX) \
                float T1_##XX##_0, T1_##XX##_1, T1_##XX##_2, T1_##XX##_3, T1_##XX##_4, T1_##XX##_5;
                DECL_T1(0)
                DECL_T1(1)
                DECL_T1(2)
                DECL_T1(3)
                DECL_T1(4)
                DECL_T1(5)

                // Calculates transpose(B).x.B
#define MULTIPLY_BT(XX) \
                multiply_bt( \
                    T1_0_##XX, T1_1_##XX, T1_2_##XX, T1_3_##XX, T1_4_##XX, T1_5_##XX, \
                    in_pad[yin + 0][xin + XX], \
                    in_pad[yin + 1][xin + XX], \
                    in_pad[yin + 2][xin + XX], \
                    in_pad[yin + 3][xin + XX], \
                    in_pad[yin + 4][xin + XX], \
                    in_pad[yin + 5][xin + XX] \
                );
                MULTIPLY_BT(0)
                MULTIPLY_BT(1)
                MULTIPLY_BT(2)
                MULTIPLY_BT(3)
                MULTIPLY_BT(4)
                MULTIPLY_BT(5)

#define MULTIPLY_B(XX) \
                multiply_bt( \
                    buffer[buffersize * (XX * WINOGRAD_ALPHA + 0) + buffer_entries], \
                    buffer[buffersize * (XX * WINOGRAD_ALPHA + 1) + buffer_entries], \
                    buffer[buffersize * (XX * WINOGRAD_ALPHA + 2) + buffer_entries], \
                    buffer[buffersize * (XX * WINOGRAD_ALPHA + 3) + buffer_entries], \
                    buffer[buffersize * (XX * WINOGRAD_ALPHA + 4) + buffer_entries], \
                    buffer[buffersize * (XX * WINOGRAD_ALPHA + 5) + buffer_entries], \
                    T1_##XX##_0, T1_##XX##_1, T1_##XX##_2, T1_##XX##_3, T1_##XX##_4, T1_##XX##_5 \
                );
                MULTIPLY_B(0)
                MULTIPLY_B(1)
                MULTIPLY_B(2)
                MULTIPLY_B(3)
                MULTIPLY_B(4)
                MULTIPLY_B(5)

                if (buffer_entries == 0) {
                    buffer_offset = ch * P + block_y * WTILES + block_x;
                }
                buffer_entries++;

                if (buffer_entries >= buffersize ||
                    (ch == C - 1 && block_x == WTILES - 1 && block_y == WTILES - 1)) {

                    for (auto i = 0; i < WINOGRAD_ALPHA * WINOGRAD_ALPHA; i++) {
                        for (auto entry = 0; entry < buffer_entries; entry++) {
                            V[i*C*P + buffer_offset + entry] = buffer[i*buffersize + entry];
                        }
                    }
                    buffer_entries = 0;
                }
            }
        }
    }
}

void CPUPipe::winograd_sgemm(const std::vector<float>& U,
                             const std::vector<float>& V,
                             std::vector<float>& M,
                             const int C, const int K) {
    constexpr auto P = WINOGRAD_P;

    for (auto b = 0; b < WINOGRAD_TILE; b++) {
        const auto offset_u = b * K * C;
        const auto offset_v = b * C * P;
        const auto offset_m = b * K * P;
#ifdef USE_BLAS
        cblas_sgemm(CblasRowMajor, CblasTrans, CblasNoTrans,
                    K, P, C,
                    1.0f,
                    &U[offset_u], K,
                    &V[offset_v], P,
                    0.0f,
                    &M[offset_m], P);
#else
        auto C_mat = EigenMatrixMap<float>(M.data() + offset_m, P, K);
        C_mat.noalias() =
           ConstEigenMatrixMap<float>(V.data() + offset_v, P, C)
            * ConstEigenMatrixMap<float>(U.data() + offset_u, K, C).transpose();
#endif
    }
}

void CPUPipe::winograd_transform_out(const std::vector<float>& M,
                                     std::vector<float>& Y,
                                     const int K) {
    constexpr auto W = BOARD_SIZE;
    constexpr auto H = BOARD_SIZE;
    constexpr auto WTILES = WINOGRAD_WTILES;
    constexpr auto P = WINOGRAD_P;

    // multiple vector [i0..i5] by At and produce [o0..o3]
    // const auto At = std::array<float, WINOGRAD_ALPHA * WINOGRAD_M>
    //       {1.0f, 1.0f,      1.0f,       1.0f,      1.0f,     0.0f,
    //        0.0f, SQ2/2.0f, -SQ2/2.0f,   SQ2,      -SQ2,      0.0f,
    //        0.0f, 1.0f/2.0f, 1.0f/2.0f,  2.0f,      2.0f,     0.0f,
    //        0.0f, SQ2/4.0f, -SQ2/4.0f,   2.0f*SQ2, -2.0f*SQ2, 1.0f};
    auto multiply_at = [](
        float & o0, float & o1, float & o2, float & o3,
        float i0, float i1, float i2, float i3, float i4, float i5
    ) {
        auto t1p2 = (i1 + i2) * (1.0f / 2.0f);
        auto t1m2 = (i1 - i2) * (SQ2/4.0f);
        auto t3p4 = i3 + i4;
        auto t3m4 = (i3 - i4) * (SQ2);

        o0 = i0 + t1p2 + t1p2 + t3p4;
        o1 = t1m2 + t1m2 + t3m4;
        o2 = t1p2 + t3p4 + t3p4;
        o3 = t1m2 + t3m4 + t3m4 + i5;
    };

    for (auto k = 0; k < K; k++) {
        for (auto block_x = 0; block_x < WTILES; block_x++) {
            const auto x = WINOGRAD_M * block_x;
            for (auto block_y = 0; block_y < WTILES; block_y++) {
                const auto y = WINOGRAD_M * block_y;

                const auto b = block_y * WTILES + block_x;
                using WinogradTile =
                    std::array<std::array<float, WINOGRAD_ALPHA>, WINOGRAD_ALPHA>;
                WinogradTile temp_m;
                for (auto xi = 0; xi < WINOGRAD_ALPHA; xi++) {
                    for (auto nu = 0; nu < WINOGRAD_ALPHA; nu++) {
                        temp_m[xi][nu] =
                            M[(xi*WINOGRAD_ALPHA + nu)*K*P + k*P + b];
                    }
                }
                std::array<std::array<float, WINOGRAD_ALPHA>, WINOGRAD_M> temp;
                std::array<std::array<float, WINOGRAD_M>, WINOGRAD_M> o;

                // Calculates transpose(A).temp_m.A
                for (auto j = 0; j < WINOGRAD_ALPHA; j++){
                    multiply_at(
                        temp[0][j], temp[1][j], temp[2][j], temp[3][j],
                        temp_m[0][j], temp_m[1][j], temp_m[2][j], temp_m[3][j], temp_m[4][j], temp_m[5][j]
                    );
                }

                for (auto i = 0; i < WINOGRAD_M; i++){
                    multiply_at(
                        o[i][0], o[i][1], o[i][2], o[i][3],
                        temp[i][0], temp[i][1], temp[i][2], temp[i][3], temp[i][4], temp[i][5]
                    );
                }

                const auto y_ind = k * H * W + y * W + x;
                for (auto i = 0; i < WINOGRAD_M; i++) {
                    for (auto j = 0; j < WINOGRAD_M; j++) {
                        if (y + i < H && x + j < W) {
                            Y[y_ind + i * W + j] = o[i][j];
                        }
                    }
                }
            }
        }
    }
}

void CPUPipe::winograd_convolve3(const int outputs,
                                 const std::vector<float>& input,
                                 const std::vector<float>& U,
                                 std::vector<float>& V,
                                 std::vector<float>& M,
                                 std::vector<float>& output) {

    constexpr unsigned int filter_len = WINOGRAD_ALPHA * WINOGRAD_ALPHA;
    const auto input_channels = U.size() / (outputs * filter_len);

    winograd_transform_in(input, V, input_channels);
    winograd_sgemm(U, V, M, input_channels, outputs);
    winograd_transform_out(M, output, outputs);
}

template<unsigned int filter_size>
void convolve(const size_t outputs,
              const std::vector<float>& input,
              const std::vector<float>& weights,
              const std::vector<float>& biases,
              std::vector<float>& output) {
    // The size of the board is defined at compile time
    constexpr unsigned int width = BOARD_SIZE;
    constexpr unsigned int height = BOARD_SIZE;
    constexpr auto num_intersections = width * height;
    constexpr auto filter_len = filter_size * filter_size;
    const auto input_channels = weights.size() / (biases.size() * filter_len);
    const auto filter_dim = filter_len * input_channels;
    assert(outputs * num_intersections == output.size());

    std::vector<float> col(filter_dim * width * height);
    im2col<filter_size>(input_channels, input, col);

    // Weight shape (output, input, filter_size, filter_size)
    // 96 18 3 3
    // C←αAB + βC
    // outputs[96,19x19] = weights[96,18x3x3] x col[18x3x3,19x19]
    // M Number of rows in matrices A and C.
    // N Number of columns in matrices B and C.
    // K Number of columns in matrix A; number of rows in matrix B.
    // lda The size of the first dimention of matrix A; if you are
    // passing a matrix A[m][n], the value should be m.
    //    cblas_sgemm(CblasRowMajor, TransA, TransB, M, N, K, alpha, A, lda, B,
    //                ldb, beta, C, N);
#ifdef USE_BLAS
    cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans,
                // M        N            K
                outputs, num_intersections, filter_dim,
                1.0f, &weights[0], filter_dim,
                &col[0], num_intersections,
                0.0f, &output[0], num_intersections);
#else
    auto C_mat = EigenMatrixMap<float>(output.data(),
                                       num_intersections, outputs);
    C_mat.noalias() =
        ConstEigenMatrixMap<float>(col.data(), num_intersections, filter_dim)
        * ConstEigenMatrixMap<float>(weights.data(), filter_dim, outputs);
#endif

    for (unsigned int o = 0; o < outputs; o++) {
        for (unsigned int b = 0; b < num_intersections; b++) {
            output[(o * num_intersections) + b] += biases[o];
        }
    }
}

template <size_t spatial_size>
void batchnorm(const size_t channels,
               std::vector<float>& data,
               const float* const means,
               const float* const stddevs,
               const float* const eltwise = nullptr) {
    const auto lambda_ReLU = [](const auto val) { return (val > 0.0f) ?
                                                          val : 0.0f; };
    for (auto c = size_t{0}; c < channels; ++c) {
        const auto mean = means[c];
        const auto scale_stddev = stddevs[c];
        const auto arr = &data[c * spatial_size];

        if (eltwise == nullptr) {
            // Classical BN
            for (auto b = size_t{0}; b < spatial_size; b++) {
                arr[b] = lambda_ReLU(scale_stddev * (arr[b] - mean));
            }
        } else {
            // BN + residual add
            const auto res = &eltwise[c * spatial_size];
            for (auto b = size_t{0}; b < spatial_size; b++) {
                arr[b] = lambda_ReLU((scale_stddev * (arr[b] - mean)) + res[b]);
            }
        }
    }
}

void CPUPipe::forward(const std::vector<float>& input,
                      std::vector<float>& output_pol,
                      std::vector<float>& output_val) {
    // Input convolution
    constexpr auto P = WINOGRAD_P;
    // Calculate output channels
    const auto output_channels = m_input_channels;
    // input_channels is the maximum number of input channels of any
    // convolution. Residual blocks are identical, but the first convolution
    // might be bigger when the network has very few filters
    const auto input_channels = std::max(static_cast<size_t>(output_channels),
                                         static_cast<size_t>(Network::INPUT_CHANNELS));
    auto conv_out = std::vector<float>(output_channels * NUM_INTERSECTIONS);

    auto V = std::vector<float>(WINOGRAD_TILE * input_channels * P);
    auto M = std::vector<float>(WINOGRAD_TILE * output_channels * P);

    winograd_convolve3(output_channels, input, m_weights->m_conv_weights[0], V, M, conv_out);
    batchnorm<NUM_INTERSECTIONS>(output_channels, conv_out,
                                 m_weights->m_batchnorm_means[0].data(),
                                 m_weights->m_batchnorm_stddevs[0].data());

    // Residual tower
    auto conv_in = std::vector<float>(output_channels * NUM_INTERSECTIONS);
    auto res = std::vector<float>(output_channels * NUM_INTERSECTIONS);
    for (auto i = size_t{1}; i < m_weights->m_conv_weights.size(); i += 2) {
        auto output_channels = m_input_channels;
        std::swap(conv_out, conv_in);
        winograd_convolve3(output_channels, conv_in,
                           m_weights->m_conv_weights[i], V, M, conv_out);
        batchnorm<NUM_INTERSECTIONS>(output_channels, conv_out,
                                     m_weights->m_batchnorm_means[i].data(),
                                     m_weights->m_batchnorm_stddevs[i].data());

        std::swap(conv_in, res);
        std::swap(conv_out, conv_in);
        winograd_convolve3(output_channels, conv_in,
                           m_weights->m_conv_weights[i + 1], V, M, conv_out);
        batchnorm<NUM_INTERSECTIONS>(output_channels, conv_out,
                                     m_weights->m_batchnorm_means[i + 1].data(),
                                     m_weights->m_batchnorm_stddevs[i + 1].data(),
                                     res.data());
    }
    convolve<1>(Network::OUTPUTS_POLICY, conv_out, m_conv_pol_w, m_conv_pol_b, output_pol);
    convolve<1>(Network::OUTPUTS_VALUE, conv_out, m_conv_val_w, m_conv_val_b, output_val);
}

void CPUPipe::push_weights(unsigned int /*filter_size*/,
                           unsigned int /*channels*/,
                           unsigned int outputs,
                           std::shared_ptr<const ForwardPipeWeights> weights) {

    m_weights = weights;

    // Output head convolutions
    m_conv_pol_w = weights->m_conv_pol_w;
    m_conv_pol_b.resize(m_conv_pol_w.size() / outputs, 0.0f);
    m_conv_val_w = weights->m_conv_val_w;
    m_conv_val_b.resize(m_conv_val_w.size() / outputs, 0.0f);
}


/*
  CreateImageFromRegion outputImage indexBegin size insideValue outsideValue referenceImage
*/
#include <iostream>
#include <string>
#include <sstream>

#include <itkImage.h>
#include <itkImageRegionIterator.h>

#include <itkImageFileReader.h>
#include <itkImageFileWriter.h>

#include <boost/program_options.hpp>

#undef __FUNCT__
#define __FUNCT__ "main"
int main(int argc, char **argv)
{
    std::string outImgFile;
    std::string refImgFile;
    std::string indexStartString;
    std::string regionSizeString;
    std::string directionString;

    double      insideValue;
    double      outsideValue;

    //------------------- Set up the command line options database and parse them -----------------------------//
    boost::program_options::options_description optionsDescription("binarize the input image by setting one to the pixels having values within the given range");
    optionsDescription.add_options()
            ("help,h", "display help message")
            ("reference,r", boost::program_options::value< std::string >(&refImgFile), "reference image to get the origin, spacing and orientation")
            ("output,o", boost::program_options::value< std::string >(&outImgFile), "output filename")
            ("indexStart,s",boost::program_options::value< std::string >(&indexStartString), "Start index of the region in the format x,y,z")
            ("regionSize,z",boost::program_options::value< std::string >(&regionSizeString), "Size of the region in the format x,y,z")
            ("direction,d",boost::program_options::value< std::string >(&directionString), "direction (-1 or 1) in each of the axes to move in the format x,y,z")
            ("inside,x",boost::program_options::value< double >(&insideValue)->default_value(1), "Inside value")
            ("outside,y",boost::program_options::value< double >(&outsideValue)->default_value(0), "Outside Value")
            ;

    boost::program_options::variables_map options;
    boost::program_options::store(boost::program_options::parse_command_line(argc,argv,optionsDescription),options);
    boost::program_options::notify(options);

    //If help is asked!
    if(options.count("help")) {
        std::cout<<optionsDescription<<std::endl;
        return EXIT_SUCCESS;
    }

    //Confirm all the options required are given.
    if(!options.count("reference") || !options.count("output")
            || !options.count("indexStart") || !options.count("regionSize") || !options.count("direction")) {
        std::cerr<<"invalid options! run with --help or -h to see the proper options."<<std::endl;
        return EXIT_FAILURE;
    }

    std::vector< int > direction;
    {
        std::stringstream ss(directionString);
        int size;
        while(ss>>size) {
            direction.push_back(size);
            if(ss.peek() == ',')
                ss.ignore();
        }
    }

    //---------------------  Read the image ----------------------//
    typedef itk::Image<double, 3>              ImageType;
    typedef itk::ImageFileReader<ImageType>     ImageReaderType;

    ImageReaderType::Pointer reader = ImageReaderType::New();
    reader->SetFileName(refImgFile);
    reader->Update();
    ImageType::Pointer refImg = reader->GetOutput();

    //Parse the region:
    ImageType::IndexType indexStart;
    {
        std::stringstream ss(indexStartString);
        unsigned int coord;
        unsigned int i=0;
        while(ss>>coord) {
            indexStart[i] = coord;
            if(ss.peek() == ',')
                ss.ignore();
            ++i;
        }
    }

    ImageType::SizeType regionSize;
    {
        std::stringstream ss(regionSizeString);
        int size;
        unsigned int i = 0;
        while(ss>>size) {
            regionSize[i] = size;
            if(ss.peek() == ',')
                ss.ignore();
            ++i;
        }
    }

    std::cout<<"Selected start index: "<<indexStart<<std::endl;
    std::cout<<"Selected size: "<<regionSize<<std::endl;
    std::cout<<"Selected direction: "<<direction.at(0)<<","<<direction.at(1)<<","<<direction.at(2)<<std::endl;

    std::cout<<"Largest possible region from the reference image: "<<std::endl;
    ImageType::RegionType fullRegion = refImg->GetLargestPossibleRegion();
    std::cout<<"Index: "<<fullRegion.GetIndex()<<std::endl;
    std::cout<<"Size: "<<fullRegion.GetSize()<<std::endl;

    ImageType::Pointer outImg;
    outImg = ImageType::New();
    outImg->SetRegions(refImg->GetLargestPossibleRegion());
    outImg->SetOrigin(refImg->GetOrigin());
    outImg->SetSpacing(refImg->GetSpacing());
    outImg->SetDirection(refImg->GetDirection());
    outImg->Allocate();
    //fill all pixels with zero;
    outImg->FillBuffer(outsideValue);

    //set outsideValue to all the pixels lying in the given region.
//    typedef itk::ImageRegionIterator< ImageType >       IteratorType;
//    IteratorType    out_it(outImg,outImg->GetLargestPossibleRegion());
//    IteratorType    in_it(refImg,refImg->GetLargestPossibleRegion());

//    for(out_it.GoToBegin(), in_it.GoToBegin(); !out_it.IsAtEnd(); ++out_it, ++in_it) {
//        {
//            double curr_val = in_it.Get();
//            if(curr_val > lowerThr && curr_val < upperThr )
//                out_it.Set(1);
//        }
//    }

    for(int x = indexStart[0]; x < signed(indexStart[0] + regionSize[0] - 1); (direction.at(0)>0) ? ++x : --x){
        for(int y = indexStart[1]; y < signed(indexStart[1] + regionSize[1] - 1); (direction.at(1)>0) ? ++y : --y){
            for(int z = indexStart[2]; z < signed(indexStart[2] + regionSize[2] - 1); (direction.at(2)>0 ? ++z : --z)){
                ImageType::IndexType pixelIndex;
                if(x<0 || y<0 || z<0) {
                    std::cout<<"Negative index: "<<x<<","<<y<<","<<z<<std::endl;
                    return EXIT_FAILURE;
                }
                pixelIndex[0] = x;
                pixelIndex[1] = y;
                pixelIndex[2] = z;
                if(!fullRegion.IsInside(pixelIndex)) {
                    std::cout<<"Invalid region, this position caused the error: ";
                    std::cout<<pixelIndex<<std::endl;
                    return EXIT_FAILURE;
                }
                outImg->SetPixel(pixelIndex,insideValue);
            }
        }
    }

    //------------------------ Write the output Image -------------------------//
    typedef itk::ImageFileWriter< ImageType >       ImageWriterType;
    ImageWriterType::Pointer    writer = ImageWriterType::New();
    writer->SetFileName(outImgFile);
    writer->SetInput(outImg);
    writer->Update();
    return EXIT_SUCCESS;
}


#include <boost/test/unit_test.hpp>

#include "../test_utils.hh"
#include "../fixture.hh"
#include "preset_matrix_value_generator.hh"

namespace {

template<class T>
void trim_to_triangular(auto matrix, scylla_blas::index_t K = -1) {
    using namespace scylla_blas;

    for (index_t i = 1; i <= matrix->get_row_count(); i++) {
        scylla_blas::vector_segment<T> old_row = matrix->get_row(i);
        scylla_blas::vector_segment<T> new_row;
        for (auto &val : old_row) {
            auto new_val = val;

            /* Skip unnecessary values */
            if (K != -1 && val.index > i + K) new_val.value = 0;
            if (val.index < i) new_val.value = 0;

            if (val.index == i) new_val.value = 0; /* We will insert those later */
            new_row.emplace_back(new_val);
        }
        matrix->update_row(i, new_row);
    }

    for (index_t i = 1; i <= matrix->get_row_count(); i++) {
        matrix->insert_value(i, i, 10);
    }
}
}

BOOST_FIXTURE_TEST_CASE(triangular_solver, mixed_fixture)
{
    using namespace scylla_blas;

    {
        auto float_B = getScyllaVector(test_const::float_vector_3_id);
        auto float_B2 = getScyllaVector(test_const::float_vector_4_id);
        auto old_vec = float_B->get_whole();

        trim_to_triangular<float>(float_BxB);
        scheduler->strsv(Upper, NoTrans, NonUnit, *float_BxB, *float_B);
        scheduler->sgemv(NoTrans, 1,  *float_BxB, *float_B, 0, *float_B2);

        auto new_vec = float_B2->get_whole();
        LogDebug("value {0}", (old_vec + new_vec * (-1)).nrminf() / new_vec.nrminf());
        BOOST_REQUIRE_LE((old_vec + new_vec * (-1)).nrminf() / new_vec.nrminf(), EPSILON);
    }

    {
        auto double_B = getScyllaDoubleVector(test_const::double_vector_3_id);
        auto double_B2 = getScyllaDoubleVector(test_const::double_vector_4_id);
        auto old_vec = double_B->get_whole();

        trim_to_triangular<double>(double_BxB);
        scheduler->dtrsv(Upper, NoTrans, NonUnit, *double_BxB, *double_B);
        scheduler->dgemv(NoTrans, 1,  *double_BxB, *double_B, 0, *double_B2);

        auto new_vec = double_B2->get_whole();
        BOOST_REQUIRE_LE((old_vec + new_vec * (-1)).nrminf() / new_vec.nrminf(), EPSILON);
    }
}

BOOST_FIXTURE_TEST_CASE(triangular_solver_preset, mixed_fixture)
{
    using namespace scylla_blas;

    {
        auto B = test_const::matrix_B;
        std::vector<float> values;
        for (scylla_blas::index_t i = 1; i <= B; i++) {
            values.emplace_back(1.0f);
        }
        auto float_B = getScyllaVectorOf(test_const::float_vector_3_id, values);
        auto float_B2 = getScyllaVector(test_const::float_vector_4_id);
        auto old_vec = float_B->get_whole();

        // generate diagonal
        float_BxB->clear_all();
        std::vector<scylla_blas::matrix_value<float>> diag;
        for (scylla_blas::index_t i = 1; i <= B; i++) {
            diag.emplace_back(i, i, 1.0f);
        }
        preset_matrix_value_generator<float> gen(B, B, diag);
        load_matrix_from_generator(session, gen, *float_BxB);

//        trim_to_triangular<float>(float_BxB);

        print_matrix(*float_BxB);
        scheduler->strsv(Upper, NoTrans, NonUnit, *float_BxB, *float_B);
        scheduler->sgemv(NoTrans, 1,  *float_BxB, *float_B, 0, *float_B2);

        auto new_vec = float_B2->get_whole();
        LogDebug("value {0}", (old_vec + new_vec * (-1)).nrminf() / new_vec.nrminf());
        BOOST_REQUIRE_LE((old_vec + new_vec * (-1)).nrminf() / new_vec.nrminf(), EPSILON);
    }
}

BOOST_FIXTURE_TEST_CASE(triangular_banded_solver, mixed_fixture)
{
    using namespace scylla_blas;
    index_t K = 2;

    {
        auto float_B = getScyllaVector(test_const::float_vector_3_id);
        auto float_B2 = getScyllaVector(test_const::float_vector_4_id);
        auto old_vec = float_B->get_whole();

        trim_to_triangular<float>(float_BxB, K);
        scheduler->stbsv(Upper, NoTrans, NonUnit, K, *float_BxB, *float_B);
        scheduler->sgbmv(NoTrans, 0, K, 1, *float_BxB, *float_B, 0, *float_B2);

        auto new_vec = float_B2->get_whole();
        BOOST_REQUIRE_LE((old_vec + new_vec * (-1)).nrminf() / new_vec.nrminf(), EPSILON);
    }

    {
        auto double_B = getScyllaDoubleVector(test_const::double_vector_3_id);
        auto double_B2 = getScyllaDoubleVector(test_const::double_vector_4_id);
        auto old_vec = double_B->get_whole();

        trim_to_triangular<double>(double_BxB, K);
        scheduler->dtbsv(Upper, NoTrans, NonUnit, K, *double_BxB, *double_B);
        scheduler->dgbmv(NoTrans, 0, K, 1, *double_BxB, *double_B, 0, *double_B2);

        auto new_vec = double_B2->get_whole();
        BOOST_REQUIRE_LE((old_vec + new_vec * (-1)).nrminf() / new_vec.nrminf(), EPSILON);
    }
}


//
// Copyright (c) 2009-2016 The Bitcoin Core developers
// Copyright (c) 2017-2020 Project PAI Foundation
// Distributed under the MIT software license, see the accompanying
// file COPYING or http://www.opensource.org/licenses/mit-license.php.
//

#include "test/test_bwscoin.h"
#include <boost/test/unit_test.hpp>
#include "uint256.h"
#include "stake/hash256prng.h"
#include "stake/treap/tickettreap.h"

BOOST_FIXTURE_TEST_SUITE(lottery_tests, BasicTestingSetup)

BOOST_AUTO_TEST_CASE(basic_prng)
{
    std::string strMsg = "Very deterministic message";
    uint256 hashMsg = Hash(strMsg.begin(), strMsg.end());
    auto prng = Hash256PRNG(hashMsg);

    for (int i = 0; i < 10000;++i)
        prng.Hash256Rand();

    auto lastHash = prng.StateHash();

    auto lastHashExpected = uint256S("c2caca97080f1eb3b9de9c755587c81813c40865c11a24de2fd13991fc39e479");
    BOOST_CHECK(lastHash == lastHashExpected);
}

BOOST_AUTO_TEST_CASE(lottery_num_selection)
{
    static const char a_seed[] = {0x01};
    uint256 seed = Hash(std::begin(a_seed), std::end(a_seed));
    auto prng = Hash256PRNG(seed);
    auto ticketsInPool = 56789;
    auto tooFewTickets = 4;
    auto justEnoughTickets = 5;
    auto ticketsPerBlock = uint16_t(5);

    {
        auto tickets = prng.FindTicketIdxs(ticketsInPool, ticketsPerBlock);
        decltype(tickets) expectedTickets;
        expectedTickets.push_back(51344u);
        expectedTickets.push_back(39045u);
        expectedTickets.push_back(10107u);
        expectedTickets.push_back(31489u);
        expectedTickets.push_back(44461u);
        BOOST_CHECK(expectedTickets == tickets);
    }

    {
        auto tickets = prng.FindTicketIdxs(justEnoughTickets, ticketsPerBlock);
        decltype(tickets) expectedTickets;
        expectedTickets.push_back(2u);
        expectedTickets.push_back(4u);
        expectedTickets.push_back(3u);
        expectedTickets.push_back(0u);
        expectedTickets.push_back(1u);
        BOOST_CHECK(expectedTickets == tickets);
    }

    auto lastHash = prng.StateHash();
    auto lastHashExpected = uint256S("7b49de6cd86694985718b6ba5f815a2fdab7630f4eb1b75560a9b1589adcb41a");
    BOOST_CHECK(lastHash == lastHashExpected);
}

BOOST_AUTO_TEST_CASE(fetch_winners_errors)
{
    auto ticketTreap = TicketTreap();
    for (int i = 0; i < 0xff; ++i) {
        const char a_seed[] = {char(i)};
        uint256 hash = Hash(std::begin(a_seed), std::end(a_seed));
        auto value = Value(
            uint32_t(i),
            (i%2 == 0),
            (i%2 != 0),
            (i%2 == 0),
            (i%2 != 0)
        );
        ticketTreap = ticketTreap.put(hash, value);
    }

    // Bad index too big.
    prevector<64,uint32_t> idxs;
    idxs.push_back(1u);
    idxs.push_back(2u);
    idxs.push_back(3u);
    idxs.push_back(4u);
    idxs.push_back(256u);
    BOOST_CHECK_THROW(
        ticketTreap.fetchWinners(idxs),
         std::runtime_error);
}

BOOST_AUTO_TEST_SUITE_END()