/** @file INPUTS/parallelInputs.cpp
 *  @brief Translates the input file into a structure that can be used by parallelFDTDField
 *
 *  Takes in a boost property tree and converts that information into a structure used by parallelFDTDField
 *
 *  @author Thomas A. Purcell (tpurcell90)
 *  @author Joshua E. Szekely (jeszekely)
 *  @bug No known bugs.
 */

#ifndef PRALLEL_FDTD_INPUTS
#define PRALLEL_FDTD_INPUTS

#include <src/OBJECTS/Obj.hpp>
#include <src/UTIL/ml_consts.hpp>
#include <src/UTIL/dielectric_params.hpp>
#include <src/UTIL/utilityFxns.hpp>
#include <boost/filesystem.hpp>
#include <boost/property_tree/ptree.hpp>
#include <boost/property_tree/json_parser.hpp>
#include <UTIL/ML_Dist_Fxn.hpp>
#include <iterator>

struct EnergyLevelDiscriptor
{
    DISTRIBUTION EDist_; //!< Distribution type for the energy level
    int nstates_; //!< number of states in the level
    int levDescribed_; //!< how many levels are described by this distribution
    std::vector<double> energyStates_; //!< energy of all the states in the level
    std::vector<double> weights_; //!< weights for the levels
};

/**
 * @brief input parameters class for the parallel FDTD class
 * @details sores all the values necessary to convert the input files into a fdtd grid
 */
class parallelProgramInputs
{
public:
    bool periodic_; //!< if true use PBC
    POLARIZATION pol_; //!< polarization of the grid (EX, EY, HZ = TE; HX, HY, EZ = TM)

    std::string filename_; //!< filename of the input file
    int res_; //!< number of grid points per unit length

    double courant_; //!< Courant factor of the cell
    double a_; //!< the unit length of the calculations
    double tMax_; //!< the max time of the calculation
    double I0_; //!< the unit current of the cell

    bool cplxFields_; //!< if true use complex fields
    bool saveFreqField_; //!< if true save the flux fields

    std::array<int,3> pmlThickness_; //!< thickness of the PMLs in all directions
    std::array<double,3> k_point_; //!< k_point vector of the light
    std::array<double,3> size_; //!< size of the cell in units of the unit length
    std::array<double,3> d_; //!< size of the cell in units of the unit length

    double dt_; //!< time step of the algorithm
    double pmlSigOptRat_; //!< Ratio of SigMax to SigOpt
    double pmlKappaMax_; //!< max a value for CPML
    double pmlAMax_; //!< max a value for CPML
    double pmlMa_; //!< scaling factor of a for the CPML
    double pmlM_; //!< scaling factor for the CPML

    std::vector<double> inputMapSlicesX_; //!< list of slices in the YZ plane
    std::vector<double> inputMapSlicesY_; //!< list of slices in the XZ plane
    std::vector<double> inputMapSlicesZ_; //!< list of slices in the XY plane

    std::vector<POLARIZATION> srcPol_; //!<polarization of all sources
    std::vector<std::vector<std::vector<cplx>>> srcFxn_; //!< pulse function parameters for the sources
    std::vector<std::array<int,3>> srcLoc_; //!<location of all sources
    std::vector<std::array<int,3>> srcSz_; //!<sizes of all all sources
    std::vector<double> srcPhi_; //!< angle of incidence for the source detector
    std::vector<double> srcTheta_; //!< angle of incidence for the source detector
    std::vector<std::vector<PLSSHAPE>> srcPulShape_; //!< pulse shape of all sources
    std::vector<std::vector<double>> srcEmax_; //!< max E incd field of all sources
    std::vector<double> srcEllipticalKratio_; //!< ratio between the long and short axis for current source surfaces
    std::vector<double> srcPsi_; //!< angle of polarization for circular and elliptical light

    std::vector<std::array<int,3>> tfsfLoc_;//!< location of the tfsf lower left point
    std::vector<std::array<int,3>> tfsfSize_; //!< size of the TFSF region
    std::vector<std::array<int,3>> tfsfM_; //!< vector describing the m value for tfsf surfaces
    std::vector<double> tfsfTheta_; //!< Phi angle of the TFSF surface
    std::vector<double> tfsfPhi_; //!< Phi angle of the TFSF surface
    std::vector<double> tfsfPsi_; //!< Phi angle of the TFSF surface
    std::vector<std::vector<std::vector<cplx>>> tfsfPulFxn_; //!< pulse function parameters for the sources
    std::vector<std::vector<PLSSHAPE>> tfsfPulShape_; //!< pulse shape of all sources
    std::vector<std::vector<double>> tfsfEmax_; //!< max E incd field of all sources
    std::vector<POLARIZATION> tfsfCircPol_; //!< sets the TFSF surface to be L or R polarized
    std::vector<double> tfsfEllipticalKratio_; //!< ratio between the long and short axis for TFSF surfaces
    std::vector<int> tfsfPMLThick_; //!< thickness of the incd PML's in number of grid points
    std::vector<double> tfsfPMLAMax_; //!< max a value for CPML
    std::vector<double> tfsfPMLM_; //!< scaling factor for the CPML
    std::vector<double> tfsfPMLMa_; //!< scaling factor of a for the CPML

    std::vector<std::shared_ptr<Obj>> objArr_; //!< array of all objects in the grid

    std::vector<std::vector<std::array<int,2>>> qeBasis_; //!< vector describing the basis functions for each state
    std::vector<EnergyLevelDiscriptor> qeELevs_; //!< vector storing the level descriptor for each level
    std::vector<std::vector<double>> qeCouplings_; //!< dipole coupling matrix for all systems
    std::vector<std::vector<std::array<int,3>>> qeLoc_; //!< location of all QE points
    std::vector<std::vector<std::vector<double>>> qeGam_; //!< relaxation matrix for all QEs
    std::vector<double> qeDen_; //!< molecular density of all QEs
    std::vector<std::vector<std::array<int,2>>> qeRelaxTransitonStates_; //!< Transitions for Lindblad Operator
    std::vector<std::vector<double>> qeRelaxRates_; //!< Rate of relaxation for Linblad operators
    std::vector<std::vector<double>> qeRelaxDephasingRate_; //!< Rate of dephasing between both states
    std::vector<bool> qeAccumP_; //!< bool to determine if output Polarization vector at every time
    std::vector<std::string> qeOutPolFname_; //!< filename for output
    std::vector<int> qeDtcPopTimeInt_; //!< Time interval for level dtc
    std::vector<std::vector<int>> qePopDtcLevs_; //!< vector storing the levels outputted at each step
    std::vector<std::vector<std::string>> qeDtcPopOutFile_; //!< filename for the pop output

    std::vector<int> qeContinuumStates_; //!< vector containing all of the Continuum states
    std::vector<double> qeContinuumNa_; //!< vector containing all of the Continuum band gap energy
    std::vector<double> qeContinuumOmgGap_; //!< vector containing all of the Continuum band gap energy
    std::vector<double> qeContinuum_dOmg_; //!< vector containing all of the Continuum inter-continuum state energy separation
    std::vector<double> qeContinuumMu_; //!< vector containing all of the Continuum transition dipole moments
    std::vector<double> qeContinuumGam1_; //!< vector containing all of the Continuum relaxation to the ground state (rationalness)
    std::vector<double> qeContinuumGamk_; //!< vector containing all of the Continuum relaxation to the conduction band edge (raditionless)
    std::vector<double> qeContinuumGamP_; //!< vector containing all of the Continuum relaxation to the conduction band edge (raditionless)
    std::vector<std::vector<std::array<int,2>>> qeContinuumLocs_; //!< vector containing all of the Continuum locations

    std::vector<DTCCLASS> dtcClass_; //!< describer of output file type for all detectors
    std::vector<bool> dtcSI_; //!< if true use SI units
    std::vector<std::array<int,3>> dtcLoc_; //!< location of all detectors
    std::vector<std::array<int,3>> dtcSz_; //!< sizes of all detectors
    std::vector<std::string> dtcName_; //!< file name for all detectors
    std::vector<DTCTYPE> dtcType_; //!< whether the dtc stores the Ex, Ey, Ez, Hx, Hy, Hz fields or the H or E power
    std::vector<double> dtcTimeInt_; //!< time interval for all detectors
    std::vector<GRIDOUTFXN> dtcOutBMPFxnType_; //!< what function should bmp converter use
    std::vector<GRIDOUTTYPE> dtcOutBMPOutType_; //!< how to output the values for the detector ina text file
    std::vector<std::vector<double>> dtcFreqList_; //!< center frequency
    std::vector<double> dtcTStart_; //!< time to start collecting fields
    std::vector<double> dtcTEnd_; //!< time to end collecting fields
    std::vector<bool> dtcOutputAvg_; //!< True if only outputting time_integrated fields
    std::vector<bool> dtcOutputMaps_; //!< True if only outputting time_integrated fields

    std::vector<int> fluxXOff_; //!< the x location offset of the fields
    std::vector<int> fluxYOff_; //!< the y location offset of the fields
    std::vector<int> fluxTimeInt_; //!< time interval used to intake the fields
    std::vector<std::array<int,3>> fluxLoc_; //!< Location of the lower left corner of the flux surface
    std::vector<std::array<int,3>> fluxSz_; //!< size of the flux surface
    std::vector<double> fluxWeight_; //!< weight of the flux surface
    std::vector<std::string> fluxName_; //!< file name for the output file
    std::vector<std::vector<double>> fluxFreqList_; //!<  center frequency
    std::vector<bool> fluxSI_; //!<  use SI units
    std::vector<bool> fluxCrossSec_; //!< calculate the cross-section?
    std::vector<bool> fluxSave_; //!< save the fields?
    std::vector<bool> fluxLoad_; //!< load the fields?
    std::vector<std::string> fluxIncdFieldsFilename_; //!< incident file names

    /**
     * @brief      Constructs the input parameter object
     *
     * @param[in]  IP    boost property tree generated from the input json file
     * @param[in]  fn    The filename of the input file
     */
    parallelProgramInputs(boost::property_tree::ptree IP,std::string fn);

    /**
     * @brief      Gets the dielectric parameters for a material.
     *
     * @param[in]  mat   String identifier of the material
     *
     * @return     The dielectric parameters.
     */
    std::vector<double> getDielectricParams(std::string mat);

    /**
     * @brief      converts a string to GRIDOUTFXN
     *
     * @param[in]  f     String identifier to a GRIDOUTFXN
     *
     * @return     GRIDOUTFXN from that input string
     */
    GRIDOUTFXN string2GRIDOUTFXN (std::string f);

    /**
     * @brief      converts a string to GRIDOUTTYPE
     *
     * @param[in]  t     String identifier to a GRIDOUTTYPE
     *
     * @return     GRIDOUTTYPE from that input string
     */
    GRIDOUTTYPE string2GRIDOUTTYPE (std::string t);

    /**
     * @brief      converts a string to POLARIZATION
     *
     * @param[in]  p     String identifier to a POLARIZATION
     *
     * @return     POLARIZATION from that input string
     */
    POLARIZATION string2pol(std::string p);

    /**
     * @brief      converts a string to SHAPE
     *
     * @param[in]  s     String identifier to a SHAPE
     *
     * @return     SHAPE from that input string
     */
    SHAPE string2shape(std::string s);

    /**
     * @brief      converts a string to DTCTYPE
     *
     * @param[in]  t     String identifier to a DTCTYPE
     *
     * @return     DTCTYPE from that input string
     */
    DTCTYPE string2out(std::string t);

    /**
     * @brief      converts a string to DTCCLASS
     *
     * @param[in]  c     String identifier to a DTCCLASS
     *
     * @return     DTCCLASS from that input string
     */
    DTCCLASS string2dtcclass(std::string c);

    /**
     * @brief      converts a string to PLSSHAPE
     *
     * @param[in]  p     String identifier to a PLSSHAPE
     *
     * @return     PLSSHAPE from that input string
     */
    PLSSHAPE string2prof(std::string p);

    /**
     * @brief      converts a string to DIRECTION
     *
     * @param[in]  dir   String identifier to a DIRECTION
     *
     * @return     DIRECTION from that input string
     */
    DIRECTION string2dir(std::string dir);

    /**
     * @brief      converts a string to DISTRIBUTION
     *
     * @param[in]  f     String identifier to a DISTRIBUTION
     *
     * @return     DISTRIBUTION from that input string
     */
    DISTRIBUTION string2dist(std::string dist);

    /**
     * @brief      Converts a string to MAT_DIP_ORIENTAITON
     *
     * @param[in]  dipOr  String identifier for the MAT_DIP_ORIENTAITON
     *
     * @return     The MAT_DIP_ORIENTAITON from the input string
     */
    MAT_DIP_ORIENTAITON string2dipor(std::string dipOr);

    /**
     * @brief      Gets the material parameters for a given material
     *
     * @param[in]  mat   String identifying the material
     *
     * @return     The material parameters
     */
    std::tuple<double,double,double, std::vector<LorenzDipoleOscillator> > getMater(std::string mat);

    /**
     * @brief      Converts the metal parameters from eV based units to FDTD based units
     *
     * @param[in]  params  The parameters for the metallic material in eV
     *
     * @return     The parameters for the metallic material in FDTD units
     */
    std::vector<LorenzDipoleOscillator> getMetal(std::vector<double> params);

    std::vector<LorenzDipoleOscillator> getMetalJM(std::vector<double> params);

    /**
     * @brief      Converts eV units to FDTD units
     *
     * @param[in]  eV    The value in eV
     *
     * @return     The value in FDTD frequency units
     */
    double ev2FDTD(double eV);

    /**
     * @brief      Converts the point from real space to grid points
     *
     * @param[in]  pt    Real space value
     *
     * @return     Corresponding grid point value
     */
    inline int find_pt(double pt, double d) {return int( floor(pt/d + 0.5) );}

    /**
     * @brief      Accessor function to tMax_
     *
     * @return     Maximum time of the simulation
     */
    inline double tMax() {return tMax_;}

    /**
     * @brief      Constructs an object from the object list child tree
     *
     * @param      iter  boost::ptree child corresponding to the object
     *
     * @return     shared_ptr to the object
     */
    std::shared_ptr<Obj> ptreeToObject(boost::property_tree::ptree::value_type &iter);
};
/**
 * @brief      strips comments from the input file
 *
 * @param      filename  The filename of the file to strip
 */
void stripComments(std::string& filename);

/**
 * @brief      boost json to std::vector<T>
 *
 * @param[in]  pt          property tree
 * @param[in]  key         property tree key
 *
 * @tparam     T           double, int
 *
 * @return     json input as a std::vector<T>
 */
template <typename T>
std::vector<T> as_vector(boost::property_tree::ptree const &pt, boost::property_tree::ptree::key_type const &key)
{
    std::vector<T> r;
    for (auto& item : pt.get_child(key))
        r.push_back(item.second.get_value<T>());
    return r;
}

/**
 * @brief      boost json to std::vector<T>
 *
 * @param[in]  pt          property tree
 * @param[in]  key         property tree key
 * @param[in]  defaultVal  The default value
 * @param[in]  szVec       The size of the vector
 *
 * @tparam     T           double, int
 *
 * @return     json input as a std::vector<T>
 */
template <typename T>
std::vector<T> as_vector(boost::property_tree::ptree const &pt, boost::property_tree::ptree::key_type const &key, T defaultVal, int szVec)
{
    std::vector<T> r;
    try
    {
        for (auto& item : pt.get_child(key))
            r.push_back(item.second.get_value<T>());
    }
    catch(std::exception& e)
    {
        r = std::vector<T>(szVec, defaultVal);
    }
    return r;
}

/**
 * @brief      boost json to std::array<T,3>
 *
 * @param[in]  pt          property tree
 * @param[in]  key         property tree key
 * @param[in]  defaultVal  The default value
 *
 * @tparam     T           double, int
 *
 * @return     json input as a std::array<T,3>
 */
template <typename T>
std::array<T,3> as_ptArr(boost::property_tree::ptree const &pt, boost::property_tree::ptree::key_type const &key, T defaultVal=0)
{
    std::array<T, 3> r = {0,0,0};
    try
    {
        int ii = 0;
        for (auto& item : pt.get_child(key))
        {
            r[ii] = item.second.get_value<T>();
            ++ii;
        }
    }
    catch(std::exception& e)
    {
        r = {{ defaultVal, defaultVal, defaultVal}};
    }
    return r;
}

/**
 * @brief      boost json to std::array<T,3>
 *
 * @param[in]  pt          property tree
 * @param[in]  key         property tree key
 * @param[in]  defaultVal  The default value
 *
 * @tparam     T           double, int
 *
 * @return     json input as a std::array<T,3>
 */
template <typename T>
std::array<T,2> as_ptArr2(boost::property_tree::ptree const &pt, boost::property_tree::ptree::key_type const &key, T defaultVal=0)
{
    std::array<T, 2> r = {0,0};
    try
    {
        int ii = 0;
        for (auto& item : pt.get_child(key))
        {
            r[ii] = item.second.get_value<T>();
            ++ii;
        }
    }
    catch(std::exception& e)
    {
        r = {{ defaultVal, defaultVal}};
    }
    return r;
}
#endif


// Copyright (c) 2017 Daniel Cooke
// Use of this source code is governed by the MIT license that can be found in the LICENSE file.

#include <boost/test/unit_test.hpp>

#include <stdexcept>
#include <vector>
#include <iterator>
#include <numeric>
#include <algorithm>
#include <cstdint>
#include <cmath>
#include <functional>

#include <boost/test/floating_point_comparison.hpp>

#include "basics/phred.hpp"

namespace octopus { namespace test {

namespace { constexpr double tolerance {0.000001}; }

BOOST_AUTO_TEST_SUITE(basics)
BOOST_AUTO_TEST_SUITE(phred)

BOOST_AUTO_TEST_CASE(phreds_must_be_positive)
{
    BOOST_CHECK_THROW((Phred<> {-1}), std::exception);
    BOOST_CHECK_NO_THROW((Phred<> {0}));
    BOOST_CHECK_NO_THROW((Phred<> {-0}));
}

BOOST_AUTO_TEST_CASE(phred_are_the_size_of_their_underlying_type)
{
    BOOST_CHECK_EQUAL(sizeof(Phred<signed char> {}),   sizeof(signed char));
    BOOST_CHECK_EQUAL(sizeof(Phred<unsigned char> {}), sizeof(unsigned char));
    BOOST_CHECK_EQUAL(sizeof(Phred<int> {}),           sizeof(int));
    BOOST_CHECK_EQUAL(sizeof(Phred<unsigned> {}),      sizeof(unsigned));
    BOOST_CHECK_EQUAL(sizeof(Phred<float> {}),         sizeof(float));
    BOOST_CHECK_EQUAL(sizeof(Phred<double> {}),        sizeof(double));
    BOOST_CHECK_EQUAL(sizeof(Phred<std::uint8_t> {}),  sizeof(std::uint8_t));
    BOOST_CHECK_EQUAL(sizeof(Phred<std::int8_t> {}),   sizeof(std::int8_t));
    BOOST_CHECK_EQUAL(sizeof(Phred<std::uint16_t> {}), sizeof(std::uint16_t));
    BOOST_CHECK_EQUAL(sizeof(Phred<std::int16_t> {}),  sizeof(std::int16_t));
}

BOOST_AUTO_TEST_CASE(phreds_can_be_converted_to_probabilities)
{
    std::vector<double> v(100);

    std::iota(std::begin(v), std::end(v), 0.0);

    for (auto x : v) {
        Phred<> phred {x};
        BOOST_CHECK_EQUAL(phred.score(), x);
        //BOOST_CHECK_CLOSE(phred.probability_false(), std::pow(10.0, -x / 10.0), tolerance);
    }
}

BOOST_AUTO_TEST_CASE(phreds_can_be_constructed_with_probabilities)
{
    std::vector<double> v(100, 0.1);

    v.front() = 1;
    std::partial_sum(std::next(std::cbegin(v)), std::cend(v), std::next(std::begin(v)), std::multiplies<> {});

    double c {0};
    for (auto p : v) {
        auto phred = probability_false_to_phred(p);
        BOOST_CHECK_CLOSE(phred.score(), 10 * c++, tolerance);
        //BOOST_CHECK_CLOSE(phred.probability_false(), p, tolerance);
    }
}

BOOST_AUTO_TEST_SUITE_END()
BOOST_AUTO_TEST_SUITE_END()

} // namespace test
} // namespace octopus


//
// Created by jonas on 06.04.20.
//

#include <boost/program_options.hpp>
#include <iostream>

#include "../src/graph/GraphIO.h"

template <typename T>
std::vector<std::vector<T>> to_matrix(const VertexPairMap<T> &map) {
    auto n = map.size();
    auto out = std::vector<std::vector<T>>(n, std::vector<T>(n));

    for (auto uv : Graph::VertexPairs(n)) {
        out[uv.u][uv.v] = map[uv];
        out[uv.v][uv.u] = map[uv];
    }

    return out;
}

std::vector<std::vector<int>> to_matrix(const Graph &graph) {
    auto n = graph.size();
    auto out = std::vector<std::vector<int>>(n, std::vector<int>(n));

    for (auto uv : graph.edges()) {
        out[uv.u][uv.v] = 1;
        out[uv.v][uv.u] = 1;
    }

    return out;
}

void to_yaml(const Instance &instance, const std::string& input_path, double multiplier, int permutation, std::ostream& os) {
    using namespace YAML;

    Emitter out;
    out << BeginDoc << BeginMap;
    out << Key << "instance" << Value << input_path;
    out << Key << "multiplier" << Value << multiplier;
    out << Key << "permutation" << Value << permutation;
    out << Key << "adj" << Value << Flow << to_matrix(instance.graph);
    out << Key << "costs" << Value << Flow << to_matrix(instance.costs);
    out << EndMap << EndDoc;

    os << out.c_str();
}

// Note: Other formats may include metis.
enum class OutputFormat {
    YAML
};

std::istream& operator>>(std::istream& in, OutputFormat& selector) {
    std::string token;
    in >> token;
    if (token == "yaml")
        selector = OutputFormat::YAML;
    else
        in.setstate(std::ios_base::failbit);
    return in;
}

std::ostream &operator<<(std::ostream &os, OutputFormat selector) {
    switch (selector) {
        case OutputFormat::YAML:
            return os << "yaml";
        default:
            return os;
    }
}

int main(int argc, char* argv[]) {
    namespace po = boost::program_options;

    try {

        std::string input_path;
        double multiplier = 1;
        int permutation = 0;
        std::string output_path;

        auto output_type = OutputFormat::YAML;

        po::options_description desc("Allowed options");
        desc.add_options()
            ("help", "Produce help message.")
            ("input", po::value<std::string>(&input_path)->required(),
                "Path to input instance.\nMay be provided as named option or as the first positional "
                "parameter.")
            ("multiplier", po::value<double>(&multiplier)->default_value(multiplier),
                "Multiplier for multiplying editing costs before discretization.")
            ("permutation", po::value<int>(&permutation)->default_value(permutation),
                "Permutation seed for the instance. The value 0 represents the identity.")
            ("output", po::value<std::string>(&output_path)->default_value(output_path),
                "Path to output file. If no path is specified the output is written to stdout.\n"
                "May be provided as named option or as the second positional parameter.")
            ("format", po::value<OutputFormat>(&output_type)->default_value(output_type),
                "Output format. Currently only 'yaml' is possible.")
            ;

        po::positional_options_description pd;
        pd.add("input", 1);
        pd.add("output", 2);

        po::variables_map vm;
        po::store(po::command_line_parser(argc, argv).options(desc).positional(pd).run(), vm);
        po::notify(vm);

        if (vm.count("help")) {
            std::cout << "Usage: " << argv[0] << " [options...] input [output]\n";
            std::cout << desc << "\n";
            return EXIT_FAILURE;
        }

        std::ostream* os = &std::cout;
        std::ofstream output_file;

        if (!output_path.empty()) {
            output_file.open(output_path);
            if (!output_file)
                throw std::runtime_error("could not open output_file");
            os = &output_file;
        }

        auto instance = GraphIO::read_instance(input_path, multiplier, permutation);

        switch (output_type) {
            case OutputFormat::YAML:
                to_yaml(instance, input_path, multiplier, permutation, *os);
                break;
            default:
                throw std::runtime_error("unknown value for output_type");
        }

    } catch (const std::exception &e) {
        std::cerr << e.what() << "\n";
        return EXIT_FAILURE;
    } catch (...) {
        std::cerr << "unknown exception\n";
        return EXIT_FAILURE;
    }

    return EXIT_SUCCESS;
}


#ifndef SKYLARK_SPARSE_MATRIX_HPP
#define SKYLARK_SPARSE_MATRIX_HPP

#include <set>
#include <vector>
#include <boost/tuple/tuple.hpp>
#include <boost/unordered_map.hpp>

namespace skylark { namespace base {

/**
 *  This implements a very crude CSC sparse matrix container only intended to
 *  hold local sparse matrices.
 *
 *  Row indices are not sorted.
 *  Structure is always constants, and can only be attached by Attached.
 *  Values of non-zeros can be modified.
 */
template<typename ValueType=double>
struct sparse_matrix_t {

    typedef int index_type;
    typedef ValueType value_type;

    typedef boost::tuple<index_type, index_type, value_type> coord_tuple_t;
    typedef std::vector<coord_tuple_t> coords_t;

    sparse_matrix_t()
        : _ownindptr(false), _ownindices(false), _ownvalues(false),
          _dirty_struct(false), _height(0), _width(0), _nnz(0),
          _indptr(nullptr), _indices(nullptr), _values(nullptr)
    {}

    // The following relies on C++11
    sparse_matrix_t(sparse_matrix_t<ValueType>&& A) :
        _ownindptr(A._ownindptr), _ownindices(A._ownindices),
        _ownvalues(A._ownvalues), _dirty_struct(A._dirty_struct),
        _height(A._height), _width(A._width), _nnz(A._nnz),
        _indptr(A._indptr), _indices(A._indices), _values(A._values)
    {
        A._ownindptr = false;
        A._ownindices = false;
        A._ownvalues = false;
    }

    ~sparse_matrix_t() {
        _free_data();
    }

    bool struct_updated() const { return _dirty_struct; }
    void reset_update_flag()    { _dirty_struct = false; }

    /**
     * Copy data to external buffers.
     */
    template<typename IdxType, typename ValType>
    void detach(IdxType *indptr, IdxType *indices, ValType *values) const {

        for(size_t i = 0; i <= _width; ++i)
            indptr[i] = static_cast<IdxType>(_indptr[i]);

        for(size_t i = 0; i < _nnz; ++i) {
            indices[i] = static_cast<IdxType>(_indices[i]);
            values[i] = static_cast<ValType>(_values[i]);
        }
    }

    /**
     * Attach new structure and values.
     */
    void attach(const index_type *indptr, const index_type *indices, double *values,
        int nnz, int n_rows, int n_cols, bool _own = false) {
        attach(indptr, indices, values, nnz, n_rows, n_cols, _own, _own, _own);
    }

    /**
     * Attach new structure and values.
     */
    void attach(const index_type *indptr, const index_type *indices, double *values,
        int nnz, int n_rows, int n_cols,
        bool ownindptr, bool ownindices, bool ownvalues) {
        _free_data();

        _indptr = indptr;
        _indices = indices;
        _values = values;
        _nnz = nnz;
        _width = n_cols;
        _height = n_rows;

        _ownindptr = ownindptr;
        _ownindices = ownindices;
        _ownvalues = ownvalues;

        _dirty_struct = true;
    }

    // attaching a coordinate structure facilitates going from distributed
    // input to local output.
    void set(coords_t coords, int n_rows = 0, int n_cols = 0) {

        sort(coords.begin(), coords.end(), &sparse_matrix_t::_sort_coords);

        n_cols = std::max(n_cols, boost::get<1>(coords.back()) + 1);
        index_type *indptr = new index_type[n_cols + 1];

        // Count non-zeros
        int nnz = 0;
        for(size_t i = 0; i < coords.size(); ++i) {
            nnz++;
            index_type cur_row = boost::get<0>(coords[i]);
            index_type cur_col = boost::get<1>(coords[i]);
            while(i + 1 < coords.size() &&
                  cur_row == boost::get<0>(coords[i + 1]) &&
                  cur_col == boost::get<1>(coords[i + 1]))
                i++;
        }

        index_type *indices = new index_type[nnz];
        value_type *values = new value_type[nnz];

        nnz = 0;
        int indptr_idx = 0;
        indptr[indptr_idx] = 0;
        for(size_t i = 0; i < coords.size(); ++i) {
            index_type cur_row = boost::get<0>(coords[i]);
            index_type cur_col = boost::get<1>(coords[i]);
            value_type cur_val = boost::get<2>(coords[i]);

            for(; indptr_idx < cur_col; ++indptr_idx)
                indptr[indptr_idx + 1] = nnz;
            nnz++;

            // sum duplicates
            while(i + 1 < coords.size() &&
                  cur_row == boost::get<0>(coords[i + 1]) &&
                  cur_col == boost::get<1>(coords[i + 1])) {

                cur_val += boost::get<2>(coords[i + 1]);
                i++;
            }

            indices[nnz - 1] = cur_row;
            values[nnz - 1] = cur_val;

            n_rows = std::max(cur_row + 1, n_rows);
        }

        for(; indptr_idx < n_cols; ++indptr_idx)
            indptr[indptr_idx + 1] = nnz;

        attach(indptr, indices, values, nnz, n_rows, n_cols, true);
    }

    int height() const {
        return _height;
    }

    int width() const {
        return _width;
    }

    int nonzeros() const {
        return _nnz;
    }

    const index_type* indptr() const {
        return _indptr;
    }

    const index_type* indices() const {
        return _indices;
    }

    value_type* values() {
        return _values;
    }

    const value_type* locked_values() const {
        return _values;
    }

    bool operator==(const sparse_matrix_t &rhs) const {

        // column pointer arrays have to be exactly the same
        if (std::vector<int>(_indptr, _indptr+_width) !=
            std::vector<int>(rhs._indptr, rhs._indptr + rhs._width))
            return false;

        // check more carefully for unordered row indices
        const int* indptr  = _indptr;
        const int* indices = _indices;
        const double* values = _values;

        const int* indices_rhs   = rhs.indices();
        const double* values_rhs = rhs.locked_values();

        for(int col = 0; col < width(); col++) {

            boost::unordered_map<int, double> col_values;

            for(int idx = indptr[col]; idx < indptr[col + 1]; idx++)
                col_values.insert(std::make_pair(indices[idx], values[idx]));

            for(int idx = indptr[col]; idx < indptr[col + 1]; idx++) {
                if(col_values[indices_rhs[idx]] != values_rhs[idx])
                    return false;
            }
        }

        return true;
    }

private:
    bool _ownindptr;
    bool _ownindices;
    bool _ownvalues;

    bool _dirty_struct;

    int _height;
    int _width;
    int _nnz;

    const index_type* _indptr;
    const index_type* _indices;
    value_type* _values;

    // TODO add the following
    sparse_matrix_t(const sparse_matrix_t&);
    void operator=(const sparse_matrix_t&);

    void _free_data() {
        if (_ownindptr)
            delete[] _indptr;
        if (_ownindices)
            delete[] _indices;
        if (_ownvalues)
            delete[] _values;
    }

    static bool _sort_coords(coord_tuple_t lhs, coord_tuple_t rhs) {
        if(boost::get<1>(lhs) != boost::get<1>(rhs))
            return boost::get<1>(lhs) < boost::get<1>(rhs);
        else
            return boost::get<0>(lhs) < boost::get<0>(rhs);
    }
};

template<typename T>
void Transpose(const sparse_matrix_t<T>& A, sparse_matrix_t<T>& B) {
    const int* aindptr = A.indptr();
    const int* aindices = A.indices();
    const double* avalues = A.locked_values();

    int m = A.width();
    int n = A.height();
    int nnz = A.nonzeros();

    int *indptr = new int[n + 1];
    int *indices = new int[nnz];
    double *values = new double[nnz];

    // Count nonzeros in each row
    int *nzrow = new int[n];
    std::fill(nzrow, nzrow + n, 0);
    for(int col = 0; col < m; col++)
        for(int idx = aindptr[col]; idx < aindptr[col + 1]; idx++)
            nzrow[aindices[idx]]++;

    // Set indptr
    indptr[0] = 0;
    for(int col = 1; col <= n; col++)
        indptr[col] = indptr[col - 1] + nzrow[col - 1];

    // Fill values
    std::fill(nzrow, nzrow + n, 0);
    for(int col = 0; col < m; col++)
        for(int idx = aindptr[col]; idx < aindptr[col + 1]; idx++) {
            int row = aindices[idx];
            double val = avalues[idx];
            indices[indptr[row] + nzrow[row]] = col;
            values[indptr[row] + nzrow[row]] = val;
            nzrow[row]++;
        }

    delete[] nzrow;

    B.attach(indptr, indices, values, nnz, m, n, true);
}

} }

#endif


// Copyright 2019 the Autoware Foundation
//
// 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.
//
// Co-developed by Tier IV, Inc. and Apex.AI, Inc.

#include <common/types.hpp>
#include <gtest/gtest.h>
#include <optimization/utils.hpp>
#include <Eigen/Core>
#include <vector>
#include <limits>
#include "test_cache_states.hpp"

using autoware::common::types::float64_t;

namespace autoware
{
namespace common
{
namespace optimization
{
TEST_F(CacheStateMachineTest, numeric_test) {
  auto eigen_comparator = [](const auto & lhs, const auto & rhs) {
      return lhs.isApprox(rhs, std::numeric_limits<float64_t>::epsilon());
    };

  for (const auto & elem : m_mode_map) {
    // Each element in the map has a term (score, hessian etc.) and all the modes
    // that enable computation of that term
    const auto term = elem.first;
    const auto modes = elem.second;
    for (auto & mode : modes) {
      do_test(mode, 1, 3, term);  // int
      do_test(mode, 1.5F, 2.7F, term);  // float
      do_test(mode, 1.5, 2.7, term);  // double
      do_test(mode, false, true, term);  // bool
      do_test(
        mode, Eigen::Vector3d{1.2, 2.5, 3.6},
        Eigen::Vector3d{7.9, 12.6, 100.9}, term, eigen_comparator);  // eigen vector
    }
  }
}

}  // namespace optimization
}  // namespace common
}  // namespace autoware


#include <boost/geometry/algorithms/difference.hpp>


#include <iostream>

#include <El.hpp>
#include <boost/mpi.hpp>
#include <boost/format.hpp>

// Some tricks to make compilation faster
#define SKYLARK_NO_ANY
#define SKYLARK_WITH_JLT_ANY
#define SKYLARK_WITH_UST_ANY
#define SKYLARK_WITH_GAUSSIAN_RFT_ANY
#define SKYLARK_WITH_LAPLACIAN_RFT_ANY
#define SKYLARK_WITH_PPT_ANY
#define SKYLARK_WITH_FAST_GAUSSIAN_RFT_ANY

#include <skylark.hpp>

// Algorithms constants
#define CLASSIC_KRR                      0
#define FASTER_KRR                       1
#define APPROXIMATE_KRR                  2
#define SKETCHED_APPROXIMATE_KRR         3
#define FAST_SKETCHED_APPROXIMATE_KRR    4
#define LARGE_SCALE_KRR                  5
#define EXPERIMENTAL_1                 100
#define EXPERIMENTAL_2                 101

// Kernels constants
#define GAUSSIAN_KERNEL   0
#define LAPLACIAN_KERNEL  1
#define POLYNOMIAL_KERNEL 2
#define LINEAR_KERNEL   100

std::string cmdline;
int seed = 38734, algorithm = FASTER_KRR, kernel_type = GAUSSIAN_KERNEL;
skylark::utility::io::fileformat_t fileformat =
    skylark::utility::io::FORMAT_LIBSVM;
int s = 2000, partial = -1, sketch_size = -1, sample = -1, maxit = 0, maxsplit = 0;
std::string fname, testname, modelname = "model.dat", logfile = "",
    outputfile = "";
double kp1 = 10.0, kp2 = 0.0, kp3 = 1.0, lambda = 0.01, tolerance=0;
bool use_single = false, use_fast = false, regression = false;
bool predict = false, decisionvals = false;
boost::property_tree::ptree pt;

#ifndef SKYLARK_AVOID_BOOST_PO

#include <boost/program_options.hpp>
namespace bpo = boost::program_options;

int parse_program_options(int argc, char* argv[]) {

    bpo::options_description desc("Options");
    desc.add_options()
        ("help,h", "produce a help message")
        ("trainfile",
            bpo::value<std::string>(&fname)->default_value(""),
            "Data to train on. For predict - data to predict on. ")
        ("testfile",
            bpo::value<std::string>(&testname)->default_value(""),
            "Test data (libsvm format).")
        ("outputfile",
            bpo::value<std::string>(&outputfile)->default_value(""),
            "Output file (for predicition). Will not output if empty.")
        ("predict", "Predict mode -- load model file and use it.")
        ("model",
            bpo::value<std::string>(&modelname)->default_value("model.dat"),
            "Name of model file.")
        ("logfile",
            bpo::value<std::string>(&logfile)->default_value(""),
            "File to write log (standard output if empty).")
        ("kernel,k",
             bpo::value<int>(&kernel_type)->default_value(GAUSSIAN_KERNEL),
            "Kernel to use (0: Gaussian, 1: Laplacian, 2: Polynomial).")
        ("algorithm,a",
             bpo::value<int>(&algorithm)->default_value(FASTER_KRR),
            "Algorithm to use (0: Classic, 1: Faster (Precond), "
            "2: Approximate (Random Features)), "
            "3: Sketched Approximate (Piecemeal Random Features + Sketch)), "
            "4: Sketched Approximate w/ Faster Sketching, "
            "5: Large Scale. OPTIONAL.")
        ("seed,s",
            bpo::value<int>(&seed)->default_value(38734),
            "Seed for random number generation. OPTIONAL.")
        ("kernelparam,g",
            bpo::value<double>(&kp1),
            "Kernel parameter. REQUIRED.")
        ("kernelparam2,x",
            bpo::value<double>(&kp2)->default_value(0.0),
            "If Applicable - Second Kernel Parameter (Polynomial Kernel: c).")
        ("kernelparam3,y",
            bpo::value<double>(&kp3)->default_value(1.0),
            "If Applicable - Third Kernel Parameter (Polynomial Kernel: gamma).")
        ("lambda,l",
            bpo::value<double>(&lambda)->default_value(0.01),
            "Lambda regularization parameter.")
        ("tolerance,t",
            bpo::value<double>(&tolerance)->default_value(0),
            "Tolerance for the iterative method (when used). "
            "0 will default based on algorithm (1e-3 for -a 1, 1e-1 for -a 5)")
        ("maxsplit,c",
            bpo::value<int>(&maxsplit)->default_value(0),
            "Maximum number of random features in a split for large scale "
            "algorithms (-a 3 to 5). 0 will default to 2 * input dimension. ")
        ("maxit,i",
            bpo::value<int>(&maxit)->default_value(0),
            "Maximum number of iterations for the iterative method (when used). "
            "0 will default based on algorithm (1000 for -a 1, 20 for -a 5)")
        ("partial,p",
            bpo::value<int>(&partial)->default_value(-1),
            "Load only specified quantity examples from training. "
            "Will read all if -1.")
        ("sample,z",
            bpo::value<int>(&sample)->default_value(-1),
            "Sample the input data. Will use all if -1. ")
        ("decisionvals",
            "In predict mode, for classification, output the "
            "decision values instead of class.")
        ("single", "Whether to use single precision instead of double.")
        ("fast", "Try using a fast feature transform.")
        ("regression", "Build a regression model"
            "(default is classification).")
        ("numfeatures,f",
            bpo::value<int>(&s),
            "Number of random features (if relevant).")
        ("sketchsize,r",
            bpo::value<int>(&sketch_size)->default_value(-1),
            "Sketch size (for regression problem; if relevant (i.e., -a 3). "
            "-1 - will be determined by software. ")
        ("fileformat",
            po::value<char>((char *)&fileformat)->
            default_value(skylark::utility::io::FORMAT_LIBSVM),
            "Fileformat (default: 0 (libsvm), 1 (hdf5)");

    bpo::positional_options_description positional;
    positional.add("trainfile", 1);
    positional.add("testfile", 2);

    bpo::variables_map vm;
    try {
        bpo::store(bpo::command_line_parser(argc, argv)
            .options(desc).positional(positional).run(), vm);

        if (vm.count("help")) {
            std::cout << "Usage: " << argv[0]
                      << " [options] input-file-name [test-file-name]"
                      << std::endl;
            std::cout << desc;
            return 0;
        }

        bpo::notify(vm);

        use_single = vm.count("single");
        use_fast = vm.count("fast");
        regression = vm.count("regression");
        predict = vm.count("predict");
        decisionvals = vm.count("decisionvals");

        if (!vm.count("trainfile")) {
            std::cout << "Input trainfile file is required! "
                      << "(In predict mode, it is the test data.)"
                      << std::endl;
            return -1;
        }

    } catch(bpo::error& e) {
        std::cerr << e.what() << std::endl;
        std::cerr << desc << std::endl;
        return -1;
    }

    return 1000;
}

#else

int parse_program_options(int argc, char* argv[]) {

    int poscount = 0;
    for (int i = 1; i < argc; i += 2) {
        std::string flag = argv[i];
        std::string value = i + 1 < argc ? argv[i+1] : "";

        if (flag == "--seed" || flag == "-s")
            seed = boost::lexical_cast<int>(value);

        if (flag == "--lambda" || flag == "-l")
            lambda = boost::lexical_cast<double>(value);

        if (flag == "--tolerance" || flag == "-t")
            tolerance = boost::lexical_cast<double>(value);

        if (flag == "--maxit" || flag == "-i")
            maxit = boost::lexical_cast<int>(value);

        if (flag == "--maxsplit" || flag == "-c")
            maxsplit = boost::lexical_cast<int>(value);

        if (flag == "--partial" || flag == "-p")
            partial = boost::lexical_cast<int>(value);

        if (flag == "--sample" || flag == "-z")
            sample = boost::lexical_cast<int>(value);

        if (flag == "--kernelparam" || flag == "-g")
            kp1 = boost::lexical_cast<double>(value);

        if (flag == "--kernelparam2" || flag == "-x")
            kp2 = boost::lexical_cast<double>(value);

        if (flag == "--kernelparam3" || flag == "-y")
            kp3 = boost::lexical_cast<double>(value);

        if (flag == "--kernel" || flag == "-k")
            kernel_type = boost::lexical_cast<int>(value);

        if (flag == "--algorithm" || flag == "-a")
            algorithm = boost::lexical_cast<int>(value);

        if (flag == "--fileformat")
            fileformat = boost::lexical_cast<int>(value);

        if (flag == "--nunmfeatures" || flag == "-f")
            s = boost::lexical_cast<int>(value);

        if (flag == "--sketchsize" || flag == "-r")
            sketch_size = boost::lexical_cast<int>(value);

        if (flag == "--single") {
            use_single = true;
            i--;
        }

        if (flag == "--decisionvals") {
            decisionvals = true;
            i--;
        }

        if (flag == "--regression") {
            regression = true;
            i--;
        }

        if (flag == "--predict") {
            predict = true;
            i--;
        }

        if (flag == "--fast") {
            use_fast = true;
            i--;
        }

        if (flag == "--trainfile")
            fname = value;

        if (flag == "--logfile")
            logfile = value;

        if (flag == "--model")
            modelname = value;

        if (flag == "--testfile")
            testname = value;

        if (flag == "--outputfile")
            outputfile = value;

        if (flag[0] != '-' && poscount != 0)
            testname = flag;

        if (flag[0] != '-' && poscount == 0) {
            fname = flag;
            poscount++;
        }

        if (flag[0] != '-')
            i--;
    }

    return 1000;
}

#endif

template<typename T>
int execute_classification(skylark::base::context_t &context) {

    boost::mpi::communicator world;
    int rank = world.rank();

    std::ostream *log_stream = &std::cout;
    if (rank == 0 && logfile != "") {
        log_stream = new std::ofstream();
        ((std::ofstream *)log_stream)->open(logfile);
    }

    El::DistMatrix<T> X0, X;
    El::DistMatrix<El::Int> L0, L;

    boost::mpi::timer timer;

    if (rank == 0) {
        *log_stream << "# Generated using kernel_regression ";
        *log_stream << "using the following command-line: " << std::endl;
        *log_stream << "#\t" << cmdline << std::endl;
        *log_stream << "# Number of ranks is " << world.size() << std::endl;
    }

    // Load X and L
    if (rank == 0) {
        *log_stream << "Reading the matrix... ";
        log_stream->flush();
        timer.restart();
    }

    switch (fileformat) {
    case skylark::utility::io::FORMAT_LIBSVM:
        skylark::utility::io::ReadLIBSVM(fname, X0, L0, skylark::base::COLUMNS,
            0, partial);
        break;

#ifdef SKYLARK_HAVE_HDF5
    case skylark::utility::io::FORMAT_HDF5: {
        H5::H5File in(fname, H5F_ACC_RDONLY);
        skylark::utility::io::ReadHDF5(in, "X", X0, -1, partial);
        skylark::utility::io::ReadHDF5(in, "Y", L0, -1, partial);
        in.close();
    }
        break;
#endif

    default:
        *log_stream << "Invalid file format specified." << std::endl;
        return -1;
    }

    if (rank == 0)
        *log_stream <<"took " << boost::format("%.2e") % timer.elapsed()
                  << " sec\n";

    skylark::sketch::generic_sketch_container_t SampT;
    if (sample == -1) {

        El::View(X, X0);
        El::View(L, L0);

    } else {
        // Sample X and L
        if (rank == 0) {
            *log_stream << "Sampling the data... ";
            log_stream->flush();
            timer.restart();
        }

        SampT =
            skylark::sketch::create_sketch<skylark::sketch::UST_t>(X0.Width(),
                sample, skylark::sketch::UST_data_t::params_t(false), context);

        X.Resize(X0.Height(), sample);
        SampT.apply(&X0, &X, skylark::sketch::rowwise_tag());

        L.Resize(1, sample);
        SampT.apply(&L0, &L, skylark::sketch::rowwise_tag());

        if (rank == 0)
            *log_stream <<"took " << boost::format("%.2e") % timer.elapsed()
                        << " sec\n";
    }

    // Training
    if (rank == 0) {
        *log_stream << "Training... " << std::endl;
        timer.restart();
    }

    std::shared_ptr<skylark::ml::kernel_t> k_ptr;

    switch (kernel_type) {
    case GAUSSIAN_KERNEL:
        k_ptr.reset(new skylark::ml::gaussian_t(X.Height(), kp1));
        break;

    case LAPLACIAN_KERNEL:
        k_ptr.reset(new skylark::ml::laplacian_t(X.Height(), kp1));
        break;

    case POLYNOMIAL_KERNEL:
        k_ptr.reset(new skylark::ml::polynomial_t(X.Height(), kp1, kp2, kp3));
        break;

    case LINEAR_KERNEL:
        k_ptr.reset(new skylark::ml::linear_t(X.Height()));
        break;

    default:
        *log_stream << "Invalid kernel specified." << std::endl;
        return -1;
    }

    skylark::ml::kernel_container_t k(k_ptr);

    El::DistMatrix<T> A, W;
    std::vector<El::Int> rcoding;

    skylark::sketch::sketch_transform_container_t<El::DistMatrix<T>,
                                                  El::DistMatrix<T> >  S;
    bool scale_maps = true;
    std::vector<
        skylark::sketch::sketch_transform_container_t<El::DistMatrix<T>,
                                                      El::DistMatrix<T> > > transforms;

    skylark::ml::rlsc_params_t rlsc_params(rank == 0, 4, *log_stream, "\t");
    rlsc_params.use_fast = use_fast;

    skylark::ml::model_t<El::Int, T> *model;

    switch(algorithm) {
    case CLASSIC_KRR:
        skylark::ml::KernelRLSC(skylark::base::COLUMNS, k, X, L,
            T(lambda), A, rcoding, rlsc_params);
        model =
            new skylark::ml::kernel_model_t<skylark::ml::kernel_container_t,
              El::Int, T>(k, skylark::base::COLUMNS, X, fname, partial, fileformat,
                  SampT, A, rcoding);
        break;

    case FASTER_KRR:
        rlsc_params.iter_lim = (maxit == 0) ? 1000 : maxit;
        rlsc_params.tolerance = (tolerance == 0) ? 1e-3 : tolerance;
        skylark::ml::FasterKernelRLSC(skylark::base::COLUMNS, k, X, L,
            T(lambda), A, rcoding, s, context, rlsc_params);
        model =
            new skylark::ml::kernel_model_t<skylark::ml::kernel_container_t,
              El::Int, T>(k, skylark::base::COLUMNS, X, fname, partial,
                  fileformat, SampT, A, rcoding);
        break;

    case APPROXIMATE_KRR:
        skylark::ml::ApproximateKernelRLSC(skylark::base::COLUMNS, k, X, L,
            T(lambda), S, W, rcoding, s, context, rlsc_params);
        model =
            new skylark::ml::feature_expansion_model_t<
                skylark::sketch::sketch_transform_container_t, El::Int, T>
            (S, W, rcoding);
        break;

    case SKETCHED_APPROXIMATE_KRR:
    case FAST_SKETCHED_APPROXIMATE_KRR:
        rlsc_params.sketched_rls = true;
        rlsc_params.sketch_size = sketch_size;
        rlsc_params.fast_sketch = algorithm == FAST_SKETCHED_APPROXIMATE_KRR;
        rlsc_params.max_split = maxsplit;
        skylark::ml::SketchedApproximateKernelRLSC(skylark::base::COLUMNS, k, X, L,
            T(lambda), scale_maps, transforms, W, rcoding, s, sketch_size,
            context, rlsc_params);
        model =
            new skylark::ml::feature_expansion_model_t<
                skylark::sketch::sketch_transform_container_t, El::Int, T>
            (scale_maps, transforms, W, rcoding);
        break;

    case LARGE_SCALE_KRR:
        rlsc_params.iter_lim = (maxit == 0) ? 20 : maxit;
        rlsc_params.tolerance = (tolerance == 0) ? 1e-1 : tolerance;
        rlsc_params.max_split = maxsplit;
        skylark::ml::LargeScaleKernelRLSC(skylark::base::COLUMNS, k, X, L,
            T(lambda), scale_maps, transforms, W, rcoding, s,
            context, rlsc_params);
        model =
            new skylark::ml::feature_expansion_model_t<
                skylark::sketch::sketch_transform_container_t, El::Int, T>
            (scale_maps, transforms, W, rcoding);
        break;

    case EXPERIMENTAL_1:
    case EXPERIMENTAL_2:
        rlsc_params.sketched_rls = true;
        rlsc_params.fast_sketch = algorithm == EXPERIMENTAL_2;
        skylark::ml::ApproximateKernelRLSC(skylark::base::COLUMNS, k, X, L,
            T(lambda), S, W, rcoding, s, context, rlsc_params);
        model =
            new skylark::ml::feature_expansion_model_t<
                skylark::sketch::sketch_transform_container_t, El::Int, T>
            (S, W, rcoding);
        break;

    default:
        *log_stream << "Invalid algorithm value specified." << std::endl;
        return -1;
    }

    if (rank == 0)
        *log_stream << "Training took " << boost::format("%.2e") % timer.elapsed()
                  << " sec\n";

    if (modelname != "NOSAVE") {
        // Save model
        if (rank == 0) {
            *log_stream << "Saving model... ";
            log_stream->flush();
            timer.restart();
        }

        pt = model->to_ptree();

        if (rank == 0) {
            std::ofstream of(modelname);
            of << "# Generated using kernel_regression ";
            of << "using the following command-line: " << std::endl;
            of << "#\t" << cmdline << std::endl;
            of << "# Number of ranks is " << world.size() << std::endl;
            boost::property_tree::write_json(of, pt);
            of.close();
        }

        if (rank == 0)
            *log_stream <<"took " << boost::format("%.2e") % timer.elapsed()
                        << " sec\n";
    }

    // Test
    if (!testname.empty()) {
        if (rank == 0) {
            *log_stream << "Predicting... ";
            log_stream->flush();
            timer.restart();
        }

        El::DistMatrix<T> XT;
        El::DistMatrix<El::Int> LT;

        switch (fileformat) {
        case skylark::utility::io::FORMAT_LIBSVM:
            skylark::utility::io::ReadLIBSVM(testname, XT, LT,
                skylark::base::COLUMNS, X.Height());
            break;

#ifdef SKYLARK_HAVE_HDF5
        case skylark::utility::io::FORMAT_HDF5: {
            H5::H5File in(testname, H5F_ACC_RDONLY);
            skylark::utility::io::ReadHDF5(in, "X", XT);
            skylark::utility::io::ReadHDF5(in, "Y", LT);
            in.close();
        }
            break;
#endif

        default:
            *log_stream << "Invalid file format specified." << std::endl;
            return -1;
        }

        El::DistMatrix<T> DV;
        El::DistMatrix<El::Int> LP;
        model->predict(skylark::base::COLUMNS, XT, LP, DV);

        if (rank == 0)
            *log_stream << "took " << boost::format("%.2e") % timer.elapsed()
                      << " sec\n";

        int errs = 0;
        if (LT.LocalHeight() > 0)
            for(int i = 0; i < LT.LocalWidth(); i++)
                if (LT.GetLocal(0, i) != LP.GetLocal(0, i))
                    errs++;

        errs = El::mpi::AllReduce(errs, MPI_SUM, LT.DistComm());

        if (rank == 0)
            *log_stream << "Error rate: "
                      << boost::format("%.2f") % ((errs * 100.0) / LT.Width())
                      << "%" << std::endl;
    }

    if (rank == 0 && logfile != "") {
        ((std::ofstream *)log_stream)->close();
        delete log_stream;
    }

    delete model;

    return 0;
}

template<typename T>
int execute_regression(skylark::base::context_t &context) {

    boost::mpi::communicator world;
    int rank = world.rank();

    std::ostream *log_stream = &std::cout;
    if (rank == 0 && logfile != "") {
        log_stream = new std::ofstream();
        ((std::ofstream *)log_stream)->open(logfile);
    }

    El::DistMatrix<T> X0, X, Y0, Y;

    boost::mpi::timer timer;

    if (rank == 0) {
        *log_stream << "# Generated using kernel_regression ";
        *log_stream << "using the following command-line: " << std::endl;
        *log_stream << "#\t" << cmdline << std::endl;
        *log_stream << "# Number of ranks is " << world.size() << std::endl;
    }

    // Load X and Y
    if (rank == 0) {
        *log_stream << "Reading the matrix... ";
        log_stream->flush();
        timer.restart();
    }

    switch (fileformat) {
    case skylark::utility::io::FORMAT_LIBSVM:
        skylark::utility::io::ReadLIBSVM(fname, X0, Y0, skylark::base::COLUMNS,
            0, partial);
        break;

#ifdef SKYLARK_HAVE_HDF5
    case skylark::utility::io::FORMAT_HDF5: {
        H5::H5File in(fname, H5F_ACC_RDONLY);
        skylark::utility::io::ReadHDF5(in, "X", X0, -1, partial);
        skylark::utility::io::ReadHDF5(in, "Y", Y0, -1, partial);
        in.close();
    }
        break;
#endif

    default:
        *log_stream << "Invalid file format specified." << std::endl;
        return -1;
    }

    if (rank == 0)
        *log_stream <<"took " << boost::format("%.2e") % timer.elapsed()
                  << " sec\n";

    skylark::sketch::generic_sketch_container_t SampT;
    if (sample == -1) {

        El::View(X, X0);
        El::View(Y, Y0);

    } else {
        // Sample X and Y
        if (rank == 0) {
            *log_stream << "Sampling the data... ";
            log_stream->flush();
            timer.restart();
        }

        SampT =
            skylark::sketch::create_sketch<skylark::sketch::UST_t>(X0.Width(),
                sample, skylark::sketch::UST_data_t::params_t(false), context);

        X.Resize(X0.Height(), sample);
        SampT.apply(&X0, &X, skylark::sketch::rowwise_tag());

        Y.Resize(1, sample);
        SampT.apply(&Y0, &Y, skylark::sketch::rowwise_tag());

        if (rank == 0)
            *log_stream <<"took " << boost::format("%.2e") % timer.elapsed()
                        << " sec\n";
    }

    // Training
    if (rank == 0) {
        *log_stream << "Training... " << std::endl;
        timer.restart();
    }

    std::shared_ptr<skylark::ml::kernel_t> k_ptr;

    switch (kernel_type) {
    case GAUSSIAN_KERNEL:
        k_ptr.reset(new skylark::ml::gaussian_t(X.Height(), kp1));
        break;

    case LAPLACIAN_KERNEL:
        k_ptr.reset(new skylark::ml::laplacian_t(X.Height(), kp1));
        break;

    case POLYNOMIAL_KERNEL:
        k_ptr.reset(new skylark::ml::polynomial_t(X.Height(), kp1, kp2, kp3));
        break;

    default:
        *log_stream << "Invalid kernel specified." << std::endl;
        return -1;
    }

    skylark::ml::kernel_container_t k(k_ptr);

    El::DistMatrix<T> A, W;

    skylark::sketch::sketch_transform_container_t<El::DistMatrix<T>,
                                                  El::DistMatrix<T> >  S;
    bool scale_maps = true;
    std::vector<
        skylark::sketch::sketch_transform_container_t<El::DistMatrix<T>,
                                                      El::DistMatrix<T> > > transforms;

    skylark::ml::krr_params_t krr_params(rank == 0, 4, *log_stream, "\t");
    krr_params.use_fast = use_fast;

    // Transpose Y since KernelRidge expects it to be a column vector (TODO ?)
    El::DistMatrix<T> Ytransp;
    El::Transpose(Y, Ytransp, true);

    skylark::ml::model_t<T, T> *model;

    switch(algorithm) {
    case CLASSIC_KRR:
        skylark::ml::KernelRidge(skylark::base::COLUMNS, k, X, Ytransp,
            T(lambda), A, krr_params);
        model =
            new skylark::ml::kernel_model_t<skylark::ml::kernel_container_t,
                T, T>(k, skylark::base::COLUMNS, X, fname, partial, fileformat,
                    SampT, A);
        break;

    case FASTER_KRR:
        krr_params.iter_lim = (maxit == 0) ? 1000 : maxit;
        krr_params.tolerance = (tolerance == 0) ? 1e-3 : tolerance;
        skylark::ml::FasterKernelRidge(skylark::base::COLUMNS, k, X, Ytransp,
            T(lambda), A, s, context, krr_params);
        model =
            new skylark::ml::kernel_model_t<skylark::ml::kernel_container_t,
             T, T>(k, skylark::base::COLUMNS, X, fname, partial, fileformat,
                 SampT, A);
        break;

    case APPROXIMATE_KRR:
        skylark::ml::ApproximateKernelRidge(skylark::base::COLUMNS, k, X, Ytransp,
            T(lambda), S, W, s, context, krr_params);
        model =
            new skylark::ml::feature_expansion_model_t<
                skylark::sketch::sketch_transform_container_t, T, T>
            (S, W);
        break;

    case SKETCHED_APPROXIMATE_KRR:
    case FAST_SKETCHED_APPROXIMATE_KRR:
        krr_params.sketched_rr = true;
        krr_params.sketch_size = sketch_size;
        krr_params.fast_sketch = algorithm == FAST_SKETCHED_APPROXIMATE_KRR;
        krr_params.max_split = maxsplit;
        skylark::ml::SketchedApproximateKernelRidge(skylark::base::COLUMNS, k,
            X, Ytransp,
            T(lambda), scale_maps, transforms, W, s, sketch_size,
            context, krr_params);
        model =
            new skylark::ml::feature_expansion_model_t<
                skylark::sketch::sketch_transform_container_t, T, T>
            (scale_maps, transforms, W);
        break;

    case LARGE_SCALE_KRR:
        krr_params.iter_lim = (maxit == 0) ? 20 : maxit;
        krr_params.tolerance = (tolerance == 0) ? 1e-1 : tolerance;
        krr_params.max_split = maxsplit;
        skylark::ml::LargeScaleKernelRidge(skylark::base::COLUMNS, k, X, Ytransp,
            T(lambda), scale_maps, transforms, W, s,
            context, krr_params);
        model =
            new skylark::ml::feature_expansion_model_t<
                skylark::sketch::sketch_transform_container_t, T, T>
            (scale_maps, transforms, W);
        break;

    case EXPERIMENTAL_1:
    case EXPERIMENTAL_2:
        krr_params.sketched_rr = true;
        krr_params.fast_sketch = algorithm == EXPERIMENTAL_2;
        skylark::ml::ApproximateKernelRidge(skylark::base::COLUMNS, k, X, Ytransp,
            T(lambda), S, W, s, context, krr_params);
        model =
            new skylark::ml::feature_expansion_model_t<
                skylark::sketch::sketch_transform_container_t, T, T>
            (S, W);
        break;

    default:
        *log_stream << "Invalid algorithm value specified." << std::endl;
        return -1;
    }

    if (rank == 0)
        *log_stream << "Training took " << boost::format("%.2e") % timer.elapsed()
                  << " sec\n";

    if (modelname != "NOSAVE") {
        // Save model
        if (rank == 0) {
            *log_stream << "Saving model... ";
            log_stream->flush();
            timer.restart();
        }

        pt = model->to_ptree();

        if (rank == 0) {
            std::ofstream of(modelname);
            of << "# Generated using kernel_regression ";
            of << "using the following command-line: " << std::endl;
            of << "#\t" << cmdline << std::endl;
            of << "# Number of ranks is " << world.size() << std::endl;
            boost::property_tree::write_json(of, pt);
            of.close();
        }

        if (rank == 0)
            *log_stream <<"took " << boost::format("%.2e") % timer.elapsed()
                        << " sec\n";
    }

    // Test
    if (!testname.empty()) {
        if (rank == 0) {
            *log_stream << "Predicting... ";
            log_stream->flush();
            timer.restart();
        }

        El::DistMatrix<T> XT;
        El::DistMatrix<T> YT;

        switch (fileformat) {
        case skylark::utility::io::FORMAT_LIBSVM:
            skylark::utility::io::ReadLIBSVM(testname, XT, YT,
                skylark::base::COLUMNS, X.Height());
            break;

#ifdef SKYLARK_HAVE_HDF5
        case skylark::utility::io::FORMAT_HDF5: {
            H5::H5File in(testname, H5F_ACC_RDONLY);
            skylark::utility::io::ReadHDF5(in, "X", XT);
            skylark::utility::io::ReadHDF5(in, "Y", YT);
            in.close();
        }
            break;
#endif

        default:
            *log_stream << "Invalid file format specified." << std::endl;
            return -1;
        }

        El::DistMatrix<T> YP;
        model->predict(skylark::base::COLUMNS, XT, YP);

        if (rank == 0)
            *log_stream << "took " << boost::format("%.2e") % timer.elapsed()
                      << " sec\n";

        T nrm_Yt = El::Nrm2(YT);
        El::Axpy(T(-1.0), YP, YT);
        T nrm_E = El::Nrm2(YT);

        if (rank == 0)
            *log_stream << "Error rate: "
                        << boost::format("%.4e") % (nrm_E / nrm_Yt)
                        << std::endl;
        world.barrier();
    }

    if (rank == 0 && logfile != "") {
        ((std::ofstream *)log_stream)->close();
        delete log_stream;
    }

    delete model;

    return 0;
}

template<typename T>
int predict_regression(skylark::base::context_t &context) {
    boost::mpi::timer timer;

    boost::mpi::communicator world;
    int rank = world.rank();

    std::ostream *log_stream = &std::cout;
    if (rank == 0 && logfile != "") {
        log_stream = new std::ofstream();
        ((std::ofstream *)log_stream)->open(logfile);
    }

    if (rank == 0) {
        *log_stream << "Reading model... ";
        log_stream->flush();
        timer.restart();
    }

    skylark::ml::model_container_t<T, T> model(pt);

    if (rank == 0)
        *log_stream <<"took " << boost::format("%.2e") % timer.elapsed()
                    << " sec\n";

    if (rank == 0) {
        *log_stream << "Predicting... ";
        log_stream->flush();
        timer.restart();
    }

    El::DistMatrix<T> XT, YT;

    switch (fileformat) {
    case skylark::utility::io::FORMAT_LIBSVM:
        skylark::utility::io::ReadLIBSVM(fname, XT, YT,
            skylark::base::COLUMNS, model.get_input_size());
        break;

#ifdef SKYLARK_HAVE_HDF5
    case skylark::utility::io::FORMAT_HDF5: {
        H5::H5File in(fname, H5F_ACC_RDONLY);
        skylark::utility::io::ReadHDF5(in, "X", XT);
        skylark::utility::io::ReadHDF5(in, "Y", YT);
        in.close();
    }
        break;
#endif

    default:
        *log_stream << "Invalid file format specified." << std::endl;
        return -1;
    }

    El::DistMatrix<T> YP;
    model.predict(skylark::base::COLUMNS, XT, YP);

    if (rank == 0)
        *log_stream << "took " << boost::format("%.2e") % timer.elapsed()
                    << " sec\n";

    T nrm_Yt = El::Nrm2(YT);
    El::Axpy(T(-1.0), YP, YT);
    T nrm_E = El::Nrm2(YT);

    if (rank == 0)
        *log_stream << "Error rate: "
                    << boost::format("%.4e") % (nrm_E / nrm_Yt)
                    << std::endl;
    world.barrier();

    if (!outputfile.empty()) {
        if (rank == 0) {
            *log_stream << "Writing output... ";
            log_stream->flush();
            timer.restart();
        }

        El::DistMatrix<T> YPT;
        El::Transpose(YP, YPT);

        El::Write(YPT, outputfile, El::ASCII);

        if (rank == 0)
            *log_stream << "took " << boost::format("%.2e") % timer.elapsed()
                    << " sec\n";
    }

    if (rank == 0 && logfile != "") {
        ((std::ofstream *)log_stream)->close();
        delete log_stream;
    }

    return 0;
}

template<typename T>
int predict_classification(skylark::base::context_t &context) {

    boost::mpi::timer timer;

    boost::mpi::communicator world;
    int rank = world.rank();

    std::ostream *log_stream = &std::cout;
    if (rank == 0 && logfile != "") {
        log_stream = new std::ofstream();
        ((std::ofstream *)log_stream)->open(logfile);
    }

    if (rank == 0) {
        *log_stream << "Reading model... ";
        log_stream->flush();
        timer.restart();
    }

    skylark::ml::model_container_t<El::Int, T> model(pt);

    if (rank == 0)
        *log_stream <<"took " << boost::format("%.2e") % timer.elapsed()
                    << " sec\n";

    if (rank == 0) {
        *log_stream << "Predicting... ";
        log_stream->flush();
        timer.restart();
    }

    El::DistMatrix<T> XT;
    El::DistMatrix<El::Int> LT;

    switch (fileformat) {
    case skylark::utility::io::FORMAT_LIBSVM:
        skylark::utility::io::ReadLIBSVM(fname, XT, LT,
            skylark::base::COLUMNS, model.get_input_size());
        break;

#ifdef SKYLARK_HAVE_HDF5
    case skylark::utility::io::FORMAT_HDF5: {
        H5::H5File in(fname, H5F_ACC_RDONLY);
        skylark::utility::io::ReadHDF5(in, "X", XT);
        skylark::utility::io::ReadHDF5(in, "Y", LT);
        in.close();
    }
        break;
#endif

    default:
        *log_stream << "Invalid file format specified." << std::endl;
        return -1;
    }

    El::DistMatrix<T> DV;
    El::DistMatrix<El::Int> LP;
    model.predict(skylark::base::COLUMNS, XT, LP, DV);

    if (rank == 0)
        *log_stream << "took " << boost::format("%.2e") % timer.elapsed()
                    << " sec\n";

    int errs = 0;
    if (LT.LocalHeight() > 0)
        for(int i = 0; i < LT.LocalWidth(); i++)
            if (LT.GetLocal(0, i) != LP.GetLocal(0, i))
                errs++;

    errs = El::mpi::AllReduce(errs, MPI_SUM, LT.DistComm());

    if (rank == 0)
        *log_stream << "Error rate: "
                    << boost::format("%.2f") % ((errs * 100.0) / LT.Width())
                    << "%" << std::endl;

    if (!outputfile.empty()) {
        if (rank == 0) {
            *log_stream << "Writing output... ";
            log_stream->flush();
            timer.restart();
        }

        if (!decisionvals) {
            El::DistMatrix<El::Int> LPT;
            El::Transpose(LP, LPT);
            El::Write(LPT, outputfile, El::ASCII);
        } else {
            El::DistMatrix<T> DVT;
            El::Transpose(DV, DVT);
            std::stringstream stream;
            stream << "# Column order for decision values:" << std::endl
                   << "# ";
            std::vector<El::Int> colcoding;
            model.get_column_coding(colcoding);
            for(int i = 0; i < colcoding.size(); i++)
                stream << colcoding[i] << " ";
            El::Write(DVT, outputfile, El::ASCII, stream.str());
        }
        if (rank == 0)
            *log_stream << "took " << boost::format("%.2e") % timer.elapsed()
                    << " sec\n";
    }

    if (rank == 0 && logfile != "") {
        ((std::ofstream *)log_stream)->close();
        delete log_stream;
    }

    return 0;
}

int main(int argc, char* argv[]) {

    for(int i = 0; i < argc; i++) {
        cmdline.append(argv[i]);
        if (i < argc - 1)
            cmdline.append(" ");
    }

    El::Initialize(argc, argv);

    boost::mpi::communicator world;
    int rank = world.rank();

    int flag = parse_program_options(argc, argv);

    if (flag != 1000)
        return flag;

    skylark::base::context_t context(seed);

    int ret = -1;

    SKYLARK_BEGIN_TRY()

        // If in predict mode, we need to read the model to see if regression
        // or classification mode.
        if (predict) {
            std::ifstream is(modelname);

            // Skip all lines begining with "#"
            while(is.peek() == '#')
                is.ignore(std::numeric_limits<std::streamsize>::max(), '\n');

            boost::property_tree::read_json(is, pt);
            is.close();

            regression = pt.get<bool>("regression");

            if (regression) {
                if (use_single)
                    ret = predict_regression<float>(context);
                else
                    ret = predict_regression<double>(context);
            } else {
                if (use_single)
                    ret = predict_classification<float>(context);
                else
                    ret = predict_classification<double>(context);
            }
        } else {

            if (regression) {
                if (use_single)
                    ret = execute_regression<float>(context);
                else
                    ret = execute_regression<double>(context);
            } else {
                if (use_single)
                    ret = execute_classification<float>(context);
                else
                    ret = execute_classification<double>(context);
            }
        }
    SKYLARK_END_TRY() SKYLARK_CATCH_AND_PRINT((rank == 0))

        catch (const std::exception& ex) {
            if (rank == 0) SKYLARK_PRINT_EXCEPTION_DETAILS(ex);
        }

    El::Finalize();

    return ret;
}


// Author(s): Wieger Wesselink
// Copyright: see the accompanying file COPYING or copy at
// https://github.com/mCRL2org/mCRL2/blob/master/COPYING
//
// Distributed under 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)
//
/// \file rewriter_test.cpp
/// \brief Add your file description here.

#include "mcrl2/data/detail/data_functional.h"
#include "mcrl2/data/detail/one_point_rule_preprocessor.h"
#include "mcrl2/data/detail/parse_substitution.h"
#include "mcrl2/data/detail/test_rewriters.h"
#include "mcrl2/data/find.h"
#include "mcrl2/data/function_sort.h"
#include "mcrl2/data/nat.h"
#include "mcrl2/data/parse.h"
#include "mcrl2/data/print.h"
#include "mcrl2/data/rewriter.h"
#include "mcrl2/data/rewriters/simplify_rewriter.h"
#include "mcrl2/utilities/text_utility.h"
#include <boost/test/minimal.hpp>
#include <iostream>
#include <memory>
#include <set>
#include <string>

using namespace mcrl2;
using namespace mcrl2::core;
using namespace mcrl2::data;
using namespace mcrl2::data::detail;

data::rewriter make_data_rewriter(const data_specification& data_spec)
{
  data::rewriter datar(data_spec);
  return datar;
}

struct A
{
  data::rewriter& r_;

  A(data::rewriter& r)
    : r_(r)
  { }
};

A make_A(data::rewriter& d)
{
  A result(d);
  return result;
}

void test1()
{
  using namespace mcrl2::data::sort_nat;

  std::string DATA_SPEC1 =
    "sort D = struct d1(Nat)?is_d1 | d2(arg2:Nat)?is_d2;\n"
    ;
  data_specification data = parse_data_specification(DATA_SPEC1);
  rewriter datar(data);
  variable x("x", sort_nat::nat());
  variable y("y", sort_nat::nat());
  variable z("z", sort_nat::nat());
  data_expression t = datar(greater(minimum(x,y), z));

  BOOST_CHECK(datar(plus(parse_data_expression("1"),
                         parse_data_expression("2"))) == datar(parse_data_expression("3")));

  // copy a rewriter
  data::rewriter datar1 = datar;
  t = datar1(greater(minimum(x,y), z));

  // rewriter as return value
  data::rewriter datar2 = make_data_rewriter(data);
  t = datar2(greater(minimum(x,y), z));

  A a(datar);
  data_expression qa = a.r_(t);

  A b = a;
  data_expression qb = b.r_(t);

  A c = make_A(datar);
  data_expression qc = c.r_(t);
}

void test2()
{
  using namespace mcrl2::data::sort_nat;

  data_specification data_spec;

  data_spec.add_context_sort(nat());

  rewriter r(data_spec);
  data_expression d1 = parse_data_expression("2+7");
  data_expression d2 = parse_data_expression("4+5");
  BOOST_CHECK(r(d1) == r(d2));

  data::variable_list variables = parse_variables("m, n: Pos;");
  data::rewriter::substitution_type sigma;
  sigma[atermpp::down_cast<variable>(parse_data_expression("m", variables))] = r(parse_data_expression("3"));
  sigma[atermpp::down_cast<variable>(parse_data_expression("n", variables))] = r(parse_data_expression("4"));

  // Rewrite two data expressions, and check if they are the same
  d1 = parse_data_expression("m+n", variables);
  d2 = parse_data_expression("7");
  BOOST_CHECK(r(d1, sigma) == r(d2));
}

template <typename Rewriter>
void test_expressions(Rewriter R, std::string const& expr1, std::string const& expr2, std::string const& declarations, const data_specification& data_spec, const std::string& substitution_text)
{
  data::rewriter::substitution_type sigma;
  data::detail::parse_substitution(substitution_text, sigma, data_spec);
  data_expression d1 = parse_data_expression(expr1, parse_variables(declarations), data_spec);
  data_expression d2 = parse_data_expression(expr2, parse_variables(declarations), data_spec);
  data_expression rd1 = R(d1, sigma);
  data_expression rd2 = R(d2);
  if (rd1 != rd2)
  {
    BOOST_CHECK(rd1 == rd2);
    std::cout << "--- failed test --- " << expr1 << " -> " << expr2 << std::endl;
    std::cout << "d1           " << d1 << std::endl;
    std::cout << "d2           " << d2 << std::endl;
    std::cout << "sigma\n      " << sigma << std::endl;
    std::cout << "R(d1, sigma) " << rd1 << std::endl;
    std::cout << "R(d2)        " << rd2 << std::endl;
  }
}

void test4()
{
  data_specification data_spec;

  data::rewriter R(data_spec);

  std::string expr1 = "exists b: Bool. if(c, c, b)";
  std::string expr2 = "exists b: Bool. if(true, true, b)";
  std::string sigma = "[c: Bool := true]";
  test_expressions(R, expr1, expr2, "c: Bool;", data_spec, sigma);
}

void test5() // Test set difference for finite sets.
{
  std::string DATA_SPEC1 =
    "map f,g:FSet(Bool);\n"
    "eqn f=({false}-{true})-{false};\n"
    "    g={};\n"
    ;
  data_specification data_spec = parse_data_specification(DATA_SPEC1);
  data::rewriter R(data_spec);

  std::string expr1 = "f";
  std::string expr2 = "g";
  std::string sigma = "[c: Bool := true]";
  test_expressions(R, expr1, expr2, "c: Bool;", data_spec, sigma);
}

void allocation_test()
{
  data_specification data_spec;
  std::shared_ptr< data::rewriter > R_heap(new data::rewriter(data_spec));
  data::rewriter                    R_stack(data_spec);

  R_stack(parse_data_expression("1 == 2"));
  R_stack(parse_data_expression("1 == 2"));

  (*R_heap)(parse_data_expression("1 == 2"));
  (*R_heap)(parse_data_expression("1 == 2"));
}

void one_point_rule_preprocessor_test()
{
  using namespace data::detail;

  test_rewriters(N(one_point_rule_preprocessor()), N(I), "!(true && false)", "!(true) || !(false)");
  test_rewriters(N(one_point_rule_preprocessor()), N(I), "!!b", "b");
  test_rewriters(N(one_point_rule_preprocessor()), N(I), "!!!!b", "b");
  test_rewriters(N(one_point_rule_preprocessor()), N(I), "!(b1 && b2 && b3)", "!b1 || !b2 || !b3");
  test_rewriters(N(one_point_rule_preprocessor()), N(I), "!(b1 || b2 || b3)", "!b1 && !b2 && !b3");
  test_rewriters(N(one_point_rule_preprocessor()), N(I), "!(b1 == b2)", "b1 != b2");
  test_rewriters(N(one_point_rule_preprocessor()), N(I), "!(b1 != b2)", "b1 == b2");
  test_rewriters(N(one_point_rule_preprocessor()), N(I), "!(b1 == b2 && b1 == b3)", "b1 != b2 || b1 != b3");
  test_rewriters(N(one_point_rule_preprocessor()), N(I), "!(b1 != b2 || b1 != b3)", "b1 == b2 && b1 == b3");
  test_rewriters(N(one_point_rule_preprocessor()), N(I), "!(n1 != n2 || n1 != n2 + 1)", "n1 == n2 && n1 == n2 + 1");
}

void simplify_rewriter_test()
{
  data::simplify_rewriter R;
  test_rewriters(N(R), N(I), "!(true && false)", "true");
  test_rewriters(N(R), N(I), "exists n:Nat, b:Bool. b", "exists b:Bool. b");
  test_rewriters(N(R), N(I), "forall b:Bool. !!b", "forall b:Bool. b");
}

// The testcase below corresponds to ticket #1426.
void test_lambda_expression()
{
  data_specification data_spec;
  data::rewriter R(data_spec);

  std::string expr1 = "(lambda t: Real. true)(t) && (lambda s1: Pos, t: Real. 1 == s1)(2, u)";
  std::string expr2 = "false";
  std::string sigma = "[]";
  test_expressions(R, expr1, expr2, "t,u: Real;", data_spec, sigma);
}

// The testcase below corresponds to ticket #1428 which indicated
// that the generated rules for equality were erroneous. Their bound variables were equal
// f == g = forall x0,x0: Nat. f(x0, x0) == g(x0, x0)
void test_equality_on_functions()
{
  std::string DATA_SPEC1 =
    "map f,g:Nat#Nat->Nat;\n"
    "var x,y:Nat; \n"
    "eqn f(x,y) = x+y;\n"
    "    g(x,y) = x+x;\n"
    ;

  data_specification data_spec = parse_data_specification(DATA_SPEC1);
  data::rewriter R(data_spec);

  std::string expr1 = "f==g";
  std::string expr2 = "false";
  std::string sigma = "[]";
  test_expressions(R, expr1, expr2, "", data_spec, sigma);
}

// The testcase below corresponds to ticket #1461 which indicated
// that rewriting enumerated functions failed as the if function would be removed
// from the rewriters if it does not occur explicitly in the specification.
void test_enumeration_of_functions()
{
  std::string DATA_SPEC1 =
    "sort Hat = struct black | white;"
    ;

  data_specification data_spec = parse_data_specification(DATA_SPEC1);
  // For this test it is essential that unnecessary equations are removed.
  data::rewriter R(data_spec,used_data_equation_selector(data_spec,std::set<function_symbol>()));

  std::string expr1 = "exists f:Hat->Hat.f(black)!=f(white)";
  std::string expr2 = "true";
  std::string sigma = "[]";
  test_expressions(R, expr1, expr2, "", data_spec, sigma);
}

int test_main(int argc, char** argv)
{
  test1();
  test2();
  test4();
  test5();
  one_point_rule_preprocessor_test();
  simplify_rewriter_test();
  test_lambda_expression();
  test_equality_on_functions();
  test_enumeration_of_functions();

  return 0;
}


// this must define the same symbol as the main module file (numpy requirement)
#define PY_ARRAY_UNIQUE_SYMBOL andres_graph_PyArray_API
#define NO_IMPORT_ARRAY

#include <Python.h>
#include <boost/python.hpp>
#include <vigra/numpy_array.hxx>
#include <vigra/numpy_array_converters.hxx>

#include <vector>

#include "andres/graph/grid-graph.hxx"
#include "andres/graph/graph.hxx"
#include "andres/graph/multicut-lifted/lifted_mc_model.hxx"
#include "andres/graph/multicut-lifted/greedy-additive.hxx"

#include "lifted_mc_solver_exporter.hxx"

namespace bp = boost::python;
namespace agraph = andres::graph;

template<class MODEL, class SETTINGS>
SETTINGS * gaSettingsFactory(
    const MODEL & model
){
    return new SETTINGS();
}

template<class MODEL, class SETTINGS, class SOLVER>
SOLVER * gaSolverFactory(
    const MODEL & model,
    const SETTINGS & settings
){
    return new SOLVER(model, settings);
}

template<class GRAPH>
void exportLiftedGaT(const std::string graphName){

    typedef agraph::multicut_lifted::LiftedMcModel<GRAPH, float> LiftedModel;
    typedef agraph::multicut_lifted::GreedyAdditiveEdgeContraction<LiftedModel> Solver;
    typedef typename Solver::Settings Settings;

    // settings
    const auto settingsClsName = std::string("LiftedGreedyAdditiveSettings") + graphName;
    bp::class_<Settings>(settingsClsName.c_str(),bp::init<>())
        .def_readwrite("nodeStopCond", &Settings::nodeStopCond)
    ;
    // factory
    bp::def("settingsLiftedGreedyAdditive", &gaSettingsFactory<LiftedModel,Settings>,
            bp::return_value_policy<bp::manage_new_object>());

    // solver
    const auto solverClsName = std::string("LiftedGreedyAdditive") + graphName;
    bp::class_<Solver>(solverClsName.c_str(),bp::no_init)
        .def(LiftedMcSolverExporter<Solver>("LiftedGreedyAdditive"))
    ;
        // factory
    bp::def("liftedGreedyAdditive", &gaSolverFactory<LiftedModel,Settings, Solver>,
        bp::with_custodian_and_ward_postcall< 0,1 ,
            bp::return_value_policy<   bp::manage_new_object>  >()
    );
}

void exportLiftedGa(){
    typedef agraph::GridGraph<2> GridGraph2D;
    typedef agraph::GridGraph<3> GridGraph3D;
    typedef agraph::Graph<> Graph;

    exportLiftedGaT<GridGraph2D>("GridGraph2D");
    exportLiftedGaT<GridGraph3D>("GridGraph3D");
    exportLiftedGaT<Graph>("Graph");

}


// -*- mode: c++; c-basic-offset: 4; indent-tabs-mode: nil -*-

#include <cmath>
#include <vector>
#include <boost/assign/std/vector.hpp>
#include "lin_time.h" // GetElapsedTime
#include "cmd/unit_generic.h"
#include "cmd/unit_util.h"
#include "gfxlib.h"
#include "viewarea.h"
#include "bubble_display.h"

namespace
{

float GetDangerRate(Radar::Sensor::ThreatLevel::Value threat)
{
    using namespace Radar;

    switch (threat)
    {
    case Sensor::ThreatLevel::High:
        return 20.0; // Fast pulsation

    case Sensor::ThreatLevel::Medium:
        return 7.5; // Slow pulsation

    default:
        return 0.0; // No pulsation
    }
}

} // anonymous namespace

namespace Radar
{

BubbleDisplay::BubbleDisplay()
    : innerSphere(0.45),
      outerSphere(1.0),
      sphereZoom(1.0),
      radarTime(0.0),
      currentTargetMarkerSize(0.0),
      lastAnimationTime(0.0)
{
    using namespace boost::assign; // vector::operator+=
    explodeSequence += 0.0, 0.0001, 0.0009, 0.0036, 0.0100, 0.0225, 0.0441, 0.0784, 0.1296, 0.2025, 0.3025, 0.4356, 0.6084, 0.8281, 1.0, 0.8713, 0.7836, 0.7465, 0.7703, 0.8657, 1.0, 0.9340, 0.9595, 1.0, 0.9659, 1.0;
    implodeSequence += 1.0, 0.9999, 0.9991, 0.9964, 0.9900, 0.9775, 0.9559, 0.9216, 0.8704, 0.7975, 0.6975, 0.5644, 0.3916, 0.1719, 0.0, 0.1287, 0.2164, 0.2535, 0.2297, 0.1343, 0.0, 0.0660, 0.0405, 0.0, 0.0341, 0.0;
}

void BubbleDisplay::PrepareAnimation(const ZoomSequence& sequence)
{
    AnimationItem firstItem;
    firstItem.duration = 0.0;
    firstItem.sphereZoom = sequence[0];
    animation.push(firstItem);

    float duration = 2.0;
    for (ZoomSequence::size_type i = 1; i < sequence.size(); ++i)
    {
        AnimationItem item;
        item.duration = duration;
        item.sphereZoom = sequence[i];
        animation.push(item);
        duration = 0.05;
    }
}

void BubbleDisplay::OnDockEnd()
{
    PrepareAnimation(explodeSequence);
}

void BubbleDisplay::OnJumpBegin()
{
    PrepareAnimation(implodeSequence);
}

void BubbleDisplay::OnJumpEnd()
{
    PrepareAnimation(explodeSequence);
}

void BubbleDisplay::Animate()
{
    if (!animation.empty())
    {
        if (radarTime > lastAnimationTime + animation.front().duration)
        {
            sphereZoom = animation.front().sphereZoom;
            animation.pop();
            lastAnimationTime = radarTime;
        }
    }
}

void BubbleDisplay::Draw(const Sensor& sensor,
                         VSSprite *frontSprite,
                         VSSprite *rearSprite)
{
    assert(frontSprite || rearSprite); // There should be at least one radar display

    radarTime += GetElapsedTime();

    leftRadar.SetSprite(frontSprite);
    rightRadar.SetSprite(rearSprite);

    if (frontSprite)
        frontSprite->Draw();
    if (rearSprite)
        rearSprite->Draw();

    Sensor::TrackCollection tracks = sensor.FindTracksInRange();

    Animate();

    GFXEnable(DEPTHTEST);
    GFXEnable(DEPTHWRITE);
    GFXEnable(SMOOTH);

    for (Sensor::TrackCollection::const_iterator it = tracks.begin(); it != tracks.end(); ++it)
    {
        if (it->GetPosition().z < 0)
        {
            // Draw tracks behind the ship
            DrawTrack(sensor, rightRadar, *it);
        }
        else
        {
            // Draw tracks in front of the ship
            DrawTrack(sensor, leftRadar, *it);
        }
    }

    DrawBackground(leftRadar, currentTargetMarkerSize);
    DrawBackground(rightRadar, currentTargetMarkerSize);

    GFXPointSize(1);
    GFXDisable(DEPTHTEST);
    GFXDisable(DEPTHWRITE);
    GFXDisable(SMOOTH);
}

void BubbleDisplay::DrawTrack(const Sensor& sensor,
                              const ViewArea& radarView,
                              const Track& track)
{
    if (!radarView.IsActive())
        return;

    GFXColor color = sensor.GetColor(track);

    Vector position = track.GetPosition();
    if (position.z < 0)
        position.z = -position.z;

    float magnitude = position.Magnitude();
    float scaleFactor = 0.0; // [0; 1] where 0 = border, 1 = center
    float maxRange = sensor.GetMaxRange();
    if (magnitude <= maxRange)
    {
        // [innerSphere; outerSphere]
        scaleFactor = (outerSphere - innerSphere) * ((maxRange - magnitude) / maxRange);
        magnitude /= (1.0 - scaleFactor);
    }

    if (sensor.InsideNebula())
    {
        magnitude /= (1.0 - 0.04 * Jitter(0.0, 1.0));
    }
    Vector scaledPosition = sphereZoom * Vector(-position.x, position.y, position.z) / magnitude;

    Vector head = radarView.Scale(scaledPosition);

    GFXColor headColor = color;

    headColor.a *= 0.2 + scaleFactor * (1.0 - 0.2); // [0;1] => [0.1;1]
    if (sensor.UseThreatAssessment())
    {
        float dangerRate = GetDangerRate(sensor.IdentifyThreat(track));
        if (dangerRate > 0.0)
        {
            // Blinking blip
            headColor.a *= cosf(dangerRate * radarTime);
        }
    }

    // Fade out dying ships
    if (track.IsExploding())
    {
        headColor.a *= (1.0 - track.ExplodingProgress());
    }

    float trackSize = std::max(1.0f, std::log10(track.GetSize()));
    if (track.GetType() != Track::Type::Cargo)
        trackSize += 1.0;

    GFXColorf(headColor);
    if (sensor.IsTracking(track))
    {
        currentTargetMarkerSize = trackSize;
        DrawTargetMarker(head, trackSize);
    }

    const bool isNebula = (track.GetType() == Track::Type::Nebula);
    const bool isEcmActive = track.HasActiveECM();
    if (isNebula || isEcmActive)
    {
        // Vary size between 50% and 150%
        trackSize *= Jitter(0.5, 1.0);
    }

    GFXPointSize(trackSize);
    GFXBegin(GFXPOINT);
    GFXVertexf(head);
    GFXEnd();
}

void BubbleDisplay::DrawTargetMarker(const Vector& position, float trackSize)
{
    // Split octagon
    float size = 3.0 * std::max(trackSize, 3.0f);
    float xsize = size / g_game.x_resolution;
    float ysize = size / g_game.y_resolution;
    GFXLineWidth(1);
    GFXBegin(GFXLINESTRIP);
    GFXVertex3f(position.x - xsize / 2, position.y - ysize, position.z);
    GFXVertex3f(position.x - xsize, position.y - ysize / 2, position.z);
    GFXVertex3f(position.x - xsize, position.y + ysize / 2, position.z);
    GFXVertex3f(position.x - xsize / 2, position.y + ysize, position.z);
    GFXEnd();
    GFXBegin(GFXLINESTRIP);
    GFXVertex3f(position.x + xsize / 2, position.y - ysize, position.z);
    GFXVertex3f(position.x + xsize, position.y - ysize / 2, position.z);
    GFXVertex3f(position.x + xsize, position.y + ysize / 2, position.z);
    GFXVertex3f(position.x + xsize / 2, position.y + ysize, position.z);
    GFXEnd();
}

void BubbleDisplay::DrawBackground(const ViewArea& radarView, float trackSize)
{
    if (!radarView.IsActive())
        return;

    GFXColor groundColor = radarView.GetColor();

    // Split octagon
    float size = 3.0 * std::max(trackSize, 3.0f);
    float xground = size / g_game.x_resolution;
    float yground = size / g_game.y_resolution;
    Vector center = radarView.Scale(Vector(0.0, 0.0, 0.0));

    GFXColorf(groundColor);
    GFXLineWidth(1);
    GFXBegin(GFXLINESTRIP);
    GFXVertex3f(center.x - xground, center.y - yground / 2, center.z);
    GFXVertex3f(center.x - xground / 2, center.y - yground, center.z);
    GFXVertex3f(center.x + xground / 2, center.y - yground, center.z);
    GFXVertex3f(center.x + xground, center.y - yground / 2, center.z);
    GFXEnd();
    GFXBegin(GFXLINESTRIP);
    GFXVertex3f(center.x - xground, center.y + yground / 2, center.z);
    GFXVertex3f(center.x - xground / 2, center.y + yground, center.z);
    GFXVertex3f(center.x + xground / 2, center.y + yground, center.z);
    GFXVertex3f(center.x + xground, center.y + yground / 2, center.z);
    GFXEnd();
}

} // namespace Radar


//          Copyright Gavin Band 2008 - 2012.
// Distributed under 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)

#include <string>
#include <memory>
#include <boost/noncopyable.hpp>
#include <boost/function.hpp>
#include <Eigen/Core>
#include "genfile/VariantIdentifyingData.hpp"
#include "genfile/VariantEntry.hpp"
#include "genfile/VariantDataReader.hpp"
#include "appcontext/OptionProcessor.hpp"
#include "components/SNPSummaryComponent/InfoComputation.hpp"

namespace stats {
	void InfoComputation::operator()(
		VariantIdentifyingData const& snp,
		Genotypes const& genotypes,
		Ploidy const& ploidy,
		genfile::VariantDataReader&,
		ResultCallback callback
	) {
		m_computation.compute( snp, genotypes, ploidy ) ;
		callback( "info", m_computation.info() ) ;
		callback( "impute_info", m_computation.impute_info() ) ;
	}

	std::string InfoComputation::get_summary( std::string const& prefix, std::size_t column_width ) const {
		return prefix + "InfoComputation" ;
	}

	namespace impl {
		namespace {
			double compute_variance( Eigen::VectorXd const& levels, Eigen::VectorXd const& levels_squared, Eigen::VectorXd const& probs ) {
				double const mean = ( probs.transpose() * levels )(0) ;
				double const variance = ( probs.transpose() * levels_squared )(0) - ( mean * mean ) ;
				return variance ;
			}

			// treat the rows of the probabilities matrix as probabilities.
			// distribution for individual i is taken as a mixture of the distribution given by row i of probabilities,
			// and the fallback distribution if the sum of row i < 1, appropriately weighted.
			// Only individuals with inclusion = 1 are used.
			template< typename Probabilities, typename InclusionIndicator >
			double compute_sum_of_variances( Eigen::VectorXd const& levels, Probabilities const& probabilities, Eigen::VectorXd const& fallback, InclusionIndicator const& inclusion ) {
				assert( levels.size() == fallback.size() ) ;
				assert( levels.size() == probabilities.cols() ) ;
				assert( inclusion.size() == probabilities.rows() ) ;
				Eigen::VectorXd levels_squared = ( levels.array() * levels.array() ) ;

				double result = 0.0 ;
				for( int i = 0; i < probabilities.rows(); ++i ) {
					if( inclusion( i ) == 1 ) {
						double const c = probabilities.row( i ).sum() ;
						result += compute_variance( levels, levels_squared, probabilities.row( i ).transpose() + ( 1 - c ) * fallback ) ;
					}
				}
				return result ;
			}

		}

		InfoComputation::InfoComputation():
			m_diploid_fallback_distribution( Eigen::VectorXd::Zero( 3 )),
			m_haploid_fallback_distribution( Eigen::VectorXd::Zero( 2 )),
			m_diploid_levels( Eigen::VectorXd::LinSpaced( 3, 0, 2 )),
			m_haploid_levels( Eigen::VectorXd::LinSpaced( 2, 0, 1 ))
		{}

		void InfoComputation::compute(
			VariantIdentifyingData const& snp,
			Genotypes const& genotypes,
			Ploidy const& ploidy
		) {
			compute(
				snp, genotypes, ploidy,
				std::vector< metro::SampleRange >( 1, metro::SampleRange( 0, genotypes.rows()) )
			) ;
		}

		void InfoComputation::compute(
			VariantIdentifyingData const& snp,
			Genotypes const& genotypes,
			Ploidy const& ploidy,
			std::vector< metro::SampleRange > const& included_samples
		) {
			m_info = std::numeric_limits< double >::quiet_NaN() ;
			m_impute_info = std::numeric_limits< double >::quiet_NaN() ;

			// we don't compute for multiallelics currently
			if( snp.number_of_alleles() == 2 ) {
				compute_impl( snp, genotypes, ploidy, included_samples ) ;
			}
		}

		void InfoComputation::compute_impl(
			VariantIdentifyingData const& snp,
			Genotypes const& genotypes,
			Ploidy const& ploidy,
			std::vector< metro::SampleRange > const& included_samples
		) {
			typedef Eigen::Block< Genotypes const > GenotypeBlock ;
			typedef Eigen::VectorBlock< Eigen::VectorXd const > PloidyBlock ;

			// This function assumes every sample is haploid or diploid.
			Eigen::VectorXd const haploids = ( ploidy.array() == 1 ).cast< double >() ;
			Eigen::VectorXd const diploids = ( ploidy.array() == 2 ).cast< double >() ;

			double b_allele_count = 0 ;
			double total_allele_count = 0 ;
			double autosomal_b_allele_count = 0 ;		// for computations that pretend ploidy == 2 everywhere.
			double autosomal_total_allele_count = 0 ;	// for computations that pretend ploidy == 2 everywhere.

			double number_of_haploid_info_samples = 0 ;
			double number_of_diploid_info_samples = 0 ;
			double total_probability = 0 ;
			for( std::size_t i = 0; i < included_samples.size(); ++i ) {
				GenotypeBlock const block = genotypes.block( included_samples[i].begin(), 0, included_samples[i].end() - included_samples[i].begin(), genotypes.cols() ) ;
				PloidyBlock const haploid_block = haploids.segment( included_samples[i].begin(), included_samples[i].size() ) ;
				PloidyBlock const diploid_block = diploids.segment( included_samples[i].begin(), included_samples[i].size() ) ;

				b_allele_count += (( block.col( 1 ) + 2.0 * block.col( 2 ) ).array() * diploid_block.array() ).sum() ;
				b_allele_count += (( block.col( 1 ) ).array() * haploid_block.array() ).sum() ;

				total_allele_count += 2 * ( block.rowwise().sum().array() * diploid_block.array() ).sum() ;
				total_allele_count += ( block.rowwise().sum().array() * haploid_block.array() ).sum() ;

				autosomal_b_allele_count += ( block.col( 1 ) + 2.0 * block.col( 2 ) ).array().sum() ;
				autosomal_total_allele_count += 2 * block.sum() ;

				number_of_haploid_info_samples += haploid_block.sum() ;
				number_of_diploid_info_samples += diploid_block.sum() ;
				total_probability += block.sum() ;
			}

			// MLE estimate of allele frequency
			double const theta_mle = b_allele_count / total_allele_count ;
			double const autosomal_theta_mle = autosomal_b_allele_count / autosomal_total_allele_count ;
			double const theta_est = theta_mle ;

			m_diploid_fallback_distribution( 0 ) = ( 1 - theta_mle ) * ( 1 - theta_mle ) ;
			m_diploid_fallback_distribution( 1 ) = 2.0 * theta_mle * ( 1 - theta_mle ) ;
			m_diploid_fallback_distribution( 2 ) = theta_mle * theta_mle ;

			m_haploid_fallback_distribution( 0 ) = 1 - theta_mle ;
			m_haploid_fallback_distribution( 1 ) = theta_mle ;

			//std::cerr << "theta = " << theta_mle << ", fallback_distribution = " << fallback_distribution.transpose() << ".\n" ;

			double diploid_info_term = 0.0 ;
			double haploid_info_term = 0.0 ;
			double impute_info_term = 0.0 ;
			for( std::size_t i = 0; i < included_samples.size(); ++i ) {
				GenotypeBlock const diploid_genotype_block = genotypes.block( included_samples[i].begin(), 0, included_samples[i].size(), genotypes.cols() ) ;
				GenotypeBlock const haploid_genotype_block = genotypes.block( included_samples[i].begin(), 0, included_samples[i].size(), 2 ) ;
				PloidyBlock const haploid_block = haploids.segment( included_samples[i].begin(), included_samples[i].size() ) ;
				PloidyBlock const diploid_block = diploids.segment( included_samples[i].begin(), included_samples[i].size() ) ;

				haploid_info_term -= compute_sum_of_variances(
					m_haploid_levels, haploid_genotype_block,
					m_haploid_fallback_distribution, haploid_block
				)
				/ ( theta_mle * ( 1 - theta_mle ) ) ;

				diploid_info_term -= compute_sum_of_variances(
					m_diploid_levels, diploid_genotype_block,
					m_diploid_fallback_distribution, diploid_block
				)
				/ ( 2.0 * theta_mle * ( 1 - theta_mle ) ) ;

				impute_info_term -= compute_sum_of_variances(
					m_diploid_levels, diploid_genotype_block,
					Eigen::VectorXd::Zero( 3 ), Eigen::VectorXd::Constant( diploid_genotype_block.rows(), 1 )
				)
				/ ( 2.0 * autosomal_theta_mle * ( 1 - autosomal_theta_mle ) ) ;
			}

			double info = std::numeric_limits< double >::quiet_NaN() ;
			double impute_info = std::numeric_limits< double >::quiet_NaN() ;

			if(( number_of_haploid_info_samples + number_of_diploid_info_samples ) > 0 ) {
				info = 1.0 + (haploid_info_term + diploid_info_term) / ( number_of_haploid_info_samples + number_of_diploid_info_samples ) ;
			}

			if( total_probability > 0 ) {
				impute_info = 1.0 + impute_info_term / total_probability ;
			}

			m_info = info ;
			m_impute_info = impute_info ;
		}

	}
}