#!/bin/bash
# CNN v2 Validation - End-to-end pipeline

set -e
PROJECT_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
BUILD_DIR="$PROJECT_ROOT/build"
WORKSPACE="main"

usage() {
    echo "Usage: $0 <checkpoint.pth> [options]"
    echo "Options:"
    echo "  -i DIR    Test images (default: training/validation)"
    echo "  -o DIR    Output (default: validation_results)"
    echo "  --skip-build    Skip rebuild"
    exit 1
}

[ $# -eq 0 ] && usage
CHECKPOINT="$1"
shift

TEST_IMAGES="$PROJECT_ROOT/training/validation"
OUTPUT="$PROJECT_ROOT/validation_results"
SKIP_BUILD=false

while [[ $# -gt 0 ]]; do
    case $1 in
        -i) TEST_IMAGES="$2"; shift 2 ;;
        -o) OUTPUT="$2"; shift 2 ;;
        --skip-build) SKIP_BUILD=true; shift ;;
        -h) usage ;;
        *) usage ;;
    esac
done

echo "=== CNN v2 Validation ==="
echo "Checkpoint: $CHECKPOINT"

# Export
echo "[1/3] Exporting shaders..."
python3 "$PROJECT_ROOT/training/export_cnn_v2_shader.py" "$CHECKPOINT" \
    --output-dir "$PROJECT_ROOT/workspaces/$WORKSPACE/shaders"

# Build
if [ "$SKIP_BUILD" = false ]; then
    echo "[2/3] Building..."
    cmake --build "$BUILD_DIR" -j4 --target cnn_test >/dev/null 2>&1
fi

# Process
echo "[3/3] Processing images..."
mkdir -p "$OUTPUT"
count=0
for img in "$TEST_IMAGES"/*.png; do
    [ -f "$img" ] || continue
    name=$(basename "$img" .png)
    "$BUILD_DIR/cnn_test" "$img" "$OUTPUT/${name}_output.png" 2>/dev/null && count=$((count+1))
done

echo "Done! Processed $count images → $OUTPUT"