library(stats)
library(tree)
library(rpart)
library(rpart.plot)
library(mgcv)
library(lattice)

ozone <- read.table("Downloads/ozone.data", header = T)

cuberootozone <- ozone$ozone^(1/3)
radiation <- ozone$radiation
temp <- ozone$temperature
wind <- ozone$wind

total_dat <- cbind(cuberootozone, radiation, temp, wind)
pairs(total_dat)

alpha_hat <- mean(cuberootozone)

rad_spline <- smooth.spline(radiation, cuberootozone - alpha_hat, all.knots = T)
temp_spline <- smooth.spline(temp, cuberootozone - alpha_hat, all.knots = T)
wind_spline <- smooth.spline(wind, cuberootozone - alpha_hat, all.knots = T)
  
rad_out <- predict(rad_spline, x = radiation)$y - mean(predict(rad_spline, x = radiation)$y)
temp_out <- predict(temp_spline, x = temp)$y - mean(predict(temp_spline, x = temp)$y)
wind_out <- predict(wind_spline, x = wind)$y - mean(predict(wind_spline, x = wind)$y)

change <- 1
eps <- 1e-5

while(change > eps){
  rad_spline <- smooth.spline(radiation, cuberootozone - 
                                alpha_hat - temp_out - wind_out, all.knots = T)
  rad_out_old <- rad_out
  rad_out <- predict(rad_spline, x = radiation)$y - mean(predict(rad_spline, x = radiation)$y)
  
  wind_spline <- smooth.spline(wind, cuberootozone - 
                                 alpha_hat - temp_out - rad_out, all.knots = T)
  wind_out_old <- wind_out
  wind_out <- predict(wind_spline, x = wind)$y - mean(predict(wind_spline, x = wind)$y)
  
  temp_spline <- smooth.spline(temp, cuberootozone - 
                                 alpha_hat - rad_out - wind_out, all.knots = T)
  temp_out_old <- temp_out
  temp_out <- predict(temp_spline, x = temp)$y - mean(predict(temp_spline, x = temp)$y)
  
  change <- max(abs(rad_out - rad_out_old), abs(temp_out - temp_out_old),
                abs(wind_out - wind_out_old))
  print(change)
}

par(mfrow=c(1,3))
plot(rad_spline$x, rad_spline$y, ylab='Fitted spline', xlab='radiation', type='l')
plot(temp_spline$x, temp_spline$y, ylab="Fitted spline", xlab="temperature", type='l')
plot(wind_spline$x, wind_spline$y, ylab="Fitted spline", xlab="wind", type='l')

# We could set nknots=5 for the temperature spline to get a smoother plot

# Trellis display for radiation
xyplot(cuberootozone ~ radiation | equal.count(temp, number=4) * equal.count(wind, number=4),
       panel = function(x, y, ...) {
         panel.xyplot(x, y, ...)
         fit <- gam(y ~ s(x))
         xs <- seq(min(x), max(x), length = 100)
         panel.lines(xs, predict(fit, newdata = data.frame(x = xs)),
                     col = "blue", lwd = 2)})

xyplot(cuberootozone ~ radiation |
         equal.count(temp, 4) * equal.count(wind, 4),
       
       data = ozone,
       
       panel = function(x, y, subscripts, ...) {
         panel.xyplot(x, y, ...)
         
         panel.lines(x, ozone$fit[subscripts],
                     col = "blue", lwd = 2)
       })

# Compute squared prediction error on data
mean((cuberootozone - alpha_hat - rad_out - temp_out - wind_out)^2)

# Now, we go to a tree approach
tree_fit <- tree(cuberootozone ~ radiation + temp + wind)

cv_tree <- cv.tree(tree_fit)
cv_tree

plot(cv_tree$size, cv_tree$dev, type='l', xlab="size", ylab="test error")

optimal_size <- cv_tree$size[which.min(cv_tree$dev)]
optimal_size

mean((predict(tree_fit, newdata=data.frame(cbind(radiation, temp, wind))) 
      - cuberootozone)^2)

par(mfrow=c(1,1))
plot(tree_fit)
text(tree_fit, pretty = 0)