#include "graphCommon.h"

#include <sstream>
#include <map>
#include <utility> // make_pair
#include <boost/program_options.hpp>

#include <boost/tokenizer.hpp>

#include <boost/lexical_cast.hpp>

/**
 * VisabilityGrapMap.
 */
class VGMap
{

public:

	typedef double ValueType;

	/**
	 * Convert all mean label signal to visability matrices.
	 */
	void run( const std::string& timeSeriesFileName, const std::string& outputRoot, const std::string& labelFileName,
			std::vector< int > labelRange, bool weighted, std::vector< unsigned int > ignore, bool verbose, const std::string& labelNames )
	{

		// [ 1 ]: dimension checks ...
		unsigned int dims = graph::Graph< ValueType >::GetImageDimensions( timeSeriesFileName );

		if ( dims != 4 )
		{
			std::cerr << "*** ERROR ***: File " << timeSeriesFileName << " is not 4D!\n" << std::endl;
			exit( EXIT_FAILURE );
		}
		dims = graph::Graph< ValueType >::GetImageDimensions( labelFileName );

		if ( dims != 3 )
		{
			std::cerr << "*** ERROR ***: File " << labelFileName << " is not 3D!\n" << std::endl;
			exit( EXIT_FAILURE );
		}

		// [ 2 ]: create visability matrix for each label ...
		process( timeSeriesFileName, outputRoot, labelFileName, labelRange, weighted, ignore, verbose, labelNames );
	}

protected:

	/**
	 *
	 */
	void process( const std::string& timeSeriesFileName, const std::string& outputRoot, const std::string& labelFileName,
			std::vector< int > labelRange, bool weighted, std::vector< unsigned int > ignore, bool verbose, const std::string& labelNames )
	{

		typedef itk::Image< ValueType, 4 > InputType;
		typedef itk::Image< ValueType, 4 > OutputType;
		typedef itk::Image< int, 3 > LabelType;

		typedef itk::ImageFileReader< InputType > InputReaderType;
		typedef itk::ImageFileReader< LabelType > LabelReaderType;
		typedef itk::ImageFileWriter< OutputType > WriterType;

		InputReaderType::Pointer inputReader = InputReaderType::New();
		inputReader->SetFileName( timeSeriesFileName.c_str() );
		inputReader->Update();

		InputType::Pointer input = inputReader->GetOutput();
		inputReader = 0;

		LabelReaderType::Pointer labelReader = LabelReaderType::New();
		labelReader->SetFileName( labelFileName.c_str() );
		labelReader->Update();

		LabelType::Pointer label = labelReader->GetOutput();
		labelReader = 0;

		std::map< int, int > labels;

		if ( labelRange.size() == 2 && labelRange[1] >= labelRange[0] )
		{
			int count = 0;
			for ( int i = labelRange[0]; i <= labelRange[1]; i++ )
			{
				labels[i] = count++;
			}
		} else
		{
			int count = 0;
			itk::ImageRegionIterator< LabelType > it( label, label->GetLargestPossibleRegion() );
			while ( !it.IsAtEnd() )
			{
				if ( it.Value() > 0 && labels.find( it.Value() ) == labels.end() )
				{
					labels[it.Value()] = count++;
				}
				++it;
			}
		}

		int numberOfLabels = labels.size();

		InputType::RegionType region = input->GetLargestPossibleRegion();
		InputType::SizeType size = region.GetSize();
		int voxels = size[0] * size[1] * size[2];
		int timepoints = size[3];

		vnl_matrix_ref< ValueType > matrix( timepoints, voxels, input->GetPixelContainer()->GetBufferPointer() );
		const int* mask = label->GetPixelContainer()->GetBufferPointer();

		vnl_matrix< ValueType > mean( timepoints, numberOfLabels );
		mean.fill( 0 );

		vnl_vector< int > count( numberOfLabels );

		count.fill( 0 );

		for ( int i = 0; i < voxels; ++i )
		{
			if ( mask[i] < 1 )
			{
				continue;
			}

			int column = labels[mask[i]];

			mean.set_column( column, mean.get_column( column ) + matrix.get_column( i ) );
			++count[column];
		}

		for ( int i = 0; i < numberOfLabels; ++i )
		{
			mean.set_column( i, mean.get_column( i ) / static_cast< ValueType > ( count( i ) ) );
		}

		// extract ignored samples ...
		int vectorLength = timepoints;
		int vectorStart = 0;

		if ( ignore.size() == 2 )
		{
			vectorStart = ignore[0];
			vectorLength = timepoints - vectorStart - ignore[1];
		}

		mean = mean.extract( vectorLength, numberOfLabels, vectorStart, 0 );

		graph::Graph< ValueType >::MapType labelMap;
		graph::Graph< ValueType >::GetLabelNames( labelMap, labelNames, "\t" );

		if ( verbose )
			std::cout << "Constructing visability graph for label: " << std::endl;

		for ( int i = 1; i <= numberOfLabels; i++ )
		{
			vnl_vector< ValueType > labelTimeSeries = mean.get_column( i );

			std::string name;
			graph::Graph< ValueType >::GetLabelName( outputRoot, ".nii.gz", labelMap, name, i );

			if ( verbose )
			{
				std::cout << i;
				std::cout.flush();
				std::cout << '\r';
			}

			graph::VisabilityGraph< ValueType >::WriteVisabilityMatrix( labelTimeSeries, name, false, weighted );
		}
	}
};

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

/**
 * Throw error if required option is not specified.
 */
void required_option( const boost::program_options::variables_map& vm, const std::string& required_option )
{
	if ( vm.count( required_option ) == 0 )
		throw std::logic_error( "Option: '" + required_option + "' is required!" );
}

/**
 * Visability graph mapping using labels.
 */
int main( int argc, char* argv[] )
{
	namespace po = boost::program_options;

	// application description ...
	std::string description = "Create visability graphs from 4D time-series using labels.\n";

	// options ...
	std::string input;
	std::string output;
	std::string labels;
	std::string labelNames;

	std::vector< int > labelRange;
	std::vector< unsigned int > ignore;

	bool weighted;
	bool verbose;

	try
	{
		po::options_description desc( "Available options" );

		desc.add_options()

		( "input,i", po::value< std::string >( &input ), "string: 4D image time-series input file." )

		( "output,o", po::value< std::string >( &output ), "string: output root." )

		( "labels,l", po::value< std::string >( &labels ), "string: 3D labels file." )

		( "label-names,n", po::value< std::string >( &labelNames ),
				"string: 3D labels file with id and name (column format: e.g. 3 S1FL le)." )

		( "weighted,w", po::value< bool >( &weighted ) ->default_value( false ) ->zero_tokens(),
				"bool: use weighted visability graph construction (slope)." )

		( "verbose,v", po::value< bool >( &verbose ) ->default_value( false ) ->zero_tokens(), "bool: print progress." )

		( "ignore,d", po::value< std::vector< unsigned int > >( &ignore )->multitoken(), "vector: samples to ignore at begin and end." )

		( "label-range,r", po::value< std::vector< int > >( &labelRange )->multitoken(), "vector: specify a label range." )

		( "help,h", "bool: produce help message." );

		po::variables_map vm;
		po::store( po::parse_command_line( argc, argv, desc ), vm );
		po::notify( vm );

		// help message ...
		if ( vm.count( "help" ) )
		{
			std::cout << argv[0] << ": " << description << std::endl;
			std::cout << desc << "\n";

			return EXIT_SUCCESS;
		}

		// required options ...
		required_option( vm, "input" );
		required_option( vm, "output" );
		required_option( vm, "labels" );

		// run application ...
		VGMap vgMap;

		vgMap.run( input, output, labels, labelRange, weighted, ignore, verbose, labelNames );

	} catch ( std::exception& e )
	{
		std::cerr << "*** ERROR ***: " << e.what() << "\n";
		std::cerr << "Use \"" << argv[0] << " --help\" for information about application usage." << std::endl;

		return EXIT_FAILURE;
	}

	return EXIT_SUCCESS;
}


//==================================================================================================
/*!
  @file

    @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_FUNCTION_EXP2_HPP_INCLUDED
#define BOOST_SIMD_FUNCTION_EXP2_HPP_INCLUDED

#if defined(DOXYGEN_ONLY)
namespace boost { namespace simd
{

  /*!
    @ingroup group-exponential
    This function object returns the base two exponential of its argument: \f$2^{x}\f$

    @par Header <boost/simd/function/exp2.hpp>

    @par Note:

    - provisions are made to obtain a flint result from a flint argument

    @par Decorators

      - std_ for floating entries calls @c std::exp2

    @see exp, exp10, pow, pow2

    @par Example:

      @snippet exp2.cpp exp2

    @par Possible output:

      @snippet exp2.txt exp2

  **/
  IEEEValue exp2(IEEEValue const& x);
} }
#endif

#include <boost/simd/function/scalar/exp2.hpp>
#include <boost/simd/function/simd/exp2.hpp>

#endif


#include <fstream>
#include <iostream>
#include <algorithm>

#include <Eigen/Dense>

#include <boost/program_options/options_description.hpp>
#include <boost/program_options/parsers.hpp>
#include <boost/program_options/variables_map.hpp>

namespace po = boost::program_options;

#include "csv.hpp"
#include "datatypes.hpp"
#include "interdataset.hpp"
#include "stats.hpp"
#include "cuda/cudasampler.hpp"
#include "utility.hpp"

#include <sys/time.h>

std::default_random_engine generator;

static const Eigen::IOFormat CSVFMT(Eigen::FullPrecision, Eigen::DontAlignCols, ",", "\n", "", "", "", "");

class csvwriter {
  bool _firstcol, _colbreak;
  std::ostream &_out;

  void nextvalue() {
    if(_firstcol) {
      assert(!_colbreak);
      _firstcol = false;
      return;
    }
    _out << ',';
    if(_colbreak) {
      _out << ' ';
      _colbreak = false;
    }
  }

  template<typename T>
  void writeraw(const T&x) {
    nextvalue();
    _out << x;
  }

public:
  csvwriter(std::ostream &out) : _firstcol(true), _colbreak(false), _out(out) {}

  void endrow() {
    // make sure we don't output an empty row
    assert(!_firstcol);
    _out << std::endl;
    _firstcol = true;
    _colbreak = false;
  }

  void columnbreak() {
    if (!_firstcol)
      _colbreak = true;
  }

  void write(short i) { writeraw(i); }
  void write(int   i) { writeraw(i); }
  void write(long  i) { writeraw(i); }

  void write(float  f) { writeraw(f); }
  void write(double f) { writeraw(f); }

  void write(const std::string &s) {
    nextvalue();

    _out << '"';
    for (std::string::const_iterator i = s.begin(), end = s.end(); i != end; ++i) {
      unsigned char c = *i;
      if (' ' <= c && c <= '~' && c != '\\' && c != '"') {
	_out << c;
      } else {
	_out << '\\';
	switch(c) {
	case '"':  _out << '"';  break;
	case '\\': _out << '\\'; break;
	case '\t': _out << 't';  break;
	case '\r': _out << 'r';  break;
	case '\n': _out << 'n';  break;
	default:
	  char const* const hexdig = "0123456789ABCDEF";
	  _out << 'x';
	  _out << hexdig[c >> 4];
	  _out << hexdig[c & 0xF];
	}
      }
    }
    _out << '"';
  }
};

void
printHeader(const shared& shared,
            const std::vector<std::string> &ss,
            csvwriter &out)
{
  std::stringstream val;
  for (int k = 0; k < shared.nfiles(); k++) {
    val.str("");
    val << "MassParameter_" << k+1;
    out.write(val.str());
  }
  out.columnbreak();
  for (int k = 0; k < shared.nfiles()-1; k++) {
    for (int l = k+1; l < shared.nfiles(); l++) {
      val.str("");
      val << "Phi_" << k+1 << l+1;
      out.write(val.str());
    }
  }
  out.columnbreak();
  for (int k = 0, a = 0; k < shared.nfiles(); k++) {
    for (int i = 0; i < shared.nitems(); i++) {
      val.str("");
      val << "Dataset" << k+1 << '_' << ss[a++];
      out.write(val.str());
    }
  }
  out.endrow();
}

void
printSample(const shared &shared, const mdisampler &mdi,
	    csvwriter &out)
{
  for (int k = 0; k < shared.nfiles(); k++) {
    out.write(mdi.dpmass(k));
  }
  out.columnbreak();
  for (int k = 0; k < shared.nfiles()-1; k++) {
    for (int l = k+1; l < shared.nfiles(); l++) {
      out.write(mdi.phi(k, l));
    }
  }
  for (int k = 0; k < shared.nfiles(); k++) {
    out.columnbreak();
    for (int i = 0; i < shared.nitems(); i++) {
      out.write(shared.alloc(k, i)+1);
    }
  }
  out.endrow();
}

void
printFeatureHeader(const std::vector<std::string> &ss,
                   csvwriter &out)
{
  for (size_t i = 0; i < ss.size(); i++) {
    out.write(ss[i]);
  }
  out.endrow();
}

void
printFeatureSample(const std::vector<std::unique_ptr<sampler> > &ss,
                   csvwriter &out)
{
  for (const auto &s : ss) {
    auto fs = s->featureState();
    out.columnbreak();
    for (int i = 0; i < fs.size(); i++) {
      out.write(fs[i]);
    }
  }
  out.endrow();
}

static gamma_distribution<> phiprior(1.0, 0.2), dpprior(2, 4);

double
sampleMdiNu (const mdisampler &mdi, const shared &shared)
{
  auto d = gamma_distribution<>(shared.nitems(), mdi.nuConditionalBeta());
  return d(generator);
}

double
sampleMdiPhi (const mdisampler &mdi, const shared &shared, int k, int l)
{
  auto d = gamma_distribution<>(phiprior.shape() + shared.phiConditionalAlpha(k, l),
				phiprior.rate() + mdi.phiConditionalBeta(k, l));
  return d(generator);
}

Eigen::MatrixXf
sampleMdiWeights(const mdisampler &mdi, const shared &shared)
{
  Eigen::MatrixXi clussize(Eigen::MatrixXi::Zero(shared.nfiles(), shared.nclus()));
  for (int i = 0; i < shared.nitems(); i++) {
    for (int k = 0; k < shared.nfiles(); k++) {
      clussize(k, shared.alloc(k, i)) += 1;
    }
  }

  Eigen::MatrixXf weight(mdi.nclus(),mdi.nfiles());

  // sample weights
  for (int k = 0; k < mdi.nfiles(); k++) {
    const gamma_distribution<> prior(mdi.dpmass(k) / mdi.nclus(), 1.0);
    for (int j = 0; j < mdi.nclus(); j++) {
      auto d = gamma_distribution<>(prior.shape() + shared.gammaConditionalAlpha(k,j),
				    prior.rate() + mdi.gammaConditionalBeta(k, j));
      weight(j,k) = d(generator);
    }
  }

  return weight;
}

Eigen::VectorXi
sampleSharedAlloc(const mdisampler &mdi, const shared &shared, int k, const sampler::item &is)
{
  Eigen::VectorXi alloc(shared.nitems());
  Eigen::VectorXf
    prob(shared.nclus()),
    lweight(mdi.weight(k).array().log().cast<float>()),
    lrho((mdi.phis(k).array().unaryExpr(std::ptr_fun(log1p)).cast<float>()));

  for (int i = 0; i < shared.nitems(); i++) {
    prob = lweight + is(i);

    for (int l = 0; l < shared.nfiles(); l++) {
      if (k == l) continue;
      prob(shared.alloc(l, i)) += lrho(l);
    }

    prob = (prob.array() - prob.maxCoeff()).exp();

    alloc(i) = std::discrete_distribution<>(prob.data(), prob.data() + prob.size()-1)(generator);
  }

  return alloc;
}

double
sampleDpMass (const mdisampler &mdi, int k)
{
  const double
    dp1 = mdi.dpmass(k),
    dp2 = std::abs(std::normal_distribution<>(dp1, 1)(generator));

  double ll1 = dpprior.logpdf(dp1), ll2 = dpprior.logpdf(dp2);

  const gamma_distribution<>
    d1(dp1 / mdi.nclus(), 1.0),
    d2(dp2 / mdi.nclus(), 1.0);
  for (int j = 0; j < mdi.nclus(); j++) {
    ll1 += d1.logpdf(mdi.weight(j,k));
    ll2 += d2.logpdf(mdi.weight(j,k));
  }

  return mcmcAcceptProposalLogProb(ll1, ll2) ?
    dp2 : dp1;
}

std::vector<int>
sampleAllocBlockSwap (const shared &shared, const mdisampler &mdi, int k)
{
  const Eigen::VectorXd lrho = mdi.phis(k).array().unaryExpr(std::ptr_fun(log1p));

  Eigen::VectorXi alloc  = shared.alloc(k);
  Eigen::VectorXf weight = mdi.weight(k);

  std::vector<int> changed(shared.nclus());
  for (int i = 0; i < shared.nclus(); i++) {
    changed[i] = i;
  }

  for (int c1 = 0; c1 < shared.nclus(); c1++) {
    int c2 = std::uniform_int_distribution<int>(0, (int)shared.nclus()-2)(generator);
    c2 += c2 >= c1;

    // for keeping track of the accept/reject ratio
    double llratio = 0;

    for (int i = 0; i < alloc.size(); i++) {
      if (alloc(i) == c1 || alloc(i) == c2) {
        for (int l = 0; l < shared.nfiles(); l++) {
          if (k == l) continue;
          if (shared.alloc(l,i) == c1) {
            llratio += alloc(i) == c2 ? -lrho(l) : lrho(l);
          }
          if (shared.alloc(l,i) == c2) {
            llratio += alloc(i) == c1 ? -lrho(l) : lrho(l);
          }
        }
      }
    }

    if (mcmcAcceptProposalLogProb(llratio)) {
      for (int i = 0; i < alloc.size(); i++) {
        if (alloc(i) == c1)
          alloc(i) = c2;
        else if(alloc(i) == c2) {
          alloc(i) = c1;
        }
      }
      std::swap(weight(c1),  weight(c2));
      std::swap(changed[c1], changed[c2]);
    }
  }

  return changed;
}

int
main(int argc, char **argv)
{
  // small pause to let the debugger attach
  usleep(10*1000);

  int nsamples, thinby, nclusters, nitems = -1;
  po::options_description genopt("General Options");
  genopt.add_options()
    ("help,h",                                                     "Display this help message")
    ("samples,n", po::value<int>(&nsamples)->default_value(10000), "Number of MCMC samples to generate")
    ("thin,t", po::value<int>(&thinby)->default_value(1),          "Thin output to one sample every 'arg' samples")
    ("clusters,c",po::value<int>(&nclusters)->default_value(50),   "Assume a maximum of 'arg' clusters")
    ("featuresel,f",po::value<std::string>(),                      "Perform feature selection")
#ifndef NOCUDA
    ("gpu,g",                                                      "Enable GPU support")
#endif
    ("benchmark",                                                  "Display wall-clock time taken to run computation")
    ("seed,s",po::value<unsigned>(),                               "Seed for Random Number Generation")
    ("version,v",                                                  "Display the version number")
    ;

  std::vector<std::string> inputs;
  po::options_description hiddenopt("Hidden options");
  hiddenopt.add_options()
    ("input-files", po::value<std::vector<std::string>>(&inputs), "datatypes and files")
    ;

  po::variables_map vm;
  try {
    po::options_description cmdopts;
    cmdopts.add(genopt).add(hiddenopt);
    po::positional_options_description p;
    p.add("input-files", -1);
    po::store(po::command_line_parser(argc, argv).options(cmdopts).positional(p).run(), vm);
    po::notify(vm);
  } catch (const std::logic_error &err) {
    std::cerr << "Error parsing command line: " << err.what() << '\n';
    return 1;
  }

  if(vm.count("help")) {
    po::options_description opts("MDI++ Command Line Help");
    opts.add(genopt);
    std::cout << opts << "\n"
      "  N  filename1 [filename2..n] : load files as independent Gaussian features\n"
      "  GP filename1 [filename2..n] : load files as Gaussian process data\n"
      "  M  filename1 [filename2..n] : load files as Multinomial features\n"
      "  BW filename1 [filename2..n] : load files as Bag-of-Words features \n";
    return 0;
  }

  if(vm.count("version")) {
    std::cout << "MDI++ version 0.2\n";
    return 0;
  }

  if (vm.count("seed")) {
    generator.seed(vm["seed"].as<unsigned>());
  } else {
    unsigned seed = (std::random_device())();
    std::cerr << "RNG Seed = " << seed << '\n';
    generator.seed(seed);
  }

  std::unique_ptr<std::ofstream> featurefile;
  std::unique_ptr<csvwriter> featuresel;
  if (vm.count("featuresel")) {
    featurefile.reset(new std::ofstream(vm["featuresel"].as<std::string>(), std::ofstream::out));
    featuresel.reset(new csvwriter(*featurefile));
  }

  std::vector<std::unique_ptr<datatype> > datasets;
  std::vector<std::string> itemnames, featurenames;

  enum {
    MODE_UNKNOWN,
    MODE_GAUSSIAN,
    MODE_GAUSSIANPROCESS,
    MODE_MULTINOMIAL,
    MODE_BAGOFWORDS
  } mode = MODE_UNKNOWN;
  for (auto const &name : inputs) {
    if (name == "N")  { mode = MODE_GAUSSIAN;        continue; }
    if (name == "GP") { mode = MODE_GAUSSIANPROCESS; continue; }
    if (name == "M")  { mode = MODE_MULTINOMIAL;     continue; }
    if (name == "BW") { mode = MODE_BAGOFWORDS;      continue; }

    std::unique_ptr<datatype> dst;
    switch(mode) {
    case MODE_UNKNOWN:
      std::cerr << "Error: Unknown datatype!\n";
      exit(1);
    case MODE_GAUSSIAN:        dst.reset(new gaussianDatatype  (name.c_str())); break;
    case MODE_GAUSSIANPROCESS: dst.reset(new gpDatatype        (name.c_str())); break;
    case MODE_MULTINOMIAL:     dst.reset(new multinomDatatype  (name.c_str())); break;
    case MODE_BAGOFWORDS:      dst.reset(new bagofwordsDatatype(name.c_str())); break;
    }
    int i = dst->items().size();
    if (datasets.size() == 0) {
      nitems = i;
    } else if (nitems != i) {
      std::cerr << "Error datafile " << name << " doesn't have the same number of items to cluster.\n";
      exit(1);
    }
    {
      auto n = dst->items();
      std::copy(n.begin(), n.end(), std::back_inserter(itemnames));
    }
    {
      auto n = dst->features();
      std::copy(n.begin(), n.end(), std::back_inserter(featurenames));
    }
    // finally(!) move here
    datasets.push_back(std::move(dst));
  }

  if (datasets.size() == 0) {
    std::cerr << "Error: Need to load at least one datafile.  Run with -h for help\n";
    exit(1);
  }

  const int nfiles = datasets.size();
  shared shared(nfiles, nclusters, nitems);
  const interdataset inter(nfiles);
  mdisampler mdi(inter, nclusters);

  shared.sampleFromPrior();
  mdi.sampleFromPrior();

  class cuda::sampler *gpu = 0;

  if (vm.count("gpu")) {
#ifdef NOCUDA
    std::cerr << "error: CUDA support not enabled in this build!";
    exit(1);
#else
    gpu = new cuda::sampler
      (nfiles, nitems, nclusters,
       inter.getPhiOrd(), inter.getWeightOrd());
    gpu->setIttr(generator());
    gpu->setNu(mdi.nu());
    gpu->setAlloc(eigenMatrixToStdVector(shared.allocs()));
    gpu->setPhis(eigenMatrixToStdVector(mdi.phis()));

    if (0) {
      std::cerr << mdi.weights() << "\n";
      std::vector<float> weights(eigenMatrixToStdVector(mdi.weights()));

      std::cerr << "weights: ";
      for (auto &x:weights) {
	std::cerr << x << ", ";
      }
      std::cerr << "\n";
    }

    gpu->setWeights(eigenMatrixToStdVector(mdi.weights()));
    gpu->setDpMass(eigenMatrixToStdVector(mdi.dpmass()));

#endif
  }

  std::vector<std::unique_ptr<sampler> > samplers;
  for (auto &dst : datasets) {
    samplers.emplace_back(dst->newSampler(nclusters, gpu));
    // TODO: add an explicit "sample from prior" call in here
  }

  csvwriter mcmcwriter(std::cout);
  printHeader (shared, itemnames, mcmcwriter);
  if(featuresel)
    printFeatureHeader(featurenames, *featuresel);

  // record the start time now initialisation is complete
  struct timeval starttime;
  gettimeofday (&starttime, 0);

  for (int ittr = 0; ittr < nsamples; ittr++) {
    if (gpu) {
#ifndef NOCUDA
      for (int k = 0; k < nfiles; k++) {
	auto &s = samplers[k];
	s->cudaSampleParameters(shared.alloc(k));
	s->cudaAccumAllocProbs();
      }

      gpu->sampleAlloc();
      gpu->sampleNuConditional();
      gpu->samplePhiConditionals();
      gpu->sampleGammaConditionals();
      gpu->sampleDpMass();

      shared.setAlloc(stdVectorToEigenMatrix(gpu->getAlloc(), nitems, nfiles));
      mdi.setWeight(stdVectorToEigenMatrix(gpu->getWeights(), nclusters, nfiles));
      mdi.setPhis(stdVectorToEigenVector(gpu->getPhis(), mdi.nphis()));
      mdi.setDpMass(stdVectorToEigenVector(gpu->getDpMass(), nfiles));
#endif
    } else {
      mdi.setNu(sampleMdiNu(mdi, shared));

      for (int k = 0; k < nfiles; k++) {
	mdi.setDpMass(k, sampleDpMass(mdi, k));
      }

      for (int k = 0; k < nfiles; k++) {
	auto &s = samplers[k];
	s->sampleParams(shared.alloc(k));

	std::unique_ptr<sampler::item> is(s->newItemSampler());
	const Eigen::VectorXi alloc(sampleSharedAlloc(mdi, shared, k, *is.get()));

	shared.setAlloc(k, alloc);
	if (featuresel)
	  s->sampleFeatureSelection(alloc);
      }

      // sample phi
      for (int k = 0; k < nfiles-1; k++) {
	for (int l = k+1; l < nfiles; l++) {
	  mdi.setPhi(k, l, sampleMdiPhi(mdi, shared, k, l));
	}
      }

      mdi.setWeight(sampleMdiWeights(mdi, shared));
    }

    // block update of cluster labels (e.g. swap all 1s and 3s in a given dataset)
    // the reason for doing this is to get labels to agree between datasets
    if (nfiles > 1) {
      for (int k = 0; k < nfiles; k++) {
	const std::vector<int> changed = sampleAllocBlockSwap (shared, mdi, k);

	std::vector<int> order(changed.size());
	for (size_t j = 0; j < order.size(); j++)
	  order[j] = j;

        std::sort(order.begin(), order.end(), [changed] (int i, int j) {
	    return changed[i] < changed[j]; });

        auto a1 = shared.alloc(k);
        Eigen::VectorXi a2(a1.size());
        for (int i = 0; i < a1.size(); i++) {
          a2(i) = order[a1(i)];
        }

        auto w1 = mdi.weight(k);
        Eigen::VectorXf w2(w1.size());
        for (int j = 0; j < w1.size(); j++) {
          w2(j) = w1(changed[j]);
        }

        shared.setAlloc(k, a2);
        mdi.setWeight(k, w2);
        samplers[k]->swap(changed);
      }

      if (gpu) {
#ifndef NOCUDA
	gpu->setAlloc(eigenMatrixToStdVector(shared.allocs()));
	gpu->setWeights(eigenMatrixToStdVector(mdi.weights()));
#endif
      }
    }

    // TODO: perform split merge steps, ala Savage et al. 2013 ICML

    // Output MCMC samples to the user, thinned as requested
    if (ittr % thinby == 0) {
      printSample(shared, mdi, mcmcwriter);
      if (featuresel)
        printFeatureSample(samplers, *featuresel);
    }
  }

  // print out benchmarking info
  if (vm.count("benchmark")) {
    struct timeval endtime;
    gettimeofday (&endtime, 0);

    std::cerr
      << nsamples << ','
      << (endtime.tv_sec - starttime.tv_sec) * 1000 + (endtime.tv_usec - starttime.tv_usec) / 1000 << " ms"
      << '\n';
  }

#ifndef NOCUDA
  if (gpu) {
    delete gpu;
  }
#endif

  return 0;
}


#ifndef EDGE_PLANE_PRIOR_HPP
#define EDGE_PLANE_PRIOR_HPP

#include <Eigen/Dense>
#include <g2o/core/base_unary_edge.h>
#include <g2o/types/slam3d_addons/vertex_plane.h>

namespace g2o {
class EdgePlanePriorNormal : public g2o::BaseUnaryEdge<3, Eigen::Vector3d, g2o::VertexPlane> {
public:
  EIGEN_MAKE_ALIGNED_OPERATOR_NEW
  EdgePlanePriorNormal() : g2o::BaseUnaryEdge<3, Eigen::Vector3d, g2o::VertexPlane>() {}

  void computeError() override {
    const g2o::VertexPlane* v1 = static_cast<const g2o::VertexPlane*>(_vertices[0]);
    Eigen::Vector3d normal = v1->estimate().normal();

    if(normal.dot(_measurement) < 0.0) {
      normal = -normal;
    }

    _error = normal - _measurement;
  }

  void setMeasurement(const Eigen::Vector3d& m) override {
    _measurement = m;
  }

  virtual bool read(std::istream& is) override {
    Eigen::Vector3d v;
    is >> v(0) >> v(1) >> v(2);
    setMeasurement(v);
    for(int i = 0; i < information().rows(); ++i)
      for(int j = i; j < information().cols(); ++j) {
        is >> information()(i, j);
        if(i != j) information()(j, i) = information()(i, j);
      }
    return true;
  }
  virtual bool write(std::ostream& os) const override {
    Eigen::Vector3d v = _measurement;
    os << v(0) << " " << v(1) << " " << v(2) << " ";
    for(int i = 0; i < information().rows(); ++i)
      for(int j = i; j < information().cols(); ++j) os << " " << information()(i, j);
    return os.good();
  }
};

class EdgePlanePriorDistance : public g2o::BaseUnaryEdge<1, double, g2o::VertexPlane> {
public:
  EIGEN_MAKE_ALIGNED_OPERATOR_NEW
  EdgePlanePriorDistance() : g2o::BaseUnaryEdge<1, double, g2o::VertexPlane>() {}

  void computeError() override {
    const g2o::VertexPlane* v1 = static_cast<const g2o::VertexPlane*>(_vertices[0]);
    _error[0] = _measurement - v1->estimate().distance();
  }

  void setMeasurement(const double& m) override {
    _measurement = m;
  }

  virtual bool read(std::istream& is) override {
    is >> _measurement;
    for(int i = 0; i < information().rows(); ++i)
      for(int j = i; j < information().cols(); ++j) {
        is >> information()(i, j);
        if(i != j) information()(j, i) = information()(i, j);
      }
    return true;
  }
  virtual bool write(std::ostream& os) const override {
    os << _measurement;
    for(int i = 0; i < information().rows(); ++i)
      for(int j = i; j < information().cols(); ++j) os << " " << information()(i, j);
    return os.good();
  }
};
}  // namespace g2o

#endif  // EDGE_SE3_PRIORXY_HPP


///////////////////////////////////////////////////////////////////////////////
// distribution::model::function::likelihood::iid::include.hpp               //
//                                                                           //
//  Copyright 2009 Erwann Rogard. Distributed under the Boost                //
//  Software License, Version 1.0. (See accompanying file                    //
//  LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)         //
///////////////////////////////////////////////////////////////////////////////
#ifndef BOOST_STATISTICS_DETAIL_DISTRIBUTION_MODEL_MODELS_LIKELIHOOD_IID_INCLUDE_HPP_ER_2009
#define BOOST_STATISTICS_DETAIL_DISTRIBUTION_MODEL_MODELS_LIKELIHOOD_IID_INCLUDE_HPP_ER_2009
#include <boost/statistics/detail/distribution/model/models/likelihood/iid/iid.hpp>
#include <boost/statistics/detail/distribution/model/models/likelihood/iid/log_unnormalized_pdf.hpp>

#endif


#pragma once

#include <polyfem/Common.hpp>

#include <polyfem/ElementAssemblyValues.hpp>
#include <polyfem/ElementBases.hpp>
#include <polyfem/ElasticityUtils.hpp>

#include <polyfem/AutodiffTypes.hpp>

#include <Eigen/Dense>
#include <functional>

//local assembler for linear elasticity
namespace polyfem
{
	class LinearElasticity
	{
	public:
		///computes local stiffness matrix is R^{dim²} for bases i,j
		//vals stores the evaluation for that element
		//da contains both the quadrature weight and the change of metric in the integral
		Eigen::Matrix<double, Eigen::Dynamic, 1, 0, 9, 1>
		assemble(const ElementAssemblyValues &vals, const int i, const int j, const QuadratureVector &da) const;

		//neccessary for mixing linear model with non-linear collision response
		Eigen::MatrixXd assemble_hessian(const ElementAssemblyValues &vals, const Eigen::MatrixXd &displacement, const QuadratureVector &da) const;
		//compute gradient of elastic energy, as assembler
		Eigen::VectorXd assemble_grad(const ElementAssemblyValues &vals, const Eigen::MatrixXd &displacement, const QuadratureVector &da) const;
		//compute elastic energy
		double compute_energy(const ElementAssemblyValues &vals, const Eigen::MatrixXd &displacement, const QuadratureVector &da) const;

		//kernel of the pde, used in kernel problem
		Eigen::Matrix<AutodiffScalarGrad, Eigen::Dynamic, 1, 0, 3, 1> kernel(const int dim, const AutodiffGradPt &r) const;

		//uses autodiff to compute the rhs for a fabbricated solution
		//uses autogenerated code to compute div(sigma)
		//pt is the evaluation of the solution at a point
		Eigen::Matrix<double, Eigen::Dynamic, 1, 0, 3, 1>
		compute_rhs(const AutodiffHessianPt &pt) const;

		//compute von mises stress for an element at the local points
		void compute_von_mises_stresses(const int el_id, const ElementBases &bs, const ElementBases &gbs, const Eigen::MatrixXd &local_pts, const Eigen::MatrixXd &displacement, Eigen::MatrixXd &stresses) const;
		//compute stress tensor for an element at the local points
		void compute_stress_tensor(const int el_id, const ElementBases &bs, const ElementBases &gbs, const Eigen::MatrixXd &local_pts, const Eigen::MatrixXd &displacement, Eigen::MatrixXd &tensor) const;

		//size of the problem, this is a tensor problem so the size is the size of the mesh
		inline int &size() { return size_; }
		inline int size() const { return size_; }

		//inialize material parameter
		void set_parameters(const json &params);
		//initialize material param per element
		void init_multimaterial(const bool is_volume, const Eigen::MatrixXd &Es, const Eigen::MatrixXd &nus);

		//class that stores and compute lame parameters per point
		const LameParameters &lame_params() const { return params_; }

	private:
		int size_ = 2;
		//class that stores and compute lame parameters per point
		LameParameters params_;

		void assign_stress_tensor(const int el_id, const ElementBases &bs, const ElementBases &gbs, const Eigen::MatrixXd &local_pts, const Eigen::MatrixXd &displacement, const int all_size, Eigen::MatrixXd &all, const std::function<Eigen::MatrixXd(const Eigen::MatrixXd &)> &fun) const;

		//aux function that computes energy
		//double compute_energy is the same with T=double
		//assemble_grad is the same with T=DScalar1 and return .getGradient()
		template <typename T>
		T compute_energy_aux(const ElementAssemblyValues &vals, const Eigen::MatrixXd &displacement, const QuadratureVector &da) const;
	};
} // namespace polyfem


/* Copyright (c) 2012, Julian Straub <jstraub@csail.mit.edu>
 * Licensed under the MIT license. See LICENSE.txt or
 * http://www.opensource.org/licenses/mit-license.php */

#include <armadillo>
#include <iostream>

#include <stdint.h>

using namespace arma;
using namespace std;

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

  Col<uint32_t> a;
  a<<1<<2<<4<<3<<1<<6<<7<<2<<4;

  cout<<a.t()<<endl;
  cout<<(a==1).t()<<endl;
  cout<<sum(a==1)<<endl;

  mat A;
  A.randn(10,2);
  A.rows(0,4) += 4;
  cout<<A<<endl;

  Col<uint32_t> b=a;
  b.resize(b.n_elem+1);
  b(b.n_elem-1)=99;
  cout<<a.t()<<endl;
  cout<<b.t()<<endl;

  colvec c(10);
  c.zeros();
  c(4)=math::inf();
  c(5)=math::nan();
  cout<<c<<endl;
  cout<<is_finite(c(3))<<endl;
  cout<<is_finite(c(4))<<endl;
  cout<<is_finite(c(5))<<endl;
  cout<<is_finite(c)<<endl;

  colvec d(2);
  d.ones();
  c.insert_rows(c.n_elem,d);
  cout<<c<<endl;

  colvec l(2);
  l << -8.0 << -16.0;
  cout<<l.t()<<endl;
  cout<<exp(l).t()<<endl;

  return 0;
}


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

#include <string>
#include <sstream>

#include <dlib/control.h>
#include <dlib/optimization.h>
#include "tester.h"

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

    logger dlog("test.mpc");

    template <
        typename EXP1,
        typename EXP2,
        typename T, long NR, long NC, typename MM, typename L
        >
    unsigned long solve_qp_box_using_smo (
        const matrix_exp<EXP1>& _Q,
        const matrix_exp<EXP2>& _b,
        matrix<T,NR,NC,MM,L>& alpha,
        matrix<T,NR,NC,MM,L>& lower,
        matrix<T,NR,NC,MM,L>& upper,
        T eps,
        unsigned long max_iter
    )
    /*!
        ensures
            - solves: 0.5*trans(x)*Q*x + trans(b)*x where x is box constrained.
    !*/
    {
        const_temp_matrix<EXP1> Q(_Q);
        const_temp_matrix<EXP2> b(_b);
        //cout << "IN QP SOLVER" << endl;
        //cout << "max eig: " << max(real_eigenvalues(Q)) << endl;
        //cout << "min eig: " << min(real_eigenvalues(Q)) << endl;
        //cout << "Q: \n" << Q << endl;
        //cout << "b: \n" << b << endl;

        // make sure requires clause is not broken
        DLIB_ASSERT(Q.nr() == Q.nc() &&
                     alpha.size() == lower.size() &&
                     alpha.size() == upper.size() &&
                     is_col_vector(b) &&
                     is_col_vector(alpha) &&
                     is_col_vector(lower) &&
                     is_col_vector(upper) &&
                     b.size() == alpha.size() &&
                     b.size() == Q.nr() &&
                     alpha.size() > 0 &&
                     0 <= min(alpha-lower) &&
                     0 <= max(upper-alpha) &&
                     eps > 0 &&
                     max_iter > 0,
                     "\t unsigned long solve_qp_box_using_smo()"
                     << "\n\t Invalid arguments were given to this function"
                     << "\n\t Q.nr():               " << Q.nr()
                     << "\n\t Q.nc():               " << Q.nc()
                     << "\n\t is_col_vector(b):     " << is_col_vector(b)
                     << "\n\t is_col_vector(alpha): " << is_col_vector(alpha)
                     << "\n\t is_col_vector(lower): " << is_col_vector(lower)
                     << "\n\t is_col_vector(upper): " << is_col_vector(upper)
                     << "\n\t b.size():             " << b.size()
                     << "\n\t alpha.size():         " << alpha.size()
                     << "\n\t lower.size():         " << lower.size()
                     << "\n\t upper.size():         " << upper.size()
                     << "\n\t Q.nr():               " << Q.nr()
                     << "\n\t min(alpha-lower):     " << min(alpha-lower)
                     << "\n\t max(upper-alpha):     " << max(upper-alpha)
                     << "\n\t eps:                  " << eps
                     << "\n\t max_iter:             " << max_iter
        );

        // Compute f'(alpha) (i.e. the gradient of f(alpha)) for the current alpha.
        matrix<T,NR,NC,MM,L> df = Q*alpha + b;
        matrix<T,NR,NC,MM,L> QQ = reciprocal_max(diag(Q));

        unsigned long iter = 0;
        for (; iter < max_iter; ++iter)
        {
            T max_df = 0;
            long best_r =0;
            for (long r = 0; r < Q.nr(); ++r)
            {
                if (alpha(r) <= lower(r) && df(r) > 0)
                    ;//alpha(r) = lower(r);
                else if (alpha(r) >= upper(r) && df(r) < 0)
                    ;//alpha(r) = upper(r);
                else if (std::abs(df(r)) > max_df)
                {
                    best_r = r;
                    max_df = std::abs(df(r));
                }
            }

            //for (long r = 0; r < Q.nr(); ++r)
            long r = best_r;
            {

                const T old_alpha = alpha(r);
                alpha(r) = -(df(r)-Q(r,r)*alpha(r))*QQ(r);
                if (alpha(r) < lower(r))
                    alpha(r) = lower(r);
                else if (alpha(r) > upper(r))
                    alpha(r) = upper(r);

                const T delta = old_alpha-alpha(r);

                // Now update the gradient. We will perform the equivalent of: df = Q*alpha + b;
                for(long k = 0; k < df.nr(); ++k)
                    df(k) -= Q(r,k)*delta;
            }

            if (max_df < eps)
                break;
        }
        //cout << "df: \n" << trans(df) << endl;
        //cout << "objective value: " << 0.5*trans(alpha)*Q*alpha + trans(b)*alpha << endl;

        return iter+1;
    }

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

    namespace impl_mpc
    {
        template <long N>
        void pack(
            matrix<double,0,1>& out,
            const std::vector<matrix<double,N,1> >& item
        )
        {
            DLIB_CASSERT(item.size() != 0,"");
            out.set_size(item.size()*item[0].size());
            long j = 0;
            for (unsigned long i = 0; i < item.size(); ++i)
                for (long r = 0; r < item[i].size(); ++r)
                    out(j++) = item[i](r);
        }

        template <long N>
        void pack(
            matrix<double,0,1>& out,
            const matrix<double,N,1>& item,
            const long num
        )
        {
            out.set_size(item.size()*num);
            long j = 0;
            for (long r = 0; r < num; ++r)
                for (long i = 0; i < item.size(); ++i)
                    out(j++) = item(i);
        }

        template <long N>
        void unpack(
            std::vector<matrix<double,N,1> >& out,
            const matrix<double,0,1>& item
        )
        {
            DLIB_CASSERT(out.size() != 0,"");
            DLIB_CASSERT((long)out.size()*out[0].size() == item.size(),"");
            long j = 0;
            for (unsigned long i = 0; i < out.size(); ++i)
                for (long r = 0; r < out[i].size(); ++r)
                    out[i](r) = item(j++);
        }
    }

    template <long S, long I>
    unsigned long solve_linear_mpc (
        const matrix<double,S,S>& A,
        const matrix<double,S,I>& B,
        const matrix<double,S,1>& C,
        const matrix<double,S,1>& Q,
        const matrix<double,I,1>& R,
        const matrix<double,I,1>& _lower,
        const matrix<double,I,1>& _upper,
        const std::vector<matrix<double,S,1> >& target,
        const matrix<double,S,1>& initial_state,
        std::vector<matrix<double,I,1> >& controls // input and output
    )
    {
        using namespace impl_mpc;
        DLIB_CASSERT(target.size() == controls.size(),"");

        matrix<double> K(B.nr()*controls.size(), B.nc()*controls.size());
        matrix<double,0,1> M(B.nr()*controls.size());

        // compute powers of A: Apow[i] == A^i
        std::vector<matrix<double,S,S> > Apow(controls.size());
        Apow[0] = identity_matrix(A);
        for (unsigned long i = 1; i < Apow.size(); ++i)
            Apow[i] = A*Apow[i-1];

        // fill in K
        K = 0;
        for (unsigned long r = 0; r < controls.size(); ++r)
            for (unsigned long c = 0; c <= r; ++c)
                set_subm(K,r*B.nr(),c*B.nc(), B.nr(), B.nc()) = Apow[r-c]*B;

        // fill in M
        set_subm(M,0*A.nr(),0,A.nr(),1) = A*initial_state + C;
        for (unsigned long i = 1; i < controls.size(); ++i)
            set_subm(M,i*A.nr(),0,A.nr(),1) = A*subm(M,(i-1)*A.nr(),0,A.nr(),1) + C;

        //cout << "M: \n" << M << endl;
        //cout << "K: \n" << K << endl;

        matrix<double,0,1> t, v, lower, upper;
        pack(t, target);
        pack(v, controls);
        pack(lower, _lower, controls.size());
        pack(upper, _upper, controls.size());

        matrix<double> QQ(K.nr(),K.nr()), RR(K.nc(),K.nc());
        QQ = 0;
        RR = 0;
        for (unsigned long c = 0; c < controls.size(); ++c)
        {
            set_subm(QQ,c*Q.nr(),c*Q.nr(),Q.nr(),Q.nr()) = diagm(Q);
            set_subm(RR,c*R.nr(),c*R.nr(),R.nr(),R.nr()) = diagm(R);
        }

        matrix<double> m1 = trans(K)*QQ*K+RR;
        matrix<double> m2 = trans(K)*QQ*(M-t);

        // run the solver...
        unsigned long iter;
        iter = solve_qp_box_using_smo(
            m1,
            m2,
            v,
            lower,
            upper,
            0.00000001,
            100000);

        //cout << "iterations: " << iter << endl;

        unpack(controls, v);
        return iter;
    }

    class test_mpc : public tester
    {
    public:
        test_mpc (
        ) :
            tester ("test_mpc",
                    "Runs tests on the mpc object.")
        {}

        void perform_test (
        )
        {
            // a basic position + velocity model
            matrix<double,2,2> A;
            A = 1, 1,
            0, 1;
            matrix<double,2,1> B, C;
            B = 0,
            1;

            C = 0.02,0.1; // no constant bias

            matrix<double,2,1> Q;
            Q = 2, 0; // only care about getting the position right
            matrix<double,1,1> R, lower, upper;
            R = 1;

            lower = -0.2;
            upper =  0.2;

            std::vector<matrix<double,1,1> > controls(30);
            std::vector<matrix<double,2,1> > target(30);
            for (unsigned long i = 0; i < controls.size(); ++i)
            {
                controls[i] = 0;
                target[i] = 0;
            }

            mpc<2,1,30> solver(A,B,C,Q,R,lower,upper);
            solver.set_epsilon(0.00000001);
            solver.set_max_iterations(10000);
            matrix<double,2,1> initial_state;
            initial_state = 0;
            initial_state(0) = 5;
            for (int i = 0; i < 30; ++i)
            {
                print_spinner();
                matrix<double,1,1> control = solver(initial_state);

                for (unsigned long i = 1; i < controls.size(); ++i)
                    controls[i-1] = controls[i];

                // Compute the correct control via SMO and make sure it matches.
                solve_linear_mpc(A,B,C,Q,R,lower,upper, target, initial_state, controls);
                dlog << LINFO << "ERROR: " << length(control-controls[0]);
                DLIB_TEST(length(control-controls[0]) < 1e-7);

                initial_state = A*initial_state + B*control + C;
                //cout << control(0) << "\t" << trans(initial_state);
            }

            {
                // also just generally test our QP solver.
                matrix<double,20,20> Q = gaussian_randm(20,20,5);
                Q = Q*trans(Q);

                matrix<double,20,1> b = randm(20,1)-0.5;
                matrix<double,20,1> alpha, lower, upper, alpha2;
                alpha = 0;
                alpha2 = 0;
                lower = -4;
                upper = 3;

                solve_qp_box_using_smo(Q,b,alpha,lower, upper, 1e-12, 500000);
                solve_qp_box_constrained(Q,b,alpha2,lower, upper, 1e-12, 50000);
                dlog << LINFO << trans(alpha);
                dlog << LINFO << trans(alpha2);
                dlog << LINFO << "objective value:  " << 0.5*trans(alpha)*Q*alpha + trans(b)*alpha;
                dlog << LINFO << "objective value2: " << 0.5*trans(alpha2)*Q*alpha + trans(b)*alpha2;
                DLIB_TEST_MSG(max(abs(alpha-alpha2)) < 1e-7, max(abs(alpha-alpha2)));
            }
        }
    } a;

}


#ifndef CX_DQMC_GREENS_REPLICA_RENYI_FT_HPP
#define CX_DQMC_GREENS_REPLICA_RENYI_FT_HPP

#include <complex>
#include <libdqmc/cx_dqmc_abstract_greens.hpp>
#include <libdqmc/la.hpp>
#include <libdqmc/cx_density.hpp>
#include <libdqmc/parameters.hpp>
#include <libdqmc/cx_workspace.hpp>
#include <libdqmc/cx_checkerboard.hpp>
#include <libdqmc/cx_interactions.hpp>
#include <libdqmc/cx_calculate_greens_ft.hpp>
#include <alps/alea.h>
#include <boost/multi_array.hpp>
#include <exception>
#include <cstring>

typedef boost::multi_array<double, 3> aux_spin_3t;

namespace cx_dqmc {

    class greens_replica_renyi_ft : public cx_dqmc::abstract_greens {

    public:
	bool first_initialization, dir_switch;

	cx_mat_t u_temp, t_temp;
	cx_mat ket, ket_sqr, bra, bra_sqr;
	vec ket_d, bra_d, d_temp;

	cx_mat_t Ul, Um, Ur, Tl, Tm, Tr;
	vec_t Dl, Dm, Dr;

	int det_method;
	bool fresh_sign;

	greens_replica_renyi_ft(dqmc::parameters* p_,
				cx_dqmc::workspace* ws_,
				cx_dqmc::workspace* reg_ws_,
				aux_spin_3t* s)
	    : cx_dqmc::abstract_greens(p_, ws_, reg_ws_, s) {};

	void initialize();
	void build_stack();
	int propagate(alps::Observable& stability);

	void log_weight();
	void calculate_greens(int direction, int slice) {
	    using namespace std;
	    // boost::timer::auto_cpu_timer t;
	    // cout << "greens_replica_renyi_ft::calculate_greens(int, int)\t";
	    calculate_greens_general();
	};
	void calculate_greens_exact(int slice);

	void update_remove_interaction();
	void update_add_interaction();

	void ket_update(int start, int stop, int i, int j);
	void bra_update(int start, int stop, int i, int j);
	void slice_sequence(int start, int stop) {};
	void slice_sequence_left(int start, int stop, cx_mat& M);
	void slice_sequence_left_renyi(int start, int stop, cx_mat& M);
	void slice_sequence_left_t(int start, int stop, cx_mat& M);
	void slice_sequence_left_renyi_t(int start, int stop, cx_mat& M);
	void slice_sequence_right(int start, int stop, cx_mat& M);

	void slice_matrix_left(int slice, cx_mat& M, bool inv = false);
	void slice_matrix_left_t(int slice, cx_mat& M, bool inv = false);
	void slice_matrix_left_renyi_t(int slice, cx_mat& M, bool inv = false);
	void slice_matrix_right(int slice, cx_mat& M, bool inv = false);

	void slice_matrix_left_renyi(int slice, cx_mat& M, bool inv = false);
	void slice_matrix_right_renyi(int slice, cx_mat& M, bool inv = false);

	void hopping_matrix_left(int slice, cx_mat& M, bool inv = false);
	void hopping_matrix_right(int slice, cx_mat& M, bool inv = false);
	void hopping_matrix_left_renyi(int slice, cx_mat& M, bool inv = false);
	void hopping_matrix_right_renyi(int slice, cx_mat& M, bool inv = false);

	void interaction_matrix_left(int slice, cx_mat& M,
				     int bond_group, bool inv = false);
	void interaction_matrix_right(int slice, cx_mat& M,
				      int bond_group, bool inv = false);

	void regularize_svd(vec& in, vec& out);

	void calculate_greens(int method = 0);
	void calculate_greens_general();
	void calculate_greens(cx_mat U, cx_mat& G);

    	double check_stability();

	void enlarge(int replica, cx_mat_t& in, cx_mat_t& out);
	void enlarge(int replica, vec_t& in, vec_t& out);
    };
}
#endif


#include <acrolib/data_loggers.h>
#include <acrolib/extern/npy.h>

#include <cassert>
#include <iostream>
#include <fstream>
#include <Eigen/Dense>

const std::string TEST_FILE = R"(one,two,three
1,2,3
4,5,6
hello,3,0
)";

using namespace acro;
using namespace std;

bool testFileLogger()
{
  Logger log("test.csv");

  log.header("one,two,three");

  Eigen::VectorXd q1(3);
  q1 << 1, 2, 3;
  log.log(q1);
  log.nextLine();

  vector<double> q2{ 4, 5, 6 };
  log.log(q2);
  log.nextLine();

  log.log("hello");
  log.log(3);
  log.log(false);
  log.nextLine();

  return true;
}

bool checkLoggedFile()
{
  // read the file again and compare result
  std::ifstream file("test.csv");
  assert(file.is_open());

  std::stringstream buffer;
  buffer << file.rdbuf();
  assert(buffer.str() == TEST_FILE);

  return true;
}

bool testNumPyLogger()
{
  // log a 5 by 3 matrix to a file
  NumPyLogger log("mat.npy", 5, 3);

  for (int row{ 0 }; row < 5; ++row)
  {
    Eigen::Vector3d v(9, 9, (double)row);
    log.log(v);
  }

  log.finish();

  return true;
}

bool checkLoggedNumPyFile()
{
  // read the file again and compare result
  // you can also read the same file using NumPy and hopefully get the same result
  bool fortran_order;
  vector<unsigned long> shape;
  vector<double> data;
  npy::LoadArrayFromNumpy("mat.npy", shape, fortran_order, data);

  assert(fortran_order == false);
  assert(shape.at(0) == 5);
  assert(shape.at(1) == 3);

  vector<double> expected_data{ 9, 9, 0, 9, 9, 1, 9, 9, 2, 9, 9, 3, 9, 9, 4 };

  for (size_t k{ 0 }; k < data.size(); ++k)
  {
    assert(data[k] == expected_data[k]);
  }

  return true;
}

int main()
{
  assert(testFileLogger());
  assert(checkLoggedFile());

  assert(testNumPyLogger());
  assert(checkLoggedNumPyFile());

  return 0;
}


//
// Created by kerin on 2019-02-26.
//
#include <map>
#include <string>
#include "data.hpp"
#include "tools/eigen3.3/Dense"

#include <boost/iostreams/filtering_stream.hpp>
#include <boost/iostreams/device/file.hpp>
#include <boost/iostreams/filter/gzip.hpp>

void Data::read_txt_file(const std::string &filename, Eigen::MatrixXd &M, unsigned long &n_cols,
                         std::vector<std::string> &col_names, std::map<int, bool> &incomplete_row) {
	// pass top line of txt file filename to col_names, and body to M.
	// TODO: Implement how to deal with missing values.

	boost_io::filtering_istream fg;
	std::string gz_str = ".gz";
	if (filename.find(gz_str) != std::string::npos) {
		fg.push(boost_io::gzip_decompressor());
	}
	fg.push(boost_io::file_source(filename));
	if (!fg) {
		throw std::runtime_error(filename+" could not be opened.");
	}

	// Reading column names
	std::string line;
	if (!getline(fg, line)) {
		throw std::runtime_error(filename+" has zero lines.");
	}
	std::stringstream ss;
	std::string s;
	n_cols = 0;
	ss.clear();
	ss.str(line);
	while (ss >> s) {
		++n_cols;
		col_names.push_back(s);
	}
	std::cout << " Reading matrix of size " << n_samples << " x " << n_cols << " from " << filename << std::endl;

	// Write remainder of file to Eigen matrix M
	incomplete_row.clear();
	M.resize(n_samples, n_cols);
	int i = 0;
	double tmp_d;
	try {
		while (getline(fg, line)) {
			if (i >= n_samples) {
				throw std::runtime_error("ERROR: could not convert txt file (too many lines).");
			}
			ss.clear();
			ss.str(line);
			for (int k = 0; k < n_cols; k++) {
				std::string sss;
				ss >> sss;
				/// NA
				if (sss == "NA" || sss == "NAN" || sss == "NaN" || sss == "nan") {
					tmp_d = params.missing_code;
				} else {
					try{
						tmp_d = stod(sss);
					} catch (const std::invalid_argument &exc) {
						std::cout << sss << " on line " << i << std::endl;
						throw;
					}
				}

				if(tmp_d != params.missing_code) {
					M(i, k) = tmp_d;
				} else {
					M(i, k) = params.missing_code;
					incomplete_row[i] = true;
				}
			}
			i++;             // loop should end at i == n_samples
		}
		if (i < n_samples) {
			throw std::runtime_error("ERROR: could not convert txt file (too few lines).");
		}
	} catch (const std::exception &exc) {
		throw;
	}
}

template<typename EigenMat>
EigenMat reduce_mat_to_complete_cases(EigenMat &M, bool &matrix_reduced, const unsigned long &n_cols,
                                      const std::map<std::size_t, bool> &incomplete_cases) {
	// Remove rows contained in incomplete_cases
	long nn = M.rows();
	EigenMat M_tmp;
	if (matrix_reduced) {
		throw std::runtime_error("ERROR: Trying to remove incomplete cases twice...");
	}

	// Create temporary matrix of complete cases
	unsigned long n_incomplete = incomplete_cases.size();
	M_tmp.resize(nn - n_incomplete, n_cols);

	// Fill M_tmp with non-missing entries of M
	int ii_tmp = 0;
	for (std::size_t ii = 0; ii < nn; ii++) {
		if (incomplete_cases.count(ii) == 0) {
			for (int kk = 0; kk < n_cols; kk++) {
				M_tmp(ii_tmp, kk) = M(ii, kk);
			}
			ii_tmp++;
		}
	}

	// Assign new values to reference variables
	matrix_reduced = true;
	return M_tmp;
}


/*
@copyright Louis Dionne 2014
Distributed under the Boost Software License, Version 1.0.
(See accompanying file LICENSE.md or copy at http://boost.org/LICENSE_1_0.txt)
 */

#include <boost/hana/detail/assert.hpp>
#include <boost/hana/detail/constexpr.hpp>
#include <boost/hana/lazy.hpp>
using namespace boost::hana;

int main() {
    //! [main]
    BOOST_HANA_CONSTEXPR_LAMBDA auto double_ = [](auto x) {
        return x * 2;
    };

    BOOST_HANA_CONSTEXPR_LAMBDA auto one_over = [](auto x) {
        return 1 / x;
    };

    BOOST_HANA_CONSTEXPR_ASSERT(eval(fmap(double_, lazy(one_over)(4))) == 1 / 2);

    fmap(double_, lazy(one_over)(0)); // never evaluated
    //! [main]
}


#include <stdexcept>
#include <iostream>

#include <boost/safe_numerics/safe_integer.hpp>

int main(int, const char *[]){
    // problem: cannot recover from arithmetic errors
    std::cout << "example 8: ";
    std::cout << "cannot detect compile time arithmetic errors" << std::endl;
    std::cout << "Not using safe numerics" << std::endl;

    try{
        const int x = 1;
        const int y = 0;
        // will emit warning at compile time
        // will leave an invalid result at runtime.
        std::cout << x / y; // will display "0"!
        std::cout << "error NOT detected!" << std::endl;
    }
    catch(const std::exception &){
        std::cout << "error detected!" << std::endl;
    }
    // solution: replace int with safe<int>
    std::cout << "Using safe numerics" << std::endl;
    try{
        using namespace boost::safe_numerics;
        const safe<int> x = 1;
        const safe<int> y = 0;
        // constexpr const safe<int> z = x / y; // note constexpr here!
        std::cout << x / y; // error would be detected at runtime
        std::cout << " error NOT detected!" << std::endl;
    }
    catch(const std::exception & e){
        std::cout << "error detected:" << e.what() << std::endl;
    }
    return 0;
}


//=======================================================================
// Copyright 2014 Alexander Lauser.
// Authors: Alexander Lauser
//
// Distributed under the Boost Software License, Version 1.0. (See
// accompanying file LICENSE_1_0.txt or copy at
// http://www.boost.org/LICENSE_1_0.txt)
//=======================================================================

// This test is an adapted version of the MWE for Bug #10231 (showing
// incorrect root_map computation).

/* Output should be:
The example graph:
a --> b
b --> a c
c --> b

Vertex a is in component 0 and has root 0
Vertex b is in component 0 and has root 0
Vertex c is in component 0 and has root 0
*/
#include <boost/config.hpp>
#include <iostream>
#include <vector>
#include <boost/graph/strong_components.hpp>
#include <boost/graph/adjacency_list.hpp>
#include <boost/graph/graph_utility.hpp>

int main(int, char*[])
{
  using namespace boost;

  adjacency_list<vecS, vecS, directedS> G;

  typedef graph_traits<adjacency_list<vecS, vecS, directedS> >::vertex_descriptor Vertex;
  Vertex a = add_vertex(G);
  Vertex b = add_vertex(G);
  Vertex c = add_vertex(G);

  add_edge(a, b, G);
  add_edge(b, a, G);

  add_edge(c, b, G);
  add_edge(b, c, G);

#if VERBOSE
  std::cout << "The example graph:" << std::endl;
  const char* name = "abc";
  print_graph(G, name);
  std::cout << std::endl;
#endif

  std::vector<int> component(num_vertices(G)), discover_time(num_vertices(G));
  std::vector<default_color_type> color(num_vertices(G));
  std::vector<Vertex> root(num_vertices(G));
  strong_components(G, make_iterator_property_map(component.begin(), get(vertex_index, G)),
		    root_map(make_iterator_property_map(root.begin(), get(vertex_index, G))).
		    color_map(make_iterator_property_map(color.begin(), get(vertex_index, G))).
		    discover_time_map(make_iterator_property_map(discover_time.begin(), get(vertex_index, G))));

#if VERBOSE
  for (std::vector<int>::size_type i = 0; i != component.size(); ++i)
    std::cout << "Vertex " << name[i]
	      << " is in component " << component[i]
	      << " and has root " << root[i] << std::endl;
#endif

#if VERBOSE
  bool test_failed;
#endif
  int ret = 0;

  //////////
#if VERBOSE
  test_failed = false;
  std::cerr << "Testing component-computation of strong_components ..." << std::endl;
#endif
  for (std::vector<int>::size_type i = 0; i != component.size(); ++i) {
    if(component[i] != 0) {
#if VERBOSE
      test_failed = true;
#endif
      ret = -1;
      break;
    }
  }
#if VERBOSE
  std::cerr << (test_failed ? "   **** Failed." :  "   Passed.") << std::endl;
#endif
  //////////

  //////////
#if VERBOSE
  test_failed = false;
  std::cerr << "Testing root_map-computation of strong_components ..." << std::endl;
#endif
  for (std::vector<int>::size_type i = 0; i != component.size(); ++i) {
    if(root[i] != 0) {
#if VERBOSE
      test_failed = true;
#endif
      ret = -1;
      break;
    }
  }
#if VERBOSE
  std::cerr << (test_failed ? "   **** Failed." :  "   Passed.") << std::endl;
#endif
  //////////

  if (ret == 0)
    std::cout << "tests passed" << std::endl;
  return ret;
}


// smooth: Lie Theory for Robotics
// https://github.com/pettni/smooth
//
// 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.

#ifndef SMOOTH__MANIFOLD_VECTOR_HPP_
#define SMOOTH__MANIFOLD_VECTOR_HPP_

#include <Eigen/Sparse>
#include <numeric>

#include "manifold.hpp"

namespace smooth {

/**
 * @brief \p std::vector based Manifold container.
 *
 * Convenient to treat a collection
 * \f[
 *   m = (m_1, m_2, ..., m_k) \in M \times M \times ... \times M
 * \f]
 * of Manifold elements as a single Manifold element \f$m\f$.
 */
template<Manifold M>
class ManifoldVector : public std::vector<M>
{
private:
  using Base = std::vector<M>;

public:
  //! Default constructor of empty ManifoldVector
  ManifoldVector() = default;
  //! Copy constructor
  ManifoldVector(const ManifoldVector & o) = default;
  //! Move constructor
  ManifoldVector(ManifoldVector && o) = default;
  //! Copy assignment operator
  ManifoldVector & operator=(const ManifoldVector & o) = default;
  //! Move assignment operator
  ManifoldVector & operator=(ManifoldVector && o) = default;
  ~ManifoldVector()                               = default;

  /**
   * Forwarding constructor to std::vector.
   */
  template<typename... Ts>
  ManifoldVector(Ts &&... ts) : Base(std::forward<Ts>(ts)...)
  {}

  /**
   * @brief Cast to different scalar type.
   */
  template<typename NewScalar>
  ManifoldVector<CastT<NewScalar, M>> cast() const
  {
    ManifoldVector<CastT<NewScalar, M>> ret;
    ret.reserve(size());
    std::transform(this->begin(), this->end(), std::back_insert_iterator(ret), [](const auto & x) {
      return ::smooth::cast<NewScalar>(x);
    });
    return ret;
  }

  /**
   * @brief Number of elements in ManifoldVector.
   */
  std::size_t size() const { return Base::size(); }

  /**
   * @brief Runtime degrees of freedom.
   *
   * Sum of the degrees of freedom of constituent elements.
   */
  Eigen::Index dof() const
  {
    if constexpr (Dof < M >> 0) {
      return size() * Dof<M>;
    } else {
      return std::accumulate(this->begin(), this->end(), 0u, [](auto s, const auto & item) {
        return s + ::smooth::dof<M>(item);
      });
    }
  }

  /**
   * @brief In-place addition.
   *
   * @note It must hold that dof() == a.dof()
   */
  template<typename Derived>
  ManifoldVector<M> & operator+=(const Eigen::MatrixBase<Derived> & a)
  {
    Eigen::Index dof_cntr = 0;
    for (auto i = 0u; i != this->size(); ++i) {
      const auto dof_i = ::smooth::dof<M>(this->operator[](i));
      this->operator[](i) =
        rplus<M>(this->operator[](i), a.template segment<Dof<M>>(dof_cntr, dof_i));
      dof_cntr += dof_i;
    }
    return *this;
  }

  /**
   * @brief Addition.
   *
   * @note It must hold that `dof() == a.dof()`
   */
  template<typename Derived>
  ManifoldVector<M> operator+(const Eigen::MatrixBase<Derived> & a) const
  {
    ManifoldVector<M> ret = *this;
    ret += a;
    return ret;
  }

  /**
   * @brief Subtraction.
   *
   * @note It must hold that `dof() == o.dof()`
   */
  Eigen::VectorX<Scalar<M>> operator-(const ManifoldVector<M> & o) const
  {
    std::size_t dof_cnts = 0;
    if (Dof < M >> 0) {
      dof_cnts = Dof<M> * size();
    } else {
      for (auto i = 0u; i != size(); ++i) { dof_cnts += ::smooth::dof<M>(this->operator[](i)); }
    }

    Eigen::VectorX<Scalar<M>> ret(dof_cnts);
    Eigen::Index idx = 0;
    for (auto i = 0u; i != size(); ++i) {
      const auto & size_i                       = ::smooth::dof<M>(this->operator[](i));
      ret.template segment<Dof<M>>(idx, size_i) = rminus<M>(this->operator[](i), o[i]);
      idx += size_i;
    }

    return ret;
  }
};

/**
 * @brief Manifold interface for ManifoldVector
 */
template<Manifold M>
struct traits::man<ManifoldVector<M>>
{
  // \cond
  using Scalar      = ::smooth::Scalar<M>;
  using PlainObject = ManifoldVector<M>;
  template<typename NewScalar>
  using CastT = ManifoldVector<typename man<M>::template CastT<NewScalar>>;

  static constexpr Eigen::Index Dof = -1;

  static inline Eigen::Index dof(const ManifoldVector<M> & m) { return m.dof(); }

  static inline PlainObject Default(Eigen::Index dof)
  {
    /// @note If underlying M has dynamic size, mdof not uniquely defined
    const Eigen::Index mdof = ::smooth::Dof<M> != -1 ? ::smooth::Dof<M> : 1;
    const Eigen::Index size = dof / mdof;

    return PlainObject(size, Default<M>(mdof));
  }

  template<typename NewScalar>
  static inline auto cast(const ManifoldVector<M> & m)
  {
    return m.template cast<NewScalar>();
  }

  template<typename Derived>
  static inline ManifoldVector<M>
  rplus(const ManifoldVector<M> & m, const Eigen::MatrixBase<Derived> & a)
  {
    return m + a;
  }

  static inline Eigen::Matrix<Scalar, Dof, 1>
  rminus(const ManifoldVector<M> & m1, const ManifoldVector<M> & m2)
  {
    return m1 - m2;
  }
  // \endcond
};

}  // namespace smooth

template<typename Stream, smooth::Manifold M>
Stream & operator<<(Stream & s, const smooth::ManifoldVector<M> & g)
{
  s << "ManifoldVector with " << g.size() << " elements:\n";
  for (auto i = 0u; i != g.size(); ++i) {
    s << i << ": " << g[i];
    if (i != g.size() - 1) { s << '\n'; }
  }
  return s;
}

#endif  // SMOOTH__MANIFOLD_VECTOR_HPP_


//==================================================================================================
/*!
  @file

  @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_SIMD_FUNCTION_NTHROOT_HPP_INCLUDED
#define BOOST_SIMD_ARCH_COMMON_SIMD_FUNCTION_NTHROOT_HPP_INCLUDED
#include <boost/simd/detail/overload.hpp>

#include <boost/simd/constant/mone.hpp>
#include <boost/simd/constant/one.hpp>
#include <boost/simd/function/abs.hpp>
#include <boost/simd/function/bitofsign.hpp>
#include <boost/simd/function/bitwise_or.hpp>
#include <boost/simd/function/divides.hpp>
#include <boost/simd/function/if_nan_else.hpp>
#include <boost/simd/function/if_else.hpp>
#include <boost/simd/function/if_else_zero.hpp>
#include <boost/simd/function/if_zero_else.hpp>
#include <boost/simd/function/is_equal.hpp>
#include <boost/simd/function/is_eqz.hpp>
#include <boost/simd/function/is_even.hpp>
#include <boost/simd/function/is_inf.hpp>
#include <boost/simd/function/is_ltz.hpp>
#include <boost/simd/function/is_nez.hpp>
#include <boost/simd/function/logical_and.hpp>
#include <boost/simd/function/logical_andnot.hpp>
#include <boost/simd/function/logical_or.hpp>
#include <boost/simd/function/minus.hpp>
#include <boost/simd/function/dec.hpp>
#include <boost/simd/function/multiplies.hpp>
#include <boost/simd/function/pow.hpp>
#include <boost/simd/function/rec.hpp>
#include <boost/simd/function/if_plus.hpp>
#include <boost/simd/function/tofloat.hpp>
#include <boost/simd/function/unary_minus.hpp>
#include <boost/simd/function/fast.hpp>
#include <boost/simd/function/log.hpp>
#include <boost/simd/function/exp.hpp>

namespace boost { namespace simd { namespace ext
{
   namespace bd = boost::dispatch;
   namespace bs = boost::simd;
   BOOST_DISPATCH_OVERLOAD_IF( nthroot_
                          , (typename A0, typename A1, typename X)
                          , (detail::is_native<X>)
                          , bd::cpu_
                          , bs::pack_<bd::floating_<A0>, X>
                          , bs::pack_<bd::integer_<A1>, X>
                          )
   {
      BOOST_FORCEINLINE A0 operator()( const A0& a0, const  A1&  a1) const BOOST_NOEXCEPT
      {
      using bA0 = bs::as_logical_t<A0>;
      A0 x =  bs::abs(a0);
      A0 aa1 = bs::tofloat(a1);
      A0 y = bs::fast_(bs::pow_abs)(x,rec(aa1));
      bA0 nul_a1 =  bs::is_eqz(bitwise_cast<A0>(a1));
      bA0 is_ltza0 = is_ltz(a0);
      auto is_odda1 = is_odd(a1);
      A0 p = fast_(bs::pow_abs)(y, aa1);
      y = bs::if_plus( bs::logical_or(bs::is_nez(y), nul_a1)
                     , y
                     , -(p - x)/(aa1*p/y)
                     );
      // Correct numerical errors (since, e.g., 64^(1/3) is not exactly 4)
      // by one iteration of Newton's method
      bA0 test =  logical_andnot(is_ltza0, is_odda1);
      bA0 done =  test;
      y = if_nan_else(test, y);  // a0 < O and a1 is not odd
      bA0 newtest =  is_equal(x, One<A0>());
      test  = logical_andnot(newtest, done);
      done  = logical_or(done, newtest);
      y = if_else(test, a0, y); // 1^a1 or (-1)^a1
      newtest =  nul_a1;
      test  = logical_andnot(newtest, done);
      done  = logical_or(done, newtest);
      y =  if_else(test,
                   if_zero_else(is_less(x, One<A0>()),
                                sign(a0)*Inf<A0>()
                               ),
                   y);
      newtest =  is_eqz(a0);
      test  = logical_andnot(newtest, done);
      done  = logical_or(done, newtest);
      y =  if_zero_else(test, y);
      #ifndef BOOST_SIMD_NO_INFINITIES
      newtest =  is_inf(a0);
      test  = logical_andnot(newtest, done);
      done  = logical_or(done, newtest);
      y =  if_else(test, if_else(is_nez(a1), a0, One<A0>()), y);
      #endif
      return bs::bitwise_or(y, bs::bitofsign(a0));
      }
   };

   BOOST_DISPATCH_OVERLOAD( nthroot_
                          , (typename A0, typename A1, typename X)
                          , bd::cpu_
                          , bs::fast_tag
                          , bs::pack_<bd::floating_<A0>, X>
                          , bs::pack_<bd::integer_<A1>, X>
                          )
   {
     BOOST_FORCEINLINE A0 operator()(const fast_tag &, const A0& a0, const  A1&  a1) const BOOST_NOEXCEPT
     {
       auto aa1 =  abs(a1);
       A0 aa0 = abs(a0);
       A0 y = sign(a0)*bs::exp(bs::log(aa0)/tofloat(aa1));
       auto l =  is_ltz(aa1);
       y =  if_nan_else(logical_and(l, is_even(a1)), y);
       return if_else(is_ltz(a1), rec(y), y);
     }
   };

} } }

#endif


// Copyright (c) 2018 Graphcore Ltd. All rights reserved.
#define BOOST_TEST_MODULE SpatialSoftmaxTest

#include <boost/test/unit_test.hpp>
#include <poplar/Engine.hpp>
#include <poplar/OptionFlags.hpp>
#include <poplibs_support/TestDevice.hpp>
#include <poplin/codelets.hpp>
#include <popnn/SpatialSoftMax.hpp>
#include <popnn/codelets.hpp>
#include <popops/codelets.hpp>
#include <poputil/TileMapping.hpp>

#include <vector>
using namespace poplibs_support;

const poplar::OptionFlags options;

BOOST_AUTO_TEST_CASE(SpatialSoftmax) {
  auto device = createTestDevice(TEST_TARGET);
  poplar::Graph g(device.getTarget());
  poplar::program::Sequence prog;
  popops::addCodelets(g);
  poplin::addCodelets(g);
  popnn::addCodelets(g);

  // Create some inputs:
  const auto width = 4u;
  const auto height = 3u;

  auto fields = g.addVariable(poplar::FLOAT, {1, height, width}, "input");
  poputil::mapTensorLinearly(g, fields);

  // Set the normalising factor to one so that the result is the mean
  // coordinate:
  const auto initialTemp = 1.f;

  // Disable adding the softmax operation itself so that computing the expected
  // values for the coordinates can be done by hand arithmetic below. (The
  // softmax is a vanilla call to Poplib's nonLinearity so independently
  // tested).
  auto ssm = popnn::spatialSoftMax2D(g, prog, fields, initialTemp, true, "ssm");
  // Check we will get a coordinate for every field down the rows:
  BOOST_CHECK_EQUAL(ssm.first.shape()[0], fields.shape()[0]);
  // Check the coords will be 2D (x,y) across columns
  BOOST_CHECK_EQUAL(ssm.first.shape()[1], 2u);

  g.createHostWrite("fields", fields);
  g.createHostRead("coords", ssm.first);
  g.createHostRead("temp", ssm.second);

  poplar::Engine e(g, prog, options);
  device.bind([&](const poplar::Device &d) {
    e.load(d);

    const std::vector<float> activations = {
        // x = -1, -1/3, 1/3,  1
        0.f, 0.f, 0.f, 0.f, // y = -1
        .5f, 0.f, 0.f, 0.f, // y =  0
        0.f, .5f, 0.f, 0.f  // y =  1
    };
    e.writeTensor("fields", activations.data(),
                  activations.data() + activations.size());
    // Compute expected coord for the 'centre of mass' of the above activations.
    // (NOTE: we disabled the softmax normalisaiton above, otherwise the
    // expected result is painful to compute!)
    // The columns coordinates map from indeices [0,3] -> [-1,1]
    // The row coordinates map from indeices [0,2] -> [-1,1]
    // Hence, mean activated index gets mapped like this:
    // [ux, uy] -> [(ux/3)*2-1, (uy/2)*2-1]
    const std::vector<float> expected = {1.f / 2,
                                         -2.f / 3}; // {ey, ex} => {row, col}

    e.run();

    std::vector<float> actual(fields.shape()[0] * 2u, -100.f);
    e.readTensor("coords", actual.data(), actual.data() + actual.size());
    for (auto i = 0u; i < actual.size(); ++i) {
      BOOST_CHECK_CLOSE(actual.at(i), expected.at(i), 0.00001);
    }

    float checkTemp = 0.f;
    e.readTensor("temp", &checkTemp, &checkTemp + 1);
    BOOST_CHECK_EQUAL(checkTemp, initialTemp);
  });
}


/*

PICCANTE
The hottest HDR imaging library!
http://vcg.isti.cnr.it/piccante

Copyright (C) 2014
Visual Computing Laboratory - ISTI CNR
http://vcg.isti.cnr.it
First author: Francesco Banterle

This Source Code Form is subject to the terms of the Mozilla Public
License, v. 2.0. If a copy of the MPL was not distributed with this
file, You can obtain one at http://mozilla.org/MPL/2.0/.

*/

#ifndef PIC_COMPUTER_VISION_CAMERA_MARTIX_HPP
#define PIC_COMPUTER_VISION_CAMERA_MARTIX_HPP

#include <vector>
#include <random>
#include <stdlib.h>

#include "../base.hpp"

#include "../util/math.hpp"
#include "../util/eigen_util.hpp"

#ifndef PIC_DISABLE_EIGEN

#ifndef PIC_EIGEN_NOT_BUNDLED
    #include "../externals/Eigen/Dense"
    #include "../externals/Eigen/SVD"
    #include "../externals/Eigen/Geometry"
    #include "../externals/Eigen/Geometry"
    #include "../externals/Eigen/QR"
#else
    #include <Eigen/Dense>
    #include <Eigen/SVD>
    #include <Eigen/Geometry>
    #include <Eigen/QR>
#endif

#endif

namespace pic {

#ifndef PIC_DISABLE_EIGEN

/**
 * @brief computeEpipole computes the epipole of a fundamental matrix F.
 * @param F is a fundamental matrix.
 * @return It returns the epipole of F.
 */
PIC_INLINE Eigen::Vector3d computeEpipole(Eigen::Matrix3d &F)
{
    Eigen::JacobiSVD< Eigen::Matrix3d > svdF(F, Eigen::ComputeFullV);
    Eigen::Matrix3d V = svdF.matrixV();

    Eigen::Vector3d e;

    e[0] = V(0, 2);
    e[1] = V(1, 2);
    e[2] = V(2, 2);

    return e;
}

/**
 * @brief getCameraMatrixFromHomography
 * @param H is 3x3 homography matrix.
 * @param K
 * @return
 */
PIC_INLINE Eigen::Matrix34d getCameraMatrixFromHomography(Eigen::Matrix3d &H, Eigen::Matrix3d &K)
{
    Eigen::Matrix34d m;
    m.setZero();

    Eigen::Matrix3d K_inv = K.inverse();

    Eigen::Matrix3d H_p = K_inv * H;

    Eigen::Vector3d r_0(H_p(0, 0), H_p(1, 0), H_p(2, 0));
    Eigen::Vector3d r_1(H_p(0, 1), H_p(1, 1), H_p(2, 1));

    r_0.normalize();
    r_1.normalize();
    Eigen::Vector3d r_2 = r_0.cross(r_1);

    Eigen::Vector3d t(H_p(0, 2), H_p(1, 2), H_p(2, 2));

    m(0, 0) = r_0[0];
    m(1, 0) = r_0[1];
    m(2, 0) = r_0[2];

    m(0, 1) = r_1[0];
    m(1, 1) = r_1[1];
    m(2, 1) = r_1[2];

    m(0, 2) = r_2[0];
    m(1, 2) = r_2[1];
    m(2, 2) = r_2[2];

    m(0 , 3) = t[0];
    m(1 , 3) = t[1];
    m(2 , 3) = t[2];

    return K * m;
}

/**
 * @brief getCameraMatrixIdentity
 * @param K
 * @return
 */
PIC_INLINE Eigen::Matrix34d getCameraMatrixIdentity(Eigen::Matrix3d &K)
{
    Eigen::Matrix34d m;
    m.setIdentity();
    return K * m;
}

/**
 * @brief getCameraMatrix
 * @param K
 * @param R
 * @param t
 * @return
 */
PIC_INLINE Eigen::Matrix34d getCameraMatrix(Eigen::Matrix3d &K, Eigen::Matrix3d &R, Eigen::Vector3d &t)
{
    Eigen::Matrix34d m;

    m(0, 0) = R(0, 0);
    m(1, 0) = R(1, 0);
    m(2, 0) = R(2, 0);

    m(0, 1) = R(0, 1);
    m(1, 1) = R(1, 1);
    m(2, 1) = R(2, 1);

    m(0, 2) = R(0, 2);
    m(1, 2) = R(1, 2);
    m(2, 2) = R(2, 2);

    m(0, 3) = t[0];
    m(1, 3) = t[1];
    m(2, 3) = t[2];

    return K * m;
}

/**
 * @brief decomposeCameraMatrix
 * @param P
 * @param K
 * @param R
 * @param t
 */
PIC_INLINE void decomposeCameraMatrix(Eigen::Matrix34d &P,
                                      Eigen::Matrix3d  &K,
                                      Eigen::Matrix3d  &R,
                                      Eigen::Vector3d  &t)
{
    Eigen::Matrix3d matrix = P.block<3, 3>(0, 0).inverse();

    //QR decomposition
    Eigen::HouseholderQR<Eigen::Matrix3d> qr(matrix.rows(), matrix.cols());
    qr.compute(matrix);

    Eigen::Matrix3d Q = qr.householderQ();
    Eigen::Matrix3d U = qr.matrixQR().triangularView<Eigen::Upper>();

    auto U_d = getDiagonalFromMatrix(U);
    Eigen::Vector3d d = U_d;
    for(int i = 0; i < 3; i++) {
        if(d[i] != 0.0) {
            d[i] = U_d[i] > 0.0 ? 1.0 : -1.0;
        }
    }
    auto D = DiagonalMatrix(d);

    Q = Q * D;
    U = D * U;

    //compute K, R, and t
    auto Q_t = Eigen::Transpose< Eigen::Matrix3d >(Q);
    auto s = Q.determinant();

    R = s * Q_t;
    t = s * U * P.col(3);

    if(U(2, 2) > 0.0) {
        U /= U(2, 2);
    }

    K = U.inverse();
}

/**
 * @brief cameraMatrixProject projects a point, p, using the camera
 * matrix, M.
 * @param M
 * @param p is a 3D point encoded in homogenous coordinate (4D vector)
 * @return
 */
PIC_INLINE Eigen::Vector2i cameraMatrixProject(Eigen::Matrix34d &M, Eigen::Vector4d &p)
{
    Eigen::Vector3d proj = M * p;
    proj[0] /= proj[2];
    proj[1] /= proj[2];

    return Eigen::Vector2i(int(proj[0]), int(proj[1]));
}

/**
 * @brief cameraMatrixProject projects a point, p, using the camera
 * matrix, M.
 * @param M
 * @param p is a 3D point (3D vector)
 * @return
 */
PIC_INLINE Eigen::Vector2i cameraMatrixProject(Eigen::Matrix34d &M, Eigen::Vector3d &p)
{
    Eigen::Vector4d p4d(p[0], p[1], p[2], 1.0);
    return cameraMatrixProject(M, p4d);
}

/**
 * @brief cameraMatrixProjection
 * @param M
 * @param p
 * @param cx
 * @param cy
 * @param fx
 * @param fy
 * @param lambda
 * @return
 */
PIC_INLINE Eigen::Vector2i cameraMatrixProjection(Eigen::Matrix34d &M, Eigen::Vector3d &p, double cx, double cy, double fx, double fy, double lambda)
{
    Eigen::Vector4d p_t = Eigen::Vector4d(p[0], p[1], p[2], 1.0);
    Eigen::Vector2i out;
    Eigen::Vector3d proj = M * p_t;
    proj[0] /= proj[2];
    proj[1] /= proj[2];

    double x_cx =  (proj[0] - cx);
    double y_cy =  (proj[1] - cy);

    double dx = x_cx / fx;
    double dy = y_cy / fy;
    double rho_sq = dx * dx + dy * dy;

    double factor = 1.0 / (1.0 + rho_sq * lambda);

    proj[0] = x_cx * factor + cx;
    proj[1] = y_cy * factor + cy;

    out[0] = int(proj[0]);
    out[1] = int(proj[1]);

    return out;
}

/**
 * @brief getOpticalCenter
 * @param P the camera matrix of a view
 * @return it returns the camera center of P
 */
PIC_INLINE Eigen::Vector3d getOpticalCenter(Eigen::Matrix34d &P)
{
    Eigen::Matrix3d Q = P.block<3, 3>(0, 0);
    auto Q_inv = Q.inverse();
    return - Q_inv * P.col(3);
}

/**
 * @brief cameraRectify
 * @param K0 intrisic matrix of view0
 * @param R0 rotation matrix of view0
 * @param t0 translation vector of view0
 * @param K1 intrisic matrix of view1
 * @param R1 rotation matrix of view1
 * @param t1 translation vector of view1
 * @param P0 new camera matrix of view0
 * @param P1 new camera matrix of view1
 * @param T0 transformation matrix for view0
 * @param T1 transformation matrix for view1
 */
PIC_INLINE void cameraRectify(Eigen::Matrix3d &K0, Eigen::Matrix3d &R0, Eigen::Vector3d &t0,
                              Eigen::Matrix3d &K1, Eigen::Matrix3d &R1, Eigen::Vector3d &t1,
                              Eigen::Matrix34d &P0_out, Eigen::Matrix34d &P1_out,
                              Eigen::Matrix3d &T0, Eigen::Matrix3d &T1)
{
    auto P0_in = getCameraMatrix(K0, R0, t0);
    auto P1_in = getCameraMatrix(K1, R1, t1);

    /*
    auto K0_i = K0.inverse();
    auto K1_i = K1.inverse();

    auto R0_t = Eigen::Transpose< Eigen::Matrix3d >(R0);
    auto R1_t = Eigen::Transpose< Eigen::Matrix3d >(R1);
    auto c0 = -R0_t * K0_i * P0t.col(3);
    auto c1 = -R1_t * K1_i * P1t.col(3);
    */

    //compute optical centers

    auto c0 = getOpticalCenter(P0_in);
    auto c1 = getOpticalCenter(P1_in);

    //compute new rotation matrix
    Eigen::Vector3d x_axis = c1 - c0;
    Eigen::Vector3d tmp = R1.row(2);
    Eigen::Vector3d y_axis = tmp.cross(x_axis);
    Eigen::Vector3d z_axis = x_axis.cross(y_axis);

    x_axis.normalize();
    y_axis.normalize();
    z_axis.normalize();

    Eigen::Matrix3d R;

    R(0, 0) = x_axis[0];
    R(0, 1) = x_axis[1];
    R(0, 2) = x_axis[2];

    R(1, 0) = y_axis[0];
    R(1, 1) = y_axis[1];
    R(1, 2) = y_axis[2];

    R(2, 0) = z_axis[0];
    R(2, 1) = z_axis[1];
    R(2, 2) = z_axis[2];

    //new camera matrices
    Eigen::Matrix3d K;
    K.setZero();
    K(0, 0) = K0(0, 0);
    K(1, 1) = K0(1, 1);
    K(0, 2) = (K0(0, 2) + K1(0, 2)) * 0.5;
    K(1, 2) = (K0(1, 2) + K1(1, 2)) * 0.5;
    K(2, 2) = 1.0;

    Eigen::Vector3d t0n = -R * c0;
    P0_out = getCameraMatrix(K, R, t0n);

    Eigen::Vector3d t1n = -R * c1;
    P1_out = getCameraMatrix(K, R, t1n);

    //transformations
    auto Q0o = P0_in.block<3, 3>(0, 0);
    auto Q0n = P0_out.block<3, 3>(0, 0);
    T0 = Q0n * Q0o.inverse();

    auto Q1o = P1_in.block<3, 3>(0, 0);
    auto Q1n = P1_out.block<3, 3>(0, 0);
    T1 = Q1n * Q1o.inverse();
}

/**
 * @brief cameraRectify
 * @param P0_in
 * @param P1_in
 * @param P0_out
 * @param P1_out
 * @param T0
 * @param T1
 */
PIC_INLINE void cameraRectify(Eigen::Matrix34d &P0_in, Eigen::Matrix34d &P1_in,
                              Eigen::Matrix34d &P0_out, Eigen::Matrix34d &P1_out,
                              Eigen::Matrix3d &T0, Eigen::Matrix3d &T1)
{
    Eigen::Matrix3d K0, K1, R0, R1;
    Eigen::Vector3d t0, t1;

    decomposeCameraMatrix(P0_in, K0, R0, t0);

    decomposeCameraMatrix(P1_in, K1, R1, t1);

    cameraRectify(K0, R0, t0,
                  K1, R1, t1,
                  P0_out, P1_out,
                  T0, T1);
}

#endif // PIC_DISABLE_EIGEN

} // end namespace pic

#endif // PIC_COMPUTER_VISION_CAMERA_MARTIX_HPP


/*
    Copyright (c) 2013, Taiga Nomi
    All rights reserved.

    Redistribution and use in source and binary forms, with or without
    modification, are permitted provided that the following conditions are met:
    * Redistributions of source code must retain the above copyright
    notice, this list of conditions and the following disclaimer.
    * Redistributions in binary form must reproduce the above copyright
    notice, this list of conditions and the following disclaimer in the
    documentation and/or other materials provided with the distribution.
    * Neither the name of the <organization> nor the
    names of its contributors may be used to endorse or promote products
    derived from this software without specific prior written permission.

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

using namespace tiny_cnn;
using namespace tiny_cnn::activation;

TEST(convolutional, fprop) {
    typedef network<mse, gradient_descent_levenberg_marquardt> CNN;
    CNN nn;

    convolutional_layer<network<mse, gradient_descent_levenberg_marquardt>, sigmoid> l(5, 5, 3, 1, 2);

    vec_t in(25);

    ASSERT_EQ(l.weight().size(), 18);

    std::fill(l.bias().begin(), l.bias().end(), 0.0);
    std::fill(l.weight().begin(), l.weight().end(), 0.0);

    uniform_rand(in.begin(), in.end(), -1.0, 1.0);

    {
        const vec_t& out = l.forward_propagation(in, 0);

        for (auto o: out)
            EXPECT_DOUBLE_EQ(o, (tiny_cnn::float_t)0.5);

    }

    l.weight()[0] = 0.3;  l.weight()[1] = 0.1; l.weight()[2] = 0.2;
    l.weight()[3] = 0.0;  l.weight()[4] =-0.1; l.weight()[5] =-0.1;
    l.weight()[6] = 0.05; l.weight()[7] =-0.2; l.weight()[8] = 0.05;

    l.weight()[9]  = 0.0; l.weight()[10] =-0.1; l.weight()[11] = 0.1;
    l.weight()[12] = 0.1; l.weight()[13] =-0.2; l.weight()[14] = 0.3;
    l.weight()[15] = 0.2; l.weight()[16] =-0.3; l.weight()[17] = 0.2;

    in[0] = 3;  in[1] = 2;  in[2] = 1;  in[3] = 5; in[4] = 2;
    in[5] = 3;  in[6] = 0;  in[7] = 2;  in[8] = 0; in[9] = 1;
    in[10] = 0; in[11] = 6; in[12] = 1; in[13] = 1; in[14] = 10;
    in[15] = 3; in[16] =-1; in[17] = 2; in[18] = 9; in[19] = 0;
    in[20] = 1; in[21] = 2; in[22] = 1; in[23] = 5; in[24] = 5;

    {
        const vec_t& out = l.forward_propagation(in, 0);

        EXPECT_NEAR(0.4875026, out[0], 1E-5);
        EXPECT_NEAR(0.8388910, out[1], 1E-5);
        EXPECT_NEAR(0.8099984, out[2], 1E-5);
        EXPECT_NEAR(0.7407749, out[3], 1E-5);
        EXPECT_NEAR(0.5000000, out[4], 1E-5);
        EXPECT_NEAR(0.1192029, out[5], 1E-5);
        EXPECT_NEAR(0.5986877, out[6], 1E-5);
        EXPECT_NEAR(0.7595109, out[7], 1E-5);
        EXPECT_NEAR(0.6899745, out[8], 1E-5);
    }

}

TEST(convolutional, bprop) {
    network<cross_entropy, gradient_descent_levenberg_marquardt> nn;
    convolutional_layer<network<cross_entropy, gradient_descent_levenberg_marquardt>, sigmoid> layer(5, 5, 3, 1, 1);

    nn.add(&layer);

    vec_t a(25, 0.0), t(9, 0.0);
    std::vector<vec_t> data, train;

    for (int y = 0; y < 5; y++) {
        a[5*y+3] = 1.0;
    }

    for (int y = 0; y < 3; y++) {
        t[3*y+0] = 0.0;
        t[3*y+1] = 0.5;
        t[3*y+2] = 1.0;
    }

    for (int i = 0; i < 100; i++) {
        data.push_back(a);
        train.push_back(t);
    }

    nn.train(data, train);

    vec_t predicted;
    nn.predict(a, &predicted);
}

TEST(convolutional, gradient_check) { // tanh - mse
    network<mse, gradient_descent_levenberg_marquardt> nn;
    convolutional_layer<network<mse, gradient_descent_levenberg_marquardt>, tan_h> layer(5, 5, 3, 1, 1);
    nn.add(&layer);

    vec_t a(25, 0.0);
    label_t t = 3;

    uniform_rand(a.begin(), a.end(), -1, 1);
    nn.init_weight();
    EXPECT_TRUE(nn.gradient_check(&a, &t, 1, 1e-4, GRAD_CHECK_ALL));
}

TEST(convolutional, gradient_check2) { // sigmoid - mse
    network<mse, gradient_descent_levenberg_marquardt> nn;
    convolutional_layer<network<mse, gradient_descent_levenberg_marquardt>, sigmoid> layer(5, 5, 3, 1, 1);
    nn.add(&layer);

    vec_t a(25, 0.0);
    label_t t = 3;

    uniform_rand(a.begin(), a.end(), -1, 1);
    nn.init_weight();
    EXPECT_TRUE(nn.gradient_check(&a, &t, 1, 1e-4, GRAD_CHECK_ALL));
}

TEST(convolutional, gradient_check3) { // rectified - mse
    network<mse, gradient_descent_levenberg_marquardt> nn;
    convolutional_layer<network<mse, gradient_descent_levenberg_marquardt>, rectified_linear> layer(5, 5, 3, 1, 1);
    nn.add(&layer);

    vec_t a(25, 0.0);
    label_t t = 3;

    uniform_rand(a.begin(), a.end(), -1, 1);
    nn.init_weight();
    EXPECT_TRUE(nn.gradient_check(&a, &t, 1, 1e-4, GRAD_CHECK_ALL));
}

TEST(convolutional, gradient_check4) { // identity - mse
    network<mse, gradient_descent_levenberg_marquardt> nn;
    convolutional_layer<network<mse, gradient_descent_levenberg_marquardt>, identity> layer(5, 5, 3, 1, 1);
    nn.add(&layer);

    vec_t a(25, 0.0);
    label_t t = 3;

    uniform_rand(a.begin(), a.end(), -1, 1);
    nn.init_weight();
    EXPECT_TRUE(nn.gradient_check(&a, &t, 1, 1e-4, GRAD_CHECK_ALL));
}

TEST(convolutional, gradient_check5) { // sigmoid - cross-entropy
    network<cross_entropy, gradient_descent_levenberg_marquardt> nn;
    convolutional_layer<network<cross_entropy, gradient_descent_levenberg_marquardt>, sigmoid> layer(5, 5, 3, 1, 1);
    nn.add(&layer);

    vec_t a(25, 0.0);
    label_t t = 3;

    uniform_rand(a.begin(), a.end(), -1, 1);
    nn.init_weight();
    EXPECT_TRUE(nn.gradient_check(&a, &t, 1, 1e-4, GRAD_CHECK_ALL));
}

TEST(convolutional, bprop2) {
    network<cross_entropy, gradient_descent_levenberg_marquardt> nn;
    convolutional_layer<network<cross_entropy, gradient_descent_levenberg_marquardt>, sigmoid> layer(5, 5, 3, 1, 1);
    fully_connected_layer<network<cross_entropy, gradient_descent_levenberg_marquardt>, sigmoid> layer2(9, 3);

    nn.add(&layer);
    nn.add(&layer2);

    vec_t a(25, 0.0), t(3, 0.0), a2(25, 0.0), t2(3, 0.0);

    for (int y = 0; y < 5; y++) {
        a[5*y+3] = 1.0;
    }

    t[0] = 0.0;
    t[1] = 0.5;
    t[2] = 1.0;

    uniform_rand(a2.begin(), a2.end(), -3, 3);
    uniform_rand(t2.begin(), t2.end(), 0, 1);

    std::vector<vec_t> data, train;

    for (int i = 0; i < 300; i++) {
        data.push_back(a);
        data.push_back(a2);
        train.push_back(t);
        train.push_back(t2);
    }
    nn.train(data, train);

    vec_t predicted;
    nn.predict(a, &predicted);
    nn.predict(a2, &predicted);
}

TEST(convolutional, serialize) {
    typedef network<mse, gradient_descent_levenberg_marquardt> NN;
    convolutional_layer<NN, tan_h> layer1(14, 14, 5, 1, 2);
    convolutional_layer<NN, tan_h> layer2(14, 14, 5, 1, 2);

    vec_t v(14*14);

    uniform_rand(v.begin(), v.end(), -1.0, 1.0);
    layer1.init_weight();

    std::ostringstream os;
    layer1.save(os);

    std::istringstream is(os.str());
    layer2.load(is);

    const vec_t& out1 = layer1.forward_propagation(v, 0);
    const vec_t& out2 = layer2.forward_propagation(v, 0);

    for (size_t i = 0; i < out1.size(); i++)
        EXPECT_NEAR(out1[i], out2[i], 1e-4);
}

TEST(convolutional, serialize2) {
    typedef network<mse, gradient_descent_levenberg_marquardt> NN;
#define O true
#define X false
    static const bool connection[] = {
        O, X, X, X, O, O,
        O, O, X, X, X, O,
        O, O, O, X, X, X
    };
#undef O
#undef X
    convolutional_layer<NN, tan_h> layer1(14, 14, 5, 3, 6, connection_table(connection, 3, 6));
    convolutional_layer<NN, tan_h> layer2(14, 14, 5, 3, 6, connection_table(connection, 3, 6));

    vec_t v(14*14*3);

    uniform_rand(v.begin(), v.end(), -1.0, 1.0);
    layer1.init_weight();

    std::ostringstream os;
    layer1.save(os);

    std::istringstream is(os.str());
    layer2.load(is);

    const vec_t& out1 = layer1.forward_propagation(v, 0);
    const vec_t& out2 = layer2.forward_propagation(v, 0);

    for (size_t i = 0; i < out1.size(); i++)
        EXPECT_NEAR(out1[i], out2[i], 1e-4);
}

TEST(fully_connected, bprop) {
    network<cross_entropy, gradient_descent_levenberg_marquardt> nn;
    fully_connected_layer<network<cross_entropy, gradient_descent_levenberg_marquardt>, sigmoid> layer(3, 2);
    nn.add(&layer);

    vec_t a(3), t(2), a2(3), t2(2);

    a[0] = 3.0; a[1] = 0.0; a[2] = -1.0;
    t[0] = 0.3; t[1] = 0.7;

    a2[0] = 0.2; a2[1] = 0.5; a2[2] = 4.0;
    t2[0] = 0.5; t2[1] = 0.1;

    std::vector<vec_t> data, train;

    for (int i = 0; i < 100; i++) {
        data.push_back(a);
        data.push_back(a2);
        train.push_back(t);
        train.push_back(t2);
    }
    nn.optimizer().alpha = 0.1;
    nn.train(data, train, 1, 10);

    vec_t predicted;
    nn.predict(a, &predicted);

    EXPECT_NEAR(predicted[0], t[0], 1E-5);
    EXPECT_NEAR(predicted[1], t[1], 1E-5);

    nn.predict(a2, &predicted);

    EXPECT_NEAR(predicted[0], t2[0], 1E-5);
    EXPECT_NEAR(predicted[1], t2[1], 1E-5);
}

TEST(fully_connected, serialize) {
    typedef network<mse, gradient_descent_levenberg_marquardt> NN;
    fully_connected_layer<NN, tan_h> layer1(10, 10);
    fully_connected_layer<NN, tan_h> layer2(10, 10);

    vec_t v(10);

    uniform_rand(v.begin(), v.end(), -1.0, 1.0);
    layer1.init_weight();

    std::ostringstream os;
    layer1.save(os);

    std::istringstream is(os.str());
    layer2.load(is);

    const vec_t& out1 = layer1.forward_propagation(v, 0);
    const vec_t& out2 = layer2.forward_propagation(v, 0);

    for (size_t i = 0; i < out1.size(); i++)
        EXPECT_NEAR(out1[i], out2[i], 1e-4);
}

TEST(fully_connected, bprop2) {
    network<mse, gradient_descent> nn;
    fully_connected_layer<network<mse, gradient_descent>, tan_h> layer(4, 6);
    fully_connected_layer<network<mse, gradient_descent>, tan_h> layer2(6, 3);

    nn.add(&layer);
    nn.add(&layer2);

    vec_t a(4, 0.0), t(3, 0.0), a2(4, 0.0), t2(3, 0.0);

    a[0] = 3.0; a[1] = 1.0; a[2] = -1.0; a[3] = 4.0;
    t[0] = 0.3; t[1] = 0.7; t[2] = 0.3;

    a2[0] = 1.0; a2[1] = 0.0; a2[2] = 4.0; a2[3] = 2.0;
    t2[0] = 0.6; t2[1] = 0.0; t2[2] = 0.1;

    std::vector<vec_t> data, train;

    for (int i = 0; i < 100; i++) {
        data.push_back(a);
        data.push_back(a2);
        train.push_back(t);
        train.push_back(t2);
    }
    nn.optimizer().alpha = 0.01;
    nn.train(data, train, 1, 10);

    vec_t predicted;
    nn.predict(a, &predicted);

    EXPECT_NEAR(predicted[0], t[0], 1E-4);
    EXPECT_NEAR(predicted[1], t[1], 1E-4);

    nn.predict(a2, &predicted);

    EXPECT_NEAR(predicted[0], t2[0], 1E-4);
    EXPECT_NEAR(predicted[1], t2[1], 1E-4);
}

TEST(multi_layer, gradient_check) { // sigmoid - cross-entropy
    typedef cross_entropy loss_func;
    typedef sigmoid activation;
    typedef network<loss_func, gradient_descent_levenberg_marquardt> network;

    network nn;
    fully_connected_layer<network, activation> l1(10, 14*14*3);
    convolutional_layer<network, activation>   l2(14, 14, 5, 3, 6);
    average_pooling_layer<network, activation> l3(10, 10, 6, 2);
    fully_connected_layer<network, activation> l4(5*5*6, 3);

    nn.add(&l1);
    nn.add(&l2);
    nn.add(&l3);
    nn.add(&l4);

    vec_t a(10, 0.0);
    label_t t = 2;

    uniform_rand(a.begin(), a.end(), -1, 1);
    nn.init_weight();
    EXPECT_TRUE(nn.gradient_check(&a, &t, 1, 1e-4, GRAD_CHECK_RANDOM));
}

TEST(multi_layer, gradient_check2) { // tan_h - mse
    typedef mse loss_func;
    typedef tan_h activation;
    typedef network<loss_func, gradient_descent_levenberg_marquardt> network;

    network nn;
    fully_connected_layer<network, activation> l1(10, 14*14*3);
    convolutional_layer<network, activation>   l2(14, 14, 5, 3, 6);
    average_pooling_layer<network, activation> l3(10, 10, 6, 2);
    fully_connected_layer<network, activation> l4(5*5*6, 3);

    nn.add(&l1);
    nn.add(&l2);
    nn.add(&l3);
    nn.add(&l4);

    vec_t a(10, 0.0);
    label_t t = 2;

    uniform_rand(a.begin(), a.end(), -1, 1);
    nn.init_weight();
    EXPECT_TRUE(nn.gradient_check(&a, &t, 1, 1e-4, GRAD_CHECK_RANDOM));
}

TEST(multi_layer, gradient_check3) { // mixture - mse
    typedef mse loss_func;
    typedef network<loss_func, gradient_descent_levenberg_marquardt> network;

    network nn;
    fully_connected_layer<network, tan_h>     l1(10, 14*14*3);
    convolutional_layer<network, sigmoid>    l2(14, 14, 5, 3, 6);
    average_pooling_layer<network, rectified_linear>    l3(10, 10, 6, 2);
    fully_connected_layer<network, identity> l4(5*5*6, 3);

    nn.add(&l1);
    nn.add(&l2);
    nn.add(&l3);
    nn.add(&l4);

    vec_t a(10, 0.0);
    label_t t = 2;

    uniform_rand(a.begin(), a.end(), -1, 1);
    nn.init_weight();
    EXPECT_TRUE(nn.gradient_check(&a, &t, 1, 1e-4, GRAD_CHECK_RANDOM));
}

TEST(multi_layer4, gradient_check) { // sigmoid - cross-entropy
    typedef cross_entropy loss_func;
    typedef sigmoid activation;
    typedef network<loss_func, gradient_descent_levenberg_marquardt> network;

    network nn;
    fully_connected_layer<network, activation> l1(10, 14 * 14 * 3);
    convolutional_layer<network, activation>   l2(14, 14, 5, 3, 6);
    average_pooling_layer<network, activation> l3(10, 10, 6, 2);
    fully_connected_dropout_layer<network, activation> l4(5 * 5 * 6, 3);

    nn.add(&l1);
    nn.add(&l2);
    nn.add(&l3);
    nn.add(&l4);

    vec_t a(10, 0.0);
    label_t t = 2;

    uniform_rand(a.begin(), a.end(), -1, 1);
    nn.init_weight();
    EXPECT_TRUE(nn.gradient_check(&a, &t, 1, 1e-4, GRAD_CHECK_RANDOM));
}

template <typename N>
void serialization_test(const layer_base<N>& src, layer_base<N>& dst)
{
    EXPECT_FALSE(src.has_same_weights(dst, 1E-5));

    boost::filesystem::path tmp_path = boost::filesystem::unique_path();

    if (boost::filesystem::exists(tmp_path))
        throw nn_error("file exists");

    std::string tmp_file_path = tmp_path.string();

    // write
    {
        std::ofstream ofs(tmp_file_path.c_str());
        ofs << src;
    }

    // read
    {
        std::ifstream ifs(tmp_file_path.c_str());
        ifs >> dst;
    }

    boost::filesystem::remove(tmp_path); // remove temporary file

    EXPECT_TRUE(src.has_same_weights(dst, 1E-5));
}

TEST(read_write, fully_connected)
{
	typedef mse loss_func;
	typedef network<loss_func, gradient_descent_levenberg_marquardt> network;

	fully_connected_layer<network, tan_h> l1(100, 100);
	fully_connected_layer<network, tan_h> l2(100, 100);

	l1.init_weight();
	l2.init_weight();

    serialization_test(l1, l2);
}

TEST(read_write, convolutional)
{
    typedef mse loss_func;
    typedef network<loss_func, gradient_descent_levenberg_marquardt> network;

    convolutional_layer<network, tan_h> l1(5, 5, 3, 1, 1);
    convolutional_layer<network, tan_h> l2(5, 5, 3, 1, 1);

    l1.init_weight();
    l2.init_weight();

    serialization_test(l1, l2);
}

int main(void) {
    RUN_ALL_TESTS();
}