"""
@author:    wangguanan
@contact:   guan.wang0706@gmail.com
"""

import os, copy
from .base import ReIDSamples


class DukeMTMCreID(ReIDSamples):
    """DukeMTMC-reID.
    Reference:
        - Ristani et al. Performance Measures and a Data Set for Multi-Target, Multi-Camera Tracking. ECCVW 2016.
        - Zheng et al. Unlabeled Samples Generated by GAN Improve the Person Re-identification Baseline in vitro. ICCV 2017.
    URL: `<https://github.com/layumi/DukeMTMC-reID_evaluation>`_
    Dataset statistics:
        - identities: 1404 (train + query).
        - images:16522 (train) + 2228 (query) + 17661 (gallery).
        - cameras: 8.
    Args:
        data_path(str): path to DukeMTMC-reID dataset
        combineall(bool): combine train and test sets as train set if True
    """

    def __init__(self, data_path, combineall=False):
        super(DukeMTMCreID, self).__init__()

        # parameters
        self.duke_path = data_path
        self.combineall = combineall

        # paths of train, query and gallery
        train_path = os.path.join(self.duke_path, 'bounding_box_train/')
        query_path = os.path.join(self.duke_path, 'query/')
        gallery_path = os.path.join(self.duke_path, 'bounding_box_test/')

        # load
        train = self._load_samples(train_path)
        query = self._load_samples(query_path)
        gallery = self._load_samples(gallery_path)
        if self.combineall:
            train += copy.deepcopy(query) + copy.deepcopy(gallery)
        train = self.relabel(train)
        self.statistics(train, query, gallery)

        # return
        self.train, self.query, self.gallery = train, query, gallery

    def _load_samples(self, folder_dir):
        '''return (img_path, identity_id, camera_id)'''
        samples = []
        root_path, _, files_name = self.os_walk(folder_dir)
        for file_name in files_name:
            if '.jpg' in file_name:
                person_id, camera_id = self._analysis_file_name(file_name)
                samples.append([root_path + file_name, person_id, camera_id])
        return samples

    def _analysis_file_name(self, file_name):
        '''
        :param file_name: format like 0002_c1_f0044158.jpg
        :return:
        '''
        split_list = file_name.replace('.jpg', '').replace('c', '').split('_')
        person_id, camera_id = int(split_list[0]), int(split_list[1])
        return person_id, camera_id
