#!/usr/bin/env python
from glob import glob
import argparse
import cv2
import numpy as np
import os

# modifiable constants
""""
---------- MODIFIABLE CONSTANTS ----------
IMG_PATH[0]: Training images (Inputs to the Neural Network)
IMG_PATH[1]: Ground truth of training images
OUT_PATH: Directory where processed images will be made
TXT_PATH: 


"""
#IMG_PATH = [
#    os.path.expanduser("/home/pganti/datasets/cityscapes/leftImg8bit/val/*/*.png"),
#    os.path.expanduser("/home/pganti/datasets/cityscapes/gtFine/val/*/*gtFine_labelTrainIds.png")]
#OUT_PATH = os.path.expanduser("/home/pganti/datasets/bayesian_segnet_data/")

#IMG_PATH = [
#    os.path.expanduser("/home/pganti/datasets/kitti/segmentation/training/image_2/val/*.png"),
#    os.path.expanduser("/home/pganti/datasets/kitti/segmentation/training/semantic_labeled/val/*labelTrainIds.png")]
#OUT_PATH = os.path.expanduser("/home/pganti/datasets/kitti_bayesian_segnet/")

IMG_PATH = [
    os.path.expanduser("/home/pganti/Downloads/dataset_watF/image_2/*.png")
]
OUT_PATH = os.path.expanduser("/home/pganti/datasets/waterloo_f/odometry/sequences/00/image_2/")

#IMG_PATH = [
#    os.path.expanduser("/home/pganti/Downloads/dataset_watF/image_3/*.png")
#]
#OUT_PATH = os.path.expanduser("/home/pganti/datasets/waterloo_f/odometry/sequences/00/image_3/

DATA_TYPE = "test"
TXT_PATH = ""
RESIZE_IMGS = True
WIDTH, HEIGHT = 1024, 352
INTERPOLATION = [cv2.INTER_CUBIC, cv2.INTER_NEAREST, cv2.INTER_AREA]
CROP_TO_ASPECT_RATIO = True
CROP_HEIGHT_POSITION = 'middle'
CROP_WIDTH_POSITION = 'middle'

def get_args(arg_vals=None):
    """
    parse any relevant runtime arguments
    :param arg_vals:    args to check (None defaults to command-line args)
    :return:            argparse.Namespace object
    """

    global WIDTH, HEIGHT, IMG_PATH, OUT_PATH, DATA_TYPE, TXT_PATH, RESIZE_IMGS
    parser = argparse.ArgumentParser(
        description="analyze SegNet output vs ground truth using \
            global accuracy, mean class accuracy, mIoU")
    parser.add_argument("--width", default=WIDTH,
                        help="desired width for images")
    parser.add_argument("--height", default=HEIGHT,
                        help="desired height for images")
    parser.add_argument("-iim", "--in_img_dir", default=IMG_PATH[0],
                        help="input path for camera images (default is {})"
                        .format(IMG_PATH[0]))
    parser.add_argument("-igt", "--in_gt_dir", default=IMG_PATH[0],
                        help="input path for ground truth (default is {})"
                        .format(IMG_PATH[0]))
    parser.add_argument("-o", "--out_dir", default=OUT_PATH,
                        help="output path (default is {})"
                        .format(OUT_PATH))
    parser.add_argument("-d", "--data_type", default=DATA_TYPE,
                        help="type of data being handled (test/val/train)")
    parser.add_argument("-nr", "--noresize", action="store_false",
                        help="flag set skips images resizing")
    parser.add_argument("-ctar", "--crop_to_aspect_ratio", default=CROP_TO_ASPECT_RATIO,
                        help="flag set crops input image to the aspect ratio of specified height x width and then resizes")
    parser.add_argument("-ch", "--crop_height_position", default=CROP_HEIGHT_POSITION, choices=['top', 'middle', 'bottom'],
                        help="position to crop on height from if crops_to_aspect_ratio is set")
    parser.add_argument("-cw", "--crop_width_position", default=CROP_WIDTH_POSITION, choices=['left', 'middle', 'right'],
                        help="position to crop on width from if crops_to_aspect_ratio is set")

    ret_args = parser.parse_args(arg_vals)

    WIDTH = int(ret_args.width) - int(ret_args.width) % 32
    HEIGHT = int(ret_args.height) - int(ret_args.height) % 32
    IMG_PATH[0] = os.path.expanduser(ret_args.in_img_dir)
    #IMG_PATH[1] = os.path.expanduser(ret_args.in_gt_dir)
    OUT_PATH = ret_args.out_dir
    DATA_TYPE = ret_args.data_type
    RESIZE_IMGS = ret_args.noresize
    if not TXT_PATH:
        TXT_PATH = "{}{}.txt".format(OUT_PATH, DATA_TYPE)

    print ("width set to: {}; height set to: {}".format(WIDTH, HEIGHT))
    return ret_args


def crop_image(img, img_dim, position):
    """
    Crops the image based on the desired dimension and position to start cropping
    :param img: openCV image
    :param img_dim: (desired_height, desired_width)
    :param position: (crop_height_position, crop_width_position)
    :return:
    """
    if position[0] == "top":
        h_t = img.shape[0] - img_dim[0]
        h_b = None
    elif position[0] == "middle":
        h_t = int((img.shape[0] - img_dim[0]) / 2)
        h_b = h_t + img_dim[0]
    elif position[0] == "bottom":
        h_t = None
        h_b = img_dim[0]
    else:
        raise ValueError("Invalid position, expecting (top, middle, bottom)")

    if position[1] == "left":
        w_l = img.shape[1] - img_dim[1]
        w_r = None
    elif position[1] == "middle":
        w_l = int((img.shape[1] - img_dim[1]) / 2)
        w_r = w_l + img_dim[1]
    elif position[1] == "right":
        w_l = None
        w_r = img_dim[1]
    else:
        raise ValueError("Invalid position, expecting (left, middle, right)")

    return img[h_t:h_b, w_l:w_r]


def main():
    get_args()

    # get lists of files
    images = [sorted(glob(IMG_PATH[0]))]#, sorted(glob(IMG_PATH[1]))]

    # check for folders' integrity
    assert len(images[0]) > 0, "no images in: %s" % (IMG_PATH[0])
    #assert len(images[1]) > 0, "no images in: %s" % (IMG_PATH[1])
    #assert len(images[0]) == len(images[1]), "number of images mismatch"

    # determine image size
    src_shape = cv2.imread(images[0][0]).shape
    
    """
    for image in np.array(images).flatten():
        assert cv2.imread(image).shape == src_shape, \
            "size mismatch: %s" % image
    """
    
    assert preprocessor(src_shape, images[0]), \
        "images not resized"


def preprocessor(src_shape, imgs):
    # verify output directories
    out_path = [os.path.join(OUT_PATH, "test/"),
                os.path.join(OUT_PATH, "testannot/")]
    if DATA_TYPE == "train":
        out_path = [os.path.join(OUT_PATH, "train/"),
                    os.path.join(OUT_PATH, "trainannot/")]
    elif DATA_TYPE == "val":
        out_path = [os.path.join(OUT_PATH, "val/"),
                    os.path.join(OUT_PATH, "valannot/")]
    #for i in range(2):
        #assert os.path.isdir(out_path[i]), \
        #    "non-existent output directory {}".format(out_path[i])
    assert os.path.isdir(OUT_PATH), \
        "non-existent output directory {}".format(OUT_PATH)


    dst_shape = (HEIGHT, WIDTH, 3)
    resize_imgs = RESIZE_IMGS and src_shape != dst_shape
    if CROP_TO_ASPECT_RATIO:
        height_multiple = int(src_shape[0] / dst_shape[0])
        width_multiple = int(src_shape[1] / dst_shape[1])
        min_multiple = min(height_multiple, width_multiple)
        height_crop_dim = dst_shape[0] * min_multiple
        width_crop_dim = dst_shape[1] * min_multiple

    count, freqs, class_counts = 0, np.zeros(256), np.zeros(256, np.int)

    with open(TXT_PATH, "w") as test_file:
        print TXT_PATH
        for image0 in imgs:
            path = ""

            img = cv2.imread(image0)
            if CROP_TO_ASPECT_RATIO:
                img = crop_image(img, (height_crop_dim, width_crop_dim), (CROP_HEIGHT_POSITION, CROP_WIDTH_POSITION))

            if resize_imgs:
                img = cv2.resize(img, (WIDTH, HEIGHT),
                                        interpolation=INTERPOLATION[0])


            path = OUT_PATH + os.path.basename(image0)
            cv2.imwrite(path, img)

            # write test.txt for SegNet
            test_file.write("%s\n" % (path))
            count += 1
            print("%d: %s" % (count, path))

    if DATA_TYPE == "train":
        n_classes = np.nonzero(class_counts)
        freqs = freqs[n_classes] / class_counts[n_classes]
        freqs = np.median(freqs) / freqs
        print("softmax class weights:\n{}".format(freqs))

    return resize_imgs


if __name__ == "__main__":
    try:
        main()
    except AssertionError as err:
        print(err.args[0])
