# -------------------------------------------------------------
# Mxnet in R - LeNet
# -------------------------------------------------------------

suppressMessages(library(mxnet))

# load train.csv
DstTrain <- read.csv('../input/train.csv', stringsAsFactors = FALSE, na.strings = c("NA", ""))
# load test.csv
DstTest  <- read.csv('../input/test.csv', stringsAsFactors = FALSE, na.strings = c("NA", ""))

# convert to matrix
DstTrain <- data.matrix(DstTrain)
DstTest <- data.matrix(DstTest)

# split x and y data
xDstTrain <- DstTrain[,-1]
yDstTrain <- DstTrain[,1]

# Normalization
xDstTrain <- t(xDstTrain/255)
DstTest <- t(DstTest/255)

# reshape the matrices into arrays
arrayDstTrain <- xDstTrain
dim(arrayDstTrain) <- c(28, 28, 1, ncol(xDstTrain))
arrayDstTest <- DstTest
dim(arrayDstTest) <- c(28, 28, 1, ncol(DstTest))

# network structure
# input
data <- mx.symbol.Variable('data')
# first conv
conv1 <- mx.symbol.Convolution(data=data, kernel=c(5,5), num_filter=20)
tanh1 <- mx.symbol.Activation(data=conv1, act_type="tanh")
pool1 <- mx.symbol.Pooling(data=tanh1, pool_type="max", kernel=c(2,2), stride=c(2,2))
# second conv
conv2 <- mx.symbol.Convolution(data=pool1, kernel=c(5,5), num_filter=50)
tanh2 <- mx.symbol.Activation(data=conv2, act_type="tanh")
pool2 <- mx.symbol.Pooling(data=tanh2, pool_type="max", kernel=c(2,2), stride=c(2,2))
# first fullc
flatten <- mx.symbol.Flatten(data=pool2)
fc1 <- mx.symbol.FullyConnected(data=flatten, num_hidden=500)
tanh3 <- mx.symbol.Activation(data=fc1, act_type="tanh")
# second fullc
fc2 <- mx.symbol.FullyConnected(data=tanh3, num_hidden=10)
# loss
lenet <- mx.symbol.SoftmaxOutput(data=fc2)


# cpu only in kaggle
device.cpu <- mx.cpu()


# Train on a CPU
mx.set.seed(0)
tic <- proc.time()
mxModel <- mx.model.FeedForward.create(
    lenet, 
    X=arrayDstTrain, 
    y=yDstTrain,
    ctx=device.cpu, 
    num.round=4, 
    array.batch.size=100,
    learning.rate=0.1, 
    momentum=0.9, 
    wd=0.00001,
    eval.metric=mx.metric.accuracy,
    epoch.end.callback=mx.callback.log.train.metric(100)
)

# print time
print(proc.time() - tic)

    
# Prediction and Submitting
Prediction <- predict(mxModel, arrayDstTest)

# get the result labels 
labelPrediction <- max.col(t(Prediction)) - 1
table(labelPrediction)

# create the submition file
result <- data.frame(ImageId=1:ncol(DstTest), Label=labelPrediction)
write.csv(result, file='result.csv', row.names=FALSE,  quote=FALSE)