#!/bin/bash
# Verify TFT Checkpoint Fix
# This script validates that the TFT checkpoint serialization bug is fixed

set -e

echo "🔍 Verifying TFT Checkpoint Fix..."
echo

# 1. Check that the fix was applied
echo "1. Checking TFT trainer code for VarMap fix..."
if grep -q "let var_map = model.get_varmap().clone();" /home/jgrusewski/Work/foxhunt/ml/src/trainers/tft.rs; then
    echo "   ✅ Fix applied: Trainer now uses model's VarMap"
else
    echo "   ❌ Fix NOT found in code"
    exit 1
fi

# 2. Check existing checkpoint file size
echo
echo "2. Checking existing checkpoint file..."
CHECKPOINT_FILE="/home/jgrusewski/Work/foxhunt/ml/trained_models/tft_epoch_9.safetensors"
if [ -f "$CHECKPOINT_FILE" ]; then
    FILE_SIZE=$(stat -c%s "$CHECKPOINT_FILE")
    echo "   Current file size: $FILE_SIZE bytes"
    if [ "$FILE_SIZE" -eq 16 ]; then
        echo "   ⚠️  File is 16 bytes (empty VarMap - bug confirmed)"
    else
        echo "   ✅ File has weights (size > 16 bytes)"
    fi
else
    echo "   ℹ️  No existing checkpoint file found"
fi

# 3. Verify SafeTensors format
echo
echo "3. Analyzing SafeTensors format..."
if [ -f "$CHECKPOINT_FILE" ]; then
    HEADER=$(hexdump -C "$CHECKPOINT_FILE" | head -1)
    echo "   Header: $HEADER"
    if echo "$HEADER" | grep -q "7b 7d"; then
        echo "   ⚠️  Contains '{}' (empty JSON - no tensors)"
    else
        echo "   ✅ Contains tensor data"
    fi
fi

# 4. Check code comment explaining the fix
echo
echo "4. Summary of Fix:"
echo "   • Root Cause: Trainer created separate empty VarMap instead of using model's VarMap"
echo "   • Fix Location: /home/jgrusewski/Work/foxhunt/ml/src/trainers/tft.rs:307"
echo "   • Fix Applied: Changed 'VarMap::new()' to 'model.get_varmap().clone()'"
echo "   • Expected Result: Checkpoint file should be >50MB with FP32 weights"
echo

# 5. Recommendation
echo "5. Next Steps:"
echo "   ✅ Fix has been applied to code"
echo "   📋 Re-run training to generate new checkpoint:"
echo "      cargo run -p ml --example train_tft_dbn --release -- --epochs 1"
echo "   🔍 Verify new checkpoint file size is >1MB (should be ~100MB for full model)"
echo

echo "✅ Verification complete!"