#!/usr/bin/env bash
# Validation script: Train all 4 models for 2 epochs each
# Wave 152: Agent 20 - Quick validation of all training pipelines
# Dependencies: Agent 19 (train_all_models_fixed.sh)
# Success criteria: All 4 models produce .safetensors files, exit 0 if pass, exit 1 if fail

set -euo 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

# Script directory
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)"
cd "$PROJECT_ROOT"

# Configuration
EPOCHS=2
DATA_DIR="test_data/real"
MODEL_OUTPUT_DIR="test_data/models"
TIMESTAMP=$(date +%Y%m%d_%H%M%S)

# Test data files (using 3-month dataset from Agent 19)
BTC_DATA="${DATA_DIR}/BTC-USD_20231001-20231231_databento_ohlcv-1s.parquet"
ETH_DATA="${DATA_DIR}/ETH-USD_20231001-20231231_databento_ohlcv-1s.parquet"

# Model names
MODELS=("DQN" "PPO" "MAMBA" "TFT")

# Output files to check
declare -A MODEL_FILES=(
    ["DQN"]="${MODEL_OUTPUT_DIR}/dqn_${TIMESTAMP}"
    ["PPO"]="${MODEL_OUTPUT_DIR}/ppo_${TIMESTAMP}"
    ["MAMBA"]="${MODEL_OUTPUT_DIR}/mamba_${TIMESTAMP}"
    ["TFT"]="${MODEL_OUTPUT_DIR}/tft_${TIMESTAMP}"
)

# Results tracking
declare -A MODEL_STATUS
declare -A MODEL_TIME
declare -A MODEL_OUTPUT

# Create output directory
mkdir -p "${MODEL_OUTPUT_DIR}"

# Banner
echo -e "${BLUE}========================================${NC}"
echo -e "${BLUE}ML Training Validation Script${NC}"
echo -e "${BLUE}Wave 152 Agent 20${NC}"
echo -e "${BLUE}========================================${NC}"
echo ""
echo -e "${YELLOW}Configuration:${NC}"
echo "  Epochs: ${EPOCHS}"
echo "  Data: ${DATA_DIR}/"
echo "  Output: ${MODEL_OUTPUT_DIR}/"
echo "  Models: ${MODELS[*]}"
echo ""

# Check prerequisites
echo -e "${BLUE}Checking prerequisites...${NC}"

# Check data files exist
if [[ ! -f "$BTC_DATA" ]]; then
    echo -e "${RED}ERROR: BTC data not found: ${BTC_DATA}${NC}"
    echo "Please run Agent 19 (train_all_models_fixed.sh) first to download data"
    exit 1
fi

if [[ ! -f "$ETH_DATA" ]]; then
    echo -e "${RED}ERROR: ETH data not found: ${ETH_DATA}${NC}"
    echo "Please run Agent 19 (train_all_models_fixed.sh) first to download data"
    exit 1
fi

echo -e "${GREEN}✓ Data files found${NC}"

# Check cargo is available
if ! command -v cargo &> /dev/null; then
    echo -e "${RED}ERROR: cargo not found${NC}"
    exit 1
fi

echo -e "${GREEN}✓ Cargo available${NC}"
echo ""

# Function to train a single model
train_model() {
    local model_name=$1
    local model_type=$2
    local output_file=$3

    echo -e "${BLUE}Training ${model_name} (${EPOCHS} epochs)...${NC}"
    local start_time=$(date +%s)

    # Build training command
    local cmd="cargo run --release --bin ml_training_cli -- train-model"
    cmd+=" --model-type ${model_type}"
    cmd+=" --data-path ${BTC_DATA}"
    cmd+=" --output-path ${output_file}"
    cmd+=" --epochs ${EPOCHS}"
    cmd+=" --batch-size 32"
    cmd+=" --learning-rate 0.001"

    # Run training and capture output
    local log_file="${MODEL_OUTPUT_DIR}/${model_name}_${TIMESTAMP}.log"
    if $cmd > "$log_file" 2>&1; then
        local end_time=$(date +%s)
        local duration=$((end_time - start_time))

        MODEL_STATUS["$model_name"]="SUCCESS"
        MODEL_TIME["$model_name"]="$duration"
        MODEL_OUTPUT["$model_name"]="$log_file"

        echo -e "${GREEN}✓ ${model_name} training completed (${duration}s)${NC}"
        return 0
    else
        local end_time=$(date +%s)
        local duration=$((end_time - start_time))

        MODEL_STATUS["$model_name"]="FAILED"
        MODEL_TIME["$model_name"]="$duration"
        MODEL_OUTPUT["$model_name"]="$log_file"

        echo -e "${RED}✗ ${model_name} training failed (${duration}s)${NC}"
        echo "  Log: $log_file"
        return 1
    fi
}

# Function to check model output
check_model_output() {
    local model_name=$1
    local output_path=$2

    # Check for .safetensors file
    if [[ -f "${output_path}.safetensors" ]]; then
        local file_size=$(du -h "${output_path}.safetensors" | cut -f1)
        echo -e "${GREEN}✓ ${model_name} model saved: ${file_size}${NC}"
        return 0
    else
        echo -e "${RED}✗ ${model_name} model NOT saved (no .safetensors file)${NC}"
        return 1
    fi
}

# Train all models
echo -e "${BLUE}========================================${NC}"
echo -e "${BLUE}Training Phase${NC}"
echo -e "${BLUE}========================================${NC}"
echo ""

success_count=0
fail_count=0

# 1. Train DQN
if train_model "DQN" "dqn" "${MODEL_FILES["DQN"]}"; then
    ((success_count++))
else
    ((fail_count++))
fi
echo ""

# 2. Train PPO
if train_model "PPO" "ppo" "${MODEL_FILES["PPO"]}"; then
    ((success_count++))
else
    ((fail_count++))
fi
echo ""

# 3. Train MAMBA-2
if train_model "MAMBA" "mamba" "${MODEL_FILES["MAMBA"]}"; then
    ((success_count++))
else
    ((fail_count++))
fi
echo ""

# 4. Train TFT
if train_model "TFT" "tft" "${MODEL_FILES["TFT"]}"; then
    ((success_count++))
else
    ((fail_count++))
fi
echo ""

# Validation phase
echo -e "${BLUE}========================================${NC}"
echo -e "${BLUE}Validation Phase${NC}"
echo -e "${BLUE}========================================${NC}"
echo ""

validation_success=0
validation_fail=0

for model_name in "${MODELS[@]}"; do
    if [[ "${MODEL_STATUS[$model_name]}" == "SUCCESS" ]]; then
        if check_model_output "$model_name" "${MODEL_FILES[$model_name]}"; then
            ((validation_success++))
        else
            ((validation_fail++))
        fi
    else
        echo -e "${YELLOW}⊘ ${model_name} model skipped (training failed)${NC}"
        ((validation_fail++))
    fi
done

echo ""

# Summary
echo -e "${BLUE}========================================${NC}"
echo -e "${BLUE}Summary${NC}"
echo -e "${BLUE}========================================${NC}"
echo ""

echo -e "${YELLOW}Training Results:${NC}"
for model_name in "${MODELS[@]}"; do
    local status="${MODEL_STATUS[$model_name]}"
    local time="${MODEL_TIME[$model_name]}"
    local log="${MODEL_OUTPUT[$model_name]}"

    if [[ "$status" == "SUCCESS" ]]; then
        echo -e "  ${GREEN}✓${NC} ${model_name}: ${status} (${time}s)"
    else
        echo -e "  ${RED}✗${NC} ${model_name}: ${status} (${time}s)"
        echo "    Log: $log"
    fi
done

echo ""
echo -e "${YELLOW}Validation Results:${NC}"
echo "  Success: ${validation_success}/4 models"
echo "  Failed:  ${validation_fail}/4 models"
echo ""

# Final verdict
if [[ $validation_success -eq 4 ]]; then
    echo -e "${GREEN}========================================${NC}"
    echo -e "${GREEN}✓ ALL TESTS PASSED${NC}"
    echo -e "${GREEN}========================================${NC}"
    echo ""
    echo "All 4 models trained successfully and saved .safetensors files"
    echo ""
    echo "Model files:"
    for model_name in "${MODELS[@]}"; do
        echo "  - ${MODEL_FILES[$model_name]}.safetensors"
    done
    exit 0
else
    echo -e "${RED}========================================${NC}"
    echo -e "${RED}✗ TESTS FAILED${NC}"
    echo -e "${RED}========================================${NC}"
    echo ""
    echo "Failed: ${validation_fail}/4 models"
    echo ""
    echo "Check logs for details:"
    for model_name in "${MODELS[@]}"; do
        if [[ "${MODEL_STATUS[$model_name]}" != "SUCCESS" ]]; then
            echo "  - ${MODEL_OUTPUT[$model_name]}"
        fi
    done
    exit 1
fi