#!/bin/bash

# Script to run fish video weight evaluation with fixed model paths

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

# Fixed model paths
YOLO_MODEL="/home/ubuntu/projects/FishMeasure/runs/train/fish_detection_20251127_104658/weights/best.pt"
POINTCLOUD_CLASSIFIER="${SCRIPT_DIR}/pointcloud_classifier/Pointnet_Pointnet2_pytorch/log/classification/fish_pointnet2_finetune/checkpoints/best_model.pth"

# Default parameters
CONF=0.5
IMGSZ=640
SAM_DEVICE="cuda"
MAX_FRAMES=0
SAVE_IMAGES=false
FILTER_POINTCLOUD=true
USE_CLUSTERING_FILTER=false
USE_DENSITY_FILTER=true
USE_POINTCLOUD_CLASSIFIER=true
POINTCLOUD_CLASSIFIER_THRESHOLD=0.5

# Parse arguments
SVO_FILE=""
IMAGE_FOLDER=""
OUTPUT_DIR=""
SCALE=1.0

# Function to print usage
usage() {
    echo "Usage: $0 [OPTIONS]"
    echo ""
    echo "Required (choose one):"
    echo "  --svo PATH              Path to SVO2 file"
    echo "  --image-folder PATH     Path to folder containing images"
    echo ""
    echo "Optional:"
    echo "  --output PATH           Output directory (default: output_preview)"
    echo "  --conf FLOAT            Confidence threshold (default: 0.25)"
    echo "  --imgsz INT             Image size (default: 640)"
    echo "  --max-frames INT        Maximum frames to process (0 = all, default: 0)"
    echo "  --save-images           Save individual images instead of video"
    echo "  --filter-pointcloud     Apply point cloud filtering"
    echo "  --use-clustering-filter Use clustering filter (requires --filter-pointcloud)"
    echo "  --use-density-filter    Use density filter (requires --filter-pointcloud)"
    echo "  --use-pointcloud-classifier  Use point cloud quality classifier to filter bad clouds"
    echo "  --pointcloud-classifier-threshold FLOAT  Confidence threshold for classifier (default: 0.7)"
    echo "  --scale FLOAT           Display scale (default: 1.0)"
    echo "  --sam-device DEVICE     SAM device: cuda or cpu (default: cuda)"
    echo ""
    echo "Examples:"
    echo "  $0 --svo /path/to/video.svo2 --output output_preview --use-pointcloud-classifier"
    echo "  $0 --image-folder /path/to/images --output output_preview --filter-pointcloud"
}

# Parse command line arguments
while [[ $# -gt 0 ]]; do
    case $1 in
        --svo)
            SVO_FILE="$2"
            shift 2
            ;;
        --image-folder)
            IMAGE_FOLDER="$2"
            shift 2
            ;;
        --output)
            OUTPUT_DIR="$2"
            shift 2
            ;;
        --conf)
            CONF="$2"
            shift 2
            ;;
        --imgsz)
            IMGSZ="$2"
            shift 2
            ;;
        --max-frames)
            MAX_FRAMES="$2"
            shift 2
            ;;
        --save-images)
            SAVE_IMAGES=true
            shift
            ;;
        --filter-pointcloud)
            FILTER_POINTCLOUD=true
            shift
            ;;
        --use-clustering-filter)
            USE_CLUSTERING_FILTER=true
            shift
            ;;
        --use-density-filter)
            USE_DENSITY_FILTER=true
            shift
            ;;
        --use-pointcloud-classifier)
            USE_POINTCLOUD_CLASSIFIER=true
            shift
            ;;
        --pointcloud-classifier-threshold)
            POINTCLOUD_CLASSIFIER_THRESHOLD="$2"
            shift 2
            ;;
        --scale)
            SCALE="$2"
            shift 2
            ;;
        --sam-device)
            SAM_DEVICE="$2"
            shift 2
            ;;
        -h|--help)
            usage
            exit 0
            ;;
        *)
            echo "Unknown option: $1"
            usage
            exit 1
            ;;
    esac
done

# Validate required arguments
if [ -z "$SVO_FILE" ] && [ -z "$IMAGE_FOLDER" ]; then
    echo "Error: Either --svo or --image-folder must be provided"
    usage
    exit 1
fi

if [ -n "$SVO_FILE" ] && [ -n "$IMAGE_FOLDER" ]; then
    echo "Error: Cannot specify both --svo and --image-folder"
    usage
    exit 1
fi

# Set default output if not provided
if [ -z "$OUTPUT_DIR" ]; then
    OUTPUT_DIR="output_preview"
fi

# Build command
CMD="python3 fish_video_weight_evaluation.py"
CMD="$CMD --yolo-model \"${YOLO_MODEL}\""
CMD="$CMD --conf ${CONF}"
CMD="$CMD --imgsz ${IMGSZ}"
CMD="$CMD --sam-device ${SAM_DEVICE}"
CMD="$CMD --scale ${SCALE}"
CMD="$CMD --max-frames ${MAX_FRAMES}"

if [ -n "$SVO_FILE" ]; then
    CMD="$CMD --svo \"${SVO_FILE}\""
fi

if [ -n "$IMAGE_FOLDER" ]; then
    CMD="$CMD --image-folder \"${IMAGE_FOLDER}\""
fi

CMD="$CMD --save-output \"${OUTPUT_DIR}\""

if [ "$SAVE_IMAGES" = true ]; then
    CMD="$CMD --save-images"
fi

if [ "$FILTER_POINTCLOUD" = true ]; then
    CMD="$CMD --filter-pointcloud"
fi

if [ "$USE_CLUSTERING_FILTER" = true ]; then
    CMD="$CMD --use-clustering-filter"
fi

if [ "$USE_DENSITY_FILTER" = true ]; then
    CMD="$CMD --use-density-filter"
fi

if [ "$USE_POINTCLOUD_CLASSIFIER" = true ]; then
    CMD="$CMD --pointcloud-classifier \"${POINTCLOUD_CLASSIFIER}\""
    CMD="$CMD --use-pointcloud-classifier"
    CMD="$CMD --pointcloud-classifier-threshold ${POINTCLOUD_CLASSIFIER_THRESHOLD}"
fi

# Print configuration
echo "=========================================="
echo "Fish Video Weight Evaluation"
echo "=========================================="
echo "YOLO Model: ${YOLO_MODEL}"
if [ "$USE_POINTCLOUD_CLASSIFIER" = true ]; then
    echo "Point Cloud Classifier: ${POINTCLOUD_CLASSIFIER}"
fi
if [ -n "$SVO_FILE" ]; then
    echo "SVO2 File: ${SVO_FILE}"
else
    echo "Image Folder: ${IMAGE_FOLDER}"
fi
echo "Output Directory: ${OUTPUT_DIR}"
echo "Confidence: ${CONF}"
echo "Image Size: ${IMGSZ}"
echo "Max Frames: ${MAX_FRAMES}"
echo "SAM Device: ${SAM_DEVICE}"
if [ "$USE_POINTCLOUD_CLASSIFIER" = true ]; then
    echo "Point Cloud Classifier Threshold: ${POINTCLOUD_CLASSIFIER_THRESHOLD}"
fi
echo "=========================================="
echo ""

# Execute command
eval $CMD

echo ""
echo "Processing completed!"