"""Task: Physical Property Estimation."""
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 {"overall_error": 0.0}


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


EXCEL_SUMMARY_SPECS = {
    "N μ": {"source": "newtonian_log10_mu", "scale": 100.0, "digits": 2},
    "N κ": {"source": "newtonian_log10_kappa", "scale": 100.0, "digits": 2},
    "N v": {"source": "newtonian_v", "scale": 100.0, "digits": 2},
    "NN μ": {"source": "non_newtonian_log10_mu", "scale": 100.0, "digits": 2},
    "NN κ": {"source": "non_newtonian_log10_kappa", "scale": 100.0, "digits": 2},
    "NN y": {"source": "non_newtonian_log10_yield", "scale": 100.0, "digits": 2},
    "NN η": {"source": "non_newtonian_log10_eta", "scale": 100.0, "digits": 2},
    "NN v": {"source": "non_newtonian_v", "scale": 100.0, "digits": 2},
    "E E": {"source": "elastic_log10_E", "scale": 100.0, "digits": 2},
    "E ν": {"source": "elastic_nu", "scale": 100.0, "digits": 2},
    "E v": {"source": "elastic_v", "scale": 100.0, "digits": 2},
    "P E": {"source": "plasticine_log10_E", "scale": 100.0, "digits": 2},
    "P y": {"source": "plasticine_log10_yield", "scale": 100.0, "digits": 2},
    "P ν": {"source": "plasticine_nu", "scale": 100.0, "digits": 2},
    "P v": {"source": "plasticine_v", "scale": 100.0, "digits": 2},
    "S φ": {"source": "sand_friction", "scale": 1.0, "digits": 2},
    "S v": {"source": "sand_v", "scale": 100.0, "digits": 2},
}


TASK = TaskPlugin(
    name="propertyEstimation",
    display_name="Properties Estimation",
    description=(
        "Predict MPM material parameters and initial velocity. "
        "This task jointly evaluates MPM Parameter Estimation, Reinitialized Simulation "
        "Rendering, and Novel View Synthesis. "
        "Abbreviations: N=newtonian, NN=non-newtonian, E=elastic, P=plasticine, S=sand."
    ),
    expected_scene_layout=(
        "```\n"
        "<scene_id>.zip\n"
        "`-- property_estimation/\n"
        "    |-- materials_configs.json\n"
        "    |-- initvelocity.json   # {\"vx\": float, \"vy\": float, \"vz\": float}\n"
        "    |-- novel/\n"
        "    |   `-- CineCamera_*/rgb/<frame>.jpg\n"
        "    `-- reinit/\n"
        "        `-- CineCamera_*/rgb/<frame>.jpg\n"
        "```"
    ),
    validate_scene_fn=validate_scene,
    evaluate_scene_fn=evaluate_scene,
    load_gt_scene_fn=load_gt_scene,
    primary_metric="rendering_pmf",
    higher_is_better=True,
    leaderboard_columns=[
        "rendering_pmf",
        "novel_pmf",
        "reinit_pmf",
        "novel_psnr",
        "novel_ssim",
        "novel_lpips",
        "reinit_psnr",
        "reinit_ssim",
        "reinit_lpips",
        "overall_error",
        *EXCEL_SUMMARY_SPECS.keys(),
    ],
)

TASK.excel_summary_specs = EXCEL_SUMMARY_SPECS
TASK.column_display_names = {
    "rendering_pmf": "rendering pmf",
    "novel_pmf": "novel pmf",
    "reinit_pmf": "reinit pmf",
    "novel_psnr": "novel psnr",
    "novel_ssim": "novel ssim",
    "novel_lpips": "novel lpips",
    "reinit_psnr": "reinit psnr",
    "reinit_ssim": "reinit ssim",
    "reinit_lpips": "reinit lpips",
    "overall_error": "param overall",
}

TASKS = {
    TASK.name: TASK,
}