//=======================================================================
// Copyright (c) 2013 Piotr Wygocki
//               2014 Piotr Godlewski
// 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 iterative_rounding.hpp
 * @brief
 * @author Piotr Wygocki, Piotr Godlewski
 * @version 1.0
 * @date 2013-05-06
 */
#ifndef PAAL_ITERATIVE_ROUNDING_HPP
#define PAAL_ITERATIVE_ROUNDING_HPP

#include "paal/iterative_rounding/ir_components.hpp"
#include "paal/lp/glp.hpp"
#include "paal/utils/floating.hpp"
#include "paal/utils/type_functions.hpp"
#include "paal/utils/irange.hpp"

#include <boost/optional.hpp>

#include <cstdlib>
#include <unordered_map>

namespace paal {
namespace ir {

/**
 * @brief Default Iterative Rounding visitor.
 */
struct trivial_visitor {
    /**
     * @brief Method called after (re)solving the LP.
     */
    template <typename Problem, typename LP>
    void solve_lp(Problem &problem, LP &lp) {}

    /**
     * @brief Method called after rounding a column of the LP.
     */
    template <typename Problem, typename LP>
    void round_col(Problem &problem, LP &lp, lp::col_id col, double val) {}

    /**
     * @brief Method called after relaxing a row of the LP.
     */
    template <typename Problem, typename LP>
    void relax_row(Problem &problem, LP &lp, lp::row_id row) {}
};

///default solve lp for row_generation,
///at first call PRIMAL, and DUAL on the next calls
template <typename Problem, typename LP>
class default_solve_lp_in_row_generation {
    bool m_first;
    LP & m_lp;
public:
    ///constructor
    default_solve_lp_in_row_generation(Problem &, LP & lp) : m_first(true), m_lp(lp) {}

    ///operator()
    lp::problem_type operator()()
    {
        if (m_first) {
            m_first = false;
            return m_lp.solve_simplex(lp::PRIMAL);
        }
        return m_lp.resolve_simplex(lp::DUAL);
    }
};

/// default row_generation for lp,
/// one can customize LP solving, by setting SolveLP
template <template <class, class> class SolveLP = default_solve_lp_in_row_generation>
struct row_generation_solve_lp {
///operator()
template <class Problem, class LP>
    auto operator()(Problem & problem, LP & lp) {
        return row_generation(problem.get_find_violation(lp), SolveLP<Problem, LP>(problem, lp));
    }
};

namespace detail {

    /**
     * @brief This class solves an iterative rounding problem.
 *
 * @tparam Problem
 * @tparam IRcomponents
 * @tparam Visitor
 * @tparam LP
 */
template <typename Problem, typename IRcomponents, typename Visitor = trivial_visitor, typename LP = lp::glp>
class iterative_rounding  {
    using RoundedCols = std::unordered_map<lp::col_id, std::pair<double, double>>;

    /**
     * @brief Returns the current value of the LP column.
     */
    double get_val(lp::col_id col) const {
        auto i = m_rounded.find(col);
        if (i == m_rounded.end()) {
            return m_lp.get_col_value(col);
        } else {
            return i->second.first;
        }
    }

  public:
    /**
     * @brief Constructor.
     */
    iterative_rounding(Problem &problem, IRcomponents e,
                       Visitor vis = Visitor())
        : m_ir_components(std::move(e)), m_visitor(std::move(vis)),
          m_problem(problem) {
        call<Init>(m_problem, m_lp);
    }

    /**
     * @brief Finds solution to the LP.
     *
     * @return LP solution status
     */
    lp::problem_type solve_lp() {
        auto prob_type = call<SolveLP>(m_problem, m_lp);
        assert(prob_type != lp::UNDEFINED);
        m_visitor.solve_lp(m_problem, m_lp);
        return prob_type;
    }

    /**
     * @brief Finds solution to the LP.
     *
     * @return LP solution status
     */
    lp::problem_type resolve_lp() {
        auto prob_type = call<ResolveLP>(m_problem, m_lp);
        assert(prob_type != lp::UNDEFINED);
        m_visitor.solve_lp(m_problem, m_lp);
        return prob_type;
    }

    /**
     * @brief Returns the solution cost based on the LP values.
     */
    double get_solution_cost() {
        double sol_cost(0);
        for (auto col : m_lp.get_columns()) {
            sol_cost += m_lp.get_col_value(col) * m_lp.get_col_coef(col);
        }
        for (auto rounded_col : m_rounded) {
            sol_cost += rounded_col.second.first * rounded_col.second.second;
        }
        return sol_cost;
    }

    /**
     * @brief Rounds the LP columns (independently) using the RoundCondition
     * component.
     *
     * @return true iff at least one column was rounded
     */
    bool round() {
        int deleted(0);
        auto &&cols = m_lp.get_columns();
        auto cbegin = std::begin(cols);
        auto cend = std::end(cols);

        while (cbegin != cend) {
            lp::col_id col = *cbegin;
            auto do_round = call<RoundCondition>(m_problem, m_lp, col);
            if (do_round) {
                ++deleted;
                m_rounded.insert(std::make_pair(col,
                    std::make_pair(*do_round, m_lp.get_col_coef(col))));
                m_visitor.round_col(m_problem, m_lp, col, *do_round);
                cbegin = delete_column(cbegin, *do_round);
            }
            else {
                ++cbegin;
            }
        }

        return deleted > 0;
    }

    /**
     * @brief Relaxes the LP rows using the RelaxCondition component.
     *
     * @return true iff at least one row was relaxed
     */
    bool relax() {
        int deleted(0);
        auto &&rows = m_lp.get_rows();
        auto rbegin = std::begin(rows);
        auto rend = std::end(rows);

        while (rbegin != rend) {
            lp::row_id row = *rbegin;
            if (call<RelaxCondition>(m_problem, m_lp, row)) {
                ++deleted;
                m_visitor.relax_row(m_problem, m_lp, row);
                rbegin = m_lp.delete_row(rbegin);
                if (call<RelaxationsLimit>(deleted)) {
                    break;
                }
            } else {
                ++rbegin;
            }
        }

        return deleted > 0;
    }

    /**
     * @brief Returns the LP object used to solve the IR.
     */
    LP &get_lp() { return m_lp; }

    /**
     * @brief Returns the IR components.
     */
    IRcomponents &get_ir_components() { return m_ir_components; }

    /**
     * @brief Sets the solution to the problem using SetSolution component.
     */
    void set_solution() {
        call<SetSolution>(m_problem, boost::bind(&iterative_rounding::get_val,
                                               this, _1));
    }

    /**
     * @brief Rounds the LP using the RoundCondition component.
     */
    void dependent_round() { call<RoundCondition>(m_problem, m_lp); }

    /**
     * @brief Checks if the IR problem has been solved, using the StopCondition
    * component.
     *
     * @return true iff the problem has been solved
     */
    bool stop_condition() { return call<StopCondition>(m_problem, m_lp); }

  private:
    template <typename Action, typename... Args>
    auto call(Args &&... args)
        ->decltype(std::declval<IRcomponents>().template call<Action>(
              std::forward<Args>(args)...)) {
        return m_ir_components.template call<Action>(
            std::forward<Args>(args)...);
    }

    /// Deletes a column from the LP and adjusts the row bounds.
    typename LP::ColIter
    delete_column(typename LP::ColIter col_iter, double value) {
        auto column = m_lp.get_rows_in_column(*col_iter);
        lp::row_id row;
        double coef;
        for (auto const &c : column) {
            boost::tie(row, coef) = c;
            double ub = m_lp.get_row_upper_bound(row);
            double lb = m_lp.get_row_lower_bound(row);
            double diff = coef * value;
            m_lp.set_row_upper_bound(row, ub - diff);
            m_lp.set_row_lower_bound(row, lb - diff);
        }
        return m_lp.delete_col(col_iter);
    };

    LP m_lp;
    IRcomponents m_ir_components;
    Visitor m_visitor;
    utils::compare<double> m_compare;
    RoundedCols m_rounded;
    Problem &m_problem;
};

} //detail

/// Iterative Rounding solution cost type. Solution cost only makes sense if the LP has been solved to optimal value.
using IRSolutionCost = boost::optional<double>;
/// Iterative Rounding result type: Pair consisting of LP problem type and IR
/// solution cost.
using IRResult = std::pair<lp::problem_type, IRSolutionCost>;

/**
 * @brief Solves an Iterative Rounding problem.
 *
 * @tparam Problem
 * @tparam IRcomponents
 * @tparam Visitor
 * @tparam LP
 * @param problem IR problem
 * @param components IR problem components
 * @param visitor visitor object used for logging progress of the algoithm
 */
template <typename Problem, typename IRcomponents, typename Visitor = trivial_visitor, typename LP = lp::glp>
IRResult solve_iterative_rounding(Problem & problem, IRcomponents components, Visitor visitor = Visitor()) {
    detail::iterative_rounding<Problem, IRcomponents, Visitor, LP> ir(problem, std::move(components), std::move(visitor));

    auto prob_type = ir.solve_lp();
    if (prob_type != lp::OPTIMAL) {
        return IRResult(prob_type, IRSolutionCost{});
    }

    while (!ir.stop_condition()) {
        bool rounded{ ir.round() };
        bool relaxed{ ir.relax() };
        assert(rounded || relaxed);

        prob_type = ir.resolve_lp();
        if (prob_type != lp::OPTIMAL) {
            return IRResult(prob_type, IRSolutionCost{});
        }
    }
    ir.set_solution();
    return IRResult(lp::OPTIMAL, IRSolutionCost(ir.get_solution_cost()));
}

/**
 * @brief Solves an Iterative Rounding problem with dependent rounding.
 *
 * @tparam Problem
 * @tparam IRcomponents
 * @tparam Visitor
 * @tparam LP
 * @param problem IR problem
 * @param components IR problem components
 * @param visitor visitor object used for logging progress of the algoithm
 */
template <typename Problem, typename IRcomponents, typename Visitor = trivial_visitor, typename LP = lp::glp>
IRResult solve_dependent_iterative_rounding(Problem & problem, IRcomponents components, Visitor visitor = Visitor()) {
    detail::iterative_rounding<Problem, IRcomponents, Visitor, LP> ir(problem, std::move(components), std::move(visitor));

    auto prob_type = ir.solve_lp();
    if (prob_type != lp::OPTIMAL) {
        return IRResult(prob_type, IRSolutionCost{});
    }

    while (!ir.stop_condition()) {
        ir.dependent_round();
        ir.relax();

        prob_type = ir.resolve_lp();
        if (prob_type != lp::OPTIMAL) {
            return IRResult(prob_type, IRSolutionCost{});
        }
    }
    ir.set_solution();
    return IRResult(lp::OPTIMAL, IRSolutionCost(ir.get_solution_cost()));
}

} // ir
} // paal

#endif // PAAL_ITERATIVE_ROUNDING_HPP


/*
 * Copyright (c) 2011, Mattia Penati <mattia.penati@gmail.com>
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without modification,
 * are permitted provided that the following conditions are met:
 *
 *     * Redistributions of source code must retain the above copyright notice,
 *       this list of conditions and the following disclaimer.
 *     * Redistributions in binary form must reproduce the above copyright notice,
 *       this list of conditions and the following disclaimer in the documentation
 *       and/or other materials provided with the distribution.
 *     * Neither the name of the Politecnico di Milano nor the names of its
 *       contributors may be used to endorse or promote products derived from
 *       this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
 * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR
 * ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
 * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
 * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
 * ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
 * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

#ifndef AMA_TENSOR_MUL_IEXP_OUTER_HPP
#define AMA_TENSOR_MUL_IEXP_OUTER_HPP 1

#include <ama/tensor/iexp/index_reorder.hpp>
#include <ama/tensor/iexp/iexp_base.hpp>
#include <ama/tensor/mul/mul_calculator.hpp>
#include <boost/mpl/bool.hpp>
#include <boost/mpl/plus.hpp>

namespace ama
{
  namespace tensor_
  {

    /* forward declaration*/
    template <typename LEFT, typename RIGHT> class mul_outer;

    /* traits definition */
    template <typename LEFT, typename RIGHT>
    struct iexp_traits< mul_outer<LEFT, RIGHT> >
    {
      typedef typename LEFT::value_type value_type;

      typedef typename LEFT::dimension_type dimension_type;

      typedef typename mul_index<LEFT, RIGHT>::controvariant_list controvariant_list;
      typedef typename mul_index<LEFT, RIGHT>::covariant_list covariant_list;

      typedef ::boost::mpl::false_ is_assignable;
    };

    /* class declaration */
    template <typename LEFT, typename RIGHT>
    class mul_outer:
        public iexp_base< mul_outer<LEFT, RIGHT> >
    {
    protected:
      typedef iexp_base< mul_outer<LEFT, RIGHT> > base_type;
      typedef mul_outer<LEFT, RIGHT> derived_type;

    protected:
      typedef LEFT left_operand_type;
      typedef RIGHT right_operand_type;

    public:
      typedef typename base_type::value_type value_type;

    public:
      /* costructor */
      mul_outer(left_operand_type const & left,
                 right_operand_type const & right)
          : m_left(left),
            m_right(right) { }

    public:
      /* retrieve the value */
      template <typename IMAP>
      value_type at() const
      {
        return (m_left.template at<IMAP>()) * (m_right.template at<IMAP>());
      }

    protected:
      /* members */
      left_operand_type const m_left;
      right_operand_type const m_right;
    };

  }
}

#endif /* AMA_TENSOR_MUL_IEXP_OUTER_HPP */


/** @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


#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;
}


//          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 ;
		}

	}
}


/*=========================================================================

  Program:   Small Body Geophysical Analysis
  Module:    SBGATObsRadar.hpp

  Class derived from VTK's vtkPolyDataAlgorithm by Benjamin Bercovici

  Copyright (c) Ken Martin, Will Schroeder, Bill Lorensen
  All rights reserved.
  See Copyright.txt or http://www.kitware.com/Copyright.htm for details.

     This software is distributed WITHOUT ANY WARRANTY; without even
     the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
     PURPOSE.  See the above copyright notice for more information.

=========================================================================*/
/**
@file SBGATObsRadar.hpp
@class  SBGATObsRadar
@author Benjamin Bercovici
@author Jay McMahon
@date October 2018
@brief  Computes range/range-rate Doppler images over the surface of
provided small body
@details Computes range/range-rate Doppler images over the surface of
provided small body
@copyright MIT License, Benjamin Bercovici and Jay McMahon

*/

#ifndef SBGATObsRadar_H
#define SBGATObsRadar_H

#include <vtkFiltersCoreModule.h> // For export macro
#include <vtkPolyDataAlgorithm.h>

#include <vtkModifiedBSPTree.h>
#include <vtkImageData.h>

#include <armadillo>
#include <array>

#include <SBGATObs.hpp>

/**
Vector of vector of radar observations. Each observation is comprised of (range,range-rate,incidence)
This incidence is not strictly speaking a measured quantity but is used to penalize the binned
observations, blurring out returns collected at a high incidence
*/
typedef typename std::vector<std::vector<std::array<double, 3> > > SBGATRadarObsSequence;

class VTKFILTERSCORE_EXPORT SBGATObsRadar : public SBGATObs{
public:
  /**
   * Constructs with initial values of zero.
   */
  static SBGATObsRadar *New();

  vtkTypeMacro(SBGATObsRadar,vtkPolyDataAlgorithm);
  void PrintSelf(std::ostream& os, vtkIndent indent) override;
  void PrintHeader(std::ostream& os, vtkIndent indent) override;
  void PrintTrailer(std::ostream& os, vtkIndent indent) override;

  /**
  Collects range/range-rate samples over the surface of the specified system of small bodies at a given time after the epoch. The radar source is positionned at 1e6 * l meters from the target's center of mass, where
  l is a measure of the first object's diagonal.
  @param measurements_sequence reference to MeasurementsSequence, holding collected range/range-rate measurements at each observation time
  @param time observation timestamp
  @param N minimum number of measurements to produce over the smallest facet in the shape. The number of samples for any other facet will be equal to N * facet_surface_area / smallest_facet_surface_area
  @param radar_dir unit direction towards radar from target in inertial frame
  @param positions_vec vector of positions of each target's center-of-mass expressed in the primary's body frame
  @param velocities_vec vector of inertial velocities of each target's center-of-mass expressed in the primary's body frame
  @param mrps_vec vector of MRPs defining the inertial-to-body DCM [BN]
  @param omegas_vec vector of angular velocities of each body expressed in the inertial frame
  @param penalize_incidence if true, each measurement will be weighed by the cos(incidence) angle between
  the sampled point and the radar squared. If false, all measurements (in view of the radar and not blocked) are weighed equally
  have the same weight
  */
  void CollectMeasurements(SBGATRadarObsSequence & measurements_sequence,
  const double & time,
  const int & N,
  const arma::vec & radar_dir,
  const std::vector<arma::vec> & positions_vec,
  const std::vector<arma::vec> & velocities_vec,
  const std::vector<arma::vec> & mrps_vec,
  const std::vector<arma::vec> & omegas_vec,
  const bool & penalize_indicence);

  /**
  Bins the provided measurements sequence into a series of 2d-histogram
  Will throw an std::runtime_error exception if
  - either of the provided bin sizes are invalid (i.e <= 0)
  - the provided bin sizes yield empty histogram dimensions
  @param measurements_sequence reference to MeasurementsSequence, holding collected range/range-rate measurements at each observation time
  @param r_bin range bin size (m)
  @param rr_bin range-rate bin size (m/s)
  */
  void BinObservations(
    const SBGATRadarObsSequence & measurements_sequence,
    const double & r_bin,
    const double & rr_bin);

  /**
  Save the binned radar images to PNGs in the prescribed folder.
  The images will be normalized by the largest value in the observation sequence
  @param savepath path to folder where images will be saved (ex: "output/")
  */
  void SaveImages(std::string savepath);

  /**
  Return vector holding the computed images
  @return vector of radar images
  */
  std::vector<vtkSmartPointer<vtkImageData>> GetImages() const;

  /**
  Clears images, if any
  */

  void ClearImages() { this -> images.clear();}

protected:
  SBGATObsRadar();
  ~SBGATObsRadar() override;

  /**
  Ray traces the facets in view to the radar and measurements sequence with range/range-rate data
  if sampled point was in view of the radar
  @param measurements_sequence Sequence of measurements to add new image data to
  @param facets_in_view reference to a vector of vector holding indices of (maybe) illuminated facets for all considered bodies
  @param radar_dir radar direction expressed in inertial frame
  @param BN_dcms_vec vector holding the DCMs orienting the body frame of each body w/r to inertial
  @param positions_vec vector holding the position vector of the CM of each body w/r to the primary
  @param velocities_vec vector holding the velocities vector of the CM of each body
  @param omega_vec vector holding the angular velocities vector of each body
  */
  void reverse_ray_trace(SBGATRadarObsSequence & measurements_sequence,
    const std::vector<std::vector<int> > & facets_in_view,
    const arma::vec & radar_dir,
    const int N,
    const bool penalize_indicence,
    const std::vector<arma::mat> & BN_dcms_vec,
    const std::vector<arma::vec> & positions_vec,
    const std::vector<arma::vec> & velocities_vec,
    const std::vector<arma::vec> & omega_vec);

  std::vector<vtkSmartPointer<vtkImageData>> images;

  double max_value;

private:
  SBGATObsRadar(const SBGATObsRadar&) = delete;
  void operator=(const SBGATObsRadar&) = delete;
};

#endif