#!/bin/bash

# Training Test Script for Suno GPT Environment
# Runs a minimal training test to verify the environment works correctly

set -e

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
LOG_DIR="${SCRIPT_DIR}/test_logs"
mkdir -p "$LOG_DIR"

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m'

log() {
    echo -e "${GREEN}[$(date +'%Y-%m-%d %H:%M:%S')]${NC} $1"
}

error() {
    echo -e "${RED}[ERROR]${NC} $1"
    exit 1
}

warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

# Check if we're in a conda or virtualenv environment
check_environment() {
    log "Checking Python environment..."
    
    if [ -n "$CONDA_DEFAULT_ENV" ]; then
        log "Conda environment detected: $CONDA_DEFAULT_ENV"
        ENV_TYPE="conda"
    elif [ -n "$VIRTUAL_ENV" ]; then
        log "Virtual environment detected: $VIRTUAL_ENV"
        ENV_TYPE="venv"
    else
        warning "No active Python environment detected!"
        info "Please activate an environment first:"
        info "  Conda: source ${SCRIPT_DIR}/activate_suno_build.sh"
        info "  Venv:  source ${SCRIPT_DIR}/activate_suno_venv.sh"
        exit 1
    fi
    
    log "Python: $(which python)"
    log "Python version: $(python --version)"
}

# Quick import test
quick_import_test() {
    log "Running quick import test..."
    
    python -c "
import torch
import transformers
import wandb
import deepspeed
print('✓ Core packages imported successfully')
print(f'  PyTorch: {torch.__version__}')
print(f'  Transformers: {transformers.__version__}')
print(f'  CUDA available: {torch.cuda.is_available()}')
" || error "Failed to import core packages"
}

# Run minimal training test
run_minimal_training() {
    log "Running minimal training test..."
    
    # Check if the original training script exists
    TRAIN_SCRIPT="/home/vibert/projects/neon/sunoGPT/train.py"
    if [ ! -f "$TRAIN_SCRIPT" ]; then
        error "Training script not found at $TRAIN_SCRIPT"
    fi
    
    # Create a test configuration
    TEST_CONFIG="${LOG_DIR}/test_config_$(date +%Y%m%d_%H%M%S).sh"
    
    cat > "$TEST_CONFIG" << 'EOF'
#!/bin/bash

echo "============================================================"
echo "Testing Suno GPT Training - Minimal Configuration"
echo "============================================================"
echo ""
echo "Configuration:"
echo "- Task: text2music (default)"
echo "- Iterations: 5 (quick test)"
echo "- Single GPU mode"
echo "- No FSDP, no compilation"
echo ""

cd /home/vibert/projects/neon/sunoGPT

# Use GPU 0 by default, or whatever is available
export CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-0}

python -u train.py \
    --out_dir=/tmp/suno_test_checkpoints \
    --data_dir=/app2/suno/data/auk_v0 \
    --train_metas_filename=metas_v1_tr_mini.jsonl \
    --val_metas_filename=metas_v3_val.jsonl \
    --step_save_iters=100000 \
    --eval_interval=25 \
    --eval_iters=2 \
    --batch_store_size=1 \
    --n_layer=2 \
    --n_head=16 \
    --d_head=64 \
    --learning_rate=1e-4 \
    --max_iters=5 \
    --warmup_iters=1 \
    --batch_size=1 \
    --grad_checkpointing=False \
    --compile=False \
    --use_text_loss=True \
    --prob_text_loss=1.0 \
    --use_hoot=False \
    --use_ditto=False \
    --fsdp=False \
    --checkpoint_save_old_format=False \
    --wandb_log=False \
    2>&1
EOF
    
    chmod +x "$TEST_CONFIG"
    
    log "Starting training test (this may take a few minutes)..."
    LOG_FILE="${LOG_DIR}/train_test_$(date +%Y%m%d_%H%M%S).log"
    
    if bash "$TEST_CONFIG" > "$LOG_FILE" 2>&1; then
        log "✓ Training test completed successfully!"
        info "Log saved to: $LOG_FILE"
        
        # Check if loss decreased
        if grep -q "loss" "$LOG_FILE"; then
            log "✓ Training losses recorded"
            tail -n 20 "$LOG_FILE" | grep -i "loss" || true
        fi
    else
        error "Training test failed! Check log at: $LOG_FILE"
    fi
}

# Run SLURM compatibility test (dry run)
test_slurm_compatibility() {
    log "Testing SLURM script compatibility..."
    
    # Check if the SLURM script exists
    SLURM_SCRIPT="${SCRIPT_DIR}/refs/train_slurm_original.sh"
    if [ ! -f "$SLURM_SCRIPT" ]; then
        warning "SLURM script not found, skipping SLURM test"
        return
    fi
    
    # Extract the python command and test it exists
    PYTHON_PATH=$(grep -o "/home/vibert/anaconda3/envs/[^/]*/bin/python" "$SLURM_SCRIPT" | head -1)
    
    if [ -n "$PYTHON_PATH" ] && [ -f "$PYTHON_PATH" ]; then
        log "✓ SLURM Python path exists: $PYTHON_PATH"
        $PYTHON_PATH --version || warning "Could not run SLURM Python"
    else
        warning "SLURM Python path not found or doesn't exist"
    fi
    
    info "For actual SLURM submission, use: sbatch train_slurm_original.sh"
}

# Main execution
main() {
    log "Starting Suno GPT Training Test"
    
    check_environment
    quick_import_test
    
    # Ask user if they want to run the actual training test
    read -p "Do you want to run the minimal training test? This will use GPU and may take a few minutes (y/n): " -n 1 -r
    echo
    if [[ $REPLY =~ ^[Yy]$ ]]; then
        run_minimal_training
    else
        info "Skipping training test"
    fi
    
    test_slurm_compatibility
    
    log "============================================================"
    log "Training test completed!"
    log "Environment is ready for training."
    log "============================================================"
}

main "$@"