#!/bin/bash
# Test script for single GPU encoding with a small sample
# This creates a test subset and runs encoding on a single GPU

set -e

echo "=========================================="
echo "Testing Semantic Encoding (Single GPU)"
echo "=========================================="
echo ""

# Configuration
FULL_METADATA="/home/tony/Work/tony/RealGen/metas_v5_val_sampled_diverse.jsonl"
TEST_METADATA="/home/tony/Work/tony/RealGen/test_metadata_10samples.jsonl"
TEST_OUTPUT_DIR="/home/tony/Work/tony/RealGen/test_semantic_codes"
BATCH_SIZE=4

echo "Creating test metadata with 10 samples..."
head -n 10 $FULL_METADATA > $TEST_METADATA
echo "Test metadata created: $TEST_METADATA"
echo ""

echo "Configuration:"
echo "  Metadata: $TEST_METADATA"
echo "  Output:   $TEST_OUTPUT_DIR"
echo "  Batch:    $BATCH_SIZE"
echo "  Samples:  10"
echo ""

# Create output directory
mkdir -p $TEST_OUTPUT_DIR

# Run encoding on single GPU
echo "Running encoding on single GPU..."
echo ""

python /home/tony/Work/tony/RealGen/encode_semantic_codes.py \
    --metadata_path $TEST_METADATA \
    --output_dir $TEST_OUTPUT_DIR \
    --batch_size $BATCH_SIZE

echo ""
echo "=========================================="
echo "Test completed!"
echo "=========================================="
echo ""

# Check results
NUM_NPZ=$(find $TEST_OUTPUT_DIR -name "*.npz" | wc -l)
echo "Results:"
echo "  Created $NUM_NPZ .npz files"
echo "  Expected: 10 files (or fewer if some failed)"
echo ""

if [ $NUM_NPZ -gt 0 ]; then
    echo "Sample output file:"
    ls -lh $TEST_OUTPUT_DIR/*.npz | head -n 1
    echo ""
    echo "✓ Single GPU test passed!"
else
    echo "✗ No output files created. Check for errors."
    exit 1
fi