"""File to transform the 10K cats dataset to a set of example objects that
are saved as a pickle file and later used by the training."""
from __future__ import print_function, division
from scipy import misc, ndimage
import numpy as np
import os
import argparse
import cStringIO as StringIO
import cPickle as pickle
import math
import imgaug as ia
from common import to_aspect_ratio_add

def main():
    """Generate datasets."""

    parser = argparse.ArgumentParser(description="Process some integers.")
    parser.add_argument("--dataset_dir", help="")
    parser.add_argument("--out_dir", default="", help="")
    parser.add_argument("--img_size", default=224, type=int, help="")
    args = parser.parse_args()

    assert args.dataset_dir is not None, "Expected 10k cats dataset directory via --dataset_dir"

    # load images and their facial keypoints
    print("Loading filepaths...")
    fps = find_image_filepaths(args.dataset_dir)

    print("Loading images...")
    examples = []
    for i, fp in enumerate(fps):
        if (i+1) % 100 == 0:
            print("Image %d of %d..." % (i+1, len(fps)))
        img = ndimage.imread(fp, mode="RGB")
        kps = load_keypoints(fp, image_height=img.shape[0], image_width=img.shape[1])
        img_square, (added_top, added_right, added_bottom, added_left) = to_aspect_ratio_add(img, 1.0, return_paddings=True)
        kps = [(x+added_left, y+added_top) for (x, y) in kps]
        #rs_factor = args.img_size / img_square.shape[0]
        img_square_rs = misc.imresize(img_square, (args.img_size, args.img_size))
        kps_rs = [(int((x/img_square.shape[1])*args.img_size), int((y/img_square.shape[0])*args.img_size)) for (x, y) in kps]
        examples.append(Example(fp, img_square_rs, kps_rs))

    # save datasets
    print("Saving pickle file...")
    with open(os.path.join(args.out_dir, "cats-dataset.pkl"), "w") as f:
        pickle.dump(examples, f)

def find_image_filepaths(dataset_dir):
    """Load image filepaths from the 10k cats dataset."""
    result = []
    for root, dirs, files in os.walk(dataset_dir):
        if "/CAT_" in root:
            for name in files:
                fp = os.path.join(root, name)
                if name.endswith(".jpg") and os.path.isfile("%s.cat" % (fp,)):
                    result.append(fp)
    return result

def load_keypoints(image_filepath, image_height, image_width):
    """Load facial keypoints of one image."""
    fp_keypoints = "%s.cat" % (image_filepath,)
    if not os.path.isfile(fp_keypoints):
        raise Exception("Could not find keypoint coordinates for image '%s'." \
                        % (image_filepath,))
    else:
        coords_raw = open(fp_keypoints, "r").readlines()[0].strip().split(" ")
        coords_raw = [abs(int(coord)) for coord in coords_raw]
        keypoints = []
        #keypoints_arr = np.zeros((9*2,), dtype=np.int32)
        for i in range(1, len(coords_raw), 2): # first element is the number of coords
            x = np.clip(coords_raw[i], 0, image_width-1)
            y = np.clip(coords_raw[i+1], 0, image_height-1)
            keypoints.append((x, y))

        return keypoints

def compress_to_jpg(img):
    """Compress an image (numpy array) to jpg."""
    return compress_img(img, method="JPEG")

def compress_img(img, method):
    """Compress an image (numpy array) using the provided image compression
    method."""
    img_compressed_buffer = StringIO.StringIO()
    im = misc.toimage(img)
    im.save(img_compressed_buffer, format=method)
    img_compressed = img_compressed_buffer.getvalue()
    img_compressed_buffer.close()
    return img_compressed

def decompress_img(img_compressed):
    """Decompress a compressed image to a numpy array."""
    img_compressed_buffer = StringIO.StringIO()
    img_compressed_buffer.write(img_compressed)
    img = ndimage.imread(img_compressed_buffer, mode="RGB")
    img_compressed_buffer.close()
    return img

class Example(object):
    """Class that encapsulates an example from the 10k cats dataset.
    It contains an example image (padded and resized) for the training
    and the facial keypoints."""

    def __init__(self, fp, image, keypoints):
        self.fp = fp
        # save images jpg-compressed to save memory
        self._image = compress_to_jpg(image)
        self.keypoints = keypoints

    @property
    def image(self):
        """Return the decompressed image of the example (i.e. numpy array)."""
        return decompress_img(self._image)

    def get_bb_coords_keypoints(self, augseq=None):
        """Get the coordinates of the bounding box appropiate for the
        cat's face."""
        # augment the coordinates if an augmentation sequence is provided
        if augseq is not None:
            keypoints = ia.KeypointsOnImage([], shape=self.image.shape)
            for (x, y) in self.keypoints:
                keypoints.keypoints.append(ia.Keypoint(x=x, y=y))
            keypoints_aug = augseq.augment_keypoints([keypoints])[0]
            keypoints_aug = [(kp.x, kp.y) for kp in keypoints_aug.keypoints]
        else:
            keypoints_aug = self.keypoints

        # the facial keypoints include eyes, nose and ears.
        # drawing a bounding box around them misses the area below the nose.
        # so we extend the BB based on the vectors pointing from the eye to the
        # nose (and extending these a bit to generate new pseudo keypoints).
        left_eye = keypoints_aug[0]
        right_eye = keypoints_aug[1]
        mouth = keypoints_aug[2]

        left_eye_vec = (left_eye[0] - mouth[0], left_eye[1] - mouth[1])
        right_eye_vec = (right_eye[0] - mouth[0], right_eye[1] - mouth[1])

        l = 0.5
        le_point = (int(mouth[0] - l*left_eye_vec[0]), int(mouth[1] - l*left_eye_vec[1]))
        re_point = (int(mouth[0] - l*right_eye_vec[0]), int(mouth[1] - l*right_eye_vec[1]))

        keypoints_aug = keypoints_aug + [le_point, re_point]

        # draw bounding box around facial keypoints and pseudo
        # keypoints (i.e. mouth area)
        bb_x1 = min([x for (x, y) in keypoints_aug])
        bb_x2 = max([x for (x, y) in keypoints_aug])
        bb_y1 = min([y for (x, y) in keypoints_aug])
        bb_y2 = max([y for (x, y) in keypoints_aug])
        #return {"x1": bb_x1, "x2": bb_x2, "y1": bb_y1, "y2": bb_y2}
        return [
            ia.Keypoint(x=bb_x1, y=bb_y1),
            ia.Keypoint(x=bb_x2, y=bb_y1),
            ia.Keypoint(x=bb_x2, y=bb_y2),
            ia.Keypoint(x=bb_x1, y=bb_y2)
        ]

    #def get_bb_coords(self, augseq=None):
    #    kps = self.get_bb_coords_keypoints(augseq=augseq)
    #    return (kps[0].x, kps[0].y, kps[2].x, kps[2].y)

if __name__ == "__main__":
    main()
