#!/usr/bin/env bash
# Train (or resume) one of the IC-LoRAs in configs/.
#
# Usage:
#   scripts/train_ic_lora.sh --config configs/v2v_reference_ic_lora.yaml
#   scripts/train_ic_lora.sh --config configs/ref_image_ic_lora.yaml
#
# Resolves the __REPO_ROOT__ placeholder in the chosen config against this
# repo's actual location (so the config works no matter where you cloned it),
# writes the resolved copy to a temp file, then launches training via
# `accelerate launch` (auto-detects GPU count) from packages/ltx-trainer/.
set -euo pipefail

REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
CONFIG="configs/v2v_reference_ic_lora.yaml"

while [[ $# -gt 0 ]]; do
  case "$1" in
    --config) CONFIG="$2"; shift 2 ;;
    *) echo "Unknown argument: $1" >&2; exit 1 ;;
  esac
done

CONFIG_ABS="$REPO_ROOT/$CONFIG"
if [[ "$CONFIG" == /* ]]; then CONFIG_ABS="$CONFIG"; fi
if [[ ! -f "$CONFIG_ABS" ]]; then
  echo "Config not found: $CONFIG_ABS" >&2
  exit 1
fi

RESOLVED_CONFIG="$(mktemp --suffix=.yaml)"
sed "s|__REPO_ROOT__|$REPO_ROOT|g" "$CONFIG_ABS" > "$RESOLVED_CONFIG"
trap 'rm -f "$RESOLVED_CONFIG"' EXIT

echo "[train_ic_lora] repo root:      $REPO_ROOT"
echo "[train_ic_lora] config:         $CONFIG_ABS"
echo "[train_ic_lora] resolved to:    $RESOLVED_CONFIG"

NUM_GPUS="$(python3 -c 'import torch; print(torch.cuda.device_count())' 2>/dev/null || echo 0)"
echo "[train_ic_lora] detected GPUs:  $NUM_GPUS"

cd "$REPO_ROOT/packages/ltx-trainer"

if [[ "$NUM_GPUS" -le 1 ]]; then
  echo "[train_ic_lora] single-GPU/CPU run -> python scripts/train.py"
  exec python scripts/train.py "$RESOLVED_CONFIG"
else
  echo "[train_ic_lora] $NUM_GPUS GPUs -> accelerate launch (DDP)"
  exec accelerate launch \
    --multi_gpu \
    --num_processes "$NUM_GPUS" \
    --mixed_precision bf16 \
    scripts/train.py "$RESOLVED_CONFIG"
fi