foxhunt/ml/src/feature_cache.rs

//! Feature Cache for DQN Training
//!
//! This module provides a disk-based caching system for computed feature vectors,
//! eliminating redundant MBP-10 OFI calculations across hyperopt trials.
//!
//! ## Cache Architecture
//!
//! - **Storage Format**: Apache Parquet (columnar, Snappy compression)
//! - **Cache Key**: SHA256 hash of (parquet path, mbp10 dir, warmup, extraction.rs)
//! - **Metadata**: JSON index mapping cache keys to file locations
//! - **Expected Savings**: 8-25 minutes per 50-trial hyperopt campaign
//!
//! ## Usage
//!
//! ```rust
//! use ml::feature_cache::{calculate_cache_key, load_features_from_cache};
//!
//! let cache_key = calculate_cache_key(
//!     Path::new("test_data/ES_FUT_180d.parquet"),
//!     Path::new("test_data/mbp10"),
//!     50, // warmup period
//! )?;
//!
//! if let Some(features) = load_features_from_cache(cache_dir, &cache_key).await? {
//!     // Use cached features
//! } else {
//!     // Compute features and save to cache
//! }
//! ```

use anyhow::{Context, Result};
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::time::SystemTime;
use tracing::{info, warn};

/// Cache metadata for a single cached feature set
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CacheMetadata {
    /// Cache format version
    pub version: String,
    /// SHA256 cache key
    pub cache_key: String,
    /// Parquet file path
    pub parquet_path: String,
    /// Parquet file modification time (Unix timestamp)
    pub parquet_mtime: u64,
    /// MBP-10 directory path
    pub mbp10_dir: String,
    /// Number of .dbn files in MBP-10 directory
    pub mbp10_file_count: usize,
    /// MBP-10 file modification times (Unix timestamps)
    pub mbp10_mtimes: Vec<u64>,
    /// Number of features per sample (51)
    pub feature_count: usize,
    /// Number of bars cached
    pub bar_count: usize,
    /// Warmup period used
    pub warmup_period: usize,
    /// Cache creation timestamp (ISO 8601)
    pub created_at: String,
}

/// Calculate cache key from input parameters
///
/// Generates a SHA256 hash that uniquely identifies a cached feature set.
/// The hash includes:
/// 1. Parquet file path + modification time
/// 2. MBP-10 directory + all .dbn files + modification times
/// 3. Feature extraction code (extraction.rs hash)
/// 4. Warmup period constant
///
/// # Arguments
///
/// * `parquet_path` - Path to OHLCV parquet file
/// * `mbp10_dir` - Path to MBP-10 snapshot directory
/// * `warmup_period` - Warmup period for feature extraction (typically 50)
///
/// # Returns
///
/// Hex-encoded SHA256 hash (64 characters)
pub fn calculate_cache_key(
    parquet_path: &Path,
    mbp10_dir: &Path,
    warmup_period: usize,
) -> Result<String> {
    let mut hasher = Sha256::new();

    // 1. Parquet file path + mtime
    hasher.update(parquet_path.to_string_lossy().as_bytes());
    let parquet_mtime = parquet_path
        .metadata()
        .context("Failed to read parquet metadata")?
        .modified()
        .context("Failed to get parquet mtime")?
        .duration_since(SystemTime::UNIX_EPOCH)
        .context("Failed to convert parquet mtime")?
        .as_secs();
    hasher.update(&parquet_mtime.to_le_bytes());

    // 2. MBP-10 directory + all .dbn files + mtimes
    hasher.update(mbp10_dir.to_string_lossy().as_bytes());
    let mut dbn_files: Vec<_> = std::fs::read_dir(mbp10_dir)
        .context("Failed to read mbp10 directory")?
        .filter_map(|e| e.ok())
        .filter(|e| {
            e.path()
                .extension()
                .and_then(|s| s.to_str())
                == Some("dbn")
        })
        .collect();
    dbn_files.sort_by_key(|e| e.path());

    for entry in &dbn_files {
        hasher.update(entry.path().to_string_lossy().as_bytes());
        let mtime = entry
            .metadata()
            .context("Failed to read dbn metadata")?
            .modified()
            .context("Failed to get dbn mtime")?
            .duration_since(SystemTime::UNIX_EPOCH)
            .context("Failed to convert dbn mtime")?
            .as_secs();
        hasher.update(&mtime.to_le_bytes());
    }

    // 3. Feature extraction version (hash of extraction.rs source)
    let extraction_src_path = PathBuf::from("ml/src/features/extraction.rs");
    if extraction_src_path.exists() {
        let extraction_src = std::fs::read_to_string(&extraction_src_path)
            .context("Failed to read extraction.rs")?;
        hasher.update(extraction_src.as_bytes());
    } else {
        warn!(
            "extraction.rs not found at {:?}, cache key may be incomplete",
            extraction_src_path
        );
    }

    // 4. Warmup period
    hasher.update(&warmup_period.to_le_bytes());

    Ok(format!("{:x}", hasher.finalize()))
}

/// Load features from cache
///
/// Reads cached feature vectors from Parquet file if cache hit.
/// Returns None if cache miss or validation fails.
///
/// # Arguments
///
/// * `cache_dir` - Directory containing cache files
/// * `cache_key` - SHA256 cache key from calculate_cache_key()
///
/// # Returns
///
/// * `Ok(Some(features))` - Cache hit, features loaded (Vec<[f64; 51]>)
/// * `Ok(None)` - Cache miss or validation failed
/// * `Err(e)` - I/O or parsing error
pub async fn load_features_from_cache(
    cache_dir: &Path,
    cache_key: &str,
) -> Result<Option<Vec<[f64; 51]>>> {
    use arrow::array::Float64Array;
    use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
    use std::fs::File;

    // Check metadata exists
    let metadata_path = cache_dir.join("cache_metadata.json");
    if !metadata_path.exists() {
        info!("⚠️  Cache miss: No metadata file found");
        return Ok(None);
    }

    // Load and verify metadata
    let file = File::open(&metadata_path).context("Failed to open cache metadata")?;
    let metadata_map: HashMap<String, CacheMetadata> =
        serde_json::from_reader(file).context("Failed to parse cache metadata")?;

    let metadata = match metadata_map.get(cache_key) {
        Some(m) => m,
        None => {
            warn!("⚠️  Cache miss: No cached features for cache key {}", cache_key);
            return Ok(None);
        }
    };

    // Check cache file exists
    let cache_path = cache_dir.join(format!("features_{}.parquet", cache_key));
    if !cache_path.exists() {
        warn!("⚠️  Cache file missing: {:?}", cache_path);
        return Ok(None);
    }

    info!("🔍 Cache hit! Loading features from {:?}", cache_path);

    // Load features from Parquet
    let file = File::open(&cache_path).context("Failed to open cache file")?;
    let builder = ParquetRecordBatchReaderBuilder::try_new(file)
        .context("Failed to create Parquet reader")?;
    let reader = builder.build().context("Failed to build Parquet reader")?;

    let mut all_features = Vec::new();

    for batch_result in reader {
        let batch = batch_result.context("Failed to read Parquet batch")?;
        let num_rows = batch.num_rows();

        for row_idx in 0..num_rows {
            let mut feature_vec = [0.0; 51];

            for col_idx in 0..51 {
                let column = batch.column(col_idx);
                let array = column
                    .as_any()
                    .downcast_ref::<Float64Array>()
                    .context("Invalid column type in cache")?;
                feature_vec[col_idx] = array.value(row_idx);
            }

            all_features.push(feature_vec);
        }
    }

    info!(
        "✅ Loaded {} feature vectors from cache (estimated savings: ~20.0s)",
        all_features.len()
    );
    info!("   Cache metadata: {} bars × {} features", metadata.bar_count, metadata.feature_count);

    Ok(Some(all_features))
}

/// Save features to cache
///
/// Writes computed feature vectors to Parquet cache file and updates metadata.
///
/// # Arguments
///
/// * `cache_dir` - Directory to store cache files
/// * `cache_key` - SHA256 cache key
/// * `features` - Feature vectors to cache ([f64; 51])
/// * `metadata` - Cache metadata
///
/// # Returns
///
/// * `Ok(())` - Cache saved successfully
/// * `Err(e)` - I/O or serialization error
pub async fn save_features_to_cache(
    cache_dir: &Path,
    cache_key: &str,
    features: &[[f64; 51]],
    metadata: CacheMetadata,
) -> Result<()> {
    use arrow::array::Float64Array;
    use arrow::datatypes::{DataType, Field, Schema};
    use arrow::record_batch::RecordBatch;
    use parquet::arrow::ArrowWriter;
    use parquet::file::properties::WriterProperties;
    use std::fs::File;
    use std::sync::Arc;

    // Create cache directory if it doesn't exist
    std::fs::create_dir_all(cache_dir).context("Failed to create cache directory")?;

    // Define Parquet schema (51 float64 columns)
    let mut fields = Vec::with_capacity(51);
    for i in 0..51 {
        fields.push(Field::new(format!("feature_{}", i), DataType::Float64, false));
    }
    let schema = Arc::new(Schema::new(fields));

    // Create Parquet writer with Snappy compression
    let cache_path = cache_dir.join(format!("features_{}.parquet", cache_key));
    let file = File::create(&cache_path).context("Failed to create cache file")?;
    let props = WriterProperties::builder()
        .set_compression(parquet::basic::Compression::SNAPPY)
        .build();
    let mut writer = ArrowWriter::try_new(file, schema.clone(), Some(props))
        .context("Failed to create Parquet writer")?;

    // Convert features to Arrow arrays
    let mut columns: Vec<Arc<dyn arrow::array::Array>> = Vec::with_capacity(51);
    for col_idx in 0..51 {
        let col_data: Vec<f64> = features.iter().map(|row| row[col_idx]).collect();
        columns.push(Arc::new(Float64Array::from(col_data)));
    }

    // Write batch to Parquet
    let batch = RecordBatch::try_new(schema, columns).context("Failed to create record batch")?;
    writer.write(&batch).context("Failed to write Parquet batch")?;
    writer.close().context("Failed to close Parquet writer")?;

    info!("💾 Saved {} features to cache: {:?}", features.len(), cache_path);

    // Update metadata index
    let metadata_path = cache_dir.join("cache_metadata.json");
    let mut metadata_map: HashMap<String, CacheMetadata> = if metadata_path.exists() {
        let file = File::open(&metadata_path).context("Failed to open metadata for update")?;
        serde_json::from_reader(file).unwrap_or_default()
    } else {
        HashMap::new()
    };

    metadata_map.insert(cache_key.to_string(), metadata);

    let file = File::create(&metadata_path).context("Failed to create metadata file")?;
    serde_json::to_writer_pretty(file, &metadata_map)
        .context("Failed to write metadata")?;

    info!("✅ Cache metadata updated: {:?}", metadata_path);

    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_cache_key_deterministic() {
        // Cache key should be deterministic for same inputs
        let parquet = PathBuf::from("test_data/ES_FUT_180d.parquet");
        let mbp10 = PathBuf::from("test_data/mbp10");

        if parquet.exists() && mbp10.exists() {
            let key1 = calculate_cache_key(&parquet, &mbp10, 50).unwrap();
            let key2 = calculate_cache_key(&parquet, &mbp10, 50).unwrap();
            assert_eq!(key1, key2, "Cache keys should be deterministic");
            assert_eq!(key1.len(), 64, "SHA256 hash should be 64 hex chars");
        }
    }

    #[test]
    fn test_cache_key_varies_with_inputs() {
        let parquet = PathBuf::from("test_data/ES_FUT_180d.parquet");
        let mbp10 = PathBuf::from("test_data/mbp10");

        if parquet.exists() && mbp10.exists() {
            let key1 = calculate_cache_key(&parquet, &mbp10, 50).unwrap();
            let key2 = calculate_cache_key(&parquet, &mbp10, 100).unwrap();
            assert_ne!(key1, key2, "Cache keys should vary with warmup period");
        }
    }
}