#!/usr/bin/env bash
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)"
EXP_ROOT="${ROOT}/experiments/ablation"
RUNNER="${ROOT}/run_bioagent_bench.py"
EVALUATOR="${ROOT}/evaluate_bioagent_bench.py"
BIOAGENT_BENCH_ROOT="$(cd "${ROOT}/.." && pwd)/bioagent-bench"
PYTHON_BIN="${BIOMANUS_PYTHON:-/225040511/miniconda3/envs/biomni_e1/bin/python}"
GRAPH_DIR="${ROOT}/graph_outputs/mcp_generated_graph_all_20260522_124110"

BACKGROUND=0
VARIANT="all"
FORCE=0
declare -a TASK_ARGS=("--all")
TASK_MODE_SET=0

while [[ $# -gt 0 ]]; do
  case "$1" in
    --background) BACKGROUND=1; shift ;;
    --foreground) BACKGROUND=0; shift ;;
    --variant) VARIANT="$2"; shift 2 ;;
    --task)
      if [[ "${TASK_MODE_SET}" -eq 0 ]]; then
        TASK_ARGS=()
        TASK_MODE_SET=1
      fi
      TASK_ARGS+=("--task" "$2")
      shift 2
      ;;
    --force-rerun) FORCE=1; shift ;;
    *) echo "Unknown argument: $1" >&2; exit 2 ;;
  esac
done

if [[ "${BACKGROUND}" -eq 1 ]]; then
  mkdir -p "${EXP_ROOT}/logs"
  LOG="${EXP_ROOT}/logs/bioagentbench_ablation_$(date -u +%Y%m%d_%H%M%S).log"
  SCRIPT="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)/$(basename "${BASH_SOURCE[0]}")"
  CMD=("${SCRIPT}" "--foreground" "--variant" "${VARIANT}")
  if [[ "${FORCE}" -eq 1 ]]; then CMD+=("--force-rerun"); fi
  nohup "${CMD[@]}" > "${LOG}" 2>&1 < /dev/null &
  echo "Started BioAgentBench ablation in background."
  echo "PID: $!"
  echo "Log: ${LOG}"
  exit 0
fi

if [[ -f "${ROOT}/.env" ]]; then
  # shellcheck disable=SC1090
  source "${ROOT}/.env"
fi

if [[ -n "${DEEPSEEK_API_KEY:-}" ]]; then
  export BIOMNI_LLM_PROVIDER="${BIOMNI_LLM_PROVIDER:-deepseek}"
  export DEEPSEEK_BASE_URL="${DEEPSEEK_BASE_URL:-https://api.deepseek.com/v1}"
  export DEEPSEEK_MODEL_NAME="${DEEPSEEK_MODEL_NAME:-deepseek-chat}"
  export BIOMNI_SOURCE="${BIOMNI_SOURCE:-Custom}"
  export BIOMNI_LLM="${BIOMNI_LLM:-${DEEPSEEK_MODEL_NAME}}"
  export BIOMNI_CUSTOM_BASE_URL="${BIOMNI_CUSTOM_BASE_URL:-${DEEPSEEK_BASE_URL}}"
  export BIOMNI_CUSTOM_API_KEY="${BIOMNI_CUSTOM_API_KEY:-${DEEPSEEK_API_KEY}}"
fi
if [[ -n "${BIOMNI_CUSTOM_BASE_URL:-}" && ! "${BIOMNI_CUSTOM_BASE_URL}" =~ ^https?:// ]]; then
  export BIOMNI_CUSTOM_BASE_URL="${DEEPSEEK_BASE_URL:-https://api.deepseek.com/v1}"
fi
if [[ -n "${DEEPSEEK_BASE_URL:-}" && ! "${DEEPSEEK_BASE_URL}" =~ ^https?:// ]]; then
  export DEEPSEEK_BASE_URL="https://api.deepseek.com/v1"
fi

run_variant() {
  local key="$1"
  local label="$2"
  shift 2
  local runs_root="${EXP_ROOT}/results/${key}/bioagentbench"
  local eval_json="${EXP_ROOT}/results/${key}/bioagentbench_evaluation.json"
  mkdir -p "${runs_root}"

  echo "=== BioAgentBench ${label} ==="
  "${PYTHON_BIN}" "${RUNNER}" \
    "${TASK_ARGS[@]}" \
    --output-root "${runs_root}" \
    --mcp-graph "${GRAPH_DIR}" \
    --executable-mcp-only \
    --timeout-seconds "${BIOMANUS_BIOAGENT_TIMEOUT_SECONDS:-1200}" \
    "$@"

  "${PYTHON_BIN}" "${EVALUATOR}" \
    --all \
    --runs-root "${runs_root}" \
    --dataset-root "${BIOAGENT_BENCH_ROOT}/dataset" \
    --judge-mode rule \
    --output "${eval_json}"
}

case "${VARIANT}" in
  biomanus)
    run_variant "biomanus" "full" --use-graph-retriever --use-mcp --mcp-retrieval-mode graph
    ;;
  mcp_flat)
    run_variant "mcp_flat" "MCP + flat retrieval" --use-graph-retriever --use-mcp --mcp-retrieval-mode flat
    ;;
  mcp_metadata)
    run_variant "mcp_metadata" "MCP + metadata retrieval" --use-graph-retriever --use-mcp --mcp-retrieval-mode metadata
    ;;
  minus_graph)
    run_variant "minus_graph" "minus graph" --no-graph-retriever --use-mcp
    ;;
  minus_mcp)
    run_variant "minus_mcp" "minus MCP" --use-graph-retriever --no-mcp
    ;;
  minus_mcp_graph)
    run_variant "minus_mcp_graph" "minus MCP and graph" --no-graph-retriever --no-mcp
    ;;
  all)
    run_variant "biomanus" "full" --use-graph-retriever --use-mcp --mcp-retrieval-mode graph
    run_variant "mcp_flat" "MCP + flat retrieval" --use-graph-retriever --use-mcp --mcp-retrieval-mode flat
    run_variant "mcp_metadata" "MCP + metadata retrieval" --use-graph-retriever --use-mcp --mcp-retrieval-mode metadata
    run_variant "minus_graph" "minus graph" --no-graph-retriever --use-mcp
    run_variant "minus_mcp" "minus MCP" --use-graph-retriever --no-mcp
    run_variant "minus_mcp_graph" "minus MCP and graph" --no-graph-retriever --no-mcp
    ;;
  *)
    echo "Unknown variant: ${VARIANT}" >&2
    exit 2
    ;;
esac