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


class LaSOTExtDataset(BaseDataset):
    """
    LaSOT test set consisting of 280 videos (see Protocol-II in the LaSOT paper)

    Publication:
        LaSOT: A High-quality Benchmark for Large-scale Single Object Tracking
        Heng Fan, Liting Lin, Fan Yang, Peng Chu, Ge Deng, Sijia Yu, Hexin Bai, Yong Xu, Chunyuan Liao and Haibin Ling
        CVPR, 2019
        https://arxiv.org/pdf/1809.07845.pdf

    Download the dataset from https://cis.temple.edu/lasot/download.html
    """
    def __init__(self):
        super().__init__()
        self.base_path = self.env_settings.lasot_ext_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', 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-2',
            'cosplay-3',
            'cosplay-4',
            'cosplay-5',
            'cosplay-6',
            'cosplay-7',
            'cosplay-8',
            'cosplay-9',
            'cosplay-10',
            '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
