# The following code runs in the R programming language. It is installed on most machines in the School

# To run on Unix in the School:
# > setup UMS
# > R

# Before this code works you might need to execute from within R
#
# install.packages('scatterplot3d')
# install.packages('mnormt')


require('scatterplot3d')
require('mnormt')

# Means and covariance matrices for the red and blue cans

red.mean  = c(1, 0, 0)
blue.mean = c(0, 0, 1)

red.cov  = diag(.2,3)
blue.cov = diag(.2,3)


# Some samples from the 3D Gaussian

red.data  = rmnorm(n=100, mean=red.mean,  varcov=red.cov)
blue.data = rmnorm(100, blue.mean, blue.cov)


# Plot it

xlim = c(max(max(red.data[,1]), max(blue.data[,1])),
         min(min(red.data[,1]), min(blue.data[,1])))
ylim = c(max(max(red.data[,2]), max(blue.data[,2])),
         min(min(red.data[,2]), min(blue.data[,2])))
zlim = c(max(max(red.data[,3]), max(blue.data[,3])),
         min(min(red.data[,3]), min(blue.data[,3])))

graph = scatterplot3d(red.data, color='red', xlim=xlim, ylim=ylim, zlim=zlim)
graph$points3d(blue.data, col='blue')

# Save the graph to a PDF
pdf('colour-detector.pdf')
graph = scatterplot3d(red.data, color='red', xlim=xlim, ylim=ylim, zlim=zlim)
graph$points3d(blue.data, col='blue')
dev.off()


# Define training and test set

red.train = red.data[1:70,]
red.test  = red.data[71:100,]

blue.train = blue.data[1:70,]
blue.test  = blue.data[71:100,]


# Estimate parameters from training data

red.sample.mean  = colMeans(red.train)
red.sample.cov   = cov(red.train)

blue.sample.mean = colMeans(blue.train)
blue.sample.cov  = cov(blue.train)


# Test on test sets!

red.red  = dmnorm(red.test, red.sample.mean, red.sample.cov)
red.blue = dmnorm(red.test, blue.sample.mean, blue.sample.cov)

red.prediction = max.col(cbind(red.red, red.blue))

blue.red  = dmnorm(blue.test, red.sample.mean, red.sample.cov)
blue.blue = dmnorm(blue.test, blue.sample.mean, blue.sample.cov)

blue.prediction = max.col(cbind(blue.red, blue.blue))

error = sum(red.prediction == 2) + sum(blue.prediction == 1)