library(MASS) # required by 'dr'
library(dr)   # for `dr`
library(MAVE) # for `mave`, `coef.mave`.
library(CVE)

################################################################################
###                       simulation utility functions                       ###
################################################################################

#' Calc distance between two subspaces.
#'
#'     || P(A) - P(B) ||_F / sqrt(2 * ncol(A))
#'
#' such that `P(A) = A (A^T A)^-1 A^T` the projection onto
#' the subspace spanned by `A`.
#'
#' @param A matrix which rows span a subspace.
#' @param B matrix of the same dim. as `A`-
#'
#' @returns numeric metric of spanning subspaces distance.
subspace_dist <- function(A, B) {
    P <- A %*% solve(t(A) %*% A, t(A))
    Q <- B %*% solve(t(B) %*% B, t(B))
    norm(P - Q, 'F') / sqrt(2 * ncol(A))
}

#' Creats a progress logger.
#'
#' @param len number of calls to the creted rpogress logger till finished.
#'
#' @returns progress logger function:
#'     \code{logger(info)}
#' which is to be called \code{len} times with a string \code{info}.
progress_logger <- function(len) {
    start <- Sys.time()
    count <- 0
    return(function(info) {
        cat("\rRunning (", info, "):",
            (count <<- count + 1), "/", len,
            " - elapsed:", format(Sys.time() - start), "\033[K")
        if (count == len) { cat("\n") }
    })
}

################################################################################
###                                cve wrapper                               ###
################################################################################
CVE <- function(...) {
    .cve <- match.call(expand.dots = TRUE)
    .cve[[1L]] <- quote(cve.call)
    .cve$method <- "simple"
    eval(.cve)
}
wCVE <- function(...) {
    .cve <- match.call(expand.dots = TRUE)
    .cve[[1L]] <- quote(cve.call)
    .cve$method <- "weighted"
    eval(.cve)
}
rCVE <- function(...) {
    .cve <- match.call(expand.dots = TRUE)
    .cve[[1L]] <- quote(cve.call)
    .cve$method <- "simple"
    dr <- eval(.cve)

    for (name in names(dr$res)) {
        .cve$method <- "weighted"
        .cve$k <- dr$res[[name]]$k
        .cve$V.init <- dr$res[[name]]$V
        dr$res[[name]] <- eval(.cve)$res[[name]]
    }
    dr$call <- match.call()
    dr$call[[1L]] <- quote(cve)
    dr$call$method <- "refined"
    return(dr)
}

wls <- function(x,y,w){ 
    x <- as.matrix(x); y <- as.vector(y)
  n=dim(x)[1];
  p=dim(x)[2]-1 
  out=c(solve(t(x*w)%*%x/n)%*%apply(x*y*w,2,mean))
  return(list(a=out[1],b=out[2:(p+1)]))
}

#' Compute kernel matrix.
#'
#' @param X matrix of dim. `n x p` (`n` samples of dimension `p`)
#' @param h numeric bandwidth.
#'
#' @returns `n x n` matrix with kernel of X_i to X_j interaction
#'      in its `i, j` component.
kern <- function(X, h) {
    X <- as.matrix(X)
    n <- nrow(X)

    k2 <- X %*% t(X)                # (X_i^T X_j)_ij
    k3 <- matrix(diag(k2), n, n)    # diag(k2) in rows
    k1 <- t(k3)                     # diag(k2) in cols
    K <- k1 - 2 * k2 + k3           # (|X_i|^2 - 2 * X_i^T X_j - |X_j|^2)_ij
    return(exp((-0.5 / h^2) * K))
}

coef.opg <- function(object, ...) {
    return(object$B)
}

OPG <- function(X, Y, k, c0 = 2.34) {
    # Cast parameters.
    X <- as.matrix(X); Y <- as.vector(Y)
    # Get Dimensions
    n <- nrow(X); p <- ncol(X)

    p0 = max(p, 3)
    rn <- n^(-1 / (2 * (p0 + 6)))
    h <- c0 * n^(-1 / (p0 + 6))
    sig <- diag(var(X))
    X <- scale(X)
    kmat <- kern(X, h)
    bmat <- numeric()
    for (i in 1:n) {
        wi <- kmat[, i]
        xi <- cbind(1, t(t(X) - X[i, ]))
        bmat <- cbind(bmat, wls(xi, Y, wi)$b)
    }
    beta <- eigen(bmat %*% t(bmat))$vectors[, 1:k]
    B <- diag(sig^(-1 / 2)) %*% beta

    return(structure(list(B = B), class = "opg"))
}

rOPG <- function(X, Y, k, nr.iter = 25L, c0 = 2.34) {
    # Cast parameters.
    X <- as.matrix(X); Y <- as.vector(Y)
    # Get Dimensions
    n <- nrow(X); p <- ncol(X)
  
    sig <- diag(var(X))
    X <- scale(X)

    p0 = max(p, 3)
    rn <- n^(-1 / (2 * (p0 + 6)))
    h <- c0 * n^(-1 / (p0 + 6)) 
    beta <- diag(p)
    for(. in seq_len(nr.iter)) {
        kmat <- kern(X %*% beta, h)
        bmat <- numeric()
        for (i in 1:n) {
            wi <- kmat[, i]
            xi <- cbind(1, t(t(X) - X[i, ])) 
            bmat <- cbind(bmat, wls(xi, Y, wi)$b)
        }
        beta <- eigen(bmat %*% t(bmat))$vectors[, 1:k, drop = F]
        h <- max(rn * h, c0 * n^((-1/(k + 4)))) 
    }
    B <- diag(sig^(-1/2)) %*% beta 

    return(structure(list(B = B), class = "opg"))
}

rmave <- function(X, Y, k, nr.iter = 25L, c0 = 2.34) {
    X <- as.matrix(X); Y <- as.vector(Y)
    # Get dimensions.
    n <- nrow(X); p <- ncol(X)

    X <- scale(X)
    sig <- attr(X, "scaled:scale")^2 # diag(var(X))

    p0 <- max(p,3)
    h <- c0 * n^(-1 / (p0 + 6))
    rn <- n^(-1 / (2 * (p0 + 6)))

    beta <- OPG(X, Y, k)$B
    for (. in seq_len(nr.iter)) {
        kermat <- kern(X %*% beta, h)
        mkermat <- colMeans(kermat)

        a <- numeric()
        b <- numeric()
        for (i in 1:n) {
            wi <- kermat[ ,i] / mkermat[i]
            ui <- cbind(1, t(t(X) - X[i, ]) %*% beta)
            out <- wls(ui, Y, wi)
            a <- c(a, out$a)
            b <- cbind(b, out$b)
        }
        out <- 0
        out1 <- 0
        for (i in 1:n) {
            xi <- kronecker(t(t(X) - X[i, ]), t(b[, i]))
            yi <- Y - a[i]
            wi <- kermat[, i] / mkermat[i]
            out <- out + colMeans(xi * yi * wi)
            out1 <- out1 + t(xi * wi) %*% xi / n
        }
        beta <- t(matrix(solve(out1, out), k, p))
        h <- max(rn * h, c0 * n^(-1 / (k + 4)))
    }
    B <- diag(sig^-0.5) %*% beta
    
    return(structure(list(B = B), class = "opg"))
}

meanOPG <- function(X, Y, max.dim = 10L, ...) {
    mave(Y ~ X, method = "meanOPG", max.dim = max.dim)
}

meanMAVE <- function(X, Y, max.dim = 10L, ...) {
    mave(Y ~ X, method = "meanMAVE", max.dim = max.dim)
}

coef.phdy <- coef.sir <- coef.save <- function(object, k, ...) {
    matrix(object$evectors[, 1:k], ncol = k)
}

phdy <- function(X, Y, k, ...) {
    dr(Y ~ X, method = 'phdy', numdir = k)
}

sir <- function(X, Y, k, ...) {
    dr(Y ~ X, method = 'sir', numdir = k)
}

save <- function(X, Y, k, ...) {
    dr(Y ~ X, method = 'save', numdir = k)
}