// Quick verification that DbnSequenceLoader produces 256-dimensional features

use ml::data_loaders::DbnSequenceLoader;

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    println!("🔍 Verifying DbnSequenceLoader feature dimensions...\n");

    // Create loader with 256 feature dimensions
    let mut loader = DbnSequenceLoader::new(60, 256).await?;
    println!("✅ Loader created: seq_len=60, d_model=256\n");

    // Load sequences from test data
    let data_dir = "test_data/real/databento/ml_training_small";
    println!("📂 Loading sequences from: {}", data_dir);

    let (train_data, val_data) = loader.load_sequences(data_dir, 0.9).await?;

    println!("\n📊 Results:");
    println!("   Training sequences: {}", train_data.len());
    println!("   Validation sequences: {}", val_data.len());

    // Check first sequence dimensions
    if let Some((input, target)) = train_data.first() {
        let input_shape = input.shape();
        let target_shape = target.shape();

        println!("\n🔢 Tensor Shapes:");
        println!("   Input:  {:?} (expected: [1, 60, 256])", input_shape.dims());
        println!("   Target: {:?} (expected: [1, 1, 256])", target_shape.dims());

        // Verify dimensions
        assert_eq!(input_shape.dims(), &[1, 60, 256], "Input shape mismatch!");
        assert_eq!(target_shape.dims(), &[1, 1, 256], "Target shape mismatch!");

        println!("\n✅ SUCCESS: All feature dimensions are correct!");
        println!("   - Extract features produces exactly 256 dimensions");
        println!("   - No zero-padding needed");
        println!("   - Ready for MAMBA-2 training");
    } else {
        println!("\n❌ ERROR: No training sequences found!");
        return Err(anyhow::anyhow!("No training data"));
    }

    Ok(())
}