foxhunt/crates/ml/examples/cache_dqn_features.rs

//! DQN Feature Caching CLI Tool
//!
//! Pre-computes and caches DQN features to eliminate redundant computation during hyperopt.
//!
//! ## Purpose
//!
//! During hyperopt, each trial wastes 10-30s recalculating identical MBP-10 OFI features.
//! This tool pre-computes features once and saves them to disk for reuse across all trials.
//!
//! ## Time Savings
//!
//! - 50-trial hyperopt: **8-25 minutes saved** (87-96% reduction)
//! - Cache creation: 30s (one-time)
//! - Cache loading: <1s per trial
//!
//! ## Usage
//!
//! ```bash
//! # Pre-compute features once
//! cargo run -p ml --example cache_dqn_features --release -- \
//!   --parquet-file test_data/ES_FUT_180d.parquet \
//!   --mbp10-dir test_data/mbp10 \
//!   --cache-dir ~/.cache/dqn_features
//!
//! # Run hyperopt with cache
//! cargo run -p ml --example hyperopt_dqn_demo --release -- \
//!   --parquet-file test_data/ES_FUT_180d.parquet \
//!   --trials 50 \
//!   --feature-cache-dir ~/.cache/dqn_features
//! ```
//!
//! ## Cache Key Strategy
//!
//! Cache key is SHA256 hash of:
//! - Parquet file path + modification time
//! - MBP-10 directory contents + modification times
//! - Warmup period (50 bars)
//! - Feature extraction version (manual constant)
//!
//! Cache automatically invalidates when any input changes.

use anyhow::{Context, Result};
use arrow::array::{ArrayRef, Float64Array};
use arrow::datatypes::{DataType, Field, Schema};
use arrow::record_batch::RecordBatch;
use clap::Parser;
use parquet::arrow::ArrowWriter;
use parquet::file::properties::WriterProperties;
use parquet::basic::Compression;
use sha2::{Digest, Sha256};
use std::fs::{self, File};
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::time::UNIX_EPOCH;
use tracing::{info, warn};

use ml::trainers::dqn::{DQNHyperparameters, DQNTrainer};

/// Command-line arguments
#[derive(Parser, Debug)]
#[command(author, version, about, long_about = None)]
struct Args {
    /// Path to input Parquet file containing OHLCV data
    #[arg(long)]
    parquet_file: PathBuf,

    /// Path to MBP-10 snapshots directory (optional, will use zeros for OFI if missing)
    #[arg(long, default_value = "test_data/mbp10")]
    mbp10_dir: PathBuf,

    /// Output cache directory
    #[arg(long, default_value = "/tmp/ml_training/feature_cache")]
    cache_dir: PathBuf,

    /// Warmup period in bars (must match training configuration)
    #[arg(long, default_value = "50")]
    warmup_period: usize,

    /// Force rebuild cache even if it exists
    #[arg(long)]
    force: bool,
}

type FeatureVector51 = [f64; 51];

/// Calculate cache key from input parameters
///
/// Cache key is SHA256 hash of:
/// - Parquet file path + modification time
/// - MBP-10 directory contents + modification times
/// - Warmup period
/// - Feature extraction version
fn calculate_cache_key(
    parquet_path: &Path,
    mbp10_dir: &Path,
    warmup_period: usize,
) -> Result<String> {
    let mut hasher = Sha256::new();

    // 1. Hash parquet file path + mtime
    hasher.update(parquet_path.to_string_lossy().as_bytes());
    let parquet_metadata = fs::metadata(parquet_path)
        .with_context(|| format!("Failed to read metadata for {}", parquet_path.display()))?;
    let parquet_mtime = parquet_metadata
        .modified()?
        .duration_since(UNIX_EPOCH)?
        .as_secs();
    hasher.update(parquet_mtime.to_le_bytes());

    // 2. Hash MBP-10 directory + all .dbn files + mtimes
    if mbp10_dir.exists() {
        hasher.update(mbp10_dir.to_string_lossy().as_bytes());
        
        let mut dbn_files: Vec<_> = fs::read_dir(mbp10_dir)?
            .filter_map(|e| e.ok())
            .filter(|e| {
                e.path()
                    .extension()
                    .and_then(|s| s.to_str())
                    .map(|s| s == "dbn")
                    .unwrap_or(false)
            })
            .collect();
        
        dbn_files.sort_by_key(|e| e.path());
        
        for entry in &dbn_files {
            hasher.update(entry.path().to_string_lossy().as_bytes());
            let metadata = entry.metadata()?;
            let mtime = metadata
                .modified()?
                .duration_since(UNIX_EPOCH)?
                .as_secs();
            hasher.update(mtime.to_le_bytes());
        }
    }

    // 3. Hash feature extraction version (from extraction.rs module constant)
    hasher.update(ml::features::extraction::FEATURE_EXTRACTION_VERSION.as_bytes());

    // 4. Hash warmup period
    hasher.update(warmup_period.to_le_bytes());

    // Return first 16 chars of hex hash for readability
    Ok(format!("{:x}", hasher.finalize())[..16].to_string())
}

/// Write features and targets to Parquet with Snappy compression
fn write_features_to_parquet(
    features: &[(FeatureVector51, Vec<f64>)],
    output_path: &Path,
) -> Result<()> {
    // Create schema: 51 feature columns + 1 target column
    let mut fields = Vec::new();
    for i in 0..51 {
        fields.push(Field::new(&format!("feature_{}", i), DataType::Float64, false));
    }
    fields.push(Field::new("target_close", DataType::Float64, false));
    let schema = Arc::new(Schema::new(fields));

    // Configure Parquet writer with Snappy compression
    let props = WriterProperties::builder()
        .set_compression(Compression::SNAPPY)
        .build();

    let file = File::create(output_path)
        .with_context(|| format!("Failed to create cache file at {}", output_path.display()))?;
    let mut writer = ArrowWriter::try_new(file, schema.clone(), Some(props))?;

    // Convert features to Arrow RecordBatch
    let mut columns: Vec<ArrayRef> = Vec::new();
    
    // Add 51 feature columns
    for i in 0..51 {
        let values: Vec<f64> = features.iter()
            .map(|(state, _)| state[i])
            .collect();
        columns.push(Arc::new(Float64Array::from(values)));
    }
    
    // Add target_close column (preprocessed_current from targets[0])
    let target_values: Vec<f64> = features.iter()
        .map(|(_, targets)| targets[0])
        .collect();
    columns.push(Arc::new(Float64Array::from(target_values)));

    let batch = RecordBatch::try_new(schema, columns)?;
    writer.write(&batch)?;
    writer.close()?;

    Ok(())
}

/// Main cache generation logic
async fn generate_cache(args: &Args) -> Result<()> {
    // Validate inputs
    if !args.parquet_file.exists() {
        anyhow::bail!("Parquet file not found: {}", args.parquet_file.display());
    }

    if !args.mbp10_dir.exists() {
        warn!("⚠️  MBP-10 directory not found at {:?}. OFI features will be zeros.", args.mbp10_dir);
    }

    // Create cache directory if needed
    fs::create_dir_all(&args.cache_dir)
        .with_context(|| format!("Failed to create cache directory: {}", args.cache_dir.display()))?;

    info!("🚀 Starting DQN feature cache generation...");
    info!("");
    info!("📂 Input:");
    info!("  • Parquet: {}", args.parquet_file.display());
    info!("  • MBP-10 dir: {}", args.mbp10_dir.display());
    info!("  • Warmup period: {} bars", args.warmup_period);
    info!("");

    // Calculate cache key
    let cache_key = calculate_cache_key(&args.parquet_file, &args.mbp10_dir, args.warmup_period)?;
    info!("🔑 Cache key: {}", cache_key);
    info!("");

    // Check if cache already exists
    let cache_file = args.cache_dir.join(format!("features_{}.parquet", cache_key));
    if cache_file.exists() && !args.force {
        info!("✅ Cache already exists: {}", cache_file.display());
        info!("   Use --force to rebuild");
        return Ok(());
    }

    // Create DQN trainer to use existing feature extraction logic
    // Use conservative parameters (only needed for initialization, not used during feature extraction)
    let params = DQNHyperparameters::conservative();

    let mut trainer = DQNTrainer::new(params)
        .with_context(|| "Failed to create DQN trainer")?;

    // Load and extract features using existing DQN trainer logic
    info!("📊 Extracting features...");
    let parquet_path = args.parquet_file.to_string_lossy().to_string();
    let (train_data, val_data) = trainer
        .load_training_data_from_parquet(&parquet_path)
        .await
        .with_context(|| "Failed to load training data from parquet")?;

    let total_samples = train_data.len() + val_data.len();
    info!("  • Extracted {} training + {} validation samples", train_data.len(), val_data.len());
    info!("");

    // Combine train and val for caching (trainer will re-split later)
    let mut all_samples = train_data;
    all_samples.extend(val_data);

    // Write cache
    info!("💾 Writing cache...");
    info!("  • Output: {}", cache_file.display());
    
    let start = std::time::Instant::now();
    write_features_to_parquet(&all_samples, &cache_file)?;
    let write_duration = start.elapsed();
    
    let file_size = fs::metadata(&cache_file)?.len();
    info!("  • Size: {:.2} MB (compressed with Snappy)", file_size as f64 / 1_048_576.0);
    info!("  • Write time: {:.1}s", write_duration.as_secs_f64());
    info!("");

    info!("✅ Cache generation complete!");
    info!("  • Total samples: {}", total_samples);
    info!("  • Cache key: {}", cache_key);
    info!("  • Ready for hyperopt!");

    Ok(())
}

#[tokio::main]
async fn main() -> Result<()> {
    // Initialize logging
    tracing_subscriber::fmt()
        .with_env_filter(
            tracing_subscriber::EnvFilter::try_from_default_env()
                .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")),
        )
        .init();

    let args = Args::parse();
    
    generate_cache(&args).await?;

    Ok(())
}