#This file holds example model code for a feedforward neural
#network. You may want to do better weight initialization since the
#current weight initialization may not be suitable for all network
#architectures.

import numpy as num
import gnumpy as gnp
import itertools



#the interface to the model:
#the model must know how to get gradEtc (gradient, gradient squared, loss)
#the model must also be able to compute Gv
# as well as having a method to get the fprop state used for Gv (getFpropState)
# and finally
#the model must compute loss by itself
#####@TODO: rethink this part of the interface
#model.getParamVect() must return a copy of the model parameters as a single packed monolithic vector
#model.setParamVect(W) must set the model parameters to be equal to
#  the vector W (not a reference to W) which may involve unpacking W



def initWeightMatrix(shape, maxNonZeroPerColumn = None):
    smallVal = 0
    fanOut = shape[1] if maxNonZeroPerColumn==None else min(maxNonZeroPerColumn, shape[1])
    W = num.random.randn(*shape)
    for j in range(shape[1]):
        perm = num.random.permutation(shape[0])
        W[perm[fanOut:],j] *= smallVal
    return W

class FeedforwardNeuralNet(object):
    def __init__(self, layerSizes, hidActFuncts, outputActFunct, L2Cost = 0):
        """
        layerSizes should be a list of positive integers with
        layerSizes[0] being the input dimensionality and
        layerSizes[-1] being the output dimensionality.
        """
        assert(len(layerSizes)>=2)
        assert(len(hidActFuncts)+2 == len(layerSizes))
        #Note that we do L2 regularization on ALL parameters, including all biases! This is also what Ilya's code does.
        #It might not always be advisable.
        self.L2Cost = L2Cost 
        self.outputActFunct = outputActFunct
        self.hidActFuncts = hidActFuncts
        shapes = [(layerSizes[i-1],layerSizes[i]) for i in range(1, len(layerSizes))]
        
        #@TODO: make the weight initialization more flexible later?
        maxNonZeroPerColumn = 15
        self.layerWeights = [gnp.garray(initWeightMatrix(s, maxNonZeroPerColumn)) for s in shapes]
        self.layerBiases = [gnp.garray(0*num.random.randn(1, layerSizes[i])) for i in range(1, len(layerSizes))]
        
        #we cache the fprop state in this variable so we don't have to
        #do a new fprop constantly if Gv is being called repeatedly
        self.fpropState = None
        
    #these methods are the model interface
    def getParamVect(self):
        """
        return a copy of the model parameters as a single packed monolithic vector
        """
        allWeights = itertools.chain(self.layerWeights, self.layerBiases)
        return gnp.concatenate([w.ravel() for w in allWeights])
    
    def setParamVect(self, vect):
        allWeights = itertools.chain(self.layerWeights, self.layerBiases)
        offset = 0
        for w in allWeights:
            #Tijmen says both the w[:]= and w[:,:]= syntax work as in numpy
            #w[:] = vect[offset:offset+w.size].reshape(w.shape) 
            w[:,:] = vect[offset:offset+w.size].reshape(w.shape)
            offset += w.size
    
    def gradEtc(self, (inputs, targets), computeGradSquared = False):
        """
        Return gradient, training loss, and optionally the gradient
        squared. We expect inputs and targets to be CPU arrays.
        """
        state = self.fprop(gnp.garray(inputs))
        #technically we don't need to do this here, but if we don't self.outputActFunct.error will compute it anyway
        outputActs = self.outputActFunct.activation(state[-1])
        targs = gnp.garray(targets)
        outputErrSignal = self.outputActFunct.dErrordNetInput(targs, state[-1], outputActs)
        #@TODO: It could be considered a bug not to include the L2 cost term in error below
        error = self.outputActFunct.error(targs, state[-1], outputActs)
        WGrads, biasGrads, WGradsSq, biasGradsSq = self.bprop(state, outputErrSignal, computeGradSquared)
        grad = gnp.concatenate([w.ravel() for w in itertools.chain(WGrads, biasGrads)])
        if self.L2Cost > 0:
            grad += self.L2Cost*inputs.shape[0]*self.getParamVect()
        #grad squared doesn't include the L2 cost because that will get added in later by HF
        #we do things this way to match Ilya's old code, probably #@TODO: it should be fixed later
        gradSquared = gnp.concatenate([w.ravel() for w in itertools.chain(WGradsSq, biasGradsSq)]) \
                      if computeGradSquared else None
        return grad, gradSquared, error
        
        
    def loss(self, (inputs, targets)):
        state = self.fprop(gnp.garray(inputs))
        error = self.outputActFunct.error(gnp.garray(targets), state[-1])
        sumSqWeights = sum((w*w).sum() for w in itertools.chain(self.layerWeights, self.layerBiases))
        return error + inputs.shape[0]*sumSqWeights*0.5*self.L2Cost
    
    def getFpropState(self, (inputs, unusedTargets)):
        """
        inputs should be a numpy array 
        """
        return self.fprop(gnp.garray(inputs))
        
    
    def Gv(self, (inputs, unusedTargets), vect, fpropState = None):
        """
        Multiply vect by the Gauss-Newton matrix defined with our
        current weights and the inputs parameter. We expect vect to be
        the same 'shape' as our weights and packed into a single
        monlithic vector. We expect vect to be a gpu array and inputs
        to be on the CPU.
        """
        if fpropState == None:
            fpropState = self.fprop(gnp.garray(inputs))
        rfpState = self.rfprop(fpropState, vect)
        outputActs = self.outputActFunct.activation(fpropState[-1])
        HJv = self.outputActFunct.HProd(rfpState[-1], outputActs)
        WParts, biasParts, unused1, unused2 = self.bprop(fpropState, HJv)
        GvNoL2 = gnp.concatenate([w.ravel() for w in itertools.chain(WParts, biasParts)])
        return GvNoL2 + vect*self.L2Cost*inputs.shape[0]
    #model interface ends
    
    def unpack(self, vect):
        """
        Interprets vect as a packed set of matrices that mirror
        self.layerWeights and self.layerBiases. Returns two lists of
        matrices mirroring self.layerWeights and self.layerBiases.
        """
        layerWeightsMirror = []
        layerBiasesMirror = []
        offset = 0
        for W in self.layerWeights:
            layerWeightsMirror.append(vect[offset:offset+W.size].reshape(W.shape))
            offset += W.size
        for b in self.layerBiases:
            layerBiasesMirror.append(vect[offset:offset+b.size].reshape(b.shape))
            offset += b.size
        return layerWeightsMirror, layerBiasesMirror
        
    def fprop(self, inputBatch):
        """
        Do a forward pass through the network and find the total input
        incident on the output units, but stop right before the output
        activation function.
        """
        assert(len(self.hidActFuncts) + 1 == len(self.layerWeights))
        state = [inputBatch]
        for i in range(len(self.hidActFuncts)):
            curActs = self.hidActFuncts[i].activation(gnp.dot(state[-1], self.layerWeights[i]) + self.layerBiases[i])
            state.append(curActs)
        #we don't apply the output nonlinearity, but we do everything except that
        state.append(gnp.dot(state[-1], self.layerWeights[-1]) + self.layerBiases[-1])
        return state
    
    def rfprop(self, fpropState, vect):
        """
        Perform an f1 pass (notation from Schraudolph) with vector
        vect which is the shape of the model weights packed into a
        single vector. We use intermediate results from fprop.
        """
        vectLayersW, vectLayersBias = self.unpack(vect)
        rfpState = [0*fpropState[0]]
        for i in range(len(self.hidActFuncts)):
            curRActs = self.hidActFuncts[i].dEdNetInput(fpropState[i+1]) * \
                       (gnp.dot(fpropState[i], vectLayersW[i]) + vectLayersBias[i] + \
                        gnp.dot(rfpState[-1], self.layerWeights[i]))
            rfpState.append(curRActs)
        rfpState.append(gnp.dot(fpropState[-2], vectLayersW[-1]) + vectLayersBias[-1] + \
                        gnp.dot(rfpState[-1], self.layerWeights[-1]))
        
        #unlike Ilya's code we return the whole list, although we only really need rfpState[-1]
        return rfpState 
    
    def predictions(self, inputBatch):
        """
        Do a full forward pass, no shortcuts. Just return the output
        activations.
        """
        state = self.fprop(inputBatch)
        return self.outputActFunct.activation(state[-1])
    
    def bprop(self, fpropState, outputErrSignal, computeGradSquared = False):
        """
        Perform a backward pass through the network. fpropState is the
        return value of fprop and outputErrSignal should be
        self.outputActFunct.dErrordNetInput(...).
        """
        #if len(self.hidActFuncts) == h then
        #len(errSignals)==len(self.layerWeights)==len(self.layerBiases)==h+1 and
        # len(fpropState) == h+2
        #fpropState[-2] is the activation of the penultimate hidden layer (or the input if there are no hidden layers)
        WGrads = [None for i in range(len(self.layerWeights))]
        biasGrads = [None for i in range(len(self.layerBiases))]
        errSignals = [None for i in range(len(self.layerWeights))]
        
        errSignals[-1] = outputErrSignal
        WGrads[-1] = gnp.dot(fpropState[-2].T, outputErrSignal)
        biasGrads[-1] = outputErrSignal.sum(axis=0)
        
        for i in reversed(range(len(self.hidActFuncts))):
            errSignals[i] = gnp.dot(errSignals[i+1], self.layerWeights[i+1].T)*self.hidActFuncts[i].dEdNetInput(fpropState[i+1])
            WGrads[i] = gnp.dot(fpropState[i].T, errSignals[i])
            biasGrads[i] = errSignals[i].sum(axis=0)
            
        if computeGradSquared:
            WGradsSq = []
            biasGradsSq = []
            for i in range(len(errSignals)):
                errSigSq = errSignals[i]**2
                WGradsSq.append(gnp.dot((fpropState[i]*fpropState[i]).T, errSigSq))
                biasGradsSq.append(errSigSq.sum(axis=0))
            return WGrads, biasGrads, WGradsSq, biasGradsSq
        
        return WGrads, biasGrads, None, None

    

    



