# Returns a similarity matrix K
# x should be a matrix
gauss_kernel <- function(x, h) {
  h_sq <- h^2
  n <- nrow(x)
  row_norm_sq <- rowSums(x^2)/2
  exp((x %*% t(x) - matrix(row_norm_sq, nrow=n, ncol=n) - matrix(row_norm_sq, byrow=TRUE, nrow=n, ncol=n))/(h_sq))
}

# x should be a vector with entries in [0,1]
First_Sobolev_kernel <- function(x) {
  outer(x, x, FUN = function(a, b) pmin(a, b))
}

Second_Sobolev_kernel <- function(x) {
  outer(x, x, FUN = function(a, b) {
    pmin(a, b)^3/3 - pmin(a, b)^2*(a+b)/2 + a*b*pmin(a,b)
  })
}

# A list to help with performing ridge regression in kernelised form
# scale_mat is D^2 / (lambda + D^2) - an n by n_lambda matrix
# lambda gives the vector of lambda values
ridge_mat <- function(k, nlambda=100) {
  out <- svd(k, nv=0)
  d_sq <- out$d
  n <- length(d_sq)
  lambda_vec <- exp(seq(from=log(50), to=log(0.000001), length.out=nlambda))
  scale_mat <- d_sq / (d_sq + matrix(lambda_vec, nrow=length(d_sq), ncol=length(lambda_vec), byrow=TRUE))
  return(list("u"=out$u,
              "t_u"=t(out$u),
              "scale_mat"=scale_mat,
              "lambda"=lambda_vec))
}

# returns a matrix of fitted values corresponding to different lambda values
fit_gauss_ridge <- function(x, y, sigma, nlambda=100) {
  K <- gauss_kernel(x, sigma)
  ridge_out <- ridge_mat(K, nlambda)
  return(list("fit"=ridge_out$u %*% (ridge_out$scale_mat * as.numeric(ridge_out$t_u %*% y)),
              "lambda"=ridge_out$lambda))
}

fit_First_Sobolev_ridge <- function(x, y, nlambda=100) {
  K <- First_Sobolev_kernel(x)
  ridge_out <- ridge_mat(K, nlambda)
  return(list("fit"=ridge_out$u %*% (ridge_out$scale_mat * as.numeric(ridge_out$t_u %*% y)),
              "lambda"=ridge_out$lambda))
}

fit_Second_Sobolev_ridge <- function(x, y, nlambda=100) {
  K <- Second_Sobolev_kernel(x)
  ridge_out <- ridge_mat(K, nlambda)
  return(list("fit"=ridge_out$u %*% (ridge_out$scale_mat * as.numeric(ridge_out$t_u %*% y)),
              "lambda"=ridge_out$lambda))
}

# plots the fit
plot_fit <- function(x, order_x, fit_out, lambda_ind, col="blue") {
  x <- as.numeric(x)
  lines(x[order_x], fit_out$fit[order_x, lambda_ind], col=col)
}


###############
# Generate data
###############
# Generate data
set.seed(101)
sigma <- 0.2
n <- 200
# Generate equally spaced points between [0,1]
x <- seq(from=0, to=1, length.out=n)
# Signal
sig <- sin(5*x)^2 / (2.5+(cos(2*x)))
y <- sig + sigma*rnorm(n)

# To help with plotting
order_x <- order(x)

# plot
plot(x, y)
lines(x[order_x], sig[order_x])

###############
# Gaussian
###############
fit_out <- fit_gauss_ridge(as.matrix(x), y, 0.1)
plot_fit(x, order_x, fit_out, 5, "blue") # can try different lambda indices

op <- par(mfrow=c(3, 3), mar=c(2, 2, 3, 1))

h_vec <- c(0.05, 0.5, 2)
lambda_vec <- c(10, 50, 90) # a vector of lambda indices

for (h in h_vec) {
  fit_out <- fit_gauss_ridge(as.matrix(x), y, h)
  for (lambda_ind in lambda_vec) {
    plot(x, y, main = paste("lambda index =", lambda_ind, " h =", h), xlab="",
         ylab="")
    lines(x[order_x], sig[order_x])
    plot_fit(x, order_x, fit_out, lambda_ind)
  }
}
par(op)

###############
# First_Sobolev
###############
fit_out <- fit_First_Sobolev_ridge(x, y)

lambda_vec <- c(10, 20, 35, 60) # a vector of lambda indices

op <- par(mfrow=c(2, 2), mar=c(2, 2, 3, 1))

for (lambda_ind in lambda_vec) {
  plot(x, y, main = paste("lambda index =", lambda_ind), xlab="", ylab="")
  lines(x[order_x], sig[order_x])
  plot_fit(x, order_x, fit_out, lambda_ind)
}

par(op)

###############
# Second_Sobolev
###############
fit_out <- fit_Second_Sobolev_ridge(x, y)

lambda_vec <- c(35, 60, 70, 90) # a vector of lambda indices

op <- par(mfrow=c(2, 2), mar=c(2, 2, 3, 1))

for (lambda_ind in lambda_vec) {
  plot(x, y, main = paste("lambda index =", lambda_ind), xlab="", ylab="")
  lines(x[order_x], sig[order_x])
  plot_fit(x, order_x, fit_out, lambda_ind)
}

par(op)