import os
import subprocess
import shutil
from pathlib import Path
from typing import Optional

from dataclasses import dataclass
from benchmark import run_command, get_output_dir, autocompute_tsdf_params
from create_table import compare_mesh, group_images

import tyro

@dataclass
class CompareConfig:
    scene_name: str 
    """Name of the scene that is being run"""

    dataset: str
    """Name of the dataset being used."""

    run_a : Path
    """Path to first run on given dataset"""

    run_b : Optional[Path] = None
    """Path to second run on given dataset"""

    dataset_path: Path = "/Datasets"
    """Name of the dataset being used."""

    output_dir: Path = Path("./compare_runs")
    """Path to output"""


def run_render(config, config_path : Path):
    assert config_path.exists(), "Make sure that config exists."
    timestamp = config_path.parent.name
    model_name = config_path.parent.parent.name

    command = [
            "python", "benchmark_utils/generate_render.py"
        , "--output_scene", config.scene_name
        , "--dataset", config.dataset
        , "--timestamp", timestamp
        , "--model_name", model_name
        , "--config_path", str(config_path)
    ]

    run_name = ""
    if model_name == "raytracingfacto":
        run_name = "qr"
    elif model_name == "dn-splatter":
        run_name = "nerfstudio"
    else:
        raise NotImplementedError("Currently not supported.")

    run_command({ run_name: command })


def generate_mesh_and_render(config, config_path: Path):
    assert config_path.exists()
    timestamp = config_path.parent.name
    model_name = config_path.parent.parent.name
    
    sparse_dir = Path(config.dataset_path).joinpath(config.dataset, config.scene_name, "processed/sparse/0")

    dns_output = Path("./mesh_exports/Open3dTSDFfusion_mesh.ply")
    if dns_output.exists():
        os.remove(dns_output)

    mesh_cmd = []
    if model_name == "dn-splatter":
        run_name = "nerfstudio"
        mesh_name = f"{config.dataset}_{config.scene_name}_DNS"
        depth_trunc = 8.0
        voxel_size, sdf_truc = autocompute_tsdf_params(sparse_dir, depth_trunc)
        mesh_cmd = [
            "gs-mesh", "o3dtsdf",
            "--voxel-size", str(voxel_size),
            "--sdf-truc", str(sdf_truc),
            "--depth-trunc", str(depth_trunc),
        ]

    elif model_name == "raytracingfacto":
        run_name = "qr"
        mesh_name = f"{config.dataset}_{config.scene_name}_ERT"
        mesh_cmd = [
            "python", "mesh_extraction.py", "o3dtsdf",
            "--sparse_dir", sparse_dir.as_posix(),
            "--depth-trunc", "3.0",
            "--name", mesh_name
        ]
    else:
        raise NotImplementedError("Currently not supported")

    mesh_cmd = mesh_cmd + ["--load-config", config_path.as_posix()]
    run_command({ run_name: mesh_cmd })

    if model_name == "dn-splatter":
        shutil.copy(dns_output, f"./mesh_exports/{mesh_name}.ply")

    result_path = Path().home().joinpath(f"benchmark/{config.scene_name}")
    output_path = result_path.joinpath("outputs", model_name, "camera.json")

    command = [
        "python", "benchmark_utils/render_mesh.py",
        "--mesh_path", Path(f"./mesh_exports/{mesh_name}.ply").absolute().as_posix(),
        "--camera_path", output_path.as_posix()
    ]
    run_command( { "qr" : command})


def compare(config: CompareConfig):
    if not config.output_dir.exists():
        config.output_dir.mkdir()
    compare_mesh(config.scene_name, config.output_dir)
    group_images(config.scene_name, config.output_dir)


def main(config: CompareConfig):
    run_render(config, config.run_a)
    generate_mesh_and_render(config, config.run_a)

    if config.run_b is not None:
        run_render(config, config.run_b)
        generate_mesh_and_render(config, config.run_b)



if __name__ == "__main__":
    config = tyro.cli(CompareConfig)
    main(config)
    compare(config)




