#!/bin/bash
if [[ -f .env ]]; then
    set -a
    source .env
    set +a
fi
# --- Fixed config (edit if you like) ---
MODEL_DEFAULT="mistralai/Mistral-7B-Instruct-v0.2"
PROMPT_TEMPLATES=("io_prompt" "zsr_prompt")
ENCODINGS=("many_to_one" "one_to_one" "one_to_many")  # will expand to one_to_many_2..5 for dirs
SEED=42

# --- Args: only dataset_dir is required; model is optional override ---
DATASET_DIR=""
MODEL="$MODEL_DEFAULT"

while [[ $# -gt 0 ]]; do
  case "$1" in
    --dataset_dir) DATASET_DIR="$2"; shift 2 ;;
    --model) MODEL="$2"; shift 2 ;;
    -h|--help)
      echo "Usage: $0 --dataset_dir <path> [--model <hf-model>]"
      exit 0
      ;;
    *)
      echo "Unknown arg: $1"
      echo "Usage: $0 --dataset_dir <path> [--model <hf-model>]"
      exit 1
      ;;
  esac
done

if [[ -z "$DATASET_DIR" ]]; then
  echo "Error: --dataset_dir is required"
  exit 1
fi

# --- Discover languages from dataset dir (needs 4 files present) ---
mapfile -t LANGUAGES < <(
  find "$DATASET_DIR" -mindepth 1 -maxdepth 1 -type d \
  | while read -r d; do
      if [[ -f "$d/data.train" && -f "$d/labels.train" && -f "$d/data.test" && -f "$d/labels.test" ]]; then
        basename "$d"
      fi
    done | sort
)

if [[ ${#LANGUAGES[@]} -eq 0 ]]; then
  echo "No languages found in $DATASET_DIR"
  exit 1
fi

# --- Create experiment folders ---
SAFE_MODEL="${MODEL//\//_}"
for pt in "${PROMPT_TEMPLATES[@]}"; do
  for enc in "${ENCODINGS[@]}"; do
    if [[ "$enc" == "one_to_many" ]]; then
      for k in 2 3 4 5; do
        for lang in "${LANGUAGES[@]}"; do
          mkdir -p "experiments/$SAFE_MODEL/$pt/one_to_many_${k}/$lang"
        done
      done
    else
      for lang in "${LANGUAGES[@]}"; do
        mkdir -p "experiments/$SAFE_MODEL/$pt/$enc/$lang"
      done
    fi
  done
done

# --- Run the Python driver (it loops over everything internally) ---
# Ensure local imports work
# export PYTHONPATH=.

# # Optionally pass dataset dir to your script via env var (if your code reads DATASET_DIR)
# export DATASET_DIR="$DATASET_DIR"

echo "Running Python with model=$MODEL, seed=$SEED"
eval python -m scripts.experiment_all_in_one \
  --model "$MODEL" \
  --seed "$SEED"
