import numpy as np
from pytracking.evaluation.data import Sequence, BaseDataset, SequenceList
from pytracking.utils.load_text import load_text


class LaSOTExtensionSubsetDataset(BaseDataset):
    """
    LaSOT Extension Subset dataset consisting of 150 videos

    Publication:
        LaSOT: A High-quality Large-scale Single Object Tracking Benchmark
        Heng Fan, Hexin Bai, Liting Lin, Fan Yang, Peng Chu, Ge Deng, Sijia Yu, Harshit, Mingzhen Huang, Juehuan Liu,
        Yong Xu, Chunyuan Liao, Lin Yuan, Haibin Ling
        IJCV, 2020
        https://arxiv.org/pdf/2009.03465.pdf

    Download the dataset from http://vision.cs.stonybrook.edu/~lasot/download.html
    """
    def __init__(self):
        super().__init__()
        self.base_path = self.env_settings.lasot_extension_subset_path
        self.sequence_list = self._get_sequence_list()
        self.clean_list = self.clean_seq_list()

    def clean_seq_list(self):
        clean_lst = []
        for i in range(len(self.sequence_list)):
            cls, _ = self.sequence_list[i].split('-')
            clean_lst.append(cls)
        return  clean_lst

    def get_sequence_list(self):
        return SequenceList([self._construct_sequence(s) for s in self.sequence_list])

    def _construct_sequence(self, sequence_name):
        class_name = sequence_name.split('-')[0]
        anno_path = '{}/{}/{}/groundtruth.txt'.format(self.base_path, class_name, sequence_name)

        ground_truth_rect = load_text(str(anno_path), delimiter=',', dtype=np.float64)

        occlusion_label_path = '{}/{}/{}/full_occlusion.txt'.format(self.base_path, class_name, sequence_name)

        # NOTE: pandas backed seems super super slow for loading occlusion/oov masks
        full_occlusion = load_text(str(occlusion_label_path), delimiter=',', dtype=np.float64, backend='numpy')

        out_of_view_label_path = '{}/{}/{}/out_of_view.txt'.format(self.base_path, class_name, sequence_name)
        out_of_view = load_text(str(out_of_view_label_path), delimiter=',', dtype=np.float64, backend='numpy')

        target_visible = np.logical_and(full_occlusion == 0, out_of_view == 0)

        frames_path = '{}/{}/{}/img'.format(self.base_path, class_name, sequence_name)

        frames_list = ['{}/{:08d}.jpg'.format(frames_path, frame_number) for frame_number in range(1, ground_truth_rect.shape[0] + 1)]

        target_class = class_name
        return Sequence(sequence_name, frames_list, 'lasot_extension_subset', ground_truth_rect.reshape(-1, 4),
                        object_class=target_class, target_visible=target_visible)

    def __len__(self):
        return len(self.sequence_list)

    def _get_sequence_list(self):
        sequence_list = ['atv-1',
                         'atv-2',
                         'atv-3',
                         'atv-4',
                         'atv-5',
                         'atv-6',
                         'atv-7',
                         'atv-8',
                         'atv-9',
                         'atv-10',
                         'badminton-1',
                         'badminton-2',
                         'badminton-3',
                         'badminton-4',
                         'badminton-5',
                         'badminton-6',
                         'badminton-7',
                         'badminton-8',
                         'badminton-9',
                         'badminton-10',
                         'cosplay-1',
                         'cosplay-10',
                         'cosplay-2',
                         'cosplay-3',
                         'cosplay-4',
                         'cosplay-5',
                         'cosplay-6',
                         'cosplay-7',
                         'cosplay-8',
                         'cosplay-9',
                         'dancingshoe-1',
                         'dancingshoe-2',
                         'dancingshoe-3',
                         'dancingshoe-4',
                         'dancingshoe-5',
                         'dancingshoe-6',
                         'dancingshoe-7',
                         'dancingshoe-8',
                         'dancingshoe-9',
                         'dancingshoe-10',
                         'footbag-1',
                         'footbag-2',
                         'footbag-3',
                         'footbag-4',
                         'footbag-5',
                         'footbag-6',
                         'footbag-7',
                         'footbag-8',
                         'footbag-9',
                         'footbag-10',
                         'frisbee-1',
                         'frisbee-2',
                         'frisbee-3',
                         'frisbee-4',
                         'frisbee-5',
                         'frisbee-6',
                         'frisbee-7',
                         'frisbee-8',
                         'frisbee-9',
                         'frisbee-10',
                         'jianzi-1',
                         'jianzi-2',
                         'jianzi-3',
                         'jianzi-4',
                         'jianzi-5',
                         'jianzi-6',
                         'jianzi-7',
                         'jianzi-8',
                         'jianzi-9',
                         'jianzi-10',
                         'lantern-1',
                         'lantern-2',
                         'lantern-3',
                         'lantern-4',
                         'lantern-5',
                         'lantern-6',
                         'lantern-7',
                         'lantern-8',
                         'lantern-9',
                         'lantern-10',
                         'misc-1',
                         'misc-2',
                         'misc-3',
                         'misc-4',
                         'misc-5',
                         'misc-6',
                         'misc-7',
                         'misc-8',
                         'misc-9',
                         'misc-10',
                         'opossum-1',
                         'opossum-2',
                         'opossum-3',
                         'opossum-4',
                         'opossum-5',
                         'opossum-6',
                         'opossum-7',
                         'opossum-8',
                         'opossum-9',
                         'opossum-10',
                         'paddle-1',
                         'paddle-2',
                         'paddle-3',
                         'paddle-4',
                         'paddle-5',
                         'paddle-6',
                         'paddle-7',
                         'paddle-8',
                         'paddle-9',
                         'paddle-10',
                         'raccoon-1',
                         'raccoon-2',
                         'raccoon-3',
                         'raccoon-4',
                         'raccoon-5',
                         'raccoon-6',
                         'raccoon-7',
                         'raccoon-8',
                         'raccoon-9',
                         'raccoon-10',
                         'rhino-1',
                         'rhino-2',
                         'rhino-3',
                         'rhino-4',
                         'rhino-5',
                         'rhino-6',
                         'rhino-7',
                         'rhino-8',
                         'rhino-9',
                         'rhino-10',
                         'skatingshoe-1',
                         'skatingshoe-2',
                         'skatingshoe-3',
                         'skatingshoe-4',
                         'skatingshoe-5',
                         'skatingshoe-6',
                         'skatingshoe-7',
                         'skatingshoe-8',
                         'skatingshoe-9',
                         'skatingshoe-10',
                         'wingsuit-1',
                         'wingsuit-2',
                         'wingsuit-3',
                         'wingsuit-4',
                         'wingsuit-5',
                         'wingsuit-6',
                         'wingsuit-7',
                         'wingsuit-8',
                         'wingsuit-9',
                         'wingsuit-10']
        return sequence_list
