#!/usr/bin/python
# -*- coding: utf-8 -*-
#
"""
    Runs ImageNet Convolutional Neural Network implemented in software Caffe.
    This module only implements the classification. The network must be trained previously using caffe.
    
    Krizhevsky, Alex, Ilya Sutskever, and Geoffrey E. Hinton, Imagenet classification with deep convolutional neural networks, Advances in neural information processing systems, 2012.
    Jia, Yangqing and Shelhamer, Evan and Donahue, Jeff and Karayev, Sergey and Long, Jonathan and Girshick, Ross and Guadarrama, Sergio and Darrell, Trevor, Caffe: Convolutional Architecture for Fast Feature Embedding, arXiv preprint arXiv:1408.5093, 2014.
    
    Name: cnn_caffe.py
    Author: Alessandro dos Santos Ferreira ( santosferreira.alessandro@gmail.com )
"""

# Make sure that caffe is on the python path:
caffe_root = '/var/tmp/caffe/' 
import sys
sys.path.insert(0, caffe_root + 'python')
import caffe

import cv2
import numpy as np
import os

from collections import OrderedDict

from util.config import Config
from util.file_utils import File
from util.utils import TimeUtils

from classifier import Classifier

class CNNCaffe(Classifier):

    def __init__(self):
        self._create_net()
        
    def _create_net(self):
        self.model_def = Config("ModelDef", caffe_root + 'models/bvlc_reference_caffenet/deploy.prototxt', str)
        self.model_weights = Config("ModelWeights", caffe_root + 'models/bvlc_reference_caffenet/bvlc_reference_caffenet.caffemodel', str)
                       
        self.net = caffe.Net(self.model_def.value,          # defines the structure of the model
                                self.model_weights.value,   # contains the trained weights
                                caffe.TEST)                 # use test mode (e.g., don't perform dropout)
        
        self.mean_image = Config("MeanImage", caffe_root + 'python/caffe/imagenet/ilsvrc_2012_mean.npy', str)
        # load the mean image for subtraction
        mu = np.load(self.mean_image.value)
        mu = mu.mean(1).mean(1)  # average over pixels to obtain the mean (BGR) pixel values

        # create transformer for the input called 'data'
        self.transformer = caffe.io.Transformer({'data': self.net.blobs['data'].data.shape})

        self.transformer.set_transpose('data', (2,0,1))  # move image channels to outermost dimension
        self.transformer.set_mean('data', mu)            # subtract the dataset-mean value in each channel
        self.transformer.set_raw_scale('data', 255)      # rescale from [0, 1] to [0, 255]
        self.transformer.set_channel_swap('data', (2,1,0))  # swap channels from RGB to BGR

        self.labels_file = Config("LabelsFile", caffe_root + 'data/ilsvrc12/synset_words.txt', str)
    
    def get_config(self):
        caffe_config = OrderedDict()
        
        caffe_config["model_def"] = self.model_def
        caffe_config["model_weights"] = self.model_weights
        caffe_config["mean_image"] = self.mean_image
        caffe_config["labels_file"] = self.labels_file
        
        return caffe_config
        
    def set_config(self, configs):
        self.model_def = Config.nvl_config(configs["model_def"], self.model_def)
        self.model_weights = Config.nvl_config(configs["model_weights"], self.model_weights)
        self.mean_image = Config.nvl_config(configs["mean_image"], self.mean_image)
        self.labels_file = Config.nvl_config(configs["labels_file"], self.labels_file)
        
        self._create_net()

    def get_summary_config(self):
        caffe_config = OrderedDict()
        
        caffe_config[self.model_def.label] = self.model_def.value
        caffe_config[self.model_weights.label] = self.model_weights.value
        caffe_config[self.mean_image.label] = self.mean_image.value
        caffe_config[self.labels_file.label] = self.labels_file.value

        summary = ''
        for config in caffe_config:
            summary += "%s: %s\n" % (config, str(caffe_config[config]))
        
        return summary

    
    def classify(self, dataset, test_dir, test_data):
        test_dir = File.make_path(dataset, test_dir)
            
        classes = []
        labels = np.loadtxt(self.labels_file.value, str)

        images = sorted(os.listdir(File.make_path(test_dir)))

        for image in images:
            filepath = File.make_path(test_dir, image)
            image = cv2.imread(filepath)
                
            # resize the segment
            resized_image = np.zeros((512, 512, image.shape[2]), dtype="uint8")
            resized_image[0:image.shape[0], 0:image.shape[1]] = image[:,:]
            resized_image = resized_image[0:256, 0:256]
            cv2.imwrite(filepath.replace('.tif', '.jpeg'), resized_image)
            
            # load the image
            input_image = caffe.io.load_image(filepath)
            transformed_image = self.transformer.preprocess('data', input_image)
            
            # copy the image data into the memory allocated for the net
            self.net.blobs['data'].data[...] = transformed_image

            # perform classification
            output = self.net.forward()
            
            # the output probability vector for the first image in the batch
            prediction = output['prob'][0] 
            print(["%0.4f" % pr for pr in prediction ])

            classes.append(labels[prediction.argmax()])

        return classes