#!/usr/bin/env bash
# generalist_train.sh — adaptive mixed-opponent training for a generalist PPO_Bot.
#
# Problem it solves: the sequential curriculum killed earlier skills (catastrophic
# forgetting). Instead we sample ONE opponent per chunk, weighted by how weak the
# bot currently is against it (max(FLOOR, 1 - winrate)), and keep the best
# generalist checkpoint by mean eval winrate across all 8 opponents.
#
# State / logs (all plain files, resumable):
#   /tmp/generalist_wr.txt        — one line per opponent: "Name winrate"
#   /tmp/generalist_seed.txt      — optional starting winrates (same format);
#                                    wins over wr.txt when both exist
#   /tmp/generalist_best.txt      — "mean rounds" of the best generalist eval so far
#   /tmp/generalist_train.log     — full run log (tee'd)
#   /tmp/generalist_eval_<Opp>.jsonl — per-opponent eval battles (fresh per eval)
#   /tmp/generalist_train_<Opp>.jsonl — per-opponent training battles (accumulated)
#
# Env-overridable config: TOTAL_ROUNDS CHUNK_SIZE EVAL_INTERVAL EVAL_ROUNDS
#   PASS_RATE FLOOR ALPHA (all declared in the config block below).
#
# Usage:
#   ./generalist_train.sh            — run the full training loop
#   ./generalist_train.sh --selftest — verify the weighted sampler against the
#                        expected distribution (no battles, no compilation)
#
# Dead-bot resilience is inherited from the harness: RunTraining.java aborts
# (exit 1) when PPO_Bot's round counter is frozen or short at battle end, and
# this script retries; a chunk that makes no counter progress across 3 attempts
# aborts loudly instead of silently burning rounds.
set -uo pipefail
export LC_ALL=C

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"
LOG="/tmp/generalist_train.log"
WR_STATE_FILE="/tmp/generalist_wr.txt"
WR_SEED_FILE="/tmp/generalist_seed.txt"
BEST_FILE="/tmp/generalist_best.txt"
SELFTEST_LOG="/tmp/generalist_sampling_selftest.log"

# ── config (env-overridable) ──────────────────────────────────────────────────
TOTAL_ROUNDS="${TOTAL_ROUNDS:-40000}"
CHUNK_SIZE="${CHUNK_SIZE:-20}"
EVAL_INTERVAL="${EVAL_INTERVAL:-1000}"
EVAL_ROUNDS="${EVAL_ROUNDS:-50}"
PASS_RATE="${PASS_RATE:-0.90}"
FLOOR="${FLOOR:-0.05}"        # minimum sampling weight — keeps maintenance sampling alive
ALPHA="${ALPHA:-0.3}"         # winrate EMA smoothing: new = ALPHA*old + (1-ALPHA)*eval

OPPONENTS=(Fire MyFirstDroid Target MyFirstLeader Crazy Corners PaintingBot MyFirstBot)

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

# read_counter — PPO_Bot's persisted round counter; robust to a mid-write read
# (partial/empty file → 0 instead of a garbage/arithmetic error).
read_counter() {
  local v=0
  [ -f "$ROUND_COUNTER_FILE" ] || { echo 0; return; }
  v=$(cat "$ROUND_COUNTER_FILE" 2>/dev/null || echo 0)
  case "$v" in
    ''|*[!0-9]*) echo 0 ;;
    *) echo "$v" ;;
  esac
}

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

# sample_opponent — weighted draw: P(opp) ∝ max(FLOOR, 1 - winrate).
# Cumulative distribution built in awk; random draw via awk's rand() seeded from
# bash $RANDOM. Deterministic when $RANDOM is pre-seeded (the self-test relies on
# this). Pass an explicit seed as $1 to override.
sample_opponent() {
  local seed="${1:-$RANDOM}"
  awk -v seed="$seed" -v floor="$FLOOR" \
      -v names="${OPPONENTS[*]}" -v rates="${WR[*]}" '
    BEGIN {
      srand(seed);
      n = split(names, N, " ");
      split(rates, R, " ");
      total = 0;
      for (i = 1; i <= n; i++) {
        w = 1.0 - R[i] + 0.0;
        if (w < floor) w = floor;
        W[i] = w; total += w;
      }
      r = rand() * total;
      acc = 0;
      for (i = 1; i <= n; i++) {
        acc += W[i];
        if (r < acc) { print N[i]; exit }
      }
      print N[n];
    }'
}

# run_chunk <opponent> <rounds> [train|eval] — one battle chunk, with the
# 3-attempt retry / loud-abort guard from curriculum_retrain.sh. Prints the
# JSONL path on success; exits the script on a zero-progress chunk.
run_chunk() {
  local opp="$1" want="$2" mode="${3:-train}"
  local logfile
  if [ "$mode" = "eval" ]; then
    logfile="/tmp/generalist_eval_${opp}.jsonl"
    > "$logfile"                    # fresh snapshot per eval
    export PPOB_EVAL_ONLY=1         # freeze training (pure evaluation)
  else
    logfile="/tmp/generalist_train_${opp}.jsonl"
    unset PPOB_EVAL_ONLY
  fi
  export TRAINING_OPPONENT="$opp"
  export TRAINING_ROUNDS="$want"
  export PPOB_LOG_FILE="$logfile"

  # Reset the counter so RunTraining's expected-end check (start+rounds) is exact.
  echo 0 > "$ROUND_COUNTER_FILE"

  local remaining="$want" attempts=0 done=0
  while [ "$remaining" -gt 0 ]; do
    attempts=$(( attempts + 1 ))
    [ "$attempts" -gt 1 ] && sleep 2
    # RunTraining.java exits 1 on its own when the counter is frozen ("process
    # dead") or short at battle end ("N rounds never trained ... aborting for
    # restart") — treat that as a retry signal, then re-check the counter.
    java -cp "$SCRIPT_DIR:$JAR" RunTraining "$opp" "$remaining" \
      || echo ">>> Java exited non-zero (attempt $attempts/3) — crash/restart expected"
    done=$(read_counter)
    remaining=$(( want - done ))
    # Zero counter progress across 3 attempts = dead bot / broken harness.
    # The bot crashed constantly before; a crash must fail loudly now.
    if [ "$attempts" -ge 3 ] && [ "$done" -eq 0 ]; then
      echo "ERROR: round_counter still 0 after $attempts attempts vs $opp — aborting for restart" >&2
      exit 1
    fi
  done
  unset PPOB_EVAL_ONLY
  echo "$logfile"
}

persist_wr() {
  local i opp
  : > "$WR_STATE_FILE"
  for i in "${!OPPONENTS[@]}"; do
    printf "%s %.4f\n" "${OPPONENTS[$i]}" "${WR[$i]}" >> "$WR_STATE_FILE"
  done
}

load_wr() {
  local src=""
  if [ -f "$WR_SEED_FILE" ]; then
    src="$WR_SEED_FILE"
    echo ">>> Seeding winrates from $WR_SEED_FILE"
  elif [ -f "$WR_STATE_FILE" ]; then
    src="$WR_STATE_FILE"
    echo ">>> Loading winrates from $WR_STATE_FILE"
  else
    echo ">>> No winrate state — starting all at 1.0 (uniform floor sampling until first eval)"
  fi
  local -A from_file=()
  if [ -n "$src" ]; then
    local name val
    while read -r name val _; do
      [ -n "$name" ] || continue
      from_file[$name]="${val:-1.0}"
    done < "$src"
  fi
  local i opp
  for i in "${!OPPONENTS[@]}"; do
    opp="${OPPONENTS[$i]}"
    WR[$i]="${from_file[$opp]:-1.0}"
  done
}

load_best() {
  BEST_MEAN=0.0
  BEST_ROUNDS=0
  if [ -f "$BEST_FILE" ]; then
    read -r BEST_MEAN BEST_ROUNDS < "$BEST_FILE" || :
    [ -n "$BEST_MEAN" ] || BEST_MEAN=0.0
    [ -n "$BEST_ROUNDS" ] || BEST_ROUNDS=0
    echo ">>> Best-so-far: mean=$BEST_MEAN at round $BEST_ROUNDS"
  fi
}

snap_weights() {
  local rounds="$1"
  local dst="$SNAPSHOTS_DIR/generalist_best_r${rounds}"
  cp -a "$PPO_BOT_SRC/weights/latest" "$dst"   # fresh dir per save (rounds strictly increase)
  echo "  saved checkpoint -> $dst"
}

# do_eval <trained_rounds> — fresh battle per opponent, PPOB_EVAL_ONLY=1,
# EVAL_ROUNDS each. Updates winrates (eval-only!) and the best generalist
# checkpoint by mean winrate.
do_eval() {
  local trained_rounds="$1"
  local i opp w logfile oldwr evwr newwr alert
  local total_wins=0
  local total_possible=$(( ${#OPPONENTS[@]} * EVAL_ROUNDS ))
  local mean
  echo ""
  echo "=== EVAL @ round $trained_rounds ==="
  for i in "${!OPPONENTS[@]}"; do
    opp="${OPPONENTS[$i]}"
    logfile=$(run_chunk "$opp" "$EVAL_ROUNDS" eval) || { echo "  $opp: ERR"; continue; }
    w=$(count_wins "$logfile")
    total_wins=$(( total_wins + w ))
    evwr=$(awk -v w="$w" -v n="$EVAL_ROUNDS" 'BEGIN{printf "%.4f", w/n}')
    oldwr="${WR[$i]}"
    newwr=$(awk -v a="$ALPHA" -v o="$oldwr" -v e="$evwr" 'BEGIN{printf "%.4f", a*o + (1-a)*e}')
    WR[$i]="$newwr"
    persist_wr   # state survives a crash mid-eval
    alert=""
    if awk -v e="$evwr" -v p="$PASS_RATE" 'BEGIN{exit !(e < p)}'; then
      alert="  <-- BELOW PASS_RATE $PASS_RATE"
    fi
    echo "  $opp: $w/$EVAL_ROUNDS (wr $oldwr -> $newwr)$alert"
  done
  mean=$(awk -v n="${#OPPONENTS[@]}" -v rates="${WR[*]}" '
    BEGIN { split(rates, R, " "); s = 0; for (i = 1; i <= n; i++) s += R[i]; printf "%.4f", s / n }')
  echo "  mean winrate: $mean (total $total_wins/$total_possible)"
  if awk -v m="$mean" -v b="$BEST_MEAN" 'BEGIN{exit !(m > b)}'; then
    BEST_MEAN="$mean"
    BEST_ROUNDS="$trained_rounds"
    snap_weights "$trained_rounds"
    printf '%s %s\n' "$BEST_MEAN" "$BEST_ROUNDS" > "$BEST_FILE"
    echo "  >>> NEW BEST generalist: mean=$mean at $trained_rounds rounds"
  else
    echo "  best remains: $BEST_MEAN at $BEST_ROUNDS rounds"
  fi
}

# ── self-test ─────────────────────────────────────────────────────────────────
# Mock winrates: MyFirstBot 0.0, PaintingBot 0.1, 6 strong bots 1.0.
# Weights: 1.0, 0.9, and 6x FLOOR 0.05 => total 2.2.
# Expected: MyFirstBot 1.0/2.2=45.5%, PaintingBot 40.9%, each strong 2.27%.
selftest() {
  local rc=0 i opp d frac
  local -A draws=()
  WR=(1.0 1.0 1.0 1.0 1.0 1.0 0.1 0.0)   # order matches OPPONENTS
  RANDOM=12345                            # fixed seed -> reproducible draws
  for ((i = 0; i < 10000; i++)); do
    opp=$(sample_opponent)
    draws[$opp]=$(( ${draws[$opp]:-0} + 1 ))
  done
  echo "=== sampling self-test: 10000 draws, seed RANDOM=12345 ==="
  echo "mock winrates: MyFirstBot 0.0, PaintingBot 0.1, strong x6 1.0"
  echo "weights -> MyFirstBot 1.0, PaintingBot 0.9, strong 0.05 each (total 2.2)"
  echo ""
  echo "actual draw distribution:"
  for opp in "${OPPONENTS[@]}"; do
    d=${draws[$opp]:-0}
    frac=$(awk -v d="$d" 'BEGIN{printf "%.4f", d/10000}')
    printf "  %-14s %5d draws (share %.4f)\n" "$opp" "$d" "$frac"
  done
  echo ""
  local mfb
  mfb=$(awk -v d="${draws[MyFirstBot]:-0}" 'BEGIN{printf "%.4f", d/10000}')
  assert_range() {
    # assert_range <name> <frac> <lo> <hi>
    if awk -v f="$2" -v lo="$3" -v hi="$4" 'BEGIN{exit !(f >= lo && f <= hi)}'; then
      echo "  OK   $1 share $2 in [$3,$4]"
    else
      echo "  FAIL $1 share $2 not in [$3,$4]"; rc=1
    fi
  }
  assert_range "MyFirstBot (weak focus)" "$mfb" 0.35 0.50
  for opp in Fire MyFirstDroid Target MyFirstLeader Crazy Corners; do
    frac=$(awk -v d="${draws[$opp]:-0}" 'BEGIN{printf "%.4f", d/10000}')
    assert_range "$opp maintenance" "$frac" 0.015 0.05
  done
  echo ""
  [ "$rc" -eq 0 ] && echo "SELF-TEST PASS" || echo "SELF-TEST FAIL"
  exit "$rc"
}

# ── main training loop ────────────────────────────────────────────────────────
main() {
  echo "=== PPO_Bot generalist mixed training — $(date) ==="
  echo "Config: TOTAL_ROUNDS=$TOTAL_ROUNDS CHUNK_SIZE=$CHUNK_SIZE EVAL_INTERVAL=$EVAL_INTERVAL EVAL_ROUNDS=$EVAL_ROUNDS PASS_RATE=$PASS_RATE FLOOR=$FLOOR ALPHA=$ALPHA"
  echo "Opponents: ${OPPONENTS[*]}"

  echo ">>> Compiling PPO_Bot..."
  (cd "$PPO_BOT_SRC" && nimble build -d:release) || { echo "ERROR: PPO_Bot build failed" >&2; exit 1; }
  echo ">>> Compiling RunTraining.java..."
  (cd "$SCRIPT_DIR" && javac -cp "$JAR" RunTraining.java) || { echo "ERROR: RunTraining compile failed" >&2; exit 1; }
  echo ">>> Compilation done."

  # ── tmux-friendly session header ────────────────────────────────────────────
  echo ""
  echo "┌─────────────────────────────────────────────────────────────────┐"
  printf "│  Start time   : %-48s│\n" "$(date '+%Y-%m-%d %H:%M:%S %Z')"
  printf "│  Total rounds : %-48s│\n" "$TOTAL_ROUNDS"
  printf "│  Chunk size   : %-48s│\n" "$CHUNK_SIZE  ($(( TOTAL_ROUNDS / CHUNK_SIZE )) chunks total)"
  printf "│  Eval interval: %-48s│\n" "every $EVAL_INTERVAL rounds  ($(( TOTAL_ROUNDS / EVAL_INTERVAL )) evals planned)"
  printf "│  Eval rounds  : %-48s│\n" "$EVAL_ROUNDS per opponent  ($(( ${#OPPONENTS[@]} * EVAL_ROUNDS )) max wins)"
  printf "│  Opponents    : %-48s│\n" "${#OPPONENTS[@]} bots"
  echo "└─────────────────────────────────────────────────────────────────┘"
  echo ""

  load_best
  load_wr
  persist_wr   # materialize state now so an early crash still leaves a readable file

  # fresh per-run train JSONLs; cross-chunk state lives in the wr/best files
  local opp
  for opp in "${OPPONENTS[@]}"; do
    : > "/tmp/generalist_train_${opp}.jsonl"
  done

  local trained_rounds=0 chunk_rounds=0
  while [ "$trained_rounds" -lt "$TOTAL_ROUNDS" ]; do
    chunk_rounds=$CHUNK_SIZE
    if [ $(( trained_rounds + CHUNK_SIZE )) -gt "$TOTAL_ROUNDS" ]; then
      chunk_rounds=$(( TOTAL_ROUNDS - trained_rounds ))
    fi
    opp=$(sample_opponent)
    echo ""
    echo "Chunk $trained_rounds: $opp ($chunk_rounds rounds expected)"
    run_chunk "$opp" "$chunk_rounds" train || :
    trained_rounds=$(( trained_rounds + chunk_rounds ))

    if [ $(( trained_rounds % EVAL_INTERVAL )) -eq 0 ]; then
      do_eval "$trained_rounds"
    fi
  done

  # final eval at the very end (skip when the loop already evaluated at TOTAL_ROUNDS)
  if [ $(( TOTAL_ROUNDS % EVAL_INTERVAL )) -ne 0 ]; then
    echo ""
    echo "=== FINAL eval — $(date) ==="
    do_eval "$trained_rounds"
  fi

  echo ""
  echo "=== generalist training complete — $(date) ==="
  echo "Final winrates in $WR_STATE_FILE"
  echo "Best generalist checkpoint: $SNAPSHOTS_DIR/generalist_best_r${BEST_ROUNDS} (mean=$BEST_MEAN at $BEST_ROUNDS rounds)"
}

if [ "${1:-}" = "--selftest" ]; then
  selftest 2>&1 | tee "$SELFTEST_LOG"
  exit "${PIPESTATUS[0]}"
fi

main "$@" 2>&1 | tee -a "$LOG"
exit "${PIPESTATUS[0]}"