"""
Copyright 2017-2018 lvaleriu (https://github.com/lvaleriu/)

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

import csv
import json
import os
import warnings

import numpy as np
from PIL import Image

from .generator import Generator
from ..utils.image import read_image_bgr


def load_hierarchy(metadata_dir, version='v4'):
    hierarchy = None
    if version == 'challenge2018':
        hierarchy = 'bbox_labels_500_hierarchy.json'
    elif version == 'v4':
        hierarchy = 'bbox_labels_600_hierarchy.json'
    elif version == 'v3':
        hierarchy = 'bbox_labels_600_hierarchy.json'

    hierarchy_json = os.path.join(metadata_dir, hierarchy)
    with open(hierarchy_json) as f:
        hierarchy_data = json.loads(f.read())

    return hierarchy_data


def load_hierarchy_children(hierarchy):
    res = [hierarchy['LabelName']]

    if 'Subcategory' in hierarchy:
        for subcategory in hierarchy['Subcategory']:
            children = load_hierarchy_children(subcategory)

            for c in children:
                res.append(c)

    return res


def find_hierarchy_parent(hierarchy, parent_cls):
    if hierarchy['LabelName'] == parent_cls:
        return hierarchy
    elif 'Subcategory' in hierarchy:
        for child in hierarchy['Subcategory']:
            res = find_hierarchy_parent(child, parent_cls)
            if res is not None:
                return res

    return None


def get_labels(metadata_dir, version='v4'):
    if version == 'v4' or version == 'challenge2018':
        csv_file = 'class-descriptions-boxable.csv' if version == 'v4' else 'challenge-2018-class-descriptions-500.csv'

        boxable_classes_descriptions = os.path.join(metadata_dir, csv_file)
        id_to_labels = {}
        cls_index    = {}

        i = 0
        with open(boxable_classes_descriptions) as f:
            for row in csv.reader(f):
                # make sure the csv row is not empty (usually the last one)
                if len(row):
                    label       = row[0]
                    description = row[1].replace("\"", "").replace("'", "").replace('`', '')

                    id_to_labels[i]  = description
                    cls_index[label] = i

                    i += 1
    else:
        trainable_classes_path = os.path.join(metadata_dir, 'classes-bbox-trainable.txt')
        description_path = os.path.join(metadata_dir, 'class-descriptions.csv')

        description_table = {}
        with open(description_path) as f:
            for row in csv.reader(f):
                # make sure the csv row is not empty (usually the last one)
                if len(row):
                    description_table[row[0]] = row[1].replace("\"", "").replace("'", "").replace('`', '')

        with open(trainable_classes_path, 'rb') as f:
            trainable_classes = f.read().split('\n')

        id_to_labels = dict([(i, description_table[c]) for i, c in enumerate(trainable_classes)])
        cls_index = dict([(c, i) for i, c in enumerate(trainable_classes)])

    return id_to_labels, cls_index


def generate_images_annotations_json(main_dir, metadata_dir, subset, cls_index, version='v4'):
    validation_image_ids = {}

    if version == 'v4':
        annotations_path = os.path.join(metadata_dir, subset, '{}-annotations-bbox.csv'.format(subset))
    elif version == 'challenge2018':
        validation_image_ids_path = os.path.join(metadata_dir, 'challenge-2018-image-ids-valset-od.csv')

        with open(validation_image_ids_path, 'r') as csv_file:
            reader = csv.DictReader(csv_file, fieldnames=['ImageID'])
            next(reader)
            for line, row in enumerate(reader):
                image_id = row['ImageID']
                validation_image_ids[image_id] = True

        annotations_path = os.path.join(metadata_dir, 'challenge-2018-train-annotations-bbox.csv')
    else:
        annotations_path = os.path.join(metadata_dir, subset, 'annotations-human-bbox.csv')

    fieldnames = ['ImageID', 'Source', 'LabelName', 'Confidence',
                  'XMin', 'XMax', 'YMin', 'YMax',
                  'IsOccluded', 'IsTruncated', 'IsGroupOf', 'IsDepiction', 'IsInside']

    id_annotations = dict()
    with open(annotations_path, 'r') as csv_file:
        reader = csv.DictReader(csv_file, fieldnames=fieldnames)
        next(reader)

        images_sizes = {}
        for line, row in enumerate(reader):
            frame = row['ImageID']

            if version == 'challenge2018':
                if subset == 'train':
                    if frame in validation_image_ids:
                        continue
                elif subset == 'validation':
                    if frame not in validation_image_ids:
                        continue
                else:
                    raise NotImplementedError('This generator handles only the train and validation subsets')

            class_name = row['LabelName']

            if class_name not in cls_index:
                continue

            cls_id = cls_index[class_name]

            if version == 'challenge2018':
                # We recommend participants to use the provided subset of the training set as a validation set.
                # This is preferable over using the V4 val/test sets, as the training set is more densely annotated.
                img_path = os.path.join(main_dir, 'images', 'train', frame + '.jpg')
            else:
                img_path = os.path.join(main_dir, 'images', subset, frame + '.jpg')

            if frame in images_sizes:
                width, height = images_sizes[frame]
            else:
                try:
                    with Image.open(img_path) as img:
                        width, height = img.width, img.height
                        images_sizes[frame] = (width, height)
                except Exception as ex:
                    if version == 'challenge2018':
                        raise ex
                    continue

            x1 = float(row['XMin'])
            x2 = float(row['XMax'])
            y1 = float(row['YMin'])
            y2 = float(row['YMax'])

            x1_int = int(round(x1 * width))
            x2_int = int(round(x2 * width))
            y1_int = int(round(y1 * height))
            y2_int = int(round(y2 * height))

            # Check that the bounding box is valid.
            if x2 <= x1:
                raise ValueError('line {}: x2 ({}) must be higher than x1 ({})'.format(line, x2, x1))
            if y2 <= y1:
                raise ValueError('line {}: y2 ({}) must be higher than y1 ({})'.format(line, y2, y1))

            if y2_int == y1_int:
                warnings.warn('filtering line {}: rounding y2 ({}) and y1 ({}) makes them equal'.format(line, y2, y1))
                continue

            if x2_int == x1_int:
                warnings.warn('filtering line {}: rounding x2 ({}) and x1 ({}) makes them equal'.format(line, x2, x1))
                continue

            img_id = row['ImageID']
            annotation = {'cls_id': cls_id, 'x1': x1, 'x2': x2, 'y1': y1, 'y2': y2}

            if img_id in id_annotations:
                annotations = id_annotations[img_id]
                annotations['boxes'].append(annotation)
            else:
                id_annotations[img_id] = {'w': width, 'h': height, 'boxes': [annotation]}
    return id_annotations


class OpenImagesGenerator(Generator):
    def __init__(
            self, main_dir, subset, version='v4',
            labels_filter=None, annotation_cache_dir='.',
            parent_label=None,
            **kwargs
    ):
        if version == 'challenge2018':
            metadata = 'challenge2018'
        elif version == 'v4':
            metadata = '2018_04'
        elif version == 'v3':
            metadata = '2017_11'
        else:
            raise NotImplementedError('There is currently no implementation for versions older than v3')

        if version == 'challenge2018':
            self.base_dir     = os.path.join(main_dir, 'images', 'train')
        else:
            self.base_dir     = os.path.join(main_dir, 'images', subset)

        metadata_dir          = os.path.join(main_dir, metadata)
        annotation_cache_json = os.path.join(annotation_cache_dir, subset + '.json')

        self.hierarchy          = load_hierarchy(metadata_dir, version=version)
        id_to_labels, cls_index = get_labels(metadata_dir, version=version)

        if os.path.exists(annotation_cache_json):
            with open(annotation_cache_json, 'r') as f:
                self.annotations = json.loads(f.read())
        else:
            self.annotations = generate_images_annotations_json(main_dir, metadata_dir, subset, cls_index, version=version)
            json.dump(self.annotations, open(annotation_cache_json, "w"))

        if labels_filter is not None or parent_label is not None:
            self.id_to_labels, self.annotations = self.__filter_data(id_to_labels, cls_index, labels_filter, parent_label)
        else:
            self.id_to_labels = id_to_labels

        self.id_to_image_id = dict([(i, k) for i, k in enumerate(self.annotations)])

        super(OpenImagesGenerator, self).__init__(**kwargs)

    def __filter_data(self, id_to_labels, cls_index, labels_filter=None, parent_label=None):
        """
        If you want to work with a subset of the labels just set a list with trainable labels
        :param labels_filter: Ex: labels_filter = ['Helmet', 'Hat', 'Analog television']
        :param parent_label: If parent_label is set this will bring you the parent label
        but also its children in the semantic hierarchy as defined in OID, ex: Animal
        hierarchical tree
        :return:
        """

        children_id_to_labels = {}

        if parent_label is None:
            # there is/are no other sublabel(s) other than the labels itself

            for label in labels_filter:
                for i, lb in id_to_labels.items():
                    if lb == label:
                        children_id_to_labels[i] = label
                        break
        else:
            parent_cls = None
            for i, lb in id_to_labels.items():
                if lb == parent_label:
                    parent_id = i
                    for c, index in cls_index.items():
                        if index == parent_id:
                            parent_cls = c
                    break

            if parent_cls is None:
                raise Exception('Couldnt find label {}'.format(parent_label))

            parent_tree = find_hierarchy_parent(self.hierarchy, parent_cls)

            if parent_tree is None:
                raise Exception('Couldnt find parent {} in the semantic hierarchical tree'.format(parent_label))

            children = load_hierarchy_children(parent_tree)

            for cls in children:
                index = cls_index[cls]
                label = id_to_labels[index]
                children_id_to_labels[index] = label

        id_map = dict([(ind, i) for i, ind in enumerate(children_id_to_labels.keys())])

        filtered_annotations = {}
        for k in self.annotations:
            img_ann = self.annotations[k]

            filtered_boxes = []
            for ann in img_ann['boxes']:
                cls_id = ann['cls_id']
                if cls_id in children_id_to_labels:
                    ann['cls_id'] = id_map[cls_id]
                    filtered_boxes.append(ann)

            if len(filtered_boxes) > 0:
                filtered_annotations[k] = {'w': img_ann['w'], 'h': img_ann['h'], 'boxes': filtered_boxes}

        children_id_to_labels = dict([(id_map[i], l) for (i, l) in children_id_to_labels.items()])

        return children_id_to_labels, filtered_annotations

    def size(self):
        return len(self.annotations)

    def num_classes(self):
        return len(self.id_to_labels)

    def has_label(self, label):
        """ Return True if label is a known label.
        """
        return label in self.id_to_labels

    def has_name(self, name):
        """ Returns True if name is a known class.
        """
        raise NotImplementedError()

    def name_to_label(self, name):
        raise NotImplementedError()

    def label_to_name(self, label):
        return self.id_to_labels[label]

    def image_aspect_ratio(self, image_index):
        img_annotations = self.annotations[self.id_to_image_id[image_index]]
        height, width = img_annotations['h'], img_annotations['w']
        return float(width) / float(height)

    def image_path(self, image_index):
        path = os.path.join(self.base_dir, self.id_to_image_id[image_index] + '.jpg')
        return path

    def load_image(self, image_index):
        return read_image_bgr(self.image_path(image_index))

    def load_annotations(self, image_index):
        image_annotations = self.annotations[self.id_to_image_id[image_index]]

        labels = image_annotations['boxes']
        height, width = image_annotations['h'], image_annotations['w']

        annotations = {'labels': np.empty((len(labels),)), 'bboxes': np.empty((len(labels), 4))}
        for idx, ann in enumerate(labels):
            cls_id = ann['cls_id']
            x1 = ann['x1'] * width
            x2 = ann['x2'] * width
            y1 = ann['y1'] * height
            y2 = ann['y2'] * height

            annotations['bboxes'][idx, 0] = x1
            annotations['bboxes'][idx, 1] = y1
            annotations['bboxes'][idx, 2] = x2
            annotations['bboxes'][idx, 3] = y2
            annotations['labels'][idx] = cls_id

        return annotations