#include <algorithm>

#include <Eigen/Core>

#include "algorithm.hpp"
#include "drawing.hpp"
#include "drawing_template.hpp"

Light makeLight()
{
    auto light = Light{};
    light.position_world = Vector4d{ 0, -10, 0, 1 };
    light.power = Vector4d{ 1, 1, 1, 0 };
    return light;
}

bool isBehindCamera(const Vector4d& v0, const Vector4d& v1, const Vector4d& v2)
{
	return v0(2) <= 0 || v1(2) <= 0 || v2(2) <= 0;
}

void vertexShader(Vertices& vertices, const Environment& environment)
{
	const auto num_vertices = vertices.size();
    const auto image_from_camera = imageFromCamera(environment.intrinsics);
    const auto camera_from_world = cameraFromWorld(environment.extrinsics);
    const auto image_from_world = Matrix4d{ image_from_camera * camera_from_world };

	for (size_t i = 0; i < num_vertices; ++i)
	{
		const auto position_world = vertices.positions_world[i];
		const auto position_image = Vector4d{image_from_world * position_world};
		vertices.positions_image[i] = position_image / position_image(3);
	}
}

Uint32 packColorArgb(Uint32 a, Uint32 r, Uint32 g, Uint32 b)
{
	return (a << 24) | (r << 16) | (g << 8) | (b << 0);
}

void basicPixelShader(const Vertex& vertex, Pixels& pixels, size_t index)
{
	using namespace vertex_index;

	if (vertex(BARY0) < 0.0 || 1.0 < vertex(BARY0)) return;
	if (vertex(BARY1) < 0.0 || 1.0 < vertex(BARY1)) return;
	if (vertex(BARY2) < 0.0 || 1.0 < vertex(BARY2)) return;

	const double disparity = vertex(DISPARITY);
	// TODO: try if defered rendering is faster.
    if (disparity <= pixels.disparities[index]) return;

	const Uint32 c = static_cast<Uint32>(clamp(500 * disparity, 0.0, 255.0));
	pixels.colors[index] = packColorArgb(255, c, c, c);
	pixels.disparities[index] = disparity;
}

Uint32 clampColor(double c)
{
    return static_cast<Uint32>(clamp(c, 0.0, 255.0));
}

struct PixelEnvironment
{
    EIGEN_MAKE_ALIGNED_OPERATOR_NEW
    const Texture* surface_texture;
    Vector4d surface_normal_world;
    Vector4d light_position_world;
    Vector4d light_power;
};

struct PixelShader
{
    Pixels* pixels;
    PixelEnvironment pixel_environment;
    void operator()(const Vertex& vertex, size_t index) const
    {
        using namespace vertex_index;

        if (vertex(BARY0) < 0.0 || 1.0 < vertex(BARY0)) return;
        if (vertex(BARY1) < 0.0 || 1.0 < vertex(BARY1)) return;
        if (vertex(BARY2) < 0.0 || 1.0 < vertex(BARY2)) return;

        const double disparity = vertex(DISPARITY);
        // TODO: try if defered rendering is faster.
        if (disparity <= pixels->disparities[index]) return;

        const Uint32 c = clampColor(255 * 2 * disparity);
        pixels->colors[index] = packColorArgb(255, c, c, c);
        pixels->disparities[index] = disparity;

        if (pixel_environment.surface_texture->empty()) return;

        const auto x = vertex(X) / disparity;
        const auto y = vertex(Y) / disparity;
        const auto z = vertex(Z) / disparity;

        const auto u = vertex(U) / disparity;
        const auto v = vertex(V) / disparity;
        const auto position_world = Vector4d{x, y, z, 1};

        //const auto light = disparity;
        const auto light = 16.0 / (position_world - pixel_environment.light_position_world).squaredNorm();

        const auto color = pixel_environment.surface_texture->sample(u, v);
        const auto red   = clampColor(light * color(RED));
        const auto green = clampColor(light * color(GREEN));
        const auto blue  = clampColor(light * color(BLUE));
        pixels->colors[index] = packColorArgb(255, red, green, blue);
    }
};

void drawPoint(Pixels& pixels, const Vector4d& vertex_image)
{
    const auto x = static_cast<int>(vertex_image.x());
    const auto y = static_cast<int>(vertex_image.y());
    const auto disparity = vertex_image.z();

    if (0 <= x && x < pixels.width && 0 <= y && y < pixels.height && disparity > 0.0)
    {
		pixels.colors[y * pixels.width + x] = 0xFFFF00FF;
    }
}

void drawPoints(Pixels& pixels, const Vectors4d& vertices_image)
{
    for (const auto& vertex_image : vertices_image)
    {
        drawPoint(pixels, vertex_image);
    }
}

Vector4d computeNormal(const Vector4d& v0, const Vector4d& v1, const Vector4d& v2)
{
    const auto a = Eigen::Vector3d(v1(0) - v0(0), v1(1) - v0(1), v1(2) - v0(2));
    const auto b = Eigen::Vector3d(v2(0) - v0(0), v2(1) - v0(1), v2(2) - v0(2));
    //const auto c = Eigen::Vector3d{a.cross(b).normalized()};
    //return {c(0), c(1), c(2), 0};
    return Vector4d::Zero();
}

void drawTriangles(Pixels& pixels, const Vertices& vertices,
    const Triangles& triangles, const Textures& textures, const Environment& environment)
{
	const auto num_triangles = triangles.size();

	fill(pixels.disparities, 0.0);
	fill(pixels.colors, 0);

    auto pixel_shader = PixelShader{};
    auto vertex0 = Vertex();
    auto vertex1 = Vertex();
    auto vertex2 = Vertex();

    for (size_t i = 0; i < num_triangles; ++i)
    {
        const auto i0 = triangles.indices0[i];
        const auto i1 = triangles.indices1[i];
        const auto i2 = triangles.indices2[i];

        const auto& v0 = vertices.positions_image[i0];
        const auto& v1 = vertices.positions_image[i1];
        const auto& v2 = vertices.positions_image[i2];

        const auto& t0 = vertices.positions_texture[i0];
        const auto& t1 = vertices.positions_texture[i1];
        const auto& t2 = vertices.positions_texture[i2];

        const auto& p0 = vertices.positions_world[i0];
        const auto& p1 = vertices.positions_world[i1];
        const auto& p2 = vertices.positions_world[i2];

		if (isBehindCamera(v0, v1, v2)) continue;

		using namespace vertex_index;

		vertex0(BARY0) = 1.0;
		vertex0(BARY1) = 0.0;
		vertex0(BARY2) = 0.0;
		vertex0(DISPARITY) = v0(2);
        vertex0(U) = t0(0) * v0(2);
        vertex0(V) = t0(1) * v0(2);
        vertex0(X) = p0(0) * v0(2);
        vertex0(Y) = p0(1) * v0(2);
        vertex0(Z) = p0(2) * v0(2);

		vertex1(BARY0) = 0.0;
		vertex1(BARY1) = 1.0;
		vertex1(BARY2) = 0.0;
		vertex1(DISPARITY) = v1(2);
        vertex1(U) = t1(0) * v1(2);
        vertex1(V) = t1(1) * v1(2);
        vertex1(X) = p1(0) * v1(2);
        vertex1(Y) = p1(1) * v1(2);
        vertex1(Z) = p1(2) * v1(2);

		vertex2(BARY0) = 0.0;
		vertex2(BARY1) = 0.0;
		vertex2(BARY2) = 1.0;
		vertex2(DISPARITY) = v2(2);
        vertex2(U) = t2(0) * v2(2);
        vertex2(V) = t2(1) * v2(2);
        vertex2(X) = p2(0) * v2(2);
        vertex2(Y) = p2(1) * v2(2);
        vertex2(Z) = p2(2) * v2(2);

        const auto texture_index = triangles.texture_indices[i];
        pixel_shader.pixels = &pixels;
        pixel_shader.pixel_environment.surface_texture = &textures[texture_index];
        pixel_shader.pixel_environment.light_position_world = environment.light.position_world;
        pixel_shader.pixel_environment.light_power = environment.light.power;

        //renderTriangleTemplate(pixels, basicPixelShader, v0, v1, v2, vertex0, vertex1, vertex2);
        renderTriangleTemplate(
            v0, v1, v2, vertex0, vertex1, vertex2,
            pixels.width, pixels.height, pixel_shader);
    }
}


#include <iostream>
#include <pybind11/pybind11.h>
#include <pybind11/numpy.h>

#include <Eigen/Core>
#include <igl/polar_dec.h>
#include <igl/polar_svd.h>

namespace py = pybind11;

template<typename float_t>
py::array_t<float_t> inv3(py::array_t<float_t, py::array::c_style> & Ts)
{
    auto Ts_buf = Ts.request();
    float_t *pT = (float_t*)Ts_buf.ptr;

    auto result = py::array_t<float_t, py::array::c_style>(Ts_buf.size);
    auto result_buf = result.request();
    float_t *pR = (float_t*)result_buf.ptr;

    for (size_t idx = 0; idx < Ts_buf.shape[0]; idx++) {
        const float_t T00 = pT[0], T01 = pT[1], T02 = pT[2];
        const float_t T10 = pT[3], T11 = pT[4], T12 = pT[5];
        const float_t T20 = pT[6], T21 = pT[7], T22 = pT[8];
        const float_t det = T00 * (T22 * T11 - T21 * T12) \
                         - T10 * (T22 * T01 - T21 * T02) \
                         + T20 * (T12 * T01 - T11 * T02);
        const float_t invDet = 1. / det;
        pR[0] =  (T11 * T22 - T21 * T12) * invDet;
        pR[1] = -(T01 * T22 - T02 * T21) * invDet;
        pR[2] =  (T01 * T12 - T02 * T11) * invDet;
        pR[3] = -(T10 * T22 - T12 * T20) * invDet;
        pR[4] =  (T00 * T22 - T02 * T20) * invDet;
        pR[5] = -(T00 * T12 - T10 * T02) * invDet;
        pR[6] =  (T10 * T21 - T20 * T11) * invDet;
        pR[7] = -(T00 * T21 - T20 * T01) * invDet;
        pR[8] =  (T00 * T11 - T10 * T01) * invDet;

        pT += 3*3;
        pR += 3*3;
    }

    return result;
}

template<typename float_t>
py::array_t<float_t> inv2(py::array_t<float_t, py::array::c_style> & Ts)
{
    auto Ts_buf = Ts.request();
    float_t *pT = (float_t*)Ts_buf.ptr;

    auto result = py::array_t<float_t, py::array::c_style>(Ts_buf.size);
    auto result_buf = result.request();
    float_t *pR = (float_t*)result_buf.ptr;

    for (size_t idx = 0; idx < Ts_buf.shape[0]; idx++) {
        const float_t T00 = pT[0], T01 = pT[1];
        const float_t T10 = pT[2], T11 = pT[3];
        const float_t det = T00 * T11 - T01 * T10;
        const float_t invDet = 1. / det;
        pR[0] =  T11 * invDet;
        pR[1] = -1 * T01 * invDet;
        pR[2] = -1 * T10 * invDet;
        pR[3] = T00 * invDet;

        pT += 2*2;
        pR += 2*2;
    }

    return result;
}

template<typename float_t>
py::array_t<float_t> matmat(
        py::array_t<float_t, py::array::c_style> & a,
        py::array_t<float_t, py::array::c_style> & b
    )
{
    auto a_buf = a.request();
    float_t *p_a = (float_t*)a_buf.ptr;
    auto b_buf = b.request();
    float_t *p_b = (float_t*)b_buf.ptr;

    auto result = py::array_t<float_t, py::array::c_style>(
            {a_buf.shape[0], a_buf.shape[1], b_buf.shape[2]});
    auto result_buf = result.request();
    float_t *p_res = (float_t*)result_buf.ptr;

    const size_t n_rows_a = a_buf.shape[1];
    const size_t n_cols_a = a_buf.shape[2];
    const size_t n_rows_b = b_buf.shape[1];
    const size_t n_cols_b = b_buf.shape[2];
    assert(n_cols_a == n_rows_b);
    for (size_t idx = 0; idx < a_buf.shape[0]; idx++) {
        for (size_t row_a = 0; row_a < n_rows_a; row_a++) {
            for (size_t col_b = 0; col_b < n_cols_b; col_b++) {
                float_t sum = 0.0;
                for (size_t k = 0; k < n_cols_a; k++) {
                    const float_t ai = p_a[row_a * n_cols_a + k];
                    const float_t bi = p_b[k * n_cols_b + col_b];
                    sum += ai * bi;
                }
                *p_res = sum;
                p_res++;
            }
        }
        p_a += n_cols_a * n_rows_a;
        p_b += n_cols_b * n_rows_b;
    }

    return result;
}

template<typename float_t>
py::array_t<float_t> matvec(
        py::array_t<float_t, py::array::c_style> & mats,
        py::array_t<float_t, py::array::c_style> & vecs
    )
{
    auto mats_buf = mats.request();
    float_t *p_mats = (float_t*)mats_buf.ptr;
    auto vecs_buf = vecs.request();
    float_t *p_vecs = (float_t*)vecs_buf.ptr;

    auto result = py::array_t<float_t, py::array::c_style>({mats_buf.shape[0], mats_buf.shape[1]});
    auto result_buf = result.request();
    float_t *p_res = (float_t*)result_buf.ptr;

    const size_t mat_stride1 = mats_buf.strides[1] / sizeof(float_t);
    for (size_t idx = 0; idx < mats_buf.shape[0]; idx++) {
        for (size_t row = 0; row < mats_buf.shape[1]; row++) {
            float_t sum = 0.0;
            for (size_t k = 0; k < mats_buf.shape[2]; k++) {
                sum += *(p_mats++) * p_vecs[k];
            }
            *p_res = sum;
            p_res++;
        }
        p_vecs += vecs_buf.shape[1];
    }

    return result;
}

template<typename float_t>
py::array_t<float_t> cross3(
        py::array_t<float_t, py::array::c_style> & a,
        py::array_t<float_t, py::array::c_style> & b
    )
{
    auto a_buf = a.request();
    float_t *p_a = (float_t*)a_buf.ptr;
    auto b_buf = b.request();
    float_t *p_b = (float_t*)b_buf.ptr;

    auto result = py::array_t<float_t, py::array::c_style>(
            {a_buf.shape[0], a_buf.shape[1]});
    auto result_buf = result.request();
    float_t *p_res = (float_t*)result_buf.ptr;

    for (size_t idx = 0; idx < a_buf.shape[0]; idx++) {
        const double ax = p_a[0];
        const double ay = p_a[1];
        const double az = p_a[2];
        const double bx = p_b[0];
        const double by = p_b[1];
        const double bz = p_b[2];
        p_res[0] = ay * bz - az * by;
        p_res[1] = az * bx - ax * bz;
        p_res[2] = ax * by - ay * bx;
        p_res += 3;
        p_a += 3;
        p_b += 3;
    }

    return result;
}

template<typename float_t>
py::array_t<float_t> multikron(
        py::array_t<float_t, py::array::c_style> & a,
        py::array_t<float_t, py::array::c_style> & b
    )
{
    auto a_buf = a.request();
    float_t *p_a = (float_t*)a_buf.ptr;
    auto b_buf = b.request();
    float_t *p_b = (float_t*)b_buf.ptr;

    const auto n_rows_a = a_buf.shape[1];
    const auto n_cols_a = a_buf.shape[2];
    const auto n_rows_b = b_buf.shape[1];
    const auto n_cols_b = b_buf.shape[2];

    auto result = py::array_t<float_t, py::array::c_style>(
            {a_buf.shape[0], n_rows_a * n_rows_b, n_cols_a * n_cols_b});
    auto result_buf = result.request();
    float_t *p_res = (float_t*)result_buf.ptr;

    for (size_t idx = 0; idx < a_buf.shape[0]; idx++) {
        // iterate over rows of a
        for (size_t row_a = 0; row_a < n_rows_a; row_a++) {
            // iterate over rows of b
            for (size_t row_b = 0; row_b < n_rows_b; row_b++) {
                // iterate over columns of a
                for (size_t col_a = 0; col_a < n_cols_a; col_a++) {
                    // iterate over columns of b
                    for (size_t col_b = 0; col_b < n_cols_b; col_b++) {
                        const float_t ai = p_a[row_a * n_cols_a + col_a];
                        const float_t bi = p_b[row_b * n_cols_b + col_b];
                        *p_res = ai * bi;
                        p_res++;
                    }
                }
            }
        }
        // next matrix
        p_a += n_cols_a * n_rows_a;
        p_b += n_cols_b * n_rows_b;
    }

    return result;
}

std::tuple<py::array_t<double>, py::array_t<double>>
polarDecompose(py::array_t<double, py::array::c_style> Ms)
{
    using RowMat3d = Eigen::Matrix<double, 3, 3, Eigen::RowMajor>;

    auto Ms_raw = Ms.unchecked<3>();
    // allocate output matrices
    auto Rs = py::array_t<double>({Ms_raw.shape(0), 3l, 3l});
    auto Rs_raw = Rs.mutable_unchecked<3>();
    auto Ss = py::array_t<double>({Ms_raw.shape(0), 3l, 3l});
    auto Ss_raw = Ss.mutable_unchecked<3>();

    RowMat3d U;
    RowMat3d V;
    Eigen::Matrix<double,3,1> S;

    for (int i = 0; i < Ms_raw.shape(0); ++i) {
        const RowMat3d Mi = Eigen::Map<const RowMat3d>(Ms_raw.data(i, 0, 0));
        Eigen::Map<RowMat3d> Ri_map(Rs_raw.mutable_data(i, 0, 0));
        Eigen::Map<RowMat3d> Si_map(Ss_raw.mutable_data(i, 0, 0));
        RowMat3d Ri = Ri_map; // performs copy, since igl::polar_dec does not take Eigen::Map
        RowMat3d Si = Si_map;
        //igl::polar_dec(Mi, Ri, Si);
        igl::polar_svd(Mi, Ri, Si, U, S, V);
        // TODO: use igl::polar_svd3x3?
        // write back results
        Ri_map = Ri;
        Si_map = Si;
    }

    return std::make_tuple(Rs, Ss);
}

PYBIND11_MODULE(_fastmath_ext, m) {
    m.def("inv3", &inv3<float>);
    m.def("inv3", &inv3<double>);
    m.def("inv2", &inv2<float>);
    m.def("inv2", &inv2<double>);
    m.def("matmat", &matmat<double>);
    m.def("matvec", &matvec<double>);
    m.def("cross3", &cross3<double>);
    m.def("multikron", &multikron<double>);
    m.def("polar_dec", &polarDecompose);
}


#ifndef _NETWORK_H_
#define _NETWORK_H_

#include <boost/multi_array.hpp>
#include <algorithm>
#include <cassert>
#include <chrono>
#include <cmath>
#include <cstdlib>
#include <cstdint>
#include <ctime>
#include <limits>
#include <fstream>
#include <iostream>
#include <numeric>
#include <memory>
#include <random>
#include <vector>
#include "tbb/tbb.h"
#include "Data.hpp"
#include "Params.hpp"

#ifdef NDEBUG
#define UNREACHABLE() __builtin_unreachable()
#else
#define UNREACHABLE() __builtin_trap()
#endif

/// Sigmoid activation function.
struct Sigmoid {
  static float compute(float z) { return 1.0f / (1.0f + std::exp(-z)); }
  static float deriv(float z) { return compute(z) * (1.0f - compute(z)); }
};

/// Rectified linear activation function.
struct ReLU {
  static float compute(float z) { return std::max(0.0f, z); }
  static float deriv(float z) { return z > 0.0f ? 1.0f : 0.0f; }
};

template<float (*activationFnDeriv)(float)>
struct QuadraticCost {
  static float compute(float activation, float label) {
    return 0.5f * std::pow(std::abs(activation - label), 2);
  }
  static float delta(float z, float activation, float label) {
    return (activation - label) * activationFnDeriv(z);
  }
};

struct CrossEntropyCost {
  static float compute(float activation, float label) {
    return (-label * std::log(activation))
             - ((1.0f - label) * std::log(1.0f - activation));
  }
  static float delta(float z, float activation, float label) {
    return activation - label;
  }
};

/// Helper functions for conversions between 1D and 3D coordinates.
static inline unsigned getX(unsigned index, unsigned dimX) {
  return index % dimX;
}
static inline unsigned getY(unsigned index, unsigned dimX, unsigned dimY) {
  return (index / dimX) % dimY;
}
static inline unsigned getZ(unsigned index, unsigned dimX, unsigned dimY) {
  return (index / (dimX * dimY)) /* % dimZ */;
}
static inline unsigned getIndex(unsigned x, unsigned y, unsigned z,
                                unsigned dimX, unsigned dimY) {
  return ((dimX * dimY) * z) + (dimX * y) + x;
}

template<unsigned mbSize>
struct Neuron {
  /// Each neuron in the network can be indexed by a one- or three-dimensional
  /// coordinate, and stores a weighted input, an activation and an error.
  /// x and y are coordinates in the 2D image plane, z indexes depth.
  unsigned index, x, y, z;
  float weightedInputs[mbSize];
  float activations[mbSize];
  float errors[mbSize];
  Neuron(unsigned index) : index(index) {}
  Neuron(unsigned x, unsigned y, unsigned z) : x(x), y(y), z(z) {}
};

template <unsigned mbSize>
struct Layer {
  virtual void initialiseDefaultWeights(std::default_random_engine&) = 0;
  virtual void feedForward(unsigned mb) = 0;
  virtual void calcBwdError(unsigned mb) = 0;
  virtual void backPropogate(unsigned mb) = 0;
  virtual void endBatch(unsigned numTrainingImages) = 0;
  virtual void setInputs(Layer<mbSize> *layer) = 0;
  virtual void setOutputs(Layer<mbSize> *layer) = 0;
  virtual float getBwdError(unsigned index, unsigned mb) = 0;
  virtual float getBwdError(unsigned x, unsigned y, unsigned z, unsigned mb) = 0;
  virtual Neuron<mbSize> &getNeuron(unsigned index) = 0;
  virtual Neuron<mbSize> &getNeuron(unsigned x, unsigned y, unsigned z) = 0;
  virtual unsigned getNumDims() = 0;
  virtual unsigned getDim(unsigned i) = 0;
  virtual unsigned size() = 0;
};

///===--------------------------------------------------------------------===///
/// Input layer.
///===--------------------------------------------------------------------===///
template <unsigned mbSize,
          unsigned imageX,
          unsigned imageY>
class InputLayer : public Layer<mbSize> {
  // x, y, z dimensions of input image.
  boost::multi_array<Neuron<mbSize>*, 3> neurons;

public:
  InputLayer() :
    neurons(boost::extents[imageX][imageY][1]) {
    for (unsigned x = 0; x < imageX; ++x) {
      for (unsigned y = 0; y < imageY; ++y) {
        neurons[x][y][0] = new Neuron<mbSize>(x, y, 0);
      }
    }
  }
  void setImage(Image &image, unsigned mb) {
    assert(image.size() == neurons.num_elements() && "invalid image size");
    for (unsigned i = 0; i < image.size(); ++i) {
      neurons[i % imageX][i / imageX][0]->activations[mb] = image[i];
    }
  }
  void initialiseDefaultWeights(std::default_random_engine&) override {
    UNREACHABLE();
  }
  virtual void calcBwdError(unsigned) override {
    UNREACHABLE();
  }
  void feedForward(unsigned) override {
    UNREACHABLE();
  }
  void backPropogate(unsigned) override {
    UNREACHABLE();
  }
  void endBatch(unsigned) override {
    UNREACHABLE();
  }
  void setInputs(Layer<mbSize>*) override {
    UNREACHABLE();
  }
  void setOutputs(Layer<mbSize>*) override {
    UNREACHABLE();
  }
  float getBwdError(unsigned, unsigned) override {
    UNREACHABLE();
  }
  float getBwdError(unsigned, unsigned, unsigned, unsigned) override {
    UNREACHABLE();
  }
  Neuron<mbSize> &getNeuron(unsigned i) override {
    assert(i < neurons.num_elements() && "Neuron index out of range.");
    return *neurons[i % imageX][i / imageX][0];
  }
  Neuron<mbSize> &getNeuron(unsigned x, unsigned y, unsigned z) override {
    assert(z == 0 && "Input image has depth 1");
    return *neurons[x][y][z];
  }
  unsigned getNumDims() override { return neurons.num_dimensions(); }
  unsigned getDim(unsigned i) override { return neurons.shape()[i]; }
  unsigned size() override { return neurons.num_elements(); }
};

///===--------------------------------------------------------------------===///
/// Fully-connected neuron.
///===--------------------------------------------------------------------===///
template <unsigned mbSize,
          float (*activationFn)(float) = nullptr,
          float (*activationFnDeriv)(float) = nullptr>
class FullyConnectedNeuron : public Neuron<mbSize> {
protected:
  float learningRate;
  float lambda;
  Layer<mbSize> *inputs;
  Layer<mbSize> *outputs;
  std::vector<float> weights;
  float bias;

public:
  FullyConnectedNeuron(unsigned index, float learningRate, float lambda) :
    Neuron<mbSize>(index),
    learningRate(learningRate), lambda(lambda),
    inputs(nullptr), outputs(nullptr) {}

  void initialiseDefaultWeights(std::default_random_engine &gen) {
    // Initialise all weights with random values from normal distribution with
    // mean 0 and stdandard deviation 1, divided by the square root of the
    // number of input connections.
    std::normal_distribution<float> distribution(0, 1.0f);
    for (unsigned i = 0; i < inputs->size(); ++i) {
      float weight = distribution(gen) / std::sqrt(inputs->size());
      weights.push_back(weight);
    }
    bias = distribution(gen);
  }

  void feedForward(unsigned mb) {
    float weightedInput = 0.0f;
    for (unsigned i = 0; i < inputs->size(); ++i) {
      weightedInput += inputs->getNeuron(i).activations[mb] * weights[i];
    }
    weightedInput += bias;
    this->weightedInputs[mb] = weightedInput;
    this->activations[mb] = activationFn(weightedInput);
  }

  void backPropogate(unsigned mb) {
    // Get the weight-error sum component from the next layer, then multiply by
    // the sigmoid derivative to get the error for this neuron.
    float error = outputs->getBwdError(this->index, mb);
    error *= activationFnDeriv(this->weightedInputs[mb]);
    this->errors[mb] = error;
  }

  void endBatch(unsigned numTrainingImages) {
    // For each weight.
    for (unsigned i = 0; i < inputs->size(); ++i) {
      float weightDelta = 0.0f;
      // For each batch element, average input activation x error (rate of
      // change of cost w.r.t. weight) and multiply by learning rate.
      // Note that FC layers can only be followed by FC layers.
      for (unsigned j = 0; j < mbSize; ++j) {
        weightDelta += inputs->getNeuron(i).activations[j] * this->errors[j];
      }
      weightDelta *= learningRate / mbSize;
      float reg = 1.0f - (learningRate * (lambda / numTrainingImages));
      weights[i] *= reg; // Regularisation term.
      weights[i] -= weightDelta;
    }
    // For each batch element, average the errors (error is equal to rate of
    // change of cost w.r.t. bias) and multiply by learning rate.
    float biasDelta = 0.0f;
    for (unsigned j = 0; j < mbSize; ++j) {
      biasDelta += this->errors[j];
    }
    biasDelta *= learningRate / mbSize;
    bias -= biasDelta;
  }

  void setInputs(Layer<mbSize> *inputs) { this->inputs = inputs; }
  void setOutputs(Layer<mbSize> *outputs) { this->outputs = outputs; }
  unsigned numWeights() { return weights.size(); }
  float getWeight(unsigned i) { return weights.at(i); }
};

///===--------------------------------------------------------------------===///
/// Fully-connected layer.
///===--------------------------------------------------------------------===///
template <unsigned mbSize,
          unsigned layerSize,
          unsigned prevSize,
          float (*activationFn)(float),
          float (*activationFnDeriv)(float)>
class FullyConnectedLayer : public Layer<mbSize> {
  using FullyConnectedNeuronTy =
      FullyConnectedNeuron<mbSize, activationFn, activationFnDeriv>;
  Layer<mbSize> *inputs;
  Layer<mbSize> *outputs;
  std::vector<FullyConnectedNeuronTy> neurons;
  boost::multi_array<float, 2> bwdErrors; // [mb][i]

public:
  FullyConnectedLayer(Params params) :
      bwdErrors(boost::extents[mbSize][prevSize]) {
    for (unsigned i = 0; i < layerSize; ++i) {
      auto n = FullyConnectedNeuronTy(i, params.learningRate, params.lambda);
      neurons.push_back(n);
    }
  }

  void setInputs(Layer<mbSize> *layer) override {
    inputs = layer;
    for (auto &neuron : neurons) {
      neuron.setInputs(layer);
    }
  }

  void setOutputs(Layer<mbSize> *layer) override {
    outputs = layer;
    for (auto &neuron : neurons) {
      neuron.setOutputs(layer);
    }
  }

  void initialiseDefaultWeights(std::default_random_engine &gen) override {
    for (auto &neuron : neurons) {
      neuron.initialiseDefaultWeights(gen);
    }
  }

  void feedForward(unsigned mb) override {
    for (auto &neuron : neurons) {
      neuron.feedForward(mb);
    }
  }

  /// Calculate the l+1 component of the error for each neuron in prev layer.
  void calcBwdError(unsigned mb) override {
    for (unsigned i = 0; i < inputs->size(); ++i) {
      float error = 0.0f;
      for (auto &neuron : neurons) {
        error += neuron.getWeight(i) * neuron.errors[mb];
      }
      bwdErrors[mb][i] = error;
    }
  }

  /// Update errors from next layer.
  void backPropogate(unsigned mb) override {
    for (auto &neuron : neurons) {
      neuron.backPropogate(mb);
    }
  }

  void endBatch(unsigned numTrainingImages) override {
    for (auto &neuron : neurons) {
      neuron.endBatch(numTrainingImages);
    }
  }

  float getBwdError(unsigned index, unsigned mb) override {
    return bwdErrors[mb][index];
  }

  float getBwdError(unsigned, unsigned, unsigned, unsigned) override {
    UNREACHABLE();
  }

  Neuron<mbSize> &getNeuron(unsigned index) override {
    return neurons.at(index);
  }

  Neuron<mbSize> &getNeuron(unsigned, unsigned, unsigned) override {
    UNREACHABLE();
  }

  unsigned getNumDims() override { return 1; }

  unsigned getDim(unsigned i) override {
    assert(i == 0 && "Layer is 1D");
    return neurons.size();
  }

  unsigned size() override { return neurons.size(); }
};

///===--------------------------------------------------------------------===///
/// Softmax neuron.
///===--------------------------------------------------------------------===///
template <unsigned mbSize,
          float (*costFn)(float, float),
          float (*costDelta)(float, float, float)>
class SoftMaxNeuron : public FullyConnectedNeuron<mbSize> {

public:
  SoftMaxNeuron(unsigned index, float learningRate, float lambda) :
      FullyConnectedNeuron<mbSize>(index, learningRate, lambda) {}

  void feedForward(unsigned mb) {
    // Only calculate weighted inputs.
    float weightedInput = 0.0f;
    for (unsigned i = 0; i < this->inputs->size(); ++i) {
      weightedInput +=
        this->inputs->getNeuron(i).activations[mb] * this->weights[i];
    }
    weightedInput += this->bias;
    this->weightedInputs[mb] = weightedInput;
  }

  void backPropogate(unsigned) { UNREACHABLE(); }

  void computeOutputError(uint8_t label, unsigned mb) {
    float y = label == this->index ? 1.0f : 0.0f;
    float error = costDelta(this->weightedInputs[mb], this->activations[mb], y);
    this->errors[mb] = error;
  }

  float computeOutputCost(uint8_t label, unsigned mb) {
    return costFn(this->activations[mb], label);
  }

  float sumSquaredWeights() {
    float result = 0.0f;
    for (auto weight : this->weights) {
      result += std::pow(weight, 2.0f);
    }
    return result;
  }
};

///===--------------------------------------------------------------------===///
/// Softmax layer.
///===--------------------------------------------------------------------===///
template <unsigned mbSize,
          unsigned layerSize,
          unsigned prevSize,
          float (*costFn)(float, float),
          float (*costDelta)(float, float, float)>
class SoftMaxLayer : public Layer<mbSize> {
  using SoftMaxNeuronTy = SoftMaxNeuron<mbSize, costFn, costDelta>;
  Layer<mbSize> *inputs;
  Layer<mbSize> *outputs;
  std::vector<SoftMaxNeuronTy> neurons;
  boost::multi_array<float, 2> bwdErrors; // [mb][i]

public:
  SoftMaxLayer(float learningRate, float lambda) :
      bwdErrors(boost::extents[mbSize][prevSize]) {
    for (unsigned i = 0; i < layerSize; ++i) {
      this->neurons.push_back(SoftMaxNeuronTy(i, learningRate, lambda));
    }
  }

  void setInputs(Layer<mbSize> *layer) override {
    inputs = layer;
    for (auto &neuron : neurons) {
      neuron.setInputs(layer);
    }
  }

  void setOutputs(Layer<mbSize>*) override {
    UNREACHABLE();
  }

  void initialiseDefaultWeights(std::default_random_engine &gen) override {
    for (auto &neuron : neurons) {
      neuron.initialiseDefaultWeights(gen);
    }
  }

  void feedForward(unsigned mb) override {
    // Calculate weighted inputs for each neuron.
    // Sum the exponential values of the weighted inputs across neurons.
    float sum = 0.0f;
    for (auto &neuron : neurons) {
      neuron.feedForward(mb);
      sum += std::exp(neuron.weightedInputs[mb]);
    }
    // Calculate each of the neuron's activations.
    for (auto &neuron : neurons) {
      neuron.activations[mb] = std::exp(neuron.weightedInputs[mb]) / sum;
    }
  }

  /// Calculate the l+1 component of the error for each neuron in prev layer.
  void calcBwdError(unsigned mb) override {
    for (unsigned i = 0; i < inputs->size(); ++i) {
      float error = 0.0f;
      for (auto &neuron : neurons) {
        error += neuron.getWeight(i) * neuron.errors[mb];
      }
      bwdErrors[mb][i] = error;
    }
  }

  /// Update errors from next layer.
  void backPropogate(unsigned mb) override {
    for (auto &neuron : neurons) {
      neuron.backPropogate(mb);
    }
  }

  void endBatch(unsigned numTrainingImages) override {
    for (auto &neuron : neurons) {
      neuron.endBatch(numTrainingImages);
    }
  }

  /// Determine the index of the highest output activation.
  unsigned readOutput(unsigned mb) {
    unsigned result = 0;
    float max = std::numeric_limits<float>::min();
    for (unsigned i = 0; i < neurons.size(); ++i) {
      float output = neurons[i].activations[mb];
      if (output > max) {
        result = i;
        max = output;
      }
    }
    return result;
  }

  void computeOutputError(uint8_t label, unsigned mb) {
    for (auto &neuron : neurons) {
      neuron.computeOutputError(label, mb);
    }
  }

  float computeOutputCost(uint8_t label, unsigned mb) {
    float outputCost = 0.0f;
    for (auto &neuron : neurons) {
      neuron.computeOutputCost(label, mb);
    }
    return outputCost;
  }

  float sumSquaredWeights() {
    float result = 0.0f;
    for (auto &neuron : neurons) {
      result += neuron.sumSquaredWeights();
    }
    return result;
  }

  float getBwdError(unsigned index, unsigned mb) override {
    return bwdErrors[mb][index];
  }

  float getBwdError(unsigned, unsigned, unsigned, unsigned) override {
    UNREACHABLE();
  }

  Neuron<mbSize> &getNeuron(unsigned index) override {
    return neurons.at(index);
  }

  Neuron<mbSize> &getNeuron(unsigned, unsigned, unsigned) override {
    UNREACHABLE();
  }

  unsigned getNumDims() override { return 1; }

  unsigned getDim(unsigned i) override {
    assert(i == 0 && "Layer is 1D");
    return neurons.size();
  }

  unsigned size() override { return neurons.size(); }
};

///===--------------------------------------------------------------------===///
/// Convolutional neuron.
///===--------------------------------------------------------------------===///
template <unsigned mbSize,
          float (*activationFn)(float),
          float (*activationFnDeriv)(float)>
class ConvNeuron : public Neuron<mbSize> {
  Layer<mbSize> *inputs;
  Layer<mbSize> *outputs;
  unsigned dimX;
  unsigned dimY;

public:
  ConvNeuron(unsigned x, unsigned y, unsigned z, unsigned dimX, unsigned dimY) :
      Neuron<mbSize>(x, y, z), dimX(dimX), dimY(dimY) {}

  void feedForward(boost::multi_array_ref<float, 4> &weights,
                   boost::multi_array_ref<float, 1> &bias,
                   unsigned mb) {
    // Convolve using each weight.
    // (z is the index of the feature map.)
    float weightedInput = 0.0f;
    for (unsigned a = 0; a < weights.shape()[1]; ++a) {
      for (unsigned b = 0; b < weights.shape()[2]; ++b) {
        for (unsigned c = 0; c < weights.shape()[3]; ++c) {
          unsigned inputX = this->x + a;
          unsigned inputY = this->y + b;
          float input = inputs->getNeuron(inputX, inputY, c).activations[mb];
          weightedInput += input * weights[this->z][a][b][c];
        }
      }
    }
    // Add bias and apply non linerarity.
    weightedInput += bias[this->z];
    this->weightedInputs[mb] = weightedInput;
    this->activations[mb] = activationFn(weightedInput);
  }

  void backPropogate(unsigned mb) {
    // If next layer is 1D, map the x, y, z coordinates onto it.
    unsigned index =
      getIndex(this->x, this->y, this->z, this->dimX, this->dimY);
    float error = outputs->getNumDims() == 1
                    ? outputs->getBwdError(index, mb)
                    : outputs->getBwdError(this->x, this->y, this->z, mb);
    error *= activationFnDeriv(this->weightedInputs[mb]);
    this->errors[mb] = error;
  }

  void setInputs(Layer<mbSize> *inputs) { this->inputs = inputs; }
  void setOutputs(Layer<mbSize> *outputs) { this->outputs = outputs; }
};

///===--------------------------------------------------------------------===///
/// Convolutional layer
///
/// kernelX is num cols
/// kernelY is num rows
/// neuron(x, y) is row y, col x
/// weights(a, b) is row b, col a
///===--------------------------------------------------------------------===///
template <unsigned mbSize,
          unsigned kernelX,
          unsigned kernelY,
          unsigned kernelZ,
          unsigned inputX,
          unsigned inputY,
          unsigned inputZ,
          unsigned numFMs,
          float (*activationFn)(float),
          float (*activationFnDeriv)(float)>
class ConvLayer : public Layer<mbSize> {
  using ConvNeuronTy = ConvNeuron<mbSize, activationFn, activationFnDeriv>;
  float learningRate;
  float lambda;
  Layer<mbSize> *inputs;
  Layer<mbSize> *outputs;
  boost::multi_array<float, 1> bias;            // [fm]
  boost::multi_array<float, 4> weights;         // [fm][x][y][z]
  boost::multi_array<ConvNeuronTy*, 3> neurons; // [fm][x][y]
  boost::multi_array<float, 4> bwdErrors;       // [mb][x][y][z]

public:
  ConvLayer(Params params) :
      learningRate(params.learningRate), lambda(params.lambda),
      inputs(nullptr), outputs(nullptr),
      bias(boost::extents[numFMs]),
      weights(boost::extents[numFMs][kernelX][kernelY][kernelZ]),
      neurons(boost::extents[numFMs][inputX-kernelX+1][inputY-kernelY+1]),
      bwdErrors(boost::extents[mbSize][inputX][inputY][inputZ]) {
    static_assert(inputZ == kernelZ, "Kernel depth should match input depth");
    unsigned dimX = neurons.shape()[1];
    unsigned dimY = neurons.shape()[2];
    for (unsigned fm = 0; fm < numFMs; ++fm) {
      for (unsigned x = 0; x < dimX; ++x) {
        for (unsigned y = 0; y < dimY; ++y) {
          neurons[fm][x][y] = new ConvNeuronTy(x, y, fm, dimX, dimY);
        }
      }
    }
  }

  void initialiseDefaultWeights(std::default_random_engine &gen) override {
    // Initialise weights random distribution of mean 0 and standard deviation
    // 1, then scale it by 1/sqrt(number of inputs).
    std::normal_distribution<float> distribution(0, 1.0f);
    float scale =
        std::sqrt(weights.shape()[1] * weights.shape()[2] * weights.shape()[3]);
    for (unsigned fm = 0; fm < weights.shape()[0]; ++fm) {
      for (unsigned a = 0; a < weights.shape()[1]; ++a) {
        for (unsigned b = 0; b < weights.shape()[2]; ++b) {
          for (unsigned c = 0; c < weights.shape()[3]; ++c) {
            weights[fm][a][b][c] = distribution(gen) / scale;
          }
        }
      }
      bias[fm] = distribution(gen);
    }
  }

  void feedForward(unsigned mb) override {
    for (unsigned fm = 0; fm < neurons.shape()[0]; ++fm) {
      for (unsigned x = 0; x < neurons.shape()[1]; ++x) {
        for (unsigned y = 0; y < neurons.shape()[2]; ++y) {
          neurons[fm][x][y]->feedForward(weights, bias, mb);
        }
      }
    }
  }

  void calcBwdError(unsigned mb) override {
    // Calculate the l+1 component of the error for each neuron in prev layer.
    for (unsigned x = 0; x < inputX; ++x) {
      for (unsigned y = 0; y < inputY; ++y) {
        for (unsigned z = 0; z < inputZ; ++z) {
          // Sum over all feature maps.
          float error = 0.0f;
          for (unsigned fm = 0; fm < numFMs; ++fm) {
            for (unsigned a = 0; a < weights.shape()[1]; ++a) {
              for (unsigned b = 0; b < weights.shape()[2]; ++b) {
                if (a <= x && b <= y &&
                    x - a < neurons.shape()[1] &&
                    y - b < neurons.shape()[2]) {
                  float ne = neurons[fm][x - a][y - b]->errors[mb];
                  error += weights[fm][a][b][z] * ne;
                }
              }
            }
          }
          bwdErrors[mb][x][y][z] = error;
        }
      }
    }
  }

  void backPropogate(unsigned mb) override {
    // Update errors from next layer.
    for (unsigned fm = 0; fm < neurons.shape()[0]; ++fm) {
      for (unsigned x = 0; x < neurons.shape()[1]; ++x) {
        for (unsigned y = 0; y < neurons.shape()[2]; ++y) {
          neurons[fm][x][y]->backPropogate(mb);
        }
      }
    }
  }

  void endBatch(unsigned numTrainingImages) override {
    // For each feature map.
    for (unsigned fm = 0; fm < numFMs; ++fm) {
      // For each weight, calculate the delta and update the weight.
      for (unsigned a = 0; a < weights.shape()[1]; ++a) {
        for (unsigned b = 0; b < weights.shape()[2]; ++b) {
          for (unsigned c = 0; c < weights.shape()[3]; ++c) {
            float weightDelta = 0.0f;
            // For each item of the minibatch.
            for (unsigned mb = 0; mb < mbSize; ++mb) {
              // For each neuron.
              for (unsigned x = 0; x < neurons.shape()[1]; ++x) {
                for (unsigned y = 0; y < neurons.shape()[2]; ++y) {
                  float i = inputs->getNeuron(x + a, y + b, c).activations[mb];
                  weightDelta += i * neurons[fm][x][y]->errors[mb];
                }
              }
            }
            weightDelta *= learningRate / mbSize;
            float reg = 1.0f - (learningRate * (lambda / numTrainingImages));
            weights[fm][a][b][c] *= reg; // Regularisation term.
            weights[fm][a][b][c] -= weightDelta;
          }
        }
      }
      // Calculate bias delta and update it.
      float biasDelta = 0.0f;
      // For each item of the minibatch.
      for (unsigned mb = 0; mb < mbSize; ++mb) {
        // For each neuron.
        for (unsigned x = 0; x < neurons.shape()[1]; ++x) {
          for (unsigned y = 0; y < neurons.shape()[2]; ++y) {
            biasDelta += neurons[fm][x][y]->errors[mb];
          }
        }
      }
      biasDelta *= learningRate / mbSize;
      bias[fm] -= biasDelta;
    }
  }

  float getBwdError(unsigned x, unsigned y, unsigned z, unsigned mb) override {
    return bwdErrors[mb][x][y][z];
  }

  void setInputs(Layer<mbSize> *layer) override {
    assert(layer->size() == inputX * inputY * inputZ &&
           "Invalid input layer size");
    inputs = layer;
    std::for_each(neurons.data(), neurons.data() + neurons.num_elements(),
                  [layer](ConvNeuronTy *n){ n->setInputs(layer); });
  }

  void setOutputs(Layer<mbSize> *layer) override {
    outputs = layer;
    std::for_each(neurons.data(), neurons.data() + neurons.num_elements(),
                  [layer](ConvNeuronTy *n){ n->setOutputs(layer); });
  }

  float getBwdError(unsigned, unsigned) override {
    UNREACHABLE(); // No FC layers preceed conv layers.
  }

  Neuron<mbSize> &getNeuron(unsigned index) override {
    // Map a 1D index onto the 3D neurons (for Conv <- FC connections).
    unsigned dimX = neurons.shape()[1];
    unsigned dimY = neurons.shape()[2];
    unsigned x = getX(index, dimX);
    unsigned y = getY(index, dimX, dimY);
    unsigned z = getZ(index, dimX, dimY);
    return *neurons[z][x][y];
  }

  Neuron<mbSize> &getNeuron(unsigned x, unsigned y, unsigned z) override {
    // Feature maps is inner dimension but corresponds to z.
    return *neurons[z][x][y];
  }

  unsigned getDim(unsigned i) override {
    assert(i <= 2 && "Dimension out of range.");
    // Feature maps is inner dimension but corresponds to z.
    return i == 2 ? neurons.shape()[0] : neurons.shape()[i + 1];
  }

  unsigned getNumDims() override { return neurons.num_dimensions(); }
  unsigned size() override { return neurons.num_elements(); }
};

///===--------------------------------------------------------------------===///
/// Max pool layer
///===--------------------------------------------------------------------===///
template <unsigned mbSize,
          unsigned poolX,
          unsigned poolY,
          unsigned inputX,
          unsigned inputY,
          unsigned inputZ>
class MaxPoolLayer : public Layer<mbSize> {
  Layer<mbSize> *inputs;
  Layer<mbSize> *outputs;
  boost::multi_array<Neuron<mbSize>*, 3> neurons; // [x][y][z]

public:
  MaxPoolLayer() :
      inputs(nullptr), outputs(nullptr),
      neurons(boost::extents[inputX / poolX][inputY / poolY][inputZ]) {
    static_assert(inputX % poolX == 0, "Dimension x mismatch with pooling");
    static_assert(inputY % poolY == 0, "Dimension y mismatch with pooling");
    for (unsigned x = 0; x < neurons.shape()[0]; ++x) {
      for (unsigned y = 0; y < neurons.shape()[1]; ++y) {
        for (unsigned z = 0; z < neurons.shape()[2]; ++z) {
          neurons[x][y][z] = new Neuron<mbSize>(x, y, z);
        }
      }
    }
  }

  void initialiseDefaultWeights(std::default_random_engine&) override {
    /* Skip */
  }

  void feedForward(unsigned mb) override {
    // For each neuron in this layer.
    for (unsigned x = 0; x < neurons.shape()[0]; ++x) {
      for (unsigned y = 0; y < neurons.shape()[1]; ++y) {
        for (unsigned z = 0; z < neurons.shape()[2]; ++z) {
          // Take maximum activation over pool area.
          float weightedInput = std::numeric_limits<float>::min();
          for (unsigned a = 0; a < poolX; ++a) {
            for (unsigned b = 0; b < poolY; ++b) {
              unsigned nX = (x * poolX) + a;
              unsigned nY = (y * poolY) + b;
              float input = inputs->getNeuron(nX, nY, z).activations[mb];
              float max = std::max(weightedInput, input);
              neurons[x][y][z]->activations[mb] = max;
            }
          }
        }
      }
    }
  }

  void calcBwdError(unsigned) override { /* Skip */ }
  void backPropogate(unsigned) override { /* Skip */ }

  float getBwdError(unsigned x, unsigned y, unsigned z, unsigned mb) override {
    // Forward the backwards error component from the next layer.
    unsigned nX = x / poolX;
    unsigned nY = y / poolY;
    unsigned nZ = z;
    unsigned dimX = neurons.shape()[0];
    unsigned dimY = neurons.shape()[1];
    // If next layer is 1D, map the x, y, z coordinates onto it.
    unsigned index = getIndex(nX, nY, nZ, dimX, dimY);
    return outputs->getNumDims() == 1
             ? outputs->getBwdError(index, mb)
             : outputs->getBwdError(nX, nY, nZ, mb);
  }

  void endBatch(unsigned) override { /* Skip */ }

  float getBwdError(unsigned, unsigned) override {
    UNREACHABLE(); // No FC layers preceed max-pooling layers.
  }

  void setInputs(Layer<mbSize> *layer) override {
    assert(layer->size() == poolX * poolY * neurons.num_elements() &&
           "invalid input layer size");
    inputs = layer;
  }

  void setOutputs(Layer<mbSize> *layer) override { outputs = layer; }

  Neuron<mbSize> &getNeuron(unsigned index) override {
    // Map a 1D index onto the 3D neurons (for Conv <- FC connections).
    unsigned dimX = neurons.shape()[0];
    unsigned dimY = neurons.shape()[1];
    unsigned x = getX(index, dimX);
    unsigned y = getY(index, dimX, dimY);
    unsigned z = getZ(index, dimX, dimY);
    return *neurons[x][y][z];
  }

  Neuron<mbSize> &getNeuron(unsigned x, unsigned y, unsigned z) override {
    return *neurons[x][y][z];
  }

  unsigned getNumDims() override { return neurons.num_dimensions(); }
  unsigned getDim(unsigned i) override { return neurons.shape()[i]; }
  unsigned size() override { return neurons.num_elements(); }
};

///===--------------------------------------------------------------------===///
/// The network.
///===--------------------------------------------------------------------===///
template <unsigned mbSize,
          unsigned inputX,
          unsigned inputY,
          unsigned softMaxSize,
          unsigned lastLayerSize,
          float (*costFn)(float, float),
          float (*costDelta)(float, float, float)>
class Network {
  using SoftMaxLayerTy = SoftMaxLayer<mbSize, softMaxSize, lastLayerSize,
                                      costFn, costDelta>;
  using LayerTy = Layer<mbSize>;
  Params params;
  InputLayer<mbSize, inputX, inputY> inputLayer;
  SoftMaxLayerTy *softMaxLayer;
  std::vector<LayerTy*> layers;
  std::default_random_engine generator;

public:
  Network(Params params, std::vector<LayerTy*> layers_) :
      params(params), layers(layers_), generator(params.seed) {
    softMaxLayer = new SoftMaxLayerTy(params.learningRate, params.lambda);
    layers.push_back(softMaxLayer);
    // Set neuron inputs.
    layers[0]->setInputs(&inputLayer);
    layers[0]->initialiseDefaultWeights(generator);
    for (unsigned i = 1; i < layers.size(); ++i) {
      layers[i]->setInputs(layers[i - 1]);
      layers[i]->initialiseDefaultWeights(generator);
    }
    // Set neuron outputs.
    for (unsigned i = 0; i < layers.size() - 1; ++i) {
      layers[i]->setOutputs(layers[i + 1]);
    }
  }

  /// The forward pass.
  void feedForward(unsigned mb) {
    for (auto layer : layers) {
      layer->feedForward(mb);
    }
  }

  /// The backward pass.
  void backPropogate(Image &image, uint8_t label, unsigned mb) {
    // Set input.
    inputLayer.setImage(image, mb);
    // Feed forward.
    feedForward(mb);
    // Compute output error in last layer.
    softMaxLayer->computeOutputError(label, mb);
    softMaxLayer->calcBwdError(mb);
    // Backpropagate the error and calculate component for next layer.
    for (int i = layers.size() - 2; i > 0; --i) {
      layers[i]->backPropogate(mb);
      layers[i]->calcBwdError(mb);
    }
    layers[0]->backPropogate(mb);
  }

  void updateMiniBatch(std::vector<Image>::iterator trainingImagesIt,
                       std::vector<uint8_t>::iterator trainingLabelsIt,
                       unsigned numTrainingImages) {
    // For each training image and label, back propogate.
    // Parallelise over the elements of the minibatch to improve performance.
    tbb::parallel_for(size_t(0), size_t(mbSize), [=](size_t i) {
      backPropogate(*(trainingImagesIt + i), *(trainingLabelsIt + i), i);
    });
    // Gradient descent: for every neuron, compute the new weights and biases.
    for (int i = layers.size() - 1; i >= 0; --i) {
      layers[i]->endBatch(numTrainingImages);
    }
  }

  float imageCost(Image &image, uint8_t label,
                  unsigned numImages, float regularisation, unsigned mb) {
    inputLayer.setImage(image, mb);
    feedForward(mb);
    float cost = softMaxLayer->computeOutputCost(label, mb) / numImages;
    cost += regularisation;
    return cost;
  }

  /// Calculate the total cost for a dataset.
  /// Parallelise over the test images (up to the minibatch size).
  float evaluateTotalCost(std::vector<Image> &testImages,
                          std::vector<uint8_t> &testLabels) {
    float regularisation = 0.5f * (params.lambda / testImages.size())
                            * softMaxLayer->sumSquaredWeights();
    float cost = 0.0f;
    for (unsigned i = 0, end = testImages.size(); i < end; i += mbSize) {
      auto mbStart = std::chrono::high_resolution_clock::now();
      // Parallel reduce over the minibatch.
      cost +=
        tbb::parallel_reduce(
          tbb::blocked_range<size_t>(0, mbSize), 0.0f,
          [&](const tbb::blocked_range<size_t> &r, float total) {
            for (size_t mb = r.begin(); mb < r.end(); ++mb) {
              total += imageCost(*(testImages.begin() + i + mb),
                                 *(testLabels.begin() + i + mb),
                                 testImages.size(), regularisation, mb);
            }
            return total;
          }, std::plus<float>());
      auto mbEnd = std::chrono::high_resolution_clock::now();
      auto ms =
        std::chrono::duration_cast<std::chrono::milliseconds>(mbEnd-mbStart);
      float imagesPerSec = (float(mbSize) / ms.count()) * 1000.0f;
      std::cout << "\rEvaluate cost " << i << " / " << end
                << " (" << imagesPerSec << " imgs/s)";
    }
    return cost;
  }

  bool testImage(Image &image, uint8_t label, unsigned mb) {
    inputLayer.setImage(image, mb);
    feedForward(mb);
    return softMaxLayer->readOutput(mb) == label;
  }

  /// Evaluate the test set and return the number of correct classifications.
  /// Parallelise over the test images (up to the minibatch size).
  unsigned evaluateAccuracy(std::vector<Image> &testImages,
                            std::vector<uint8_t> &testLabels) {
    unsigned result = 0;
    for (unsigned i = 0, end = testImages.size(); i < end; i += mbSize) {
      auto mbStart = std::chrono::high_resolution_clock::now();
      // Parallel reduce over the minibatch.
      result +=
        tbb::parallel_reduce(
          tbb::blocked_range<size_t>(0, mbSize), 0,
          [&](const tbb::blocked_range<size_t> &r, unsigned total) {
            for (size_t mb = r.begin(); mb < r.end(); ++mb) {
              total += testImage(*(testImages.begin() + i + mb),
                                 *(testLabels.begin() + i + mb), mb);
            }
            return total;
          }, std::plus<unsigned>());
      auto mbEnd = std::chrono::high_resolution_clock::now();
      auto ms =
        std::chrono::duration_cast<std::chrono::milliseconds>(mbEnd-mbStart);
      float imagesPerSec = (float(mbSize) / ms.count()) * 1000.0f;
      std::cout << "\rEvaluate accuracy " << i << " / " << end
                << " (" << imagesPerSec << " imgs/s)";
    }
    return result;
  }

  void SGD(Data &data) {
    // For each epoch.
    for (unsigned epoch = 0; epoch < params.numEpochs; ++epoch) {
      auto epochStart = std::chrono::high_resolution_clock::now();
      // Identically randomly shuffle the training images and labels.
      std::uniform_int_distribution<unsigned> distribution;
      unsigned seed = distribution.operator ()(generator);
      std::shuffle(data.getTrainingLabels().begin(),
                   data.getTrainingLabels().end(),
                   std::default_random_engine(seed));
      std::shuffle(data.getTrainingImages().begin(),
                   data.getTrainingImages().end(),
                   std::default_random_engine(seed));
      // For each mini batch.
      unsigned numTrainingImages = data.getTrainingImages().size();
      for (unsigned i = 0; i < numTrainingImages; i += mbSize) {
        auto mbStart = std::chrono::high_resolution_clock::now();
        updateMiniBatch(data.getTrainingImages().begin() + i,
                        data.getTrainingLabels().begin() + i,
                        mbSize);
        auto mbEnd = std::chrono::high_resolution_clock::now();
        auto ms =
          std::chrono::duration_cast<std::chrono::milliseconds>(mbEnd-mbStart);
        float imagesPerSec = (float(mbSize) / ms.count()) * 1000.0f;
        std::cout << "\rMinibatch " << i << " / " << numTrainingImages
                  << " (" << imagesPerSec << " imgs/s)";
        if (i % params.monitorInterval == 0) {
          std::cout << '\r' << std::string(100, ' ');
          // Evaluate the test set.
          if (params.monitorEvaluationAccuracy) {
            unsigned result = evaluateAccuracy(data.getValidationImages(),
                                               data.getValidationLabels());
            std::cout << '\r' << std::string(100, ' ');
            std::cout << "\rAccuracy on evaluation data: "
                      << result << " / " << data.getValidationImages().size()
                      << '\n';
          }
          if (params.monitorEvaluationCost) {
            float cost = evaluateTotalCost(data.getValidationImages(),
                                           data.getValidationLabels());
            std::cout << '\r' << std::string(100, ' ');
            std::cout << "\rCost on evaluation data: " << cost << "\n";
          }
          if (params.monitorTrainingAccuracy) {
            unsigned result = evaluateAccuracy(data.getTestImages(),
                                               data.getTestLabels());
            std::cout << '\r' << std::string(100, ' ');
            std::cout << "\rAccuracy on test data: "
                      << result << " / " << data.getTestImages().size()
                      << '\n';
          }
          if (params.monitorTrainingCost) {
            float cost = evaluateTotalCost(data.getTestImages(),
                                           data.getTestLabels());
            std::cout << '\r' << std::string(100, ' ');
            std::cout << "\rCost on test data: " << cost << "\n";
          }
        }
      }
      std::cout << '\n';
      // Display end of epoch and time.
      auto epochEnd = std::chrono::high_resolution_clock::now();
      auto s =
        std::chrono::duration_cast<std::chrono::seconds>(epochEnd-epochStart);
      std::cout << "Epoch " << epoch << " complete in " << s.count() << " s.\n";
    }
  }
};

#endif


// Flatten.cpp : Defines the entry point for the console application.
//
#ifdef WIN32
#define NOMINMAX
#include <windows.h>
#endif

#if defined (__APPLE__) || defined (OSX)
	#include <OpenGL/gl.h>
	#include <GLUT/glut.h>
#else
	#include <GL/gl.h>
	#include <GL/glut.h>
#endif

#include "GA/c3ga.h"
#include "GA/c3ga_util.h"
#include "GA/gl_util.h"

#include "primitivedraw.h"
#include "gahelper.h"
#include "Laplacian.h"

#include <memory>

#include <vector>
#include <map>
#include "numerics.h"
#include "HalfEdge/Mesh.h"
#include "GARotorEstimator.h"

#include <Eigen/Core>
#include <Eigen/Dense>
#include <Eigen/Sparse>
#include <Eigen/Geometry>

const char *WINDOW_TITLE = "Interactive 3D Shape Deformation using Conformal Geometric Algebra";

// GLUT state information
int g_viewportWidth = 800;
int g_viewportHeight = 600;

void display();
void reshape(GLint width, GLint height);
void MouseButton(int button, int state, int x, int y);
void MouseMotion(int x, int y);
void KeyboardUpFunc(unsigned char key, int x, int y);
void SpecialFunc(int key, int x, int y);
void SpecialUpFunc(int key, int x, int y);
void Idle();
void DestroyWindow();

using namespace c3ga;
using namespace std;
using namespace numerics;

class Camera
{
public:
	float		pos[3];
	float		fw[3];
	float		up[3];
	float		translateVel;
	float		rotateVel;

	Camera()
	{
		float		_pos[] = { 0, 0, -2};
		float		_fw[] = { 0, 0, 1 };
		float		_up[] = { 0, 1, 0 };

		translateVel = 0.005;
		rotateVel = 0.005;
		memcpy(pos, _pos, sizeof(float)*3);
		memcpy(fw, _fw, sizeof(float)*3);
		memcpy(up, _up, sizeof(float)*3);
	}

	void glLookAt()
	{
		gluLookAt( pos[0], pos[1], pos[2], fw[0],  fw[1],  fw[2], up[0],  up[1],  up[2] );
	}
};

class VertexBuffer
{
public:
	std::vector<Eigen::Vector3d> deformedPositions; //deformed mesh positions
	std::map<int,Eigen::Vector3d> constrainedPositions; //positional constraints
	std::vector<Eigen::Vector3d> laplacianCoordinates; //laplacian Coordinates
	std::vector<Eigen::Vector3d> normals; //for rendering (lighting)
	std::vector<Eigen::Vector3d> normalsOrig; //original normals
	std::vector<Eigen::Quaterniond> rotors;
	int size;

	VertexBuffer() : size(0)
	{
	}

	void resize(int size)
	{
		this->size = size;
		deformedPositions.resize(size);
		laplacianCoordinates.resize(size);
		normals.resize(size);
		normalsOrig.resize(size);
		rotors.resize(size);
	}
	int get_size() { return size; }
};

class IndexBuffer{
public:
	std::vector<int> faces;
	int size;

	IndexBuffer() : size(0)
	{
	}

	void resize(int size)
	{
		this->size = size;
		faces.resize(size);
	}
	int get_size() { return size; }

};

class Handle
{
public:
	rotor R;
	translator T;
	translator Tcenter;
	dualSphere dS;
	std::set<int> constraints;

	Handle() {
		T = _rotor(1.0);
		R = _translator(1.0);
		Tcenter = _translator(1.0);
	}
	Handle(dualSphere dS)
	{
		normalizedPoint x = DualSphereCenter(dS);
		Tcenter = exp( -0.5*_vectorE3GA(x)*ni );
		T = _translator(1.0);
		R = _rotor(1.0);
		this->dS = dS;
	}

	TRversor GetTRVersor()
	{
		return _TRversor(T * _TRversor(Tcenter * R * inverse(Tcenter)));
	}
};

Camera g_camera;
Mesh mesh;
vectorE3GA g_prevMousePos;
bool g_rotateModel = false;
bool g_rotateModelOutOfPlane = false;
rotor g_modelRotor = _rotor(1.0);
bool g_rotateKeyRotors = false;
bool g_translateKeyRotors = false;
bool g_computeBasis = false;
float g_dragDistance = -1.0f;
int g_dragObject;
std::shared_ptr<SparseMatrix> A;
Eigen::SparseLU<Eigen::SparseMatrix<double>, Eigen::COLAMDOrdering<int>> solver;
int systemType = LaplaceBeltrami; //MeanValue; //LaplaceBeltrami
bool g_showSpheres = true;
bool g_showWires = false;
bool g_iterateManyTimes = false;
Eigen::MatrixXd b3;
Eigen::MatrixXd xyz;

VertexBuffer vertexDescriptors;
IndexBuffer triangles;
std::vector<Handle> handles;

Eigen::Affine3d MotorToMatrix(const TRversor &R) {
	TRversor Ri = inverse(R);

	// compute images of basis vectors:
	c3ga::flatPoint imageOfE1NI = _flatPoint(R * c3ga::e1ni * Ri);
	c3ga::flatPoint imageOfE2NI = _flatPoint(R * c3ga::e2ni * Ri);
	c3ga::flatPoint imageOfE3NI = _flatPoint(R * c3ga::e3ni * Ri);
	c3ga::flatPoint imageOfNONI = _flatPoint(R * c3ga::noni * Ri);

	// create matrix representation:
	Eigen::Affine3d M;
	M(0, 0) = imageOfE1NI.m_c[0];
	M(1, 0) = imageOfE1NI.m_c[1];
	M(2, 0) = imageOfE1NI.m_c[2];
	M(3, 0) = imageOfE1NI.m_c[3];
	M(0, 1) = imageOfE2NI.m_c[0];
	M(1, 1) = imageOfE2NI.m_c[1];
	M(2, 1) = imageOfE2NI.m_c[2];
	M(3, 1) = imageOfE2NI.m_c[3];
	M(0, 2) = imageOfE3NI.m_c[0];
	M(1, 2) = imageOfE3NI.m_c[1];
	M(2, 2) = imageOfE3NI.m_c[2];
	M(3, 2) = imageOfE3NI.m_c[3];
	M(0, 3) = imageOfNONI.m_c[0];
	M(1, 3) = imageOfNONI.m_c[1];
	M(2, 3) = imageOfNONI.m_c[2];
	M(3, 3) = imageOfNONI.m_c[3];
	return M;
}

void ComputeLaplacianCoordinates(std::shared_ptr<SparseMatrix> A, Mesh* mesh, std::vector<Eigen::Vector3d>& laplacianCoordinates)
{
	std::fill(laplacianCoordinates.begin(), laplacianCoordinates.end(), Eigen::Vector3d(0,0,0));

	auto numRows = A->numRows();

	for( int i = 0; i < numRows ; ++i)
	{
		SparseMatrix::RowIterator aIter = A->iterator(i);
		for( ; !aIter.end() ; ++aIter )
		{
			auto j = aIter.columnIndex();
			laplacianCoordinates[i] += mesh->vertexAt(j).p * aIter.value();
		}
	}
}

bool is_constrained(std::vector<Handle>& handles, int vertex)
{
	for(Handle& handle : handles) {
		if(handle.constraints.find(vertex) != handle.constraints.end())
			return true;
	}
	return false;
}

void PreFactor(std::shared_ptr<SparseMatrix> A, std::vector<Handle>& handles)
{
	Eigen::SparseMatrix<double> Lc = Eigen::SparseMatrix<double>(A->numRows(), A->numColumns());

	auto numRows = A->numRows();
	for( int i = 0; i < numRows ; ++i)
	{
		if(!is_constrained(handles, i))
		{
			SparseMatrix::RowIterator aIter = A->iterator(i);
			for( ; !aIter.end() ; ++aIter )
			{
				auto j = aIter.columnIndex();
				Lc.insert(i, j) = (*A)(i,j);
			}
		}
		else
		{
			Lc.insert(i, i) = 1.0;
		}
	}
	Lc.makeCompressed();
	solver.compute(Lc);
	if(solver.info() != Eigen::Success) {
		// TODO: error handling
	}
}

int main(int argc, char* argv[])
{
	mesh.readOBJ("cactus2.obj"); //armadillo-5k-smooth.obj female.obj david.obj rabbit.obj tyra.obj horse.obj cylinder.obj bar.obj planewithpeaks.obj dragon.obj catHead.obj  cactus.obj  bunny.obj  764_hand-olivier-10kf.obj armadillo.obj

	mesh.CenterAndNormalize();

	mesh.computeNormals();

	// GLUT Window Initialization:
	glutInit (&argc, argv);
	glutInitWindowSize(g_viewportWidth, g_viewportHeight);
	glutInitDisplayMode( GLUT_RGB | GLUT_ALPHA | GLUT_DOUBLE | GLUT_DEPTH);
	glutCreateWindow(WINDOW_TITLE);

	// Register callbacks:
	glutDisplayFunc(display);
	glutReshapeFunc(reshape);
	glutMouseFunc(MouseButton);
	glutMotionFunc(MouseMotion);
	glutKeyboardUpFunc(KeyboardUpFunc);
	glutSpecialFunc(SpecialFunc);
	glutSpecialUpFunc(SpecialUpFunc);
	glutIdleFunc(Idle);
	atexit(DestroyWindow);

	InitializeDrawing();

	//cactus.obj
	handles.push_back(Handle(_dualSphere(c3gaPoint(.0, .04, 0.7) - 0.5*SQR(0.15)*ni)));
	handles.push_back(Handle(_dualSphere(c3gaPoint(.0, .04, -0.8) - 0.5*SQR(0.15)*ni)));

	////cylinder.obj
	//handles.push_back( boost::shared_ptr<Handle>( new Handle( _dualSphere(c3gaPoint(.0, 0.9,.0) - 0.5*SQR(0.25)*ni), false, P1 ) ) );
	//handles.push_back( boost::shared_ptr<Handle>( new Handle( _dualSphere(c3gaPoint(.0,-0.9,.0) - 0.5*SQR(0.25)*ni), true, P2 ) ) );

	//Armadillo Pie y Mano
	//handles.push_back( boost::shared_ptr<Handle>( new Handle( _dualSphere(c3gaPoint(-.5, 0.45,-.3) - 0.5*SQR(0.15)*ni), false, P1 ) ) );
	//handles.push_back( boost::shared_ptr<Handle>( new Handle( _dualSphere(c3gaPoint(.2,-0.6,.1) - 0.5*SQR(0.15)*ni), true, P2 ) ) );
	//Armadillo Pubis y Cabeza
	//handles.push_back( boost::shared_ptr<Handle>( new Handle( _dualSphere(c3gaPoint(.0, 0.4,-.2) - 0.5*SQR(0.15)*ni), false, P1 ) ) );
	//handles.push_back( boost::shared_ptr<Handle>( new Handle( _dualSphere(c3gaPoint(.0,-0.05,.1) - 0.5*SQR(0.15)*ni), true, P2 ) ) );

	vertexDescriptors.resize(mesh.numVertices());
	triangles.resize(mesh.numFaces()*3);
	int n = vertexDescriptors.get_size();

	for(Vertex& vertex : mesh.getVertices()) {
		vertexDescriptors.normalsOrig[vertex.ID] = vertex.n;

		normalizedPoint position = c3gaPoint( vertex.p.x(), vertex.p.y(), vertex.p.z() );

		for( Handle& handle : handles)
		{
			TRversor TR = handle.GetTRVersor();

			if( _double(position << (TR * handle.dS * inverse(TR))) > 0 ) //inside the sphere
			{
				handle.constraints.insert(vertex.ID);
				vertexDescriptors.constrainedPositions[vertex.ID] = vertex.p;
			}
		}
	}

	for(Face& face : mesh.getFaces()) {
		int i = face.ID;
		int	v1 = face.edge->vertex->ID;
		int	v2 = face.edge->next->vertex->ID;
		int	v3 = face.edge->next->next->vertex->ID;
		triangles.faces[i*3 + 0] = v1;
		triangles.faces[i*3 + 1] = v2;
		triangles.faces[i*3 + 2] = v3;
	}

	A = CreateLaplacianMatrix( &mesh, systemType );

	ComputeLaplacianCoordinates(A, &mesh, vertexDescriptors.laplacianCoordinates);

	b3 = Eigen::MatrixXd(A->numRows(), 3);

	PreFactor(A, handles);

	glutMainLoop();

	return 0;
}

void SolveLinearSystem(VertexBuffer& vertexDescriptors)
{
	int n = vertexDescriptors.get_size();
	for( int i = 0 ; i < n ; ++i ) {
		b3.row(i) = vertexDescriptors.laplacianCoordinates[i];
	}

	for( Handle& handle : handles) {
		Eigen::Affine3d M = MotorToMatrix(handle.GetTRVersor());
		for(int i : handle.constraints) {
			b3.row(i) = M * vertexDescriptors.constrainedPositions[i];
		}
	}

	xyz = solver.solve(b3);

	for( int i = 0 ; i < n ; ++i )
	{
		vertexDescriptors.deformedPositions[i] = xyz.row(i);
		vertexDescriptors.normals[i] = vertexDescriptors.rotors[i]._transformVector(vertexDescriptors.normalsOrig[i]);
	}
}

void UpdateLaplaciansRotation(Mesh *mesh, std::shared_ptr<SparseMatrix> A, VertexBuffer& vertexDescriptors)
{
	Eigen::Matrix3d m;
	for( Vertex& vertex : mesh->getVertices() )
	{
		m.setZero();
		int i = vertex.ID;
		const Eigen::Vector3d &pi = mesh->vertexAt(i).p;
		const Eigen::Vector3d &tpi = vertexDescriptors.deformedPositions[i];
		double S = 0;
		for(Vertex::EdgeAroundIterator edgeAroundIter = vertex.iterator() ; !edgeAroundIter.end() ; edgeAroundIter++)
		{
			int j = edgeAroundIter.edge_out()->pair->vertex->ID;
			const Eigen::Vector3d &pj = mesh->vertexAt(j).p;
			const Eigen::Vector3d &tpj = vertexDescriptors.deformedPositions[j];
			const double wij = (*A)(i, j);
			Eigen::Vector3d eij = (pj - pi);
			Eigen::Vector3d teij = tpj - tpi;
			m += (wij * eij) * teij.transpose();
			S += eij.dot(eij) * wij;
			S += teij.dot(teij) * wij;
		}
		vertexDescriptors.rotors[i] = GARotorEstimator(m, S);
	}

	std::fill(vertexDescriptors.laplacianCoordinates.begin(), vertexDescriptors.laplacianCoordinates.end(), Eigen::Vector3d(0,0,0));
	for(Vertex& vertex : mesh->getVertices())
	{
		int i = vertex.ID;
		for(Vertex::EdgeAroundIterator edgeAroundIter = vertex.iterator() ; !edgeAroundIter.end() ; edgeAroundIter++)
		{
			int j = edgeAroundIter.edge_out()->pair->vertex->ID;
			double wij = (*A)(i, j);
			Eigen::Quaterniond &Ri = vertexDescriptors.rotors[i];
			Eigen::Quaterniond &Rj = vertexDescriptors.rotors[j];
			Eigen::Vector3d V = mesh->vertexAt(j).p - mesh->vertexAt(i).p;
			vertexDescriptors.laplacianCoordinates[i] += 0.5 * wij * (Ri._transformVector(V) + Rj._transformVector(V));
		}
	}
}

void display()
{
	/*
	 *	matrices
	 */
	glViewport( 0, 0, g_viewportWidth, g_viewportHeight );
	glMatrixMode(GL_PROJECTION);
	glLoadIdentity();
	pickLoadMatrix();
	GLpick::g_frustumFar = 1000.0;
	GLpick::g_frustumNear = .1;
	gluPerspective( 60.0, (double)g_viewportWidth/(double)g_viewportHeight, GLpick::g_frustumNear, GLpick::g_frustumFar );
	glMatrixMode(GL_MODELVIEW);
	glLoadIdentity();

	glShadeModel(GL_SMOOTH);	//gouraud shading
	glClearDepth(1.0f);
	glClearColor( .75f, .75f, .75f, .0f );
	glHint( GL_PERSPECTIVE_CORRECTION_HINT, GL_NICEST );

	/*
	 *	estados
	 */
	glEnable(GL_CULL_FACE);		//face culling
	glCullFace( GL_BACK );
	glFrontFace( GL_CCW );
	glEnable(GL_DEPTH_TEST);	//z-buffer
	glDepthFunc(GL_LEQUAL);

	/*
	 *	iluminacion
	 */
	float		ambient[] = { .3f, .3f, .3f, 1.f };
	float		diffuse[] = { .3f, .3f, .3f, 1.f };
	float		position[] = { .0f, 0.f, -150.f, 1.f };
	float		specular[] = { 1.f, 1.f, 1.f };

	glLightfv( GL_LIGHT0, GL_AMBIENT, ambient );
	glLightfv( GL_LIGHT0, GL_DIFFUSE, diffuse );
	glLightf(GL_LIGHT0, GL_CONSTANT_ATTENUATION, 0);
	glLightf(GL_LIGHT0, GL_LINEAR_ATTENUATION, 0.0125);
	glEnable(  GL_LIGHT0   );
	glEnable(  GL_LIGHTING );
	glMaterialfv( GL_FRONT_AND_BACK, GL_SPECULAR, specular );
	glMaterialf( GL_FRONT_AND_BACK, GL_SHININESS, 50.f );

	glClear( GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT );

	glLoadIdentity();

	g_camera.glLookAt();

	glLightfv( GL_LIGHT0, /*GL_POSITION*/GL_SPOT_DIRECTION, position );
	glPushMatrix();

	rotorGLMult(g_modelRotor);

	if(!g_computeBasis)

	{
		static bool oneTime = true;
		if(oneTime || (g_rotateKeyRotors || g_translateKeyRotors))
		{
			if(oneTime == true)
			{
				SolveLinearSystem(vertexDescriptors);
			}

			oneTime = false;

			for(int i = 0 ; i < 10 ; ++i)
			{
				UpdateLaplaciansRotation(&mesh, A, vertexDescriptors);
				SolveLinearSystem(vertexDescriptors);
			}
		}
	}
	if(g_iterateManyTimes)
	{
		for(int i = 0 ; i < 400 ; ++i)
		{
			UpdateLaplaciansRotation(&mesh, A, vertexDescriptors);
			SolveLinearSystem(vertexDescriptors);
		}
		g_iterateManyTimes = false;
	}

	if (GLpick::g_pickActive) glLoadName((GLuint)-1);

	double alpha = 1.0;

	//glEnable (GL_BLEND);
	//glBlendFunc (GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA);
	//alpha = 0.5;

	//Mesh-Faces Rendering
	glPolygonMode( GL_FRONT_AND_BACK, GL_FILL /*GL_LINE GL_FILL GL_POINT*/);
	glEnable (GL_POLYGON_OFFSET_FILL);
	glPolygonOffset (1., 1.);
	glColorMaterial(GL_FRONT_AND_BACK, GL_AMBIENT_AND_DIFFUSE);
	glEnable( GL_COLOR_MATERIAL );
	if (GLpick::g_pickActive) glLoadName((GLuint)10);

	glColor4d( 1, 1, 1, alpha );
	glEnableClientState(GL_NORMAL_ARRAY);
	glEnableClientState(GL_VERTEX_ARRAY);
	glVertexPointer(3, GL_DOUBLE, 0, &vertexDescriptors.deformedPositions[0]);
	glNormalPointer(GL_DOUBLE, 0, &vertexDescriptors.normals[0]);
	// draw the model
	glDrawElements(GL_TRIANGLES, triangles.get_size(), GL_UNSIGNED_INT, &triangles.faces[0]);
	// deactivate vertex arrays after drawing
	glDisableClientState(GL_VERTEX_ARRAY);
	glDisableClientState(GL_NORMAL_ARRAY);

	if(g_showWires)
	{
		if (!GLpick::g_pickActive)
		{
			//Mesh-Edges Rendering (superimposed to faces)
			glPolygonMode( GL_FRONT_AND_BACK, GL_LINE /*GL_LINE GL_FILL GL_POINT*/);
			glColor4d( .5, .5, .5, alpha );
			glDisable( GL_LIGHTING );
				glEnableClientState(GL_VERTEX_ARRAY);
				glVertexPointer(3, GL_DOUBLE, 0, &vertexDescriptors.deformedPositions[0]);
				// draw the model
				glDrawElements(GL_TRIANGLES, triangles.get_size(), GL_UNSIGNED_INT, &triangles.faces[0]);
				// deactivate vertex arrays after drawing
				glDisableClientState(GL_VERTEX_ARRAY);
			glEnable( GL_LIGHTING );
		}
	}
	glDisable( GL_COLOR_MATERIAL );
	glDisable(GL_POLYGON_OFFSET_FILL);

	//glDisable (GL_BLEND);

	if(g_showSpheres)
	{
		//Handles rendering
		glPolygonMode( GL_FRONT_AND_BACK, GL_FILL /*GL_LINE GL_FILL GL_POINT*/);

		float	turcoise[] = { .0f, .5f, .5f, 0.3f };
		float	red[] = { .5f, .0f, .0f, 0.3f };

		for( int k = 0 ; k < handles.size() ; ++k)
		{
			if (GLpick::g_pickActive) glLoadName((GLuint)k);
			TRversor R = handles[k].GetTRVersor();
			DrawTransparentDualSphere( _dualSphere( R * handles[k].dS * inverse(R) ), turcoise );
		}
	}

	glPopMatrix();

	glutSwapBuffers();
}

void reshape(GLint width, GLint height)
{
	g_viewportWidth = width;
	g_viewportHeight = height;

	// redraw viewport
	glutPostRedisplay();
}

vectorE3GA mousePosToVector(int x, int y) {
	x -= g_viewportWidth / 2;
	y -= g_viewportHeight / 2;
	return _vectorE3GA((float)-x * e1 - (float)y * e2);
}

void MouseButton(int button, int state, int x, int y)
{
	g_rotateModel = false;
	g_rotateKeyRotors = false;
	g_translateKeyRotors = false;

	if (button == GLUT_LEFT_BUTTON)
	{
		g_prevMousePos = mousePosToVector(x, y);

		GLpick::g_pickWinSize = 1;
		g_dragObject = pick(x, g_viewportHeight - y, display, &g_dragDistance);

		if(g_dragObject == -1 || g_dragObject == 10 )
		{
			vectorE3GA mousePos = mousePosToVector(x, y);
			g_rotateModel = true;

			if ((_Float(norm_e(mousePos)) / _Float(norm_e(g_viewportWidth * e1 + g_viewportHeight * e2))) < 0.2)
				g_rotateModelOutOfPlane = true;
			else g_rotateModelOutOfPlane = false;
		}
		else if(g_dragObject >= 0 && g_dragObject < handles.size())
		{
			g_rotateKeyRotors = true;
		}
	}

	if (button == GLUT_RIGHT_BUTTON)
	{
		g_prevMousePos = mousePosToVector(x, y);

		GLpick::g_pickWinSize = 1;
		g_dragObject = pick(x, g_viewportHeight - y, display, &g_dragDistance);

		if(g_dragObject >= 0 && g_dragObject < handles.size())
			g_translateKeyRotors = true;
	}
}

void MouseMotion(int x, int y)
{
	if (g_rotateModel || g_rotateKeyRotors || g_translateKeyRotors )
	{
		// get mouse position, motion
		vectorE3GA mousePos = mousePosToVector(x, y);
		vectorE3GA motion = mousePos - g_prevMousePos;

		if (g_rotateModel)
		{
			// update rotor
			if (g_rotateModelOutOfPlane)
				g_modelRotor = exp(g_camera.rotateVel * (motion ^ e3) ) * g_modelRotor;
			else
				g_modelRotor = exp(0.00001f * (motion ^ mousePos) ) * g_modelRotor;
		}
		if(g_rotateKeyRotors)
		{
			//rotor R1 =  _rotor( inverse(g_modelRotor) * exp(-g_camera.rotateVel * (motion ^ e3) ) * g_modelRotor);
			rotor R1 =  _rotor( exp(-g_camera.rotateVel * (motion ^ e3) ) );
			if(g_dragObject < handles.size())
			{
				rotor R = handles[g_dragObject].R;
				handles[g_dragObject].R = normalize(_TRversor( R1 * R  ) );
			}
		}

		if(g_translateKeyRotors)
		{
			normalizedTranslator T1 = _normalizedTranslator(inverse(g_modelRotor) * exp( _freeVector(-g_camera.translateVel*motion*ni) ) * g_modelRotor);
			if(g_dragObject < handles.size())
			{
				translator T = handles[g_dragObject].T;
				handles[g_dragObject].T = normalize(_TRversor( T1 * T ));
			}
		}

		// remember mouse pos for next motion:
		g_prevMousePos = mousePos;

		// redraw viewport
		glutPostRedisplay();
	}
}

void SpecialFunc(int key, int x, int y)
{
	switch(key) {
		case GLUT_KEY_F1 :
			{
				int mod = glutGetModifiers();
				if(mod == GLUT_ACTIVE_CTRL || mod == GLUT_ACTIVE_SHIFT )
				{
				}
			}
			break;
		case GLUT_KEY_UP:
			{
				if(g_rotateKeyRotors)
				{
					handles[g_dragObject].dS = ChangeDualSphereRadiusSize(handles[g_dragObject].dS, 0.025);

					// redraw viewport
					glutPostRedisplay();
				}
			}
			break;
		case GLUT_KEY_DOWN:
			{
				if(g_rotateKeyRotors)
				{
					handles[g_dragObject].dS = ChangeDualSphereRadiusSize(handles[g_dragObject].dS, -0.025);

					// redraw viewport
					glutPostRedisplay();
				}
			}
			break;
	}
}

void SpecialUpFunc(int key, int x, int y)
{
}

void KeyboardUpFunc(unsigned char key, int x, int y)
{
	if(key == 'w' || key == 'W')
	{
		g_showWires = !g_showWires;
		glutPostRedisplay();
	}

	if( key == 'h' || key == 'H' )
	{
		g_showSpheres = !g_showSpheres;
		glutPostRedisplay();
	}

	if( key == 'x' || key == 'X' )
	{
		g_iterateManyTimes = true;
		glutPostRedisplay();
	}
}

void Idle()
{
	// redraw viewport
	//glutPostRedisplay();
}

void DestroyWindow()
{
	ReleaseDrawing();
}


#include <iostream>
#include <armadillo>

using namespace std;
using namespace arma;

extern "C" {

  int *createMatrix(int n, int arr[], int size) {
    sp_mat A = zeros<sp_mat>(n, n);
    sp_mat B = zeros<sp_mat>(n, n);
    for (int i =0; i<size; ++i) {
      A(i, i) = arr[i]+1;
      B(i, i) = arr[i];
    }
    sp_mat C = A*B;
    cout << C.t() << endl;
    for (int i=0; i< size; ++i) {
      arr[i] = C(i, i);
    }
    return arr;
  }

}


//////////////////////////////////////////////////////////////////////////////
// random::poisson_ext::devroye::detail::step4::standard.hpp         		//
//                                                                          //
//                                                                          //
//  (C) Copyright 2010 Erwann Rogard                                        //
//  Use, modification and distribution are subject to the                   //
//  Boost Software License, Version 1.0. (See accompanying file             //
//  LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)        //
//////////////////////////////////////////////////////////////////////////////
#ifndef BOOST_RANDOM_POISSON_EXT_DEVROYE_DETAIL_STEP4_STANDARD_HPP_ER_2010
#define BOOST_RANDOM_POISSON_EXT_DEVROYE_DETAIL_STEP4_STANDARD_HPP_ER_2010
// --- For testing only ---
#include <string>
#include <boost/timer.hpp>
#include <boost/format.hpp>
// ------------------------

#include <boost/random/poisson_ext/devroye/detail/q.hpp>
#include <boost/random/poisson_ext/devroye/detail/int_mean.hpp>

namespace boost{
namespace random{
namespace poisson{
namespace devroye{
namespace detail{
namespace step4{

	// Poisson sampler by the method of Devroye with a standard step4
	//
	// Complexity O(1+sqrt(mean)) as mean -> infinity, p.203
	//
    // Properties:
    // Int = int fails for mean >= c, 256 < c <= 512
    // Int = long int fails for mean >= c, 65536 < c <= 262144
    //
    // Warning : the properties need to be checked.
	template<typename Int,typename T,typename P>
    class standard :
        detail::q<Int,T,P>,
        public detail::int_mean<step4::standard<Int,T,P>,Int,T,P>
    {
		typedef detail::q<Int,T,P> q_fun_;
		typedef detail::int_mean<step4::standard<Int,T,P>,Int,T,P> super_;

		public:
		typedef typename super_::result_type result_type;
		typedef typename super_::input_type input_type;

		standard():super_(){}
		explicit standard(const result_type& mean)
        :super_(mean)
//      ,elapsed(0)
        {}

		bool accept(){
//        	t.restart();
            typename super_::input_type q_val = this->q_fun(
            	this->int_mean_val(),
                this->y()
            );
//            elapsed += t.elapsed();
			if(this->v()<q_val)
            {
            	this->y_ += this->int_mean_val();
                return true;
            }else{
				return false;
            }
		}

// Only for testing :
//		~standard(){
//        	std::string str = "elapsed for q() at mean = %1% : %2%";
//            boost::format f(str);
//            f%this->int_mean_val()%elapsed;
//        	std::cout << f << std::endl;
//        }
//        boost::timer t;
//        typename super_::input_type elapsed;

	};

}// step4
}// detail
}// devroye
}// poisson
}// random
}// boost

#endif


#ifndef SHIFT_MATH_FIXED_HPP
#define SHIFT_MATH_FIXED_HPP

#include <type_traits>
#include <boost/operators.hpp>
#include "shift/core/types.hpp"

namespace shift::math
{
template <typename K, std::size_t F>
class fixed;
}

namespace std
{
/// Returns the squared magnitude of the scalar value (|value|^2).
template <typename K, std::size_t F>
shift::math::fixed<K, F> norm(const shift::math::fixed<K, F>& value)
{
  return value * value;
}

/// Returns the magnitude of the scalar value (|value|).
template <typename K, std::size_t F>
shift::math::fixed<K, F> abs(shift::math::fixed<K, F> value)
{
  value.data() = std::abs(value.data());
  return value;
}

///
template <typename K, std::size_t F>
shift::math::fixed<K, F> floor(shift::math::fixed<K, F> value)
{
  if ((value.data() & shift::math::fixed<K, F>::fractional_mask) != 0)
  {
    if (value.data() >= 0)
      return shift::math::fixed<K, F>(static_cast<K>(value));
    else
      return shift::math::fixed<K, F>(static_cast<K>(value) - 1);
  }
  else
    return value;
}

///
template <typename K, std::size_t F>
shift::math::fixed<K, F> ceil(shift::math::fixed<K, F> value)
{
  if (value.data() & shift::math::fixed<K, F>::fractional_mask != 0)
  {
    if (value.data() >= 0)
      return shift::math::fixed<K, F>(static_cast<K>(value) + 1);
    else
      return shift::math::fixed<K, F>(static_cast<K>(value));
  }
  else
    return value;
}

///
template <typename K, std::size_t F>
shift::math::fixed<K, F> round(shift::math::fixed<K, F> value)
{
  return shift::math::fixed<K, F>(
    static_cast<K>(value + shift::math::fixed<K, F>::one / 2));
}
}

namespace shift::math
{
/// A fixed-point arithmetic class
template <typename K, std::size_t F>
class fixed : boost::operators<fixed<K, F>>, boost::shiftable<fixed<K, F>, int>
{
public:
  static_assert(std::is_integral<K>::value,
                "K in fixed<K, F> must be a signed integral type.");

  using storage_type = K;

  static const std::size_t total_bits = sizeof(K);
  static const std::size_t fractional_bits = F;
  static const std::size_t fractional_mask = (1 << F) - 1;
  static const std::size_t integer_mask = ~fractional_mask;
  static const std::size_t integer_bits = total_bits - F;
  static const K one = K(1) << fractional_bits;

  struct direct_copy
  {
  };

public:
  /// Default constructor.
  fixed() : _number(0)
  {
  }

  /// Copy constructor.
  fixed(const fixed& other) : _number(other._number)
  {
  }

  /// Copy constructor.
  /// ToDo...
  // template <typename K2, std::size_t F2>
  // fixed(const fixed<K2, F2>& other) : _number(other._number)
  //{
  //}

  /// Constructor which interprets the passed value as correctly formatted
  /// fixed point number.
  fixed(K value, direct_copy&) : _number(value)
  {
  }

  /// Constructor from a signed integral type T.
  template <typename T, ENABLE_IF(std::is_integral<T>::value)>
  fixed(T number) : _number(static_cast<K>(number) << fractional_bits)
  {
  }

  /// Constructor from a floating-point type.
  template <typename T, ENABLE_IF(std::is_floating_point<T>::value)>
  fixed(T number) : _number(static_cast<K>(one * number))
  {
  }

  /// Conversion operator to signed integral type T.
  template <typename T, ENABLE_IF(std::is_integral<T>::value)>
  operator T() const
  {
    static_assert(sizeof(T) * 8 >= sizeof(K) * 8 - fractional_bits,
                  "Cannot convert fixed-point number to type T because of "
                  "possible loss of data.");
    return static_cast<T>(_number >> fractional_bits);
  }

  /// Conversion operator to floating-point type T.
  template <typename T, ENABLE_IF(std::is_floating_point<T>::value)>
  operator T() const
  {
    return static_cast<T>(_number) / (1 << fractional_bits);
  }

  /// Returns a reference to the internal number store.
  storage_type& data()
  {
    return _number;
  }

  /// Assignment operator.
  fixed& operator=(const fixed& other)
  {
    _number = other._number;
    return *this;
  }

  ///
  bool operator==(const fixed& other) const
  {
    return _number == other._number;
  }

  ///
  bool operator<(const fixed& other) const
  {
    return _number < other._number;
  }

  ///
  bool operator!() const
  {
    return !_number;
  }

  ///
  fixed operator~() const
  {
    return fixed(~_number, direct_copy());
  }

  ///
  fixed operator-() const
  {
    return fixed(-t._number, direct_copy());
  }

  ///
  fixed& operator++()
  {
    _number += one;
    return *this;
  }

  ///
  fixed& operator--()
  {
    _number -= one;
    return *this;
  }

  ///
  fixed& operator+=(const fixed& other)
  {
    _number += other._number;
    return *this;
  }

  ///
  fixed& operator-=(const fixed& other)
  {
    _number -= other._number;
    return *this;
  }

  ///
  fixed& operator&=(const fixed& other)
  {
    _number &= other._number;
    return *this;
  }

  ///
  fixed& operator|=(const fixed& other)
  {
    _number |= other._number;
    return *this;
  }

  ///
  fixed& operator^=(const fixed& other)
  {
    _number ^= other._number;
    return *this;
  }

  ///
  fixed& operator>>=(const int& other)
  {
    _number >>= other;
    return *this;
  }

  ///
  fixed& operator<<=(const int& other)
  {
    _number <<= other;
    return *this;
  }

  ///
  fixed& operator>>=(const fixed& other)
  {
    _number >>= static_cast<K>(other);
    return *this;
  }

  ///
  fixed& operator<<=(const fixed& other)
  {
    _number <<= static_cast<K>(other);
    return *this;
  }

private:
  storage_type _number;
};
}

#endif


/// General Matrix Multiplication
#include <HElib/FHE.h>
#include <HElib/FHEContext.h>
#include <HElib/EncryptedArray.h>
#include <HElib/NumbTh.h>

#include "SMP/DoublePacking.hpp"
#include "SMP/Matrix.hpp"
#include "SMP/Timer.hpp"
#include <boost/asio.hpp>
#include <boost/asio/ip/tcp.hpp>
#include <iostream>
#include <numeric>
#include <list>
using boost::asio::ip::tcp;
static void print_time(std::string const& doc, Duration_t const& dur) {
    std::cout << doc << " " << time_as_millsecond(dur) << std::endl;
}

inline long round_div(long a, long b) {
    return (a + b - 1) / b;
}

void zero(Matrix &mat) {
    for (long i = 0; i < mat.NumRows(); i++)
        for (long j = 0; j < mat.NumCols(); j++)
            mat[i][j] = 0;
}

void randomize(Matrix &mat) {
    for (long i = 0; i < mat.NumRows(); i++)
        for (long j = 0; j < mat.NumCols(); j++)
            mat[i][j] = NTL::RandomBnd(2L);
}

void mod_to_plain(Matrix &mat, long prime) {
    for (long i = 0; i < mat.NumRows(); i++)
        for (long j = 0; j < mat.NumCols(); j++)
            mat[i][j] %= prime;
}

struct Duplication {
    long blk_id;
    long dup;
    long slots_left;
};

Duplication compute_duplications(const long total_rows,
                                 const long x_blk_id,
                                 const long num_slots) {
    long N = std::min((x_blk_id  + 1) * num_slots, total_rows) - x_blk_id * num_slots;
    Duplication d;
    d.blk_id = x_blk_id;
    d.dup = num_slots / N;
    d.slots_left = num_slots - d.dup * N;
    return d;
}

std::vector<long> compute_rotations(long x_blk, long total_rows_X,
                                    long y_blk, long total_rows_Y,
                                    long num_slots) {
    auto dupA = compute_duplications(total_rows_X, x_blk, num_slots);
    auto dupB = compute_duplications(total_rows_Y, y_blk, num_slots);
    std::vector<long> rotations;
    long sze;
    // fully packed
    if (dupA.dup == 1 and dupB.dup == 1) {
        sze = num_slots;
    } else if (dupA.slots_left != 0 or dupB.slots_left != 0) {
        // cumble to implement, so just skip this case
        std::cout << "cumble\n" << std::endl;
        sze = num_slots;
    } else {
        long NA = (num_slots - dupA.slots_left) / dupA.dup;
        long NB = (num_slots - dupB.slots_left) / dupB.dup;
        sze = std::max(NA, NB) / (dupA.dup * dupB.dup);
        sze = std::max(1L, sze);
    }
    rotations.resize(sze);
    std::iota(rotations.begin(), rotations.end(), 0);
    return rotations;
}

void fill_compute(Matrix& mat,
                  int x, int y, int k,
                  const std::vector<NTL::zz_pX> &slots,
                  const EncryptedArray *ea) {
    const long l = ea->size();
    const long d = ea->getDegree();
    long row_start = x * l;
    long row_end = row_start + l;
    long col_start = y * l;
    long col_end = col_start + l;

    assert(slots.size() == l);
    for (long ll = 0; ll < l; ll++) {
        long computed = NTL::coeff(slots[ll], d - 1)._zz_p__rep;
        long row = row_start + ll;
        long col = col_start + ll + k;
        if (col >= col_end)
            col = col_start + col % l;
        mat.put(row, col, computed);
    }
}

long ceil_round(long a, long l) {
    return (a + l - 1) / l * l;
}

FHEcontext receive_context(std::istream &s) {
    unsigned long m, p, r;
    std::vector<long> gens, ords;
    readContextBase(s, m, p, r, gens, ords);
    FHEcontext context(m, p, r, gens, ords);
    s >> context;
    return context;
}

void send_context(std::ostream &s, FHEcontext const& context) {
    writeContextBase(s, context);
    s << context;
}

void play_client(tcp::iostream &conn,
                 FHESecKey &sk,
                 FHEcontext &context,
                 const long n1,
                 const long n2,
                 const long n3) {
    FHEPubKey ek(sk);
    ek.makeSymmetric();
    conn << ek;
    const EncryptedArray *ea = context.ea;
    const long l = ea->size();
    const long d = ea->getDegree();

    NTL::SetSeed(NTL::to_ZZ(123));
    Matrix A, B, ground_truth;
    A.SetDims(n1, n2);
    B.SetDims(n2, n3);
    randomize(A);
    randomize(B);
    ground_truth = mul(A, B);
    mod_to_plain(ground_truth, context.alMod.getPPowR());
    /// print grouth truth for debugging
    // save_matrix(std::cout, ground_truth);
    const long MAX_X1 = round_div(A.NumRows(), l);
    const long MAX_Y1 = round_div(A.NumCols(), d);
    const long MAX_X2 = round_div(B.NumCols(), l);

    std::vector<std::vector<Ctxt>> uploading;
    uploading.resize(MAX_X1, std::vector<Ctxt>(MAX_Y1, sk));
    auto start_time = Clock::now();
    /// encrypt matrix
    for (int x = 0; x < MAX_X1; x++) {
        for (int k = 0; k < MAX_Y1; k++) {
            internal::BlockId blk = {x, k};
            auto packed_rows = internal::partition(A, blk, *ea, false);
            //ea->skEncrypt(uploading[x][k], sk, packed_rows.polys);
        }
    }
    auto end_time = Clock::now();
    print_time("encryption", end_time - start_time);

    /// send ciphertexts of matrix
    start_time = Clock::now();
    for (auto const& row : uploading) {
        for (auto const& ctx : row)
            conn << ctx;
    }
    std::cout << MAX_X1 * MAX_Y1 << " ciphertexts sent" << std::endl;
    end_time = Clock::now();
    print_time("client->server", end_time - start_time);

    /// waiting results
    long rows_of_A = A.NumRows();
    long rows_of_Bt = B.NumCols(); // Bt::Rows = B::Cols
    std::list<Ctxt> ret_ctxs;
    for (int x = 0; x < MAX_X1; x++) {
        for (int y = 0; y < MAX_X2; y++) {
            size_t num_rots = compute_rotations(x, rows_of_A,
                                                y, rows_of_Bt, l).size();
            for (size_t k = 0; k < num_rots; k++) {
                Ctxt result(sk);
                conn >> result;
                ret_ctxs.emplace_back(result);
            }
        }
    }
    /// decrypt
    Matrix computed;
    computed.SetDims(A.NumRows(), B.NumCols());
    zero(computed);
    int x = 0;
    int y = 0;
    auto itr = ret_ctxs.begin();
    std::vector<NTL::zz_pX> slots;
    start_time = Clock::now();
    for (int x = 0; x < MAX_X1; x++) {
        for (int y = 0; y < MAX_X2; y++) {
            size_t num_rots  = compute_rotations(x, rows_of_A,
                                                 y, rows_of_Bt, l).size();
            for (size_t k = 0; k < num_rots; k++) {
                //ea->decrypt(*itr++, sk, slots);
                fill_compute(computed, x, y, k, slots, ea);
            }
        }
    }
    end_time = Clock::now();
    print_time("decryption", end_time - start_time);
    if (!::is_same(ground_truth, computed))
        std::cerr << "The computation seems wrong " << std::endl;
    else
        std::cout << "passed" << std::endl;
}

void my_encode(EncryptedArray const& ea,
               zzX &ret,
               std::vector<NTL::zz_pX> &slots) {
    auto const& encoder = ea.getContext().alMod.getDerived(PA_zz_p());
    NTL::zz_pX tmp;
    encoder.CRT_reconstruct(tmp, slots);
    // NTL::conv(ret, tmp);
}

void play_server(tcp::iostream &conn,
                 const long n1,
                 const long n2,
                 const long n3) {
    auto context = receive_context(conn);
    const EncryptedArray *ea = context.ea;
    const long l = ea->size();
    const long d = ea->getDegree();
    FHEPubKey ek(context);
    conn >> ek;
    NTL::SetSeed(NTL::to_ZZ(123));

    Matrix A, B;
    A.SetDims(n1, n2);
    B.SetDims(n2, n3);
    randomize(A);
    randomize(B);

    const long MAX_X1 = round_div(A.NumRows(), l);
    const long MAX_Y1 = round_div(A.NumCols(), d);
    Matrix Bt;
    transpose(&Bt, B);
    const long MAX_X2 = round_div(Bt.NumRows(), l);
    const long MAX_Y2 = round_div(Bt.NumCols(), d);
    auto start_time = Clock::now();
    std::vector<std::vector<NewPlaintextArray>> precomputed;
    precomputed.resize(MAX_X2, std::vector<NewPlaintextArray>(MAX_Y1, *ea));
    for (int y = 0; y < MAX_X2; y++) {
        for (int k = 0; k < MAX_Y1; k++) {
            internal::BlockId blk = {y, k};
            auto packed_rows = internal::partition(Bt, blk, *ea, true);
            encode(*ea, precomputed[y][k], packed_rows.polys);
        }
    }
    auto end_time = Clock::now();
    print_time("precomputation", end_time - start_time);

    /// receving ciphertexts from the client
    std::vector<std::vector<Ctxt>> received;
    received.resize(MAX_X1, std::vector<Ctxt>(MAX_Y1, ek));
    for (int x = 0; x < MAX_X1; x++) {
        for (int k = 0; k < MAX_Y1; k++)
            conn >> received[x][k];
    }

    /// compute the matrix mulitplication
    long rows_of_A = A.NumRows();
    long rows_of_Bt = Bt.NumRows();
    size_t send_ctx = 0;
    zzX packed_polys;
    Duration_t computation(0), network(0);
    for (int x = 0; x < MAX_X1; x++) {
        for (int y = 0; y < MAX_X2; y++) {
            start_time = Clock::now();
            std::vector<long> rotations = compute_rotations(x, rows_of_A,
                                                            y, rows_of_Bt, l);
            size_t num_rotations = rotations.size();
            std::vector<Ctxt> summations(num_rotations, ek);
            for (int k = 0; k < MAX_Y1; k++) {
                for (size_t rot = 0; rot < num_rotations; rot++) {
                    auto rotated(precomputed[y][k]);
                    rotate(*ea, rotated, -rotations[rot]);
                    /// pack the polys into one poly
                    //my_encode(*ea, packed_polys, rotated.getData());
                    Ctxt client_ctx(received[x][k]);
                    /// ctx * plain
                    client_ctx.multByConstant(packed_polys);
                    summations[rot] += client_ctx;
                }
            }
            for (auto &sm : summations)
                sm.modDownToLevel(1);
            end_time = Clock::now();
            computation += (end_time - start_time);

            start_time = Clock::now();
            for (auto const&sm : summations)
                conn << sm;
            end_time = Clock::now();
            network += (end_time - start_time);
            send_ctx += summations.size();
        }
    }
    print_time("computation", computation);
    print_time("server->client", network);
    std::cout << "Sent " << send_ctx << " ciphertexts" << std::endl;
}

int run_client(long n1, long n2, long n3) {
    tcp::iostream conn("127.0.0.1", "12345");
    if (!conn) {
        std::cerr << "Can not connect to server!" << std::endl;
        return -1;
    }
    const long m = 8192;
    //const long p = 641;
    const long p = 3329;
    const long r = 1;
    const long L = 3;
    FHEcontext context(m, p, r);
    context.bitsPerLevel = 14 + std::ceil(std::log(m)/2 + r* std::log(p));
    buildModChain(context, L);
    std::cout << "kappa = " << context.securityLevel() << std::endl;
    FHESecKey sk(context);
    sk.GenSecKey(64);
    /// send FHEcontext obj
    send_context(conn, context);
    /// send the evaluation key
    play_client(conn, sk, context, n1, n2, n3);
    conn.close();
    return 1;
}

int run_server(long n1, long n2, long n3) {
    boost::asio::io_service ios;
    tcp::endpoint endpoint(tcp::v4(), 12345);
    tcp::acceptor acceptor(ios, endpoint);

    for (;;) {
        tcp::iostream conn;
        boost::system::error_code err;
        acceptor.accept(*conn.rdbuf(), err);
        if (!err) {
            std::cout << "Connected!" << std::endl;
            play_server(conn, n1, n2, n3);
            break;
        }
    }
}

int main(int argc, char *argv[]) {
    ArgMapping argmap;
    long role;
    long n1 = 8;
    long n2 = 8;
    long n3 = 8;
    argmap.arg("N", n1, "n1");
    argmap.arg("M", n2, "n2");
    argmap.arg("D", n3, "n3");
    argmap.arg("R", role, "role");
    argmap.parse(argc, argv);
    if (role == 0) {
        return run_server(n1, n2, n3);
    } else if (role == 1) {
        return run_client(n1, n2, n3);
    }
}


#include <cstdlib>
#include <iostream>

#include <boost/units/systems/si/electric_potential.hpp>
#include <boost/units/systems/si/current.hpp>
#include <boost/units/systems/si/resistance.hpp>
#include <boost/units/systems/si/io.hpp>

int main()
{
    using namespace boost::units;
    using namespace boost::units::si;

    quantity<current> I = 5 * amperes;
    quantity<resistance> R = 10 * ohms;

    std::cout << "I = " << I << std::endl;
    std::cout << "R = " << R << std::endl;
    std::cout << "U = I * R = " << I * R << std::endl;
    return EXIT_SUCCESS;
}