#include "Camera.h"
#include "TinyRaytracer.h"

#include <omp.h>
#include <iostream>
#include <string>
#include <numbers>
#include <vector>
#include <sstream>
#include <filesystem>

#include <stb_image.h>
#include <stb_image_write.h>
#include <Eigen/Dense>
#include <GLFW/glfw3.h>

using Eigen::Vector3f;

struct Sphere
{
    Vector3f center;
    float radius;

    Sphere(const Vector3f& c, const float& r) : center(c), radius(r) {}

    bool ray_intersect(const Vector3f& orig, const Vector3f& dir, float& t0) const
    {
        Vector3f L = center - orig;
        float tca = L.dot(dir);
        float d2 = L.dot(L) - tca * tca;
        if (d2 > radius * radius) return false;
        float thc = sqrtf(radius * radius - d2);
        t0 = tca - thc;
        float t1 = tca + thc;
        if (t0 < 0) t0 = t1;
        if (t0 < 0) return false;
        return true;
    }
};

Vector3f cast_ray(const Vector3f& orig, const Vector3f& dir, const Sphere& sphere)
{
    float sphere_dist = std::numeric_limits<float>::max();
    if (!sphere.ray_intersect(orig, dir, sphere_dist))
    {
        return Vector3f(0.2f, 0.7f, 0.8f); // background color
    }
    return Vector3f(0.4f, 0.4f, 0.3f);
}

void RenderToBuffer(const Camera& cam, const Sphere& sphere, std::stringstream& outBuffer)
{
    const std::string ProjPath(PROJECT_PATH);
    const std::string ResoucePath = ProjPath + std::string("/resource");

    if (!std::filesystem::exists(ResoucePath))
        std::filesystem::create_directories(ResoucePath);

    std::string OutPath = ResoucePath + std::string("/out.png");
    float tanFov2 = (float)tan(cam.GetFov() / 2.);
    std::vector<Vector3f> framebuffer(width * height);

#pragma omp parallel for
    for (size_t j = 0; j < height; j++)
    {
        for (size_t i = 0; i < width; i++)
        {
            float x = (2 * (i + 0.5f) / (float)width - 1.0f) * tanFov2 * (float)width / (float)height;
            float y = -(2 * (j + 0.5f) / (float)height - 1.0f) * tanFov2;
            // Vector3f dir = Eigen::AngleAxisf(std::numbers::pi / 4, Vector3f(0.0f, 0.0f, 1.0f)).matrix().transpose() * Vector3f(x, y, -1.0f).normalized();
            Vector3f dir = Vector3f(x, y, -1.0f).normalized();
            framebuffer[i + j * width] = cast_ray(cam.GetPos(), dir, sphere);
        }
    }

    for (size_t i = 0; i < height * width; i++)
    {
        for (size_t j = 0; j < 3; j++)
        {
            outBuffer << (char)(255 * std::max(0.f, std::min(1.f, framebuffer[i][j])));
        }
    }

    //glDrawPixels(width, height, GL_RGB, GL_UNSIGNED_BYTE, ss.str().c_str());

    //stbi_write_png(OutPath.c_str(), width, height, 3, ss.str().c_str(), 0);
}

void display(std::stringstream& buffer)
{
    glDrawPixels(width, height, GL_RGB, GL_UNSIGNED_BYTE, buffer.str().c_str());
}

int main()
{
    GLFWwindow* window;
    Camera cam;
    Sphere sphere(Vector3f(-3.0f, 0.0f, -16.0f), 2);
    std::stringstream buffer;

    /* Initialize the library */
    if (!glfwInit())
        return -1;

    /* Create a windowed mode window and its OpenGL context */
    window = glfwCreateWindow(width, height, "Hello World", NULL, NULL);
    if (!window)
    {
        glfwTerminate();
        return -1;
    }

    /* Make the window's context current */
    glfwMakeContextCurrent(window);

    /* Loop until the user closes the window */
    while (!glfwWindowShouldClose(window))
    {
        /* Render here */
        RenderToBuffer(cam, sphere, buffer);
        display(buffer);

        /* Swap front and back buffers */
        glfwSwapBuffers(window);

        /* Poll for and process events */
        glfwPollEvents();
    }

    glfwTerminate();
}


#include "hgd.hh"

#include <NTL/RR.h>

using namespace std;
using namespace NTL;

static RR
AFC(const RR &I)
{
    /*
     * FUNCTION TO EVALUATE LOGARITHM OF THE FACTORIAL I
     * IF (I .GT. 7), USE STIRLING'S APPROXIMATION
     * OTHERWISE,  USE TABLE LOOKUP
     */
    double AL[8] =
    { 0.0, 0.0, 0.6931471806, 1.791759469, 3.178053830, 4.787491743,
      6.579251212, 8.525161361 };

    if (I <= 7) {
        return to_RR(AL[to_int(round(I))]);
    } else {
        RR LL = log(I);
        return (I+0.5) * LL - I + 0.399089934;
    }
}

static RR
RAND(PRNG *prng, long precision)
{
    ZZ div = to_ZZ(1) << precision;
    ZZ rzz = prng->rand_zz_mod(div);
    return to_RR(rzz) / to_RR(div);
}

ZZ
HGD(const ZZ &KK, const ZZ &NN1, const ZZ &NN2, PRNG *prng)
{
    /*
     * XXX
     * NTL is single-threaded by design: there is a global precision
     * setting, which gets switched back and forth all over the place
     * (see NTL's RR.c).  We should hold a lock around any RR usage,
     * or re-implement the relevant parts of NTL::RR with a scoped
     * precision parameter..
     */
    long precision = NumBits(NN1 + NN2 + KK) + 10;
    RR::SetPrecision(precision);

    RR JX;      // the result
    RR TN, N1, N2, K;
    RR P, U, V, A, IX, XL, XR, M;
    RR KL, KR, LAMDL, LAMDR, NK, NM, P1, P2, P3;

    bool REJECT;
    RR MINJX, MAXJX;

    double CON = 57.56462733;
    double DELTAL = 0.0078;
    double DELTAU = 0.0034;
    double SCALE = 1.0e25;

    /*
     * CHECK PARAMETER VALIDITY
     */
    if ((NN1 < 0) || (NN2 < 0) || (KK < 0) || (KK > NN1 + NN2))
        throw_c(false);

    /*
     * INITIALIZE
     */
    REJECT = true;

    if (NN1 >= NN2) {
        N1 = to_RR(NN2);
        N2 = to_RR(NN1);
    } else {
        N1 = to_RR(NN1);
        N2 = to_RR(NN2);
    }

    TN = N1 + N2;

    if (to_RR(KK + KK) >= TN)  {
        K = TN - to_RR(KK);
    } else {
        K = to_RR(KK);
    }

    M = (K+1) * (N1+1) / to_RR(TN+2);

    if (K-N2 < 0) {
        MINJX = 0;
    } else {
        MINJX = K-N2;
    }

    if (N1 < K) {
        MAXJX = N1;
    } else {
        MAXJX = K;
    }

    /*
     * GENERATE RANDOM VARIATE
     */
    if (MINJX == MAXJX)  {
        /*
         * ...DEGENERATE DISTRIBUTION...
         */
        IX = MAXJX;
    } else if (M-MINJX < 10) {
        /*
         * ...INVERSE TRANSFORMATION...
         * Shouldn't really happen in OPE because M will be on the order of N1.
         * In practice, this does get invoked.
         */
        RR W;
        if (K < N2) {
            W = exp(CON + AFC(N2) + AFC(N1+N2-K) - AFC(N2-K) - AFC(N1+N2));
        } else {
            W = exp(CON + AFC(N1) + AFC(K) - AFC(K-N2) - AFC(N1+N2));
        }

 label10:
        P  = W;
        IX = MINJX;
        U  = RAND(prng, precision) * SCALE;

 label20:
        if (U > P) {
            U  = U - P;
            P  = P * (N1-IX)*(K-IX);
            IX = IX + 1;
            P  = P / IX / (N2-K+IX);
            if (IX > MAXJX)
                goto label10;
            goto label20;
        }
    } else {
        /*
         * ...H2PE...
         */
        RR S;
        SqrRoot(S, (TN-K) * K * N1 * N2 / (TN-1) / TN /TN);

        /*
         * ...REMARK:  D IS DEFINED IN REFERENCE WITHOUT INT.
         * THE TRUNCATION CENTERS THE CELL BOUNDARIES AT 0.5
         */
        RR D = trunc(1.5*S) + 0.5;
        XL = trunc(M - D + 0.5);
        XR = trunc(M + D + 0.5);
        A = AFC(M) + AFC(N1-M) + AFC(K-M) + AFC(N2-K+M);
        RR expon = A - AFC(XL) - AFC(N1-XL)- AFC(K-XL) - AFC(N2-K+XL);

        KL = exp(expon);
        KR = exp(A - AFC(XR-1) - AFC(N1-XR+1) - AFC(K-XR+1) - AFC(N2-K+XR-1));
        LAMDL = -log(XL * (N2-K+XL) / (N1-XL+1) / (K-XL+1));
        LAMDR = -log((N1-XR+1) * (K-XR+1) / XR / (N2-K+XR));
        P1 = 2*D;
        P2 = P1 + KL / LAMDL;
        P3 = P2 + KR / LAMDR;

 label30:
        U = RAND(prng, precision) * P3;
        V = RAND(prng, precision);

        if (U < P1)  {
            /* ...RECTANGULAR REGION... */
            IX    = XL + U;
        } else if  (U <= P2)  {
            /* ...LEFT TAIL... */
            IX = XL + log(V)/LAMDL;
            if (IX < MINJX) {
                goto label30;
            }
            V = V * (U-P1) * LAMDL;
        } else  {
            /* ...RIGHT TAIL... */
            IX = XR - log(V)/LAMDR;
            if (IX > MAXJX)  {
                goto label30;
            }
            V = V * (U-P2) * LAMDR;
        }

        /*
         * ...ACCEPTANCE/REJECTION TEST...
         */
        RR F;
        if ((M < 100) || (IX <= 50))  {
            /* ...EXPLICIT EVALUATION... */
            F = to_RR(1.0);
            if (M < IX) {
                for (RR I = M+1; I < IX; I++) {
                    /*40*/ F = F * (N1-I+1) * (K-I+1) / (N2-K+I) / I;
                }
            } else if (M > IX) {
                for (RR I = IX+1; I < M; I++) {
                    /*50*/ F = F * I * (N2-K+I) / (N1-I) / (K-I);
                }
            }
            if (V <= F)  {
                REJECT = false;
            }
        } else {
            /* ...SQUEEZE USING UPPER AND LOWER BOUNDS... */

            RR Y   = IX;
            RR Y1  = Y + 1.;
            RR YM  = Y - M;
            RR YN  = N1 - Y + 1.;
            RR YK  = K - Y + 1.;
            NK     = N2 - K + Y1;
            RR R   = -YM / Y1;
            S      = YM / YN;
            RR T   = YM / YK;
            RR E   = -YM / NK;
            RR G   = YN * YK / (Y1*NK) - 1.;
            RR DG  = to_RR(1.0);
            if (G < 0)  { DG = 1.0 + G; }
            RR GU  = G * (1.+G*(-0.5+G/3.0));
            RR GL  = GU - 0.25 * sqr(sqr(G)) / DG;
            RR XM  = M + 0.5;
            RR XN  = N1 - M + 0.5;
            RR XK  = K - M + 0.5;
            NM     = N2 - K + XM;
            RR UB  = Y * GU - M * GL + DELTAU +
                     XM * R * (1.+R*(-.5+R/3.)) +
                     XN * S * (1.+S*(-.5+S/3.)) +
                     XK * T * (1.+T*(-.5+T/3.)) +
                     NM * E * (1.+E*(-.5+E/3.));

            /* ...TEST AGAINST UPPER BOUND... */

            RR ALV = log(V);
            if (ALV > UB) {
                REJECT = true;
            } else {
                /* ...TEST AGAINST LOWER BOUND... */

                RR DR = XM * sqr(sqr(R));
                if (R < 0) {
                    DR = DR / (1.+R);
                }
                RR DS = XN * sqr(sqr(S));
                if (S < 0) {
                    DS = DS / (1.+S);
                }
                RR DT = XK * sqr(sqr(T));
                if (T < 0) {
                    DT = DT / (1.+T);
                }
                RR DE = NM * sqr(sqr(E));
                if (E < 0) {
                    DE = DE / (1.+E);
                }
                if (ALV < UB-0.25*(DR+DS+DT+DE) + (Y+M)*(GL-GU) - DELTAL) {
                    REJECT = false;
                } else {
                    /* ...STIRLING'S FORMULA TO MACHINE ACCURACY... */

                    if (ALV <=
                        (A - AFC(IX) -
                         AFC(N1-IX)  - AFC(K-IX) - AFC(N2-K+IX)) ) {
                        REJECT = false;
                    } else {
                        REJECT = true;
                    }
                }
            }
        }
        if (REJECT)  {
            goto label30;
        }
    }

    /*
     * RETURN APPROPRIATE VARIATE
     */

    if (KK + KK >= to_ZZ(TN)) {
        if (NN1 > NN2) {
            IX = to_RR(KK - NN2) + IX;
        } else {
            IX = to_RR(NN1) - IX;
        }
    } else {
        if (NN1 > NN2) {
            IX = to_RR(KK) - IX;
        }
    }
    JX = IX;
    return to_ZZ(JX);
}


#define NO_IMPORT_ARRAY

// debug
#include <iostream>
using namespace std;

#include <pybindings.h>

#include <assert.h>
#include <math.h>

#include <omp.h>
#include <boost/math/quaternion.hpp>

#include <container_pybindings.h>

#include "so3g_numpy.h"
#include "Projection.h"
#include "Ranges.h"
#include "exceptions.h"
#include "so_linterp.h"

// TRIG_TABLE_SIZE
//
// Set this macro to enable trig interpolation tables.  Studies show
// that a table sizeof 16k gives errors below .1 arcsecond.  To use
// math library asin and atan2, leave the macro undefined.
//
#define TRIG_TABLE_SIZE (1 << 14)

#ifdef TRIG_TABLE_SIZE
static asinTable asin_lookup(TRIG_TABLE_SIZE);
static atan2Table atan2_lookup(TRIG_TABLE_SIZE);
#  define ASIN asin_lookup.get
#  define ATAN2 atan2_lookup.get
#else
#  define ASIN asin
#  define ATAN2 atan2
#endif

typedef boost::math::quaternion<double> quatd;

inline bool isNone(const bp::object &pyo)
{
    return (pyo.ptr() == Py_None);
}

// ProjEng template system
//
// ProjEng classes will be templated like:
//
//     template<typename C, typename P, typename S>
//
// Parameters C, P, S are shorthands for:
//
// - C: CoordSys, determines how the boresight position and detector
//   offsets are interpreted and combined.  The most obvious cases
//   (e.g. ProjTAN, ProjCAR), labeled with FITS projection codes,
//   imply that boresight and detector offsets should be treated as
//   rotation quaternions, multiplied together, and interpreted as
//   celestial positions within a particular projection.  Other cases
//   include ProjFlat (where boresight and detector offsets live in a
//   flat 2-d space) and ProjQuat (where requests for coordinates will
//   simply return the composed quaternion, without celestial
//   interpretation).  See Pointer<C> implementation for more details.
//
// - P: PixelSys, determining how computed coordinates (in a
//   projection plane, say) should be turned into pixel indices.  This
//   is the bit that actually needs to know the footprint of a
//   specific map -- the equivalent of CRVAL, CRPIX and CDELT.
//
// - S: SpinSys, determining the number of spin-components and the
//   equations for their weights.  This provides the abstraction for
//   T-only, TQU, and QU-only mapping, by specifying how to compute
//   "pointing matrix" projection factors for each active map
//   component.  The input here is the polarized detector's position
//   angle, gamma, relative to the reference axis of the active
//   projection.  This system can support mapping of other spin fields
//   quite easily; the most obvious candidate being spin-1, to carry
//   information such as in-plane pointing displacements.

// State the template<CoordSys> options
class ProjQuat;
class ProjFlat;
class ProjCEA;
class ProjCAR;
class ProjARC;
class ProjTAN;
class ProjZEA;

// State the template<TilingSys> options
class Tiled;
class NonTiled;

// Helper template for SpinSys...
template <int N>
class SpinClass {
public:
    static const int comp_count = N;
};

// State the template<SpinSys> options
class SpinT : public SpinClass<1> {};
class SpinQU : public SpinClass<2> {};
class SpinTQU : public SpinClass<3> {};

template <typename CoordSys>
class Pointer {
public:
    bool TestInputs(bp::object &map, bp::object &pbore, bp::object &pdet,
                    bp::object &signal, bp::object &det_weights);
    void InitPerDet(int i_det, double *dofs);
    int DetCount() { return n_det; }
    int TimeCount() { return n_time; }
    void GetCoords(int i_det, int i_time, const double *dofs, double *coords);
private:
    BufferWrapper<double> _pborebuf;
    BufferWrapper<double> _pdetbuf;
    int n_det;
    int n_time;
};

template <typename CoordSys>
bool Pointer<CoordSys>::TestInputs(
    bp::object &map, bp::object &pbore, bp::object &pdet,
    bp::object &signal, bp::object &det_weights)
{
    // Boresight and Detector must present and inter-compatible.
    _pborebuf = BufferWrapper<double>("boresight", pbore, false,
                                      vector<int>{-1, 4});
    _pdetbuf = BufferWrapper<double>("detectors", pdet, false,
                                     vector<int>{-1, 4});
    n_time = _pborebuf->shape[0];
    n_det = _pdetbuf->shape[0];
    return true;
}

template <typename CoordSys>
inline
void Pointer<CoordSys>::InitPerDet(int i_det, double *dofs)
{
    const char *det = (char*)_pdetbuf->buf
        + _pdetbuf->strides[0] * i_det;
    for (int ic = 0; ic < 4; ++ic)
        dofs[ic] = *(double*)(det + _pdetbuf->strides[1] * ic);
}

// ProjQuat: Not a projection -- returns the quaternion rotation
// components.

template <>
inline
void Pointer<ProjQuat>::GetCoords(int i_det, int i_time,
                                  const double *dofs, double *coords)
{
    double _qbore[4];
    for (int ic=0; ic<4; ic++)
        _qbore[ic] = *(double*)((char*)_pborebuf->buf +
                            _pborebuf->strides[0] * i_time +
                            _pborebuf->strides[1] * ic);

    // What could possibly go wrong.
    const quatd *qbore = reinterpret_cast<const quatd*>(_qbore);
    const quatd *qofs = reinterpret_cast<const quatd*>(dofs);
    quatd *qdet = reinterpret_cast<quatd*>(coords);
    *qdet =(*qbore) * (*qofs);
}

// ProjFlat: Not a spherical projection -- assumes flat space (as in
// FITS X,Y type coordinates).

template <>
inline
void Pointer<ProjFlat>::GetCoords(int i_det, int i_time,
                                  const double *dofs, double *coords)
{
    for (int ic=0; ic<4; ic++)
        coords[ic] = *(double*)((char*)_pborebuf->buf +
                                _pborebuf->strides[0] * i_time +
                                _pborebuf->strides[1] * ic);
    coords[0] += dofs[0];
    coords[1] += dofs[1];
    const double coords_2_ = coords[2];
    coords[2] = coords[2] * dofs[2] - coords[3] * dofs[3];
    coords[3] = coords[3] * dofs[2] + coords_2_ * dofs[3];
}

// ProjARC: the zenithal equidistant projection.
//
// The first two coordinates are R(lat)*sin(lon) and -R(lat)*cos(lon)
// [see FITS-II].  Then cos and sin of parallactic angle.  For ARC,
// R(lat) = 90 - lat = theta.

template <>
inline
void Pointer<ProjARC>::GetCoords(int i_det, int i_time,
                                 const double *dofs, double *coords)
{
    double _qbore[4];
    for (int ic=0; ic<4; ic++)
        _qbore[ic] = *(double*)((char*)_pborebuf->buf +
                            _pborebuf->strides[0] * i_time +
                            _pborebuf->strides[1] * ic);

    // What could possibly go wrong.
    const quatd *qbore = reinterpret_cast<const quatd*>(_qbore);
    const quatd *qofs = reinterpret_cast<const quatd*>(dofs);
    quatd qdet = (*qbore) * (*qofs);

    const double a = qdet.R_component_1();
    const double b = qdet.R_component_2();
    const double c = qdet.R_component_3();
    const double d = qdet.R_component_4();

    const double cos_theta2_sq = a*a + d*d;

    const double sc = c*a + d*b;
    const double ss = a*b - c*d;
    const double half_sin_theta = sqrt(sc*sc + ss*ss);
    double R_factor;
    if (half_sin_theta < 1e-8)
        R_factor = 2 + 1.33333333333*half_sin_theta*half_sin_theta;
    else
        R_factor = ASIN(half_sin_theta*2) / half_sin_theta;

    coords[0] = ss * R_factor;
    coords[1] = sc * R_factor;
    coords[2] = (a*a - d*d) / cos_theta2_sq;
    coords[3] = (2*a*d) / cos_theta2_sq;
}

// ProjTAN: the tangent plane (gnomonic) projection.  It is zenithal.
//
// The first two coordinates are R(lat)*sin(lon) and -R(lat)*cos(lon)
// [see FITS-II].  Then cos and sin of parallactic angle.  For TAN,
// R(lat) = tan(lat) = cot(theta).

template <>
inline
void Pointer<ProjTAN>::GetCoords(int i_det, int i_time,
                                 const double *dofs, double *coords)
{
    double _qbore[4];
    for (int ic=0; ic<4; ic++)
        _qbore[ic] = *(double*)((char*)_pborebuf->buf +
                            _pborebuf->strides[0] * i_time +
                            _pborebuf->strides[1] * ic);

    // What could possibly go wrong.
    const quatd *qbore = reinterpret_cast<const quatd*>(_qbore);
    const quatd *qofs = reinterpret_cast<const quatd*>(dofs);
    quatd qdet = (*qbore) * (*qofs);

    const double a = qdet.R_component_1();
    const double b = qdet.R_component_2();
    const double c = qdet.R_component_3();
    const double d = qdet.R_component_4();

    const double cos_theta2_sq = a*a + d*d;
    const double cos_theta = 2*cos_theta2_sq - 1;

    coords[0] = (a*b - c*d) * 2 / cos_theta;
    coords[1] = (a*c + b*d) * 2 / cos_theta;
    coords[2] = (a*a - d*d) / cos_theta2_sq;
    coords[3] = (2*a*d) / cos_theta2_sq;
}

// ProjZEA: the zenithal equal area projection.
//
// The first two coordinates are R(lat)*sin(lon) and -R(lat)*cos(lon)
// [see FITS-II].  Then cos and sin of parallactic angle.  For ZEA,
// R(lat) = sqrt(2(1-cos(lat))) = sqrt(2(1-sin(theta))).

template <>
inline
void Pointer<ProjZEA>::GetCoords(int i_det, int i_time,
                                 const double *dofs, double *coords)
{
    double _qbore[4];
    for (int ic=0; ic<4; ic++)
        _qbore[ic] = *(double*)((char*)_pborebuf->buf +
                            _pborebuf->strides[0] * i_time +
                            _pborebuf->strides[1] * ic);

    // What could possibly go wrong.
    const quatd *qbore = reinterpret_cast<const quatd*>(_qbore);
    const quatd *qofs = reinterpret_cast<const quatd*>(dofs);
    quatd qdet = (*qbore) * (*qofs);

    const double a = qdet.R_component_1();
    const double b = qdet.R_component_2();
    const double c = qdet.R_component_3();
    const double d = qdet.R_component_4();

    const double cos_theta2_sq = a*a + d*d;
    const double cos_theta = 2*cos_theta2_sq - 1;
    const double cos_theta2 = sqrt(cos_theta2_sq);

    coords[0] = (a*b - c*d) * 2 / cos_theta2;
    coords[1] = (a*c + b*d) * 2 / cos_theta2;
    coords[2] = (a*a - d*d) / cos_theta2_sq;
    coords[3] = (2*a*d) / cos_theta2_sq;
}

// ProjCEA: Cylindrical projection.
//
// First two coordinates are lon (in radians) and sin(lat).  Then cos
// and sin of parallactic angle.

template <>
inline
void Pointer<ProjCEA>::GetCoords(int i_det, int i_time,
                                 const double *dofs, double *coords)
{
    double _qbore[4];
    for (int ic=0; ic<4; ic++)
        _qbore[ic] = *(double*)((char*)_pborebuf->buf +
                            _pborebuf->strides[0] * i_time +
                            _pborebuf->strides[1] * ic);

    // What could possibly go wrong.
    const quatd *qbore = reinterpret_cast<const quatd*>(_qbore);
    const quatd *qofs = reinterpret_cast<const quatd*>(dofs);
    quatd qdet = (*qbore) * (*qofs);

    const double a = qdet.R_component_1();
    const double b = qdet.R_component_2();
    const double c = qdet.R_component_3();
    const double d = qdet.R_component_4();

    const double cos_theta = a*a - b*b - c*c + d*d;
    const double half_sin_theta = 0.5 * sqrt(1 - cos_theta*cos_theta);

    coords[0] = atan2(c*d - a*b, c*a + d*b);
    coords[1] = cos_theta; // Yes, cos(theta) = sin(lat).
    coords[2] = (a*c - b*d) / half_sin_theta;
    coords[3] = (c*d + a*b) / half_sin_theta;
}

// ProjCAR: Cylindrical projection.
//
// First two coordinates are lon and lat (in radians).  Then cos and
// sin of parallactic angle.

template <>
inline
void Pointer<ProjCAR>::GetCoords(int i_det, int i_time,
                                      const double *dofs, double *coords)
{
    double _qbore[4];
    for (int ic=0; ic<4; ic++)
        _qbore[ic] = *(double*)((char*)_pborebuf->buf +
                            _pborebuf->strides[0] * i_time +
                            _pborebuf->strides[1] * ic);

    // What could possibly go wrong.
    const quatd *qbore = reinterpret_cast<const quatd*>(_qbore);
    const quatd *qofs = reinterpret_cast<const quatd*>(dofs);
    quatd qdet = (*qbore) * (*qofs);

    const double a = qdet.R_component_1();
    const double b = qdet.R_component_2();
    const double c = qdet.R_component_3();
    const double d = qdet.R_component_4();

    const double cos_theta = a*a - b*b - c*c + d*d;
    const double half_sin_theta = 0.5 * sqrt(1 - cos_theta*cos_theta);

    coords[0] = ATAN2(c*d - a*b, c*a + d*b);
    coords[1] = ASIN(cos_theta);   // Yes, cos(theta) = sin(lat).
    coords[2] = (a*c - b*d) / half_sin_theta;
    coords[3] = (c*d + a*b) / half_sin_theta;
}

//! Pixelizors
//
template <typename TilingSys>
class Pixelizor2_Flat;

template <>
class Pixelizor2_Flat<NonTiled> {
public:
    static const int index_count = 2;
    Pixelizor2_Flat(int ny, int nx,
                    double dy=1., double dx=1.,
                    double iy0=0., double ix0=0.) {
        naxis[0] = ny;
        naxis[1] = nx;
        cdelt[0] = dy;
        cdelt[1] = dx;
        crpix[0] = iy0;
        crpix[1] = ix0;
    };
    Pixelizor2_Flat() : naxis{1,1} {};
    Pixelizor2_Flat(bp::object args) {
        bp::tuple args_tuple = bp::extract<bp::tuple>(args);
        naxis[0] = bp::extract<int>(args_tuple[0])();
        naxis[1] = bp::extract<int>(args_tuple[1])();
        cdelt[0] = bp::extract<double>(args_tuple[2])();
        cdelt[1] = bp::extract<double>(args_tuple[3])();
        crpix[0] = bp::extract<double>(args_tuple[4])();
        crpix[1] = bp::extract<double>(args_tuple[5])();
    }
    ~Pixelizor2_Flat() {};

    bp::object zeros(vector<int> shape) {
        shape.push_back(naxis[0]);
        shape.push_back(naxis[1]);
        int ndim = 0;
        npy_intp dims[32];
        for (auto d: shape)
            dims[ndim++] = d;

        int dtype = NPY_FLOAT64;
        PyObject *v = PyArray_ZEROS(ndim, dims, dtype, 0);
        return bp::object(bp::handle<>(v));
    }

    inline
    void GetPixel(int i_det, int i_time, const double *coords, int *pixel_index) {
        //double ix = coords[0] / cdelt[1] + crpix[1] + 0.5;
        double ix = coords[0] / cdelt[1] + crpix[1] - 1 + 0.5;
        if (ix < 0 || ix >= naxis[1]) {
            *pixel_index = -1;
            return;
        }

        //double iy = coords[1] / cdelt[0] + crpix[0] + 0.5;
        double iy = coords[1] / cdelt[0] + crpix[0] - 1 + 0.5;
        if (iy < 0 || iy >= naxis[0]) {
            *pixel_index = -1;
            return;
        }

        pixel_index[0] = int(iy);
        pixel_index[1] = int(ix);
    }

    bool TestInputs(bp::object &map, bool need_map, bool need_weight_map, int comp_count) {
        if (need_map) {
            // The map is mandatory, and the leading axis must match the
            // component count.  It can have 1+ other dimensions.
            mapbuf = BufferWrapper<double>("map", map, false,
                                           vector<int>{comp_count,-1,-3});
        } else if (need_weight_map) {
            // The map is mandatory, and the two leading axes must match
            // the component count.  It can have 1+ other dimensions.
            mapbuf = BufferWrapper<double>("map", map, false,
                                           vector<int>{comp_count,comp_count,-1,-3});
        }
        return true;
    }
    double *pix(int imap, const int pixel_index[]) {
        return (double*)((char*)mapbuf->buf +
                         mapbuf->strides[0]*imap +
                         mapbuf->strides[1]*pixel_index[0] +
                         mapbuf->strides[2]*pixel_index[1]);
    }
    double *wpix(int imap, int jmap, const int pixel_index[]) {
        return (double*)((char*)mapbuf->buf +
                         mapbuf->strides[0]*imap +
                         mapbuf->strides[1]*jmap +
                         mapbuf->strides[2]*pixel_index[0] +
                         mapbuf->strides[3]*pixel_index[1]);
    }
    int stripe(const int pixel_index[], int thread_count) {
        if (pixel_index[0] < 0)
            return -1;
        return pixel_index[1] * thread_count / naxis[1];
    }

    int crpix[2];
    double cdelt[2];
    int naxis[2];
    BufferWrapper<double> mapbuf;
};

template <>
class Pixelizor2_Flat<Tiled> {
public:
    static const int index_count = 3;
    Pixelizor2_Flat() {};
    Pixelizor2_Flat(int ny, int nx, double dy, double dx,
                    double iy0, double ix0, int tiley, int tilex) {
        parent_pix = Pixelizor2_Flat<NonTiled>(ny, nx, dy, dx, iy0, ix0);
        tile_shape[0] = tiley;
        tile_shape[1] = tilex;
    };
    Pixelizor2_Flat(Pixelizor2_Flat<NonTiled> _parent_pix, int tiley, int tilex) {
        parent_pix = _parent_pix;
        tile_shape[0] = tiley;
        tile_shape[1] = tilex;
    };
    Pixelizor2_Flat(bp::object args) {
        parent_pix = Pixelizor2_Flat<NonTiled>(args); // first 6...
        bp::tuple args_tuple = bp::extract<bp::tuple>(args);

        tile_shape[0] = bp::extract<int>(args_tuple[6])();
        tile_shape[1] = bp::extract<int>(args_tuple[7])();
        // Check for tile list as arg[8].
        if (bp::len(args) >= 9) {
            int n_ty = (parent_pix.naxis[0] + tile_shape[0] - 1) / tile_shape[0];
            int n_tx = (parent_pix.naxis[1] + tile_shape[1] - 1) / tile_shape[1];
            populate = vector<bool>(n_ty * n_tx, false);
            bp::object active_tiles = bp::extract<bp::object>(args_tuple[8])();
            for (int i=0; i<bp::len(active_tiles); i++) {
                // We're using C API instead of boost because this:
                //   bp::extract<int>(active_tiles[i])
                // does not work on an ndarray, nor even on a list of
                // elements extracted from an array, unless one carefully
                // casts them to int first.
                int idx = PyLong_AsLong(bp::object(active_tiles[i]).ptr());
                if (idx >= 0 && idx < n_tx*n_ty)
                    populate[idx] = true;
            }
        }
    }
    ~Pixelizor2_Flat() {};

    bp::object zeros(vector<int> shape) {
        int dtype = NPY_FLOAT64;
        int ndim = 0;
        npy_intp dims[32];
        for (auto d: shape)
            dims[ndim++] = d;
        ndim += 2;

        int n_ty = (parent_pix.naxis[0] + tile_shape[0] - 1) / tile_shape[0];
        int n_tx = (parent_pix.naxis[1] + tile_shape[1] - 1) / tile_shape[1];

        if (populate.size() == 0)
            throw shape_exception("zeros", "Cannot create blank tiled map unless "
                                  "user has specified what tiles to populate.");

        bp::list maps_out;
        auto pop_iter = populate.begin();
        for (int i_ty = 0; i_ty < n_ty; i_ty++) {
            for (int i_tx = 0; i_tx < n_tx; i_tx++) {
                bool pop_this = (pop_iter != populate.end()) && *(pop_iter++);
                if (pop_this) {
                    dims[ndim-2] = min(tile_shape[0], parent_pix.naxis[0] - i_ty * tile_shape[0]);
                    dims[ndim-1] = min(tile_shape[1], parent_pix.naxis[1] - i_tx * tile_shape[1]);
                    PyObject *v = PyArray_ZEROS(ndim, dims, dtype, 0);
                    maps_out.append(bp::handle<>(v));
                } else
                    maps_out.append(bp::object());
            }
        }
        return maps_out;
    }

    inline
    void GetPixel(int i_det, int i_time, const double *coords, int *pixel_index) {
        //double ix = coords[0] / parent_pix.cdelt[1] + parent_pix.crpix[1] + 0.5;
        double ix = coords[0] / parent_pix.cdelt[1] + parent_pix.crpix[1] - 1 + 0.5;
        if (ix < 0 || ix >= parent_pix.naxis[1]) {
            pixel_index[0] = -1;
            return;
        }

        //double iy = coords[1] / parent_pix.cdelt[0] + parent_pix.crpix[0] + 0.5;
        double iy = coords[1] / parent_pix.cdelt[0] + parent_pix.crpix[0] - 1 + 0.5;
        if (iy < 0 || iy >= parent_pix.naxis[0]) {
            *pixel_index = -1;
            return;
        }

        int sub_y = int(iy) / tile_shape[0];
        int sub_x = int(ix) / tile_shape[1];
        pixel_index[0] = sub_y * ((parent_pix.naxis[1] + tile_shape[1] - 1) / tile_shape[1]) + sub_x;
        //pixel_index[1] = strides[0]*int(iy) + strides[1]*int(ix);
        pixel_index[1] = int(iy) - sub_y * tile_shape[0];
        pixel_index[2] = int(ix) - sub_x * tile_shape[1];
    }

    bool TestInputs(bp::object &map, bool need_map, bool need_weight_map, int comp_count) {
        vector<int> map_shape_req;
        if (need_map) {
            // The map is mandatory, and the leading axis must match the
            // component count.  It can have 1+ other dimensions.
            map_shape_req = {comp_count,-1,-3};
        } else if (need_weight_map) {
            // The map is mandatory, and the two leading axes must match
            // the component count.  It can have 1+ other dimensions.
            map_shape_req = {comp_count,comp_count,-1,-3};
        }
        if (map_shape_req.size() == 0)
            return true;

        mapbufs.clear();
        for (int i_tile = 0; i_tile < bp::len(map); i_tile++) {
            if (isNone(map[i_tile])) {
                mapbufs.push_back(BufferWrapper<double>());
            } else {
                // You should be checking that presence and size matches expectation.
                mapbufs.push_back(
                    BufferWrapper<double>("map", map[i_tile], false, map_shape_req));
            }
        }

        return true;
    }
    double *pix(int imap, const int pixel_index[]) {
        const BufferWrapper<double> &mapbuf = mapbufs[pixel_index[0]];
        // This assertion is needed in case the user did not populate
        // the right set of tiles.
        assert(mapbuf->buf != nullptr);
        return (double*)((char*)mapbuf->buf +
                         mapbuf->strides[0]*imap +
                         mapbuf->strides[1]*pixel_index[1] +
                         mapbuf->strides[2]*pixel_index[2]);
    }
    double *wpix(int imap, int jmap, const int pixel_index[]) {
        // Expensive shared_ptr copy?
        const BufferWrapper<double> &mapbuf = mapbufs[pixel_index[0]];
        // This assertion is needed in case the user did not populate
        // the right set of tiles.
        assert(mapbuf->buf != nullptr);
        return (double*)((char*)mapbuf->buf +
                         mapbuf->strides[0]*imap +
                         mapbuf->strides[1]*jmap +
                         mapbuf->strides[2]*pixel_index[1] +
                         mapbuf->strides[3]*pixel_index[2]);
    }
    int stripe(const int pixel_index[], int thread_count) {
        return pixel_index[0] % thread_count;
    }

    Pixelizor2_Flat<NonTiled> parent_pix;
    vector<bool> populate;
    int tile_shape[2];

    vector<BufferWrapper<double>> mapbufs;
};

template <typename SpinSys>
static inline
void spin_proj_factors(const double* coords, FSIGNAL *projfacs);

template <>
inline void spin_proj_factors<SpinT>(const double* coords, FSIGNAL *projfacs)
{
    projfacs[0] = 1.;
}

template <>
inline
void spin_proj_factors<SpinQU>(const double* coords, FSIGNAL *projfacs)
{
    const double c = coords[2];
    const double s = coords[3];
    projfacs[0] = c*c - s*s;
    projfacs[1] = 2*c*s;
}

template <>
inline
void spin_proj_factors<SpinTQU>(const double* coords, FSIGNAL *projfacs)
{
     const double c = coords[2];
     const double s = coords[3];
     projfacs[0] = 1.;
     projfacs[1] = c*c - s*s;
     projfacs[2] = 2*c*s;
}

template <typename DTYPE>
bool SignalSpace<DTYPE>::_Validate(bp::object input, std::string var_name,
                                   int dtype)
{
    // We want a list of arrays here.
    bp::list sig_list;
    auto list_extractor = bp::extract<bp::list>(input);
    if (isNone(input)) {
        npy_intp _dims[dims.size()];
        for (int d=0; d<dims.size(); ++d) {
            if (dims[d] < 0)
                throw shape_exception(var_name, "Cannot create space with wildcard dimensons.");
            _dims[d] = dims[d];
        }
        for (int i=0; i<dims[0]; ++i) {
            PyObject *v = PyArray_ZEROS(dims.size()-1, _dims+1, dtype, 0);
            sig_list.append(bp::object(bp::handle<>(v)));
        }
    } else if (list_extractor.check()) {
        sig_list = list_extractor();
    } else {
        // Probably an array... listify it.
        for (int i=0; i<bp::len(input); ++i)
            sig_list.append(input[i]);
    }
    ret_val = sig_list;

    if (dims[0] == -1) {
        dims[0] = bp::len(sig_list);
        if (dims[0] == 0)
            throw shape_exception(var_name, "has not been tested on shape 0 objects");
    } else if (bp::len(sig_list) != dims[0])
        throw shape_exception(var_name, "must contain (n_det) vectors");

    const int n_det = dims[0];

    // Now extract each list member into our vector of BufferWrappers..
    data_ptr = (DTYPE**)calloc(n_det, sizeof(*data_ptr));

    bw.reserve(n_det);

    // Copy dims[1:] into sub_dims; potentially update sub_dims during
    // extraction, then copy those back into dims.
    vector<int> sub_dims(dims.begin()+1, dims.end());

    for (int i=0; i<n_det; i++) {
        bp::object item = bp::extract<bp::object>(sig_list[i])();
        bw.push_back(BufferWrapper<DTYPE>(var_name, item, false, sub_dims));
        if (i == 0) {
            sub_dims.clear();
            for (int d=0; d<bw[0]->ndim; d++)
                sub_dims.push_back(bw[0]->shape[d]);
        } else {
            for (int d=0; d<sub_dims.size(); ++d) {
                if (bw[i]->strides[d] != bw[0]->strides[d])
                    throw shape_exception(var_name, "[all elements must have same stride]");
            }
        }
        data_ptr[i] = (DTYPE*)bw[i]->buf;
    }

    // Store the step in units of the itemsize; update dims from sub_dims.
    for (int d=0; d<dims.size()-1; d++) {
        dims[d+1] = sub_dims[d];
        if (bw[0]->strides[d] % bw[0]->itemsize != 0)
            throw shape_exception(var_name, "stride is non-integral; realign.");
        steps[d] = bw[0]->strides[d] / bw[0]->itemsize;
    }
    return true;
}

template <typename DTYPE>
SignalSpace<DTYPE>::SignalSpace(
    bp::object input, std::string var_name, int dtype, int n_det, int n_time)
{
    dims = {n_det, n_time};
    _Validate(input, var_name, dtype);
}

template <typename DTYPE>
SignalSpace<DTYPE>::SignalSpace(
    bp::object input, std::string var_name, int dtype, int n_det, int n_time,
    int n_thirdaxis)
{
    dims = {n_det, n_time, n_thirdaxis};
    _Validate(input, var_name, dtype);
}

template<typename C, typename P, typename S>
ProjectionEngine<C,P,S>::ProjectionEngine(bp::object pix_args)
{
    _pixelizor = P(pix_args);
}

template<typename C, typename P, typename S>
int ProjectionEngine<C,P,S>::index_count() const {
    return P::index_count;
}

template<typename C, typename P, typename S>
int ProjectionEngine<C,P,S>::comp_count() const {
    return S::comp_count;
}

template<typename C, typename P, typename S>
bp::object ProjectionEngine<C,P,S>::coords(
    bp::object pbore, bp::object pofs, bp::object coord)
{
    auto _none = bp::object();
    auto pointer = Pointer<C>();
    pointer.TestInputs(_none, pbore, pofs, _none, _none);

    int n_det = pointer.DetCount();
    int n_time = pointer.TimeCount();

    const int n_coord = 4;
    auto coord_buf_man = SignalSpace<double>(
        coord, "coord", NPY_FLOAT64, n_det, n_time, n_coord);

#pragma omp parallel for
    for (int i_det = 0; i_det < n_det; ++i_det) {
        double dofs[4];
        pointer.InitPerDet(i_det, dofs);

        double* const coords_det = coord_buf_man.data_ptr[i_det];
        const int step0 = coord_buf_man.steps[0];
        const int step1 = coord_buf_man.steps[1];

        for (int i_time = 0; i_time < n_time; ++i_time) {
            double coords[4];
            pointer.GetCoords(i_det, i_time, (double*)dofs, (double*)coords);
            for (int ic=0; ic<4; ic++)
                *(coords_det + step0 * i_time + step1 * ic) = coords[ic];
        }
    }

    return coord_buf_man.ret_val;
}

template<typename C, typename P, typename S>
bp::object ProjectionEngine<C,P,S>::pixels(
    bp::object pbore, bp::object pofs, bp::object pixel)
{
    auto _none = bp::object();

    auto pointer = Pointer<C>();
    pointer.TestInputs(_none, pbore, pofs, _none, _none);
    //_pixelizor.TestInputs(_none, _none, _none, _none, _none);

    int n_det = pointer.DetCount();
    int n_time = pointer.TimeCount();

    auto pixel_buf_man = SignalSpace<int32_t>(
        pixel, "pixel", NPY_INT32, n_det, n_time, P::index_count);

#pragma omp parallel for
    for (int i_det = 0; i_det < n_det; ++i_det) {
        double dofs[4];
        pointer.InitPerDet(i_det, dofs);
        int* const pix_buf = pixel_buf_man.data_ptr[i_det];
        int pixel_offset[P::index_count] = {-1};
        for (int i_time = 0; i_time < n_time; ++i_time) {
            double coords[4];
            pointer.GetCoords(i_det, i_time, (double*)dofs, (double*)coords);
            _pixelizor.GetPixel(i_det, i_time, (double*)coords, pixel_offset);
            for (int i_dim = 0; i_dim < P::index_count; i_dim++)
                pix_buf[i_time * pixel_buf_man.steps[0] +
                        i_dim * pixel_buf_man.steps[1]] = pixel_offset[i_dim];
            // pix_buf[i_time * pixel_buf_man.steps[0]] = pixel_offset;
        }
    }

    return pixel_buf_man.ret_val;
}

template<typename C, typename P, typename S>
bp::object ProjectionEngine<C,P,S>::pointing_matrix(
    bp::object pbore, bp::object pofs, bp::object pixel, bp::object proj)
{
    auto _none = bp::object();

    auto pointer = Pointer<C>();
    pointer.TestInputs(_none, pbore, pofs, _none, _none);
    //_pixelizor.TestInputs(_none, _none, _none, _none, _none);

    int n_det = pointer.DetCount();
    int n_time = pointer.TimeCount();

    auto pixel_buf_man = SignalSpace<int32_t>(
        pixel, "pixel", NPY_INT32, n_det, n_time, P::index_count);

    auto proj_buf_man = SignalSpace<FSIGNAL>(
        proj, "proj", FSIGNAL_NPY_TYPE, n_det, n_time, S::comp_count);

#pragma omp parallel for
    for (int i_det = 0; i_det < n_det; ++i_det) {
        double dofs[4];
        pointer.InitPerDet(i_det, dofs);
        int* const pix_buf = pixel_buf_man.data_ptr[i_det];
        FSIGNAL* const proj_buf = proj_buf_man.data_ptr[i_det];
        const int step = pixel_buf_man.steps[0];
        int pixel_offset[P::index_count] = {-1};
        for (int i_time = 0; i_time < n_time; ++i_time) {
            double coords[4];
            FSIGNAL pf[S::comp_count];
            pointer.GetCoords(i_det, i_time, (double*)dofs, (double*)coords);
            _pixelizor.GetPixel(i_det, i_time, (double*)coords, pixel_offset);
            spin_proj_factors<S>(coords, pf);

            for (int i_dim = 0; i_dim < P::index_count; i_dim++)
                pix_buf[i_time * pixel_buf_man.steps[0] +
                        i_dim * pixel_buf_man.steps[1]] = pixel_offset[i_dim];
            for (int i_spin = 0; i_spin < S::comp_count; i_spin++)
                proj_buf[i_time * proj_buf_man.steps[0] +
                         i_spin * proj_buf_man.steps[1]] = pf[i_spin];
        }
    }

    return bp::make_tuple(pixel_buf_man.ret_val,
                          proj_buf_man.ret_val);
}

template<typename C, typename P, typename S>
bp::object ProjectionEngine<C,P,S>::pixel_ranges(
    bp::object pbore, bp::object pofs, bp::object map)
{
    auto _none = bp::object();

    auto pointer = Pointer<C>();
    pointer.TestInputs(map, pbore, pofs, _none, _none);
    int n_det = pointer.DetCount();
    int n_time = pointer.TimeCount();

    bool from_map = !isNone(map);
    if (from_map)
        _pixelizor.TestInputs(map, true, false, S::comp_count);

    vector<vector<RangesInt32>> ranges;

#pragma omp parallel
    {
        int n_domain = omp_get_num_threads();
#pragma omp single
        {
            for (int i=0; i<n_domain; ++i) {
                vector<RangesInt32> v(n_det);
                for (auto &_v: v)
                    _v.count = n_time;
                ranges.push_back(v);
            }
        }

#pragma omp for
        for (int i_det = 0; i_det < n_det; ++i_det) {
            double dofs[4];
            pointer.InitPerDet(i_det, dofs);
            int last_slice = -1;
            int slice_start = 0;
            int pixel_offset[P::index_count] = {-1};
            for (int i_time = 0; i_time < n_time; ++i_time) {
                double coords[4];
                pointer.GetCoords(i_det, i_time, (double*)dofs, (double*)coords);
                _pixelizor.GetPixel(i_det, i_time, (double*)coords, pixel_offset);

                int this_slice = -1;
                if (from_map) {
                    if (pixel_offset[0] >= 0)
                        this_slice = *_pixelizor.pix(0, pixel_offset);
                } else {
                    this_slice = _pixelizor.stripe(pixel_offset, n_domain);
                }
                if (this_slice != last_slice) {
                    if (last_slice >= 0)
                        ranges[last_slice][i_det].append_interval_no_check(
                            slice_start, i_time);
                    slice_start = i_time;
                    last_slice = this_slice;
                }
            }
            if (last_slice >= 0)
                ranges[last_slice][i_det].append_interval_no_check(
                    slice_start, n_time);
        }
    }

    // Convert super vector to a list and return
    auto ivals_out = bp::list();
    for (int j=0; j<ranges.size(); j++) {
        auto ivals = bp::list();
        for (int i_det=0; i_det<n_det; i_det++) {
            auto iv = ranges[j][i_det];
            ivals.append(bp::object(iv));
        }
        ivals_out.append(bp::extract<bp::object>(ivals)());
    }
    return bp::extract<bp::object>(ivals_out);
}

template<typename C, typename P, typename S>
bp::object ProjectionEngine<C,P,S>::zeros(bp::object shape)
{
    vector<int> dims;
    bp::extract<int> int_ex(shape);
    if (int_ex.check()) {
        dims.push_back(int_ex());
        return _pixelizor.zeros(dims);
    }

    bp::extract<bp::tuple> tuple_ex(shape);
    if (tuple_ex.check()) {
        auto tuple = tuple_ex();
        for (int i=0; i<bp::len(tuple); i++)
            dims.push_back(bp::extract<int>(tuple[i])());
        return _pixelizor.zeros(dims);
    }

    return bp::object();  //None on fall-through
}

template<typename C, typename P, typename S>
bp::object ProjectionEngine<C,P,S>::from_map(
    bp::object map, bp::object pbore, bp::object pofs, bp::object signal)
{
    auto _none = bp::object();

    // Initialize pointer and _pixelizor.
    auto pointer = Pointer<C>();
    pointer.TestInputs(map, pbore, pofs, signal, _none);
    int n_det = pointer.DetCount();
    int n_time = pointer.TimeCount();

    // Confirm that map has the right meta-shape.
    _pixelizor.TestInputs(map, true, false, S::comp_count);

    // Get pointers to the signal and (optional) per-det weights.
    auto _signalspace = SignalSpace<FSIGNAL>(
            signal, "signal", FSIGNAL_NPY_TYPE, n_det, n_time);

    #pragma omp parallel for
    for (int i_det = 0; i_det < n_det; ++i_det) {
        double dofs[4];
        pointer.InitPerDet(i_det, dofs);
        int pixel_offset[P::index_count] = {-1};
        FSIGNAL pf[S::comp_count];
        for (int i_time = 0; i_time < n_time; ++i_time) {
            double coords[4];
            pointer.GetCoords(i_det, i_time, (double*)dofs, (double*)coords);
            _pixelizor.GetPixel(i_det, i_time, (double*)coords, pixel_offset);
            if (pixel_offset[0] < 0) continue;
            spin_proj_factors<S>(coords, pf);
            FSIGNAL *sig = (_signalspace.data_ptr[i_det] +
                            _signalspace.steps[0]*i_time);
            for (int imap=0; imap<S::comp_count; ++imap)
                *sig += *_pixelizor.pix(imap, pixel_offset) * pf[imap];
        }
    }

    return _signalspace.ret_val;
}

template<typename C, typename P, typename S>
static
void to_map_single_thread(Pointer<C> &pointer,
                          P &_pixelizor,
                          vector<RangesInt32> ivals,
                          BufferWrapper<FSIGNAL> &_det_weights,
                          SignalSpace<FSIGNAL> *_signalspace)
{
    int n_det = pointer.DetCount();
    for (int i_det = 0; i_det < n_det; ++i_det) {
        FSIGNAL det_wt = 1.;
        if (_det_weights->obj != NULL)
            det_wt = *(FSIGNAL*)((char*)_det_weights->buf +
                                 _det_weights->strides[0]*i_det);
        double dofs[4];
        double coords[4];
        int pixel_offset[P::index_count] = {-1};
        FSIGNAL pf[S::comp_count];
        pointer.InitPerDet(i_det, dofs);
        for (auto const &rng: ivals[i_det].segments) {
            for (int i_time = rng.first; i_time < rng.second; ++i_time) {
                pointer.GetCoords(i_det, i_time, (double*)dofs, (double*)coords);
                _pixelizor.GetPixel(i_det, i_time, (double*)coords, pixel_offset);
                if (pixel_offset[0] < 0) continue;
                const FSIGNAL sig = *(_signalspace->data_ptr[i_det] +
                                      _signalspace->steps[0]*i_time);
                spin_proj_factors<S>(coords, pf);
                for (int imap=0; imap<S::comp_count; ++imap)
                    *_pixelizor.pix(imap, pixel_offset) += sig * pf[imap] * det_wt;
            }
        }
    }
}

template<typename C, typename P, typename S>
static
void to_weight_map_single_thread(Pointer<C> &pointer,
                                 P &_pixelizor,
                                 vector<RangesInt32> ivals,
                                 BufferWrapper<FSIGNAL> &_det_weights)
{
    int n_det = pointer.DetCount();
    for (int i_det = 0; i_det < n_det; ++i_det) {
        FSIGNAL det_wt = 1.;
        if (_det_weights->obj != NULL)
            det_wt = *(FSIGNAL*)((char*)_det_weights->buf +
                                 _det_weights->strides[0]*i_det);
        double dofs[4];
        double coords[4];
        int pixel_offset[P::index_count] = {-1};
        FSIGNAL pf[S::comp_count];
        pointer.InitPerDet(i_det, dofs);
        for (auto const &rng: ivals[i_det].segments) {
            for (int i_time = rng.first; i_time < rng.second; ++i_time) {
                pointer.GetCoords(i_det, i_time, (double*)dofs, (double*)coords);
                _pixelizor.GetPixel(i_det, i_time, (double*)coords, pixel_offset);
                if (pixel_offset[0] < 0) continue;
                spin_proj_factors<S>(coords, pf);
                for (int imap=0; imap<S::comp_count; ++imap)
                    for (int jmap=imap; jmap<S::comp_count; ++jmap)
                        *_pixelizor.wpix(imap, jmap, pixel_offset) += pf[imap] * pf[jmap] * det_wt;
            }
        }
    }
}

static
vector<vector<RangesInt32>> derive_ranges(bp::object thread_intervals,
                                          int n_det, int n_time)
{
    // The first index of the returned object should correspond to
    // (OMP) execution thread; the second index is over detectors.
    // The input object can be any of the following:
    //
    // - None: Result will be full coverage for all detectors in a
    //   single thread.
    // - List of Ranges: This will be promoted to a single thread.
    // - List containing N lists of Ranges: N threads, with each list
    //   used by corresponding thread.

    vector<vector<RangesInt32>> ivals;
    if (isNone(thread_intervals)) {
        auto r = RangesInt32(n_time).add_interval(0, n_time);
        vector<RangesInt32> v(n_det, r);
        ivals.push_back(v);
    } else {
        // Don't convert to a list... just use indexing.
        if (bp::extract<RangesInt32>(thread_intervals[0]).check()) {
            ivals.push_back(extract_ranges<int32_t>(thread_intervals));
        } else {
            // Assumed it is list of lists of ranges.
            for (int i=0; i<bp::len(thread_intervals); i++)
                ivals.push_back(extract_ranges<int32_t>(thread_intervals[i]));
        }
    }
    return ivals;
}

template<typename C, typename P, typename S>
bp::object ProjectionEngine<C,P,S>::to_map(
    bp::object map, bp::object pbore, bp::object pofs, bp::object signal, bp::object det_weights,
    bp::object thread_intervals)
{
    //Initialize it / check inputs.
    auto pointer = Pointer<C>();
    pointer.TestInputs(map, pbore, pofs, signal, det_weights);
    int n_det = pointer.DetCount();
    int n_time = pointer.TimeCount();

    //Do we need a map?  Now is the time.
    if (isNone(map))
        map = _pixelizor.zeros(vector<int>{S::comp_count});

    // Confirm that map has the right meta-shape.
    _pixelizor.TestInputs(map, true, false, S::comp_count);

    // Get pointers to the signal and (optional) per-det weights.
    auto _signalspace = SignalSpace<FSIGNAL>(
            signal, "signal", FSIGNAL_NPY_TYPE, n_det, n_time);
    auto _det_weights = BufferWrapper<FSIGNAL>(
         "det_weights", det_weights, true, vector<int>{n_det});

    // For multi-threading, the principle here is that all threads
    // loop over all detectors, but the sample ranges encoded in ivals
    // are disjoint.
    auto ivals = derive_ranges(thread_intervals, n_det, n_time);

    if (ivals.size() <= 1) {
        // This block may also be used if OMP is disabled.
        for (int i_thread=0; i_thread < ivals.size(); i_thread++)
            to_map_single_thread<C,P,S>(
                pointer, _pixelizor, ivals[i_thread],_det_weights, &_signalspace);
    } else {
#pragma omp parallel
        {
            // This loop construction handles the (sub-optimal) case
            // that the number of ivals is not equal to the number of
            // OMP threads.
            for (int i_thread=0; i_thread < ivals.size(); i_thread++) {
                if (i_thread % omp_get_num_threads() == omp_get_thread_num())
                    to_map_single_thread<C,P,S>(
                        pointer, _pixelizor, ivals[i_thread], _det_weights, &_signalspace);
            }
        }
    }
    return map;
}

template<typename C, typename P, typename S>
bp::object ProjectionEngine<C,P,S>::to_weight_map(
    bp::object map, bp::object pbore, bp::object pofs, bp::object det_weights,
    bp::object thread_intervals)
{
    auto _none = bp::object();

    //Initialize it / check inputs.
    auto pointer = Pointer<C>();
    pointer.TestInputs(map, pbore, pofs, _none, det_weights);
    int n_det = pointer.DetCount();
    int n_time = pointer.TimeCount();

    //Do we need a map?  Now is the time.
    if (isNone(map))
        map = _pixelizor.zeros(vector<int>{S::comp_count,S::comp_count});

    // Confirm that map has the right meta-shape.
    _pixelizor.TestInputs(map, false, true, S::comp_count);

    // Get pointer to (optional) per-det weights.
    auto _det_weights = BufferWrapper<FSIGNAL>(
         "det_weights", det_weights, true, vector<int>{n_det});

    // For multi-threading, the principle here is that all threads
    // loop over all detectors, but the sample ranges encoded in ivals
    // are disjoint.
    auto ivals = derive_ranges(thread_intervals, n_det, n_time);

    if (ivals.size() <= 1) {
        // This block may also be used if OMP is disabled.
        for (int i_thread=0; i_thread < ivals.size(); i_thread++)
            to_weight_map_single_thread<C,P,S>(
                pointer, _pixelizor, ivals[i_thread],_det_weights);
    } else {
#pragma omp parallel
        {
            // This loop construction handles the (sub-optimal) case
            // that the number of ivals is not equal to the number of
            // OMP threads.
            for (int i_thread=0; i_thread < ivals.size(); i_thread++) {
                if (i_thread % omp_get_num_threads() == omp_get_thread_num())
                    to_weight_map_single_thread<C,P,S>(
                        pointer, _pixelizor, ivals[i_thread], _det_weights);
            }
        }
    }
    return map;
}

// ProjEng_Precomp
//
// The ProjEng_Precomp may be used for very fast projection
// operations, provided you have precomputed pixel index and spin
// projection factors for every sample.  This is agnostic of
// coordinate system, and so not crazily templated.
//

template<typename TilingSys>
class ProjEng_Precomp {
public:
    ProjEng_Precomp() {};
    bp::object from_map(bp::object map, bp::object pixel_index, bp::object spin_proj,
                        bp::object signal);
    bp::object to_map(bp::object map, bp::object pixel_index, bp::object spin_proj,
                      bp::object signal, bp::object weights, bp::object thread_intervals);
    bp::object to_weight_map(bp::object map, bp::object pixel_index, bp::object spin_proj,
                             bp::object weights, bp::object thread_intervals);
};

template<typename TilingSys>
bp::object ProjEng_Precomp<TilingSys>::from_map(
    bp::object map, bp::object pixel_index, bp::object spin_proj,
    bp::object signal)
{
    // You won't get far without pixel_index, so use that to nail down
    // the n_time and n_det.
    auto pixel_buf_man = SignalSpace<int32_t>(
        pixel_index, "pixel_index", NPY_INT32, -1, -1, -1);
    int n_det = pixel_buf_man.dims[0];
    int n_time = pixel_buf_man.dims[1];
    int n_pix = pixel_buf_man.dims[2];

    // Similarly, spin_proj tells you the number of components.
    auto spin_proj_man = SignalSpace<FSIGNAL>(
        spin_proj, "spin_proj", FSIGNAL_NPY_TYPE, n_det, n_time, -1);
    int n_spin = spin_proj_man.dims[2];

    auto tiling = Pixelizor2_Flat<TilingSys>();
    tiling.TestInputs(map, true, false, n_spin);

    // Check the signal.
    auto signal_man = SignalSpace<FSIGNAL>(
        signal, "signal", FSIGNAL_NPY_TYPE, n_det, n_time);

    // Below we assume the pixel sub-indices are close-packed.
    if (pixel_buf_man.steps[1] != 1)
        throw shape_exception("pixel_index",
                              "Fast dimension of pixel indices must be close-packed.");

#pragma omp parallel for
    for (int i_det = 0; i_det < n_det; ++i_det) {
        for (int i_time = 0; i_time < n_time; ++i_time) {
            const int *pixel_ofs = pixel_buf_man.data_ptr[i_det] +
                pixel_buf_man.steps[0]*i_time;
            if (pixel_ofs[0] < 0)
                continue;
            FSIGNAL sig = 0;
            for (int i_spin = 0; i_spin < n_spin; ++i_spin) {
                FSIGNAL pf = *(spin_proj_man.data_ptr[i_det] +
                               spin_proj_man.steps[0] * i_time +
                               spin_proj_man.steps[1] * i_spin);
                sig += pf * *tiling.pix(i_spin, pixel_ofs);
            }
            *(signal_man.data_ptr[i_det] + signal_man.steps[0]*i_time) += sig;
        }
    }

    return signal_man.ret_val;
}

template<typename TilingSys>
static
void precomp_to_map_single_thread(Pixelizor2_Flat<TilingSys> &tiling,
                                  const SignalSpace<int32_t> &pixel_buf_man,
                                  const SignalSpace<FSIGNAL> &spin_proj_man,
                                  vector<RangesInt32> ivals,
                                  BufferWrapper<FSIGNAL> &_det_weights,
                                  SignalSpace<FSIGNAL> *_signalspace)
{
    int n_det = pixel_buf_man.dims[0];
    int n_spin = spin_proj_man.dims[2];
    assert(spin_proj_man.steps[1] == 1); // assumed below

    for (int i_det = 0; i_det < n_det; ++i_det) {
        FSIGNAL det_wt = 1.;
        if (_det_weights->obj != NULL)
            det_wt = *(FSIGNAL*)((char*)_det_weights->buf + _det_weights->strides[0]*i_det);

        for (auto const &rng: ivals[i_det].segments) {
            for (int i_time = rng.first; i_time < rng.second; ++i_time) {
                const int *pixel_offset = pixel_buf_man.data_ptr[i_det] +
                    pixel_buf_man.steps[0]*i_time;
                if (pixel_offset[0] < 0)
                    continue;
                const FSIGNAL sig = *(_signalspace->data_ptr[i_det] +
                                      _signalspace->steps[0]*i_time);
                const FSIGNAL *pf = (spin_proj_man.data_ptr[i_det] +
                                     spin_proj_man.steps[0] * i_time);
                for (int i_spin = 0; i_spin < n_spin; ++i_spin) {
                    *tiling.pix(i_spin, pixel_offset) += det_wt * sig * pf[i_spin];
                }
            }
        }
    }
}

template<typename TilingSys>
static
void precomp_to_weight_map_single_thread(Pixelizor2_Flat<TilingSys> &tiling,
                                         const SignalSpace<int32_t> &pixel_buf_man,
                                         const SignalSpace<FSIGNAL> &spin_proj_man,
                                         vector<RangesInt32> ivals,
                                         BufferWrapper<FSIGNAL> &_det_weights)
{
    int n_det = pixel_buf_man.dims[0];
    int n_spin = spin_proj_man.dims[2];
    assert(spin_proj_man.steps[1] == 1); // assumed below
    for (int i_det = 0; i_det < n_det; ++i_det) {
        FSIGNAL det_wt = 1.;
        if (_det_weights->obj != NULL)
            det_wt = *(FSIGNAL*)((char*)_det_weights->buf + _det_weights->strides[0]*i_det);

        for (auto const &rng: ivals[i_det].segments) {
            for (int i_time = rng.first; i_time < rng.second; ++i_time) {
                const int *pixel_offset = pixel_buf_man.data_ptr[i_det] +
                    pixel_buf_man.steps[0]*i_time;
                if (pixel_offset[0] < 0)
                    continue;
                const FSIGNAL *pf = (spin_proj_man.data_ptr[i_det] +
                                     spin_proj_man.steps[0] * i_time);
                for (int imap=0; imap < n_spin; ++imap)
                    for (int jmap=imap; jmap < n_spin; ++jmap)
                        *tiling.wpix(imap, jmap, pixel_offset) += pf[imap] * pf[jmap] * det_wt;
            }
        }
    }
}

template<typename TilingSys>
bp::object ProjEng_Precomp<TilingSys>::to_map(
    bp::object map, bp::object pixel_index, bp::object spin_proj,
    bp::object signal, bp::object det_weights, bp::object thread_intervals)
{
    // You won't get far without pixel_index, so use that to nail down
    // the n_time and n_det.
    auto pixel_buf_man = SignalSpace<int32_t>(
        pixel_index, "pixel_index", NPY_INT32, -1, -1, -1);
    int n_det = pixel_buf_man.dims[0];
    int n_time = pixel_buf_man.dims[1];
    int n_pix = pixel_buf_man.dims[2];

    // Similarly, spin_proj tells you the number of components.
    auto spin_proj_man = SignalSpace<FSIGNAL>(
        spin_proj, "spin_proj", FSIGNAL_NPY_TYPE, n_det, n_time, -1);
    int n_spin = spin_proj_man.dims[2];

    // Unlike the on-the-fly class, we aren't able to make the map
    // here, because there's no initialized pixelizor.
    auto pixelizor = Pixelizor2_Flat<TilingSys>();
    pixelizor.TestInputs(map, true, false, n_spin);

    // Check the signal.
    auto _signalspace = SignalSpace<FSIGNAL>(
        signal, "signal", FSIGNAL_NPY_TYPE, n_det, n_time);

    BufferWrapper<FSIGNAL> _det_weights("det_weights", det_weights, true,
                                        vector<int>{n_det});

    // Below we assume the pixel sub-indices are close-packed.
    if (pixel_buf_man.steps[1] != 1)
        throw shape_exception("pixel_index",
                              "Fast dimension of pixel indices must be close-packed.");

    auto ivals = derive_ranges(thread_intervals, n_det, n_time);

    if (ivals.size() <= 1) {
        // This block may also be used if OMP is disabled.
        for (int i_thread=0; i_thread < ivals.size(); i_thread++)
            precomp_to_map_single_thread<TilingSys>(
                pixelizor, pixel_buf_man, spin_proj_man,
                ivals[i_thread],_det_weights, &_signalspace);
    } else {
#pragma omp parallel
        {
            // This loop construction handles the (sub-optimal) case
            // that the number of ivals is not equal to the number of
            // OMP threads.
            for (int i_thread=0; i_thread < ivals.size(); i_thread++) {
                if (i_thread % omp_get_num_threads() == omp_get_thread_num())
                    precomp_to_map_single_thread<TilingSys>(
                        pixelizor, pixel_buf_man, spin_proj_man,
                        ivals[i_thread], _det_weights, &_signalspace);
            }
        }
    }
    return map;
}

template<typename TilingSys>
bp::object ProjEng_Precomp<TilingSys>::to_weight_map(
    bp::object map, bp::object pixel_index, bp::object spin_proj,
    bp::object det_weights, bp::object thread_intervals)
{
    // You won't get far without pixel_index, so use that to nail down
    // the n_time and n_det.
    auto pixel_buf_man = SignalSpace<int32_t>(
        pixel_index, "pixel_index", NPY_INT32, -1, -1, -1);
    int n_det = pixel_buf_man.dims[0];
    int n_time = pixel_buf_man.dims[1];
    int n_pix = pixel_buf_man.dims[2];

    // Similarly, spin_proj tells you the number of components.
    auto spin_proj_man = SignalSpace<FSIGNAL>(
        spin_proj, "spin_proj", FSIGNAL_NPY_TYPE, n_det, n_time, -1);
    int n_spin = spin_proj_man.dims[2];

    // Unlike the on-the-fly class, we aren't able to make the map
    // here, because there's no initialized pixelizor.
    auto pixelizor = Pixelizor2_Flat<TilingSys>();
    pixelizor.TestInputs(map, true, false, n_spin);

    BufferWrapper<FSIGNAL> _det_weights("det_weights", det_weights, true,
                                        vector<int>{n_det});

    // Below we assume the pixel sub-indices are close-packed.
    if (pixel_buf_man.steps[1] != 1)
        throw shape_exception("pixel_index",
                              "Fast dimension of pixel indices must be close-packed.");

    auto ivals = derive_ranges(thread_intervals, n_det, n_time);

    if (ivals.size() <= 1) {
        // This block may also be used if OMP is disabled.
        for (int i_thread=0; i_thread < ivals.size(); i_thread++)
            precomp_to_weight_map_single_thread<TilingSys>(
                pixelizor, pixel_buf_man, spin_proj_man,
                ivals[i_thread],_det_weights);
    } else {
#pragma omp parallel
        {
            // This loop construction handles the (sub-optimal) case
            // that the number of ivals is not equal to the number of
            // OMP threads.
            for (int i_thread=0; i_thread < ivals.size(); i_thread++) {
                if (i_thread % omp_get_num_threads() == omp_get_thread_num())
                    precomp_to_weight_map_single_thread<TilingSys>(
                        pixelizor, pixel_buf_man, spin_proj_man,
                        ivals[i_thread], _det_weights);
            }
        }
    }
    return map;
}

typedef ProjEng_Precomp<NonTiled> ProjEng_Precomp_NonTiled;
typedef ProjEng_Precomp<Tiled> ProjEng_Precomp_Tiled;

#define PROJENG(PIX, SPIN, TILING) ProjEng_##PIX##_##SPIN##_##TILING

#define TYPEDEF_TILING(PIX, SPIN, TILING)                               \
    typedef ProjectionEngine<Proj ## PIX,Pixelizor2_Flat<TILING>,Spin##SPIN> \
        PROJENG(PIX, SPIN, TILING);

#define TYPEDEF_SPIN(PIX, SPIN)                 \
    TYPEDEF_TILING(PIX, SPIN, Tiled)            \
    TYPEDEF_TILING(PIX, SPIN, NonTiled)

#define TYPEDEF_PIX(PIX)                         \
    TYPEDEF_SPIN(PIX, T)                         \
    TYPEDEF_SPIN(PIX, QU)                        \
    TYPEDEF_SPIN(PIX, TQU)

TYPEDEF_PIX(Flat)
TYPEDEF_PIX(Quat)
TYPEDEF_PIX(CAR)
TYPEDEF_PIX(CEA)
TYPEDEF_PIX(ARC)
TYPEDEF_PIX(TAN)
TYPEDEF_PIX(ZEA)

#define STRINGIFY(X) #X

#define EXPORT_ENGINE(CLASSNAME)                                        \
    bp::class_<CLASSNAME>(STRINGIFY(CLASSNAME), bp::init<bp::object>()) \
    .add_property("index_count", &CLASSNAME::index_count)               \
    .add_property("comp_count", &CLASSNAME::comp_count)                 \
    .def("coords", &CLASSNAME::coords)                                  \
    .def("pixels", &CLASSNAME::pixels)                                  \
    .def("pointing_matrix", &CLASSNAME::pointing_matrix)                \
    .def("pixel_ranges", &CLASSNAME::pixel_ranges)                      \
    .def("zeros", &CLASSNAME::zeros)                                    \
    .def("from_map", &CLASSNAME::from_map)                              \
    .def("to_map", &CLASSNAME::to_map)                                  \
    .def("to_weight_map", &CLASSNAME::to_weight_map)                    \
    ;

#define EXPORT_TILING(PIX, SPIN, TILING)        \
    EXPORT_ENGINE(PROJENG(PIX, SPIN, TILING))

#define EXPORT_SPIN(PIX, SPIN)                                          \
    EXPORT_TILING(PIX, SPIN, Tiled)                                     \
    EXPORT_TILING(PIX, SPIN, NonTiled)

#define EXPORT_PIX(PIX) \
    EXPORT_SPIN(PIX, T)                                                 \
    EXPORT_SPIN(PIX, QU)                                                \
    EXPORT_SPIN(PIX, TQU)

#define EXPORT_PRECOMP(CLASSNAME)                                       \
    bp::class_<CLASSNAME>(#CLASSNAME)                                   \
    .def("from_map", &CLASSNAME::from_map)                              \
    .def("to_map", &CLASSNAME::to_map)                                  \
    .def("to_weight_map", &CLASSNAME::to_weight_map)                    \
;

// There are probably better ways to do this..
template<typename T>
inline
int _index_count(const T &) { return T::index_count; }

PYBINDINGS("so3g")
{
    EXPORT_PIX(Flat);
    EXPORT_PIX(Quat);
    EXPORT_PIX(CAR);
    EXPORT_PIX(CEA);
    EXPORT_PIX(ARC);
    EXPORT_PIX(TAN);
    EXPORT_PIX(ZEA);

    EXPORT_PRECOMP(ProjEng_Precomp_NonTiled);
    EXPORT_PRECOMP(ProjEng_Precomp_Tiled);
}


#include "dynet/nodes.h"
#include "dynet/dynet.h"
#include "dynet/training.h"
#include "dynet/timing.h"
#include "dynet/rnn.h"
#include "dynet/gru.h"
#include "dynet/lstm.h"
#include "dynet/dict.h"
#include "dynet/expr.h"

// Platform-independent header to allow calls to getpid()
#if _WINDOWS
    #include <process.h>
#endif

#include <iostream>
#include <fstream>
#include <sstream>

#include <boost/archive/text_iarchive.hpp>
#include <boost/archive/text_oarchive.hpp>

#include "ExternalSpines.pb.h"

using namespace std;
using namespace dynet;

unsigned LAYERS = 2;
unsigned INPUT_DIM = 128;
unsigned HIDDEN_DIM = 256;
unsigned TAG_HIDDEN_DIM = 64;
unsigned TAG_SIZE = 0;
unsigned VOCAB_SIZE = 0;

double CUTOFF_RATIO = 0.5;

dynet::Dict d;
dynet::Dict td;
int kNONE;
int kSOS;
int kEOS;
vector<int> pruning_info;
vector<int> pruning_info_gold;
const int PRUNING_BUCKETS = 1000;
const double PRUNING_MAX = 1.0;
const double PRUNING_DELTA = -0.001;

template <class Builder>
struct BiLSTMTagger {
  LookupParameter p_w;
  Parameter p_l2th;
  Parameter p_r2th;
  Parameter p_thbias;

  Parameter p_th2t;
  Parameter p_tbias;
  Builder l2rbuilder;
  Builder r2lbuilder;
  explicit BiLSTMTagger(Model& model) :
      l2rbuilder(LAYERS, INPUT_DIM, HIDDEN_DIM, model),
      r2lbuilder(LAYERS, INPUT_DIM, HIDDEN_DIM, model) {
    p_w = model.add_lookup_parameters(VOCAB_SIZE, {INPUT_DIM});
    p_l2th = model.add_parameters({TAG_HIDDEN_DIM, HIDDEN_DIM});
    p_r2th = model.add_parameters({TAG_HIDDEN_DIM, HIDDEN_DIM});
    p_thbias = model.add_parameters({TAG_HIDDEN_DIM});

    p_th2t = model.add_parameters({TAG_SIZE, TAG_HIDDEN_DIM});
    p_tbias = model.add_parameters({TAG_SIZE});
  }

  // return Expression of total loss
  Expression BuildTaggingGraph(
      const vector<int>& sent, const vector<int>& tags,
      ComputationGraph& cg, double* cor = 0, unsigned* ntagged = 0,
      bool print_labels = false, bool eval = false,
      graphparser::Sentence* nsent = nullptr
  ) {
    const unsigned slen = sent.size();
    l2rbuilder.new_graph(cg);  // reset RNN builder for new graph
    l2rbuilder.start_new_sequence();
    r2lbuilder.new_graph(cg);  // reset RNN builder for new graph
    r2lbuilder.start_new_sequence();
    Expression i_l2th = parameter(cg, p_l2th);
    Expression i_r2th = parameter(cg, p_r2th);
    Expression i_thbias = parameter(cg, p_thbias);
    Expression i_th2t = parameter(cg, p_th2t);
    Expression i_tbias = parameter(cg, p_tbias);
    vector<Expression> errs;
    vector<Expression> i_words(slen);
    vector<Expression> fwds(slen);
    vector<Expression> revs(slen);

    // set words, adding noise during training (non-eval)
    for (unsigned t = 0; t < slen; ++t) {
///      if (print_probs)
///        cout << t << " " << sent[t] << " " << d.convert(sent[t]) << endl;
      i_words[t] = lookup(cg, p_w, sent[t]);
      if (!eval) { i_words[t] = noise(i_words[t], 0.1); }
    }

    // read sequence from left to right
    l2rbuilder.add_input(lookup(cg, p_w, kSOS));
    for (unsigned t = 0; t < slen; ++t)
      fwds[t] = l2rbuilder.add_input(i_words[t]);
    // read sequence from right to left
    r2lbuilder.add_input(lookup(cg, p_w, kEOS));
    for (unsigned t = 0; t < slen; ++t)
      revs[slen - t - 1] = r2lbuilder.add_input(i_words[slen - t - 1]);

    for (unsigned t = 0; t < slen; ++t) {
      if (tags[t] != kNONE) {
        if (ntagged) (*ntagged)++;
        Expression i_th = tanh(affine_transform({i_thbias, i_l2th, fwds[t], i_r2th, revs[t]}));
        Expression i_t = affine_transform({i_tbias, i_th2t, i_th});
        if (cor) {
          vector<float> dist = as_vector(cg.incremental_forward(i_t));

          // Find best tag according to the distribution and also collect info
          // on the score for the correct tag.
          double best = -9e99;
          double worst = 9e99;
          double corr_score = 0.0;
          int besti = -1;
          if (print_labels)
            cout << "Word " << t << " " << d.convert(sent[t]);
          for (unsigned int i = 0; i < dist.size(); ++i) {
///            if (print_probs) cout << " " << dist[i];
            if (i == tags[t]) corr_score = dist[i];
            if (dist[i] > best) {
              best = dist[i];
              besti = i;
            }
            if (dist[i] < worst) worst = dist[i];
          }
          if (print_labels)
            cout << " " << besti << " correct: " << tags[t] << " ( " << td.convert(besti) << " vs. " << td.convert(tags[t]) << " )" << endl;
          if (tags[t] == besti) (*cor)++;

          if (eval) {
            // Record information for pruning
            double range = best - worst;

            // Save scores
            graphparser::ScoreSet* saved_scores = nullptr;
            if (nsent != nullptr) {
              saved_scores = nsent->add_scores();
              saved_scores->set_min_score(worst);
              double scaled_ratio = (range * CUTOFF_RATIO) + worst;
              for (unsigned int i = 0; i < dist.size(); ++i) {
                if (dist[i] > scaled_ratio) {
                  graphparser::Score* cscore = saved_scores->add_scores();
                  cscore->set_id(i);
                  cscore->set_score(dist[i]);
                }
              }
            }

            // Sort dist to make the next stage more efficient
            sort(dist.begin(), dist.end());
            int cpos = 0;
            for (unsigned int i = dist.size() - 1; i > 0; i--) {
              int count = dist.size() - i - 1;
              double scaled = (dist[i] - worst) / range;
              while (
                  cpos * PRUNING_DELTA + PRUNING_MAX > scaled &&
                  cpos < PRUNING_BUCKETS - 2
              ) {
                pruning_info[cpos] += count;
                if (dist[i] < corr_score) pruning_info_gold[cpos] += 1;
                cpos += 1;
              }
            }
            pruning_info[PRUNING_BUCKETS - 1] += dist.size();
            pruning_info_gold[PRUNING_BUCKETS - 1] += 1;
          }
        }

        Expression i_err = pickneglogsoftmax(i_t, tags[t]);
        errs.push_back(i_err);
      }
    }
    if (print_labels) cout << endl;
    return sum(errs);
  }
};

string get_string_arg(int argc, char** argv, string arg, string default_val="") {
  for (int i = 0; i < argc; i++) {
    if (arg.compare(argv[i]) == 0 && i < (argc - 1)) {
      return string(argv[i + 1]);
    }
  }
  return default_val;
}

double get_double_arg(int argc, char** argv, string arg, double default_val) {
  for (int i = 0; i < argc; i++) {
    if (arg.compare(argv[i]) == 0 && i < (argc - 1)) {
      return stod(string(argv[i + 1]));
    }
  }
  return default_val;
}

int get_int_arg(int argc, char** argv, string arg, int default_val) {
  for (int i = 0; i < argc; i++) {
    if (arg.compare(argv[i]) == 0 && i < (argc - 1)) {
      return stoi(string(argv[i + 1]));
    }
  }
  return default_val;
}

bool get_bool_arg(int argc, char** argv, string arg) {
  for (int i = 0; i < argc; i++) {
    if (arg.compare(argv[i]) == 0) return true;
  }
  return false;
}

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

  dynet::initialize(argc, argv);
  if (argc < 3) {
    cerr << "Usage: "
      << argv[0] << endl
      << "  [-train corpus.txt]" << endl
      << "  [-dev dev.txt]" << endl
      << "  [-test test.txt]" << endl
      << "  [-model model.params]" << endl
      << "  [-word-dict dict.words]" << endl
      << "  [-tag-dict dict.tags]" << endl
      << "  [-prefix name (default=tagger)]" << endl
      << "  [-cutoff-ratio CUTOFF_RATIO]" << endl
      << "  [-layers LAYERS]" << endl
      << "  [-input-dim INPUT_DIM]" << endl
      << "  [-hidden-dim HIDDEN_DIM]" << endl
      << "  [-tag-hidden-dim TAG_HIDDEN_DIM]" << endl
      << "  [-print-labels]" << endl
      << endl;
    return 1;
  }
  bool doing_train = get_bool_arg(argc, argv, "-train");
  bool doing_test = get_bool_arg(argc, argv, "-test");
  bool min_loss = get_bool_arg(argc, argv, "-min_loss");
  string prefix = get_string_arg(argc, argv, "-prefix", "tagger");

  CUTOFF_RATIO = get_double_arg(argc, argv, "-cutoff-ratio", CUTOFF_RATIO);
  LAYERS = get_int_arg(argc, argv, "-layers", LAYERS);
  INPUT_DIM = get_int_arg(argc, argv, "-input-dim", INPUT_DIM);
  HIDDEN_DIM = get_int_arg(argc, argv, "-hidden-dim", HIDDEN_DIM);
  TAG_HIDDEN_DIM = get_int_arg(argc, argv, "-tag-hidden-dim", TAG_HIDDEN_DIM);
  bool print_labels = get_bool_arg(argc, argv, "-print-labels");

  vector<pair<vector<int>,vector<int>>> training, dev;
  vector<pair<int, pair<vector<int>,vector<int>>>> test;

  if (doing_train) {
    kNONE = td.convert("*");
    kSOS = d.convert("<s>");
    kEOS = d.convert("</s>");
    string line;
    int tlc = 0;
    int ttoks = 0;
    string train_in = get_string_arg(argc, argv, "-train");
    assert(train_in.compare("") != 0);
    cerr << "Reading training data from " << train_in << "...\n";
    {
      ifstream in(train_in);
      assert(in);
      while(getline(in, line)) {
        auto found = line.find("# SentID");
        if (found == std::string::npos) {
          ++tlc;
          int nc = 0;
          vector<int> x,y;
          read_sentence_pair(line, x, d, y, td);
          assert(x.size() == y.size());
          if (x.size() == 0) { cerr << line << endl; abort(); }
          training.push_back(make_pair(x,y));
          for (unsigned i = 0; i < y.size(); ++i) {
            if (y[i] != kNONE) { ++nc; }
          }
          if (nc == 0) {
            cerr << "No tagged tokens in line " << tlc << endl;
            abort();
          }
          ttoks += x.size();
        }
      }
      cerr << tlc << " lines, " << ttoks << " tokens, " << d.size() << " types\n";
      cerr << "Tags: " << td.size() << endl;
    }

    d.freeze(); // no new word types allowed
    td.freeze(); // no new tag types allowed
    d.set_unk("UNKNOWN_WORD");
    td.set_unk("UNKNOWN_TAG");

    ostringstream file_dw;
    ostringstream file_dt;
    file_dw << prefix << ".dict.words";
    file_dt << prefix << ".dict.tags";
    ofstream ows(file_dw.str());
    ofstream ots(file_dt.str());
    boost::archive::text_oarchive ow(ows);
    boost::archive::text_oarchive ot(ots);
    ow << d;
    ot << td;
  } else {
    string words_infile = get_string_arg(argc, argv, "-word-dict");
    string tags_infile = get_string_arg(argc, argv, "-tag-dict");
    cerr << "Reading existing word dictionary from: " << words_infile << endl;
    cerr << "Reading existing tag dictionary from: " << tags_infile << endl;
    ifstream words_in(words_infile);
    ifstream tags_in(tags_infile);
    boost::archive::text_iarchive words_in_archive(words_in);
    boost::archive::text_iarchive tags_in_archive(tags_in);
    words_in_archive >> d;
    tags_in_archive >> td;

    kNONE = td.convert("*");
    kSOS = d.convert("<s>");
    kEOS = d.convert("</s>");
  }

  VOCAB_SIZE = d.size();
  TAG_SIZE = td.size();

  if (! doing_test) {
    int dlc = 0;
    int dtoks = 0;
    string dev_in = get_string_arg(argc, argv, "-dev");
    assert(dev_in.compare("") != 0);
    cerr << "Reading dev data from " << dev_in << "...\n";
    {
      ifstream in(dev_in);
      assert(in);
      string line;
      while(getline(in, line)) {
        auto found = line.find("# SentID");
        if (found == std::string::npos) {
          ++dlc;
          vector<int> x,y;
          read_sentence_pair(line, x, d, y, td);
          assert(x.size() == y.size());
          dev.push_back(make_pair(x,y));
          dtoks += x.size();
        }
      }
      cerr << dlc << " lines, " << dtoks << " tokens\n";
    }
  }

  if (doing_test) {
    int tlc = 0;
    int ttoks = 0;
    int sentence_id = -1;
    string test_fname = get_string_arg(argc, argv, "-test");
    cerr << "Reading test data from " << test_fname << "...\n";
    {
      ifstream in(test_fname);
      assert(in);
      string line;
      while(getline(in, line)) {
        // Get the ID
        auto found = line.find("# SentID");
        if (found == std::string::npos) {
          ++tlc;
          vector<int> x,y;
          read_sentence_pair(line, x, d, y, td);
          assert(x.size() == y.size());
          test.push_back(make_pair(sentence_id, make_pair(x,y)));
          ttoks += x.size();
        } else {
          sentence_id = stoi(line.substr(8));
        }
      }
      cerr << tlc << " lines, " << ttoks << " tokens\n";
    }
  }

  Model model;
///  Trainer* sgd = new MomentumSGDTrainer(model);
  Trainer* sgd = new AdamTrainer(model);
  BiLSTMTagger<LSTMBuilder> tagger(model);
  string model_fname = get_string_arg(argc, argv, "-model");
  if (model_fname.compare("") != 0) {
    cerr << "Reading existing model params from: " << model_fname << endl;
    ifstream in(model_fname);
    boost::archive::text_iarchive ia(in);
    ia >> model;
  }

  for (int i = 0; i < PRUNING_BUCKETS; i++) {
    pruning_info.push_back(0);
    pruning_info_gold.push_back(0);
  }

  if (doing_test) {
    ostringstream os;
    os << prefix << ".data";
    const string fname = os.str();

    graphparser::Scores scores_for_saving;

    // First, print the tag mapping
    for (unsigned int i = 0; i < td.size(); ++i) {
      string spine_text = td.convert(i);
      graphparser::Spine* spine = scores_for_saving.add_spines();
      spine->set_id(i);
      spine->set_text(spine_text);
///      cout << "Tag " << i << " " << spine << endl;
    }

    double tloss = 0;
    unsigned ttags = 0;
    double tcorr = 0;
    for (auto& id_sent : test) {
      int sentence_id = id_sent.first;
      auto sent = id_sent.second;
      ComputationGraph cg;
      graphparser::Sentence* nsent = scores_for_saving.add_sentences();
      nsent->set_id(sentence_id);
      tagger.BuildTaggingGraph(sent.first, sent.second, cg, &tcorr, &ttags,
          print_labels, true, nsent);
    }

    // Write the new address book back to disk.
    fstream output(fname, ios::out | ios::trunc | ios::binary);
    if (!scores_for_saving.SerializeToOstream(&output)) {
      cerr << "Failed to write scores for saving." << endl;
      return -1;
    }
  } else {
    double best = 9e+99; // Tracks the loss for the best model so far
    double best_prune = 1.0; // Tracks the best pruning by a model so far

    // Prepare filename for model storage, and write dictionaries
    ostringstream os;
    os << prefix << ".params";
    const string fname = os.str();
    cerr << "Parameters will be written to: " << fname << endl;

    unsigned report_every_i = get_int_arg(argc, argv, "-report-freq", 1000);
    unsigned dev_every_i_reports = get_int_arg(argc, argv, "-dev-report-freq", 10);
    unsigned si = training.size();
    vector<unsigned> order(training.size());
    for (unsigned i = 0; i < order.size(); ++i) order[i] = i;
    bool first = true;
    int report = 0;
    unsigned lines = 0;
    while(1) {
      Timer iteration("completed in");
      double loss = 0;
      unsigned ttags = 0;
      double correct = 0;
      for (unsigned i = 0; i < report_every_i; ++i) {
        if (si == training.size()) {
          si = 0;
          if (first) { first = false; } else { sgd->update_epoch(); }
          cerr << "**SHUFFLE\n";
          shuffle(order.begin(), order.end(), *dynet::rndeng);
        }

        // build graph for this instance
        ComputationGraph cg;
        auto& sent = training[order[si]];
        ++si;
        Expression loss_expr = tagger.BuildTaggingGraph(sent.first, sent.second, cg, &correct, &ttags, false, false);

        // Run forward pass, backpropagate, and do an update
        loss += as_scalar(cg.forward(loss_expr));
        cg.backward(loss_expr);
        sgd->update(1.0);
        ++lines;
      }
      sgd->status();
      cerr << " E = " << (loss / ttags) << " ppl=" << exp(loss / ttags) << " (acc=" << (correct / ttags) << ") ";

      // show score on dev data?
      report++;
      if (report % dev_every_i_reports == 0) {

        // Prepare storage for pruning information
        pruning_info.clear();
        pruning_info_gold.clear();
        for (int i = 0; i < PRUNING_BUCKETS; i++) {
          pruning_info.push_back(0);
          pruning_info_gold.push_back(0);
        }

        double dloss = 0;
        unsigned dtags = 0;
        double dcorr = 0;
        for (auto& sent : dev) {
          ComputationGraph cg;
          Expression loss_expr = tagger.BuildTaggingGraph(sent.first, sent.second, cg, &dcorr, &dtags, false, true);
          dloss += as_scalar(cg.forward(loss_expr));
        }
        // Return pruning information
        double goal = 0.97;
        double gold_tag_count = pruning_info_gold.back();
        double tag_count = pruning_info.back();
        double score_at_99 = 1.0;
        for (int i = 0; i < PRUNING_BUCKETS; i++) {
          // Identify if the goal is met
          if (goal * gold_tag_count < pruning_info_gold[i]) {
            double gold_pruned = pruning_info_gold[i] / gold_tag_count;
            double pruned = pruning_info[i] / tag_count;
            if ((goal - 0.99) < 0.001) score_at_99 = pruned;
            cerr << "\n***DEV Pruning " << gold_pruned << " gold, " << pruned << " = " << pruning_info[i] << " at " << (PRUNING_MAX + i * PRUNING_DELTA);
            goal += 0.005;
          }
          // Note the result at 99 (to be used to determine model saving)
        }
        if (min_loss) {
          if (dloss < best) {
            best = dloss;
            ofstream out(fname);
            boost::archive::text_oarchive oa(out);
            oa << model;
          }
        } else {
          if (score_at_99 < best_prune) {
            best_prune = score_at_99;
            ofstream out(fname);
            boost::archive::text_oarchive oa(out);
            oa << model;
          }
        }

        cerr << "\n***DEV [epoch=" << (lines / (double)training.size()) << "] E = " << (dloss / dtags) << " ppl=" << exp(dloss / dtags) << " acc=" << (dcorr / dtags) << ' ';
      }
    }
  }

  google::protobuf::ShutdownProtobufLibrary();

  delete sgd;
}


//---------------------------------------------------------------------------//
// Copyright (c) 2018-2020 Mikhail Komarov <nemo@nil.foundation>
// Copyright (c) 2020 Nikita Kaskov <nbering@nil.foundation>
//
// MIT License
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.
//---------------------------------------------------------------------------//

#define BOOST_TEST_MODULE kasumi_cipher_test

#include <iostream>
#include <unordered_map>

#include <boost/test/unit_test.hpp>
#include <boost/test/data/test_case.hpp>
#include <boost/test/data/monomorphic.hpp>

#include <nil/crypto3/block/algorithm/encrypt.hpp>
#include <nil/crypto3/block/algorithm/decrypt.hpp>

#include <nil/crypto3/block/kasumi.hpp>

using namespace nil::crypto3;

BOOST_AUTO_TEST_SUITE(kasumi_test_suite)

BOOST_AUTO_TEST_CASE(kasumi_1) {

    std::vector<char> input = {'\xea', '\x02', '\x47', '\x14', '\xad', '\x5c', '\x4d', '\x84'};
    std::vector<char> key = {'\x2b', '\xd6', '\x45', '\x9f', '\x82', '\xc5', '\xb3', '\x00', '\x95', '\x2c', '\x49', '\x10', '\x48', '\x81', '\xff', '\x48'};

    std::string out = encrypt<block::kasumi>(input, key);

    BOOST_CHECK_EQUAL(out, "df1f9b251c0bf45f");
}

BOOST_AUTO_TEST_SUITE_END()


#define BOOST_TEST_MODULE TensorDot
#include <boost/test/unit_test.hpp>
#include <vexcl/vector.hpp>
#include <vexcl/vector_view.hpp>
#include <vexcl/tensordot.hpp>
#include <vexcl/random.hpp>
#include <vexcl/reductor.hpp>
#include <vexcl/element_index.hpp>
#include <vexcl/function.hpp>
#include "context_setup.hpp"

BOOST_AUTO_TEST_CASE(axes_pairs) {
    auto ax = vex::axes_pairs(1, 2, 3, 4, 5, 6);

    BOOST_CHECK_EQUAL(ax[0][0], 1);
    BOOST_CHECK_EQUAL(ax[0][1], 2);
    BOOST_CHECK_EQUAL(ax[1][0], 3);
    BOOST_CHECK_EQUAL(ax[1][1], 4);
    BOOST_CHECK_EQUAL(ax[2][0], 5);
    BOOST_CHECK_EQUAL(ax[2][1], 6);
}

BOOST_AUTO_TEST_CASE(mat_mat) {
    using vex::_;
    using vex::extents;

    std::vector<vex::command_queue> queue(1, ctx.queue(0));

    const size_t N = 32;

    vex::vector<double> a(queue, N * N);
    vex::vector<double> b(queue, N * N);

    vex::Random<double> rnd;
    vex::Reductor<double, vex::SUM> sum(queue);

    auto idx = vex::element_index();

    a = rnd(idx, 0);
    b = rnd(idx, 1);

    vex::slicer<2> dim(extents[N][N]);

    BOOST_CHECK_SMALL(
            sum(
                fabs(
                    vex::tensordot(
                        dim[_](a), dim[_](b), vex::axes_pairs(1, 0)
                        )
                    -
                    vex::reduce<vex::SUM>(
                        vex::extents[N][N][N],
                        vex::reshape(a, extents[N][N][N], extents[0][1]) *
                        vex::reshape(b, extents[N][N][N], extents[1][2]),
                        1)
                    )
               ),
            1e-8);
}

BOOST_AUTO_TEST_CASE(mat_vec) {
    using vex::_;
    using vex::extents;

    std::vector<vex::command_queue> queue(1, ctx.queue(0));

    const size_t N = 32;

    vex::vector<double> a(queue, N * N);
    vex::vector<double> b(queue, N);

    vex::Random<double> rnd;
    vex::Reductor<double, vex::SUM> sum(queue);

    auto idx = vex::element_index();

    a = rnd(idx, 0);
    b = rnd(idx, 1);

    vex::slicer<2> mat(extents[N][N]);
    vex::slicer<1> vec(extents[N]);

    BOOST_CHECK_SMALL(
            sum(
                fabs(
                    vex::tensordot(
                        mat[_](a), vec[_](b), vex::axes_pairs(1, 0)
                        )
                    -
                    vex::reduce<vex::SUM>(
                        vex::extents[N][N],
                        a * vex::reshape(b, extents[N][N], extents[1]),
                        1)
                    )
               ),
            1e-8);
}

BOOST_AUTO_TEST_CASE(vec_mat) {
    using vex::_;
    using vex::extents;

    std::vector<vex::command_queue> queue(1, ctx.queue(0));

    const size_t N = 32;

    vex::vector<double> a(queue, N);
    vex::vector<double> b(queue, N * N);

    vex::Random<double> rnd;
    vex::Reductor<double, vex::SUM> sum(queue);

    auto idx = vex::element_index();

    a = rnd(idx, 0);
    b = rnd(idx, 1);

    vex::slicer<2> mat(extents[N][N]);
    vex::slicer<1> vec(extents[N]);

    BOOST_CHECK_SMALL(
            sum(
                fabs(
                    vex::tensordot(
                        vec[_](a), mat[_](b), vex::axes_pairs(0, 0)
                        )
                    -
                    vex::reduce<vex::SUM>(
                        vex::extents[N][N],
                        vex::reshape(a, extents[N][N], extents[0]) * b,
                        0)
                    )
               ),
            1e-8);
}

BOOST_AUTO_TEST_SUITE_END()