#!/usr/bin/env bash
# mixed_train.sh — mixed-opponent training for generalist PPO_Bot policy.
# Randomly picks an opponent each chunk to prevent catastrophic forgetting.
set -uo pipefail

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

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}"

SNAPSHOTS_DIR="$SCRIPT_DIR/snapshots"
mkdir -p "$SNAPSHOTS_DIR"

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

# ── config ────────────────────────────────────────────────────────────────────
OPPONENTS=(Fire MyFirstDroid Target MyFirstLeader Crazy Corners PaintingBot MyFirstBot)
TOTAL_ROUNDS=10000
CHUNK_SIZE=60
EVAL_INTERVAL=500
EVAL_ROUNDS=50
PASS_THRESHOLD=45

# ── compile once ──────────────────────────────────────────────────────────────
echo "=== PPO_Bot mixed training — $(date) ==="
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 ───────────────────────────────────────────────────────────────────

count_wins() {
  local f="$1"
  [ -f "$f" ] || { echo 0; return; }
  awk '/"win":true/{c++} END{print c+0}' "$f"
}

# run_chunk <opponent> <rounds> [eval]
# Sets PPOB_EVAL_ONLY if third arg is "eval". Retries up to 3 times on failure.
run_chunk() {
  local opp="$1" rounds="$2" mode="${3:-train}"
  local logfile
  if [ "$mode" = "eval" ]; then
    logfile="/tmp/${opp}_eval.jsonl"
    export PPOB_EVAL_ONLY=1
  else
    logfile="/tmp/${opp}_mixed_train.jsonl"
    unset PPOB_EVAL_ONLY
  fi
  export TRAINING_OPPONENT="$opp"
  export TRAINING_ROUNDS="$rounds"
  export PPOB_LOG_FILE="$logfile"
  > "$logfile"

  # Reset counter so java knows when it's done
  echo 0 > "$ROUND_COUNTER_FILE"

  local attempt=0
  while [ "$attempt" -lt 3 ]; do
    attempt=$(( attempt + 1 ))
    [ "$attempt" -gt 1 ] && sleep 2
    java -cp "$SCRIPT_DIR:$JAR" RunTraining "$opp" "$rounds" && break
    echo ">>> Java exited non-zero (attempt $attempt/3)"
  done

  if [ "$attempt" -ge 3 ]; then
    local done=0
    [ -f "$ROUND_COUNTER_FILE" ] && done=$(cat "$ROUND_COUNTER_FILE" 2>/dev/null || echo 0)
    if [ "$done" -eq 0 ]; then
      echo "WARNING: chunk $opp made no progress after 3 attempts — skipping"
      unset PPOB_EVAL_ONLY
      return 1
    fi
  fi

  unset PPOB_EVAL_ONLY
  echo "$logfile"
}

evaluate_all() {
  local trained_rounds="$1"
  local total_wins=0
  local total_possible=$(( ${#OPPONENTS[@]} * EVAL_ROUNDS ))
  local scoreboard="[EVAL @ round $trained_rounds]"

  for opp in "${OPPONENTS[@]}"; do
    local logfile
    logfile=$(run_chunk "$opp" "$EVAL_ROUNDS" eval) || { scoreboard="$scoreboard ${opp}:ERR"; continue; }
    local wins
    wins=$(count_wins "$logfile")
    total_wins=$(( total_wins + wins ))
    scoreboard="$scoreboard ${opp}:${wins}/${EVAL_ROUNDS}"
  done

  scoreboard="$scoreboard Total:${total_wins}/${total_possible}"
  echo "$scoreboard"
  echo "$total_wins"
}

snap_weights() {
  local trained_rounds="$1"
  local dst="$SNAPSHOTS_DIR/mixed_r${trained_rounds}"
  cp -a "$PPO_BOT_SRC/weights/latest" "$dst"
  echo "  snapshot → $dst"
}

random_opponent() {
  local idx=$(( RANDOM % ${#OPPONENTS[@]} ))
  echo "${OPPONENTS[$idx]}"
}

# ── training loop ─────────────────────────────────────────────────────────────
trained_rounds=0
best_total_wins=-1
n_opponents=${#OPPONENTS[@]}
total_possible=$(( n_opponents * EVAL_ROUNDS ))

echo "Config: TOTAL_ROUNDS=$TOTAL_ROUNDS CHUNK_SIZE=$CHUNK_SIZE EVAL_INTERVAL=$EVAL_INTERVAL EVAL_ROUNDS=$EVAL_ROUNDS PASS_THRESHOLD=$PASS_THRESHOLD"
echo "Opponents: ${OPPONENTS[*]}"

while [ "$trained_rounds" -lt "$TOTAL_ROUNDS" ]; do
  opp=$(random_opponent)
  echo "--- [round $trained_rounds] training vs $opp ($CHUNK_SIZE rounds)"
  run_chunk "$opp" "$CHUNK_SIZE" train || true
  trained_rounds=$(( trained_rounds + CHUNK_SIZE ))

  if [ $(( trained_rounds % EVAL_INTERVAL )) -eq 0 ]; then
    # evaluate_all prints scoreboard then total_wins on separate lines
    mapfile -t eval_out < <(evaluate_all "$trained_rounds")
    echo "${eval_out[0]}"
    total_wins="${eval_out[1]:-0}"
    if [ "$total_wins" -gt "$best_total_wins" ]; then
      best_total_wins="$total_wins"
      snap_weights "$trained_rounds"
    fi
  fi
done

echo ""
echo "=== Final evaluation — $(date) ==="
mapfile -t eval_out < <(evaluate_all "$trained_rounds")
echo "${eval_out[0]}"
total_wins="${eval_out[1]:-0}"
if [ "$total_wins" -gt "$best_total_wins" ]; then
  snap_weights "$trained_rounds"
fi
echo "=== mixed training complete — $(date) ==="