#!/bin/bash
set -e
set -o pipefail  # Ensure pipeline returns non-zero if any command fails
# 스크립트 시작시 출력 버퍼링 비활성화
export PYTHONUNBUFFERED=1
export HYDRA_FULL_ERROR=1
export ENABLE_AMP=1

# Usage: 
#   bash run.sh <config_path> [additional_args]
#
# Environment variables (optional):
#   ROOT_BASE: Dataset root path (default: /ssd/datasets/TimeSeries)
#   NUM_WORKERS: Number of workers for data loading (default: 12)
#
# Examples:
#   bash run.sh imputation_pypots/ETTh1/T1/9062_0000
#   ROOT_BASE=/data/datasets NUM_WORKERS=8 bash run.sh imputation_pypots/ETTh1/T1/9062_0000
#   bash run.sh imputation_pypots/ETTh1/T1/9062_0000 learning_rate=0.001

# Get the absolute path of the script directory
SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )"

# Remove --config-name from the input if present
FULL_PATH=${1#"--config-name "}
FULL_PATH=${FULL_PATH#"--config-name="}
shift

# Use absolute paths
DEFAULT="$SCRIPT_DIR/lab/configs/default.yaml"
BASE_CONFIG_DIR="$SCRIPT_DIR/lab/configs"
RESULTS_DIR="$SCRIPT_DIR/lab/results"

# Default values (can be overridden by environment variables)
ROOT_BASE="${ROOT_BASE:-/ssd/datasets/TimeSeries}"
NUM_WORKERS="${NUM_WORKERS:-12}"

# Check if config file exists first
CONFIG_FILE="$BASE_CONFIG_DIR/$FULL_PATH.yaml"
if [ ! -f "$CONFIG_FILE" ]; then
    echo "Error: Config file does not exist at: $CONFIG_FILE"
    exit 1
fi

# Process the path and create configs
if [[ $FULL_PATH == *"/"* ]]; then
    TARGET_DIR="$BASE_CONFIG_DIR/$(dirname "$FULL_PATH")"
    CONFIG_NAME=$(basename "$FULL_PATH")
    TARGET_FILE="$TARGET_DIR/$CONFIG_NAME.yaml"
    DEFAULT_TEMP_FILE="$TARGET_DIR/default_temp.yaml"
    
    # Create directory if it doesn't exist
    mkdir -p "$TARGET_DIR"
    
    # Copy default config to both files, using -f to force overwrite
    if [ ! -f "$TARGET_FILE" ]; then
        cp -f "$DEFAULT" "$TARGET_FILE"
        echo "Created new config file at: $TARGET_FILE"
    fi
    
    # Always use -f flag to force overwrite existing files
    cp -f "$DEFAULT" "$DEFAULT_TEMP_FILE"
    echo "Created/Updated default_temp.yaml at: $DEFAULT_TEMP_FILE"
    
    # Create results directory
    LOG_DIR="$RESULTS_DIR/$FULL_PATH"
    mkdir -p "$LOG_DIR"
    
    # Run the python script with correct config path and additional arguments
    CONFIG_DIR=$(dirname "$FULL_PATH")
    CONFIG_NAME=$(basename "$FULL_PATH")
    echo "Using config: $FULL_PATH"
    echo "Additional arguments: $@"
    echo "Logging to: $LOG_DIR/log.txt"
    
    # Modified command to ensure full error output
    cd "$SCRIPT_DIR/lab/configs" && python "$SCRIPT_DIR/run.py" --config-dir=$CONFIG_DIR --config-name=$CONFIG_NAME root_base=$ROOT_BASE num_workers=$NUM_WORKERS $@ 2>&1 | tee "$LOG_DIR/log.txt"
    exit_code=${PIPESTATUS[0]}
    exit $exit_code
else
    # Create results directory
    LOG_DIR="$RESULTS_DIR/$FULL_PATH"
    mkdir -p "$LOG_DIR"
    
    # Always use -f flag for any potential cp operations
    DEFAULT_TEMP_FILE="$BASE_CONFIG_DIR/default_temp.yaml"
    cp -f "$DEFAULT" "$DEFAULT_TEMP_FILE"
    echo "Created/Updated default_temp.yaml at: $DEFAULT_TEMP_FILE"
    
    echo "Using config: $FULL_PATH"
    echo "Additional arguments: $@"
    echo "Logging to: $LOG_DIR/log.txt"
    
    # Modified command to ensure full error output
    cd "$SCRIPT_DIR/lab/configs" && python "$SCRIPT_DIR/run.py" --config-name=$FULL_PATH root_base=$ROOT_BASE num_workers=$NUM_WORKERS $@ 2>&1 | tee "$LOG_DIR/log.txt"
    exit_code=${PIPESTATUS[0]}
    exit $exit_code
fi