/* ----------------------------------------------------------------------- */
/**
 *
 * @file cox_proportional_hazards.cpp
 *
 * @brief Cox proportional hazards
 *
 * ----------------------------------------------------------------------- */

#include <dbconnector/dbconnector.hpp>
#include <modules/shared/HandleTraits.hpp>
#include <modules/prob/boost.hpp>
#include <math.h>

#include "marginal_cox.hpp"

namespace madlib {
namespace modules {
namespace stats {

// Use Eigen
using namespace dbal::eigen_integration;
// Import names from other MADlib modules
using dbal::NoSolutionFoundException;

using namespace std;

// ---------------------------------------------------------------------------
//            Marginal Effects Cox Proportional Hazard State
// ---------------------------------------------------------------------------
/**
 * @brief State for marginal effects calculation for cox proportional hazards
 *
 * TransitionState encapsualtes the transition state during the
 * marginal effects calculation.
 * 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 MarginsCoxPropHazardState {
    template <class OtherHandle>
    friend class MarginsCoxPropHazardState;

  public:
    MarginsCoxPropHazardState(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]));
    }

    /**
     * @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>
    MarginsCoxPropHazardState &operator=(
        const MarginsCoxPropHazardState<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>
    MarginsCoxPropHazardState &operator+=(
        const MarginsCoxPropHazardState<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;
        baseline_hazard = 0;
        marginal_effects.fill(0);
        training_data_vcov.fill(0);
        delta.fill(0);
        if (numCategoricalVarsInSubset > 0)
            categorical_basis_indices.fill(0);
    }

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

    /**
     * @brief Rebind to a new storage array
     *
     * @param inWidthOfX The number of independent variables.
     *
     */
    void rebind(const uint16_t inWidthOfX,
                const uint16_t inNumBasis,
                const uint16_t inNumCategoricalVars) {
        widthOfX.rebind(&mStorage[0]);
        numBasis.rebind(&mStorage[1]);
        numCategoricalVarsInSubset.rebind(&mStorage[2]);
        numRows.rebind(&mStorage[3]);
        baseline_hazard.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 (inNumCategoricalVars > 0)
            categorical_basis_indices.rebind(&mStorage[5 + inNumBasis +
                                            (inWidthOfX + inNumBasis)*inWidthOfX],
                                            inNumCategoricalVars);
    }
    Handle mStorage;

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

/**
 * @brief Perform the marginal effects transition step
 */
AnyType
margins_coxph_int_transition::run(AnyType &args) {

    MarginsCoxPropHazardState<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.");

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


    // basis_indices represents the indices (of beta) for 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.size());
    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");
        }
        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.");

        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));
        state.reset();

        // coxph stores the Hessian matrix in it's output table. This is
        // different from the various regression methods that provides the
        // training vcov matrix directly.
        Matrix hessian = args[3].getAs<MappedMatrix>();
        // Computing pseudo inverse of a PSD matrix
        SymmetricPositiveDefiniteEigenDecomposition<Matrix> decomposition(
            hessian, EigenvaluesOnly, ComputePseudoInverse);
        state.training_data_vcov = decomposition.pseudoInverse();

        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);
    state.baseline_hazard += f_beta;
    double exp_f_beta = 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 = exp_f_beta * J_trans_beta;
    Matrix curr_delta = exp_f_beta * (trans(J) + J_trans_beta * trans(f));

    // 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 h_set = exp(dot(f_set, beta));
        double h_unset = exp(dot(f_unset, beta));

        curr_margins(static_cast<uint16_t>(state.categorical_basis_indices(i))) = h_set - h_unset;
        curr_delta.row(static_cast<uint16_t>(state.categorical_basis_indices(i))) =
            (h_set * trans(f_set) - h_unset * trans(f_unset));
    }

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


/**
 * @brief Cox Proportional Hazards marginal effects: Merge transition states
 */
AnyType
margins_coxph_int_merge::run(AnyType &args) {
    MarginsCoxPropHazardState<MutableArrayHandle<double> > stateLeft = args[0];
    MarginsCoxPropHazardState<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 Cox Proportional Hazards marginal effects: Final step
 */
AnyType
margins_coxph_int_final::run(AnyType &args) {
    // We request a mutable object.
    // Depending on the backend, this might perform a deep copy.
    MarginsCoxPropHazardState<MutableArrayHandle<double> > state = args[0];
    // Aggregates that haven't seen any data just return Null.
    if (state.numRows == 0)
        return Null();

    state.baseline_hazard = exp(state.baseline_hazard / static_cast<double>(state.numRows));
    state.baseline_hazard = 1;
    state.marginal_effects /= static_cast<double>(state.numRows) / state.baseline_hazard;

    // 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 std_err =
        variance.cwiseProduct(state.delta).rowwise().sum();
    std_err = std_err.array().sqrt() / static_cast<double>(state.numRows);

    MutableNativeColumnVector tStats(this->allocateArray<double>(state.numBasis));
    MutableNativeColumnVector pValues(this->allocateArray<double>(state.numBasis));

    for (Index i = 0; i < state.numBasis; ++i) {
        tStats(i) = state.marginal_effects(i) / std_err(i);

        // p-values only make sense if numRows > coef.size()
        if (state.numRows > state.numBasis)
            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 << state.marginal_effects
          << std_err
          << tStats
          << (state.numRows > state.numBasis ? pValues: Null());
    return tuple;
}


/**
 * @brief Cox Proportional Hazards marginal effects: Statistics function
 */
AnyType
margins_compute_stats::run(AnyType &args) {
    // NULL input returns a NULL output
    if (args[0].isNull() || args[1].isNull()) {
        return Null();
    }

    MappedColumnVector marginal_effects = args[0].getAs<MappedColumnVector>();
    MappedColumnVector std_err = args[1].getAs<MappedColumnVector>();
    uint16_t n_basis_terms = static_cast<uint16_t>(marginal_effects.size());
    MutableNativeColumnVector tStats(
        (*this).allocateArray<double>(n_basis_terms));
    MutableNativeColumnVector pValues(
        (*this).allocateArray<double>(n_basis_terms));

    for (Index i = 0; i < n_basis_terms; ++i) {
        tStats(i) = marginal_effects(i) / std_err(i);
        pValues(i) = 2. * prob::cdf( prob::normal(),
                                         -std::abs(tStats(i)));
    }

    // Return all standard errors and p_values in a tuple
    AnyType tuple;
    tuple << marginal_effects << std_err << tStats << pValues;
    return tuple;
}
// ----------------- End of Marginal Effects for Cox proportional hazards ------

} // namespace stats
} // namespace modules
} // namespace madlib