#!/bin/bash

# Submit REPPO jobs to SLURM
# Usage: ./submit_job.sh [environment_type] [env_name] [experiment_type]

set -e

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
cd "$SCRIPT_DIR"

# Default values
ENV_TYPE=${1:-maniskill}
ENV_NAME=${2:-PickCube-v1}
EXPERIMENT_TYPE=${3:-mjx_dmc_medium_data}

echo "Submitting REPPO job..."
echo "Environment type: $ENV_TYPE"
echo "Environment name: $ENV_NAME"
echo "Experiment type: $EXPERIMENT_TYPE"

case $ENV_TYPE in
    maniskill)
        echo "Submitting ManiSkill job..."
        ENV_NAME="$ENV_NAME" EXPERIMENT_TYPE="$EXPERIMENT_TYPE" sbatch slurm/run_reppo_maniskill.sh
        ;;
    brax)
        echo "Submitting Brax job..."
        ENV_NAME="$ENV_NAME" EXPERIMENT_TYPE="$EXPERIMENT_TYPE" sbatch slurm/run_reppo_brax.sh
        ;;
    *)
        echo "Unknown environment type: $ENV_TYPE"
        echo "Supported types: maniskill, brax"
        exit 1
        ;;
esac

echo ""
echo "Job submitted! Check status with:"
echo "  squeue -u $USER"
echo ""
echo "Check logs in: logs/ directory"
echo ""
echo "Available ManiSkill environments:"
echo "  PickCube-v1, PlaceApple-v1, StackCube-v1, PegInsertionSide-v1"
echo ""
echo "Available Brax environments:"  
echo "  ant, cheetah, hopper, walker2d, humanoid"
echo ""
echo "Available experiment types:"
echo "  mjx_dmc_small_data (32k), mjx_dmc_medium_data (512k), mjx_dmc_large_data (1M)"