#include <iostream>
#include <iomanip>
#include <fstream>
#include <iterator>
#include <cmath>
#include <regex>
#include <numeric>

#include <boost/program_options.hpp>
#include <boost/log/sources/logger.hpp>
#include <boost/archive/binary_iarchive.hpp>
#include <boost/archive/binary_oarchive.hpp>
#include <boost/serialization/map.hpp>
#include <boost/mpi.hpp>
#include <boost/core/demangle.hpp>
#include <boost/numeric/conversion/cast.hpp>
#include <boost/algorithm/string/trim.hpp>
#include <boost/algorithm/string/split.hpp>
#include <boost/range/adaptors.hpp>
#include <boost/range/algorithm.hpp>
#include <boost/range/combine.hpp>

#include <benchmark.hpp>
#include <logging.hpp>

#include <blas.hpp>
#include <matrix_serialization.hpp>

namespace po = boost::program_options;

auto init_options() -> po::options_description
{
	po::options_description description;
	description.add_options()
		( "help", "Produce the help message." )
		( "type"
		, po::value<std::string>()->required()
		, "Type of an element of the epsilon matrix. It may be "
		  "either cfloat or cdouble." )
		( "epsilon"
		, po::value<std::string>()->required()
		, "File to where epsilon matrix was saved to." )
		( "positions"
		, po::value<std::string>()->required()
		, "File to where atomic site positions were saved to." )
		( "q"
		, po::value<std::string>()->required()
		, "List of |q|s separated by commas. MIND YOU: "
		  "no spaces!" )
		( "direction"
		, po::value<std::string>()->required()
		, "Direction of q as (x,y,z). It is automatically"
		  "normalized." );
	return description;
}

auto element_type(std::string input) -> std::type_index
{
	using namespace std::string_literals;
	static std::unordered_map<std::string, std::type_index> const types =
		{ { "cfloat"s,  std::type_index(typeid(std::complex<float>))  }
		, { "cdouble"s, std::type_index(typeid(std::complex<double>)) }
		};

	boost::to_lower(input);
	try {
		return types.at(input);
	} catch(std::out_of_range & e) {
		std::cerr << "Invalid element type `" + input + "`!\n";
		throw;
	}
}

template <class _R>
auto parse_direction(std::string str) -> std::array<_R, 3>
{
	using namespace boost;
	using namespace boost::algorithm;

	auto const normalize = [](_R const x, _R const y, _R const z)
		-> std::array<_R, 3> {
		auto const _abs = _R{1.0} / std::sqrt(std::norm(x) + std::norm(y) + std::norm(z));
		return {_abs * x, _abs * y, _abs * z};
	};

	std::regex vector_re{"\\((.*),(.*),(.*)\\)"};
	std::smatch results;

	trim(str);
	if (std::regex_match(str, results, vector_re)) {
		assert(results.size() == 4);
		return normalize( lexical_cast<_R>(trim_copy(results[1].str()))
		                , lexical_cast<_R>(trim_copy(results[2].str()))
			            , lexical_cast<_R>(trim_copy(results[3].str())) );
	}
	throw std::invalid_argument{ "Could not convert '" + str
	                           + "' to a 3D vector."};
}

template <class _R>
auto parse_qs(std::string const& str) -> std::vector<_R>
{
	using namespace boost;
	using namespace boost::adaptors;
	using namespace boost::algorithm;

	auto const parse_number = [](auto const& s) {
		return lexical_cast<_R>(trim_copy(s));
	};

	std::vector<std::string> tokens;
	split(tokens, str, [](auto const ch) { return ch == ','; });

	std::vector<_R> qs;
	qs.reserve(tokens.size());
	copy( tokens | transformed(std::cref(parse_number))
	    , std::back_inserter(qs) );
	return qs;
}

template<class _Help, class _Run>
auto process_command_line( int argc, char** argv
                         , _Help&& help
						 , _Run&& run ) -> void
{
	auto const description = init_options();
	po::variables_map vm;

	po::store( po::command_line_parser(argc, argv)
	              .options(description)
	              .run()
	         , vm );

	if (vm.count("help")) {
		help(description);
		return;
	}

	po::notify(vm);
	run(vm);
}

template<class _T>
auto load_matrix(std::string const& file_name) -> tcm::Matrix<_T>
{
	std::ifstream in_stream{file_name};
	if (not in_stream)
		throw std::runtime_error{"Failed to open `" + file_name + "`."};
	boost::archive::binary_iarchive in_archive{in_stream};

	tcm::Matrix<_T> A;
	in_archive >> A;
	return A;
}

template<class _T>
auto read_positions(std::string const& file_name)
	-> std::vector<std::array<_T, 3>>
{
	std::ifstream in_stream{file_name};
	if (not in_stream)
		throw std::runtime_error{"Failed to open `" + file_name + "`."};

	std::vector<std::array<_T, 3>> positions;
	auto i = std::istream_iterator<_T>{in_stream};
	while(i != std::istream_iterator<_T>{}) {
		const auto x = *i++;
		const auto y = *i++;
		const auto z = *i++;
		positions.push_back({x, y, z});
	}

	return positions;
}

template < std::size_t _Dim = 3
         , class _R1 = double
         , class _R2 = double
         , class _R3 = std::common_type_t<_R1, _R2>
         >
auto make_momentum_eigenvector( std::array<_R1, 3> const wavevector
                              , std::vector<std::array<_R2, 3>> const& positions
                              , _R3 const pi = _R3{M_PI} )
{
	using _R = std::common_type_t<_R1, _R2, _R3>;
	using _C = std::complex<_R>;

	auto constexpr I = _C{0, 1};
	auto const _norm =
		std::pow(_R{2} * pi, -boost::numeric_cast<_R>(_Dim) / _R{2});
	auto const _dot = [] (auto const& x, auto const& y) noexcept {
		return x[0]*y[0] + x[1]*y[1] + x[2]*y[2];
	};

	tcm::Matrix<_C> q{positions.size(), 1};
	std::transform( std::begin(positions), std::end(positions)
	              , q.data()
				  , [ _norm
	                , &wavevector
	                , _dot = std::cref(_dot)
	                ] (auto const& r) {
	                    return _norm * std::exp<_R>(I * _dot(wavevector, r));
	                } );
	return q;
}

template <class _R, class _Matrix>
auto loss_function( std::array<_R, 3> const& direction
                  , std::vector<_R> const& qs
                  , _Matrix const& epsilon
                  , std::vector<std::array<_R, 3>> const& positions )
{
	using namespace boost;
	using namespace boost::adaptors;

	assert(tcm::is_square(epsilon));
	assert(epsilon.height() == positions.size());
	using _C = typename _Matrix::value_type;
	static_assert( std::is_same<typename _C::value_type, _R>::value
	             , "Types mismatch." );

	auto const make_wavevector = [&direction](auto const q) noexcept
		-> std::array<_R, 3> {
		return {q * direction[0], q * direction[1], q * direction[2]};
	};

	auto const make_epsilon_q = [&epsilon,&positions](auto const& wavevector) {
		auto const q_state = make_momentum_eigenvector(wavevector, positions);
		tcm::Matrix<_C> _temp{epsilon.height(), 1};
		tcm::blas::gemv( tcm::blas::Operator::H
		               , _C{1}, epsilon, q_state
		               , _C{0}, _temp );
		return tcm::blas::dot(_temp, q_state);
	};

	std::vector<_C> epsilon_q;
	epsilon_q.reserve(qs.size());
	copy( qs | transformed(std::cref(make_wavevector))
	         | transformed(std::cref(make_epsilon_q))
	    , std::back_inserter(epsilon_q) );
	return epsilon_q;
}

template <class _C>
auto run(boost::program_options::variables_map const& vm) -> void
{
	using namespace boost;
	using _R = typename _C::value_type;

	auto const direction =
		parse_direction<_R>(vm["direction"].as<std::string>());
	auto const qs =
		parse_qs<_R>(vm["q"].as<std::string>());
	auto const positions =
		read_positions<_R>(vm["positions"].as<std::string>());
	auto const epsilon =
		load_matrix<_C>(vm["epsilon"].as<std::string>());

	auto const epsilon_q = loss_function( direction, qs
	                                    , epsilon
	                                    , positions );

	auto const print = [](auto const& t) {
		std::cout << tuples::get<0>(t) << '\t'
		          << std::real(tuples::get<1>(t)) << '\t'
		          << std::imag(tuples::get<1>(t)) << '\n';
	};

	std::cout << std::scientific << std::setprecision(15);
	for_each(combine(qs, epsilon_q), std::cref(print));
}

auto dispatch(boost::program_options::variables_map const& vm) -> void
{
	auto const type = element_type(vm["type"].as<std::string>());
	if (type == typeid(std::complex<float>)) {
		run<std::complex<float>>(vm);
	}
	else if (type == typeid(std::complex<double>)) {
		run<std::complex<double>>(vm);
	}
	else throw std::invalid_argument{"Wrong type."};
}

int main(int argc, char** argv)
{
	process_command_line
		( argc, argv
		, [](auto desc) { std::cout << desc << '\n'; }
		, &dispatch
		);
	return EXIT_SUCCESS;
}


/* Copyright (C) 2015 Ion Torrent Systems, Inc. All Rights Reserved */

/* Author: Alex Artyomenko <aartyomenko@cs.gsu.edu> */

#ifndef ION_ANALYSIS_VITERBI_HPP
#define ION_ANALYSIS_VITERBI_HPP

#include <stdlib.h>
#include <vector>
#include <algorithm>
#include <numeric>
#include <boost/math/distributions/poisson.hpp>
#include <boost/algorithm/minmax_element.hpp>

template<typename T>
template<typename _ForwardIterator>
markov_chain<T>::markov_chain(_ForwardIterator begin, _ForwardIterator end) {
  pair<_ForwardIterator, _ForwardIterator> m = boost::minmax_element(begin, end);
  initialize(begin, end, *(m.first), *(m.second));
}

template<typename T>
template<typename _ForwardIterator>
markov_chain<T>::markov_chain(_ForwardIterator begin, _ForwardIterator end, T min, T max) {
  initialize(begin, end, min, max);
}

template<typename T>
template<typename _ForwardIterator>
void markov_chain<T>::initialize(_ForwardIterator begin, _ForwardIterator end, T min, T max) {
  // initialize a constructor delegation method
  size_t s = (size_t)distance(begin, end);
  T dmax = (T) max < MAX_STATE_VALUE ? max : MAX_STATE_VALUE,
    dmin = (T) min < MAX_STATE_VALUE ? min : MAX_STATE_VALUE;
  if (dmin == dmax) {
    T avg = (T) round(accumulate(begin, end, 0.0) / s);
    items.push_back(make_pair(depth_info<T>(avg, min, max), s - 1));
    return;
  }
  values.reserve(s);
  reserve(s);
  generate_states(dmin, dmax);
  for (_ForwardIterator it = begin; it != end; it++)
    process_next_value(*it);
  optimal_path();
}

template<typename T>
void markov_chain<T>::generate_states(T min, T max) {
  // prepare state values and state objects for viterbi execution
  double dmax = (double) max,
      dmin = (double) min,
      alpha = 2,
      val;
  val = ceil(pow(dmin, 1 / alpha));

  states.push_back(dmin);
  for (double st = pow(val, alpha); st < dmax; st = pow(val, alpha), val++) {
    states.push_back(st);
  }
  states.push_back(dmax);
}

template<typename T>
void markov_chain<T>::process_next_value(T val) {
  // accept next value and fill out column in table of states
  vector<long> states_trace;
  states_trace.reserve(states.size());
  for (vector<markov_state>::iterator st = states.begin(); st != states.end(); st++) {
    markov_state& current = *st;
    markov_chain_comparator mccmp = markov_chain_comparator(current);
    vector<markov_state>::iterator min_prev = min_element(states.begin(), states.end(), mccmp);
    states_trace.push_back(min_prev - states.begin());
    double new_penalty = min_prev->prev + current.cost(*min_prev) + current.cost(val);
    st->current = new_penalty;
  }
  for (vector<markov_state>::iterator st = states.begin(); st != states.end(); st++) st->step_forward();
  push_back(states_trace);
  values.push_back(val);
}

template<typename T>
void markov_chain<T>::optimal_path() {
  // find optimal state switches and generate items for a chain

  double mean = 0;
  int i = 0;

  items.reserve(size());
  long state_index = min_element(states.begin(), states.end(), markov_chain_comparator()) - states.begin();

  markov_chain::reverse_iterator it = rbegin();
  typename vector<T>::reverse_iterator vit = values.rbegin();
  items.push_back(make_pair(depth_info<T>(*vit), size() - 1));
  for (; it != rend() && vit != values.rend(); it++, vit++) {
    mean += *vit;
    i++;
    if (vit == values.rend() - 1) {
      items.back().first.dp = (T)round(mean / i);
    } else if ((*it)[state_index] != state_index){
      items.back().first.dp = (T)round(mean / i);
      mean = 0;
      i = 0;
      items.push_back(make_pair(depth_info<T>(*(vit+1)), rend() - it - 2));
    } else {
      depth_info<T>& cdi = items.back().first;
      cdi.max_dp = max(*vit, cdi.max_dp);
      cdi.min_dp = min(*vit, cdi.min_dp);
    }
    state_index = (*it)[state_index];
  }
}

#endif //ION_ANALYSIS_VITERBI_HPP


#include <boost/program_options.hpp>
#include <string>
#include <d3d_base/grid.h>
#include <iostream>
#include <d3d_base/configFile.h>
#include <d3d_fusion/volumetricFusionTools.h>
#include <algorithm>

struct Hist {
    float counts[8];
};

D3D::Grid<float> runTVHist(D3D::Grid<Hist> dataCost, Eigen::Vector3f& minCorner, Eigen::Vector3f& size, Eigen::Matrix4f boxToGlobal, Eigen::Vector3f color, std::string vrmlOutputFile,
                                const int numIter,
                                const float lambda, const float theta)
{
    D3D::Grid<float> u(dataCost.getWidth(), dataCost.getHeight(), dataCost.getDepth());
    for (unsigned z = 0; z < u.getDepth(); z++)
        for (unsigned y = 0; y < u.getHeight(); y++)
            for (unsigned x = 0; x < u.getWidth(); x++)
            {
                u(x,y,z) = 0.f;
            }

    D3D::Grid<float> uBar = u.clone();

    D3D::Grid<float> p1(dataCost.getWidth(), dataCost.getHeight(), dataCost.getDepth(), 0);
    D3D::Grid<float> p2(dataCost.getWidth(), dataCost.getHeight(), dataCost.getDepth(), 0);
    D3D::Grid<float> p3(dataCost.getWidth(), dataCost.getHeight(), dataCost.getDepth(), 0);

    const int xDim = u.getWidth();
    const int yDim = u.getHeight();
    const int zDim = u.getDepth();

    // uses diagonal preconditioned first order pirmal dual

    const float sigma = 0.99f/2.f;
    const float tau = 0.99f/6.f;
    const float centers[8] = {-1.f, -5.f/7.f, -3.f/7.f, -1.f/7.f, 1.f/7.f, 3.f/7.f, 5.f/7.f, 1.f};

    for (int c = 0; c < numIter; c++)
    {
        std::cout << ".";
        std::cout.flush();

        // update dual variables
#pragma omp parallel for
        for (int z = 0; z < zDim; z++)
        {
            int const Z1 = (z < zDim-1) ? (z+1) : z;
            for (int y = 0; y < yDim; y++)
            {
                int const Y1 = (y < yDim-1) ? (y+1) : y;
                for (int x = 0; x < xDim; x++)
                {
                    float const u0 = uBar(x,y,z);

                    int const X1 = (x < xDim -1) ? (x+1) : x;

                    float const u_x = (uBar(X1, y, z) - u0);
                    float const u_y = (uBar(x, Y1, z) - u0);
                    float const u_z = (uBar(x, y, Z1) - u0);

                    float const d_x = p1(x,y,z) + sigma*u_x;
                    float const d_y = p2(x,y,z) + sigma*u_y;
                    float const d_z = p3(x,y,z) + sigma*u_z;

                    float const tv = sqrtf(d_x*d_x + d_y*d_y + d_z*d_z);

                    float const denom = std::max(1.f, tv);

                    p1(x,y,z) = d_x/denom;
                    p2(x,y,z) = d_y/denom;
                    p3(x,y,z) = d_z/denom;
                }
            }
        }

        // update primal
#pragma omp parallel for
        for (int z = 0; z < zDim; z++)
        {
            int const Z0 = z-1;
            for (int y = 0; y < yDim; y++)
            {
                int const Y0 = y-1;
                for (int x = 0; x < xDim; x++)
                {
                    int const X0 = x-1;

                    float const div = (X0 < xDim -1 ? p1(x,y,z) : 0.0f) - (X0 >= 0 ? p1(X0, y, z) : 0.0f) +
                                      (Y0 < yDim -1 ? p2(x,y,z) : 0.0f) - (Y0 >= 0 ? p2(x, Y0, z) : 0.0f) +
                                      (Z0 < zDim -1 ? p3(x,y,z) : 0.0f) - (Z0 >= 0 ? p3(x, y, Z0) : 0.0f);

                    float const u0 = u(x,y,z);
                    float U = u(x,y,z) + tau*div;

//                    // check most likely bin, is in [1, 8] as u0 in [-1, 1]
//                    int mlBin = int(7.f * (u0 + 1.f) / 2.f + 1.f);

//                    float step = 0.f;
//                    for (int b = 0; b <= mlBin; ++b) {
//                        step -= float(dataCost(x, y, z).counts[b]);
//                    }
//                    for (int b = mlBin+1; b < 10; ++b) {
//                        step += float(dataCost(x, y, z).counts[b]);
//                    }

//                    float proxMin = u0 + tau*lambda*step;

                    float proxMin, step;

//                    if (proxMin <= centers[mlBin] || centers[mlBin+1] <= proxMin) {
                        // outside of range
                        bool found = false;
                        for (int bin = 0; bin < 7; bin++) {
//                            if (bin == mlBin)
//                                continue;
                            step = 0.f;
                            for (int b = 0; b <= bin; ++b) {
                                step -= dataCost(x, y, z).counts[b];
                            }
                            for (int b = bin+1; b < 8; ++b) {
                                step += dataCost(x, y, z).counts[b];
                            }

                            proxMin = U + tau*lambda*step;

                            if (centers[bin] < proxMin && proxMin < centers[bin+1]) {
                                // inside of range, val is the minimum
                                found = true;
                                break;
                            }
                        }
                        if (!found) {
                            // minimum must be one of the bin centers
                            float proxMinEnergy = std::numeric_limits<float>::max();
                            for (int bin = 0; bin < 8; bin++) {
                                float dataterm = 0.f;
                                for (int i = 0; i < 8; ++i) {
                                    dataterm += dataCost(x, y, z).counts[i] * fabs(centers[bin]-centers[i]);
                                }
                                float proxEnergy = (U-centers[bin])*(U-centers[bin])/(2.f*tau) + lambda*dataterm;
                                if (proxEnergy < proxMinEnergy) {
                                    proxMinEnergy = proxEnergy;
                                    proxMin = centers[bin];
                                }
                            }
                        }
//                    }

                    U = proxMin;

                    u(x,y,z) = U;
                    uBar(x,y,z) = U + theta*(U - u0);

                }
            }
        }

        if ((c+1) % 50 == 0) {
            D3D::saveVolumeAsVRMLMesh(u, 0.0f, minCorner, size, boxToGlobal, color, vrmlOutputFile, true);
            std::cout << std::endl;
        }
    }
    return u;
}

int main(int argc, char* argv[])
{
    std::string configFile;
    std::string intDepthGridFile;
    std::string vrmlOutputFile;

    boost::program_options::options_description desc("Allowed options");
    desc.add_options()
            ("help", "Produce help message")
            ("intDepthGridFile", boost::program_options::value<std::string>(&intDepthGridFile)->default_value("intCSDFGrid.dat"), "Depth map list file")
            ("configFile", boost::program_options::value<std::string>(&configFile)->default_value("conf.txt"), "Config file")
            ("vrmlOutputFile", boost::program_options::value<std::string>(&vrmlOutputFile)->default_value("tvHistModel.wrl"), "Output data file with integrated depth data.")
            ;

    boost::program_options::variables_map vm;
    boost::program_options::store(boost::program_options::command_line_parser(argc, argv).options(desc).run(), vm);
    boost::program_options::notify(vm);

    if (vm.count("help"))
    {
        std::cout << desc << std::endl;
        return 1;
    }

    // config file
    D3D::ConfigFile conf(configFile);

    if (!conf.isFileRead())
    {
        D3D_THROW_EXCEPTION("Could not open config file.")
    }

    // load the box parameters
    std::string bpMinXStr = conf.get("BP_MIN_X");
    if (bpMinXStr.empty())
    {
        D3D_THROW_EXCEPTION("Parameter BP_MIN_X not specified in the config file.")
    }
    float bpMinX = atof(bpMinXStr.c_str());

    std::string bpSizeXStr = conf.get("BP_SIZE_X");
    if (bpSizeXStr.empty())
    {
        D3D_THROW_EXCEPTION("Parameter BP_SIZE_X not specified in the config file.")
    }
    float bpSizeX = atof(bpSizeXStr.c_str());

    std::string bpMinYStr = conf.get("BP_MIN_Y");
    if (bpMinYStr.empty())
    {
        D3D_THROW_EXCEPTION("Parameter BP_MIN_Y not specified in the config file.")
    }
    float bpMinY = atof(bpMinYStr.c_str());

    std::string bpSizeYStr = conf.get("BP_SIZE_Y");
    if (bpSizeYStr.empty())
    {
        D3D_THROW_EXCEPTION("Parameter BP_SIZE_Y not specified in the config file.")
    }
    float bpSizeY = atof(bpSizeYStr.c_str());

    std::string bpMinZStr = conf.get("BP_MIN_Z");
    if (bpMinZStr.empty())
    {
        D3D_THROW_EXCEPTION("Parameter BP_MIN_Z not specified in the config file.")
    }
    float bpMinZ = atof(bpMinZStr.c_str());

    std::string bpSizeZStr = conf.get("BP_SIZE_Z");
    if (bpSizeZStr.empty())
    {
        D3D_THROW_EXCEPTION("Parameter BP_SIZE_Z not specified in the config file.")
    }
    float bpSizeZ = atof(bpSizeZStr.c_str());

    std::string bpResXStr = conf.get("BP_RES_X");
    if (bpResXStr.empty())
    {
        D3D_THROW_EXCEPTION("Parameter BP_RES_X not specified in the config file.")
    }
    int bpResX = atoi(bpResXStr.c_str());

    std::string bpResYStr = conf.get("BP_RES_Y");
    if (bpResYStr.empty())
    {
        D3D_THROW_EXCEPTION("Parameter BP_RES_Y not specified in the config file.")
    }
    int bpResY = atoi(bpResYStr.c_str());

    std::string bpResZStr = conf.get("BP_RES_Z");
    if (bpResZStr.empty())
    {
        D3D_THROW_EXCEPTION("Parameter BP_RES_Z not specified in the config file.")
    }
    int bpResZ = atoi(bpResZStr.c_str());

    std::string bpOriginStr = conf.get("BP_ORIGIN");
    Eigen::Vector3f bpOrigin;
    if (bpOriginStr.empty())
    {
        bpOrigin = Eigen::Vector3f::Zero();
    }
    else
    {
        std::stringstream bpOriginStream(bpOriginStr);
        bpOriginStream >> bpOrigin(0);
        bpOriginStream >> bpOrigin(1);
        bpOriginStream >> bpOrigin(2);

        if (bpOriginStream.fail())
        {
            D3D_THROW_EXCEPTION("Error reading BP_ORIGIN");
        }
    }

    std::string bpXAxisStr = conf.get("BP_X_AXIS");
    Eigen::Vector3f bpXAxis;
    bool xAxisSet = false;
    if (!bpXAxisStr.empty())
    {
        xAxisSet = true;

        std::stringstream bpXAxisStream(bpXAxisStr);
        bpXAxisStream >> bpXAxis(0);
        bpXAxisStream >> bpXAxis(1);
        bpXAxisStream >> bpXAxis(2);

        if (bpXAxisStream.fail())
        {
            D3D_THROW_EXCEPTION("Error reading BP_X_AXIS");
        }
    }

    std::string bpYAxisStr = conf.get("BP_Y_AXIS");
    Eigen::Vector3f bpYAxis;
    bool yAxisSet = false;
    if (!bpYAxisStr.empty())
    {
        yAxisSet = true;

        std::stringstream bpYAxisStream(bpYAxisStr);
        bpYAxisStream >> bpYAxis(0);
        bpYAxisStream >> bpYAxis(1);
        bpYAxisStream >> bpYAxis(2);

        if (bpYAxisStream.fail())
        {
            D3D_THROW_EXCEPTION("Error reading BP_Y_AXIS");
        }
    }

    std::string bpZAxisStr = conf.get("BP_Z_AXIS");
    Eigen::Vector3f bpZAxis;
    bool zAxisSet = false;
    if (!bpZAxisStr.empty())
    {
        zAxisSet = true;

        std::stringstream bpZAxisStream(bpZAxisStr);
        bpZAxisStream >> bpZAxis(0);
        bpZAxisStream >> bpZAxis(1);
        bpZAxisStream >> bpZAxis(2);

        if (bpZAxisStream.fail())
        {
            D3D_THROW_EXCEPTION("Error reading BP_Z_AXIS");
        }
    }

    if (xAxisSet == false && yAxisSet == false && zAxisSet == false)
    {
        bpXAxis(0) = 1; bpXAxis(1) = 0; bpXAxis(2) = 0;
        bpYAxis(0) = 0; bpYAxis(1) = 1; bpYAxis(2) = 0;
        bpZAxis(0) = 0; bpZAxis(1) = 0; bpZAxis(2) = 1;
    }
    else if (!(xAxisSet && yAxisSet && zAxisSet))
    {
        D3D_THROW_EXCEPTION("Either all or none out of BP_X_AXIS, BP_Y_AXIS and BP_Z_AXIS have to be specified.")
    }

    // rectify the axes
    // normlaize
    bpXAxis /= bpXAxis.norm();
    bpYAxis /= bpYAxis.norm();
    bpZAxis /= bpZAxis.norm();

    Eigen::Vector3f rightAngleZAxis = bpXAxis.cross(bpYAxis);
    if (rightAngleZAxis.dot(bpZAxis) < 0.95)
    {
        D3D_THROW_EXCEPTION("The given axes are not orthogonal to each other.")
    }

    bpZAxis = rightAngleZAxis;

    Eigen::Vector3f rightAngleYAxis = bpZAxis.cross(bpXAxis);
    if (rightAngleYAxis.dot(bpYAxis) < 0.95)
    {
        D3D_THROW_EXCEPTION("The given axes are not orthogonal to each other.");
    }

    bpYAxis = rightAngleYAxis;

    Eigen::Matrix4f boxToGlobal = Eigen::Matrix4f::Identity();
    boxToGlobal.topLeftCorner(3,3).col(0) = bpXAxis;
    boxToGlobal.topLeftCorner(3,3).col(1) = bpYAxis;
    boxToGlobal.topLeftCorner(3,3).col(2) = bpZAxis;
    boxToGlobal.topRightCorner(3,1).col(0) = bpOrigin;

    std::string hfNumIterStr = conf.get("HF_NUM_ITER");
    if (hfNumIterStr.empty())
    {
        D3D_THROW_EXCEPTION("Paramter HF_NUM_ITER is not specified in the config file.")
    }
    int ffNumIter = atoi(hfNumIterStr.c_str());

    std::string hfLambdaStr = conf.get("HF_LAMBDA");
    if (hfLambdaStr.empty())
    {
        D3D_THROW_EXCEPTION("Parameter HF_LAMBDA is not specified in the config file.")
    }
    float hfLambda = atof(hfLambdaStr.c_str());

    // load grid
    D3D::Grid<Hist> dataCost;
    dataCost.loadFromDataFile(intDepthGridFile);

    std::cout << "Data Cost Loaded: width = " << dataCost.getWidth() << ", height = " << dataCost.getHeight() << ", depth = " << dataCost.getDepth() << std::endl;

    if (bpResX != (int) dataCost.getWidth() || bpResY != (int) dataCost.getHeight() || bpResZ != (int) dataCost.getDepth())
    {
        D3D_THROW_EXCEPTION("Resolution specified in the config file does not match the dimension of the grid with integrated depth maps.")
    }

    Eigen::Vector3f minCorner;
    minCorner(0) = bpMinX;
    minCorner(1) = bpMinY;
    minCorner(2) = bpMinZ;
    Eigen::Vector3f size;
    size(0) = bpSizeX;
    size(1) = bpSizeY;
    size(2) = bpSizeZ;
    Eigen::Vector3f color;
    color(0) = 0.5;
    color(1) = 0.5;
    color(2) = 0.5;

    runTVHist(dataCost, minCorner, size, boxToGlobal, color, vrmlOutputFile, ffNumIter, hfLambda, 1);
}


/** tsne.cc
    Jeremy Barnes, 16 December 2014
    Copyright (c) 2014 mldb.ai inc.  All rights reserved.

    This file is part of MLDB. Copyright 2015 mldb.ai inc. All rights reserved.

    Implementation of an TSNE algorithm for embedding of a dataset.
*/

#include "tsne_interface.h"
#include "mldb/builtin/matrix.h"
#include "mldb/core/mldb_engine.h"
#include "mldb/core/dataset.h"
#include "mldb/utils/distribution.h"
#include <boost/multi_array.hpp>
#include "mldb/base/scope.h"
#include "mldb/base/parallel.h"
#include "mldb/utils/pair_utils.h"
#include "mldb/arch/timers.h"
#include "mldb/arch/simd_vector.h"
#include "mldb/utils/vector_utils.h"
#include "mldb/types/basic_value_descriptions.h"
#include "mldb/types/any_impl.h"
#include "mldb/builtin/sql_config_validator.h"
#include "mldb/utils/vantage_point_tree.h"
#include "mldb/plugins/jml/tsne/tsne.h"
#include "mldb/sql/sql_expression.h"
#include "mldb/core/analytics.h"
#include "mldb/vfs/fs_utils.h"
#include "mldb/vfs/filter_streams.h"
#include "mldb/types/jml_serialization.h"
#include "mldb/utils/log.h"

using namespace std;

namespace MLDB {

DEFINE_STRUCTURE_DESCRIPTION(TsneConfig);

TsneConfigDescription::
TsneConfigDescription()
{
    addField("trainingData", &TsneConfig::trainingData,
             "Specification of the data for input to the TSNE procedure.  This should be "
             "organized as an embedding, with each selected row containing the same "
             "set of columns with numeric values to be used as coordinates.  The select statement "
             "does not support groupby and having clauses.");
    addField("rowOutputDataset", &TsneConfig::output,
             "Dataset for TSNE output, with embeddings of training data. "
             "One row will be added for each row in the input dataset, "
             "with a list of coordinates.",
             PolyConfigT<Dataset>().withType("embedding"));
    addField("numInputDimensions", &TsneConfig::numInputDimensions,
             "Number of dimensions from the input to use.  This will limit "
             "the columns to the n first columns in the alphabetical "
             "sorting of the columns (-1 = all).",
             -1);
    addField("numOutputDimensions", &TsneConfig::numOutputDimensions,
             "Number of dimensions to produce in t-SNE space.  Normally "
             "this will be 2 or 3, depending upon the number of dimensions "
             "in the visualization.",
             2);
    addField("tolerance", &TsneConfig::tolerance,
             "Tolerance of perplexity calculation.  This is an internal "
             "parameter that only needs to be changed in rare circumstances.");
    addField("perplexity", &TsneConfig::perplexity,
             "Perplexity to aim for; higher means more spread out.  This "
             "controls how hard t-SNE tries to spread the points out.  If "
             "the resulting output looks more like a ball or a sphere than "
             "individual clusters, you should reduce this number.  If it "
             "looks like a dot or star, you should increase it.",
             30.0);
    addField("learningRate", &TsneConfig::learningRate,
             "The learning rate specifies the gradient descent step size during "
             "optimization of the cost function.  A learning rate that is too small "
             "may hold optimization in a local minimum.  A learning rate that is too high "
             "may jump over the best optimal point. In general, the learning rate "
             "should be between 100 and 1000.",
             500.0);
    addField("modelFileUrl", &TsneConfig::modelFileUrl,
             "URL where the model file (with extension '.tsn') should be saved. "
             "This file can be loaded by the ![](%%doclink tsne.embedRow function). "
             "This parameter is optional unless the `functionName` parameter is used.");
    addField("functionName", &TsneConfig::functionName,
             "If specified, an instance of the ![](%%doclink tsne.embedRow function) of this name will be created using "
             "the trained model. Note that to use this parameter, the `modelFileUrl` must "
             "also be provided.");
    addAuto("minIterations", &TsneConfig::minIterations,
            "Minimum number of t-SNE iterations to run.  Making this too "
            "small may lead to spurious early stopping.");
    addAuto("maxIterations", &TsneConfig::maxIterations,
            "Maximum number of t-SNE iterations to run (it will stop if "
            "convergance is obtained before this number of iterations). ");
    addParent<ProcedureConfig>();

    onPostValidate = chain(validateQuery(&TsneConfig::trainingData,
                                         MustContainFrom(),
                                         NoGroupByHaving()),
                           validateFunction<TsneConfig>());
}

/*****************************************************************************/
/* TSNE PROCEDURE                                                             */
/*****************************************************************************/

struct TsneItl {
    TsneItl()
    {
    }

    TsneItl(const Url & filename)
    {
        filter_istream stream(filename);
        MLDB::DB::Store_Reader store(stream);

        reconstitute(store);

        stream.close();
    }

    ML::TSNE_Params params;
    boost::multi_array<float, 2> inputPath;
    boost::multi_array<float, 2> outputPath;
    std::unique_ptr<MLDB::VantagePointTree> vpTree;
    std::unique_ptr<MLDB::Quadtree> qtree;
    std::vector<Utf8String> inputColumnNames;
    std::vector<Utf8String> outputColumnNames;
    std::shared_ptr<const std::vector<ColumnPath> > outputColumnNamesShared;

    size_t numOutputDimensions() const { return outputPath.shape()[1]; }

    int64_t memusage() const
    {
        int64_t result = sizeof(*this);
        result += sizeof(float) * inputPath.shape()[0] * inputPath.shape()[1];
        result += sizeof(float) * outputPath.shape()[0] * outputPath.shape()[1];
        result += vpTree->memusage();
        result += qtree->root->memusage();
        return result;
    }

    void save(const std::string & filename) const
    {
        filter_ostream stream(filename);
        MLDB::DB::Store_Writer store(stream);
        serialize(store);
    }

    void serialize(MLDB::DB::Store_Writer & store) const
    {
        using namespace MLDB::DB;

        store << string("TSNE") << compact_size_t(2);

        size_t rows = inputPath.shape()[0];
        size_t dimsIn = inputPath.shape()[1];
        size_t dimsOut = outputPath.shape()[1];

        store << compact_size_t(rows)
              << compact_size_t(dimsIn)
              << compact_size_t(dimsOut);

        store << inputPath << outputPath;

        store << inputColumnNames << outputColumnNames;

        MLDB::VantagePointTree::serializePtr(store, vpTree.get());
        qtree->serialize(store);
    }

    void reconstitute(MLDB::DB::Store_Reader & store)
    {
        using namespace MLDB::DB;

        string tag;
        store >> tag;
        if (tag != "TSNE")
            throw MLDB::Exception("Expected TSNE tag");
        compact_size_t version(store);
        if (version != 2)
            throw MLDB::Exception("Unknown version for t-SNE");

        compact_size_t rows(store), dimsIn(store), dimsOut(store);

        store >> inputPath >> outputPath;

        store >> inputColumnNames >> outputColumnNames;

        outputColumnNamesShared.reset(new vector<ColumnPath>(outputColumnNames.begin(), outputColumnNames.end()));

        vpTree.reset(MLDB::VantagePointTree::reconstitutePtr(store));
        qtree.reset(new MLDB::Quadtree(store));
    }

    distribution<float> reembed(const distribution<float> & v) const
    {
        return retsneApproxFromCoords(v, inputPath, outputPath,
                                      *qtree, *vpTree, params);
    }

};

TsneProcedure::
TsneProcedure(MldbEngine * owner,
            PolyConfig config,
            const std::function<bool (const Json::Value &)> & onProgress)
    : Procedure(owner)
{
    tsneConfig = config.params.convert<TsneConfig>();
}

Any
TsneProcedure::
getStatus() const
{
    return Any();
}

RunOutput
TsneProcedure::
run(const ProcedureRunConfig & run,
      const std::function<bool (const Json::Value &)> & onProgress) const
{
    auto runProcConf = applyRunConfOverProcConf(tsneConfig, run);

    if (!runProcConf.modelFileUrl.empty()) {
        checkWritability(runProcConf.modelFileUrl.toString(), "modelFileUrl");
    }

    auto onProgress2 = [&] (const Json::Value & progress)
        {
            Json::Value value;
            value["dataset"] = progress;
            return onProgress(value);
        };

    // 1.  Get a numsubjects by numdimensions input matrix to train the
    //     t-SNE algorithm on
    auto itl = std::make_shared<TsneItl>();

    itl->params.perplexity = runProcConf.perplexity;
    itl->params.tolerance = runProcConf.tolerance;
    itl->params.eta = runProcConf.learningRate;
    itl->params.min_iter = runProcConf.minIterations;
    itl->params.max_iter = runProcConf.maxIterations;

    DEBUG_MSG(logger) << "perplexity = " << itl->params.perplexity;
    DEBUG_MSG(logger) << "tolerance = " << itl->params.tolerance;
    DEBUG_MSG(logger) << "learningRate = " << itl->params.eta;
    DEBUG_MSG(logger) << "minIterations = " << itl->params.min_iter;
    DEBUG_MSG(logger) << "maxIterations = " << itl->params.max_iter;
    DEBUG_MSG(logger) << "doing t-SNE";

    SqlExpressionMldbScope context(engine);

    ConvertProgressToJson convertProgressToJson(onProgress);
    auto embeddingOutput = getEmbedding(*runProcConf.trainingData.stm,
                                        context,
                                        runProcConf.numInputDimensions,
                                        convertProgressToJson);

    std::vector<std::tuple<RowHash, RowPath, std::vector<double>,
                           std::vector<ExpressionValue> > > & rows
        = embeddingOutput.first;
    std::vector<KnownColumn> & vars = embeddingOutput.second;

    size_t numDims = vars.size();

    DEBUG_MSG(logger) << "numDims = " << numDims;

    boost::multi_array<float, 2> coords
        (boost::extents[rows.size()][numDims]);

    for (unsigned i = 0;  i < rows.size();  ++i) {
        for (auto & e: std::get<2>(rows[i]))
            ExcAssert(isfinite(e));
        std::copy(std::get<2>(rows[i]).begin(), std::get<2>(rows[i]).end(),
                  &coords[i][0]);
    }

    if (coords.size() == 0)
        throw AnnotatedException(400, "t-sne training requires at least 1 datapoint. "
                                  "Make sure your dataset is not empty and that your WHERE, offset "
                                  "and limit expressions do not filter all the rows");

    DEBUG_MSG(logger) << "copied into matrix";

//     DEBUG_MSG(logger) << "rows[0] = " << rows[0].second;
//     DEBUG_MSG(logger) << "rows[1] = " << rows[1].second;
//     DEBUG_MSG(logger) << "rows[2] = " << rows[1].second;
//     DEBUG_MSG(logger) << "rows[0] dot rows[1] = " << rows[0].second.dotprod(rows[1].second);
//     DEBUG_MSG(logger) << "rows[0] dot rows[2] = " << rows[0].second.dotprod(rows[2].second);
//     DEBUG_MSG(logger) << "rows[0] dist rows[1] = " << (rows[0].second - rows[1].second).two_norm();
//     DEBUG_MSG(logger) << "rows[0] dist rows[2] = " << (rows[0].second - rows[2].second).two_norm();
//     DEBUG_MSG(logger) << "rows[1] dist rows[2] = " << (rows[1].second - rows[2].second).two_norm();

    itl->inputPath.resize(boost::extents[rows.size()][numDims]);
    itl->inputPath = coords;

    ML::TSNE_Callback callback = [&] (int iter, float cost,
                                      std::string phase)
        {
            if (iter == 1 || iter % 10 == 0)
                INFO_MSG(logger) << "phase " << phase << " iter " << iter
                     << " cost " << cost;
            return true;
        };

    ExcAssertGreaterEqual(runProcConf.numOutputDimensions, 1);

    itl->outputPath.resize(boost::extents[rows.size()][runProcConf.numOutputDimensions]);
    itl->outputPath
        = ML::tsneApproxFromCoords(coords, runProcConf.numOutputDimensions,
                                   itl->params, callback, &itl->vpTree,
                                   &itl->qtree);

    ExcAssert(itl->qtree);
    ExcAssert(itl->vpTree);

    vector<ColumnPath> names = { ColumnPath("x"), ColumnPath("y"), ColumnPath("z") };
    if (runProcConf.numOutputDimensions <= 3)
        names.resize(runProcConf.numOutputDimensions);
    else {
        names.clear();
        for (unsigned i = 0; i < runProcConf.numOutputDimensions;  ++i)
            names.push_back(ColumnPath(MLDB::format("dim%04d", i)));
    }

    // Record the column names for later
    for (auto & c: vars) {
        itl->inputColumnNames.emplace_back(c.columnName.toUtf8String());
    }
    for (auto & c: names) {
        itl->outputColumnNames.emplace_back(c.toUtf8String());
    }

    itl->outputColumnNamesShared
        .reset(new vector<ColumnPath>(itl->outputColumnNames.begin(),
                                      itl->outputColumnNames.end()));

    if (!runProcConf.modelFileUrl.empty()) {
        makeUriDirectory(runProcConf.modelFileUrl.toString());
        itl->save(runProcConf.modelFileUrl.toString());
    }

    // Create a dataset to contain the output embedding if we ask for it
    if (!runProcConf.output.type.empty()) {
        auto output = createDataset(engine, runProcConf.output, onProgress2, true /*overwrite*/);

        for (unsigned i = 0;  i < rows.size();  ++i) {
            TRACE_MSG(logger) << "row " << i << " had coords " << itl->outputPath[i][0] << ","
                 << itl->outputPath[i][1];
            std::vector<std::tuple<ColumnPath, CellValue, Date> > cols;
            for (unsigned j = 0;  j < runProcConf.numOutputDimensions;  ++j) {
                ExcAssert(isfinite(itl->outputPath[i][j]));
                cols.emplace_back(names[j], itl->outputPath[i][j], Date());
            }

            output->recordRow(std::get<1>(rows[i]), cols);
        }

        output->commit();
    }

    if(!runProcConf.functionName.empty()) {
        PolyConfig tsneFuncPC;
        tsneFuncPC.type = "tsne.embedRow";
        tsneFuncPC.id = runProcConf.functionName;
        tsneFuncPC.params = TsneEmbedConfig(runProcConf.modelFileUrl);

        createFunction(engine, tsneFuncPC, onProgress, true);
    }

    return Any();
}

DEFINE_STRUCTURE_DESCRIPTION(TsneEmbedConfig);

TsneEmbedConfigDescription::
TsneEmbedConfigDescription()
{
    addField("modelFileUrl", &TsneEmbedConfig::modelFileUrl,
             "URL of the model file (with extension '.tns') to load. "
             "This file is created by the ![](%%doclink tsne.train procedure).");
}

/*****************************************************************************/
/* TSNE EMBED ROW                                                            */
/*****************************************************************************/

DEFINE_STRUCTURE_DESCRIPTION(TsneInput);

TsneInputDescription::TsneInputDescription()
{
    addField("embedding", &TsneInput::embedding,
             "Undocumented.");
}

DEFINE_STRUCTURE_DESCRIPTION(TsneOutput);

TsneOutputDescription::TsneOutputDescription()
{
    addField("cluster", &TsneOutput::tsne,
             "Undocumented.");
}

TsneEmbed::
TsneEmbed(MldbEngine * owner,
          PolyConfig config,
          const std::function<bool (const Json::Value &)> & onProgress)
    : BaseT(owner, config)
{
    functionConfig = config.params.convert<TsneEmbedConfig>();
    itl = std::make_shared<TsneItl>(functionConfig.modelFileUrl);
}

TsneOutput
TsneEmbed::
call(TsneInput input) const
{
    throw AnnotatedException(500, "t-SNE Embed apply function is not yet implemented");
#if 0
    ExpressionValue result;

    ExpressionValue storage;
    const ExpressionValue & inputVal = context.get("embedding", storage);

    distribution<float> input = inputVal.getEmbedding(itl->inputColumnNames.size());

    Date ts = Date::negativeInfinity();
    auto embedding = itl->reembed(input);

    result.set("tsne", ExpressionValue(embedding, ts));

    return result;
#endif
}

namespace {

RegisterProcedureType<TsneProcedure, TsneConfig>
regTsne(builtinPackage(),
        "Project a high dimensional space into a low-dimensional space suitable for visualization",
        "procedures/TsneProcedure.md.html");

RegisterFunctionType<TsneEmbed, TsneEmbedConfig>
regTsneEmbed(builtinPackage(),
             "tsne.embedRow",
             "Embed a pre-trained t-SNE algorithm to new data points",
             "functions/TsneEmbed.md.html",
             nullptr,
             {MldbEntity::INTERNAL_ENTITY});

} // file scope

} // namespace MLDB


#ifdef USE_EIGEN_BACKEND

/** Eigen3 backend.
 *
 * Only supports float32 computation with NHWC memory layout (at runtime and as input).
 */

// CR lpuchallafiore: Add multi-threading support (see "Evaluating with a Thread Pool" in the Eigen Tensor docs).

#include "../neuralnet/nninterface.h"

#include <Eigen/Dense>
#include <unsupported/Eigen/CXX11/Tensor>
#include <zstr/src/zstr.hpp>

#include "../neuralnet/desc.h"
#include "../neuralnet/modelversion.h"
#include "../neuralnet/nninputs.h"
#include "../neuralnet/nneval.h"

using namespace std;
using Eigen::Tensor;
using Eigen::TensorMap;

//Eigen doesn't seem to have a way to make a const tensor map out of a const float* ??
//So we have to cast away qualifiers to build it.
#pragma GCC diagnostic ignored "-Wcast-qual"

// Eigen tensors are stored in column-major order, so an NHWC memory layout is given by Tensor<4>(C,W,H,N).

#define SCALAR float
#define TENSOR2 Tensor<SCALAR, 2>
#define TENSOR3 Tensor<SCALAR, 3>
#define TENSOR4 Tensor<SCALAR, 4>
#define TENSORMAP2 TensorMap<Tensor<SCALAR, 2>>
#define TENSORMAP3 TensorMap<Tensor<SCALAR, 3>>
#define TENSORMAP4 TensorMap<Tensor<SCALAR, 4>>

#define CONSTTENSOR2 const Tensor<SCALAR, 2>
#define CONSTTENSOR3 const Tensor<SCALAR, 3>
#define CONSTTENSOR4 const Tensor<SCALAR, 4>
#define CONSTTENSORMAP2 const TensorMap<Tensor<SCALAR, 2>>
#define CONSTTENSORMAP3 const TensorMap<Tensor<SCALAR, 3>>
#define CONSTTENSORMAP4 const TensorMap<Tensor<SCALAR, 4>>

// Debugging -----------------------------------------------------------------------------------------------------------
// #define DEBUG true

template <typename T>
void printTensorShape(const string& name, const T* t) {
  auto d = t->dimensions();
  cout << name << " rank=" << d.size() << " - (";
  for (int i = 0; i < d.size(); i++) {
    cout << d[i] << ",";
  }
  cout << ")" << endl;
}

#if DEBUG
#define DSHAPE(n, x) printTensorShape(n,x)
#define DTENSOR(n, x) cout << n << *x << endl
#else
#define DSHAPE(n, x)
#define DTENSOR(n, x)
#endif

// LoadedModel / ModelDesc ---------------------------------------------------------------------------------------------

struct LoadedModel {
  ModelDesc modelDesc;

  LoadedModel(const string& fileName) {
    ModelDesc::loadFromFileMaybeGZipped(fileName,modelDesc);
  }

  LoadedModel() = delete;
  LoadedModel(const LoadedModel&) = delete;
  LoadedModel& operator=(const LoadedModel&) = delete;
};

LoadedModel* NeuralNet::loadModelFile(const string& file) {
  LoadedModel* loadedModel = new LoadedModel(file);
  return loadedModel;
}

void NeuralNet::freeLoadedModel(LoadedModel* loadedModel) {
  delete loadedModel;
}

string NeuralNet::getModelName(const LoadedModel* loadedModel) {
  return loadedModel->modelDesc.name;
}

int NeuralNet::getModelVersion(const LoadedModel* loadedModel) {
  return loadedModel->modelDesc.version;
}

Rules NeuralNet::getSupportedRules(const LoadedModel* loadedModel, const Rules& desiredRules, bool& supported) {
  return loadedModel->modelDesc.getSupportedRules(desiredRules, supported);
}

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

struct ComputeContext {
  int nnXLen;
  int nnYLen;
};

ComputeContext* NeuralNet::createComputeContext(
  const std::vector<int>& gpuIdxs,
  Logger* logger,
  int nnXLen,
  int nnYLen,
  const string& openCLTunerFile,
  const string& homeDataDirOverride,
  bool openCLReTunePerBoardSize,
  enabled_t useFP16Mode,
  enabled_t useNHWCMode,
  const LoadedModel* loadedModel
) {
  (void)gpuIdxs;
  (void)logger;
  (void)openCLTunerFile;
  (void)homeDataDirOverride;
  (void)openCLReTunePerBoardSize;
  (void)loadedModel;

  bool useFP16 = useFP16Mode == enabled_t::True ? true : false;
  bool useNHWC = useNHWCMode == enabled_t::False ? false : true;

  if(useFP16)
    throw StringError("Eigen backend: useFP16 = true not supported");
  if(!useNHWC)
    throw StringError("Eigen backend: useNHWC = false not supported");

  ComputeContext* context = new ComputeContext();
  context->nnXLen = nnXLen;
  context->nnYLen = nnYLen;
  return context;
}

void NeuralNet::freeComputeContext(ComputeContext* computeContext) {
  delete computeContext;
}

// Helpers --------------------------------------------------------------------------------------------------------------

static void computeMaskSum(CONSTTENSORMAP3* mask, float* maskSum) {
  for (int n = 0; n < mask->dimension(2); n++) {
    float s = 0.f;
    for (int h = 0; h < mask->dimension(1); h++) {
      for (int w = 0; w < mask->dimension(0); w++) {
        s += (*mask)(w, h, n);
      }
    }
    maskSum[n] = s;
  }
}

// in NxHxWxC, bias NxC
static void addNCBiasInplace(TENSORMAP4* in, CONSTTENSORMAP2* bias) {
  assert(in->dimension(0) == bias->dimension(0) && in->dimension(3) == bias->dimension(1));
  for (int n = 0; n < in->dimension(3); n++) {
    for (int h = 0; h < in->dimension(2); h++) {
      for (int w = 0; w < in->dimension(1); w++) {
        for (int c = 0; c < in->dimension(0); c++) {
          (*in)(c,w,h,n) += (*bias)(c,n);
        }
      }
    }
  }
}

static void poolRowsGPool(CONSTTENSORMAP4* in, TENSORMAP2* out, const float* maskSum) {
  for (int n = 0; n < in->dimension(3); n++) {
    for (int c = 0; c < in->dimension(0); c++) {
      float s = 0.f;
      float m = 0.f;
      for (int h = 0; h < in->dimension(2); h++) {
        for (int w = 0; w < in->dimension(1); w++) {
          float x = (*in)(c, w, h, n);
          s += x;
          m = max(m, x);
        }
      }
      float div = maskSum[n];
      float sqrtdiv = sqrt(div);
      float mean = s / div;
      (*out)(c, n) = mean;
      (*out)(c + in->dimension(0), n) = mean * (sqrtdiv - 14.f) * 0.1f;
      (*out)(c + 2*in->dimension(0), n) = m;
    }
  }
}

// // Given input [n,w,h,c] fills output of shape [n,c] with sum over c.
// static void poolRowsSum(CONSTTENSORMAP4* in, TENSORMAP2* out, float scaleSum) {
//   for (int n = 0; n < in->dimension(3); n++) {
//     for (int c = 0; c < in->dimension(0); c++) {
//       float s = 0.f;
//       for (int h = 0; h < in->dimension(2); h++) {
//         for (int w = 0; w < in->dimension(1); w++) {
//           float x = in(c, w, h, n);
//           s += x;
//         }
//       }
//       out(c, n) = s * scaleSum;
//     }
//   }
// }

static void poolRowsValueHead(CONSTTENSORMAP4* in, TENSORMAP2* out, const float* maskSum) {
  for (int n = 0; n < in->dimension(3); n++) {
    for (int c = 0; c < in->dimension(0); c++) {
      float s = 0.f;
      for (int h = 0; h < in->dimension(2); h++) {
        for (int w = 0; w < in->dimension(1); w++) {
          float x = (*in)(c, w, h, n);
          s += x;
        }
      }
      float div = maskSum[n];
      float sqrtdiv = sqrt(div);
      float mean = s / div;
      (*out)(c, n) = mean;
      (*out)(c + in->dimension(0), n) = mean * (sqrtdiv - 14.f) * 0.1f;
      (*out)(c + 2*in->dimension(0), n) = mean * ((sqrtdiv - 14.0f) * (sqrtdiv - 14.0f) * 0.01f - 0.1f);
    }
  }
}

// Layers --------------------------------------------------------------------------------------------------------------

// Convolution layer with zero-padding.
struct ConvLayer {
  string name;

  Eigen::array<pair<int, int>, 4> paddings;
  TENSOR4 kernel;
  int inChannels, outChannels;

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

  ConvLayer(const ConvLayerDesc& desc) {
    name = desc.name;
    int convYSize = desc.convYSize;
    int convXSize = desc.convXSize;
    inChannels = desc.inChannels;
    outChannels = desc.outChannels;
    //Currently eigen impl doesn't support dilated convs
    int dilationY = desc.dilationY;
    int dilationX = desc.dilationX;
    int paddingX = (convXSize / 2) * dilationX;
    int paddingY = (convYSize / 2) * dilationY;

    if(dilationX != 1 || dilationY != 1)
      throw StringError("Eigen backend: Encountered convolution dilation factors other than 1, not supported");

    assert(convXSize % 2 == 1);
    assert(convYSize % 2 == 1);

    paddings[0] = make_pair(0, 0);                // C
    paddings[1] = make_pair(paddingX, paddingX);  // W
    paddings[2] = make_pair(paddingY, paddingY);  // H
    paddings[3] = make_pair(0, 0);                // N

    // CR-someday lpuchallafiore: optimize NHWC vs NCHW, etc.
    kernel = TensorMap<const Tensor<const SCALAR, 4>>(
      desc.weights.data(), convXSize, convYSize, inChannels, outChannels);
  }

  void apply(CONSTTENSORMAP4* input, TENSORMAP4* output, bool accumulate) const {
    auto padded = input->pad(paddings);
    assert(output->dimension(0) == outChannels);
    for(int n = 0; n < input->dimension(3); n++) {
      auto inN = padded.chip(n, 3);
      for(int oc = 0; oc < outChannels; oc++) {
        TENSOR2 sum(input->dimension(1), input->dimension(2));
        sum.setZero();

        for(int ic = 0; ic < inChannels; ic++) {
          Eigen::array<ptrdiff_t, 2> dims({0, 1});
          auto kChip = kernel.chip(oc, 3).chip(ic, 2);
          auto inNC = inN.chip(ic, 0);
          sum += inNC.convolve(kChip, dims);
        }

        if(accumulate)
          output->chip(n, 3).chip(oc, 0) += sum;
        else
          output->chip(n, 3).chip(oc, 0) = sum;
      }
    }
  }
};

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

struct BatchNormLayer {
  string name;

  vector<float> mergedScale;
  vector<float> mergedBias;

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

  BatchNormLayer(const BatchNormLayerDesc& desc) {
    name = desc.name;
    int numChannels = desc.numChannels;
    float epsilon = desc.epsilon;

    mergedScale.resize(numChannels);
    mergedBias.resize(numChannels);
    for(int c = 0; c < numChannels; c++) {
      mergedScale[c] = desc.scale[c] / sqrt(desc.variance[c] + epsilon);
      mergedBias[c] = desc.bias[c] - mergedScale[c] * desc.mean[c];
    }
  }

  // Mask should be in 'NHW' format (no "C" channel).
  void apply(
    bool applyRelu,
    CONSTTENSORMAP4* input,
    TENSORMAP4* output,
    CONSTTENSORMAP3* mask
  ) const {
    for(int c = 0; c < input->dimension(0); c++) {
      auto inC = input->chip(c, 0);
      auto x = inC * mergedScale[c] + mergedBias[c];
      auto z = TENSOR3(mask->dimension(0), mask->dimension(1), mask->dimension(2)).setZero();
      if(applyRelu)
        output->chip(c, 0) = (*mask == 1.f).select(x.cwiseMax(0.f), z);
      else
        output->chip(c, 0) = (*mask == 1.f).select(x, z);
    }
  }
};

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

struct ActivationLayer {
  string name;

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

  ActivationLayer(const ActivationLayerDesc& desc) { name = desc.name; }

  template <int N>
  void apply(const Tensor<SCALAR, N>* input, Tensor<SCALAR, N>* output) const { *output = input->cwiseMax(0.f); }
  template <int N>
  void apply(const TensorMap<Tensor<SCALAR, N>>* input, TensorMap<Tensor<SCALAR, N>>* output) const { *output = input->cwiseMax(0.f); }
};

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

struct MatMulLayer {
  string name;
  TENSOR2 weights;

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

  MatMulLayer(const MatMulLayerDesc& desc)
    : name(desc.name)
  {
    weights = TENSOR2(desc.outChannels, desc.inChannels);
    memcpy(weights.data(), desc.weights.data(), sizeof(SCALAR) * weights.size());
  }

  void apply(CONSTTENSORMAP2* in, TENSORMAP2* out) const {
    Eigen::array<Eigen::IndexPair<int>, 1> product_dims = { Eigen::IndexPair<int>(1, 0) };
    *out = weights.contract(*in, product_dims);
  }
};

struct MatBiasLayer {
  string name;
  std::vector<float> weights;

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

  MatBiasLayer(const MatBiasLayerDesc& desc)
    : name(desc.name),
      weights(desc.weights) {}

  void apply(TENSORMAP2* mat) const {
    for(int n = 0; n < mat->dimension(1); n++) {
      for(int c = 0; c < mat->dimension(0); c++) {
        (*mat)(c, n) += weights[c];
      }
    }
  }
};

// Blocks
// --------------------------------------------------------------------------------------------------------------

struct ResidualBlockIntf {
  virtual ~ResidualBlockIntf(){}

  virtual void apply(
    TENSORMAP4* trunk,
    TENSORMAP4* trunkScratch,
    TENSORMAP4* regularOut,
    TENSORMAP4* regularScratch,
    TENSORMAP4* midIn,
    TENSORMAP4* midScratch,
    TENSORMAP4* gpoolOut,
    TENSORMAP4* gpoolOut2,
    TENSORMAP2* gpoolConcat,
    TENSORMAP2* gpoolBias,
    CONSTTENSORMAP3* mask,
    const float* maskSum
  ) const = 0;
};

struct ResidualBlock final : public ResidualBlockIntf {
  string name;
  BatchNormLayer preBN;
  ConvLayer regularConv;
  BatchNormLayer midBN;
  ConvLayer finalConv;

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

  ~ResidualBlock(){}

  ResidualBlock(const ResidualBlockDesc& desc)
    : name(desc.name),
      preBN(desc.preBN),
      regularConv(desc.regularConv),
      midBN(desc.midBN),
      finalConv(desc.finalConv) {}

  void apply(
    TENSORMAP4* trunk,
    TENSORMAP4* trunkScratch,
    TENSORMAP4* regularOut,
    TENSORMAP4* regularScratch,
    TENSORMAP4* midIn,
    TENSORMAP4* midScratch,
    TENSORMAP4* gpoolOut,
    TENSORMAP4* gpoolOut2,
    TENSORMAP2* gpoolConcat,
    TENSORMAP2* gpoolBias,
    CONSTTENSORMAP3* mask,
    const float* maskSum
  ) const override {
    (void)regularOut;
    (void)regularScratch;
    (void)gpoolOut;
    (void)gpoolOut2;
    (void)gpoolConcat;
    (void)gpoolBias;
    (void)maskSum;
    const bool applyBNRelu = true;
    preBN.apply(applyBNRelu, trunk, trunkScratch, mask);
    regularConv.apply(trunkScratch, midIn, false);
    midBN.apply(applyBNRelu, midIn, midScratch, mask);
    finalConv.apply(midScratch, trunk, true);
  }
};

// // Given two tensors with shapes inA: [n, h, w, cA] and inB: [n, h, w, cB]
// // Copy them into a single tensor out: [n, h, w, cA + cB]
// TENSOR4 concatTensors(CONSTTENSOR4& a, CONSTTENSOR4& b) {
//   assert(a->dimension(1) == b->dimension(1) && a->dimension(2) == b->dimension(2) && a->dimension(3) == b->dimension(3));
//   TENSOR4 x = TENSOR4(/* C */ a->dimension(0) + b->dimension(0),
//                                           /* W */ a->dimension(1),
//                                           /* H */ a->dimension(2),
//                                           /* N */ a->dimension(3));
//   for (int n = 0; n < a->dimension(3); n++) {
//     for (int h = 0; h < a->dimension(2); h++) {
//       for (int w = 0; w < a->dimension(1); w++) {
//         int c = 0;
//         for (int ca = 0; a->dimension(0); ca++, c++) {
//           x(c,w,h,n) = a(ca,w,h,n);
//         }
//         for (int cb = 0; b->dimension(0); cb++, c++) {
//           x(c,w,h,n) = b(cb,w,h,n);
//         }
//       }
//     }
//   }
//   return x;
// }

struct GlobalPoolingResidualBlock final : public ResidualBlockIntf {
  string name;
  BatchNormLayer preBN;
  ActivationLayer preActivation;
  ConvLayer regularConv;
  ConvLayer gpoolConv;
  BatchNormLayer gpoolBN;
  ActivationLayer gpoolActivation;
  MatMulLayer gpoolToBiasMul;
  BatchNormLayer midBN;
  ActivationLayer midActivation;
  ConvLayer finalConv;

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

  ~GlobalPoolingResidualBlock(){}

  GlobalPoolingResidualBlock(const GlobalPoolingResidualBlockDesc& desc)
    : name(desc.name),
      preBN(desc.preBN),
      preActivation(desc.preActivation),
      regularConv(desc.regularConv),
      gpoolConv(desc.gpoolConv),
      gpoolBN(desc.gpoolBN),
      gpoolActivation(desc.gpoolActivation),
      gpoolToBiasMul(desc.gpoolToBiasMul),
      midBN(desc.midBN),
      midActivation(desc.midActivation),
      finalConv(desc.finalConv) {}

  void apply(
    TENSORMAP4* trunk,
    TENSORMAP4* trunkScratch,
    TENSORMAP4* regularOut,
    TENSORMAP4* regularScratch,
    TENSORMAP4* midIn,
    TENSORMAP4* midScratch,
    TENSORMAP4* gpoolOut,
    TENSORMAP4* gpoolOut2,
    TENSORMAP2* gpoolConcat,
    TENSORMAP2* gpoolBias,
    CONSTTENSORMAP3* mask,
    const float* maskSum
  ) const override {
    (void)midIn;
    (void)midScratch;
    const bool applyBNRelu = true;
    DTENSOR("trunk", trunk);
    DTENSOR("mask", mask);
    preBN.apply(applyBNRelu, trunk, trunkScratch, mask);
    DTENSOR("trunkScratch", trunkScratch);
    regularConv.apply(trunkScratch, regularOut, false);
    DTENSOR("regularOut", regularOut);
    gpoolConv.apply(trunkScratch, gpoolOut, false);
    DTENSOR("gpoolOut", gpoolOut);
    gpoolBN.apply(applyBNRelu, gpoolOut, gpoolOut2, mask);
    DTENSOR("gpoolOut2", gpoolOut2);
    poolRowsGPool(gpoolOut2, gpoolConcat, maskSum);
    gpoolToBiasMul.apply(gpoolConcat, gpoolBias);
    addNCBiasInplace(regularOut, gpoolBias);
    midBN.apply(applyBNRelu, regularOut, regularScratch, mask);
    finalConv.apply(regularScratch, trunk, true);
    DSHAPE("trunk", trunk);
    DSHAPE("trunkScratch", trunkScratch);
    DSHAPE("regularOut", regularOut);
    DSHAPE("gpoolOut", gpoolOut);
    DSHAPE("gpoolOut2", gpoolOut2);
    DSHAPE("gpoolConcat", gpoolConcat);
    DSHAPE("gpoolBias", gpoolBias);
    DSHAPE("mask", mask);
  }
};

struct Trunk {
  string name;
  int version;
  int numBlocks;

  ConvLayer initialConv;
  MatMulLayer initialMatMul;
  vector<pair<int, ResidualBlockIntf*>> blocks;
  BatchNormLayer trunkTipBN;
  ActivationLayer trunkTipActivation;

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

  Trunk(const TrunkDesc& desc)
    : name(desc.name),
      version(desc.version),
      numBlocks(desc.numBlocks),
      initialConv(desc.initialConv),
      initialMatMul(desc.initialMatMul),
      trunkTipBN(desc.trunkTipBN),
      trunkTipActivation(desc.trunkTipActivation)
  {
    for (int i = 0; i < numBlocks; ++i) {
      if (desc.blocks[i].first == ORDINARY_BLOCK_KIND) {
        ResidualBlockDesc* blockDesc = (ResidualBlockDesc*)desc.blocks[i].second;
        ResidualBlockIntf* block = new ResidualBlock(*blockDesc);
        blocks.push_back(make_pair(ORDINARY_BLOCK_KIND, block));
      }
      else if (desc.blocks[i].first == DILATED_BLOCK_KIND) {
        throw StringError("Eigen backend: Dilated residual blocks are not supported right now");
      }
      else if (desc.blocks[i].first == GLOBAL_POOLING_BLOCK_KIND) {
        GlobalPoolingResidualBlockDesc* blockDesc = (GlobalPoolingResidualBlockDesc*)desc.blocks[i].second;
        GlobalPoolingResidualBlock* block = new GlobalPoolingResidualBlock(*blockDesc);
        blocks.push_back(make_pair(GLOBAL_POOLING_BLOCK_KIND, block));
      }
      else {
        ASSERT_UNREACHABLE;
      }
    }
  }

  virtual ~Trunk() {
    for (auto p : blocks) {
      delete p.second;
    }
  }

  void apply(
    CONSTTENSORMAP4* input,
    CONSTTENSORMAP2* inputGlobal,
    TENSORMAP2* inputMatMulOut,
    TENSORMAP4* trunk,
    TENSORMAP4* trunkScratch,
    TENSORMAP4* regularOut,
    TENSORMAP4* regularScratch,
    TENSORMAP4* midIn,
    TENSORMAP4* midScratch,
    TENSORMAP4* gpoolOut,
    TENSORMAP4* gpoolOut2,
    TENSORMAP2* gpoolConcat,
    TENSORMAP2* gpoolBias,
    CONSTTENSORMAP3* mask,
    const float* maskSum
  ) const {

    initialConv.apply(input, trunkScratch, false);
    initialMatMul.apply(inputGlobal, inputMatMulOut);
    addNCBiasInplace(trunkScratch, inputMatMulOut);

    // apply blocks
    // Flip trunkBuf and trunkScratchBuf so that the result gets accumulated in trunkScratchBuf
    for (auto block : blocks) {
      block.second->apply(
        trunkScratch,
        trunk,
        regularOut,
        regularScratch,
        midIn,
        midScratch,
        gpoolOut,
        gpoolOut2,
        gpoolConcat,
        gpoolBias,
        mask,
        maskSum
      );
    }

    // And now with the final BN port it from trunkScratchBuf to trunkBuf.
    const bool applyBNRelu = true;
    trunkTipBN.apply(applyBNRelu, trunkScratch, trunk, mask);
  }
};

struct PolicyHead {
  string name;
  int version;

  ConvLayer p1Conv;
  ConvLayer g1Conv;
  BatchNormLayer g1BN;
  ActivationLayer g1Activation;
  MatMulLayer gpoolToBiasMul;
  BatchNormLayer p1BN;
  ActivationLayer p1Activation;
  ConvLayer p2Conv;
  MatMulLayer gpoolToPassMul;

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

  PolicyHead(const PolicyHeadDesc& desc)
    : name(desc.name),
      version(desc.version),
      p1Conv(desc.p1Conv),
      g1Conv(desc.g1Conv),
      g1BN(desc.g1BN),
      g1Activation(desc.g1Activation),
      gpoolToBiasMul(desc.gpoolToBiasMul),
      p1BN(desc.p1BN),
      p1Activation(desc.p1Activation),
      p2Conv(desc.p2Conv),
      gpoolToPassMul(desc.gpoolToPassMul) {}

  void apply(
    CONSTTENSORMAP4* trunk,
    TENSORMAP4* p1Out,
    TENSORMAP4* p1Out2,
    TENSORMAP4* g1Out,
    TENSORMAP4* g1Out2,
    TENSORMAP2* g1Concat,
    TENSORMAP2* g1Bias,
    TENSORMAP2* policyPass,
    TENSORMAP4* policy,
    CONSTTENSORMAP3* mask,
    const float* maskSum
  ) const {
    const bool applyBNRelu = true;
    p1Conv.apply(trunk, p1Out, false);
    g1Conv.apply(trunk, g1Out, false);
    g1BN.apply(applyBNRelu, g1Out, g1Out2, mask);
    poolRowsGPool(g1Out2, g1Concat, maskSum);
    gpoolToBiasMul.apply(g1Concat, g1Bias);
    addNCBiasInplace(p1Out, g1Bias);
    p1BN.apply(true, p1Out, p1Out2, mask);
    p2Conv.apply(p1Out2, policy, false);
    gpoolToPassMul.apply(g1Concat, policyPass);
  }
};

struct ValueHead {
  string name;
  int version;

  ConvLayer v1Conv;
  BatchNormLayer v1BN;
  ActivationLayer v1Activation;
  MatMulLayer v2Mul;
  MatBiasLayer v2Bias;
  ActivationLayer v2Activation;
  MatMulLayer v3Mul;
  MatBiasLayer v3Bias;
  MatMulLayer sv3Mul;
  MatBiasLayer sv3Bias;
  ConvLayer vOwnershipConv;

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

  ValueHead(const ValueHeadDesc& desc)
    : name(desc.name),
      version(desc.version),
      v1Conv(desc.v1Conv),
      v1BN(desc.v1BN),
      v1Activation(desc.v1Activation),
      v2Mul(desc.v2Mul),
      v2Bias(desc.v2Bias),
      v2Activation(desc.v2Activation),
      v3Mul(desc.v3Mul),
      v3Bias(desc.v3Bias),
      sv3Mul(desc.sv3Mul),
      sv3Bias(desc.sv3Bias),
      vOwnershipConv(desc.vOwnershipConv) {}

  void apply(
    CONSTTENSORMAP4* trunk,
    TENSORMAP4* v1Out,
    TENSORMAP4* v1Out2,
    TENSORMAP2* v1Mean,
    TENSORMAP2* v2Out,
    TENSORMAP2* value,
    TENSORMAP2* scoreValue,
    TENSORMAP4* ownership,
    CONSTTENSORMAP3* mask,
    const float* maskSum
  ) const {
    bool applyBNRelu = true;
    v1Conv.apply(trunk, v1Out, false);
    v1BN.apply(applyBNRelu, v1Out, v1Out2, mask);
    poolRowsValueHead(v1Out2, v1Mean, maskSum);
    v2Mul.apply(v1Mean, v2Out);
    v2Bias.apply(v2Out);
    v2Activation.apply(v2Out, v2Out);
    v3Mul.apply(v2Out, value);
    v3Bias.apply(value);

    sv3Mul.apply(v2Out, scoreValue);
    sv3Bias.apply(scoreValue);

    vOwnershipConv.apply(v1Out2, ownership, false);
  }
};

// Model and Buffer I/O ------------------------------------------------------------------------------------------------

struct Model {
  string name;
  int version;
  int numInputChannels;
  int numInputGlobalChannels;
  int numValueChannels;
  int numScoreValueChannels;
  int numOwnershipChannels;

  Trunk trunk;
  PolicyHead policyHead;
  ValueHead valueHead;

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

  Model(const ModelDesc& desc)
    : name(desc.name), version(desc.version), numInputChannels(desc.numInputChannels),
      numInputGlobalChannels(desc.numInputGlobalChannels),
      numValueChannels(desc.numValueChannels),
      numScoreValueChannels(desc.numScoreValueChannels),
      numOwnershipChannels(desc.numOwnershipChannels),
      trunk(desc.trunk),
      policyHead(desc.policyHead),
      valueHead(desc.valueHead) {}

  void apply(
    CONSTTENSORMAP4* input,
    CONSTTENSORMAP2* inputGlobal,
    TENSORMAP2* inputMatMulOut,
    TENSORMAP4* trunkBuf,
    TENSORMAP4* trunkScratch,
    TENSORMAP4* regularOut,
    TENSORMAP4* regularScratch,
    TENSORMAP4* midIn,
    TENSORMAP4* midScratch,
    TENSORMAP4* gpoolOut,
    TENSORMAP4* gpoolOut2,
    TENSORMAP2* gpoolConcat,
    TENSORMAP2* gpoolBias,

    TENSORMAP4* p1Out,
    TENSORMAP4* p1Out2,
    TENSORMAP4* g1Out,
    TENSORMAP4* g1Out2,
    TENSORMAP2* g1Concat,
    TENSORMAP2* g1Bias,
    TENSORMAP2* policyPass,
    TENSORMAP4* policy,

    TENSORMAP4* v1Out,
    TENSORMAP4* v1Out2,
    TENSORMAP2* v1Mean,
    TENSORMAP2* v2Out,
    TENSORMAP2* value,
    TENSORMAP2* scoreValue,
    TENSORMAP4* ownership,

    TENSORMAP3* mask,
    float* maskSum
  ) const {
    *mask = input->chip(0,0);
    computeMaskSum(mask,maskSum);

    trunk.apply(
      input,
      inputGlobal,
      inputMatMulOut,
      trunkBuf,
      trunkScratch,
      regularOut,
      regularScratch,
      midIn,
      midScratch,
      gpoolOut,
      gpoolOut2,
      gpoolConcat,
      gpoolBias,
      mask,
      maskSum
    );
    policyHead.apply(
      trunkBuf,
      p1Out,
      p1Out2,
      g1Out,
      g1Out2,
      g1Concat,
      g1Bias,
      policyPass,
      policy,
      mask,
      maskSum
    );
    valueHead.apply(
      trunkBuf,
      v1Out,
      v1Out2,
      v1Mean,
      v2Out,
      value,
      scoreValue,
      ownership,
      mask,
      maskSum
    );
  }
};

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

struct Buffers {
  TENSOR2 inputMatMulOut;
  TENSOR4 trunk;
  TENSOR4 trunkScratch;
  TENSOR4 regularOut;
  TENSOR4 regularScratch;
  TENSOR4 midIn;
  TENSOR4 midScratch;
  TENSOR4 gpoolOut;
  TENSOR4 gpoolOut2;
  TENSOR2 gpoolConcat;
  TENSOR2 gpoolBias;

  TENSOR4 p1Out;
  TENSOR4 p1Out2;
  TENSOR4 g1Out;
  TENSOR4 g1Out2;
  TENSOR2 g1Concat;
  TENSOR2 g1Bias;
  TENSOR2 policyPass;
  TENSOR4 policy;

  TENSOR4 v1Out;
  TENSOR4 v1Out2;
  TENSOR2 v1Mean;
  TENSOR2 v2Out;
  TENSOR2 value;
  TENSOR2 scoreValue;
  TENSOR4 ownership;

  TENSOR3 mask;
  vector<float> maskSum;

  Buffers(
    const ModelDesc& desc,
    int maxBatchSize,
    int nnXLen,
    int nnYLen
  ) :
    inputMatMulOut(desc.trunk.trunkNumChannels, maxBatchSize),
    trunk(desc.trunk.trunkNumChannels, nnXLen, nnYLen, maxBatchSize),
    trunkScratch(desc.trunk.trunkNumChannels, nnXLen, nnYLen, maxBatchSize),
    regularOut(desc.trunk.regularNumChannels, nnXLen, nnYLen, maxBatchSize),
    regularScratch(desc.trunk.regularNumChannels, nnXLen, nnYLen, maxBatchSize),
    midIn(desc.trunk.midNumChannels, nnXLen, nnYLen, maxBatchSize),
    midScratch(desc.trunk.midNumChannels, nnXLen, nnYLen, maxBatchSize),
    gpoolOut(desc.trunk.gpoolNumChannels, nnXLen, nnYLen, maxBatchSize),
    gpoolOut2(desc.trunk.gpoolNumChannels, nnXLen, nnYLen, maxBatchSize),
    gpoolConcat(desc.trunk.gpoolNumChannels*3, maxBatchSize),
    gpoolBias(desc.trunk.regularNumChannels, maxBatchSize),

    p1Out(desc.policyHead.p1Conv.outChannels, nnXLen, nnYLen, maxBatchSize),
    p1Out2(desc.policyHead.p1Conv.outChannels, nnXLen, nnYLen, maxBatchSize),
    g1Out(desc.policyHead.g1Conv.outChannels, nnXLen, nnYLen, maxBatchSize),
    g1Out2(desc.policyHead.g1Conv.outChannels, nnXLen, nnYLen, maxBatchSize),
    g1Concat(desc.policyHead.g1Conv.outChannels*3, maxBatchSize),
    g1Bias(desc.policyHead.gpoolToBiasMul.outChannels, maxBatchSize),
    policyPass(desc.policyHead.gpoolToPassMul.outChannels, maxBatchSize),
    policy(desc.policyHead.p2Conv.outChannels, nnXLen, nnYLen, maxBatchSize),

    v1Out(desc.valueHead.v1Conv.outChannels, nnXLen, nnYLen, maxBatchSize),
    v1Out2(desc.valueHead.v1Conv.outChannels, nnXLen, nnYLen, maxBatchSize),
    v1Mean(desc.valueHead.v1Conv.outChannels*3, maxBatchSize),
    v2Out(desc.valueHead.v2Mul.outChannels, maxBatchSize),
    value(desc.valueHead.v3Mul.outChannels, maxBatchSize),
    scoreValue(desc.valueHead.sv3Mul.outChannels, maxBatchSize),
    ownership(desc.valueHead.vOwnershipConv.outChannels, nnXLen, nnYLen, maxBatchSize),

    mask(nnXLen, nnYLen, maxBatchSize),
    maskSum(maxBatchSize)
  {}
};

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

struct InputBuffers {
  int maxBatchSize;

  size_t singleInputElts;
  size_t singleInputGlobalElts;

  size_t singlePolicyPassResultElts;
  size_t singlePolicyResultElts;
  size_t singleValueResultElts;
  size_t singleScoreValueResultElts;
  size_t singleOwnershipResultElts;

  std::vector<float> spatialInput;
  std::vector<float> globalInput;

  InputBuffers(const LoadedModel* loadedModel, int maxBatchSz, int nnXLen, int nnYLen) {
    const ModelDesc& m = loadedModel->modelDesc;

    int xSize = nnXLen;
    int ySize = nnYLen;

    maxBatchSize = maxBatchSz;
    singleInputElts = m.numInputChannels * xSize * ySize;
    singleInputGlobalElts = m.numInputGlobalChannels;

    singlePolicyPassResultElts = (size_t)(1);
    singlePolicyResultElts = (size_t)(xSize * ySize);
    singleValueResultElts = (size_t)m.numValueChannels;
    singleScoreValueResultElts = (size_t)m.numScoreValueChannels;
    singleOwnershipResultElts = (size_t)m.numOwnershipChannels * xSize * ySize;

    assert(NNModelVersion::getNumSpatialFeatures(m.version) == m.numInputChannels);
    assert(NNModelVersion::getNumGlobalFeatures(m.version) == m.numInputGlobalChannels);

    spatialInput = vector<float>(m.numInputChannels * xSize * ySize * maxBatchSize);
    globalInput = vector<float>(m.numInputGlobalChannels * maxBatchSize);
  }

  ~InputBuffers() { }

  InputBuffers() = delete;
  InputBuffers(const InputBuffers&) = delete;
  InputBuffers& operator=(const InputBuffers&) = delete;
};

InputBuffers* NeuralNet::createInputBuffers(const LoadedModel* loadedModel, int maxBatchSize, int nnXLen, int nnYLen) {
  return new InputBuffers(loadedModel, maxBatchSize, nnXLen, nnYLen);
}
void NeuralNet::freeInputBuffers(InputBuffers* inputBuffers) {
  delete inputBuffers;
}

// float* NeuralNet::getBatchEltSpatialInplace(InputBuffers* inputBuffers, int nIdx) {
//   assert(nIdx < inputBuffers->maxBatchSize);
//   return inputBuffers->spatialInput.data() + (inputBuffers->singleInputElts * nIdx);
// }

// float* NeuralNet::getBatchEltGlobalInplace(InputBuffers* inputBuffers, int rowIdx) {
//   assert(rowIdx < inputBuffers->maxBatchSize);
//   return inputBuffers->globalInput.data() + (inputBuffers->singleInputGlobalElts * rowIdx);
// }

// int NeuralNet::getBatchEltSpatialLen(const InputBuffers* inputBuffers) {
//   return inputBuffers->singleInputElts;
// }
// int NeuralNet::getBatchEltGlobalLen(const InputBuffers* inputBuffers) {
//   return inputBuffers->singleInputGlobalElts;
// }

// bool* NeuralNet::getSymmetriesInplace(InputBuffers* inputBuffers) {
//   return inputBuffers->symmetriesBuffer;
// }

// NeuralNet -----------------------------------------------------------------------------------------------------------

void NeuralNet::globalInitialize() {
  // no-op for cpu
}

void NeuralNet::globalCleanup() {
  // no-op for cpu
}

struct ComputeHandle {
  const ComputeContext* context;
  int maxBatchSize;
  bool inputsUseNHWC;
  Model model;
  Buffers buffers;

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

  ComputeHandle(const ComputeContext* ctx, const LoadedModel& loadedModel, int maxBSize, bool iNHWC)
    : context(ctx),
      maxBatchSize(maxBSize),
      inputsUseNHWC(iNHWC),
      model(loadedModel.modelDesc),
      buffers(loadedModel.modelDesc,maxBSize,ctx->nnXLen,ctx->nnYLen)
  {}
};

ComputeHandle* NeuralNet::createComputeHandle(
  ComputeContext* context,
  const LoadedModel* loadedModel,
  Logger* logger,
  int maxBatchSize,
  bool requireExactNNLen,
  bool inputsUseNHWC,
  int gpuIdxForThisThread
) {
  if(logger != NULL) {
    logger->write("Eigen (CPU) backend: Model version " + Global::intToString(loadedModel->modelDesc.version));
    logger->write("Eigen (CPU) backend: Model name: " + loadedModel->modelDesc.name);
  }

  (void)requireExactNNLen; //We don't bother with mask optimizations if we know exact sizes right now.
  (void)gpuIdxForThisThread; //Doesn't matter

  if(!inputsUseNHWC)
    throw StringError("Eigen backend: inputsUseNHWC = false unsupported");
  return new ComputeHandle(context, *loadedModel, maxBatchSize, inputsUseNHWC);
}

void NeuralNet::freeComputeHandle(ComputeHandle* gpuHandle) {
  delete gpuHandle;
}

void NeuralNet::getOutput(
  ComputeHandle* computeHandle,
  InputBuffers* inputBuffers,
  int numBatchEltsFilled,
  NNResultBuf** inputBufs,
  int symmetry,
  vector<NNOutput*>& outputs
) {
  assert(numBatchEltsFilled <= inputBuffers->maxBatchSize);
  assert(numBatchEltsFilled > 0);
  int batchSize = numBatchEltsFilled;
  int nnXLen = computeHandle->context->nnXLen;
  int nnYLen = computeHandle->context->nnYLen;
  int version = computeHandle->model.version;

  int numSpatialFeatures = NNModelVersion::getNumSpatialFeatures(version);
  int numGlobalFeatures = NNModelVersion::getNumGlobalFeatures(version);
  assert(numSpatialFeatures == computeHandle->model.numInputChannels);
  assert(numSpatialFeatures * nnXLen * nnYLen == inputBuffers->singleInputElts);
  assert(numGlobalFeatures == inputBuffers->singleInputGlobalElts);

  for(int nIdx = 0; nIdx<batchSize; nIdx++) {
    float* rowSpatialInput = inputBuffers->spatialInput.data() + (inputBuffers->singleInputElts * nIdx);
    float* rowGlobalInput = inputBuffers->globalInput.data() + (inputBuffers->singleInputGlobalElts * nIdx);

    const float* rowGlobal = inputBufs[nIdx]->rowGlobal;
    const float* rowSpatial = inputBufs[nIdx]->rowSpatial;
    std::copy(rowGlobal,rowGlobal+numGlobalFeatures,rowGlobalInput);
    SymmetryHelpers::copyInputsWithSymmetry(rowSpatial, rowSpatialInput, 1, nnYLen, nnXLen, numSpatialFeatures, computeHandle->inputsUseNHWC, symmetry);
  }

  Buffers& buffers = computeHandle->buffers;

  CONSTTENSORMAP4 input(inputBuffers->spatialInput.data(), numSpatialFeatures, nnXLen, nnYLen, batchSize);
  CONSTTENSORMAP2 inputGlobal(inputBuffers->globalInput.data(), numGlobalFeatures, batchSize);

#define MAP4(NAME) TENSORMAP4 NAME(buffers.NAME.data(), buffers.NAME.dimension(0), buffers.NAME.dimension(1), buffers.NAME.dimension(2), batchSize)
#define MAP3(NAME) TENSORMAP3 NAME(buffers.NAME.data(), buffers.NAME.dimension(0), buffers.NAME.dimension(1), batchSize)
#define MAP2(NAME) TENSORMAP2 NAME(buffers.NAME.data(), buffers.NAME.dimension(0), batchSize)

  MAP2(inputMatMulOut);
  MAP4(trunk);
  MAP4(trunkScratch);
  MAP4(regularOut);
  MAP4(regularScratch);
  MAP4(midIn);
  MAP4(midScratch);
  MAP4(gpoolOut);
  MAP4(gpoolOut2);
  MAP2(gpoolConcat);
  MAP2(gpoolBias);
  MAP4(p1Out);
  MAP4(p1Out2);
  MAP4(g1Out);
  MAP4(g1Out2);
  MAP2(g1Concat);
  MAP2(g1Bias);
  MAP2(policyPass);
  MAP4(policy);
  MAP4(v1Out);
  MAP4(v1Out2);
  MAP2(v1Mean);
  MAP2(v2Out);
  MAP2(value);
  MAP2(scoreValue);
  MAP4(ownership);
  MAP3(mask);
  vector<float>& maskSum = buffers.maskSum;

  computeMaskSum(&mask,maskSum.data());

  computeHandle->model.apply(
    &input,
    &inputGlobal,
    &inputMatMulOut,
    &trunk,
    &trunkScratch,
    &regularOut,
    &regularScratch,
    &midIn,
    &midScratch,
    &gpoolOut,
    &gpoolOut2,
    &gpoolConcat,
    &gpoolBias,
    &p1Out,
    &p1Out2,
    &g1Out,
    &g1Out2,
    &g1Concat,
    &g1Bias,
    &policyPass,
    &policy,
    &v1Out,
    &v1Out2,
    &v1Mean,
    &v2Out,
    &value,
    &scoreValue,
    &ownership,
    &mask,
    maskSum.data()
  );

  assert(outputs.size() == batchSize);

  float* policyData = policy.data();
  float* policyPassData = policyPass.data();
  float* valueData = value.data();
  float* scoreValueData = scoreValue.data();
  float* ownershipData = ownership.data();

  for(int row = 0; row < batchSize; row++) {
    NNOutput* output = outputs[row];
    assert(output->nnXLen == nnXLen);
    assert(output->nnYLen == nnYLen);

    const float* policySrcBuf = policyData + row * inputBuffers->singlePolicyResultElts;
    float* policyProbs = output->policyProbs;

    //These are not actually correct, the client does the postprocessing to turn them into
    //policy probabilities and white game outcome probabilities
    //Also we don't fill in the nnHash here either
    SymmetryHelpers::copyOutputsWithSymmetry(policySrcBuf, policyProbs, 1, nnYLen, nnXLen, symmetry);
    policyProbs[inputBuffers->singlePolicyResultElts] = policyPassData[row];

    int numValueChannels = computeHandle->model.numValueChannels;
    assert(numValueChannels == 3);
    output->whiteWinProb = valueData[row * numValueChannels];
    output->whiteLossProb = valueData[row * numValueChannels + 1];
    output->whiteNoResultProb = valueData[row * numValueChannels + 2];

    //As above, these are NOT actually from white's perspective, but rather the player to move.
    //As usual the client does the postprocessing.
    if(output->whiteOwnerMap != NULL) {
      const float* ownershipSrcBuf = ownershipData + row * nnXLen * nnYLen;
      assert(computeHandle->model.numOwnershipChannels == 1);
      SymmetryHelpers::copyOutputsWithSymmetry(ownershipSrcBuf, output->whiteOwnerMap, 1, nnYLen, nnXLen, symmetry);
    }

    if(version >= 8) {
      int numScoreValueChannels = computeHandle->model.numScoreValueChannels;
      assert(numScoreValueChannels == 4);
      output->whiteScoreMean = scoreValueData[row * numScoreValueChannels];
      output->whiteScoreMeanSq = scoreValueData[row * numScoreValueChannels + 1];
      output->whiteLead = scoreValueData[row * numScoreValueChannels + 2];
      output->varTimeLeft = scoreValueData[row * numScoreValueChannels + 3];
    }
    else if(version >= 4) {
      int numScoreValueChannels = computeHandle->model.numScoreValueChannels;
      assert(numScoreValueChannels == 2);
      output->whiteScoreMean = scoreValueData[row * numScoreValueChannels];
      output->whiteScoreMeanSq = scoreValueData[row * numScoreValueChannels + 1];
      output->whiteLead = output->whiteScoreMean;
      output->varTimeLeft = 0;
    }
    else if(version >= 3) {
      int numScoreValueChannels = computeHandle->model.numScoreValueChannels;
      assert(numScoreValueChannels == 1);
      output->whiteScoreMean = scoreValueData[row * numScoreValueChannels];
      //Version 3 neural nets don't have any second moment output, implicitly already folding it in, so we just use the mean squared
      output->whiteScoreMeanSq = output->whiteScoreMean * output->whiteScoreMean;
      output->whiteLead = output->whiteScoreMean;
      output->varTimeLeft = 0;
    }
    else {
      ASSERT_UNREACHABLE;
    }
  }
}

void NeuralNet::printDevices() {
}

// FOR TESTING ---------------------------------------------------------------------------------------------------------
bool NeuralNet::testEvaluateConv(
  const ConvLayerDesc* desc,
  int batchSize,
  int nnXLen,
  int nnYLen,
  bool useFP16,
  bool useNHWC,
  const std::vector<float>& inputBuffer,
  std::vector<float>& outputBuffer
) {
  if(!useNHWC || useFP16)
    return false;
  ConvLayer layer(*desc);
  TENSORMAP4 inTensor(
    (float*)inputBuffer.data(), desc->inChannels, nnXLen, nnYLen, batchSize);
  TENSOR4 outTensorBuf(desc->outChannels, nnXLen, nnYLen, batchSize);
  TENSORMAP4 outTensor(outTensorBuf);

  layer.apply(&inTensor, &outTensor, false);

  outputBuffer.resize(outTensorBuf.size());
  memcpy(outputBuffer.data(), outTensorBuf.data(), sizeof(SCALAR) * outTensorBuf.size());
  return true;
}

// Mask should be in 'NHW' format (no "C" channel).
bool NeuralNet::testEvaluateBatchNorm(
  const BatchNormLayerDesc* desc,
  int batchSize,
  int nnXLen,
  int nnYLen,
  bool useFP16,
  bool useNHWC,
  const std::vector<float>& inputBuffer,
  const std::vector<float>& maskBuffer,
  std::vector<float>& outputBuffer
) {
  if(!useNHWC || useFP16)
    return false;
  BatchNormLayer layer(*desc);
  TENSORMAP4 inTensor((float*)inputBuffer.data(), desc->numChannels, nnXLen, nnYLen, batchSize);
  TENSORMAP3 mask((float*)maskBuffer.data(), nnXLen, nnYLen, batchSize);
  TENSOR4 outTensorBuf(desc->numChannels, nnXLen, nnYLen, batchSize);
  TENSORMAP4 outTensor(outTensorBuf);

  layer.apply(false, &inTensor, &outTensor, &mask);

  outputBuffer.resize(outTensorBuf.size());
  memcpy(outputBuffer.data(), outTensorBuf.data(), sizeof(SCALAR) * outTensorBuf.size());
  return true;
}

// CR lpuchallafiore: test evaluate activation layer.
// CR lpuchallafiore: test evaluate matmul layer.
// CR lpuchallafiore: test evaluate matbias layer.

bool NeuralNet::testEvaluateResidualBlock(
  const ResidualBlockDesc* desc,
  int batchSize,
  int nnXLen,
  int nnYLen,
  bool useFP16,
  bool useNHWC,
  const std::vector<float>& inputBuffer,
  const std::vector<float>& maskBuffer,
  std::vector<float>& outputBuffer
) {
  if(!useNHWC || useFP16)
    return false;
  ResidualBlock block(*desc);
  TENSORMAP4 inTensor((float*)inputBuffer.data(), desc->preBN.numChannels, nnXLen, nnYLen, batchSize);
  TENSORMAP3 mask((float*)maskBuffer.data(), nnXLen, nnYLen, batchSize);

  TENSOR4 trunkBuf(desc->preBN.numChannels, nnXLen, nnYLen, batchSize);
  TENSOR4 trunkScratchBuf(desc->preBN.numChannels, nnXLen, nnYLen, batchSize);
  TENSOR4 midBuf(desc->finalConv.inChannels, nnXLen, nnYLen, batchSize);
  TENSOR4 midScratchBuf(desc->finalConv.inChannels, nnXLen, nnYLen, batchSize);

  TENSORMAP4 trunk(trunkBuf);
  TENSORMAP4 trunkScratch(trunkScratchBuf);
  TENSORMAP4 mid(midBuf);
  TENSORMAP4 midScratch(midScratchBuf);

  trunk = inTensor;

  block.apply(
    &trunk,
    &trunkScratch,
    NULL,
    NULL,
    &mid,
    &midScratch,
    NULL,
    NULL,
    NULL,
    NULL,
    &mask,
    NULL
  );

  outputBuffer.resize(trunk.size());
  memcpy(outputBuffer.data(), trunk.data(), sizeof(SCALAR) * trunk.size());
  return true;
}

bool NeuralNet::testEvaluateGlobalPoolingResidualBlock(
  const GlobalPoolingResidualBlockDesc* desc,
  int batchSize,
  int nnXLen,
  int nnYLen,
  bool useFP16,
  bool useNHWC,
  const std::vector<float>& inputBuffer,
  const std::vector<float>& maskBuffer,
  std::vector<float>& outputBuffer) {
  if(!useNHWC || useFP16)
    return false;

  GlobalPoolingResidualBlock block(*desc);

  TENSORMAP4 inTensor((float*)inputBuffer.data(), desc->preBN.numChannels, nnXLen, nnYLen, batchSize);
  TENSORMAP3 mask((float*)maskBuffer.data(), nnXLen, nnYLen, batchSize);

  TENSOR4 trunkBuf(desc->preBN.numChannels, nnXLen, nnYLen, batchSize);
  TENSOR4 trunkScratchBuf(desc->preBN.numChannels, nnXLen, nnYLen, batchSize);
  TENSOR4 regularOutBuf(desc->finalConv.inChannels, nnXLen, nnYLen, batchSize);
  TENSOR4 regularScratchBuf(desc->finalConv.inChannels, nnXLen, nnYLen, batchSize);
  TENSOR4 gpoolOutBuf(desc->gpoolConv.outChannels, nnXLen, nnYLen, batchSize);
  TENSOR4 gpoolOut2Buf(desc->gpoolConv.outChannels, nnXLen, nnYLen, batchSize);
  TENSOR2 gpoolConcatBuf(desc->gpoolConv.outChannels*3, batchSize);
  TENSOR2 gpoolBiasBuf(desc->gpoolToBiasMul.outChannels, batchSize);

  TENSORMAP4 trunk(trunkBuf);
  TENSORMAP4 trunkScratch(trunkScratchBuf);
  TENSORMAP4 regularOut(regularOutBuf);
  TENSORMAP4 regularScratch(regularScratchBuf);
  TENSORMAP4 gpoolOut(gpoolOutBuf);
  TENSORMAP4 gpoolOut2(gpoolOut2Buf);
  TENSORMAP2 gpoolConcat(gpoolConcatBuf);
  TENSORMAP2 gpoolBias(gpoolBiasBuf);

  std::vector<float> maskSum(batchSize);
  computeMaskSum(&mask,maskSum.data());

  trunk = inTensor;

  block.apply(
    &trunk,
    &trunkScratch,
    &regularOut,
    &regularScratch,
    NULL,
    NULL,
    &gpoolOut,
    &gpoolOut2,
    &gpoolConcat,
    &gpoolBias,
    &mask,
    maskSum.data()
  );

  outputBuffer.resize(trunk.size());
  memcpy(outputBuffer.data(), trunk.data(), sizeof(SCALAR) * trunk.size());

  return true;
}

#endif  // USE_EIGEN_BACKEND

// CR lpuchallafiore: test evaluate Trunk
// CR lpuchallafiore: test evaluate Policy Head
// CR lpuchallafiore: test evaluate Value Head