##-----------------------------------------------------------------------------
## ST790 -- R codes for Lecture 10b
## Author: Ryan Martin (rgmarti3@ncsu.edu)                                      
##-----------------------------------------------------------------------------

# Conformal prediction -- response only

conformal.y <- function(y, ygrid) {
  
  n <- length(y)
  pl <- 0 * ygrid
  for(i in seq_along(ygrid)) {
    
    yy <- c(ygrid[i], y)
    med.yy <- median(yy)
    e <- abs(yy - med.yy)
    pl[i] <- mean(e >= e[1])
    
  }
  return(list(y=ygrid, pl=pl))
  
}

Y <- log(rgamma(200, shape=3, rate=2))
#hist(Y, col="gray", border="white", main="")
o <- conformal.y(Y, seq(min(Y), max(Y), len=100))
plot(o$y, o$pl, type="l", xlab=expression(tilde(y)), ylab=expression(pi[y](tilde(y))))
abline(h=0.05, lty=2)
abline(v=median(Y), lty=3)


# Conformal prediction -- response and covariate

library(splines)

conformal.xy <- function(x, y, xnew, df=7, alpha=0.05) {
  
  hi <- lo <- 0 * xnew
  n <- length(y)
  B <- bs(x, df)
  oo <- lm(y ~ B - 1)
  sig <- summary(oo)$sigma 
  ynew <- as.numeric(predict(B, xnew) %*% oo$coefficients)
  get.pl <- function(x0, y0) {
    
    xx <- c(x0, x)
    yy <- c(y0, y)
    o <- lm(yy ~ bs(xx, df) - 1)
    e <- abs(o$residuals)
    return(mean(e >= e[1]))
    
  }
  for(i in seq_along(xnew)) {
    
    f <- function(u) get.pl(xnew[i], u) - alpha
    lo[i] <- uniroot(f, c(ynew[i] - 5 * sig, ynew[i]))$root #, extendInt="downX")$root
    hi[i] <- uniroot(f, c(ynew[i], ynew[i] + 5 * sig))$root #, extendInt="upX")$root
    
  }
  return(cbind(xnew=xnew, ynew=ynew, lo=lo, hi=hi))
  
}

n <- 200
x <- runif(n) 
m <- function(x) sin(2 * pi * x**3)**3
y <- m(x) + 0.1 * rt(n, df=5)
xnew <- seq(0.01, 0.99, len=50)
out <- conformal.xy(x, y, xnew, df=7)
plot(x, y, ylim=range(c(y, out[,3:4])), type="n")
polygon(x=c(xnew, rev(xnew)), y=c(out[,4], rev(out[,3])), col="lightgray", border=NA)
lines(out[,1], out[,2])
points(x, y)
curve(m, add=TRUE, lwd=2)