#!/bin/bash
# Memory Profiling Script for ML Training
# Usage: ./profile_memory.sh <MODEL> <PARQUET_FILE> [EPOCHS]
#
# Example: ./profile_memory.sh dqn test_data/ES_FUT_small.parquet 3

set -e

MODEL=$1
PARQUET=$2
EPOCHS=${3:-3}

if [ -z "$MODEL" ] || [ -z "$PARQUET" ]; then
    echo "Usage: $0 <MODEL> <PARQUET_FILE> [EPOCHS]"
    echo "  MODEL: dqn, ppo, mamba2, tft"
    echo "  PARQUET_FILE: path to training data"
    echo "  EPOCHS: number of epochs (default: 3)"
    exit 1
fi

echo "=========================================="
echo "Memory Profiling: train_${MODEL}"
echo "Data: $PARQUET"
echo "Epochs: $EPOCHS"
echo "=========================================="
echo

# Get base directory (foxhunt root)
SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"
FOXHUNT_ROOT="$( cd "$SCRIPT_DIR/../.." && pwd )"

cd "$FOXHUNT_ROOT"

# Check if parquet file exists
if [ ! -f "$PARQUET" ]; then
    echo "Error: Parquet file not found: $PARQUET"
    exit 1
fi

# Run training with memory profiling
echo "Starting profiling run..."
echo

/usr/bin/time -v cargo run -p ml --example train_${MODEL} --release -- \
    --parquet-file "$PARQUET" \
    --epochs "$EPOCHS" \
    2>&1 | tee /tmp/profile_${MODEL}_$$.log

echo
echo "=========================================="
echo "Memory Profile Summary"
echo "=========================================="
grep -E "(Maximum resident|User time|System time|Percent of CPU|Elapsed|Minor.*page faults|Major.*page faults)" /tmp/profile_${MODEL}_$$.log || true

echo
echo "Full log saved to: /tmp/profile_${MODEL}_$$.log"