#!/bin/bash
# Complete CNN v2 Training Pipeline
# Train → Export → Build → Validate
# Usage: ./train_cnn_v2_full.sh [OPTIONS]
#
# MODES:
#   (none)                    Run complete pipeline: train → export → build → validate
#   --validate                Validate only (skip training, use existing weights)
#   --validate CHECKPOINT     Validate with specific checkpoint file
#   --export-only CHECKPOINT  Export weights only (skip training, build, validation)
#
# TRAINING PARAMETERS:
#   --epochs N                Training epochs (default: 200)
#   --batch-size N            Batch size (default: 16)
#   --checkpoint-every N      Checkpoint interval (default: 50)
#   --kernel-sizes K          Comma-separated kernel sizes (default: 3,3,3)
#   --num-layers N            Number of layers (default: 3)
#   --mip-level N             Mip level for p0-p3 features: 0-3 (default: 0)
#   --grayscale-loss          Compute loss on grayscale instead of RGBA
#
# PATCH PARAMETERS:
#   --patch-size N            Patch size (default: 8)
#   --patches-per-image N     Patches per image (default: 256)
#   --detector TYPE           Detector: harris|fast|shi-tomasi|gradient (default: harris)
#   --full-image              Use full-image training (disables patch mode)
#   --image-size N            Image size for full-image mode (default: 256)
#
# DIRECTORIES:
#   --input DIR               Input directory (default: training/input)
#   --target DIR              Target directory (default: training/target_1)
#   --checkpoint-dir DIR      Checkpoint directory (default: checkpoints)
#   --validation-dir DIR      Validation directory (default: validation_results)
#
# OTHER:
#   --help                    Show this help message
#
# Examples:
#   ./train_cnn_v2_full.sh
#   ./train_cnn_v2_full.sh --epochs 500 --batch-size 32
#   ./train_cnn_v2_full.sh --validate
#   ./train_cnn_v2_full.sh --validate checkpoints/checkpoint_epoch_50.pth
#   ./train_cnn_v2_full.sh --export-only checkpoints/checkpoint_epoch_100.pth
#   ./train_cnn_v2_full.sh --mip-level 1 --kernel-sizes 3,5,3

set -e

PROJECT_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
cd "$PROJECT_ROOT"

# Default configuration
INPUT_DIR="training/input"
TARGET_DIR="training/target_1"
CHECKPOINT_DIR="checkpoints"
VALIDATION_DIR="validation_results"
EPOCHS=200
CHECKPOINT_EVERY=50
BATCH_SIZE=16
PATCH_SIZE=8
PATCHES_PER_IMAGE=256
DETECTOR="harris"
KERNEL_SIZES="3,3,3"
NUM_LAYERS=3
MIP_LEVEL=0
GRAYSCALE_LOSS=false
FULL_IMAGE_MODE=false
IMAGE_SIZE=256

# Parse arguments
VALIDATE_ONLY=false
VALIDATE_CHECKPOINT=""
EXPORT_ONLY=false
EXPORT_CHECKPOINT=""

if [ "$1" = "--help" ] || [ "$1" = "-h" ]; then
    head -47 "$0" | grep "^#" | grep -v "^#!/" | sed 's/^# *//'
    exit 0
fi

# Parse all arguments
while [[ $# -gt 0 ]]; do
    case "$1" in
        --export-only)
            EXPORT_ONLY=true
            if [ -z "$2" ]; then
                echo "Error: --export-only requires a checkpoint file argument"
                exit 1
            fi
            EXPORT_CHECKPOINT="$2"
            shift 2
            ;;
        --validate)
            VALIDATE_ONLY=true
            if [ -n "$2" ] && [[ ! "$2" =~ ^-- ]]; then
                VALIDATE_CHECKPOINT="$2"
                shift 2
            else
                shift
            fi
            ;;
        --epochs)
            if [ -z "$2" ]; then
                echo "Error: --epochs requires a number argument"
                exit 1
            fi
            EPOCHS="$2"
            shift 2
            ;;
        --batch-size)
            if [ -z "$2" ]; then
                echo "Error: --batch-size requires a number argument"
                exit 1
            fi
            BATCH_SIZE="$2"
            shift 2
            ;;
        --checkpoint-every)
            if [ -z "$2" ]; then
                echo "Error: --checkpoint-every requires a number argument"
                exit 1
            fi
            CHECKPOINT_EVERY="$2"
            shift 2
            ;;
        --kernel-sizes)
            if [ -z "$2" ]; then
                echo "Error: --kernel-sizes requires a comma-separated list"
                exit 1
            fi
            KERNEL_SIZES="$2"
            shift 2
            ;;
        --num-layers)
            if [ -z "$2" ]; then
                echo "Error: --num-layers requires a number argument"
                exit 1
            fi
            NUM_LAYERS="$2"
            shift 2
            ;;
        --mip-level)
            if [ -z "$2" ]; then
                echo "Error: --mip-level requires a level argument (0-3)"
                exit 1
            fi
            MIP_LEVEL="$2"
            shift 2
            ;;
        --grayscale-loss)
            GRAYSCALE_LOSS=true
            shift
            ;;
        --patch-size)
            if [ -z "$2" ]; then
                echo "Error: --patch-size requires a number argument"
                exit 1
            fi
            PATCH_SIZE="$2"
            shift 2
            ;;
        --patches-per-image)
            if [ -z "$2" ]; then
                echo "Error: --patches-per-image requires a number argument"
                exit 1
            fi
            PATCHES_PER_IMAGE="$2"
            shift 2
            ;;
        --detector)
            if [ -z "$2" ]; then
                echo "Error: --detector requires a type argument"
                exit 1
            fi
            DETECTOR="$2"
            shift 2
            ;;
        --full-image)
            FULL_IMAGE_MODE=true
            shift
            ;;
        --image-size)
            if [ -z "$2" ]; then
                echo "Error: --image-size requires a number argument"
                exit 1
            fi
            IMAGE_SIZE="$2"
            shift 2
            ;;
        --input)
            if [ -z "$2" ]; then
                echo "Error: --input requires a directory argument"
                exit 1
            fi
            INPUT_DIR="$2"
            shift 2
            ;;
        --target)
            if [ -z "$2" ]; then
                echo "Error: --target requires a directory argument"
                exit 1
            fi
            TARGET_DIR="$2"
            shift 2
            ;;
        --checkpoint-dir)
            if [ -z "$2" ]; then
                echo "Error: --checkpoint-dir requires a directory argument"
                exit 1
            fi
            CHECKPOINT_DIR="$2"
            shift 2
            ;;
        --validation-dir)
            if [ -z "$2" ]; then
                echo "Error: --validation-dir requires a directory argument"
                exit 1
            fi
            VALIDATION_DIR="$2"
            shift 2
            ;;
        *)
            echo "Unknown option: $1"
            exit 1
            ;;
    esac
done

# Build training arguments
if [ "$FULL_IMAGE_MODE" = true ]; then
    TRAINING_MODE_ARGS="--full-image --image-size $IMAGE_SIZE"
else
    TRAINING_MODE_ARGS="--patch-size $PATCH_SIZE --patches-per-image $PATCHES_PER_IMAGE --detector $DETECTOR"
fi

# Handle export-only mode
if [ "$EXPORT_ONLY" = true ]; then
    echo "=== CNN v2 Export Weights Only ==="
    echo "Checkpoint: $EXPORT_CHECKPOINT"
    echo ""

    if [ ! -f "$EXPORT_CHECKPOINT" ]; then
        echo "Error: Checkpoint file not found: $EXPORT_CHECKPOINT"
        exit 1
    fi

    python3 training/export_cnn_v2_weights.py "$EXPORT_CHECKPOINT" \
      --output-weights workspaces/main/weights/cnn_v2_weights.bin

    if [ $? -ne 0 ]; then
        echo "Error: Export failed"
        exit 1
    fi

    echo ""
    echo "=== Export Complete ==="
    echo "Output: workspaces/main/weights/cnn_v2_weights.bin"
    exit 0
fi

if [ "$VALIDATE_ONLY" = true ]; then
    echo "=== CNN v2 Validation Only ==="
    echo "Skipping training, using existing weights"
    echo ""
else
    echo "=== CNN v2 Complete Training Pipeline ==="
    echo "Input: $INPUT_DIR"
    echo "Target: $TARGET_DIR"
    echo "Epochs: $EPOCHS"
    echo "Checkpoint interval: $CHECKPOINT_EVERY"
    echo "Mip level: $MIP_LEVEL (p0-p3 features)"
    echo ""
fi

if [ "$VALIDATE_ONLY" = false ]; then
    # Step 1: Train model
    echo "[1/4] Training CNN v2 model..."

# Build optional flags
OPTIONAL_FLAGS=""
if [ "$GRAYSCALE_LOSS" = true ]; then
    OPTIONAL_FLAGS="$OPTIONAL_FLAGS --grayscale-loss"
fi

python3 training/train_cnn_v2.py \
  --input "$INPUT_DIR" \
  --target "$TARGET_DIR" \
  $TRAINING_MODE_ARGS \
  --kernel-sizes $KERNEL_SIZES \
  --num-layers $NUM_LAYERS \
  --mip-level $MIP_LEVEL \
  --epochs $EPOCHS \
  --batch-size $BATCH_SIZE \
  --checkpoint-dir "$CHECKPOINT_DIR" \
  --checkpoint-every $CHECKPOINT_EVERY \
  $OPTIONAL_FLAGS

if [ $? -ne 0 ]; then
    echo "Error: Training failed"
    exit 1
fi

echo ""
echo "Training complete!"
echo ""

# Step 2: Export final checkpoint to shaders
FINAL_CHECKPOINT="$CHECKPOINT_DIR/checkpoint_epoch_${EPOCHS}.pth"

if [ ! -f "$FINAL_CHECKPOINT" ]; then
    echo "Warning: Final checkpoint not found, using latest available..."
    FINAL_CHECKPOINT=$(ls -t "$CHECKPOINT_DIR"/checkpoint_epoch_*.pth | head -1)
fi

echo "[2/4] Exporting final checkpoint to binary weights..."
echo "Checkpoint: $FINAL_CHECKPOINT"
python3 training/export_cnn_v2_weights.py "$FINAL_CHECKPOINT" \
  --output-weights workspaces/main/weights/cnn_v2_weights.bin

if [ $? -ne 0 ]; then
    echo "Error: Shader export failed"
    exit 1
fi

echo ""
fi  # End of training/export section

# Determine which checkpoint to use
if [ "$VALIDATE_ONLY" = true ]; then
    if [ -n "$VALIDATE_CHECKPOINT" ]; then
        FINAL_CHECKPOINT="$VALIDATE_CHECKPOINT"
    else
        # Use latest checkpoint
        FINAL_CHECKPOINT=$(ls -t "$CHECKPOINT_DIR"/checkpoint_epoch_*.pth | head -1)
    fi
    echo "Using checkpoint: $FINAL_CHECKPOINT"
    echo ""
fi

# Step 3: Rebuild with new shaders
if [ "$VALIDATE_ONLY" = false ]; then
    echo "[3/4] Rebuilding demo with new shaders..."
    cmake --build build -j4 --target demo64k > /dev/null 2>&1

    if [ $? -ne 0 ]; then
        echo "Error: Build failed"
        exit 1
    fi

    echo "  → Build complete"
    echo ""
fi

# Step 4: Visual assessment - process final checkpoint only
if [ "$VALIDATE_ONLY" = true ]; then
    echo "Validation on all input images (using existing weights)..."
else
    echo "[4/4] Visual assessment on all input images..."
fi

mkdir -p "$VALIDATION_DIR"
echo "  Using checkpoint: $FINAL_CHECKPOINT"

# Export weights for validation mode (already exported in step 2 for training mode)
if [ "$VALIDATE_ONLY" = true ]; then
    python3 training/export_cnn_v2_weights.py "$FINAL_CHECKPOINT" \
      --output-weights workspaces/main/weights/cnn_v2_weights.bin > /dev/null 2>&1
fi

# Build cnn_test
cmake --build build -j4 --target cnn_test > /dev/null 2>&1

# Process all input images
for input_image in "$INPUT_DIR"/*.png; do
    basename=$(basename "$input_image" .png)
    echo "    Processing $basename..."
    build/cnn_test "$input_image" "$VALIDATION_DIR/${basename}_output.png" 2>/dev/null
done

# Build demo only if not in validate mode
if [ "$VALIDATE_ONLY" = false ]; then
    cmake --build build -j4 --target demo64k > /dev/null 2>&1
fi

echo ""
if [ "$VALIDATE_ONLY" = true ]; then
    echo "=== Validation Complete ==="
else
    echo "=== Training Pipeline Complete ==="
fi
echo ""
echo "Results:"
if [ "$VALIDATE_ONLY" = false ]; then
    echo "  - Checkpoints: $CHECKPOINT_DIR"
    echo "  - Final weights: workspaces/main/weights/cnn_v2_weights.bin"
fi
echo "  - Validation outputs: $VALIDATION_DIR"
echo ""
echo "Opening results directory..."
open "$VALIDATION_DIR" 2>/dev/null || xdg-open "$VALIDATION_DIR" 2>/dev/null || true

if [ "$VALIDATE_ONLY" = false ]; then
    echo ""
    echo "Run demo to see final result:"
    echo "  ./build/demo64k"
fi