"""Task: Video Generation (scene-level).

Per-scene metrics: PSNR / SSIM / LPIPS / PMF . The evaluator is initialized on CPU by
default; export ``PHYS_METRIC_DEVICE=cuda`` to opt in to GPU when one is
available to the worker.
"""
from __future__ import annotations

from typing import Any, Dict

from src.tasks.base import TaskPlugin


def validate_scene(_scene_dir: str) -> None:
    return None


def evaluate_scene(_scene_dir: str, _gt_scene: Any) -> Dict[str, float]:
    return {"psnr": 0.0, "ssim": 0.0, "lpips": 0.0, "pmf": 0.0}


def load_gt_scene(scene_id: str, gt_dir: str) -> Any:
    return {"scene_id": scene_id, "gt_dir": gt_dir}

TASK_VIDEO_GEN_STATIC = TaskPlugin(
    name="videoGeneration_static",
    display_name="Video Generation on Static Camera",
    description="Given first frame and caption, generate a physical plausible video on a static camera. " \
                "Metrics: PSNR / SSIM / LPIPS / PMF.",
    expected_scene_layout=(
        "```\n"
        "<scene_id>.zip\n"
        "├── CineCamera_n/\n"
        "│   └── rgb/\n"
        "│       ├── ...\n"
        "│       ├── 0075.jpg\n"
        "│       └── ...\n"
        "\n"
        "```"
    ),
    validate_scene_fn=validate_scene,
    evaluate_scene_fn=evaluate_scene,
    load_gt_scene_fn=load_gt_scene,
    primary_metric="pmf",
    higher_is_better=True,
    leaderboard_columns=["psnr", "ssim", "lpips", "pmf"],
)

TASK_VIDEO_GEN_MOVING = TaskPlugin(
    name="videoGeneration_moving",
    display_name="Video Generation on Moving Camera",
    description="Given first frame and caption, generate a physical plausible video a static camera. " \
                "Metrics: PSNR / SSIM / LPIPS / PMF.",
    expected_scene_layout=(
        "```\n"
        "<scene_id>.zip\n"
        "├── CineCamera_Moving/\n"
        "    └── rgb/\n"
        "│       ├── ...\n"
        "│       ├── 0075.jpg\n"
        "│       └── ...\n"
        "\n"
        "```"
    ),
    validate_scene_fn=validate_scene,
    evaluate_scene_fn=evaluate_scene,
    load_gt_scene_fn=load_gt_scene,
    primary_metric="pmf",
    higher_is_better=True,
    leaderboard_columns=["psnr", "ssim", "lpips", "pmf"],
)

TASKS = {
    TASK_VIDEO_GEN_STATIC.name:  TASK_VIDEO_GEN_STATIC,
    TASK_VIDEO_GEN_MOVING.name:  TASK_VIDEO_GEN_MOVING,
}