# Source code for torchreid.data.datasets.video.mars

    from __future__ import absolute_import
    from __future__ import print_function
    from __future__ import division

    import sys
    import os
    import os.path as osp
    from scipy.io import loadmat
    import warnings

    from torchreid.data.datasets import VideoDataset

    class Mars(VideoDataset):
        """MARS.

        Reference:
            Zheng et al. MARS: A Video Benchmark for Large-Scale Person Re-identification. ECCV 2016.

        URL: `<http://www.liangzheng.com.cn/Project/project_mars.html>`_

        Dataset statistics:
            - identities: 1261.
            - tracklets: 8298 (train) + 1980 (query) + 9330 (gallery).
            - cameras: 6.
        """
        dataset_dir = 'mars'
        dataset_url = None

        def __init__(self, root='', **kwargs):
            self.root = osp.abspath(osp.expanduser(root))
            self.dataset_dir = osp.join(self.root, self.dataset_dir)
            self.download_dataset(self.dataset_dir, self.dataset_url)

            self.train_name_path = osp.join(self.dataset_dir, 'info/train_name.txt')
            self.test_name_path = osp.join(self.dataset_dir, 'info/test_name.txt')
            self.track_train_info_path = osp.join(self.dataset_dir, 'info/tracks_train_info.mat')
            self.track_test_info_path = osp.join(self.dataset_dir, 'info/tracks_test_info.mat')
            self.query_IDX_path = osp.join(self.dataset_dir, 'info/query_IDX.mat')

            required_files = [
                self.dataset_dir,
                self.train_name_path,
                self.test_name_path,
                self.track_train_info_path,
                self.track_test_info_path,
                self.query_IDX_path
            ]
            self.check_before_run(required_files)

            train_names = self.get_names(self.train_name_path)
            test_names = self.get_names(self.test_name_path)
            track_train = loadmat(self.track_train_info_path)['track_train_info'] # numpy.ndarray (8298, 4)
            track_test = loadmat(self.track_test_info_path)['track_test_info'] # numpy.ndarray (12180, 4)
            query_IDX = loadmat(self.query_IDX_path)['query_IDX'].squeeze() # numpy.ndarray (1980,)
            query_IDX -= 1 # index from 0
            track_query = track_test[query_IDX,:]
            gallery_IDX = [i for i in range(track_test.shape[0]) if i not in query_IDX]
            track_gallery = track_test[gallery_IDX,:]

            train = self.process_data(train_names, track_train, home_dir='bbox_train', relabel=True)
            query = self.process_data(test_names, track_query, home_dir='bbox_test', relabel=False)
            gallery = self.process_data(test_names, track_gallery, home_dir='bbox_test', relabel=False)

            super(Mars, self).__init__(train, query, gallery, **kwargs)

        def get_names(self, fpath):
            names = []
            with open(fpath, 'r') as f:
                for line in f:
                    new_line = line.rstrip()
                    names.append(new_line)
            return names

        def process_data(self, names, meta_data, home_dir=None, relabel=False, min_seq_len=0):
            assert home_dir in ['bbox_train', 'bbox_test']
            num_tracklets = meta_data.shape[0]
            pid_list = list(set(meta_data[:,2].tolist()))
            num_pids = len(pid_list)

            if relabel: pid2label = {pid:label for label, pid in enumerate(pid_list)}
            tracklets = []

            for tracklet_idx in range(num_tracklets):
                data = meta_data[tracklet_idx,...]
                start_index, end_index, pid, camid = data
                if pid == -1:
                    continue # junk images are just ignored
                assert 1 <= camid <= 6
                if relabel: pid = pid2label[pid]
                camid -= 1 # index starts from 0
                img_names = names[start_index - 1:end_index]

                # make sure image names correspond to the same person
                pnames = [img_name[:4] for img_name in img_names]
                assert len(set(pnames)) == 1, 'Error: a single tracklet contains different person images'

                # make sure all images are captured under the same camera
                camnames = [img_name[5] for img_name in img_names]
                assert len(set(camnames)) == 1, 'Error: images are captured under different cameras!'

                # append image names with directory information
                img_paths = [osp.join(self.dataset_dir, home_dir, img_name[:4], img_name) for img_name in img_names]
                if len(img_paths) >= min_seq_len:
                    img_paths = tuple(img_paths)
                    tracklets.append((img_paths, pid, camid))

            return tracklets

        def combine_all(self):
            warnings.warn('Some query IDs do not appear in gallery. Therefore, combineall '
                          'does not make any difference to Mars')