#!/usr/bin/env python3
"""
GS-only eval for H3D mini139-style bench: run sparse-structure + SLAT inference, decode Gaussians,
and save multiview GS renders under a *separate* tree (`work_dirs/eval_gs_only/...` by default).

Does not write GLB/mesh, Blender `render/`, or the standard `work_dirs/eval_results/<save_name>/` layout.

Typical "h3d_only + EMA + 60k" — point checkpoint roots at your **h3d_only** training outputs:

  export TRAINED_SS_DIR=/path/to/h3d_only_img_to_voxel
  export TRAINED_SLAT_DIR=/path/to/h3d_only_voxel_to_texture

  /root/miniconda3/envs/vinedresser3d/bin/python scripts/test_gs_only_h3d_bench.py \
    --ss_latents_load_dir "${TRAINED_SS_DIR}" \
    --latents_load_dir "${TRAINED_SLAT_DIR}" \
    --ss_latents_load_ckpt 60000 \
    --latents_load_ckpt 60000 \
    --load_ema_model_for_inference \
    --run_name h3d_only_ema_step60000_mini139_gs \
    --end_idx 20
"""

from __future__ import annotations

import argparse
import json
import math
import os
import sys
import time
from pathlib import Path

_REPO_ROOT = Path(__file__).resolve().parents[1]
if str(_REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(_REPO_ROOT))
os.chdir(_REPO_ROOT)

import tqdm  # noqa: E402

from eval_3d_editing import TrellisEdititngModel, set_seed  # noqa: E402


def _default_from_env(key: str, fallback: str | None) -> str | None:
    v = os.environ.get(key)
    return v if v else fallback


def build_argparser() -> argparse.ArgumentParser:
    bench = "/mnt/zsn/zsn_workspace/PartCraft3D/H3D_v1_hf/data/test_data/h3d_manual_keep_mini139_unique_3deditformer_style"

    p = argparse.ArgumentParser(description="H3D bench — GS multiview only (separate output_root).")
    p.add_argument(
        "--dataset_root_dir",
        type=str,
        default=bench,
        help="Bench root (contains dataset_info.json + select json).",
    )
    p.add_argument("--data_info_json_path", type=str, default="dataset_info.json")
    p.add_argument(
        "--select_json_path",
        type=str,
        default="test_data_info_3deditformer.json",
        help="Same as full-bench eval (`run_h3d_manual_mini140_compare_tmux.sh`).",
    )
    p.add_argument("--flux_edit_root_path", type=str, default="flux_edit")
    p.add_argument("--alpaca_root_path", type=str, default="alpaca")
    p.add_argument("--mixamo_root_path", type=str, default="mixamo")

    p.add_argument(
        "--output_root",
        type=str,
        default="./work_dirs/eval_gs_only",
        help="Root for outputs (does not touch work_dirs/eval_results/).",
    )
    p.add_argument(
        "--run_name",
        type=str,
        default="h3d_only_ema_step60000_mini139_gs",
        help="Subfolder under output_root.",
    )

    p.add_argument(
        "--ss_latents_config",
        type=str,
        default="ss_flow_img_dit_L_16l8_fp16_h3d_add_del_mod_scale_plus_3deditverse_clean_40k.json",
    )
    p.add_argument(
        "--latents_config",
        type=str,
        default="slat_flow_img_dit_L_64l8p2_fp16_h3d_add_del_mod_scale_plus_3deditverse_clean_40k.json",
    )

    d_ss = _default_from_env("TRAINED_SS_DIR", "/root/workspace/outputs/h3d_clean_img_to_voxel_40k_bsz8x2")
    d_slat = _default_from_env("TRAINED_SLAT_DIR", "/root/workspace/outputs/h3d_clean_voxel_to_texture_40k_bsz16x2")
    p.add_argument(
        "--ss_latents_load_dir",
        type=str,
        default=d_ss,
        help="Sparse-structure ckpt root (contains ckpts/). Override for h3d_only training dirs.",
    )
    p.add_argument(
        "--latents_load_dir",
        type=str,
        default=d_slat,
        help="SLAT ckpt root (contains ckpts/). Override for h3d_only training dirs.",
    )
    p.add_argument("--ss_latents_load_ckpt", type=int, default=60000)
    p.add_argument("--latents_load_ckpt", type=int, default=60000)
    p.add_argument("--load_ema_model_for_inference", action="store_true")

    p.add_argument("--trellis_pipeline_path", type=str, default="/mnt/zsn/ckpts/TRELLIS-image-large")

    p.add_argument("--total_render_view_num", type=int, default=300)
    p.add_argument("--eval_view_num", type=int, default=10)
    p.add_argument("--gs_render_resolution", type=int, default=512)

    p.add_argument("--cuda_idx", type=int, nargs="*", default=[0])
    p.add_argument("--world_size", type=int, default=1)
    p.add_argument("--rank", type=int, default=0)

    p.add_argument("--start_idx", type=int, default=None)
    p.add_argument("--end_idx", type=int, default=None)

    p.add_argument("--debug", action="store_true")
    p.add_argument("--print_time", action="store_true")
    p.add_argument(
        "--empty_structure_fallback",
        type=str,
        default="original",
        choices=["error", "original"],
    )
    p.add_argument(
        "--force",
        action="store_true",
        help="Recompute GS views even if render_gs/transforms.json exists.",
    )
    return p


def main() -> None:
    args = build_argparser().parse_args()

    args.data_info_json_path = os.path.join(args.dataset_root_dir, args.data_info_json_path)
    args.select_json_path = os.path.join(args.dataset_root_dir, args.select_json_path)
    args.flux_edit_root_path = os.path.join(args.dataset_root_dir, args.flux_edit_root_path)
    args.alpaca_root_path = os.path.join(args.dataset_root_dir, args.alpaca_root_path)
    args.mixamo_root_path = os.path.join(args.dataset_root_dir, args.mixamo_root_path)

    set_seed(42)
    assert len(args.cuda_idx) == 1

    save_path = os.path.join(os.path.abspath(args.output_root), args.run_name)

    with open(args.select_json_path, "r") as f:
        select_data = json.load(f)
    with open(args.data_info_json_path, "r") as f:
        data_info = json.load(f)

    processed_keys = []
    for key, values in select_data.items():
        for value in values:
            processed_keys.append((key, value))

    if args.start_idx is not None and args.end_idx is not None:
        start_idx, end_idx = args.start_idx, args.end_idx
        processed_keys = processed_keys[start_idx:end_idx]
    else:
        chunk_size = math.ceil(len(processed_keys) / args.world_size)
        start_idx = args.rank * chunk_size
        end_idx = min((args.rank + 1) * chunk_size, len(processed_keys))
        processed_keys = processed_keys[start_idx:end_idx]

    print("processed keys:", len(processed_keys), "->", save_path)
    if args.debug:
        processed_keys = processed_keys[:6]

    model = TrellisEdititngModel(
        ss_latents_config_path=f"./configs/editing/{args.ss_latents_config}",
        ss_latents_load_dir=args.ss_latents_load_dir,
        ss_latents_load_ckpt=args.ss_latents_load_ckpt,
        latents_config_path=f"./configs/editing/{args.latents_config}",
        latents_load_dir=args.latents_load_dir,
        latents_load_ckpt=args.latents_load_ckpt,
        load_ema_model_for_inference=args.load_ema_model_for_inference,
        trellis_pipeline_path=args.trellis_pipeline_path,
    )
    print("init TrellisEdititngModel done")

    for data in tqdm.tqdm(processed_keys):
        dataset_type = data[0]
        if dataset_type == "alpaca":
            key = data[1]
            ori_ss_latents_path = os.path.join(
                args.alpaca_root_path, data_info[dataset_type][key]["ori_ss_latents_path"]
            )
            ori_latents_path = os.path.join(
                args.alpaca_root_path, data_info[dataset_type][key]["ori_latents_path"]
            )
            ori_img_path = os.path.join(args.alpaca_root_path, data_info[dataset_type][key]["ori_img_path"])
            edit_img_path = os.path.join(args.alpaca_root_path, data_info[dataset_type][key]["edit_img_path"])
        elif dataset_type == "flux_edit":
            key = data[1]
            ori_ss_latents_path = os.path.join(
                args.flux_edit_root_path, data_info[dataset_type][key]["ori_ss_latents_path"]
            )
            ori_latents_path = os.path.join(
                args.flux_edit_root_path, data_info[dataset_type][key]["ori_latents_path"]
            )
            ori_img_path = os.path.join(args.flux_edit_root_path, data_info[dataset_type][key]["ori_img_path"])
            edit_img_path = os.path.join(args.flux_edit_root_path, data_info[dataset_type][key]["edit_img_path"])
        elif dataset_type == "mixamo":
            character_name, ori_idx, edit_idx = data[1]
            key = f"{character_name}_{ori_idx}_{edit_idx}"
            ori_ss_latents_path = os.path.join(
                args.mixamo_root_path,
                data_info[dataset_type][character_name][ori_idx]["ss_latents_path"],
            )
            ori_latents_path = os.path.join(
                args.mixamo_root_path,
                data_info[dataset_type][character_name][ori_idx]["latents_path"],
            )
            ori_img_path = os.path.join(
                args.mixamo_root_path, data_info[dataset_type][character_name][ori_idx]["img_path"]
            )
            edit_img_path = os.path.join(
                args.mixamo_root_path, data_info[dataset_type][character_name][edit_idx]["img_path"]
            )
        else:
            raise ValueError(f"Invalid dataset type: {dataset_type}")

        mesh_save_path = os.path.join(save_path, dataset_type, key, "edit.glb")
        gs_render_marker = os.path.join(os.path.dirname(mesh_save_path), "render_gs", "transforms.json")
        os.makedirs(os.path.dirname(mesh_save_path), exist_ok=True)

        if os.path.exists(gs_render_marker) and not args.force:
            continue

        t0 = time.time()
        model.editing_inference(
            ori_ss_latents_path=ori_ss_latents_path,
            ori_latents_path=ori_latents_path,
            edited_img_path=edit_img_path,
            ori_img_path=ori_img_path,
            output_video=False,
            output_mesh=False,
            render_gs_views=True,
            gs_render_total_view_num=args.total_render_view_num,
            gs_render_eval_view_num=args.eval_view_num,
            gs_render_resolution=args.gs_render_resolution,
            video_save_path=None,
            mesh_save_path=mesh_save_path,
            slat_save_path=None,
            ori_latents_norm=True if dataset_type == "mixamo" else False,
            print_time=args.print_time,
            empty_structure_fallback=args.empty_structure_fallback,
        )
        if args.print_time:
            print(f"sample {key} time: {time.time() - t0:.3f}s")


if __name__ == "__main__":
    main()