import caffe
import numpy as np
import lmdb
import os
import csv
import gdal
import Extract_images
import math
import WriteTiff
from Config import SINGLE_LABEL, IMAGE_HEIGHT, IMAGE_WIDTH, \
    NUM_CHANNELS, NUM_SLICES, PATCH_SIZE, HALF_WIDTH, HALF_HEIGHT, \
    STRIDE, OUTPUT_FOLDER, MASK_SIZE, CONTROL_IMAGE_SIZE, CONTROL_IMAGE_OUTPUT, \
    TEST_STRIDE
from random import randrange
from itertools import izip
from itertools import islice
import matplotlib.pyplot as plt
from caffe.proto import caffe_pb2
from caffe.io import datum_to_array, array_to_datum
import tqdm
import Tests


def make_datum(array):
    return caffe.io.array_to_datum(array)


def find_all_files(root_dir):
    train_files = [os.path.join(root, name)
                   for root, dirs, files in os.walk(root_dir)
                   for name in files
                   if name.endswith(".tif") and not name.startswith("fmask")]
    with open("train_files", "w") as wf:
        for in_idx, img_path in enumerate(train_files):
            wf.write(str(in_idx))
            wf.write(",")
            wf.write(img_path)
            wf.write("\n")


def grouper(iterable, n):
    it = iter(iterable)
    while True:
       chunk = tuple(islice(it, n))
       if not chunk:
           return
       yield chunk

def slice_img_stride(img, stride):
    for i in xrange(0, img.shape[1], stride):
        if i + IMAGE_HEIGHT > img.shape[1]:
            i = img.shape[1] - IMAGE_HEIGHT
        for j in xrange(0, img.shape[2], stride):
            if j + IMAGE_WIDTH > img.shape[2]:
                j = img.shape[2] - IMAGE_WIDTH
            yield (i, j)


def slice_img_stride_count(img, stride):
    count = 0
    for i in xrange(0, img.shape[1], stride):
        if i + IMAGE_HEIGHT > img.shape[1]:
            i = img.shape[1] - IMAGE_HEIGHT
        for j in xrange(0, img.shape[2], stride):
            if j + IMAGE_WIDTH > img.shape[2]:
                j = img.shape[2] - IMAGE_WIDTH
            count += 1


def classify_patch_votes_batch(img, net, stride, group):
    output = np.zeros((img.shape[1], img.shape[2]))
    votes = np.zeros((4, img.shape[1], img.shape[2]))
    mean_val = np.array([32.97, 35.68, 40.7])
    mean_val = np.rollaxis(np.tile(mean_val, (IMAGE_HEIGHT, IMAGE_WIDTH, 1)), 2, start=0)
    for group in tqdm.tqdm(grouper(slice_img_stride(img, stride), group), total=slice_img_stride_count(img, stride)):
        iter = 0
        for i, j in group:
            sub_img = img[:, i:i+IMAGE_HEIGHT, j:j+IMAGE_WIDTH]        
            net.blobs['data'].data[iter] = sub_img - mean_val
            iter += 1
        
        net.forward()
        out = net.blobs['prob'].data[:len(group)]
        for (i,j), prediction in izip(group, out):
            votes[:, i:i+IMAGE_HEIGHT, j:j+IMAGE_WIDTH] += prediction
            #print votes[:, i, j]

    for i in range(img.shape[1]):
        for j in range(img.shape[2]):
            output[i,j] = np.argmax(votes[:, i, j])
    return output.astype('uint8'), votes


def classify_patch_votes(img, net, stride):
    output = np.zeros((img.shape[1], img.shape[2]), np.uint8)
    votes = np.zeros((4, img.shape[1], img.shape[2]), np.float64)
    mean_val = np.array([32.97, 35.68, 40.7])
    mean_val = np.rollaxis(np.tile(mean_val, (IMAGE_HEIGHT, IMAGE_WIDTH, 1)), 2, start=0)
    for i in tqdm.tqdm(xrange(0, img.shape[1], stride)):
        if i + IMAGE_HEIGHT > img.shape[1]:
            i = img.shape[1] - IMAGE_HEIGHT
        for j in xrange(0, img.shape[2], stride):
            if j + IMAGE_WIDTH > img.shape[2]:
                j = img.shape[2] - IMAGE_WIDTH
            sub_img = img[:, i:i+IMAGE_HEIGHT, j:j+IMAGE_WIDTH]
            net.blobs['data'].data[...] = sub_img - mean_val
            out = net.forward()['prob'][0]
            votes[:, i:i+IMAGE_HEIGHT, j:j+IMAGE_WIDTH] += out;

    for i in range(img.shape[1]):
        for j in range(img.shape[2]):
            output[i,j] = np.argmax(votes[:, i, j])
    return output.astype('uint8'), votes


def classify_image(image, net):
    mean_val = np.array([32.97, 35.68, 40.7])
    mean_val = np.rollaxis(np.tile(mean_val, (IMAGE_HEIGHT, IMAGE_WIDTH, 1)), 2, start=0)
    net.blobs['data'].data[...] = image - mean_val
    out = net.forward()['prob'][0]
    out_votes = np.zeros((out.shape[1], out.shape[2]))
    for i in range(out.shape[1]):
        for j in range(out.shape[2]):
            out_votes[i,j] = np.argmax(out[:, i, j])
            #sorted_out = np.argsort(out[:,i,j])
            #if (out[sorted_out[-1],i , j]/out[sorted_out[-2], i, j]) > 1.08:
            #    best = sorted_out[-1]
            #else:
            #    best = 4 #No observation
            #out_votes[i,j] = best
    print out_votes.shape
    return out_votes


def classify_all_images(train_data, net, train_labels, img_name='test'):
    with open(train_data) as f, open(train_labels) as labels:
        my_iter = 0
        for img_path, fmask_path in izip(f, labels):
            img_path = img_path.strip()
            fmask_path = fmask_path.strip()

            img = gdal.Open(img_path, gdal.GDT_UInt16)
            img = img.ReadAsArray()

            mask = gdal.Open(fmask_path, gdal.GDT_UInt16)
            mask = mask.ReadAsArray()

            mask = Extract_images.transform_mask(mask, fmask_path)
            img = Extract_images.transform_img(img, NUM_CHANNELS, img_path)
           
            WriteTiff.write_tiff_mask(CONTROL_IMAGE_OUTPUT + img_name + str(my_iter) + "Orig.tif", mask)
            voted, lab_votes = classify_patch_votes_batch(img, net, TEST_STRIDE, 10)
            voted[np.where(mask == 3)] = 3
            acc = (mask == voted).astype(np.uint8)
            accuracy = float(np.sum(acc)) / (img.shape[1]*img.shape[2])
            print accuracy
            print Tests.intersection_over_union_metric(voted, mask)
            WriteTiff.write_tiff_mask(CONTROL_IMAGE_OUTPUT + img_name +str( my_iter) +  "Voted" + str(TEST_STRIDE) + ".tif", voted)
            #WriteTiff.write_tiff(CONTROL_IMAGE_OUTPUT + img_name +str( my_iter) +  "Votes" + str(TEST_STRIDE) + ".tif", lab_votes, type='float')
            my_iter += 1
        return accuracy, acc, voted



def classify_all_images_patches(train_data, net, train_labels, img_name='test'):
    with open(train_data) as f, open(train_labels) as labels:
        my_iter = 0
        for img_path, fmask_path in izip(f, labels):
            img_path = img_path.strip()
            fmask_path = fmask_path.strip()

            img = gdal.Open(img_path, gdal.GDT_UInt16)
            img = img.ReadAsArray()

            mask = gdal.Open(fmask_path, gdal.GDT_UInt16)
            mask = mask.ReadAsArray()
            
            mask = Extract_images.transform_mask(mask, fmask_path)
            img = Extract_images.transform_img(img, NUM_CHANNELS, img_path)

            WINDOW_SIZE = CONTROL_IMAGE_SIZE/2

            for i in range(10):
                hrand = randrange(WINDOW_SIZE + 1, img.shape[1] - WINDOW_SIZE - 1)
                wrand = randrange(WINDOW_SIZE + 1, img.shape[2] - WINDOW_SIZE - 1)

                sub_img = img[:, hrand-WINDOW_SIZE: hrand + WINDOW_SIZE,\
                       wrand-WINDOW_SIZE:wrand+WINDOW_SIZE]

                output = classify_image(sub_img, net)
                sub_mask = mask[\
                       hrand-WINDOW_SIZE: hrand+WINDOW_SIZE,\
                       wrand-WINDOW_SIZE:wrand+WINDOW_SIZE]
                output[np.where(sub_mask == 3)] = 3

                WriteTiff.write_tiff_mask(CONTROL_IMAGE_OUTPUT + "net_output_mask" + img_name + str(my_iter) + str(i) + ".tif", output)
                
                WriteTiff.write_tiff(CONTROL_IMAGE_OUTPUT + "temp" + img_name + str(my_iter) + str(i) + ".tif", sub_img)
                WriteTiff.write_tiff_mask(CONTROL_IMAGE_OUTPUT + "temp_mask" + img_name + str(my_iter) + str(i) + ".tif", sub_mask)

                acc = (sub_mask == output).astype(np.uint8)
                accuracy = float(np.sum(acc)) / (CONTROL_IMAGE_SIZE*CONTROL_IMAGE_SIZE)
                print accuracy
                print Tests.intersection_over_union_metric(output, sub_mask)
            my_iter += 1
            break
        return accuracy, acc, output


def create_lmdb(data_filename, db_filename, db_im_size, RNG_STATE):
    in_db = lmdb.open(db_filename, map_size=int(db_im_size))
    miterator = 0
    with in_db.begin(write=True) as in_txn:
        with open(data_filename) as f:
            for img_path in f:
                img_path = img_path.strip()
                print img_path
                img = gdal.Open(img_path, gdal.GDT_UInt16)
                img = img.ReadAsArray()

                if "mask" in img_path:
                    img = Extract_images.transform_mask(img, img_path)

                    img = img.reshape(img.shape + (1,))
                    # Switch to Channel x Height x Width order
                    img = img.transpose((2, 0, 1))
                else:
                    img = Extract_images.transform_img(img, NUM_CHANNELS, img_path)

                images = Extract_images.sliding_window(img, STRIDE,
                                                       PATCH_SIZE)  # slice_imgs(img, mask, img_width=IMAGE_WIDTH, img_height=IMAGE_HEIGHT)

                if SINGLE_LABEL and "mask" in img_path:
                    new_images = []
                    for img in images:
                        new_mask = img[:, img.shape[1] / 2, img.shape[2] / 2]
                        new_mask = new_mask.reshape(new_mask.shape + (1, 1,))
                        new_images.append(new_mask)
                    images = new_images

                np.random.set_state(RNG_STATE)
                np.random.shuffle(images)
                for img in images:
                    datum = make_datum(img)
                    in_idx = miterator
                    miterator += 1
                    in_txn.put('{:0>10d}'.format(in_idx), datum.SerializeToString())
                    print '{:0>10d}'.format(in_idx) + ':' + img_path
    in_db.close()


def calc_mean(data_filename):
    miterator = 0
    all_means = [0, 0, 0]
    with open(data_filename) as f:
        for img_path in f:
            img_path = img_path.strip()
            img = gdal.Open(img_path, gdal.GDT_UInt16)
            img = img.ReadAsArray()

            img = Extract_images.transform_img(img, NUM_CHANNELS, img_path)
            all_means += np.mean(img, axis=(1, 2))
            miterator += 1
    all_means /= miterator
    return all_means

if __name__ == "__main__":
    rng_state = np.random.get_state()
    data = "../gisat_data_tr"
    #mean_data = calc_mean(data)
    #with open("Data/mean_generated", "w") as f:
    #    for i in mean_data:
    #        f.write(str(i) + "\n")

    TRAINING_DATA_NUM = 70
    TESTING_DATA_NUM = 20
    db_size = NUM_CHANNELS * IMAGE_WIDTH * IMAGE_HEIGHT * NUM_SLICES * TRAINING_DATA_NUM * 8 + NUM_SLICES * TRAINING_DATA_NUM
    create_lmdb(data, OUTPUT_FOLDER + "gisat_data_tr.lmdb", db_size, rng_state)

    data = "../gisat_mask_tr"
    db_size = MASK_SIZE * MASK_SIZE * 1 * NUM_SLICES * TRAINING_DATA_NUM * 8 + NUM_SLICES * TRAINING_DATA_NUM
    create_lmdb(data, OUTPUT_FOLDER + "gisat_mask_tr.lmdb", db_size, rng_state)

    #data = "../testing_data"
    #db_size = NUM_CHANNELS * IMAGE_WIDTH * IMAGE_HEIGHT * NUM_SLICES * TESTING_DATA_NUM * 8 + NUM_SLICES * TESTING_DATA_NUM
    #create_lmdb(data, OUTPUT_FOLDER + "rgbgenval.lmdb", db_size, rng_state)

    #data = "../testing_mask"
    #db_size = MASK_SIZE * MASK_SIZE * 1 * NUM_SLICES * TESTING_DATA_NUM * 8 + NUM_SLICES * TESTING_DATA_NUM
    #create_lmdb(data, OUTPUT_FOLDER + "rgbgenvallabel.lmdb", db_size, rng_state)
