@dataclass
class LSFBContConfig:
    """
        Simple configuration class for the LSFB_CONT dataset.
        Each instance of the LSFB_CONT dataset is a video with corresponding annotations.
        Each annotation is made up of a segment (start, end) and a label.

        LSFB: French Belgian Sign Language
        CONT: Continuous videos in LSFB

        See `LSFBContLandmarks` to load the dataset.

        Args:
            root: Root directory of the LSFB_CONT dataset.
                The dataset must already be downloaded !
            landmarks: Select which landmarks (features) to use.
                'face' for face mesh (468 landmarks);
                'pose' for pose skeleton (23 landmarks);
                'left_hand' for left hand skeleton (21 landmarks);
                'right_hand' for right hand skeleton (21 landmarks);
                Default = ('pose', 'left_hand', 'right_hand').
            use_3d: If true, use 3D landmarks. Otherwise, use 2D landmarks.
                Default=False.
            use_raw: If true, use raw landmarks. Otherwise, use preprocessed landmarks where:
                - missing landmarks are interpolated;
                - vibrations have been reduced by using smoothing (Savitchy Golay filter).
                Default=False.

            features_transform: Callable object used to transform the features.
            target_transform: Callable object used to transform the targets.
            transform: Callable object used to transform both the features and the targets.

            split: Specify which subset of the dataset is used.
                'fold_x' where x is in {0, 1, 2, 3, 4} for a specific fold;
                'train' for training set (folds 2, 3, 4);
                'test' for the test set (folds 0 and 1);
                'all' for all the instances of the dataset (all folds);
                'mini_sample' for a tiny set of instances (10 instances).
                Default = 'all'.

            hands: Only load the sign of a specific hand, or both of them.
                'right' for the signs from the right hand of the signer;
                'left' for the signs from the left hand of the signer;
                'both' for the signs from both hands of the signer.
                Default = 'both'.
            segment_level: Specifies the level at which annotations are extracted.
                'signs'
                'subtitles'
                Default = 'signs'
            segment_label: Specify which label to use for each segment.
                'sign_gloss': the gloss of the signs is used. Example: HAND
                'sign_index': the index (class) of the signs is used. Example: 45
                'text': the text of the sign, or subtitle, is used. This can be useful when annotations are subtitles.
                If the annotations are subtitles, multiple labels are assigned to a segment.
                For example, a sequence of glosses, indices or event the full text of the subtitles.
                Default='sign_index'
            segment_unit: Specify which unit is used for the boundaries (start, end) of the annotations.
                'frame': frame indices in annotations boundaries.
                'ms': milliseconds in annotations boundaries.
                default='ms'

            n_labels: If this parameter is an integer `x`, then `x+1` labels are used for the `x` most
                frequent signs and the background label. If none, the number of labels is the number of different signs
                in the dataset no matter their occurrences.
                This parameter is used to filter out signs with very few examples.
                Default=750

            window: Optional argument (window_size, window_stride) to use fixed-size windows instead of
                variable length sequences.
                If specified, the dataset is windowed with a window size and a window stride.
                Be careful, this argument changes the number of instances in the dataset !
                Default = None.

            show_progress: If true, shows a progress bar while the dataset is loading. Default = True.

        Author:
            ppoitier (v 2.0)
    """

    root: str
    landmarks: Optional[tuple[str, ...]] = ('pose', 'left_hand', 'right_hand')
    use_3d: bool = False
    use_raw: bool = False

    features_transform: callable = None
    target_transform: callable = None
    transform: callable = None

    split: str = 'all'

    hands: str = 'both'
    segment_level: str = 'signs'
    segment_label: str = 'sign_index'
    segment_unit: str = 'ms'

    n_labels: Union[int, None] = 750

    window: Optional[tuple[int, int]] = None

    show_progress: bool = True

class LSFBContLandmarks(LSFBContBase):
    """
        Utility class to load the LSFB CONT Landmarks dataset.
        The dataset must be already downloaded!

        All the landmarks and targets are loaded in memory.
        Therefore, iterating over all the instances is fast but consumes a lot of RAM.
        If you don't have enough RAM, use the `LSFBContLandmarksGenerator` class instead.

        Example:
            ```python
            my_dataset_config = LSFBContConfig(
                root="./my_dataset",
                landmarks=['pose', 'left_hand', 'right_hand'],
                split="fold_1",
                n_labels=750,
                segment_level='signs',
                segment_unit='frame',
                segment_label='sign_gloss',
                use_3d=True,
                window=(1500, 1200),
            )

            my_dataset = LSFBContLandmarks(my_dataset_config)
            features, target_annotations = dataset[10]
            ```

        If you did not download the dataset, see `lsfb_dataset.Downloader`.

        Args:
            config: The configuration object (see `LSFBContConfig`).

        Author:
            ppoitier (v 2.0)
    """
    # TODO: add class properties to docstring

    def __init__(self, config: LSFBContConfig):
        super().__init__(config)
        self.features: dict[str, dict[str, np.ndarray]] = {}
        self._load_features()

    def __get_instance__(self, index):
        instance_id = self.instances[index]
        features = self.features[instance_id]
        annotations = self.annotations[instance_id].values
        features, annotations = self._apply_transforms(features, annotations)
        return features, annotations

    def __get_window__(self, index):
        instance_id, start, end = self.windows[index]
        features = {lm: lm_feat[start:end] for lm, lm_feat in self.features[instance_id].items()}

        annotations = self.annotations[instance_id]
        if self.config.segment_unit == 'ms':
            annotations = annotations.loc[
                ((annotations['end'] / 20) >= start) &
                ((annotations['start'] / 20) <= end)
            ]
            annotations.loc[:, 'start'] = annotations['start'] - start * 20
            annotations.loc[:, 'end'] = annotations['end'] - start * 20
        else:
            annotations = annotations.loc[
                (annotations['end'] >= start) &
                (annotations['start'] <= end)
            ]
            annotations.loc[:, 'start'] = annotations['start'] - start
            annotations.loc[:, 'end'] = annotations['end'] - start
        features, annotations = self._apply_transforms(features, annotations)
        return features, annotations

    def _load_features(self):
        pose_folder = 'poses_raw' if self.config.use_raw else 'poses'
        coordinate_indices = [0, 1, 2] if self.config.use_3d else [1, 2]
        for instance_id in tqdm(self.instances, disable=(not self.config.show_progress)):
            instance_feat = {}
            for landmark_set in self.config.landmarks:
                filepath = f"{self.config.root}/{pose_folder}/{landmark_set}/{instance_id}.npy"
                instance_feat[landmark_set] = np.load(filepath)[:, :, coordinate_indices]
            self.features[instance_id] = instance_feat
        gc.collect()

class LSFBContLandmarksGenerator(LSFBContBase):
    """
        Utility class to load the LSFB CONT Landmarks dataset.
        The dataset must be already downloaded!

        All the landmarks and targets are lazily loaded.
        Therefore, iterating over all the instances can be a bit slow.
        If you have enough RAM and want faster iterations, use the `LSFBContLandmarks` class instead.

        Example:
            ```python
            my_dataset_config = LSFBContConfig(
                root="./my_dataset",
                landmarks=['pose', 'left_hand', 'right_hand'],
                split="fold_1",
                n_labels=750,
                segment_level='signs',
                segment_unit='frame',
                segment_label='sign_gloss',
                use_3d=True,
                window=(1500, 1200),
            )

            my_dataset = LSFBContLandmarksGenerator(my_dataset_config)
            features, target_annotations = dataset[10]
            ```

        If you did not download the dataset, see `lsfb_dataset.Downloader`.

        Args:
            config: The configuration object (see `LSFBContConfig`).

        Author:
            ppoitier (v 2.0)
    """

    def __init__(self, config: LSFBContConfig):
        super(LSFBContLandmarksGenerator, self).__init__(config)

    def __get_instance__(self, index):
        instance_id = self.instances[index]
        features = self._load_instance_features(instance_id)
        annotations = self.annotations[instance_id]
        features, annotations = self._apply_transforms(features, annotations)
        return features, annotations

    def __get_window__(self, index):
        instance_id, start, end = self.windows[index]
        features = self._load_instance_features(instance_id)
        features = {lm: lm_feat[start:end] for lm, lm_feat in features[instance_id].items()}

        annotations = self.annotations[instance_id]
        if self.config.segment_unit == 'ms':
            annotations = annotations.loc[
                ((annotations['end'] / 20) >= start) &
                ((annotations['start'] / 20) <= end)
            ]
        else:
            annotations = annotations.loc[
                (annotations['end'] >= start) &
                (annotations['start'] <= end)
            ]
        annotations.loc[:, 'start'] = annotations['start'] - start
        annotations.loc[:, 'end'] = annotations['end'] - start
        features, annotations = self._apply_transforms(features, annotations)
        return features, annotations

    def _load_instance_features(self, instance_id: str):
        pose_folder = 'poses_raw' if self.config.use_raw else 'poses'
        coordinate_indices = [0, 1, 2] if self.config.use_3d else [1, 2]
        features = {}
        for landmark_set in self.config.landmarks:
            filepath = f"{self.config.root}/{pose_folder}/{landmark_set}/{instance_id}.npy"
            features[landmark_set] = np.load(filepath)[:, :, coordinate_indices]
        return features

@dataclass
class LSFBIsolConfig:
    """
    Simple configuration class for the LSFB ISOL Dataset.

    LSFB: French Belgian Sign Language
    ISOL: Isolated signs (videos and poses) in LSFB

    If you did not download the dataset, see `lsfb_dataset.Downloader`.

    Args:
        root: Root directory of the LSFB_ISOL dataset.
            The dataset must already be downloaded!
        landmarks: Select which landmarks (features) to use:
            'face' for face mesh (468 landmarks);
            'pose' for pose skeleton (23 landmarks);
            'left_hand' for left hand skeleton (21 landmarks);
            'right_hand' for right hand skeleton (21 landmarks);
            Default = ['pose', 'left_hand', 'right_hand'].
        use_3d: If true, use 3D landmarks. Otherwise, use 2D landmarks. Default = False.
        use_raw: If true, use raw landmarks. Otherwise, use preprocessed landmarks where:
            - missing landmarks are interpolated (linear interpolation);
            - vibrations have been reduced by using smoothing (Savitchy Golay filter).
            Default = False.

        target: Specify which target is returned with the features for each instance:
            'sign_index': the index of the sign is used;
            'sign_gloss': the gloss of the sign is used;
            Default = 'sign_index'.

        transform: Callable object used to transform the features.

        split: Specify which subset of the dataset is used.
            'fold_x' where x is in {0, 1, 2, 3, 4} for a specific fold;
            'train' for training set (folds 2, 3, 4);
            'test' for the test set (folds 0 and 1);
            'all' for all the instances of the dataset (all folds);
            'mini_sample' for a tiny set of instances (10 instances).
            Default = 'all'.

        sequence_max_length: (Optional) Max length of the clip sequence. Default=50.

        n_labels: If this parameter is an integer `x`, then `x+1` labels are used for the `x` most
            frequent signs and the background label. If none, the number of labels is the number of different signs
            in the dataset no matter their occurrences.
            This parameter is used to filter out signs with very few examples.
            Default=750.

        show_progress: If true, shows a progress bar while the dataset is loading. Default = True.

    Authors:
        jfink (v 1.0)
        ppoitier (v 2.0)
    """

    root: str
    landmarks: Optional[tuple[str, ...]] = ("pose", "left_hand", "right_hand")
    use_3d: bool = False
    use_raw: bool = False

    target: str = "sign_index"

    transform: callable = None

    split: str = "all"
    n_labels: int = 750
    sequence_max_length: int = 50

    show_progress: bool = True

class LSFBIsolLandmarks(LSFBIsolBase):
    """
    Utility class to load the LSFB ISOL Landmarks dataset.
    The dataset must be already downloaded!

    All the landmarks and targets are loaded in memory.
    Therefore, iterating over all the instances is fast but consumes a lot of RAM.
    If you don't have enough RAM, use the `LSFBIsolLandmarksGenerator` class instead.

    Example:
        ```python
        my_dataset_config = LSFBIsolConfig(
            root="./my_dataset",
            split="fold_1",
            n_labels=750,
            target='sign_gloss',
            sequence_max_length=10,
            use_3d=True,
        )

        my_dataset = LSFBIsolLandmarks(my_dataset_config)
        features, target = dataset[30]
        ```

    If you did not download the dataset, see `lsfb_dataset.Downloader`.

    Args:
        config: The configuration object (see `LSFBContConfig`).

    Author:
        ppoitier (v 2.0)
    """
    # TODO: add class properties to docstring

    def __init__(self, config: LSFBIsolConfig):
        super().__init__(config)
        self.features: dict[str, dict[str, np.ndarray]] = self._load_features()

    def __getitem__(self, index):
        instance_id = self.instances[index]
        features = self.features[instance_id]
        target = self.targets[instance_id]

        if self.config.transform is not None:
            features = self.config.transform(features)

        return features, target

    def _load_features(self):
        pose_folder = "poses_raw" if self.config.use_raw else "poses"
        coordinate_indices = [0, 1, 2] if self.config.use_3d else [0, 1]
        all_features = {}
        max_len = self.config.sequence_max_length

        for instance_id in tqdm(self.instances, disable=(not self.config.show_progress)):
            instance_features = {}
            for landmark_set in self.config.landmarks:
                filepath = f"{self.config.root}/{pose_folder}/{landmark_set}/{instance_id}.npy"
                lm_set_features = np.load(filepath)[:, :, coordinate_indices]
                if max_len is not None:
                    lm_set_features = lm_set_features[:max_len]
                instance_features[landmark_set] = lm_set_features
            all_features[instance_id] = instance_features
        gc.collect()
        return all_features

class LSFBIsolLandmarksGenerator(LSFBIsolBase):
    """
        Utility class to load the LSFB ISOL Landmarks dataset.
        The dataset must be already downloaded!

        All the landmarks and targets are lazily loaded.
        Therefore, iterating over all the instances can be a bit slow.
        If you have enough RAM and want faster iterations, use the `LSFBIsolLandmarks` class instead.

        Example:
            ```python
            my_dataset_config = LSFBIsolConfig(
                root="./my_dataset",
                split="fold_1",
                n_labels=750,
                target='sign_gloss',
                sequence_max_length=10,
                use_3d=True,
            )

            my_dataset = LSFBIsolLandmarksGenerator(my_dataset_config)
            features, target = dataset[30]
            ```

        Args:
            config: The configuration object (see `LSFBContConfig`).

        Author:
            ppoitier (v 2.0)
        """

    def __init__(self, config: LSFBIsolConfig):
        super().__init__(config)

    def __getitem__(self, index):
        instance_id = self.instances[index]
        features = self._load_instance_features(instance_id)
        target = self.targets[instance_id]

        if self.config.transform is not None:
            features = self.config.transform(features)

        return features, target

    def _load_instance_features(self, instance_id):
        pose_folder = 'poses_raw' if self.config.use_raw else 'poses'
        coordinate_indices = [0, 1, 2] if self.config.use_3d else [0, 1]
        max_len = self.config.sequence_max_length
        instance_features = {}
        for landmark_set in self.config.landmarks:
            pose_path = f"{self.config.root}/{pose_folder}/{landmark_set}/{instance_id}.npy"
            lm_set_features = np.load(pose_path)[:, :, coordinate_indices]
            if max_len is not None:
                lm_set_features = lm_set_features[:max_len]
            instance_features[instance_id] = lm_set_features
        return instance_features