# THE MNIST DATABASE of handwritten digits 
# Identification using http://yann.lecun.com/exdb/mnist/
print "Download MNIST DATABASE of handwritten digits"
from jhplot import *
print Web.get("http://datamelt.org/examples/data/mnist_bin.zip")
print IO.unzip("mnist_bin.zip")

NMax=10 # Total runs. Reduce this number to get results faster
from org.ea.javacnn.data import DataBlock,OutputDefinition,TrainResult
from org.ea.javacnn.layers import DropoutLayer,FullyConnectedLayer,InputLayer,LocalResponseNormalizationLayer
from org.ea.javacnn.layers import ConvolutionLayer,RectifiedLinearUnitsLayer,PoolingLayer
from org.ea.javacnn.losslayers import SoftMaxLayer
from org.ea.javacnn.readers import ImageReader,MnistReader,PGMReader,Reader
from org.ea.javacnn.trainers import AdaGradTrainer
from org.ea.javacnn import JavaCNN
from java.util import ArrayList,Arrays
from java.lang import System



c1 = SPlot()
c1.visible()
c1.setAutoRange()
c1.setMarksStyle('various')
c1.setConnected(1, 0)
c1.setNameX('Epoch')
c1.setNameY('Loss') 


layers = ArrayList(); de = OutputDefinition() 
print "Total number of runs=", NMax 
print "Reading train sample.."
mr = MnistReader("mnist_bin/train-labels-idx1-ubyte", "mnist_bin/train-images-idx3-ubyte");
print "Total number of trainning images=",mr.size()," Nr of types=",mr.numOfClasses()
print "Read test sample .."
mrTest = MnistReader("mnist_bin/t10k-labels-idx1-ubyte", "mnist_bin/t10k-images-idx3-ubyte");
 
print "Total number of test images=",mrTest.size()," Nr of types=",mrTest.numOfClasses()
modelName = "model.ser" # save NN to this file  

layers.add(InputLayer(de, mr.getSizeX(), mr.getSizeY(), 1))
layers.add(ConvolutionLayer(de, 5, 32, 1, 2)) # uses different filters 
layers.add(RectifiedLinearUnitsLayer())       # applies the non-saturating activation function 
layers.add(PoolingLayer(de, 2,2, 0))          # creats a smaller zoomed out version
layers.add(ConvolutionLayer(de, 5, 64, 1, 2))
layers.add(RectifiedLinearUnitsLayer())
layers.add(PoolingLayer(de, 2,2, 0))
layers.add(FullyConnectedLayer(de, 1024))
layers.add(LocalResponseNormalizationLayer())
layers.add(DropoutLayer(de))
layers.add(FullyConnectedLayer(de, mr.numOfClasses()))
layers.add(SoftMaxLayer(de))

print "Training.."
net = JavaCNN(layers)
trainer = AdaGradTrainer(net, 20, 0.001) # 20 batch size, 0.001 weight decay 

from jarray import zeros
numberDistribution,correctPredictions = zeros(10, "i"),zeros(10, "i") 

tot=0
start = System.currentTimeMillis()
db = DataBlock(mr.getSizeX(), mr.getSizeY(), 1, 0)
for j in range(NMax):
  loss = 0
  for i in range(mr.size()):
    db.addImageData(mr.readNextImage(), mr.getMaxvalue())
    tr = trainer.train(db, mr.readNextLabel())
    loss = loss + tr.getLoss()
    tot +=1
    if (i != 0 and i % 100 == 0):
       print "Nr of images: ",i," Loss: ",(loss/float(i))
       c1.addPoint(0,tot,(loss/float(i)),1)
       c1.update()
  print "Loss: ", (loss / float(mr.size())), " for run=",j 
  mr.reset()
  Arrays.fill(correctPredictions, 0)
  Arrays.fill(numberDistribution, 0)
  for i in range(mrTest.size()):
            db.addImageData(mrTest.readNextImage(), mr.getMaxvalue())
            net.forward(db, False)
            correct = mrTest.readNextLabel()
            prediction = net.getPrediction()
            if(correct == prediction): correctPredictions[correct] +=1 
            numberDistribution[correct] +=1
  mrTest.reset()
  print " -> Testing time: ",int(0.001*(System.currentTimeMillis() - start))," s"
  print " -> Current run:",j
  print net.getPredictions(correctPredictions, numberDistribution, mrTest.size(), mrTest.numOfClasses())
  print " -> Save current state to ",modelName
  net.saveModel(modelName)

print "Read trained network from ",modelName," and make the final test"
cnn =net.loadModel(modelName)
Arrays.fill(correctPredictions, 0)
Arrays.fill(numberDistribution, 0)
for i in range(mrTest.size()):
            db.addImageData(mrTest.readNextImage(), mr.getMaxvalue())
            net.forward(db, False)
            correct = mrTest.readNextLabel()
            prediction = net.getPrediction()
            if(correct == prediction): correctPredictions[correct] +=1
            numberDistribution[correct] +=1
print "Final test:"
print net.getPredictions(correctPredictions, numberDistribution, mrTest.size(), mrTest.numOfClasses())