#!/bin/bash

# Function to echo with timestamp
log_message() {
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] $1"
}

# Define the scenes and difficulties
# scenes=('grass' 'stair' 'hydrant' 'road' 'sky' 'pillar' 'lab')
# difficulties=('easy' 'median' 'hard')

scenes=('stair')
difficulties=('median')

# KDE bandwidth (fixed at 7.5)
kde_bandwidth=7.5

# Fixed parameters
noise_scheduler="linear"
start_sigma=350

# Base port number (will be incremented for each experiment)
base_port=7900

# Base path for source data
base_source_path="/project2/hentci/ours_data/free-dataset"

# Function to run the training script
run_training() {
    local scene=$1
    local difficulty=$2
    local target_view=$3
    local port=$4
    
    # Base paths
    local source_path="${base_source_path}/poison_${scene}/${difficulty}"
    local output_path="/project2/hentci/Metrics/ours/single-view/freedataset/poison_${scene}/${difficulty}"
    local points_dir="${source_path}/sparse/0"
    local original_points="${points_dir}/points3D.ply"
    local kde_points="${points_dir}/points3D_KDE${kde_bandwidth}.ply"
    
    log_message "----------------------------------------"
    log_message "Starting experiment:"
    log_message "- Scene: ${scene}"
    log_message "- Difficulty: ${difficulty}"
    log_message "- Target view: ${target_view}"
    log_message "- KDE bandwidth: ${kde_bandwidth}"
    log_message "- Noise scheduler: ${noise_scheduler} (fixed)"
    log_message "- Start sigma: ${start_sigma} (fixed)"
    log_message "- Port: ${port}"
    
    # Save the current points3D.ply if it exists
    if [ -f "$original_points" ]; then
        log_message "Backing up existing points3D.ply"
        mv "$original_points" "${original_points}.backup"
        log_message "Backup created at: ${original_points}.backup"
    fi
    
    # Copy the KDE points file to points3D.ply
    if [ -f "$kde_points" ]; then
        log_message "Copying points3D_KDE${kde_bandwidth}.ply to points3D.ply"
        cp "$kde_points" "$original_points"
        
        if [ $? -eq 0 ]; then
            log_message "File copied successfully"
        else
            log_message "ERROR: Failed to copy file"
            return 1
        fi
    else
        log_message "WARNING: KDE points file not found at: ${kde_points}"
        log_message "Will continue with existing points file if any"
    fi
    
    log_message "Starting training script with these parameters:"
    log_message "- Source path: ${source_path}"
    log_message "- Output path: ${output_path}"
    log_message "- Target view: ${target_view}"
    log_message "- Noise scheduler: ${noise_scheduler}"
    log_message "- Start sigma: ${start_sigma}"
    
    # Create output directory if it doesn't exist
    mkdir -p "$output_path"
    
    # Run the training script
    python3 ../train_scheduling_noisev2.py \
        -s "$source_path" \
        -i images \
        -m "$output_path" \
        --target_view "$target_view" \
        --noise_scheduler "$noise_scheduler" \
        --start_sigma "$start_sigma" \
        --eval \
        --port "$port"
    
    # Check if training script executed successfully
    if [ $? -eq 0 ]; then
        log_message "Training script completed successfully"
    else
        log_message "ERROR: Training script failed"
    fi
    
    # Restore the original points3D.ply
    if [ -f "${original_points}.backup" ]; then
        log_message "Restoring original points3D.ply from backup"
        mv "${original_points}.backup" "$original_points"
        log_message "Backup restored successfully"
    fi
    
    log_message "Experiment completed"
}

# Main execution
log_message "Starting execution of all experiments"
log_message "Will run experiments for:"
log_message "- Scenes: ${scenes[*]}"
log_message "- Difficulties: ${difficulties[*]}"
log_message "- Fixed KDE bandwidth: ${kde_bandwidth}"
log_message "- Fixed noise scheduler: ${noise_scheduler}"
log_message "- Fixed start sigma: ${start_sigma}"
log_message "----------------------------------------"

# Initialize port counter
port=$base_port

# Run for each combination of scene, difficulty, and kde_bandwidth
for scene in "${scenes[@]}"; do
    for difficulty in "${difficulties[@]}"; do
        # Construct the attack images path
        attack_images_path="${base_source_path}/poison_${scene}/${difficulty}/attack/images"
        
        # Check if the attack/images directory exists
        if [ -d "$attack_images_path" ]; then
            # Get the target view (first image name in attack/images directory without extension)
            target_view=$(ls "$attack_images_path" | head -n 1 | sed 's/\.[^.]*$//')
            
            if [ -n "$target_view" ]; then
                log_message "Found target view: ${target_view} for ${scene}/${difficulty}"
                
                # Increment port for each experiment
                port=$((port + 1))
                
                # Run the training
                run_training "$scene" "$difficulty" "$target_view" "$port"
            else
                log_message "WARNING: No target view found in ${attack_images_path}. Skipping this experiment."
            fi
        else
            log_message "WARNING: Attack images directory not found at: ${attack_images_path}. Skipping this experiment."
        fi
    done
done

log_message "All experiments completed"