#!/bin/bash
set -e

DATA_DIR="${DATA_DIR:-/data/synthetic}"
MODEL_DIR="${MODEL_DIR:-/models}"
NUM_SAMPLES="${NUM_SAMPLES:-50000}"
EPOCHS="${EPOCHS:-100}"
BATCH_SIZE="${BATCH_SIZE:-32}"
LEARNING_RATE="${LEARNING_RATE:-0.001}"
DETECTOR="${DETECTOR:-radiacode_103}"
MIN_DURATION="${MIN_DURATION:-30}"
MAX_DURATION="${MAX_DURATION:-300}"
SEED="${SEED:-42}"
MEASURED_BACKGROUND_PATH="${MEASURED_BACKGROUND_PATH:-}"

echo "============================================"
echo "  Radiacode 103 — Pipeline d'entraînement"
echo "============================================"
echo "  Data dir    : $DATA_DIR"
echo "  Model dir   : $MODEL_DIR"
echo "  Samples     : $NUM_SAMPLES"
echo "  Detector    : $DETECTOR"
echo "  Duration    : $MIN_DURATION-$MAX_DURATION s"
echo "  Epochs      : $EPOCHS"
echo "  Batch size  : $BATCH_SIZE"
echo "  Learning rate: $LEARNING_RATE"
echo "============================================"

MEASURED_BG_ARG=""
if [ -n "$MEASURED_BACKGROUND_PATH" ] && [ -f "$MEASURED_BACKGROUND_PATH" ]; then
    MEASURED_BG_ARG="--measured_background $MEASURED_BACKGROUND_PATH"
    echo "Using measured background: $MEASURED_BACKGROUND_PATH"
fi

echo ""
echo "=== Phase 1 : Génération des spectres synthétiques ==="
python -m vega_ml.synthetic_spectra.generate_spectra \
  --num_samples "$NUM_SAMPLES" \
  --output_dir "$DATA_DIR" \
  --detector "$DETECTOR" \
  --min_duration "$MIN_DURATION" \
  --max_duration "$MAX_DURATION" \
  --seed "$SEED" \
  $MEASURED_BG_ARG

echo ""
echo "=== Phase 2 : Entraînement du VegaModel ==="
python -m vega_ml.training.vega.run_training \
  --data-dir "$DATA_DIR" \
  --model-dir "$MODEL_DIR" \
  --epochs "$EPOCHS" \
  --batch-size "$BATCH_SIZE" \
  --learning-rate "$LEARNING_RATE"

echo ""
echo "=== Entraînement terminé ==="
echo "Fichiers modèle :"
ls -lh "$MODEL_DIR/"

echo ""
echo "Copie de l'index des isotopes..."
if [ -f "$MODEL_DIR/vega_isotope_index.txt" ]; then
    echo "  vega_isotope_index.txt présent"
else
    echo "  ATTENTION : vega_isotope_index.txt absent"
fi