#!/bin/bash
#
# Quarterly Model Retraining Script
#
# This script automates the quarterly retraining of all ML models
# in the Foxhunt HFT trading system. It should be run on the first
# Sunday of each quarter (Jan, Apr, Jul, Oct).
#
# Usage:
#   ./scripts/quarterly_retrain.sh [--dry-run] [--parallel] [--models DQN,PPO]
#
# Environment Variables:
#   FOXHUNT_ROOT: Project root directory (default: current directory)
#   SLACK_WEBHOOK: Slack webhook URL for notifications
#   EMAIL_RECIPIENTS: Comma-separated list of email addresses
#
# Exit Codes:
#   0 - Success (all models trained and passed quality gates)
#   1 - Partial failure (some models failed)
#   2 - Complete failure (no models succeeded)
#   3 - Prerequisites not met

set -e
set -o pipefail

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# Configuration
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="${FOXHUNT_ROOT:-$(cd "$SCRIPT_DIR/.." && pwd)}"
LOG_DIR="$PROJECT_ROOT/logs"
DATA_DIR="$PROJECT_ROOT/test_data/real/databento/ml_training"
OUTPUT_DIR="$PROJECT_ROOT/ml/trained_models/quarterly"
HYPERPARAMS_FILE="$PROJECT_ROOT/ml/config/best_hyperparameters.yaml"

# Generate version tag (e.g., 2024Q4_v1)
YEAR=$(date +"%Y")
MONTH=$(date +"%m")
QUARTER=$(( (MONTH - 1) / 3 + 1 ))
VERSION_TAG="${YEAR}Q${QUARTER}_v1"
OUTPUT_SUBDIR="$OUTPUT_DIR/$YEAR/Q$QUARTER"
LOG_FILE="$LOG_DIR/retraining_${VERSION_TAG}_$(date +%Y%m%d_%H%M%S).log"

# Default options
DRY_RUN=false
PARALLEL=false
MODELS="DQN,PPO,MAMBA2,TFT"
MIN_SHARPE=1.5
MIN_WIN_RATE=0.55
LATEST_DAYS=90

# Parse arguments
while [[ $# -gt 0 ]]; do
    case $1 in
        --dry-run)
            DRY_RUN=true
            shift
            ;;
        --parallel)
            PARALLEL=true
            shift
            ;;
        --models)
            MODELS="$2"
            shift 2
            ;;
        --min-sharpe)
            MIN_SHARPE="$2"
            shift 2
            ;;
        --min-win-rate)
            MIN_WIN_RATE="$2"
            shift 2
            ;;
        --latest-days)
            LATEST_DAYS="$2"
            shift 2
            ;;
        --help)
            echo "Usage: $0 [OPTIONS]"
            echo ""
            echo "Options:"
            echo "  --dry-run           Validate without training"
            echo "  --parallel          Train models in parallel (may OOM)"
            echo "  --models LIST       Comma-separated models (default: DQN,PPO,MAMBA2,TFT)"
            echo "  --min-sharpe N      Minimum Sharpe ratio (default: 1.5)"
            echo "  --min-win-rate N    Minimum win rate (default: 0.55)"
            echo "  --latest-days N     Days of data to use (default: 90)"
            echo "  --help              Show this help message"
            exit 0
            ;;
        *)
            echo -e "${RED}Unknown option: $1${NC}"
            exit 3
            ;;
    esac
done

# Logging function
log() {
    local level=$1
    shift
    local message="$@"
    local timestamp=$(date +"%Y-%m-%d %H:%M:%S")

    case $level in
        INFO)
            echo -e "${BLUE}[$timestamp] INFO:${NC} $message" | tee -a "$LOG_FILE"
            ;;
        SUCCESS)
            echo -e "${GREEN}[$timestamp] SUCCESS:${NC} $message" | tee -a "$LOG_FILE"
            ;;
        WARNING)
            echo -e "${YELLOW}[$timestamp] WARNING:${NC} $message" | tee -a "$LOG_FILE"
            ;;
        ERROR)
            echo -e "${RED}[$timestamp] ERROR:${NC} $message" | tee -a "$LOG_FILE"
            ;;
        *)
            echo "[$timestamp] $message" | tee -a "$LOG_FILE"
            ;;
    esac
}

# Send notification (Slack + Email)
notify() {
    local status=$1
    local message=$2

    # Slack notification
    if [ -n "$SLACK_WEBHOOK" ]; then
        local emoji
        case $status in
            SUCCESS) emoji=":white_check_mark:" ;;
            WARNING) emoji=":warning:" ;;
            ERROR) emoji=":x:" ;;
            *) emoji=":information_source:" ;;
        esac

        curl -X POST "$SLACK_WEBHOOK" \
            -H 'Content-Type: application/json' \
            -d "{\"text\":\"$emoji Foxhunt Quarterly Retraining\\n$message\"}" \
            2>/dev/null || true
    fi

    # Email notification
    if [ -n "$EMAIL_RECIPIENTS" ]; then
        echo "$message" | mail -s "Foxhunt Quarterly Retraining - $status" "$EMAIL_RECIPIENTS" 2>/dev/null || true
    fi
}

# Check prerequisites
check_prerequisites() {
    log INFO "Checking prerequisites..."

    # Check Rust installation
    if ! command -v cargo &> /dev/null; then
        log ERROR "Cargo not found. Please install Rust."
        return 1
    fi

    # Check GPU availability
    if command -v nvidia-smi &> /dev/null; then
        log SUCCESS "GPU detected: $(nvidia-smi --query-gpu=name --format=csv,noheader | head -1)"
    else
        log WARNING "No GPU detected. Training will use CPU (slower)."
    fi

    # Check data directory
    if [ ! -d "$DATA_DIR" ]; then
        log ERROR "Data directory not found: $DATA_DIR"
        return 1
    fi

    local dbn_count=$(find "$DATA_DIR" -name "*.dbn" | wc -l)
    if [ "$dbn_count" -eq 0 ]; then
        log ERROR "No DBN files found in $DATA_DIR"
        return 1
    fi
    log INFO "Found $dbn_count DBN files in data directory"

    # Check hyperparameters file
    if [ ! -f "$HYPERPARAMS_FILE" ]; then
        log WARNING "Hyperparameters file not found: $HYPERPARAMS_FILE"
        log WARNING "Will use default hyperparameters"
    else
        log SUCCESS "Hyperparameters file found: $HYPERPARAMS_FILE"
    fi

    # Check Docker services
    if command -v docker-compose &> /dev/null; then
        if docker-compose ps | grep -q "Up"; then
            log SUCCESS "Docker services are running"
        else
            log WARNING "Some Docker services may not be running"
            log WARNING "Starting services with: docker-compose up -d"
            cd "$PROJECT_ROOT" && docker-compose up -d
        fi
    fi

    # Create output directories
    mkdir -p "$LOG_DIR"
    mkdir -p "$OUTPUT_SUBDIR"

    log SUCCESS "Prerequisites check passed"
    return 0
}

# Run retraining pipeline
run_retraining() {
    log INFO "═══════════════════════════════════════════════════════"
    log INFO "Starting quarterly model retraining"
    log INFO "═══════════════════════════════════════════════════════"
    log INFO "Version: $VERSION_TAG"
    log INFO "Models: $MODELS"
    log INFO "Mode: $([ "$PARALLEL" = true ] && echo "parallel" || echo "sequential")"
    log INFO "Data: Latest $LATEST_DAYS days"
    log INFO "Output: $OUTPUT_SUBDIR"
    log INFO "Dry run: $DRY_RUN"
    log INFO ""

    cd "$PROJECT_ROOT"

    # Build command
    local cmd="cargo run -p ml --example retrain_all_models --release --features cuda --"
    cmd="$cmd --models $MODELS"
    cmd="$cmd --data-dir $DATA_DIR"
    cmd="$cmd --output-dir $OUTPUT_SUBDIR"
    cmd="$cmd --hyperparams-file $HYPERPARAMS_FILE"
    cmd="$cmd --latest-days $LATEST_DAYS"
    cmd="$cmd --min-sharpe $MIN_SHARPE"
    cmd="$cmd --min-win-rate $MIN_WIN_RATE"
    cmd="$cmd --version-tag $VERSION_TAG"

    if [ "$PARALLEL" = true ]; then
        cmd="$cmd --parallel"
        log WARNING "⚠️  Parallel training may cause GPU OOM on RTX 3050 Ti"
    fi

    if [ "$DRY_RUN" = true ]; then
        cmd="$cmd --dry-run"
    fi

    log INFO "Executing: $cmd"
    log INFO ""

    # Execute retraining
    local start_time=$(date +%s)

    if eval "$cmd" 2>&1 | tee -a "$LOG_FILE"; then
        local end_time=$(date +%s)
        local duration=$((end_time - start_time))
        local hours=$((duration / 3600))
        local minutes=$(((duration % 3600) / 60))

        log SUCCESS "Retraining completed successfully"
        log INFO "Duration: ${hours}h ${minutes}m"

        return 0
    else
        local exit_code=$?
        log ERROR "Retraining failed with exit code $exit_code"
        return $exit_code
    fi
}

# Analyze results
analyze_results() {
    log INFO "Analyzing retraining results..."

    local summary_file="$OUTPUT_SUBDIR/retraining_summary_${VERSION_TAG}.json"

    if [ ! -f "$summary_file" ]; then
        log ERROR "Summary file not found: $summary_file"
        return 1
    fi

    # Parse results using jq (if available)
    if command -v jq &> /dev/null; then
        local attempted=$(jq -r '.models_attempted' "$summary_file")
        local succeeded=$(jq -r '.models_succeeded' "$summary_file")
        local failed=$(jq -r '.models_failed' "$summary_file")
        local passed=$(jq -r '.models_passed_quality_gate' "$summary_file")

        log INFO "Results:"
        log INFO "  • Models attempted: $attempted"
        log INFO "  • Models succeeded: $succeeded"
        log INFO "  • Models failed: $failed"
        log INFO "  • Quality gate passed: $passed"
        log INFO ""

        # Print per-model results
        log INFO "Per-model results:"
        jq -r '.results[] | "  • \(.model_type): Sharpe=\(.validation_metrics.sharpe_ratio), WinRate=\(.validation_metrics.win_rate*100)%, QualityGate=\(if .quality_gate_passed then "✅" else "❌" end)"' "$summary_file" | tee -a "$LOG_FILE"

        # Determine overall status
        if [ "$passed" -eq "$attempted" ]; then
            log SUCCESS "All models passed quality gates! ✅"
            return 0
        elif [ "$passed" -gt 0 ]; then
            log WARNING "Some models passed quality gates ($passed/$attempted)"
            return 1
        else
            log ERROR "No models passed quality gates"
            return 2
        fi
    else
        log WARNING "jq not installed, cannot parse results"
        log INFO "Summary file: $summary_file"
        return 0
    fi
}

# Main execution
main() {
    log INFO "Foxhunt Quarterly Model Retraining"
    log INFO "Started at: $(date)"
    log INFO ""

    # Check prerequisites
    if ! check_prerequisites; then
        log ERROR "Prerequisites check failed"
        notify ERROR "Prerequisites check failed. Aborting retraining."
        exit 3
    fi

    # Run retraining
    notify INFO "Starting quarterly retraining for $VERSION_TAG"

    local retrain_exit_code=0
    if ! run_retraining; then
        retrain_exit_code=$?
        log ERROR "Retraining execution failed"
        notify ERROR "Retraining failed with exit code $retrain_exit_code. Check logs: $LOG_FILE"

        if [ "$DRY_RUN" = true ]; then
            exit $retrain_exit_code
        fi
    fi

    # Analyze results (skip if dry run)
    if [ "$DRY_RUN" = false ]; then
        local analysis_exit_code=0
        if ! analyze_results; then
            analysis_exit_code=$?
        fi

        # Send final notification
        case $analysis_exit_code in
            0)
                notify SUCCESS "Quarterly retraining completed successfully! All models passed quality gates. Version: $VERSION_TAG"
                ;;
            1)
                notify WARNING "Quarterly retraining completed with warnings. Some models failed quality gates. Version: $VERSION_TAG. Review: $summary_file"
                ;;
            2)
                notify ERROR "Quarterly retraining failed. No models passed quality gates. Version: $VERSION_TAG. Review: $summary_file"
                ;;
        esac

        log INFO ""
        log INFO "Completed at: $(date)"
        log INFO "Log file: $LOG_FILE"
        log INFO "Summary: $OUTPUT_SUBDIR/retraining_summary_${VERSION_TAG}.json"

        exit $analysis_exit_code
    else
        log INFO "Dry run completed successfully"
        exit 0
    fi
}

# Execute main
main "$@"