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

#

# The MIT License (MIT)

# 

# Peter Moss Acute Myeloid/Lymphoblastic Leukemia AI Research Project

# Copyright (C) 2018 Adam Milton-Barker (AdamMiltonBarker.com)

# 

# Permission is hereby granted, free of charge, to any person obtaining a copy

# of this software and associated documentation files (the "Software"), to deal

# in the Software without restriction, including without limitation the rights

# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell

# copies of the Software, and to permit persons to whom the Software is

# furnished to do so, subject to the following conditions:

# 

# The above copyright notice and this permission notice shall be included in

# all copies or substantial portions of the Software.

# 

# 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.

#

# Title:         Caffe Acute Lymphoblastic Leukemia CNN Info

# Description:   Used to view info Caffe Acute Lymphoblastic Leukemia CNN

# Configuration: Required/Confs.json

# Last Modified: 2019-03-10

#

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

import os, sys, cv2

sys.path.append('/home/upsquared/caffe/python')

import caffe

import numpy as np

from Classes.Helpers import Helpers

class allCNN():

    def __init__(self):

        """

        Sets up all default requirements and placeholders 

        needed for the Caffe Acute Lymphoblastic Leukemia CNN.

        """

        self.Helpers = Helpers()

        self.confs = self.Helpers.loadConfs()

        self.logFile = self.Helpers.setLogFile(self.confs["Settings"]["Logs"]["allCNN"])

        self.Helpers.logMessage(self.logFile, "allCNN", "Status", "Init complete")

    def loadCaffeNet(self):

        """

        Loads the Caffe network using prototxt layer definition.

        """

        self.net = caffe.Net(self.confs["Settings"]["Classifier"]["Caffe"]["layerFile"], caffe.TEST)

        print("")

        self.Helpers.logMessage(self.logFile, "allCNN", "Status", "Caffe net initialized")

    def printDetails(self):

        """

        Prints and logs input, blob and parameter info.

        """

        # Prints the Net Inputs

        self.Helpers.logMessage(self.logFile, "allCNN", "Net Inputs", str(self.net.inputs))

        # Prints the Net Blobs

        self.Helpers.logMessage(self.logFile, "allCNN", "Net Blobs", str(self.net.blobs))

        # Prints the Net Blob shapes

        self.Helpers.logMessage(self.logFile, "allCNN", "Net Blob shapes", str([(k, v.data.shape) for k, v in self.net.blobs.items()]))

        # Prints the Net Params

        self.Helpers.logMessage(self.logFile, "allCNN", "Net Params", str(self.net.params))

        # Prints the Net Params shapes

        self.Helpers.logMessage(self.logFile, "allCNN", "Net Params", str([(k, v[0].data.shape, v[1].data.shape) for k, v in self.net.params.items()]))

        print("")

    def writeOutputImages(self, image):

        """

        Writes the output images for each neuron in the first convolution layer.

        """

        # Transposes the input (50,50,3) -> (3,50,50)

        inp = np.transpose(cv2.imread(image))

        # Reshape the data blob

        self.net.blobs['data'].reshape(1, *inp.shape)

        self.net.blobs['data'].data[...] = inp

        # Passes the input data through the network to compute the output

        self.net.forward()

        # Loops through each neuron in the first convolution layer and saves the images in that neuron

        for i in range(30):

            cv2.imwrite(self.confs["Settings"]["Classifier"]["Data"]["dir"] + self.confs["Settings"]["Classifier"]["Info"]["outDir"] + 'conv1/out_' + str(i) + '.jpg', 

                        255 * self.net.blobs['conv1'].data[0,i])

        # Loops through each neuron in the second convolution layer and saves the images in that neuron

        for i in range(30):

            cv2.imwrite(self.confs["Settings"]["Classifier"]["Data"]["dir"] + self.confs["Settings"]["Classifier"]["Info"]["outDir"] + 'conv2/out_' + str(i) + '.jpg', 

                        255 * self.net.blobs['conv2'].data[0,i])

        self.Helpers.logMessage(self.logFile, 

                                "allCNN", 

                                "Output Images", 

                                "Output images written to " + self.confs["Settings"]["Classifier"]["Data"]["dir"] + self.confs["Settings"]["Classifier"]["Info"]["outDir"])

    def saveCaffeNet(self):

        """

        Saves our Caffe network.

        """

        self.net.save(self.confs["Settings"]["Classifier"]["Model"]["file"])

        self.Helpers.logMessage(self.logFile, 

                                "allCNN", 

                                "Status", 

                                "Caffe net saved")

allCNN = allCNN()

def main(argv):

    if(len(argv) < 1):

        """

        Incorrect arguments size.

        """

        allCNN.Helpers.logMessage(allCNN.logFile, 

                                  "allCNN", 

                                  "Arguments", 

                                  "Please provide NetworkInfo or Outputs argument")

    elif argv[0] == "NetworkInfo":

        """

        Provides information about our Caffe network.

        """

        allCNN.loadCaffeNet()

        allCNN.printDetails()

    elif argv[0] == "Outputs":

        """

        Plots the outputs of each neuron as images.

        """

        allCNN.loadCaffeNet()

        allCNN.writeOutputImages(allCNN.confs["Settings"]["Classifier"]["Data"]["dir"] + allCNN.confs["Settings"]["Classifier"]["Data"]["dirTest"] + allCNN.confs["Settings"]["Classifier"]["Info"]["testImage"])

    elif argv[0] == "Save":

        """

        Saves our Caffe network.

        """

        allCNN.loadCaffeNet()

        allCNN.saveCaffeNet()

if __name__ == "__main__":

    main(sys.argv[1:])