#!/usr/bin/env bash
# Verify our PRIME-RL composer-loss adapter matches UPSTREAM default_loss_fn
# byte-for-byte, WITHOUT installing the full prime-rl package (which pulls vLLM,
# pydantic config trees, flash-attn, etc.). We clone prime-rl, build a throwaway
# venv with only torch+beartype+jaxtyping+numpy, load upstream loss.py by path
# with stubbed config/utils modules, and run identical inputs through both.
#
# Usage:
#   bash composer_replication/recipes/prime_rl/verify_parity.sh
#
# Exit 0 = byte-for-byte parity confirmed; non-zero = mismatch or setup failure.
#
# This is the reproducible counterpart to the skip-marked
# test_parity_with_prime_rl_default_loss_fn unit test: that test only runs when
# prime-rl is importable in the framework venv (it usually isn't, by design —
# we don't want prime-rl's heavy deps in our test env). This script provides the
# real upstream check out-of-band.
set -euo pipefail

PRIME_RL_REPO="${PRIME_RL_REPO:-https://github.com/PrimeIntellect-ai/prime-rl.git}"
WORK="${WORK:-/tmp/prime-rl-parity-check}"
FRAMEWORK="$(cd "$(dirname "${BASH_SOURCE[0]}")/../../.." && pwd)"
CLONE="$WORK/prime-rl"
VENV="$WORK/venv"
HARNESS="$WORK/harness.py"

mkdir -p "$WORK"

echo "==> Cloning prime-rl (shallow) into $CLONE"
if [ ! -d "$CLONE/.git" ]; then
    git clone --depth 1 "$PRIME_RL_REPO" "$CLONE"
fi
PRIME_REV="$(cd "$CLONE" && git rev-parse --short HEAD)"
echo "    upstream rev: $PRIME_REV"

echo "==> Building isolated venv (torch+beartype+jaxtyping+numpy only)"
if [ ! -x "$VENV/bin/python" ]; then
    python3 -m venv "$VENV"
    "$VENV/bin/pip" install --quiet --upgrade pip
    # CPU torch is plenty for a loss-numerics parity check.
    "$VENV/bin/pip" install --quiet torch --index-url https://download.pytorch.org/whl/cpu
    "$VENV/bin/pip" install --quiet beartype jaxtyping numpy
fi

echo "==> Writing parity harness"
cp "$FRAMEWORK/composer_replication/recipes/prime_rl/_parity_harness.py" "$HARNESS"

echo "==> Running parity sweep"
"$VENV/bin/python" "$HARNESS" "$CLONE" "$FRAMEWORK"