#!/bin/bash
# run_vis.sh - Run visualization experiments for datasets and link functions

# Configuration
BASE_DIR="../datasets-orreview/ordinal-regression"
VIS_DIR="vis"
VIS_COMBINED_DIR="vis/combined"
LOGS_DIR="logs_vis"
GPU=0  # Use GPU 0 by default

# Available datasets and link functions
DATASETS=("ERA" "LEV" "SWD" "car" "winequality-red")
LINKS=("logit" "probit")

# Create necessary directories
mkdir -p "$VIS_DIR"
mkdir -p "$VIS_COMBINED_DIR"
mkdir -p "$LOGS_DIR"

# Function to run visualization for a dataset and link function
run_visualization() {
    local dataset=$1
    local link=$2
    
    echo "=========================================="
    echo "Starting visualization for dataset: $dataset, link function: $link"
    echo "=========================================="
    
    # Check if dataset exists
    if [ ! -d "$BASE_DIR/$dataset/weka" ]; then
        echo "Error: Dataset directory '$BASE_DIR/$dataset/weka' does not exist!"
        return 1
    fi
    
    # Create log file for this run
    local log_file="$LOGS_DIR/${dataset}_${link}_vis.log"
    
    echo "Running visualization experiment..."
    
    # Run the visualization experiment
    CUDA_VISIBLE_DEVICES=$GPU \
        python -u experiment_vis.py \
            --dataset "$dataset" \
            --link "$link" \
        > "$log_file" 2>&1
    
    local status=$?
    if [ $status -eq 0 ]; then
        echo "Visualization completed successfully for $dataset with $link"
        
        # Immediately run the combine_plots_vis.py script for this dataset and link function
        echo "Combining visualization images for $dataset with $link..."
        python -u combine_plots_vis.py \
            --dataset "$dataset" \
            --link "$link" \
            --vis-dir "$VIS_DIR" \
            --output-dir "$VIS_COMBINED_DIR" \
            >> "$log_file" 2>&1
        
        if [ $? -eq 0 ]; then
            echo "Combined visualization completed successfully for $dataset with $link"
        else
            echo "Error: Combined visualization failed for $dataset with $link"
            echo "Check log file: $log_file"
        fi
    else
        echo "Error: Visualization failed for $dataset with $link"
        echo "Check log file: $log_file"
    fi
    
    return $status
}

# Process command line arguments
if [ $# -eq 0 ]; then
    # No arguments: process ALL datasets with ALL link functions
    echo "Processing all datasets with all link functions"
    
    for DATASET in "${DATASETS[@]}"; do
        for LINK in "${LINKS[@]}"; do
            run_visualization "$DATASET" "$LINK"
        done
    done
    
    echo "All processing completed!"
    exit 0
fi

# Process specific dataset(s) and link function(s)
if [ $# -ge 1 ] && [ $# -le 2 ]; then
    # Parse datasets (comma-separated)
    IFS=',' read -ra SELECTED_DATASETS <<< "$1"
    
    # Validate each dataset
    for DS in "${SELECTED_DATASETS[@]}"; do
        VALID=0
        for VALID_DS in "${DATASETS[@]}"; do
            if [ "$DS" == "$VALID_DS" ]; then
                VALID=1
                break
            fi
        done
        
        if [ "$VALID" -eq 0 ]; then
            echo "Error: Dataset '$DS' is not valid. Must be one of: ${DATASETS[*]}"
            exit 1
        fi
    done
    
    # Determine link functions to use
    SELECTED_LINKS=("${LINKS[@]}")  # Default to all links
    if [ $# -eq 2 ]; then
        # Parse link functions (comma-separated)
        IFS=',' read -ra SELECTED_LINKS <<< "$2"
        
        # Validate each link function
        for LK in "${SELECTED_LINKS[@]}"; do
            if [ "$LK" != "logit" ] && [ "$LK" != "probit" ]; then
                echo "Error: Link function '$LK' is not valid. Must be 'logit' or 'probit'"
                exit 1
            fi
        done
    fi
    
    # Process each combination
    for DATASET in "${SELECTED_DATASETS[@]}"; do
        for LINK in "${SELECTED_LINKS[@]}"; do
            run_visualization "$DATASET" "$LINK"
        done
    done
    
    echo "Processing completed!"
    exit 0
fi

# Wrong number of arguments
echo "Usage: $0 [dataset1,dataset2,...] [link1,link2,...]"
echo "If no arguments are provided, processes all datasets with all link functions"
echo "If only datasets are provided, processes those datasets with all link functions"
echo "Available datasets: ${DATASETS[*]}"
echo "Available link functions: ${LINKS[*]}"
exit 1