#!/usr/bin/env bash
# curriculum_retrain.sh — sequential curriculum retraining for PPO_Bot v2.
# Bots in order: Fire → MyFirstDroid → Target → MyFirstLeader → Crazy → Corners → PaintingBot → MyFirstBot
set -uo pipefail

SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
REPO_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)"
PPO_BOT_SRC="$REPO_ROOT/PPO_Bot_garage"

# Source hyperparams (but we'll override opponent/rounds/log per phase)
ENV_FILE="$SCRIPT_DIR/training.env"
[ -f "$ENV_FILE" ] && set -a && . "$ENV_FILE" && set +a

JAR="${TANK_ROYALE_JAR:-/home/davide/Projects/tank-royale/runner/examples/lib/robocode-tankroyale-runner.jar}"
export PPO_BOT_DIR="$PPO_BOT_SRC"
export SAMPLE_BOTS_DIR="${SAMPLE_BOTS_DIR:-/home/davide/Projects/tank-royale/sample-bots/java/build/archive}"

LOGS_DIR="$SCRIPT_DIR/logs"
SNAPSHOTS_DIR="$SCRIPT_DIR/snapshots"
SUMMARY="$LOGS_DIR/curriculum_retrain_summary.txt"
mkdir -p "$LOGS_DIR" "$SNAPSHOTS_DIR"

ROUND_COUNTER_FILE="$PPO_BOT_SRC/weights/round_counter.txt"

echo "=== PPO_Bot v2 curriculum retrain — $(date) ===" | tee -a "$SUMMARY"

# ── compile once up front (run.sh does this; curriculum skipped it) ───────────
echo ">>> Compiling PPO_Bot..."
(cd "$PPO_BOT_SRC" && nimble build -d:release)
echo ">>> Compiling RunTraining.java..."
(cd "$SCRIPT_DIR" && javac -cp "$JAR" RunTraining.java)
echo ">>> Compilation done."

# ── helpers ──────────────────────────────────────────────────────────────────

run_java() {
  # run_java <opponent> <rounds>
  local opp="$1" rounds="$2"
  local remaining="$rounds"
  local attempts=0
  while [ "$remaining" -gt 0 ]; do
    attempts=$(( attempts + 1 ))
    [ "$attempts" -gt 1 ] && sleep 2
    java -cp "$SCRIPT_DIR:$JAR" RunTraining "$opp" "$remaining" || \
      echo ">>> Java exited non-zero (crash/restart expected)"
    local done=0
    [ -f "$ROUND_COUNTER_FILE" ] && done=$(cat "$ROUND_COUNTER_FILE" 2>/dev/null || echo 0)
    remaining=$(( rounds - done ))
    # Guard: if counter didn't advance at all on this attempt, something is
    # fundamentally broken (bad binary, missing bots dir, etc.) — abort rather
    # than loop forever.
    if [ "$attempts" -ge 3 ] && [ "$done" -eq 0 ]; then
      echo "ERROR: round_counter still 0 after $attempts attempts — aborting" >&2
      exit 1
    fi
  done
}

count_wins() {
  # count_wins <jsonl_file> → prints win count
  # grep -c exits 1 on zero matches (not an error) which was triggering || echo 0
  # making WINS="0\n0". Use awk which always exits 0.
  local f="$1"
  [ -f "$f" ] || { echo 0; return; }
  awk '/"win":true/{c++} END{print c+0}' "$f"
}

snap_weights() {
  local label="$1"
  local counter
  counter=$(cat "$ROUND_COUNTER_FILE" 2>/dev/null || echo 0)
  local dst="$SNAPSHOTS_DIR/v2_${label}_r${counter}"
  cp -a "$PPO_BOT_SRC/weights/latest" "$dst"
  echo "  snapshot → $dst"
}

# ── curriculum ────────────────────────────────────────────────────────────────

declare -A TRAIN_ROUNDS=( [Fire]=200 [MyFirstDroid]=200 [Target]=200 [MyFirstLeader]=200 [Crazy]=200
                          [Corners]=400 [PaintingBot]=400 [MyFirstBot]=1000 )
BOTS=(Fire MyFirstDroid Target MyFirstLeader Crazy Corners PaintingBot MyFirstBot)
GATE=150
GATE_PASS=148

for BOT in "${BOTS[@]}"; do
  TR="${TRAIN_ROUNDS[$BOT]}"
  echo "" | tee -a "$SUMMARY"
  echo "══ $BOT  (train=${TR} gate=${GATE}) ══" | tee -a "$SUMMARY"

  # Reset round counter so run_java terminates correctly for each phase
  echo 0 > "$ROUND_COUNTER_FILE"

  # ── training phase ──────────────────────────────────────────────────────────
  TRAIN_LOG="$LOGS_DIR/${BOT}_train.jsonl"
  export TRAINING_OPPONENT="$BOT"
  export TRAINING_ROUNDS="$TR"
  export PPOB_LOG_FILE="$TRAIN_LOG"
  unset PPOB_EVAL_ONLY
  > "$TRAIN_LOG"
  echo "  [train] $TR rounds..." | tee -a "$SUMMARY"
  run_java "$BOT" "$TR"
  echo "  [train] done" | tee -a "$SUMMARY"

  # ── gate eval ───────────────────────────────────────────────────────────────
  echo 0 > "$ROUND_COUNTER_FILE"
  EVAL_LOG="$LOGS_DIR/${BOT}_eval.jsonl"
  export TRAINING_ROUNDS="$GATE"
  export PPOB_LOG_FILE="$EVAL_LOG"
  export PPOB_EVAL_ONLY=1
  > "$EVAL_LOG"
  echo "  [eval]  $GATE rounds..." | tee -a "$SUMMARY"
  run_java "$BOT" "$GATE"

  WINS=$(count_wins "$EVAL_LOG")
  if [ "$WINS" -ge "$GATE_PASS" ]; then
    STATUS="PASS ($WINS/$GATE)"
    snap_weights "$BOT"
  else
    STATUS="FAIL ($WINS/$GATE)"
  fi
  echo "  [gate]  $STATUS" | tee -a "$SUMMARY"
  unset PPOB_EVAL_ONLY
done

echo "" | tee -a "$SUMMARY"
echo "=== curriculum complete — $(date) ===" | tee -a "$SUMMARY"