#!/bin/bash
# TFT Training Restart Script (Agent 117)
# After Agent 112 TLOB Decoder fix

set -e

echo "=== TFT Training Restart ==="
echo "Timestamp: $(date)"
echo "GPU Status:"
nvidia-smi --query-gpu=name,memory.total,memory.used,temperature.gpu,utilization.gpu --format=csv,noheader
echo ""

echo "=== Pre-flight Checks ==="
echo "Checking for running training processes..."
if ps aux | grep -E "(train_tft|train_dqn|train_ppo)" | grep -v grep; then
    echo "WARNING: Found running training processes!"
    exit 1
fi

echo "Checking compilation status..."
cargo check -p ml 2>&1 | grep -E "(error|Finished)" | tail -5

echo "Checking output directory..."
mkdir -p ml/trained_models/production/tft_real_data
ls -la ml/trained_models/production/tft_real_data/

echo "Checking training data..."
echo "Available DBN files: $(find test_data/real/databento/ml_training -name '*.dbn' | wc -l)"
echo "Total size: $(du -sh test_data/real/databento/ml_training/)"

echo ""
echo "=== Starting TFT Training ==="
echo "Configuration:"
echo "  - Epochs: 200"
echo "  - Learning Rate: 0.001"
echo "  - Batch Size: 32"
echo "  - Lookback Window: 60"
echo "  - Forecast Horizon: 10"
echo "  - GPU: Enabled (CUDA_VISIBLE_DEVICES=0)"
echo "  - Output: ml/trained_models/production/tft_real_data"
echo ""

# Log file
LOG_FILE="tft_training_$(date +%Y%m%d_%H%M%S).log"
echo "Logging to: $LOG_FILE"
echo ""

# Start training
RUST_LOG=info \
RUST_BACKTRACE=1 \
CUDA_VISIBLE_DEVICES=0 \
cargo run --release -p ml --example train_tft_dbn -- \
  --data-path test_data/real/databento/ml_training \
  --epochs 200 \
  --learning-rate 0.001 \
  --batch-size 32 \
  --lookback-window 60 \
  --forecast-horizon 10 \
  --use-gpu \
  --output-dir ml/trained_models/production/tft_real_data \
  --early-stopping-patience 20 \
  --early-stopping-threshold 0.0001 \
  --verbose 2>&1 | tee "$LOG_FILE"

echo ""
echo "=== Training Complete ==="
echo "Final GPU status:"
nvidia-smi --query-gpu=memory.used,memory.total,temperature.gpu --format=csv,noheader
echo "Log saved to: $LOG_FILE"