from datasets.video import (
    MinecraftVideoDataset,
    MinecraftVideoDeMemWMLatentDataset,
)

from algorithms.worldmem import WorldMemMinecraft
from algorithms.dememwm import DeMemWMMinecraft
from .exp_base import BaseLightningExperiment


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

    compatible_algorithms = dict(
        df_video_worldmemminecraft=WorldMemMinecraft,
        dememwm_base=DeMemWMMinecraft,
    )

    compatible_datasets = dict(
        # video datasets
        video_minecraft=MinecraftVideoDataset,
        video_minecraft_dememwm_latent=MinecraftVideoDeMemWMLatentDataset,
    )