#!/bin/bash
# Example hyperparameter tuning job
# Demonstrates how to launch the Optuna tuner subprocess

set -e

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
SERVICE_DIR="$(dirname "$SCRIPT_DIR")"

echo "==> Example Hyperparameter Tuning Job"
echo ""

# Configuration
JOB_ID="example_$(date +%s)"
MODEL_TYPE="TLOB"
NUM_TRIALS=10
CONFIG_PATH="$SERVICE_DIR/tuning_config.yaml"
STORAGE_PATH="/tmp/study_${JOB_ID}.log"
DATA_SOURCE_JSON='{"file_path": "/tmp/test_data.parquet", "start_time": 1704067200, "end_time": 1704672000}'
USE_GPU="--use-gpu"  # Remove if no GPU available

echo "Job ID: $JOB_ID"
echo "Model Type: $MODEL_TYPE"
echo "Trials: $NUM_TRIALS"
echo "Storage: $STORAGE_PATH"
echo ""

# Check if ML Training Service is running
if ! grpc_health_probe -addr=localhost:50054 2>/dev/null; then
    echo "WARNING: ML Training Service (port 50054) is not running"
    echo "Start with: cargo run -p ml_training_service"
    echo ""
    read -p "Continue anyway? (y/n) " -n 1 -r
    echo
    if [[ ! $REPLY =~ ^[Yy]$ ]]; then
        exit 1
    fi
fi

# Check Python dependencies
echo "==> Checking Python dependencies..."
if ! python3 -c "import optuna, grpc, yaml; import pynvml" 2>/dev/null; then
    echo "Missing dependencies. Install with:"
    echo "  pip3 install -r requirements-tuner.txt"
    exit 1
fi

# Generate Python proto stubs if needed
if [ ! -f "$SERVICE_DIR/proto/ml_training_pb2.py" ]; then
    echo "==> Generating Python gRPC stubs..."
    "$SCRIPT_DIR/generate_python_proto.sh"
fi

# Run tuner
echo "==> Starting hyperparameter tuning..."
echo ""

cd "$SERVICE_DIR"

python3 hyperparameter_tuner.py \
    --job-id "$JOB_ID" \
    --model-type "$MODEL_TYPE" \
    --num-trials "$NUM_TRIALS" \
    --config "$CONFIG_PATH" \
    --data-source-json "$DATA_SOURCE_JSON" \
    $USE_GPU \
    --storage-path "$STORAGE_PATH" \
    --grpc-host localhost \
    --grpc-port 50054

echo ""
echo "==> Tuning completed!"
echo ""

# Load results
echo "==> Best hyperparameters:"
python3 -c "
import optuna
from optuna.storages import JournalStorage, JournalFileStorage

storage = JournalStorage(JournalFileStorage('$STORAGE_PATH'))
study = optuna.load_study(study_name='study_${JOB_ID}', storage=storage)

print(f'Completed trials: {len([t for t in study.trials if t.state == optuna.trial.TrialState.COMPLETE])}')
print(f'Best Sharpe ratio: {study.best_value:.4f}')
print(f'Best parameters:')
for param, value in study.best_params.items():
    print(f'  {param}: {value}')
"

echo ""
echo "Study saved to: $STORAGE_PATH"