# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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 os
import numpy as np
from imdb import Imdb


class YoloFormat(Imdb):
    """
    Base class for loading datasets as used in YOLO

    Parameters:
    ----------
    name : str
        name for this dataset
    classes : list or tuple of str
        class names in this dataset
    list_file : str
        filename of the image list file
    image_dir : str
        image directory
    label_dir : str
        label directory
    extension : str
        by default .jpg
    label_extension : str
        by default .txt
    shuffle : bool
        whether to shuffle the initial order when loading this dataset,
        default is True
    """
    def __init__(self, name, classes, list_file, image_dir, label_dir, \
                 extension='.jpg', label_extension='.txt', shuffle=True):
        if isinstance(classes, list) or isinstance(classes, tuple):
            num_classes = len(classes)
        elif isinstance(classes, str):
            with open(classes, 'r') as f:
                classes = [l.strip() for l in f.readlines()]
                num_classes = len(classes)
        else:
            raise ValueError("classes should be list/tuple or text file")
        assert num_classes > 0, "number of classes must > 0"
        super(YoloFormat, self).__init__(name + '_' + str(num_classes))
        self.classes = classes
        self.num_classes = num_classes
        self.list_file = list_file
        self.image_dir = image_dir
        self.label_dir = label_dir
        self.extension = extension
        self.label_extension = label_extension

        self.image_set_index = self._load_image_set_index(shuffle)
        self.num_images = len(self.image_set_index)
        self.labels = self._load_image_labels()


    def _load_image_set_index(self, shuffle):
        """
        find out which indexes correspond to given image set (train or val)

        Parameters:
        ----------
        shuffle : boolean
            whether to shuffle the image list
        Returns:
        ----------
        entire list of images specified in the setting
        """
        assert os.path.exists(self.list_file), 'Path does not exists: {}'.format(self.list_file)
        with open(self.list_file, 'r') as f:
            image_set_index = [x.strip() for x in f.readlines()]
        if shuffle:
            np.random.shuffle(image_set_index)
        return image_set_index

    def image_path_from_index(self, index):
        """
        given image index, find out full path

        Parameters:
        ----------
        index: int
            index of a specific image
        Returns:
        ----------
        full path of this image
        """
        assert self.image_set_index is not None, "Dataset not initialized"
        name = self.image_set_index[index]
        image_file = os.path.join(self.image_dir, name) + self.extension
        assert os.path.exists(image_file), 'Path does not exist: {}'.format(image_file)
        return image_file

    def label_from_index(self, index):
        """
        given image index, return preprocessed ground-truth

        Parameters:
        ----------
        index: int
            index of a specific image
        Returns:
        ----------
        ground-truths of this image
        """
        assert self.labels is not None, "Labels not processed"
        return self.labels[index]

    def _label_path_from_index(self, index):
        """
        given image index, find out annotation path

        Parameters:
        ----------
        index: int
            index of a specific image

        Returns:
        ----------
        full path of annotation file
        """
        label_file = os.path.join(self.label_dir, index + self.label_extension)
        assert os.path.exists(label_file), 'Path does not exist: {}'.format(label_file)
        return label_file

    def _load_image_labels(self):
        """
        preprocess all ground-truths

        Returns:
        ----------
        labels packed in [num_images x max_num_objects x 5] tensor
        """
        temp = []

        # load ground-truths
        for idx in self.image_set_index:
            label_file = self._label_path_from_index(idx)
            with open(label_file, 'r') as f:
                label = []
                for line in f.readlines():
                    temp_label = line.strip().split()
                    assert len(temp_label) == 5, "Invalid label file" + label_file
                    cls_id = int(temp_label[0])
                    x = float(temp_label[1])
                    y = float(temp_label[2])
                    half_width = float(temp_label[3]) / 2
                    half_height = float(temp_label[4]) / 2
                    xmin = x - half_width
                    ymin = y - half_height
                    xmax = x + half_width
                    ymax = y + half_height
                    label.append([cls_id, xmin, ymin, xmax, ymax])
                temp.append(np.array(label))
        return temp
