#!/bin/bash

# UNIFIED TRAINING MONITORING DASHBOARD - Agent 134
# Monitors all 5 model training processes with GPU metrics, epoch progress, and alerting

set -e

# Color definitions
RED='\033[0;31m'
YELLOW='\033[1;33m'
GREEN='\033[0;32m'
BLUE='\033[0;34m'
CYAN='\033[0;36m'
MAGENTA='\033[0;35m'
NC='\033[0m' # No Color

# Monitoring configuration
REFRESH_INTERVAL=30  # Refresh every 30 seconds
LOG_DIR="/home/jgrusewski/Work/foxhunt"
ALERT_LOG="/tmp/training_alerts.log"
STATUS_FILE="/tmp/training_dashboard_status.txt"

# Training process definitions (name, log file, expected epochs, PID file)
declare -A TRAINING_PROCESSES=(
    ["TFT"]="tft_training_output.log:200:/tmp/tft_training.pid"
    ["MAMBA2"]="mamba2_training_output.log:200:/tmp/mamba2_training.pid"
    ["Liquid"]="liquid_training_output.log:200:/tmp/liquid_training.pid"
    ["DQN"]="/tmp/tuning_run.log:50:/tmp/dqn_tuning.pid"
    ["PPO"]="/tmp/ppo_tuning_run.log:50:/tmp/ppo_tuning.pid"
)

# Function to log alerts
log_alert() {
    local level=$1
    local model=$2
    local message=$3
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] [$level] [$model] $message" >> "$ALERT_LOG"
}

# Function to get GPU metrics
get_gpu_metrics() {
    if command -v nvidia-smi &> /dev/null; then
        nvidia-smi --query-gpu=index,name,utilization.gpu,memory.used,memory.total,temperature.gpu,power.draw --format=csv,noheader,nounits 2>/dev/null || echo "0,N/A,0,0,0,0,0"
    else
        echo "0,N/A,0,0,0,0,0"
    fi
}

# Function to parse GPU metrics into readable format
format_gpu_metrics() {
    local metrics=$1
    local gpu_util=$(echo "$metrics" | cut -d',' -f3 | xargs)
    local mem_used=$(echo "$metrics" | cut -d',' -f4 | xargs)
    local mem_total=$(echo "$metrics" | cut -d',' -f5 | xargs)
    local temp=$(echo "$metrics" | cut -d',' -f6 | xargs)
    local power=$(echo "$metrics" | cut -d',' -f7 | xargs)

    # Calculate memory percentage
    local mem_percent=0
    if [ "$mem_total" -gt 0 ] 2>/dev/null; then
        mem_percent=$(echo "scale=1; $mem_used * 100 / $mem_total" | bc 2>/dev/null || echo "0")
    fi

    # Color code based on utilization
    local util_color=$GREEN
    if [ "$gpu_util" -gt 70 ] 2>/dev/null; then
        util_color=$YELLOW
    fi
    if [ "$gpu_util" -gt 90 ] 2>/dev/null; then
        util_color=$RED
    fi

    echo -e "${util_color}GPU: ${gpu_util}%${NC} | VRAM: ${mem_used}/${mem_total}MB (${mem_percent}%) | Temp: ${temp}°C | Power: ${power}W"
}

# Function to get process status
get_process_status() {
    local pid_file=$1

    if [ ! -f "$pid_file" ]; then
        echo "NOT_STARTED"
        return
    fi

    local pid=$(cat "$pid_file" 2>/dev/null)
    if [ -z "$pid" ]; then
        echo "NOT_STARTED"
        return
    fi

    if ps -p "$pid" > /dev/null 2>&1; then
        echo "RUNNING:$pid"
    else
        echo "STOPPED"
    fi
}

# Function to extract epoch progress from log file
get_epoch_progress() {
    local log_file=$1
    local model_name=$2

    if [ ! -f "$log_file" ]; then
        echo "0/0|0|N/A"
        return
    fi

    # Different parsing logic for different models
    case "$model_name" in
        "TFT"|"MAMBA2"|"Liquid")
            # Look for epoch completion patterns like "Epoch 45/200" or "Epoch 45 complete"
            local last_epoch=$(grep -oP "Epoch \K\d+" "$log_file" 2>/dev/null | tail -1 || echo "0")
            local total_epochs=$(grep -oP "Epoch \d+/\K\d+" "$log_file" 2>/dev/null | head -1 || echo "0")
            local last_loss=$(grep -oP "loss: \K[\d\.]+" "$log_file" 2>/dev/null | tail -1 || echo "N/A")

            # If total_epochs not found, use expected
            [ "$total_epochs" == "0" ] && total_epochs=$(echo "${TRAINING_PROCESSES[$model_name]}" | cut -d':' -f2)

            # Calculate percentage
            local percent=0
            if [ "$total_epochs" -gt 0 ] 2>/dev/null && [ "$last_epoch" -gt 0 ] 2>/dev/null; then
                percent=$(echo "scale=1; $last_epoch * 100 / $total_epochs" | bc 2>/dev/null || echo "0")
            fi

            echo "${last_epoch}/${total_epochs}|${percent}|${last_loss}"
            ;;

        "DQN"|"PPO")
            # Look for trial completion patterns like "Trial 23 completed"
            local trials=0
            if [ -f "$log_file" ]; then
                trials=$(grep -c "Trial .* completed" "$log_file" 2>/dev/null | tr -d '\n' || echo "0")
            fi
            [ -z "$trials" ] && trials=0
            [ "$trials" == "" ] && trials=0

            local total_trials=$(echo "${TRAINING_PROCESSES[$model_name]}" | cut -d':' -f2 | tr -d '\n')
            [ -z "$total_trials" ] && total_trials=50

            # Calculate percentage - use awk for safer arithmetic
            local percent="0"
            if [ "$total_trials" -gt 0 ] 2>/dev/null && [ "$trials" -ge 0 ] 2>/dev/null; then
                percent=$(awk "BEGIN {printf \"%.1f\", ($trials * 100.0 / $total_trials)}" 2>/dev/null || echo "0")
            fi

            # Get last trial's best value
            local best_value="N/A"
            if [ -f "$log_file" ]; then
                best_value=$(grep -oP "Best value: \K[\d\.]+" "$log_file" 2>/dev/null | tail -1 | tr -d '\n' || echo "N/A")
            fi
            [ -z "$best_value" ] && best_value="N/A"

            echo "${trials}/${total_trials}|${percent}|${best_value}"
            ;;
    esac
}

# Function to estimate time remaining
estimate_time_remaining() {
    local pid=$1
    local current_epoch=$2
    local total_epochs=$3

    if [ "$total_epochs" -le "$current_epoch" ] || [ "$current_epoch" -eq 0 ]; then
        echo "N/A"
        return
    fi

    # Get process elapsed time in seconds
    local elapsed_seconds=$(ps -p "$pid" -o etimes= 2>/dev/null | xargs || echo "0")

    if [ "$elapsed_seconds" -eq 0 ] || [ "$current_epoch" -eq 0 ]; then
        echo "N/A"
        return
    fi

    # Calculate time per epoch
    local seconds_per_epoch=$((elapsed_seconds / current_epoch))

    # Calculate remaining epochs
    local remaining_epochs=$((total_epochs - current_epoch))

    # Calculate remaining time
    local remaining_seconds=$((seconds_per_epoch * remaining_epochs))

    # Format as HH:MM:SS
    local hours=$((remaining_seconds / 3600))
    local minutes=$(((remaining_seconds % 3600) / 60))
    local seconds=$((remaining_seconds % 60))

    printf "%02d:%02d:%02d" "$hours" "$minutes" "$seconds"
}

# Function to get process runtime
get_runtime() {
    local pid=$1
    if ps -p "$pid" > /dev/null 2>&1; then
        ps -p "$pid" -o etime= | xargs
    else
        echo "N/A"
    fi
}

# Function to check for OOM or crashes
check_for_errors() {
    local log_file=$1
    local model_name=$2

    if [ ! -f "$log_file" ]; then
        return
    fi

    # Check for recent errors (last 100 lines)
    local errors=$(tail -100 "$log_file" 2>/dev/null | grep -iE "error|panic|killed|out of memory|oom|cuda error|segmentation fault" || true)

    if [ -n "$errors" ]; then
        log_alert "ERROR" "$model_name" "Errors detected in log file"
        echo -e "${RED}⚠️  ERRORS DETECTED${NC}"
        return 1
    fi

    return 0
}

# Function to display model status
display_model_status() {
    local model_name=$1
    local config="${TRAINING_PROCESSES[$model_name]}"

    local log_file=$(echo "$config" | cut -d':' -f1)
    local expected_epochs=$(echo "$config" | cut -d':' -f2)
    local pid_file=$(echo "$config" | cut -d':' -f3)

    # Resolve full log path
    if [[ ! "$log_file" =~ ^/ ]]; then
        log_file="${LOG_DIR}/${log_file}"
    fi

    # Get process status
    local status=$(get_process_status "$pid_file")
    local status_type=$(echo "$status" | cut -d':' -f1)
    local pid=$(echo "$status" | cut -d':' -f2 2>/dev/null || echo "")

    # Get epoch progress
    local progress=$(get_epoch_progress "$log_file" "$model_name")
    local epochs=$(echo "$progress" | cut -d'|' -f1)
    local percent=$(echo "$progress" | cut -d'|' -f2)
    local metric=$(echo "$progress" | cut -d'|' -f3)

    local current_epoch=$(echo "$epochs" | cut -d'/' -f1)
    local total_epochs=$(echo "$epochs" | cut -d'/' -f2)

    # Color code status
    local status_color=$YELLOW
    local status_icon="⏸️ "
    case "$status_type" in
        "RUNNING")
            status_color=$GREEN
            status_icon="🟢"
            ;;
        "STOPPED")
            status_color=$RED
            status_icon="🔴"
            ;;
        "NOT_STARTED")
            status_color=$YELLOW
            status_icon="⚪"
            ;;
    esac

    # Display header
    echo -e "${CYAN}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
    echo -e "${MAGENTA}${model_name}${NC} ${status_icon} ${status_color}${status_type}${NC}"

    # Display PID and runtime
    if [ -n "$pid" ]; then
        local runtime=$(get_runtime "$pid")
        echo -e "  PID: ${pid} | Runtime: ${runtime}"

        # Get process memory usage
        local mem_mb=$(ps -p "$pid" -o rss= 2>/dev/null | awk '{printf "%.1f", $1/1024}' || echo "0")
        echo -e "  Memory: ${mem_mb}MB"
    fi

    # Display progress
    echo -e "  Progress: ${epochs} (${percent}%)"

    # Display progress bar
    local bar_length=40
    # Handle decimal percentage
    local percent_int=$(echo "$percent" | cut -d'.' -f1)
    [ -z "$percent_int" ] && percent_int=0
    local filled=$((percent_int * bar_length / 100))
    local empty=$((bar_length - filled))

    local bar_color=$GREEN
    if [ "$percent_int" -lt 30 ] 2>/dev/null; then
        bar_color=$YELLOW
    fi
    if [ "$percent_int" -lt 10 ] 2>/dev/null; then
        bar_color=$RED
    fi

    printf "  ["
    printf "${bar_color}%${filled}s${NC}" | tr ' ' '█'
    printf "%${empty}s" | tr ' ' '░'
    printf "]\n"

    # Display metric
    if [ "$metric" != "N/A" ]; then
        if [[ "$model_name" == "DQN" || "$model_name" == "PPO" ]]; then
            echo -e "  Best Value: ${metric}"
        else
            echo -e "  Last Loss: ${metric}"
        fi
    fi

    # Display time remaining estimate
    if [ "$status_type" == "RUNNING" ] && [ -n "$pid" ]; then
        local time_remaining=$(estimate_time_remaining "$pid" "$current_epoch" "$total_epochs")
        if [ "$time_remaining" != "N/A" ]; then
            echo -e "  ETA: ${time_remaining}"
        fi
    fi

    # Check for errors
    check_for_errors "$log_file" "$model_name" || true

    # Display log file location
    echo -e "  Log: ${log_file}"

    echo ""
}

# Function to display summary statistics
display_summary() {
    local running=0
    local stopped=0
    local not_started=0
    local total=0
    local total_progress=0

    for model_name in "${!TRAINING_PROCESSES[@]}"; do
        local config="${TRAINING_PROCESSES[$model_name]}"
        local pid_file=$(echo "$config" | cut -d':' -f3)
        local status=$(get_process_status "$pid_file")
        local status_type=$(echo "$status" | cut -d':' -f1)

        case "$status_type" in
            "RUNNING") ((running++)) ;;
            "STOPPED") ((stopped++)) ;;
            "NOT_STARTED") ((not_started++)) ;;
        esac

        # Get progress for running processes
        if [ "$status_type" == "RUNNING" ]; then
            local log_file=$(echo "$config" | cut -d':' -f1)
            if [[ ! "$log_file" =~ ^/ ]]; then
                log_file="${LOG_DIR}/${log_file}"
            fi
            local progress=$(get_epoch_progress "$log_file" "$model_name")
            local percent=$(echo "$progress" | cut -d'|' -f2)
            total_progress=$(echo "scale=1; $total_progress + $percent" | bc)
        fi

        ((total++))
    done

    # Calculate average progress
    local avg_progress=0
    if [ "$running" -gt 0 ]; then
        avg_progress=$(echo "scale=1; $total_progress / $running" | bc)
    fi

    echo -e "${CYAN}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
    echo -e "${BLUE}SUMMARY${NC}"
    echo -e "  Total Models: ${total}"
    echo -e "  ${GREEN}Running: ${running}${NC} | ${RED}Stopped: ${stopped}${NC} | ${YELLOW}Not Started: ${not_started}${NC}"
    if [ "$running" -gt 0 ]; then
        echo -e "  Average Progress: ${avg_progress}%"
    fi
    echo ""
}

# Function to display consolidated log viewer commands
display_log_commands() {
    echo -e "${CYAN}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
    echo -e "${BLUE}LOG VIEWER COMMANDS${NC}"
    echo ""

    for model_name in "${!TRAINING_PROCESSES[@]}"; do
        local config="${TRAINING_PROCESSES[$model_name]}"
        local log_file=$(echo "$config" | cut -d':' -f1)

        if [[ ! "$log_file" =~ ^/ ]]; then
            log_file="${LOG_DIR}/${log_file}"
        fi

        echo -e "  ${MAGENTA}${model_name}${NC}: tail -f ${log_file}"
    done

    echo ""
    echo -e "  ${MAGENTA}All Alerts${NC}: tail -f ${ALERT_LOG}"
    echo ""
}

# Function to check system resources
check_system_resources() {
    # Check memory
    local mem_percent=$(free | grep Mem | awk '{printf "%.0f", $3*100/$2}')
    local mem_color=$GREEN
    if [ "$mem_percent" -gt 70 ] 2>/dev/null; then
        mem_color=$YELLOW
    fi
    if [ "$mem_percent" -gt 90 ] 2>/dev/null; then
        mem_color=$RED
    fi

    # Check disk
    local disk_percent=$(df -h / | tail -1 | awk '{print $5}' | tr -d '%')
    local disk_color=$GREEN
    if [ "$disk_percent" -gt 70 ] 2>/dev/null; then
        disk_color=$YELLOW
    fi
    if [ "$disk_percent" -gt 85 ] 2>/dev/null; then
        disk_color=$RED
    fi

    echo -e "${CYAN}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
    echo -e "${BLUE}SYSTEM RESOURCES${NC}"
    echo -e "  Memory: ${mem_color}${mem_percent}%${NC}"
    echo -e "  Disk: ${disk_color}${disk_percent}%${NC}"

    # Get GPU metrics
    local gpu_metrics=$(get_gpu_metrics)
    echo -e "  $(format_gpu_metrics "$gpu_metrics")"

    # Alert if resources critical
    if [ "$mem_percent" -gt 90 ]; then
        log_alert "CRITICAL" "SYSTEM" "Memory usage critical: ${mem_percent}%"
        echo -e "  ${RED}⚠️  CRITICAL: Memory usage >90%${NC}"
    fi

    if [ "$disk_percent" -gt 85 ]; then
        log_alert "WARNING" "SYSTEM" "Disk usage high: ${disk_percent}%"
        echo -e "  ${YELLOW}⚠️  WARNING: Disk usage >85%${NC}"
    fi

    echo ""
}

# Main monitoring loop
monitor_training() {
    echo -e "${GREEN}Starting unified training monitoring dashboard...${NC}"
    echo -e "${GREEN}Press Ctrl+C to stop${NC}"
    echo ""

    # Initialize alert log
    touch "$ALERT_LOG"

    while true; do
        # Clear screen
        clear

        # Display header
        echo -e "${CYAN}╔════════════════════════════════════════════════════════╗${NC}"
        echo -e "${CYAN}║     ${MAGENTA}UNIFIED TRAINING MONITORING DASHBOARD${CYAN}           ║${NC}"
        echo -e "${CYAN}╚════════════════════════════════════════════════════════╝${NC}"
        echo -e "${BLUE}Updated: $(date '+%Y-%m-%d %H:%M:%S')${NC}"
        echo ""

        # Check system resources
        check_system_resources

        # Display each model status
        for model_name in TFT MAMBA2 Liquid DQN PPO; do
            if [ -n "${TRAINING_PROCESSES[$model_name]}" ]; then
                display_model_status "$model_name"
            fi
        done

        # Display summary
        display_summary

        # Display log commands
        display_log_commands

        # Display footer
        echo -e "${CYAN}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
        echo -e "${BLUE}CONTROLS${NC}"
        echo -e "  Refresh interval: ${REFRESH_INTERVAL}s"
        echo -e "  Press Ctrl+C to exit"
        echo -e "  Alert log: ${ALERT_LOG}"
        echo -e "${CYAN}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"

        # Save status to file
        {
            echo "TRAINING_DASHBOARD_STATUS"
            echo "Updated: $(date '+%Y-%m-%d %H:%M:%S')"
            echo ""
            for model_name in "${!TRAINING_PROCESSES[@]}"; do
                local config="${TRAINING_PROCESSES[$model_name]}"
                local pid_file=$(echo "$config" | cut -d':' -f3)
                local status=$(get_process_status "$pid_file")
                echo "$model_name: $status"
            done
        } > "$STATUS_FILE"

        # Wait for next refresh
        sleep "$REFRESH_INTERVAL"
    done
}

# Function to display one-time status
display_status() {
    clear

    # Display header
    echo -e "${CYAN}╔════════════════════════════════════════════════════════╗${NC}"
    echo -e "${CYAN}║     ${MAGENTA}TRAINING STATUS SNAPSHOT${CYAN}                       ║${NC}"
    echo -e "${CYAN}╚════════════════════════════════════════════════════════╝${NC}"
    echo -e "${BLUE}Generated: $(date '+%Y-%m-%d %H:%M:%S')${NC}"
    echo ""

    # Check system resources
    check_system_resources

    # Display each model status
    for model_name in TFT MAMBA2 Liquid DQN PPO; do
        if [ -n "${TRAINING_PROCESSES[$model_name]}" ]; then
            display_model_status "$model_name"
        fi
    done

    # Display summary
    display_summary

    # Display recent alerts
    if [ -f "$ALERT_LOG" ]; then
        echo -e "${CYAN}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
        echo -e "${BLUE}RECENT ALERTS (Last 10)${NC}"
        tail -10 "$ALERT_LOG" 2>/dev/null || echo "No alerts"
        echo ""
    fi

    echo -e "${CYAN}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
}

# Main command handler
case "${1:-monitor}" in
    monitor)
        monitor_training
        ;;
    status)
        display_status
        ;;
    alerts)
        if [ -f "$ALERT_LOG" ]; then
            cat "$ALERT_LOG"
        else
            echo "No alerts logged yet"
        fi
        ;;
    clear-alerts)
        > "$ALERT_LOG"
        echo "Alert log cleared"
        ;;
    *)
        echo "Usage: $0 {monitor|status|alerts|clear-alerts}"
        echo ""
        echo "Commands:"
        echo "  monitor       - Start continuous monitoring dashboard (default)"
        echo "  status        - Show one-time status snapshot"
        echo "  alerts        - Display all logged alerts"
        echo "  clear-alerts  - Clear alert log"
        echo ""
        echo "Examples:"
        echo "  $0 monitor              # Start live dashboard"
        echo "  $0 status               # Quick status check"
        echo "  watch -n 30 $0 status   # Auto-refresh status every 30s"
        exit 1
        ;;
esac