#!/usr/bin/env bash

set -euo pipefail

if [[ $# -ne 4 ]]; then
    echo "Usage: $0 <res_dir> <dataset> <frame_stride> <n_iter>" >&2
    exit 1
fi

res_dir="$1"
dataset="$2"
frame_stride="$3"
n_iter="$4"

repo_root="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
cd "${repo_root}"

gpu_devices="${CUDA_VISIBLE_DEVICES:-1}"
python_bin="${PYTHON_BIN:-}"
if [[ -z "${python_bin}" ]]; then
    if [[ -n "${CONDA_PREFIX:-}" ]] && [[ -x "${CONDA_PREFIX}/bin/python" ]]; then
        python_bin="${CONDA_PREFIX}/bin/python"
    elif [[ -n "${VIRTUAL_ENV:-}" ]] && [[ -x "${VIRTUAL_ENV}/bin/python" ]]; then
        python_bin="${VIRTUAL_ENV}/bin/python"
    else
        python_bin="$(command -v python || true)"
    fi
fi
if [[ -z "${python_bin}" ]] || [[ ! -x "${python_bin}" ]]; then
    echo "Unable to resolve a usable Python interpreter. Set PYTHON_BIN explicitly." >&2
    exit 1
fi
save_dir="${res_dir}/backbone_profile_dense"
log_path="${save_dir}/run.log"

mkdir -p "${save_dir}"

interaction_cmd=(
    "${python_bin}" scripts/evaluation/world_model_interaction.py
    --seed 123
    --ckpt_path ckpts/unifolm_wma_dual.ckpt
    --config configs/inference/world_model_interaction.yaml
    --savedir "${save_dir}"
    --bs 1
    --height 320
    --width 512
    --unconditional_guidance_scale 1.0
    --ddim_steps 50
    --ddim_eta 1.0
    --prompt_dir "${res_dir}/world_model_interaction_prompts"
    --dataset "${dataset}"
    --video_length 16
    --frame_stride "${frame_stride}"
    --n_action_steps 16
    --exe_steps 16
    --n_iter "${n_iter}"
    --timestep_spacing uniform_trailing
    --guidance_rescale 0.7
    --perframe_ae
    --analysis_log_metrics
    --analysis_reference_steps 50
    --analysis_profile_backbone_blocks
)

metrics_analysis_cmd=(
    "${python_bin}" scripts/evaluation/analyze_metrics.py
    --input_dir "${save_dir}/inference"
    --output_dir "${save_dir}/inference/analysis"
)

backbone_analysis_cmd=(
    "${python_bin}" scripts/evaluation/analyze_backbone_blocks.py
    --input_csv "${save_dir}/inference/backbone_block_log.csv"
    --output_dir "${save_dir}/inference/backbone_analysis"
)

{
    echo "============================================================"
    echo "Running dense backbone profiling for ${res_dir} on GPU ${gpu_devices}"
    echo "============================================================"
    echo "Using Python: ${python_bin}"
    time env CUDA_VISIBLE_DEVICES="${gpu_devices}" "${interaction_cmd[@]}"

    "${metrics_analysis_cmd[@]}"
    "${backbone_analysis_cmd[@]}"
} 2>&1 | tee "${log_path}"