import numpy as np
import lmdb
import Utilities
import caffe
import struct
from tqdm import tqdm
from sklearn import metrics
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
import Tests
from itertools import izip
from caffe.proto import caffe_pb2
from caffe.io import datum_to_array, array_to_datum
from Config import CONTROL_IMAGE_OUTPUT, CATEGORIES, IMAGE_HEIGHT, IMAGE_WIDTH, NUM_CHANNELS, OUTPUT_FOLDER, GPU_PROC


def evaluate(net, db_name):
    env = lmdb.open(db_name)
    entries = env.stat()['entries']
    batchsize = 25

    datum = caffe_pb2.Datum()
    iters = entries / batchsize

    outfile = open('Temp.bin', 'w+b')
    struct_fmt = 'B'*16384
    struct_len = struct.calcsize(struct_fmt)
    struct_unpack = struct.Struct(struct_fmt).unpack_from

    with env.begin() as in_txn:
        cursor = in_txn.cursor()
        for i in tqdm(range(iters)):
            out = net.forward()
            out = out["conv1_1_D"]

            for j in range(out.shape[0]):
                cursor.next()
                datum.ParseFromString(cursor.value())
                true_val = datum_to_array(datum).astype(int)
                out_j = out[j]
                predicted_val = np.zeros((out_j.shape[1], out.shape[2]))
                for k in range(out_j.shape[1]):
                    for l in range(out_j.shape[2]):
                        indices = np.argsort(out_j[:, k, l])
                        predicted_val[k,l] = indices[-1]
                predicted_val = predicted_val.flatten()
                true_val = true_val.flatten()
                predicted_val[np.where(true_val == 3)] = 3
                predicted_val = predicted_val.astype(int)
                s_truth = struct.pack(struct_fmt, *true_val.tolist())
                s_pred = struct.pack(struct_fmt, *predicted_val.tolist())
                outfile.write(s_truth + s_pred)
            outfile.flush()
    outfile.close()
    return iters, batchsize, struct_fmt


def evaluate_thresholds(net, db_name, thresholds):
    env = lmdb.open(db_name)
    entries = env.stat()['entries']
    batchsize = 25

    datum = caffe_pb2.Datum()
    iters = entries / batchsize

    outfiles = []
    for i in range(len(thresholds)):
    	outfiles.append(open('Temp' + str(i) + '.bin', 'w+b'))

    struct_fmt = 'B'*16384
    struct_len = struct.calcsize(struct_fmt)
    struct_unpack = struct.Struct(struct_fmt).unpack_from

    with env.begin() as in_txn:
        cursor = in_txn.cursor()
        for i in tqdm(range(iters)):
            out = net.forward()
            out = out["conv1_1_D"]

            for j in range(out.shape[0]):
                cursor.next()
                datum.ParseFromString(cursor.value())
                true_val = datum_to_array(datum).astype(int)
                out_j = out[j]
                for threshold_index in range(len(thresholds)):
                    predicted_val = np.zeros((out_j.shape[1], out.shape[2]))
                    for k in range(out_j.shape[1]):
                        for l in range(out_j.shape[2]):
                            indices = np.argsort(out_j[:, k, l])
                            sorted_vals = out_j[indices, k, l]
                            if (sorted_vals[-2]/sorted_vals[-1]) < thresholds[threshold_index]:
                                predicted_val[k,l] = indices[-1]
                            else:
                                predicted_val[k,l] = 0
                    predicted_val = predicted_val.flatten()
                    true_val = true_val.flatten()
                    predicted_val[np.where(true_val == 3)] = 3
                    predicted_val = predicted_val.astype(int)
                    s_truth = struct.pack(struct_fmt, *true_val.tolist())
                    s_pred = struct.pack(struct_fmt, *predicted_val.tolist())
                    outfiles[threshold_index].write(s_truth + s_pred)
                outfiles[threshold_index].flush()
    for i in range(len(thresholds)):
        outfiles[i].close()
    return iters, batchsize, struct_fmt


def make_report(file, iters, batchsize, struct_fmt, report_name):
    max_number = iters*batchsize*len(struct_fmt)
    struct_len = struct.calcsize(struct_fmt)
    struct_unpack = struct.Struct(struct_fmt).unpack_from

    struct_batch_len = len(struct_fmt)

    truth = np.zeros(max_number, dtype=np.uint8)
    predicted = np.zeros(max_number, dtype=np.uint8)
    with open(file) as tempfile:
        for m_iter in tqdm(range(iters*batchsize)):
            data = tempfile.read(struct_len)
            if not data: break
            s = struct_unpack(data)
            truth[m_iter*struct_batch_len:(m_iter+1)*struct_batch_len] = s

            data = tempfile.read(struct_len)
            s = struct_unpack(data)
            predicted[m_iter*struct_batch_len:(m_iter+1)*struct_batch_len] = s

    with open(report_name, 'w') as report_file:
        acc = metrics.accuracy_score(truth, predicted)
        report_file.write("Accuracy\n")
        report_file.write(str(acc))
        report_file.write("\n")
        print acc
        prec = metrics.precision_score(truth, predicted, average='macro')
        report_file.write("Precision\n")
        report_file.write(str(prec))
        report_file.write("\n")
        print prec
        rec = metrics.recall_score(truth, predicted, average='macro')
        report_file.write("Recall\n")
        report_file.write(str(rec))
        report_file.write("\n")
        print rec
        report = metrics.classification_report(truth, predicted)
        report_file.write("Report\n")
        report_file.write(report)
        report_file.write("\n")
        print report
        conf = metrics.confusion_matrix(truth, predicted)
        report_file.write("Confusion Matrix\n")
        report_file.write(str(conf))
        report_file.write("\n")
        print conf
        #iou = Tests.intersection_over_union_metric(truth, predicted)
        #report_file.write("IoU\n")
        #report_file.write(str(iou))
        #report_file.write("\n")
        #print iou


if __name__ == "__main__":
    caffe.set_device(GPU_PROC)
    caffe.set_mode_gpu()
    test_net_file = "Training/test_segnet.prototxt"
    weights = "Training/snap_hard_low2.caffemodel"
    db_name = OUTPUT_FOLDER + "gisat_mask_tst.lmdb"

    my_net = caffe.Net(test_net_file, weights=weights, phase=caffe.TEST)
    iters, batchsize, struct_fmt = evaluate_thresholds(my_net, db_name, [0.1, 0.25])
    print 'evaluated'
    #iters = 3745
    batchsize = 25
    struct_fmt = 'B' * 16384
    make_report('Temp0.bin', iters, batchsize, struct_fmt, 'Report_Gisat_hard_low_t_01.txt')
    make_report('Temp1.bin', iters, batchsize, struct_fmt, 'Report_Gisat_hard_low_t_025.txt')
