# encoding: utf-8
"""
@author:  yuntaej
@contact: jyt0131@gmail.com
"""


import glob
import os.path as osp
import re
import warnings

from .bases import ImageDataset
from ..datasets import DATASET_REGISTRY

import time
import os


@DATASET_REGISTRY.register()
class AIC23_seg(ImageDataset):
    """MOT17.

    Reference:
        Milan, A., Leal-Taixé, L., Reid, I., Roth, S. & Schindler, K. MOT16: A Benchmark for Multi-Object Tracking. arXiv:1603.00831 [cs], 2016., (arXiv: 1603.00831)

    URL: `<https://motchallenge.net/data/MOT17/>`_

    Dataset statistics:
        - identities: ?
        - images: ?
    """
    dataset_dir = '/data/cvprw/AIC23/dataset/converted_seg_dataset/train'
    dataset_url = ''  # 'https://motchallenge.net/data/MOT17.zip'
    dataset_name = "AIC23_seg"

    def __init__(self, root='datasets', **kwargs):
        # self.root = osp.abspath(osp.expanduser(root))
#         self.root = root
#         self.dataset_dir = osp.join(self.root, self.dataset_dir)
#         self.dataset_dir = dataset_dir
        dataset_dir = '/data/cvprw/AIC23/dataset/converted_seg_dataset/train'
        # allow alternative directory structure
        self.data_dir = dataset_dir
        if osp.isdir(self.data_dir):
            self.data_dir = self.data_dir
        else:
            warnings.warn('The current data structure is deprecated. Please '
                          'put data folders such as "bounding_box_train" under '
                          '"AIC23_bbox".')

        self.train_dir = osp.join(self.data_dir, 'detect_result')
        self.query_dir = osp.join(self.data_dir, 'query')
        self.gallery_dir = osp.join(self.data_dir, 'bounding_box_test')
        self.extra_gallery_dir = osp.join(self.data_dir, 'images')
        self.extra_gallery = False

        required_files = [
            self.data_dir,
            self.train_dir,
            # self.query_dir,
            # self.gallery_dir,
        ]

        self.check_before_run(required_files)

        train = lambda: self.process_dir(self.train_dir)
        query = lambda: self.process_dir(self.query_dir, is_train=False)
        gallery = lambda: self.process_dir(self.gallery_dir, is_train=False) + \
                          (self.process_dir(self.extra_gallery_dir, is_train=False) if self.extra_gallery else [])

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

    def process_dir(self, dir_path, is_train=True):
        
        data = []
        
        site_list = [file for file in os.listdir(dir_path) if not file.startswith('.')]
        for site in site_list:
            channel_list = [file for file in os.listdir(os.path.join(dir_path,site)) if not file.startswith('.')]
            for channel in channel_list:
                img_list = [file for file in os.listdir(os.path.join(dir_path,site,channel)) if  file.endswith('.png')]
                for img in img_list:
                    file_name=img.split('.')[0]
                    pid = file_name.split('_')[1]
                    camid = str(site)+"-"+str(channel)
                    
                    if is_train:
                        pid = self.dataset_name + "_" + str(pid)
                        camid = self.dataset_name + "_" + str(camid)
                        
                    img_path = os.path.join(dir_path,site,channel,img)
                    
                    data.append((img_path, pid, camid))

        return data