#!/bin/bash
#
# MAMBA-2 Production Training Launch Script
#
# Configuration:
# - 200 epochs (early stopping at ~150)
# - 32 batch size (4GB VRAM optimized)
# - 0.0001 learning rate
# - 60 sequence length
# - 128 hidden dim (memory efficient)
# - 64 state dim
# - GPU acceleration (RTX 3050 Ti)
# - 360 training files (665,483 bars)
#
# Expected Duration: 2-3 hours
#

set -e  # Exit on error

# Configuration
EPOCHS=200
BATCH_SIZE=32
LEARNING_RATE=0.0001
SEQ_LENGTH=60
HIDDEN_DIM=128
STATE_DIM=64
DATA_DIR="test_data/real/databento/ml_training"
OUTPUT_DIR="ml/trained_models/production/mamba2"

# Create output directory
mkdir -p "$OUTPUT_DIR"

# Print configuration
echo "╔═══════════════════════════════════════════════════════════╗"
echo "║       MAMBA-2 Production Training Launch                 ║"
echo "╚═══════════════════════════════════════════════════════════╝"
echo ""
echo "Configuration:"
echo "  Epochs: $EPOCHS"
echo "  Batch Size: $BATCH_SIZE"
echo "  Learning Rate: $LEARNING_RATE"
echo "  Sequence Length: $SEQ_LENGTH"
echo "  Hidden Dim: $HIDDEN_DIM"
echo "  State Dim: $STATE_DIM"
echo "  Data Directory: $DATA_DIR"
echo "  Output Directory: $OUTPUT_DIR"
echo ""
echo "GPU: RTX 3050 Ti (4GB VRAM)"
echo "Data: 360 files, 665,483 bars"
echo "Expected Duration: 2-3 hours"
echo ""
echo "Starting training..."
echo ""

# Launch training with CUDA
CUDA_VISIBLE_DEVICES=0 cargo run --release -p ml --example train_mamba2_dbn --features cuda -- \
  --epochs "$EPOCHS" \
  --batch-size "$BATCH_SIZE" \
  --learning-rate "$LEARNING_RATE" \
  --sequence-length "$SEQ_LENGTH" \
  --hidden-dim "$HIDDEN_DIM" \
  --state-dim "$STATE_DIM" \
  --data-dir "$DATA_DIR" \
  --output-dir "$OUTPUT_DIR" \
  --use-gpu \
  2>&1 | tee "$OUTPUT_DIR/training.log"

echo ""
echo "╔═══════════════════════════════════════════════════════════╗"
echo "║           MAMBA-2 Training Complete                       ║"
echo "╚═══════════════════════════════════════════════════════════╝"
echo ""
echo "Checkpoints saved to: $OUTPUT_DIR"
echo "Training log: $OUTPUT_DIR/training.log"
echo ""