#!/bin/bash

# AUK Mini Test - Permutation Tag Fix Validation
# Based on train_sample_permutation_test.sh
# Tests the fix for permutation tag behavior to ensure source_type includes _permuted suffix
# Uses CUDA device 0 for testing

# Check if we're in the sunoGPT directory
if [[ ! "$(basename $(pwd))" == "sunoGPT" ]]; then
    echo "Error: This script must be run from the sunoGPT directory"
    echo "Please run: cd sunoGPT && ./auk_mini_test.sh"
    exit 1
fi

# Check if train.py exists
if [[ ! -f "train.py" ]]; then
    echo "Error: train.py not found in current directory"
    echo "Please make sure you're in the correct sunoGPT directory"
    exit 1
fi

echo "Starting AUK Mini Test - Permutation Tag Fix Validation..."
echo "CUDA Device: 0,1"
echo "Testing sample permutation tag fix with sample_permutation_prob=0.8"
echo "Expected: source_type should always be stem_permuted/full_mix_permuted when permutation enabled"

CUDA_VISIBLE_DEVICES=0,1 /home/vibert/anaconda3/envs/suno_3/bin/python -u train.py \
    --out_dir="/app2/suno/checkpoints/auk_mini_test" \
    --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=500 \
    --eval_interval=250 \
    --eval_iters=3 \
    --batch_store_size=1 \
    \
    --n_layer=1 \
    --n_head=16 \
    --d_head=64 \
    \
    --learning_rate=1e-4 \
    --max_iters=10 \
    --warmup_iters=0 \
    --batch_size=1 \
    \
    --fsdp=False \
    --grad_checkpointing=False \
    --compile=False \
    \
    --wandb_log=False \
    --preload_optimizer=False \
    --allow_skip=True \
    --allow_sample=True \
    --sample_permutation_prob=0.8 \
    --prob_sample=0.5 \
    --prob_sample_from_stems=0.9 \
    2>&1 | tee auk_mini_test.log

echo "AUK Mini Test completed!"
echo "Check logs for permutation tag validation:"
echo "- All samples with permutation enabled should have source_type as stem_permuted or full_mix_permuted"
echo "- This includes cases where no actual effects were applied (effects_info=none)"
echo "- Log location: ./auk_mini_test.log"