####################################
# DATASET 2: Motorcycle data
####################################
rm(list = ls())

library(MASS)
library(tree)

data(mcyle)
?mcycle

str(mcycle)

# OLS polynomial regression
ols <- lm(accel ~ poly(times,12), data = mcycle)
yhat <- predict(ols, newdata = mcycle)
plot(mcycle)
points(mcycle$times,yhat,col=2)

tree <- tree(accel ~ ., data = mcycle)
summary(tree)
par(mfrow=c(1,2))
plot(tree, type = "uniform")
text(tree, cex = 0.75)
yhat.tree <- predict(tree, newdata = mcycle)
plot(mcycle)
points(mcycle$times,yhat.tree,col=3,pch=16)

pruned <- prune.tree(tree, best = 6)
summary(pruned)
yhat.tree1 <- predict(pruned, newdata = mcycle)

par(mfrow=c(1,2))
plot(tree, type = "uniform")
text(tree, cex = 0.75)
plot(pruned, type = "uniform")
text(pruned, cex = 0.75)

par(mfrow=c(1,1))
plot(mcycle)
points(mcycle$times,yhat,col=2,pch=16)
points(mcycle$times,yhat.tree,col=3,pch=16)
points(mcycle$times,yhat.tree1,col=4)
legend("topleft",legend=c("Data","OLS poly reg","CART","Pruned CART"),col=1:4,pch=16)

sqrt(mean((yhat-mcycle$accel)^2))
sqrt(mean((yhat.tree-mcycle$accel)^2))
sqrt(mean((yhat.tree1-mcycle$accel)^2))