import numpy as num
from scipy.io import loadmat, savemat
from hfOpt import HF
from ffnnet import *
import activationFunctions as Act
import sys
import cPickle as pk
import time
from hfDataInterface import HFDataSource
from hfModelInterface import HFModel

dataPath = "./digs3pts_1.mat"

def sampleMinibatch(mbsz, inputs, targets):
    if mbsz == inputs.shape[0]:
        return inputs, targets
    assert(inputs.shape[0] == targets.shape[0])
    rows = num.random.random_integers(0, inputs.shape[0]-1, mbsz)
    return inputs[rows], targets[rows]

def calcTestLoss(net, testInputs, testTargets):
    testPreds = net.predictions(testInputs)
    diff = (testPreds - testTargets)
    testLoss = (diff*diff).sum()/testTargets.shape[0]
    return testLoss

def main():
    logFile = sys.stdout
    
    
    d = loadmat(dataPath, struct_as_record = True)
    trainInputs = d["bdata"]
    testInputs = d["bdatatest"]
    trainTargets = trainInputs[:,:]
    testTargets = testInputs[:,:]

    num.random.seed(130)
    perm = num.random.permutation(20000)
    trainInputs = trainInputs[perm]
    trainTargets = trainTargets[perm]
    
    initDamping = 15.
    
    optimizer = HF(logFile, initDamping)
    optimizer.doCGBacktracking = True
    
    layerSizes = [trainInputs.shape[1], 400, 200, 100, 50, 25, 6, 25, 50, 100, 200, 400, trainTargets.shape[1]]
    hidActFuncts = [Act.Sigmoid() for i in range(len(layerSizes)-2)]
    hidActFuncts[5] = Act.Linear()
    outputActFunct = Act.Sigmoid()
    L2Cost = 2e-5
    
    net = FeedforwardNeuralNet(layerSizes, hidActFuncts, outputActFunct, L2Cost)
    hfnet = HFModel(net)
    
    mbsz = 5000
    curvatureBatchSize = mbsz
    gradBatchSize = 4*mbsz
    
    iters = 500
    
    data = ((trainInputs[(i%4)*mbsz:(i%4+1)*mbsz,:], trainTargets[(i%4)*mbsz:(i%4+1)*mbsz,:]) for i in range(4))

    data = HFDataSource(data, mbsz, gradBatchSize = 20000, curvatureBatchSize = 5000, lineSearchBatchSize = 20000, \
                        gradBatchIsTrainingSet = True)
    
    hfIters = optimizer.optimize(hfnet, data, iters, otherPrecondDampingTerm=net.L2Cost)
    
    print "Initial test loss:", calcTestLoss(net, testInputs, testTargets)
    print "Starting HF"
    print >>logFile, "gradBatchSize=", gradBatchSize, "curvatureBatchSize=", curvatureBatchSize
    print >>logFile, "L2Cost=%.8f" % (L2Cost)
    t = time.time()
    for j,w in enumerate(hfIters):
        testPreds = net.predictions(testInputs)
        diff = (testPreds - testTargets)
        testLoss = (diff*diff).sum()/testTargets.shape[0] #should eventually reach 0.2?
        trainPreds = net.predictions(trainInputs)
        diff = trainPreds-trainTargets
        trainLoss = (diff*diff).sum()/trainTargets.shape[0]
        print >>logFile, "\n\nHF:", j+1, "train sq loss=", trainLoss, "test sq loss=", testLoss
        logFile.flush()
    print >>logFile, "total time elapsed %f hours" % ((time.time()-t)/3600.)
    logFile.close()

if __name__ == "__main__":
    main()
