from datasets.video import (
    MinecraftVideoDataset,
    MinecraftVideoPoseDataset
)

from algorithms.worldmem import WorldMemMinecraft
from algorithms.worldmem import PosePrediction
from .exp_base import BaseLightningExperiment


class VideoPredictionExperiment(BaseLightningExperiment):
    """
    A video prediction experiment
    """

    compatible_algorithms = dict(
        df_video_worldmemminecraft=WorldMemMinecraft,
        pose_prediction=PosePrediction
    )

    compatible_datasets = dict(
        # video datasets
        video_minecraft=MinecraftVideoDataset,
        video_minecraft_pose=MinecraftVideoPoseDataset
    )