#include "TensorLines.hh"

#include "TensorProductBezierTriangles.hh"
#include "TensorCoreLinesEvaluator.hh"
#include "TensorTopologyEvaluator.hh"
#include "ParallelEigenvectorsEvaluator.hh"

#include <Eigen/Eigenvalues>
#include <Eigen/LU>

#include <boost/range/algorithm/min_element.hpp>
#include <boost/range/algorithm_ext/insert.hpp>
#include <boost/optional.hpp>

#include <stack>
#include <queue>
#include <iterator>
#include <complex>
#include <type_traits>

using namespace cpp_utils;

namespace
{
using namespace tl;

/**
 * Representative Solution in a cluster of similar solutions
 */
template <typename Evaluator>
struct ClusterRepr
{
    std::size_t cluster_size;
    Evaluator eval;
};

/**
 * Cluster all triangles in a list that are closer than a given distance
 *
 * @param cands List of candidates generated by parallelEigenvectorSearch()
 * @param epsilon Maximum distance of triangles in a cluster
 *
 * @return List of clusters (each cluster is a list of candidates)
 */
template <typename CandList>
std::vector<CandList>
clusterTris(const CandList& cands, double epsilon)
{
    auto classes = std::vector<CandList>{};
    for(const auto& t : cands)
    {
        classes.push_back({t});
    }

    auto has_close_elements = [&](const CandList& c1,
                                  const CandList& c2) {
        if(c1 == c2) return false;
        for(const auto& t1 : c1)
        {
            for(const auto& t2 : c2)
            {
                if(distance(t1, t2) <= epsilon)
                {
                    return true;
                }
            }
        }
        return false;
    };

    auto changed = true;
    while(changed)
    {
        changed = false;
        for(auto it = std::begin(classes); it != std::end(classes); ++it)
        {
            auto jt = it;
            ++jt;
            for(; jt != std::end(classes); ++jt)
            {
                if(has_close_elements(*it, *jt))
                {
                    it->insert(std::end(*it), std::begin(*jt), std::end(*jt));
                    jt = classes.erase(jt);
                    --jt;
                    changed = true;
                }
            }
        }
    }
    return classes;
}

/**
 * @brief Select representative solutions for each cluster.
 * @details For each cluster, the solution candidate with the smallest error
 *      estimate is chosen as a representative. Also returns the number of
 *      elements in the cluster.
 *
 * @param clusters Vector of vectors of Evaluators representing solution
 *     clusters as produced by clusterTris().
 * @return A cluster representative for each input cluster.
 */
template <typename Evaluator>
std::vector<ClusterRepr<Evaluator>>
findRepresentatives(const std::vector<std::vector<Evaluator>>& clusters)
{
    static_assert(is_evaluator_v<Evaluator>,
                  "findRepresentatives requires a valid Evaluator!");
    auto result = std::vector<ClusterRepr<Evaluator>>{};
    for(const auto& c : clusters)
    {
        using namespace boost;
        using namespace boost::adaptors;
        result.push_back(
                {c.size(),
                 *min_element(c, [](const auto& c1, const auto& c2) {
                     return c1.error() < c2.error();
                 })});
    }
    return result;
}

/**
 * @brief Compute context info for representatives
 * @details Computes global point position, eigenvalue order, presence of other
 *      imaginary eigenvalues, and packs into result list together with point
 *      position, eigenvector direction, eigenvalues.
 *
 * @param representatives TriPairs selected by findRepresentatives()
 * @param s_interp First tensor field on the triangle
 * @param t_interp Second tensor field on the triangle
 * @param tri Spatial triangle
 * @return List of TLPoints with context info
 */
PointList
computeContextInfoPEV(
        const std::vector<
            ClusterRepr<ParallelEigenvectorsEvaluator>>& representatives,
        const TensorInterp& s_interp,
        const TensorInterp& t_interp,
        const Triangle& tri)
{
    auto points = PointList{};
    points.reserve(representatives.size());

    for(const auto& r : representatives)
    {
        const auto& pos_tri = r.eval.tris().pos_tri;
        const auto& dir_tri = r.eval.tris().dir_tri;

        auto result_center = pos_tri({1. / 3., 1. / 3., 1. / 3.});
        auto result_dir = dir_tri({1. / 3., 1. / 3., 1. / 3.}).normalized();

        // We want to know which eigenvector of each tensor field we have
        // found (i.e. corresponding to largest, middle, or smallest
        // eigenvalue)
        // Therefore we explicitly compute the eigenvalues at the result
        // position and check which ones the found eigenvector direction
        // corresponds to.
        // @todo: make this step optional

        auto s = s_interp(result_center);
        auto t = t_interp(result_center);

        // Get eigenvalues from our computed direction
        auto s_real_eigv = (s * result_dir).dot(result_dir);
        auto t_real_eigv = (t * result_dir).dot(result_dir);

        // Compute all eigenvalues using Eigen
        auto s_eigvs = s.eigenvalues().eval();
        auto t_eigvs = t.eigenvalues().eval();

        // Find index of eigenvalue that is closest to the one we computed
        using Vec3c = decltype(s_eigvs);
        auto s_closest_index = Vec3d::Index{0};
        (s_eigvs - Vec3c::Ones() * s_real_eigv)
                .cwiseAbs()
                .minCoeff(&s_closest_index);

        auto t_closest_index = Vec3d::Index{0};
        (t_eigvs - Vec3c::Ones() * t_real_eigv)
                .cwiseAbs()
                .minCoeff(&t_closest_index);

        // Find which of the (real) eigenvalues ours is
        auto count_larger_real = [](double ref,
                                    const std::complex<double>& val) {
            if(val.imag() != 0) return 0;
            if(std::abs(ref) >= std::abs(val.real())) return 0;
            return 1;
        };
        auto s_order = s_eigvs.unaryExpr([&](const std::complex<double>& val) {
                                  return count_larger_real(
                                          s_eigvs[s_closest_index].real(), val);
                              })
                               .sum();
        auto t_order = t_eigvs.unaryExpr([&](const std::complex<double>& val) {
                                  return count_larger_real(
                                          t_eigvs[t_closest_index].real(), val);
                              })
                               .sum();

        points.push_back(
                TLPoint{tri(result_center),
                         ERank(s_order),
                         ERank(t_order),
                         result_dir,
                         s_real_eigv,
                         t_real_eigv,
                         s_eigvs.sum().imag() != 0,
                         t_eigvs.sum().imag() != 0,
                         r.cluster_size,
                         (pos_tri[1] - pos_tri[0]).norm(),
                         (dir_tri[1] - dir_tri[0]).norm(),
                         0.});
    }
    return points;
}

PointList
computeContextInfoTCL(
        const std::vector<ClusterRepr<TensorCoreLinesEvaluator>>& representatives,
        const TensorInterp& t_interp,
        const TensorInterp& tx_interp,
        const TensorInterp& ty_interp,
        const TensorInterp& tz_interp,
        const Triangle& tri)
{
    auto points = PointList{};
    points.reserve(representatives.size());

    for(const auto& r : representatives)
    {
        const auto& pos_tri = r.eval.tris().pos_tri;
        const auto& dir_tri = r.eval.tris().dir_tri;

        auto result_center = pos_tri({1. / 3., 1. / 3., 1. / 3.});
        auto result_dir = dir_tri({1. / 3., 1. / 3., 1. / 3.}).normalized();

        // We want to know which eigenvector of each tensor field we have
        // found (i.e. corresponding to largest, middle, or smallest
        // eigenvalue)
        // Therefore we explicitly compute the eigenvalues at the result
        // position and check which ones the found eigenvector direction
        // corresponds to.
        // @todo: make this step optional

        auto t = t_interp(result_center);
        auto tx = tx_interp(result_center);
        auto ty = ty_interp(result_center);
        auto tz = tz_interp(result_center);
        auto dt = (tx * result_dir[0] + ty * result_dir[1] + tz * result_dir[2])
                          .eval();

        // Get eigenvalues from our computed direction
        auto t_real_eigv = (t * result_dir).dot(result_dir);
        auto dt_real_eigv = (dt * result_dir).dot(result_dir);

        // Compute all eigenvalues using Eigen
        auto t_eigvs = t.eigenvalues().eval();
        auto dt_eigvs = dt.eigenvalues().eval();

        // Find index of eigenvalue that is closest to the one we computed
        using Vec3c = decltype(t_eigvs);
        auto t_closest_index = Vec3d::Index{0};
        (t_eigvs - Vec3c::Ones() * t_real_eigv)
                .cwiseAbs()
                .minCoeff(&t_closest_index);

        auto dt_closest_index = Vec3d::Index{0};
        (dt_eigvs - Vec3c::Ones() * dt_real_eigv)
                .cwiseAbs()
                .minCoeff(&dt_closest_index);

        // Find which of the (real) eigenvalues ours is
        auto count_larger_real = [](double ref,
                                    const std::complex<double>& val) {
            if(val.imag() != 0) return 0;
            if(std::abs(ref) >= std::abs(val.real())) return 0;
            return 1;
        };
        auto t_order = t_eigvs.unaryExpr([&](const std::complex<double>& val) {
                                  return count_larger_real(
                                          t_eigvs[t_closest_index].real(), val);
                              })
                               .sum();
        auto dt_order =
                dt_eigvs.unaryExpr([&](const std::complex<double>& val) {
                            return count_larger_real(
                                    dt_eigvs[dt_closest_index].real(), val);
                        })
                        .sum();

        // det( (NablaT*R1)*R  ,  (NablaT*R2)*R ,  R )
        auto r2 = Vec3d::Random().normalized().eval();
        while(result_dir.cross(r2).norm() < 0.1)
        {
            r2 = Vec3d::Random().normalized().eval();
        }
        auto r1 = result_dir.cross(r2).normalized().eval();
        r2 = r1.cross(result_dir).normalized().eval();
        auto scale = t.operatorNorm();
        auto stability = std::log(std::abs(
                (Mat3d{} << ((tx * r1[0] + ty * r1[1] + tz * r1[2])
                             * result_dir)
                                    / scale,
                 ((tx * r2[0] + ty * r2[1] + tz * r2[2]) * result_dir) / scale,
                 result_dir)
                        .finished()
                        .determinant()));

        points.push_back(
                TLPoint{tri(result_center),
                         ERank(t_order),
                         ERank(dt_order),
                         result_dir,
                         t_real_eigv,
                         dt_real_eigv,
                         t_eigvs.sum().imag() != 0,
                         dt_eigvs.sum().imag() != 0,
                         r.cluster_size,
                         (pos_tri[1] - pos_tri[0]).norm(),
                         (dir_tri[1] - dir_tri[0]).norm(),
                         stability});
    }
    return points;
}

PointList computeContextInfoTopo(
        const std::vector<ClusterRepr<TensorTopologyEvaluator>>& representatives,
        const Triangle& tri)
{
    auto points = PointList{};
    points.reserve(representatives.size());

    for(const auto& r : representatives)
    {
        const auto& pos_tri = r.eval.tris().pos_tri;

        auto result_center = pos_tri({1. / 3., 1. / 3., 1. / 3.});

        points.push_back(
                TLPoint{tri(result_center),
                         ERank::First,
                         ERank::First,
                         Vec3d::Zero(),
                         0,
                         0,
                         false,
                         false,
                         r.cluster_size,
                         (pos_tri[1] - pos_tri[0]).norm(),
                         0,
                         0});
    }
    return points;
}

/**
 * @brief Perform the recursive root search using an evaluator.
 * @details Performs a breadth-first recursive search for solutions of the
 *      starting evaluator. Terminates when all solutions have been found or
 *      when more than @a max_candidates are in the queue. In the latter case,
 *      @c boost::none is returned.
 *
 * @param start_ev Starting evaluator
 * @param max_candidates Maximum number of triangles produced during subdivision
 *     before early termination
 * @param num_splits Optional output parameter for storing the number of split
 *     operations performed
 * @param max_level Optional output parameter for storing the maximum
 *     subdivision level reached
 * @return A vector of solution candidates represented by Evaluators at the lowest
 *      subdivision level, or boost::none if the search was terminated early.
 */
template <typename Evaluator>
boost::optional<std::vector<Evaluator>>
rootSearch(const Evaluator& start_ev,
           std::size_t max_candidates,
           uint64_t* num_splits = nullptr,
           uint64_t* max_level = nullptr)
{
    static_assert(is_evaluator_v<Evaluator>,
                  "rootSearch requires a valid Evaluator!");

    auto work_lst = std::queue<Evaluator>{};
    work_lst.push(start_ev);
    auto result = std::vector<Evaluator>{};

    while(!work_lst.empty())
    {
        if(work_lst.size() > max_candidates) return boost::none;
        auto ev = work_lst.front();
        work_lst.pop();
        if(num_splits) *num_splits += 1;
        if(max_level && *max_level < ev.splitLevel())
        {
            *max_level = ev.splitLevel();
        }

        switch(ev.eval())
        {
            case Result::Split:
                for(const auto& p : ev.split())
                {
                    work_lst.push(p);
                }
                break;
            case Result::Accept:
                result.push_back(ev);
                break;
            case Result::Discard:
                break;
        }
    }

    return result;
}

/**
 * Search for parallel eigenvector intersections with a triangle.
 *
 * @param s First tensor field (linear on a triangle)
 * @param t Second tensor field (linear on a triangle)
 * @param tri Physical location of the triangle
 * @param tolerance Error tolerance for subdivision
 * @param max_candidates Maximum number of triangles produced during subdivision
 *     before early termination
 * @param num_splits Optional output parameter for storing the number of split
 *     operations performed
 * @param max_level Optional output parameter for storing the maximum
 *     subdivision level reached
 * @return A vector of all found solution candidates and a vector of the
 *     rough eigenvector directions that resulted in early termination
 */
std::pair<std::vector<ParallelEigenvectorsEvaluator>, std::vector<Vec3d>>
parallelEigenvectorSearch(const TensorInterp& s,
                          const TensorInterp& t,
                          const Triangle& tri,
                          double tolerance,
                          std::size_t max_candidates,
                          uint64_t* num_splits = nullptr,
                          uint64_t* max_level = nullptr)
{
    auto result = std::vector<ParallelEigenvectorsEvaluator>{};
    // Stores directions for which the search was terminated because of
    // too many splits
    auto failed_dirs = std::vector<Vec3d>{};

    auto compute_tri = [&](const Triangle& r) {
        auto start_ev = ParallelEigenvectorsEvaluator(
                {tri, r}, s, t, {tolerance});
        auto solutions =
                rootSearch(start_ev, max_candidates, num_splits, max_level);
        if(solutions)
        {
            boost::insert(
                    result,
                    result.end(),
                    solutions.value());
        }
        else
        {
            failed_dirs.push_back(r({1./3, 1./3, 1./3}));
        }
    };

    // Four triangles covering hemisphere
    auto dir_tris = std::array<Triangle, 4>{
            Triangle{{Vec3d{1, 0, 0}, Vec3d{0, 1, 0}, Vec3d{0, 0, 1}}},
            Triangle{{Vec3d{0, 1, 0}, Vec3d{-1, 0, 0}, Vec3d{0, 0, 1}}},
            Triangle{{Vec3d{-1, 0, 0}, Vec3d{0, -1, 0}, Vec3d{0, 0, 1}}},
            Triangle{{Vec3d{0, -1, 0}, Vec3d{1, 0, 0}, Vec3d{0, 0, 1}}}};

    for(const auto& tri : dir_tris)
    {
        for(const auto& t : tri.split())
        {
            compute_tri(t);
        }
    }

    return {result, failed_dirs};
}

/**
 * Search for tensor core line intersections with a triangle.
 *
 * @param t Tensor field (linear on a triangle)
 * @param dt derivatives of the tensor field (constant on a triangle)
 * @param tri Physical location of the triangle
 * @param tolerance Error tolerance for subdivision
 * @param max_candidates Maximum number of triangles produced during subdivision
 *     before early termination
 * @param num_splits Optional output parameter for storing the number of split
 *     operations performed
 * @param max_level Optional output parameter for storing the maximum
 *     subdivision level reached
 * @return A vector of all found solution candidates and a vector of the
 *     rough eigenvector directions that resulted in early termination
 */
std::pair<std::vector<TensorCoreLinesEvaluator>, std::vector<Vec3d>>
tensorCoreLinesSearch(const TensorInterp& t,
                         const std::array<TensorInterp, 3>& dt,
                         const Triangle& tri,
                         double tolerance,
                         std::size_t max_candidates,
                         uint64_t* num_splits = nullptr,
                         uint64_t* max_level = nullptr)
{
    auto result = std::vector<TensorCoreLinesEvaluator>{};
    // Stores directions for which the search was terminated because of
    // too many splits
    auto failed_dirs = std::vector<Vec3d>{};

    auto compute_tri = [&](const Triangle& r) {
        auto start_ev = TensorCoreLinesEvaluator(
                {tri, r}, t, dt, {tolerance});
        auto solutions =
                rootSearch(start_ev, max_candidates, num_splits, max_level);
        if(solutions)
        {
            boost::insert(result, result.end(), solutions.value());
        }
        else
        {
            failed_dirs.push_back(r({1. / 3, 1. / 3, 1. / 3}));
        }
    };

    auto dir_tris = std::array<Triangle, 4>{
            Triangle{{Vec3d{1, 0, 0}, Vec3d{0, 1, 0}, Vec3d{0, 0, 1}}},
            Triangle{{Vec3d{0, 1, 0}, Vec3d{-1, 0, 0}, Vec3d{0, 0, 1}}},
            Triangle{{Vec3d{-1, 0, 0}, Vec3d{0, -1, 0}, Vec3d{0, 0, 1}}},
            Triangle{{Vec3d{0, -1, 0}, Vec3d{1, 0, 0}, Vec3d{0, 0, 1}}}};

    for(const auto& tri : dir_tris)
    {
        for(const auto& t : tri.split())
        {
            compute_tri(t);
        }
    }

    return {result, failed_dirs};
}

/**
 * Search for degenerate line intersections with a triangle.
 *
 * @param t Tensor field (linear on a triangle)
 * @param tri Physical location of the triangle
 * @param tolerance Error tolerance for subdivision
 * @param max_candidates Maximum number of triangles produced during subdivision
 *     before early termination
 * @param num_splits Optional output parameter for storing the number of split
 *     operations performed
 * @param max_level Optional output parameter for storing the maximum
 *     subdivision level reached
 * @return A vector of all found solution candidates and a vector of the
 *     rough eigenvector directions that resulted in early termination
 */
std::pair<std::vector<TensorTopologyEvaluator>, std::vector<Vec3d>>
tensorTopologySearch(const TensorInterp& t,
                     const Triangle& tri,
                     double tolerance,
                     std::size_t max_candidates,
                     uint64_t* num_splits = nullptr,
                     uint64_t* max_level = nullptr)
{
    auto start_ev = TensorTopologyEvaluator(
            {tri, Triangle{{Vec3d::Zero(), Vec3d::Zero(), Vec3d::Zero()}}},
            t,
            {tolerance});
    auto solutions =
            rootSearch(start_ev, max_candidates, num_splits, max_level);

    if(solutions)
    {
        return {solutions.value(), std::vector<Vec3d>{}};
    }
    else
    {
        return {std::vector<TensorTopologyEvaluator>{},
                std::vector<Vec3d>{1, Vec3d::Zero()}};
    }
}

} // namespace

namespace tl
{

TLResult findParallelEigenvectors(const std::array<Mat3d, 3>& s,
                                   const std::array<Mat3d, 3>& t,
                                   const std::array<Vec3d, 3>& x,
                                   const TLOptions& opts)
{
    auto start_tri =
            Triangle{{Vec3d{1., 0., 0.}, Vec3d{0., 1., 0.}, Vec3d{0., 0., 1.}}};

    auto st = TensorInterp{{s[0], s[1], s[2]}};
    auto tt = TensorInterp{{t[0], t[1], t[2]}};
    auto xt = Triangle{{x[0], x[1], x[2]}};

    auto num_splits = uint64_t{0};
    auto max_level = uint64_t{0};
    auto tris = parallelEigenvectorSearch(st,
                                          tt,
                                          start_tri,
                                          opts.tolerance,
                                          opts.max_candidates,
                                          &num_splits,
                                          &max_level);

    auto clustered_tris = clusterTris(tris.first, opts.cluster_epsilon);

    auto representatives = findRepresentatives(clustered_tris);

    return {computeContextInfoPEV(representatives, st, tt, xt),
            tris.second};
}

TLResult findParallelEigenvectors(const std::array<Mat3d, 3>& s,
                                   const std::array<Mat3d, 3>& t,
                                   const TLOptions& opts)
{
    return findParallelEigenvectors(
            s,
            t,
            {Vec3d{1., 0., 0.}, Vec3d{0., 1., 0.}, Vec3d{0., 0., 1.}},
            opts);
}

TLResult findTensorCoreLines(const std::array<Mat3d, 3>& t,
                                 const std::array<Mat3d, 3>& dt,
                                 const std::array<Vec3d, 3>& x,
                                 const TLOptions& opts)
{
    auto start_tri =
            Triangle{{Vec3d{1., 0., 0.}, Vec3d{0., 1., 0.}, Vec3d{0., 0., 1.}}};

    auto tt = TensorInterp{{t[0], t[1], t[2]}};
    auto tx = TensorInterp{{dt[0], dt[0], dt[0]}};
    auto ty = TensorInterp{{dt[1], dt[1], dt[1]}};
    auto tz = TensorInterp{{dt[2], dt[2], dt[2]}};
    auto xt = Triangle{{x[0], x[1], x[2]}};

    auto tolerance_scale = std::max({tt[0].operatorNorm(),
                                     tt[1].operatorNorm(),
                                     tt[2].operatorNorm()});

    auto num_splits = uint64_t{0};
    auto max_level = uint64_t{0};
    auto tris = tensorCoreLinesSearch(tt,
                                         {tx, ty, tz},
                                         start_tri,
                                         opts.tolerance*tolerance_scale,
                                         opts.max_candidates,
                                         &num_splits,
                                         &max_level);

    auto clustered_tris = clusterTris(tris.first, opts.cluster_epsilon);

    auto representatives = findRepresentatives(clustered_tris);

    return {computeContextInfoTCL(representatives, tt, tx, ty, tz, xt),
            tris.second};
}

TLResult findTensorCoreLines(const std::array<Mat3d, 3>& t,
                                 const std::array<Mat3d, 3>& dt,
                                 const TLOptions& opts)
{
    return findTensorCoreLines(
            t,
            dt,
            {Vec3d{1., 0., 0.}, Vec3d{0., 1., 0.}, Vec3d{0., 0., 1.}},
            opts);
}

TLResult findTensorTopology(const std::array<Mat3d, 3>& t,
                             const std::array<Vec3d, 3>& x,
                             const TLOptions& opts)
{
    auto start_tri =
            Triangle{{Vec3d{1., 0., 0.}, Vec3d{0., 1., 0.}, Vec3d{0., 0., 1.}}};

    auto tt = TensorInterp{{t[0], t[1], t[2]}};
    auto xt = Triangle{{x[0], x[1], x[2]}};

    auto tolerance_scale = std::max({tt[0].operatorNorm(),
                                     tt[1].operatorNorm(),
                                     tt[2].operatorNorm()});

    auto num_splits = uint64_t{0};
    auto max_level = uint64_t{0};
    auto tris = tensorTopologySearch(tt,
                                     start_tri,
                                     opts.tolerance * tolerance_scale,
                                     opts.max_candidates,
                                     &num_splits,
                                     &max_level);

    auto clustered_tris = clusterTris(tris.first, opts.cluster_epsilon);

    auto representatives = findRepresentatives(clustered_tris);

    return {computeContextInfoTopo(representatives, xt),
            tris.second};
}

TLResult findTensorTopology(const std::array<Mat3d, 3>& t,
                             const TLOptions& opts)
{
    return findTensorTopology(
            t,
            {Vec3d{1., 0., 0.}, Vec3d{0., 1., 0.}, Vec3d{0., 0., 1.}},
            opts);
}

} // namespace tl


/* ------------------------------------------------------
 *
 * @file marginal.cpp
 *
 * @brief Marginal effects for regression functions
 *
 *
 *//* ----------------------------------------------------------------------- */
#include <limits>
#include <dbconnector/dbconnector.hpp>
#include <modules/shared/HandleTraits.hpp>
#include <modules/prob/boost.hpp>
#include <boost/math/distributions.hpp>
#include <modules/prob/student.hpp>
#include "marginal.hpp"

namespace madlib {

// Use Eigen
using namespace dbal::eigen_integration;

namespace modules {

// Import names from other MADlib modules
using dbal::NoSolutionFoundException;

namespace regress {

inline double logistic(double x) {
    return 1. / (1. + std::exp(-x));
}

/**
 * @brief Helper function that computes the final statistics for the marginal variance
 */

AnyType margins_stateToResult(
    const Allocator &inAllocator,
    const ColumnVector &diagonal_of_variance_matrix,
    const ColumnVector &inmarginal_effects_per_observation,
    const double numRows) {

    uint16_t n_basis_terms = static_cast<uint16_t>(inmarginal_effects_per_observation.size());
    MutableNativeColumnVector marginal_effects(
        inAllocator.allocateArray<double>(n_basis_terms));
    MutableNativeColumnVector stdErr(
        inAllocator.allocateArray<double>(n_basis_terms));
    MutableNativeColumnVector tStats(
        inAllocator.allocateArray<double>(n_basis_terms));
    MutableNativeColumnVector pValues(
        inAllocator.allocateArray<double>(n_basis_terms));

    for (Index i = 0; i < n_basis_terms; ++i) {
        marginal_effects(i) = inmarginal_effects_per_observation(i) / numRows;
        stdErr(i) = std::sqrt(diagonal_of_variance_matrix(i));
        tStats(i) = marginal_effects(i) / stdErr(i);

        // p-values only make sense if numRows > coef.size()
        if (numRows > n_basis_terms)
            pValues(i) = 2. * prob::cdf( prob::normal(),
                                         -std::abs(tStats(i)));
    }

    // Return all coefficients, standard errors, etc. in a tuple
    // Note: p-values will return NULL if numRows <= coef.size
    AnyType tuple;
    tuple << marginal_effects
          << stdErr
          << tStats
          << (numRows > n_basis_terms? pValues: Null());
    return tuple;
}
// -------------------------------------------------------------------------

// ---------------------------------------------------------------------------
//             Marginal Effects Linear Regression States
// ---------------------------------------------------------------------------

/**
 * @brief State for marginal effects calculation for logistic regression
 *
 * TransitionState encapsualtes the transition state during the
 * marginal effects calculation for the logistic-regression aggregate function.
 * To the database, the state is exposed as a single DOUBLE PRECISION array,
 * to the C++ code it is a proper object containing scalars and vectors.
 *
 * Note: We assume that the DOUBLE PRECISION array is initialized by the
 * database with length at least 5, and all elemenets are 0.
 *
 */
template <class Handle>
class MarginsLinregrInteractionState {
    template <class OtherHandle>
    friend class MarginsLinregrInteractionState;

  public:
    MarginsLinregrInteractionState(const AnyType &inArray)
        : mStorage(inArray.getAs<Handle>()) {
        rebind(static_cast<uint16_t>(mStorage[1]),
               static_cast<uint16_t>(mStorage[2]));
    }

    /**
     * @brief Convert to backend representation
     *
     * We define this function so that we can use State in the
     * argument list and as a return type.
     */
    inline operator AnyType() const {
        return mStorage;
    }

    /**
     * @brief Initialize the marginal variance calculation state.
     *
     * This function is only called for the first iteration, for the first row.
     */
    inline void initialize(const Allocator &inAllocator,
                           const uint16_t inWidthOfX,
                           const uint16_t inNumBasis) {
        mStorage = inAllocator.allocateArray<double, dbal::AggregateContext,
                                             dbal::DoZero, dbal::ThrowBadAlloc>(
                arraySize(inWidthOfX, inNumBasis));
        rebind(inWidthOfX, inNumBasis);
        widthOfX = inWidthOfX;
        numBasis = inNumBasis;
    }

    /**
     * @brief We need to support assigning the previous state
     */
    template <class OtherHandle>
    MarginsLinregrInteractionState &operator=(
        const MarginsLinregrInteractionState<OtherHandle> &inOtherState) {

        for (size_t i = 0; i < mStorage.size(); i++)
            mStorage[i] = inOtherState.mStorage[i];
        return *this;
    }

    /**
     * @brief Merge with another State object by copying the intra-iteration
     *     fields
     */
    template <class OtherHandle>
    MarginsLinregrInteractionState &operator+=(
        const MarginsLinregrInteractionState<OtherHandle> &inOtherState) {

        if (mStorage.size() != inOtherState.mStorage.size() ||
            widthOfX != inOtherState.widthOfX)
            throw std::logic_error("Internal error: Incompatible transition "
                                   "states");
        numRows += inOtherState.numRows;
        marginal_effects += inOtherState.marginal_effects;
        delta += inOtherState.delta;
        return *this;
    }

    /**
     * @brief Reset the inter-iteration fields.
     */
    inline void reset() {
        numRows = 0;
        training_data_vcov.fill(0);
        marginal_effects.fill(0);
        delta.fill(0);
    }

  private:
    static inline size_t arraySize(const uint16_t inWidthOfX,
                                   const uint16_t inNumBasis) {
        return 4 + inNumBasis + (inWidthOfX + inNumBasis) * inWidthOfX;
    }

    /**
     * @brief Rebind to a new storage array
     *
     * @param inWidthOfX The number of independent variables.
     *
     */
    void rebind(uint16_t inWidthOfX, uint16_t inNumBasis) {
        iteration.rebind(&mStorage[0]);
        widthOfX.rebind(&mStorage[1]);
        numBasis.rebind(&mStorage[2]);
        numRows.rebind(&mStorage[3]);
        marginal_effects.rebind(&mStorage[4], inNumBasis);
        training_data_vcov.rebind(&mStorage[4 + inNumBasis], inWidthOfX, inWidthOfX);
        delta.rebind(&mStorage[4 + inNumBasis + inWidthOfX * inWidthOfX],
                     inNumBasis, inWidthOfX);
    }
    Handle mStorage;

  public:
    typename HandleTraits<Handle>::ReferenceToUInt32 iteration;
    typename HandleTraits<Handle>::ReferenceToUInt16 widthOfX;
    typename HandleTraits<Handle>::ReferenceToUInt16 numBasis;
    typename HandleTraits<Handle>::ReferenceToUInt64 numRows;
    typename HandleTraits<Handle>::ColumnVectorTransparentHandleMap marginal_effects;
    typename HandleTraits<Handle>::MatrixTransparentHandleMap training_data_vcov;
    typename HandleTraits<Handle>::MatrixTransparentHandleMap delta;
};
// ----------------------------------------------------------------------

/**
 * @brief Perform the marginal effects transition step
 */
AnyType
margins_linregr_int_transition::run(AnyType &args) {
    // Early return because of an exception has been "thrown"
    // (actually "warning") in the previous invocations
    if (args[0].isNull())
        return Null();
    MarginsLinregrInteractionState<MutableArrayHandle<double> > state = args[0];
    if (args[1].isNull() || args[2].isNull() ||
            args[3].isNull() || args[4].isNull()) {
        return args[0];
    }
    MappedColumnVector x;
    try {
        // an exception is raised in the backend if args[2] contains nulls
        MappedColumnVector xx = args[1].getAs<MappedColumnVector>();
        // x is a const reference, we can only rebind to change its pointer
        x.rebind(xx.memoryHandle(), xx.size());
    } catch (const ArrayWithNullException &e) {
        return args[0];
    }

    // The following check was added with MADLIB-138.
    if (!dbal::eigen_integration::isfinite(x)) {
        //throw std::domain_error("Design matrix is not finite.");
        warning("Design matrix is not finite.");
        return Null();
     }

    MappedColumnVector beta = args[2].getAs<MappedColumnVector>();

    Matrix J_trans = args[4].getAs<MappedMatrix>();
    J_trans.transposeInPlace(); // we actually pass-in J but transpose since
                                // we only require J^T in our equations

    if (state.numRows == 0) {
        if (x.size() > std::numeric_limits<uint16_t>::max()) {
            //throw std::domain_error("Number of independent variables cannot be "
            //                        "larger than 65535.");
            warning("Number of independent variables cannot be larger than 65535.");
            return Null();
        }
        state.initialize(*this,
                         static_cast<uint16_t>(beta.size()),
                         static_cast<uint16_t>(J_trans.rows()));
        Matrix training_data_vcov = args[3].getAs<MappedMatrix>();
        state.training_data_vcov = training_data_vcov;
    }

    // Now do the transition step
    state.numRows++;

    // compute marginal effects and delta using 1st and 2nd derivatives
    state.marginal_effects += J_trans * beta;
    state.delta += J_trans;
    return state;
}

/**
 * @brief Marginal effects: Merge transition states
 */
AnyType
margins_linregr_int_merge::run(AnyType &args) {
    // In case the aggregator should be terminated because
    // an exception has been "thrown" in the transition function
    if (args[0].isNull() || args[1].isNull())
        return Null();
    MarginsLinregrInteractionState<MutableArrayHandle<double> > stateLeft = args[0];
    MarginsLinregrInteractionState<ArrayHandle<double> > stateRight = args[1];
    // We first handle the trivial case where this function is called with one
    // of the states being the initial state
    if (stateLeft.numRows == 0)
        return stateRight;
    else if (stateRight.numRows == 0)
        return stateLeft;

    // Merge states together and return
    stateLeft += stateRight;
    return stateLeft;
}

/**
 * @brief Marginal effects: Final step
 */
AnyType
margins_linregr_int_final::run(AnyType &args) {
    // In case the aggregator should be terminated because
    // an exception has been "thrown" in the transition function
    if (args[0].isNull())
        return Null();
    // We request a mutable object.
    // Depending on the backend, this might perform a deep copy.
    MarginsLinregrInteractionState<ArrayHandle<double> > state = args[0];
    // Aggregates that haven't seen any data just return Null.
    if (state.numRows == 0)
        return Null();

    // Variance of the marginal effects (computed by delta method)
    Matrix variance;
    variance = state.delta * state.training_data_vcov;
    // we only need the diagonal elements of the variance, so we perform a dot
    // product of each row with itself to compute each diagonal element.
    ColumnVector variance_diagonal =
        variance.cwiseProduct(state.delta).rowwise().sum() / static_cast<double>(state.numRows * state.numRows);

    // Computing the marginal effects
    return margins_stateToResult(*this, variance_diagonal,
                                 state.marginal_effects, static_cast<double>(state.numRows));
}

// ---------------------------------------------------------------------------
//             Marginal Effects Logistic Regression States
// ---------------------------------------------------------------------------
/**
 * @brief State for marginal effects calculation for logistic regression
 *
 * TransitionState encapsualtes the transition state during the
 * marginal effects calculation for the logistic-regression aggregate function.
 * To the database, the state is exposed as a single DOUBLE PRECISION array,
 * to the C++ code it is a proper object containing scalars and vectors.
 *
 * Note: We assume that the DOUBLE PRECISION array is initialized by the
 * database with length at least 5, and all elemenets are 0.
 *
 */
template <class Handle>
class MarginsLogregrInteractionState {
    template <class OtherHandle>
    friend class MarginsLogregrInteractionState;

  public:
    MarginsLogregrInteractionState(const AnyType &inArray)
        : mStorage(inArray.getAs<Handle>()) {

        rebind(static_cast<uint16_t>(mStorage[1]),
               static_cast<uint16_t>(mStorage[2]),
               static_cast<uint16_t>(mStorage[3]));
    }

    /**
     * @brief Convert to backend representation
     *
     * We define this function so that we can use State in the
     * argument list and as a return type.
     */
    inline operator AnyType() const {
        return mStorage;
    }

    /**
     * @brief Initialize the marginal variance calculation state.
     *
     * This function is only called for the first iteration, for the first row.
     */
    inline void initialize(const Allocator &inAllocator,
                           const uint16_t inWidthOfX,
                           const uint16_t inNumBasis,
                           const uint16_t inNumCategoricals) {
        mStorage = inAllocator.allocateArray<double, dbal::AggregateContext,
                                             dbal::DoZero, dbal::ThrowBadAlloc>(
                arraySize(inWidthOfX, inNumBasis, inNumCategoricals));
        rebind(inWidthOfX, inNumBasis, inNumCategoricals);
        widthOfX = inWidthOfX;
        numBasis = inNumBasis;
        numCategoricalVarsInSubset = inNumCategoricals;
    }

    /**
     * @brief We need to support assigning the previous state
     */
    template <class OtherHandle>
    MarginsLogregrInteractionState &operator=(
        const MarginsLogregrInteractionState<OtherHandle> &inOtherState) {

        for (size_t i = 0; i < mStorage.size(); i++)
            mStorage[i] = inOtherState.mStorage[i];
        return *this;
    }

    /**
     * @brief Merge with another State object by copying the intra-iteration
     *     fields
     */
    template <class OtherHandle>
    MarginsLogregrInteractionState &operator+=(
        const MarginsLogregrInteractionState<OtherHandle> &inOtherState) {

        if (mStorage.size() != inOtherState.mStorage.size() ||
            widthOfX != inOtherState.widthOfX)
            throw std::logic_error("Internal error: Incompatible transition "
                                   "states");

        numRows += inOtherState.numRows;
        marginal_effects += inOtherState.marginal_effects;
        delta += inOtherState.delta;
        return *this;
    }

    /**
     * @brief Reset the inter-iteration fields.
     */
    inline void reset() {
        numRows = 0;
        marginal_effects.fill(0);
        categorical_basis_indices.fill(0);
        training_data_vcov.fill(0);
        delta.fill(0);
    }

  private:
    static inline size_t arraySize(const uint16_t inWidthOfX,
                                   const uint16_t inNumBasis,
                                   const uint16_t inNumCategoricals) {
        return 5 + inNumBasis + inNumCategoricals + (inWidthOfX + inNumBasis) * inWidthOfX;
    }

    /**
     * @brief Rebind to a new storage array
     *
     * @param inWidthOfX The number of independent variables.
     *
     */
    void rebind(uint16_t inWidthOfX, uint16_t inNumBasis, uint16_t inNumCategoricals) {
        iteration.rebind(&mStorage[0]);
        widthOfX.rebind(&mStorage[1]);
        numBasis.rebind(&mStorage[2]);
        numCategoricalVarsInSubset.rebind(&mStorage[3]);
        numRows.rebind(&mStorage[4]);
        marginal_effects.rebind(&mStorage[5], inNumBasis);
        training_data_vcov.rebind(&mStorage[5 + inNumBasis], inWidthOfX, inWidthOfX);
        delta.rebind(&mStorage[5 + inNumBasis + inWidthOfX * inWidthOfX],
                     inNumBasis, inWidthOfX);
        if (inNumCategoricals > 0)
            categorical_basis_indices.rebind(&mStorage[5 + inNumBasis +
                                            (inWidthOfX + inNumBasis) * inWidthOfX],
                                            inNumCategoricals);
    }
    Handle mStorage;

  public:

    typename HandleTraits<Handle>::ReferenceToUInt32 iteration;
    typename HandleTraits<Handle>::ReferenceToUInt16 widthOfX;
    typename HandleTraits<Handle>::ReferenceToUInt16 numBasis;
    typename HandleTraits<Handle>::ReferenceToUInt16 numCategoricalVarsInSubset;
    typename HandleTraits<Handle>::ReferenceToUInt64 numRows;
    typename HandleTraits<Handle>::ColumnVectorTransparentHandleMap marginal_effects;
    typename HandleTraits<Handle>::ColumnVectorTransparentHandleMap categorical_basis_indices;
    typename HandleTraits<Handle>::MatrixTransparentHandleMap training_data_vcov;
    typename HandleTraits<Handle>::MatrixTransparentHandleMap delta;
};
// ----------------------------------------------------------------------

/**
 * @brief Perform the marginal effects transition step
 */
AnyType
margins_logregr_int_transition::run(AnyType &args) {
    // Early return because of an exception has been "thrown"
    // (actually "warning") in the previous invocations
    if (args[0].isNull())
        return Null();
    MarginsLogregrInteractionState<MutableArrayHandle<double> > state = args[0];
    if (args[1].isNull() || args[2].isNull() ||
            args[3].isNull() || args[4].isNull()) {
        return args[0];
    }    MappedColumnVector f;
    try {
        // an exception is raised in the backend if args[2] contains nulls
        MappedColumnVector xx = args[1].getAs<MappedColumnVector>();
        // x is a const reference, we can only rebind to change its pointer
        f.rebind(xx.memoryHandle(), xx.size());
    } catch (const ArrayWithNullException &e) {
        return args[0];
    }

    // The following check was added with MADLIB-138.
    if (!dbal::eigen_integration::isfinite(f)) {
        //throw std::domain_error("Design matrix is not finite.");
        warning("Design matrix is not finite.");
        return Null();
    }

    // beta is the coefficient vector from logistic regression
    MappedColumnVector beta = args[2].getAs<MappedColumnVector>();

    // basis_indices represents the indices (from beta) of which we want to
    // compute the marginal effect. We don't need to compute the ME for all
    // variables, thus basis_indices could be a subset of all indices.
    MappedColumnVector basis_indices = args[4].getAs<MappedColumnVector>();

    // below symbols match the ones used in the design doc
    const uint16_t N = static_cast<uint16_t>(beta.size());
    const uint16_t M = static_cast<uint16_t>(basis_indices.size());
    assert(N >= M);

    Matrix J;  // J: N * M
    if (args[5].isNull()){
        J = Matrix::Zero(N, M);
        for (Index i = 0; i < M; ++i)
            J(static_cast<Index>(basis_indices(i)), i) = 1;
    } else{
        J = args[5].getAs<MappedMatrix>();
    }
    assert(J.rows() == N && J.cols() == M);

    MappedColumnVector categorical_indices;
    uint16_t numCategoricalVars = 0;

    if (!args[6].isNull()) {
        // categorical_indices represents which indices (from beta) are
        // categorical variables
        try {
            MappedColumnVector xx = args[6].getAs<MappedColumnVector>();
            categorical_indices.rebind(xx.memoryHandle(), xx.size());
        } catch (const ArrayWithNullException &e) {
             //throw std::runtime_error("The categorical indices contain NULL values");
             warning("The categorical indices contain NULL values");
             return Null();
        }
        numCategoricalVars = static_cast<uint16_t>(categorical_indices.size());
    }

    if (state.numRows == 0) {
        if (f.size() > std::numeric_limits<uint16_t>::max()) {
            //throw std::domain_error("Number of independent variables cannot be "
            //                        "larger than 65535.");
            warning("Number of independent variables cannot be larger than 65535.");
            return Null();
        }
        std::vector<uint16_t> tmp_cat_basis_indices;
        if (numCategoricalVars > 0){
            // find which of the variables in basis_indices are categorical
            // and store only the indices for these variables in our state
            for (Index i = 0; i < basis_indices.size(); ++i){
                for (Index j = 0; j < categorical_indices.size(); ++j){
                    if (basis_indices(i) == categorical_indices(j)){
                        tmp_cat_basis_indices.push_back(static_cast<uint16_t>(i));
                        continue;
                    }
                }
            }
            state.numCategoricalVarsInSubset = static_cast<uint16_t>(tmp_cat_basis_indices.size());
        }
        state.initialize(*this,
                         static_cast<uint16_t>(N),
                         static_cast<uint16_t>(M),
                         static_cast<uint16_t>(state.numCategoricalVarsInSubset));

        Matrix training_data_vcov = args[3].getAs<MappedMatrix>();
        state.training_data_vcov = training_data_vcov;

        if (state.numCategoricalVarsInSubset > 0){
            for (int i=0; i < state.numCategoricalVarsInSubset; ++i){
                state.categorical_basis_indices(i) = tmp_cat_basis_indices[i];
            }
        }
    }

    // Now do the transition step
    state.numRows++;
    double f_beta = dot(f, beta);
    double p = std::exp(f_beta)/ (1 + std::exp(f_beta));

    // compute marginal effects and delta using 1st and 2nd derivatives
    ColumnVector J_trans_beta;
    J_trans_beta = trans(J) * beta;
    ColumnVector curr_margins = J_trans_beta * p * (1 - p);

    Matrix curr_delta;
    curr_delta = p * (1 - p) * (trans(J) +
                                (1 - 2 * p) * J_trans_beta * trans(f));

    // margins and delta using discrete differences for categoricals variables

    // f_set_mat and f_unset_mat are matrices where each row corresponds to a
    //  categorical basis variable and the columns correspond to all terms
    // (basis and interaction)
    Matrix f_set_mat;  // numCategoricalVarsInSubset x N
    Matrix f_unset_mat;  // numCategoricalVarsInSubset x N
    if (!args[7].isNull() && !args[8].isNull()){
        // the matrix is read in column-order but passed in row-order
        f_set_mat = args[7].getAs<MappedMatrix>();
        f_set_mat.transposeInPlace();

        f_unset_mat = args[8].getAs<MappedMatrix>();
        f_unset_mat.transposeInPlace();
    }

    // PERFORMANCE TWEAK: for the no interaction case, f_set_mat and f_unset_mat
    // only need column entries for the categorical variables (others are same
    // as f). Since, passing a smaller matrix into the transition function is
    // faster, for the no interaction case, we don't need to input the whole
    // matrix into this function. For the interaction case, since it is unknown
    // which indices have interaction, we require entries for all columns in the
    // matrices.
    bool no_interactions = (f_set_mat.cols() < N);
    for (Index i = 0; i < state.numCategoricalVarsInSubset; ++i) {
        // Note: categorical_indices are assumed to be zero-based
        ColumnVector f_set;
        ColumnVector f_unset;
        ColumnVector shortened_f_set = f_set_mat.row(i);
        ColumnVector shortened_f_unset = f_unset_mat.row(i);

        if (no_interactions){
            f_set = f;
            f_unset = f;
            for (Index j=0; j < shortened_f_set.size(); ++j){
                f_set(static_cast<Index>(categorical_indices(j))) = shortened_f_set(j);
                f_unset(static_cast<Index>(categorical_indices(j))) = shortened_f_unset(j);
            }
        } else {
            f_set = shortened_f_set;
            f_unset = shortened_f_unset;
        }
        double p_set = logistic(dot(f_set, beta));
        double p_unset = logistic(dot(f_unset, beta));

        curr_margins(static_cast<uint16_t>(state.categorical_basis_indices(i))) = p_set - p_unset;
        curr_delta.row(static_cast<uint16_t>(state.categorical_basis_indices(i))) = (
            p_set * (1 - p_set) * f_set - p_unset * (1 - p_unset) * f_unset);
    }
    state.marginal_effects += curr_margins;
    state.delta += curr_delta;
    return state;
}

/**
 * @brief Marginal effects: Merge transition states
 */
AnyType
margins_logregr_int_merge::run(AnyType &args) {
    // In case the aggregator should be terminated because
    // an exception has been "thrown" in the transition function
    if (args[0].isNull() || args[1].isNull())
        return Null();
    MarginsLogregrInteractionState<MutableArrayHandle<double> > stateLeft = args[0];
    MarginsLogregrInteractionState<ArrayHandle<double> > stateRight = args[1];
    // We first handle the trivial case where this function is called with one
    // of the states being the initial state
    if (stateLeft.numRows == 0)
        return stateRight;
    else if (stateRight.numRows == 0)
        return stateLeft;

    // Merge states together and return
    stateLeft += stateRight;
    return stateLeft;
}

/**
 * @brief Marginal effects: Final step
 */
AnyType
margins_logregr_int_final::run(AnyType &args) {
    // In case the aggregator should be terminated because
    // an exception has been "thrown" in the transition function
    if (args[0].isNull())
        return Null();
    // We request a mutable object.
    // Depending on the backend, this might perform a deep copy.
    MarginsLogregrInteractionState<MutableArrayHandle<double> > state = args[0];
    // Aggregates that haven't seen any data just return Null.
    if (state.numRows == 0)
        return Null();

    // Variance for marginal effects according to the delta method
    Matrix variance;
    variance = state.delta * state.training_data_vcov;

    // we only need the diagonal elements of the variance, so we perform a dot
    // product of each row with state.delta to compute each diagonal element.
    // We divide by numRows^2 since we need the average variance
    ColumnVector variance_diagonal =
        variance.cwiseProduct(state.delta).rowwise().sum() / static_cast<double>(state.numRows * state.numRows);

    // Computing the final results
    return margins_stateToResult(*this, variance_diagonal,
                                 state.marginal_effects, static_cast<double>(state.numRows));
}
// ------------------------ End of Logistic Marginal ---------------------------

// ---------------------------------------------------------------------------
//             Marginal Effects Multilogistic Regression
// ---------------------------------------------------------------------------
/**
 * @brief State for marginal effects calculation for multilogistic regression
 *
 * TransitionState encapsualtes the transition state during the
 * marginal effects calculation for the logistic-regression aggregate function.
 * To the database, the state is exposed as a single DOUBLE PRECISION array,
 * to the C++ code it is a proper object containing scalars and vectors.
 *
 * Note: We assume that the DOUBLE PRECISION array is initialized by the
 * database with length at least 5, and all elemenets are 0.
 *
 */
template <class Handle>
class MarginsMLogregrInteractionState {
    template <class OtherHandle>
    friend class MarginsMLogregrInteractionState;

  public:
    MarginsMLogregrInteractionState(const AnyType &inArray)
        : mStorage(inArray.getAs<Handle>()) {

        rebind(static_cast<uint16_t>(mStorage[0]),
               static_cast<uint16_t>(mStorage[1]),
               static_cast<uint16_t>(mStorage[2]),
               static_cast<uint16_t>(mStorage[3]));
    }

    /**
     * @brief Convert to backend representation
     *
     * We define this function so that we can use State in the
     * argument list and as a return type.
     */
    inline operator AnyType() const {
        return mStorage;
    }

    /**
     * @brief Initialize the marginal variance calculation state.
     *
     * This function is only called for the first iteration, for the first row.
     */
    inline void initialize(const Allocator &inAllocator,
                           const uint16_t inWidthOfX,
                           const uint16_t inNumCategories,
                           const uint16_t inNumBasis,
                           const uint16_t inNumCategoricalVars) {
        mStorage = inAllocator.allocateArray<double, dbal::AggregateContext,
                                             dbal::DoZero, dbal::ThrowBadAlloc>(
                arraySize(inWidthOfX, inNumCategories, inNumBasis, inNumCategoricalVars));
        rebind(inWidthOfX,  inNumCategories, inNumBasis, inNumCategoricalVars);
        widthOfX = inWidthOfX;
        numCategories = inNumCategories;
        numBasis = inNumBasis;
        numCategoricalVarsInSubset = inNumCategoricalVars;
    }

    /**
     * @brief We need to support assigning the previous state
     */
    template <class OtherHandle>
    MarginsMLogregrInteractionState &operator=(
        const MarginsMLogregrInteractionState<OtherHandle> &inOtherState) {

        for (size_t i = 0; i < mStorage.size(); i++)
            mStorage[i] = inOtherState.mStorage[i];
        return *this;
    }

    /**
     * @brief Merge with another State object by copying the intra-iteration
     *     fields
     */
    template <class OtherHandle>
    MarginsMLogregrInteractionState &operator+=(
        const MarginsMLogregrInteractionState<OtherHandle> &inOtherState) {

        if (mStorage.size() != inOtherState.mStorage.size() ||
            widthOfX != inOtherState.widthOfX)
            throw std::logic_error("Internal error: Incompatible transition "
                                   "states");

        numRows += inOtherState.numRows;
        marginal_effects += inOtherState.marginal_effects;
        delta += inOtherState.delta;
        return *this;
    }

    /**
     * @brief Reset the inter-iteration fields.
     */
    inline void reset() {
        numRows = 0;
        marginal_effects.fill(0);
        categorical_basis_indices.fill(0);
        training_data_vcov.fill(0);
        delta.fill(0);
    }

  private:
    static inline size_t arraySize(const uint16_t inWidthOfX,
                                   const uint16_t inNumCategories,
                                   const uint16_t inNumBasis,
                                   const uint16_t inNumCategoricalVars) {
        return 5 + (inNumCategories - 1) * (inNumBasis +
                    inNumBasis*inWidthOfX*(inNumCategories - 1) +
                    inWidthOfX*inWidthOfX*(inNumCategories - 1)) +
                inNumCategoricalVars;
    }

    /**
     * @brief Rebind to a new storage array
     *
     * @param inWidthOfX The number of independent variables.
     *
     */
    void rebind(const uint16_t inWidthOfX,
                const uint16_t inNumCategories,
                const uint16_t inNumBasis,
                const uint16_t inNumCategoricalVars) {

        const uint16_t & L = inNumCategories;
        const uint16_t & N = inWidthOfX;
        const uint16_t & M = inNumBasis;

        widthOfX.rebind(&mStorage[0]);
        numCategories.rebind(&mStorage[1]);
        numBasis.rebind(&mStorage[2]);
        numCategoricalVarsInSubset.rebind(&mStorage[3]);
        numRows.rebind(&mStorage[4]);

        if (L == 0) { return; }

        marginal_effects.rebind(&mStorage[5], M, L-1);

        int current_length = 5 + M * (L - 1);

        training_data_vcov.rebind(&mStorage[current_length], N*(L-1), N*(L-1));
        current_length += N * (L - 1) * N * (L - 1);

        delta.rebind(&mStorage[current_length], M*(L-1), N*(L-1));
        current_length += N * (L - 1) * M * (L - 1);

        if (inNumCategoricalVars > 0)
            categorical_basis_indices.rebind(&mStorage[static_cast<uint16_t>(current_length)], inNumCategoricalVars);
    }
    Handle mStorage;

  public:
    // symbols in comments correspond to the design document
    typename HandleTraits<Handle>::ReferenceToUInt16 widthOfX;       // N
    typename HandleTraits<Handle>::ReferenceToUInt16 numCategories;  // L
    typename HandleTraits<Handle>::ReferenceToUInt16 numBasis;       // M
    typename HandleTraits<Handle>::ReferenceToUInt16 numCategoricalVarsInSubset;
    typename HandleTraits<Handle>::ReferenceToUInt64 numRows;
    typename HandleTraits<Handle>::MatrixTransparentHandleMap marginal_effects; // ME
    typename HandleTraits<Handle>::ColumnVectorTransparentHandleMap categorical_basis_indices;
    typename HandleTraits<Handle>::MatrixTransparentHandleMap training_data_vcov;
    typename HandleTraits<Handle>::MatrixTransparentHandleMap delta;   // S
};
// ----------------------------------------------------------------------

namespace {
inline Index
reindex(Index outer, Index inner, Index block) { return outer * block + inner; }
}

/**
 * @brief Perform the marginal effects transition step
 */
AnyType
margins_mlogregr_int_transition::run(AnyType &args) {
    // Early return because of an exception has been "thrown"
    // (actually "warning") in the previous invocations
    if (args[0].isNull())
        return Null();
    MarginsMLogregrInteractionState<MutableArrayHandle<double> > state = args[0];
    if (args[1].isNull() || args[2].isNull() || args[3].isNull() ||
        args[4].isNull()) {
        return args[0];
    }

    MappedColumnVector f;
    try {
        // an exception is raised in the backend if args[1] contains nulls
        MappedColumnVector xx = args[1].getAs<MappedColumnVector>();
        // x is a const reference, we can only rebind to change its pointer
        f.rebind(xx.memoryHandle(), xx.size());
    } catch (const ArrayWithNullException &e) {
        return args[0];
    }

    // The following check was added with MADLIB-138.
    if (!dbal::eigen_integration::isfinite(f)) {
        //throw std::domain_error("Design matrix is not finite.");
        warning("Design matrix is not finite.");
        return Null();
    }

    // coefficients are arranged in a matrix
    MappedMatrix beta = args[2].getAs<MappedMatrix>();  // beta: N x (L - 1)

    // basis_indices represents the indices (from beta) of which we want to
    // compute the marginal effect. We don't need to compute the ME for all
    // variables, thus basis_indices could be a subset of all indices.
    MappedColumnVector basis_indices = args[4].getAs<MappedColumnVector>();

    // all variable symbols correspond to the design document
    const uint16_t N = static_cast<uint16_t>(beta.rows());
    const uint16_t M = static_cast<uint16_t>(basis_indices.size());
    assert(N >= M);

    Matrix J;  // J: N x M
    if (args[5].isNull()){
        J = Matrix::Zero(N, M);
        for (Index i = 0; i < M; ++i)
            J(static_cast<Index>(basis_indices(i)), i) = 1;
    } else{
        J = args[5].getAs<MappedMatrix>();
    }
    assert(J.rows() == N && J.cols() == M);

    MappedColumnVector categorical_indices;
    uint16_t numCategoricalVars = 0;

    if (!args[6].isNull()) {
        // categorical_indices represents which indices (from beta) are
        // categorical variables
        try {
            MappedColumnVector xx = args[6].getAs<MappedColumnVector>();
            categorical_indices.rebind(xx.memoryHandle(), xx.size());
        } catch (const ArrayWithNullException &e) {
             //throw std::runtime_error("The categorical indices contain NULL values");
             warning("The categorical indices contain NULL values");
             return Null();
        }
        numCategoricalVars = static_cast<uint16_t>(categorical_indices.size());
    }

    if (state.numRows == 0) {
        if (f.size() > std::numeric_limits<uint16_t>::max()) {
            //throw std::domain_error("Number of independent variables cannot be "
            //                        "larger than 65535.");
            warning("Number of independent variables cannot be larger than 65535.");
            return Null();
        }
        std::vector<uint16_t> tmp_cat_basis_indices;
        if (numCategoricalVars > 0){
            // find which of the variables in basis_indices are categorical
            // and store only the indices for these variables in our state
            for (Index i = 0; i < basis_indices.size(); ++i){
                for (Index j = 0; j < categorical_indices.size(); ++j){
                    if (basis_indices(i) == categorical_indices(j)){
                        tmp_cat_basis_indices.push_back(static_cast<uint16_t>(i));
                        continue;
                    }
                }
            }
            state.numCategoricalVarsInSubset = static_cast<uint16_t>(tmp_cat_basis_indices.size());
        }
        state.initialize(*this,
                         static_cast<uint16_t>(J.rows()),
                         static_cast<uint16_t>(beta.cols() + 1),
                         static_cast<uint16_t>(J.cols()),
                         static_cast<uint16_t>(state.numCategoricalVarsInSubset));

        Matrix training_data_vcov = args[3].getAs<MappedMatrix>();
        state.training_data_vcov = training_data_vcov;
        if (state.numCategoricalVarsInSubset > 0){
            for (int i=0; i < state.numCategoricalVarsInSubset; ++i){
                state.categorical_basis_indices(i) = tmp_cat_basis_indices[i];
            }
        }
    }

    state.numRows++;

    // all variable symbols correspond to the design document
    const uint16_t & L = state.numCategories;
    ColumnVector prob(trans(beta) * f);
    Matrix J_trans_beta(trans(J) * beta);

    // Calculate the odds ratio
    prob = prob.array().exp();
    double prob_sum = prob.sum();
    prob = prob / (1 + prob_sum);

    ColumnVector JBP(J_trans_beta * prob);
    Matrix curr_margins = J_trans_beta * prob.asDiagonal() - JBP * trans(prob);

    // compute delta using 2nd derivatives
    // delta matrix is 2-D of size (L-1)M x (L-1)N:
    //      row_index = [0, (L-1)M), col_index = [0, (L-1)N)
    // row_index(m, l) = m * (L-1) + l
    // col_index(n, l1) = n * (L-1) + l1
    Index row_index, col_index;
    int delta_l_l1;
    for (int m = 0; m < M; m++){
        // Skip the categorical variables
        if (state.numCategoricalVarsInSubset > 0) {
            bool is_categorical = false;
            for(int i = 0; i < state.categorical_basis_indices.size(); i++)
                if (m == state.categorical_basis_indices(i)) {
                    is_categorical = true;
                    break;
                }
             if (is_categorical)
                continue;
        }

        for (int l=0; l < (L-1); l++){
            row_index = reindex(m, l, L-1);
            for (int n=0; n < N; n++){
                for (int l1=0; l1 < (L-1); l1++){
                    delta_l_l1 = (l==l1) ? 1 : 0;
                    col_index = reindex(n, l1, L-1);
                    state.delta(row_index, col_index) +=
                        f(n)*(delta_l_l1 - prob(l1)) * curr_margins(m, l) +
                        prob(l) * (delta_l_l1 * J(n, m) -
                                   f(n) * curr_margins(m, l1) -
                                   prob(l1) * J(n, m));
                }
            }
        }
    }

    // update marginal effects and delta using discrete differences just for
    // categorical variables
    Matrix f_set_mat;   // numCategoricalVarsInSubset * N
    Matrix f_unset_mat; // numCategoricalVarsInSubset * N
    // the above matrices contain the f_set and f_unset for all categorical variables
    if (!args[7].isNull() && !args[8].isNull()){
        // the matrix is read in column-order but passed in row-order
        f_set_mat = args[7].getAs<MappedMatrix>();
        f_set_mat.transposeInPlace();

        f_unset_mat = args[8].getAs<MappedMatrix>();
        f_unset_mat.transposeInPlace();
    }

    // PERFORMANCE TWEAK: for the no interaction case, f_set_mat and f_unset_mat
    // only need column entries for the categorical variables (others are same
    // as f). Since, passing a smaller matrix into the transition function is
    // faster, for the no interaction case, we don't need to input the whole
    // matrix into this function. For the interaction case, since it is unknown
    // which indices have interaction, we require entries for all columns in the
    // matrices.
    bool no_interactions = (f_set_mat.cols() < N);
    for (Index i = 0; i < state.numCategoricalVarsInSubset; ++i) {
        // Note: categorical_indices are assumed to be zero-based
        ColumnVector f_set;
        ColumnVector f_unset;
        ColumnVector shortened_f_set = f_set_mat.row(i);
        ColumnVector shortened_f_unset = f_unset_mat.row(i);

        if (no_interactions){
            f_set = f;
            f_unset = f;
            for (Index j=0; j < shortened_f_set.size(); ++j){
                f_set(static_cast<Index>(categorical_indices(j))) = shortened_f_set(j);
                f_unset(static_cast<Index>(categorical_indices(j))) = shortened_f_unset(j);
            }
        } else {
            f_set = shortened_f_set;
            f_unset = shortened_f_unset;
        }

        RowVector p_set(trans(f_set) * beta);
        {
            p_set = p_set.array().exp();
            double p_sum = p_set.sum();
            p_set = p_set / (1 + p_sum);
        }

        RowVector p_unset(trans(f_unset) * beta);
        {
            p_unset = p_unset.array().exp();
            double p_sum = p_unset.sum();
            p_unset = p_unset / (1 + p_sum);
        }
        // Compute the marginal effect using difference method
        curr_margins.row(static_cast<uint16_t>(state.categorical_basis_indices(i))) = p_set - p_unset;

        // Compute the delta using difference method
        int m = static_cast<uint16_t>(state.categorical_basis_indices(i));
        for (int l = 0; l < L - 1; l++) {
            row_index = reindex(m, l, L - 1);
            for (int n = 0; n < N; n++) {
                for (int l1 = 0; l1 < L - 1; l1++) {
                    double delta = - p_set(l) * p_set(l1) * f_set(n) + p_unset(l) * p_unset(l1) * f_unset(n);
                    if (l1 == l)
                        delta += p_set(l) * f_set(n) - p_unset(l) * f_unset(n);
                    col_index = reindex(n, l1, L - 1);
                    state.delta(row_index, col_index) += delta;
                }
            }
        }

    }
    state.marginal_effects += curr_margins;
    return state;
}

/**
 * @brief Marginal effects: Merge transition states
 */
AnyType
margins_mlogregr_int_merge::run(AnyType &args) {
    // In case the aggregator should be terminated because
    // an exception has been "thrown" in the transition function
    if (args[0].isNull() || args[1].isNull())
        return Null();
    MarginsMLogregrInteractionState<MutableArrayHandle<double> > stateLeft = args[0];
    MarginsMLogregrInteractionState<ArrayHandle<double> > stateRight = args[1];
    // We first handle the trivial case where this function is called with one
    // of the states being the initial state
    if (stateLeft.numRows == 0)
        return stateRight;
    else if (stateRight.numRows == 0)
        return stateLeft;

    // Merge states together and return
    stateLeft += stateRight;
    return stateLeft;
}

/**
 * @brief Marginal effects: Final step
 */
AnyType
margins_mlogregr_int_final::run(AnyType &args) {
    // In case the aggregator should be terminated because
    // an exception has been "thrown" in the transition function
    if (args[0].isNull())
        return Null();
    // We request a mutable object.
    // Depending on the backend, this might perform a deep copy.
    MarginsMLogregrInteractionState<MutableArrayHandle<double> > state = args[0];
    // Aggregates that haven't seen any data just return Null.
    if (state.numRows == 0)
        return Null();

    state.marginal_effects /= static_cast<double>(state.numRows);
    Matrix marginal_effects_trans = trans(state.marginal_effects);
    AnyType tuple;
    tuple << marginal_effects_trans;

    // Variance for marginal effects according to the delta method
    Matrix variance(state.delta * state.training_data_vcov);
    // // we only need the diagonal elements of the variance, so we perform a dot
    // // product of each row with itself to compute each diagonal element.
    // // We divide by numRows^2 since we need the average variance
    Matrix std_err = variance.cwiseProduct(state.delta).rowwise().sum() / static_cast<double>(state.numRows * state.numRows);
    std_err = std_err.array().sqrt();
    std_err.resize(state.numCategories-1, state.numBasis);
    tuple << std_err;

    Matrix t_stats = marginal_effects_trans.cwiseQuotient(std_err);
    tuple << t_stats;

    // Note: p-values will return NULL if numRows <= coef.size
    if (state.numRows > state.numBasis) {
        MutableNativeMatrix p_values(
                this->allocateArray<double>(state.numBasis * (state.numCategories-1)),
                state.numCategories-1,
                state.numBasis);
        for (Index l = 0; l < p_values.rows(); l ++) {
            for (Index m = 0; m < p_values.cols(); m ++) {
                p_values(l,m) = 2. * prob::cdf(prob::normal(), -std::abs(t_stats(l,m)));
            }
        }
        tuple << static_cast<Matrix>(p_values);
    }

    return tuple;
}
// ------------------------ End of Logistic Marginal ---------------------------

} // namespace regress

} // namespace modules

} // namespace madlib