#include "teca_cartesian_mesh_source.h"
#include "teca_dataset.h"
#include "teca_cartesian_mesh.h"
#include "teca_coordinate_util.h"

#include <string>
#include <vector>

#if defined(TECA_HAS_BOOST)
#include <boost/program_options.hpp>
#endif

using std::cerr;
using std::endl;

struct teca_cartesian_mesh_source::internals_t
{
    // given the index space i0 to i1 spanning the world space
    // x0 to x1, generate an equally spaced coordinate axis
    template<typename num_t>
    static
    void initialize_axis(p_teca_variant_array_impl<num_t> x,
        unsigned long i0, unsigned long i1, num_t x0, num_t x1);

    // given an index space [i0 i1 j0 j1 k0 k1 q0 q1] spanning
    // the world space [x0 x1 y0 y1 z0 z1 t0 t1] generate
    // equally spaced coordinate axes x,y,z,t
    static
    void initialize_axes(int type_code, const unsigned long *extent,
        const double *bounds, p_teca_variant_array &x_axis,
        p_teca_variant_array &y_axis, p_teca_variant_array &z_axis,
        p_teca_variant_array &t_axis);

    static
    void initialize_axes(int type_code, const unsigned long *extent,
        const double *bounds, p_teca_variant_array &x_axis,
        p_teca_variant_array &y_axis, p_teca_variant_array &z_axis);

    // cached metadata
    teca_metadata metadata;
    p_teca_variant_array t_axis;
};

// --------------------------------------------------------------------------
template<typename num_t>
void teca_cartesian_mesh_source::internals_t::initialize_axis(
    p_teca_variant_array_impl<num_t> x, unsigned long i0, unsigned long i1,
    num_t x0, num_t x1)
{
    // generate an equally spaced coordinate axes
    unsigned long nx = i1 - i0 + 1;
    x->resize(nx);

    num_t *px = x->get();

    // avoid divide by zero
    if (nx < 2)
    {
        px[0] = x0;
        return;
    }

    num_t dx = (x1 - x0)/(nx - 1l);
    num_t xx = x0 + i0*dx;

    for (unsigned long i = 0; i < nx; ++i)
        px[i] = xx + dx*i;
}

// --------------------------------------------------------------------------
void teca_cartesian_mesh_source::internals_t::initialize_axes(int type_code,
    const unsigned long *extent, const double *bounds, p_teca_variant_array &x_axis,
    p_teca_variant_array &y_axis, p_teca_variant_array &z_axis,
    p_teca_variant_array &t_axis)
{
    // gernate equally spaced coordinate axes x,y,z,t
    x_axis = teca_variant_array_factory::New(type_code);
    y_axis = x_axis->new_instance();
    z_axis = x_axis->new_instance();
    t_axis = x_axis->new_instance();

    TEMPLATE_DISPATCH(teca_variant_array_impl,
        x_axis.get(),

        internals_t::initialize_axis<NT>(std::static_pointer_cast<TT>(x_axis),
            extent[0], extent[1], bounds[0], bounds[1]);

        internals_t::initialize_axis<NT>(std::static_pointer_cast<TT>(y_axis),
            extent[2], extent[3], bounds[2], bounds[3]);

        internals_t::initialize_axis<NT>(std::static_pointer_cast<TT>(z_axis),
            extent[4], extent[5], bounds[4], bounds[5]);

        internals_t::initialize_axis<NT>(std::static_pointer_cast<TT>(t_axis),
            extent[6], extent[7], bounds[6], bounds[7]);
        )
}

// --------------------------------------------------------------------------
void teca_cartesian_mesh_source::internals_t::initialize_axes(int type_code,
    const unsigned long *extent, const double *bounds, p_teca_variant_array &x_axis,
    p_teca_variant_array &y_axis, p_teca_variant_array &z_axis)
{
    // gernate equally spaced coordinate axes x,y,z,t
    x_axis = teca_variant_array_factory::New(type_code);
    y_axis = x_axis->new_instance();
    z_axis = x_axis->new_instance();

    TEMPLATE_DISPATCH(teca_variant_array_impl,
        x_axis.get(),

        internals_t::initialize_axis<NT>(std::static_pointer_cast<TT>(x_axis),
            extent[0], extent[1], bounds[0], bounds[1]);

        internals_t::initialize_axis<NT>(std::static_pointer_cast<TT>(y_axis),
            extent[2], extent[3], bounds[2], bounds[3]);

        internals_t::initialize_axis<NT>(std::static_pointer_cast<TT>(z_axis),
            extent[4], extent[5], bounds[4], bounds[5]);
        )
}

// --------------------------------------------------------------------------
teca_cartesian_mesh_source::teca_cartesian_mesh_source() :
    coordinate_type_code(teca_variant_array_code<double>::get()),
    field_type_code(teca_variant_array_code<double>::get()),
    x_axis_variable("lon"), y_axis_variable("lat"), z_axis_variable("plev"),
    t_axis_variable("time"),
    whole_extents{0l, 359l, 0l, 179l, 0l, 0l, 0l, 0l},
    bounds{0., 360, -90., 90., 0., 0., 0., 0.},
    internals(new internals_t)
{

    this->set_number_of_input_connections(0);
    this->set_number_of_output_ports(1);
}

// --------------------------------------------------------------------------
teca_cartesian_mesh_source::~teca_cartesian_mesh_source()
{
    delete this->internals;
}

#if defined(TECA_HAS_BOOST)
// --------------------------------------------------------------------------
void teca_cartesian_mesh_source::get_properties_description(
    const std::string &prefix, options_description &global_opts)
{
    (void)prefix;
    (void)global_opts;
}

// --------------------------------------------------------------------------
void teca_cartesian_mesh_source::set_properties(const std::string &prefix,
    variables_map &opts)
{
    (void)prefix;
    (void)opts;
}
#endif

// --------------------------------------------------------------------------
void teca_cartesian_mesh_source::set_modified()
{
    // clear cached metadata before forwarding on to
    // the base class.
    this->clear_cached_metadata();
    teca_algorithm::set_modified();
}

// --------------------------------------------------------------------------
void teca_cartesian_mesh_source::clear_cached_metadata()
{
    this->internals->metadata.clear();
    teca_algorithm::set_modified();
}

// --------------------------------------------------------------------------
void teca_cartesian_mesh_source::set_x_axis_variable(const std::string &name)
{
    this->x_axis_variable = name;
    this->x_axis_attributes.clear();
    teca_algorithm::set_modified();
}

// --------------------------------------------------------------------------
void teca_cartesian_mesh_source::set_x_axis_variable(const std::string &name,
    const teca_metadata &atts)
{
    this->x_axis_variable = name;
    this->x_axis_attributes = atts;
    teca_algorithm::set_modified();
}

// --------------------------------------------------------------------------
int teca_cartesian_mesh_source::set_x_axis_variable(const teca_metadata &md)
{
    teca_metadata coords;
    if (md.get("coordinates", coords))
        return -1;

    if (coords.get("x_variable", this->x_axis_variable))
        return -1;

    teca_metadata atts;
    if (md.get("attributes", atts))
        return -1;

    if (atts.get(this->x_axis_variable, this->x_axis_attributes))
        return -1;

    return 0;
}

// --------------------------------------------------------------------------
void teca_cartesian_mesh_source::set_y_axis_variable(const std::string &name)
{
    this->y_axis_variable = name;
    this->y_axis_attributes.clear();
    teca_algorithm::set_modified();
}

// --------------------------------------------------------------------------
void teca_cartesian_mesh_source::set_y_axis_variable(const std::string &name,
    const teca_metadata &atts)
{
    this->y_axis_variable = name;
    this->y_axis_attributes = atts;
    teca_algorithm::set_modified();
}

// --------------------------------------------------------------------------
int teca_cartesian_mesh_source::set_y_axis_variable(const teca_metadata &md)
{
    teca_metadata coords;
    if (md.get("coordinates", coords))
        return -1;

    if (coords.get("y_variable", this->y_axis_variable))
        return -1;

    teca_metadata atts;
    if (md.get("attributes", atts))
        return -1;

    if (atts.get(this->y_axis_variable, this->y_axis_attributes))
        return -1;

    return 0;
}

// --------------------------------------------------------------------------
void teca_cartesian_mesh_source::set_z_axis_variable(const std::string &name)
{
    this->z_axis_variable = name;
    this->z_axis_attributes.clear();
    teca_algorithm::set_modified();
}

// --------------------------------------------------------------------------
void teca_cartesian_mesh_source::set_z_axis_variable(const std::string &name,
    const teca_metadata &atts)
{
    this->z_axis_variable = name;
    this->z_axis_attributes = atts;
    teca_algorithm::set_modified();
}

// --------------------------------------------------------------------------
int teca_cartesian_mesh_source::set_z_axis_variable(const teca_metadata &md)
{
    // get coordinates and attributes, fail if either are missing
    teca_metadata coords;
    if (md.get("coordinates", coords))
        return -1;

    if (coords.get("x_variable", this->z_axis_variable))
        return -1;

    teca_metadata atts;
    if (md.get("attributes", atts))
        return -1;

    if (atts.get(this->z_axis_variable, this->z_axis_attributes))
        return -1;

    return 0;
}

// --------------------------------------------------------------------------
void teca_cartesian_mesh_source::set_t_axis_variable(const std::string &name)
{
    this->t_axis_variable = name;
    this->t_axis_attributes.clear();
    teca_algorithm::set_modified();
}

// --------------------------------------------------------------------------
void teca_cartesian_mesh_source::set_calendar(
    const std::string &calendar, const std::string &units)
{
    this->t_axis_attributes.clear();
    this->t_axis_attributes.set("calendar", calendar);
    this->t_axis_attributes.set("units", units);
    teca_algorithm::set_modified();
}

// --------------------------------------------------------------------------
void teca_cartesian_mesh_source::set_t_axis_variable(const std::string &name,
    const teca_metadata &atts)
{
    this->t_axis_variable = name;
    this->t_axis_attributes = atts;
    teca_algorithm::set_modified();
}

// --------------------------------------------------------------------------
int teca_cartesian_mesh_source::set_t_axis_variable(const teca_metadata &md)
{
    teca_metadata coords;
    if (md.get("coordinates", coords))
        return -1;

    if (coords.get("t_variable", this->t_axis_variable))
        return -1;

    teca_metadata atts;
    if (md.get("attributes", atts))
        return -1;

    if (atts.get(this->t_axis_variable, this->t_axis_attributes))
        return -1;

    return 0;
}

// --------------------------------------------------------------------------
int teca_cartesian_mesh_source::set_t_axis(const teca_metadata &md)
{
    teca_metadata coords;
    if (md.get("coordinates", coords))
        return -1;

    this->internals->t_axis = coords.get("t");

    return 0;
}

// --------------------------------------------------------------------------
void teca_cartesian_mesh_source::set_t_axis(const p_teca_variant_array &t)
{
    this->internals->t_axis = t;
}

// --------------------------------------------------------------------------
int teca_cartesian_mesh_source::set_output_metadata(const teca_metadata &md)
{
    teca_metadata coords;
    if (md.get("coordinates", coords))
        return -1;

    teca_metadata atts;
    if (md.get("attributes", atts))
        return -1;

    // get the coordinate axes.
    const_p_teca_variant_array x = coords.get("x");
    const_p_teca_variant_array y = coords.get("y");
    const_p_teca_variant_array z = coords.get("z");
    const_p_teca_variant_array t = coords.get("t");

    // because of assumptions made in execute, all must be provided
    if (!x || !y || !z || !t)
        return -1;

    unsigned long nx = x->size();
    unsigned long ny = y->size();
    unsigned long nz = z->size();
    unsigned long nxyz = nx*ny*nz;

    // clear out any variables, and replace with those that we provide.
    std::vector<std::string> vars;
    std::vector<field_generator_t>::iterator it = this->field_generators.begin();
    std::vector<field_generator_t>::iterator end = this->field_generators.end();
    for (; it != end; ++it)
    {
        vars.push_back(it->name);

        // correct size
        teca_metadata var_atts = it->attributes;
        var_atts.set("size", nxyz);

        atts.set(it->name, var_atts);
    }

    // copy the metadata
    this->set_modified();

    this->internals->metadata = md;
    this->internals->metadata.set("variables", vars);

    return 0;
}

// --------------------------------------------------------------------------
int teca_cartesian_mesh_source::set_spatial_extents(const teca_metadata &md,
    bool three_d)
{
    // get coordinates and attributes, fail if either are missing
    teca_metadata coords;
    if (md.get("coordinates", coords))
        return -1;

    teca_metadata attributes;
    if (md.get("attributes", attributes))
        return -1;

    // get the coordinate axes
    p_teca_variant_array x = coords.get("x");
    p_teca_variant_array y = coords.get("y");
    p_teca_variant_array z = coords.get("z");

    // verify
    if (!x || !y || (three_d && !z))
        return -1;

    // set the extents
    this->whole_extents[0] = 0;
    this->whole_extents[1] = x->size() - 1;
    this->whole_extents[2] = 0;
    this->whole_extents[3] = y->size() - 1;
    this->whole_extents[4] = 0;
    this->whole_extents[5] = three_d ? z->size() - 1 : 0;

    return 0;
}
// --------------------------------------------------------------------------
int teca_cartesian_mesh_source::set_spatial_bounds(const teca_metadata &md,
    bool three_d)
{
    // get coordinates and attributes, fail if either are missing
    teca_metadata coords;
    if (md.get("coordinates", coords))
        return -1;

    teca_metadata attributes;
    if (md.get("attributes", attributes))
        return -1;

    // get the coordinate axes
    p_teca_variant_array x = coords.get("x");
    p_teca_variant_array y = coords.get("y");
    p_teca_variant_array z = coords.get("z");

    // verify
    if (!x || !y || (three_d && !z))
        return -1;

    // get the bounds
    x->get(0lu, this->bounds[0]);
    x->get(x->size() - 1lu, this->bounds[1]);
    y->get(0lu, this->bounds[2]);
    y->get(y->size() - 1lu, this->bounds[3]);
    z->get(0lu, this->bounds[4]);

    unsigned long khi = three_d ? z->size() - 1lu : 0lu;
    z->get(khi, this->bounds[5]);

    // set the coordinate type
    this->set_coordinate_type_code(x->type_code());

    return 0;
}

// --------------------------------------------------------------------------
void teca_cartesian_mesh_source::append_field_generator(
    const std::string &name, const teca_metadata &atts,
    field_generator_callback &callback)
{
    this->append_field_generator({name, atts, callback});
}

// --------------------------------------------------------------------------
teca_metadata teca_cartesian_mesh_source::get_output_metadata(
    unsigned int port, const std::vector<teca_metadata> &input_md)
{
#ifdef TECA_DEBUG
    cerr << teca_parallel_id()
        << "teca_cartesian_mesh_source::get_output_metadata" << endl;
#endif
    (void)port;
    (void)input_md;
    if (this->internals->metadata)
        return this->internals->metadata;

    if (this->whole_extents.size() != 8)
    {
        TECA_FATAL_ERROR("invalid whole extents were specified")
        return teca_metadata();
    }

    if (this->bounds.size() != 8)
    {
        TECA_FATAL_ERROR("invalid bounds were specified")
        return teca_metadata();
    }

    // generate cooridnate axes
    p_teca_variant_array x_axis, y_axis, z_axis, t_axis;

    if (this->internals->t_axis)
    {
        // generate x,y,z axes but use cached time axis
        internals_t::initialize_axes(this->coordinate_type_code,
            this->whole_extents.data(), this->bounds.data(), x_axis,
            y_axis, z_axis);

        t_axis = this->internals->t_axis;
    }
    else
    {
        // generate x,y,z and t axes
        internals_t::initialize_axes(this->coordinate_type_code,
            this->whole_extents.data(), this->bounds.data(), x_axis,
            y_axis, z_axis, t_axis);
    }

    size_t nx = x_axis->size();
    size_t ny = y_axis->size();
    size_t nz = z_axis->size();
    size_t nt = t_axis->size();

    // construct attributes
    teca_metadata x_atts = this->x_axis_attributes;
    x_atts.set("type_code", x_axis->type_code());
    x_atts.set("size", nx);

    teca_metadata y_atts = this->y_axis_attributes;
    y_atts.set("type_code", y_axis->type_code());
    y_atts.set("size", ny);

    teca_metadata z_atts = this->z_axis_attributes;
    z_atts.set("type_code", z_axis->type_code());
    z_atts.set("size", nz);

    teca_metadata t_atts = this->t_axis_attributes;
    t_atts.set("type_code", t_axis->type_code());
    t_atts.set("size", nt);

    teca_metadata atts;
    atts.set(this->x_axis_variable, x_atts);
    atts.set(this->y_axis_variable, y_atts);
    atts.set(this->z_axis_variable, z_atts);

    if (!this->t_axis_variable.empty())
        atts.set(this->t_axis_variable, t_atts);

    // construct dataset metadata
    teca_metadata coords;
    coords.set("x_variable", this->x_axis_variable);
    coords.set("y_variable", this->y_axis_variable);
    coords.set("z_variable", this->z_axis_variable);
    coords.set("t_variable", this->t_axis_variable);

    coords.set("x", x_axis);
    coords.set("y", y_axis);
    coords.set("z", z_axis);
    coords.set("t", t_axis);

    this->internals->metadata.set("whole_extent", this->whole_extents);
    this->internals->metadata.set("coordinates", coords);

    size_t nxyz = nx*ny*nz;
    std::vector<std::string> vars;
    std::vector<field_generator_t>::iterator it = this->field_generators.begin();
    std::vector<field_generator_t>::iterator end = this->field_generators.end();
    for (; it != end; ++it)
    {
        vars.push_back(it->name);

        // correct size
        teca_metadata var_atts = it->attributes;
        var_atts.set("size", nxyz);

        atts.set(it->name, var_atts);
    }

    this->internals->metadata.set("variables", vars);
    this->internals->metadata.set("attributes", atts);

    // setup the execution control keys
    this->internals->metadata.set("number_of_time_steps",
        t_axis->size());

    this->internals->metadata.set("index_initializer_key",
        std::string("number_of_time_steps"));

    this->internals->metadata.set("index_request_key",
        std::string("time_step"));

    return this->internals->metadata;
}

// --------------------------------------------------------------------------
const_p_teca_dataset teca_cartesian_mesh_source::execute(unsigned int port,
    const std::vector<const_p_teca_dataset> &input_data,
    const teca_metadata &request)
{
#ifdef TECA_DEBUG
    cerr << teca_parallel_id()
        << "teca_cartesian_mesh_source::execute" << endl;
#endif
    (void)port;
    (void)input_data;

    // get coordinates
    teca_metadata coords;
    if (this->internals->metadata.get("coordinates", coords))
    {
        TECA_FATAL_ERROR("metadata is missing \"coordinates\"")
        return nullptr;
    }

    p_teca_variant_array in_x, in_y, in_z, in_t;
    if (!(in_x = coords.get("x")) || !(in_y = coords.get("y"))
        || !(in_z = coords.get("z")) || !(in_t = coords.get("t")))
    {
        TECA_FATAL_ERROR("metadata is missing coordinate arrays")
        return nullptr;
    }

    // get the extent of the dataset we could generate
    unsigned long md_whole_extent[6] = {0};
    if (this->internals->metadata.get("whole_extent", md_whole_extent, 6))
    {
        TECA_FATAL_ERROR("metadata is missing \"whole_extent\"")
        return nullptr;
    }

    // determine the subset we are being requested, fall back to the
    // whole extent if a subset isn't specified
    unsigned long req_extent[6] = {0};
    double req_bounds[6] = {0.0};
    if (request.get("bounds", req_bounds, 6))
    {
        // bounds key not present, check for extent key
        // if not present use whole_extent
        if (request.get("extent", req_extent, 6))
        {
            memcpy(req_extent, md_whole_extent, 6*sizeof(unsigned long));
        }
    }
    else
    {
        // bounds key was present, convert the bounds to an
        // an extent that covers them.
        if (teca_coordinate_util::bounds_to_extent(req_bounds,
                in_x, in_y, in_z, req_extent) ||
            teca_coordinate_util::validate_extent(in_x->size(),
                in_y->size(), in_z->size(), req_extent, true))
        {
            TECA_FATAL_ERROR("invalid bounds requested.")
            return nullptr;
        }
    }

    // get the timestep, no matter what the key is named we treat it as
    // a time step. this is to support metadata provided by another source
    // eg. a different reader.
    std::string request_key;
    if (request.get("index_request_key", request_key))
    {
        TECA_FATAL_ERROR("Request is missing the \"index_request_key\"")
        return nullptr;
    }

    unsigned long req_index = 0;
    if (request.get(request_key, req_index))
    {
        TECA_FATAL_ERROR("Request is missing \"" << request_key << "\"")
        return nullptr;
    }

    // check that the we have a time value for the requested index.
    if (req_index >= in_t->size())
    {
        TECA_FATAL_ERROR("The requested index " << req_index
            << " is out of bounds [0, " << in_t->size() << "]")
        return nullptr;
    }

    // get the time
    double t = 0.;
    in_t->get(req_index, t);

    // slice axes on the requested extent
    p_teca_variant_array out_x = in_x->new_copy(req_extent[0], req_extent[1]);
    p_teca_variant_array out_y = in_y->new_copy(req_extent[2], req_extent[3]);
    p_teca_variant_array out_z = in_z->new_copy(req_extent[4], req_extent[5]);

    // create output dataset
    p_teca_cartesian_mesh mesh = teca_cartesian_mesh::New();

    std::string x_variable = this->x_axis_variable.empty() ? "x" : this->x_axis_variable;
    std::string y_variable = this->y_axis_variable.empty() ? "y" : this->y_axis_variable;
    std::string z_variable = this->z_axis_variable.empty() ? "z" : this->z_axis_variable;

    mesh->set_x_coordinates(x_variable, out_x);
    mesh->set_y_coordinates(y_variable, out_y);
    mesh->set_z_coordinates(z_variable, out_z);

    // get the calendar
    std::string calendar;
    std::string units;
    teca_metadata atts;
    this->internals->metadata.get("attributes", atts);
    atts.get("calendar", calendar);
    atts.get("units", units);

    // set metadata
    mesh->set_whole_extent(md_whole_extent);
    mesh->set_extent(req_extent);
    mesh->set_time_step(req_index);
    mesh->set_time(t);
    mesh->set_calendar(calendar);
    mesh->set_time_units(units);

    teca_metadata &mesh_md = mesh->get_metadata();
    mesh_md.set("index_request_key", request_key);

    // generate fields over the requested subset
    std::vector<field_generator_t>::iterator it = this->field_generators.begin();
    std::vector<field_generator_t>::iterator end = this->field_generators.end();
    for (; it != end; ++it)
    {
        p_teca_variant_array f_xyzt = it->generator(out_x, out_y, out_z, t);
        mesh->get_point_arrays()->append(it->name, f_xyzt);
    }

    // pass the attributes
    mesh_md.set("attributes", atts);

    return mesh;
}


/**
 * @file tests/main_tests/pca_test.cpp
 * @author Ryan Curtin
 *
 * Test mlpackMain() of pca_main.cpp.
 *
 * mlpack is free software; you may redistribute it and/or modify it under the
 * terms of the 3-clause BSD license.  You should have received a copy of the
 * 3-clause BSD license along with mlpack.  If not, see
 * http://www.opensource.org/licenses/BSD-3-Clause for more information.
 */
#include <string>

#define BINDING_TYPE BINDING_TYPE_TEST
static const std::string testName = "PrincipalComponentAnalysis";

#include <mlpack/core.hpp>
#include <mlpack/core/util/mlpack_main.hpp>
#include "test_helper.hpp"
#include <mlpack/methods/pca/pca_main.cpp>

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

using namespace mlpack;

struct PCATestFixture
{
 public:
  PCATestFixture()
  {
    // Cache in the options for this program.
    CLI::RestoreSettings(testName);
  }

  ~PCATestFixture()
  {
    // Clear the settings.
    bindings::tests::CleanMemory();
    CLI::ClearSettings();
  }
};

BOOST_FIXTURE_TEST_SUITE(PCAMainTest, PCATestFixture);

/**
 * Make sure that if we ask for a dataset in 3 dimensions back, we get it.
 */
BOOST_AUTO_TEST_CASE(PCADimensionTest)
{
  arma::mat x = arma::randu<arma::mat>(5, 5);

  // Random input, new dimensionality of 3.
  SetInputParam("input", std::move(x));
  SetInputParam("new_dimensionality", (int) 3);

  mlpackMain();

  // Now check that the output has 3 dimensions.
  BOOST_REQUIRE_EQUAL(CLI::GetParam<arma::mat>("output").n_rows, 3);
  BOOST_REQUIRE_EQUAL(CLI::GetParam<arma::mat>("output").n_cols, 5);
}

/**
 * Ensure that if we retain all variance, we get back a matrix with the same
 * dimensionality.
 */
BOOST_AUTO_TEST_CASE(PCAVarRetainTest)
{
  arma::mat x = arma::randu<arma::mat>(4, 5);

  SetInputParam("input", std::move(x));
  SetInputParam("var_to_retain", (double) 1.0);
  SetInputParam("scale", true);
  SetInputParam("new_dimensionality", (int) 3); // Should be ignored.

  mlpackMain();

  // Check that the output has 5 dimensions.
  BOOST_REQUIRE_EQUAL(CLI::GetParam<arma::mat>("output").n_rows, 4);
  BOOST_REQUIRE_EQUAL(CLI::GetParam<arma::mat>("output").n_cols, 5);
}

/**
 * Ensure that if we retain no variance, we get back no dimensions.
 */
BOOST_AUTO_TEST_CASE(PCANoVarRetainTest)
{
  arma::mat x = arma::randu<arma::mat>(5, 5);

  SetInputParam("input", std::move(x));
  SetInputParam("var_to_retain", (double) 0.01);
  SetInputParam("scale", true);
  SetInputParam("new_dimensionality", (int) 3); // Should be ignored.

  mlpackMain();

  // Check that the output has 1 dimensions.
  BOOST_REQUIRE_EQUAL(CLI::GetParam<arma::mat>("output").n_rows, 1);
  BOOST_REQUIRE_EQUAL(CLI::GetParam<arma::mat>("output").n_cols, 5);
}

/**
 * Check that we can't specify an invalid new dimensionality.
 */
BOOST_AUTO_TEST_CASE(PCATooHighNewDimensionalityTest)
{
  arma::mat x = arma::randu<arma::mat>(5, 5);

  SetInputParam("input", std::move(x));
  SetInputParam("new_dimensionality", (int) 7); // Invalid.

  Log::Fatal.ignoreInput = true;
  BOOST_REQUIRE_THROW(mlpackMain(), std::runtime_error);
  Log::Fatal.ignoreInput = false;
}

BOOST_AUTO_TEST_SUITE_END();


/*
@copyright Louis Dionne 2014
Distributed under the Boost Software License, Version 1.0.
(See accompanying file LICENSE.md or copy at http://boost.org/LICENSE_1_0.txt)
 */

#include <boost/hana/ext/std/tuple.hpp>

#include <boost/hana/detail/assert.hpp>
#include <boost/hana/detail/constexpr.hpp>
#include <boost/hana/integral.hpp>

#include <tuple>
using namespace boost::hana;

BOOST_HANA_CONSTEXPR_LAMBDA auto tuple = [](auto ...xs) {
    return std::make_tuple(xs...);
};

BOOST_HANA_CONSTEXPR_LAMBDA auto f = [](auto ...xs) {
    return std::make_tuple(1, xs...);
};

BOOST_HANA_CONSTEXPR_LAMBDA auto g = [](auto ...xs) {
    return std::make_tuple(2, xs...);
};

int main() {
    BOOST_HANA_CONSTEXPR_ASSERT(lift<StdTuple>(int_<0>) == tuple(int_<0>));
    BOOST_HANA_CONSTEXPR_ASSERT(lift<StdTuple>(int_<1>) == tuple(int_<1>));
    BOOST_HANA_CONSTEXPR_ASSERT(lift<StdTuple>(3) == tuple(3));

    BOOST_HANA_CONSTEXPR_ASSERT(ap(tuple(), tuple()) == tuple());
    BOOST_HANA_CONSTEXPR_ASSERT(ap(tuple(), tuple(1)) == tuple());
    BOOST_HANA_CONSTEXPR_ASSERT(ap(tuple(), tuple(1, '2')) == tuple());
    BOOST_HANA_CONSTEXPR_ASSERT(ap(tuple(), tuple(1, '2', 3.3)) == tuple());

    BOOST_HANA_CONSTEXPR_ASSERT(ap(tuple(f), tuple()) == tuple());
    BOOST_HANA_CONSTEXPR_ASSERT(ap(tuple(f), tuple(1)) == tuple(f(1)));
    BOOST_HANA_CONSTEXPR_ASSERT(ap(tuple(f), tuple(1, '2')) == tuple(f(1), f('2')));
    BOOST_HANA_CONSTEXPR_ASSERT(ap(tuple(f), tuple(1, '2', 3.3)) == tuple(f(1), f('2'), f(3.3)));

    BOOST_HANA_CONSTEXPR_ASSERT(ap(tuple(f, g), tuple()) == tuple());
    BOOST_HANA_CONSTEXPR_ASSERT(ap(tuple(f, g), tuple(1)) == tuple(f(1), g(1)));
    BOOST_HANA_CONSTEXPR_ASSERT(ap(tuple(f, g), tuple(1, '2')) == tuple(f(1), f('2'), g(1), g('2')));
    BOOST_HANA_CONSTEXPR_ASSERT(ap(tuple(f, g), tuple(1, '2', 3.3)) == tuple(f(1), f('2'), f(3.3), g(1), g('2'), g(3.3)));
}


#include <numeric>

#include "Defaults.hpp"
#include "Container.hpp"

#include <iostream>
#include <ctime>
#include <array>

#include <boost/multi_array.hpp>

bool equal(const double& a, const double& b)
{
  return a == b;
}

template <typename ArrayA, typename ArrayB>
bool equal(const ArrayA& A, const ArrayB& B)
{
  typename ArrayA::const_iterator ia;
  typename ArrayB::const_iterator ib = B.begin();
  for (ia = A.begin(); ia != A.end(); ++ia, ++ib)
    if (!equal(*ia, *ib))
      return false;
  return true;
}

using namespace GFTools;

int main()
{
    std::cout << "Hi!" << std::endl;
    std::array<size_t,1> Ar1 {{5}};
    std::array<size_t,2> Ar2 {{3,3}};
    std::array<size_t,3> Ar3 {{2,3,3}};

    INFO("Container")

    Container<ComplexType,2> A1(Ar2), A2(Ar2), A3(Ar2);
    A1 = 1.;
    A2 = 2.;
    DEBUG(A1);
    A1+=A2;
    DEBUG(A1);
    A3 = 3.0*(A1+0.5)*2+A2;
    DEBUG(A3);
    DEBUG(A1);
    auto a01 = 3*(A1[1]+A1[0])*2;
    DEBUG(a01);
    DEBUG(A1);

    Container<ComplexType,1> E1(A1[1]), E2(A1[2]);
    E1*=3;
    DEBUG(E1);
    DEBUG(A1);

    Container<ComplexType,2> B(Ar2);
    Container<ComplexType,2> C(Ar2);
    Container<ComplexType,3> D(Ar3);

    B[0][2]=3.0;
    C[0][2]=-2.0;

    std::array<size_t, 2> coord1 {{0,1}};
    DEBUG(B._data(coord1));

    Container<ComplexType,2>::iterator t1;
    auto it2 = B.begin();
    auto it3 = B[0].begin();

    DEBUG(B[0][2]);
    DEBUG(B[0][0]);

    //DEBUG(B._data.shape()[0]);
    //DEBUG(B[0]._data.shape()[0]);

    DEBUG(B);
    DEBUG(B[0]);

    D[0][0][0]=-1.0;
    D[0][0][1]=1.0;
    DEBUG(D);
    DEBUG(D._data.num_elements());

    DEBUG(C);
    C=B;
    DEBUG(C);

    C[0][0]=1.0;
    C[0][2]=-5;
    DEBUG(C);
    C+=B;
    DEBUG(C);
    C+=1.0;
    DEBUG(C);
    DEBUG(C+B);
    Container<ComplexType,3> E(D);

    DEBUG(B);
    B+=2.0;
    DEBUG(B);
    DEBUG(B[0]+C[0]);
    DEBUG((B+C)[0]);
    DEBUG((B-C)[0]);

    B*=4.;
    DEBUG(B);
    DEBUG(B[0]*2.0);
    DEBUG(D+2.0);
    DEBUG(D[0]+B);
    //DEBUG((B[0]*3.0));
    E*=(-1);
    DEBUG(D);
    DEBUG(E);
    DEBUG(D+E);
    DEBUG(D*5+2.0);

    INFO("Matrix test");
    INFO("2-dim");
    std::array<size_t,2> dim1 {{4,5}};
    Container<RealType,2> Vals (dim1);
    Vals[1][0] = 1.0;
    Vals[0][3] = -0.4;
    Vals[2][1] = 1.5;
    INFO(Vals);
    INFO(Vals.getAsMatrix());
    DEBUG("!");
    Container<RealType,2> Vals_2(Vals.getAsMatrix());
    INFO(Vals_2);

    decltype(Vals) Vals_3(dim1);
    Vals_3 = Vals.getAsMatrix();
    INFO(Vals_3);

    if ((Vals-Vals_2).sum()!=0 || (Vals-Vals_3).sum()!=0) return EXIT_FAILURE;
    INFO("1-dim test");
    Container<RealType,1> Vals_1d(std::array<size_t,1>({{5}}));
    Vals_1d[2]=1.6;
    Vals_1d[4]=3;
    INFO(Vals_1d);
    INFO(Vals_1d.getAsVector());
    INFO(Vals_1d.getAsDiagonalMatrix());

    INFO("View test");
    Container<RealType, 3> C1(std::array<size_t,3>({{2,1,3}}));
    C1[0][0][1] = 1.0; C1[0][0][2] = 2.0; C1[1][0][0] = 3.0; C1[1][0][2] = 4.0;
    INFO_NONEWLINE("Data in memory: "); for (size_t i = 0; i<C1._data.num_elements(); ++i) { INFO_NONEWLINE(*(C1._data.origin()+i) << " "); }; INFO("");
    INFO("Container in : " << C1);
    typedef boost::multi_array<RealType, 3>::array_view<3>::type myview_3_t;
    typedef boost::multi_array_types::index_range range;
    myview_3_t boost_view1 = C1._data[boost::indices[range(0,1)][range(0,1)][range(0,2)]];
    ContainerView<RealType, 3, 3> C1_v2(myview_3_t(C1._data[boost::indices[range().stride(2)][range()][range().stride(1)]]));
    ContainerView<RealType, 3, 3> C1_v3(C1_v2);
    INFO(C1_v2);
    DEBUG(C1_v2*C1_v3);
    DEBUG(C1);
    C1_v2 *= C1_v3;
    DEBUG(C1_v2);
    DEBUG(C1);
    ContainerView<RealType, 3> C1_v(C1._data[boost::indices[range()][range()][range().stride(3)]]);
    INFO("Container view : " << C1_v);

    return EXIT_SUCCESS;
}


//  Copyright (c) 2006 Xiaogang Zhang, 2015 John Maddock.
//  Use, modification and distribution are subject to the
//  Boost Software License, Version 1.0. (See accompanying file
//  LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)
//
//  History:
//  XZ wrote the original of this file as part of the Google
//  Summer of Code 2006.  JM modified it slightly to fit into the
//  Boost.Math conceptual framework better.
//  Updated 2015 to use Carlson's latest methods.

#ifndef BOOST_MATH_ELLINT_RD_HPP
#define BOOST_MATH_ELLINT_RD_HPP

#ifdef _MSC_VER
#pragma once
#endif

#include <boost/math/special_functions/math_fwd.hpp>
#include <boost/math/special_functions/ellint_rc.hpp>
#include <boost/math/special_functions/pow.hpp>
#include <boost/math/tools/config.hpp>
#include <boost/math/policies/error_handling.hpp>

// Carlson's elliptic integral of the second kind
// R_D(x, y, z) = R_J(x, y, z, z) = 1.5 * \int_{0}^{\infty} [(t+x)(t+y)]^{-1/2} (t+z)^{-3/2} dt
// Carlson, Numerische Mathematik, vol 33, 1 (1979)

namespace boost { namespace math { namespace detail{

template <typename T, typename Policy>
T ellint_rd_imp(T x, T y, T z, const Policy& pol)
{
   BOOST_MATH_STD_USING
   using std::swap;

   static const char* function = "boost::math::ellint_rd<%1%>(%1%,%1%,%1%)";

   if(x < 0)
   {
      return policies::raise_domain_error<T>(function,
         "Argument x must be >= 0, but got %1%", x, pol);
   }
   if(y < 0)
   {
      return policies::raise_domain_error<T>(function,
         "Argument y must be >= 0, but got %1%", y, pol);
   }
   if(z <= 0)
   {
      return policies::raise_domain_error<T>(function,
         "Argument z must be > 0, but got %1%", z, pol);
   }
   if(x + y == 0)
   {
      return policies::raise_domain_error<T>(function,
         "At most one argument can be zero, but got, x + y = %1%", x + y, pol);
   }
   //
   // Special cases from http://dlmf.nist.gov/19.20#iv
   //
   using std::swap;
   if(x == z)
      swap(x, y);
   if(y == z)
   {
      if(x == y)
      {
         return 1 / (x * sqrt(x));
      }
      else if(x == 0)
      {
         return 3 * constants::pi<T>() / (4 * y * sqrt(y));
      }
      else
      {
         if((std::min)(x, y) / (std::max)(x, y) > 1.3)
            return 3 * (ellint_rc_imp(x, y, pol) - sqrt(x) / y) / (2 * (y - x));
         // Otherwise fall through to avoid cancellation in the above (RC(x,y) -> 1/x^0.5 as x -> y)
      }
   }
   if(x == y)
   {
      if((std::min)(x, z) / (std::max)(x, z) > 1.3)
         return 3 * (ellint_rc_imp(z, x, pol) - 1 / sqrt(z)) / (z - x);
      // Otherwise fall through to avoid cancellation in the above (RC(x,y) -> 1/x^0.5 as x -> y)
   }
   if(y == 0)
      swap(x, y);
   if(x == 0)
   {
      //
      // Special handling for common case, from
      // Numerical Computation of Real or Complex Elliptic Integrals, eq.47
      //
      T xn = sqrt(y);
      T yn = sqrt(z);
      T x0 = xn;
      T y0 = yn;
      T sum = 0;
      T sum_pow = 0.25f;

      while(fabs(xn - yn) >= 2.7 * tools::root_epsilon<T>() * fabs(xn))
      {
         T t = sqrt(xn * yn);
         xn = (xn + yn) / 2;
         yn = t;
         sum_pow *= 2;
         sum += sum_pow * boost::math::pow<2>(xn - yn);
      }
      T RF = constants::pi<T>() / (xn + yn);
      //
      // This following calculation suffers from serious cancellation when y ~ z
      // unless we combine terms.  We have:
      //
      // ( ((x0 + y0)/2)^2 - z ) / (z(y-z))
      //
      // Substituting y = x0^2 and z = y0^2 and simplifying we get the following:
      //
      T pt = (x0 + 3 * y0) / (4 * z * (x0 + y0));
      //
      // Since we've moved the demoninator from eq.47 inside the expression, we
      // need to also scale "sum" by the same value:
      //
      pt -= sum / (z * (y - z));
      return pt * RF * 3;
   }

   T xn = x;
   T yn = y;
   T zn = z;
   T An = (x + y + 3 * z) / 5;
   T A0 = An;
   // This has an extra 1.2 fudge factor which is really only needed when x, y and z are close in magnitude:
   T Q = pow(tools::epsilon<T>() / 4, -T(1) / 8) * (std::max)((std::max)(An - x, An - y), An - z) * 1.2f;
   BOOST_MATH_INSTRUMENT_VARIABLE(Q);
   T lambda, rx, ry, rz;
   unsigned k = 0;
   T fn = 1;
   T RD_sum = 0;

   for(; k < policies::get_max_series_iterations<Policy>(); ++k)
   {
      rx = sqrt(xn);
      ry = sqrt(yn);
      rz = sqrt(zn);
      lambda = rx * ry + rx * rz + ry * rz;
      RD_sum += fn / (rz * (zn + lambda));
      An = (An + lambda) / 4;
      xn = (xn + lambda) / 4;
      yn = (yn + lambda) / 4;
      zn = (zn + lambda) / 4;
      fn /= 4;
      Q /= 4;
      BOOST_MATH_INSTRUMENT_VARIABLE(k);
      BOOST_MATH_INSTRUMENT_VARIABLE(RD_sum);
      BOOST_MATH_INSTRUMENT_VARIABLE(Q);
      if(Q < An)
         break;
   }

   policies::check_series_iterations<T, Policy>(function, k, pol);

   T X = fn * (A0 - x) / An;
   T Y = fn * (A0 - y) / An;
   T Z = -(X + Y) / 3;
   T E2 = X * Y - 6 * Z * Z;
   T E3 = (3 * X * Y - 8 * Z * Z) * Z;
   T E4 = 3 * (X * Y - Z * Z) * Z * Z;
   T E5 = X * Y * Z * Z * Z;

   T result = fn * pow(An, T(-3) / 2) *
      (1 - 3 * E2 / 14 + E3 / 6 + 9 * E2 * E2 / 88 - 3 * E4 / 22 - 9 * E2 * E3 / 52 + 3 * E5 / 26 - E2 * E2 * E2 / 16
      + 3 * E3 * E3 / 40 + 3 * E2 * E4 / 20 + 45 * E2 * E2 * E3 / 272 - 9 * (E3 * E4 + E2 * E5) / 68);
   BOOST_MATH_INSTRUMENT_VARIABLE(result);
   result += 3 * RD_sum;

   return result;
}

} // namespace detail

template <class T1, class T2, class T3, class Policy>
inline typename tools::promote_args<T1, T2, T3>::type
   ellint_rd(T1 x, T2 y, T3 z, const Policy& pol)
{
   typedef typename tools::promote_args<T1, T2, T3>::type result_type;
   typedef typename policies::evaluation<result_type, Policy>::type value_type;
   return policies::checked_narrowing_cast<result_type, Policy>(
      detail::ellint_rd_imp(
         static_cast<value_type>(x),
         static_cast<value_type>(y),
         static_cast<value_type>(z), pol), "boost::math::ellint_rd<%1%>(%1%,%1%,%1%)");
}

template <class T1, class T2, class T3>
inline typename tools::promote_args<T1, T2, T3>::type
   ellint_rd(T1 x, T2 y, T3 z)
{
   return ellint_rd(x, y, z, policies::policy<>());
}

}} // namespaces

#endif // BOOST_MATH_ELLINT_RD_HPP


#include "parameters.hpp"
#include "eigen_utils.hpp"
#include "mpi_utils.hpp"
#include "eigen_utils.hpp"
#include "rhe_reg_component.hpp"

#include "tools/eigen3.3/Dense"

#include <boost/iostreams/filtering_stream.hpp>

#include <random>

RHEreg_Component::RHEreg_Component(const parameters &myparams,
                                   const Eigen::VectorXd &myY,
                                   const Eigen::MatrixXd &myWzz,
                                   const Eigen::MatrixXd &myC,
                                   const Eigen::MatrixXd &myCtC_inv,
                                   const long &myNJacknifeLocal) : params(myparams), Y(myY),
	zz(myWzz), C(myC), CtC_inv(myCtC_inv), n_jacknife_local(myNJacknifeLocal) {
	assert(n_jacknife_local > 0);
	n_covar = C.cols();
	n_samples = zz.rows();
	n_draws = zz.cols();

	label = "";
	is_gxe = false;
	is_active = true;
	is_finalised = false;
	rm_jacknife_block = -1;

	ytXXty = 0;
	ytXXtys.resize(n_jacknife_local, 0);

	n_var_local = 0;
	n_vars_local.resize(n_jacknife_local, 0);

	_XXtzs.resize(n_jacknife_local);
	for(long ii = 0; ii < n_jacknife_local; ii++) {
		_XXtzs[ii] = Eigen::MatrixXd::Zero(n_samples, n_draws);
	}

	if(params.debug) {
		std::string ram = mpiUtils::currentUsageRAM();
		std::cout << "(" << ram << ")" << std::endl;
	}
}

void RHEreg_Component::set_env_var(const Eigen::Ref<const Eigen::VectorXd>& my_env_var) {
	assert(is_active);
	is_gxe = true;
	env_var = my_env_var;
	Y.array() *= env_var.array();
	if(zz.rows() > 0) {
		zz.array().colwise() *= env_var.array();
	}
}

void RHEreg_Component::change_env_var(const Eigen::Ref<const Eigen::VectorXd>& new_env_var) {
	assert(is_active);
	assert(is_gxe);
	Y.array() /= env_var.array();
	Y.array() *= new_env_var.array();
	if(zz.rows() > 0) {
		zz.array().colwise() /= env_var.array();
		zz.array().colwise() *= new_env_var.array();
	}
	env_var = new_env_var;
}

void RHEreg_Component::set_inactive() {
	// The inactive component corresponds to sigma_e
	// Ie the 'noise' component
	assert(!is_gxe);
	is_active = false;
	_XXtWz = zz;
	n_var_local = 1;
	ytXXty = mpiUtils::mpiReduce_inplace(Y.squaredNorm());
}

void RHEreg_Component::add_to_trace_estimator(Eigen::Ref <Eigen::MatrixXd> X,
                                              long jacknife_index) {
	assert(jacknife_index < n_jacknife_local);
	if(is_active) {
		Eigen::MatrixXd Xty = X.transpose() * Y;
		Xty = mpiUtils::mpiReduce_inplace(Xty);
		ytXXtys[jacknife_index] += Xty.squaredNorm();
		if(n_covar > 0) {
			Eigen::MatrixXd XtWz = X.transpose() * zz;
			XtWz = mpiUtils::mpiReduce_inplace(XtWz);
			_XXtzs[jacknife_index] += X * XtWz;
		}
		n_vars_local[jacknife_index] += X.cols();
	}
}

void RHEreg_Component::finalise() {
	// Sum over the different jacknife blocks;
	if(is_active) {
		_XXtWz = Eigen::MatrixXd::Zero(n_samples, n_draws);
		for (auto& mm : _XXtzs) {
			_XXtWz += mm;
		}

		n_var_local = std::accumulate(n_vars_local.begin(), n_vars_local.end(), 0.0);
		ytXXty = std::accumulate(ytXXtys.begin(), ytXXtys.end(), 0.0);
		is_finalised = true;
	}
}

Eigen::MatrixXd RHEreg_Component::getXXtz() const {
	Eigen::MatrixXd res;
	if(rm_jacknife_block >= 0) {
		res = _XXtWz - _XXtzs[rm_jacknife_block];
	} else {
		res = _XXtWz;
	}
	if(is_gxe) {
		res.array().colwise() *= env_var.array();
	}
	return res;
}

double RHEreg_Component::get_bb_trace() const {
	if(rm_jacknife_block >= 0) {
		return (ytXXty - ytXXtys[rm_jacknife_block]) / get_n_var_local();
	} else {
		return ytXXty / get_n_var_local();
	}
}

double RHEreg_Component::get_n_var_local() const {
	if(rm_jacknife_block >= 0) {
		return n_var_local - n_vars_local[rm_jacknife_block];
	} else {
		return n_var_local;
	}
}

double RHEreg_Component::operator*(const RHEreg_Component &other) const {
	double res;
	if(n_covar == 0) {
		res = getXXtz().cwiseProduct(other.getXXtz()).sum();
	} else if (n_covar > 0) {
		if(label == "noise" || other.label == "noise") {
			res = getXXtz().cwiseProduct(other.getXXtz()).sum();
		} else {
			Eigen::MatrixXd XXtWz = getXXtz();
			Eigen::MatrixXd WXXtWz = project_out_covars(XXtWz);
			res = WXXtWz.cwiseProduct(other.getXXtz()).sum();
		}
	} else {
		throw std::runtime_error("Error in PVE_Component");
	}
	res = mpiUtils::mpiReduce_inplace(res);
	return res / get_n_var_local() / other.get_n_var_local() / (double) n_draws;
}

Eigen::MatrixXd RHEreg_Component::project_out_covars(Eigen::Ref <Eigen::MatrixXd> rhs) const {
	return EigenUtils::project_out_covars(rhs, C, CtC_inv);
}

void get_GxE_collapsed_component(const std::vector<RHEreg_Component> &vec_of_components,
                                 RHEreg_Component &new_comp,
                                 const Eigen::Ref<const Eigen::MatrixXd> &E,
                                 const Eigen::Ref<const Eigen::VectorXd> &env_weights,
                                 const std::vector<Eigen::MatrixXd> &ytEXXtEy,
                                 const bool& copy_jacknife_partitions) {
	long n_components = vec_of_components.size();
	long n_samples = vec_of_components[0].n_samples;
	long n_draws = vec_of_components[0].n_draws;
	long n_jacknife = new_comp.params.n_jacknife;

	Eigen::VectorXd eta = E * env_weights;
	Eigen::MatrixXd sum_WlVbl = Eigen::MatrixXd::Zero(n_samples, n_draws);
	long GxE_comp_index;
	for (int ii = 0; ii < n_components; ii++) {
		if (vec_of_components[ii].effect_type == "GxE") {
			assert(vec_of_components[ii].rm_jacknife_block == -1);
			GxE_comp_index = ii;
			long ll = vec_of_components[ii].env_var_index;
			sum_WlVbl += env_weights[ll] * E.col(ll).asDiagonal().inverse() * vec_of_components[ii].getXXtz();
		}
	}
	new_comp._XXtWz = sum_WlVbl;
	new_comp.ytXXty = 0;
	for (long jj = 0; jj < n_jacknife; jj++) {
		new_comp.ytXXty += env_weights.transpose() * ytEXXtEy[jj] * env_weights;
	}
	new_comp.is_finalised = true;

	if(copy_jacknife_partitions) {
		new_comp.n_vars_local.resize(n_jacknife);
		new_comp.ytXXtys.resize(n_jacknife);
		new_comp._XXtzs.resize(n_jacknife);
		for(long ii = 0; ii < n_jacknife; ii++) {
			new_comp._XXtzs[ii] = Eigen::MatrixXd::Zero(n_samples, n_draws);
		}

		for(long jj = 0; jj < n_jacknife; jj++) {
			for (int ii = 0; ii < n_components; ii++) {
				if (vec_of_components[ii].effect_type == "GxE") {
					long ll = vec_of_components[ii].env_var_index;
					new_comp._XXtzs[jj] += env_weights[ll] * vec_of_components[ii]._XXtzs[jj];
				}
			}
			new_comp.n_vars_local[jj] = vec_of_components[0].n_vars_local[jj];
			new_comp.ytXXtys[jj] = env_weights.transpose() * ytEXXtEy[jj] * env_weights;
		}
	}

	new_comp.set_env_var(eta);
	new_comp.n_var_local = vec_of_components[GxE_comp_index].n_var_local;
	new_comp.n_draws = vec_of_components[GxE_comp_index].n_draws;
	new_comp.n_covar = vec_of_components[GxE_comp_index].n_covar;
	new_comp.Y = vec_of_components[GxE_comp_index].Y;
	new_comp.C = vec_of_components[GxE_comp_index].C;
	new_comp.CtC_inv = vec_of_components[GxE_comp_index].CtC_inv;
	new_comp.label = "GxE";
	new_comp.effect_type = "GxE";
}

void get_GxE_collapsed_system(const std::vector<RHEreg_Component> &vec_of_components,
                              std::vector<RHEreg_Component> &new_components,
                              const Eigen::Ref<const Eigen::MatrixXd> &E,
                              const Eigen::Ref<const Eigen::VectorXd> &env_weights,
                              const std::vector<Eigen::MatrixXd> &ytEXXtEy,
                              const bool& copy_jacknife_partitions){

	long n_components = vec_of_components.size();
	long n_samples = vec_of_components[0].n_samples;
	long n_draws = vec_of_components[0].n_draws;

	// Create component for \diag{Ew} \left( \sum_l w_l \bm{v}_{b, l} \right)
	Eigen::MatrixXd placeholder = Eigen::MatrixXd::Zero(n_samples, n_draws);
	RHEreg_Component combined_comp(vec_of_components[0].params, vec_of_components[0].Y,
	                               vec_of_components[0].C, vec_of_components[0].CtC_inv,
	                               placeholder);
	get_GxE_collapsed_component(vec_of_components, combined_comp, E, env_weights, ytEXXtEy,
	                            copy_jacknife_partitions);

	// Get main and noise components
	new_components.clear();
	new_components.reserve(3);
	bool transfered_gxe = false;
	for (int ii = 0; ii < n_components; ii++) {
		if(vec_of_components[ii].effect_type != "GxE") {
			new_components.push_back(vec_of_components[ii]);
		} else if (!transfered_gxe) {
			transfered_gxe = true;
			new_components.push_back(std::move(combined_comp));
		}
	}
}


#include <iostream>

#include <boost/array.hpp>
#include <boost/numeric/odeint.hpp>

#include "Dynamics.h"

using std::cout;
using std::endl;

using namespace boost::numeric::odeint;

Dynamics::Dynamics() : PI(3.14159265359), rng_(new boost::mt19937()),
                       nd_(0, 1.0), var_nor(rng_, nd_),
                       process_noise_(0), measurement_noise_(0)
{
     for (int i = 0; i < 5; i++) {
          u_[i] = 0;
          x0_[i] = 0;
          state_[i] = 0;
     }

     rng_->seed((++seed_) + time(NULL));
}

std::vector<double> Dynamics::time_vector(double t0, double dt, double tend)
{
     std::vector<double> tt;
     for( double t = t0; t < tend ; t += dt ) {
          tt.push_back(t);
     }
     return tt;
}

void Dynamics::set_time(double t0, double dt, double tend)
{
     t0_ = t0;
     dt_ = dt;
     tend_ = tend;

     tt_.clear();
     for( double t = t0_ ; t < tend_ ; t += dt_ ) {
          tt_.push_back(t);
     }
}

void Dynamics::set_x0(state_5d_type x0)
{
     for(int i = 0; i < 5; i++) {
          x0_[i] = x0[i];
          state_[i] = x0[i];
          measurement_[i] = 0;
     }

     measurement_[0] = x0[0] + var_nor() * measurement_noise_;
     measurement_[1] = x0[1] + var_nor() * measurement_noise_;
}

void Dynamics::compute_trajectory()
{
     headings_.clear();
     truth_points_.clear();

     if (model_ == cart || model_ == roomba) {
          state_5d_type x = {x0_[0], x0_[1], x0_[2], x0_[3], x0_[4]};
          runge_kutta4< state_5d_type > stepper;

          std::vector<double>::iterator it = tt_.begin();
          for (; it != tt_.end(); it++) {
               Eigen::Vector2d truth, measured;

               truth = Eigen::Vector2d(x[0], x[1]);

               measured(0,0) = truth(0,0) + var_nor() * measurement_noise_;
               measured(1,0) = truth(1,0) + var_nor() * measurement_noise_;

               headings_.push_back(x[2]*180.0/PI);
               truth_points_.push_back(truth);
               measured_points_.push_back(measured);

               state_5d_type state;
               state[0] = x[0];
               state[1] = x[1];
               state[2] = x[2];
               state[3] = 0;
               state[4] = 0;
               states_.push_back(state);

               if (model_ == cart) {
                    stepper.do_step(make_ode_wrapper( *this , &Dynamics::cart_model ), x , *it , dt_ );
               } else {
                    stepper.do_step(make_ode_wrapper( *this , &Dynamics::roomba_model ), x , *it , dt_ );
               }
          }
     } else if (model_ == constant_velocity) {
     }
}

void Dynamics::step_motion_model(double dt, double time)
{
     if (input_sequence_.size() > 0 && input_sequence_.front().time <= time) {
          cout << "Using input: " << endl;
          for (int i = 0; i < 5; i++) {
               u_[i] = input_sequence_.front().input[i];

               printf("%d : %f\n", i, u_[i]);
          }
          input_sequence_.erase(input_sequence_.begin());
     }

     if (model_ == cart || model_ == roomba) {
          if (model_ == cart) {
               stepper_.do_step(make_ode_wrapper( *this , &Dynamics::cart_model ), state_ , 0 , dt );
          } else {
               stepper_.do_step(make_ode_wrapper( *this , &Dynamics::roomba_model ), state_ , 0 , dt );
          }

          Eigen::Vector2d truth, measured;
          truth = Eigen::Vector2d(state_[0], state_[1]);

          measured(0,0) = truth(0,0) + var_nor() * measurement_noise_;
          measured(1,0) = truth(1,0) + var_nor() * measurement_noise_;

          measurement_[0] = measured(0,0);
          measurement_[1] = measured(1,0);
          measurement_[2] = 0;
          measurement_[3] = 0;
          measurement_[4] = 0;

          headings_.push_back(state_[2]*180.0/PI);
          truth_points_.push_back(truth);
          measured_points_.push_back(measured);

     } else if (model_ == constant_velocity) {
     }
}

void Dynamics::set_input(state_5d_type input)
{
     input_sequence_.clear();

     ControlInput control_input;
     control_input.time = 0;
     for(int i = 0; i < 5; i++) {
          control_input.input[i] = input[i];
     }

     input_sequence_.push_back(control_input);
}

void Dynamics::set_input(state_5d_type input, double time)
{
     ControlInput control_input;
     control_input.time = time;
     for(int i = 0; i < 5; i++) {
          control_input.input[i] = input[i];
     }

     input_sequence_.push_back(control_input);
}

void Dynamics::set_input_sequence(std::vector<ControlInput> &input_sequence)
{
     input_sequence_.clear();

     for (std::vector<ControlInput>::iterator it = input_sequence.begin();
          it != input_sequence.end(); it++) {

          ControlInput control_input;
          control_input.time = it->time;
          for(int i = 0; i < 5; i++) {
               control_input.input[i] = it->input[i];
          }
          input_sequence_.push_back(control_input);
     }
}

void Dynamics::cart_model(const state_5d_type &x , state_5d_type &dxdt , double t)
{
     /// 0 : x-position
     /// 1 : y-position
     /// 2 : theta

     //u_[0] = 2;
     //u_[1] = 3.14159265359/20;

     double u = u_[0];
     double u_theta = u_[1];
     double L = 3;

     dxdt[0] = u*cos(x[2]) + var_nor() * process_noise_;
     dxdt[1] = u*sin(x[2]) + var_nor() * process_noise_;
     dxdt[2] = u/L*tan(u_theta) + var_nor() * process_noise_;
     dxdt[3] = 0;
     dxdt[4] = 0;
}

void Dynamics::roomba_model(const state_5d_type &x , state_5d_type &dxdt , double t)
{
     /// 0 : x-position
     /// 1 : y-position
     /// 2 : theta

     double R = 1;
     double b = 0.5;

     double w_l = u_[0];
     double w_r = u_[1];

     dxdt[0] = R*cos(x[2])*(w_r + w_l)/2;
     dxdt[1] = R*sin(x[2])*(w_r + w_l)/2;
     dxdt[2] = R*(w_r-w_l)/(2*b);
     dxdt[3] = 0;
     dxdt[4] = 0;
}


/* Copyright (C) 2012-2017 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 <NTL/BasicThreadPool.h>

#include "recryption.h"
#include "EncryptedArray.h"
#include "EvalMap.h"
#include "powerful.h"
#include "CtPtrs.h"
#include "intraSlot.h"
#include "norms.h"
#include "sample.h"
#include "debugging.h"

NTL_CLIENT

#ifdef DEBUG_PRINTOUT
#include "debugging.h"
long printFlag = FLAG_PRINT_VEC;
#endif

/************************ Some local functions ***********************/
/*********************************************************************/
static void
printSizesPowerful(const vector<ZZX>& zzParts, const DoubleCRT& sKey,
                   const RecryptData& rcData, long q, double noise);

// Return in poly a polynomial with X^i encoded in all the slots
static void x2iInSlots(ZZX& poly, long i,
		       vector<ZZX>& xVec, const EncryptedArray& ea)
{
  xVec.resize(ea.size());
  ZZX x2i = ZZX(i,1);
  for (long j=0; j<(long)xVec.size(); j++) xVec[j] = x2i;
  ea.encode(poly, xVec);
}

// Make every entry of vec divisible by p2e by adding/subtracting
// multiples of p2r and q, while keeping the added multiples small.
// Specifically, for q = 1 mod p2e and any a < p2e/(2*p2r), any
// integer z can be made divisible by p2e via z' = z + u*p2r + v*q,
// with |u|*p2r <= a and |v| <= p2e/2 -a.
// Returns the largest absolute values of the u's and the new entries.

#if 1
static long makeDivisible(vec_ZZ& vec, long p2e, long p2r, long q, long a,
                          double& U_norm, const PAlgebra& palg, PowerfulDCRT* p2dConv = 0)
{
  assert(q>0 && p2e>0 && p2r>0 && a>=0
         && q % p2e == 1 && a % p2r == 0 && a*2 < p2e);
  long aa = a / p2r;

#ifdef DEBUG_PRINTOUT
  zzX uVec(INIT_SIZE, vec.length());
  zzX vVec(INIT_SIZE, vec.length());
#endif

  long maxU = 0;
  for (long i=0; i<vec.length(); i++) {
    ZZ& z = vec[i];
    long u, v;

    // What to add to z to make it divisible by p2e?
    long zMod = rem(z, p2e); // zMod is in [0,p2e-1]
    if (zMod > p2e/2) { // need to add a positive number
      zMod = p2e - zMod;
      u = zMod/p2r;
      if (u > aa) u = aa;
    }
    else {              // need to add a negative number
      u = -(zMod/p2r);
      if (u < -aa) u = -aa;
      zMod = -zMod;
    }
    v = zMod - u*p2r;
    z += u*p2r + to_ZZ(q)*v; // make z divisible by p2e

    if (rem(z,p2e) != 0) { // sanity check
      cerr << "**error: original z["<<i<<"]=" << (z-(u*p2r+to_ZZ(q)*v))
	   << std::dec << ", p^r="<<p2r << ", p^e="<<p2e << endl;
      cerr << "z' = z + "<<u<<"*p^r +"<<v<<"*q = "<<z<<endl;
      exit(1);
    }
    if (abs(u) > maxU) maxU = abs(u);
#ifdef DEBUG_PRINTOUT
    uVec[i] = u;
    vVec[i] = v;
#endif
  }

#ifdef DEBUG_PRINTOUT
  if (dbgEa) {
    const PAlgebra& palg = dbgEa->getPAlgebra();
    double V_norm;
    if (p2dConv) {
       ZZX poly;
       p2dConv->powerfulToZZX(poly, conv<Vec<ZZ>>(uVec));
       U_norm = conv<double>(embeddingLargestCoeff(poly, palg)) *p2r;
       p2dConv->powerfulToZZX(poly, conv<Vec<ZZ>>(vVec));
       V_norm = conv<double>(embeddingLargestCoeff(poly, palg));
    }
    else {
      U_norm = conv<double>(embeddingLargestCoeff(uVec, palg)) *p2r;
      V_norm = conv<double>(embeddingLargestCoeff(vVec, palg));
    }
    cerr << "  makeDivisible: maxU=" << (maxU*p2r)
         << ", U_norm=" << U_norm
         << ", V_norm=" << V_norm << endl;
  }
#endif

  return maxU;
}
#else
// experimental, randomized version
static long makeDivisible(vec_ZZ& vec, long p2e, long p2r, long q, long a,
                          double& U_norm, const PAlgebra& palg)
{
  assert(q>0 && p2e>0 && p2r>0 && a>=0
         && q % p2e == 1 && a % p2r == 0 && a*2 < p2e);
  long aa = a / p2r;

  vec_ZZ orig_vec = vec;
  double sum_norm;
  long maxU;

  zzX uVec(INIT_SIZE, vec.length());
  zzX vVec(INIT_SIZE, vec.length());

  for (long trials = 0; trials < 10; trials++) {
    vec_ZZ try_vec = orig_vec;

    long maxU1 = 0;
    for (long i=0; i<vec.length(); i++) {
      ZZ& z = try_vec[i];
      long u, v;

      // What to add to z to make it divisible by p2e?
      long zMod = rem(z, p2e); // zMod is in [0,p2e-1]
      if (zMod > p2e/2) { // need to add a positive number
	zMod = p2e - zMod;
	u = zMod/p2r;
	if (u > aa) u = aa;
        if (u > 0) {
          long ran = RandomBnd(u+1); //0..u
          u -= ran;
        }
      }
      else {              // need to add a negative number
	u = -(zMod/p2r);
	if (u < -aa) u = -aa;
        if (-u > 0) {
          long ran = RandomBnd(-u+1); //0..|u|
          u += ran;
        }
	zMod = -zMod;
      }
      v = zMod - u*p2r;
      z += u*p2r + to_ZZ(q)*v; // make z divisible by p2e

      if (rem(z,p2e) != 0) { // sanity check
	cerr << "**error: original z["<<i<<"]=" << (z-(u*p2r+to_ZZ(q)*v))
	     << std::dec << ", p^r="<<p2r << ", p^e="<<p2e << endl;
	cerr << "z' = z + "<<u<<"*p^r +"<<v<<"*q = "<<z<<endl;
	exit(1);
      }
      if (abs(u) > maxU1) maxU1 = abs(u);

      uVec[i] = u;
      vVec[i] = v;
    }

    double U_norm1 = conv<double>(embeddingLargestCoeff(uVec, palg)) *p2r;
    double V_norm1 = conv<double>(embeddingLargestCoeff(vVec, palg));
    double sum_norm1 = U_norm1 + V_norm1;
#ifdef DEBUG_PRINTOUT
      cerr << "  makeDivisible: maxU1=" << (maxU1*p2r)
	   << ", U_norm1=" << U_norm1
	   << ", V_norm1=" << V_norm1
           << ", sum=" << (U_norm1+V_norm1)
           << "\n";
#endif

    if (trials == 0 || sum_norm1 < sum_norm) {
      sum_norm = sum_norm1;
      maxU = maxU1;
      U_norm = U_norm1;
      vec = try_vec;
    }
  }

  return maxU;
}

#endif

static inline double pow(long a, long b) {return pow(double(a), double(b));}
/*********************************************************************/
/*********************************************************************/

RecryptData::~RecryptData()
{
  if (alMod!=NULL)     delete alMod;
  if (ea!=NULL)        delete ea;
  if (firstMap!=NULL)  delete firstMap;
  if (secondMap!=NULL) delete secondMap;
  if (p2dConv!=NULL)   delete p2dConv;
}

/**
 * Fix the "ring constant" cM, a target norm tau for the secret key,
 * and plaintext space mod p^r. We want to find e,e' that minimize
 * e-e', subject to the constraint
 *
 *    (1) (p^{e'}/2 + 2(p^r+1))(tau+1)*cM <= (q-1)/2  = p^e/2
 *
 * Note that as we let e,e' tend to infinity the constraint above
 * degenerates to (tau+1)*cM < p^{e-e'}, so the smallest value
 * of e-e' that we can hope for is
 *
 *    (2) e-e' = 1 + floor( log_p( (tau+1)*cM) )
 *
 * The setAE procedure tries to minimize e-e' subject to (1), and
 * in addition subject to the constraint that e is "not too big".
 * Specifically, it tries to ensure p^e<2^{30}, and failing that it
 * uses the smallest e for which 2(p^r+1)(tau+1)*cM*2 <= p^e, and the
 * largest e' for that value of e.
 *
 * Once e,e' are set, it splits p^{e'}/2=a+b with a,b about equal and
 * a divisible by p^r. Then it computes and returns the largest Hamming
 * weight for the key (that implies the norm tau') for which constraint
 * (1) still holds.
 *
 * NOTE: setAE returns the Hamming weight, *not* the norm tau'. The norm
 * can be computed from the weight using sampleHWtBoundedEffectiveBound.
 **/
long RecryptData::setAE(long& a, long& e, long& ePrime,
                    const FHEcontext& context, long targetWeight)
{
  bool default_target=false;
  if (targetWeight<=0) {
    targetWeight = RecryptData::defSkHwt;
    default_target=true;
  }

  double skSize = sampleHWtBoundedEffectiveBound(context, targetWeight);

  double magicConst = context.zMStar.get_cM();
  double factor = (skSize+1) * magicConst * 2;
  long p = context.zMStar.getP();
  long p2r = context.alMod.getPPowR();
  long frstTerm = 2*(p2r+1);
  double logp = log(p);    // log p

  // Start with the smallest e s.t. p^e > (2p^r+1)(skSize+1)*magicCons
  ePrime = 0;
  e = ceil( log(frstTerm*factor) /logp );
  assert(e*logp < log(NTL_SP_BOUND));
         // p^e for this smallest e must be single precision

  // Loop to try and find better solutions, that still satisfy
  //          (p^{e'}/2 + 2*p^r+1) * (s+1)*magicConst*2 <= p^e
  //          <=> p^{e'}/2 + 2*p^r +1 <= p^e/factor
  //          <=> p^{e'} <=  2*(p^e/factor - 2*p^r -1)
  // stop when you hit p^e > 2^30
  long eTry = e;
  long eMinusEprime = e; // want to minimize e-e'
  do {
    double p2e = pow(p,eTry);

    // Solve for the largest e' satisfying constraints (1)
    // p^{e'} <=  2*(p^e/factor - 2*p^r -1)
    long ePrimeTry = floor( log(2*(p2e/factor -frstTerm)) / logp);
    if (eTry - ePrimeTry < eMinusEprime) {
      e = eTry;
      ePrime = ePrimeTry;
      eMinusEprime = e - ePrime;
    }
  } while ((++eTry)*log(p) <= log(1L<<30));

  // Split p^{e'}/2 into a+b with a divisible by p^r
  long pToEprimeOver2 = floor(pow(p,ePrime)/2.0);
  a = 2*pToEprimeOver2 /5; // empirically we want a little below a half
  a -= (a % p2r);

  if (default_target) {
    // Try to increase t, while maintaining constraint (1)
    double bound = (frstTerm + pToEprimeOver2)*magicConst * 2;
    long p2e = pow(p,e);
    while (targetWeight++) {
      skSize = sampleHWtBoundedEffectiveBound(context, targetWeight);
      if (targetWeight>256 || bound*(skSize+1) >= p2e)
	break;
    }
    targetWeight--;
  }

#ifdef DEBUG_PRINTOUT
  long b = pToEprimeOver2 - a;
  cerr << "RecryptData::setAE(): e="<<e<<", e'="<<ePrime
       << ", a="<<a<<", b="<<b<<", sk-hwt="<<targetWeight
       << " (size="<<skSize<<"), magicConst="<<magicConst<<endl;
#endif
  return targetWeight;
}

bool RecryptData::operator==(const RecryptData& other) const
{
  if (mvec != other.mvec) return false;
  if (skHwt != other.skHwt) return false;

  return true;
}

// The main method
void RecryptData::init(const FHEcontext& context, const Vec<long>& mvec_,
                  bool enableThick, long t, bool build_cache_, bool minimal)
{
  if (alMod != NULL) { // were we called for a second time?
    cerr << "@Warning: multiple calls to RecryptData::init\n";
    return;
  }
  assert(computeProd(mvec_) == (long)context.zMStar.getM()); // sanity check

  // Record the arguments to this function
  mvec = mvec_;
  build_cache = build_cache_;

  skHwt = setAE(a, e, ePrime, context, t);
  long p = context.zMStar.getP();
  long r = context.alMod.getR();

  // First part of Bootstrapping works wrt plaintext space p^{r'}
  alMod = new PAlgebraMod(context.zMStar, e-ePrime+r);
  ea = new EncryptedArray(context, *alMod);
         // Polynomial defaults to F0, PAlgebraMod explicitly given

  p2dConv = new PowerfulDCRT(context, mvec);

  if (!enableThick) return;

  // Initialize the linear polynomial for unpacking the slots
  zz_pBak bak; bak.save(); ea->getAlMod().restoreContext();
  long nslots = ea->size();
  long d = ea->getDegree();

  const Mat<zz_p>& CBi=ea->getDerived(PA_zz_p()).getNormalBasisMatrixInverse();

  vector<ZZX> LM;
  LM.resize(d);
  for (long i = 0; i < d; i++) // prepare the linear polynomial
    LM[i] = rep(CBi[i][0]);

  vector<ZZX> C;
  ea->buildLinPolyCoeffs(C, LM); // "build" the linear polynomial

  unpackSlotEncoding.resize(d);  // encode the coefficients

  for (long j = 0; j < d; j++) {
    vector<ZZX> v(nslots);
    for (long k = 0; k < nslots; k++) v[k] = C[j];
    ea->encode(unpackSlotEncoding[j], v);
  }
  firstMap = new EvalMap(*ea, minimal, mvec, true, build_cache);
  secondMap = new EvalMap(*context.ea, minimal, mvec, false, build_cache);
}

/********************************************************************/
/********************************************************************/

// Extract digits from fully packed slots
void extractDigitsPacked(Ctxt& ctxt, long botHigh, long r, long ePrime,
			 const vector<ZZX>& unpackSlotEncoding);

// Extract digits from unpacked slots
void extractDigitsThin(Ctxt& ctxt, long botHigh, long r, long ePrime);

// bootstrap a ciphertext to reduce noise
void FHEPubKey::reCrypt(Ctxt &ctxt)
{
  FHE_TIMER_START;

  // Some sanity checks for dummy ciphertext
  long ptxtSpace = ctxt.getPtxtSpace();
  if (ctxt.isEmpty()) return;
  if (ctxt.parts.size()==1 && ctxt.parts[0].skHandle.isOne()) {
    // Dummy encryption, just ensure that it is reduced mod p
    ZZX poly = to_ZZX(ctxt.parts[0]);
    for (long i=0; i<poly.rep.length(); i++)
      poly[i] = to_ZZ( rem(poly[i],ptxtSpace) );
    poly.normalize();
    ctxt.DummyEncrypt(poly);
    return;
  }

  assert(recryptKeyID>=0); // check that we have bootstrapping data

  long p = getContext().zMStar.getP();
  long r = getContext().alMod.getR();
  long p2r = getContext().alMod.getPPowR();

  long intFactor = ctxt.intFactor;

  // the bootstrapping key is encrypted relative to plaintext space p^{e-e'+r}.
  const RecryptData& rcData = getContext().rcData;
  long e = rcData.e;
  long ePrime = rcData.ePrime;
  long p2ePrime = power_long(p,ePrime);
  long q = power_long(p,e)+1;
  assert(e>=r);

#ifdef DEBUG_PRINTOUT
  cerr << "reCrypt: p="<<p<<", r="<<r<<", e="<<e<<" ePrime="<<ePrime
       << ", q="<<q<<endl;
  CheckCtxt(ctxt, "init");
#endif

  // can only bootstrap ciphertext with plaintext-space dividing p^r
  assert(p2r % ptxtSpace == 0);

  ctxt.dropSmallAndSpecialPrimes();

#ifdef DEBUG_PRINTOUT
  CheckCtxt(ctxt, "after mod down");
#endif

  FHE_NTIMER_START(AAA_preProcess);

  // Make sure that this ciphertxt is in canonical form
  if (!ctxt.inCanonicalForm()) ctxt.reLinearize();

  // Mod-switch down if needed
  IndexSet s = ctxt.getPrimeSet() / context.specialPrimes;
  assert(s <= context.ctxtPrimes);
  if (s.card()>2) { // leave only bottom two primes
    long frst = s.first();
    long scnd = s.next(frst);
    IndexSet s2(frst,scnd);
    s.retain(s2); // retain only first two primes
  }
  ctxt.modDownToSet(s);

  // key-switch to the bootstrapping key
  ctxt.reLinearize(recryptKeyID);

  // "raw mod-switch" to the bootstrapping mosulus q=p^e+1.
  vector<ZZX> zzParts; // the mod-switched parts, in ZZX format
  double noise = ctxt.rawModSwitch(zzParts, q);

#ifdef DEBUG_PRINTOUT
  if (dbgKey) {
    cerr << "  before makeDivisible (recryption modulus q="<<q
         << "), noise_bnd=" << noise<<endl;
    printSizesPowerful(zzParts, dbgKey->sKeys[recryptKeyID],
                       ctxt.getContext().rcData, q, noise);
  }
#endif

  // Add multiples of p2r and q to make the zzParts divisible by p^{e'}
  long maxU=0;
  double maxU_norm = 0;
  for (long i=0; i<(long)zzParts.size(); i++) {
    // make divisible by p^{e'}
    double U_norm;
    long newMax = makeDivisible(zzParts[i].rep, p2ePrime, p2r, q,
				rcData.a, U_norm, context.zMStar);
    zzParts[i].normalize();   // normalize after working directly on the rep
    if (maxU < newMax)  maxU = newMax;
    if (maxU_norm < U_norm)  maxU_norm = U_norm;
  }
#ifdef DEBUG_PRINTOUT
  double newNoise = noise + maxU_norm*p2r*(skBounds[recryptKeyID]+1);
  cerr << "  after makeDivisible, maxU=" << maxU
       << ", maxU_norm="<<maxU_norm<<", p2r="<<p2r
       << ", noise_bnd="<<newNoise<<", sk_bnd="<< skBounds[recryptKeyID]
       << endl;
   if (dbgKey)
     printSizesPowerful(zzParts, dbgKey->sKeys[recryptKeyID],
                        ctxt.getContext().rcData, q, newNoise);
#endif

  for (long i=0; i<(long)zzParts.size(); i++)
    zzParts[i] /= p2ePrime;   // divide by p^{e'}

  // Multiply the post-processed cipehrtext by the encrypted sKey

  double p0size = to_double(embeddingLargestCoeff(zzParts[0], context.zMStar));
  double p1size = to_double(embeddingLargestCoeff(zzParts[1], context.zMStar));
  // FIXME: This might be slow without Armadillo

  // NOTE: here we lose the intFactor associated with ctxt.
  // We will restore it below.
  ctxt = recryptEkey;

  ctxt.multByConstant(zzParts[1]);
  ctxt.addConstant(zzParts[0]);

#ifdef DEBUG_PRINTOUT
  CheckCtxt(ctxt, "after preProcess");
#endif
  FHE_NTIMER_STOP(AAA_preProcess);

  // Move the powerful-basis coefficients to the plaintext slots
  FHE_NTIMER_START(AAA_LinearTransform1);
  ctxt.getContext().rcData.firstMap->apply(ctxt);
  FHE_NTIMER_STOP(AAA_LinearTransform1);

#ifdef DEBUG_PRINTOUT
  CheckCtxt(ctxt, "after LinearTransform1");
#endif

  // Extract the digits e-e'+r-1,...,e-e' (from fully packed slots)
  FHE_NTIMER_START(AAA_extractDigitsPacked);
  extractDigitsPacked(ctxt, e-ePrime, r, ePrime,
		      context.rcData.unpackSlotEncoding);
  FHE_NTIMER_STOP(AAA_extractDigitsPacked);

#ifdef DEBUG_PRINTOUT
  CheckCtxt(ctxt, "after extractDigitsPacked");
#endif

  // Move the slots back to powerful-basis coefficients
  FHE_NTIMER_START(AAA_LinearTransform2);
  ctxt.getContext().rcData.secondMap->apply(ctxt);
  FHE_NTIMER_STOP(AAA_LinearTransform2);

#ifdef DEBUG_PRINTOUT
  CheckCtxt(ctxt, "after linearTransform2");
#endif

  // restore intFactor
  if (intFactor != 1)
    ctxt.intFactor = MulMod(ctxt.intFactor, intFactor, ptxtSpace);
}

#ifdef FHE_BOOT_THREADS

// Extract digits from fully packed slots, multithreaded version
void extractDigitsPacked(Ctxt& ctxt, long botHigh, long r, long ePrime,
			 const vector<ZZX>& unpackSlotEncoding)
{
  FHE_TIMER_START;

  // Step 1: unpack the slots of ctxt
  FHE_NTIMER_START(unpack);
  ctxt.cleanUp();

  // Apply the d automorphisms and store them in scratch area
  long d = ctxt.getContext().zMStar.getOrdP();

  vector<Ctxt> unpacked(d, Ctxt(ZeroCtxtLike, ctxt));
  { // explicit scope to force all temporaries to be released
    vector< shared_ptr<DoubleCRT> > coeff_vector;
    coeff_vector.resize(d);

    FHE_NTIMER_START(unpack1);
    for (long i = 0; i < d; i++)
      coeff_vector[i] = shared_ptr<DoubleCRT>(new
        DoubleCRT(unpackSlotEncoding[i], ctxt.getContext(), ctxt.getPrimeSet()) );
    FHE_NTIMER_STOP(unpack1);

    FHE_NTIMER_START(unpack2);
    vector<Ctxt> frob(d, Ctxt(ZeroCtxtLike, ctxt));

    NTL_EXEC_RANGE(d, first, last)
    // FIXME: implement using hoisting!
        for (long j = first; j < last; j++) { // process jth Frobenius
          frob[j] = ctxt;
          frob[j].frobeniusAutomorph(j);
          frob[j].cleanUp();
          // FIXME: not clear if we should call cleanUp here
        }
    NTL_EXEC_RANGE_END

    FHE_NTIMER_STOP(unpack2);

    FHE_NTIMER_START(unpack3);
    Ctxt tmp1(ZeroCtxtLike, ctxt);
    for (long i = 0; i < d; i++) {
      for (long j = 0; j < d; j++) {
        tmp1 = frob[j];
        tmp1.multByConstant(*coeff_vector[mcMod(i+j, d)]);
        unpacked[i] += tmp1;
      }
    }
    FHE_NTIMER_STOP(unpack3);
  }
  FHE_NTIMER_STOP(unpack);

  //#ifdef DEBUG_PRINTOUT
  //  CheckCtxt(unpacked[0], "after unpack");
  //#endif

  NTL_EXEC_RANGE(d, first, last)
  for (long i = first; i < last; i++) {
    extractDigitsThin(unpacked[i], botHigh, r, ePrime);
  }
  NTL_EXEC_RANGE_END

  //#ifdef DEBUG_PRINTOUT
  //CheckCtxt(unpacked[0], "before repack");
  //#endif

  // Step 3: re-pack the slots
  FHE_NTIMER_START(repack);
  const EncryptedArray& ea2 = *ctxt.getContext().ea;
  ZZX xInSlots;
  vector<ZZX> xVec(ea2.size());
  ctxt = unpacked[0];
  for (long i=1; i<d; i++) {
    x2iInSlots(xInSlots, i, xVec, ea2);
    unpacked[i].multByConstant(xInSlots);
    ctxt += unpacked[i];
  }
  FHE_NTIMER_STOP(repack);
  //#ifdef DEBUG_PRINTOUT
  //CheckCtxt(ctxt, "after repack");
  //#endif
}

#else

// Extract digits from fully packed slots
void extractDigitsPacked(Ctxt& ctxt, long botHigh, long r, long ePrime,
			 const vector<ZZX>& unpackSlotEncoding)
{
  FHE_TIMER_START;

  // Step 1: unpack the slots of ctxt
  FHE_NTIMER_START(unpack);
  ctxt.cleanUp();

  // Apply the d automorphisms and store them in scratch area
  long d = ctxt.getContext().zMStar.getOrdP();

  vector<Ctxt> unpacked(d, Ctxt(ZeroCtxtLike, ctxt));
  { // explicit scope to force all temporaries to be released
    vector< shared_ptr<DoubleCRT> > coeff_vector;
    coeff_vector.resize(d);
    for (long i = 0; i < d; i++)
      coeff_vector[i] = shared_ptr<DoubleCRT>(new
        DoubleCRT(unpackSlotEncoding[i], ctxt.getContext(), ctxt.getPrimeSet()) );
    Ctxt tmp1(ZeroCtxtLike, ctxt);
    Ctxt tmp2(ZeroCtxtLike, ctxt);

    // FIXME: implement using hoisting!
    for (long j = 0; j < d; j++) { // process jth Frobenius
      tmp1 = ctxt;
      tmp1.frobeniusAutomorph(j);
      tmp1.cleanUp();
      // FIXME: not clear if we should call cleanUp here

      for (long i = 0; i < d; i++) {
        tmp2 = tmp1;
        tmp2.multByConstant(*coeff_vector[mcMod(i+j, d)]);
        unpacked[i] += tmp2;
      }
    }
  }
  FHE_NTIMER_STOP(unpack);

  //#ifdef DEBUG_PRINTOUT
  //  CheckCtxt(unpacked[0], "after unpack");
  //#endif

  for (long i=0; i<(long)unpacked.size(); i++) {
    extractDigitsThin(unpacked[i], botHigh, r, ePrime);
  }

  //#ifdef DEBUG_PRINTOUT
  //  CheckCtxt(unpacked[0], "before repack");
  //#endif

  // Step 3: re-pack the slots
  FHE_NTIMER_START(repack);
  const EncryptedArray& ea2 = *ctxt.getContext().ea;
  ZZX xInSlots;
  vector<ZZX> xVec(ea2.size());
  ctxt = unpacked[0];
  for (long i=1; i<d; i++) {
    x2iInSlots(xInSlots, i, xVec, ea2);
    unpacked[i].multByConstant(xInSlots);
    ctxt += unpacked[i];
  }
  FHE_NTIMER_STOP(repack);
}

#endif

// Use packed bootstrapping, so we can bootstrap all in just one go.
void packedRecrypt(const CtPtrs& cPtrs,
                   const std::vector<zzX>& unpackConsts,
                   const EncryptedArray& ea)
{
  FHEPubKey& pKey = (FHEPubKey&)cPtrs[0]->getPubKey();

  // Allocate temporary ciphertexts for the recryption
  int nPacked = divc(cPtrs.size(), ea.getDegree()); // ceil(totoalNum/d)
  std::vector<Ctxt> cts(nPacked, Ctxt(pKey));

  repack(CtPtrs_vectorCt(cts), cPtrs, ea);  // pack ciphertexts
  //  cout << "@"<< lsize(cts)<<std::flush;
  for (Ctxt& c: cts) {     // then recrypt them
    c.reducePtxtSpace(2);  // we only have recryption data for binary ctxt
    pKey.reCrypt(c);
  }
  unpack(cPtrs, CtPtrs_vectorCt(cts), ea, unpackConsts);
}

// recrypt all ctxt at level < belowLvl
void packedRecrypt(const CtPtrs& array,
                   const std::vector<zzX>& unpackConsts,
                   const EncryptedArray& ea, long belowLvl)
{
  std::vector<Ctxt*> v;
  for (long i=0; i<array.size(); i++)
    if ( array.isSet(i) && !array[i]->isEmpty()
         && array[i]->bitCapacity()<belowLvl*(array[i]->getContext().BPL()) )
      v.push_back(array[i]);
  packedRecrypt(CtPtrs_vectorPt(v), unpackConsts, ea);
}
void packedRecrypt(const CtPtrMat& m,
                   const std::vector<zzX>& unpackConsts,
                   const EncryptedArray& ea, long belowLvl)
{
  std::vector<Ctxt*> v;
  for (long i=0; i<m.size(); i++)
    for (long j=0; j<m[i].size(); j++)
      if ( m[i].isSet(j) && !m[i][j]->isEmpty()
           && m[i][j]->bitCapacity()<belowLvl*(m[i][j]->getContext().BPL()) )
        v.push_back(m[i][j]);
  packedRecrypt(CtPtrs_vectorPt(v), unpackConsts, ea);
}

//===================== Thin Bootstrapping stuff ==================

ThinRecryptData::~ThinRecryptData()
{
  if (coeffToSlot!=NULL)  delete coeffToSlot;
  if (slotToCoeff!=NULL) delete slotToCoeff;
}

// This code was copied from RecryptData::init, and is mostly
// the same, except for the linear-map-related stuff.
// FIXME: There is really too much code (and data!) duplication here.
void ThinRecryptData::init(const FHEcontext& context, const Vec<long>& mvec_,
                      bool alsoThick, long t, bool build_cache_, bool minimal)
{
  RecryptData::init(context, mvec_, alsoThick, t, build_cache_, minimal);
  coeffToSlot = new ThinEvalMap(*ea, minimal, mvec, true, build_cache);
  slotToCoeff = new ThinEvalMap(*context.ea, minimal, mvec, false, build_cache);
}

// Extract digits from thinly packed slots

long fhe_force_chen_han = 0;

void extractDigitsThin(Ctxt& ctxt, long botHigh, long r, long ePrime)
{
  FHE_TIMER_START;

  Ctxt unpacked(ctxt);
  unpacked.cleanUp();

  vector<Ctxt> scratch;

  long p = ctxt.getContext().zMStar.getP();
  long p2r = power_long(p,r);
  long topHigh = botHigh + r-1;

  // degree Chen/Han technique is p^{bot-1}(p-1)r
  // degree of basic technique is p^{bot-1}p^r,
  //     or p^{bot-1}p^{r-1} if p==2, r > 1, and bot+r > 2

  bool use_chen_han = false;
  if (r > 1) {
    double chen_han_cost = log(p-1) + log(r);
    double basic_cost;
    if (p == 2 && botHigh + r > 2)
       basic_cost = (r-1)*log(p);
    else
       basic_cost = r*log(p);

    //cerr << "*** basic: " << basic_cost << "\n";
    //cerr << "*** chen/han: " << chen_han_cost << "\n";

    double thresh = 1.5;
    if (p == 2) thresh = 1.75;
    // increasing thresh makes chen_han less likely to be chosen.
    // For p == 2, the basic algorithm is just squaring,
    // and so is a bit cheaper, so we raise thresh a bit.
    // This is all a bit heuristic.

    if (basic_cost > thresh*chen_han_cost)
      use_chen_han = true;
  }

  if (fhe_force_chen_han > 0)
    use_chen_han = true;
  else if (fhe_force_chen_han < 0)
    use_chen_han = false;

  if (use_chen_han) {
    // use Chen and Han technique

    extendExtractDigits(scratch, unpacked, botHigh, r);

#if 0
    for (long i: range(scratch.size())) {
      CheckCtxt(scratch[i], "**");
    }
#endif

    for (long j = 0; j < botHigh; j++) {
      unpacked -= scratch[j];
      unpacked.divideByP();
    }

    if (p==2 && botHigh>0)   // For p==2, subtract also the previous bit
      unpacked += scratch[botHigh-1];
    unpacked.negate();

    if (r>ePrime) {          // Add in digits from the bottom part, if any
      long topLow = r-1 - ePrime;
      Ctxt tmp = scratch[topLow];
      for (long j=topLow-1; j>=0; --j) {
	tmp.multByP();
	tmp += scratch[j];
      }
      if (ePrime>0)
	tmp.multByP(ePrime); // multiply by p^e'
      unpacked += tmp;
    }
    unpacked.reducePtxtSpace(p2r); // Our plaintext space is now mod p^r

    ctxt = unpacked;
  }
  else {

    if (p==2 && r>1 && topHigh+1 > 2)
      topHigh--; // For p==2 we sometime get a bit for free

    extractDigits(scratch, unpacked, topHigh+1);

    // set upacked = -\sum_{j=botHigh}^{topHigh} scratch[j] * p^{j-botHigh}
    if (topHigh >= LONG(scratch.size())) {
      topHigh = scratch.size() -1;
      cerr << " @ suspect: not enough digits in extractDigitsPacked\n";
    }

    unpacked = scratch[topHigh];
    for (long j=topHigh-1; j>=botHigh; --j) {
      unpacked.multByP();
      unpacked += scratch[j];
    }
    if (p==2 && botHigh>0)   // For p==2, subtract also the previous bit
      unpacked += scratch[botHigh-1];
    unpacked.negate();

    if (r>ePrime) {          // Add in digits from the bottom part, if any
      long topLow = r-1 - ePrime;
      Ctxt tmp = scratch[topLow];
      for (long j=topLow-1; j>=0; --j) {
	tmp.multByP();
	tmp += scratch[j];
      }
      if (ePrime>0)
	tmp.multByP(ePrime); // multiply by p^e'
      unpacked += tmp;
    }
    unpacked.reducePtxtSpace(p2r); // Our plaintext space is now mod p^r
    ctxt = unpacked;
  }

}

// Hack to get at private fields of public key
struct FHEPubKeyHack { // The public key
  const FHEcontext& context; // The context

  //! @var Ctxt pubEncrKey
  //! The public encryption key is an encryption of 0,
  //! relative to the first secret key
  Ctxt pubEncrKey;

  std::vector<long> skHwts; // The Hamming weight of the secret keys
  std::vector<KeySwitch> keySwitching; // The key-switching matrices

  // The keySwitchMap structure contains pointers to key-switching matrices
  // for re-linearizing automorphisms. The entry keySwitchMap[i][n] contains
  // the index j such that keySwitching[j] is the first matrix one needs to
  // use when re-linearizing s_i(X^n).
  std::vector< std::vector<long> > keySwitchMap;

  NTL::Vec<int> KS_strategy; // NTL Vec's support I/O, which is
                             // more convenient

  // bootstrapping data

  long recryptKeyID; // index of the bootstrapping key
  Ctxt recryptEkey;  // the key itself, encrypted under key #0

};

// bootstrap a ciphertext to reduce noise
void FHEPubKey::thinReCrypt(Ctxt &ctxt)
{
  FHE_TIMER_START;

  // Some sanity checks for dummy ciphertext
  long ptxtSpace = ctxt.getPtxtSpace();
  if (ctxt.isEmpty()) return;

  if (ctxt.parts.size()==1 && ctxt.parts[0].skHandle.isOne()) {
    // Dummy encryption, just ensure that it is reduced mod p
    ZZX poly = to_ZZX(ctxt.parts[0]);
    for (long i=0; i<poly.rep.length(); i++)
      poly[i] = to_ZZ( rem(poly[i],ptxtSpace) );
    poly.normalize();
    ctxt.DummyEncrypt(poly);
    return;
  }

  assert(recryptKeyID>=0); // check that we have bootstrapping data

  long p = ctxt.getContext().zMStar.getP();
  long r = ctxt.getContext().alMod.getR();
  long p2r = ctxt.getContext().alMod.getPPowR();

  long intFactor = ctxt.intFactor;

  const ThinRecryptData& trcData = ctxt.getContext().rcData;

  // the bootstrapping key is encrypted relative to plaintext space p^{e-e'+r}.
  long e = trcData.e;
  long ePrime = trcData.ePrime;
  long p2ePrime = power_long(p,ePrime);
  long q = power_long(p,e)+1;
  assert(e>=r);

  // can only bootstrap ciphertext with plaintext-space dividing p^r
  assert(p2r % ptxtSpace == 0);

#ifdef DEBUG_PRINTOUT
  CheckCtxt(ctxt, "init");
#endif

  ctxt.dropSmallAndSpecialPrimes();

#ifdef DROP_BEFORE_THIN_RECRYPT
  // experimental code...we should drop down to a reasonably low level
  // before doing the first linear map.
  long first = context.ctxtPrimes.first();
  long last = min(context.ctxtPrimes.last(),
                  first + thinRecrypt_initial_level - 1);
  ctxt.bringToSet(IndexSet(first, last));
#endif

#ifdef DEBUG_PRINTOUT
  CheckCtxt(ctxt, "after mod down");
#endif

  // Move the slots to powerful-basis coefficients
  FHE_NTIMER_START(AAA_slotToCoeff);
  trcData.slotToCoeff->apply(ctxt);
  FHE_NTIMER_STOP(AAA_slotToCoeff);

#ifdef DEBUG_PRINTOUT
  CheckCtxt(ctxt, "after slotToCoeff");
#endif

  FHE_NTIMER_START(AAA_bootKeySwitch);

  // Make sure that this ciphertxt is in canonical form
  if (!ctxt.inCanonicalForm()) ctxt.reLinearize();

  // Mod-switch down if needed
  IndexSet s = ctxt.getPrimeSet() / context.specialPrimes;
  assert(s <= context.ctxtPrimes);
  if (s.card()>2) { // leave only bottom two primes
    long frst = s.first();
    long scnd = s.next(frst);
    IndexSet s2(frst,scnd);
    s.retain(s2); // retain only first two primes
  }
  ctxt.modDownToSet(s);

  // key-switch to the bootstrapping key
  ctxt.reLinearize(recryptKeyID);

  // "raw mod-switch" to the bootstrapping mosulus q=p^e+1.
  vector<ZZX> zzParts; // the mod-switched parts, in ZZX format
  double noise = ctxt.rawModSwitch(zzParts, q);
  assert(zzParts.size() == 2);

#ifdef DEBUG_PRINTOUT
  if (dbgKey) {
    cerr << "  before makeDivisible (recryption modulus q="<<q
         << "), noise_bnd=" << noise<<endl;
    printSizesPowerful(zzParts, dbgKey->sKeys[recryptKeyID],
                       ctxt.getContext().rcData, q, noise);
  }
#endif

#if 1

  // Add multiples of p2r and q to make the zzParts divisible by p^{e'}
  double maxU_norm = 0;
  long maxU=0;
  for (long i=0; i<(long)zzParts.size(); i++) {
    // make divisible by p^{e'}
    double U_norm;

    long newMax = makeDivisible(zzParts[i].rep, p2ePrime, p2r, q,
				trcData.a, U_norm, context.zMStar);

    zzParts[i].normalize(); // normalize after working directly on the rep
    if (maxU < newMax)  maxU = newMax;
    if (maxU_norm < U_norm)  maxU_norm = U_norm;
  }
#ifdef DEBUG_PRINTOUT
  double newNoise = noise + maxU_norm*(skBounds[recryptKeyID]+1);
  cerr << "  after makeDivisible, maxU=" << maxU
       << ", maxU_norm="<<maxU_norm<<", p2r="<<p2r
       << ", noise_bnd="<<newNoise<<", sk_bnd="<< skBounds[recryptKeyID]
       << endl;
   if (dbgKey)
     printSizesPowerful(zzParts, dbgKey->sKeys[recryptKeyID],
                        ctxt.getContext().rcData, q, newNoise);
#endif

  for (long i=0; i<(long)zzParts.size(); i++)
    zzParts[i] /= p2ePrime;   // divide by p^{e'}

#else
  // Experimental version
  // Add multiples of p2r and q to make the zzParts divisible by p^{e'}
  double maxU_norm = 0;
  long maxU=0;
  for (long i=0; i<(long)zzParts.size(); i++) {
    // make divisible by p^{e'}
    double U_norm;

    Vec<ZZ> pwrfl;
    trcData.p2dConv->ZZXtoPowerful(pwrfl, zzParts[i]);

    long newMax = makeDivisible(pwrfl, p2ePrime, p2r, q,
				trcData.a, U_norm, context.zMStar);

    trcData.p2dConv->powerfulToZZX(zzParts[i], pwrfl);

    zzParts[i].normalize(); // normalize after working directly on the rep
    if (maxU < newMax)  maxU = newMax;
    if (maxU_norm < U_norm)  maxU_norm = U_norm;
  }
#ifdef DEBUG_PRINTOUT
  double newNoise = noise + maxU_norm*(skBounds[recryptKeyID]+1);
  cerr << "  after makeDivisible, maxU=" << maxU
       << ", maxU_norm="<<maxU_norm<<", p2r="<<p2r
       << ", noise_bnd="<<newNoise<<", sk_bnd="<< skBounds[recryptKeyID]
       << endl;
   if (dbgKey)
     printSizesPowerful(zzParts, dbgKey->sKeys[recryptKeyID],
                        ctxt.getContext().rcData, q, newNoise);
#endif

  for (long i=0; i<(long)zzParts.size(); i++)
    zzParts[i] /= p2ePrime;   // divide by p^{e'}

#endif

  // Multiply the post-processed cipehrtext by the encrypted sKey

  double p0size = to_double(embeddingLargestCoeff(zzParts[0], context.zMStar));
  double p1size = to_double(embeddingLargestCoeff(zzParts[1], context.zMStar));
  // FIXME: This might be slow without Armadillo

  // NOTE: here we lose the intFactor associated with ctxt.
  // We will restore it below.
  ctxt = recryptEkey;

  ctxt.multByConstant(zzParts[1], p1size);
  ctxt.addConstant(zzParts[0], p0size);

#ifdef DEBUG_PRINTOUT
   CheckCtxt(ctxt, "after bootKeySwitch");
#endif

  FHE_NTIMER_STOP(AAA_bootKeySwitch);

  // Move the powerful-basis coefficients to the plaintext slots
  FHE_NTIMER_START(AAA_coeffToSlot);
  trcData.coeffToSlot->apply(ctxt);
  FHE_NTIMER_STOP(AAA_coeffToSlot);

#ifdef DEBUG_PRINTOUT
   CheckCtxt(ctxt, "after coeffToSlot");
#endif

  // Extract the digits e-e'+r-1,...,e-e' (from fully packed slots)
  FHE_NTIMER_START(AAA_extractDigitsThin);
  extractDigitsThin(ctxt, e-ePrime, r, ePrime);
  FHE_NTIMER_STOP(AAA_extractDigitsThin);

#ifdef DEBUG_PRINTOUT
   CheckCtxt(ctxt, "after extractDigitsThin");
#endif

  // restore intFactor
  if (intFactor != 1)
    ctxt.intFactor = MulMod(ctxt.intFactor, intFactor, ptxtSpace);
}

static void
printSizesPowerful(const vector<ZZX>& zzParts, const DoubleCRT& sKey,
                   const RecryptData& rcData, long q, double noise)
{
  ZZX ptxt;
  //  const RecryptData& rcData = ctxt.getContext().rcData;
  const PAlgebra& palg = rcData.ea->getPAlgebra();
  rawDecrypt(ptxt, zzParts, sKey); // no mod q

  Vec<ZZ> powerful;
  rcData.p2dConv->ZZXtoPowerful(powerful, ptxt);
  xdouble max_pwrfl = conv<xdouble>(largestCoeff(powerful));
  xdouble max_canon = embeddingLargestCoeff(ptxt, palg);
  double ratio = log(max_pwrfl/max_canon)/log(2.0);
  xdouble critical_value = (max_pwrfl/q)/q;
  cerr << "                     max_pwrfl/q^2=" << ((max_pwrfl/q)/q)
       << ", log2(max_pwrfl/max_canon)=" << ratio;
  if (ratio > 0) cerr << " BAD-BOUND";
  cerr << endl;

  vecRed(powerful, powerful, q, false);
  max_pwrfl = conv<xdouble>(largestCoeff(powerful));
  rcData.p2dConv->powerfulToZZX(ptxt, powerful);
  max_canon = embeddingLargestCoeff(ptxt, palg);
  ratio = log(max_pwrfl/max_canon)/log(2.0);
  critical_value += max_pwrfl/q;
  cerr << "        after mod q, max_pwrfl/q=" << (max_pwrfl/q);
  if (critical_value > 0.5) cerr << " BAD-BOUND";
  cerr << ", log2(max_pwrfl/max_canon)=" << ratio;
  if (ratio > 0) cerr << " BAD-BOUND";

  ratio = log(max_canon/noise)/log(2.0);
  cerr << "\n        log2(max_canon/noiseEst)=" << ratio;
  if (ratio > 0) cerr << " BAD-BOUND";

  ratio = log(noise/q)/log(2.0);
  cerr << ", log2(noiseEst/q)=" << ratio;

  cerr << endl;
}


/*======================================================================
Copyright 2019 Google LLC

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

    https://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.
======================================================================*/

#define PCV_STATE_GLOBAL_CREATE

#include <Eigen/Dense>
#include "PCV_Types.h"

#include "PCV_State.h"

using namespace Eigen;

// CREATE THESE GLOBAL VARIABLES HERE
Float hw_freq;
Float cur_time;

// PREFIX: d=desired, g=global, r=raw ; SUFFIX d=dot, dd=dot.dot
//       pos   vel   acc
Vector8d   q              = Vector8d::Zero()
Vector8d        qd        = Vector8d::Zero()

Vector3d       rxd        = Vector3d::Zero();
Vector3d      rgxd        = Vector3d::Zero();
Vector3d   x              = Vector3d::Zero();
Vector3d        xd        = Vector3d::Zero();
Vector3d  gx              = Vector3d::Zero();
Vector3d       gxd        = Vector3d::Zero();
Vector3d            gxdd  = Vector3d::Zero();
Vector3d  dx              = Vector3d::Zero();
Vector3d       dxd        = Vector3d::Zero();
Vector3d            dxdd  = Vector3d::Zero();

Traj3  *traj;


//==================================================================================================
/*!
  @file

  @copyright 2016 NumScale SAS

  Distributed under the Boost Software License, Version 1.0.
  (See accompanying file LICENSE.md or copy at http://boost.org/LICENSE_1_0.txt)
*/
//==================================================================================================
#ifndef BOOST_SIMD_ARCH_COMMON_SCALAR_FUNCTION_FREXP_HPP_INCLUDED
#define BOOST_SIMD_ARCH_COMMON_SCALAR_FUNCTION_FREXP_HPP_INCLUDED
#include <boost/simd/function/fast.hpp>
#include <boost/simd/function/std.hpp>

#ifndef BOOST_SIMD_NO_DENORMALS
#include <boost/simd/constant/twotonmb.hpp>
#include <boost/simd/function/is_eqz.hpp>
#endif
#include <boost/simd/detail/constant/limitexponent.hpp>
#include <boost/simd/detail/constant/mask1frexp.hpp>
#include <boost/simd/detail/constant/mask2frexp.hpp>
#include <boost/simd/detail/constant/maxexponentm1.hpp>
#include <boost/simd/constant/nbmantissabits.hpp>
#include <boost/simd/constant/zero.hpp>
#include <boost/simd/function/bitwise_and.hpp>
#include <boost/simd/function/bitwise_andnot.hpp>
#include <boost/simd/function/bitwise_cast.hpp>
#include <boost/simd/function/bitwise_or.hpp>
#include <boost/simd/function/is_invalid.hpp>
#include <boost/simd/function/shr.hpp>
#include <boost/simd/detail/dispatch/function/overload.hpp>
#include <boost/config.hpp>
#include <utility>
#include <cmath>

namespace boost { namespace simd { namespace ext
{
  namespace bd = boost::dispatch;
  BOOST_DISPATCH_OVERLOAD ( frexp_
                          , (typename A0)
                          , bd::cpu_
                          , bd::scalar_< bd::floating_<A0> >
                          )
  {
    using i_t = bd::as_integer_t<A0, signed>;
    BOOST_FORCEINLINE std::pair<A0,i_t> operator() ( A0 a0) const BOOST_NOEXCEPT
    {
      if (a0 == 0 || is_invalid(a0))
      {
        return {a0, Zero<i_t>()};
      }
      else
      {
        i_t r1 = simd::bitwise_cast<i_t>(bitwise_and(Mask1frexp<A0>(), a0));  // extract exp.
#ifndef BOOST_SIMD_NO_DENORMALS
        i_t t = Zero<i_t>();
        if(is_eqz(r1)) // denormal
        {
          a0 *= Twotonmb<A0>();
          r1 = simd::bitwise_cast<i_t>(bitwise_and(Mask1frexp<A0>(), a0));  // extract exp. again
          t = Nbmantissabits<A0>();
        }
#endif
        A0 x  = bitwise_andnot(a0, Mask1frexp<A0>());                    // clear exp. in a0
        r1 = shr(r1,Nbmantissabits<A0>())- Maxexponentm1<A0>();         // compute exp.
        if (r1 > Limitexponent<A0>())
        {
          return {a0, Zero<i_t>()};
        }
#ifndef BOOST_SIMD_NO_DENORMALS
        r1 -= t;
#endif
        return {bitwise_or(x,Mask2frexp<A0>()), r1};
      }
    }
  };

  BOOST_DISPATCH_OVERLOAD ( frexp_
                          , (typename A0)
                          , bd::cpu_
                          , boost::simd::fast_tag
                          , bd::scalar_< bd::floating_<A0> >
                          )
  {
    using i_t = bd::as_integer_t<A0, signed>;
    BOOST_FORCEINLINE std::pair<A0,i_t> operator() (const fast_tag &
                                                   , A0 a0 ) const BOOST_NOEXCEPT
    {
      i_t r1  = bitwise_cast<i_t>(bitwise_and(Mask1frexp<A0>(), a0));
      A0  x = bitwise_andnot(a0, Mask1frexp<A0>());
      return {bitwise_or(x,Mask2frexp<A0>()), shr(r1,Nbmantissabits<A0>()) - Maxexponentm1<A0>()};
    }
  };
  BOOST_DISPATCH_OVERLOAD ( frexp_
                          , (typename A0)
                          , bd::cpu_
                          , boost::simd::std_tag
                          , bd::scalar_< bd::floating_<A0> >
                          )
  {
    using i_t = bd::as_integer_t<A0, signed>;
    BOOST_FORCEINLINE std::pair<A0,i_t> operator() (const std_tag &,  A0 a0
                                                   ) const BOOST_NOEXCEPT
    {
      int e = Zero<int>();
      A0 r = std::frexp(a0, &e);
      return {r, i_t(e)};
    }
  };
} } }

#endif


/**
 * @file test_finalSelectionSampleExpQ.cpp
 * @brief This file defines the test cases for the final selection policy based on the sample
 * exponential Q-function.
 *
 * @copyright Copyright (c) 2022
 *
 */

#include <algorithm>
#include <boost/test/unit_test.hpp>
#include <boost/test/unit_test_suite.hpp>
#include <cstddef>
#include <map>
#include <memory>
#include <string>
#include <vector>

#include "proseco_planning/config/configuration.h"
#include "proseco_planning/config/defaultConfiguration.h"
#include "proseco_planning/math/mathlib.h"
#include "proseco_planning/policies/final_selection/finalSelectionSampleExpQ.h"
#include "proseco_planning/util/alias.h"

namespace utf = boost::unit_test;
using namespace proseco_planning;

struct ConfigFixture {
  ConfigFixture() { Config::create(config::scenarioSimple, config::optionsSimple); }
  ~ConfigFixture() { Config::get()->reset(); }
};

struct FinalSelectionSampleExpQFixture : ConfigFixture, FinalSelectionSampleExpQ {
  FinalSelectionSampleExpQFixture() : FinalSelectionSampleExpQ("") {}
};

BOOST_FIXTURE_TEST_SUITE(FinalSelectionSampleExpQ, FinalSelectionSampleExpQFixture)

BOOST_AUTO_TEST_CASE(exponentiated_action_value, *utf::tolerance(0.00001f)) {
  BOOST_TEST(FinalSelectionSampleExpQ::calculateActionWeight(1000.0f) ==
             std::numeric_limits<float>::max());
  BOOST_TEST(FinalSelectionSampleExpQ::calculateActionWeight(10.0f) ==
             std::numeric_limits<float>::max());
  BOOST_TEST(FinalSelectionSampleExpQ::calculateActionWeight(0.1f) == 22026.46579480f);
  BOOST_TEST(FinalSelectionSampleExpQ::calculateActionWeight(-0.1f) == 4.5399931e-05f);
  BOOST_TEST(FinalSelectionSampleExpQ::calculateActionWeight(-1.0f) == 0.0f);
}

BOOST_AUTO_TEST_CASE(sampleActionFromWeights, *utf::tolerance(0.00001f)) {
  math::Random::setRandomSeed(99);
  std::vector<float> weights{0.1f, 0.1f, 0.999f, 0.1f, 0.1f};
  auto [index, probability] = FinalSelectionSampleExpQ::sampleActionFromWeights(weights);
  BOOST_TEST(index == 2);
  BOOST_TEST(probability == 0.714081466f);
  weights                      = {0.1f, 0.2f, 0.3f, 0.4f, 0.5f};
  std::tie(index, probability) = FinalSelectionSampleExpQ::sampleActionFromWeights(weights);
  BOOST_TEST(index == 4);
  BOOST_TEST(probability == 0.333333343f);
}

BOOST_AUTO_TEST_SUITE_END()


/*
@copyright Louis Dionne 2014
Distributed under the Boost Software License, Version 1.0.
(See accompanying file LICENSE.md or copy at http://boost.org/LICENSE_1_0.txt)
 */

#include <boost/hana/map.hpp>

#include <boost/hana/detail/assert.hpp>
#include <boost/hana/detail/constexpr.hpp>
#include <boost/hana/detail/minimal/comparable.hpp>
#include <boost/hana/detail/minimal/list.hpp>
#include <boost/hana/detail/minimal/product.hpp>
using namespace boost::hana;

template <int i>
constexpr auto key = detail::minimal::comparable<>(i);

template <int i>
constexpr auto value = detail::minimal::comparable<>(i + 100);

template <int i, int j>
constexpr auto p = detail::minimal::product<>(key<i>, value<j>);

int main() {
    using L = detail::minimal::List<>;
    constexpr auto list = detail::minimal::list<>;
    BOOST_HANA_CONSTANT_ASSERT(to<L>(map()) == list());
    BOOST_HANA_CONSTEXPR_ASSERT(to<L>(map(p<1, 1>)) == list(p<1, 1>));

    BOOST_HANA_CONSTEXPR_ASSERT(elem(permutations(list(p<1, 1>, p<2, 2>)), to<L>(map(p<1, 1>, p<2, 2>))));
    BOOST_HANA_CONSTEXPR_ASSERT(elem(permutations(list(p<1, 1>, p<2, 2>, p<3, 3>)), to<L>(map(p<1, 1>, p<2, 2>, p<3, 3>))));
}


#ifndef PWALK_PUBLIC_INTERFACE_HPP_
#define PWALK_PUBLIC_INTERFACE_HPP_

#include <Eigen/Dense>

Eigen::MatrixXd generateDikinWalkSamples(const Eigen::VectorXd& initialization, const Eigen::MatrixXd& cons_A, const Eigen::VectorXd& cons_b, const double r, const int nb_samples);

Eigen::MatrixXd generateVaidyaWalkSamples(const Eigen::VectorXd& initialization, const Eigen::MatrixXd& cons_A, const Eigen::VectorXd& cons_b, const double r, const int nb_samples);

Eigen::MatrixXd generateJohnWalkSamples(const Eigen::VectorXd& initialization, const Eigen::MatrixXd& cons_A, const Eigen::VectorXd& cons_b, const double r, const int nb_samples);

#endif // PWALK_PUBLIC_INTERFACE_HPP_