/******************************************************************************
*       SOFA, Simulation Open-Framework Architecture, development version     *
*                (c) 2006-2017 INRIA, USTL, UJF, CNRS, MGH                    *
*                                                                             *
* This program is free software; you can redistribute it and/or modify it     *
* under the terms of the GNU Lesser General Public License as published by    *
* the Free Software Foundation; either version 2.1 of the License, or (at     *
* your option) any later version.                                             *
*                                                                             *
* This program 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 Lesser General Public License *
* for more details.                                                           *
*                                                                             *
* You should have received a copy of the GNU Lesser General Public License    *
* along with this program. If not, see <http://www.gnu.org/licenses/>.        *
*******************************************************************************
* Authors: The SOFA Team and external contributors (see Authors.txt)          *
*                                                                             *
* Contact information: contact@sofa-framework.org                             *
******************************************************************************/
// Author: François Faure, INRIA-UJF, (C) 2011
//
// Copyright: See COPYING file that comes with this distribution
#include <SofaEigen2Solver/SVDLinearSolver.h>
#include <sofa/core/visual/VisualParams.h>
#include <SofaBaseLinearSolver/FullMatrix.h>
#include <SofaBaseLinearSolver/SparseMatrix.h>
#include <SofaBaseLinearSolver/CompressedRowSparseMatrix.h>
#include <sofa/simulation/MechanicalVisitor.h>
#include <sofa/helper/system/thread/CTime.h>
#include <sofa/helper/AdvancedTimer.h>
#include <sofa/core/ObjectFactory.h>
#include <iostream>
#include <Eigen/Dense>
#include <Eigen/Core>

namespace sofa
{
namespace component
{
namespace linearsolver
{
using core::VecId;
using namespace sofa::defaulttype;
using namespace sofa::core::behavior;
using namespace sofa::simulation;
#ifdef DISPLAY_TIME
using sofa::helper::system::thread::CTime;
#endif

template<class TMatrix, class TVector>
SVDLinearSolver<TMatrix,TVector>::SVDLinearSolver()
    : f_verbose( initData(&f_verbose,false,"verbose","Dump system state at each iteration") )
    , f_minSingularValue( initData(&f_minSingularValue,(Real)1.0e-6,"minSingularValue","Thershold under which a singular value is set to 0, for the stabilization of ill-conditioned system.") )
    , f_conditionNumber( initData(&f_conditionNumber,(Real)0.0,"conditionNumber","Condition number of the matrix: ratio between the largest and smallest singular values. Computed in method solve.") )
{
#ifdef DISPLAY_TIME
    timeStamp = 1.0 / (double)CTime::getRefTicksPerSec();
#endif
}

/// Solve Mx=b
template<class TMatrix, class TVector>
void SVDLinearSolver<TMatrix,TVector>::solve(Matrix& M, Vector& x, Vector& b)
{
#ifdef SOFA_DUMP_VISITOR_INFO
    simulation::Visitor::printComment("SVD");
#endif
#ifdef DISPLAY_TIME
    CTime timer;
    double time1 = (double) timer.getTime();
#endif
    const bool printLog = this->f_printLog.getValue();
    const bool verbose  = f_verbose.getValue();

    /// Convert the matrix and the right-hand vector to Eigen objects
    Eigen::MatrixXd m(M.rowSize(),M.colSize());
    Eigen::VectorXd rhs(M.rowSize());
    for(unsigned i=0; i<(unsigned)M.rowSize(); i++ )
    {
        for( unsigned j=0; j<(unsigned)M.colSize(); j++ )
            m(i,j) = M[i][j];
        rhs(i) = b[i];
    }
    if(verbose)
    {
        serr << "SVDLinearSolver<TMatrix,TVector>::solve, Here is the matrix m:" << sendl << m << sendl;
    }

    /// Compute the SVD decomposition and the condition number
    Eigen::JacobiSVD<Eigen::MatrixXd> svd(m, Eigen::ComputeThinU | Eigen::ComputeThinV);
    f_conditionNumber.setValue( (Real)(svd.singularValues()(0) / svd.singularValues()(M.rowSize()-1)) );
    if(printLog)
    {
        serr << "SVDLinearSolver<TMatrix,TVector>::solve, the singular values are:" << sendl << svd.singularValues() << sendl;
    }
    if(verbose)
    {
        serr << "Its left singular vectors are the columns of the thin U matrix:" << sendl << svd.matrixU() << sendl;
        serr << "Its right singular vectors are the columns of the thin V matrix:" << sendl << svd.matrixV() << sendl;
    }

    /// Solve the equation system and copy the solution to the SOFA vector
//    Eigen::VectorXd solution = svd.solve(rhs);
//    for(unsigned i=0; i<M.rowSize(); i++ ){
//        x[i] = solution(i);
//    }
    Eigen::VectorXd Ut_b = svd.matrixU().transpose() *  rhs;
    Eigen::VectorXd S_Ut_b(M.colSize());
    for( unsigned i=0; i<(unsigned)M.colSize(); i++ )   /// product with the diagonal matrix, using the threshold for near-null values
    {
        if( svd.singularValues()[i] > f_minSingularValue.getValue() )
            S_Ut_b[i] = Ut_b[i]/svd.singularValues()[i];
        else
            S_Ut_b[i] = (Real)0.0 ;
    }
    Eigen::VectorXd solution = svd.matrixV() * S_Ut_b;
    for(unsigned i=0; i<(unsigned)M.rowSize(); i++ )
    {
        x[i] = (Real) solution(i);
    }

    if( printLog )
    {
#ifdef DISPLAY_TIME
        time1 = (double)(((double) timer.getTime() - time1) * timeStamp / (nb_iter-1));
        std::cerr<<"SVDLinearSolver::solve, SVD = "<<time1<<std::endl;
#endif
        serr << "SVDLinearSolver<TMatrix,TVector>::solve, rhs vector = " << sendl << rhs.transpose() << sendl;
        serr << "SVDLinearSolver<TMatrix,TVector>::solve, solution = " << sendl << x << sendl;
        serr << "SVDLinearSolver<TMatrix,TVector>::solve, verification, mx - b = " << sendl << (m * solution - rhs ).transpose() << sendl;
    }
}

SOFA_DECL_CLASS(SVDLinearSolver)

int SVDLinearSolverClass = core::RegisterObject("Linear system solver using the conjugate gradient iterative algorithm")
        .add< SVDLinearSolver< FullMatrix<double>, FullVector<double> > >()
        .add< SVDLinearSolver< FullMatrix<float>, FullVector<float> > >()
        .addAlias("SVDLinear")
        .addAlias("SVD")
        ;

} // namespace linearsolver

} // namespace component

} // namespace sofa


#include "Random.hpp"

#include <boost/random.hpp>
#include <boost/random/random_device.hpp>

//
// class Random
//

template <class T>
Random<T>::Random() {

    boost::random::random_device seed_gen;
    this->engine = new boost::random::mt19937_64(seed_gen());

}

template <class T>
Random<T>::~Random() {

    delete(this->engine);

}

template <class T>
T Random<T>::operator()() const {

    return get();

}

//
// class Poison
//

Poisson::Poisson(double arg) {

    this->poisson = new boost::random::poisson_distribution<int, double>(arg);

}

Poisson::~Poisson() {

    delete(this->poisson);

}

int Poisson::get() const {

    return (*poisson)(*engine);

}

//
// class Exponential
//

Exponential::Exponential(double arg) {

    this->exponential = new boost::random::exponential_distribution<double>(arg);

}

Exponential::~Exponential() {

    delete(this->exponential);

}

double Exponential::get() const {

    return (*exponential)(*engine);

}


/*=========================================================================
 *
 * Copyright Universitat Pompeu Fabra, Department of Information and
 * Comunication Technologies.
 *
 *  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.txt
 *
 *  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 itkRWSegmentationFilter_hxx
#define itkRWSegmentationFilter_hxx

#include <itkImageRegionIterator.h>
#include <itkImageRegionConstIterator.h>

#include <vector>
#include <Eigen/Dense>
#include <Eigen/Sparse>

namespace itk
{
template <typename TInputImage, typename TOutputImage>
void RWSegmentationFilter<TInputImage, TOutputImage>::GenerateData()
{
  if (m_LabelImage == nullptr) // Exit if SetLabelImage has not been called
  {
    std::cout << "Label image has not been set" << std::endl;
    return;
  }
  if (OutputImageType::ImageDimension != 2 && OutputImageType::ImageDimension != 3)
  {
    std::cout << "Exit segmentation. Image dimension must be 2 or 3 but is " << OutputImageType::ImageDimension << std::endl;
    return;
  }

  // Get regions to iterate for original and label images
  typename OutputImageType::RegionType regionLabel = m_LabelImage->GetLargestPossibleRegion();

  // Set label image iterator to get the bounding box size and the quantity of marked nodes
  typedef itk::ImageRegionIterator<OutputImageType> IteratorLabelType;
  IteratorLabelType itLabel(m_LabelImage, regionLabel);

  // Crop image to bounds containin labels
  typename OutputImageType::RegionType regionLabelCrop;
  typename InputImageType::RegionType regionCrop;

  for (int i = 0; i != OutputImageType::ImageDimension; ++i)
  {
    regionLabelCrop.SetSize(i, regionLabel.GetIndex()[i]);
    regionLabelCrop.SetIndex(i, regionLabel.GetSize()[i]);
  }

  int markedLength = 0;

  itLabel.GoToBegin();
  while (!itLabel.IsAtEnd())
  {
    // Get image boundaries to crop image
    if (itLabel.Get() != 0)
    {
      ++markedLength;
      // Get bounding region
      for (int i = 0; i != OutputImageType::ImageDimension; ++i)
      {
        if (itLabel.GetIndex()[i] < regionLabelCrop.GetIndex()[i])
        {
          regionCrop.SetIndex(i, itLabel.GetIndex()[i]);
          regionLabelCrop.SetIndex(i, itLabel.GetIndex()[i]);
        }
        else if (itLabel.GetIndex()[i] > regionLabelCrop.GetSize()[i])
        {
          regionCrop.SetSize(i, itLabel.GetIndex()[i]);
          regionLabelCrop.SetSize(i, itLabel.GetIndex()[i]);
        }
      }
    }
    ++itLabel;
  }

  int totalNodes = 1;
  for (int i = 0; i != OutputImageType::ImageDimension; ++i)
  {
    regionCrop.SetIndex(i, regionCrop.GetIndex()[i]);
    regionLabelCrop.SetIndex(i, regionLabelCrop.GetIndex()[i]);
    regionCrop.SetSize(i, regionCrop.GetSize()[i] - regionCrop.GetIndex()[i] + 1);
    regionLabelCrop.SetSize(i, regionLabelCrop.GetSize()[i] - regionLabelCrop.GetIndex()[i] + 1);

    totalNodes *= regionCrop.GetSize()[i];
  }

  /////////////////////// Build Graph /////////////////////////////
  /*
    // Create graph. Each node will correspond to a pixel connected to its neighbors by an edge.
    // Neighbors are those nodes which distance is 1 ( d = sqrt((x-xi)² + (y-yi)² + (z-zi)²) = 1 ),
    // where x, y, and z are the indices of a pixel and xi, yi, and zi are the indeces of a neighboring pixel
    // Iterate through all pixels of input and label images
    */

  int unmarkedLength = totalNodes - markedLength; /* Quantity of unmarked nodes */

  // Define vectors to store graph data
  std::vector<float> *nodes = new std::vector<float>(totalNodes);  /* Pixel intensity for all nodes/pixels */
  std::vector<float> *labels = new std::vector<float>(totalNodes); /* Label of each node/pixel */
  /*  For node 'i'. If 'i' is a marked node, previousFound->at(i) is how many unmarked nodes there are before node 'i'.
        If 'i' is an unmarked node, previousFound->at(i) is how many marked nodes there are before node 'i'.
        This values are needed to build ordered Lu and BT matrices.
    */
  std::vector<int> *previousFound = new std::vector<int>(totalNodes);
  std::vector<int> *unmarked = new std::vector<int>(unmarkedLength); /* Indices of unmarked nodes ordered */
  std::vector<int> *marked = new std::vector<int>(markedLength);     /* Store labels of marked nodes */
  std::vector<int> *markedIdx = new std::vector<int>(markedLength);  /* Indices of marked nodes ordered */
  std::vector<int> *nameLabels = new std::vector<int>();             /* Values of the different labels of the prior */

  // Set bounding box image iterators
  typedef itk::ImageRegionConstIterator<InputImageType> ConstIteratorImageType;
  ConstIteratorImageType itImageCrop(this->GetInput(), regionCrop);
  IteratorLabelType itLabelCrop(m_LabelImage, regionLabelCrop);

  int foundMarked = 0;
  int foundUnmarked = 0;
  int NodeIdx = 0;

  itImageCrop.GoToBegin();
  itLabelCrop.GoToBegin();
  while (!itImageCrop.IsAtEnd())
  {
    // Store intensity in a std::vector
    nodes->at(NodeIdx) = itImageCrop.Get();
    // Store labels in a std::vector. Each index of 'nodes' and 'labels' correspond to the same pixel
    labels->at(NodeIdx) = itLabelCrop.Get();

    if (itLabelCrop.Get() == 0)
    {
      // Get the index of each node that is unmarked
      unmarked->at(foundUnmarked) = NodeIdx;
      // Store how many marked points have been found before this unmarked node
      previousFound->at(NodeIdx) = foundMarked;
      ++foundUnmarked;
    }
    else
    {
      // Set to label value if label != 0
      marked->at(foundMarked) = itLabelCrop.Get();
      // Get the index of each node that is marked
      markedIdx->at(foundMarked) = NodeIdx;
      // Store how many unmarked points have been found before this marked node
      previousFound->at(NodeIdx) = foundUnmarked;
      ++foundMarked;
      bool found;

      // Store the different labels in a std::vector
      found = std::find(nameLabels->begin(), nameLabels->end(), itLabelCrop.Get()) != nameLabels->end();
      if (!found)
      {
        nameLabels->push_back(itLabelCrop.Get());
      }
    }
    ++itImageCrop;
    ++itLabelCrop;
    ++NodeIdx;
  }

  // Sort labels
  std::sort(nameLabels->begin(), nameLabels->end());

  int totalLabels = nameLabels->size();

  if (m_SolveForAllLabels)
  {
    totalLabels += 1;
  }

  // Linear system: Lu * X = -BT * M
  // Convert marked (M) into a Eigen::Sparse matrix markedRHS. Needed for the computation of -BT * M (Eigen::SparseMatrix * Eigen::SparseMatrix)
  Eigen::SparseMatrix<float, Eigen::ColMajor> *markedRHS = new Eigen::SparseMatrix<float, Eigen::ColMajor>(markedLength, totalLabels - 1);
  for (int i = 0; i != totalLabels - 1; ++i)
  {
    int pos = 0;
    for (auto itMarked = marked->begin(); itMarked != marked->end(); ++itMarked)
    {
      if (*itMarked == nameLabels->at(i))
        markedRHS->insert(pos, i) = 1;
      ++pos;
    }
  }

  marked->clear();
  delete marked;

  /////////////////////// Build Laplacian matrix /////////////////////////////
  // Normalize intensity gradient over image spacing
  std::vector<float> spacing;
  typename InputImageType::SpacingType space = this->GetInput()->GetSpacing();
  if (OutputImageType::ImageDimension == 2)
    spacing = {space[1], space[0], space[0], space[1]};
  else if (OutputImageType::ImageDimension == 3)
    spacing = {space[2], space[1], space[0], space[0], space[1], space[2]};

  std::vector<int> neighbors;
  int x = regionCrop.GetSize()[0];
  int y = regionCrop.GetSize()[1];
  //  Right hand of the equation: -BT * M
  Eigen::SparseMatrix<float, Eigen::ColMajor> *BTxM = new Eigen::SparseMatrix<float, Eigen::ColMajor>(markedLength, totalLabels - 1);

  // Build BT
  // Compare whether NumCols < NumRows for less iterations during BT building
  if (markedIdx->size() < unmarked->size())
  {
    Eigen::SparseMatrix<float, Eigen::ColMajor> *BT = new Eigen::SparseMatrix<float, Eigen::ColMajor>(unmarkedLength, markedLength);
    BT->reserve(Eigen::VectorXi::Constant(markedLength, 6));

    int node;
    float valNode, valNeighbor, w;

    // Iterate through marked nodes to build BT. Rows correspond to unmarked nodes, columns to marked nodes
    for (auto itMarked = markedIdx->begin(); itMarked != markedIdx->end(); ++itMarked)
    {
      valNode = nodes->at(*itMarked); // Intensity of node
      // Obtain neighbors indexes. right and left, top and bottom, front and back.
      node = *itMarked;
      if (OutputImageType::ImageDimension == 2)
        neighbors = {node - x, node - 1, node + 1, node + x};
      else if (OutputImageType::ImageDimension == 3)
        neighbors = {node - x * y, node - x, node - 1, node + 1, node + x, node + x * y};
      for (int i = 0; i != neighbors.size(); ++i)
      {
        // Make sure all the neighbors computed fall within the bounding box dimension
        if (neighbors.at(i) >= 0 && neighbors.at(i) < totalNodes && labels->at(neighbors.at(i)) == 0)
        {
          valNeighbor = nodes->at(neighbors.at(i));                                    // Intensity of neighbor pixel
          w = (exp(-m_Beta * pow((valNode - valNeighbor) / spacing.at(i), 2)) + 1e-6); // Intensity gradient following a Gaussian function
          //  Columns of BT correspond to marked nodes, rows to unmarked
          BT->insert(neighbors.at(i) - previousFound->at(neighbors.at(i)), node - previousFound->at(node)) = -w;
        }
      }
    }
    markedIdx->clear();
    //  Right hand of the equation: -BT * M
    *BTxM = -*BT * *markedRHS;
    BT->resize(0, 0);
    BT->data().squeeze();
    markedRHS->resize(0, 0);
    markedRHS->data().squeeze();
    delete markedIdx, markedRHS, BT;
  }
  else
  {
    markedIdx->clear();
    delete markedIdx;

    Eigen::SparseMatrix<float, Eigen::RowMajor> *BT = new Eigen::SparseMatrix<float, Eigen::RowMajor>(unmarkedLength, markedLength);
    BT->reserve(Eigen::VectorXi::Constant(markedLength, 6));

    int node;
    float valNode, valNeighbor, w;

    // Iterate through unmarked nodes to build BT. Rows correspond to unmarked nodes, columns to marked nodes
    for (auto itUnmarked = unmarked->begin(); itUnmarked != unmarked->end(); ++itUnmarked)
    {
      valNode = nodes->at(*itUnmarked); // Intensity of node
      // Obtain neighbors indexes. right and left, top and bottom, front and back.
      node = *itUnmarked;
      if (OutputImageType::ImageDimension == 2)
        neighbors = {node - x, node - 1, node + 1, node + x};
      else if (OutputImageType::ImageDimension == 3)
        neighbors = {node - x * y, node - x, node - 1, node + 1, node + x, node + x * y};
      for (int i = 0; i != neighbors.size(); ++i)
      {
        // Make sure all the neighbors computed fall within the bounding box dimension
        if (neighbors.at(i) >= 0 && neighbors.at(i) < totalNodes && labels->at(neighbors.at(i)) != 0)
        {
          valNeighbor = nodes->at(neighbors.at(i));                                    // Intensity of neighbor pixel
          w = (exp(-m_Beta * pow((valNode - valNeighbor) / spacing.at(i), 2)) + 1e-6); // Intensity gradient following a Gaussian function
          //  Columns of BT correspond to marked nodes, rows to unmarked
          BT->insert(node - previousFound->at(node), neighbors.at(i) - previousFound->at(neighbors.at(i))) = -w;
        }
      }
    }
    //  Right hand of the equation: -BT * M
    *BTxM = -*BT * *markedRHS;
    BT->resize(0, 0);
    BT->data().squeeze();
    markedRHS->resize(0, 0);
    markedRHS->data().squeeze();
    delete markedRHS, BT;
  }

  // Build Lu. LHS of the equation
  Eigen::SparseMatrix<float, Eigen::RowMajor> *Lu = new Eigen::SparseMatrix<float, Eigen::RowMajor>(unmarkedLength, unmarkedLength);
  Lu->reserve(Eigen::VectorXi::Constant(unmarkedLength, 7));

  // Iterate through unmarked nodes to build Lu. Rows correspond to unmarked nodes, columns to unmarked nodes
  for (auto itUnmarked = unmarked->begin(); itUnmarked != unmarked->end(); ++itUnmarked)
  {
    int node;
    float valNode, valNeighbor, w, degree;

    valNode = nodes->at(*itUnmarked); // Intensity of node
    // Obtain neighbors indexes. right and left, top and bottom, front and back.
    node = *itUnmarked;
    if (OutputImageType::ImageDimension == 2)
      neighbors = {node - x, node - 1, node + 1, node + x};
    else if (OutputImageType::ImageDimension == 3)
      neighbors = {node - x * y, node - x, node - 1, node + 1, node + x, node + x * y};
    degree = 0.0;
    for (int i = 0; i != neighbors.size(); ++i)
    {
      // Make sure all the neighbors computed fall within the bounding box dimension
      if (neighbors.at(i) >= 0 && neighbors.at(i) < totalNodes)
      {
        valNeighbor = nodes->at(neighbors.at(i));                                    // Intensity of neighbor pixel
        w = (exp(-m_Beta * pow((valNode - valNeighbor) / spacing.at(i), 2)) + 1e-6); // Intensity gradient following a Gaussian function
        degree += w;                                                                 // Sum of the weights
        //  Columns of Lu correspond to unmarked nodes
        if (labels->at(neighbors.at(i)) == 0) // If neighbor is an unmarked node, build Lu
        {
          Lu->insert(node - previousFound->at(node), neighbors.at(i) - previousFound->at(neighbors.at(i))) = -w;
        }
      }
    }
    // Add node degree to diagonal of Lu
    Lu->insert(node - previousFound->at(node), node - previousFound->at(node)) = degree;
  }
  Lu->makeCompressed();

  nodes->clear();
  unmarked->clear();
  previousFound->clear();
  neighbors.clear();
  delete nodes, unmarked, previousFound;

  /////////////////////// Solve linear system /////////////////////////////
  // Lu * X = -BT * M
  // Set vector to store the result of the solver
  Eigen::Matrix<float, Eigen::Dynamic, Eigen::Dynamic> *probabilities = new Eigen::Matrix<float, Eigen::Dynamic, Eigen::Dynamic>(unmarkedLength, totalLabels - 1);

  // Allow multithreading. For the moment limited to 8 threads.
  if (m_NumberOfThreads > 0 && m_NumberOfThreads < 9)
  {
    omp_set_num_threads(m_NumberOfThreads);
    Eigen::setNbThreads(m_NumberOfThreads);
  }
  else
  {
    omp_set_num_threads(1);
    Eigen::setNbThreads(1);
  }

  // Set the solver for the problem with LHS Lu. BiCGStab.
  // Select either Eigen BiCGSTAB or ConjugateGradient
  Eigen::BiCGSTAB<Eigen::SparseMatrix<float, Eigen::RowMajor>, Eigen::DiagonalPreconditioner<float>> *solver = new Eigen::BiCGSTAB<Eigen::SparseMatrix<float, Eigen::RowMajor>, Eigen::DiagonalPreconditioner<float>>(*Lu); // Usually faster
  // Eigen::ConjugateGradient< Eigen::SparseMatrix< float ,Eigen::RowMajor> , Eigen::Lower|Eigen::Upper , Eigen::DiagonalPreconditioner<float> > solver(Lu);

  // Set solver parameters
  solver->setTolerance(m_Tolerance);
  solver->setMaxIterations(m_MaximumNumberOfIterations);
  // Compute probabilities with RHS BTxM
  *probabilities = solver->solve(*BTxM);

  m_SolverIterations = solver->iterations();
  m_SolverError = solver->error();

  Lu->resize(0, 0);
  Lu->data().squeeze();
  BTxM->resize(0, 0);
  BTxM->data().squeeze();
  delete Lu, BTxM;

  std::vector<int> *RWLabels = new std::vector<int>(unmarkedLength);

  /*  Assign a label to each unmarked node according to the result of the solver.
        The label that is assigned is that one corresponding to the highest probability.
        Since we solver for S-1 systems, last label probability is computed by subtraction */
  for (int i = 0; i != unmarkedLength; ++i)
  {
    float maxProbability = probabilities->coeffRef(i, 0);
    int maxLabelPos = 0;
    float accumulatedProbability = maxProbability;
    for (int j = 1; j != totalLabels - 1; ++j)
    {
      accumulatedProbability += probabilities->coeffRef(i, j);
      if (probabilities->coeffRef(i, j) > maxProbability)
      {
        maxProbability = probabilities->coeffRef(i, j);
        maxLabelPos = j;
      }
    }
    RWLabels->at(i) = 0.95 - accumulatedProbability > maxProbability ? nameLabels->back() : nameLabels->at(maxLabelPos);
  }
  probabilities->resize(0, 0);
  delete probabilities;

  int valBackground;
  if (!m_WriteBackground)
    valBackground = 0;
  else
    valBackground = nameLabels->back();

  typename OutputImageType::Pointer outputLabels = this->GetOutput();
  outputLabels->Graft(m_LabelImage);
  outputLabels->FillBuffer(valBackground);

  // Iterate through original label image to create segmentation image
  // Assign labels known from label image to their original value. Assign unmarked labels
  // according to the result from the solver

  IteratorLabelType itOut1(outputLabels, regionLabelCrop);

  // Set computed labels to segmentation image
  int unmarkedIdx = 0;
  int idxOutput = 0;

  itOut1.GoToBegin();
  while (!itOut1.IsAtEnd())
  {
    if (labels->at(idxOutput) == 0)
    {
      if (RWLabels->at(unmarkedIdx) != nameLabels->back())
        itOut1.Set(RWLabels->at(unmarkedIdx));
      else
        itOut1.Set(valBackground);
      ++unmarkedIdx;
    }
    else if (labels->at(idxOutput) != nameLabels->back())
      itOut1.Set(labels->at(idxOutput));

    ++idxOutput;
    ++itOut1;
  }

  labels->clear();
  RWLabels->clear();
  nameLabels->clear();
  delete labels, RWLabels, nameLabels;

  return;
}

template <typename TInputImage, typename TOutputImage>
void RWSegmentationFilter<TInputImage, TOutputImage>::PrintSelf(std::ostream &os, Indent indent) const
{
  Superclass::PrintSelf(os, indent);
  os << indent << "Beta: " << m_Beta << std::endl;
  os << indent << "NumberOfThreads: " << m_NumberOfThreads << std::endl;
  os << indent << "Tolerance: " << m_Tolerance << std::endl;
  os << indent << "MaximumNumberOfIterations: " << m_MaximumNumberOfIterations << std::endl;
  os << indent << "WriteBackground: " << m_WriteBackground << std::endl;
}

} // namespace itk
#endif


// Copyright (C) 2011  Davis E. King (davis@dlib.net)
// License: Boost Software License   See LICENSE.txt for the full license.

#include <sstream>
#include <string>
#include <cstdlib>
#include <ctime>
#include "tester.h"
#include <dlib/svm_threaded.h>
#include <dlib/rand.h>

typedef dlib::matrix<double,3,1> lhs_element;
typedef dlib::matrix<double,3,1> rhs_element;

namespace
{
    using namespace test;
    using namespace dlib;
    using namespace std;

    logger dlog("test.assignment_learning");

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

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

    struct feature_extractor_dense
    {
        typedef matrix<double,3,1> feature_vector_type;

        typedef ::lhs_element lhs_element;
        typedef ::rhs_element rhs_element;

        unsigned long num_features() const
        {
            return 3;
        }

        void get_features (
            const lhs_element& left,
            const rhs_element& right,
            feature_vector_type& feats
        ) const
        {
            feats = squared(left - right);
        }

    };

    void serialize   (const feature_extractor_dense& , std::ostream& ) {}
    void deserialize (feature_extractor_dense&       , std::istream& ) {}

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

    struct feature_extractor_sparse
    {
        typedef std::vector<std::pair<unsigned long,double> > feature_vector_type;

        typedef ::lhs_element lhs_element;
        typedef ::rhs_element rhs_element;

        unsigned long num_features() const
        {
            return 3;
        }

        void get_features (
            const lhs_element& left,
            const rhs_element& right,
            feature_vector_type& feats
        ) const
        {
            feats.clear();
            feats.push_back(make_pair(0,squared(left-right)(0)));
            feats.push_back(make_pair(1,squared(left-right)(1)));
            feats.push_back(make_pair(2,squared(left-right)(2)));
        }

    };

    void serialize   (const feature_extractor_sparse& , std::ostream& ) {}
    void deserialize (feature_extractor_sparse&       , std::istream& ) {}

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

    typedef std::pair<std::vector<lhs_element>, std::vector<rhs_element> > sample_type;
    typedef std::vector<long> label_type;

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

    void make_data (
        std::vector<sample_type>& samples,
        std::vector<label_type>& labels
    )
    {
        lhs_element a, b, c, d;
        a = 1,0,0;
        b = 0,1,0;
        c = 0,0,1;
        d = 0,1,1;

        std::vector<lhs_element> lhs;
        std::vector<rhs_element> rhs;
        label_type label;

        lhs.push_back(a);
        lhs.push_back(b);
        lhs.push_back(c);

        rhs.push_back(b);
        rhs.push_back(a);
        rhs.push_back(c);

        label.push_back(1);
        label.push_back(0);
        label.push_back(2);

        samples.push_back(make_pair(lhs,rhs));
        labels.push_back(label);

        lhs.clear();
        rhs.clear();
        label.clear();

        lhs.push_back(a);
        lhs.push_back(b);
        lhs.push_back(c);

        rhs.push_back(c);
        rhs.push_back(b);
        rhs.push_back(a);
        rhs.push_back(d);

        label.push_back(2);
        label.push_back(1);
        label.push_back(0);

        samples.push_back(make_pair(lhs,rhs));
        labels.push_back(label);

        lhs.clear();
        rhs.clear();
        label.clear();

        lhs.push_back(a);
        lhs.push_back(b);
        lhs.push_back(c);

        rhs.push_back(c);
        rhs.push_back(a);
        rhs.push_back(d);

        label.push_back(1);
        label.push_back(-1);
        label.push_back(0);

        samples.push_back(make_pair(lhs,rhs));
        labels.push_back(label);

        lhs.clear();
        rhs.clear();
        label.clear();

        lhs.push_back(d);
        lhs.push_back(b);
        lhs.push_back(c);

        label.push_back(-1);
        label.push_back(-1);
        label.push_back(-1);

        samples.push_back(make_pair(lhs,rhs));
        labels.push_back(label);

        lhs.clear();
        rhs.clear();
        label.clear();

        samples.push_back(make_pair(lhs,rhs));
        labels.push_back(label);

    }

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

    void make_data_force (
        std::vector<sample_type>& samples,
        std::vector<label_type>& labels
    )
    {
        lhs_element a, b, c, d;
        a = 1,0,0;
        b = 0,1,0;
        c = 0,0,1;
        d = 0,1,1;

        std::vector<lhs_element> lhs;
        std::vector<rhs_element> rhs;
        label_type label;

        lhs.push_back(a);
        lhs.push_back(b);
        lhs.push_back(c);

        rhs.push_back(b);
        rhs.push_back(a);
        rhs.push_back(c);

        label.push_back(1);
        label.push_back(0);
        label.push_back(2);

        samples.push_back(make_pair(lhs,rhs));
        labels.push_back(label);

        lhs.clear();
        rhs.clear();
        label.clear();

        lhs.push_back(a);
        lhs.push_back(b);
        lhs.push_back(c);

        rhs.push_back(c);
        rhs.push_back(b);
        rhs.push_back(a);
        rhs.push_back(d);

        label.push_back(2);
        label.push_back(1);
        label.push_back(0);

        samples.push_back(make_pair(lhs,rhs));
        labels.push_back(label);

        lhs.clear();
        rhs.clear();
        label.clear();

        lhs.push_back(a);
        lhs.push_back(c);

        rhs.push_back(c);
        rhs.push_back(a);

        label.push_back(1);
        label.push_back(0);

        samples.push_back(make_pair(lhs,rhs));
        labels.push_back(label);

        lhs.clear();
        rhs.clear();
        label.clear();

        samples.push_back(make_pair(lhs,rhs));
        labels.push_back(label);

    }

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

    template <typename fe_type, typename F>
    void test1(F make_data, bool force_assignment)
    {
        print_spinner();

        std::vector<sample_type> samples;
        std::vector<label_type> labels;

        make_data(samples, labels);
        make_data(samples, labels);
        make_data(samples, labels);

        randomize_samples(samples, labels);

        structural_assignment_trainer<fe_type> trainer;

        DLIB_TEST(trainer.forces_assignment() == false);
        DLIB_TEST(trainer.get_c() == 100);
        DLIB_TEST(trainer.get_num_threads() == 2);
        DLIB_TEST(trainer.get_max_cache_size() == 5);

        trainer.set_forces_assignment(force_assignment);
        trainer.set_num_threads(3);
        trainer.set_c(50);

        DLIB_TEST(trainer.get_c() == 50);
        DLIB_TEST(trainer.get_num_threads() == 3);
        DLIB_TEST(trainer.forces_assignment() == force_assignment);

        assignment_function<fe_type> ass = trainer.train(samples, labels);

        for (unsigned long i = 0; i < samples.size(); ++i)
        {
            std::vector<long> out = ass(samples[i]);
            dlog << LINFO << "true labels: " << trans(mat(labels[i]));
            dlog << LINFO << "pred labels: " << trans(mat(out));
            DLIB_TEST(trans(mat(labels[i])) == trans(mat(out)));
        }

        double accuracy;

        dlog << LINFO << "samples.size(): "<< samples.size();
        accuracy = test_assignment_function(ass, samples, labels);
        dlog << LINFO << "accuracy: "<< accuracy;
        DLIB_TEST(accuracy == 1);

        accuracy = cross_validate_assignment_trainer(trainer, samples, labels, 3);
        dlog << LINFO << "cv accuracy: "<< accuracy;
        DLIB_TEST(accuracy == 1);

        ostringstream sout;
        serialize(ass, sout);
        istringstream sin(sout.str());
        assignment_function<fe_type> ass2;
        deserialize(ass2, sin);

        DLIB_TEST(ass2.forces_assignment() == ass.forces_assignment());
        DLIB_TEST(length(ass2.get_weights() - ass.get_weights()) < 1e-10);

        for (unsigned long i = 0; i < samples.size(); ++i)
        {
            std::vector<long> out = ass2(samples[i]);
            dlog << LINFO << "true labels: " << trans(mat(labels[i]));
            dlog << LINFO << "pred labels: " << trans(mat(out));
            DLIB_TEST(trans(mat(labels[i])) == trans(mat(out)));
        }
    }

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

    class test_assignment_learning : public tester
    {
    public:
        test_assignment_learning (
        ) :
            tester ("test_assignment_learning",
                    "Runs tests on the assignment learning code.")
        {}

        void perform_test (
        )
        {
            test1<feature_extractor_dense>(make_data, false);
            test1<feature_extractor_sparse>(make_data, false);

            test1<feature_extractor_dense>(make_data_force, false);
            test1<feature_extractor_sparse>(make_data_force, false);
            test1<feature_extractor_dense>(make_data_force, true);
            test1<feature_extractor_sparse>(make_data_force, true);
        }
    } a;

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

}


//---------------------------------------------------------------------------//
//!
//! \file   tstWattDistribution.cpp
//! \author Aaron Tumulak
//! \brief  Watt distribution unit tests.
//!
//---------------------------------------------------------------------------//

// Std Lib Includes
#include <iostream>

// Boost Includes
#include <boost/units/systems/si.hpp>
#include <boost/units/systems/cgs.hpp>
#include <boost/units/io.hpp>

// Trilinos Includes
#include <Teuchos_UnitTestHarness.hpp>
#include <Teuchos_RCP.hpp>
#include <Teuchos_Array.hpp>
#include <Teuchos_ParameterList.hpp>
#include <Teuchos_XMLParameterListCoreHelpers.hpp>
#include <Teuchos_VerboseObject.hpp>

// FRENSIE Includes
#include "Utility_UnitTestHarnessExtensions.hpp"
#include "Utility_OneDDistribution.hpp"
#include "Utility_WattDistribution.hpp"
#include "Utility_PhysicalConstants.hpp"
#include "Utility_RandomNumberGenerator.hpp"
#include "Utility_UnitTraits.hpp"
#include "Utility_QuantityTraits.hpp"
#include "Utility_ElectronVoltUnit.hpp"

using boost::units::quantity;
using namespace Utility::Units;
namespace si = boost::units::si;
namespace cgs = boost::units::cgs;

//---------------------------------------------------------------------------//
// Testing Variables
//---------------------------------------------------------------------------//

Teuchos::RCP<Teuchos::ParameterList> test_dists_list;

Teuchos::RCP<Utility::OneDDistribution> distribution(
				 new Utility::WattDistribution( 1.0, 1.0, 1.0, 0.1 ) );

Teuchos::RCP<Utility::UnitAwareOneDDistribution<MegaElectronVolt,si::amount> >
unit_aware_distribution( new Utility::UnitAwareWattDistribution<MegaElectronVolt,si::amount>( 1e6*eV, 1e3*keV, 1e-6/eV, 0.1*MeV ) );

//---------------------------------------------------------------------------//
// Tests.
//---------------------------------------------------------------------------//
// Check that the distribution can be evaluated
TEUCHOS_UNIT_TEST( WattDistribution, evaluate )
{
  double test_value_1 = 0.0 ;
  double test_value_2 = exp( -1.0 ) * sinh( 1.0 );

  TEST_EQUALITY_CONST( distribution->evaluate( 0.0 ), test_value_1 );
  TEST_EQUALITY_CONST( distribution->evaluate( 1.0 ), test_value_2 );
}

//---------------------------------------------------------------------------//
// Check that the unit-aware distribution can be evaluated
TEUCHOS_UNIT_TEST( UnitAwareWattDistribution, evaluate )
{
  double scale_factor = exp( -1.0 )*sinh( 1.0 );

  TEST_EQUALITY_CONST( unit_aware_distribution->evaluate( 0.0*MeV ),
		       0.0*si::mole );
  TEST_EQUALITY_CONST( unit_aware_distribution->evaluate( 1.0*MeV ),
		       scale_factor*si::mole );
}

//---------------------------------------------------------------------------//
// Check that the PDF can be evaluated
TEUCHOS_UNIT_TEST( WattDistribution, evaluatePDF )
{
  double test_value_1 = 0.0 ;
  double test_value_2 = 0.25 * sqrt( Utility::PhysicalConstants::pi ) * exp( 0.25 ) * ( erf( sqrt(0.9) - sqrt(0.25) ) + erf( sqrt(0.9) + sqrt(0.25) ) ) - exp( -0.9 ) * sinh( sqrt(0.9) );
  test_value_2 = pow( test_value_2, -1.0 ) * exp( -1.0 ) * sinh( 1.0 );

  TEST_EQUALITY_CONST( distribution->evaluatePDF( 0.0 ), test_value_1 );
  TEST_EQUALITY_CONST( distribution->evaluatePDF( 1.0 ), test_value_2 );
}

//---------------------------------------------------------------------------//
// Check that the unit-aware PDF can be evaluated
TEUCHOS_UNIT_TEST( UnitAwareWattDistribution, evaluatePDF )
{
  double scale_factor = exp( -1.0 )*sinh( 1.0 );
  scale_factor /= 0.25 * sqrt( Utility::PhysicalConstants::pi ) * exp( 0.25 ) * ( erf( sqrt(0.9) - sqrt(0.25) ) + erf( sqrt(0.9) + sqrt(0.25) ) ) - exp( -0.9 ) * sinh( sqrt(0.9) );

  TEST_EQUALITY_CONST( unit_aware_distribution->evaluatePDF( 0.0*MeV ),
		       0.0/MeV );
  TEST_EQUALITY_CONST( unit_aware_distribution->evaluatePDF( 1.0*MeV ),
		       scale_factor/MeV );
}

//---------------------------------------------------------------------------//
// Check that the distribution can be sampled using OpenMC method
TEUCHOS_UNIT_TEST( WattDistribution, sample )
{
  std::vector<double> fake_stream( 15 );
  fake_stream[0] = 0.8110855833521807; // Maxwellian Distribution
  fake_stream[1] = 0.9603231091455754; // Sample Accepted. Sample is:
  fake_stream[2] = 0.3888878000049402; // 0.23654793157394
  fake_stream[3] = 0.1790067705163392; // Watt Distribution Sample Accepted. Sample is: 0.17431015530718
  fake_stream[4] = 0.7971990719536760; // Maxwellian Distribution
  fake_stream[5] = 0.0448247918373053; // Sample Rejected. Sample is:
  fake_stream[6] = 0.0375590391821071; // 3.32084968701263
  fake_stream[7] = 0.9522705508082248; // Maxwellian Distribution
  fake_stream[8] = 0.5691586660769004; // Sample Accepted. Sample is:
  fake_stream[9] = 0.0083308067369354; // 0.61240561891848
  fake_stream[10] = 0.6415677450291360; // Watt Distribution Sample Rejected. Sample is: 1.08397711671084
  fake_stream[11] = 0.8589987123720327; // Maxwellian Distribution
  fake_stream[12] = 0.4981431293148418; // Sample Accepted. Sample is:
  fake_stream[13] = 0.5932835342621010; // 0.39976547875177
  fake_stream[14] = 0.6731968005067009; // Watt Distribution Sample Accepted. Sample is: 0.86877979533891

  Utility::RandomNumberGenerator::setFakeStream( fake_stream );

  double sample = distribution->sample();
  TEST_FLOATING_EQUALITY( sample, 0.17431015530718, 1e-13 );

  sample = distribution->sample();
  TEST_FLOATING_EQUALITY( sample, 0.86877979533891, 1e-13 );

  Utility::RandomNumberGenerator::unsetFakeStream();
}

//---------------------------------------------------------------------------//
// Check that the unit-aware distribution can be sampled using OpenMC method
TEUCHOS_UNIT_TEST( UnitAwareWattDistribution, sample )
{
  std::vector<double> fake_stream( 15 );
  fake_stream[0] = 0.8110855833521807; // Maxwellian Distribution
  fake_stream[1] = 0.9603231091455754; // Sample Accepted. Sample is:
  fake_stream[2] = 0.3888878000049402; // 0.23654793157394
  fake_stream[3] = 0.1790067705163392; // Watt Distribution Sample Accepted. Sample is: 0.17431015530718
  fake_stream[4] = 0.7971990719536760; // Maxwellian Distribution
  fake_stream[5] = 0.0448247918373053; // Sample Rejected. Sample is:
  fake_stream[6] = 0.0375590391821071; // 3.32084968701263
  fake_stream[7] = 0.9522705508082248; // Maxwellian Distribution
  fake_stream[8] = 0.5691586660769004; // Sample Accepted. Sample is:
  fake_stream[9] = 0.0083308067369354; // 0.61240561891848
  fake_stream[10] = 0.6415677450291360; // Watt Distribution Sample Rejected. Sample is: 1.08397711671084
  fake_stream[11] = 0.8589987123720327; // Maxwellian Distribution
  fake_stream[12] = 0.4981431293148418; // Sample Accepted. Sample is:
  fake_stream[13] = 0.5932835342621010; // 0.39976547875177
  fake_stream[14] = 0.6731968005067009; // Watt Distribution Sample Accepted. Sample is: 0.86877979533891

  Utility::RandomNumberGenerator::setFakeStream( fake_stream );

  quantity<MegaElectronVolt> sample = unit_aware_distribution->sample();
  UTILITY_TEST_FLOATING_EQUALITY( sample, 0.17431015530718*MeV, 1e-13 );

  sample = unit_aware_distribution->sample();
  UTILITY_TEST_FLOATING_EQUALITY( sample, 0.86877979533891*MeV, 1e-13 );

  Utility::RandomNumberGenerator::unsetFakeStream();
}

//---------------------------------------------------------------------------//
// Check that the distribution can be sampled using OpenMC method, passing in
// parameters
TEUCHOS_UNIT_TEST( WattDistribution, sample_pass_parameters )
{
  std::vector<double> fake_stream( 15 );
  fake_stream[0] = 0.6617443503056450; // Maxwellian Distribution
  fake_stream[1] = 0.8510592242616175; // Sample Accepted. Sample is:
  fake_stream[2] = 0.5401745197210969; // 0.09667248521245
  fake_stream[3] = 0.2005969462099806; // Watt Distribution Sample Accepted. Sample is: 0.07927727029875
  fake_stream[4] = 0.5189418543931951; // Maxwellian Distribution
  fake_stream[5] = 0.7484231272861934; // Sample Rejected. Sample is:
  fake_stream[6] = 0.8345865134048199; // 0.13501824430785
  fake_stream[7] = 0.8704061668810904; // Maxwellian Distribution
  fake_stream[8] = 0.5575773053122431; // Sample Accepted. Sample is:
  fake_stream[9] = 0.7452925729418971; // 0.04548390223293
  fake_stream[10] = 0.6556974287666129; // Watt Distribution Sample Rejected. Sample is: 0.05068410532744
  fake_stream[11] = 0.8140855348977614; // Maxwellian Distribution
  fake_stream[12] = 0.5750169688125915; // Sample Accepted. Sample is:
  fake_stream[13] = 0.8372435273691630; // 0.04821527540845
  fake_stream[14] = 0.4721727279140688; // Watt Distribution Sample Accepted. Sample is: 0.04844237604136

  Utility::RandomNumberGenerator::setFakeStream( fake_stream );

  double incident_energy, a_parameter, b_parameter, restriction_energy, sample;

  incident_energy = 0.5;
  a_parameter = 0.2;
  b_parameter = 0.3;
  restriction_energy = 0.4;

  sample = Utility::WattDistribution::sample( incident_energy, a_parameter, b_parameter, restriction_energy );
  TEST_FLOATING_EQUALITY( sample, 0.07927727029875, 1e-13 );

  incident_energy = 0.3;
  a_parameter = 0.2;
  b_parameter = 0.1;
  restriction_energy = 0.25;

  sample = Utility::WattDistribution::sample( incident_energy, a_parameter, b_parameter, restriction_energy );
  TEST_FLOATING_EQUALITY( sample, 0.04844237604136, 1e-13 );

  Utility::RandomNumberGenerator::unsetFakeStream();
}

//---------------------------------------------------------------------------//
// Check that the unit-aware distribution can be sampled using OpenMC method,
// passing in parameters
TEUCHOS_UNIT_TEST( UnitAwareWattDistribution,
		   sample_pass_parameters )
{
  std::vector<double> fake_stream( 15 );
  fake_stream[0] = 0.6617443503056450; // Maxwellian Distribution
  fake_stream[1] = 0.8510592242616175; // Sample Accepted. Sample is:
  fake_stream[2] = 0.5401745197210969; // 0.09667248521245
  fake_stream[3] = 0.2005969462099806; // Watt Distribution Sample Accepted. Sample is: 0.07927727029875
  fake_stream[4] = 0.5189418543931951; // Maxwellian Distribution
  fake_stream[5] = 0.7484231272861934; // Sample Rejected. Sample is:
  fake_stream[6] = 0.8345865134048199; // 0.13501824430785
  fake_stream[7] = 0.8704061668810904; // Maxwellian Distribution
  fake_stream[8] = 0.5575773053122431; // Sample Accepted. Sample is:
  fake_stream[9] = 0.7452925729418971; // 0.04548390223293
  fake_stream[10] = 0.6556974287666129; // Watt Distribution Sample Rejected. Sample is: 0.05068410532744
  fake_stream[11] = 0.8140855348977614; // Maxwellian Distribution
  fake_stream[12] = 0.5750169688125915; // Sample Accepted. Sample is:
  fake_stream[13] = 0.8372435273691630; // 0.04821527540845
  fake_stream[14] = 0.4721727279140688; // Watt Distribution Sample Accepted. Sample is: 0.04844237604136

  Utility::RandomNumberGenerator::setFakeStream( fake_stream );

  Utility::UnitAwareWattDistribution<MegaElectronVolt>::IndepQuantity incident_energy, a_parameter, restriction_energy, sample;

  Utility::UnitAwareWattDistribution<MegaElectronVolt>::InverseIndepQuantity b_parameter;

  incident_energy = 0.5*MeV;
  a_parameter = 0.2*MeV;
  b_parameter = 0.3/MeV;
  restriction_energy = 0.4*MeV;

  sample = Utility::UnitAwareWattDistribution<MegaElectronVolt>::sample( incident_energy, a_parameter, b_parameter, restriction_energy );
  UTILITY_TEST_FLOATING_EQUALITY( sample, 0.07927727029875*MeV, 1e-13 );

  incident_energy = 0.3*MeV;
  a_parameter = 0.2*MeV;
  b_parameter = 0.1/MeV;
  restriction_energy = 0.25*MeV;

  sample = Utility::UnitAwareWattDistribution<MegaElectronVolt>::sample( incident_energy, a_parameter, b_parameter, restriction_energy );
  UTILITY_TEST_FLOATING_EQUALITY( sample, 0.04844237604136*MeV, 1e-13 );

  Utility::RandomNumberGenerator::unsetFakeStream();
}

//---------------------------------------------------------------------------//
// Check that the distribution can be sampled using OpenMC method, passing in
// parameters
TEUCHOS_UNIT_TEST( WattDistribution, sampleAndRecordTrials_pass_parameters )
{
  std::vector<double> fake_stream( 15 );
  fake_stream[0] = 0.6617443503056450; // Maxwellian Distribution
  fake_stream[1] = 0.8510592242616175; // Sample Accepted. Sample is:
  fake_stream[2] = 0.5401745197210969; // 0.09667248521245
  fake_stream[3] = 0.2005969462099806; // Watt Distribution Sample Accepted. Sample is: 0.07927727029875
  fake_stream[4] = 0.5189418543931951; // Maxwellian Distribution
  fake_stream[5] = 0.7484231272861934; // Sample Rejected. Sample is:
  fake_stream[6] = 0.8345865134048199; // 0.13501824430785
  fake_stream[7] = 0.8704061668810904; // Maxwellian Distribution
  fake_stream[8] = 0.5575773053122431; // Sample Accepted. Sample is:
  fake_stream[9] = 0.7452925729418971; // 0.04548390223293
  fake_stream[10] = 0.6556974287666129; // Watt Distribution Sample Rejected. Sample is: 0.05068410532744
  fake_stream[11] = 0.8140855348977614; // Maxwellian Distribution
  fake_stream[12] = 0.5750169688125915; // Sample Accepted. Sample is:
  fake_stream[13] = 0.8372435273691630; // 0.04821527540845
  fake_stream[14] = 0.4721727279140688; // Watt Distribution Sample Accepted. Sample is: 0.04844237604136

  Utility::RandomNumberGenerator::setFakeStream( fake_stream );

  double incident_energy, a_parameter, b_parameter, restriction_energy, sample;
  unsigned trials = 0;

  incident_energy = 0.5;
  a_parameter = 0.2;
  b_parameter = 0.3;
  restriction_energy = 0.4;

  sample = Utility::WattDistribution::sampleAndRecordTrials(incident_energy, a_parameter, b_parameter, restriction_energy, trials);
  TEST_FLOATING_EQUALITY( sample, 0.07927727029875, 1e-13 );
  TEST_EQUALITY_CONST( trials, 1.0 );

  incident_energy = 0.3;
  a_parameter = 0.2;
  b_parameter = 0.1;
  restriction_energy = 0.25;

  sample = Utility::WattDistribution::sampleAndRecordTrials(incident_energy, a_parameter, b_parameter, restriction_energy, trials);
  TEST_FLOATING_EQUALITY( sample, 0.04844237604136, 1e-13 );
  TEST_EQUALITY_CONST( trials, 4.0 );

  Utility::RandomNumberGenerator::unsetFakeStream();
}

//---------------------------------------------------------------------------//
// Check that the unit-aware distribution can be sampled using OpenMC method,
// passing in parameters
TEUCHOS_UNIT_TEST( UnitAwareWattDistribution,
		   sampleAndRecordTrials_pass_parameters )
{
  std::vector<double> fake_stream( 15 );
  fake_stream[0] = 0.6617443503056450; // Maxwellian Distribution
  fake_stream[1] = 0.8510592242616175; // Sample Accepted. Sample is:
  fake_stream[2] = 0.5401745197210969; // 0.09667248521245
  fake_stream[3] = 0.2005969462099806; // Watt Distribution Sample Accepted. Sample is: 0.07927727029875
  fake_stream[4] = 0.5189418543931951; // Maxwellian Distribution
  fake_stream[5] = 0.7484231272861934; // Sample Rejected. Sample is:
  fake_stream[6] = 0.8345865134048199; // 0.13501824430785
  fake_stream[7] = 0.8704061668810904; // Maxwellian Distribution
  fake_stream[8] = 0.5575773053122431; // Sample Accepted. Sample is:
  fake_stream[9] = 0.7452925729418971; // 0.04548390223293
  fake_stream[10] = 0.6556974287666129; // Watt Distribution Sample Rejected. Sample is: 0.05068410532744
  fake_stream[11] = 0.8140855348977614; // Maxwellian Distribution
  fake_stream[12] = 0.5750169688125915; // Sample Accepted. Sample is:
  fake_stream[13] = 0.8372435273691630; // 0.04821527540845
  fake_stream[14] = 0.4721727279140688; // Watt Distribution Sample Accepted. Sample is: 0.04844237604136

  Utility::RandomNumberGenerator::setFakeStream( fake_stream );

  Utility::UnitAwareOneDDistribution<MegaElectronVolt,si::amount>::IndepQuantity incident_energy, a_parameter, restriction_energy, sample;

  Utility::UnitAwareOneDDistribution<MegaElectronVolt,si::amount>::InverseIndepQuantity b_parameter;

  unsigned trials = 0;

  incident_energy = 0.5*MeV;
  a_parameter = 0.2*MeV;
  b_parameter = 0.3/MeV;
  restriction_energy = 0.4*MeV;

  sample = Utility::UnitAwareWattDistribution<MegaElectronVolt,si::amount>::sampleAndRecordTrials(incident_energy, a_parameter, b_parameter, restriction_energy, trials);
  UTILITY_TEST_FLOATING_EQUALITY( sample, 0.07927727029875*MeV, 1e-13 );
  TEST_EQUALITY_CONST( trials, 1.0 );

  incident_energy = 0.3*MeV;
  a_parameter = 0.2*MeV;
  b_parameter = 0.1/MeV;
  restriction_energy = 0.25*MeV;

  sample = Utility::UnitAwareWattDistribution<MegaElectronVolt,si::amount>::sampleAndRecordTrials(incident_energy, a_parameter, b_parameter, restriction_energy, trials);
  UTILITY_TEST_FLOATING_EQUALITY( sample, 0.04844237604136*MeV, 1e-13 );
  TEST_EQUALITY_CONST( trials, 4.0 );

  Utility::RandomNumberGenerator::unsetFakeStream();
}

//---------------------------------------------------------------------------//
// Check that the upper bound of the distribution independent variable can be
// returned
TEUCHOS_UNIT_TEST( WattDistribution, getUpperBoundOfIndepVar )
{
  TEST_EQUALITY_CONST( distribution->getUpperBoundOfIndepVar(), 0.9 );
}

//---------------------------------------------------------------------------//
// Check that the upper bound of the unit-aware distribution independent
// variable can be returned
TEUCHOS_UNIT_TEST( UnitAwareWattDistribution, getUpperBoundOfIndepVar )
{
  TEST_EQUALITY_CONST( unit_aware_distribution->getUpperBoundOfIndepVar(),
		       0.9*MeV );
}

//---------------------------------------------------------------------------//
// Check that the lower bound of the distribution independent variable can be
// returned
TEUCHOS_UNIT_TEST( WattDistribution, getLowerBoundOfIndepVar )
{
  TEST_EQUALITY_CONST( distribution->getLowerBoundOfIndepVar(), 0.0 );
}

//---------------------------------------------------------------------------//
// Check that the lower bound of the unit-aware distribution independent
// variable can be returned
TEUCHOS_UNIT_TEST( UnitAwareWattDistribution, getLowerBoundOfIndepVar )
{
  TEST_EQUALITY_CONST( unit_aware_distribution->getLowerBoundOfIndepVar(),
		       0.0*MeV );
}

//---------------------------------------------------------------------------//
// Check that the distribution type can be returned
TEUCHOS_UNIT_TEST( WattDistribution, getDistributionType )
{
  TEST_EQUALITY_CONST( distribution->getDistributionType(),
		       Utility::WATT_DISTRIBUTION );
}

//---------------------------------------------------------------------------//
// Check that the unit-aware distribution type can be returned
TEUCHOS_UNIT_TEST( UnitAwareWattDistribution, getDistributionType )
{
  TEST_EQUALITY_CONST( unit_aware_distribution->getDistributionType(),
		       Utility::WATT_DISTRIBUTION );
}

//---------------------------------------------------------------------------//
// Check if the distribution is tabular
TEUCHOS_UNIT_TEST( WattDistribution, isTabular )
{
  TEST_ASSERT( !distribution->isTabular() );
}

//---------------------------------------------------------------------------//
// Check if the unit-aware distribution is tabular
TEUCHOS_UNIT_TEST( UnitAwareWattDistribution, isTabular )
{
  TEST_ASSERT( !unit_aware_distribution->isTabular() );
}

//---------------------------------------------------------------------------//
// Check if the distribution is continuous
TEUCHOS_UNIT_TEST( WattDistribution, isContinuous )
{
  TEST_ASSERT( distribution->isContinuous() );
}

//---------------------------------------------------------------------------//
// Check if the unit-aware distribution is continuous
TEUCHOS_UNIT_TEST( UnitAwareWattDistribution, isContinuous )
{
  TEST_ASSERT( unit_aware_distribution->isContinuous() );
}

//---------------------------------------------------------------------------//
// Check that the distribution can be written to an xml file
TEUCHOS_UNIT_TEST( WattDistribution, toParameterList )
{
  Teuchos::RCP<Utility::WattDistribution> true_distribution =
  Teuchos::rcp_dynamic_cast<Utility::WattDistribution>( distribution );

  Teuchos::ParameterList parameter_list;

  parameter_list.set<Utility::WattDistribution>( "test distribution",
                                                  *true_distribution );

  Teuchos::writeParameterListToXmlFile( parameter_list,
                                       "watt_dist_test_list.xml" );

  Teuchos::RCP<Teuchos::ParameterList> read_parameter_list =
  Teuchos::getParametersFromXmlFile( "watt_dist_test_list.xml" );

  TEST_EQUALITY( parameter_list, *read_parameter_list );

  Teuchos::RCP<Utility::WattDistribution>
  copy_distribution( new Utility::WattDistribution );

  *copy_distribution = read_parameter_list->get<Utility::WattDistribution>(
                                                                             "test distribution");

  TEST_EQUALITY( *copy_distribution, *true_distribution );
}

//---------------------------------------------------------------------------//
// Check that the unit-aware distribution can be written to an xml file
TEUCHOS_UNIT_TEST( UnitAwareWattDistribution, toParameterList )
{
  typedef Utility::UnitAwareWattDistribution<MegaElectronVolt,si::amount> UnitAwareWattDistribution;

  Teuchos::RCP<UnitAwareWattDistribution> true_distribution =
  Teuchos::rcp_dynamic_cast<UnitAwareWattDistribution>( unit_aware_distribution );

  Teuchos::ParameterList parameter_list;

  parameter_list.set<UnitAwareWattDistribution>( "test distribution",
                                                  *true_distribution );

  Teuchos::writeParameterListToXmlFile( parameter_list,
                                       "unit_aware_watt_dist_test_list.xml" );

  Teuchos::RCP<Teuchos::ParameterList> read_parameter_list =
  Teuchos::getParametersFromXmlFile( "unit_aware_watt_dist_test_list.xml" );

  TEST_EQUALITY( parameter_list, *read_parameter_list );

  Teuchos::RCP<UnitAwareWattDistribution>
  copy_distribution( new UnitAwareWattDistribution );

  *copy_distribution = read_parameter_list->get<UnitAwareWattDistribution>(
							 "test distribution" );

  TEST_EQUALITY( *copy_distribution, *true_distribution );
}

//---------------------------------------------------------------------------//
// Check that the distribution can be read from an xml file
TEUCHOS_UNIT_TEST( WattDistribution, fromParameterList )
{
  double test_value_1;
  double test_value_2;

  Utility::WattDistribution read_distribution =
    test_dists_list->get<Utility::WattDistribution>( "Watt Distribution A" );

  test_value_1 = 0.0 ;
  test_value_2 = 0.25 * sqrt( Utility::PhysicalConstants::pi ) * exp( 0.25 ) * ( erf( sqrt(0.9) - sqrt(0.25) ) + erf( sqrt(0.9) + sqrt(0.25) ) ) - exp( -0.9 ) * sinh( sqrt(0.9) );
  test_value_2 = pow( test_value_2, -1.0 ) * exp( -1.0 ) * sinh( 1.0 );

  TEST_EQUALITY_CONST( read_distribution.evaluatePDF( 0.0 ), test_value_1 );
  TEST_FLOATING_EQUALITY( read_distribution.evaluatePDF( 1.0 ),
			  test_value_2,
			  1e-15 );

   read_distribution =
    test_dists_list->get<Utility::WattDistribution>( "Watt Distribution B" );

  test_value_1 = 0.0 ;
  test_value_2 = 0.25 * sqrt( Utility::PhysicalConstants::pi * pow( 2.0, 3.0 ) * 1.0 )
    * exp( 0.25 * 2.0 ) * ( erf( sqrt( 1.5 ) - sqrt( 0.25 * 2.0 ) )
    + erf( sqrt( 1.5 ) + sqrt( 0.25 * 2.0 ) ) )
    - 2.0 * exp( - 1.5 ) * sinh( sqrt( 3.0 ) );
  test_value_2 = pow( test_value_2, -1.0 ) * exp( -0.5 ) * sinh( 1.0 );

  TEST_EQUALITY_CONST( read_distribution.evaluatePDF( 0.0 ), test_value_1 );
  TEST_FLOATING_EQUALITY( read_distribution.evaluatePDF( 1.0 ),
			  test_value_2,
			  1e-15 );

}

//---------------------------------------------------------------------------//
// Check that the unit-aware distribution can be read from an xml file
TEUCHOS_UNIT_TEST( UnitAwareWattDistribution, fromParameterList )
{
  typedef Utility::UnitAwareWattDistribution<MegaElectronVolt,si::amount> UnitAwareWattDistribution;

  double scale_factor;
  scale_factor = exp( -1.0 )*sinh( 1.0 );
  scale_factor /= 0.25*sqrt( Utility::PhysicalConstants::pi )*exp( 0.25 )*
    ( erf( sqrt(0.9) - sqrt(0.25) ) +
      erf( sqrt(0.9) + sqrt(0.25) ) ) - exp( -0.9 ) * sinh( sqrt(0.9) );

  UnitAwareWattDistribution read_distribution =
    test_dists_list->get<UnitAwareWattDistribution>( "Unit-Aware Watt Distribution A" );

  TEST_EQUALITY_CONST( read_distribution.evaluate( 0.0*MeV ), 0.0*si::mole );
  UTILITY_TEST_FLOATING_EQUALITY( read_distribution.evaluate( 1.0*MeV ),
				  exp( -1.0 )*sinh( 1.0 )*si::mole,
				  1e-15 );
  UTILITY_TEST_FLOATING_EQUALITY( read_distribution.evaluatePDF( 1.0*MeV ),
				  scale_factor/MeV,
				  1e-15 );

  read_distribution =
    test_dists_list->get<UnitAwareWattDistribution>( "Unit-Aware Watt Distribution B" );

   scale_factor = exp( -0.5 )*sinh( 1.0 );
   scale_factor /= 0.25*sqrt( Utility::PhysicalConstants::pi*pow( 2.0, 3.0 ) )*
     exp( 0.25 * 2.0 )*( erf( sqrt( 1.5 ) - sqrt( 0.25 * 2.0 ) ) +
			 erf( sqrt( 1.5 ) + sqrt( 0.25 * 2.0 ) ) ) -
     2.0*exp( -1.5 )*sinh( sqrt( 3.0 ) );

   TEST_EQUALITY_CONST( read_distribution.evaluate( 0.0*MeV ), 0.0*si::mole );
   UTILITY_TEST_FLOATING_EQUALITY( read_distribution.evaluate( 1.0*MeV ),
				   2*exp( -0.5 )*sinh( 1.0 )*si::mole,
				   1e-15 );
   UTILITY_TEST_FLOATING_EQUALITY( read_distribution.evaluatePDF( 1.0*MeV ),
				   scale_factor/MeV,
				   1e-15 );

}

//---------------------------------------------------------------------------//
// Check that distributions can be scaled
TEUCHOS_UNIT_TEST_TEMPLATE_4_DECL( UnitAwareWattDistribution,
				   explicit_conversion,
				   IndepUnitA,
				   DepUnitA,
				   IndepUnitB,
				   DepUnitB )
{
  typedef typename Utility::UnitTraits<IndepUnitA>::template GetQuantityType<double>::type IndepQuantityA;
  typedef typename Utility::UnitTraits<typename Utility::UnitTraits<IndepUnitA>::InverseUnit>::template GetQuantityType<double>::type InverseIndepQuantityA;

  typedef typename Utility::UnitTraits<IndepUnitB>::template GetQuantityType<double>::type IndepQuantityB;
  typedef typename Utility::UnitTraits<typename Utility::UnitTraits<IndepUnitB>::InverseUnit>::template GetQuantityType<double>::type InverseIndepQuantityB;

  typedef typename Utility::UnitTraits<DepUnitA>::template GetQuantityType<double>::type DepQuantityA;
  typedef typename Utility::UnitTraits<DepUnitB>::template GetQuantityType<double>::type DepQuantityB;

  // Copy from unitless distribution to distribution type A (static method)
  Utility::UnitAwareWattDistribution<IndepUnitA,DepUnitA>
    unit_aware_dist_a_copy = Utility::UnitAwareWattDistribution<IndepUnitA,DepUnitA>::fromUnitlessDistribution( *Teuchos::rcp_dynamic_cast<Utility::WattDistribution>( distribution ) );

  // Copy from distribution type A to distribution type B (explicit cast)
  Utility::UnitAwareWattDistribution<IndepUnitB,DepUnitB>
    unit_aware_dist_b_copy( unit_aware_dist_a_copy );

  IndepQuantityA indep_quantity_a =
    Utility::QuantityTraits<IndepQuantityA>::initializeQuantity( 0.0 );
  InverseIndepQuantityA inv_indep_quantity_a =
    Utility::QuantityTraits<InverseIndepQuantityA>::initializeQuantity( 0.0 );
  DepQuantityA dep_quantity_a =
    Utility::QuantityTraits<DepQuantityA>::initializeQuantity( 0.0 );

  IndepQuantityB indep_quantity_b( indep_quantity_a );
  InverseIndepQuantityB inv_indep_quantity_b( inv_indep_quantity_a );
  DepQuantityB dep_quantity_b( dep_quantity_a );

  UTILITY_TEST_FLOATING_EQUALITY(
			   unit_aware_dist_a_copy.evaluate( indep_quantity_a ),
			   dep_quantity_a,
			   1e-15 );
  UTILITY_TEST_FLOATING_EQUALITY(
			unit_aware_dist_a_copy.evaluatePDF( indep_quantity_a ),
			inv_indep_quantity_a,
			1e-15 );
  UTILITY_TEST_FLOATING_EQUALITY(
			   unit_aware_dist_b_copy.evaluate( indep_quantity_b ),
			   dep_quantity_b,
			   1e-15 );
  UTILITY_TEST_FLOATING_EQUALITY(
			unit_aware_dist_b_copy.evaluatePDF( indep_quantity_b ),
			inv_indep_quantity_b,
			1e-15 );

  Utility::setQuantity( indep_quantity_a, 1.0 );
  Utility::setQuantity( inv_indep_quantity_a,
			exp(-1.0)*sinh(1.0)/
			(0.25*sqrt( Utility::PhysicalConstants::pi )*
			 exp( 0.25 )*( erf( sqrt(0.9) - sqrt(0.25) ) +
				       erf( sqrt(0.9) + sqrt(0.25) ) ) -
			 exp( -0.9 ) * sinh( sqrt(0.9) )) );
  Utility::setQuantity( dep_quantity_a,
			exp(-1.0)*sinh(1.0) );

  indep_quantity_b = IndepQuantityB( indep_quantity_a );
  inv_indep_quantity_b = InverseIndepQuantityB( inv_indep_quantity_a );
  dep_quantity_b = DepQuantityB( dep_quantity_a );

  UTILITY_TEST_FLOATING_EQUALITY(
			   unit_aware_dist_a_copy.evaluate( indep_quantity_a ),
			   dep_quantity_a,
			   1e-15 );
  UTILITY_TEST_FLOATING_EQUALITY(
			unit_aware_dist_a_copy.evaluatePDF( indep_quantity_a ),
			inv_indep_quantity_a,
			1e-15 );
  UTILITY_TEST_FLOATING_EQUALITY(
			   unit_aware_dist_b_copy.evaluate( indep_quantity_b ),
			   dep_quantity_b,
			   1e-15 );
  UTILITY_TEST_FLOATING_EQUALITY(
			unit_aware_dist_b_copy.evaluatePDF( indep_quantity_b ),
			inv_indep_quantity_b,
			1e-15 );
}

typedef si::energy si_energy;
typedef cgs::energy cgs_energy;
typedef si::amount si_amount;
typedef si::length si_length;
typedef cgs::length cgs_length;
typedef si::mass si_mass;
typedef cgs::mass cgs_mass;
typedef si::dimensionless si_dimensionless;
typedef cgs::dimensionless cgs_dimensionless;

TEUCHOS_UNIT_TEST_TEMPLATE_4_INSTANT( UnitAwareWattDistribution,
				      explicit_conversion,
				      si_energy,
				      si_amount,
				      cgs_energy,
				      si_amount );
TEUCHOS_UNIT_TEST_TEMPLATE_4_INSTANT( UnitAwareWattDistribution,
				      explicit_conversion,
				      cgs_energy,
				      si_amount,
				      si_energy,
				      si_amount );
TEUCHOS_UNIT_TEST_TEMPLATE_4_INSTANT( UnitAwareWattDistribution,
				      explicit_conversion,
				      si_energy,
				      si_length,
				      cgs_energy,
				      cgs_length );
TEUCHOS_UNIT_TEST_TEMPLATE_4_INSTANT( UnitAwareWattDistribution,
				      explicit_conversion,
				      cgs_energy,
				      cgs_length,
				      si_energy,
				      si_length );
TEUCHOS_UNIT_TEST_TEMPLATE_4_INSTANT( UnitAwareWattDistribution,
				      explicit_conversion,
				      si_energy,
				      si_mass,
				      cgs_energy,
				      cgs_mass );
TEUCHOS_UNIT_TEST_TEMPLATE_4_INSTANT( UnitAwareWattDistribution,
				      explicit_conversion,
				      cgs_energy,
				      cgs_mass,
				      si_energy,
				      si_mass );
TEUCHOS_UNIT_TEST_TEMPLATE_4_INSTANT( UnitAwareWattDistribution,
				      explicit_conversion,
				      si_energy,
				      si_dimensionless,
				      cgs_energy,
				      cgs_dimensionless );
TEUCHOS_UNIT_TEST_TEMPLATE_4_INSTANT( UnitAwareWattDistribution,
				      explicit_conversion,
				      cgs_energy,
				      cgs_dimensionless,
				      si_energy,
				      si_dimensionless );
TEUCHOS_UNIT_TEST_TEMPLATE_4_INSTANT( UnitAwareWattDistribution,
				      explicit_conversion,
				      si_energy,
				      void,
				      cgs_energy,
				      void );
TEUCHOS_UNIT_TEST_TEMPLATE_4_INSTANT( UnitAwareWattDistribution,
				      explicit_conversion,
				      cgs_energy,
				      void,
				      si_energy,
				      void );
TEUCHOS_UNIT_TEST_TEMPLATE_4_INSTANT( UnitAwareWattDistribution,
				      explicit_conversion,
				      ElectronVolt,
				      si_amount,
				      si_energy,
				      si_amount );
TEUCHOS_UNIT_TEST_TEMPLATE_4_INSTANT( UnitAwareWattDistribution,
				      explicit_conversion,
				      ElectronVolt,
				      si_amount,
				      cgs_energy,
				      si_amount );
TEUCHOS_UNIT_TEST_TEMPLATE_4_INSTANT( UnitAwareWattDistribution,
				      explicit_conversion,
				      ElectronVolt,
				      si_amount,
				      KiloElectronVolt,
				      si_amount );
TEUCHOS_UNIT_TEST_TEMPLATE_4_INSTANT( UnitAwareWattDistribution,
				      explicit_conversion,
				      ElectronVolt,
				      si_amount,
				      MegaElectronVolt,
				      si_amount );
TEUCHOS_UNIT_TEST_TEMPLATE_4_INSTANT( UnitAwareWattDistribution,
				      explicit_conversion,
				      KiloElectronVolt,
				      si_amount,
				      si_energy,
				      si_amount );
TEUCHOS_UNIT_TEST_TEMPLATE_4_INSTANT( UnitAwareWattDistribution,
				      explicit_conversion,
				      KiloElectronVolt,
				      si_amount,
				      cgs_energy,
				      si_amount );
TEUCHOS_UNIT_TEST_TEMPLATE_4_INSTANT( UnitAwareWattDistribution,
				      explicit_conversion,
				      KiloElectronVolt,
				      si_amount,
				      ElectronVolt,
				      si_amount );
TEUCHOS_UNIT_TEST_TEMPLATE_4_INSTANT( UnitAwareWattDistribution,
				      explicit_conversion,
				      KiloElectronVolt,
				      si_amount,
				      MegaElectronVolt,
				      si_amount );
TEUCHOS_UNIT_TEST_TEMPLATE_4_INSTANT( UnitAwareWattDistribution,
				      explicit_conversion,
				      MegaElectronVolt,
				      si_amount,
				      si_energy,
				      si_amount );
TEUCHOS_UNIT_TEST_TEMPLATE_4_INSTANT( UnitAwareWattDistribution,
				      explicit_conversion,
				      MegaElectronVolt,
				      si_amount,
				      cgs_energy,
				      si_amount );
TEUCHOS_UNIT_TEST_TEMPLATE_4_INSTANT( UnitAwareWattDistribution,
				      explicit_conversion,
				      MegaElectronVolt,
				      si_amount,
				      ElectronVolt,
				      si_amount );
TEUCHOS_UNIT_TEST_TEMPLATE_4_INSTANT( UnitAwareWattDistribution,
				      explicit_conversion,
				      MegaElectronVolt,
				      si_amount,
				      KiloElectronVolt,
				      si_amount );
TEUCHOS_UNIT_TEST_TEMPLATE_4_INSTANT( UnitAwareWattDistribution,
				      explicit_conversion,
				      void,
				      MegaElectronVolt,
				      void,
				      KiloElectronVolt );

//---------------------------------------------------------------------------//
// Custom main function
//---------------------------------------------------------------------------//
int main( int argc, char** argv )
{
  std::string test_dists_xml_file;

  Teuchos::CommandLineProcessor& clp = Teuchos::UnitTestRepository::getCLP();

  clp.setOption( "test_dists_xml_file",
		 &test_dists_xml_file,
		 "Test distributions xml file name" );

  const Teuchos::RCP<Teuchos::FancyOStream> out =
    Teuchos::VerboseObjectBase::getDefaultOStream();

  Teuchos::CommandLineProcessor::EParseCommandLineReturn parse_return =
    clp.parse(argc,argv);

  if ( parse_return != Teuchos::CommandLineProcessor::PARSE_SUCCESSFUL ) {
    *out << "\nEnd Result: TEST FAILED" << std::endl;
    return parse_return;
  }

  TEUCHOS_ADD_TYPE_CONVERTER( Utility::WattDistribution );
  typedef Utility::UnitAwareWattDistribution<MegaElectronVolt,si::amount> UnitAwareWattDistribution;
  TEUCHOS_ADD_TYPE_CONVERTER( UnitAwareWattDistribution );

  test_dists_list = Teuchos::getParametersFromXmlFile( test_dists_xml_file );

  // Initialize the random number generator
  Utility::RandomNumberGenerator::createStreams();

  // Run the unit tests
  Teuchos::GlobalMPISession mpiSession( &argc, &argv );

  const bool success = Teuchos::UnitTestRepository::runUnitTests(*out);

  if (success)
    *out << "\nEnd Result: TEST PASSED" << std::endl;
  else
    *out << "\nEnd Result: TEST FAILED" << std::endl;

  clp.printFinalTimerSummary(out.ptr());

  return (success ? 0 : 1);
}

//---------------------------------------------------------------------------//
// end tstWattDistribution.cpp
//---------------------------------------------------------------------------//


#ifndef STAN_MATH_PRIM_SCAL_FUN_FMAX_HPP
#define STAN_MATH_PRIM_SCAL_FUN_FMAX_HPP

#include <stan/math/prim/meta.hpp>
#include <stan/math/prim/scal/fun/is_nan.hpp>
#include <boost/math/tools/promotion.hpp>

namespace stan {
namespace math {

/**
 * Return the greater of the two specified arguments.  If one is
 * greater than the other, return not-a-number.
 *
 * @param x First argument.
 * @param y Second argument.
 * @return maximum of x or y and if one is NaN return the other
 */
template <typename T1, typename T2>
inline typename boost::math::tools::promote_args<T1, T2>::type fmax(
    const T1& x, const T2& y) {
  if (is_nan(x))
    return y;
  if (is_nan(y))
    return x;
  return x > y ? x : y;
}

}  // namespace math
}  // namespace stan
#endif


///////////////////////////////////////////////////////////////////////////////
// BSD 3-Clause License
//
// Copyright (C) 2018-2020, LAAS-CNRS, The University of Edinburgh
// Copyright note valid unless otherwise stated in individual files.
// All rights reserved.
///////////////////////////////////////////////////////////////////////////////

#ifndef CROCODDYL_CORE_DIFF_ACTION_BASE_HPP_
#define CROCODDYL_CORE_DIFF_ACTION_BASE_HPP_

#include <stdexcept>
#include <boost/shared_ptr.hpp>
#include <boost/make_shared.hpp>
#include "crocoddyl/core/state-base.hpp"
#include "crocoddyl/core/utils/to-string.hpp"

namespace crocoddyl {

struct DifferentialActionDataAbstract;  // forward declaration

/**
 * @brief This class DifferentialActionModelAbstract represents a first-order
 * ODE, i.e.
 * \f[
 * \mathbf{\dot{v}} = \mathbf{f}(\mathbf{q}, \mathbf{v}, \boldsymbol{\tau})
 * \f]
 * where \f$ xout = \mathbf{\dot{v}} \f$ and represents the  acceleration of the
 * system. Note that Jacobians Fx and Fu in the
 * DifferentialActionDataAbstract are in \f$ \mathbb{R}^{nv\times ndx} \f$ and
 * \f$ \mathbb{R}^{nv\times nu} \f$, respectively.
 *
 * Then we use the acceleration to integrate the system, and as consequence we
 * obtain:
 * \f[
 * \mathbf{\dot{x}} = (\mathbf{v}, \mathbf{\dot{v}}) = \mathbf{f}(\mathbf{x},\mathbf{u})
 * \f]
 * where this \f$ f \f$ function is different to the other one.
 * So \f$ xout \f$ is interpreted here as \f$ vdout \f$ or \f$ aout \f$.
 */
class DifferentialActionModelAbstract {
 public:
  DifferentialActionModelAbstract(boost::shared_ptr<StateAbstract> state, const std::size_t& nu,
                                  const std::size_t& nr = 0);
  virtual ~DifferentialActionModelAbstract();

  virtual void calc(const boost::shared_ptr<DifferentialActionDataAbstract>& data,
                    const Eigen::Ref<const Eigen::VectorXd>& x, const Eigen::Ref<const Eigen::VectorXd>& u) = 0;
  virtual void calcDiff(const boost::shared_ptr<DifferentialActionDataAbstract>& data,
                        const Eigen::Ref<const Eigen::VectorXd>& x, const Eigen::Ref<const Eigen::VectorXd>& u,
                        const bool& recalc = true) = 0;
  virtual boost::shared_ptr<DifferentialActionDataAbstract> createData();

  void calc(const boost::shared_ptr<DifferentialActionDataAbstract>& data, const Eigen::Ref<const Eigen::VectorXd>& x);
  void calcDiff(const boost::shared_ptr<DifferentialActionDataAbstract>& data,
                const Eigen::Ref<const Eigen::VectorXd>& x);

  const std::size_t& get_nu() const;
  const std::size_t& get_nr() const;
  const boost::shared_ptr<StateAbstract>& get_state() const;

  const Eigen::VectorXd& get_u_lb() const;
  const Eigen::VectorXd& get_u_ub() const;
  bool const& get_has_control_limits() const;

  void set_u_lb(const Eigen::VectorXd& u_lb);
  void set_u_ub(const Eigen::VectorXd& u_ub);

 protected:
  std::size_t nu_;                          //!< Control dimension
  std::size_t nr_;                          //!< Dimension of the cost residual
  boost::shared_ptr<StateAbstract> state_;  //!< Model of the state
  Eigen::VectorXd unone_;                   //!< Neutral state
  Eigen::VectorXd u_lb_;                    //!< Lower control limits
  Eigen::VectorXd u_ub_;                    //!< Upper control limits
  bool has_control_limits_;                 //!< Indicates whether any of the control limits is finite

  void update_has_control_limits();

#ifdef PYTHON_BINDINGS

 public:
  void calc_wrap(const boost::shared_ptr<DifferentialActionDataAbstract>& data, const Eigen::VectorXd& x,
                 const Eigen::VectorXd& u = Eigen::VectorXd()) {
    if (u.size() == 0) {
      calc(data, x);
    } else {
      calc(data, x, u);
    }
  }

  void calcDiff_wrap(const boost::shared_ptr<DifferentialActionDataAbstract>& data, const Eigen::VectorXd& x,
                     const Eigen::VectorXd& u, const bool& recalc) {
    calcDiff(data, x, u, recalc);
  }
  void calcDiff_wrap(const boost::shared_ptr<DifferentialActionDataAbstract>& data, const Eigen::VectorXd& x,
                     const Eigen::VectorXd& u) {
    calcDiff(data, x, u, true);
  }
  void calcDiff_wrap(const boost::shared_ptr<DifferentialActionDataAbstract>& data, const Eigen::VectorXd& x) {
    calcDiff(data, x, unone_, true);
  }
  void calcDiff_wrap(const boost::shared_ptr<DifferentialActionDataAbstract>& data, const Eigen::VectorXd& x,
                     const bool& recalc) {
    calcDiff(data, x, unone_, recalc);
  }

#endif
};

struct DifferentialActionDataAbstract {
  EIGEN_MAKE_ALIGNED_OPERATOR_NEW

  template <typename Model>
  explicit DifferentialActionDataAbstract(Model* const model)
      : cost(0.),
        xout(model->get_state()->get_nv()),
        Fx(model->get_state()->get_nv(), model->get_state()->get_ndx()),
        Fu(model->get_state()->get_nv(), model->get_nu()),
        r(model->get_nr()),
        Lx(model->get_state()->get_ndx()),
        Lu(model->get_nu()),
        Lxx(model->get_state()->get_ndx(), model->get_state()->get_ndx()),
        Lxu(model->get_state()->get_ndx(), model->get_nu()),
        Luu(model->get_nu(), model->get_nu()) {
    xout.setZero();
    r.setZero();
    Fx.setZero();
    Fu.setZero();
    Lx.setZero();
    Lu.setZero();
    Lxx.setZero();
    Lxu.setZero();
    Luu.setZero();
  }
  virtual ~DifferentialActionDataAbstract() {}

  double cost;
  Eigen::VectorXd xout;
  Eigen::MatrixXd Fx;
  Eigen::MatrixXd Fu;
  Eigen::VectorXd r;
  Eigen::VectorXd Lx;
  Eigen::VectorXd Lu;
  Eigen::MatrixXd Lxx;
  Eigen::MatrixXd Lxu;
  Eigen::MatrixXd Luu;
};

}  // namespace crocoddyl

#endif  // CROCODDYL_CORE_DIFF_ACTION_BASE_HPP_


#include "stdafx.h"

#include <cmath>

#include <boost/geometry.hpp>
#include <boost/geometry/geometries/point_xy.hpp>
#include <boost/geometry/geometries/polygon.hpp>
#include <fmt/format.h>
#include <nlohmann/json.hpp>

#include "contest_types.h"
#include "judge.h"
#include "solver_registry.h"
#include "visual_editor.h"

namespace OptunaAnnealingSolver {

  namespace bg = boost::geometry;
  using BoostPoint = bg::model::d2::point_xy<double>;
  using BoostPolygon = bg::model::polygon<BoostPoint>;
  using BoostLinestring = bg::model::linestring<BoostPoint>;

  template <typename T>
  void shrink(T& point_a, T& point_b) {
    constexpr double eps = 1e-6;
    auto shrink_a = point_a * (1.0 - eps) + point_b * eps;
    auto shrink_b = point_b * (1.0 - eps) + point_a * eps;
    point_a = shrink_a;
    point_b = shrink_b;
  }

  template <typename T>
  BoostPoint ToBoostPoint(const T& point) {
    const auto [x, y] = point;
    return BoostPoint(x, y);
  }

  template <typename T>
  BoostPolygon ToBoostPolygon(const std::vector<T>& points) {
    BoostPolygon polygon;
    for (std::size_t i = 0; i <= points.size(); ++i) {
      polygon.outer().push_back(ToBoostPoint(points[i % points.size()]));
    }
    if (bg::area(polygon) < 0.0) {
      bg::reverse(polygon);
    }
    return polygon;
  }

  template <typename T, typename U>
  double SquaredDistance(const T& vertex0, const U& vertex1) {
    const auto [x0, y0] = vertex0;
    const auto [x1, y1] = vertex1;
    return (x0 - x1) * (x0 - x1) + (y0 - y1) * (y0 - y1);
  }

  template <typename T>
  double SquaredEdgeLength(const T& vertices, const Edge& edge) {
    const auto [a, b] = edge;
    return SquaredDistance(vertices[a], vertices[b]);
  }

  class Solver : public SolverBase {
  public:
    SolverOutputs solve(const SolverArguments& args) override {
      hole_ = args.problem->hole_polygon;
      vertices_ = args.problem->vertices;
      edges_ = args.problem->edges;
      epsilon_ = args.problem->epsilon;
      hole_polygon_ = ToBoostPolygon(hole_);

      // パラメーターファイルを読み込む
      std::ifstream ifs;
      // 実行バイナリと同じフォルダから読み込む
      ifs.open(args.parameters_file_path);
      CHECK(ifs);
      ifs >> parameters_;

      SVisualEditorPtr editor;
      if (args.visualize) {
        editor = std::make_shared<SVisualEditor>(args.problem, "OptunaAnnealingSolver", "visualize");
      }

      const int N = vertices_.size();
      auto pose = vertices_;
      double cost = std::numeric_limits<double>::infinity();
      double best_feasible_cost = std::numeric_limits<double>::infinity();
      std::vector<Point> best_feasible_pose;

      const int num_iters = 100000;
      const double T0 = parameters_["T0"]; // 10.0 1.0-100.0 log
      const double T1 = parameters_["T1"]; // 0.01 0.001-0.1 log
      double progress = 0.0;

      integer ymin = INT_MAX, ymax = INT_MIN;
      integer xmin = INT_MAX, xmax = INT_MIN;
      for (auto p : hole_) {
        xmin = std::min(xmin, get_x(p));
        ymin = std::min(ymin, get_y(p));
        xmax = std::max(xmax, get_x(p));
        ymax = std::max(ymax, get_y(p));
      }

      // lesser version of tonagi's idea
      const int initialize_pose_by_hole = parameters_["initialize_pose_by_hole"]; // 0 0-1
      if (initialize_pose_by_hole) {
        for (auto& p : pose) { p = hole_[0]; }
      }

      const int prohibit_unfeasible_after_feasible = parameters_["prohibit_unfeasible_after_feasible"]; // 0 0-1
      auto evaluate_and_descide_rollback = [&]() -> bool {
        auto [feasible, updated_cost] = Evaluate(pose);

        // tonagi's idea.
        if (prohibit_unfeasible_after_feasible || !best_feasible_pose.empty()) {
          auto res = judge(*args.problem, pose);
          if (!res.fit_in_hole()) {
            feasible = false;
            updated_cost = DBL_MAX;
          }
        }

#if 0
        auto judge_valid = judge(*args.problem, pose).is_valid();
        if (feasible != judge_valid) {
          LOG(INFO) << feasible << " " << judge_valid;
          if (editor) {
            editor->set_pose(args.problem->create_solution(pose));
            while (true) {
              int c = editor->show(1);
              if (c == 27) break;
            }
          }
        }
#endif

        if (feasible && updated_cost < best_feasible_cost) {
          best_feasible_cost = updated_cost;
          best_feasible_pose = pose;
          if (editor) editor->set_persistent_custom_stat(fmt::format("best_cost = {}", best_feasible_cost));
        }
        const double T = std::pow(T0, 1.0 - progress) * std::pow(T1, progress);
        if (std::uniform_real_distribution(0.0, 1.0)(rng_) < std::exp(-(updated_cost - cost) / T)) {
          cost = updated_cost;
          return false; // accepted
        }
        else {
          return true; // rejected
        }
      };
      evaluate_and_descide_rollback();

      const int single_small_change_max_delta = parameters_["single_small_change_max_delta"]; // 1 1-5
      auto single_small_change = [&] { // ynasu87 original
        const int v = std::uniform_int_distribution(0, N - 1)(rng_);
        const int dx = std::uniform_int_distribution(-single_small_change_max_delta, single_small_change_max_delta)(rng_);
        const int dy = std::uniform_int_distribution(-single_small_change_max_delta, single_small_change_max_delta)(rng_);
        auto& [x, y] = pose[v];
        x += dx;
        y += dy;
        if (evaluate_and_descide_rollback()) {
          x -= dx;
          y -= dy;
        }
      };

      const int shift_max_delta = parameters_["shift_max_delta"]; // 1 1-5
      auto shift = [&] {
        const int dx = std::uniform_int_distribution(-shift_max_delta, shift_max_delta)(rng_);
        const int dy = std::uniform_int_distribution(-shift_max_delta, shift_max_delta)(rng_);
        auto pose_bak = pose;
        for (auto& p : pose) {
          p.first += dx;
          p.second += dy;
        }
        if (evaluate_and_descide_rollback()) {
          pose = pose_bak;
        }
      };

      const double slight_rotate_max_deg = parameters_["slight_rotate_max_deg"]; // 2.0 1.0-180.0
      auto slight_rotate = [&] {
        const double deg = std::uniform_real_distribution(-slight_rotate_max_deg, slight_rotate_max_deg)(rng_);
        auto pose_bak = pose;

        integer curr_ymin = INT_MAX, curr_ymax = INT_MIN;
        integer curr_xmin = INT_MAX, curr_xmax = INT_MIN;
        for (auto p : hole_) {
          curr_xmin = std::min(curr_xmin, get_x(p));
          curr_ymin = std::min(curr_ymin, get_y(p));
          curr_xmax = std::max(curr_xmax, get_x(p));
          curr_ymax = std::max(curr_ymax, get_y(p));
        }

        const double cx = double(curr_xmin + curr_xmax) / 2;
        const double cy = double(curr_ymin + curr_ymax) / 2;
        const double sin = std::sin(deg * 3.1415 / 180.0);
        const double cos = std::cos(deg * 3.1415 / 180.0);
        for (auto& p : pose) {
          const double dx = p.first - cx;
          const double dy = p.second - cy;
          p.first = std::round(dx * cos - dy * sin + cx);
          p.second = std::round(dx * sin + dy * cos + cy);
        }
        if (evaluate_and_descide_rollback()) {
          pose = pose_bak;
        }
      };

      auto flip = [&] { // from FlipAnnealingSolver
        const int v0 = std::uniform_int_distribution(0, N - 1)(rng_);
        const int v1 = std::uniform_int_distribution(0, N - 1)(rng_);
        if (v0 == v1) return;
        auto pose_bak = pose;
        auto reflect = [](Point a, Point c, Point v) {
          return v - 2 * double(dot(v - c, a)) / double(dot(a, a)) * a;
        };
        auto diff = pose[v1] - pose[v0];
        Point n = { get_y(diff), -get_x(diff) };
        for (int v2 = 0; v2 < pose.size(); ++v2) {
          if (ccw(pose[v0], pose[v1], pose[v2])) {
            pose[v2] = reflect(n, pose[v0], pose[v2]);
          }
        }
        if (evaluate_and_descide_rollback()) {
          pose = pose_bak;
        }
      };

      const double vote_pow = parameters_["vote_pow"]; // 5.0 1.0-5.0
      auto edges_cache = edges_from_vertex(*args.problem);
      std::vector<std::vector<int> > good_pos(ymax - ymin + 1, std::vector<int>(xmax - xmin + 1));
      std::vector<double> pow_table;
      for (int i = 0; i < 1024; ++i) {
        pow_table.push_back(std::pow(static_cast<double>(i), vote_pow));
      }

      auto hop_grid = [&] { // jump to a tolerated (by at least one edge) point.
        const int pivot = std::uniform_int_distribution(0, N - 1)(rng_);
        const auto pivot_bak = pose[pivot];
        const auto& edges = edges_cache[pivot];
        for (auto& row : good_pos) {
          fill(row.begin(), row.end(), 0);
        }
        for (auto eid : edges) {
          auto [u, v] = args.problem->edges[eid];
          const int counter_vid = u == pivot ? v : u;
          const auto org_d2 = distance2(args.problem->vertices[pivot], args.problem->vertices[counter_vid]);
          for (int y = ymin; y <= ymax; ++y) {
            for (int x = xmin; x <= xmax; ++x) {
              const auto moved_d2 = distance2({ x, y }, pose[counter_vid]);
              if (tolerate(org_d2, moved_d2, epsilon_)) {
                ++good_pos[y - ymin][x - xmin];
              }
            }
          }
        }
        // emphasize large votes.
        double total_votes = 0;
        for (int y = ymin; y <= ymax; ++y) {
          for (int x = xmin; x <= xmax; ++x) {
            total_votes += pow_table[good_pos[y - ymin][x - xmin]];
          }
        }
        const double select_accum_vote = std::uniform_real_distribution<double>(0.0, total_votes)(rng_);
        double accum_vote = 0;
        bool found = false;
        for (int y = ymin; !found && y <= ymax; ++y) {
          for (int x = xmin; !found && x <= xmax; ++x) {
            accum_vote += pow_table[good_pos[y - ymin][x - xmin]];
            if (select_accum_vote <= accum_vote) {
              pose[pivot] = {x, y};
              found = true;
            }
          }
        }
        if (evaluate_and_descide_rollback()) {
          pose[pivot] = pivot_bak;
        }
      };

      // はみ出している線分を移動させる
      const int slide_protrusion_max_delta = parameters_["slide_protrusion_max_delta"]; // 1 1-5
      auto slide_protrusion = [&] { // from OptunaAnnealingSolver
        bool found = false;
        int vertex_index_backup[2];
        Point vertex_backup[2];
        for (const auto& edge : edges_) {
          const auto [a, b] = edge;
          Point2d pa = pose[a];
          Point2d pb = pose[b];
          shrink(pa, pb);
          const auto boost_point_a = ToBoostPoint(pa);
          const auto boost_point_b = ToBoostPoint(pb);
          BoostLinestring linestring{ boost_point_a, boost_point_b };
          std::vector<BoostLinestring> differences;
          bg::difference(linestring, hole_polygon_, differences);
          if (differences.empty()) {
            continue;
          }

          found = true;
          vertex_index_backup[0] = a;
          vertex_index_backup[1] = b;
          vertex_backup[0] = pose[a];
          vertex_backup[1] = pose[b];

          pose[a].first += std::uniform_int_distribution(-slide_protrusion_max_delta, slide_protrusion_max_delta)(rng_);
          pose[a].second += std::uniform_int_distribution(-slide_protrusion_max_delta, slide_protrusion_max_delta)(rng_);
          pose[b].first += std::uniform_int_distribution(-slide_protrusion_max_delta, slide_protrusion_max_delta)(rng_);
          pose[b].second += std::uniform_int_distribution(-slide_protrusion_max_delta, slide_protrusion_max_delta)(rng_);

          break;
        }

        if (found) {
          if (evaluate_and_descide_rollback()) {
            pose[vertex_index_backup[0]] = vertex_backup[0];
            pose[vertex_index_backup[1]] = vertex_backup[1];
          }
        }
      };

      using Action = std::function<void()>;
      const double single_small_change_probability = parameters_["single_small_change_probability"]; // 0.9 0.0-1.0
      const double slight_rotate_probability = parameters_["slight_rotate_probability"]; // 0.01 0.0-1.0
      const double shift_probability = parameters_["shift_probability"]; // 0.01 0.0-1.0
      const double hop_grid_probability = parameters_["hop_grid_probability"]; // 0.01 0.0-1.0
      const double flip_probability = parameters_["flip_probability"]; // 0.03 0.0-1.0
      const double slide_protrusion_probability = parameters_["slide_protrusion_probability"]; // 0.10 0.0-1.0
      std::vector<std::pair<double, Action>> action_probs = {
        {single_small_change_probability, single_small_change},
        {slight_rotate_probability, slight_rotate},
        {shift_probability, shift},
        {hop_grid_probability, hop_grid},
        {flip_probability, flip},
        {slide_protrusion_probability, slide_protrusion},
      };
      {
        // normalize probs
        double p = 0.0;
        for (int i = 0; i < action_probs.size(); ++i) { p += action_probs[i].first; }
        for (int i = 0; i < action_probs.size(); ++i) { action_probs[i].first /= p; }
      }

      for (int iter = 0; iter < num_iters; ++iter) {
        progress = 1.0 * iter / num_iters;

        const double p_action = std::uniform_real_distribution(0.0, 1.0)(rng_);
        double p_accum = 0.0;
        for (int i = 0; i < action_probs.size(); ++i) {
          p_accum += action_probs[i].first;
          if (p_action < p_accum) {
            action_probs[i].second();
          }
        }

        if (editor && iter % 100 == 0) {
          editor->set_oneshot_custom_stat(fmt::format("iter = {}/{}", iter, num_iters));
          editor->set_pose(args.problem->create_solution(pose));
          if (auto show_result = editor->show(1); show_result.edit_result) {
            pose = show_result.edit_result->pose_after_edit->vertices;
          }
        }
      }

      SolverOutputs outputs;
      if (best_feasible_pose.empty()) {
        outputs.solution = args.problem->create_solution(pose);
      }
      else {
        outputs.solution = args.problem->create_solution(best_feasible_pose);
      }
      return outputs;
    }

    template <typename P>
    std::tuple<bool, double> Evaluate(const std::vector<P>& pose) const {
      const double protrusion_cost_coefficient0 = parameters_["protrusion_cost_coefficient0"]; // 1.0 0.01-100.0 log
      const double protrusion_cost_coefficient1 = parameters_["protrusion_cost_coefficient1"]; // 0.01 0.01-100.0 log
      const double deformation_cost_coefficient = parameters_["deformation_cost_coefficient"]; // 10.0 0.01-100.0 log
      const double dislikes_cost_coefficient = parameters_["dislikes_cost_coefficient"]; // 0.01 0.01-100.0 log

      double deformation_cost = 0.0;
      double protrusion_cost = 0.0;
      double dislikes_cost = 0.0;

      const double tolerance = epsilon_ / 1'000'000.0;
      for (const auto& vertex : pose) {
        protrusion_cost += protrusion_cost_coefficient0 * bg::distance(ToBoostPoint(vertex), hole_polygon_);
      }
      for (const auto& edge : edges_) {
        const auto [a, b] = edge;
        Point2d pa = pose[a];
        Point2d pb = pose[b];
        shrink(pa, pb);
        BoostLinestring linestring{ ToBoostPoint(pa), ToBoostPoint(pb) };
        std::vector<BoostLinestring> differences;
        bg::difference(linestring, hole_polygon_, differences);
        for (const auto& segment : differences) {
          protrusion_cost += protrusion_cost_coefficient1 * bg::length(segment);
        }

        const auto d0 = SquaredEdgeLength(vertices_, edge);
        const auto d1 = SquaredEdgeLength(pose, edge);
        deformation_cost += deformation_cost_coefficient * std::max(0.0, std::abs(d1 / d0 - 1.0) - tolerance);
      }

      for (const auto h : hole_) {
        double best = std::numeric_limits<double>::infinity();
        for (const auto v : pose) {
          best = std::min(best, SquaredDistance(h, v));
        }
        dislikes_cost += best * dislikes_cost_coefficient;
      }

      const bool feasible = deformation_cost + protrusion_cost == 0.0;
      const double cost = deformation_cost + protrusion_cost + dislikes_cost;

      return { feasible, cost };
    }

  private:
    std::mt19937 rng_;
    std::vector<Point> hole_;
    std::vector<Point> vertices_;
    std::vector<Edge> edges_;
    integer epsilon_;
    BoostPolygon hole_polygon_;
    nlohmann::json parameters_;
  };

}

REGISTER_SOLVER("OptunaAnnealingSolver", OptunaAnnealingSolver::Solver);
// vim:ts=2 sw=2 sts=2 et ci


//==================================================================================================
/**
  Copyright 2016 NumScale SAS

  Distributed under the Boost Software License, Version 1.0.
  (See accompanying file LICENSE.md or copy at http://boost.org/LICENSE_1_0.txt)
**/
//==================================================================================================
#ifndef BOOST_SIMD_ARCH_COMMON_SCALAR_FUNCTION_MEANOF_HPP_INCLUDED
#define BOOST_SIMD_ARCH_COMMON_SCALAR_FUNCTION_MEANOF_HPP_INCLUDED

#include <boost/simd/constant/half.hpp>
#include <boost/simd/function/max.hpp>
#include <boost/simd/function/min.hpp>
#include <boost/simd/function/is_finite.hpp>
#include <boost/simd/detail/dispatch/function/overload.hpp>
#include <boost/config.hpp>

namespace boost { namespace simd { namespace ext
{
   namespace bd = boost::dispatch;
  BOOST_DISPATCH_OVERLOAD ( meanof_
                          , (typename A0)
                          , bd::cpu_
                          , bd::scalar_< bd::floating_<A0> >
                          , bd::scalar_< bd::floating_<A0> >
                          )
  {
    BOOST_FORCEINLINE A0 operator() ( A0 a0, A0 a1) const BOOST_NOEXCEPT
    {
      if (is_finite(a0)&&is_finite(a1))
      {
        A0 m = min(a0, a1);
        return m + (max(a0, a1)-m)*Half<A0>();
      }
      else
        return a0+a1;
    }
  };

  BOOST_DISPATCH_OVERLOAD ( meanof_
                          , (typename A0)
                          , bd::cpu_
                          , bd::scalar_< bd::integer_<A0>>
                          , bd::scalar_< bd::integer_<A0>>
                          )
  {
    BOOST_FORCEINLINE A0 operator() ( A0 a0, A0 a1) const BOOST_NOEXCEPT
    {
      return (a0 & a1) + ((a0 ^ a1) >> 1);
    }
  };

} } }

#endif


// smooth_feedback: Control theory on Lie groups
// https://github.com/pettni/smooth_feedback
//
// Licensed under the MIT License <http://opensource.org/licenses/MIT>.
//
// Copyright (c) 2021 Petter Nilsson
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.

#include <gtest/gtest.h>

#include <Eigen/Core>

#include "smooth/feedback/collocation/dyn_error.hpp"

TEST(CollocationDyn, DynError)
{
  // given trajectory
  const auto x = [](double t) -> Eigen::Vector<double, 1> {
    return Eigen::Vector<double, 1>{{0.1 * t * t - 0.4 * t + 0.2}};
  };

  // system dynamics
  const auto f =
    []<typename T>(const T & t, const Eigen::Vector<T, 1> &, const Eigen::Vector<T, 0> &)
    -> Eigen::Vector<T, 1> { return Eigen::Vector<T, 1>{{0.2 * t - 0.4}}; };

  double t0 = 3;
  double tf = 5;

  smooth::feedback::Mesh<5, 5> m;

  // trajectory is not a polynomial, so we need a couple of intervals for a good approximation
  m.refine_ph(0, 16 * 5);
  ASSERT_EQ(m.N_ivals(), 16);

  // fill X with curve values at the two intervals
  std::size_t M = 0;
  Eigen::MatrixXd X(1, m.N_colloc() + 1);
  for (auto p = 0u; p < m.N_ivals(); ++p) {
    for (const auto & [i, tau] :
         smooth::utils::zip(std::views::iota(0u, m.N_colloc_ival(p)), m.interval_nodes(p))) {
      X.col(M + i) = x(t0 + (tf - t0) * tau);
    }
    M += m.N_colloc_ival(p);
  }
  X.col(m.N_colloc()) = x(tf);

  auto xfun = [X = X, t0 = t0, tf = tf, m = m](const double t) -> Eigen::Vector<double, 1> {
    return m.eval<Eigen::Vector<double, 1>>((t - t0) / (tf - t0), X.colwise(), 0, true);
  };
  auto ufun = [](const double) -> Eigen::Vector<double, 0> {
    return Eigen::Vector<double, 0>::Zero();
  };

  m.increase_degrees();
  auto rel_errs = smooth::feedback::mesh_dyn_error(f, m, t0, tf, xfun, ufun);
  m.decrease_degrees();

  ASSERT_LE(rel_errs.cwiseAbs().maxCoeff(), 1e-8);

  const auto Npre = m.N_ivals();
  m.refine_errors(rel_errs, 1e-8);

  ASSERT_EQ(m.N_ivals(), Npre);
}


#include <iostream>
#include <vector>
#include <string>
#include "jefflib.h"
#include <boost/tokenizer.hpp>
#include <map>
#include <set>

using namespace std;
using namespace boost;

void processNode(vector<int>& data, int& index, int& checksum);

int main()
{
    vector<string> vect;
    if(GetStringInput(vect)){
        cout << "Got data!" << endl;
        cout << endl;
    }
    else {
        cout << "Failed to read input :( " << endl;
        return -1;
    }

    char_separator<char> sep(" ");
    tokenizer< char_separator<char> > tokens(vect[0], sep);
    vector<string> v_temp(tokens.begin(), tokens.end());

    vector<int> data;
    for(auto s : v_temp){
        data.push_back(stoi(s));
    }

    cout << "Data size: " << data.size() << endl;

    int checksum = 0;
    int index = 0;
    processNode(data, index, checksum);
    cout << "Checksum: " << checksum << endl;
}

void processNode(vector<int>& data, int& index, int& checksum){
    int numChildren = data[index];
    int numMeta = data[++index];

    for(int i = 0; i < numChildren; i++){
        processNode(data, ++index, checksum);
    }

    for(int i = 0; i < numMeta; i++){
        checksum += data[++index];
    }
}