"""Task: Future Frame Prediction (scene-level)."""
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 {
        "pmf_train": 0.0,
        "pmf_novel": 0.0,
        "psnr_train": 0.0,
        "ssim_train": 0.0,
        "lpips_train": 0.0,
        "psnr_novel": 0.0,
        "ssim_novel": 0.0,
        "lpips_novel": 0.0,
    }


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


TASK = TaskPlugin(
    name="task_future_prediction",
    display_name="Future Frame Prediction",
    description="Predict future frames from historical video frames. Evaluated with PSNR / SSIM / LPIPS / PMF.",
    expected_scene_layout=(
        "```\n"
        "<scene_id>.zip\n"
        "├── CineCamera_0/\n"
        "│   └── rgb/\n"
        "│       ├── 0075.jpg\n"
        "│       └── ...\n"
        "├── CineCamera_1/\n"
        "│   └── rgb/\n"
        "└── ... (must match the GT camera set except for CineCamera_Moving)\n"
        "\n"
        "The system generates a uid for each submission (the submission_id), and the Worker extracts the archive to\n"
        "  /data/tmp_inference_output/<uid>/<benchmark>/<scene_id>_trajectory/\n"
        "GT provides transforms_train.json and transforms_val.json to define training and novel views,\n"
        "while transforms_test.json defines the exact evaluation frames for those views.\n"
        "```"
    ),
    validate_scene_fn=validate_scene,
    evaluate_scene_fn=evaluate_scene,
    load_gt_scene_fn=load_gt_scene,
    primary_metric="pmf_novel",
    higher_is_better=True,
    leaderboard_columns=[
        "pmf_novel",
        "psnr_novel",
        "ssim_novel",
        "lpips_novel",
        "pmf_train",
        "psnr_train",
        "ssim_train",
        "lpips_train",
    ],
)
TASK.column_display_names = {
    "pmf_train": "PMF (training view)",
    "pmf_novel": "PMF (novel view)",
    "psnr_train": "PSNR (training view)",
    "ssim_train": "SSIM (training view)",
    "lpips_train": "LPIPS (training view)",
    "psnr_novel": "PSNR (novel view)",
    "ssim_novel": "SSIM (novel view)",
    "lpips_novel": "LPIPS (novel view)",

}