-- BASIC MODEL
input_size = 3

hidden_size = 150
mlp = nn.Sequential()
mlp:add(nn.Linear(input_size,hidden_size))
mlp:add(nn.Sigmoid())
mlp:add(nn.Linear(hidden_size,1))

criterion = nn.MSECriterion()

trainer = nn.StochasticGradient(mlp,criterion)
trainer.maxIteration = 100

-- INITIALIZE DATA
params = mlp:getParameters()
params:uniform(-0.1,0.1)

trainer:train(dataset)

require 'torch'
require 'image'
require 'nn'


if not opt then
   print '==> processing options'
   cmd = torch.CmdLine()
   cmd:text()
   cmd:text('SVHN Dataset Preprocessing')
   cmd:text()
   cmd:text('Options:')
   cmd:option('-size', 'small', 'how many samples do we load: small | full | extra')
   cmd:option('-visualize', true, 'visualize input data and weights during training')
   cmd:text()
   opt = cmd:parse(arg or {})
end

-- This would download the dataset if its not already present
print '==> downloading dataset'
www = 'http://torch7.s3-website-us-east-1.amazonaws.com/data/housenumbers/'

train_file = 'train_32x32.t7'
test_file = 'test_32x32.t7'
extra_file = 'extra_32x32.t7'

if not paths.filep(train_file) then
   os.execute('wget ' .. www .. train_file)
end
if not paths.filep(test_file) then
   os.execute('wget ' .. www .. test_file)
end
if opt.size == 'extra' and not paths.filep(extra_file) then
   os.execute('wget ' .. www .. extra_file)   
end



-- training/test size by default small
if opt.size == 'extra' then
   print '==> using extra training data'
   trsize = 73257 + 531131
   tesize = 26032
elseif opt.size == 'full' then
   print '==> using regular, full training data'
   trsize = 73257
   tesize = 26032
elseif opt.size == 'small' then
   print '==> using reduced training data, for fast experiments'
   trsize = 10000
   tesize = 2000
end

-- load train and test data
loaded = torch.load(train_file,'ascii')
trainData = {
   data = loaded.X:transpose(3,4),
   labels = loaded.y[1],
   size = function() return trsize end -- this would be called trainData:size()
}

loaded = torch.load(test_file,'ascii')
testData = {
   data = loaded.X:transpose(3,4),
   labels = loaded.y[1],
   size = function() return tesize end
}

-- convert to float
trainData.data = trainData.data:float()
testData.data = testData.data:float()

-- preprocessing

-- Convert all images to YUV
print '==> preprocessing data: colorspace RGB -> YUV'
for i = 1,trainData:size() do
   trainData.data[i] = image.rgb2yuv(trainData.data[i])
end
for i = 1,testData:size() do
   testData.data[i] = image.rgb2yuv(testData.data[i])
end

-- Name channels for convenience
channels = {'y','u','v'}
print '==> preprocessing data: normalize each feature (channel) globally'
mean = {}
std = {}
for i,channel in ipairs(channels) do
   -- normalize each channel globally by subtracting data from mean and dividing by standard deviation
   mean[i] = trainData.data[{ {},i,{},{} }]:mean()
   std[i] = trainData.data[{ {},i,{},{} }]:std()
   trainData.data[{ {},i,{},{} }]:add(-mean[i])
   trainData.data[{ {},i,{},{} }]:div(std[i])
end

-- Normalize test data, using the training means/stds
for i,channel in ipairs(channels) do
   -- normalize each channel globally:
   testData.data[{ {},i,{},{} }]:add(-mean[i])
   testData.data[{ {},i,{},{} }]:div(std[i])
end


neighborhood = image.gaussian1D(13) -- declaring a gaussian with 13 elements.

-- Local normalization uniformizes the local mean and variance of an image
normalization = nn.SpatialContrastiveNormalization(1, neighborhood, 1):float()

-- Normalize all channels locally:
for c in ipairs(channels) do
   for i = 1,trainData:size() do
      trainData.data[{ i,{c},{},{} }] = normalization:forward(trainData.data[{ i,{c},{},{} }])
   end
   for i = 1,testData:size() do
      testData.data[{ i,{c},{},{} }] = normalization:forward(testData.data[{ i,{c},{},{} }])
   end
end