#include "cpu/image_proc.h"

#include <Eigen/Dense>
#include <map>
#include <string>
#include <cmath>

namespace helper {

    bool in_bounds(Eigen::Vector2f p, int h, int w) {
        return p.x() >= 0.0 && p.x() < w && p.y() >= 0.0 && p.y() < h;
    }

    bool valid_flow_at(const Eigen::Vector2f& p, int h, int w, py::array_t<float>& flow_image, Eigen::Vector2f& flow) {
        if (!in_bounds(p, h, w)) {
            return false;
        }

        flow.x() = *flow_image.data(p.y(), p.x(), 0);
        flow.y() = *flow_image.data(p.y(), p.x(), 1);

        if (!flow.allFinite()) {
            return false;
        }

        return true;
    }
}

namespace image_proc {

    using Vec2f = Eigen::Vector2f;

    template <class V, class L = std::less<std::string>,
              class A = Eigen::aligned_allocator<std::pair<const std::string, V>>>
    using aligned_dict = std::map<std::string, V, L, A>;

    py::array_t<float> compute_augmented_flow_from_rotation(py::array_t<float>& flow_image_rot_sa2so,
                                              py::array_t<float>& flow_image_so2to,
                                              py::array_t<float>& flow_image_rot_to2ta,
                                              const int h, const int w) {
        // TODO: change to runtime asserts
        // assert(flow_image_rot_sa2so.ndim() == 3);
        // assert(flow_image_rot_sa2so.shape(0) == 2);
        // assert(flow_image_rot_sa2so.shape(1) == h);
        // assert(flow_image_rot_sa2so.shape(2) == w);

        // assert(flow_image_so2to.ndim() == 3);
        // assert(flow_image_so2to.shape(0) == 2);
        // assert(flow_image_so2to.shape(1) == h);
        // assert(flow_image_so2to.shape(2) == w);

        // assert(flow_image_rot_to2ta.ndim() == 3);
        // assert(flow_image_rot_to2ta.shape(0) == 2);
        // assert(flow_image_rot_to2ta.shape(1) == h);
        // assert(flow_image_rot_to2ta.shape(2) == w);

        // allocate memory for output array
        py::array_t<float> flow_image_rot_sa2ta = py::array_t<float>(flow_image_rot_sa2so.request().size);

        // reshape array to match input shape
        flow_image_rot_sa2ta.resize({h, w, 2});

        for (int y = 0; y < h; y++) {
            for (int x = 0; x < w; x++) {

                // update output flow image
                *flow_image_rot_sa2ta.mutable_data(y, x, 0) = - std::numeric_limits<float>::infinity();
                *flow_image_rot_sa2ta.mutable_data(y, x, 1) = - std::numeric_limits<float>::infinity();

                Vec2f p_sa(x, y);

                /////////////////////////////////////////////////////////////////////////////////
                // 1. SOURCE AUGMENTED TO SOURCE ORIGINAL
                /////////////////////////////////////////////////////////////////////////////////

                // flow from source augmented to source original
                Vec2f flow_sa2so(*flow_image_rot_sa2so.data(y, x, 0), *flow_image_rot_sa2so.data(y, x, 1));

                // flow_sa2so should be dense and w/o any invalid value
                if (!flow_sa2so.allFinite()) {
                    throw "flow_sa2so sould be dense and w/o any invalid values!";
                }

                // compute warped location on source original (so we're going from source augmented to source original)
                Vec2f p_so = p_sa + flow_sa2so;

                // init flow_sa2ta with the first contribution, i.e, flow_sa2so
                Vec2f flow_sa2ta = flow_sa2so;

                /////////////////////////////////////////////////////////////////////////////////
                // 2. SOURCE ORIGINAL TO TARGET ORIGINAL
                /////////////////////////////////////////////////////////////////////////////////
                int u0 = std::floor(p_so.x());
                int u1 = u0 + 1;
                int v0 = std::floor(p_so.y());
                int v1 = v0 + 1;

                Vec2f p00(u0, v0);
                Vec2f p01(u0, v1);
                Vec2f p10(u1, v0);
                Vec2f p11(u1, v1);

                aligned_dict<Vec2f> valid_coords;
                aligned_dict<Vec2f> valid_flows;

                Vec2f flow_00_so2to;
                if (helper::valid_flow_at(p00, h, w, flow_image_so2to, flow_00_so2to)) {
                    valid_coords["p00"] = p00;
                    valid_flows["p00"]  = flow_00_so2to;
                }

                Vec2f flow_01_so2to;
                if (helper::valid_flow_at(p01, h, w, flow_image_so2to, flow_01_so2to)) {
                    valid_coords["p01"] = p01;
                    valid_flows["p01"]  = flow_01_so2to;
                }

                Vec2f flow_10_so2to;
                if (helper::valid_flow_at(p10, h, w, flow_image_so2to, flow_10_so2to)) {
                    valid_coords["p10"] = p10;
                    valid_flows["p10"]  = flow_10_so2to;
                }

                Vec2f flow_11_so2to;
                if (helper::valid_flow_at(p11, h, w, flow_image_so2to, flow_11_so2to)) {
                    valid_coords["p11"] = p11;
                    valid_flows["p11"]  = flow_11_so2to;
                }

                // Depending on how many valid flows we have, do bilinear interpolation or nearest neighbor:
                Vec2f flow_so2to;

                if (valid_coords.size() == 0) {
                    continue;
                } else if (valid_coords.size() == 4) {
                    // Bilinear interpolation
                    float du = p_so.x() - u0;
                    float dv = p_so.y() - v0;

                    float w00 = (1 - du) * (1 - dv);
                    float w01 = (1 - du) * dv;
                    float w10 = du * (1 - dv);
                    float w11 = du * dv;

                    flow_so2to = w00 * valid_flows["p00"] +
                                 w01 * valid_flows["p01"] +
                                 w10 * valid_flows["p10"] +
                                 w11 * valid_flows["p11"];
                } else {
                    // Nearest Neighbor
                    std::string nn = "None";
                    float min_dist = std::numeric_limits<float>::max();

                    for (const auto& valid_coord : valid_coords) {
                        const std::string k = valid_coord.first;
                        const Vec2f& p = valid_coord.second;

                        float dist = (p_so - p).norm();
                        if (dist < min_dist) {
                            min_dist = dist;
                            nn = k;
                        }
                    }

                    if (nn == "None") {
                        throw std::runtime_error("Neighrest Neighbor 'nn' was not assigned...");
                    }

                    flow_so2to = valid_flows[nn];
                }

                // compute warped location on target original (so we're going from source original to target original)
                Vec2f p_to = p_so + flow_so2to;

                // add flow_so2to to flow_sa2ta
                flow_sa2ta += flow_so2to;

                /////////////////////////////////////////////////////////////////////////////////
                // 3. TARGET ORIGINAL TO TARGET AUGMENTED
                /////////////////////////////////////////////////////////////////////////////////
                u0 = std::floor(p_to.x());
                u1 = u0 + 1;
                v0 = std::floor(p_to.y());
                v1 = v0 + 1;

                p00 = Vec2f(u0, v0);
                p01 = Vec2f(u0, v1);
                p10 = Vec2f(u1, v0);
                p11 = Vec2f(u1, v1);

                valid_coords.clear();
                valid_flows.clear();

                Vec2f flow_00_to2ta;
                if (helper::valid_flow_at(p00, h, w, flow_image_rot_to2ta, flow_00_to2ta)) {
                    valid_coords["p00"] = p00;
                    valid_flows["p00"]  = flow_00_to2ta;
                }

                Vec2f flow_01_to2ta;
                if (helper::valid_flow_at(p01, h, w, flow_image_rot_to2ta, flow_01_to2ta)) {
                    valid_coords["p01"] = p01;
                    valid_flows["p01"]  = flow_01_to2ta;
                }

                Vec2f flow_10_to2ta;
                if (helper::valid_flow_at(p10, h, w, flow_image_rot_to2ta, flow_10_to2ta)) {
                    valid_coords["p10"] = p10;
                    valid_flows["p10"]  = flow_10_to2ta;
                }

                Vec2f flow_11_to2ta;
                if (helper::valid_flow_at(p11, h, w, flow_image_rot_to2ta, flow_11_to2ta)) {
                    valid_coords["p11"] = p11;
                    valid_flows["p11"]  = flow_11_to2ta;
                }

                // Depending on how many valid flows we have, do bilinear interpolation or nearest neighbor:
                Vec2f flow_to2ta;

                if (valid_coords.size() == 0) {
                    continue;
                } else if (valid_coords.size() == 4) {
                    // Bilinear interpolation
                    float du = p_to.x() - u0;
                    float dv = p_to.y() - v0;

                    float w00 = (1 - du) * (1 - dv);
                    float w01 = (1 - du) * dv;
                    float w10 = du * (1 - dv);
                    float w11 = du * dv;

                    flow_to2ta = w00 * valid_flows["p00"] +
                                 w01 * valid_flows["p01"] +
                                 w10 * valid_flows["p10"] +
                                 w11 * valid_flows["p11"];
                } else {
                    // Nearest Neighbor
                    std::string nn = "None";
                    float min_dist = std::numeric_limits<float>::max();

                    for (const auto& valid_coord : valid_coords) {
                        const std::string k = valid_coord.first;
                        const Vec2f& p = valid_coord.second;

                        float dist = (p_to - p).norm();
                        if (dist < min_dist) {
                            min_dist = dist;
                            nn = k;
                        }
                    }

                    if (nn == "None") {
                        throw std::runtime_error("Neighrest Neighbor 'nn' was not assigned...");
                    }

                    flow_to2ta = valid_flows[nn];
                }

                // add flow_to2ta to flow_sa2ta
                flow_sa2ta += flow_to2ta;

                // update output flow image
                *flow_image_rot_sa2ta.mutable_data(y, x, 0) = flow_sa2ta.x();
                *flow_image_rot_sa2ta.mutable_data(y, x, 1) = flow_sa2ta.y();
            }
        }

        return flow_image_rot_sa2ta;
    }

    int count_tp1(py::array_t<bool> &p, py::array_t<bool> &gt) {
        assert(p.ndim() == 2);
        assert(gt.ndim() == 2);
        const int n_batch = p.shape(0);
        const int dimz = p.shape(1);

        auto& ptr = p;
        int counter = 0;
        for (int i = 0; i < n_batch; i++)
            for (int z = 0 ; z < dimz ; z++)
                if (*gt.data(i, z)) {
                    counter += CHECK1();
                }
        return counter;
    }

    int count_tp2(py::array_t<bool> &p, py::array_t<bool> &gt) {
        assert(p.ndim() == 4);
        assert(gt.ndim() == 4);
        const int n_batch = p.shape(0);
        assert(p.shape(1) == 1);
        const int height = p.shape(2);
        const int width = p.shape(3);

        auto& ptr = p;
        int counter = 0;
        for (int i = 0; i < n_batch; i++)
            for (int y = 0; y <  height; y++)
                for (int x = 0; x < width; x++) {
                    if (*gt.data(i, 0, y, x)) {
                        counter += CHECK2();
                    }
                }
        return counter;
    }

    int count_tp3(py::array_t<bool> &p, py::array_t<bool> &gt) {
        assert(p.ndim() == 5);
        assert(gt.ndim() == 5);
        const int n_batch = p.shape(0);
        assert(p.shape(1) == 1);
        const int dimz = p.shape(2);
        const int dimy = p.shape(3);
        const int dimx = p.shape(4);

        auto& ptr = p;
        int counter = 0;
        for (int i = 0; i < n_batch; i++)
            for (int z = 0 ; z < dimz ; z++)
                for (int y = 0; y <  dimy; y++)
                    for (int x = 0; x < dimx; x++) {
                        if (*gt.data(i, 0, z, y, x)) {
                            counter += CHECK3();
                            //printf("i %d x %d y %d z %d in %d\n", i, x, y, z, res);
                        }
                    }
        return counter;
    }

    void extend3(py::array_t<bool> &in, py::array_t<bool> &out) {
        assert(in.ndim() == 5);
        assert(out.ndim() == 5);
        int n_batch = in.shape(0);
        assert(in.shape(1) == 1);
        int dimz = in.shape(2);
        int dimy = in.shape(3);
        int dimx = in.shape(4);

        auto& ptr = in;
        for (int i = 0; i < n_batch; i++)
            for (int z = 1 ; z < dimz - 1; z++)
                for (int y = 1; y <  dimy - 1; y++)
                    for (int x = 1; x < dimx - 1; x++) {
                        *out.mutable_data(i, 0, z, y, x) = CHECK3();
                    }
    }

    void backproject_depth_ushort(py::array_t<unsigned short>& in, py::array_t<float>& out, float fx, float fy, float cx, float cy, float normalizer) {
        assert(in.ndim() == 2);
        assert(out.ndim() == 3);

        int width = in.shape(1);
        int height = in.shape(0);
        assert(out.shape(0) == 3);
        assert(out.shape(1) == height);
        assert(out.shape(2) == width);

        for (int y = 0; y < height; y++) {
            for (int x = 0; x < width; x++) {
                float depth = float(*in.data(y, x)) / normalizer;

                if (depth > 0) {
                    float pos_x = depth * (x - cx) / fx;
                    float pos_y = depth * (y - cy) / fy;
                    float pos_z = depth;

                    *out.mutable_data(0, y, x) = pos_x;
                    *out.mutable_data(1, y, x) = pos_y;
                    *out.mutable_data(2, y, x) = pos_z;
                }
            }
        }
    }

    void backproject_depth_float(py::array_t<float>& in, py::array_t<float>& out, float fx, float fy, float cx, float cy) {
        assert(in.ndim() == 2);
        assert(out.ndim() == 3);

        int width = in.shape(1);
        int height = in.shape(0);
        assert(out.shape(0) == 3);
        assert(out.shape(1) == height);
        assert(out.shape(2) == width);

        for (int y = 0; y < height; y++) {
            for (int x = 0; x < width; x++) {
                float depth = *in.data(y, x);

                if (depth > 0) {
                    float pos_x = depth * (x - cx) / fx;
                    float pos_y = depth * (y - cy) / fy;
                    float pos_z = depth;

                    *out.mutable_data(0, y, x) = pos_x;
                    *out.mutable_data(1, y, x) = pos_y;
                    *out.mutable_data(2, y, x) = pos_z;
                }
            }
        }
    }

    void compute_mesh_from_depth(
        const py::array_t<float>& pointImage, float maxTriangleEdgeDistance,
        py::array_t<float>& vertexPositions, py::array_t<int>& vertexPixels, py::array_t<int>& faceIndices
    ) {
        int width = pointImage.shape(2);
        int height = pointImage.shape(1);

        // Compute valid pixel vertices and faces.
        // We also need to compute the pixel -> vertex index mapping for
        // computation of faces.
        // We connect neighboring pixels on the square into two triangles.
        // We only select valid triangles, i.e. with all valid vertices and
        // not too far apart.
        // Important: The triangle orientation is set such that the normals
        // point towards the camera.
        std::vector<Eigen::Vector3f> vertices;
        std::vector<Eigen::Vector3i> faces;
        std::vector<Eigen::Vector2i> pixels;

        int vertexIdx = 0;
        std::vector<int> mapPixelToVertexIdx(width * height, -1);

        for (int y = 0; y < height - 1; y++) {
            for (int x = 0; x < width - 1; x++) {
                Eigen::Vector3f obs00(*pointImage.data(0, y, x), *pointImage.data(1, y, x), *pointImage.data(2, y, x));
                Eigen::Vector3f obs01(*pointImage.data(0, y + 1, x), *pointImage.data(1, y + 1, x), *pointImage.data(2, y + 1, x));
                Eigen::Vector3f obs10(*pointImage.data(0, y, x + 1), *pointImage.data(1, y, x + 1), *pointImage.data(2, y, x + 1));
                Eigen::Vector3f obs11(*pointImage.data(0, y + 1, x + 1), *pointImage.data(1, y + 1, x + 1), *pointImage.data(2, y + 1, x + 1));

                int idx00 = y * width + x;
                int idx01 = (y + 1) * width + x;
                int idx10 = y * width + (x + 1);
                int idx11 = (y + 1) * width + (x + 1);

                bool valid00 = obs00.z() > 0;
                bool valid01 = obs01.z() > 0;
                bool valid10 = obs10.z() > 0;
                bool valid11 = obs11.z() > 0;

                if (valid00 && valid01 && valid10) {
                    float d0 = (obs00 - obs01).norm();
                    float d1 = (obs00 - obs10).norm();
                    float d2 = (obs01 - obs10).norm();

                    if (d0 <= maxTriangleEdgeDistance && d1 <= maxTriangleEdgeDistance && d2 <= maxTriangleEdgeDistance) {
                        int vIdx0 = mapPixelToVertexIdx[idx00];
                        int vIdx1 = mapPixelToVertexIdx[idx01];
                        int vIdx2 = mapPixelToVertexIdx[idx10];

                        if (vIdx0 == -1) {
                            vIdx0 = vertexIdx;
                            mapPixelToVertexIdx[idx00] = vertexIdx;
                            vertices.push_back(obs00);
                            pixels.push_back(Eigen::Vector2i(x, y));

                            vertexIdx++;
                        }
                        if (vIdx1 == -1) {
                            vIdx1 = vertexIdx;
                            mapPixelToVertexIdx[idx01] = vertexIdx;
                            vertices.push_back(obs01);
                            pixels.push_back(Eigen::Vector2i(x, y+1));

                            vertexIdx++;
                        }
                        if (vIdx2 == -1) {
                            vIdx2 = vertexIdx;
                            mapPixelToVertexIdx[idx10] = vertexIdx;
                            vertices.push_back(obs10);
                            pixels.push_back(Eigen::Vector2i(x+1, y));
                            vertexIdx++;
                        }

                        faces.push_back(Eigen::Vector3i(vIdx0, vIdx1, vIdx2));
                    }
                }

                if (valid01 && valid10 && valid11) {
                    float d0 = (obs10 - obs01).norm();
                    float d1 = (obs10 - obs11).norm();
                    float d2 = (obs01 - obs11).norm();

                    if (d0 <= maxTriangleEdgeDistance && d1 <= maxTriangleEdgeDistance && d2 <= maxTriangleEdgeDistance) {
                        int vIdx0 = mapPixelToVertexIdx[idx11];
                        int vIdx1 = mapPixelToVertexIdx[idx10];
                        int vIdx2 = mapPixelToVertexIdx[idx01];

                        if (vIdx0 == -1) {
                            vIdx0 = vertexIdx;
                            mapPixelToVertexIdx[idx11] = vertexIdx;
                            vertices.push_back(obs11);
                            pixels.push_back(Eigen::Vector2i(x+1, y + 1));

                            vertexIdx++;
                        }
                        if (vIdx1 == -1) {
                            vIdx1 = vertexIdx;
                            mapPixelToVertexIdx[idx10] = vertexIdx;
                            vertices.push_back(obs10);
                            pixels.push_back(Eigen::Vector2i(x+1, y));
                            vertexIdx++;
                        }
                        if (vIdx2 == -1) {
                            vIdx2 = vertexIdx;
                            mapPixelToVertexIdx[idx01] = vertexIdx;
                            vertices.push_back(obs01);
                            pixels.push_back(Eigen::Vector2i(x, y+1));
                            vertexIdx++;
                        }

                        faces.push_back(Eigen::Vector3i(vIdx0, vIdx1, vIdx2));
                    }
                }
            }
        }

        // Convert to numpy array.
        int nVertices = vertices.size();
        int nFaces = faces.size();

        if (nVertices > 0 && nFaces > 0) {
            // Reference check should be set to false otherwise there is a runtime
            // error. Check why that is the case.
            vertexPositions.resize({ nVertices, 3 }, false);
            vertexPixels.resize({ nVertices, 2 }, false);
            faceIndices.resize({ nFaces, 3 }, false);

            for (int i = 0; i < nVertices; i++) {
                *vertexPositions.mutable_data(i, 0) = vertices[i].x();
                *vertexPositions.mutable_data(i, 1) = vertices[i].y();
                *vertexPositions.mutable_data(i, 2) = vertices[i].z();

                *vertexPixels.mutable_data(i, 0) = pixels[i].x();
                *vertexPixels.mutable_data(i, 1) = pixels[i].y();
            }

            for (int i = 0; i < nFaces; i++) {
                *faceIndices.mutable_data(i, 0) = faces[i].x();
                *faceIndices.mutable_data(i, 1) = faces[i].y();
                *faceIndices.mutable_data(i, 2) = faces[i].z();
            }
        }
    }

    void compute_mesh_from_depth_and_color(
        const py::array_t<float>& pointImage, const py::array_t<int>& colorImage, float maxTriangleEdgeDistance,
        py::array_t<float>& vertexPositions, py::array_t<int>& vertexColors, py::array_t<int>& faceIndices
    ) {
        int width = pointImage.shape(2);
        int height = pointImage.shape(1);

        // Compute valid pixel vertices and faces.
        // We also need to compute the pixel -> vertex index mapping for
        // computation of faces.
        // We connect neighboring pixels on the square into two triangles.
        // We only select valid triangles, i.e. with all valid vertices and
        // not too far apart.
        // Important: The triangle orientation is set such that the normals
        // point towards the camera.
        std::vector<Eigen::Vector3f> vertices;
        std::vector<Eigen::Vector3i> colors;
        std::vector<Eigen::Vector3i> faces;

        int vertexIdx = 0;
        std::vector<int> mapPixelToVertexIdx(width * height, -1);

        for (int y = 0; y < height - 1; y++) {
            for (int x = 0; x < width - 1; x++) {
                Eigen::Vector3f obs00(*pointImage.data(0, y, x), *pointImage.data(1, y, x), *pointImage.data(2, y, x));
                Eigen::Vector3f obs01(*pointImage.data(0, y + 1, x), *pointImage.data(1, y + 1, x), *pointImage.data(2, y + 1, x));
                Eigen::Vector3f obs10(*pointImage.data(0, y, x + 1), *pointImage.data(1, y, x + 1), *pointImage.data(2, y, x + 1));
                Eigen::Vector3f obs11(*pointImage.data(0, y + 1, x + 1), *pointImage.data(1, y + 1, x + 1), *pointImage.data(2, y + 1, x + 1));

                Eigen::Vector3i color00(*colorImage.data(0, y, x), *colorImage.data(1, y, x), *colorImage.data(2, y, x));
                Eigen::Vector3i color01(*colorImage.data(0, y + 1, x), *colorImage.data(1, y + 1, x), *colorImage.data(2, y + 1, x));
                Eigen::Vector3i color10(*colorImage.data(0, y, x + 1), *colorImage.data(1, y, x + 1), *colorImage.data(2, y, x + 1));
                Eigen::Vector3i color11(*colorImage.data(0, y + 1, x + 1), *colorImage.data(1, y + 1, x + 1), *colorImage.data(2, y + 1, x + 1));

                int idx00 = y * width + x;
                int idx01 = (y + 1) * width + x;
                int idx10 = y * width + (x + 1);
                int idx11 = (y + 1) * width + (x + 1);

                bool valid00 = obs00.z() > 0;
                bool valid01 = obs01.z() > 0;
                bool valid10 = obs10.z() > 0;
                bool valid11 = obs11.z() > 0;

                if (valid00 && valid01 && valid10) {
                    float d0 = (obs00 - obs01).norm();
                    float d1 = (obs00 - obs10).norm();
                    float d2 = (obs01 - obs10).norm();

                    if (d0 <= maxTriangleEdgeDistance && d1 <= maxTriangleEdgeDistance && d2 <= maxTriangleEdgeDistance) {
                        int vIdx0 = mapPixelToVertexIdx[idx00];
                        int vIdx1 = mapPixelToVertexIdx[idx01];
                        int vIdx2 = mapPixelToVertexIdx[idx10];

                        if (vIdx0 == -1) {
                            vIdx0 = vertexIdx;
                            mapPixelToVertexIdx[idx00] = vertexIdx;
                            vertices.push_back(obs00);
                            colors.push_back(color00);
                            vertexIdx++;
                        }
                        if (vIdx1 == -1) {
                            vIdx1 = vertexIdx;
                            mapPixelToVertexIdx[idx01] = vertexIdx;
                            vertices.push_back(obs01);
                            colors.push_back(color01);
                            vertexIdx++;
                        }
                        if (vIdx2 == -1) {
                            vIdx2 = vertexIdx;
                            mapPixelToVertexIdx[idx10] = vertexIdx;
                            vertices.push_back(obs10);
                            colors.push_back(color10);
                            vertexIdx++;
                        }

                        faces.push_back(Eigen::Vector3i(vIdx0, vIdx1, vIdx2));
                    }
                }

                if (valid01 && valid10 && valid11) {
                    float d0 = (obs10 - obs01).norm();
                    float d1 = (obs10 - obs11).norm();
                    float d2 = (obs01 - obs11).norm();

                    if (d0 <= maxTriangleEdgeDistance && d1 <= maxTriangleEdgeDistance && d2 <= maxTriangleEdgeDistance) {
                        int vIdx0 = mapPixelToVertexIdx[idx11];
                        int vIdx1 = mapPixelToVertexIdx[idx10];
                        int vIdx2 = mapPixelToVertexIdx[idx01];

                        if (vIdx0 == -1) {
                            vIdx0 = vertexIdx;
                            mapPixelToVertexIdx[idx11] = vertexIdx;
                            vertices.push_back(obs11);
                            colors.push_back(color11);
                            vertexIdx++;
                        }
                        if (vIdx1 == -1) {
                            vIdx1 = vertexIdx;
                            mapPixelToVertexIdx[idx10] = vertexIdx;
                            vertices.push_back(obs10);
                            colors.push_back(color10);
                            vertexIdx++;
                        }
                        if (vIdx2 == -1) {
                            vIdx2 = vertexIdx;
                            mapPixelToVertexIdx[idx01] = vertexIdx;
                            vertices.push_back(obs01);
                            colors.push_back(color01);
                            vertexIdx++;
                        }

                        faces.push_back(Eigen::Vector3i(vIdx0, vIdx1, vIdx2));
                    }
                }
            }
        }

        // Convert to numpy array.
        int nVertices = vertices.size();
        int nFaces = faces.size();

        if (nVertices > 0 && nFaces > 0) {
            // Reference check should be set to false otherwise there is a runtime
            // error. Check why that is the case.
            vertexPositions.resize({ nVertices, 3 }, false);
            vertexColors.resize({ nVertices, 3 }, false);
            faceIndices.resize({ nFaces, 3 }, false);

            for (int i = 0; i < nVertices; i++) {
                *vertexPositions.mutable_data(i, 0) = vertices[i].x();
                *vertexPositions.mutable_data(i, 1) = vertices[i].y();
                *vertexPositions.mutable_data(i, 2) = vertices[i].z();

                *vertexColors.mutable_data(i, 0) = colors[i].x();
                *vertexColors.mutable_data(i, 1) = colors[i].y();
                *vertexColors.mutable_data(i, 2) = colors[i].z();
            }

            for (int i = 0; i < nFaces; i++) {
                *faceIndices.mutable_data(i, 0) = faces[i].x();
                *faceIndices.mutable_data(i, 1) = faces[i].y();
                *faceIndices.mutable_data(i, 2) = faces[i].z();
            }
        }
    }

    void compute_mesh_from_depth_and_flow(
        const py::array_t<float>& pointImage, const py::array_t<float>& flowImage, float maxTriangleEdgeDistance,
        py::array_t<float>& vertexPositions, py::array_t<float>& vertexFlows, py::array_t<int>& vertexPixels, py::array_t<int>& faceIndices
    ) {
        int width = pointImage.shape(2);
        int height = pointImage.shape(1);

        // Compute valid pixel vertices and faces.
        // We also need to compute the pixel -> vertex index mapping for
        // computation of faces.
        // We connect neighboring pixels on the square into two triangles.
        // We only select valid triangles, i.e. with all valid vertices and
        // not too far apart.
        // Important: The triangle orientation is set such that the normals
        // point towards the camera.
        std::vector<Eigen::Vector3f> vertices;
        std::vector<Eigen::Vector3f> flows;
        std::vector<Eigen::Vector2i> pixels;
        std::vector<Eigen::Vector3i> faces;

        int vertexIdx = 0;
        std::vector<int> mapPixelToVertexIdx(width * height, -1);

        for (int y = 0; y < height - 1; y++) {
            for (int x = 0; x < width - 1; x++) {
                Eigen::Vector3f obs00(*pointImage.data(0, y, x), *pointImage.data(1, y, x), *pointImage.data(2, y, x));
                Eigen::Vector3f obs01(*pointImage.data(0, y + 1, x), *pointImage.data(1, y + 1, x), *pointImage.data(2, y + 1, x));
                Eigen::Vector3f obs10(*pointImage.data(0, y, x + 1), *pointImage.data(1, y, x + 1), *pointImage.data(2, y, x + 1));
                Eigen::Vector3f obs11(*pointImage.data(0, y + 1, x + 1), *pointImage.data(1, y + 1, x + 1), *pointImage.data(2, y + 1, x + 1));

                Eigen::Vector3f flow00(*flowImage.data(0, y, x), *flowImage.data(1, y, x), *flowImage.data(2, y, x));
                Eigen::Vector3f flow01(*flowImage.data(0, y + 1, x), *flowImage.data(1, y + 1, x), *flowImage.data(2, y + 1, x));
                Eigen::Vector3f flow10(*flowImage.data(0, y, x + 1), *flowImage.data(1, y, x + 1), *flowImage.data(2, y, x + 1));
                Eigen::Vector3f flow11(*flowImage.data(0, y + 1, x + 1), *flowImage.data(1, y + 1, x + 1), *flowImage.data(2, y + 1, x + 1));

                int idx00 = y * width + x;
                int idx01 = (y + 1) * width + x;
                int idx10 = y * width + (x + 1);
                int idx11 = (y + 1) * width + (x + 1);

                bool valid00 = obs00.z() > 0 && std::isfinite(flow00.x()) && std::isfinite(flow00.y()) && std::isfinite(flow00.z());
                bool valid01 = obs01.z() > 0 && std::isfinite(flow01.x()) && std::isfinite(flow01.y()) && std::isfinite(flow01.z());
                bool valid10 = obs10.z() > 0 && std::isfinite(flow10.x()) && std::isfinite(flow10.y()) && std::isfinite(flow10.z());
                bool valid11 = obs11.z() > 0 && std::isfinite(flow11.x()) && std::isfinite(flow11.y()) && std::isfinite(flow11.z());

                if (valid00 && valid01 && valid10) {
                    float d0 = (obs00 - obs01).norm();
                    float d1 = (obs00 - obs10).norm();
                    float d2 = (obs01 - obs10).norm();

                    if (d0 <= maxTriangleEdgeDistance && d1 <= maxTriangleEdgeDistance && d2 <= maxTriangleEdgeDistance) {
                        int vIdx0 = mapPixelToVertexIdx[idx00];
                        int vIdx1 = mapPixelToVertexIdx[idx01];
                        int vIdx2 = mapPixelToVertexIdx[idx10];

                        if (vIdx0 == -1) {
                            vIdx0 = vertexIdx;
                            mapPixelToVertexIdx[idx00] = vertexIdx;
                            vertices.push_back(obs00);
                            flows.push_back(flow00);
                            pixels.push_back(Eigen::Vector2i(x, y));
                            vertexIdx++;
                        }
                        if (vIdx1 == -1) {
                            vIdx1 = vertexIdx;
                            mapPixelToVertexIdx[idx01] = vertexIdx;
                            vertices.push_back(obs01);
                            flows.push_back(flow01);
                            pixels.push_back(Eigen::Vector2i(x, y + 1));
                            vertexIdx++;
                        }
                        if (vIdx2 == -1) {
                            vIdx2 = vertexIdx;
                            mapPixelToVertexIdx[idx10] = vertexIdx;
                            vertices.push_back(obs10);
                            flows.push_back(flow10);
                            pixels.push_back(Eigen::Vector2i(x + 1, y));
                            vertexIdx++;
                        }

                        faces.push_back(Eigen::Vector3i(vIdx0, vIdx1, vIdx2));
                    }
                }

                if (valid01 && valid10 && valid11) {
                    float d0 = (obs10 - obs01).norm();
                    float d1 = (obs10 - obs11).norm();
                    float d2 = (obs01 - obs11).norm();

                    if (d0 <= maxTriangleEdgeDistance && d1 <= maxTriangleEdgeDistance && d2 <= maxTriangleEdgeDistance) {
                        int vIdx0 = mapPixelToVertexIdx[idx11];
                        int vIdx1 = mapPixelToVertexIdx[idx10];
                        int vIdx2 = mapPixelToVertexIdx[idx01];

                        if (vIdx0 == -1) {
                            vIdx0 = vertexIdx;
                            mapPixelToVertexIdx[idx11] = vertexIdx;
                            vertices.push_back(obs11);
                            flows.push_back(flow11);
                            pixels.push_back(Eigen::Vector2i(x + 1, y + 1));
                            vertexIdx++;
                        }
                        if (vIdx1 == -1) {
                            vIdx1 = vertexIdx;
                            mapPixelToVertexIdx[idx10] = vertexIdx;
                            vertices.push_back(obs10);
                            flows.push_back(flow10);
                            pixels.push_back(Eigen::Vector2i(x + 1, y));
                            vertexIdx++;
                        }
                        if (vIdx2 == -1) {
                            vIdx2 = vertexIdx;
                            mapPixelToVertexIdx[idx01] = vertexIdx;
                            vertices.push_back(obs01);
                            flows.push_back(flow01);
                            pixels.push_back(Eigen::Vector2i(x, y + 1));
                            vertexIdx++;
                        }

                        faces.push_back(Eigen::Vector3i(vIdx0, vIdx1, vIdx2));
                    }
                }
            }
        }

        // Convert to numpy array.
        int nVertices = vertices.size();
        int nFaces = faces.size();

        if (nVertices > 0 && nFaces > 0) {
            // Reference check should be set to false otherwise there is a runtime
            // error. Check why that is the case.
            vertexPositions.resize({ nVertices, 3 }, false);
            vertexFlows.resize({ nVertices, 3 }, false);
            vertexPixels.resize({ nVertices, 2 }, false);
            faceIndices.resize({ nFaces, 3 }, false);

            for (int i = 0; i < nVertices; i++) {
                *vertexPositions.mutable_data(i, 0) = vertices[i].x();
                *vertexPositions.mutable_data(i, 1) = vertices[i].y();
                *vertexPositions.mutable_data(i, 2) = vertices[i].z();

                *vertexFlows.mutable_data(i, 0) = flows[i].x();
                *vertexFlows.mutable_data(i, 1) = flows[i].y();
                *vertexFlows.mutable_data(i, 2) = flows[i].z();

                *vertexPixels.mutable_data(i, 0) = pixels[i].x();
                *vertexPixels.mutable_data(i, 1) = pixels[i].y();
            }

            for (int i = 0; i < nFaces; i++) {
                *faceIndices.mutable_data(i, 0) = faces[i].x();
                *faceIndices.mutable_data(i, 1) = faces[i].y();
                *faceIndices.mutable_data(i, 2) = faces[i].z();
            }
        }
    }

    void filter_depth(py::array_t<unsigned short>& in, py::array_t<unsigned short>& out, int radius) {
        assert(in.ndim() == 2);
        assert(out.ndim() == 2);

        unsigned kernelSize = 2 * radius + 1;
        unsigned windowSize = kernelSize * kernelSize;

        int width = in.shape(1);
        int height = in.shape(0);
        assert(out.shape(0) == height);
        assert(out.shape(1) == width);

        // #pragma omp parallel for
        for (int y = 0; y < height; y++) {
            for (int x = 0; x < width; x++) {
                // Get all values in the median window.
                int xMin = std::max(x - radius, 0);
                int xMax = std::min(x + radius, int(width) - 1);
                int yMin = std::max(y - radius, 0);
                int yMax = std::min(y + radius, int(height) - 1);

                std::vector<unsigned short> windowValues;
                windowValues.reserve(windowSize);

                for (int yNear = yMin; yNear <= yMax; yNear++) {
                    for (int xNear = xMin; xNear <= xMax; xNear++) {
                        unsigned short depth = *in.data(yNear, xNear);
                        if (depth > 0) {
                            windowValues.push_back(depth);
                        }
                    }
                }

                // Sort the values and pick the median as the middle element.
                unsigned nElements = windowValues.size();
                std::sort(windowValues.begin(), windowValues.end());

                unsigned middleIdx = std::floor(nElements / 2);
                unsigned short median = windowValues[middleIdx];

                // Write out the median value.
                *out.mutable_data(y, x) = median;
            }
        }
    }

    py::array_t<float> warp_flow(const py::array_t<float>& image, const py::array_t<float>& flow, const py::array_t<float>& mask) {
        // We assume:
        //      image shape (3, h, w)
        //      flow shape  (2, h, w)
        //      mask shape  (2, h, w)

        int width = image.shape(2);
        int height = image.shape(1);

        py::array_t<float> imageWarped = py::array_t<float>({3, height, width});
        py::array_t<float> weightsWarped = py::array_t<float>({1, height, width});

        // Initialize to zero.
        for (int v = 0; v < height; v++) {
            for (int u = 0; u < width; u++) {
                *imageWarped.mutable_data(0, v, u) = 0.0;
                *imageWarped.mutable_data(1, v, u) = 0.0;
                *imageWarped.mutable_data(2, v, u) = 0.0;
                *weightsWarped.mutable_data(0, v, u) = 0.0;
            }
        }

        // Compute image values and interpolation weights.
        for (int v = 0; v < height; v++) {
            for (int u = 0; u < width; u++) {
                // Check if pixel is inside the mask.
                if (*mask.data(0, v, u) <= 0 || *mask.data(1, v, u) <= 0) continue;

                // Compute the warped pixel.
                float u_warped = u + *flow.data(0, v, u);
                float v_warped = v + *flow.data(1, v, u);

                int u0 = std::floor(u_warped);
                int u1 = u0 + 1;
                int v0 = std::floor(v_warped);
                int v1 = v0 + 1;

                if (u0 < 0 || u1 >= width || v0 < 0 || v1 >= height) continue;

                // Interpolate the color contributions.
                float du = u_warped - u0;
                float dv = v_warped - v0;

                float w00 = (1 - du)*(1 - dv);
                float w01 = (1 - du)*dv;
                float w10 = du*(1 - dv);
                float w11 = du*dv;

                float c0 = *image.data(0, v, u);
                float c1 = *image.data(1, v, u);
                float c2 = *image.data(2, v, u);

                *imageWarped.mutable_data(0, v0, u0) += w00 * c0;
                *imageWarped.mutable_data(1, v0, u0) += w00 * c1;
                *imageWarped.mutable_data(2, v0, u0) += w00 * c2;
                *imageWarped.mutable_data(0, v1, u0) += w01 * c0;
                *imageWarped.mutable_data(1, v1, u0) += w01 * c1;
                *imageWarped.mutable_data(2, v1, u0) += w01 * c2;
                *imageWarped.mutable_data(0, v0, u1) += w10 * c0;
                *imageWarped.mutable_data(1, v0, u1) += w10 * c1;
                *imageWarped.mutable_data(2, v0, u1) += w10 * c2;
                *imageWarped.mutable_data(0, v1, u1) += w11 * c0;
                *imageWarped.mutable_data(1, v1, u1) += w11 * c1;
                *imageWarped.mutable_data(2, v1, u1) += w11 * c2;

                *weightsWarped.mutable_data(0, v0, u0) += w00;
                *weightsWarped.mutable_data(0, v1, u0) += w01;
                *weightsWarped.mutable_data(0, v0, u1) += w10;
                *weightsWarped.mutable_data(0, v1, u1) += w11;
            }
        }

        // Normalize image.
        for (int v = 0; v < height; v++) {
            for (int u = 0; u < width; u++) {
                float w = *weightsWarped.data(0, v, u);
                if (w > 0) {
                    *imageWarped.mutable_data(0, v, u) /= w;
                    *imageWarped.mutable_data(1, v, u) /= w;
                    *imageWarped.mutable_data(2, v, u) /= w;
                }
                else {
                    *imageWarped.mutable_data(0, v, u) = 1.0;
                    *imageWarped.mutable_data(1, v, u) = 1.0;
                    *imageWarped.mutable_data(2, v, u) = 1.0;
                }
            }
        }

        return imageWarped;
    }

    py::array_t<float> warp_rigid(
        const py::array_t<float>& rgbd,
        const py::array_t<float>& rotation,
        const py::array_t<float>& translation,
        float fx, float fy, float cx, float cy
    ) {
        // We assume:
        //      rgbd shape (6, h, w)
        //      rotation shape  (9)
        //      translation shape  (2)

        int width = rgbd.shape(2);
        int height = rgbd.shape(1);

        float r00 = *rotation.data(0);
        float r01 = *rotation.data(1);
        float r02 = *rotation.data(2);
        float r10 = *rotation.data(3);
        float r11 = *rotation.data(4);
        float r12 = *rotation.data(5);
        float r20 = *rotation.data(6);
        float r21 = *rotation.data(7);
        float r22 = *rotation.data(8);
        float t0 = *translation.data(0);
        float t1 = *translation.data(1);
        float t2 = *translation.data(2);

        py::array_t<float> imageWarped = py::array_t<float>({3, height, width});
        py::array_t<float> weightsWarped = py::array_t<float>({1, height, width});

        // Initialize to zero.
        for (int v = 0; v < height; v++) {
            for (int u = 0; u < width; u++) {
                *imageWarped.mutable_data(0, v, u) = 0.0;
                *imageWarped.mutable_data(1, v, u) = 0.0;
                *imageWarped.mutable_data(2, v, u) = 0.0;
                *weightsWarped.mutable_data(0, v, u) = 0.0;
            }
        }

        // Compute image values and interpolation weights.
        for (int v = 0; v < height; v++) {
            for (int u = 0; u < width; u++) {
                // Compute the warped pixel.
                float x = *rgbd.data(3, v, u);
                float y = *rgbd.data(4, v, u);
                float z = *rgbd.data(5, v, u);
                if (z <= 0) continue;

                float x_def = r00 * x + r01 * y + r02 * z + t0;
                float y_def = r10 * x + r11 * y + r12 * z + t1;
                float z_def = r20 * x + r21 * y + r22 * z + t2;
                if (z_def <= 0) continue;

                float u_warped = fx * x_def / z_def + cx;
                float v_warped = fy * y_def / z_def + cy;

                int u0 = std::floor(u_warped);
                int u1 = u0 + 1;
                int v0 = std::floor(v_warped);
                int v1 = v0 + 1;

                if (u0 < 0 || u1 >= width || v0 < 0 || v1 >= height) continue;

                // Interpolate the color contributions.
                float du = u_warped - u0;
                float dv = v_warped - v0;

                float w00 = (1 - du)*(1 - dv);
                float w01 = (1 - du)*dv;
                float w10 = du*(1 - dv);
                float w11 = du*dv;

                float c0 = *rgbd.data(0, v, u);
                float c1 = *rgbd.data(1, v, u);
                float c2 = *rgbd.data(2, v, u);

                *imageWarped.mutable_data(0, v0, u0) += w00 * c0;
                *imageWarped.mutable_data(1, v0, u0) += w00 * c1;
                *imageWarped.mutable_data(2, v0, u0) += w00 * c2;
                *imageWarped.mutable_data(0, v1, u0) += w01 * c0;
                *imageWarped.mutable_data(1, v1, u0) += w01 * c1;
                *imageWarped.mutable_data(2, v1, u0) += w01 * c2;
                *imageWarped.mutable_data(0, v0, u1) += w10 * c0;
                *imageWarped.mutable_data(1, v0, u1) += w10 * c1;
                *imageWarped.mutable_data(2, v0, u1) += w10 * c2;
                *imageWarped.mutable_data(0, v1, u1) += w11 * c0;
                *imageWarped.mutable_data(1, v1, u1) += w11 * c1;
                *imageWarped.mutable_data(2, v1, u1) += w11 * c2;

                *weightsWarped.mutable_data(0, v0, u0) += w00;
                *weightsWarped.mutable_data(0, v1, u0) += w01;
                *weightsWarped.mutable_data(0, v0, u1) += w10;
                *weightsWarped.mutable_data(0, v1, u1) += w11;
            }
        }

        // Normalize image.
        for (int v = 0; v < height; v++) {
            for (int u = 0; u < width; u++) {
                float w = *weightsWarped.data(0, v, u);
                if (w > 0) {
                    *imageWarped.mutable_data(0, v, u) /= w;
                    *imageWarped.mutable_data(1, v, u) /= w;
                    *imageWarped.mutable_data(2, v, u) /= w;
                }
                else {
                    *imageWarped.mutable_data(0, v, u) = 1.0;
                    *imageWarped.mutable_data(1, v, u) = 1.0;
                    *imageWarped.mutable_data(2, v, u) = 1.0;
                }
            }
        }

        return imageWarped;
    }

    py::array_t<float> warp_3d(const py::array_t<float>& rgbd, const py::array_t<float>& points, const py::array_t<int>& pointValidity, float fx, float fy, float cx, float cy) {
        // We assume:
        //      image shape             (6, h, w)
        //      points shape            (3, h, w)
        //      pointValidity shape     (h, w)

        int width = rgbd.shape(2);
        int height = rgbd.shape(1);

        py::array_t<float> imageWarped = py::array_t<float>({3, height, width});
        py::array_t<float> weightsWarped = py::array_t<float>({1, height, width});

        // Initialize to zero.
        for (int v = 0; v < height; v++) {
            for (int u = 0; u < width; u++) {
                *imageWarped.mutable_data(0, v, u) = 0.0;
                *imageWarped.mutable_data(1, v, u) = 0.0;
                *imageWarped.mutable_data(2, v, u) = 0.0;
                *weightsWarped.mutable_data(0, v, u) = 0.0;
            }
        }

        // Compute image values and interpolation weights.
        for (int v = 0; v < height; v++) {
            for (int u = 0; u < width; u++) {
                // Compute the warped pixel.
                if (*pointValidity.data(v, u) <= 0) continue;

                float z = *rgbd.data(5, v, u);
                if (z <= 0) continue;

                float x_def = *points.data(0, v, u);
                float y_def = *points.data(1, v, u);
                float z_def = *points.data(2, v, u);
                if (z_def <= 0) continue;

                float u_warped = fx * x_def / z_def + cx;
                float v_warped = fy * y_def / z_def + cy;

                int u0 = std::floor(u_warped);
                int u1 = u0 + 1;
                int v0 = std::floor(v_warped);
                int v1 = v0 + 1;

                if (u0 < 0 || u1 >= width || v0 < 0 || v1 >= height) continue;

                // Interpolate the color contributions.
                float du = u_warped - u0;
                float dv = v_warped - v0;

                float w00 = (1 - du)*(1 - dv);
                float w01 = (1 - du)*dv;
                float w10 = du*(1 - dv);
                float w11 = du*dv;

                float c0 = *rgbd.data(0, v, u);
                float c1 = *rgbd.data(1, v, u);
                float c2 = *rgbd.data(2, v, u);

                *imageWarped.mutable_data(0, v0, u0) += w00 * c0;
                *imageWarped.mutable_data(1, v0, u0) += w00 * c1;
                *imageWarped.mutable_data(2, v0, u0) += w00 * c2;
                *imageWarped.mutable_data(0, v1, u0) += w01 * c0;
                *imageWarped.mutable_data(1, v1, u0) += w01 * c1;
                *imageWarped.mutable_data(2, v1, u0) += w01 * c2;
                *imageWarped.mutable_data(0, v0, u1) += w10 * c0;
                *imageWarped.mutable_data(1, v0, u1) += w10 * c1;
                *imageWarped.mutable_data(2, v0, u1) += w10 * c2;
                *imageWarped.mutable_data(0, v1, u1) += w11 * c0;
                *imageWarped.mutable_data(1, v1, u1) += w11 * c1;
                *imageWarped.mutable_data(2, v1, u1) += w11 * c2;

                *weightsWarped.mutable_data(0, v0, u0) += w00;
                *weightsWarped.mutable_data(0, v1, u0) += w01;
                *weightsWarped.mutable_data(0, v0, u1) += w10;
                *weightsWarped.mutable_data(0, v1, u1) += w11;
            }
        }

        // Normalize image.
        for (int v = 0; v < height; v++) {
            for (int u = 0; u < width; u++) {
                float w = *weightsWarped.data(0, v, u);
                if (w > 0) {
                    *imageWarped.mutable_data(0, v, u) /= w;
                    *imageWarped.mutable_data(1, v, u) /= w;
                    *imageWarped.mutable_data(2, v, u) /= w;
                }
                else {
                    *imageWarped.mutable_data(0, v, u) = 1.0;
                    *imageWarped.mutable_data(1, v, u) = 1.0;
                    *imageWarped.mutable_data(2, v, u) = 1.0;
                }
            }
        }

        return imageWarped;
    }

} //namespace image_proc


#include "HybridBot.cpp"
#include "SweepBot.cpp"
#include <kilosim/World.h>
#include <kilosim/Viewer.h>
#include <kilosim/ConfigParser.h>
#include <kilosim/Logger.h>

#include <Eigen/Dense>
#include <Eigen/Eigenvalues>

#include <unistd.h>

// using Eigen::MatrixXd;

// AGGREGATORS

std::vector<double> network_eigenvals(std::vector<Kilosim::Robot *> &robots)
{
    // Get the eigenvalues of the Laplacian of the robot network/communication graph for this:
    // https://en.wikipedia.org/wiki/Algebraic_connectivity

    // L = D - A
    // D is the diagonal matrix of the degree matrix
    // A is the adjacency matrix of the robot network
    // L_{i,j} = ...
    // deg(v_i) if v_i == v_j
    // -1       if v_i != v_j and v_i is connected to v_j
    // 0        otherwise

    // Get the IDs of all robots
    std::vector<int> robot_ids(robots.size());
    for (int i = 0; i < robots.size(); i++)
    {
        robot_ids[i] = robots[i]->id;
    }
    // Initialize a 2D matrix with the number of robots
    Eigen::MatrixXd connectivity_matrix = Eigen::MatrixXd::Zero(robots.size(), robots.size());
    for (int i = 0; i < robots.size(); i++)
    {
        // From each robot, get a list of the neighbors it has heard from on this tick
        Kilosim::BaseBot *bot = (Kilosim::BaseBot *)robots[i];
        std::vector<int> neighbor_ids = bot->get_neighbors();
        // Convert the map to a vector, sorted by neighbor ID (include self as 1!)
        // All neighbors not heard from are 0; all neighbors heard from are 1
        // std::vector<int> network_row(robots.size(), 0);
        for (int j = 0; j < neighbor_ids.size(); j++)
        {
            // Find the index of the neighbor in the list of robot IDs
            // (row order must match column order)
            auto iter = std::find(robot_ids.begin(), robot_ids.end(), neighbor_ids[j]);
            int neighbor_id_ind = iter - robot_ids.begin();
            connectivity_matrix(i, neighbor_id_ind) = -1;
        }
        connectivity_matrix(i, i) = bot->neighbor_count;
    }
    // Compute the eigenvalues of the matrix
    Eigen::VectorXcd eivals = connectivity_matrix.eigenvalues();
    // // Convert the eigenvalues to an std::vector for returning
    // Convert all eigenvalues to real numbers. Those with an imaginary part are converted to nan
    std::vector<double> ei_real(eivals.size());
    for (int i = 0; i < ei_real.size(); i++)
    {
        if (eivals[i].imag() == 0)
        {
            ei_real[i] = eivals[i].real();
        }
        else
        {
            ei_real[i] = std::nan("");
        }
    }
    return ei_real;
}

std::vector<double> decision_states(std::vector<Kilosim::Robot *> &robots)
{
    // Get the current state of the robots' decisions. (undecided, self decided, all decided)
    // This is represented by the robot's m_state
    std::vector<double> decision_states(robots.size());
    for (auto i = 0ul; i < robots.size(); i++)
    {
        Kilosim::BaseBot *bot = (Kilosim::BaseBot *)robots[i];
        decision_states[i] = bot->get_state();
    }
    return decision_states;
}

std::vector<double> has_decided_count(std::vector<Kilosim::Robot *> &robots)
{
    // Count how many robots have made a decision (ie, found value <= threshold)
    double decided_count = 0;
    for (auto i = 0ul; i < robots.size(); i++)
    {
        Kilosim::BaseBot *bot = (Kilosim::BaseBot *)robots[i];
        if (bot->has_decided())
        {
            decided_count++;
        }
    }
    return std::vector<double>{decided_count};
}

std::vector<double> neighbor_count(std::vector<Kilosim::Robot *> &robots)
{
    // Get the number of neighbors each robot has heard from on this tick
    std::vector<double> neighbor_count(robots.size());
    for (auto i = 0ul; i < robots.size(); i++)
    {
        Kilosim::BaseBot *bot = (Kilosim::BaseBot *)robots[i];
        neighbor_count[i] = bot->neighbor_count;
    }
    return neighbor_count;
}

std::vector<double> robot_coverage(std::vector<Kilosim::Robot *> &robots)
{
    // Get the total number of cells covered by each robot.
    // This includes the number of cells covered by the robot's sensors, not where the robot walked.
    // It also doesn't account for duplicates (ie, if a robot is in a cell twice, it will only count once).
    std::vector<double> coverage(robots.size());
    int i = 0;
    for (auto &robot : robots)
    {
        Kilosim::BaseBot *bot = (Kilosim::BaseBot *)robot;
        coverage[i] = bot->map_coverage_count();
        i++;
    }
    return coverage;
}

std::vector<double> robot_distance(std::vector<Kilosim::Robot *> &robots)
{
    // Get the distance each robot has traveled.
    // This is the number of cells the robot has walked, not the number of cells covered by the robot's sensors.
    std::vector<double> distances(robots.size());
    for (auto i = 0ul; i < robots.size(); i++)
    {
        Kilosim::BaseBot *bb = (Kilosim::BaseBot *)robots[i];
        distances[i] = bb->map_visited_count();
    }
    return distances;
}

std::vector<double> collective_coverage(std::vector<Kilosim::Robot *> &robots)
{
    // Add up how much of the map the robots have observed (over all robots).
    // This is very time-intensive -- AT LEAST O(nwh) (n=num_robots, w=width, h=height)
    // Probably most useful to only call at the end to check how much was covered
    std::vector<std::vector<int>> collective_map;
    for (auto i = 0ul; i < robots.size(); i++)
    {
        Kilosim::BaseBot *bb = (Kilosim::BaseBot *)robots[i];
        std::vector<std::vector<int>> map = bb->get_map();
        // On the first go, resize the map to match the robot's map and fill with 0
        if (i == 0)
        {
            collective_map.resize(map.size(), std::vector<int>(map[0].size(), 0));
            // m_map.resize(m_arena_grid_height, std::vector<int>(m_arena_grid_width, -1))
        }
        // Add the robot's map to the collective map
        for (auto j = 0ul; j < map.size(); j++)
        {
            for (auto k = 0ul; k < map[j].size(); k++)
            {
                if (map[j][k] != -1)
                {
                    // Only add 1 if the cell was observed
                    collective_map[j][k] += std::min(1, map[j][k]);
                    collective_map[j][k] = 1;
                }
            }
        }
    }
    // Add up the collective map to get the total number of cells covered
    double total_coverage = 0;
    for (auto j = 0ul; j < collective_map.size(); j++)
    {
        for (auto k = 0ul; k < collective_map[j].size(); k++)
        {
            total_coverage += collective_map[j][k];
        }
    }
    return std::vector<double>{total_coverage};
}

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

bool all_robots_home(std::vector<Kilosim::HybridBot *> robots)
{
    for (auto robot : robots)
    {
        if (!robot->is_home())
        {
            // At least one robot is not home (ie finished)
            return false;
        }
    }
    return true;
}

// Check if all of the robots in the world are finished
// (either they all found the source and returned home, or the time limit expired)
bool is_finished(Kilosim::World &world, std::vector<Kilosim::HybridBot *> robots, std::string end_condition, int end_val)
{
    if (end_condition == "time")
    {
        // Either time is up OR all of the robots are home (done)
        return (world.get_time() * world.get_tick_rate() >= end_val) || all_robots_home(robots);
    }
    else if (end_condition == "value")
    {
        return all_robots_home(robots);
    }
    else if (end_condition == "first_find")
    {
        for (auto robot : robots)
        {
            // Robot is finished if it found a target below threshold
            if (robot->is_finished())
            {
                return true;
            }
        }
        // There are no robots finished
        return false;
    }
    // Else no robots that *didn't* meet the end condition
    return true;
}

bool all_robots_home(std::vector<Kilosim::SweepBot *> &robots)
{
    // Check if all robots have returned to the origin
    for (auto robot : robots)
    {
        if (!robot->is_home())
        {
            // At least one robot is not home (ie finished)
            return false;
        }
    }
    return true;
}

bool is_finished(Kilosim::World &world, std::vector<Kilosim::SweepBot *> robots, std::string end_condition, int end_val)
{
    // HACKY: This is just a copy of the HybridBot version
    if (end_condition == "time")
    {
        return world.get_time() * world.get_tick_rate() >= end_val || all_robots_home(robots);
    }
    else if (end_condition == "value")
    {
        return all_robots_home(robots);
    }
    // Else no robots that *didn't* meet the end condition
    return true;
}

void print_pos_vec(std::vector<Pos> pos_vec)
{
    for (auto pos : pos_vec)
    {
        std::cout << pos.x << ", " << pos.y << std::endl;
    }
}

std::vector<std::vector<Pos>> compute_sweep_paths(int arena_width, int arena_height, int num_robots)
{
    // Compute the predetermined sweep paths for the robots to follow
    std::vector<std::vector<Pos>> paths;
    for (int i = 0; i < num_robots; i++)
    {
        Pos start_pos = {0, 0};
        int x_max = arena_width - i * 3 - 1;
        int x_min = num_robots * 3 - i * 3 - 2;
        int y_move_r = (num_robots * 3 - i * 3) * 2 - 3;
        int y_move_l = (i * 3) * 2 + 3;
        int x = 0;
        int y = i * 3 + 1;
        std::vector<Pos> path = {start_pos};

        // Start by going from the origin to the starting position along the left side
        // To get everyone lined up takes a fixed time dependent on the number of robots
        // Then subtract off the number of ticks it takes for a robot to get to its OWN start pos.
        // time_to_lineup = num_robots * 3 - 1  # Should be -2, but we're accounting for origin @ t=0
        int next_y = start_pos.y;
        int next_x = start_pos.x;
        while (path.back() != Pos({x, y}))
        {
            if (x > next_x)
            {
                next_x += 1;
            }
            next_y += 1;
            path.push_back({next_x, next_y});
        }

        // Go until hitting the bottom
        while (path.back().y < arena_height)
        {
            // x right
            if ((y >= arena_height - y_move_r && arena_height / num_robots % 2 == 1) || num_robots >= arena_width / 3)
            {
                // Last row-set of an odd number or multiples of number of robots
                for (int x_ = 0; x_ < arena_width - 1; x_++)
                {
                    path.push_back({x_, y});
                }
                break;
            }
            else
            {
                for (int x_ = x + 1; x_ < x_max; x_++)
                {
                    path.push_back({x_, y});
                }
                x = x_max - 1;
                // down (right)
                for (int y_ = y + 1; y_ < y_move_r + y; y_++)
                {
                    path.push_back({x, y_});
                }
                y = y + y_move_r;
                // x left
                for (int x_ = x_max - 1; x_ > x_min - 1; x_--)
                {
                    path.push_back({x_, y});
                }
                x = x_min;
                // down (left)
                for (int y_ = y + 1; y_ < y_move_l + y; y_++)
                {
                    path.push_back({x, y_});
                }
                y = y + y_move_l;
                // ... repeat
            }
        }
        // Filter out points outside of the arena
        path.erase(std::remove_if(path.begin(), path.end(), [&](Pos p)
                                  { return p.x < 0 || p.x >= arena_width || p.y < 0 || p.y >= arena_height; }),
                   path.end());
        std::reverse(path.begin(), path.end());
        paths.push_back(path);
    }
    return paths;
}

void hybrid_sim(Kilosim::World &world, Kilosim::Logger &logger, Kilosim::ConfigParser &config)
{

    int num_robots = config.get("num_robots");
    std::string end_condition = config.get("end_condition");
    int end_val;
    try
    {
        end_val = config.get("end_val");
    }
    catch (std::invalid_argument &e)
    {
        if (end_condition == "value")
        {
            end_val = config.get("end_condition_params").at("value").at("end_val");
        }
        else if (end_condition == "time")
        {
            end_val = config.get("end_condition_params").at("time").at("end_val");
        }
        logger.log_param("end_val", end_val);
    }
    int target_val = config.get("target_val");

    unsigned long int max_duration = config.get("max_trial_duration");

    // This accounts for possibility of different ways of doing weights
    double pso_self_weight;
    double pso_group_weight;
    try
    {
        // Try to use the same weight for own/others' observations
        pso_self_weight = config.get("pso_weights");
        pso_group_weight = config.get("pso_weights");
    }
    catch (std::exception &e)
    {
        // Use different weights for own/others' observations
        pso_self_weight = config.get("pso_self_weight");
        pso_group_weight = config.get("pso_group_weight");
    }

    std::vector<Kilosim::HybridBot *> robots(num_robots);
    for (int n = 0; n < num_robots; n++)
    {
        robots[n] = new Kilosim::HybridBot();
        world.add_robot(robots[n]);
        robots[n]->robot_init(0, 0, 0);
        robots[n]->max_speed = config.get("max_speed");
        // PSO/GD hybrid pre-decision movement
        robots[n]->pso_step_interval = config.get("pso_step_interval");
        robots[n]->is_pso_step_interval_constant = config.get("is_pso_step_interval_constant", true);
        robots[n]->pso_self_weight = pso_self_weight;
        robots[n]->pso_group_weight = pso_group_weight;
        // This is only used with the old init method. It's left here for legacy reasons.
        // (Now, robots, start by going to a random initial position)
        robots[n]->start_interval = (int)config.get("pso_step_interval") * 3;
        robots[n]->pso_inertia = config.get("pso_inertia");
        robots[n]->gradient_weight = config.get("gradient_weight");
        // comm_range is used by comm_criteria to determine communication range
        // In config, comm_range is in grid cells (as dimension)
        robots[n]->comm_range = (int)config.get("comm_range") * 10;
        robots[n]->num_neighbors = num_robots - 1;
        robots[n]->rx_table_timeout = config.get("rx_table_timeout");
        // End conditions
        robots[n]->end_condition = end_condition;
        robots[n]->end_val = end_val;
        robots[n]->target_val = target_val;
        // Post-decision movement options
        robots[n]->post_decision_movement = config.get("post_decision_movement");
        // Boids parameters
        robots[n]->lj_a = config.get("lj_a");
        robots[n]->lj_b = config.get("lj_b");
        robots[n]->lj_epsilon = config.get("lj_epsilon");
        robots[n]->lj_gamma = config.get("lj_gamma");
        robots[n]->boids_step_interval = config.get("boids_step_interval");
    }

    logger.add_aggregator("num_neighbors", neighbor_count);
    logger.add_aggregator("network_eigenvals", network_eigenvals);
    logger.add_aggregator("decision_states", decision_states);
    logger.add_aggregator("has_decided_count", has_decided_count);

    Kilosim::Viewer viewer(world, 1080);
    // viewer.set_show_network(true);
    // viewer.set_show_tags(true);

    sleep(2);
    while (!is_finished(world, robots, end_condition, end_val) &&
           world.get_tick() < max_duration)
    {
        viewer.draw();
        world.step();
        if (world.get_tick() % 20 == 0)
        {
            logger.log_state();
        }
        usleep(10000);
    }
}

void sweep_sim(Kilosim::World &world, Kilosim::Logger &logger, Kilosim::ConfigParser &config)
{
    // Kilosim::Viewer viewer(world, 1200);
    // viewer.set_show_network(true);
    // viewer.set_show_tags(true);

    int num_robots = config.get("num_robots");
    std::string end_condition = config.get("end_condition");
    int end_val;
    try
    {
        end_val = config.get("end_val");
    }
    catch (std::invalid_argument &e)
    {
        // Try getting it from sub-config instead
        end_val = config.get("end_condition_params").at(end_condition).at("end_val");
        logger.log_param("end_val", end_val);
    }
    unsigned long int max_duration = config.get("max_trial_duration");

    std::vector<std::vector<Pos>> sweep_paths =
        compute_sweep_paths(config.get("world_grid_width"),
                            config.get("world_grid_height"),
                            config.get("num_robots"));
    int max_path_len = 0;
    for (auto path : sweep_paths)
    {
        if (path.size() > max_path_len)
        {
            max_path_len = path.size();
        }
    }

    std::vector<Kilosim::SweepBot *> robots(num_robots);
    for (int n = 0; n < num_robots; n++)
    {
        std::vector<Pos> this_path = sweep_paths[n];

        Pos start_pos = this_path.back();
        this_path.pop_back();
        robots[n] = new Kilosim::SweepBot();
        world.add_robot(robots[n]);
        robots[n]->robot_init(start_pos.x, start_pos.y, 0);
        robots[n]->end_condition = end_condition;
        robots[n]->end_val = end_val;
        robots[n]->set_path_external(this_path);
        robots[n]->max_path_len = max_path_len;
        // In config, comm_range is in grid cells (as dimension)
        robots[n]->comm_range = (int)config.get("comm_range") * 10;
    }

    logger.add_aggregator("decision_states", decision_states);

    // sleep(2);
    while (!is_finished(world, robots, end_condition, end_val) &&
           world.get_tick() < max_duration)
    {
        // viewer.draw();
        world.step();
        if (world.get_tick() % 20 == 0)
        {
            logger.log_state();
        }
    }
}

int main(int argc, char *argv[])
{
    // Get config file name
    std::vector<std::string> args(argv, argv + argc);
    if (args.size() < 2)
    {
        std::cout << "ERROR: You must provide a config file name" << std::endl;
        exit(1);
    }
    Kilosim::ConfigParser config(args[1]);

    // Get configuration values
    const int start_trial = config.get("start_trial");
    const int num_trials = config.get("num_trials");
    const std::string img_dir = config.get("img_dir");
    const int env_octaves = config.get("env_octaves");
    const int world_grid_width = (int)config.get("world_grid_width");
    const int world_grid_height = (int)config.get("world_grid_height");
    const int world_width = world_grid_width * 10;
    const int world_height = world_grid_height * 10;
    const std::string movement_type = config.get("movement_type");

    for (auto trial = start_trial; trial < start_trial + num_trials; trial++)
    {
        std::string trial_str = std::to_string(trial);
        trial_str.insert(trial_str.begin(), 3 - trial_str.size(), '0');
        const std::string img_filename = img_dir + "/" + trial_str +
                                         "_oct=" + std::to_string(env_octaves) + ".png";
        // const std::string img_filename = img_dir + "/img_oct=" + std::to_string(env_octaves) +
        //                                  "_" + trial_str + ".png";
        // std::cout << img_filename << std::endl;

        // Create 3m x 3m world (no background image, for now)
        Kilosim::World world(
            world_width, world_height, img_filename // World image
        );

        // Leaving this out keeps the default communication rate of every 3 ticks
        // Hopefully this will speed everything up without damaging results
        // world.set_comm_rate(1);
        world.set_comm_rate(2);

        // Create log file
        std::string log_filename = (std::string)config.get("log_dir") + "data.h5";
        // False = don't overwrite logs
        Kilosim::Logger logger(world, log_filename, trial, false);
        // Kilosim::Logger logger(world, log_filename, trial, false);
        // False = don't warn about config parameters that can't be saved
        logger.log_config(config, false);

        // Run the right simulation
        std::cout << "Starting trial " << trial << " (" << config.get("log_dir") << ")" << std::endl;
        if (movement_type == "hybrid")
        {
            hybrid_sim(world, logger, config);
        }
        else if (movement_type == "sweep")
        {
            sweep_sim(world, logger, config);
        }

        // These are things that are only logged once, at the end of the experiment!
        logger.add_aggregator("robot_distance", robot_distance);
        logger.add_aggregator("robot_coverage", robot_coverage);
        logger.add_aggregator("collective_coverage", collective_coverage);
        logger.log_state();

        std::cout << "Finished trial " << trial << "\tt=" << world.get_time() * world.get_tick_rate() << std::endl;
    }

    printf("Finished\n");

    return 0;
}


// Std includes
#include <cmath>
#include <iostream>
// Thirdparties includes
#include <Eigen/Dense>
// Lib includes
#include "s0s/runge_kutta_fehlberg.h"
#include "sl0/point_inertial.h"
// Simple includes
#include "flow.h"

using TypeScalar = double;
// Linear Algebra
template<int Size>
using TypeVector = Eigen::Matrix<TypeScalar, Size, 1>;
// Space
constexpr unsigned int DIM = 3;
using TypeSpaceVector = Eigen::Matrix<TypeScalar, DIM, 1>;
// Ref and View
template<typename ...Args>
using TypeRef = Eigen::Ref<Args...>;
template<typename ...Args>
using TypeView = Eigen::Map<Args...>;
// Solver
using TypeSolver = s0s::SolverRungeKuttaFehlberg<TypeVector<Eigen::Dynamic>, TypeView>;
// Flow
using TypeFlow = Flow<TypeSpaceVector, TypeRef>;

int main () {
    TypeSpaceVector x0 = TypeSpaceVector::Constant(1.0);
    double t0 = 0.0;
    double dt = 1e-3;
    double tEnd = 1.0;
    unsigned int nt = std::round((tEnd - t0) / dt);
    // Create point
    sl0::PointInertial<TypeVector, DIM, TypeView, TypeFlow, TypeSolver> point(std::make_shared<TypeFlow>(), 1.0);
    // Set initial state
    point.sStep->x(point.state.data()) = x0;
    point.t = t0;
    // Computation
    for(std::size_t i = 0; i < nt; i++) {
        point.update(dt);
    }
    // out
    std::cout << "\n";
    std::cout << "Point advected following a an exponential flow, exp(" << point.t << ") = " << "\n";
    std::cout << "\n";
    std::cout << "Point Final Position : " << "\n" << point.sStep->x(point.state.data()) << "\n";
    std::cout << "Point Final Velocity : " << "\n" << point.sStep->u(point.state.data()) << "\n";
    std::cout << std::endl;
}


#include <math.h>
#include <stdlib.h>
#include <string>

#include <opencv2/highgui/highgui.hpp>
#include <opencv2/opencv.hpp>
//#include <opencv2/legacy/compat.hpp>

#include "dlib/opencv.h"
#include "dlib/image_processing/frontal_face_detector.h"
#include "dlib/image_processing/render_face_detections.h"
#include "dlib/gui_widgets.h"
#include <dlib/image_processing.h>

#include "util.h"
#include "faceDetection.h"

double get_conversion_factor (dlib::full_object_detection shape, FacePose* face_pose, double magnitude_normal, int mode) {
	cv::Point p1, p2;
    //mode : 1 for left eye, 2 for right eye
	if(mode == 1) {
		p1 = cv::Point(shape.part(42).x(), shape.part(42).y());
		p2 = cv::Point(shape.part(45).x(), shape.part(45).y());
	}
	else if(mode == 2) {
		p1 = cv::Point(shape.part(36).x(), shape.part(36).y());
		p2 = cv::Point(shape.part(39).x(), shape.part(39).y());
	}

	double dx = p1.x - p2.x, dy = p1.y - p2.y;
	double temp1, temp2, beta;
	double n1 = face_pose->normal[0], n2 = face_pose->normal[1], n3 = face_pose->normal[2];
	double beta_old = sqrt(dx*dx + dy*dy)/magnitude_normal;

	temp1 = dx*dx*(1.0 - n2*n2);
	temp2 = dy*dy*(1.0 - n1*n1);

	beta = sqrt(temp1 + temp2)/((double)(magnitude_normal*fabs(n3)));
	beta = 1.0/((double) beta);

	//std::cout<<"Beta : "<<beta<<std::endl;
	return beta;
}

void compute_vec_LR (cv::Point p1, cv::Point p2, FacePose* face_pose, std::vector<double>& LR) {
	double scale = 20.784/30.0;

	LR[0] = p1.x - p2.x;
	LR[1] = p1.y - p2.y;
	LR[0] = LR[0]*scale;
	LR[1] = LR[1]*scale;
	LR[2] = -(LR[0]*face_pose->normal[0] + LR[1]*face_pose->normal[1])/face_pose->normal[3];
}

void get_quadratic_solution (std::vector<double> coeff, double& solution, int mode) {
	solution = (-coeff[1] + mode*sqrt(coeff[1]*coeff[1] - 4*coeff[0]*coeff[2]))/(2*coeff[0]);
	std::cout<<"soln : "<<solution<<std::endl;
}

void get_quadratic_equation (std::vector<double> coeff, std::vector<double>& quad_eqn) {
	quad_eqn[0] = coeff[0]*coeff[0];
	quad_eqn[1] = 2*coeff[0]*coeff[1];
	quad_eqn[2] = coeff[1]*coeff[1];
}

void solve(std::vector<double> coeff_1, double const_1, std::vector<double> coeff_2, double const_2, double mag, std::vector<double>& vec, int mode) {
	double det = coeff_1[0]*coeff_2[1] - coeff_1[1]*coeff_2[0];

	std::vector<double> linear_eqn_1(2), linear_eqn_2(2);
	linear_eqn_1[0] = (coeff_1[1]*coeff_2[2] - coeff_1[2]*coeff_2[1])/det;
	linear_eqn_1[1] = (const_1*coeff_2[1] - coeff_1[1]*const_2)/det;
	linear_eqn_2[0] = (coeff_1[2]*coeff_2[0] - coeff_1[0]*coeff_2[2])/det;
	linear_eqn_2[1] = (coeff_1[0]*const_2 - coeff_2[0]*const_1)/det;

	std::vector<double> quad_eqn_1(3), quad_eqn_2(3), quad_eqn_final(3);
	get_quadratic_equation(linear_eqn_1, quad_eqn_1);
	get_quadratic_equation(linear_eqn_2, quad_eqn_2);

	quad_eqn_final[0] = quad_eqn_1[0] + quad_eqn_2[0] + 1;
	quad_eqn_final[1] = quad_eqn_1[1] + quad_eqn_2[1];
	quad_eqn_final[2] = quad_eqn_1[2] + quad_eqn_2[2] - mag*mag;

	//std::cout<<"const_1 : "<<const_1<<" const_2 : "<<const_2<<std::endl;
	std::vector<double> coeff = quad_eqn_final;
	std::cout<<"Discriminant : "<<coeff[1]*coeff[1] - 4*coeff[0]*coeff[2]<<std::endl;
	log_vec("quad_eqn_final", quad_eqn_final);
	get_quadratic_solution (quad_eqn_final, vec[2], mode);
	vec[0] = linear_eqn_1[0]*vec[2] + linear_eqn_1[1];
	vec[1] = linear_eqn_2[0]*vec[2] + linear_eqn_2[1];
}

void get_section(cv::Point p1, cv::Point p2, cv::Point pupil, double& Y1, double& Y2, double& h) {
	std::vector<double> line(3);

	line[0] = p2.y - p1.y;
	line[1] = -(p2.x - p1.x);
	line[2] =  p1.y*(p2.x - p1.x) - p1.x*(p2.y - p1.y);

	cv::Point pupil_proj;
	pupil_proj.x = -(line[0]*pupil.x + line[1]*pupil.y + line[2])*line[0]/(line[0]*line[0] + line[1]*line[1]) + pupil.x;
	pupil_proj.y = -(line[0]*pupil.x + line[1]*pupil.y + line[2])*line[1]/(line[0]*line[0] + line[1]*line[1]) + pupil.y;

	Y1 = get_distance (p1, pupil_proj);
	Y2 = get_distance (p2, pupil_proj);
	h = get_distance (pupil, pupil_proj);
}

//List : Y1, Y2 can be interchanged. Magnitudes of the vectors in real world may be wrong.
//		 mag_LR square changed to just mag_LR.

void compute_vec_CP(cv::Point p1, cv::Point p2, cv::Point pupil, cv::Rect rect, FacePose* face_pose,
	std::vector<double> vec_CR_u, double mag_CR, std::vector<double> vec_LR_u, double mag_LR,
	std::vector<double> vec_UD_u, double mag_CP, std::vector<double>& vec_CP, double S2R, int mode) {
	double Y1, Y2, H;
	get_section(p1, p2, cv::Point(pupil.x + rect.x, pupil.y + rect.y), Y1, Y2, H);

	double const_1, const_2;
	const_1 = (S2R*H);///std::cos(face_pose->pitch);
	if(mode == 1) {
		const_2 = mag_CR*(scalar_product(vec_CR_u, vec_LR_u)) + ((mag_LR*Y1)/((double) (Y1 + Y2)));
	}
	else if(mode == 2) {
		const_2 = mag_CR*(scalar_product(vec_CR_u, vec_LR_u)) + ((mag_LR*Y2)/((double) (Y1 + Y2)));
	}

	//std::cout<<"Y1 : "<<Y1<<" Y2 : "<<Y2<<" H : "<<H<<std::endl;
	std::cout<<"CP - constants : "<<const_1<<" "<<const_2<<std::endl;

	solve(vec_UD_u, const_1, vec_LR_u, const_2, mag_CP, vec_CP, 1);
}

bool vec_isnan(std::vector<double>& vec) {
	int f = 0;
	for(int i=0; i<vec.size(); i++) {
		if(std::isnan(vec[i])) {
			f=1;
			break;
		}
	}

	if(f) {
		for(int i=0; i<vec.size(); i++) {
			vec[i] = 0.0;
		}
	}

	return (1-f);
}

void compute_eye_gaze (FacePose* face_pose, dlib::full_object_detection shape, cv::Rect rect, cv::Point pupil, double mag_CP, double mag_LR, double mag_CR, double mag_CM, double theta, int mode, std::vector<double>& vec_CP) {

	std::vector<double> vec_LR_u(3), vec_RP(3), vec_CR_u(3), vec_CM_u(3), vec_UD_u(3);
	std::vector<double> vec_CP_l(3), vec_CP_r(3);
	double S2R = get_conversion_factor(shape, face_pose, mag_CM, mode);

	cv::Point p1, p2;
    //mode : 1 for left eye, 2 for right eye
	if(mode == 1) {
		p1 = cv::Point(shape.part(42).x(), shape.part(42).y());
		p2 = cv::Point(shape.part(45).x(), shape.part(45).y());
	}
	else if(mode == 2) {
		p1 = cv::Point(shape.part(36).x(), shape.part(36).y());
		p2 = cv::Point(shape.part(39).x(), shape.part(39).y());
	}

	vec_CP[0] = 1.0;
	vec_CP[1] = 1.0;
	vec_CP[2] = 1.0;

	compute_vec_LR(p1, p2, face_pose, vec_LR_u);
	make_unit_vector(vec_LR_u, vec_LR_u);

	//log_vec("LR", vec_LR_u);

	vec_CM_u[0] = face_pose->normal[0];
	vec_CM_u[1] = face_pose->normal[1];
	vec_CM_u[2] = face_pose->normal[2];

	cross_product(vec_CM_u, vec_LR_u, vec_UD_u);
	make_unit_vector(vec_UD_u, vec_UD_u);

	//log_vec("UD", vec_UD_u);

	double const_1 = std::cos(theta/2.0);
	double const_2 = 0.0;

	solve(vec_UD_u, const_1, vec_CM_u, const_2, 1.0, vec_CR_u, -1);
	make_unit_vector(vec_CR_u, vec_CR_u);

	//log_vec("CR", vec_CR_u);

	compute_vec_CP(p1, p2, pupil, rect, face_pose, vec_CR_u, mag_CR, vec_LR_u, mag_LR,
		vec_UD_u, mag_CP, vec_CP_l, S2R, 2);

	compute_vec_CP(p1, p2, pupil, rect, face_pose, vec_CR_u, mag_CR, vec_LR_u, mag_LR,
		vec_UD_u, mag_CP, vec_CP_r, S2R, 1);

	double f1 = vec_isnan(vec_CP_l);
	double f2 = vec_isnan(vec_CP_r);

	if(f1 || f2) {
			vec_CP[0] = (vec_CP_l[0] + vec_CP_r[0]);
			vec_CP[1] = (vec_CP_l[1] + vec_CP_r[1]);
			vec_CP[2] = (vec_CP_l[2] + vec_CP_r[2]);
	}
	else {
		vec_CP[0] = (vec_CP_l[0] + vec_CP_r[0])/2.0;
		vec_CP[1] = (vec_CP_l[1] + vec_CP_r[1])/2.0;
		vec_CP[2] = (vec_CP_l[2] + vec_CP_r[2])/2.0;
	}

	log_vec("CP_l", vec_CP_l);
	log_vec("CP_r", vec_CP_r);
	log_vec("CP", vec_CP);
}


// Copyright (c) 2021 FRC Team 3512. All Rights Reserved.

#pragma once

#include <Eigen/Cholesky>
#include <Eigen/Core>
#include <drake/math/discrete_algebraic_riccati_equation.h>
#include <frc/system/Discretization.h>
#include <units/time.h>

#include "controllers/DARE.hpp"

namespace frc3512 {

/**
 * Returns the LQR controller gain for the given coefficients and plant.
 *
 * @param A  Continuous system matrix of the plant being controlled.
 * @param B  Continuous input matrix of the plant being controlled.
 * @param Q  The state cost matrix.
 * @param R  The input cost matrix.
 * @param dt Discretization timestep.
 */
template <int States, int Inputs>
Eigen::Matrix<double, Inputs, States> LQR(
    const Eigen::Matrix<double, States, States>& A,
    const Eigen::Matrix<double, States, Inputs>& B,
    const Eigen::Matrix<double, States, States>& Q,
    const Eigen::Matrix<double, Inputs, Inputs>& R, units::second_t dt) {
    Eigen::Matrix<double, States, States> discA;
    Eigen::Matrix<double, States, Inputs> discB;
    frc::DiscretizeAB<States, Inputs>(A, B, dt, &discA, &discB);

    Eigen::Matrix<double, States, States> S =
        drake::math::DiscreteAlgebraicRiccatiEquation(discA, discB, Q, R);
    return (discB.transpose() * S * discB + R)
        .llt()
        .solve(discB.transpose() * S * discA);
}

/**
 * Returns the LQR controller gain for the given coefficients and plant.
 *
 * @param A Discrete system matrix of the plant being controlled.
 * @param B Discrete input matrix of the plant being controlled.
 * @param Q The state cost matrix.
 * @param R The input cost matrix.
 * @param N The state-input cross-term cost matrix.
 */
template <int States, int Inputs>
Eigen::Matrix<double, Inputs, States> LQR(
    const Eigen::Matrix<double, States, States>& A,
    const Eigen::Matrix<double, States, Inputs>& B,
    const Eigen::Matrix<double, States, States>& Q,
    const Eigen::Matrix<double, Inputs, Inputs>& R,
    const Eigen::Matrix<double, States, Inputs>& N) {
    Eigen::Matrix<double, States, States> S =
        DARE<States, Inputs>(A, B, Q, R, N);
    return (B.transpose() * S * B + R)
        .llt()
        .solve(B.transpose() * S * A + N.transpose());
}

}  // namespace frc3512


#define BOOST_TEST_DYN_LINK
#define BOOST_TEST_MODULE HeapTest
/*
* @author : Ionesio Junior
*/

#include <boost/test/unit_test.hpp>
#include "../BinaryHeap.hpp"
#include <cstdlib>
#include <iostream>
#include <time.h>
#include <vector>

struct setUp{
	BinaryHeap<int> *heap = new BinaryHeap<int>(10);
	int empty[0] = {};
	int even[10] = {9,8,7,6,5,4,3,2,1,0};
	int odd[9] = {9,8,7,6,5,4,3,2,1};
	int repeated[9] = {8,5,8,5,9,5,7,8,5};

	int equal[5] = {5,5,5,5,5};

	int repeatedOrdered[9] = {5,5,5,5,7,8,8,8,9};
	int evenOrdered[10] = {0,1,2,3,4,5,6,7,8,9};
	int oddOrdered[9] = {1,2,3,4,5,6,7,8,9};
};

BOOST_FIXTURE_TEST_SUITE(HeapTest,setUp)
BOOST_AUTO_TEST_CASE(testInit){
	BOOST_CHECK(true == heap->isEmpty());
	BOOST_CHECK(false == heap->isFull());
	BOOST_CHECK(NULL == heap->rootElement());
	BOOST_CHECK(0 == heap->size());
}

BOOST_AUTO_TEST_CASE(testIsEmpty){
	BOOST_CHECK(true == heap->isEmpty());
	heap->insert(10);
	BOOST_CHECK(false == heap->isEmpty());
	heap->extractRoot();
	BOOST_CHECK(true == heap->isEmpty());
}

BOOST_AUTO_TEST_CASE(testSize){
	BOOST_CHECK(0 == heap->size());
	for(int i = 0 ; i < 5;i++){
		heap->insert(i);
		BOOST_CHECK(i + 1 == heap->size());
	}
	BOOST_CHECK(5 == heap->size());
	for(int i = 10; i < 20 ;i = i + 2){
		heap->insert(i);
	}
	BOOST_CHECK(10 == heap->size());

	for(int i = 0 ; i < 5;i++){
		heap->extractRoot();
	}

	BOOST_CHECK(5 == heap->size());
}

BOOST_AUTO_TEST_CASE(testInsert){
	BOOST_CHECK(0 == heap->size());

	//Test insert repeated elements
	heap->insert(5);
	BOOST_CHECK(1 == heap->size());
	heap->insert(5);
	BOOST_CHECK(2 == heap->size());

	heap->extractRoot();
	heap->extractRoot();
	BOOST_CHECK(0 == heap->size());

	//Test responsive array
	for(int i = 0 ; i < 10;i++){
		heap->insert(i);
		BOOST_CHECK(i == *heap->rootElement());
	}
	BOOST_CHECK(true == heap->isFull());
	heap->insert(10);
	BOOST_CHECK(false == heap->isFull());
	BOOST_CHECK(11 == heap->size());
}

BOOST_AUTO_TEST_CASE(testExtractRoot){
	BOOST_CHECK(0 == heap->size());
	BOOST_CHECK(true == heap->isEmpty());

	for(int i = 0; i < 100;i++){
		heap->insert(i);
	}

	BOOST_CHECK(100 == heap->size());
	BOOST_CHECK(false == heap->isEmpty());
	BOOST_CHECK(99 == *heap->rootElement());

	for(int i = 99; i >= 0;i--){
		BOOST_CHECK(i == heap->extractRoot());
	}

	BOOST_CHECK(0 == heap->size());
	BOOST_CHECK(true == heap->isEmpty());

	for(int i = 0 ; i < 10;i++){
		if(i < 5){
			heap->insert(10);
		}else{
			heap->insert(20);
		}
	}

	BOOST_CHECK(10 == heap->size());

	for(int i = 0 ; i < 10;i++){
		if(i < 5){
			BOOST_CHECK(20 == heap->extractRoot());
		}else{
			BOOST_CHECK(10 == heap->extractRoot());
		}
	}
}

BOOST_AUTO_TEST_CASE(testRootElement){
	BOOST_CHECK(true == heap->isEmpty());
	heap->insert(100);
	for(int i = 0 ; i < 10;i++){
		BOOST_CHECK(100 == *heap->rootElement());
	}

	BOOST_CHECK(1 == heap->size());

	for(int i = 0 ; i < 10;i++){
		heap->insert(i);
		BOOST_CHECK(100 == *heap->rootElement());
	}

	for(int i = 0 ; i <= 10;i++){
		BOOST_CHECK(*heap->rootElement() == heap->extractRoot());
	}
}

BOOST_AUTO_TEST_CASE(testHeapSort){
	//Test Even Array
	Array<int> result = heap->heapSort(even,10);
	for(int i = 0 ; i < result.size;i++){
		BOOST_CHECK(evenOrdered[i] == result[i]);
	}

	//Test Odd Array
	result = heap->heapSort(odd,9);
	for(int i = 0 ; i < result.size;i++){
		BOOST_CHECK(oddOrdered[i] == result[i]);
	}

	//Test Empty array
	result = heap->heapSort(empty,0);
	for(int i = 0; i < result.size;i++){
		BOOST_CHECK(empty[i] == result[i]);
	}

	//Test Repeated Array
	result = heap->heapSort(repeated,9);
	for(int i = 0 ; i < result.size;i++){
		BOOST_CHECK(repeatedOrdered[i] == result[i]);
	}

	//Test equal array
	result = heap->heapSort(equal,5);
	for(int i = 0 ; i < result.size;i++){
		BOOST_CHECK(equal[i] == result[i]);
	}
}

BOOST_AUTO_TEST_CASE(testBuildHeap){

	//Test even array
	BOOST_CHECK(0 == heap->size());
	BOOST_CHECK(true == heap->isEmpty());
	heap->buildHeap(even,10);
	BOOST_CHECK(10 == heap->size());
	BOOST_CHECK(false == heap->isEmpty());
	for(int i = 9 ; i >= 0;i--){
		BOOST_CHECK(i == heap->extractRoot());
	}

	//Test odd array
        BOOST_CHECK(0 == heap->size());
        BOOST_CHECK(true == heap->isEmpty());
        heap->buildHeap(odd,9);
        BOOST_CHECK(9 == heap->size());
        BOOST_CHECK(false == heap->isEmpty());
        for(int i = 9 ; i > 0;i--){
                BOOST_CHECK(i == heap->extractRoot());
        }

	//Test empty array
        BOOST_CHECK(0 == heap->size());
        BOOST_CHECK(true == heap->isEmpty());
        heap->buildHeap(empty,0);
        BOOST_CHECK(0 == heap->size());
        BOOST_CHECK(true == heap->isEmpty());

	//Test equal array
        BOOST_CHECK(0 == heap->size());
        BOOST_CHECK(true == heap->isEmpty());
        heap->buildHeap(equal,5);
        BOOST_CHECK(5 == heap->size());
        BOOST_CHECK(false == heap->isEmpty());
        for(int i = 5 ; i > 0;i--){
                BOOST_CHECK(5 == heap->extractRoot());
        }

	//Test repeated array
	int reverseRepeated[9] = {9,8,8,8,7,5,5,5,5};
        BOOST_CHECK(0 == heap->size());
        BOOST_CHECK(true == heap->isEmpty());
        heap->buildHeap(repeated,9);
        BOOST_CHECK(9 == heap->size());
        BOOST_CHECK(false == heap->isEmpty());
        for(int i = 0 ; i < 9;i++){
		BOOST_CHECK(reverseRepeated[i] == heap->extractRoot());
        }
}

BOOST_AUTO_TEST_CASE(testException){
	BOOST_CHECK_THROW(heap->extractRoot(),HeapUnderflowException);
}

BOOST_AUTO_TEST_SUITE_END()


/*
Part of the Fluid Corpus Manipulation Project (http://www.flucoma.org/)
Copyright 2017-2019 University of Huddersfield.
Licensed under the BSD-3 License.
See license.md file in the project root for full license information.
This project has received funding from the European Research Council (ERC)
under the European Union’s Horizon 2020 research and innovation programme
(grant agreement No 725899).
*/

#pragma once

#include "../util/AlgorithmUtils.hpp"
#include "../util/FluidEigenMappings.hpp"
#include "../util/KWeightingFilter.hpp"
#include "../util/TruePeak.hpp"
#include "../../data/FluidIndex.hpp"
#include "../../data/TensorTypes.hpp"
#include <Eigen/Core>
#include <cmath>

namespace fluid {
namespace algorithm {

class Loudness
{

public:
  Loudness(index maxSize) : mTP(maxSize) {}

  void init(index size, double sampleRate)
  {
    mFilter.init(sampleRate);
    mTP.init(size, sampleRate);
    mSize = size;
    mInitialized = true;
  }

  void processFrame(const RealVectorView& input, RealVectorView output,
                    bool weighting, bool truePeak)
  {
    using namespace Eigen;
    using namespace std;
    assert(mInitialized);
    assert(output.size() == 2);
    assert(input.size() == mSize);
    ArrayXd in = _impl::asEigen<Array>(input);
    ArrayXd filtered(mSize);
    for (index i = 0; i < mSize; i++)
      filtered(i) = weighting ? mFilter.processSample(in(i)) : in(i);
    double loudness = -0.691 + 10 * log10(filtered.square().mean() + epsilon);
    double peak = truePeak ? mTP.processFrame(input) : in.abs().maxCoeff();
    peak = 20 * log10(peak + epsilon);
    output(0) = loudness;
    output(1) = peak;
  }

private:
  TruePeak         mTP;
  KWeightingFilter mFilter;
  index            mSize{1024};
  bool             mInitialized{false};
};

} // namespace algorithm
} // namespace fluid


#include <Eigen/Dense>
#include <Eigen/SparseCholesky>
#include <Eigen/SparseLU>
#include <Eigen/IterativeLinearSolvers>

using namespace Eigen;

typedef SparseMatrix<double, ColMajor> MySparseMatrix;

extern "C" {

int solve_eigen_icholt_coo(
    const double *coo_data,
    const int *row,
    const int *col,
    int nnz,
    const double *b,
    double *x,
    int n,
    double rtol,
    double initial_shift
){
    MySparseMatrix A(n, n);

    std::vector<Triplet<double> > triplets(nnz);

    for (int k = 0; k < nnz; k++){
        triplets[k] = Triplet<double>(row[k], col[k], coo_data[k]);
    }

    A.setFromTriplets(triplets.begin(), triplets.end());

    A.makeCompressed();

    VectorXd b_temp(n);

    for (int k = 0; k < n; k++){
        b_temp[k] = b[k];
    }

    typedef IncompleteCholesky<double> Preconditioner;
    ConjugateGradient<MySparseMatrix, Lower, Preconditioner> solver;

    solver.preconditioner().setInitialShift(initial_shift);

    solver.setTolerance(rtol);

    solver.compute(A);

    VectorXd x_temp = solver.solve(b_temp);

    for (int k = 0; k < n; k++){
        x[k] = x_temp[k];
    }

    return 0;
}

int solve_eigen_cholesky_coo(
    const double *coo_data,
    const int *row,
    const int *col,
    int nnz,
    const double *b,
    double *x,
    int n,
    double rtol,
    double initial_shift
){
    MySparseMatrix A(n, n);

    std::vector<Triplet<double> > triplets(nnz);

    for (int k = 0; k < nnz; k++){
        triplets[k] = Triplet<double>(row[k], col[k], coo_data[k]);
    }

    A.setFromTriplets(triplets.begin(), triplets.end());

    A.makeCompressed();

    VectorXd b_temp(n);

    for (int k = 0; k < n; k++){
        b_temp[k] = b[k];
    }

    SimplicialLDLT<MySparseMatrix> solver;

    solver.analyzePattern(A);

    solver.factorize(A);

    VectorXd x_temp = solver.solve(b_temp);

    for (int k = 0; k < n; k++){
        x[k] = x_temp[k];
    }

    return 0;
}

}


#include <iostream>
#include <cassert>
#include <vector>
#include <CGAL/Exact_predicates_inexact_constructions_kernel.h>
#include <CGAL/Delaunay_triangulation_2.h>
#include <CGAL/Triangulation_vertex_base_with_info_2.h>
#include <CGAL/Triangulation_face_base_2.h>
#include <CGAL/Triangulation_data_structure_2.h>
#include <boost/pending/disjoint_sets.hpp>

typedef CGAL::Exact_predicates_inexact_constructions_kernel K;
typedef CGAL::Triangulation_vertex_base_with_info_2<int, K> Vb;
typedef CGAL::Triangulation_face_base_2<K> Fb;
typedef CGAL::Triangulation_data_structure_2<Vb, Fb> Tds;
typedef CGAL::Delaunay_triangulation_2<K, Tds> Triangulation;

const int debug_level = 0;

#define DEBUG(min_level, x)      \
  if (debug_level >= min_level)  \
  {                              \
    std::cerr << x << std::endl; \
  }

struct TreeEdge
{
  int i1, i2;
  double sq_dist;
};

struct GraphProblem
{
  std::vector<int> nearest_tree_by_bone;
  std::vector<double> bone_tree_sq_dists;
  std::vector<TreeEdge> tree_edges;
};

typedef std::vector<std::pair<K::Point_2, int>> LocationVec;

int binary_search_first(std::function<bool(int)> is_target)
{
  int upper_power = 0;
  while (true)
  {
    int i = 1 << upper_power;
    if (is_target(i))
    {
      break;
    }
    upper_power++;
  }

  int low = upper_power == 0 ? 0 : (1 << (upper_power - 1));
  int high = 1 << upper_power;
  while (low < high)
  {
    int mid = (low + high) / 2;
    if (is_target(mid))
    {
      high = mid;
    }
    else
    {
      low = mid + 1;
    }
  }

  assert(!is_target(high - 1) && is_target(high));
  return high;
}

LocationVec read_locations(int n)
{
  LocationVec locations;
  for (int i = 0; i < n; i++)
  {
    long x, y;
    std::cin >> x >> y;
    assert(abs(x) < (1 << 24) && abs(y) < (1 << 24));
    locations.push_back(std::make_pair(K::Point_2(x, y), i));
  }
  return locations;
}

GraphProblem graph_problem_from_locations(const LocationVec &tree_locations, const LocationVec &bone_locations)
{
  Triangulation triangulation;
  triangulation.insert(tree_locations.begin(), tree_locations.end());

  GraphProblem g;

  for (const auto &bone_location : bone_locations)
  {
    Triangulation::Vertex_handle vertex = triangulation.nearest_vertex(bone_location.first);
    g.nearest_tree_by_bone.push_back(vertex->info());
    g.bone_tree_sq_dists.push_back(CGAL::squared_distance(vertex->point(), bone_location.first));
  }

  for (auto it = triangulation.finite_edges_begin(); it != triangulation.finite_edges_end(); it++)
  {
    TreeEdge e;
    e.i1 = it->first->vertex((it->second + 1) % 3)->info();
    e.i2 = it->first->vertex((it->second + 2) % 3)->info();
    if (e.i1 > e.i2)
    {
      std::swap(e.i1, e.i2);
    }
    e.sq_dist = triangulation.segment(it).squared_length();
    g.tree_edges.push_back(e);
  }

  DEBUG(2, "g.tree_edges.size() " << g.tree_edges.size());

  return g;
}

int count_reachable_bones(int n, double s, const GraphProblem &graph_problem)
{
  const int m = graph_problem.nearest_tree_by_bone.size();
  assert(int(graph_problem.bone_tree_sq_dists.size()) == m);

  std::vector<int> ds_rank(n);
  std::vector<int> ds_parent(n);
  boost::disjoint_sets<int *, int *> ds(ds_rank.data(), ds_parent.data());
  for (int i = 0; i < n; i++)
  {
    ds.make_set(i);
  }

  for (const auto &edge : graph_problem.tree_edges)
  {
    DEBUG(2, "tree edge " << edge.i1 << " " << edge.i2 << " " << edge.sq_dist);
    if (edge.sq_dist <= s)
    {
      DEBUG(2, "union");
      ds.union_set(edge.i1, edge.i2);
    }
  }

  std::vector<int> bones_per_component(n, 0);
  for (int i = 0; i < m; i++)
  {
    if (4 * graph_problem.bone_tree_sq_dists.at(i) <= s)
    {
      bones_per_component.at(ds.find_set(graph_problem.nearest_tree_by_bone.at(i)))++;
    }
  }

  return *std::max_element(bones_per_component.begin(), bones_per_component.end());
}

void testcase()
{
  int n, m, k;
  double s;
  std::cin >> n >> m >> s >> k;
  assert(n >= 1 && n <= 4e4);
  assert(m >= 1 && m <= 4e4);
  assert(s >= 1 && s <= (1L << 51));
  assert(k >= 1 && k <= m);

  auto tree_locations = read_locations(n);
  auto bone_locations = read_locations(m);

  GraphProblem graph_problem = graph_problem_from_locations(tree_locations, bone_locations);

  int a = count_reachable_bones(n, s, graph_problem);

  std::vector<double> interesting_sq_dists;
  for (const auto &e : graph_problem.tree_edges)
  {
    interesting_sq_dists.push_back(e.sq_dist);
  }
  for (const auto &sq_dist : graph_problem.bone_tree_sq_dists)
  {
    interesting_sq_dists.push_back(4 * sq_dist);
  }
  std::sort(interesting_sq_dists.begin(), interesting_sq_dists.end());

  int critical_i = binary_search_first([&interesting_sq_dists, &graph_problem, n, k](int i) {
    assert(i >= 0);
    if (i >= int(interesting_sq_dists.size()))
    {
      return true;
    }
    return count_reachable_bones(n, interesting_sq_dists.at(i), graph_problem) >= k;
  });
  double q = interesting_sq_dists.at(critical_i);

  std::cout << a << " " << q << "\n";
}

int main()
{
  std::ios_base::sync_with_stdio(false);
  std::cout << std::fixed << std::setprecision(0);

  int t;
  std::cin >> t;
  for (int i = 0; i < t; i++)
  {
    testcase();
  }

  return 0;
}