################################################################################
# (c) [2013] The Johns Hopkins University / Applied Physics Laboratory All Rights Reserved.
# Contact the JHU/APL Office of Technology Transfer for any additional rights.  www.jhuapl.edu/ott
# 
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# 
#    http://www.apache.org/licenses/LICENSE-2.0
# 
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
################################################################################

""" Solves a membrane detection/classification problem.

This module provides the toplevel interface for solving a binary
"membrane vs non-membrane" classification problem for EM data sets
(e.g. [1]) using convolutional neural networks.

The overall approach is based on Dan Ciresan's paper [2] and the code
is derived from a LeNet example included in the Theano code base for
MNIST classification.

References:
  [1] http://brainiac2.mit.edu/isbi_challenge/
  [2] Ciresan, Dan, et al. "Deep neural networks segment neuronal membranes
      in electron microscopy images." Advances in neural information
      processing systems. 2012.

December 2013, mjp
"""

import os, os.path
import sys, time
import socket
import argparse

import numpy
from PIL import Image

import pdb

import theano
import theano.tensor as T

import em_networks as EMN
from em_utils import *
from tiles import *



def load_membrane_data(trainDataFile, trainLabelsFile, 
                       tileSize,
                       trainSlices, validSlices,
                       nZeeChannels=0):
    """Loads data set and creates corresponding tile managers.  
    """
    
    # load the volume and the labels
    if trainDataFile.endswith('.tif'):
        X = load_tiff_data(trainDataFile)
        # Assumes raw conference data (i.e. not preprocessed).
        #for ii in range(X.shape[0]):
        #    X[ii,:,:] = X[ii,:,:] - numpy.mean(X[ii,:,:])
        #X = X / numpy.max(numpy.abs(X))
        print '[%s]: Warning: no longer zero-meaning and scaling data' % __name__
    elif trainDataFile.endswith('.npz'):
        # assumes volume data is stored as the tensor X and is suitably preprocessed
        X = numpy.load(trainDataFile)['X']
    else:
        raise RuntimeError('unexpected data file extension')

    Y = load_tiff_data(trainLabelsFile)

    # mirror edges
    border = numpy.floor(tileSize/2.)
    X = mirror_edges_tensor(X, border)
    Y = mirror_edges_tensor(Y, border)

    # Use 0 and 1 as class labels.  This is actually important because
    # the neural network code will use class labels as indices into
    # the outputs of the last network layer.
    #
    # 0 := non-membrane
    # 1 := membrane
    Y[Y==0] = 1;  Y[Y==255] = 0
    assert(Y.max() == 1)

    X_train = X[trainSlices,:,:]
    Y_train = Y[trainSlices,:,:]
    X_valid = X[validSlices,:,:]
    Y_valid = Y[validSlices,:,:]

    # tile managers will put the images into GPU memory via Theano shared vars.
    train = TileManager(X_train, Y_train, tileSize=tileSize, nZeeChannels=nZeeChannels)
    valid = TileManager(X_valid, Y_valid, tileSize=tileSize, nZeeChannels=nZeeChannels)

    return (train, valid, (X, Y))



def random_image_modifiers(flipProb=.6, rotProb=.6):
    """Randomly applies certain transforms to a 2d image.
    As of this writing, these transforms are some
    combination of flips and rotations.
    """
    # clip probabilities to [0,1]
    flipProb = max(min(flipProb,1),0)
    rotProb = max(min(rotProb,1),0)
    
    flipDim = 0; rotDir = 0
    if numpy.random.rand() < flipProb:
        flipDim = numpy.sign(numpy.random.rand() - .5)
    if numpy.random.rand() < rotProb:
        rotDir = numpy.sign(numpy.random.rand() - .5)
    return flipDim, rotDir


    
def train_network(nn, trainMgr, validMgr,
                  nEpochs=30, learningRate=.001, decay=.995,
                  maxNumTilesPerEpoch=sys.maxint,
                  outDir="."):
    """Learns parameters for the given neural network.
    """

    p2 = int(numpy.floor(nn.p/2.0))
    
    # compute number of minibatches 
    nTrainBatches = int(numpy.ceil(trainMgr.batchSize / nn.miniBatchSize))
    nValidBatches = int(numpy.ceil(validMgr.batchSize / nn.miniBatchSize))
    print '[%s]: # of training batches is %d' % (__name__, nTrainBatches)

    # allocate symbolic variables 
    indexT = T.lscalar()         # index to a [mini]batch
    learningRateT = T.scalar()   # learning rate, theano variable
    
    print '[%s]: initializing Theano...' % __name__

    #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    # functions for the validation data
    #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    predict_validation_data = theano.function([indexT], nn.layers[-1].p_y_given_x,
             givens={