foxhunt/ml/tests/liquid_nn_training_tests.rs

//! Liquid NN Training Pipeline TDD Test Suite
//!
//! Comprehensive E2E tests for Liquid Neural Network training with CPU-only
//! fixed-point arithmetic. Tests cover forward/backward passes, training loop
//! convergence, checkpoint persistence, inference determinism, and memory usage.
//!
//! Architecture:
//! - CPU-ONLY (fixed-point arithmetic for <100μs inference)
//! - No CUDA dependencies (by design, not a limitation)
//! - Fixed-point precision: 8 decimal places (PRECISION = 100_000_000)
//! - Training: CPU-based gradient descent with MSE loss
//! - Inference: Deterministic fixed-point computation
//!
//! Test Coverage:
//! 1. Forward pass: Fixed-point computation correctness
//! 2. Backward pass: Gradient calculation (CPU only)
//! 3. Training loop: Loss convergence over epochs
//! 4. Checkpoint save/load: Safetensors persistence
//! 5. Inference determinism: Same input → same output
//! 6. Memory usage: CPU memory within limits

#![allow(unused_crate_dependencies)]

use ml::liquid::{
    ActivationType, FixedPoint, LTCConfig, LayerConfig, LiquidNetwork, LiquidNetworkConfig,
    LiquidTrainer, LiquidTrainingConfig, NetworkType, OutputLayerConfig, SolverType, TrainingBatch,
    TrainingSample, TrainingUtils, PRECISION,
};
use std::time::Instant;

// ============================================================================
// Test 1: Forward Pass - Fixed-Point Computation
// ============================================================================

#[test]
fn test_liquid_nn_forward_pass() -> anyhow::Result<()> {
    println!("\n=== Test 1: Forward Pass - Fixed-Point Computation ===");

    // Create minimal Liquid NN (16 input → 8 hidden → 3 output)
    let ltc_config = LTCConfig {
        input_size: 16,
        hidden_size: 8,
        tau_min: FixedPoint::from_f64(0.1),
        tau_max: FixedPoint::from_f64(1.0),
        use_bias: true,
        solver_type: SolverType::Euler, // Simplest solver for testing
        activation: ActivationType::Tanh,
    };

    let network_config = LiquidNetworkConfig {
        network_type: NetworkType::LTC,
        input_size: 16,
        output_size: 3,
        layer_configs: vec![LayerConfig::LTC(ltc_config)],
        output_layer: OutputLayerConfig {
            use_linear_output: true,
            output_activation: Some(ActivationType::Linear),
            dropout_rate: None,
        },
        default_dt: FixedPoint::from_f64(0.01),
        market_regime_adaptation: false,
    };

    let mut network = LiquidNetwork::new(network_config)?;
    println!("✓ Created network: 16 inputs → 8 hidden (LTC) → 3 outputs");
    println!("  Parameters: {}", network.parameter_count());

    // Create input with fixed-point values
    let input: Vec<FixedPoint> = (0..16)
        .map(|i| FixedPoint::from_f64(0.5 + (i as f64) * 0.01))
        .collect();

    println!("\n  Input features (first 5): {:?}", &input[0..5]);

    // Forward pass
    let start = Instant::now();
    let output = network.forward(&input)?;
    let duration = start.elapsed();

    println!("  Forward pass time: {:?}", duration);
    println!("  Output shape: {} values", output.len());
    println!("  Output values: {:?}", output);

    // Assertions
    assert_eq!(output.len(), 3, "Output should have 3 values");
    assert!(
        duration.as_micros() < 1000,
        "Forward pass should be <1ms (target: <100μs in production)"
    );

    // Verify fixed-point arithmetic correctness
    for &val in &output {
        assert!(
            val.is_finite(),
            "Output values should be finite (no overflow)"
        );
    }

    println!("✓ Forward pass completed successfully");
    Ok(())
}

// ============================================================================
// Test 2: Backward Pass - Gradient Computation (CPU Only)
// ============================================================================

#[test]
fn test_liquid_nn_backward_pass() -> anyhow::Result<()> {
    println!("\n=== Test 2: Backward Pass - Gradient Computation ===");

    // Create network
    let ltc_config = LTCConfig {
        input_size: 4,
        hidden_size: 4,
        tau_min: FixedPoint::from_f64(0.1),
        tau_max: FixedPoint::from_f64(1.0),
        use_bias: true,
        solver_type: SolverType::Euler,
        activation: ActivationType::Tanh,
    };

    let network_config = LiquidNetworkConfig {
        network_type: NetworkType::LTC,
        input_size: 4,
        output_size: 2,
        layer_configs: vec![LayerConfig::LTC(ltc_config)],
        output_layer: OutputLayerConfig {
            use_linear_output: true,
            output_activation: Some(ActivationType::Linear),
            dropout_rate: None,
        },
        default_dt: FixedPoint::from_f64(0.01),
        market_regime_adaptation: false,
    };

    let mut network = LiquidNetwork::new(network_config.clone())?;
    let trainer_config = LiquidTrainingConfig::default();
    let mut trainer = LiquidTrainer::new(trainer_config);

    println!("✓ Created network: 4 → 4 (LTC) → 2");

    // Create training sample
    let input = vec![
        FixedPoint::from_f64(0.5),
        FixedPoint::from_f64(0.3),
        FixedPoint::from_f64(0.7),
        FixedPoint::from_f64(0.2),
    ];
    let target = vec![FixedPoint::one(), FixedPoint::zero()];

    let sample = TrainingSample {
        input: input.clone(),
        target: target.clone(),
        timestamp: None,
        market_regime: None,
        volatility: None,
    };

    println!("\n  Input:  {:?}", input);
    println!("  Target: {:?}", target);

    // Forward pass to get predictions
    let predictions = network.forward(&input)?;
    println!("  Predictions (before training): {:?}", predictions);

    // Calculate loss manually (MSE)
    let loss_before: f64 = predictions
        .iter()
        .zip(target.iter())
        .map(|(pred, tgt)| {
            let diff = pred.to_f64() - tgt.to_f64();
            diff * diff
        })
        .sum::<f64>()
        / predictions.len() as f64;
    println!("  Loss (before training): {:.6}", loss_before);

    // Train to verify gradient computation (use public train method)
    let batches = vec![TrainingBatch::new(vec![sample])];
    trainer.train(&mut network, &batches, None)?;

    let history = trainer.get_training_history();
    let batch_loss = history.last().map(|m| m.training_loss).unwrap_or(0.0);
    println!("\n  Final training loss: {:.6}", batch_loss);
    println!(
        "  Gradient history length: {}",
        trainer.gradient_history.len()
    );

    // Verify gradient was computed
    assert!(
        !trainer.gradient_history.is_empty(),
        "Gradient history should contain gradients after training"
    );

    // Verify gradient is finite
    let last_gradient = trainer.gradient_history.last().unwrap();
    assert!(
        last_gradient.is_finite(),
        "Gradient should be finite (no overflow)"
    );

    println!("✓ Backward pass completed successfully");
    println!("  Last gradient norm: {:.6}", last_gradient.to_f64());

    Ok(())
}

// ============================================================================
// Test 3: Training Loop Convergence - Loss Decreases Over Epochs
// ============================================================================

#[test]
fn test_training_loop_convergence() -> anyhow::Result<()> {
    println!("\n=== Test 3: Training Loop Convergence ===");

    // Create small network for fast convergence test
    let ltc_config = LTCConfig {
        input_size: 3,
        hidden_size: 4,
        tau_min: FixedPoint::from_f64(0.1),
        tau_max: FixedPoint::from_f64(1.0),
        use_bias: true,
        solver_type: SolverType::Euler,
        activation: ActivationType::Tanh,
    };

    let network_config = LiquidNetworkConfig {
        network_type: NetworkType::LTC,
        input_size: 3,
        output_size: 2,
        layer_configs: vec![LayerConfig::LTC(ltc_config)],
        output_layer: OutputLayerConfig {
            use_linear_output: true,
            output_activation: Some(ActivationType::Linear),
            dropout_rate: None,
        },
        default_dt: FixedPoint::from_f64(0.01),
        market_regime_adaptation: false,
    };

    let mut network = LiquidNetwork::new(network_config)?;
    println!("✓ Created network: 3 → 4 (LTC) → 2");

    // Create synthetic training data (simple XOR-like problem)
    let mut training_samples = Vec::new();
    for i in 0..20 {
        let x = (i % 4) as f64;
        let input = vec![
            FixedPoint::from_f64(x / 4.0),
            FixedPoint::from_f64((x * 2.0) / 4.0),
            FixedPoint::from_f64((x * 3.0) / 4.0),
        ];
        let target = if i % 2 == 0 {
            vec![FixedPoint::one(), FixedPoint::zero()]
        } else {
            vec![FixedPoint::zero(), FixedPoint::one()]
        };

        training_samples.push(TrainingSample {
            input,
            target,
            timestamp: None,
            market_regime: None,
            volatility: None,
        });
    }

    println!("  Created {} training samples", training_samples.len());

    // Create batches
    let batches = TrainingUtils::create_batches(training_samples, 4);
    println!("  Created {} batches (batch size: 4)", batches.len());

    // Configure training (10 epochs for convergence test)
    let training_config = LiquidTrainingConfig {
        learning_rate: FixedPoint(PRECISION / 100), // 0.01
        batch_size: 4,
        max_epochs: 10,
        early_stopping_patience: 5,
        gradient_clip_threshold: FixedPoint::one(),
        l2_regularization: FixedPoint::zero(),
        adaptive_learning_rate: false,
        market_regime_adaptation: false,
        validation_split: 0.0,
    };

    let mut trainer = LiquidTrainer::new(training_config);
    println!("\n  Training configuration:");
    println!("    Learning rate: 0.01");
    println!("    Max epochs: 10");
    println!("    Batch size: 4");

    // Train network
    println!("\n  Starting training...");
    let start = Instant::now();
    trainer.train(&mut network, &batches, None)?;
    let training_time = start.elapsed();

    println!("\n  Training completed in {:?}", training_time);

    // Verify loss convergence
    let history = trainer.get_training_history();
    assert!(
        history.len() >= 2,
        "Training history should have at least 2 epochs"
    );

    let first_loss = history[0].training_loss;
    let last_loss = history.last().unwrap().training_loss;

    println!("\n  Loss progression:");
    println!("    Epoch 0:  {:.6}", first_loss);
    for (i, metrics) in history.iter().enumerate().skip(1) {
        println!("    Epoch {}: {:.6}", i, metrics.training_loss);
    }
    println!("    Final:    {:.6}", last_loss);

    // Assert loss decreased (convergence)
    assert!(
        last_loss < first_loss,
        "Loss should decrease during training (first={:.6}, last={:.6})",
        first_loss,
        last_loss
    );

    let loss_reduction = ((first_loss - last_loss) / first_loss) * 100.0;
    println!("\n✓ Training converged successfully");
    println!("  Loss reduction: {:.2}%", loss_reduction);

    Ok(())
}

// ============================================================================
// Test 4: Checkpoint Save/Load - Safetensors Persistence
// ============================================================================

#[test]
fn test_checkpoint_save_load() -> anyhow::Result<()> {
    println!("\n=== Test 4: Checkpoint Save/Load ===");

    // Create network
    let ltc_config = LTCConfig {
        input_size: 5,
        hidden_size: 6,
        tau_min: FixedPoint::from_f64(0.1),
        tau_max: FixedPoint::from_f64(1.0),
        use_bias: true,
        solver_type: SolverType::Euler,
        activation: ActivationType::Tanh,
    };

    let network_config = LiquidNetworkConfig {
        network_type: NetworkType::LTC,
        input_size: 5,
        output_size: 3,
        layer_configs: vec![LayerConfig::LTC(ltc_config)],
        output_layer: OutputLayerConfig {
            use_linear_output: true,
            output_activation: Some(ActivationType::Linear),
            dropout_rate: None,
        },
        default_dt: FixedPoint::from_f64(0.01),
        market_regime_adaptation: false,
    };

    let network = LiquidNetwork::new(network_config.clone())?;
    println!("✓ Created original network: 5 → 6 (LTC) → 3");

    // Create test input
    let input: Vec<FixedPoint> = (0..5)
        .map(|i| FixedPoint::from_f64(0.1 + (i as f64) * 0.1))
        .collect();

    // Save checkpoint BEFORE running forward pass to preserve initial state
    println!("\n  Saving checkpoint...");
    let original_network = network.clone();
    let checkpoint_json = serde_json::to_string(&original_network)?;
    println!("  Checkpoint size: {} bytes", checkpoint_json.len());

    // Run forward pass on original network
    let mut original_network_mut = original_network.clone();
    let original_output = original_network_mut.forward(&input)?;
    println!("  Original predictions: {:?}", original_output);

    // Load checkpoint
    println!("\n  Loading checkpoint...");
    let mut loaded_network: LiquidNetwork = serde_json::from_str(&checkpoint_json)?;
    println!("  ✓ Checkpoint loaded successfully");

    // Verify predictions match
    let loaded_output = loaded_network.forward(&input)?;
    println!("\n  Loaded predictions: {:?}", loaded_output);

    // Compare outputs
    for (i, (&orig, &loaded)) in original_output.iter().zip(loaded_output.iter()).enumerate() {
        let diff = (orig.0 - loaded.0).abs();
        println!(
            "  Output[{}]: orig={:.6}, loaded={:.6}, diff={}",
            i,
            orig.to_f64(),
            loaded.to_f64(),
            diff
        );
        assert_eq!(
            orig, loaded,
            "Output {} should match exactly after checkpoint reload",
            i
        );
    }

    println!("\n✓ Checkpoint save/load verified (deterministic)");
    Ok(())
}

// ============================================================================
// Test 5: Inference Determinism - Same Input → Same Output
// ============================================================================

#[test]
fn test_inference_determinism() -> anyhow::Result<()> {
    println!("\n=== Test 5: Inference Determinism ===");

    // Create network
    let ltc_config = LTCConfig {
        input_size: 8,
        hidden_size: 8,
        tau_min: FixedPoint::from_f64(0.1),
        tau_max: FixedPoint::from_f64(1.0),
        use_bias: true,
        solver_type: SolverType::RK4, // Higher-order solver
        activation: ActivationType::Tanh,
    };

    let network_config = LiquidNetworkConfig {
        network_type: NetworkType::LTC,
        input_size: 8,
        output_size: 4,
        layer_configs: vec![LayerConfig::LTC(ltc_config)],
        output_layer: OutputLayerConfig {
            use_linear_output: true,
            output_activation: Some(ActivationType::Linear),
            dropout_rate: None,
        },
        default_dt: FixedPoint::from_f64(0.01),
        market_regime_adaptation: false,
    };

    let mut network = LiquidNetwork::new(network_config)?;
    println!("✓ Created network: 8 → 8 (LTC, RK4) → 4");

    // Create test input
    let input: Vec<FixedPoint> = vec![
        FixedPoint::from_f64(0.123),
        FixedPoint::from_f64(0.456),
        FixedPoint::from_f64(0.789),
        FixedPoint::from_f64(0.234),
        FixedPoint::from_f64(0.567),
        FixedPoint::from_f64(0.890),
        FixedPoint::from_f64(0.345),
        FixedPoint::from_f64(0.678),
    ];

    println!("\n  Running 10 inference passes with identical input...");

    // Run inference 10 times with same input
    let mut outputs = Vec::new();
    for i in 0..10 {
        // Reset network state before each inference
        network.reset_states();

        let output = network.forward(&input)?;
        outputs.push(output);

        if i == 0 {
            println!("  Run {}: {:?}", i, outputs[i]);
        }
    }

    // Verify all outputs are identical
    let first_output = &outputs[0];
    for (run_idx, output) in outputs.iter().enumerate().skip(1) {
        for (i, (&expected, &actual)) in first_output.iter().zip(output.iter()).enumerate() {
            assert_eq!(
                expected, actual,
                "Run {} output[{}] should match run 0 (deterministic)",
                run_idx, i
            );
        }
    }

    println!("\n✓ Inference is deterministic (10/10 runs identical)");
    println!("  First output: {:?}", first_output);
    Ok(())
}

// ============================================================================
// Test 6: Memory Usage - CPU Memory Within Limits
// ============================================================================

#[test]
fn test_memory_usage() -> anyhow::Result<()> {
    println!("\n=== Test 6: Memory Usage ===");

    // Create realistic-sized network (16 → 128 → 3)
    let ltc_config = LTCConfig {
        input_size: 16,
        hidden_size: 128,
        tau_min: FixedPoint::from_f64(0.1),
        tau_max: FixedPoint::from_f64(1.0),
        use_bias: true,
        solver_type: SolverType::RK4,
        activation: ActivationType::Tanh,
    };

    let network_config = LiquidNetworkConfig {
        network_type: NetworkType::LTC,
        input_size: 16,
        output_size: 3,
        layer_configs: vec![LayerConfig::LTC(ltc_config)],
        output_layer: OutputLayerConfig {
            use_linear_output: true,
            output_activation: Some(ActivationType::Linear),
            dropout_rate: None,
        },
        default_dt: FixedPoint::from_f64(0.01),
        market_regime_adaptation: false,
    };

    let network = LiquidNetwork::new(network_config)?;
    println!("✓ Created network: 16 → 128 (LTC) → 3");

    // Calculate memory footprint
    let param_count = network.parameter_count();
    let bytes_per_param = size_of::<FixedPoint>(); // 8 bytes (i64)
    let total_bytes = param_count * bytes_per_param;
    let kb = total_bytes as f64 / 1024.0;
    let mb = kb / 1024.0;

    println!("\n  Memory Analysis:");
    println!("    Parameters: {}", param_count);
    println!(
        "    Bytes per param: {} (FixedPoint = i64)",
        bytes_per_param
    );
    println!(
        "    Total memory: {} bytes ({:.2} KB / {:.3} MB)",
        total_bytes, kb, mb
    );

    // Parameter breakdown
    let input_weights = 16 * 128; // input_size × hidden_size
    let recurrent_weights = 128 * 128; // hidden_size × hidden_size
    let bias = 128; // hidden_size
    let output_weights = 128 * 3; // hidden_size × output_size
    let output_bias = 3; // output_size

    println!("\n  Parameter Breakdown:");
    println!("    Input weights:      {}", input_weights);
    println!("    Recurrent weights:  {}", recurrent_weights);
    println!("    Hidden bias:        {}", bias);
    println!("    Output weights:     {}", output_weights);
    println!("    Output bias:        {}", output_bias);
    println!(
        "    Total calculated:   {}",
        input_weights + recurrent_weights + bias + output_weights + output_bias
    );

    // Verify memory is reasonable (<10 MB for this size)
    assert!(
        mb < 10.0,
        "Network memory should be <10 MB (actual: {:.3} MB)",
        mb
    );

    // Create training dataset and measure memory
    println!("\n  Testing with 1000 training samples...");
    let mut samples = Vec::new();
    for i in 0..1000 {
        let input: Vec<FixedPoint> = (0..16)
            .map(|j| FixedPoint::from_f64((i * j) as f64 / 1000.0))
            .collect();
        let target = vec![
            FixedPoint::from_f64((i % 3 == 0) as u8 as f64),
            FixedPoint::from_f64((i % 3 == 1) as u8 as f64),
            FixedPoint::from_f64((i % 3 == 2) as u8 as f64),
        ];

        samples.push(TrainingSample {
            input,
            target,
            timestamp: None,
            market_regime: None,
            volatility: None,
        });
    }

    let sample_memory = samples.len() * (16 + 3) * size_of::<FixedPoint>();
    let sample_mb = sample_memory as f64 / 1024.0 / 1024.0;
    println!("  Sample dataset memory: {:.3} MB", sample_mb);

    // Total memory (network + samples)
    let total_mb = mb + sample_mb;
    println!("\n  Total memory usage: {:.3} MB", total_mb);

    // Verify total memory is reasonable (<50 MB)
    assert!(
        total_mb < 50.0,
        "Total memory (network + samples) should be <50 MB (actual: {:.3} MB)",
        total_mb
    );

    println!("\n✓ Memory usage within limits");
    println!("  Network: {:.3} MB", mb);
    println!("  Samples: {:.3} MB", sample_mb);
    println!("  Total:   {:.3} MB (<50 MB limit)", total_mb);

    Ok(())
}

// ============================================================================
// Helper Functions
// ============================================================================

#[allow(dead_code)]
fn print_test_separator() {
    println!("\n{}\n", "=".repeat(60));
}