#!/usr/bin/env sh
set -eu

export CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-0}

PYTHON_BIN=${PYTHON_BIN:-/opt/miniconda/envs/focallora/bin/python}
MODEL_PATH=${1:-${MODEL_PATH:-}}
LORA_PATH=${2:-${LORA_PATH:-}}
INPUT_PATH=${3:-${INPUT_PATH:-./Ident_IH_05_visualize_any_testmsg.txt}}
OUTPUT_PATH=${4:-${OUTPUT_PATH:-./Ident_IH_05_visualize_any_output.png}}

if [ ! -x "$PYTHON_BIN" ]; then
  echo "[Ident_IH_05_visualize_any.sh] Python not found: $PYTHON_BIN" >&2
  exit 1
fi

run_one() {
  run_model_path=$1
  run_lora_path=$2
  run_input_path=$3
  run_output_path=$4
  run_title=${5:-}
  run_head_order_path=${6:-}

  if [ ! -d "$run_model_path" ]; then
    echo "[Ident_IH_05_visualize_any.sh] Model path not found: $run_model_path" >&2
    exit 1
  fi
  if [ -n "$run_lora_path" ] && [ ! -d "$run_lora_path" ]; then
    echo "[Ident_IH_05_visualize_any.sh] LoRA path not found: $run_lora_path" >&2
    exit 1
  fi
  if [ ! -f "$run_input_path" ]; then
    echo "[Ident_IH_05_visualize_any.sh] Input file not found: $run_input_path" >&2
    exit 1
  fi
  if [ -n "$run_head_order_path" ] && [ ! -f "$run_head_order_path" ]; then
    echo "[Ident_IH_05_visualize_any.sh] Head-order file not found: $run_head_order_path" >&2
    exit 1
  fi

  echo "[run] model=$run_model_path"
  echo "[run] lora=${run_lora_path:-<off>}"
  echo "[run] input=$run_input_path"
  echo "[run] output=$run_output_path"
  echo "[run] head_order=${run_head_order_path:-<default>}"

  set -- \
    "$PYTHON_BIN" ./Ident_IH_05_visualize_any.py \
    --model-path "$run_model_path" \
    --lora-path "$run_lora_path" \
    --input-path "$run_input_path" \
    --output-path "$run_output_path"
  if [ -n "$run_title" ]; then
    set -- "$@" --title "$run_title"
  fi
  if [ -n "$run_head_order_path" ]; then
    set -- "$@" --head-order-path "$run_head_order_path"
  fi
  "$@"
}

if [ -n "$MODEL_PATH" ]; then
  run_one "$MODEL_PATH" "$LORA_PATH" "$INPUT_PATH" "$OUTPUT_PATH"
  exit 0
fi

CODE_DIR=../code
IGNORE_INPUT=$CODE_DIR/Ident_IH_05_visualize_any_ignore_attk.txt
CONV_INPUT=$CODE_DIR/Ident_IH_05_visualize_any_conv_attack_attk.txt

for required_input in "$IGNORE_INPUT" "$CONV_INPUT"; do
  if [ ! -f "$required_input" ]; then
    echo "[Ident_IH_05_visualize_any.sh] Required batch input missing: $required_input" >&2
    exit 1
  fi
done

while IFS='|' read -r model_tag model_path lora_path head_order_path; do
  [ -n "$model_tag" ] || continue

  while IFS='|' read -r lora_state effective_lora; do
    [ -n "$lora_state" ] || continue
    if [ "$lora_state" = "on" ]; then
      current_lora=$lora_path
    else
      current_lora=
    fi

    while IFS='|' read -r attack_tag attack_input; do
      [ -n "$attack_tag" ] || continue
      output_file=$CODE_DIR/Ident_IH_05_visualize_any_${model_tag}_lora-${lora_state}_${attack_tag}.png
      title="${model_tag} | lora=${lora_state} | ${attack_tag}"
      run_one "$model_path" "$current_lora" "$attack_input" "$output_file" "$title" "$head_order_path"
    done <<EOF_ATTACK
ignore|$IGNORE_INPUT
conv_attack|$CONV_INPUT
EOF_ATTACK
  done <<EOF_LORA
off|
on|$lora_path
EOF_LORA
done <<'EOF_MODEL'
qwen3-4b|../../models/Qwen3-4B|../3-2_model_training/lora/qwen3-4b_sep_tool_simple_1e-4_orig/batch_14_99|../2-2_head_identification/head_scoring/qwen3-4b_sep/heads_sorted/user_prop_inst_0.1.json
qwen3-8b|../../models/Qwen3-8B|../3-2_model_training/lora/qwen3-8b_sep_tool_simple_debug2_1e-4_orig/batch_0_562|../2-2_head_identification/head_scoring/qwen3-8b_sep/heads_sorted/all_roc_inst_0.1.json
llama31-8b|../../models/Llama-3.1-8B-Instruct|../3-2_model_training/lora/llama31-8b_sep_tool_simple_debug/batch_1_562|../2-2_head_identification/head_scoring/llama31-8b_sep/heads_sorted/all_roc_inst_0.1.json
EOF_MODEL