""" 

Copyright (c) 2016, Jose Dolz .All rights reserved.

Redistribution and use in source and binary forms, with or without modification,

are permitted provided that the following conditions are met:

    1. Redistributions of source code must retain the above copyright notice,

       this list of conditions and the following disclaimer.

    2. Redistributions in binary form must reproduce the above copyright notice,

       this list of conditions and the following disclaimer in the documentation

       and/or other materials provided with the distribution.

    THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,

    EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES

    OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND

    NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT

    HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,

    WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING

    FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR

    OTHER DEALINGS IN THE SOFTWARE.

Jose Dolz. Dec, 2016.

email: jose.dolz.upv@gmail.com

LIVIA Department, ETS, Montreal.

"""

import pdb

import os

import numpy as np

import theano

import theano.tensor as T

import gzip

import cPickle

import sys

from os.path import isfile, join

# https://github.com/Theano/Theano/issues/689

sys.setrecursionlimit(50000)

# To set a learning rate at each layer

def extendLearningRateToParams(numberParamsPerLayer,learning_rate):

    if not isinstance(learning_rate, list):

        learnRates = np.ones(sum(numberParamsPerLayer), dtype = "float32") * learning_rate

    else:

        print("")

        learnRates = []

        for p_i in range(len(numberParamsPerLayer)) :

            for lr_i in range(numberParamsPerLayer[p_i]) :

                learnRates.append(learning_rate[p_i])

    return learnRates

# TODO: Check that length of learning rate (in config ini) actually corresponds to length of layers (CNNs + FCs + SoftMax)

def computeReceptiveField(kernelsCNN, stride) :

    # To-do. Verify receptive field with stride size other than 1

    if len(kernelsCNN) == 0:

        return 0

    # Check number of ConvLayers

    numberCNNLayers = []

    for l_i in range(1,len(kernelsCNN)):

        if len(kernelsCNN[l_i]) == 3:

            numberCNNLayers = l_i + 1

    kernelDim = len(kernelsCNN[0])

    receptiveField = [stride]*kernelDim

    for d_i in xrange(kernelDim) :

        for l_i in xrange(numberCNNLayers) :

            receptiveField[d_i] += kernelsCNN[l_i][d_i] - 1

    return receptiveField

###########################################################

######## Create bias and include them on feat maps ########

###########################################################

# TODO. Remove number of FeatMaps

def addBiasParametersOnFeatureMaps( bias, featMaps, numberOfFeatMaps ) :

    output = featMaps + bias.dimshuffle('x', 0, 'x', 'x', 'x')

    return (output)

###########################################################

########         Initialize CNN weights            ########

###########################################################

def initializeWeights(filter_shape, initializationMethodType, weights) :

    # filter_shape:[#FMs in this layer, #FMs in input, KernelDim_0, KernelDim_1, KernelDim_2]

    def Classic():

        print " --- Weights initialization type: Classic "

        rng = np.random.RandomState(24575)

        stdForInitialization = 0.01

        W = theano.shared(

            np.asarray(

                rng.normal(loc=0.0, scale=stdForInitialization, size=(filter_shape[0],filter_shape[1],filter_shape[2],filter_shape[3],filter_shape[4])),

                dtype='float32'#theano.config.floatX

                ),

            borrow=True

        )

        return W

    def Delving():

       # https://arxiv.org/pdf/1502.01852.pdf

       print " --- Weights initialization type: Delving "

       rng = np.random.RandomState(24575)

       stdForInitialization = np.sqrt( 2.0 / (filter_shape[1] * filter_shape[2] * filter_shape[3] * filter_shape[4]) ) #Delving Into rectifiers suggestion.

       W = theano.shared(

           np.asarray(

               rng.normal(loc=0.0, scale=stdForInitialization, size=(filter_shape[0],filter_shape[1],filter_shape[2],filter_shape[3],filter_shape[4])),

               dtype='float32'#theano.config.floatX

               ),

           borrow=True

       )

       return W

    # TODO: Add checks so that weights and kernel have the same shape

    def Load():

       print " --- Weights initialization type: Transfer learning... "

       W = theano.shared(

           np.asarray(

               weights,

               dtype=theano.config.floatX

               ),

           borrow=True

       )

       return W

    optionsInitWeightsType = {0 : Classic,

                              1 : Delving,

                              2 : Load}

    W = optionsInitWeightsType[initializationMethodType]()

    return W

def getCentralVoxels(sampleSize, receptiveField) :

    centralVoxels = []

    for d_i in xrange(0, len(sampleSize)) :

        centralVoxels.append(sampleSize[d_i] - receptiveField[d_i] + 1)

    return centralVoxels

def extractCenterFeatMaps(featMaps, featMaps_shape, centralVoxels) :

    centerValues = []

    minValues = []

    maxValues = []

    for i in xrange(3) :

        C_v = (featMaps_shape[i + 2] - 1) / 2

        min_v = C_v - (centralVoxels[i]-1)/2

        max_v = min_v + centralVoxels[i]

        centerValues.append(C_v)

        minValues.append(min_v)

        maxValues.append(max_v)

    return featMaps[:,

                    :,

                    minValues[0] : maxValues[0],

                    minValues[1] : maxValues[1],

                    minValues[2] : maxValues[2]]

###########################################

############# Save/Load models ############

###########################################

def load_model_from_gzip_file(modelFileName) :

    f = gzip.open(modelFileName, 'rb')

    model_obj = cPickle.load(f)

    f.close()

    return model_obj

def dump_model_to_gzip_file(model, modelFileName) :

    # First release GPU memory

    model.releaseGPUData()

    f = gzip.open(modelFileName, 'wb')

    cPickle.dump(model, f, protocol=cPickle.HIGHEST_PROTOCOL)

    f.close()

    return modelFileName

def makeFolder(folderName, display_Str) :

    if not os.path.exists(folderName) :

        os.makedirs(folderName)

    strToPrint = "..Folder " + display_Str + " created..."

    print strToPrint

    from os import listdir

""" Get a set of images from a folder given an array of indexes """

def getImagesSet(imagesFolder, imageIndexes) :

   imageNamesToGetWithFullPath = []

   imageNamesToGet = []

   if os.path.exists(imagesFolder):

       imageNames = [f for f in os.listdir(imagesFolder) if isfile(join(imagesFolder, f))]

       imageNames.sort()

       # Remove corrupted files (if any)

       if '.DS_Store' in imageNames: imageNames.remove('.DS_Store')

       imageNamesToGetWithFullPath = []

       imageNamesToGet = []

       if ( len(imageNames) > 0):  

           imageNamesToGetWithFullPath = [join(imagesFolder,imageNames[imageIndexes[i]]) for i in range(0,len(imageIndexes))]

           imageNamesToGet = [imageNames[imageIndexes[i]] for i in range(0,len(imageIndexes))]

   return (imageNamesToGetWithFullPath,imageNamesToGet)

"""" Get a set of weights from a folder given an array of indexes """

def getWeightsSet(weightsFolder, weightsIndexes) :

   weightNames = [f for f in os.listdir(weightsFolder) if isfile(join(weightsFolder, f))]

   weightNames.sort()

   # Remove corrupted files (if any)

   if '.DS_Store' in weightNames: weightNames.remove('.DS_Store')

   weightNamesToGetWithFullPath = [join(weightsFolder,weightNames[weightsIndexes[i]]) for i in range(0,len(weightsIndexes))]

   return (weightNamesToGetWithFullPath)