#!/usr/bin/env bash
set -euo pipefail

# Generate a GitLab CI child pipeline YAML for ML model training.
# Reads environment variables with sensible defaults and writes
# .training-generated.yml to the repository root.

###############################################################################
# Configuration (environment variables with defaults)
###############################################################################

SYMBOLS="${SYMBOLS:-ES.FUT}"
MODELS="${MODELS:-all}"
PHASE="${PHASE:-full}"
MAX_PARALLEL="${MAX_PARALLEL:-10}"
EPOCHS="${EPOCHS:-50}"
HYPEROPT_TRIALS="${HYPEROPT_TRIALS:-20}"
RUN_ID="${RUN_ID:-$(date +%Y%m%d-%H%M%S)}"
REGISTRY="${REGISTRY:-rg.fr-par.scw.cloud/foxhunt-ci}"

OUTPUT=".training-generated.yml"

###############################################################################
# Model definitions
###############################################################################

RL_MODELS=(dqn ppo)
SUPERVISED_MODELS=(tft mamba2 tggn tlob liquid kan xlstm diffusion)
# Models with hyperopt adapters (RL: DQN/PPO, Supervised: TFT/Mamba2)
HYPEROPT_RL=(dqn ppo)
HYPEROPT_SUPERVISED=(tft mamba2)
HYPEROPT_MODELS=("${HYPEROPT_RL[@]}" "${HYPEROPT_SUPERVISED[@]}")

###############################################################################
# Resolve which models to generate jobs for
###############################################################################

resolve_models() {
  case "${MODELS}" in
    all)
      echo "${RL_MODELS[*]} ${SUPERVISED_MODELS[*]}"
      ;;
    rl)
      echo "${RL_MODELS[*]}"
      ;;
    supervised)
      echo "${SUPERVISED_MODELS[*]}"
      ;;
    *)
      # Comma-separated list
      echo "${MODELS}" | tr ',' ' '
      ;;
  esac
}

SELECTED_MODELS=$(resolve_models)

###############################################################################
# Helpers
###############################################################################

# Replace dots with hyphens for valid YAML keys
sanitize() {
  echo "$1" | tr '.' '-'
}

# Check if a model is an RL model
is_rl_model() {
  local model="$1"
  for m in "${RL_MODELS[@]}"; do
    if [[ "$m" == "$model" ]]; then
      return 0
    fi
  done
  return 1
}

# Check if a model has a hyperopt adapter
has_hyperopt() {
  local model="$1"
  for m in "${HYPEROPT_MODELS[@]}"; do
    if [[ "$m" == "$model" ]]; then
      return 0
    fi
  done
  return 1
}

# Check if a model is a supervised hyperopt model (TFT/Mamba2)
is_supervised_hyperopt() {
  local model="$1"
  for m in "${HYPEROPT_SUPERVISED[@]}"; do
    if [[ "$m" == "$model" ]]; then
      return 0
    fi
  done
  return 1
}

# Get the hyperopt binary for a model
hyperopt_binary() {
  local model="$1"
  if is_rl_model "$model"; then
    echo "hyperopt_baseline_rl"
  else
    echo "hyperopt_baseline_supervised"
  fi
}

# Get the training binary name for a model
training_binary() {
  local model="$1"
  if is_rl_model "$model"; then
    echo "train_baseline_rl"
  else
    echo "train_baseline_supervised"
  fi
}

###############################################################################
# Begin generating the pipeline YAML
###############################################################################

: > "$OUTPUT"

cat >> "$OUTPUT" <<YAML
# Auto-generated by scripts/generate-training-pipeline.sh
# RUN_ID: ${RUN_ID}
# SYMBOLS: ${SYMBOLS}
# MODELS: ${MODELS} -> ${SELECTED_MODELS}
# PHASE: ${PHASE}
# Generated: $(date -u +%Y-%m-%dT%H:%M:%SZ)

YAML

# --- Stages --------------------------------------------------------------- #

if [[ "$PHASE" == "eval" ]]; then
  cat >> "$OUTPUT" <<YAML
stages:
  - evaluate

YAML
elif [[ "$PHASE" == "hyperopt" ]]; then
  cat >> "$OUTPUT" <<YAML
stages:
  - hyperopt

YAML
elif [[ "$PHASE" == "train" ]]; then
  cat >> "$OUTPUT" <<YAML
stages:
  - train

YAML
else
  cat >> "$OUTPUT" <<YAML
stages:
  - hyperopt
  - train
  - evaluate

YAML
fi

# --- Base template --------------------------------------------------------- #

cat >> "$OUTPUT" <<YAML
.training-base:
  image: ${REGISTRY}/training:latest
  tags:
    - kapsule
    - rust
  variables:
    NVIDIA_VISIBLE_DEVICES: all
    NVIDIA_DRIVER_CAPABILITIES: compute,utility
    RUST_LOG: info
    SQLX_OFFLINE: "true"
  interruptible: true

YAML

# --- Collect job names for evaluate needs ---------------------------------- #

TRAIN_JOB_NAMES=()

# --- Parse symbols -------------------------------------------------------- #

IFS=',' read -ra SYMBOL_LIST <<< "$SYMBOLS"

###############################################################################
# Hyperopt jobs (RL + supervised models with adapters, phases: full or hyperopt)
###############################################################################

if [[ "$PHASE" == "full" || "$PHASE" == "hyperopt" ]]; then
  for model in $SELECTED_MODELS; do
    if ! has_hyperopt "$model"; then
      continue
    fi
    for symbol in "${SYMBOL_LIST[@]}"; do
      symbol=$(echo "$symbol" | xargs)  # trim whitespace
      safe_symbol=$(sanitize "$symbol")
      job_name="hyperopt-${model}-${safe_symbol}"
      hp_binary=$(hyperopt_binary "$model")

      # RL models use --data-dir + --symbol; supervised use --data-dir (symbol dir)
      if is_rl_model "$model"; then
        hp_script="    - ${hp_binary}
      --model ${model}
      --trials ${HYPEROPT_TRIALS}
      --epochs 10
      --data-dir /data
      --symbol ${symbol}
      --output /output/${RUN_ID}/hyperopt/${model}/${symbol}/results.json
      --base-dir /tmp/ml_training"
      else
        hp_script="    - ${hp_binary}
      --model ${model}
      --trials ${HYPEROPT_TRIALS}
      --epochs 10
      --data-dir /data/${symbol}
      --output /output/${RUN_ID}/hyperopt/${model}/${symbol}/results.json
      --base-dir /tmp/ml_training"
      fi

      cat >> "$OUTPUT" <<YAML
${job_name}:
  extends: .training-base
  stage: hyperopt
  resource_group: gpu-${job_name}
  parallel: 1
  script:
${hp_script}

YAML
    done
  done
fi

###############################################################################
# Train jobs (phases: full or train)
###############################################################################

if [[ "$PHASE" == "full" || "$PHASE" == "train" ]]; then
  for model in $SELECTED_MODELS; do
    for symbol in "${SYMBOL_LIST[@]}"; do
      symbol=$(echo "$symbol" | xargs)
      safe_symbol=$(sanitize "$symbol")
      job_name="train-${model}-${safe_symbol}"
      binary=$(training_binary "$model")

      TRAIN_JOB_NAMES+=("$job_name")

      # Determine needs
      needs_block="  needs: []"
      if [[ "$PHASE" == "full" ]] && has_hyperopt "$model"; then
        hyperopt_job="hyperopt-${model}-${safe_symbol}"
        needs_block="  needs:
    - ${hyperopt_job}"
      fi

      # Build the script command
      if is_rl_model "$model"; then
        script_block="    - ${binary}
      --model ${model}
      --epochs ${EPOCHS}
      --data-dir /data
      --output-dir /output/${RUN_ID}/models/${model}/${symbol}"
      else
        script_block="    - ${binary}
      --model ${model}
      --epochs ${EPOCHS}
      --data-dir /data
      --symbol ${symbol}
      --output-dir /output/${RUN_ID}/models/${model}/${symbol}"
      fi
      # Wire hyperopt results into training (both RL and supervised)
      if [[ "$PHASE" == "full" ]] && has_hyperopt "$model"; then
        script_block="${script_block}
      --hyperopt-params /output/${RUN_ID}/hyperopt/${model}/${symbol}/results.json"
      fi

      cat >> "$OUTPUT" <<YAML
${job_name}:
  extends: .training-base
  stage: train
  resource_group: gpu-${job_name}
${needs_block}
  script:
${script_block}

YAML
    done
  done
fi

###############################################################################
# Evaluate job (phases: full or eval)
###############################################################################

if [[ "$PHASE" == "full" || "$PHASE" == "eval" ]]; then
  # Build needs list
  if [[ "$PHASE" == "eval" ]] || [[ ${#TRAIN_JOB_NAMES[@]} -eq 0 ]]; then
    eval_needs="  needs: []"
  else
    eval_needs="  needs:"
    for tjob in "${TRAIN_JOB_NAMES[@]}"; do
      eval_needs="${eval_needs}
    - ${tjob}"
    done
  fi

  cat >> "$OUTPUT" <<YAML
evaluate:
  extends: .training-base
  stage: evaluate
  resource_group: gpu-evaluate
${eval_needs}
  script:
    - evaluate_baseline
      --models-dir /output/${RUN_ID}/models
      --data-dir /data
      --output /output/${RUN_ID}/eval/ensemble_report.json
  artifacts:
    paths:
      - /output/${RUN_ID}/eval/ensemble_report.json
    expire_in: 30 days

YAML
fi

echo "Generated ${OUTPUT} (RUN_ID=${RUN_ID}, ${PHASE} phase, $(echo "$SELECTED_MODELS" | wc -w) models x ${#SYMBOL_LIST[@]} symbols)"