import numpy as np
import lmdb
import Utilities
import caffe
import gdal
from random import randrange
from random import random
from random import shuffle
from Config import OUTPUT_FOLDER, CATEGORIES, IMAGE_HEIGHT, IMAGE_WIDTH, NUM_CHANNELS, GPU_PROC
from itertools import izip
from tqdm import tqdm
import matplotlib.pyplot as plt
from caffe.proto import caffe_pb2
from caffe.io import datum_to_array, array_to_datum


def shuffle_db(database_name, database_mask_name, new_db_name, new_db_mask_name):
    in_db = lmdb.open(database_name, readonly=True)
    in_db_lab = lmdb.open(database_mask_name, readonly=True)

    num_images = in_db.stat()["entries"]
    indices = []
    with in_db.begin() as in_txn:
        cursor = in_txn.cursor()
        for i in tqdm(range(num_images)):
            indices.append(cursor.key())
            cursor.next()
    print indices
    indices_repl = indices
    shuffle(indices_repl)
    print indices_repl
    out_db = lmdb.open(new_db_name)
    out_db_lab = lmdb.open(new_db_mask_name)
    with in_db.begin(write=True) as in_txn, in_db_lab.begin(write=True) as in_lab_txn, out_db.begin(write=True) as out_txn, out_lab_db.begin(write=True) as out_lab_txn:
        cursor = in_txn.cursor()
        cursor_lab = in_lab_txn.cursor()
        for i in tqdm(range(num_images)):
            val = cursor.value()
            mask_val = cursor_lab.value()
            
            out_txn.put('{:0>10d}'.format(indices_repl[indices[i]]), val)
            out_lab_txn.put('{:0>10d}'.format(indices_repl[indices[i]]), mask_val)
            
            cursor.next()
            cursor_lab.next()

def test_net_on_lmdb(database_name, deploy_net_filename, deploy_weights):
    env = lmdb.open(database_name, readonly=True)
    num_images = env.stat()["entries"]
    env.begin()

    testnet_net = caffe.Net(deploy_net_filename, weights=deploy_weights, phase=caffe.TEST)
    numClasses = 4

    # num_batches depends on the number of Images you will test and the batch_size defined on test.prototxt
    batch_size = 25
    num_batches = num_images/ batch_size

    # Initialize the accuracy reported by the net
    acc_top1 = 0
    #acc_top5 = 0

    #datum = caffe_pb2.Datum()

    for t in tqdm(range(num_batches)):
        #for j in range(batch_size):
        #    value_text = '{:0>10d}'.format(t*j+j)
        #    value = env.get(value_txt)
        #    datum.ParseFromString(value)
        #    X = datum_to_array(datum)
        #    testnet_net.blobs['data'].data[j] = X
        # Run the net for the current image batch
        #testnet_net.forward()
	testnet_net.forward()

        # Update accuracy with the average accuracy of the current batch
        output = testnet_net.blobs["accuracy"].data
        print output
        acc_top1 = acc_top1 + output
        #acc_top5 = acc_top5 + net.blobs[nameAccuracyTop5].data
    acc_top1 = acc_top1 / num_batches
    #acc_top5 = acc_top5 / num_batches
    return acc_top1



def intersection_over_union_metric(prediction, ground_truth):
    iou_all = []
    for i in range(len(CATEGORIES)):
        union_of_i = np.sum(np.bitwise_or(prediction == i, ground_truth == i))
        intersection_of_i = np.sum(np.bitwise_and(prediction == i, ground_truth == i))
        if union_of_i == 0:
            iou_all.append(None)
        else:
            iou_all.append(float(intersection_of_i) / union_of_i)
    failed = any(i < 0.4 for i in iou_all if i != None)
    return iou_all, not failed


def accuracy_metric(prediction, ground_truth):
    accuracy = float(np.sum(ground_truth == prediction)) / (prediction.shape[0] * prediction.shape[1])
    failed = False
    if accuracy < 0.75:
        failed = True
    return accuracy, not failed


def classify_lmdb(lmdb_images, lmdb_labels, net, db_filename, db_filename_label, num_iter=60000, metric=accuracy_metric):
    env = lmdb.open(lmdb_images)
    env_label = lmdb.open(lmdb_labels)
    txn = env.begin()
    txn_label = env_label.begin()
    cursor = txn.cursor()
    cursor_label = txn_label.cursor()

    num_images = env.stat()["entries"]
    indices = 0
    in_db = lmdb.open(db_filename, map_size=int(num_iter * IMAGE_HEIGHT * IMAGE_WIDTH * NUM_CHANNELS * 8 + num_images))
    in_db_lab = lmdb.open(db_filename_label, map_size=int(num_iter * IMAGE_HEIGHT * IMAGE_WIDTH * 1 * 8 + num_images))
    with in_db.begin(write=True) as in_txn, in_db_lab.begin(write=True) as in_txn_lab:
        for i in tqdm(range(num_iter)):
            rand_idx = randrange(0, num_images - 1)
            str_rand_idx = '{:0>10d}'.format(rand_idx)

            img_idx = txn.get(str_rand_idx)
            lab_idx = txn_label.get(str_rand_idx)
            value = img_idx
            label = lab_idx
            datum = caffe_pb2.Datum()
            datum.ParseFromString(value)
            X = datum_to_array(datum)
            X = np.array(X)

            datum.ParseFromString(label)
            Y = datum_to_array(datum)
            Y = np.array(Y)

            net.blobs['data'].data[...] = X
            net.forward()
            out = net.blobs['conv1_1_D'].data[0]
            
            predictions = np.zeros((out.shape[1], out.shape[2]))
            for r in range(out.shape[1]):
                for c in range(out.shape[2]):
                    predictions[r, c] = np.argmax(out[:, r, c])
            metric_val, metric_eval = metric(predictions, Y)
            print metric_val
            if not metric_eval:
                in_txn.put(str_rand_idx, img_idx)
                in_txn_lab.put(str_rand_idx, lab_idx)
                indices += 1
                print 'Hard'
            else:
                if random() > 0.8:
                    in_txn.put(str_rand_idx, img_idx)
                    in_txn_lab.put(str_rand_idx, lab_idx)
                    indices += 1
                    print 'Luck'
            print indices

def show_database_images(db_path):
	env = lmdb.open(db_path)
	txn = env.begin()
	cursor = txn.cursor()
	for idx, (key, value) in enumerate(cursor):
		datum = caffe_pb2.Datum()
		datum.ParseFromString(value)
		X = datum_to_array(datum)
		X = np.array(X)
		img = X[[1,2,0], :,:]
		img = np.rollaxis(img, 0, start=3)
		plt.imshow(img)
		y = datum.label


def small_images_vis_check():
	for i in range(10):
		fig = plt.figure()
		filenames = ["CroppedTests/net_output_mask", "CroppedTests/temp_mask", "CroppedTests/temp"]
		for j in range(0, 2):
			cmap = plt.get_cmap('jet')
			plt.subplot(1, 3, j+1)
			img = gdal.Open(filenames[j] + str(i) + ".tif", gdal.GDT_UInt16)
			img = img.ReadAsArray()
			cax = plt.imshow(img, cmap=cmap, vmin=0, vmax=3)
			cbar = fig.colorbar(cax, ticks=[0, 1, 2, 3])
			cbar.ax.set_yticklabels(['Ground', 'Shadow', 'Cloud', 'NaN'])  # vertically oriented colorbar
		plt.subplot(1, 3, 3)
		img = gdal.Open(filenames[2] + str(i) + ".tif", gdal.GDT_UInt16)
		img = img.ReadAsArray()
		img = img.astype("uint8")
		img = img[[1,2,0], :,:]
		img = np.rollaxis(img, 0, start=3)
		cax = plt.imshow(img)
		plt.show()


if __name__ == "__main__":
        caffe.set_device(GPU_PROC)
        caffe.set_mode_gpu()
        data_pick = 1
        datanames = ['training', 'testing', 'gisat']
        data = '../' + datanames[data_pick] + "_data"
        label = '../' + datanames[data_pick] + "_mask"
        
        data = '../gisat_data_tst' #'../gisat_scl_data'
        label = '../gisat_mask_tst' #'../gisat_scl_mask'

        out_name = 'test_HL21'
	net_file = "/Training/deploy_segnet.prototxt"
        test_net_file = "/Training/test_segnet.prototxt"
        weights = "/Training/snap_hard_low2.caffemodel"

	my_net = caffe.Net(net_file, weights=weights, phase=caffe.TEST)
        
      	Utilities.classify_all_images(train_data=data, train_labels=label, net=my_net, img_name=out_name)
 	#Utilities.classify_all_images_patches(train_data=data, train_labels=label, net=my_net, img_name=out_name)
        #classify_lmdb(OUTPUT_FOLDER + "rgbgendata.lmdb", OUTPUT_FOLDER + "rgbgendatalabel.lmdb", my_net, OUTPUT_FOLDER + "superrgbgendata2.lmdb", OUTPUT_FOLDER + "superrgbgendatalabel2.lmdb", num_iter=20000)