#include <cstdio>
#include <vector>
#include <algorithm>

#include <boost/compute.hpp>
#include <boost/compute/types/complex.hpp>

#include <clFFT.h>

#include <ceres/ceres.h>

#include "image.hpp"
#include "vec2.hpp"

namespace compute = boost::compute;
using boost::compute::dim;

using homography_t = std::array<double, 9>;

template <typename T>
void hann(img_t<T>& out, int w, int h, int d=1) {
    out.resize(w, h, d);
    out.set_value(0);
    // with modifications from "Burst photography for high dynamic range and low-light imaging on mobile cameras"
    // namely: half pixel offset and /w|/h instead of /(w-1)|/(h-1)
    for (int l = 0; l < d; l++) {
        for (int y = 0; y < h; y++) {
            T vy = 0.5f * (1 - std::cos(2*M_PI*(y+0.5) / h));
            for (int x = 0; x < w; x++) {
                T vx = 0.5 * (1 - std::cos(2*M_PI*(x+0.5) / w));
                out(x, y, l) = vx * vy;
            }
        }
    }
}

template <typename T>
vec2<T> homography_apply(const homography_t& H, vec2<T> x) {
    T X = H[0]*x[0] + H[1]*x[1] + H[2];
    T Y = H[3]*x[0] + H[4]*x[1] + H[5];
    T Z = H[6]*x[0] + H[7]*x[1] + H[8];
    return { X / Z, Y / Z };
}

struct TranslationResidual {
    TranslationResidual(vec2<double> x, vec2<double> y) : x_(x), y_(y) {}

    template <typename T> bool operator()(const T* const h,
                                          T* residual) const {
        T X = h[0]*x_[0] + h[1]*x_[1] + h[2];
        T Y = h[3]*x_[0] + h[4]*x_[1] + h[5];
        T Z = h[6]*x_[0] + h[7]*x_[1] + h[8];

        vec2<T> p = { X / Z, Y / Z };
        residual[0] = y_[0] - p[0];
        residual[1] = y_[1] - p[1];
        return true;
    }

    private:
    const vec2<double> x_;
    const vec2<double> y_;
};

homography_t homography_from_translations_robust(const std::vector<vec2<vec2<double>>>& translations)
{
    homography_t h = {1,0,0, 0,1,0, 0,0,1};
    ceres::Problem problem;
    double* ph = &h[0];
    for (unsigned i = 0; i < translations.size(); i++) {
        problem.AddResidualBlock(new ceres::AutoDiffCostFunction<TranslationResidual, 2, 9>(
                                    new TranslationResidual(translations[i][0], translations[i][1])),
                                 new ceres::SoftLOneLoss(1.0), ph);
    }

    ceres::Solver::Options options;
    ceres::Solver::Summary summary;
    ceres::Solve(options, &problem, &summary);
    for (int k = 0; k < 9; k++)
        h[k] /= h[8];
    return h;
}

template <typename T>
img_t<T> img_from_device(const compute::vector<T>& input,
                                 int w, int h, int d,
                                 compute::command_queue& queue) {
    queue.finish();
    img_t<T> out(w, h, d);
    compute::copy(input.begin(), input.end(), out.data.begin(), queue);
    queue.finish();
    return out;
}

compute::vector<float> img_to_device(const img_t<float>& input, compute::command_queue& queue) {
    queue.finish();
    return compute::vector<float>(input.data.begin(), input.data.end(), queue);
}

const char kernels_src[] = BOOST_COMPUTE_STRINGIZE_SOURCE(
    typedef float2 cfloat;

    inline cfloat cmult(cfloat a, cfloat b){
        return (cfloat)( a.x*b.x - a.y*b.y, a.x*b.y + a.y*b.x);
    }

    inline cfloat cconj(cfloat a){
        return (cfloat)(a.x, -a.y);
    }

    __kernel void extract_tile(__global const float* input,
                               __global float* output,
                               const int ox, const int oy,
                               const int w, const int h)
    {
        int x = get_global_id(0) + ox;
        int y = get_global_id(1) + oy;
        const int dx = get_global_id(0);
        const int dy = get_global_id(1);

        x = max(0, min(x, w-1));
        y = max(0, min(y, h-1));

        output[(dx+dy*W)*3+0] = input[(x+y*w)*3+0];
        output[(dx+dy*W)*3+1] = input[(x+y*w)*3+1];
        output[(dx+dy*W)*3+2] = input[(x+y*w)*3+2];
    }

    __kernel void fulltohalf(__global const float* input,
                             __global float* output)
    {
        const int x = get_global_id(0);
        const int y = get_global_id(1);
        int dx = (x + W/2) % W;
        int dy = (y + W/2) % W;
        if (dx >= W/4 && dx < W*3/4 && dy >= W/4 && dy < W*3/4) {
            output[(x+y*W)*3+0] = input[(dx+dy*W)*3+0];
            output[(x+y*W)*3+1] = input[(dx+dy*W)*3+1];
            output[(x+y*W)*3+2] = input[(dx+dy*W)*3+2];
        } else {
            output[(x+y*W)*3+0] = 0.f;
            output[(x+y*W)*3+1] = 0.f;
            output[(x+y*W)*3+2] = 0.f;
        }
    }

    __kernel void float2complex(__global const float* input,
                                __global cfloat* output)
    {
        const int x = get_global_id(0);

        output[x].x = input[x];
        output[x].y = 0.f;
    }

    __kernel void magnitude(__global const cfloat* input,
                            __global float* out)
    {
        const int x = get_global_id(0);

        out[x] = (fast_length(input[x*3+0])
                + fast_length(input[x*3+1])
                + fast_length(input[x*3+2])) / 3.f;
    }

    // /!\ transposed result
    __kernel void blur(__global const float* in,
                       __global float* out,
                       __global const float* gaussian, int size)
    {
        const int x = get_global_id(0);
        const int y = get_global_id(1);

        float v = 0.;
        for (int i = -size; i <= size; i++) {
            v += in[(x+i+W)%W+W*y] * gaussian[i+size];
        }
        out[y+W*x] = v;
    }

    __kernel void pow_(__global float* buf, float p)
    {
        const int x = get_global_id(0);
        buf[x] = pow(buf[x], p);
    }

    __kernel void crosscorrelation(__global const cfloat* img,
                                   __global const cfloat* ref,
                                   __global cfloat* cc)
    {
        const int x = get_global_id(0);
        cc[x] = (cmult(img[x*3+0], cconj(ref[x*3+0]))
               + cmult(img[x*3+1], cconj(ref[x*3+1]))
               + cmult(img[x*3+2], cconj(ref[x*3+2]))) / 3.f;
    }

    __kernel void l2residuals(__global const cfloat* cc,
                              __global const float* boxfiltered,
                              __global float* D)
    {
        const int x = get_global_id(0);
        D[x] = boxfiltered[x] - 2.f * cc[x].x;
    }

    __kernel void translate(__global const float2* in,
                            __global float2* out,
                            float dx, float dy)
    {
        const int x = get_global_id(0);
        const int y = get_global_id(1);
        const int wx = (x + W / 2) % W - W / 2;
        const int wy = (y + W / 2) % W - W / 2;

        const float d = 2.f * M_PI_F * (wx * dx / W + wy * dy / W);
        const cfloat phase = (cfloat)(cos(d), sin(d));

        out[(x+y*W)*3+0] = cmult(in[(x+y*W)*3+0], phase);
        out[(x+y*W)*3+1] = cmult(in[(x+y*W)*3+1], phase);
        out[(x+y*W)*3+2] = cmult(in[(x+y*W)*3+2], phase);
    }

    __kernel void accumulate(__global const cfloat* tile,
                             __global const float* hann,
                             __global float* image,
                             __global float* image_weight,
                             const int ox, const int oy,
                             const int w, const int h)
    {
        int x = get_global_id(0) + ox;
        int y = get_global_id(1) + oy;
        const int dx = get_global_id(0);
        const int dy = get_global_id(1);
        float weight = hann[dx+dy*W];

        if (x >= 0 && x < w && y >= 0 && y < h) {
            image_weight[x+y*w] += weight;
            image[(x+y*w)*3+0] += tile[(dx+dy*W)*3+0].x * weight;
            image[(x+y*w)*3+1] += tile[(dx+dy*W)*3+1].x * weight;
            image[(x+y*w)*3+2] += tile[(dx+dy*W)*3+2].x * weight;
        }
    }

    __kernel void unweight(__global float* image,
                           __global const float* image_weight)
    {
        int x = get_global_id(0);
        image[x*3+0] /= image_weight[x];
        image[x*3+1] /= image_weight[x];
        image[x*3+2] /= image_weight[x];
    }

    __kernel void cunweight(__global cfloat* image,
                            __global const float* image_weight)
    {
        int x = get_global_id(0);
        image[x*3+0] /= image_weight[x];
        image[x*3+1] /= image_weight[x];
        image[x*3+2] /= image_weight[x];
    }

    __kernel void fba(__global cfloat* accum,
                      __global float* accum_weight,
                      __global const cfloat* tile,
                      __global const float* tile_weight)
    {
        int x = get_global_id(0);
        const float weight = tile_weight[x];

        accum_weight[x] += weight + 1e-6;
        accum[x*3+0].x += tile[x*3+0].x * weight;
        accum[x*3+0].y += tile[x*3+0].y * weight;
        accum[x*3+1].x += tile[x*3+1].x * weight;
        accum[x*3+1].y += tile[x*3+1].y * weight;
        accum[x*3+2].x += tile[x*3+2].x * weight;
        accum[x*3+2].y += tile[x*3+2].y * weight;
    }

    __kernel void sqr(__global const float* in,
                      __global float* out)
    {
        const int x = get_global_id(0);
        out[x] = in[x] * in[x];
    }

    // /!\ unnormalized + transposed output
    __kernel void boxfilter(__global const float* in,
                            __global float* out)
    {
        const int y = get_global_id(0);

        float v = in[y*W];
        for (int x = 1; x <= hw; x++) {
            v += in[x+y*W] + in[(W-x)+y*W];
        }

        out[y] = v;
        for (int x = 1; x <= hw; x++) {
            v += in[(x+hw)+y*W] - in[(W+x-hw-1)+y*W];
            out[y+x*W] = v;
        }
        for (int x = hw + 1; x < W - hw; x++) {
            v += in[(x+hw)+y*W] - in[(x-hw-1)+y*W];
            out[y+x*W] = v;
        }
        for (int x = W - hw; x < W; x++) {
            v += in[(x-W+hw)+y*W] - in[(x-hw-1)+y*W];
            out[y+x*W] = v;
        }
    }

    __kernel void rgb(__global const float* input,
                      __global float* out)
    {
        const int x = get_global_id(0);
        out[x] = (input[x*3+0]
                + input[x*3+1]
                + input[x*3+2]) / 3.f;
    }
);

struct tile {

    ///////////////
    // constants //
    ///////////////

    int x, y;
    bool use_for_estimation;
    img_t<float> src;
    compute::vector<float> f; // full tiles
    compute::vector<float> f_sqr; // full tiles
    compute::vector<float> h; // half tiles
    compute::vector<std::complex<float>> ff; // fourier full tiles
    compute::vector<std::complex<float>> fh; // fourier half tiles
    compute::vector<float> boxfiltered; // W*W*1
    compute::vector<float> boxfiltered2; // W*W*1
    compute::vector<float> w; // W*W*1
    compute::vector<float> magn; // W*W*1
    compute::vector<float> wblur; // W*W*1
    compute::vector<char> tmpbuf;
    compute::vector<char> tmpbufgray;

    ///////////////////
    // time-variable //
    ///////////////////

    bool valid;
    float dx, dy;
    compute::vector<std::complex<float>> cc; // W*W*1
    compute::vector<float> l2residuals; // W*W*1
    compute::vector<std::complex<float>> rff; // registered fourier full tiles
};

struct image {

    ///////////////
    // constants //
    ///////////////

    int w, h, d, nt;
    img_t<float> src;
    compute::vector<float> dev; // w * h * d
    std::vector<tile> tiles; // nt tiles

    ///////////////////
    // time-variable //
    ///////////////////

    bool allocated;
};

struct result_tile {
    int x, y;

    compute::vector<std::complex<float>> rf;
    compute::vector<std::complex<float>> accum;
    compute::vector<float> accum_weight;
    compute::vector<char> tmpbuf;
};

struct result {
    int w, h, d, nt;
    std::vector<result_tile> tiles;

    compute::vector<float> accumulated; // w * h * d
    compute::vector<float> accumulated_weight; // w * h
};

struct things {
    int W;
    int O;
    float p;

    compute::vector<float> hann;
    compute::vector<float> gaussian;

    compute::command_queue queue;
    clfftPlanHandle ftplan;
    clfftPlanHandle ftplangray;
    compute::program prog;
};

void to_tiles_with_allocation(image& image, things& T)
{
    auto ctx = T.queue.get_context();
    static auto kernel = T.prog.create_kernel("extract_tile");

    for (int y = -T.O; y < image.h; y+=T.O) {
        for (int x = -T.O; x < image.w; x+=T.O) {
            image.tiles.push_back(tile());
            tile& t = image.tiles[image.tiles.size()-1];
            t.x = x;
            t.y = y;
            t.f = compute::vector<float>(T.W * T.W * image.d, ctx);
            kernel.set_args(image.dev, t.f, x, y, image.w, image.h);
            compute::extents<2> offset = dim(0, 0);
            compute::extents<2> ts = dim(T.W, T.W);
            T.queue.enqueue_nd_range_kernel(kernel, 2, offset.data(), ts.data(), 0);

            t.use_for_estimation  = !(t.x+T.W/2 < T.W/2 || t.x+T.W/2 > image.w - T.W/2);
            t.use_for_estimation &= !(t.y+T.W/2 < T.W/2 || t.y+T.W/2 > image.h - T.W/2);
        }
    }
}

void to_tiles_without_allocation(image& image, things& T)
{
    auto ctx = T.queue.get_context();
    static auto kernel = T.prog.create_kernel("extract_tile");

    compute::extents<2> offset = dim(0, 0);
    compute::extents<2> ts = dim(T.W, T.W);
    for (int i = 0; i < image.nt; i++) {
        tile& t = image.tiles[i];
        kernel.set_args(image.dev, t.f, t.x, t.y, image.w, image.h);
        T.queue.enqueue_nd_range_kernel(kernel, 2, offset.data(), ts.data(), 0);
    }
}

void fulltohalf(image& image, things& T)
{
    auto ctx = T.queue.get_context();
    static auto kernel = T.prog.create_kernel("fulltohalf");

    for (int i = 0; i < image.nt; i++) {
        tile& t = image.tiles[i];
        if (!t.use_for_estimation)
            continue;

        kernel.set_args(t.f, t.h);
        compute::extents<2> offset = dim(0, 0);
        compute::extents<2> ts = dim(T.W, T.W);
        T.queue.enqueue_nd_range_kernel(kernel, 2, offset.data(), ts.data(), 0);
    }
}

void fftfull(image& image, things& T)
{
    auto ctx = T.queue.get_context();
    static auto kernel = T.prog.create_kernel("float2complex");

    cl_command_queue q = T.queue;
    int err;
    for (int i = 0; i < image.nt; i++) {
        tile& t = image.tiles[i];

        // XXX: this line is necessary even though I don't understand why
        t.ff = compute::vector<std::complex<float>>(T.W*T.W*image.d, ctx);

        kernel.set_args(t.f, t.ff);
        T.queue.enqueue_1d_range_kernel(kernel, 0, t.ff.size(), 0);

        err = clfftEnqueueTransform(T.ftplan, CLFFT_FORWARD, 1, &q, 0, NULL, NULL,
                                    &t.ff.get_buffer().get(), NULL, t.tmpbuf.get_buffer().get());
        assert(!err);
    }
}

void ffthalf(image& image, things& T)
{
    static auto kernel = T.prog.create_kernel("float2complex");

    cl_command_queue q = T.queue;
    int err;
    for (int i = 0; i < image.nt; i++) {
        tile& t = image.tiles[i];
        if (!t.use_for_estimation)
            continue;

        kernel.set_args(t.h, t.fh);
        T.queue.enqueue_1d_range_kernel(kernel, 0, t.h.size(), 0);

        err = clfftEnqueueTransform(T.ftplan, CLFFT_FORWARD, 1, &q, 0, NULL, NULL,
                                    &t.fh.get_buffer().get(), NULL, t.tmpbuf.get_buffer().get());
        assert(!err);
    }
}

void backtospace(result& image, things& T)
{
    auto ctx = T.queue.get_context();

    cl_command_queue q = T.queue;
    int err;
    for (int i = 0; i < image.nt; i++) {
        result_tile& t = image.tiles[i];

        compute::copy(t.accum.begin(), t.accum.end(), t.rf.begin(), T.queue);

        err = clfftEnqueueTransform(T.ftplan, CLFFT_BACKWARD, 1, &q, 0, NULL, NULL,
                                    &t.rf.get_buffer().get(), NULL, t.tmpbuf.get_buffer().get());
        assert(!err);
    }
}

void weight(image& image, things& T)
{
    auto ctx = T.queue.get_context();
    static auto kernel_magn = T.prog.create_kernel("magnitude");
    static auto kernel_pow = T.prog.create_kernel("pow_");
    static auto kernel_blur = T.prog.create_kernel("blur");

    compute::extents<2> offset = dim(0, 0);
    compute::extents<2> ts = dim(T.W, T.W);
    for (int i = 0; i < image.nt; i++) {
        tile& t = image.tiles[i];

        kernel_magn.set_args(t.ff, t.magn);
        T.queue.enqueue_1d_range_kernel(kernel_magn, 0, t.magn.size(), 0);

        kernel_blur.set_args(t.magn, t.wblur, T.gaussian, (int)T.gaussian.size()/2);
        T.queue.enqueue_nd_range_kernel(kernel_blur, 2, offset.data(), ts.data(), 0);

        kernel_blur.set_args(t.wblur, t.w, T.gaussian, (int)T.gaussian.size()/2);
        T.queue.enqueue_nd_range_kernel(kernel_blur, 2, offset.data(), ts.data(), 0);

        kernel_pow.set_args(t.w, T.p);
        T.queue.enqueue_1d_range_kernel(kernel_pow, 0, t.w.size(), 0);
    }
}

void boxfilter(image& image, things& T)
{
    auto ctx = T.queue.get_context();
    static auto kernel_box = T.prog.create_kernel("boxfilter");
    static auto kernel_sqr = T.prog.create_kernel("sqr");
    static auto kernel_rgb = T.prog.create_kernel("rgb");

    compute::extents<2> offset = dim(0, 0);
    compute::extents<2> ts = dim(T.W, T.W);
    for (int i = 0; i < image.nt; i++) {
        tile& t = image.tiles[i];
        if (!t.use_for_estimation)
            continue;

        kernel_sqr.set_args(t.f, t.f_sqr);
        T.queue.enqueue_1d_range_kernel(kernel_sqr, 0, t.f_sqr.size(), 0);

        kernel_rgb.set_args(t.f_sqr, t.boxfiltered);
        T.queue.enqueue_1d_range_kernel(kernel_rgb, 0, t.boxfiltered.size(), 0);

        kernel_box.set_args(t.boxfiltered, t.boxfiltered2);
        T.queue.enqueue_1d_range_kernel(kernel_box, 0, T.W, 0);

        kernel_box.set_args(t.boxfiltered2, t.boxfiltered);
        T.queue.enqueue_1d_range_kernel(kernel_box, 0, T.W, 0);
    }
}

void l2residuals(struct image& image, const struct image& ref, things& T)
{
    auto ctx = T.queue.get_context();
    static auto kernel_cc = T.prog.create_kernel("crosscorrelation");
    static auto kernel = T.prog.create_kernel("l2residuals");

    for (int i = 0; i < image.nt; i++) {
        tile& t = image.tiles[i];
        const tile& tref = ref.tiles[i];
        if (!t.use_for_estimation)
            continue;

        kernel_cc.set_args(t.ff, tref.fh, t.cc);
        T.queue.enqueue_1d_range_kernel(kernel_cc, 0, t.cc.size(), 0);
    }

    cl_command_queue q = T.queue;
    int err;
    for (int i = 0; i < image.nt; i++) {
        tile& t = image.tiles[i];
        if (!t.use_for_estimation)
            continue;

        err = clfftEnqueueTransform(T.ftplangray, CLFFT_BACKWARD, 1, &q, 0, NULL, NULL,
                                    &t.cc.get_buffer().get(), NULL, t.tmpbufgray.get_buffer().get());
        assert(!err);
    }

    T.queue.finish();

    for (int i = 0; i < image.nt; i++) {
        tile& t = image.tiles[i];
        if (!t.use_for_estimation)
            continue;

        kernel.set_args(t.cc, t.boxfiltered, t.l2residuals);
        T.queue.enqueue_1d_range_kernel(kernel, 0, t.l2residuals.size(), 0);
    }
}

void fetch_translations(struct image& image, things& T)
{
    for (int i = 0; i < image.nt; i++) {
        tile& t = image.tiles[i];
        if (!t.use_for_estimation)
            continue;

        auto it = compute::min_element(t.l2residuals.begin(), t.l2residuals.end(), T.queue);
        int x = std::distance(t.l2residuals.begin(), it);
        t.dx = T.W/2 - x % T.W;
        t.dy = T.W/2 - x / T.W;
        t.dx = -t.dx;
        t.dy = -t.dy;
    }
}

void homshift(struct image& image, homography_t hom, things& T)
{
    auto ctx = T.queue.get_context();
    static auto kernel = T.prog.create_kernel("translate");

    for (int i = 0; i < image.nt; i++) {
        tile& t = image.tiles[i];

        vec2<float> p = {t.x + T.W/2.f, t.y + T.W/2.f};
        vec2<float> d = homography_apply(hom, p) - p;
        d[0] = std::round(d[0]);
        d[1] = std::round(d[1]);

        if (std::abs(d[0]) > T.W/4 || std::abs(d[1]) > T.W/4) {
            t.valid = false;
            continue;
        }
        t.valid = true;

        kernel.set_args(t.ff, t.rff, d[0], d[1]);
        compute::extents<2> offset = dim(0, 0);
        compute::extents<2> ts = dim(T.W, T.W);
        T.queue.enqueue_nd_range_kernel(kernel, 2, offset.data(), ts.data(), 0);
    }

    T.queue.finish();
}

void accumulate(result& image, things& T)
{
    auto ctx = T.queue.get_context();
    static auto kernel = T.prog.create_kernel("accumulate");

    auto barrier = T.queue.enqueue_marker();

    compute::fill(image.accumulated.begin(), image.accumulated.end(), 0.f, T.queue);
    compute::fill(image.accumulated_weight.begin(), image.accumulated_weight.end(), 0.f, T.queue);

    barrier.wait(); // wait for ifft to finish

    for (int i = 0; i < image.nt; i++) {
        result_tile& t = image.tiles[i];
        int x = t.x;
        int y = t.y;

        kernel.set_args(t.rf, T.hann, image.accumulated, image.accumulated_weight, x, y, image.w, image.h);
        compute::extents<2> offset = dim(0, 0);
        compute::extents<2> ts = dim(T.W, T.W);
        T.queue.enqueue_nd_range_kernel(kernel, 2, offset.data(), ts.data(), 0);
    }

    static auto kernel_unweight = T.prog.create_kernel("unweight");
    kernel_unweight.set_args(image.accumulated, image.accumulated_weight);
    T.queue.enqueue_1d_range_kernel(kernel_unweight, 0, image.accumulated_weight.size(), 0);
}

void fba(result& result, std::vector<image*>& images, things& T)
{
    auto ctx = T.queue.get_context();
    static auto kernel = T.prog.create_kernel("fba");
    static auto kernel_unweight = T.prog.create_kernel("cunweight");

    for (int i = 0; i < result.nt; i++) {
        result_tile& tbuf = result.tiles[i];
        compute::fill(tbuf.accum.begin(), tbuf.accum.end(), 0, T.queue);
        compute::fill(tbuf.accum_weight.begin(), tbuf.accum_weight.end(), 0, T.queue);
    }

    for (int j = 0; j < images.size(); j++) {
        image& image = *images[j];
        for (int i = 0; i < image.nt; i++) {
            result_tile& tbuf = result.tiles[i];
            tile& t = image.tiles[i];

            if (!t.valid)  {
                continue;
            }

            kernel.set_args(tbuf.accum, tbuf.accum_weight, t.rff, t.w);
            T.queue.enqueue_1d_range_kernel(kernel, 0, tbuf.accum_weight.size(), 0);
        }
    }

    for (int i = 0; i < result.nt; i++) {
        result_tile& tbuf = result.tiles[i];
        kernel_unweight.set_args(tbuf.accum, tbuf.accum_weight);
        T.queue.enqueue_1d_range_kernel(kernel_unweight, 0, tbuf.accum_weight.size(), 0);
    }
}

void register_all(image& ref, std::vector<image*>& images, things& T)
{
    for (unsigned i = 0; i < images.size(); i++) {
        if (images[i] != &ref)
            l2residuals(*images[i], ref, T);
    }
    for (unsigned i = 0; i < images.size(); i++) {
        if (images[i] != &ref)
            fetch_translations(*images[i], T);
    }
    for (unsigned i = 0; i < images.size(); i++) {
        if (images[i] == &ref)
            continue;

        int W = T.W;
        std::vector<vec2<vec2<double>>> translations;
        for (int j = 0; j < images[i]->nt; j++) {
            tile& t = images[i]->tiles[j];
            if (!t.use_for_estimation)
                continue;
            vec2<vec2<double>> tr;
            tr[0] = vec2<double>(t.x + T.W/2, t.y + T.W/2);
            tr[1] = vec2<double>(t.x + T.W/2 + t.dx, t.y + T.W/2 + t.dy);
            translations.push_back(tr);
        }

        homography_t H = homography_from_translations_robust(translations);
        homshift(*images[i], H, T);
    }

    for (int j = 0; j < ref.nt; j++) {
        tile& t = ref.tiles[j];
        t.valid = true;
        compute::copy_async(t.ff.begin(), t.ff.end(), t.rff.begin(), T.queue);
    }
}

void fuse_all(result& result, std::vector<image*>& images, things& T)
{
    fba(result, images, T);
    backtospace(result, T);
    accumulate(result, T);
}

void initialize_result(result& result, const image& image, things& T)
{
    auto ctx = T.queue.get_context();

    result.w = image.w;
    result.h = image.h;
    result.d = image.d;
    result.nt = image.nt;
    result.accumulated = compute::vector<float>(image.w*image.h*image.d, ctx);
    result.accumulated_weight = compute::vector<float>(image.w*image.h, ctx);

    result.tiles.resize(image.nt);
    for (int t = 0; t < image.nt; t++) {
        auto& tt = result.tiles[t];
        auto& ti = image.tiles[t];

        tt.x = ti.x;
        tt.y = ti.y;
        tt.rf = compute::vector<std::complex<float>>(T.W*T.W*image.d, ctx);
        tt.accum = compute::vector<std::complex<float>>(T.W*T.W*image.d, ctx);
        tt.accum_weight = compute::vector<float>(T.W*T.W, ctx);
        size_t size;
        clfftGetTmpBufSize(T.ftplan, &size);
        tt.tmpbuf = compute::vector<char>(size, ctx);
    }
}

void prepare_image(image& image, const img_t<float>& img, things& T)
{
    auto ctx = T.queue.get_context();
    image.src = img;
    image.w = img.w;
    image.h = img.h;
    image.d = img.d;

    if (!image.allocated) {
        image.dev = img_to_device(img, T.queue);
        to_tiles_with_allocation(image, T);
        image.nt = image.tiles.size();

        for (int t = 0; t < image.nt; t++) {
            auto& tt = image.tiles[t];
            tt.h = compute::vector<float>(T.W*T.W*image.d, ctx);
            tt.f_sqr = compute::vector<float>(T.W*T.W*image.d, ctx);
            tt.ff = compute::vector<std::complex<float>>(T.W*T.W*image.d, ctx);
            tt.fh = compute::vector<std::complex<float>>(T.W*T.W*image.d, ctx);
            tt.cc = compute::vector<std::complex<float>>(T.W*T.W, ctx);
            tt.l2residuals = compute::vector<float>(T.W*T.W, ctx);
            tt.rff = compute::vector<std::complex<float>>(T.W*T.W*image.d, ctx);
            tt.w = compute::vector<float>(T.W*T.W, ctx);
            tt.wblur = compute::vector<float>(T.W*T.W, ctx);
            tt.magn = compute::vector<float>(T.W*T.W, ctx);
            tt.boxfiltered = compute::vector<float>(T.W*T.W, ctx);
            tt.boxfiltered2 = compute::vector<float>(T.W*T.W, ctx);
            size_t size;
            clfftGetTmpBufSize(T.ftplan, &size);
            tt.tmpbuf = compute::vector<char>(size, ctx);
            clfftGetTmpBufSize(T.ftplangray, &size);
            tt.tmpbufgray = compute::vector<char>(size, ctx);
        }
        image.allocated = true;
    } else {
        compute::copy(img.data.begin(), img.data.end(), image.dev.begin(), T.queue);
        to_tiles_without_allocation(image, T);
    }

    fulltohalf(image, T);
    fftfull(image, T);
    ffthalf(image, T);
    weight(image, T);
    boxfilter(image, T);
}

int main(int argc, char** argv)
{
    if (argc < 2 || argc > 4) {
        return fprintf(stderr, "usage: %s <output_fmt> [file_of_inputs (stdin)]\n", argv[0]), 1;
    }

    char* output_fmt = argv[1];
    FILE* inputs = stdin;
    if (argc == 3) {
        inputs = fopen(argv[2], "r");
        if (!inputs) {
            return perror(argv[2]), 1;
        }
    }

    compute::device device = compute::system::default_device();
    std::cout << "device: " << device.name() << std::endl;

    compute::context ctx(device);

    int W = 256;

    things things;
    things.queue = compute::command_queue(ctx, device);
    things.W = W;
    things.O = W/3;
    things.p = 3;

    {
        img_t<float> _hann;
        ::hann(_hann, W/2, W/2);
        img_t<float> hann(W, W);
        hann.set_value(0);
        for (int y = 0; y < W/2; y++) {
            for (int x = 0; x < W/2; x++) {
                hann(W/4 + x, W/4 + y) = _hann(x, y);
            }
        }
        things.hann = img_to_device(hann, things.queue);
    }

    {
        float sigma = things.W / 50.f;
        std::vector<float> gaussian(21);
        float sum = 0.f;
        for (int x = 0; x < (int) gaussian.size(); x++) {
            gaussian[x] = 1.f/std::sqrt(2*M_PI*sigma*sigma)
                        * std::exp(- std::pow((float)(x-(int)gaussian.size()/2), 2.f) / (2*sigma*sigma));
            sum += gaussian[x];
        }
        for (unsigned x = 0; x < gaussian.size(); x++) {
            gaussian[x] /= sum;
        }
        things.gaussian = compute::vector<float>(gaussian.begin(), gaussian.end(), things.queue);
    }

    things.prog = compute::program::build_with_source(kernels_src, ctx,
                                                      "-D W=" + std::to_string(W)
                                                      + " -D hw=" + std::to_string(W/4));

    {
        clfftPlanHandle planHandle;
        clfftDim dim = CLFFT_2D;
        size_t clLengths[2] = {(size_t)W, (size_t)W};
        size_t strides[] = {(size_t)3, (size_t)W*3};
        int err;
        clfftSetupData fftSetup;
        err = clfftInitSetupData(&fftSetup);
        err = clfftSetup(&fftSetup);
        err = clfftCreateDefaultPlan(&things.ftplan, ctx, dim, clLengths);
        err = clfftSetPlanPrecision(things.ftplan, CLFFT_SINGLE);
        err = clfftSetLayout(things.ftplan, CLFFT_COMPLEX_INTERLEAVED, CLFFT_COMPLEX_INTERLEAVED);
        err = clfftSetResultLocation(things.ftplan, CLFFT_INPLACE);
        err = clfftSetPlanInStride(things.ftplan, dim, strides);
        err = clfftSetPlanOutStride(things.ftplan, dim, strides);
        err = clfftSetPlanBatchSize(things.ftplan, 3);
        err = clfftSetPlanDistance(things.ftplan, 1, 1);
        cl_command_queue q = things.queue;
        err = clfftBakePlan(things.ftplan, 1, &q, NULL, NULL);
        assert(!err);
    }
    {
        clfftPlanHandle planHandle;
        clfftDim dim = CLFFT_2D;
        size_t clLengths[2] = {(size_t)W, (size_t)W};
        size_t strides[] = {(size_t)1, (size_t)W};
        int err;
        clfftSetupData fftSetup;
        err = clfftInitSetupData(&fftSetup);
        err = clfftSetup(&fftSetup);
        err = clfftCreateDefaultPlan(&things.ftplangray, ctx, dim, clLengths);
        err = clfftSetPlanPrecision(things.ftplangray, CLFFT_SINGLE);
        err = clfftSetLayout(things.ftplangray, CLFFT_COMPLEX_INTERLEAVED, CLFFT_COMPLEX_INTERLEAVED);
        err = clfftSetResultLocation(things.ftplangray, CLFFT_INPLACE);
        err = clfftSetPlanInStride(things.ftplangray, dim, strides);
        err = clfftSetPlanOutStride(things.ftplangray, dim, strides);
        err = clfftSetPlanBatchSize(things.ftplangray, 1);
        err = clfftSetPlanDistance(things.ftplangray, 1, 1);
        cl_command_queue q = things.queue;
        err = clfftBakePlan(things.ftplangray, 1, &q, NULL, NULL);
        assert(!err);
    }

    std::vector<struct image*> images;
    struct result result;

    int N = 3;
    int cur = 0;
    int i = 0;
    char file[2048];
    while (fgets(file, sizeof(file), inputs) && file[0]) {
        file[strlen(file) - 1] = 0;
        img_t<float> img = img_t<float>::load(file);
        float max = img.max();
        for (auto& v : img.data) v /= max;

        if (images.size() < N) {
            images.push_back(new struct image);
            images[images.size()-1]->allocated = false;
            cur = images.size() - 1;
        } else {
            cur = (cur + 1) % images.size();
        }

        prepare_image(*images[cur], img, things);

        if (images.size() == 1)
            initialize_result(result, *images[0], things);

        register_all(*images[cur], images, things);

        fuse_all(result, images, things);

        auto accumulated = img_from_device(result.accumulated, img.w, img.h, img.d, things.queue);
        for (auto& v : accumulated.data) v *= max;
        std::string output = string_format(output_fmt, i);
        accumulated.save(output);
        printf("%s\n", output.c_str());
        i++;
    }

    int err = clfftDestroyPlan(&things.ftplan);
    err = clfftDestroyPlan(&things.ftplangray);
    clfftTeardown();
    return 0;
}


#pragma once

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

namespace northstar {
    namespace math {
        namespace types {
            typedef Eigen::Transform<double, 3, Eigen::TransformTraits::Affine, Eigen::ColMajor> AffineMatrix4d;
            typedef Eigen::Transform<double, 3, Eigen::TransformTraits::Isometry, Eigen::ColMajor> IsometryMatrix4d;
            typedef Eigen::Transform<double, 3, Eigen::TransformTraits::AffineCompact, Eigen::ColMajor> AffineMatrix4x3d;
            typedef Eigen::Transform<double, 3, Eigen::TransformTraits::Projective, Eigen::ColMajor> ProjMatrix4d;
            typedef Eigen::Translation<double, 3> Translation3d;
            typedef Eigen::Vector<double, 2> Vector2d;
            typedef Eigen::Vector<double, 3> Vector3d;
            typedef Eigen::Vector<double, 4> Vector4d;
            typedef Eigen::ParametrizedLine<double, 3> Ray3d;
            typedef Eigen::Hyperplane<double, 3> Plane3d;
            typedef Eigen::Quaternion<double> Quaterniond;
            typedef Eigen::Vector<int32_t, 2> Vector2i;
            typedef Eigen::AngleAxis<double> AngleAxisd;

            struct SPose {
                Vector3d position;
                Quaterniond orientation;
            };

            struct SSphere {
                Vector3d origin;
                double radius;
            };

            // NOTE: this hash function gets worse with the number of elements being iterated
            template<typename T>
            struct SHasher {
                std::size_t operator()(const T & IterativeEigen) const {
                    return std::transform_reduce(
                        IterativeEigen.begin(),
                        IterativeEigen.end(),
                        static_cast<double>(0.0),
                        std::bit_xor<std::size_t>(),
                        std::hash<typename T::Scalar>());
                }
            };
        }
    }
}


// Copyright (c) 2020 Chris Richardson
// FEniCS Project
// SPDX-License-Identifier:    MIT

#include "raviart-thomas.h"
#include "dof-permutations.h"
#include "lagrange.h"
#include "moments.h"
#include "polyset.h"
#include "quadrature.h"
#include <Eigen/Dense>
#include <numeric>
#include <vector>

using namespace libtab;

//----------------------------------------------------------------------------
FiniteElement libtab::create_rt(cell::type celltype, int degree,
                                const std::string& name)
{
  if (celltype != cell::type::triangle and celltype != cell::type::tetrahedron)
    throw std::runtime_error("Unsupported cell type");

  const int tdim = cell::topological_dimension(celltype);

  const cell::type facettype
      = (tdim == 2) ? cell::type::interval : cell::type::triangle;

  // The number of order (degree-1) scalar polynomials
  const int nv = polyset::dim(celltype, degree - 1);
  // The number of order (degree-2) scalar polynomials
  const int ns0 = polyset::dim(celltype, degree - 2);
  // The number of additional polnomials in the polynomial basis for
  // Raviart-Thomas
  const int ns = polyset::dim(facettype, degree - 1);

  // Evaluate the expansion polynomials at the quadrature points
  auto [Qpts, Qwts] = quadrature::make_quadrature(celltype, 2 * degree);
  Eigen::ArrayXXd Pkp1_at_Qpts
      = polyset::tabulate(celltype, degree, 0, Qpts)[0];

  // The number of order (degree) polynomials
  const int psize = Pkp1_at_Qpts.cols();

  // Create coefficients for order (degree-1) vector polynomials
  Eigen::MatrixXd wcoeffs = Eigen::MatrixXd::Zero(nv * tdim + ns, psize * tdim);
  for (int j = 0; j < tdim; ++j)
  {
    wcoeffs.block(nv * j, psize * j, nv, nv)
        = Eigen::MatrixXd::Identity(nv, nv);
  }

  // Create coefficients for additional polynomials in Raviart-Thomas
  // polynomial basis
  for (int i = 0; i < ns; ++i)
  {
    for (int k = 0; k < psize; ++k)
    {
      for (int j = 0; j < tdim; ++j)
      {
        const double w_sum = (Qwts * Pkp1_at_Qpts.col(ns0 + i) * Qpts.col(j)
                              * Pkp1_at_Qpts.col(k))
                                 .sum();
        wcoeffs(nv * tdim + i, k + psize * j) = w_sum;
      }
    }
  }

  // Dual space
  Eigen::MatrixXd dual = Eigen::MatrixXd::Zero(nv * tdim + ns, psize * tdim);

  // quadrature degree
  int quad_deg = 5 * degree;

  // Add rows to dualmat for integral moments on facets
  const int facet_count = tdim + 1;
  const int facet_dofs = ns;
  dual.block(0, 0, facet_count * facet_dofs, psize * tdim)
      = moments::make_normal_integral_moments(
          create_dlagrange(facettype, degree - 1), celltype, tdim, degree,
          quad_deg);

  // Add rows to dualmat for integral moments on interior
  if (degree > 1)
  {
    const int internal_dofs = tdim * ns0;
    // Interior integral moment
    dual.block(facet_count * facet_dofs, 0, internal_dofs, psize * tdim)
        = moments::make_integral_moments(create_dlagrange(celltype, degree - 2),
                                         celltype, tdim, degree, quad_deg);
  }

  const std::vector<std::vector<std::vector<int>>> topology
      = cell::topology(celltype);

  const int ndofs = dual.rows();
  int perm_count = 0;
  for (int i = 1; i < tdim; ++i)
    perm_count += topology[i].size() * i;

  std::vector<Eigen::MatrixXd> base_permutations(
      perm_count, Eigen::MatrixXd::Identity(ndofs, ndofs));
  if (tdim == 2)
  {
    Eigen::ArrayXi edge_ref = dofperms::interval_reflection(degree - 1);
    for (int edge = 0; edge < facet_count; ++edge)
    {
      const int start = edge_ref.size() * edge;
      for (int i = 0; i < edge_ref.size(); ++i)
      {
        base_permutations[edge](start + i, start + i) = 0;
        base_permutations[edge](start + i, start + edge_ref[i]) = -1;
      }
    }
  }
  else if (tdim == 3)
  {
    Eigen::ArrayXi face_ref = dofperms::triangle_reflection(degree - 1);
    Eigen::ArrayXi face_rot = dofperms::triangle_rotation(degree - 1);

    for (int face = 0; face < facet_count; ++face)
    {
      const int start = face_ref.size() * face;
      for (int i = 0; i < face_rot.size(); ++i)
      {
        base_permutations[2 * face](start + i, start + i) = 0;
        base_permutations[2 * face](start + i, start + face_rot[i]) = 1;
        base_permutations[2 * face + 1](start + i, start + i) = 0;
        base_permutations[2 * face + 1](start + i, start + face_ref[i]) = -1;
      }
    }
  }

  // Raviart-Thomas has ns dofs on each facet, and ns0*tdim in the interior
  std::vector<std::vector<int>> entity_dofs(topology.size());
  for (int i = 0; i < tdim - 1; ++i)
    entity_dofs[i].resize(topology[i].size(), 0);
  entity_dofs[tdim - 1].resize(topology[tdim - 1].size(), ns);
  entity_dofs[tdim] = {ns0 * tdim};

  Eigen::MatrixXd coeffs = compute_expansion_coefficients(wcoeffs, dual);
  return FiniteElement(name, celltype, degree, {tdim}, coeffs, entity_dofs,
                       base_permutations);
}
//-----------------------------------------------------------------------------


/**
 * @file init_rules_test.cpp
 * @author Marcus Edel
 *
 * Tests for the various weight initialize methods.
 *
 * mlpack is free software; you may redistribute it and/or modify it under the
 * terms of the 3-clause BSD license.  You should have received a copy of the
 * 3-clause BSD license along with mlpack.  If not, see
 * http://www.opensource.org/licenses/BSD-3-Clause for more information.
 */
#include <mlpack/core.hpp>
#include <mlpack/core/math/random.hpp>

#include <mlpack/methods/ann/layer/layer.hpp>
#include <mlpack/methods/ann/layer/layer_types.hpp>
#include <mlpack/methods/ann/ffn.hpp>

#include <mlpack/methods/ann/init_rules/kathirvalavakumar_subavathi_init.hpp>
#include <mlpack/methods/ann/init_rules/nguyen_widrow_init.hpp>
#include <mlpack/methods/ann/init_rules/oivs_init.hpp>
#include <mlpack/methods/ann/init_rules/orthogonal_init.hpp>
#include <mlpack/methods/ann/init_rules/random_init.hpp>
#include <mlpack/methods/ann/init_rules/const_init.hpp>
#include <mlpack/methods/ann/init_rules/gaussian_init.hpp>

#include <boost/test/unit_test.hpp>
#include "test_tools.hpp"

using namespace mlpack;
using namespace mlpack::ann;

BOOST_AUTO_TEST_SUITE(InitRulesTest);

/**
 * Test the RandomInitialization class with a constant value.
 */
BOOST_AUTO_TEST_CASE(ConstantInitTest)
{
  arma::mat weights;
  RandomInitialization constantInit(1, 1);
  constantInit.Initialize(weights, 100, 100);

  bool b = arma::all(arma::vectorise(weights) == 1);
  BOOST_REQUIRE_EQUAL(b, 1);
}

/**
 * Simple test of the OrthogonalInitialization class with two different
 * sizes.
 */
BOOST_AUTO_TEST_CASE(OrthogonalInitTest)
{
  arma::mat weights;
  OrthogonalInitialization orthogonalInit;
  orthogonalInit.Initialize(weights, 100, 200);

  arma::mat orthogonalWeights = arma::eye<arma::mat>(100, 100);
  weights *= weights.t();

  for (size_t i = 0; i < weights.n_rows; i++)
    for (size_t j = 0; j < weights.n_cols; j++)
      BOOST_REQUIRE_SMALL(weights.at(i, j) - orthogonalWeights.at(i, j), 1e-3);

  orthogonalInit.Initialize(weights, 200, 100);
  weights = weights.t() * weights;

  for (size_t i = 0; i < weights.n_rows; i++)
    for (size_t j = 0; j < weights.n_cols; j++)
      BOOST_REQUIRE_SMALL(weights.at(i, j) - orthogonalWeights.at(i, j), 1e-3);
}

/**
 * Test the OrthogonalInitialization class with a non default gain.
 */
BOOST_AUTO_TEST_CASE(OrthogonalInitGainTest)
{
  arma::mat weights;

  const double gain = 2;
  OrthogonalInitialization orthogonalInit(gain);
  orthogonalInit.Initialize(weights, 100, 200);

  arma::mat orthogonalWeights = arma::eye<arma::mat>(100, 100);
  orthogonalWeights *= (gain * gain);
  weights *= weights.t();

  for (size_t i = 0; i < weights.n_rows; i++)
    for (size_t j = 0; j < weights.n_cols; j++)
      BOOST_REQUIRE_SMALL(weights.at(i, j) - orthogonalWeights.at(i, j), 1e-3);
}

/**
 * Test the ConstInitialization class. If you think about it, it's kind of
 * ridiculous to test the const init rule. But at least we make sure it
 * builds without any problems.
 */
BOOST_AUTO_TEST_CASE(ConstInitTest)
{
  arma::mat weights;
  ConstInitialization zeroInit(0);
  zeroInit.Initialize(weights, 100, 100);

  bool b = arma::all(arma::vectorise(weights) == 0);
  BOOST_REQUIRE_EQUAL(b, 1);
}

/*
 * Simple test of the KathirvalavakumarSubavathiInitialization class with
 * two different sizes.
 */
BOOST_AUTO_TEST_CASE(KathirvalavakumarSubavathiInitTest)
{
  arma::mat data = arma::randu<arma::mat>(100, 1);

  arma::mat weights;
  arma::cube weights3d;

  KathirvalavakumarSubavathiInitialization kathirvalavakumarSubavathiInit(
      data, 1.5);

  kathirvalavakumarSubavathiInit.Initialize(weights, 100, 100);
  kathirvalavakumarSubavathiInit.Initialize(weights3d, 100, 100, 2);

  BOOST_REQUIRE_EQUAL(weights.n_rows, 100);
  BOOST_REQUIRE_EQUAL(weights.n_cols, 100);

  BOOST_REQUIRE_EQUAL(weights3d.n_rows, 100);
  BOOST_REQUIRE_EQUAL(weights3d.n_cols, 100);
  BOOST_REQUIRE_EQUAL(weights3d.n_slices, 2);
}

/**
 * Simple test of the NguyenWidrowInitialization class.
 */
BOOST_AUTO_TEST_CASE(NguyenWidrowInitTest)
{
  arma::mat weights;
  arma::cube weights3d;

  NguyenWidrowInitialization nguyenWidrowInit;

  nguyenWidrowInit.Initialize(weights, 100, 100);
  nguyenWidrowInit.Initialize(weights3d, 100, 100, 2);

  BOOST_REQUIRE_EQUAL(weights.n_rows, 100);
  BOOST_REQUIRE_EQUAL(weights.n_cols, 100);

  BOOST_REQUIRE_EQUAL(weights3d.n_rows, 100);
  BOOST_REQUIRE_EQUAL(weights3d.n_cols, 100);
  BOOST_REQUIRE_EQUAL(weights3d.n_slices, 2);
}

/**
 * Simple test of the OivsInitialization class with two different sizes.
 */
BOOST_AUTO_TEST_CASE(OivsInitTest)
{
  arma::mat weights;
  arma::cube weights3d;

  OivsInitialization<> oivsInit;

  oivsInit.Initialize(weights, 100, 100);
  oivsInit.Initialize(weights3d, 100, 100, 2);

  BOOST_REQUIRE_EQUAL(weights.n_rows, 100);
  BOOST_REQUIRE_EQUAL(weights.n_cols, 100);

  BOOST_REQUIRE_EQUAL(weights3d.n_rows, 100);
  BOOST_REQUIRE_EQUAL(weights3d.n_cols, 100);
  BOOST_REQUIRE_EQUAL(weights3d.n_slices, 2);
}

/**
 * Simple test of the GaussianInitialization class.
 */
BOOST_AUTO_TEST_CASE(GaussianInitTest)
{
  const size_t rows = 7;
  const size_t cols = 8;
  const size_t slices = 2;

  arma::mat weights;
  arma::cube weights3d;

  GaussianInitialization t(0, 0.2);

  t.Initialize(weights, rows, cols);
  t.Initialize(weights3d, rows, cols, slices);

  BOOST_REQUIRE_EQUAL(weights.n_rows, rows);
  BOOST_REQUIRE_EQUAL(weights.n_cols, cols);

  BOOST_REQUIRE_EQUAL(weights3d.n_rows, rows);
  BOOST_REQUIRE_EQUAL(weights3d.n_cols, cols);
  BOOST_REQUIRE_EQUAL(weights3d.n_slices, slices);
}

/**
 * Simple test of the NetworkInitialization class, we test it with every
 * implemented initialization rule and make sure the output is reasonable.
 */
BOOST_AUTO_TEST_CASE(NetworkInitTest)
{
  arma::mat input = arma::ones(5, 1);
  arma::mat response;
  NegativeLogLikelihood<> outputLayer;

  // Create a simple network and use the RandomInitialization rule to
  // initialize the network parameters.
  RandomInitialization randomInit(0.5, 0.5);

  FFN<NegativeLogLikelihood<>, RandomInitialization> randomModel(
      std::move(outputLayer), randomInit);
  randomModel.Add<IdentityLayer<> >();
  randomModel.Add<Linear<> >(5, 5);
  randomModel.Add<Linear<> >(5, 2);
  randomModel.Add<LogSoftMax<> >();
  randomModel.Predict(input, response);

  bool b = arma::all(arma::vectorise(randomModel.Parameters()) == 0.5);
  BOOST_REQUIRE_EQUAL(b, 1);
  BOOST_REQUIRE_EQUAL(randomModel.Parameters().n_elem, 42);

  // Create a simple network and use the OrthogonalInitialization rule to
  // initialize the network parameters.
  FFN<NegativeLogLikelihood<>, OrthogonalInitialization> orthogonalModel;
  orthogonalModel.Add<IdentityLayer<> >();
  orthogonalModel.Add<Linear<> >(5, 5);
  orthogonalModel.Add<Linear<> >(5, 2);
  orthogonalModel.Add<LogSoftMax<> >();
  orthogonalModel.Predict(input, response);

  BOOST_REQUIRE_EQUAL(orthogonalModel.Parameters().n_elem, 42);

  // Create a simple network and use the ZeroInitialization rule to
  // initialize the network parameters.
  FFN<NegativeLogLikelihood<>, ConstInitialization>
    zeroModel(NegativeLogLikelihood<>(), ConstInitialization(0));
  zeroModel.Add<IdentityLayer<> >();
  zeroModel.Add<Linear<> >(5, 5);
  zeroModel.Add<Linear<> >(5, 2);
  zeroModel.Add<LogSoftMax<> >();
  zeroModel.Predict(input, response);

  BOOST_REQUIRE_EQUAL(arma::accu(zeroModel.Parameters()), 0);
  BOOST_REQUIRE_EQUAL(zeroModel.Parameters().n_elem, 42);

  // Create a simple network and use the
  // KathirvalavakumarSubavathiInitialization rule to initialize the network
  // parameters.
  KathirvalavakumarSubavathiInitialization kathirvalavakumarSubavathiInit(
      input, 1.5);
  FFN<NegativeLogLikelihood<>, KathirvalavakumarSubavathiInitialization>
      ksModel(std::move(outputLayer), kathirvalavakumarSubavathiInit);
  ksModel.Add<IdentityLayer<> >();
  ksModel.Add<Linear<> >(5, 5);
  ksModel.Add<Linear<> >(5, 2);
  ksModel.Add<LogSoftMax<> >();
  ksModel.Predict(input, response);

  BOOST_REQUIRE_EQUAL(ksModel.Parameters().n_elem, 42);

  // Create a simple network and use the OivsInitialization rule to
  // initialize the network parameters.
  FFN<NegativeLogLikelihood<>, OivsInitialization<> > oivsModel;
  oivsModel.Add<IdentityLayer<> >();
  oivsModel.Add<Linear<> >(5, 5);
  oivsModel.Add<Linear<> >(5, 2);
  oivsModel.Add<LogSoftMax<> >();
  oivsModel.Predict(input, response);

  BOOST_REQUIRE_EQUAL(oivsModel.Parameters().n_elem, 42);

  // Create a simple network and use the GaussianInitialization rule to
  // initialize the network parameters.
  FFN<NegativeLogLikelihood<>, GaussianInitialization> gaussianModel;
  gaussianModel.Add<IdentityLayer<> >();
  gaussianModel.Add<Linear<> >(5, 5);
  gaussianModel.Add<Linear<> >(5, 2);
  gaussianModel.Add<LogSoftMax<> >();
  gaussianModel.Predict(input, response);

  BOOST_REQUIRE_EQUAL(gaussianModel.Parameters().n_elem, 42);
}

BOOST_AUTO_TEST_SUITE_END();


#include "fieldtrack/DimensionParser.h"
#include "argus_utils/geometry/PoseSE2.h"
#include "argus_utils/geometry/PoseSE3.h"

#include <boost/algorithm/string.hpp>

namespace argus
{

unsigned int twod_to_threed( unsigned int d )
{
	// NOTE Integer division rounding down
	unsigned int order = d / PoseSE2::TangentDimension;
	unsigned int ind = d - order * PoseSE2::TangentDimension;
	if( ind == 2 ) { ind = 5; } // 2D angle maps to 3D z angle
	return order * PoseSE3::TangentDimension + ind;
}

MatrixType promote_3d_matrix( bool twoDimensional, unsigned int order )
{
	unsigned int fullDim = PoseSE3::TangentDimension * (order + 1);
	if( !twoDimensional )
	{
		return MatrixType::Identity( fullDim, fullDim );
	}

	unsigned int dim = twoDimensional ? PoseSE2::TangentDimension
							: PoseSE3::TangentDimension;
	unsigned int inDim = dim * (order + 1);

	MatrixType p = MatrixType::Zero( fullDim, inDim );

	for( unsigned int i = 0; i < inDim; ++i )
	{
		p( twod_to_threed(i), i ) = 1;
	}
	return p;
}

unsigned int parse_dim_string( const std::string& s,
                               bool twoDimensional,
                               unsigned int maxOrder,
                               unsigned int minOrder )
{
	std::vector<std::string> splits;
	boost::split( splits, s, boost::is_any_of( "_" ), boost::token_compress_on );
	if( splits.size() != 3 )
	{
		throw std::invalid_argument( "String " + s + " invalid format" );
	}

	std::string typeStr = splits[0];
	std::string dimStr = splits[1];

	unsigned int order = std::stoi( splits[2] ); // NOTE Should really use stoul
	if( order < minOrder || order > maxOrder )
	{
		throw std::invalid_argument( "String " + s + " order exceeds order limits" );
	}
	order = order - minOrder;

	unsigned int typeInd, dimInd;
	if( typeStr == "pos" )
	{
		typeInd = 0;
	}
	else if( typeStr == "ori" )
	{
		typeInd = twoDimensional ? 0 : 3;
	}
	else
	{
		throw std::invalid_argument( "String " + s + " type invalid" );
	}

	if( dimStr == "x" )
	{
		dimInd = 0;
	}
	else if( dimStr == "y" )
	{
		dimInd = 1;
	}
	else if( dimStr == "z" )
	{
		dimInd = 2;
	}
	else
	{
		throw std::invalid_argument( "String " + s + " dim invalid" );
	}

	if( twoDimensional )
	{
		if( (typeStr == "pos" && dimStr == "z") ||
		    (typeStr == "ori" && dimStr == "x") ||
		    (typeStr == "ori" && dimStr == "y") )
		{
			throw std::invalid_argument( "String " + s + " invalid in 2D mode" );
		}
	}

	unsigned int baseDim = twoDimensional ? 3 : 6;

	return baseDim * order + typeInd + dimInd;
}

std::vector<unsigned int> parse_dim_string( const std::vector<std::string>& s,
                                            bool twoDimensional,
                                            unsigned int maxOrder,
                                            unsigned int minOrder )
{
	std::vector<unsigned int> ret;
	ret.reserve( s.size() );
	for( unsigned int i = 0; i < s.size(); ++i )
	{
		ret.push_back( parse_dim_string( s[i], twoDimensional, maxOrder, minOrder ) );
	}
	return ret;
}
}


#include "RANSAC.h"

#include "Line.h"
#include <Eigen/Core>
#include <opencv2/highgui/highgui.hpp>

#include <ctime>

namespace RANSAC
{

using Vector = Eigen::Vector2f;
using Line = localization::Line;

Vector findBestConvergencePoint(const std::vector<Line>& lines, int k)
{
    std::srand(std::time(NULL));
    float bestDistance = std::numeric_limits<float>::max();
    Vector bestConvergencePoint(0.0, 0.0);
    for (int i = 0; i < k; ++i) {
        const Line& firstLine = lines[std::rand() % lines.size()];
        const Line& otherLine = lines[std::rand() % lines.size()];

        Vector intersection;
        bool isIntersecting = firstLine.findIntersection(otherLine, intersection);
        if (isIntersecting) {
            double distance = summConvergeDistance(lines, firstLine, intersection);
            distance += summConvergeDistance(lines, otherLine, intersection);
            if (distance < bestDistance) {
                bestConvergencePoint = intersection;
                bestDistance = distance;
            }
        }
        // else handle parallel lines here.
    }
    return bestConvergencePoint;
}

float summConvergeDistance(const std::vector<Line>& lines, const Line& pivot, const Vector& convergence)
{
    float distance = 0.0;
    for (int i = 0; i < lines.size(); ++i) {
        Vector intersection;
        bool isIntersecting = lines[i].findIntersection(pivot, intersection);
        if (isIntersecting) {
            distance += (intersection - convergence).squaredNorm();
        }
    }
    return distance;
}

bool findBestOrientationSplit(const std::vector<localization::Line>& lines, int k, Vector bestOrientations[2])
{
    std::srand(std::time(NULL));
    float bestFitWeight = 0.0;
    for (int i = 0; i < k; ++i) {
        const Line& line = lines[std::rand() % lines.size()];
        Vector orientation = line.getOrientation();

        Vector orientations[2] = { orientation,
                                   { orientation.y(),
                                    -orientation.x() }};
        float fitWeight = evaluateOrientationModel(lines, orientations);
        if (fitWeight > bestFitWeight) {
            bestFitWeight = fitWeight;
            bestOrientations[0] = orientations[0];
            bestOrientations[1] = orientations[1];
        }
    }
    return true;
}

float evaluateOrientationModel(const std::vector<localization::Line>& lines, Eigen::Vector2f orientations[2])
{
    double dotSumm = 0.0;
    for (int i = 0; i < lines.size(); ++i) {
        Vector u = lines[i].getOrientation();

        double udotv = std::abs(u.dot(orientations[0]));
        double udotw = std::abs(u.dot(orientations[1]));

        bool isCloserToV = udotv > udotw;
        dotSumm += (isCloserToV) ? udotv : udotw;
        int iOrientation = (isCloserToV) ? 0 : 1;
    }
    return dotSumm;
}

}


#include "Geometry.h"
#include <cmath>
//#include <boost/polygon/polygon.hpp>
#include <boost/geometry.hpp>
#include <boost/geometry/geometries/point_xy.hpp>
#include <boost/geometry/geometries/polygon.hpp>

//////////////////////////////////////////////////////////////////////////////////////////////////////////////////
Point2D::Point2D(void) { ; }
Point2D::Point2D(coord_t _x, coord_t _y) : x(_x), y(_y) { ; }
Vector2D Point2D::operator - (const Point2D& p) const { return Vector2D(x - p.x, y - p.y); }
Point2D Point2D::operator + (const Vector2D& v) const { return Point2D(x + v.x, y + v.y); }
Point2D Point2D::operator - (const Vector2D& v) const { return Point2D(x - v.x, y - v.y); }
bool Point2D::operator == (const Point2D& v) const {
	return abs(x - v.x) <= EPSILON
		&& abs(y - v.y) <= EPSILON;
}
//////////////////////////////////////////////////////////////////////////////////////////////////////////////////
Point3D::Point3D(void) { ; }
Point3D::Point3D(coord_t _x, coord_t _y, coord_t _z) :x(_x), y(_y), z(_z) { ; }

Vector3D Point3D::to_vector(void) const {
	return Vector3D(x, y, z);
}

Point3D Point3D::operator - (void) const { return Point3D(-x, -y, -z); }
Vector3D Point3D::operator - (const Point3D& p) const { return Vector3D(x - p.x, y - p.y, z - p.z); }
Point3D Point3D::operator + (const Vector3D& v) const { return Point3D(x + v.x, y + v.y, z + v.z); }
Point3D Point3D::operator - (const Vector3D& v) const { return Point3D(x - v.x, y - v.y, z - v.z); }
//bool Point3D::operator == (const Point3D& v) const { return x == v.x && y == v.y && z == v.z;  }
bool Point3D::operator == (const Point3D& v) const {
	return abs(x - v.x) <= EPSILON
		&& abs(y - v.y) <= EPSILON
		&& abs(z - v.z) <= EPSILON;
}
//////////////////////////////////////////////////////////////////////////////////////////////////////////////////
AxisAlignedBoundingBox3D::AxisAlignedBoundingBox3D(void) { ; }
AxisAlignedBoundingBox3D::AxisAlignedBoundingBox3D(const Point3D& _m, const Point3D& _M) : min(_m), max(_M) { ; }
void AxisAlignedBoundingBox3D::add(const Point3D& p) {
	if (min.x > p.x) min.x = p.x;
	if (min.y > p.y) min.y = p.y;
	if (min.z > p.z) min.z = p.z;

	if (max.x < p.x) max.x = p.x;
	if (max.y < p.y) max.y = p.y;
	if (max.z < p.z) max.z = p.z;
}
AxisAlignedBoundingBox3D AxisAlignedBoundingBox3D::operator + (const AxisAlignedBoundingBox3D& operand) const {
	AxisAlignedBoundingBox3D ret = operand;
	ret.add(max);
	ret.add(min);
	return ret;
}

XYRotatedBoundingBox3D::XYRotatedBoundingBox3D(void)
:min(0,0,0), axis_u(1,0,0), axis_v(0,1,0), xy(0,0), height(0.0) { ; }

void XYRotatedBoundingBox3D::compute(const std::vector<Point3D>& pts){
	if (pts.empty()) {
		*this = XYRotatedBoundingBox3D();
		return;
	}
	min.z = pts[0].z;
	height = 0.0;
	for (int i = 1; i < pts.size(); ++i) {
		if (min.z > pts[i].z) {
			height += min.z - pts[i].z;
			min.z = pts[i].z;
		}
		if (height < pts[i].z - min.z) {
			height = pts[i].z - min.z;
		}
	}

	// TODO: may use the convex hull to improve the complexity
	scalar_t min_area = -1;
	Point3D  ma_p(0, 0, 0);
	Vector2D ma_xy(0, 0);
	Vector3D ma_u(1, 0, 0);
	Vector3D ma_v(0, 1, 0);
	for (int i = 0; i < pts.size(); ++i) {
		for (int j = i + 1; j < pts.size(); ++j) {
			//Vector2D delta(pts[j].x - pts[i].x, pts[j].y - pts[i].y);
			Vector3D u = pts[j] - pts[i];
			u.z = 0;
			if (u.is_zero()) continue;
			u = u.normalized();
			Vector3D v = Vector3D(0, 0, 1).cross_product(u);

			scalar_t min_u, max_u;
			scalar_t min_v, max_v;
			min_u = max_u = pts[0].to_vector().dot_product(u);
			min_v = max_v = pts[0].to_vector().dot_product(v);

			for (int k = 1; k < pts.size(); ++k) {
				scalar_t x = pts[k].to_vector().dot_product(u);
				scalar_t y = pts[k].to_vector().dot_product(v);

				if (min_u > x) min_u = x;
				if (min_v > y) min_v = y;

				if (max_u < x) max_u = x;
				if (max_v < y) max_v = y;
			}

			scalar_t area = (max_u - min_u) * (max_v - min_v);
			if (min_area < 0 || min_area > area) {
				min_area = area;
				min = Point3D(0,0,min.z) + (u * min_u + v * min_v);
				xy = Vector2D(max_u - min_u, max_v - min_v);
				axis_u = u;
				axis_v = v;
			}
		}
	}
	// Post condition:
	//  "min" is the bottom-left corner, "axis_u" and "axis_v" comprise the unit vectors indicating two sides of the bounding box.
	//  "xy" is (width, height).

};

//////////////////////////////////////////////////////////////////////////////////////////////////////////////////
Vector2D::Vector2D(void) { ; }
Vector2D::Vector2D(coord_t _x, coord_t _y) : x(_x), y(_y) { ; }

Vector2D Vector2D::normalize(void) const {
	length_t size = sqrt(dot_product(*this));
	return Vector2D(x / size, y / size);
}

scalar_t Vector2D::dot_product(const Vector2D& v) const { return x * v.x + y * v.y; }
area_t Vector2D::signed_area(const Vector2D& v) const { return x * v.y - y * v.x; }

length_t Vector2D::length(void) const {
	return sqrt(dot_product(*this));
}

Vector2D Vector2D::operator + (const Vector2D& v) const { return Vector2D(x + v.x, y + v.y); }
Vector2D Vector2D::operator - (const Vector2D& v) const { return Vector2D(x - v.x, y - v.y); }
Vector2D Vector2D::operator - (void) const { return Vector2D(-x, -y); }
Vector2D operator * (scalar_t a, const Vector2D& v) { return Vector2D(a*v.x, a*v.y); }

//////////////////////////////////////////////////////////////////////////////////////////////////////////////////
Vector3D::Vector3D(void) { ; }
Vector3D::Vector3D(coord_t _x, coord_t _y, coord_t _z) :x(_x), y(_y), z(_z) { ; }

Point3D Vector3D::to_point(void) const {
	return Point3D(x, y, z);
}

scalar_t Vector3D::dot_product(const Vector3D& v) const { return x * v.x + y * v.y + z * v.z; }
Vector3D Vector3D::cross_product(const Vector3D& v) const {
	float xx = y * v.z - z * v.y;
	float yy = z * v.x - x * v.z;
	float zz = x * v.y - y * v.x;
	return Vector3D(xx, yy, zz);
}
Vector3D Vector3D::normalized(void) const {
	length_t size = length();
	return Vector3D(x / size, y / size, z / size);
}
length_t Vector3D::length(void) const {
	return sqrt(dot_product(*this));

}
length_t Vector3D::length_square(void) const {
	return dot_product(*this);

}
Vector3D Vector3D::rotate(const Vector3D& q, radian_t theta) const {

	scalar_t rx = sin(theta / 2) * q.x;
	scalar_t ry = sin(theta / 2) * q.y;
	scalar_t rz = sin(theta / 2) * q.z;
	scalar_t rw = cos(theta / 2);

	scalar_t ix = rw * x + ry * z - rz * y;
	scalar_t iy = rw * y + rz * x - rx * z;
	scalar_t iz = rw * z + rx * y - ry * x;
	scalar_t iw = -rx * x - ry * y - rz * z;

	return Vector3D(
		ix * rw - iw * rx + ry * iz - rz * iy,
		iy * rw - iw * ry + rz * ix - rx * iz,
		iz * rw - iw * rz + rx * iy - ry * ix
	);
}
Vector3D Vector3D::get_projection_component_to(const Vector3D& x) const {
	return x.normalized() * dot_product(x) / sqrt(x.dot_product(x));
}
Vector3D Vector3D::get_perpendicular_component_to(const Vector3D& x) const {
	return *this - get_projection_component_to(x);
}

Vector3D Vector3D::operator + (const Vector3D& v) const { return Vector3D(x + v.x, y + v.y, z + v.z); }
Vector3D Vector3D::operator - (const Vector3D& v) const { return Vector3D(x - v.x, y - v.y, z - v.z); }
Vector3D operator * (scalar_t a, const Vector3D& v) { return Vector3D(a*v.x, a*v.y, a*v.z); }
Vector3D Vector3D::operator * (scalar_t a) const { return Vector3D(a*x, a*y, a*z); }
Vector3D Vector3D::operator / (scalar_t a) const { return Vector3D(x / a, y / a, z / a); }

bool Vector3D::is_zero(void) const {
	return x * x + y * y + z * z < EPSILON;
}

//////////////////////////////////////////////////////////////////////////////////////////////////////////////////
Line2D::Line2D(const Point2D& _p, const Vector2D& _v) : p(_p), v(_v) { ; }

Point2D Line2D::get_projection_of(const Point2D& _p, length_t *_alpha) const {
	scalar_t denominator = v.dot_product(v);
	if (denominator == 0) {
		if (_alpha) *_alpha = 0;
		return p;
	}
	length_t alpha = (_p - p).dot_product(v) / denominator;
	if (_alpha) *_alpha = alpha;
	return p + alpha * v;
}
//////////////////////////////////////////////////////////////////////////////////////////////////////////////////

Line3D::Line3D() { ; }
Line3D::Line3D(const Point3D& _p, const Point3D& _q) : p(_p), v(_q-_p) { ; }
Line3D::Line3D(const Point3D& _p, const Vector3D& _v) : p(_p), v(_v) { ; }

//////////////////////////////////////////////////////////////////////////////////////////////////////////////////
/*
typedef boost::polygon::polygon_with_holes_data<scalar_t> boost_polygon;
typedef boost::polygon::polygon_traits<boost_polygon>::point_type boost_point;
static boost_polygon to_boost_polygon(const Polygon2D& poly) {
	using namespace boost::polygon;
	boost_polygon p;
	std::vector<boost_point> exterior;
	for (int i = 0; i < poly.exterior.size(); ++i) {
		exterior.push_back(construct<boost_point>(poly.exterior[i].x, poly.exterior[i].y));
	}

	std::vector< std::vector<boost_point> > holes;
	for (int r = 0; r < poly.hole.size(); ++r) {
		std::vector<boost_point> hole;
		for (int i = 0; i < poly.hole[r].size(); ++i) {
			hole.push_back(construct<boost_point>(poly.hole[r][i].x, poly.hole[r][i].y));
		}
		holes.push_back(hole);
	}

	set_points(p, exterior.begin(), exterior.end());
	set_holes(p, holes.begin(), holes.end());

	return p;
}

static Polygon2D to_polygon2d(const boost_polygon& poly) {
	Polygon2D ret;

	for (auto i = poly.begin(); i != poly.end(); ++i) {
		ret.exterior.push_back(Point2D(i->x(), i->y()));
	}

	for (auto i = poly.begin_holes(); i != poly.end_holes(); ++i) {
		std::vector<Point2D> hole;
		for (auto j = i->begin(); j != i->end(); ++j) {
			hole.push_back(Point2D(j->x(), j->y()));
		}
		ret.hole.push_back(hole);
	}

	return ret;
}
std::vector<Polygon2D> Polygon2D::intersection(const Polygon2D& x) const {
	boost_polygon p = to_boost_polygon(*this);
	boost_polygon q = to_boost_polygon(x);
	using namespace boost::polygon::operators;
	//typedef  PolygonSet;
	//PolygonSet ps;
	std::vector<boost_polygon> intersection_ret;
	std::cout << "DIRTY" << (p & q).dirty() << std::endl;
	assign(intersection_ret, p & q);

	std::vector<Polygon2D> ret;
	for (int i = 0; i < intersection_ret.size(); ++i) {
		ret.push_back(to_polygon2d(intersection_ret[i]));
	}
	return ret;
}
*/
typedef boost::geometry::model::polygon<boost::geometry::model::d2::point_xy<double> > boost_polygon;
typedef boost_polygon::point_type boost_point;

static boost_polygon to_boost_polygon(const Polygon2D& poly) {
	boost_polygon ret;
	std::vector<boost_point> pts;
	for (int i = 0; i < poly.exterior.size(); ++i) {
		pts.push_back(boost_point(poly.exterior[i].x, poly.exterior[i].y));
	}
	pts.push_back(boost_point(poly.exterior[0].x, poly.exterior[0].y));
	boost::geometry::assign_points(ret, pts);
	return ret;
}
static Polygon2D to_polygon2d(const boost_polygon& poly) {
	Polygon2D ret;
	for (auto i = poly.outer().begin(); i != poly.outer().end(); ++i) {
		ret.exterior.push_back(Point2D((*i).x(), (*i).y()));
	}
	return ret;
}
std::vector<Polygon2D> Polygon2D::intersection(const Polygon2D& x) const {
	boost_polygon p = to_boost_polygon(*this);
	boost_polygon q = to_boost_polygon(x);

	std::vector<boost_polygon> intersection_ret;
	boost::geometry::intersection(p, q, intersection_ret);
	std::vector<Polygon2D> ret;
	for (int i = 0; i < intersection_ret.size(); ++i) {
		ret.push_back(to_polygon2d(intersection_ret[i]));
	}
	return ret;
}

//////////////////////////////////////////////////////////////////////////////////////////////////////////////////
Plane::Plane(void) { ; }
Plane::Plane(const Point3D& _p, const Vector3D& _h) : p(_p), h(_h) { ; }
Point3D Plane::project(const Point3D& q, length_t *d) const {
	length_t l = (p - q).dot_product(h) / h.length();
	if (d) *d = l;
	return q + l * h.normalized();
}
Line3D Plane::project(const Line3D& l) const {
	Point3D proj_p = project(l.p);
	Point3D proj_q = project(l.p + l.v);
	Vector3D proj_v = proj_q - proj_p;
	return Line3D(proj_p, proj_v);
}

CoordinatedPlane::CoordinatedPlane(void) { ; }
CoordinatedPlane::CoordinatedPlane(const Point3D& _p, const Vector3D& _x, const Vector3D& _y) : Plane(_p, _x.cross_product(_y)), x(_x), y(_y) { ; }

Point3D CoordinatedPlane::convert(const Point2D& q) const { return p + q.x*x + q.y*y; }
Point2D CoordinatedPlane::convert(const Point3D& q, length_t *d) const {
	Vector3D denominator_vector = x.cross_product(y);
	scalar_t denominator = denominator_vector.dot_product(denominator_vector);
	Point2D ret;
	ret.x = denominator_vector.dot_product((q - p).cross_product(y)) / denominator;
	ret.y = denominator_vector.dot_product((p - q).cross_product(x)) / denominator;
	if (d) { *d = (q - p).dot_product(denominator_vector) / denominator; }
	return ret;
}

bool Plane::is_parallel(const Plane& plane) const {
	return h.cross_product(plane.h).is_zero();
}

Line3D Plane::intersect(const Plane& plane) const {
	const Vector3D& u = h;
	const Vector3D& v = plane.h;

	Vector3D q_p = plane.p - p;

	Vector3D uv = u.cross_product(v);

	Vector3D vec_p_to_line = u.cross_product(uv);
	scalar_t alpha = (q_p).dot_product(v) / vec_p_to_line.dot_product(v);
	Point3D start_point = p + alpha * vec_p_to_line;

	Vector3D vec_line_to_q = v.cross_product(uv);
	scalar_t beta = (q_p).dot_product(u) / vec_line_to_q.dot_product(u);
	Point3D end_point = plane.p - beta * vec_line_to_q;

	return Line3D(start_point, end_point - start_point);
}

void Plane::get_basis(Vector3D* u, Vector3D* v) const {
	Vector3D ux(1, 0, 0);
	Vector3D uy(0, 1, 0);
	Vector3D uz(0, 0, 1);
	scalar_t dx = abs(h.dot_product(ux));
	scalar_t dy = abs(h.dot_product(uy));
	scalar_t dz = abs(h.dot_product(uz));

	Vector3D base_u, base_v;

	if (dx <= dy && dx <= dz) {
		base_u = ux.get_perpendicular_component_to(h).normalized();
		base_v = h.cross_product(base_u).normalized();
	}
	else if (dy <= dx && dy <= dz) {
		base_u = uy.get_perpendicular_component_to(h).normalized();
		base_v = h.cross_product(base_u).normalized();
	}
	else if (dz <= dx && dz <= dy) {
		base_u = uz.get_perpendicular_component_to(h).normalized();
		base_v = h.cross_product(base_u).normalized();
	}

	if (u) *u = base_u;
	if (v) *v = base_v;
}

//////////////////////////////////////////////////////////////////////////////////////////////////////////////////
bool interp(const Point3D& x, const Line3D& l, scalar_t *_alpha) {
	// x = l.p + alpha * l.v + h
	// where h is perpendicular to l.v
	scalar_t denominator = l.v.dot_product(l.v);
	if (abs(denominator) < EPSILON) return false;
	if (_alpha) {
		*_alpha = (x - l.p).dot_product(l.v) / denominator;
	}
	return true;
}
bool interp(const Line3D& l1, const Line3D& l2, scalar_t *_alpha, scalar_t *_beta) {
	// l1.p + alpha * l1.v + h = l2.p + beta * l2.v
	// where h is perpendicular to both l1.v and l2.v
	Vector3D q_p = l2.p - l1.p;
	Vector3D uv = l1.v.cross_product(l2.v);
	scalar_t denominator = uv.dot_product(uv);
	if (abs(denominator) < EPSILON) return false;
	if (_alpha) {
		*_alpha = (q_p.cross_product(l2.v)).dot_product(uv) / denominator;
	}
	if (_beta) {
		*_beta = (q_p.cross_product(l1.v)).dot_product(uv) / denominator;
	}
	return true;
}
bool interp(const Point3D& x, const Plane& p, scalar_t *_alpha) {
	// x = p.p + alpha * p.h + v
	// where v is a vector perpendicular to p.h
	if (_alpha) {
		*_alpha = (x - p.p).dot_product(p.h) / p.h.dot_product(p.h);
	}
	return true;
}
bool interp(const Line3D& l, const Plane& p, scalar_t *_alpha) {
	// l.p + alpha * l.v = p.p + v
	// where v is a vector perpendicular to p.h
	scalar_t vh = l.v.dot_product(p.h);
	if (abs(vh) < EPSILON) return false;
	if (_alpha) {
		*_alpha = (p.p - l.p).dot_product(p.h) / vh;
	}
	return true;
}

bool interp(const Plane& p, const Plane& q, Line3D *_line, Point3D *p_proj, Point3D *q_proj) {
	Vector3D uv = p.h.cross_product(q.h);
	if (uv.is_zero()) return false;
	Vector3D uuv = p.h.cross_product(uv);
	Vector3D vuv = q.h.cross_product(uv);

	scalar_t d_alpha = uuv.dot_product(q.h);
	scalar_t d_gamma = vuv.dot_product(p.h);
	if (abs(d_alpha) < EPSILON || abs(d_gamma) < EPSILON) return false;

	Vector3D q_p = q.p - p.p;
	Point3D p_p = p.p + (q_p.dot_product(q.h) / d_alpha) * uuv;
	if (p_proj) {
		*p_proj = p_p;
	}
	if (q_proj) {
		*q_proj = q.p - (q_p.dot_product(p.h) / d_gamma) * vuv;
	}
	if (_line) {
		*_line = Line3D(p_p, uv.normalized());
	}
	return true;
}

#include <iostream>
float solid_angle(const Point3D& o, const Point3D& p, const Point3D& q, const Point3D& r) {
	Vector3D a = (p - o).normalized();
	Vector3D b = (q - o).normalized();
	Vector3D c = (r - o).normalized();

	//std::cout << "VEC " << a.x << ' ' << a.y << ' ' << a.z << std::endl;
	//std::cout << "VEC " << b.x << ' ' << b.y << ' ' << b.z << std::endl;
	//std::cout << "VEC " << c.x << ' ' << c.y << ' ' << c.z << std::endl;

	scalar_t abc = a.dot_product(b.cross_product(c));
	scalar_t div = 1 + a.dot_product(b) + a.dot_product(c) + b.dot_product(c);

	scalar_t omega = atan(abc/div);

	//std::cout << "NORM " << omega << '\t' << abc << '\t' << div << '\t';

	if (div < 0) {
		if (omega < 0) omega += PI;
		else omega -= PI;
	}

	//std::cout << omega << std::endl;
	//if ((a + b + c).dot_product((b - a).cross_product(c - a)) < 0) omega = -omega;
	return 2*omega;
}


#include "eigen_ext.hpp"
#include "parameters.hpp"
#include "stiff.hpp"
#include <cassert>
#include <iostream>

#include <Eigen/Core>
#include <Eigen/SparseCore>

namespace pear {
void grad_phi(Vec &xp, Vec &yp, double &T, Vec &Dphi2, Vec &Dphi3) {

  Dphi2 << yp(1) - yp(2), yp(2) - yp(0), yp(0) - yp(1);
  Dphi3 << xp(2) - xp(1), xp(0) - xp(2), xp(1) - xp(0);

  T = xp(1) * yp(2) + xp(0) * yp(1) + xp(2) * yp(0) - xp(1) * yp(0) -
      xp(0) * yp(2) - xp(2) * yp(1);
}

Mat stiff_block(Vec &xp, Vec &yp, double Dr, double Dz) {

  Mat K_block(3, 3);

  Vec Dphi2(3);
  Vec Dphi3(3);
  double T = 1;

  grad_phi(xp, yp, T, Dphi2, Dphi3);

  for (int idx1 = 0; idx1 < 3; idx1++) {
    for (int idx2 = 0; idx2 < 3; idx2++) {
      K_block(idx1, idx2) =
          (xp(0) + xp(1) + xp(2)) *
          (Dr * Dphi2(idx1) * Dphi2(idx2) + Dz * Dphi3(idx1) * Dphi3(idx2)) /
          12 / T;
    }
  }
  return K_block;
}

void stiff(Vec &xp, Vec &yp, MatI &t, Mat &Ku, Mat &Kv) {
  // PRELIMINARIES
  int np = xp.rows();
  int nt = t.rows();

  Mat Ku_block(3, 3);
  Mat Kv_block(3, 3);

  VecI t_loc(3);
  Vec xp_loc(3);
  Vec yp_loc(3);

  for (int idxm = 0; idxm < nt; idxm++) {
    t_loc = t.row(idxm);
    xp_loc = pear::extract<Vec>(xp, t_loc);
    yp_loc = pear::extract<Vec>(yp, t_loc);

    Ku_block = stiff_block(xp_loc, yp_loc, pear::Dur, pear::Duz);
    Kv_block = stiff_block(xp_loc, yp_loc, pear::Dvr, pear::Dvz);
    for (int idx1 = 0; idx1 < 3; idx1++) {
      for (int idx2 = 0; idx2 < 3; idx2++) {
        Ku(t_loc(idx1), t_loc(idx2)) += Ku_block(idx1, idx2);
        Kv(t_loc(idx1), t_loc(idx2)) += Kv_block(idx1, idx2);
      }
    }
  }
}

} // namespace pear


#pragma once

#include <Eigen/Dense>
#include <Eigen/StdVector>

template <int ROWS>
using VecNd = Eigen::Matrix<double, ROWS, 1>;

template <int ROWS, int COLS>
using MatNd = Eigen::Matrix<double, ROWS, COLS>;

using VecXd = Eigen::VectorXd;
using MatXd = Eigen::MatrixXd;

template <int ROWS>
using SquareMatNd = Eigen::Matrix<double, ROWS, ROWS>;

template <typename Eig>
using StdVector = std::vector<Eig, Eigen::aligned_allocator<Eig>>;

namespace jcc {
using Vec1 = VecNd<1>;
using Vec2 = Eigen::Vector2d;
using Vec3 = Eigen::Vector3d;
using Vec4 = Eigen::Vector4d;
using Vec5 = VecNd<5>;
using Vec6 = VecNd<6>;

}  // namespace jcc