# ============================================================
#  krnn.R
#  k-Ranked Nearest Neighbour classifier for functional data
#  Sections 3.2 - 3.4 of Perez Ruiz (2020)
# ============================================================


# ---- Internal helpers --------------------------------------------------

#' @keywords internal
.krnn_predict_sdp <- function(sdp_new, sdp_train, y_train, k) {
  o          <- order(sdp_train)
  sdp_sorted <- sdp_train[o]
  y_sorted   <- y_train[o]
  n_tr       <- length(sdp_sorted)
  vapply(sdp_new, function(s0) {
    pos <- findInterval(s0, sdp_sorted)
    lo  <- max(1L, pos - k + 1L)
    hi  <- min(n_tr, pos + k)
    neigh_labels <- y_sorted[lo:hi]
    tab <- table(neigh_labels)
    as.integer(names(tab)[which.max(tab)])
  }, integer(1L))
}

#' @keywords internal
.smoother_prob <- function(sdp_new, sdp_train, y_train, k,
                           pre_sorted = FALSE) {
  if (!pre_sorted) {
    o         <- order(sdp_train)
    sdp_train <- sdp_train[o]
    y_train   <- y_train[o]
  }
  n_tr <- length(sdp_train)
  vapply(sdp_new, function(s0) {
    pos <- findInterval(s0, sdp_train)
    lo  <- max(1L, pos - k + 1L)
    hi  <- min(n_tr, pos + k)
    mean(y_train[lo:hi] == 0L)
  }, numeric(1L))
}

#' @keywords internal
.project_signed_depth <- function(X_new, sd_train) {
  sdi_new   <- signed_distance_integral(X_new,
                                        x_ref = sd_train$x_ref,
                                        grid  = sd_train$grid)
  depth_new <- fm_depth(X_new, grid = sd_train$grid, type = sd_train$type)
  sign(sdi_new) * depth_new
}

#' @keywords internal
.bootstrap_ci <- function(sdp_sorted, y_sorted, k, B = 200L, alpha = 0.05) {
  p_hat    <- .smoother_prob(sdp_sorted, sdp_sorted, y_sorted, k,
                             pre_sorted = TRUE)
  n        <- length(sdp_sorted)
  boot_mat <- matrix(NA_real_, nrow = n, ncol = B)
  for (b in seq_len(B)) {
    y_boot        <- stats::rbinom(n, 1L, 1 - p_hat)
    boot_mat[, b] <- .smoother_prob(sdp_sorted, sdp_sorted, y_boot, k,
                                    pre_sorted = TRUE)
  }
  se_boot <- apply(boot_mat, 1, stats::sd)
  z_boot  <- sweep(boot_mat, 1, p_hat, "-") / (se_boot + 1e-10)
  t_lo    <- apply(z_boot, 1, stats::quantile, probs = alpha / 2)
  t_hi    <- apply(z_boot, 1, stats::quantile, probs = 1 - alpha / 2)
  list(lower = pmax(0, p_hat - t_hi * se_boot),
       upper = pmin(1, p_hat - t_lo * se_boot))
}


# -----------------------------------------------------------------------
#  krnn_fit
# -----------------------------------------------------------------------

#' k-Ranked Nearest Neighbour Classifier for Functional Data
#'
#' Fits the k-RNN classifier of Perez Ruiz (2020). Curves are ranked by their
#' signed depth and each new observation is assigned to the majority class
#' among its \eqn{2k} nearest ranked neighbours (k above and k below).
#'
#' @param X_train Numeric matrix (\eqn{n_{\mathrm{train}} \times M}) of
#'   training curves (rows = curves, columns = time points).
#' @param y_train Integer vector of class labels (0/1) of length
#'   \eqn{n_{\mathrm{train}}}.
#' @param X_test Numeric matrix of test curves. If \code{NULL} (default) the
#'   training set is used for in-sample evaluation.
#' @param k Positive integer: half-neighbourhood size. The total number of
#'   neighbours per observation is \eqn{2k}. Use \code{\link{krnn_cv}} to
#'   select k.
#' @param grid Numeric vector of length \eqn{M} giving observation times.
#'   Defaults to equally spaced points on \eqn{[0, 1]}.
#' @param type Depth variant passed to \code{\link{fm_depth}}: \code{"tukey"}
#'   (default) or \code{"simplicial"}.
#' @param sd_train Optional pre-computed \code{"fd_signed_depth"} object for
#'   the training set (avoids recomputation when called repeatedly).
#'
#' @return An object of class \code{"krnn_fit"} with components:
#'   \describe{
#'     \item{predicted}{Integer vector of predicted class labels.}
#'     \item{prob}{Numeric vector of estimated probabilities
#'       \eqn{\hat{P}(y=0|\mathrm{sdp})}.}
#'     \item{sd_train}{The \code{"fd_signed_depth"} object for the training
#'       set.}
#'     \item{y_train}{The training labels.}
#'     \item{k}{The value of k used.}
#'   }
#'
#' @seealso \code{\link{krnn_cv}}, \code{\link{predict.krnn_fit}}
#'
#' @references
#' Perez Ruiz, D. A. (2020). \emph{Supervised Classification for Functional
#' Data}. PhD thesis, University of Manchester. Section 3.2.
#'
#' @examples
#' set.seed(7)
#' M <- 100; n <- 80
#' t <- seq(0, 1, length.out = M)
#' X0 <- t(replicate(n / 2, sin(2 * pi * t) + rnorm(M, sd = 0.4)))
#' X1 <- t(replicate(n / 2, cos(2 * pi * t) + rnorm(M, sd = 0.4)))
#' X  <- rbind(X0, X1)
#' y  <- rep(0:1, each = n / 2)
#' idx_tr <- sample(n, 60)
#' fit <- krnn_fit(X[idx_tr, ], y[idx_tr], X[-idx_tr, ], k = 5)
#' print(fit)
#'
#' @export
krnn_fit <- function(X_train, y_train,
                     X_test   = NULL,
                     k        = 5L,
                     grid     = NULL,
                     type     = c("tukey", "simplicial"),
                     sd_train = NULL) {
  mc      <- match.call()
  type    <- match.arg(type)
  y_train <- as.integer(y_train)
  if (!all(y_train %in% c(0L, 1L)))
    stop("'y_train' must contain only 0 and 1.")
  if (!is.matrix(X_train))
    stop("'X_train' must be a matrix.")
  k <- as.integer(k)
  if (k < 1L) stop("'k' must be a positive integer.")
  M <- ncol(X_train)
  if (is.null(grid)) grid <- seq(0, 1, length.out = M)
  if (is.null(sd_train))
    sd_train <- signed_depth(X_train, grid = grid, type = type)
  sdp_tr <- sd_train$sdp
  if (is.null(X_test)) {
    sdp_te <- sdp_tr
  } else {
    if (ncol(X_test) != M)
      stop("'X_test' must have the same number of columns as 'X_train'.")
    sdp_te <- .project_signed_depth(X_test, sd_train)
  }
  predicted <- .krnn_predict_sdp(sdp_te, sdp_tr, y_train, k)
  prob      <- .smoother_prob(sdp_te, sdp_tr, y_train, k)
  structure(
    list(predicted = predicted, prob = prob, sd_train = sd_train,
         y_train = y_train, k = k, type = type, call = mc),
    class = "krnn_fit"
  )
}

#' @param x A \code{"krnn_fit"} object.
#' @param ... Ignored.
#' @return Invisibly returns \code{x}.
#' @rdname krnn_fit
#' @export
print.krnn_fit <- function(x, ...) {
  cat("k-RNN Functional Classifier\n")
  cat("  k          :", x$k, "\n")
  cat("  Depth type :", x$type, "\n")
  cat("  Training N :", x$sd_train$N, "\n")
  invisible(x)
}


# -----------------------------------------------------------------------
#  predict.krnn_fit
# -----------------------------------------------------------------------

#' Predict Method for krnn_fit Objects
#'
#' Classifies new functional observations using a fitted
#' \code{\link{krnn_fit}} object.
#'
#' @param object A \code{"krnn_fit"} object.
#' @param newdata Numeric matrix of new curves
#'   (\eqn{n_{\mathrm{new}} \times M}).
#' @param y_true Optional integer vector of true labels for computing the
#'   misclassification rate.
#' @param ... Ignored.
#'
#' @return A list with components \code{predicted}, \code{prob}, and (if
#'   \code{y_true} is supplied) \code{misclass}.
#'
#' @examples
#' set.seed(7)
#' M <- 100; n <- 80
#' t <- seq(0, 1, length.out = M)
#' X0 <- t(replicate(n / 2, sin(2 * pi * t) + rnorm(M, sd = 0.4)))
#' X1 <- t(replicate(n / 2, cos(2 * pi * t) + rnorm(M, sd = 0.4)))
#' X  <- rbind(X0, X1)
#' y  <- rep(0:1, each = n / 2)
#' idx_tr <- sample(n, 60)
#' fit  <- krnn_fit(X[idx_tr, ], y[idx_tr], k = 5)
#' pred <- predict(fit, newdata = X[-idx_tr, ], y_true = y[-idx_tr])
#' pred$misclass
#'
#' @export
predict.krnn_fit <- function(object, newdata, y_true = NULL, ...) {
  sdp_new   <- .project_signed_depth(newdata, object$sd_train)
  predicted <- .krnn_predict_sdp(sdp_new, object$sd_train$sdp,
                                 object$y_train, object$k)
  prob      <- .smoother_prob(sdp_new, object$sd_train$sdp,
                              object$y_train, object$k)
  out <- list(predicted = predicted, prob = prob)
  if (!is.null(y_true))
    out$misclass <- mean(predicted != as.integer(y_true))
  out
}


# -----------------------------------------------------------------------
#  krnn_cv
# -----------------------------------------------------------------------

#' Cross-Validation for k-RNN Neighbourhood Size
#'
#' Selects the optimal half-neighbourhood size \eqn{k} by \eqn{R}-fold
#' cross-validation, using either the minimum-error rule or the
#' one-standard-error (1-SE) rule (Section 3.3.1 of Perez Ruiz, 2020).
#'
#' @param X Numeric matrix (\eqn{N \times M}).
#' @param y Integer vector of class labels (0/1).
#' @param k_max Maximum value of k to evaluate. Default
#'   \code{floor(nrow(X) / 4)}.
#' @param R Number of CV folds. Default \code{10}.
#' @param rule Character: \code{"min"} (default) or \code{"1se"}.
#' @param grid Numeric vector of length \eqn{M}.
#' @param type Depth variant.
#' @param seed Optional integer seed for reproducibility.
#'
#' @return An object of class \code{"krnn_cv"} with components:
#'   \describe{
#'     \item{k_opt}{Selected optimal k.}
#'     \item{cv_error}{Mean CV misclassification error for each k.}
#'     \item{cv_se}{Standard error of CV error per k.}
#'     \item{k_max}{Maximum k evaluated.}
#'     \item{R}{Number of folds used.}
#'     \item{rule}{The selection rule used.}
#'   }
#'
#' @examples
#' \donttest{
#' set.seed(3)
#' M <- 80
#' t  <- seq(0, 1, length.out = M)
#' X0 <- t(replicate(50, sin(2 * pi * t) + rnorm(M, sd = 0.4)))
#' X1 <- t(replicate(50, cos(2 * pi * t) + rnorm(M, sd = 0.4)))
#' X  <- rbind(X0, X1)
#' y  <- rep(0:1, each = 50)
#' cv <- krnn_cv(X, y, k_max = 20, R = 5, seed = 1)
#' print(cv)
#' plot_krnn_cv(cv)
#' }
#'
#' @export
krnn_cv <- function(X, y,
                    k_max = NULL,
                    R     = 10L,
                    rule  = c("min", "1se"),
                    grid  = NULL,
                    type  = c("tukey", "simplicial"),
                    seed  = NULL) {
  rule <- match.arg(rule)
  type <- match.arg(type)
  y    <- as.integer(y)
  N    <- nrow(X); M <- ncol(X)
  if (is.null(k_max)) k_max <- max(1L, floor(N / 4L))
  k_max <- as.integer(k_max); R <- as.integer(R)
  if (is.null(grid)) grid <- seq(0, 1, length.out = M)
  if (!is.null(seed)) set.seed(seed)
  folds   <- sample(rep(seq_len(R), length.out = N))
  err_mat <- matrix(NA_real_, nrow = R, ncol = k_max)
  for (r in seq_len(R)) {
    idx_te      <- which(folds == r)
    idx_tr      <- which(folds != r)
    sd_tr       <- signed_depth(X[idx_tr, , drop = FALSE],
                                grid = grid, type = type)
    sdp_te_proj <- .project_signed_depth(X[idx_te, , drop = FALSE], sd_tr)
    sdp_tr      <- sd_tr$sdp
    y_tr        <- y[idx_tr]; y_te <- y[idx_te]
    for (k in seq_len(k_max)) {
      pred          <- .krnn_predict_sdp(sdp_te_proj, sdp_tr, y_tr, k)
      err_mat[r, k] <- mean(pred != y_te)
    }
  }
  cv_error <- colMeans(err_mat)
  cv_se    <- apply(err_mat, 2, stats::sd) / sqrt(R)
  k_min    <- which.min(cv_error)
  k_opt    <- if (rule == "min") k_min else
    which(cv_error <= cv_error[k_min] + cv_se[k_min])[1L]
  structure(
    list(k_opt = k_opt, cv_error = cv_error, cv_se = cv_se,
         k_max = k_max, R = R, rule = rule),
    class = "krnn_cv"
  )
}

#' @param x A \code{"krnn_cv"} object.
#' @param ... Ignored.
#' @return Invisibly returns \code{x}.
#' @rdname krnn_cv
#' @export
print.krnn_cv <- function(x, ...) {
  cat("k-RNN Cross-Validation\n")
  cat("  Folds (R)      :", x$R,     "\n")
  cat("  Rule           :", x$rule,  "\n")
  cat("  Optimal k      :", x$k_opt, "\n")
  cat("  CV error at k* :", round(x$cv_error[x$k_opt], 4), "\n")
  invisible(x)
}


# -----------------------------------------------------------------------
#  plot_krnn_cv
# -----------------------------------------------------------------------

#' Plot a krnn_cv Object
#'
#' Plots the cross-validation error curve against the neighbourhood size k,
#' with standard-error bars and a vertical line at the selected optimum.
#'
#' @param x A \code{"krnn_cv"} object returned by \code{\link{krnn_cv}}.
#' @param ... Additional graphical parameters passed to
#'   \code{\link[graphics]{plot.default}}.
#'
#' @return Invisibly returns \code{x}.
#'
#' @examples
#' \donttest{
#' set.seed(3)
#' M <- 80
#' t  <- seq(0, 1, length.out = M)
#' X0 <- t(replicate(50, sin(2 * pi * t) + rnorm(M, sd = 0.4)))
#' X1 <- t(replicate(50, cos(2 * pi * t) + rnorm(M, sd = 0.4)))
#' X  <- rbind(X0, X1)
#' y  <- rep(0:1, each = 50)
#' cv <- krnn_cv(X, y, k_max = 20, R = 5, seed = 1)
#' plot_krnn_cv(cv)
#' }
#'
#' @export
plot_krnn_cv <- function(x, ...) {
  k_seq <- seq_len(x$k_max)
  graphics::plot(k_seq, x$cv_error,
                 type = "b", pch = 19, cex = 0.7,
                 xlab = "Number of neighbours (k)",
                 ylab = "CV misclassification error",
                 main = paste0("k-RNN cross-validation  (", x$R, "-fold)"),
                 ylim = range(c(x$cv_error - x$cv_se,
                                x$cv_error + x$cv_se)),
                 ...)
  graphics::segments(k_seq,
                     x$cv_error - x$cv_se,
                     k_seq,
                     x$cv_error + x$cv_se,
                     col = "grey60")
  graphics::abline(v = x$k_opt, lty = 2, col = "firebrick")
  graphics::text(x$k_opt,
                 max(x$cv_error + x$cv_se),
                 paste0("k* = ", x$k_opt),
                 adj = c(-0.1, 1), col = "firebrick", cex = 0.9)
  invisible(x)
}


# -----------------------------------------------------------------------
#  krnn_smoother
# -----------------------------------------------------------------------

#' k-RNN Moving-Average Smoother
#'
#' Estimates the conditional probability
#' \eqn{\hat{P}(y = 0 \mid \mathrm{sdp})} using the running-mean smoother
#' (moving average) with span \eqn{2k}. This is the regression interpretation
#' of the k-RNN (Section 3.4 of Perez Ruiz, 2020).
#'
#' @param X Numeric matrix (\eqn{N \times M}).
#' @param y Integer vector of labels (0/1).
#' @param k Half-neighbourhood size.
#' @param grid_eval Optional numeric vector of signed-depth values at which to
#'   evaluate the smoother. Defaults to the training signed depths.
#' @param boot Logical; if \code{TRUE} compute bootstrap-t confidence
#'   intervals. Default \code{FALSE}.
#' @param B Number of bootstrap replicates (only used when
#'   \code{boot = TRUE}). Default \code{200}.
#' @param alpha Nominal level: coverage is \eqn{1 - \alpha}. Default
#'   \code{0.05}.
#' @param grid_fd Numeric vector of length \eqn{M} (time grid).
#' @param type Depth variant.
#'
#' @return An object of class \code{"krnn_smoother"} with components:
#'   \describe{
#'     \item{sdp_eval}{Evaluation points (signed depths).}
#'     \item{prob}{Estimated conditional probabilities.}
#'     \item{ci_lower, ci_upper}{Bootstrap CI bounds (if
#'       \code{boot = TRUE}).}
#'     \item{k}{Half-neighbourhood size used.}
#'   }
#'
#' @examples
#' set.seed(5)
#' M <- 100; N <- 80
#' t  <- seq(0, 1, length.out = M)
#' X0 <- t(replicate(40, sin(2 * pi * t) + rnorm(M, sd = 0.4)))
#' X1 <- t(replicate(40, cos(2 * pi * t) + rnorm(M, sd = 0.4)))
#' X  <- rbind(X0, X1)
#' y  <- rep(0:1, each = 40)
#' sm <- krnn_smoother(X, y, k = 10)
#' plot_krnn_smoother(sm)
#'
#' @export
krnn_smoother <- function(X, y,
                          k         = 10L,
                          grid_eval = NULL,
                          boot      = FALSE,
                          B         = 200L,
                          alpha     = 0.05,
                          grid_fd   = NULL,
                          type      = c("tukey", "simplicial")) {
  type <- match.arg(type)
  y    <- as.integer(y)
  M    <- ncol(X)
  if (is.null(grid_fd)) grid_fd <- seq(0, 1, length.out = M)
  sd_obj <- signed_depth(X, grid = grid_fd, type = type)
  sdp    <- sd_obj$sdp
  o      <- order(sdp)
  sdp_s  <- sdp[o]; y_s <- y[o]
  if (is.null(grid_eval)) grid_eval <- sdp_s
  prob <- .smoother_prob(grid_eval, sdp_s, y_s, k, pre_sorted = TRUE)
  out  <- list(sdp_eval = grid_eval, prob = prob, k = k,
               sd_obj = sd_obj, y = y)
  if (boot) {
    ci           <- .bootstrap_ci(sdp_s, y_s, k, B = B, alpha = alpha)
    out$ci_lower <- ci$lower
    out$ci_upper <- ci$upper
  }
  structure(out, class = "krnn_smoother")
}


# -----------------------------------------------------------------------
#  plot_krnn_smoother
# -----------------------------------------------------------------------

#' Plot a krnn_smoother Object
#'
#' Plots the estimated conditional probability
#' \eqn{\hat{P}(y = 0 \mid \mathrm{sdp})} against the signed depth, with
#' optional bootstrap confidence bands.
#'
#' @param x A \code{"krnn_smoother"} object returned by
#'   \code{\link{krnn_smoother}}.
#' @param ... Additional graphical parameters passed to
#'   \code{\link[graphics]{plot.default}}.
#'
#' @return Invisibly returns \code{x}.
#'
#' @examples
#' set.seed(5)
#' M <- 100; N <- 80
#' t  <- seq(0, 1, length.out = M)
#' X0 <- t(replicate(40, sin(2 * pi * t) + rnorm(M, sd = 0.4)))
#' X1 <- t(replicate(40, cos(2 * pi * t) + rnorm(M, sd = 0.4)))
#' X  <- rbind(X0, X1)
#' y  <- rep(0:1, each = 40)
#' sm <- krnn_smoother(X, y, k = 10)
#' plot_krnn_smoother(sm)
#'
#' @export
plot_krnn_smoother <- function(x, ...) {
  graphics::plot(x$sdp_eval, x$prob,
                 type = "l", lwd = 2,
                 xlab = "Signed depth",
                 ylab = expression(hat(P)(y == 0 ~ "|" ~ sdp)),
                 main = paste0("k-RNN smoother  (k = ", x$k, ")"),
                 ylim = c(0, 1),
                 ...)
  graphics::abline(h = 0.5, lty = 2, col = "grey50")
  if (!is.null(x$ci_lower)) {
    graphics::lines(x$sdp_eval, x$ci_lower, lty = 3, col = "steelblue")
    graphics::lines(x$sdp_eval, x$ci_upper, lty = 3, col = "steelblue")
  }
  invisible(x)
}
