//
// Development file.
//
// Usage:
//  ~$ R -e "library(Rcpp); Rcpp::sourceCpp('cve_V1.cpp')"
//

// only `RcppArmadillo.h` which includes `Rcpp.h`
#include <RcppArmadillo.h>

// through the depends attribute `Rcpp` is tolled to create
// hooks for `RcppArmadillo` needed by the build process.
//
// [[Rcpp::depends(RcppArmadillo)]]

// required for `R::qchisq()` used in `estimateBandwidth()`
#include <Rmath.h>

//' Estimated bandwidth for CVE.
//'
//' Estimates a propper bandwidth \code{h} according
//' \deqn{h = \chi_{p-q}^{-1}\left(\frac{nObs - 1}{n-1}\right)\frac{2 tr(\Sigma)}{p}}{%
//'       h = qchisq( (nObs - 1)/(n - 1), p - q ) 2 tr(Sigma) / p}
//'
//' @param X data matrix of dimension (n x p) with n data points X_i of dimension
//'  q. Therefor each row represents a datapoint of dimension p.
//' @param k Guess for rank(B).
//' @param nObs Ether numeric of a function. If specified as numeric value
//'  its used in the computation of the bandwidth directly. If its a function
//'  `nObs` is evaluated as \code{nObs(nrow(x))}. The default behaviou if not
//'  supplied at all is to use \code{nObs <- nrow(x)^0.5}.
//'
//' @seealso [qchisq()]
//'
//' @export
// [[Rcpp::export]]
double estimateBandwidth(const arma::mat& X, arma::uword k, double nObs) {
    using namespace arma;

    uword n = X.n_rows; // nr samples
    uword p = X.n_cols; // dimension of rand. var. `X`

    // column mean
    mat M = mean(X);
    // center `X`
    mat C = X.each_row() - M;
    // trace of covariance matrix, `traceSigma = Tr(C' C)`
    double traceSigma = accu(C % C);
    // compute estimated bandwidth
    double qchi2 = R::qchisq((nObs - 1.0) / (n - 1), static_cast<double>(k), 1, 0);

    return 2.0 * qchi2 * traceSigma / (p * n);
}

//' Random element from Stiefel Manifold `S(p, q)`.
//'
//' Draws an element of \eqn{S(p, q)} which is the Stiefel Manifold.
//' This is done by taking the Q-component of the QR decomposition
//' from a `(p, q)` Matrix with independent standart normal entries.
//' As a semi-orthogonal Matrix the result `V` satisfies \eqn{V'V = I_q}.
//'
//' @param p Row dimension
//' @param q Column dimension
//'
//' @return Matrix of dim `(p, q)`.
//'
//' @seealso <https://en.wikipedia.org/wiki/Stiefel_manifold>
//'
//' @export
// [[Rcpp::export]]
arma::mat rStiefel(arma::uword p, arma::uword q) {
    arma::mat Q, R;
    arma::qr_econ(Q, R, arma::randn<arma::mat>(p, q));
    return Q;
}

//' Gradient computation of the loss `L_n(V)`.
//'
//' The loss is defined as
//' \deqn{L_n(V) := \frac{1}{n}\sum_{j=1}^n (y_2(V, X_j) - y_1(V, X_j)^2)}
//' with
//' \deqn{y_l(s_0) := \sum_{i=1}^n w_i(V, s_0)Y_i^l}{y_l(s_0) := sum_i(w_i(V, s_0) Y_i^l)}
//'
//' @rdname optStiefel
//' @keywords internal
double gradient(const arma::mat& X,
                const arma::mat& X_diff,
                const arma::mat& Y,
                const arma::mat& Y_rep,
                const arma::mat& V,
                const double h,
                arma::mat* G = 0
) {
    using namespace arma;

    uword n = X.n_rows;
    uword p = X.n_cols;

    // orthogonal projection matrix `Q = I - VV'` for dist computation
    mat Q = -(V * V.t());
    Q.diag() += 1;
    // calc pairwise distances as `D` with `D(i, j) = d_i(V, X_j)`
    vec D_vec = sum(square(X_diff * Q), 1);
    mat D = reshape(D_vec, n, n);
    // calc weights as `W` with `W(i, j) = w_i(V, X_j)`
    mat W = exp(D / (-2.0 * h));
    // column-wise normalization via 1-norm
    W = normalise(W, 1);
    vec W_vec = vectorise(W);
    // weighted `Y` means (first and second order)
    vec y1 = W.t() * Y;
    vec y2 = W.t() * square(Y);
    // loss for each `X_i`, meaning `L(i) = L_n(V, X_i)`
    vec L = y2 - square(y1);
    // `loss = L_n(V)`
    double loss = mean(L);
    // check if gradient as output variable is set
    if (G != 0) {
        // `G = grad(L_n(V))` a.k.a. gradient of `L` with respect to `V`
        vec scale = (repelem(L, n, 1) - square(Y_rep - repelem(y1, n, 1))) % W_vec % D_vec;
        mat X_diff_scale = X_diff.each_col() % scale;
        (*G) = X_diff_scale.t() * X_diff * V;
        (*G) *= -2.0 / (h * h * n);
    }

    return loss;
}

//' Stiefel Optimization.
//'
//' Stiefel Optimization for \code{V} given a dataset \code{X} and responces
//' \code{Y} for the model \deqn{Y\sim g(B'X) + \epsilon}{Y ~ g(B'X) + epsilon}
//' to compute the matrix `B` such that \eqn{span{B} = span(V)^{\bot}}{%
//' span(B) = orth(span(B))}.
//'
//' @param X data points
//' @param Y response
//' @param k assumed \eqn{rank(B)}
//' @param nObs parameter for bandwidth estimation, typical value
//'     \code{nObs = nrow(X)^lambda} with \code{lambda} in the range [0.3, 0.8].
//' @param tau Initial step size
//' @param tol Tolerance for update error used for stopping criterion
//'     \eqn{|| V(j) V(j)' - V(j+1) V(j+1)' ||_2 < tol}{%
//'          \| V_j V_j' - V_{j+1} V_{j+1}' \|_2 < tol}.
//' @param maxIter Upper bound of optimization iterations
//'
//' @return List containing the bandwidth \code{h}, optimization objective \code{V}
//'     and the matrix \code{B} estimated for the model as a orthogonal basis of the
//'     orthogonal space spaned by \code{V}.
//'
//' @rdname optStiefel
//' @keywords internal
double optStiefel(
        const arma::mat& X,
        const arma::vec& Y,
        const int k,
        const double h,
        const double tauInitial,
        const double tol,
        const double slack,
        const int maxIter,
        arma::mat& V, // out
        arma::vec& history // out
) {
    using namespace arma;

    // get dimensions
    const uword n = X.n_rows; // nr samples
    const uword p = X.n_cols; // dim of random variable `X`
    const uword q = p - k;    // rank(V) e.g. dim of ortho space of span{B}

    // all `X_i - X_j` differences, `X_diff.row(i * n + j) = X_i - X_j`
    mat X_diff(n * n, p);
    for (uword i = 0, k = 0; i < n; ++i) {
        for (uword j = 0; j < n; ++j) {
            X_diff.row(k++) = X.row(i) - X.row(j);
        }
    }
    const vec Y_rep = repmat(Y, n, 1);
    const mat I_p = eye<mat>(p, p);

    // initial start value for `V`
    V = rStiefel(p, q);

    // init optimization `loss`es, `error` and stepsize `tau`
    // double loss_next = datum::inf;
    double loss;
    double error = datum::inf;
    double tau = tauInitial;
    int iter;
    // main optimization loop
    for (iter = 0; iter < maxIter && error > tol; ++iter) {
        // calc gradient `G = grad_V(L)(V)`
        mat G;
        loss = gradient(X, X_diff, Y, Y_rep, V, h, &G);
        // matrix `A` for colescy-transform of the gradient
        mat A = tau * ((G * V.t()) - (V * G.t()));
        // next iteration step of `V`
        mat V_tau = inv(I_p + A) * (I_p - A) * V;
        // loss after step `L(V(tau))`
        double loss_tau = gradient(X, X_diff, Y, Y_rep, V_tau, h);

        // store `loss` in `history` and increase `iter`
        history(iter) = loss;

        // validate if loss decreased
        if ((loss_tau - loss) > slack * loss) {
            // ignore step, retry with half the step size
            tau = tau / 2.;
            error = datum::inf;
        } else {
            // compute step error (break condition)
            error = norm((V * V.t()) - (V_tau * V_tau.t()), 2) / (2 * q);
            // shift for next iteration
            V = V_tau;
            loss = loss_tau;
        }
    }

    // store final `loss`
    history(iter) = loss;

    return loss;
}

//' Conditional Variance Estimation (CVE) method.
//'
//' This version uses a "simple" stiefel optimization schema.
//'
//' @param X data points
//' @param Y response
//' @param k assumed \eqn{rank(B)}
//' @param nObs parameter for bandwidth estimation, typical value
//'     \code{nObs = nrow(X)^lambda} with \code{lambda} in the range [0.3, 0.8].
//' @param tau Initial step size (default 1)
//' @param tol Tolerance for update error used for stopping criterion (default 1e-5)
//' @param slack Ratio of small negative error allowed in loss optimization (default -1e-10)
//' @param maxIter Upper bound of optimization iterations (default 50)
//' @param attempts Number of tryes with new random optimization starting points (default 10)
//'
//' @return List containing the bandwidth \code{h}, optimization objective \code{V}
//'     and the matrix \code{B} estimated for the model as a orthogonal basis of the
//'     orthogonal space spaned by \code{V}.
//'
//' @rdname cve_cpp
//' @export
// [[Rcpp::export]]
Rcpp::List cve_cpp(
        const arma::mat& X,
        const arma::vec& Y,
        const int k,
        const double nObs,
        const double tauInitial = 1.,
        const double tol = 1e-5,
        const double slack = -1e-10,
        const int maxIter = 50,
        const int attempts = 10
) {
    using namespace arma;

    // tracker of current best results
    mat V_best;
    double loss_best = datum::inf;
    // estimate bandwidth
    double h = estimateBandwidth(X, k, nObs);

    // loss history for each optimization attempt
    // each column contaions the iteration history for the loss
    mat history = mat(maxIter + 1, attempts);

    // multiple stiefel optimization attempts
    for (int i = 0; i < attempts; ++i) {
        // declare output variables
        mat V; // estimated `V` space
        vec hist = vec(history.n_rows, fill::zeros); // optimization history
        double loss = optStiefel(X, Y, k, h, tauInitial, tol, slack, maxIter, V, hist);
        if (loss < loss_best) {
            loss_best = loss;
            V_best = V;
        }
        // write history to history collection
        history.col(i) = hist;
    }

    // get `B` as kernal of `V.t()`
    mat B = null(V_best.t());

    return Rcpp::List::create(
        Rcpp::Named("history") = history,
        Rcpp::Named("loss") = loss_best,
        Rcpp::Named("h") = h,
        Rcpp::Named("V") = V_best,
        Rcpp::Named("B") = B
    );
}

/*** R

source("CVE/R/datasets.R")
ds <- dataset()

print(system.time(
    cve.res <- cve_cpp(
        X = ds$X,
        Y = ds$Y,
        k = ncol(ds$B),
        nObs = sqrt(nrow(ds$X))
    )
))

pdf('plots/cve_V1_history.pdf')
H <- cve.res$history
H_i <- H[H[, 1] > 0, 1]
plot(1:length(H_i), H_i,
    main = "History cve_V1",
    xlab = "Iterations i",
    ylab = expression(loss == L[n](V^{(i)})),
    xlim = c(1, nrow(H)),
    ylim = c(0, max(H)),
    type = "l"
)
for (i in 2:ncol(H)) {
    H_i <- H[H[, i] > 0, i]
    lines(1:length(H_i), H_i)
}
x.ends <- apply(H, 2, function(h) { length(h[h > 0]) })
y.ends <- apply(H, 2, function(h) { min(h[h > 0]) })
points(x.ends, y.ends)

*/
