foxhunt/ml/tests/liquid_ensemble_risk_tests.rs

//! Comprehensive Tests for Liquid Neural Networks, Ensemble Methods, and Risk Models
//!
//! This test suite covers:
//! - Liquid Time-constant (LTC) cell dynamics validation
//! - ODE solver convergence and accuracy (Euler vs RK4)
//! - Ensemble voting strategies (majority, weighted, confidence-based)
//! - Kelly criterion edge cases and boundary behavior
//! - VaR calculation methods (historical vs parametric)

#![allow(unused_crate_dependencies)]

use ml::ensemble::voting::{EnsembleVoter, VotingConfig, VotingStrategy};
use ml::liquid::activation::ActivationType;
use ml::liquid::cells::{CfCCell, CfCConfig, LTCCell, LTCConfig};
use ml::liquid::ode_solvers::{
    EulerSolver, ODESolver, RK4Solver, SolverType, VolatilityAwareTimeConstants,
};
use ml::liquid::{FixedPoint, LiquidError, Result as LiquidResult, PRECISION};
use ml::risk::kelly_optimizer::{KellyCriterionOptimizer, KellyOptimizerConfig};
use ml::risk::var_models::{MarketTick, NeuralVarConfig, NeuralVarModel, VarFeatures};
use ml::MLError;

use chrono::Utc;
use common::types::{Price, Quantity, Symbol};

// ============================================================================
// LIQUID CELL TESTS
// ============================================================================

#[test]
fn test_ltc_cell_time_constant_dynamics() -> LiquidResult<()> {
    // Test that time constants affect adaptation speed
    let config_fast = LTCConfig {
        input_size: 2,
        hidden_size: 3,
        tau_min: FixedPoint(PRECISION / 100), // 0.01 (fast)
        tau_max: FixedPoint(PRECISION / 10),  // 0.1
        use_bias: true,
        solver_type: SolverType::Euler,
        activation: ActivationType::Tanh,
    };

    let config_slow = LTCConfig {
        tau_min: FixedPoint(PRECISION),     // 1.0 (slow)
        tau_max: FixedPoint(10 * PRECISION), // 10.0
        ..config_fast.clone()
    };

    let mut cell_fast = LTCCell::new(config_fast)?;
    let mut cell_slow = LTCCell::new(config_slow)?;

    let input = vec![
        FixedPoint(PRECISION / 2), // 0.5
        FixedPoint(PRECISION / 4), // 0.25
    ];
    let dt = FixedPoint(PRECISION / 100); // 0.01

    // Run multiple steps
    for _ in 0..10 {
        cell_fast.forward(&input, dt)?;
        cell_slow.forward(&input, dt)?;
    }

    // Fast cell should have adapted more (higher magnitude changes)
    let fast_state = &cell_fast.hidden_state;
    let slow_state = &cell_slow.hidden_state;

    // At least one neuron should show different adaptation
    let mut found_difference = false;
    for i in 0..3 {
        if (fast_state[i].0 - slow_state[i].0).abs() > PRECISION / 100 {
            found_difference = true;
            break;
        }
    }
    assert!(
        found_difference,
        "Time constants should affect adaptation speed"
    );

    Ok(())
}

#[test]
fn test_ltc_cell_volatility_adaptation() -> LiquidResult<()> {
    let config = LTCConfig {
        input_size: 2,
        hidden_size: 2,
        tau_min: FixedPoint(PRECISION / 100),
        tau_max: FixedPoint(PRECISION),
        use_bias: true,
        solver_type: SolverType::Euler,
        activation: ActivationType::Sigmoid,
    };

    let mut cell = LTCCell::new(config.clone())?;

    // Initial time constants
    let initial_taus = cell.get_time_constants();

    // Apply low volatility - should increase time constants (slower adaptation)
    let low_volatility = FixedPoint(PRECISION / 10); // 0.1
    cell.update_market_volatility(low_volatility)?;
    let low_vol_taus = cell.get_time_constants();

    // Apply high volatility - should decrease time constants (faster adaptation)
    let high_volatility = FixedPoint(5 * PRECISION); // 5.0
    cell.update_market_volatility(high_volatility)?;
    let high_vol_taus = cell.get_time_constants();

    // Time constants should be clamped within bounds
    for tau in &high_vol_taus {
        assert!(tau.0 >= config.tau_min.0, "Tau below minimum");
        assert!(tau.0 <= config.tau_max.0, "Tau above maximum");
    }

    // High volatility should generally produce lower time constants
    let avg_low_vol = low_vol_taus.iter().map(|t| t.0).sum::<i64>() / low_vol_taus.len() as i64;
    let avg_high_vol =
        high_vol_taus.iter().map(|t| t.0).sum::<i64>() / high_vol_taus.len() as i64;
    assert!(
        avg_high_vol <= avg_low_vol,
        "High volatility should decrease time constants"
    );

    Ok(())
}

#[test]
fn test_cfc_cell_backbone_network() -> LiquidResult<()> {
    // Test CfC with multi-layer backbone
    let config = CfCConfig {
        input_size: 4,
        hidden_size: 6,
        backbone_layers: vec![8, 6, 4], // 3-layer backbone
        mixed_memory: true,
        use_gate: true,
        solver_type: SolverType::RK4,
    };

    let mut cell = CfCCell::new(config)?;

    // Verify backbone structure
    assert_eq!(
        cell.backbone_weights.len(),
        3,
        "Should have 3 backbone layers"
    );
    assert_eq!(
        cell.backbone_weights[0].len(),
        8,
        "First layer should have 8 neurons"
    );
    assert_eq!(
        cell.backbone_weights[1].len(),
        6,
        "Second layer should have 6 neurons"
    );
    assert_eq!(
        cell.backbone_weights[2].len(),
        4,
        "Third layer should have 4 neurons"
    );

    // Test forward pass
    let input = vec![
        FixedPoint(PRECISION / 4),
        FixedPoint(PRECISION / 3),
        FixedPoint(PRECISION / 2),
        FixedPoint(PRECISION / 5),
    ];
    let dt = FixedPoint(PRECISION / 100);

    let output = cell.forward(&input, dt)?;
    assert_eq!(output.len(), 6, "Output should match hidden_size");

    // All outputs should be finite
    for &out in &output {
        assert!(out.is_finite(), "Output should be finite");
    }

    Ok(())
}

// ============================================================================
// ODE SOLVER TESTS
// ============================================================================

#[test]
fn test_euler_vs_rk4_convergence() -> LiquidResult<()> {
    let euler = EulerSolver;
    let rk4 = RK4Solver;

    // Test exponential decay: dx/dt = -x, analytical solution: x(t) = x0 * e^(-t)
    let decay = |x: FixedPoint, _t: FixedPoint| -> FixedPoint { FixedPoint(-x.0) };

    let x0 = FixedPoint(PRECISION); // 1.0
    let t0 = FixedPoint(0);
    let dt = FixedPoint(PRECISION / 100); // 0.01

    // Simulate 100 steps (t = 1.0)
    let mut x_euler = x0;
    let mut x_rk4 = x0;
    let mut t = t0;

    for _ in 0..100 {
        x_euler = euler.step(&decay, x_euler, t, dt)?;
        x_rk4 = rk4.step(&decay, x_rk4, t, dt)?;
        t = (t + dt)?;
    }

    // Analytical solution at t=1: e^(-1) ≈ 0.36788
    let analytical = FixedPoint((0.36788 * PRECISION as f64) as i64);

    // RK4 should be more accurate than Euler
    let euler_error = (x_euler.0 - analytical.0).abs();
    let rk4_error = (x_rk4.0 - analytical.0).abs();

    assert!(
        rk4_error < euler_error,
        "RK4 should be more accurate than Euler (RK4 error: {}, Euler error: {})",
        rk4_error,
        euler_error
    );

    // RK4 error should be very small (within 0.1%)
    assert!(
        rk4_error < PRECISION / 1000,
        "RK4 error should be < 0.1%"
    );

    Ok(())
}

#[test]
fn test_ode_solver_stability() -> LiquidResult<()> {
    let euler = EulerSolver;

    // Test stiff equation: dx/dt = -100*x (requires very small timestep for Euler)
    let stiff = |x: FixedPoint, _t: FixedPoint| -> FixedPoint { FixedPoint(-100 * x.0) };

    let x0 = FixedPoint(PRECISION);
    let t0 = FixedPoint(0);

    // Large timestep should cause instability for Euler on stiff problems
    let dt_large = FixedPoint(PRECISION / 10); // 0.1
    let result = euler.step(&stiff, x0, t0, dt_large);

    // Should either work or indicate overflow (both acceptable for stiff problems)
    match result {
        Ok(x) => {
            // If it works, value should still be reasonable
            assert!(x.0.abs() < 1000 * PRECISION, "Value exploded");
        }
        Err(LiquidError::Overflow(_)) => {
            // Expected for stiff problems with large timesteps
        }
        Err(e) => panic!("Unexpected error: {:?}", e),
    }

    Ok(())
}

#[test]
fn test_volatility_aware_time_constants_clamping() -> LiquidResult<()> {
    let base_tau = FixedPoint(PRECISION / 10); // 0.1
    let min_tau = FixedPoint(PRECISION / 100); // 0.01
    let max_tau = FixedPoint(PRECISION); // 1.0

    let mut vol_aware = VolatilityAwareTimeConstants::new(base_tau, min_tau, max_tau);

    // Test extreme volatility values
    let extreme_high = FixedPoint(100 * PRECISION); // 100.0 (should be capped)
    vol_aware.update_volatility(extreme_high)?;
    let tau_high = vol_aware.current_tau();

    assert!(
        tau_high.0 >= min_tau.0,
        "Should respect minimum tau: {} >= {}",
        tau_high.0,
        min_tau.0
    );
    assert!(
        tau_high.0 <= max_tau.0,
        "Should respect maximum tau: {} <= {}",
        tau_high.0,
        max_tau.0
    );

    // Zero volatility
    let zero_vol = FixedPoint(0);
    vol_aware.update_volatility(zero_vol)?;
    let tau_zero = vol_aware.current_tau();

    assert!(
        tau_zero.0 >= min_tau.0 && tau_zero.0 <= max_tau.0,
        "Zero volatility should produce valid tau"
    );

    Ok(())
}

// ============================================================================
// ENSEMBLE VOTING TESTS
// ============================================================================

#[test]
fn test_ensemble_weighted_voting() -> Result<(), MLError> {
    let config = VotingConfig {
        strategy: VotingStrategy::WeightedAverage,
        dynamic_strategy: false,
        outlier_threshold: 2.0,
        minimum_confidence: 0.1,
    };

    let mut voter = EnsembleVoter::new(config);

    // Test with empty weights (should still work)
    let signals = vec![];
    let weights = std::collections::HashMap::new();

    let result = voter.aggregate_signals(&signals, &weights)?;

    // Should return valid result even with no signals (fallback behavior)
    assert!(result.signal >= 0.0 && result.signal <= 1.0);
    assert!(result.confidence >= 0.0 && result.confidence <= 1.0);

    Ok(())
}

#[test]
fn test_ensemble_voting_strategies() -> Result<(), MLError> {
    let strategies = vec![
        VotingStrategy::WeightedAverage,
        VotingStrategy::ConfidenceWeighted,
        VotingStrategy::Adaptive,
        VotingStrategy::Robust,
        VotingStrategy::MajorityVote,
    ];

    for strategy in strategies {
        let config = VotingConfig {
            strategy: strategy.clone(),
            dynamic_strategy: false,
            outlier_threshold: 1.5,
            minimum_confidence: 0.5,
        };

        let mut voter = EnsembleVoter::new(config);
        let signals = vec![];
        let weights = std::collections::HashMap::new();

        let result = voter.aggregate_signals(&signals, &weights)?;

        // Each strategy should produce valid output
        assert!(
            result.signal >= 0.0,
            "Strategy {:?} produced invalid signal",
            strategy
        );
        assert!(
            result.confidence >= 0.0 && result.confidence <= 1.0,
            "Strategy {:?} produced invalid confidence",
            strategy
        );
    }

    Ok(())
}

#[test]
fn test_voting_config_validation() {
    // Test default configuration
    let config = VotingConfig::default();
    assert!(config.outlier_threshold > 0.0);
    assert!(config.minimum_confidence >= 0.0 && config.minimum_confidence <= 1.0);
    assert_eq!(config.strategy, VotingStrategy::WeightedAverage);
}

// ============================================================================
// KELLY CRITERION TESTS
// ============================================================================

#[test]
fn test_kelly_zero_edge() -> Result<(), MLError> {
    let config = KellyOptimizerConfig::default();
    let optimizer = KellyCriterionOptimizer::new(config)?;

    // Zero edge: 50% win probability, equal win/loss amounts
    let kelly = optimizer.calculate_basic_kelly(0.5, 1.0, 1.0)?;

    // Zero edge should produce zero Kelly fraction (clamped to min_fraction)
    assert!(
        kelly <= 0.02,
        "Zero edge should produce near-zero Kelly fraction, got {}",
        kelly
    );

    Ok(())
}

#[test]
fn test_kelly_negative_edge() -> Result<(), MLError> {
    let config = KellyOptimizerConfig::default();
    let optimizer = KellyCriterionOptimizer::new(config.clone())?;

    // Negative edge: 40% win probability, 1:1 payout
    let kelly = optimizer.calculate_basic_kelly(0.4, 1.0, 1.0)?;

    // Negative edge should produce minimum Kelly fraction
    assert_eq!(
        kelly, config.min_fraction,
        "Negative edge should produce minimum Kelly fraction"
    );

    Ok(())
}

#[test]
fn test_kelly_100_percent_confidence() -> Result<(), MLError> {
    let config = KellyOptimizerConfig::default();
    let optimizer = KellyCriterionOptimizer::new(config.clone())?;

    // Cannot test exactly 1.0 (validation error), test near-certainty
    let kelly = optimizer.calculate_basic_kelly(0.99, 2.0, 1.0)?;

    // Near-certain win with 2:1 payout should produce high Kelly (clamped to max)
    assert_eq!(
        kelly, config.max_fraction,
        "Near-certain win should produce maximum Kelly fraction"
    );

    Ok(())
}

#[test]
fn test_kelly_boundary_validation() -> Result<(), MLError> {
    let config = KellyOptimizerConfig::default();
    let optimizer = KellyCriterionOptimizer::new(config)?;

    // Test invalid probabilities
    assert!(
        optimizer.calculate_basic_kelly(0.0, 2.0, 1.0).is_err(),
        "Should reject 0% probability"
    );
    assert!(
        optimizer.calculate_basic_kelly(1.0, 2.0, 1.0).is_err(),
        "Should reject 100% probability"
    );
    assert!(
        optimizer.calculate_basic_kelly(1.1, 2.0, 1.0).is_err(),
        "Should reject >100% probability"
    );

    // Test invalid win/loss amounts
    assert!(
        optimizer.calculate_basic_kelly(0.6, -1.0, 1.0).is_err(),
        "Should reject negative win"
    );
    assert!(
        optimizer.calculate_basic_kelly(0.6, 1.0, -1.0).is_err(),
        "Should reject negative loss"
    );
    assert!(
        optimizer.calculate_basic_kelly(0.6, 0.0, 1.0).is_err(),
        "Should reject zero win"
    );

    Ok(())
}

#[test]
fn test_kelly_enhanced_with_volatility() -> Result<(), MLError> {
    let config = KellyOptimizerConfig {
        volatility_adjustment: true,
        ..Default::default()
    };
    let optimizer = KellyCriterionOptimizer::new(config)?;

    // Test enhanced Kelly combines standard and basic approaches
    let kelly = optimizer.calculate_enhanced_kelly(
        0.1,  // 10% expected return
        0.04, // 4% variance (20% volatility)
        0.6,  // 60% win probability
        0.15, // 15% average win
        0.1,  // 10% average loss
    )?;

    assert!(kelly > 0.0, "Should produce positive Kelly fraction");
    assert!(kelly <= 0.25, "Should respect maximum fraction");

    // Compare with volatility adjustment disabled
    let config_no_vol = KellyOptimizerConfig {
        volatility_adjustment: false,
        ..Default::default()
    };
    let optimizer_no_vol = KellyCriterionOptimizer::new(config_no_vol)?;

    let kelly_no_vol = optimizer_no_vol.calculate_enhanced_kelly(
        0.1, 0.04, 0.6, 0.15, 0.1,
    )?;

    // Results should differ when volatility adjustment is enabled/disabled
    // (though in practice they might be the same due to weighting)
    assert!(
        kelly >= 0.0 && kelly_no_vol >= 0.0,
        "Both should produce valid results"
    );

    Ok(())
}

#[test]
fn test_kelly_fractional_sizing() -> Result<(), MLError> {
    let config = KellyOptimizerConfig::default();
    let optimizer = KellyCriterionOptimizer::new(config.clone())?;

    let full_kelly = 0.2;

    // Test fractional Kelly
    let half_kelly = optimizer.calculate_fractional_kelly(full_kelly, 0.5);
    assert_eq!(half_kelly, 0.1, "Half Kelly should be exactly half");

    let quarter_kelly = optimizer.calculate_fractional_kelly(full_kelly, 0.25);
    assert_eq!(
        quarter_kelly, 0.05,
        "Quarter Kelly should be exactly 1/4"
    );

    // Test that excessive values are clamped
    let excessive = optimizer.calculate_fractional_kelly(1.0, 1.0);
    assert_eq!(
        excessive, config.max_fraction,
        "Excessive Kelly should be clamped to max"
    );

    Ok(())
}

#[test]
fn test_kelly_position_recommendation() -> Result<(), MLError> {
    let config = KellyOptimizerConfig::default();
    let optimizer = KellyCriterionOptimizer::new(config.clone())?;

    // Test with winning history
    let winning_returns = vec![0.1, -0.02, 0.08, 0.05, 0.12, -0.01, 0.06];
    let rec = optimizer.recommend_position("AAPL".to_string(), &winning_returns)?;

    assert!(rec.recommended_fraction > 0.0);
    assert!(rec.recommended_fraction <= config.max_fraction);
    assert!(rec.win_probability > 0.5, "Should detect winning edge");
    assert!(rec.expected_return > 0.0, "Should have positive expected return");

    // Test with losing history
    let losing_returns = vec![-0.05, -0.03, 0.02, -0.08, -0.01, -0.04];
    let rec_loss = optimizer.recommend_position("LOSS".to_string(), &losing_returns)?;

    assert!(rec_loss.win_probability < 0.5, "Should detect losing edge");
    assert!(
        rec_loss.recommended_fraction == config.min_fraction,
        "Should recommend minimum fraction for losing edge"
    );

    // Test with empty history
    let empty_returns: Vec<f64> = vec![];
    assert!(
        optimizer
            .recommend_position("EMPTY".to_string(), &empty_returns)
            .is_err(),
        "Should reject empty history"
    );

    Ok(())
}

// ============================================================================
// VAR MODEL TESTS
// ============================================================================

#[test]
fn test_var_model_creation() -> Result<(), MLError> {
    let config = NeuralVarConfig::default();
    let model = NeuralVarModel::new(config)?;

    assert!(
        !model.config.confidence_levels.is_empty(),
        "Should have confidence levels"
    );
    assert!(
        model.config.lookback_period > 0,
        "Should have positive lookback"
    );

    Ok(())
}

#[test]
fn test_var_confidence_levels() -> Result<(), MLError> {
    let config = NeuralVarConfig {
        confidence_levels: vec![0.90, 0.95, 0.99, 0.999],
        ..Default::default()
    };
    let model = NeuralVarModel::new(config)?;

    assert_eq!(
        model.config.confidence_levels.len(),
        4,
        "Should have 4 confidence levels"
    );

    // Confidence levels should be in ascending order
    let levels = &model.config.confidence_levels;
    for i in 1..levels.len() {
        assert!(
            levels[i] > levels[i - 1],
            "Confidence levels should be ascending"
        );
    }

    Ok(())
}

#[test]
fn test_var_features_extraction() -> Result<(), MLError> {
    let symbol = Symbol::from("AAPL");
    let mut market_data = Vec::new();

    // Create sample market data with increasing prices
    for i in 0..20 {
        market_data.push(MarketTick {
            symbol: symbol.clone(),
            price: Price::from_f64(100.0 + i as f64).unwrap(),
            quantity: Quantity::from_f64(1000.0 + i as f64 * 10.0).unwrap(),
            timestamp: Utc::now(),
        });
    }

    let features = VarFeatures::from_market_data(&market_data, 252)?;

    // Check returns calculation
    assert_eq!(
        features.returns.len(),
        19,
        "Should have n-1 returns for n prices"
    );

    // Returns should be positive (increasing prices)
    for &ret in &features.returns {
        assert!(ret > 0.0, "Returns should be positive for increasing prices");
    }

    // Volatility should be positive
    assert!(features.volatility > 0.0, "Volatility should be positive");

    // Volume should be positive
    assert!(features.volume > 0.0, "Volume should be positive");

    Ok(())
}

#[test]
fn test_var_features_with_volatile_data() -> Result<(), MLError> {
    let symbol = Symbol::from("VOL");
    let mut market_data = Vec::new();

    // Create volatile market data
    let prices = vec![
        100.0, 105.0, 98.0, 110.0, 95.0, 108.0, 92.0, 115.0, 90.0, 120.0,
    ];

    for (i, price) in prices.into_iter().enumerate() {
        market_data.push(MarketTick {
            symbol: symbol.clone(),
            price: Price::from_f64(price).unwrap(),
            quantity: Quantity::from_f64(1000.0).unwrap(),
            timestamp: Utc::now(),
        });
    }

    let features = VarFeatures::from_market_data(&market_data, 252)?;

    // Volatility should be high for volatile data
    assert!(
        features.volatility > 0.05,
        "Volatile data should have high volatility: {}",
        features.volatility
    );

    // Returns should alternate sign
    let mut has_positive = false;
    let mut has_negative = false;
    for &ret in &features.returns {
        if ret > 0.0 {
            has_positive = true;
        }
        if ret < 0.0 {
            has_negative = true;
        }
    }
    assert!(
        has_positive && has_negative,
        "Volatile data should have both positive and negative returns"
    );

    Ok(())
}

#[test]
fn test_var_feature_vector_conversion() -> Result<(), MLError> {
    let symbol = Symbol::from("TEST");
    let mut market_data = Vec::new();

    for i in 0..10 {
        market_data.push(MarketTick {
            symbol: symbol.clone(),
            price: Price::from_f64(100.0 + i as f64).unwrap(),
            quantity: Quantity::from_f64(1000.0).unwrap(),
            timestamp: Utc::now(),
        });
    }

    let features = VarFeatures::from_market_data(&market_data, 252)?;
    let feature_vector = features.to_feature_vector();

    // Should produce fixed-size vector
    assert_eq!(
        feature_vector.len(),
        100,
        "Feature vector should have 100 elements"
    );

    // First elements should be volatility and volume
    assert_eq!(
        feature_vector[0], features.volatility,
        "First element should be volatility"
    );
    assert_eq!(
        feature_vector[1], features.volume,
        "Second element should be volume"
    );

    Ok(())
}

#[test]
fn test_var_empty_data_handling() {
    let empty_data: Vec<MarketTick> = vec![];

    let result = VarFeatures::from_market_data(&empty_data, 252);
    assert!(
        result.is_err(),
        "Should reject empty market data"
    );
}

// ============================================================================
// INTEGRATION TESTS
// ============================================================================

#[test]
fn test_liquid_cell_parameter_count() -> LiquidResult<()> {
    let ltc_config = LTCConfig {
        input_size: 4,
        hidden_size: 8,
        tau_min: FixedPoint(PRECISION / 100),
        tau_max: FixedPoint(PRECISION),
        use_bias: true,
        solver_type: SolverType::Euler,
        activation: ActivationType::Tanh,
    };

    let ltc_cell = LTCCell::new(ltc_config)?;
    let ltc_params = ltc_cell.parameter_count();

    // Expected: (4*8 input weights) + (8*8 recurrent weights) + (8 bias) + (8 tau) = 32 + 64 + 8 + 8 = 112
    assert_eq!(
        ltc_params, 112,
        "LTC parameter count should be 112"
    );

    let cfc_config = CfCConfig {
        input_size: 4,
        hidden_size: 6,
        backbone_layers: vec![8, 4],
        mixed_memory: true,
        use_gate: true,
        solver_type: SolverType::RK4,
    };

    let cfc_cell = CfCCell::new(cfc_config)?;
    let cfc_params = cfc_cell.parameter_count();

    // Backbone: (10*8 + 8 bias) + (8*4 + 4 bias) = 88 + 36 = 124
    // Final: 6 + 6 + 6 = 18
    // Total: 142
    assert_eq!(
        cfc_params, 142,
        "CfC parameter count should be 142"
    );

    Ok(())
}

#[test]
fn test_state_reset() -> LiquidResult<()> {
    let config = LTCConfig {
        input_size: 2,
        hidden_size: 3,
        tau_min: FixedPoint(PRECISION / 100),
        tau_max: FixedPoint(PRECISION / 10),
        use_bias: true,
        solver_type: SolverType::Euler,
        activation: ActivationType::Sigmoid,
    };

    let mut cell = LTCCell::new(config)?;

    // Run forward pass to change state
    let input = vec![FixedPoint(PRECISION / 2), FixedPoint(PRECISION / 4)];
    let dt = FixedPoint(PRECISION / 100);
    cell.forward(&input, dt)?;

    // State should be non-zero
    assert!(
        cell.hidden_state.iter().any(|&s| s.0 != 0),
        "State should be non-zero after forward pass"
    );

    // Reset state
    cell.reset_state();

    // State should be all zeros
    assert!(
        cell.hidden_state.iter().all(|&s| s.0 == 0),
        "State should be zero after reset"
    );

    Ok(())
}

#[test]
fn test_inference_count_tracking() -> LiquidResult<()> {
    let config = CfCConfig {
        input_size: 3,
        hidden_size: 4,
        backbone_layers: vec![6],
        mixed_memory: false,
        use_gate: false,
        solver_type: SolverType::Euler,
    };

    let mut cell = CfCCell::new(config)?;

    assert_eq!(cell.inference_count, 0, "Initial inference count should be 0");

    // Run multiple forward passes
    let input = vec![
        FixedPoint(PRECISION / 3),
        FixedPoint(PRECISION / 2),
        FixedPoint(PRECISION / 4),
    ];
    let dt = FixedPoint(PRECISION / 100);

    for i in 1..=5 {
        cell.forward(&input, dt)?;
        assert_eq!(
            cell.inference_count, i,
            "Inference count should be {}",
            i
        );
    }

    Ok(())
}