#!/usr/bin/env python

import unittest
import os
from context import config
from context import model
from context import predictor
from context import transformer
from context import image_and_label
import sys
import argparse
import shutil
import pickle


# GPU is for program to decide whether run CPU or GPU tests.
# GPU_ID is for GPU card pick.
GPU = False
GPU_ID = None

# Use argparse to get parameters form command linbe
# This can not be done in main() because of unittest itself
# So we put it here.
ap = argparse.ArgumentParser()
ap.add_argument("-g", "--gpu", required = False,
    help = "Sepcify your gpu id if you want to test in gpu.")
args = vars(ap.parse_args())
if args['gpu'] != None:
    GPU = True
    GPU_ID = int(args['gpu'])


# Functions below can delete file and directories silently,
# and it will not report error if file do not exists.
def remove_file(filename):
    try:
        os.remove(filename)
    except OSError:
        pass

def remove_dir(dirname):
    try:
        shutil.rmtree(dirname)
    except OSError:
        pass

# This function is designed for test task, all the file gennerated by
# test will be deleted here.
# We cann't use tearDown() in test because it will delete file after
# every test, some file generated by one test will be useful to another test.
def remove_files():
    remove_file('test_list.txt')
    remove_file('image_list_save_tmp.p')


# We choose unittest for our test task because it is easy to read and maintain.
# Note if you want to add a test here, the function must be start with "test"
# (something like "test_foo()").
# unittest is designed for isolated test, so it get test function in lexicographical
# order, 0 is fisrt and z is last. Here some test has to use former test's result,
# so I put a serial number after "test" to keep test run in right order.
class TestTools(unittest.TestCase):

    # Check if pycaffe can be imported.
    def test_1_pycaffe(self):
        import imp
        sys.path.append(config.pycaffe_path)
        try:
            imp.find_module('caffe')
            found = True
        except ImportError:
            found = False
        self.assertTrue(found)

    # Check if list file can be generated.
    def test_2_labellist_create(self):
        list_config = {}

        list_config['_path'] = ['test_images/blot',
                                    'test_images/other']
        list_config['class_number'] = 2
        list_config['_label'] = [0, 1]
        list_config['file_path'] = 'test_list.txt'
        list_config['amount'] = [2, 2]


        # This is to make sure file list can be used for generating by using absolute path
        list_config['absolute_path'] = 1

        list_file = image_and_label.build_list(list_config)
        self.assertTrue(os.path.exists(list_file))

    # Check if list file can be divide correctly.
    def test_3_labellist_divide(self):
        list_file = 'test_list.txt'
        image_list, label_list = image_and_label.get(list_file)

        pickle.dump(image_list, open('image_list_save_tmp.p', 'wb'))

        self.assertEqual(len(image_list), 4)
        self.assertEqual(sum(label_list), 2)

    # Skip CPU test if command line accept GPU parameters.
    @unittest.skipIf(GPU,'Skip because it\' GPU test')
    def test_4_test_CPU(self):
        _net = model.read(False, None, 'test_deploy.prototxt',
                    'testmodel.caffemodel', 1, 3, 227, 227)
	_transformer = transformer.set(_net, 'test_mean_bin', True)
        
        image_list = pickle.load(open('image_list_save_tmp.p', 'rb'))

	predict_label_list = []
        expected_result = [0, 0, 1, 1]
	
        for image in image_list:
            prediction = predictor.predict(image, _transformer, _net)
            predict_label_list.append(prediction)

        self.assertEqual(predict_label_list, expected_result)

    # Skip GPU test if no GPU parameter was sepcified in commandline.
    @unittest.skipIf(not GPU,'Skip because it\' CPU test')
    def test_4_test_GPU(self):
        _net = model.read(True, GPU_ID, 'test_deploy.prototxt',
                    'testmodel.caffemodel', 1, 3, 227, 227)
        _transformer = transformer.set(_net, 'test_mean_bin', True)
        
        image_list = pickle.load(open('image_list_save_tmp.p', 'rb'))
	
	predict_label_list = []
        expected_result = [0, 0, 1, 1]

        for image in image_list:
            prediction = predictor.predict(image, _transformer, _net)
            predict_label_list.append(prediction)

        self.assertEqual(predict_label_list, expected_result) 

if __name__ == '__main__':
    # Clean up generated files in old time.
    remove_files()

    # Run all the tests.
    # Note that if you use unittest.main(), its own command line parameters
    # will overide GPU parameters which we wrote and wanted.
    # So we choose another way to run these tests.
    suite = unittest.TestLoader().loadTestsFromTestCase(TestTools)
    unittest.TextTestRunner(verbosity=4).run(suite)

    # Clean up the mess.
    remove_files()
