#include <Eigen/Dense>
#include <Eigen/Core>
#include <Eigen/Geometry>
#include <cmath>
#include "UKF/Types.h"
#include "UKF/Integrator.h"
#include "UKF/StateVector.h"
#include "UKF/MeasurementVector.h"
#include "UKF/Core.h"
#include "filter.h"

static double fov = 90;
static int width = 672;
static int height = 672;

enum StateFields {
    CameraPosition,
    CameraQuat,
    CameraVelocity,
    CameraAngVelocity,
    CameraAcceleration,
    CameraAngAcceleration,
    ActorPosition,
    ActorYaw,
    ActorVelocity,
    ActorYawVelocity,
    ActorAcceleration,
    ActorYawAcceleration
};

using TrajectoryStateVector = UKF::StateVector<
    UKF::Field<CameraPosition, UKF::Vector<3>>,
    UKF::Field<CameraQuat, UKF::Vector<4>>,
    UKF::Field<CameraVelocity, UKF::Vector<3>>,
    UKF::Field<CameraAngVelocity, UKF::Vector<3>>,
    UKF::Field<CameraAcceleration, UKF::Vector<3>>,
    UKF::Field<CameraAngAcceleration, UKF::Vector<3>>,
    UKF::Field<ActorPosition, UKF::Vector<3>>,
    UKF::Field<ActorYaw, real_t>,
    UKF::Field<ActorVelocity, UKF::Vector<3>>,
    UKF::Field<ActorYawVelocity, real_t>,
    UKF::Field<ActorAcceleration, UKF::Vector<3>>,
    UKF::Field<ActorYawAcceleration, real_t>
>;

namespace UKF {
    template <> template<>
    TrajectoryStateVector TrajectoryStateVector::derivative<>() const {
        TrajectoryStateVector temp;

        /* Position derivative */
        temp.set_field<CameraPosition>(get_field<CameraVelocity>());
        UKF::Vector<3> ang_vel = get_field<CameraAngVelocity>();
        float cr = std::cos(ang_vel[0] * 0.5);
        float sr = std::sin(ang_vel[0] * 0.5);
        float cp = std::cos(ang_vel[1] * 0.5);
        float sp = std::sin(ang_vel[1] * 0.5);
        float cy = std::cos(ang_vel[2] * 0.5);
        float sy = std::sin(ang_vel[2] * 0.5);
        temp.set_field<CameraQuat>(
            UKF::Vector<4>(cr * cp * cy + sr * sp * sy, //w
                            sr * cp * cy - cr * sp * sy, //x
                            cr * sp * cy + sr * cp * sy, //y
                            cr * cp * sy - sr * sp * cy  //z
                            ));

        /* Velocity derivative */
        temp.set_field<CameraVelocity>(get_field<CameraAcceleration>());
        temp.set_field<CameraAngVelocity>(get_field<CameraAngAcceleration>());

        /* Acceleration derivative */
        temp.set_field<CameraAcceleration>(Vector<3>(0,0,0));
        temp.set_field<CameraAngAcceleration>(0);

        /* Position derivative */
        temp.set_field<ActorPosition>(get_field<ActorVelocity>());
        temp.set_field<ActorYaw>(get_field<ActorYawVelocity>());

        /* Velocity derivative */
        temp.set_field<ActorVelocity>(get_field<ActorAcceleration>());
        temp.set_field<ActorYawVelocity>(get_field<ActorYawAcceleration>());

        /* Acceleration derivative */
        temp.set_field<ActorAcceleration>(Vector<3>(0,0,0));
        temp.set_field<ActorYawAcceleration>(0);

        return temp;
    }
}

enum MeasurementFields {
    BoundingBox,
    Depth,
    DronePosition,
    DroneYaw,
    DronePitch,
    HDE
};

using MeasurementVector = UKF::DynamicMeasurementVector<
    UKF::Field<BoundingBox, UKF::Vector<2>>,
    UKF::Field<Depth, real_t>,
    UKF::Field<DronePosition, UKF::Vector<3>>,
    UKF::Field<DroneYaw, real_t>,
    UKF::Field<DronePitch, real_t>,
    UKF::Field<HDE, real_t>
>;

using MotionForecastingCore = UKF::Core<
    TrajectoryStateVector,
    MeasurementVector,
    UKF::IntegratorRK4
>;

UKF::Vector<4> flatten(UKF::Vector<4> quat) {
    double length = sqrt(quat[0] * quat[0] + quat[3] * quat[3]);
    return UKF::Vector<4>(quat[0] / length, 0, 0, quat[3] / length);
}

float getYaw(UKF::Vector<4> q) {
    q = flatten(q);
    double siny_cosp = 2 * (q[0] * q[3] + q[1] * q[2]);
    double cosy_cosp = 1 - 2 * (q[2] * q[2] + q[3] * q[3]);
    return std::atan2(siny_cosp, cosy_cosp);
}

float getPitch(UKF::Vector<4> q) {
    UKF::Vector<4> f = flatten(q);

    float dot_product = f[0] * q[0] + f[1] * q[1] + f[2] * q[2] + f[3] * q[3];
    float angle = acos(2*dot_product*dot_product - 1);
    return angle;
}

namespace UKF {
    template <> template <>
    UKF::Vector<2> MeasurementVector::expected_measurement
    <TrajectoryStateVector, BoundingBox>(
            const TrajectoryStateVector& state) {
        double fx = width/fov;
        double fy = height/fov;
        int cu = width/2;
        int cv = height/2;

        UKF::Vector<3> diff = state.get_field<ActorPosition>() - state.get_field<CameraPosition>();

        float camera_yaw = getYaw(state.get_field<CameraQuat>());
        float camera_pitch = getPitch(state.get_field<CameraQuat>());

        float actor_yaw = std::atan2(diff[0], diff[1]);
        float actor_pitch = std::atan2(diff[2], std::sqrt(diff[1]*diff[1] + diff[0]*diff[0]));

        int u = fx*(actor_yaw - camera_yaw) + cu;
        int v = fy*(actor_pitch - camera_pitch) + cv;

        return UKF::Vector<2>(u, v);
    }
    template <> template <>
    real_t MeasurementVector::expected_measurement
    <TrajectoryStateVector, Depth>(
            const TrajectoryStateVector& state) {
        // Calculate vector r_q/c between drone and actor, and get the magnitude of that
        UKF::Vector<3> diff = state.get_field<CameraPosition>() - state.get_field<ActorPosition>();
        return diff.norm();
    }
    template <> template <>
    UKF::Vector<3> MeasurementVector::expected_measurement
    <TrajectoryStateVector, DronePosition>(
            const TrajectoryStateVector& state) {
        return state.get_field<CameraPosition>();
    }
    template <> template <>
    real_t MeasurementVector::expected_measurement
    <TrajectoryStateVector, DroneYaw>(
            const TrajectoryStateVector& state) {
        UKF::Vector<4> q = state.get_field<CameraQuat>();
        return getYaw(q);
    }
    template <> template <>
    real_t MeasurementVector::expected_measurement
    <TrajectoryStateVector, DronePitch>(
            const TrajectoryStateVector& state) {
        UKF::Vector<4> q = state.get_field<CameraQuat>();
        return getPitch(q);
    }
    template <> template <>
    real_t MeasurementVector::expected_measurement
    <TrajectoryStateVector, HDE>(
            const TrajectoryStateVector& state) {
        float yaw_actor = state.get_field<ActorYaw>();
        float yaw_camera = getYaw(state.get_field<CameraQuat>());
        return yaw_actor - yaw_camera;
    }
}

static MotionForecastingCore filter;
static MeasurementVector meas;

void ukf_init(float x, float y, float z, float yaw) {
    filter.state.set_field<CameraPosition>(UKF::Vector<3>(0,0,0));
    filter.state.set_field<CameraQuat>(UKF::Vector<4>(cos(yaw/2), 0, 0, sin(yaw/2)));
    filter.state.set_field<CameraVelocity>(UKF::Vector<3>(0,0,0));
    filter.state.set_field<CameraAngVelocity>(UKF::Vector<3>(0,0,0));
    filter.state.set_field<CameraAcceleration>(UKF::Vector<3>(0,0,0));
    filter.state.set_field<CameraAngAcceleration>(UKF::Vector<3>(0,0,0));
    filter.state.set_field<ActorPosition>(UKF::Vector<3>(0,0,0));
    filter.state.set_field<ActorYaw>(0);
    filter.state.set_field<ActorVelocity>(UKF::Vector<3>(0,0,0));
    filter.state.set_field<ActorYawVelocity>(0);
    filter.state.set_field<ActorAcceleration>(UKF::Vector<3>(0,0,0));
    filter.state.set_field<ActorYawAcceleration>(0);
    filter.covariance = TrajectoryStateVector::CovarianceMatrix::Identity();

    filter.process_noise_covariance = TrajectoryStateVector::CovarianceMatrix::Identity() * 0.1;
    filter.measurement_covariance << 1e-2, 1e-2, 0.1, 0.2, 0.2, 0.2, 0.1, 0.1, 1e-2;
}

cinematography_msgs::msg::MultiDOF get_state(rclcpp::Duration duration) {
    cinematography_msgs::msg::MultiDOF point;
    point.duration = duration.seconds();
    point.x = filter.state.get_field<ActorPosition>()[0];
    point.y = filter.state.get_field<ActorPosition>()[1];
    point.z = filter.state.get_field<ActorPosition>()[2];
    point.yaw = filter.state.get_field<ActorYaw>();
    point.vx = filter.state.get_field<ActorVelocity>()[0];
    point.vy = filter.state.get_field<ActorVelocity>()[1];
    point.vz = filter.state.get_field<ActorVelocity>()[2];
    point.ax = filter.state.get_field<ActorAcceleration>()[0];
    point.ay = filter.state.get_field<ActorAcceleration>()[1];
    point.az = filter.state.get_field<ActorAcceleration>()[2];

    return point;
}

cinematography_msgs::msg::MultiDOF get_state(rclcpp::Duration duration, MotionForecastingCore filter) {
    cinematography_msgs::msg::MultiDOF point;
    point.duration = duration.seconds();
    point.x = filter.state.get_field<ActorPosition>()[0];
    point.y = filter.state.get_field<ActorPosition>()[1];
    point.z = filter.state.get_field<ActorPosition>()[2];
    point.yaw = filter.state.get_field<ActorYaw>();
    point.vx = filter.state.get_field<ActorVelocity>()[0];
    point.vy = filter.state.get_field<ActorVelocity>()[1];
    point.vz = filter.state.get_field<ActorVelocity>()[2];
    point.ax = filter.state.get_field<ActorAcceleration>()[0];
    point.ay = filter.state.get_field<ActorAcceleration>()[1];
    point.az = filter.state.get_field<ActorAcceleration>()[2];

    return point;
}

// forecast length must be at least 2 (now and next point)
std::vector<cinematography_msgs::msg::MultiDOF> ukf_iterate(rclcpp::Duration point_duration, int forecast_length) {
    if (forecast_length < 2) {
        return std::vector<cinematography_msgs::msg::MultiDOF>();
    }

    std::vector<cinematography_msgs::msg::MultiDOF> path;
    path.reserve(forecast_length);

    path[0] = get_state(point_duration);

    filter.step(point_duration.seconds(), meas);
    path[1] = get_state(point_duration);

    MotionForecastingCore forecaster = filter;
    for(int i = 2; i < forecast_length; i++) {
        forecaster.step(point_duration.seconds(), MeasurementVector());
        path[i] = get_state(point_duration, forecaster);
    }

    return path;
}

void ukf_meas_clear() {
    meas = MeasurementVector();
}

void ukf_set_fov(float f) {
    fov = f;
}

void ukf_set_bb(float width, float height) {
    meas.set_field<BoundingBox>(UKF::Vector<2>(width, height));
}

void ukf_set_depth(float depth) {
    meas.set_field<Depth>(depth);
}

void ukf_set_position(float x, float y, float z) {
    meas.set_field<DronePosition>(UKF::Vector<3>(x, y, z));
}

void ukf_set_yaw(float yaw) {
    meas.set_field<DroneYaw>(yaw);
}

void ukf_set_pitch(float pitch) {
    meas.set_field<DronePitch>(pitch);
}

void ukf_set_hde(float hde) {
    meas.set_field<HDE>(hde);
}


///////////////////////////////////////////////////////////////
//  Copyright 2012 John Maddock. Distributed under the Boost
//  Software License, Version 1.0. (See accompanying file
//  LICENSE_1_0.txt or copy at https://www.boost.org/LICENSE_1_0.txt

//
// Compare arithmetic results using fixed_int to GMP results.
//

#ifdef _MSC_VER
#define _SCL_SECURE_NO_WARNINGS
#endif

#include <nil/crypto3/multiprecision/gmp.hpp>
#include <nil/crypto3/multiprecision/fixed_int.hpp>
#include <boost/random/mersenne_twister.hpp>
#include <boost/random/uniform_int.hpp>
#include "test.hpp"

template<class T>
T generate_random(unsigned bits_wanted) {
    static boost::random::mt19937 gen;
    typedef boost::random::mt19937::result_type random_type;

    T max_val;
    unsigned digits;
    if (std::numeric_limits<T>::is_bounded && (bits_wanted == std::numeric_limits<T>::digits)) {
        max_val = (std::numeric_limits<T>::max)();
        digits = std::numeric_limits<T>::digits;
    } else {
        max_val = T(1) << bits_wanted;
        digits = bits_wanted;
    }

    unsigned bits_per_r_val = std::numeric_limits<random_type>::digits - 1;
    while ((random_type(1) << bits_per_r_val) > (gen.max)())
        --bits_per_r_val;

    unsigned terms_needed = digits / bits_per_r_val + 1;

    T val = 0;
    for (unsigned i = 0; i < terms_needed; ++i) {
        val *= (gen.max)();
        val += gen();
    }
    val %= max_val;
    return val;
}

int main() {
    using namespace nil::crypto3::multiprecision;
    typedef number<fixed_int<1024, true>> packed_type;
    unsigned last_error_count = 0;
    for (int i = 0; i < 1000; ++i) {
        mpz_int a = generate_random<mpz_int>(1000);
        mpz_int b = generate_random<mpz_int>(512);
        mpz_int c = generate_random<mpz_int>(256);
        mpz_int d = generate_random<mpz_int>(32);

        int si = d.convert_to<int>();

        packed_type a1 = a.str();
        packed_type b1 = b.str();
        packed_type c1 = c.str();
        packed_type d1 = d.str();

        BOOST_CHECK_EQUAL(a.str(), a1.str());
        BOOST_CHECK_EQUAL(b.str(), b1.str());
        BOOST_CHECK_EQUAL(c.str(), c1.str());
        BOOST_CHECK_EQUAL(d.str(), d1.str());
        BOOST_CHECK_EQUAL(mpz_int(a + b).str(), packed_type(a1 + b1).str());
        BOOST_CHECK_EQUAL(mpz_int(a - b).str(), packed_type(a1 - b1).str());
        BOOST_CHECK_EQUAL(mpz_int(mpz_int(-a) + b).str(), packed_type(packed_type(-a1) + b1).str());
        BOOST_CHECK_EQUAL(mpz_int(mpz_int(-a) - b).str(), packed_type(packed_type(-a1) - b1).str());
        BOOST_CHECK_EQUAL(mpz_int(c * d).str(), packed_type(c1 * d1).str());
        BOOST_CHECK_EQUAL(mpz_int(c * -d).str(), packed_type(c1 * -d1).str());
        BOOST_CHECK_EQUAL(mpz_int(-c * d).str(), packed_type(-c1 * d1).str());
        BOOST_CHECK_EQUAL(mpz_int(b * c).str(), packed_type(b1 * c1).str());
        BOOST_CHECK_EQUAL(mpz_int(a / b).str(), packed_type(a1 / b1).str());
        BOOST_CHECK_EQUAL(mpz_int(a / -b).str(), packed_type(a1 / -b1).str());
        BOOST_CHECK_EQUAL(mpz_int(-a / b).str(), packed_type(-a1 / b1).str());
        BOOST_CHECK_EQUAL(mpz_int(a / d).str(), packed_type(a1 / d1).str());
        BOOST_CHECK_EQUAL(mpz_int(a % b).str(), packed_type(a1 % b1).str());
        BOOST_CHECK_EQUAL(mpz_int(a % -b).str(), packed_type(a1 % -b1).str());
        BOOST_CHECK_EQUAL(mpz_int(-a % b).str(), packed_type(-a1 % b1).str());
        BOOST_CHECK_EQUAL(mpz_int(a % d).str(), packed_type(a1 % d1).str());
        // bitwise ops:
        BOOST_CHECK_EQUAL(mpz_int(a | b).str(), packed_type(a1 | b1).str());
        BOOST_CHECK_EQUAL(mpz_int(a & b).str(), packed_type(a1 & b1).str());
        BOOST_CHECK_EQUAL(mpz_int(a ^ b).str(), packed_type(a1 ^ b1).str());
        // Now check operations involving integers:
        BOOST_CHECK_EQUAL(mpz_int(a + si).str(), packed_type(a1 + si).str());
        BOOST_CHECK_EQUAL(mpz_int(a + -si).str(), packed_type(a1 + -si).str());
        BOOST_CHECK_EQUAL(mpz_int(-a + si).str(), packed_type(-a1 + si).str());
        BOOST_CHECK_EQUAL(mpz_int(si + a).str(), packed_type(si + a1).str());
        BOOST_CHECK_EQUAL(mpz_int(a - si).str(), packed_type(a1 - si).str());
        BOOST_CHECK_EQUAL(mpz_int(a - -si).str(), packed_type(a1 - -si).str());
        BOOST_CHECK_EQUAL(mpz_int(-a - si).str(), packed_type(-a1 - si).str());
        BOOST_CHECK_EQUAL(mpz_int(si - a).str(), packed_type(si - a1).str());
        BOOST_CHECK_EQUAL(mpz_int(b * si).str(), packed_type(b1 * si).str());
        BOOST_CHECK_EQUAL(mpz_int(b * -si).str(), packed_type(b1 * -si).str());
        BOOST_CHECK_EQUAL(mpz_int(-b * si).str(), packed_type(-b1 * si).str());
        BOOST_CHECK_EQUAL(mpz_int(si * b).str(), packed_type(si * b1).str());
        BOOST_CHECK_EQUAL(mpz_int(a / si).str(), packed_type(a1 / si).str());
        BOOST_CHECK_EQUAL(mpz_int(a / -si).str(), packed_type(a1 / -si).str());
        BOOST_CHECK_EQUAL(mpz_int(-a / si).str(), packed_type(-a1 / si).str());
        BOOST_CHECK_EQUAL(mpz_int(a % si).str(), packed_type(a1 % si).str());
        BOOST_CHECK_EQUAL(mpz_int(a % -si).str(), packed_type(a1 % -si).str());
        BOOST_CHECK_EQUAL(mpz_int(-a % si).str(), packed_type(-a1 % si).str());
        BOOST_CHECK_EQUAL(mpz_int(a | si).str(), packed_type(a1 | si).str());
        BOOST_CHECK_EQUAL(mpz_int(a & si).str(), packed_type(a1 & si).str());
        BOOST_CHECK_EQUAL(mpz_int(a ^ si).str(), packed_type(a1 ^ si).str());
        BOOST_CHECK_EQUAL(mpz_int(si | a).str(), packed_type(si | a1).str());
        BOOST_CHECK_EQUAL(mpz_int(si & a).str(), packed_type(si & a1).str());
        BOOST_CHECK_EQUAL(mpz_int(si ^ a).str(), packed_type(si ^ a1).str());
        BOOST_CHECK_EQUAL(mpz_int(gcd(a, b)).str(), packed_type(gcd(a1, b1)).str());
        BOOST_CHECK_EQUAL(mpz_int(lcm(c, d)).str(), packed_type(lcm(c1, d1)).str());

        if (last_error_count != boost::detail::test_errors()) {
            last_error_count = boost::detail::test_errors();
            std::cout << std::hex << std::showbase;

            std::cout << "a    = " << a << std::endl;
            std::cout << "a1   = " << a1 << std::endl;
            std::cout << "b    = " << b << std::endl;
            std::cout << "b1   = " << b1 << std::endl;
            std::cout << "c    = " << c << std::endl;
            std::cout << "c1   = " << c1 << std::endl;
            std::cout << "d    = " << d << std::endl;
            std::cout << "d1   = " << d1 << std::endl;
            std::cout << "a + b   = " << a + b << std::endl;
            std::cout << "a1 + b1 = " << a1 + b1 << std::endl;
            std::cout << std::dec;
            std::cout << "a - b   = " << a - b << std::endl;
            std::cout << "a1 - b1 = " << a1 - b1 << std::endl;
            std::cout << "-a + b   = " << mpz_int(-a) + b << std::endl;
            std::cout << "-a1 + b1 = " << packed_type(-a1) + b1 << std::endl;
            std::cout << "-a - b   = " << mpz_int(-a) - b << std::endl;
            std::cout << "-a1 - b1 = " << packed_type(-a1) - b1 << std::endl;
            std::cout << "c*d    = " << c * d << std::endl;
            std::cout << "c1*d1  = " << c1 * d1 << std::endl;
            std::cout << "b*c    = " << b * c << std::endl;
            std::cout << "b1*c1  = " << b1 * c1 << std::endl;
            std::cout << "a/b    = " << a / b << std::endl;
            std::cout << "a1/b1  = " << a1 / b1 << std::endl;
            std::cout << "a/d    = " << a / d << std::endl;
            std::cout << "a1/d1  = " << a1 / d1 << std::endl;
            std::cout << "a%b    = " << a % b << std::endl;
            std::cout << "a1%b1  = " << a1 % b1 << std::endl;
            std::cout << "a%d    = " << a % d << std::endl;
            std::cout << "a1%d1  = " << a1 % d1 << std::endl;
        }
    }
    return boost::report_errors();
}


/* areagrid - a geometric data structure which splits an area into rectangles
*/
#include <iostream>
#include <tuple>
#include <array>
#include <algorithm>
#include <climits>
#include <random>
#include <string>
#include <sstream>
#include <unordered_map>
#include <unordered_set>
#include <tuple>
#include <algorithm>
#include <vector>
#include <boost/format.hpp>

struct EquiDistRange {
    EquiDistRange(int min, int max, int divisions)
    : min_{std::min(min, max)}, max_{std::max(min, max)}
    , dist_{max_ - min_}
    , divisions_{divisions}
    , d_{static_cast<double>((dist_)/divisions_)}
    {}
    int operator[](int i) const {
        return static_cast<int>(min_ + d_ * i);
    }
    struct Iterator {
        const EquiDistRange& r_;
        int index_ { 0 };
        bool operator==(const Iterator& b) const { return index_ == b.index_; }
        bool operator!=(const Iterator& b) const { return !(*this == b); }
        bool operator<(const Iterator& b) const { return index_ < b.index_; }
        bool operator<=(const Iterator& b) const { return index_ <= b.index_; }
        bool operator>(const Iterator& b) const { return index_ > b.index_; }
        bool operator>=(const Iterator& b) const { return index_ >= b.index_; }
        Iterator& operator++() { ++index_; return *this; }
        Iterator operator++(int) { auto tmp = *this; ++index_; return tmp; }
        Iterator& operator--() { --index_; return *this; }
        Iterator operator--(int) { auto tmp = *this; --index_; return tmp; }
        int operator*() { return r_[index_]; }
    };
    Iterator begin() const { return Iterator{*this}; }
    Iterator end() const { return Iterator{*this, divisions_ + 1}; }
    const int min_, max_, dist_, divisions_;
    double d_;
};

struct Point {
  int x, y;
  bool is_top_left(const Point& other) const noexcept {
    return x < other.x && y < other.y;
  }
  bool is_top_left_eq(const Point& other) const noexcept {
    return x <= other.x && y <= other.y;
  }
  bool operator==(const Point& p) const noexcept {
    return x == p.x && y == p.y;
  }
};

std::ostream& operator<<(std::ostream& o, const Point& p) {
  return o << '(' << p.x << ", " << p.y << ')';
}

struct Rect {
  Point min, max;
  void sanitize() {
    if (min.x > max.x)
      std::swap(min.x, max.x);
    if (min.y > max.y)
      std::swap(min.y, max.y);
  }
  bool operator==(const Rect& o) const noexcept { return min == o.min && max == o.max; }
  bool contains(const Point& p) const noexcept {
    return min.is_top_left_eq(p) && p.is_top_left_eq(max);
  }
  int width() const noexcept { return max.x - min.x + 1; }
  int height() const noexcept { return max.y - min.y + 1; }
  int area() const noexcept { return width() * height(); }
  bool intersects(const Rect& r, Rect* intersection = nullptr) const noexcept {
    Rect tmp{ std::max(min.x, r.min.x)
            , std::max(min.y, r.min.y)
            , std::min(max.x, r.max.x)
            , std::min(max.y, r.max.y)
            };
    auto tl = tmp.min.is_top_left_eq(tmp.max);
    if (tl && intersection != nullptr) {
      *intersection = tmp;
    }
    return tl;
  }
  Point center() const noexcept { return Point{ min.x + width()/2, min.y + height()/2};}
  template<typename T>
  static
  Rect bounding_rect(T begin, const T end) {
    Rect tmp = { INT_MAX, INT_MAX, INT_MIN, INT_MIN};
    for(;begin != end; ++begin) {
      const Rect* ptr = static_cast<const Rect*>(*begin);
      tmp.min.x = std::min(ptr->min.x, tmp.min.x);
      tmp.min.y = std::min(ptr->min.y, tmp.min.y);
      tmp.max.x = std::max(ptr->max.x, tmp.max.x);
      tmp.max.y = std::max(ptr->max.y, tmp.max.y);
    }
    return tmp;
  }
  template<typename T>
  static
  bool intersect(T begin, const T end, Rect* intersection = nullptr) {
    Rect tmp = { INT_MIN, INT_MIN, INT_MAX, INT_MAX};
    for(;begin != end; ++begin) {
      const Rect* ptr = static_cast<const Rect*>(*begin);
      if (!tmp.intersects(*ptr, &tmp))
        return false;
    }
    if (intersection != nullptr) {
      *intersection = tmp;
    }
    return true;
  }
};

class Randomizer {
public:
  Randomizer(unsigned area_width  = 1'000'000
           , unsigned area_height = 1'000'000
           , unsigned max_width   =  200'000
           , unsigned max_height  =  200'000)
  : distx(0, area_width - 1)
  , disty(0, area_height - 1)
  , max_width_{max_width}
  , max_height_{max_height}
  {
    if (area_width < max_width || area_height < max_height)
        throw "Randomizer: Illegal dimensions";
  }
  void operator()(Point& p) {
    p.x = distx(rd);
    p.y = disty(rd);
  }
  void operator()(Rect& r) {
    operator()(r.min);
    std::uniform_int_distribution<unsigned> distw(0, std::min(max_width_, distx.max() - r.min.x));
    std::uniform_int_distribution<unsigned> disth(0, std::min(max_height_, disty.max() - r.min.y));
    r.max.x = r.min.x + distw(rd);
    r.max.y = r.min.y + disth(rd);
    r.sanitize();
  }
private:
  std::random_device rd;
  std::uniform_int_distribution<unsigned> distx;
  std::uniform_int_distribution<unsigned> disty;
  unsigned max_width_, max_height_;
};

struct RectScale {
    const Rect bounds_;
    const unsigned width_, height_;
    const double fx_, fy_;
    RectScale(const Rect& bounds, unsigned width = 100, unsigned height = 100)
    : bounds_{bounds}
    , width_{width}
    , height_{height}
    , fx_{width_/static_cast<double>(bounds_.width())}
    , fy_{height_/static_cast<double>(bounds_.height())}
    {}
    Rect operator()(const Rect& r) const {
        return Rect{
            static_cast<int>(std::floor(fx_ * ( r.min.x - bounds_.min.x)))
          , static_cast<int>(std::floor(fy_ * ( r.min.y - bounds_.min.y)))
          , static_cast<int>(std::ceil(fx_ * ( r.max.x - bounds_.min.x)))
          , static_cast<int>(std::ceil(fy_ * ( r.max.y - bounds_.min.y)))
        };
    }
    Point operator()(const Point& p) const {
      return Point{
        static_cast<int>(fx_ * (p.x - bounds_.min.x))
      , static_cast<int>(fy_ * (p.y - bounds_.min.y))
      };
    }
};

std::ostream& operator<<(std::ostream& o, const Rect& r) {
  return o << '[' << r.min << ", " << r.max << ']';
}

bool operator<(const Point& a, const Point& b) {
  return std::tie(a.x, a.y) < std::tie(b.x, b.y);
}

bool operator<(const Rect& a, const Rect& b) {
  return std::tie(a.min, a.max) < std::tie(b.min, b.max);
}

struct VRect {
  Rect r;
  int value;
  operator Rect*() noexcept { return &r; }
  bool operator==(const VRect& o) const noexcept { return r == o.r && value == o.value;}
};
std::ostream& operator<<(std::ostream& o, const VRect& r) {
  return o << '[' << r.value << ", " << r.r.min << ", " << r.r.max << ']';
}
bool operator<(const VRect& a, const VRect& b) {
  return a < b;
}

template<typename T>
void print_all(const T& l) {
  for(const auto& e : l) {
    std::cout << e << std::endl;
  }
}

namespace std
{
  template<> struct hash<Point>
  {
    typedef Point argument_type;
    typedef std::size_t result_type;
    result_type operator()(argument_type const& s) const noexcept
    {
      result_type const h1 ( std::hash<int>{}(s.x) );
      result_type const h2 ( std::hash<int>{}(s.y) );
      return h1 ^ (h2 << 1);
    }
  };
  template<> struct hash<Rect>
  {
    typedef Rect argument_type;
    typedef std::size_t result_type;
    result_type operator()(argument_type const& s) const noexcept
    {
      result_type const h1 ( std::hash<Point>{}(s.min) );
      result_type const h2 ( std::hash<Point>{}(s.max) );
      return h1 ^ (h2 << 1);
    }
  };
  template<> struct hash<VRect>
  {
    typedef VRect argument_type;
    typedef std::size_t result_type;
    result_type operator()(argument_type const& s) const noexcept
    {
      result_type const h1 ( std::hash<Rect>{}(s.r) );
      result_type const h2 ( std::hash<int>{}(s.value) );
      return h1 ^ (h2 << 1);
    }
  };
}

class VRectGrid {
public:
  using it_t  = std::vector<VRect>::size_type;
  using set_t = std::unordered_set<it_t>;
  using map_t = std::unordered_map<Point, set_t>;
  template<typename T>
  VRectGrid(const RectScale& scale, T begin, const T end);
  // const set_t& operator()(int x, int y);
  // const set_t& operator()(const Point&);
  size_t rects_containing(const Point&, std::vector<VRect>& result);
  const RectScale scale_;
private:
  map_t map_;
  const std::vector<VRect> vec_;
};

template<typename T>
VRectGrid::VRectGrid(const RectScale& scale, T begin, const T end)
: scale_{scale}, vec_(begin, end)
{
  for(int y = 0; y <= scale_.height_ + 1; ++y) {
    for(int x = 0; x <= scale_.width_ + 1; ++x) {
      Point p{x, y};
      set_t set;
      for(size_t i = 0; i < vec_.size(); ++i) {
        if (scale_(vec_[i].r).contains(p)) {
          set.insert(i);
        }
      }
      map_.insert({p, set});
    }
  }
}
// const VRectGrid::set_t& VRectGrid::operator()(const Point& p) {
//   return map_[scale_(p)];
// }
// const VRectGrid::set_t& VRectGrid::operator()(int x, int y) {
//   return this->operator()(Point{x, y});
// }
size_t VRectGrid::rects_containing(const Point& p, std::vector<VRect>& result) {
  size_t ret = 0;
  decltype(auto) res = map_[scale_(p)];
  std::cout << "Checking " << res.size() << " VRects" << std::endl;
  for(const auto& e : res) {
    if (vec_[e].r.contains(p)) {
      ++ret;
      result.push_back(vec_[e]);
    }
  }
  return ret;
}

// std::ostream& operator<<(std::ostream& o, VRectGrid& g) {
//   boost::format fmt{"%3d"};
//   for(int y = g.scale_.height_ + 1; y >= 0; --y) {
//     for(int x = 0; x <= g.scale_.width_ + 1; ++x) {
//       auto s = g(x, y).size();
//       std::cout << fmt % s << " ";
//     }
//     std::cout << std::endl;
//   }
//   return o;
// }

int main() {
  using namespace std;
  std::array<VRect, 2000> rs;
  Randomizer rnd;
  int idx = 1;
  for(auto & e : rs) {
    e.value = idx++;
    rnd(e.r);
  }
  // print_all(rs);
  std::sort(rs.begin(), rs.end(), [](const VRect& a, const VRect& b) { return a.r < b.r;});
  std::cout << "***********" << std::endl;
  print_all(rs);
  Rect bounds = Rect::bounding_rect(rs.begin(), rs.end());
  std::cout << "*" << bounds << std::endl;

  const int buckets = 5;
  float rects_per_bucket = rs.size() / (float)buckets;
  float dx = static_cast<float>(bounds.width())/rs.size();
  for(size_t i = 0; i < rs.size(); ++i) {
    int current_x = dx * i;
    int bucket = i / rects_per_bucket;
    cout << i << " " << current_x << " " << bucket
         << " " << rs[i].r.min.x
         << " " << rs[i] << rs[i].r.center()
         << endl;
  }
  auto b = rs.begin();
  int bb = 0;
  for(size_t i = 0; ; ++i) {
    int bucket = i / rects_per_bucket;
    if (bb < bucket) {
      Rect intersection;
      auto intersect = Rect::intersect(b, rs.begin() + i, &intersection);
      cout << "Bucket " << bb << " intersect: "
           << boolalpha << intersect << ": " << intersection
           << static_cast<float>(intersection.area())/bounds.area()/buckets << endl;
      bb = bucket;
      b = rs.begin() + i;
    }
    if (i >= rs.size())
      break;
  }
  // EquiDistRange r{0, 8, 3};
  // for(auto i : r) {
  //   cout << i << ", ";
  // }
  // cout << endl;
  RectScale scale(bounds, 30, 30);
  VRectGrid vrg{scale, rs.begin(), rs.end()};
  Point p;
  std::vector<VRect> res;
  for(int i = 0; i < 20; ++i) {
    res.resize(0);
    rnd(p);
    cout << "**** Loopkup " << p
         << ": " << vrg.rects_containing(p, res) << "VRects:"
         << endl;
    print_all(res);
  }
  cout << "*******" << endl;
  // std::string bla;
  // cin >> bla;
}


/*
 * Copyright (C) 2018 Swift Navigation Inc.
 * Contact: Swift Navigation <dev@swiftnav.com>
 *
 * This source is subject to the license found in the file 'LICENSE' which must
 * be distributed together with this source. All other rights reserved.
 *
 * THIS CODE AND INFORMATION IS PROVIDED "AS IS" WITHOUT WARRANTY OF ANY KIND,
 * EITHER EXPRESSED OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE IMPLIED
 * WARRANTIES OF MERCHANTABILITY AND/OR FITNESS FOR A PARTICULAR PURPOSE.
 */

#include "covariance_functions/covariance_functions.h"
#include <Eigen/Core>
#include <Eigen/Dense>
#include <gtest/gtest.h>
#include <iostream>
#include <vector>

namespace albatross {

std::vector<Eigen::Vector3d> points_on_a_line(const int n) {
  std::vector<Eigen::Vector3d> xs;
  for (int i = 0; i < n; i++) {
    Eigen::Vector3d x;
    for (int j = 0; j < 3; j++)
      x[static_cast<std::size_t>(j)] = 1000 * i + j;
    xs.push_back(x);
  }
  return xs;
}

TEST(test_covariance_functions, test_build_covariance) {
  using Feature = Eigen::Vector3d;
  using Noise = IndependentNoise<Feature>;
  using SqExp = SquaredExponential<EuclideanDistance>;
  using RadialSqExp = SquaredExponential<RadialDistance>;

  CovarianceFunction<SqExp> sqexp = {SqExp()};
  CovarianceFunction<Constant> constant = {Constant()};
  CovarianceFunction<Noise> noise = {Noise()};
  CovarianceFunction<RadialSqExp> radial_sqexp = {RadialSqExp()};

  // Add and multiply covariance functions together and make sure they are
  // still capable of producing a covariance matrix.
  auto product = sqexp * radial_sqexp;
  auto covariance_function = constant + product + noise;

  auto xs = points_on_a_line(5);
  Eigen::MatrixXd C = symmetric_covariance(covariance_function, xs);
  assert(C.rows() == xs.size());
  assert(C.cols() == xs.size());
}

/*
 * In the following we test any covariance functions which should support
 * Eigen::Vector feature vectors.
 */
template <typename T>
class TestVectorCovarianceFunctions : public ::testing::Test {

public:
  typedef CovarianceFunction<T> CovFunc;
  T value_;
};

typedef ::testing::Types<
    SquaredExponential<EuclideanDistance>, SquaredExponential<RadialDistance>,
    Exponential<EuclideanDistance>, Exponential<AngularDistance>,
    Exponential<RadialDistance>>
    VectorCompatibleCovarianceFunctions;

TYPED_TEST_CASE(TestVectorCovarianceFunctions,
                VectorCompatibleCovarianceFunctions);

TYPED_TEST(TestVectorCovarianceFunctions, WorksWithEigen) {
  typename TestFixture::CovFunc covariance_function;
  auto xs = points_on_a_line(5);
  Eigen::MatrixXd C = symmetric_covariance(covariance_function, xs);
  assert(C.rows() == xs.size());
  assert(C.cols() == xs.size());
  // Make sure C is positive definite.
  auto inverse = C.inverse();
}

TYPED_TEST(TestVectorCovarianceFunctions, WorksDirectlyOnCovarianceterms) {
  typename TestFixture::CovFunc covariance_function;
  auto xs = points_on_a_line(5);
  Eigen::MatrixXd C = symmetric_covariance(covariance_function.term, xs);
  assert(C.rows() == xs.size());
  assert(C.cols() == xs.size());
  // Make sure C is positive definite.
  auto inverse = C.inverse();
}
} // namespace albatross


//==============================================================================
//         Copyright 2003 - 2011 LASMEA UMR 6602 CNRS/Univ. Clermont II
//         Copyright 2009 - 2011 LRI    UMR 8623 CNRS/Univ Paris Sud XI
//
//          Distributed under the Boost Software License, Version 1.0.
//                 See accompanying file LICENSE.txt or copy at
//                     http://www.boost.org/LICENSE_1_0.txt
//==============================================================================
#ifndef NT2_TOOLBOX_TRIGONOMETRIC_FUNCTIONS_SIMD_COMMON_SINC_HPP_INCLUDED
#define NT2_TOOLBOX_TRIGONOMETRIC_FUNCTIONS_SIMD_COMMON_SINC_HPP_INCLUDED
#include <nt2/toolbox/trigonometric/functions/sinc.hpp>
#include <nt2/sdk/meta/as_floating.hpp>
#include <nt2/include/functions/simd/sin.hpp>
#include <nt2/include/functions/simd/abs.hpp>
#include <nt2/include/functions/simd/is_inf.hpp>
#include <nt2/include/functions/simd/is_less.hpp>
#include <nt2/include/functions/simd/if_else.hpp>
#include <nt2/include/constants/eps.hpp>
#include <nt2/include/constants/one.hpp>
#include <nt2/include/constants/zero.hpp>
#include <boost/simd/sdk/config.hpp>

namespace nt2 { namespace ext
{
  NT2_FUNCTOR_IMPLEMENTATION( nt2::tag::sinc_, boost::simd::tag::simd_
                            , (A0)(X)
                            , ((simd_<arithmetic_<A0>,X>))
                            )
  {

    typedef typename meta::as_floating<A0>::type result_type;

    NT2_FUNCTOR_CALL(1)
    {
      return nt2::sinc(tofloat(a0));
    }
  };
} }

namespace nt2 { namespace ext
{
  NT2_FUNCTOR_IMPLEMENTATION( nt2::tag::sinc_, boost::simd::tag::simd_
                            , (A0)(X)
                            , ((simd_<floating_<A0>,X>))
                            )
  {

    typedef A0 result_type;

    NT2_FUNCTOR_CALL(1)
    {
      result_type r1 = nt2::if_else(nt2::lt(nt2::abs(a0), nt2::Eps<A0>()),
                                    nt2::One<A0>(),
                                    nt2::sin(a0)/a0);
      #ifdef BOOST_SIMD_NO_INFINITIES
      return r1;
      #else
      return nt2::if_else(nt2::is_inf(a0), nt2::Zero<A0>(), r1);
      #endif
    }
  };
} }

#endif


#ifndef FNN_MODEL_HPP
#define FNN_MODEL_HPP

#include <vector>
#include <list>
#include <Eigen/Dense>
#include "SFML/System.hpp"
#include "gnuplot_i.hpp"

struct Data
{
  Eigen::MatrixXd &training_sample_i;
  Eigen::MatrixXd &training_sample_o;
  Eigen::MatrixXd &eval_input;
  Eigen::MatrixXd &eval_output;
  Data(Eigen::MatrixXd &t_i, Eigen::MatrixXd &t_o,
      Eigen::MatrixXd &e_i, Eigen::MatrixXd &e_o) :
    training_sample_i(t_i),
    training_sample_o(t_o),
    eval_input(e_i),
    eval_output(e_o)
  {}
};
class FNN_Model
{
  public:
    FNN_Model(
        std::vector<unsigned int> layers);
    ~FNN_Model();

    void print_FNN();
    static double normal_distri(double input);
    void Manual_Set_FNN(
        std::vector<Eigen::MatrixXd> &weights,
        std::vector<Eigen::MatrixXd> &bias);
    void Init();
    void ResizeBatch();
    void SetInput(Eigen::MatrixXd &inputs);
    static double sigmoid(double input);
    static double sigmoidPrime(double input);
    void FeedForward();
    void ComputeError(Eigen::MatrixXd &d_outputs);
    void GradientDescent();
    void BackProgagation(
        Eigen::MatrixXd &inputs,
        Eigen::MatrixXd &d_outputs);
    void train(
        Eigen::MatrixXd &training_sample_i,
        Eigen::MatrixXd &training_sample_o,
        unsigned int nbr_epoch,
        unsigned int batch_size,
        double learning_rate,
        Eigen::MatrixXd &eval_input,
        Eigen::MatrixXd &eval_output);
    void core_train();
    double evaluate(
        Eigen::MatrixXd &eval_input,
        Eigen::MatrixXd &eval_output,
        unsigned int epoch);
    double evaluate_MNIST(
        Eigen::MatrixXd &eval_input,
        Eigen::MatrixXd &eval_output,
        unsigned int epoch);

    sf::Thread train_thread;
    sf::Mutex chart_mutex;
    bool thread_state;
    Gnuplot graph;

  private:
    unsigned int _nbr_layer;
    unsigned int _batch_size;
    unsigned int _nbr_epoch;
    double _learning_rate;
    std::vector<unsigned int> _layers;
    std::vector<Eigen::MatrixXd> _weights;
    std::vector<Eigen::MatrixXd> _bias;
    std::vector<Eigen::MatrixXd> _zs;
    std::vector<Eigen::MatrixXd> _activations;
    std::vector<Eigen::MatrixXd> _errors;
    struct Data *_mnist_data;
    std::vector<double> _x;
    std::vector<double> _y;
};

#endif


/* Copyright (C) 2012-2020 IBM Corp.
 * This program is Licensed under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance
 * with the License. You may obtain a copy of the License at
 *   http://www.apache.org/licenses/LICENSE-2.0
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License. See accompanying LICENSE file.
 */
#include <iostream>
#include <fstream>
#include <vector>
#include <cmath>
#include <algorithm>
#include <numeric>
#include <tuple>
#include <NTL/BasicThreadPool.h>

#include <helib/helib.h>

#include <helib/intraSlot.h>
#include <helib/binaryArith.h>

#include "gtest/gtest.h"
#include "test_common.h"

#include <helib/debugging.h>

// define flags FLAG_PRINT_ZZX, FLAG_PRINT_POLY, FLAG_PRINT_VEC, functions
//        decryptAndPrint(ostream, ctxt, sk, ea, flags)
//        decryptAndCompare(ctxt, sk, ea, pa);

namespace {

struct Parameters
{
  Parameters(long prm,
             long bitSize,
             long bitSize2,
             long outSize,
             bool bootstrap,
             long seed,
             long nthreads) :
      prm(prm),
      bitSize(bitSize),
      bitSize2(bitSize2),
      outSize(outSize),
      bootstrap(bootstrap),
      seed(seed),
      nthreads(nthreads){};

  long prm;       // parameter size (0-tiny,...,7-huge)
  long bitSize;   // bitSize of input integers (<=32)
  long bitSize2;  // bitSize of 2nd input integer (<=32)
  long outSize;   // bitSize of output integers, as many as needed
  bool bootstrap; // test multiplication with bootstrapping
  long seed;      // PRG seed
  long nthreads;  // number of threads

  friend std::ostream& operator<<(std::ostream& os, const Parameters& params)
  {
    return os << "{"
              << "prm=" << params.prm << ","
              << "bitSize=" << params.bitSize << ","
              << "bitSize2=" << params.bitSize2 << ","
              << "outSize=" << params.outSize << ","
              << "bootstrap=" << params.bootstrap << ","
              << "seed=" << params.seed << ","
              << "nthreads=" << params.nthreads << "}";
  };
};

class GTestBinaryArith :
    public ::testing::TestWithParam<std::tuple<Parameters, int>>
{
protected:
  static std::vector<helib::zzX> unpackSlotEncoding;
  constexpr static long mValues[8][15] = {
      // clang-format off
   // {p,phi(m),     m,  d,   m1,  m2,  m3,   g1,   g2,   g3,ord1,ord2,ord3,  B, c}
      {2,    48,   105, 12,    3,  35,   0,   71,   76,    0,   2,   2,   0, 25, 2},
      {2,   600,  1023, 10,   11,  93,   0,  838,  584,    0,  10,   6,   0, 25, 2},
      {2,  2304,  4641, 24,    7,   3, 221, 3979, 3095, 3760,   6,   2,  -8, 25, 3},
      {2,  5460,  8193, 26, 8193,   0,   0,   46,    0,    0, 210,   0,   0, 25, 3},
      {2,  8190,  8191, 13, 8191,   0,   0,   39,    0,    0, 630,   0,   0, 25, 3},
      {2, 10752, 11441, 48,   17, 673,   0, 4712, 2024,    0,  16, -14,   0, 25, 3},
      {2, 15004, 15709, 22,   23, 683,   0, 4099,13663,    0,  22,  31,   0, 25, 3},
      {2, 27000, 32767, 15,   31,   7, 151,11628,28087,25824,  30,   6, -10, 28, 4}
      // clang-format on
  };

  static long correctBitSize(long minimum, long oldBitSize)
  {
    long newBitSize;
    if (oldBitSize <= 0)
      newBitSize = minimum;
    else if (oldBitSize > 32)
      newBitSize = 32;
    else
      newBitSize = oldBitSize;
    return newBitSize;
  };

  // Validates the prm value, throwing if invalid
  static long validatePrm(long prm)
  {
    if (prm < 0 || prm >= 5)
      throw std::invalid_argument("prm must be in the interval [0, 4]");
    return prm;
  };

  static NTL::Vec<long> calculateMvec(const long* vals)
  {
    NTL::Vec<long> mvec;
    NTL::append(mvec, vals[4]);
    if (vals[5] > 1)
      NTL::append(mvec, vals[5]);
    if (vals[6] > 1)
      NTL::append(mvec, vals[6]);
    return mvec;
  };

  static std::vector<long> calculateGens(const long* vals)
  {
    std::vector<long> gens;
    gens.push_back(vals[7]);
    if (vals[8] > 1)
      gens.push_back(vals[8]);
    if (vals[9] > 1)
      gens.push_back(vals[9]);
    return gens;
  };

  static std::vector<long> calculateOrds(const long* vals)
  {
    std::vector<long> ords;
    ords.push_back(vals[10]);
    if (abs(vals[11]) > 1)
      ords.push_back(vals[11]);
    if (abs(vals[12]) > 1)
      ords.push_back(vals[12]);
    return ords;
  };

  static long calculateLevels(bool bootstrap, long outSize, long bitSize)
  {
    long L;
    if (bootstrap)
      L = 900; // that should be enough
    else {
      double nBits =
          (outSize > 0 && outSize < 2 * bitSize) ? outSize : (2 * bitSize);
      double three4twoLvls = log(nBits / 2) / log(1.5);
      double add2NumsLvls = log(nBits) / log(2.0);
      L = (5 + ceil(three4twoLvls + add2NumsLvls)) * 30;
    }
    return L;
  };

  // Returns a reference to the passed-in context once it has been modified to
  // get it ready for test. This is not static because it uses quite a lot of
  // state of the object.
  helib::Context& prepareContext(helib::Context& context)
  {
    if (helib_test::verbose) {
      std::cout << "input bitSizes=" << bitSize << ',' << bitSize2
                << ", output size bound=" << outSize << std::endl;
      if (nthreads > 1)
        std::cout << "  using " << NTL::AvailableThreads() << " threads\n";
      std::cout << "computing key-independent tables..." << std::flush;
    }
    buildModChain(context, L, c, /*willBeBootstrappable=*/bootstrap);
    if (bootstrap) {
      context.enableBootStrapping(mvec);
    }
    buildUnpackSlotEncoding(unpackSlotEncoding, *context.ea);
    if (helib_test::verbose) {
      std::cout << " done.\n";
      context.zMStar.printout();
    }
    return context;
  };

  void prepareSecKey(helib::SecKey& secKey)
  {
    if (helib_test::verbose) {
      std::cout << " L=" << L << ", B=" << B << std::endl;
      std::cout << "\ncomputing key-dependent tables..." << std::flush;
    }
    secKey.GenSecKey();
    addSome1DMatrices(secKey); // compute key-switching matrices
    addFrbMatrices(secKey);
    if (bootstrap)
      secKey.genRecryptData();
    if (helib_test::verbose)
      std::cout << " done\n";
  };

  const long prm;
  const long bitSize;
  const long bitSize2;
  const long outSize;
  const bool bootstrap;
  const long seed;
  const long nthreads;

  const long* vals;
  const long p;
  const long m;
  const NTL::Vec<long> mvec;
  const std::vector<long> gens;
  const std::vector<long> ords;
  const long B;
  const long c;
  const long L;
  helib::Context context;
  helib::SecKey secKey;

  GTestBinaryArith() :
      prm(validatePrm(std::get<0>(GetParam()).prm)),
      bitSize(correctBitSize(5, std::get<0>(GetParam()).bitSize)),
      bitSize2(correctBitSize(bitSize, std::get<0>(GetParam()).bitSize2)),
      outSize(std::get<0>(GetParam()).outSize),
      bootstrap(std::get<0>(GetParam()).bootstrap),
      seed(std::get<0>(GetParam()).seed),
      nthreads(std::get<0>(GetParam()).nthreads),
      vals(mValues[prm]),
      p(vals[0]),
      m(vals[2]),
      mvec(calculateMvec(vals)),
      gens(calculateGens(vals)),
      ords(calculateOrds(vals)),
      B(vals[13]),
      c(vals[14]),
      L(calculateLevels(bootstrap, outSize, bitSize)),
      context(m, p, /*r=*/1, gens, ords),
      secKey(prepareContext(context)){};

  void SetUp() override
  {
    if (seed)
      NTL::SetSeed(NTL::ZZ(seed));
    if (nthreads > 1)
      NTL::SetNumThreads(nthreads);

    prepareSecKey(secKey);

    helib::activeContext = &context; // make things a little easier sometimes

    helib::setupDebugGlobals(&secKey, context.ea);
  }

  virtual void TearDown() override
  {
#ifdef HELIB_DEBUG
    helib::cleanupDebugGlobals();
#endif
  }

  // This gets called once all of the tests have run
public:
  static void TearDownTestCase()
  {
    if (helib_test::verbose)
      helib::printAllTimers(std::cout);
  };
};
constexpr long GTestBinaryArith::mValues[8][15];
std::vector<helib::zzX> GTestBinaryArith::unpackSlotEncoding;

TEST_P(GTestBinaryArith, fifteenForFour)
{
  // Randomly generate up to 15 integers from {0,1} with some entries
  // randomly set to null.  We then encrypt and use the fifteenOrLess4Four
  // function to calculate the binary representation of their sum.  This is
  // then checked against the plaintext calculation.
  // In this case each ciphertext is considered to be the encryption of one
  // bit, however packing more into the slots is possible.
  // LSB is at the start (left) of the vector.

  // Note: fifteenOrLess4Four is not entirely thread safe so this test will only
  // be ran single-threaded. Save current number of threads and set to 1.
  auto numThreads = NTL::AvailableThreads();
  NTL::SetNumThreads(1);

  // vector of ciphertexts corresponding to encrypted input bit vectors.
  std::vector<helib::Ctxt> inBuf(15, helib::Ctxt(secKey));
  std::vector<helib::Ctxt*> inPtrs(15, nullptr);

  // vector of ciphertexts corresponding to the summation of the input bit
  // vectors.
  std::vector<helib::Ctxt> outBuf(5, helib::Ctxt(secKey));

  // Randomly generate and encrypt the input vectors.
  long sum = 0;
  std::string inputBits = "(";
  for (int i = 0; i < 15; i++) {
    if (NTL::RandomBnd(10) > 0) { // Leave empty (null) with small probability.
      inPtrs[i] = &(inBuf[i]);
      long bit = NTL::RandomBnd(2); // Select a randomised bit.
      secKey.Encrypt(inBuf[i], NTL::ZZX(bit));
      inputBits += std::to_string(bit) + ",";
      sum += bit; // Keep track of the plaintext sum.
    } else
      inputBits += "-,"; // This represents a null bit.
  }
  inputBits += ")";

  if (helib_test::verbose) {
    std::cout << std::endl;
    helib::CheckCtxt(inBuf[helib::lsize(inBuf) - 1], "b4 15for4");
  }
  // Add the encrypted bits.
  long numOutputs = fifteenOrLess4Four(helib::CtPtrs_vectorCt(outBuf),
                                       helib::CtPtrs_vectorPt(inPtrs));
  if (helib_test::verbose) {
    std::cout << "numOutputs: " << numOutputs << std::endl;
    helib::CheckCtxt(outBuf[helib::lsize(outBuf) - 1], "after 15for4");
  }

  // Check the result.
  long sum2 = 0;
  for (int i = 0; i < numOutputs; i++) {
    NTL::ZZX poly;
    secKey.Decrypt(poly, outBuf[i]);
    sum2 += to_long(ConstTerm(poly)) << i;
  }
  EXPECT_EQ(sum, sum2) << "inputs = " << inputBits << std::endl;
  if (helib_test::verbose) {
    std::cout << "15to4 succeeded, sum" << inputBits << "=" << sum2
              << std::endl;
  }

  // Restore the number of threads to original value.
  NTL::SetNumThreads(numThreads);
}

TEST_P(GTestBinaryArith, product)
{
  // Randomly generate a pair of numbers of a specified bit size and then
  // encrypt them in binary representation. Then use multTwoNumbers to
  // multiply the two positive binary numbers and then check against the
  // plaintext calculation. Next, use multTwoNumbers with binary numbers in
  // 2's complement to calculate the product where the multiplier is negative
  // and then check against the plaintext calculation.
  //
  // In this case each ciphertext is considered to be the encryption of one
  // bit, however packing more into the slots is possible.
  // LSB is at the start (left) of the vector.

  const helib::EncryptedArray& ea = *context.ea;
  // outSize 1's on the least significant end of mask.
  long mask = (outSize ? ((1L << outSize) - 1) : -1);

  // Choose two random integers with correct bit sizes.
  long multiplicand_data = NTL::RandomBits_long(bitSize);
  long multiplier_data = NTL::RandomBits_long(bitSize2);

  // Encrypt the individual bits.
  NTL::Vec<helib::Ctxt> encrypted_product, encrypted_multiplicand,
      encrypted_multiplier;

  // Resize the vector of ciphertexts (encrypted_bits) to match the input size.
  helib::resize(encrypted_multiplicand, bitSize, helib::Ctxt(secKey));
  for (long i = 0; i < bitSize; i++) {
    secKey.Encrypt(encrypted_multiplicand[i],
                   NTL::ZZX((multiplicand_data >> i) & 1));
    if (bootstrap) { // If bootstrapping then modulo down to a lower level.
      encrypted_multiplicand[i].bringToSet(context.getCtxtPrimes(5));
    }
  }
  helib::resize(encrypted_multiplier, bitSize2, helib::Ctxt(secKey));
  for (long i = 0; i < bitSize2; i++) {
    secKey.Encrypt(encrypted_multiplier[i],
                   NTL::ZZX((multiplier_data >> i) & 1));
    if (bootstrap) { // If bootstrapping then modulo down to a lower level.
      encrypted_multiplier[i].bringToSet(context.getCtxtPrimes(5));
    }
  }
  if (helib_test::verbose) {
    std::cout << "\n  bits-size " << bitSize << '+' << bitSize2;
    if (outSize > 0)
      std::cout << "->" << outSize;
    helib::CheckCtxt(encrypted_multiplier[0], "b4 multiplication");
  }
  std::vector<long> slots; // Vector that will hold the decrypted result.
  // Test multiplication with two positive numbers.
  // A scope which tests multTwoNumbers using wrappers around the encrypted
  // data.
  {
    helib::CtPtrs_VecCt output_wrapper(
        encrypted_product); // A wrapper around the output vector.
    helib::multTwoNumbers(output_wrapper,
                          helib::CtPtrs_VecCt(encrypted_multiplicand),
                          helib::CtPtrs_VecCt(encrypted_multiplier),
                          /*negative=*/false,
                          outSize,
                          &unpackSlotEncoding);
    helib::decryptBinaryNums(slots, output_wrapper, secKey, ea);
  } // output_wrapper is deleted once out of scope.
  if (helib_test::verbose)
    helib::CheckCtxt(encrypted_product[helib::lsize(encrypted_product) - 1],
                     "after multiplication");

  // Calculate the multiplication in the plain.
  long plaintext_product = multiplicand_data * multiplier_data;
  EXPECT_EQ(slots[0], ((multiplicand_data * multiplier_data) & mask))
      << "Positive product error: multiplicand_data=" << multiplicand_data
      << ", multiplier_data=" << multiplier_data << ", but product=" << slots[0]
      << " (should be " << plaintext_product << '&' << mask << '='
      << (plaintext_product & mask) << ")\n";

  if (helib_test::verbose) {
    std::cout << "positive product succeeded: ";
    if (outSize)
      std::cout << "bottom " << outSize << " bits of ";
    std::cout << multiplicand_data << "*" << multiplier_data << "=" << slots[0]
              << std::endl;
  }

  // Test multiplication of numbers in 2's complement where the multiplier is
  // negative.
  secKey.Encrypt(encrypted_multiplier[bitSize2 - 1], NTL::ZZX(1));
  decryptBinaryNums(slots,
                    helib::CtPtrs_VecCt(encrypted_multiplier),
                    secKey,
                    ea,
                    /*negative=*/true);
  multiplier_data = slots[0];
  encrypted_product.kill(); // Clear the data in encrypted_product.
  // A scope which tests multTwoNumbers (with a negative multiplier) using
  // wrappers around the encrypted data.
  {
    helib::CtPtrs_VecCt output_wrapper(
        encrypted_product); // A wrapper around the output vector.
    multTwoNumbers(output_wrapper,
                   helib::CtPtrs_VecCt(encrypted_multiplicand),
                   helib::CtPtrs_VecCt(encrypted_multiplier),
                   /*negative=*/true,
                   outSize,
                   &unpackSlotEncoding);
    decryptBinaryNums(slots, output_wrapper, secKey, ea, /*negative=*/true);
  } // output_wrapper is deleted once out of scope.

  if (helib_test::verbose)
    helib::CheckCtxt(encrypted_product[helib::lsize(encrypted_product) - 1],
                     "after multiplication");

  // Calculate the multiplication in the plain.
  plaintext_product = multiplicand_data * multiplier_data;
  EXPECT_EQ((slots[0] & mask), (plaintext_product & mask))
      << "Negative product error: multiplicand_data=" << multiplicand_data
      << ", multiplier_data=" << multiplier_data << ", but product=" << slots[0]
      << " (should be " << plaintext_product << '&' << mask << '='
      << (plaintext_product & mask) << ")\n";
  if (helib_test::verbose) {
    std::cout << "negative product succeeded: ";
    if (outSize)
      std::cout << "bottom " << outSize << " bits of ";
    std::cout << multiplicand_data << "*" << multiplier_data << "=" << slots[0]
              << std::endl;
  }

#ifdef HELIB_DEBUG
  // Print out the ciphertext with the lowest level after multiplication
  // if HELIB_DEBUG is defined.
  const helib::Ctxt* minCtxt = nullptr;
  long minLvl = 10000000;
  for (const helib::Ctxt& c : encrypted_product) {
    long lvl = c.logOfPrimeSet();
    if (lvl < minLvl) {
      minCtxt = &c;
      minLvl = lvl;
    }
  }
  decryptAndPrint((std::cout << " after multiplication: "),
                  *minCtxt,
                  secKey,
                  ea,
                  0);
  std::cout << std::endl;
#endif
}

TEST_P(GTestBinaryArith, add)
{
  // Randomly generate a pair of numbers of a specified bit size and then
  // encrypt them in binary representation. Then use addTwoNumbers to add the
  // two binary numbers and then check against the plaintext calculation.
  //
  // In this case each ciphertext is considered to be the encryption of one
  // bit, however packing more into the slots is possible.
  // LSB is at the start(left) of the vector.

  const helib::EncryptedArray& ea = *context.ea;
  // outSize 1's on the least significant end of mask.
  long mask = (outSize ? ((1L << outSize) - 1) : -1);

  // Choose two random n-bit integers.
  long addend_data = NTL::RandomBits_long(bitSize);
  long augend_data = NTL::RandomBits_long(bitSize2);

  // Encrypt the individual bits.
  NTL::Vec<helib::Ctxt> encrypted_sum, encrypted_addend, encrypted_augend;

  // Resize the vector of ciphertexts (encrypted_bits) to match the input size.
  helib::resize(encrypted_addend, bitSize, helib::Ctxt(secKey));
  for (long i = 0; i < bitSize; i++) {
    secKey.Encrypt(encrypted_addend[i], NTL::ZZX((addend_data >> i) & 1));
    if (bootstrap) { // If bootstrapping then modulo down to a lower level.
      encrypted_addend[i].bringToSet(context.getCtxtPrimes(5));
    }
  }
  // Resize the vector of ciphertexts (encrypted_bits) to match the input size.
  helib::resize(encrypted_augend, bitSize2, helib::Ctxt(secKey));
  for (long i = 0; i < bitSize2; i++) {
    secKey.Encrypt(encrypted_augend[i], NTL::ZZX((augend_data >> i) & 1));
    if (bootstrap) { // If bootstrapping then modulo down to a lower level.
      encrypted_augend[i].bringToSet(context.getCtxtPrimes(5));
    }
  }
  if (helib_test::verbose) {
    std::cout << "\n  bits-size " << bitSize << '+' << bitSize2;
    if (outSize > 0)
      std::cout << "->" << outSize;
    std::cout << std::endl;
    helib::CheckCtxt(encrypted_augend[0], "b4 addition");
  }

  std::vector<long>
      decrypted_result; // Vector that will hold the decrypted result.
  // Test addition.
  // A scope which tests addTwoNumbers using wrappers around the encrypted data.
  {
    helib::CtPtrs_VecCt output_wrapper(
        encrypted_sum); // A wrapper around the output vector.
    helib::addTwoNumbers(output_wrapper,
                         helib::CtPtrs_VecCt(encrypted_addend),
                         helib::CtPtrs_VecCt(encrypted_augend),
                         outSize,
                         &unpackSlotEncoding);
    helib::decryptBinaryNums(decrypted_result, output_wrapper, secKey, ea);
  } // output_wrapper is deleted once out of scope.

  if (helib_test::verbose)
    helib::CheckCtxt(encrypted_sum[helib::lsize(encrypted_sum) - 1],
                     "after addition");

  // Calculate the addition in the plain.
  long plaintext_sum = addend_data + augend_data;
  EXPECT_EQ(decrypted_result[0], ((addend_data + augend_data) & mask))
      << "addTwoNums error: addend_data=" << addend_data
      << ", augend_data=" << augend_data
      << ", but plaintext_sum=" << decrypted_result[0]
      << " (should be =" << (plaintext_sum & mask) << ")\n";
  if (helib_test::verbose) {
    std::cout << "addTwoNums succeeded: ";
    if (outSize)
      std::cout << "bottom " << outSize << " bits of ";
    std::cout << addend_data << "+" << augend_data << "=" << decrypted_result[0]
              << std::endl;
  }

#ifdef HELIB_DEBUG
  // Print out the ciphertext with the lowest level after addition if
  // HELIB_DEBUG is defined.
  const helib::Ctxt* minCtxt = nullptr;
  long minLvl = 1000;
  for (const helib::Ctxt& c : encrypted_sum) {
    long lvl = c.logOfPrimeSet();
    if (lvl < minLvl) {
      minCtxt = &c;
      minLvl = lvl;
    }
  }
  decryptAndPrint((std::cout << " after addition: "), *minCtxt, secKey, ea, 0);
  std::cout << std::endl;
#endif
}

TEST_P(GTestBinaryArith, addManyNumbers)
{
  // Randomly generate a vector of numbers of a specified bit size and then
  // encrypt them in binary representation. Then use addManyNumbers to add
  // them all up and then check against the plaintext calculation.
  //
  // In this case each ciphertext is considered to be the encryption of one
  // bit, however packing more into the slots is possible.
  // LSB is at the start(left) of the vector.

  const long num_summands = 5;
  const helib::EncryptedArray& ea = *context.ea;
  // outSize 1's on the least significant end of mask.
  long mask = (outSize ? ((1L << outSize) - 1) : -1);

  // Choose a set of random n-bit integers.
  std::vector<long> summands_data;
  for (long i = 0; i < num_summands; ++i)
    summands_data.push_back(NTL::RandomBits_long(bitSize));

  // Encrypt the individual bits.
  std::vector<helib::Ctxt> encrypted_sum;
  std::vector<std::vector<helib::Ctxt>> encrypted_summands;

  // Utility function for encrypting a number into a binary representation.
  const auto encrypt_binary_number =
      [&](const long num) -> std::vector<helib::Ctxt> {
    std::vector<helib::Ctxt> encrypted_num;
    // Resize the vector of ciphertexts (encrypted_bits) to match the input
    // size.
    helib::resize(encrypted_num, bitSize, helib::Ctxt(secKey));
    for (long i = 0; i < bitSize; i++) {
      secKey.Encrypt(encrypted_num[i], NTL::ZZX((num >> i) & 1));
      if (bootstrap) { // If bootstrapping then modulo down to a lower level.
        encrypted_num[i].bringToSet(context.getCtxtPrimes(5));
      }
    }
    return encrypted_num;
  };

  // Encrypt the set of numbers into binary representation.
  for (long i = 0; i < num_summands; ++i)
    encrypted_summands.push_back(encrypt_binary_number(summands_data[i]));

  if (helib_test::verbose)
    helib::CheckCtxt(encrypted_summands[0][0], "b4 addition");

  std::vector<long>
      decrypted_result; // Vector that will hold the decrypted result.
  // Test summation.
  // A scope which tests addManyNumbers using wrappers around the encrypted
  // data.
  {
    helib::CtPtrs_vectorCt output_wrapper(
        encrypted_sum); // A wrapper around the output vector.
    helib::CtPtrMat_vectorCt summands_wrapper(
        encrypted_summands); // A wrapper around the output vector.
    helib::addManyNumbers(output_wrapper,
                          summands_wrapper,
                          outSize,
                          &unpackSlotEncoding);
    helib::decryptBinaryNums(decrypted_result, output_wrapper, secKey, ea);
  } // output_wrapper is deleted once out of scope.

  if (helib_test::verbose)
    helib::CheckCtxt(encrypted_sum[helib::lsize(encrypted_sum) - 1],
                     "after addition");

  // Calculate the summation in the plain.
  long plaintext_sum = std::accumulate(summands_data.begin(),
                                       summands_data.end(),
                                       0l,
                                       std::plus<long>());
  EXPECT_EQ(decrypted_result[0], plaintext_sum & mask);
  if (helib_test::verbose) {
    std::cout << "addManyNums succeeded: ";
    if (outSize)
      std::cout << "bottom " << outSize << " bits of ";
    std::cout << summands_data[0];
    for (long i = 1; i < num_summands; ++i)
      std::cout << "+" << summands_data[i];
    std::cout << "=" << decrypted_result[0] << std::endl;
  }
}

TEST_P(GTestBinaryArith, negateNegatesCorrectly)
{
  // Randomly generate a number in 2's complement and negate it.

  const helib::EncryptedArray& ea = *context.ea;
  unsigned long input_data = NTL::RandomBits_long(bitSize);

  long mask = ((1L << bitSize) - 1);
  long expected_result = ((~input_data) + 1) & mask;
  expected_result = helib::bitSetToLong(expected_result, bitSize);

  std::vector<helib::Ctxt> encrypted_data(bitSize, helib::Ctxt(secKey));

  for (long i = 0; i < bitSize; i++) {
    secKey.Encrypt(encrypted_data[i], NTL::ZZX((input_data >> i) & 1));
    if (bootstrap) { // If bootstrapping then modulo down to a lower level.
      encrypted_data[i].bringToSet(context.getCtxtPrimes(5));
    }
  }

  std::vector<long> decrypted_result;
  std::vector<helib::Ctxt> result_vector(bitSize, helib::Ctxt(secKey));
  helib::CtPtrs_vectorCt output_wrapper(result_vector);
  helib::negateBinary(output_wrapper, helib::CtPtrs_vectorCt(encrypted_data));
  helib::decryptBinaryNums(decrypted_result, output_wrapper, secKey, ea, true);

  EXPECT_EQ(decrypted_result.size(), ea.size());
  for (std::size_t i = 0; i < decrypted_result.size(); ++i) {
    EXPECT_EQ(decrypted_result[i], expected_result) << "i = " << i << std::endl;
  }

  // Make sure it throws for incorrect-length args
  const auto do_negate = [&]() {
    helib::negateBinary(output_wrapper, helib::CtPtrs_vectorCt(encrypted_data));
  };
  encrypted_data.emplace_back(secKey);
  EXPECT_THROW(do_negate(), helib::LogicError);
  encrypted_data.pop_back();
  encrypted_data.pop_back();
  EXPECT_THROW(do_negate(), helib::LogicError);
}

TEST_P(GTestBinaryArith, subtractSubtractsCorrectly)
{
  // Randomly generate two numbers in 2's complement and subtract one from the
  // other.
  const helib::EncryptedArray& ea = *context.ea;
  unsigned long minuend_data = NTL::RandomBits_long(bitSize);
  unsigned long subtrahend_data = NTL::RandomBits_long(bitSize);

  long mask = ((1L << bitSize) - 1);

  // Do the bitSize-bit subtraction manually by negating the subtrahend, adding,
  // and masking.
  long expected_result =
      (minuend_data + (((~subtrahend_data) + 1) & mask)) & mask;
  expected_result = helib::bitSetToLong(expected_result, bitSize);

  std::vector<helib::Ctxt> encrypted_minuend(bitSize, helib::Ctxt(secKey));
  std::vector<helib::Ctxt> encrypted_subtrahend(bitSize, helib::Ctxt(secKey));

  for (long i = 0; i < bitSize; i++) {
    secKey.Encrypt(encrypted_minuend[i], NTL::ZZX((minuend_data >> i) & 1));
    if (bootstrap) { // If bootstrapping then modulo down to a lower level.
      encrypted_minuend[i].bringToSet(context.getCtxtPrimes(5));
    }
  }
  for (long i = 0; i < bitSize; i++) {
    secKey.Encrypt(encrypted_subtrahend[i],
                   NTL::ZZX((subtrahend_data >> i) & 1));
    if (bootstrap) { // If bootstrapping then modulo down to a lower level.
      encrypted_subtrahend[i].bringToSet(context.getCtxtPrimes(5));
    }
  }
  std::vector<long> decrypted_result;
  std::vector<helib::Ctxt> result_vector(bitSize, helib::Ctxt(secKey));
  helib::CtPtrs_vectorCt output_wrapper(result_vector);
  helib::subtractBinary(output_wrapper,
                        helib::CtPtrs_vectorCt(encrypted_minuend),
                        helib::CtPtrs_vectorCt(encrypted_subtrahend));
  helib::decryptBinaryNums(decrypted_result, output_wrapper, secKey, ea, true);

  EXPECT_EQ(decrypted_result.size(), ea.size());
  for (std::size_t i = 0; i < decrypted_result.size(); ++i) {
    EXPECT_EQ(decrypted_result[i], expected_result) << "i = " << i << std::endl;
  }

  // Make sure it throws for incorrect-length args
  result_vector.emplace_back(secKey);
  const auto do_subtract = [&]() {
    helib::subtractBinary(output_wrapper,
                          helib::CtPtrs_vectorCt(encrypted_minuend),
                          helib::CtPtrs_vectorCt(encrypted_subtrahend));
  };
  EXPECT_THROW(do_subtract(), helib::LogicError);
  encrypted_subtrahend.emplace_back(secKey);
  EXPECT_THROW(do_subtract(), helib::LogicError);
}

TEST_P(GTestBinaryArith, binaryMaskMasksCorrectly)
{
  const helib::EncryptedArray& ea = *context.ea;
  const helib::PubKey& pubKey = secKey;
  helib::Ctxt mask(secKey);
  helib::Ptxt<helib::BGV> mask_data(context);
  for (std::size_t i = 0; i < mask_data.size(); ++i)
    mask_data[i] = i % 2;

  pubKey.Encrypt(mask, mask_data);
  std::vector<helib::Ctxt> eNums(bitSize, helib::Ctxt(secKey));
  long input_number = NTL::RandomBits_long(bitSize);

  for (long i = 0; i < bitSize; ++i)
    secKey.Encrypt(eNums[i], NTL::ZZX((input_number >> i) & 1));

  std::vector<helib::Ctxt> output(eNums);
  helib::CtPtrs_vectorCt output_wrapper(output);
  helib::binaryMask(output_wrapper, mask);

  std::vector<long> decrypted_result;
  helib::decryptBinaryNums(decrypted_result, output_wrapper, secKey, ea);

  EXPECT_EQ(decrypted_result.size(), ea.size());
  for (long i = 0; i < bitSize; ++i)
    if (i % 2 == 1) {
      EXPECT_EQ(decrypted_result[i], input_number) << "i = " << i << std::endl;
    } else {
      EXPECT_EQ(decrypted_result[i], 0L) << "i = " << i << std::endl;
    }
  // No error cases to test, all sized inputs are valid
}

TEST_P(GTestBinaryArith, binaryCondWorksCorrectly)
{
  const helib::EncryptedArray& ea = *context.ea;
  const helib::PubKey& pubKey = secKey;
  helib::Ctxt cond(secKey);
  helib::Ptxt<helib::BGV> cond_data(context);
  for (std::size_t i = 0; i < cond_data.size(); ++i)
    cond_data[i] = i % 2;

  pubKey.Encrypt(cond, cond_data);
  std::vector<helib::Ctxt> lhsNums(bitSize, helib::Ctxt(secKey));
  std::vector<helib::Ctxt> rhsNums(bitSize, helib::Ctxt(secKey));
  long lhs_number = NTL::RandomBits_long(bitSize);
  long rhs_number = NTL::RandomBits_long(bitSize);

  for (long i = 0; i < bitSize; ++i) {
    secKey.Encrypt(lhsNums[i], NTL::ZZX((lhs_number >> i) & 1));
    secKey.Encrypt(rhsNums[i], NTL::ZZX((rhs_number >> i) & 1));
  }

  std::vector<helib::Ctxt> output(lhsNums);
  helib::CtPtrs_vectorCt output_wrapper(output);
  helib::binaryCond(output_wrapper,
                    cond,
                    helib::CtPtrs_vectorCt(lhsNums),
                    helib::CtPtrs_vectorCt(rhsNums));

  std::vector<long> decrypted_result;
  helib::decryptBinaryNums(decrypted_result, output_wrapper, secKey, ea);

  EXPECT_EQ(decrypted_result.size(), ea.size());
  for (long i = 0; i < bitSize; ++i) {
    EXPECT_EQ(decrypted_result[i], (i % 2) ? lhs_number : rhs_number)
        << "i = " << i << std::endl;
  }

  // All three CtPtrs args must have same size.  Check that it throws if not.
  const auto do_cond = [&]() {
    helib::binaryCond(output_wrapper,
                      cond,
                      helib::CtPtrs_vectorCt(lhsNums),
                      helib::CtPtrs_vectorCt(rhsNums));
  };
  lhsNums.emplace_back(secKey);
  EXPECT_THROW(do_cond(), /*1*/ helib::LogicError);
  rhsNums.emplace_back(secKey);
  EXPECT_THROW(do_cond(), /*2*/ helib::LogicError);
  output.emplace_back(secKey);
  output.emplace_back(secKey);
  EXPECT_THROW(do_cond(), /*3*/ helib::LogicError);
}

TEST_P(GTestBinaryArith, concatBinaryNumsConcatsCorrectly)
{
  const helib::EncryptedArray& ea = *context.ea;
  long lhs_number = NTL::RandomBits_long(bitSize);
  long rhs_number = NTL::RandomBits_long(bitSize);

  std::vector<helib::Ctxt> lhs(bitSize, helib::Ctxt(secKey));
  std::vector<helib::Ctxt> rhs(bitSize, helib::Ctxt(secKey));
  for (long i = 0; i < bitSize; ++i) {
    secKey.Encrypt(lhs[i], NTL::ZZX((lhs_number >> i) & 1));
    secKey.Encrypt(rhs[i], NTL::ZZX((rhs_number >> i) & 1));
  }

  std::vector<helib::Ctxt> output(bitSize * 2, helib::Ctxt(secKey));
  helib::CtPtrs_vectorCt output_wrapper(output);
  helib::concatBinaryNums(output_wrapper,
                          helib::CtPtrs_vectorCt(lhs),
                          helib::CtPtrs_vectorCt(rhs));

  std::vector<long> decrypted_result;
  helib::decryptBinaryNums(decrypted_result, output_wrapper, secKey, ea);

  EXPECT_EQ(decrypted_result.size(), ea.size());
  for (std::size_t i = 0; i < decrypted_result.size(); ++i) {
    EXPECT_EQ(decrypted_result[i], (rhs_number << bitSize) + lhs_number)
        << "i = " << i << std::endl;
  }

  // Make sure that incorrect sizing issues throw
  const auto do_concat = [&]() {
    helib::concatBinaryNums(output_wrapper,
                            helib::CtPtrs_vectorCt(lhs),
                            helib::CtPtrs_vectorCt(rhs));
  };
  output.emplace_back(secKey);
  EXPECT_THROW(do_concat(), helib::LogicError);
  output.pop_back();
  output.pop_back();
  EXPECT_THROW(do_concat(), helib::LogicError);
}

TEST_P(GTestBinaryArith, splitBinaryNumsSplitsCorrectly)
{
  const helib::EncryptedArray& ea = *context.ea;
  long lhs_number = NTL::RandomBits_long(bitSize + 1);
  long rhs_number = NTL::RandomBits_long(bitSize);

  std::vector<helib::Ctxt> lhs(bitSize + 1, helib::Ctxt(secKey));
  std::vector<helib::Ctxt> rhs(bitSize, helib::Ctxt(secKey));
  for (long i = 0; i < bitSize; ++i) {
    secKey.Encrypt(lhs[i], NTL::ZZX((lhs_number >> i) & 1));
    secKey.Encrypt(rhs[i], NTL::ZZX((rhs_number >> i) & 1));
  }
  secKey.Encrypt(lhs[bitSize], NTL::ZZX((lhs_number >> bitSize) & 1));

  std::vector<helib::Ctxt> concatenation((bitSize * 2) + 1,
                                         helib::Ctxt(secKey));
  helib::CtPtrs_vectorCt concatenation_wrapper(concatenation);
  helib::concatBinaryNums(concatenation_wrapper,
                          helib::CtPtrs_vectorCt(lhs),
                          helib::CtPtrs_vectorCt(rhs));

  std::vector<helib::Ctxt> lhs_output(bitSize + 1, helib::Ctxt(secKey));
  std::vector<helib::Ctxt> rhs_output(bitSize, helib::Ctxt(secKey));
  helib::CtPtrs_vectorCt lhs_output_wrapper(lhs_output);
  helib::CtPtrs_vectorCt rhs_output_wrapper(rhs_output);
  helib::splitBinaryNums(lhs_output_wrapper,
                         rhs_output_wrapper,
                         concatenation_wrapper);

  std::vector<long> decrypted_lhs;
  std::vector<long> decrypted_rhs;
  helib::decryptBinaryNums(decrypted_lhs, lhs_output_wrapper, secKey, ea);
  helib::decryptBinaryNums(decrypted_rhs, rhs_output_wrapper, secKey, ea);

  EXPECT_EQ(decrypted_lhs.size(), ea.size());
  EXPECT_EQ(decrypted_rhs.size(), ea.size());
  for (std::size_t i = 0; i < decrypted_lhs.size(); ++i) {
    EXPECT_EQ(decrypted_lhs[i], lhs_number) << "i = " << i << std::endl;
    EXPECT_EQ(decrypted_rhs[i], rhs_number) << "i = " << i << std::endl;
  }

  // Make sure it throws if the sizes don't line up
  const auto do_split = [&]() {
    helib::splitBinaryNums(lhs_output_wrapper,
                           rhs_output_wrapper,
                           concatenation_wrapper);
  };
  lhs_output.emplace_back(secKey);
  EXPECT_THROW(do_split(), helib::LogicError);
  lhs_output.pop_back();
  lhs_output.pop_back();
  EXPECT_THROW(do_split(), helib::LogicError);
}

TEST_P(GTestBinaryArith, bitwiseShiftShiftsCorrectly)
{
  const helib::EncryptedArray& ea = *context.ea;
  long number = NTL::RandomBits_long(bitSize);

  std::vector<helib::Ctxt> eNums(bitSize, helib::Ctxt(secKey));
  for (long i = 0; i < bitSize; ++i)
    secKey.Encrypt(eNums[i], NTL::ZZX((number >> i) & 1));

  unsigned long mask = (1Lu << bitSize) - 1;
  for (long shamt = 0; shamt <= bitSize; ++shamt) {
    std::vector<helib::Ctxt> output(bitSize, helib::Ctxt(secKey));
    helib::CtPtrs_vectorCt output_wrapper(output);

    helib::leftBitwiseShift(output_wrapper,
                            helib::CtPtrs_vectorCt(eNums),
                            shamt);

    std::vector<long> decrypted_result;
    helib::decryptBinaryNums(decrypted_result, output_wrapper, secKey, ea);

    EXPECT_EQ(decrypted_result.size(), ea.size());
    for (std::size_t i = 0; i < decrypted_result.size(); ++i)
      EXPECT_EQ(decrypted_result[i], (number << shamt) & mask)
          << "i = " << i << std::endl;
  }

  // Make sure it throws if output and input aren't the same size
  std::vector<helib::Ctxt> output(bitSize + 1, helib::Ctxt(secKey));
  helib::CtPtrs_vectorCt output_wrapper(output);

  const auto do_shift = [&]() {
    helib::leftBitwiseShift(output_wrapper, helib::CtPtrs_vectorCt(eNums), 0);
  };
  EXPECT_THROW(do_shift(), helib::LogicError);
  output.pop_back();
  output.pop_back();
  EXPECT_THROW(do_shift(), helib::LogicError);
}

TEST_P(GTestBinaryArith, bitwiseRotateRotatesCorrectly)
{
  const helib::EncryptedArray& ea = *context.ea;
  long input = NTL::RandomBits_long(bitSize);

  std::vector<helib::Ctxt> eNums(bitSize, helib::Ctxt(secKey));
  for (long i = 0; i < bitSize; ++i)
    secKey.Encrypt(eNums[i], NTL::ZZX((input >> i) & 1));

  const auto plaintext_rotate = [](long num, long amt, long bitSize) {
    // Make sure that amt is in the right range
    amt = ((amt % bitSize) + bitSize) % bitSize;
    long mask = (1LU << bitSize) - 1;
    // Get the left hand part
    long result = (num << amt) & mask;
    // Get the right-hand part
    result |= (num >> (bitSize - amt)) & mask;
    return result;
  };

  // Test all rotation amounts from negative values which wrap around mod
  // bitSize, up to positive values which wrap around mod bitSize.
  for (long rotamt = -bitSize - 1; rotamt <= bitSize + 1; ++rotamt) {
    std::vector<helib::Ctxt> output(bitSize, helib::Ctxt(secKey));
    helib::CtPtrs_vectorCt output_wrapper(output);
    helib::bitwiseRotate(output_wrapper, helib::CtPtrs_vectorCt(eNums), rotamt);
    std::vector<long> decrypted_result;
    helib::decryptBinaryNums(decrypted_result, output_wrapper, secKey, ea);
    EXPECT_EQ(decrypted_result.size(), ea.size());
    for (std::size_t i = 0; i < decrypted_result.size(); ++i) {
      EXPECT_EQ(decrypted_result[i], plaintext_rotate(input, rotamt, bitSize))
          << "i = " << i << std::endl;
    }
  }

  // Check that non-matching input and output sizes throw
  std::vector<helib::Ctxt> output(bitSize + 1, helib::Ctxt(secKey));
  helib::CtPtrs_vectorCt output_wrapper(output);
  const auto do_rotate = [&]() {
    helib::bitwiseRotate(output_wrapper, helib::CtPtrs_vectorCt(eNums), 0);
  };
  EXPECT_THROW(do_rotate(), helib::LogicError);
  output.pop_back();
  output.pop_back();
  EXPECT_THROW(do_rotate(), helib::LogicError);
}

TEST_P(GTestBinaryArith, binaryAndWithLongAndsCorrectly)
{
  const helib::EncryptedArray& ea = *context.ea;
  long number = NTL::RandomBits_long(bitSize);

  unsigned long long_mask = 0;
  std::vector<long> mask;
  std::vector<helib::Ctxt> eNums(bitSize, helib::Ctxt(secKey));
  for (long i = 0; i < bitSize; ++i) {
    secKey.Encrypt(eNums[i], NTL::ZZX((number >> i) & 1));
    mask.push_back(i % 2);
    long_mask |= (i % 2) << i;
  }

  std::vector<helib::Ctxt> output(bitSize, helib::Ctxt(secKey));
  helib::CtPtrs_vectorCt output_wrapper(output);

  helib::bitwiseAnd(output_wrapper, helib::CtPtrs_vectorCt(eNums), mask);

  std::vector<long> decrypted_result;
  helib::decryptBinaryNums(decrypted_result, output_wrapper, secKey, ea);

  EXPECT_EQ(decrypted_result.size(), ea.size());
  for (std::size_t i = 0; i < decrypted_result.size(); ++i) {
    EXPECT_EQ(decrypted_result[i], (number & long_mask))
        << "i = " << i << std::endl;
  }

  // Check that non-matching input and output sizes throw
  const auto do_and = [&]() {
    helib::bitwiseAnd(output_wrapper, helib::CtPtrs_vectorCt(eNums), mask);
  };
  output.emplace_back(secKey);
  EXPECT_THROW(do_and(), helib::LogicError);
  output.pop_back();
  output.pop_back();
  EXPECT_THROW(do_and(), helib::LogicError);
}

TEST_P(GTestBinaryArith, binaryXORXORsCorrectly)
{
  const helib::EncryptedArray& ea = *context.ea;
  long lhs = NTL::RandomBits_long(bitSize);
  long rhs = NTL::RandomBits_long(bitSize);

  std::vector<helib::Ctxt> encrypted_lhs(bitSize, helib::Ctxt(secKey));
  std::vector<helib::Ctxt> encrypted_rhs(bitSize, helib::Ctxt(secKey));

  for (long i = 0; i < bitSize; ++i) {
    secKey.Encrypt(encrypted_lhs[i], NTL::ZZX((lhs >> i) & 1));
    secKey.Encrypt(encrypted_rhs[i], NTL::ZZX((rhs >> i) & 1));
  }

  std::vector<helib::Ctxt> output(bitSize, helib::Ctxt(secKey));
  helib::CtPtrs_vectorCt output_wrapper(output);

  helib::bitwiseXOR(output_wrapper,
                    helib::CtPtrs_vectorCt(encrypted_lhs),
                    helib::CtPtrs_vectorCt(encrypted_rhs));

  std::vector<long> decrypted_result;
  helib::decryptBinaryNums(decrypted_result, output_wrapper, secKey, ea);

  EXPECT_EQ(decrypted_result.size(), ea.size());
  for (std::size_t i = 0; i < decrypted_result.size(); ++i) {
    EXPECT_EQ(decrypted_result[i], lhs ^ rhs) << "i = " << i << std::endl;
  }

  // Check that non-matching sizes throw
  const auto do_xor = [&]() {
    helib::bitwiseXOR(output_wrapper,
                      helib::CtPtrs_vectorCt(encrypted_lhs),
                      helib::CtPtrs_vectorCt(encrypted_rhs));
  };
  output.emplace_back(secKey);
  EXPECT_THROW(do_xor(), helib::LogicError);
  encrypted_lhs.emplace_back(secKey);
  EXPECT_THROW(do_xor(), helib::LogicError);
}

TEST_P(GTestBinaryArith, binaryAndAndsCorrectly)
{
  const helib::EncryptedArray& ea = *context.ea;
  long lhs = NTL::RandomBits_long(bitSize);
  long rhs = NTL::RandomBits_long(bitSize);

  std::vector<helib::Ctxt> encrypted_lhs(bitSize, helib::Ctxt(secKey));
  std::vector<helib::Ctxt> encrypted_rhs(bitSize, helib::Ctxt(secKey));

  for (long i = 0; i < bitSize; ++i) {
    secKey.Encrypt(encrypted_lhs[i], NTL::ZZX((lhs >> i) & 1));
    secKey.Encrypt(encrypted_rhs[i], NTL::ZZX((rhs >> i) & 1));
  }

  std::vector<helib::Ctxt> output(bitSize, helib::Ctxt(secKey));
  helib::CtPtrs_vectorCt output_wrapper(output);

  helib::bitwiseAnd(output_wrapper,
                    helib::CtPtrs_vectorCt(encrypted_lhs),
                    helib::CtPtrs_vectorCt(encrypted_rhs));

  std::vector<long> decrypted_result;
  helib::decryptBinaryNums(decrypted_result, output_wrapper, secKey, ea);

  EXPECT_EQ(decrypted_result.size(), ea.size());
  for (std::size_t i = 0; i < decrypted_result.size(); ++i) {
    EXPECT_EQ(decrypted_result[i], lhs & rhs) << "i = " << i << std::endl;
  }

  // Make sure it throws if arguments' sizes don't match
  const auto do_and = [&]() {
    helib::bitwiseAnd(output_wrapper,
                      helib::CtPtrs_vectorCt(encrypted_lhs),
                      helib::CtPtrs_vectorCt(encrypted_rhs));
  };
  output.emplace_back(secKey);
  EXPECT_THROW(do_and(), helib::LogicError);
  encrypted_lhs.emplace_back(secKey);
  EXPECT_THROW(do_and(), helib::LogicError);
}

TEST_P(GTestBinaryArith, binaryOrOrsCorrectly)
{
  const helib::EncryptedArray& ea = *context.ea;
  long lhs = NTL::RandomBits_long(bitSize);
  long rhs = NTL::RandomBits_long(bitSize);

  std::vector<helib::Ctxt> encrypted_lhs(bitSize, helib::Ctxt(secKey));
  std::vector<helib::Ctxt> encrypted_rhs(bitSize, helib::Ctxt(secKey));

  for (long i = 0; i < bitSize; ++i) {
    secKey.Encrypt(encrypted_lhs[i], NTL::ZZX((lhs >> i) & 1));
    secKey.Encrypt(encrypted_rhs[i], NTL::ZZX((rhs >> i) & 1));
  }

  std::vector<helib::Ctxt> output(bitSize, helib::Ctxt(secKey));
  helib::CtPtrs_vectorCt output_wrapper(output);

  helib::bitwiseOr(output_wrapper,
                   helib::CtPtrs_vectorCt(encrypted_lhs),
                   helib::CtPtrs_vectorCt(encrypted_rhs));

  std::vector<long> decrypted_result;
  helib::decryptBinaryNums(decrypted_result, output_wrapper, secKey, ea);

  EXPECT_EQ(decrypted_result.size(), ea.size());
  for (std::size_t i = 0; i < decrypted_result.size(); ++i) {
    EXPECT_EQ(decrypted_result[i], lhs | rhs) << "i = " << i << std::endl;
  }

  // Make sure it throws if arguments' sizes don't match
  const auto do_or = [&]() {
    helib::bitwiseOr(output_wrapper,
                     helib::CtPtrs_vectorCt(encrypted_lhs),
                     helib::CtPtrs_vectorCt(encrypted_rhs));
  };
  output.emplace_back(secKey);
  EXPECT_THROW(do_or(), helib::LogicError);
  encrypted_lhs.emplace_back(secKey);
  EXPECT_THROW(do_or(), helib::LogicError);
}

TEST_P(GTestBinaryArith, bitwiseNotNotsCorrectly)
{
  const helib::EncryptedArray& ea = *context.ea;
  long input = NTL::RandomBits_long(bitSize);

  std::vector<helib::Ctxt> eNums(bitSize, helib::Ctxt(secKey));
  long mask = (1LU << bitSize) - 1;

  for (long i = 0; i < bitSize; ++i)
    secKey.Encrypt(eNums[i], NTL::ZZX((input >> i) & 1));

  std::vector<helib::Ctxt> output(bitSize, helib::Ctxt(secKey));
  helib::CtPtrs_vectorCt output_wrapper(output);

  helib::bitwiseNot(output_wrapper, helib::CtPtrs_vectorCt(eNums));

  std::vector<long> decrypted_result;
  helib::decryptBinaryNums(decrypted_result, output_wrapper, secKey, ea);

  EXPECT_EQ(decrypted_result.size(), ea.size());
  for (std::size_t i = 0; i < decrypted_result.size(); ++i) {
    EXPECT_EQ(decrypted_result[i], (~input) & mask) << "i = " << i << std::endl;
  }

  // Make sure it throws if input and output are different sizes
  const auto do_not = [&]() {
    helib::bitwiseNot(output_wrapper, helib::CtPtrs_vectorCt(eNums));
  };
  output.emplace_back(secKey);
  EXPECT_THROW(do_not(), helib::LogicError);
  output.pop_back();
  output.pop_back();
  EXPECT_THROW(do_not(), helib::LogicError);
}

INSTANTIATE_TEST_SUITE_P(
    smallParameterSizesRepeated,
    GTestBinaryArith,
    ::testing::Combine(
        ::testing::Values(
            // SLOW
            Parameters(1, 5, 0, 0, false, 0, 1)
            // FAST
            // Parameters(0, 5, 0, 0, false, 0, 1)), ::testing::Range(0,2)

            ),
        ::testing::Range(0, 2)) // The range is for repeats
);

} // anonymous namespace


#pragma once

#include <Eigen/Dense>
#include <boost/assert.hpp>
#include <type_traits>

#include "spectral/basis/spectral_basis_factory_hermite.hpp"
#include "spectral/basis/spectral_basis_factory_ks.hpp"
#include "spectral/basis/toolbox/spectral_basis.hpp"
#include "spectral/polar_to_hermite.hpp"

namespace boltzmann {

class MQEval;  // forward declaration

namespace detail_ {
/**
 * @helper to evaluate macroscopic quantities
 *
 * @param mq
 *
 * @return
 */
struct MQEval_helper
{
 public:
  typedef Eigen::Vector3d vector_t;
  typedef Eigen::Matrix3d tensor_t;
  typedef double scalar_t;

 public:
  MQEval_helper(const MQEval &mq)
      : mq_(mq)
  { /* empty */
  }

  void operator()(const double *ptr_c, unsigned int N);

  template <typename DERIVED>
  void operator()(const Eigen::DenseBase<DERIVED> &c);

 public:
  double m;    // mass
  double e;    // energy
  vector_t v;  // velocity
  vector_t r;  // energy flow
  tensor_t P;  // pressure
  vector_t q;  // heat flow
  tensor_t M;  // momentum flow

 private:
  const MQEval &mq_;
};

}  // detail_

/**
 * @brief macroscopic quantities evaluator
 *
 */
class MQEval
{
 public:
  typedef detail_::MQEval_helper evaluator_t;

 public:
  MQEval() { /*  default constructor */}

  template <typename BASIS>
  MQEval(const BASIS &basis)
  {
    init(basis);
  }

  /**
   *  @brief evaluator to do actual compuations
   *
   *  Returns a struct storing the moments, it has a const reference to *this
   * (where the
   *  coefficients are stored).
   *
   *  Attention: MQEval_helper depends on MQEval by a reference, make sure it
   * does not run out of
   * scope as
   *  long as the MQEval_helper is in use.
   *
   */
  detail_::MQEval_helper evaluator() const { return detail_::MQEval_helper(*this); }

  /**
   *
   * @param basis  Polar-Laguerre basis
   */
  template <typename BASIS>
  void init(const BASIS &basis);

 private:
  typedef Eigen::VectorXd vec_t;
  typedef Eigen::ArrayXd array_t;
  typedef Eigen::MatrixXd mat_t;

 private:
  /// basis size
  unsigned int N_ = 0;
  const double tol_ = 1e-10;

  array_t mass_;
  array_t energy_;
  array_t ux_;
  array_t uy_;
  /// momentum flow
  array_t uxx_;
  array_t uyy_;
  array_t uxy_;
  /// energy flow
  array_t rx_;
  array_t ry_;

 public:
  //@{
  /// coefficients in Polar-Laguerre basis
  const array_t &cmass() const { return mass_; }
  const array_t &cenergy() const { return energy_; }
  const array_t &cux() const { return ux_; }
  const array_t &cuy() const { return uy_; }
  const array_t &cuxx() const { return uxx_; }
  const array_t &cuyy() const { return uyy_; }
  const array_t &cuxy() const { return uxy_; }
  const array_t &crx() const { return rx_; }
  const array_t &cry() const { return ry_; }
  //@}

  unsigned int N() const { return N_; }

 private:
  /**
   * @brief trim non zero coefficients starting from the end
   */
  template <typename DERIVED>
  void trim_coeff(Eigen::DenseBase<DERIVED> &coeffs);
};

// ---------------------------------------------------------------------------
template <typename BASIS>
void
MQEval::init(const BASIS &basis)
{
  typedef BASIS polar_basis_t;
  // compute coefficients

  int K = spectral::get_max_k(basis) + 1;
  int N = basis.n_dofs();
  N_ = N;

  typedef SpectralBasisFactoryHN::basis_type hermite_basis_t;
  hermite_basis_t hermite_basis;

  SpectralBasisFactoryHN::create(hermite_basis, K, 2);
  Polar2Hermite<polar_basis_t, hermite_basis_t> P2H(basis, hermite_basis);
  typedef typename hermite_basis_t::elem_t hermite_elem_t;

  typedef typename boost::mpl::at_c<typename hermite_elem_t::types_t, 0>::type hx_t;
  typedef typename boost::mpl::at_c<typename hermite_elem_t::types_t, 1>::type hy_t;

  typename hermite_elem_t::Acc::template get<hx_t> get_hx;
  typename hermite_elem_t::Acc::template get<hy_t> get_hy;

  // Hermite quadrature
  QHermiteW quad(0.5, K);
  // Hermite polynomials
  HermiteNW<double> hermw(K);
  hermw.compute(quad.pts());

  // test for mass
  vec_t herm_coeffs(N);
  // apply "quadrature":
  Eigen::Map<const array_t> x(quad.points_data(), K);
  Eigen::Map<const array_t> w(quad.weights_data(), K);

  // ------------------------------
  // MASS
  for (int i = 0; i < N; ++i) {
    int kx = get_hx(hermite_basis.get_elem(i)).get_id().k;
    int ky = get_hy(hermite_basis.get_elem(i)).get_id().k;
    auto hx = hermw.get_array(kx);
    auto hy = hermw.get_array(ky);

    herm_coeffs[i] = (hx * w).sum() * (hy * w).sum();
  }
  mass_.resize(N);
  P2H.to_hermite_T(mass_, herm_coeffs);

  // ------------------------------
  // ENERGY
  for (int i = 0; i < N; ++i) {
    int kx = get_hx(hermite_basis.get_elem(i)).get_id().k;
    int ky = get_hy(hermite_basis.get_elem(i)).get_id().k;

    auto hx = hermw.get_array(kx);
    auto hy = hermw.get_array(ky);

    herm_coeffs[i] =
        (x * x * hx * w).sum() * (hy * w).sum() + (x * x * hy * w).sum() * (hx * w).sum();
  }
  energy_.resize(N);
  P2H.to_hermite_T(energy_, herm_coeffs);

  // ------------------------------
  // MOMENTUM
  vec_t herm_coeffs2(N);
  for (int i = 0; i < N; ++i) {
    int kx = get_hx(hermite_basis.get_elem(i)).get_id().k;
    int ky = get_hy(hermite_basis.get_elem(i)).get_id().k;
    auto hx = hermw.get_array(kx);
    auto hy = hermw.get_array(ky);
    double sumx = (x * hx * w).sum() * (hy * w).sum();
    double sumy = (x * hy * w).sum() * (hx * w).sum();
    herm_coeffs[i] = sumx;
    herm_coeffs2[i] = sumy;
  }

  ux_.resize(N);
  P2H.to_hermite_T(ux_, herm_coeffs);
  uy_.resize(N);
  P2H.to_hermite_T(uy_, herm_coeffs2);

  // resize to non-zero contribution
  trim_coeff(mass_);
  trim_coeff(energy_);
  trim_coeff(ux_);
  trim_coeff(uy_);

  uxx_.resize(N);
  uxy_.resize(N);
  uyy_.resize(N);
  // ------------------------------
  // Momentum flow
  for (int i = 0; i < N; ++i) {
    int kx = get_hx(hermite_basis.get_elem(i)).get_id().k;
    int ky = get_hy(hermite_basis.get_elem(i)).get_id().k;
    auto hx = hermw.get_array(kx);
    auto hy = hermw.get_array(ky);
    const double tx = (x * hx * w).sum();
    const double ty = (x * hy * w).sum();
    uxx_[i] = (x * x * hx * w).sum() * (hy * w).sum();
    uxy_[i] = (x * hx * w).sum() * (x * hy * w).sum();
    uyy_[i] = (x * x * hy * w).sum() * (hx * w).sum();
  }
  herm_coeffs = uxx_;
  P2H.to_hermite_T(uxx_, herm_coeffs);
  herm_coeffs = uxy_;
  P2H.to_hermite_T(uxy_, herm_coeffs);
  herm_coeffs = uyy_;
  P2H.to_hermite_T(uyy_, herm_coeffs);

  // ------------------------------
  // Momentum flow
  uxx_.resize(N);
  uxy_.resize(N);
  uyy_.resize(N);
  for (int i = 0; i < N; ++i) {
    int kx = get_hx(hermite_basis.get_elem(i)).get_id().k;
    int ky = get_hy(hermite_basis.get_elem(i)).get_id().k;
    auto hx = hermw.get_array(kx);
    auto hy = hermw.get_array(ky);

    uxx_[i] = (x * x * hx * w).sum() * (hy * w).sum();
    uxy_[i] = (x * hx * w).sum() * (x * hy * w).sum();
    uyy_[i] = (x * x * hy * w).sum() * (hx * w).sum();
  }
  herm_coeffs = uxx_;
  P2H.to_hermite_T(uxx_, herm_coeffs);
  herm_coeffs = uxy_;
  P2H.to_hermite_T(uxy_, herm_coeffs);
  herm_coeffs = uyy_;
  P2H.to_hermite_T(uyy_, herm_coeffs);

  // ------------------------------
  // Energy flow
  rx_.resize(N);
  ry_.resize(N);
  for (int i = 0; i < N; ++i) {
    int kx = get_hx(hermite_basis.get_elem(i)).get_id().k;
    int ky = get_hy(hermite_basis.get_elem(i)).get_id().k;
    auto hx = hermw.get_array(kx);
    auto hy = hermw.get_array(ky);

    rx_[i] =
        (x * x * x * hx * w).sum() * (hy * w).sum() + (x * hx * w).sum() * (x * x * hy * w).sum();
    ry_[i] =
        (x * x * x * hy * w).sum() * (hx * w).sum() + (x * hy * w).sum() * (x * x * hx * w).sum();
  }
  herm_coeffs = rx_;
  P2H.to_hermite_T(rx_, herm_coeffs);
  herm_coeffs = ry_;
  P2H.to_hermite_T(ry_, herm_coeffs);

  trim_coeff(uxx_);
  trim_coeff(uxy_);
  trim_coeff(uyy_);

  trim_coeff(rx_);
  trim_coeff(ry_);
}

template <typename DERIVED>
void
MQEval::trim_coeff(Eigen::DenseBase<DERIVED> &coeffs)
{
  for (unsigned int i = N_ - 1; i >= 0; --i) {
    if (std::abs(coeffs[i]) > tol_) {
      DERIVED tmp = coeffs.segment(0, i + 1);
      coeffs.derived().resize(i + 1);
      coeffs = tmp;
      break;
    }
  }
}

namespace detail_ {
template <typename DERIVED>
void
MQEval_helper::operator()(const Eigen::DenseBase<DERIVED> &c)
{
  static_assert(DERIVED::RowsAtCompileTime == 1 || DERIVED::ColsAtCompileTime == 1,
                "Shape mismatch");

  BOOST_ASSERT(c.cols() * c.rows() == mq_.N());
  auto &cmass = mq_.cmass();
  auto &cux = mq_.cux();
  auto &cuy = mq_.cuy();
  auto &ce = mq_.cenergy();
  auto &cuxx = mq_.cuxx();
  auto &cuxy = mq_.cuxy();
  auto &cuyy = mq_.cuyy();
  auto &crx = mq_.crx();
  auto &cry = mq_.cry();

  auto ca = c.derived().array();

  double rho = (ca.segment(0, cmass.size()) * cmass).sum();
  m = rho;
  double rho_vx = (ca.segment(0, cux.size()) * cux).sum();
  double rho_vy = (ca.segment(0, cuy.size()) * cuy).sum();
  // total energy
  double rho_e = (ca.segment(0, ce.size()) * ce).sum();
  e = rho_e / rho;
  // energy density per unit volume
  double w = 0.5 * rho_e;
  double mxx = (ca.segment(0, cuxx.size()) * cuxx).sum();
  double mxy = (ca.segment(0, cuxy.size()) * cuxy).sum();
  double myy = (ca.segment(0, cuyy.size()) * cuyy).sum();
  M << mxx, mxy, 0, mxy, myy, 0, 0, 0, 0;
  v(0) = rho_vx / rho;
  v(1) = rho_vy / rho;
  v(2) = 0;
  double pxx = mxx - rho * v(0) * v(0);
  double pxy = mxy - rho * v(1) * v(0);
  double pyy = myy - rho * v(1) * v(1);
  P << pxx, pxy, 0, pxy, pyy, 0, 0, 0, 0;
  // energy flow
  r(0) = (ca.segment(0, crx.size()) * crx).sum();
  r(1) = (ca.segment(0, cry.size()) * cry).sum();
  r(2) = 0;
  double v2 = v.squaredNorm();
  // heat flow
  q(0) = 0.5 * (r(0) - 2 * (v * M.row(0)).sum() + rho * v2 + 2 * v(0) * w + v2 * v(0) * rho);
  q(1) = 0.5 * (r(1) - 2 * (v * M.row(1)).sum() + rho * v2 + 2 * v(1) * w + v2 * v(1) * rho);
  q(2) = 0;
}

inline void
MQEval_helper::operator()(const double *ptr_c, unsigned int N)
{
  assert(N == mq_.N());
  Eigen::Map<const Eigen::ArrayXd> c(ptr_c, N);
  this->operator()(c);
}

}  // detail_

}  // end namespace boltzmann


#include "visa.hpp"
#include <chrono>
#include <thread>
#include <iostream>
#include <armadillo>
#include <cstdlib>
#include <ctime>

using namespace std;

int main( int argc, char** argv )
{
  srand(time(0));
  arma::vec gaussian(700);
  for ( int i=0;i<gaussian.n_elem;i++ )
  {
    double x = static_cast<double>(i)-static_cast<double>(gaussian.n_elem)/2.0;
    gaussian(i) = exp(-x*x*0.0001);
  }
  visa::WindowHandler plots;
  try
  {
    plots.addLinePlot("Gaussian");
    plots.get("Gaussian").setLimits(0.0,1.2);
    plots.get("Gaussian").fillVertexArray(gaussian);
    for ( unsigned int i=0;i<10;i++ )
    {
      clog << "Closes in " << 10-i << " sec...\r";
      plots.show();
      this_thread::sleep_for(chrono::seconds(1));
    }
  }
  catch ( exception &exc )
  {
    cout << exc.what() << endl;
    return 1;
  }
  catch (...)
  {
    cout << "Unrecognized exception!\n";
    return 1;
  }
  return 0;
}


#include "max.hpp"
#include "ast.hpp"
#include "parsers.hpp"
#include "types.hpp"
#include <boost/variant.hpp>
#include <stdexcept>
#include <algorithm>
namespace HT
{
    void max(PASTNode astnode, ParsersHelper& ph)
    {
        auto myParserHelper(ph);
        if (astnode->ch.size()<2)
          throw std::runtime_error("Max should have at least one parameter");
        auto & secondCh = *astnode->ch.rbegin();
        ph.parse(secondCh);

        astnode->type = Simple;
        astnode->token.tokenType = Complex;
        astnode->token.info = boost::get<ComplexType>(secondCh->token.info);
        std::for_each( ++astnode->ch.begin(), astnode->ch.end(), [&ph, &astnode](PASTNode an)
                    {
                    ph.parse(an);
                    if (an->token.tokenType != Complex || !boost::get<ComplexType>(an->token.info).isReal())
                    throw std::runtime_error("arguments of Max must be real");
                    auto cast = boost::get<ComplexType>(an->token.info);
                    if (cast.isRational() && boost::get<ComplexType>(astnode->token.info).isRational())
                    {
                    if (cast.getRealR() > boost::get<ComplexType>(astnode->token.info).getRealR())
                    astnode->token.info = cast;
                    }
                    else
                    {
                    auto cast2= boost::get<ComplexType>(astnode->token.info);
                    cast.toinexact();
                    cast2.toinexact();
                    astnode->token.info = ComplexType( std::max(cast.getRealD(), cast2.getRealD()));
                    }
                    });
        astnode->remove();
    }
}