"""
Mask R-CNN
Configurations and data loading code for the synthetic Shapes dataset.
This is a duplicate of the code in the noteobook train_shapes.ipynb for easy
import into other notebooks, such as inspect_model.ipynb.

Copyright (c) 2017 Matterport, Inc.
Licensed under the MIT License (see LICENSE for details)
Written by Waleed Abdulla
"""

import os
import glob
import sys
import math
import random
import numpy as np
import cv2
import imageio
import matplotlib.pyplot as plt

# Root directory of the project
ROOT_DIR = os.path.abspath("../../")

from pycocotools.coco import COCO
from pycocotools.cocoeval import COCOeval
from pycocotools import mask as maskUtils


# Import Mask RCNN
sys.path.append(ROOT_DIR)  # To find local version of the library
from mrcnn.config import Config
from mrcnn import utils

# Path to trained weights file
COCO_MODEL_PATH = os.path.join(ROOT_DIR, "mask_rcnn_coco.h5")


class InteriorNetConfig(Config):
    """Configuration for training on the InteriorNet dataset.
    Derives from the base Config class and overrides values specific
    to the InterioNet dataset.
    """
    # Give the configuration a recognizable name
    NAME = "InteriorNet"

    # Train on 1 GPU and 8 images per GPU. We can put multiple images on each
    # GPU because the images are small. Batch size is 8 (GPUs * images/GPU).
    GPU_COUNT = 1
    IMAGES_PER_GPU = 2

    # Number of classes (including background)
    NUM_CLASSES = 1 + 40  # background + 40 classes

    # Use small images for faster training. Set the limits of the small side
    # the large side, and that determines the image shape.
    IMAGE_MIN_DIM = 640
    IMAGE_MAX_DIM = 640


class InteriorDataset(utils.Dataset):
    
    
    def load_Interior(self, dataset_dir, subset, class_ids=None, return_coco=False):
        """Load a subset of the InteriorNet dataset.
        dataset_dir: The root directory of the InteriorNet dataset.
        subset: What to load (train, test, val
        class_ids: If provided, only loads images that have the given classes.
        return_coco: If True, returns the COCO object.
        """
        dataset_dir = "{}/{}".format(dataset_dir, subset)
        self.dataset_dir = dataset_dir
        
        # Iterate trough all files in the subset folder
        image_dirs = glob.iglob(os.path.join(dataset_dir,'*'))
        for image_dir in image_dirs:
            # Select current cocolabe.json file
            coco = COCO("{}/cocolabel.json".format(image_dir))
            head, tail = os.path.split(image_dir)
           
            if not class_ids:
                # All classes
                class_ids = sorted(coco.getCatIds())
                print(class_ids)

            # All images or a subset?
            if class_ids:
                image_ids = []
                for id in class_ids:
                    image_ids.extend(list(coco.getImgIds(catIds=[id])))
                # Remove duplicates
                image_ids = list(set(image_ids))
            else:
                # All images
                image_ids = list(coco.imgs.keys())

            # Add classes
            for i in class_ids:
                self.add_class("interior", i, coco.loadCats(i)[0]["name"])

            # Add images
            for i in image_ids:
                id = tail + '_id' + str(i)
                self.add_image(
                    "interior", image_id=id,
                    image_sub_id = i,
                    path=os.path.join(image_dir, coco.imgs[i]['file_name']),
                    width=coco.imgs[i]["width"],
                    height=coco.imgs[i]["height"],
                    annotations=coco.loadAnns(coco.getAnnIds(
                        imgIds=[i], catIds=class_ids, iscrowd=None)),
                    subfolder=tail)
        if return_coco:
            return coco


    def image_reference(self, image_id):
        """Return the path of the image"""
        
        info = self.image_info[image_id]
        if info["source"] == "interior":
            return info["path"]
        else:
            super(self.__class__, self).image_reference(image_id)
        
    def load_mask(self, image_id):
        """Load instance masks for the given image.

        Different datasets use different ways to store masks. This
        function converts the different mask format to one format
        in the form of a bitmap [height, width, instances].

        Returns:
        masks: A bool array of shape [height, width, instance count] with
            one mask per instance.
        class_ids: a 1D array of class IDs of the instance masks.
        """
        # If not an InteriorNet image, delegate to parent class.
        image_info = self.image_info[image_id]
        if image_info["source"] != "interior":
            return super(self.__class__, self).load_mask(image_id)

        instance_masks = []
        class_ids = []
        path = self.image_info[image_id]["path"]
        # Build mask of shape [height, width, instance_count] and list
        # of class IDs that correspond to each channel of the mask.
        instance_mask_path = os.path.join(self.dataset_dir, self.image_info[image_id]["subfolder"], 
                                   'label0/data',
                                   str(self.image_info[image_id]["image_sub_id"])+'_instance.png')
        nyu_mask_path = os.path.join(self.dataset_dir, self.image_info[image_id]["subfolder"], 
                                   'label0/data',
                                   str(self.image_info[image_id]["image_sub_id"])+'_nyu.png')
        instance_im = imageio.imread(instance_mask_path)
        nyu_im = imageio.imread(nyu_mask_path)
        instance_ids = np.unique(instance_im)
        for instance_id in instance_ids:
            binary_mask = np.where(instance_im == instance_id , True, False)
            class_id = nyu_im[binary_mask]
            class_id = np.unique(class_id)
            if len(class_id==1):
                class_ids.append(class_id[0])
                instance_masks.append(binary_mask)
            else:
                print('Multiple different classes inside one instance mask.')

        # Pack instance masks into an array
        if class_ids:
            mask = np.stack(instance_masks, axis=2).astype(np.bool)
            class_ids = np.array(class_ids, dtype=np.int32)
            return mask, class_ids
        else:
            # Call super class to return an empty mask
            return super(CocoDataset, self).load_mask(image_id)

    
        
'''
############################################################
#  Training
############################################################


if __name__ == '__main__':
    import argparse

    # Parse command line arguments
    parser = argparse.ArgumentParser(
        description='Train Mask R-CNN on InteriorNet.')
    parser.add_argument("command",
                        metavar="<command>",
                        help="'train' or 'evaluate' on InteriorNet")
    parser.add_argument('--dataset', required=True,
                        metavar="/path/to/intriornet/",
                        help='Directory of the InterioNet dataset')
    parser.add_argument('--model', required=True,
                        metavar="/path/to/weights.h5",
                        help="Path to weights .h5 file")
    parser.add_argument('--logs', required=False,
                        default=DEFAULT_LOGS_DIR,
                        metavar="/path/to/logs/",
                        help='Logs and checkpoints directory (default=logs/)')
    parser.add_argument('--limit', required=False,
                        default=500,
                        metavar="<image count>",
                        help='Images to use for evaluation (default=500)')

    args = parser.parse_args()
    print("Command: ", args.command)
    print("Model: ", args.model)
    print("Dataset: ", args.dataset)
    print("Logs: ", args.logs)

    # Configurations
    if args.command == "train":
        config = InteriorNetConfig()
    else:
        class InferenceConfig(InteriorNetConfig):
            # Set batch size to 1 since we'll be running inference on
            # one image at a time. Batch size = GPU_COUNT * IMAGES_PER_GPU
            GPU_COUNT = 1
            IMAGES_PER_GPU = 1
            DETECTION_MIN_CONFIDENCE = 0
        config = InferenceConfig()
    config.display()

    # Create model
    if args.command == "train":
        model = modellib.MaskRCNN(mode="training", config=config,
                                  model_dir=args.logs)
    else:
        model = modellib.MaskRCNN(mode="inference", config=config,
                                  model_dir=args.logs)

    # Select weights file to load
    if args.model.lower() == "coco":
        model_path = COCO_MODEL_PATH
    elif args.model.lower() == "last":
        # Find last trained weights
        model_path = model.find_last()
    elif args.model.lower() == "imagenet":
        # Start from ImageNet trained weights
        model_path = model.get_imagenet_weights()
    else:
        model_path = args.model

    # Load weights
    print("Loading weights ", model_path)
    model.load_weights(model_path, by_name=True)

    # Train or evaluate
    if args.command == "train":
        # Training dataset. Use the training set and 35K from the
        # validation set, as as in the Mask RCNN paper.
        dataset_train = InteriorDataset()
        dataset_train.load_coco(args.dataset, "train", year=args.year, auto_download=args.download)
        if args.year in '2014':
            dataset_train.load_coco(args.dataset, "valminusminival", year=args.year, auto_download=args.download)
        dataset_train.prepare()

        # Validation dataset
        dataset_val = CocoDataset()
        val_type = "val" if args.year in '2017' else "minival"
        dataset_val.load_coco(args.dataset, val_type, year=args.year, auto_download=args.download)
        dataset_val.prepare()

        # Image Augmentation
        # Right/Left flip 50% of the time
        augmentation = imgaug.augmenters.Fliplr(0.5)

        # *** This training schedule is an example. Update to your needs ***

        # Training - Stage 1
        print("Training network heads")
        model.train(dataset_train, dataset_val,
                    learning_rate=config.LEARNING_RATE,
                    epochs=40,
                    layers='heads',
                    augmentation=augmentation)

        # Training - Stage 2
        # Finetune layers from ResNet stage 4 and up
        print("Fine tune Resnet stage 4 and up")
        model.train(dataset_train, dataset_val,
                    learning_rate=config.LEARNING_RATE,
                    epochs=120,
                    layers='4+',
                    augmentation=augmentation)

        # Training - Stage 3
        # Fine tune all layers
        print("Fine tune all layers")
        model.train(dataset_train, dataset_val,
                    learning_rate=config.LEARNING_RATE / 10,
                    epochs=160,
                    layers='all',
                    augmentation=augmentation)

    elif args.command == "evaluate":
        # Validation dataset
        dataset_val = CocoDataset()
        val_type = "val" if args.year in '2017' else "minival"
        coco = dataset_val.load_coco(args.dataset, val_type, year=args.year, return_coco=True, auto_download=args.download)
        dataset_val.prepare()
        print("Running COCO evaluation on {} images.".format(args.limit))
        evaluate_coco(model, dataset_val, coco, "bbox", limit=int(args.limit))
    else:
        print("'{}' is not recognized. "
              "Use 'train' or 'evaluate'".format(args.command))

'''