//! Liquid Neural Network Implementation
//!
//! Complete liquid network with multiple layers of LTC/CfC cells,
//! `MLModel` trait integration, and HFT-optimized inference.

use std::collections::HashMap;
use std::time::{Instant, SystemTime, UNIX_EPOCH};

use serde::{Deserialize, Serialize};

use super::activation::{apply_activation, ActivationType};
use super::cells::{CfCCell, CfCConfig, LTCCell, LTCConfig};
use super::{
    FixedPoint, LiquidError, MarketRegime, NetworkType, PerformanceMetrics, Result, PRECISION,
};
use ml_core::{MLError, MLResult};

/// Layer configuration for liquid networks
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum LayerConfig {
    LTC(LTCConfig),
    CfC(CfCConfig),
}

/// Output layer configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct OutputLayerConfig {
    pub use_linear_output: bool,
    pub output_activation: Option<ActivationType>,
    pub dropout_rate: Option<f32>,
}

/// Complete liquid neural network configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LiquidNetworkConfig {
    pub network_type: NetworkType,
    pub input_size: usize,
    pub output_size: usize,
    pub layer_configs: Vec<LayerConfig>,
    pub output_layer: OutputLayerConfig,
    pub default_dt: FixedPoint,
    pub market_regime_adaptation: bool,
}

/// Liquid network layer enum
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum LiquidLayer {
    LTC(LTCCell),
    CfC(CfCCell),
}

impl LiquidLayer {
    pub fn forward(&mut self, input: &[FixedPoint], dt: FixedPoint) -> Result<Vec<FixedPoint>> {
        match self {
            LiquidLayer::LTC(cell) => cell.forward(input, dt),
            LiquidLayer::CfC(cell) => cell.forward(input, dt),
        }
    }

    pub fn reset_state(&mut self) {
        match self {
            LiquidLayer::LTC(cell) => cell.reset_state(),
            LiquidLayer::CfC(cell) => cell.reset_state(),
        }
    }

    pub fn parameter_count(&self) -> usize {
        match self {
            LiquidLayer::LTC(cell) => cell.parameter_count(),
            LiquidLayer::CfC(cell) => cell.parameter_count(),
        }
    }

    pub fn update_market_volatility(&mut self, volatility: FixedPoint) -> Result<()> {
        match self {
            LiquidLayer::LTC(cell) => cell.update_market_volatility(volatility),
            LiquidLayer::CfC(_cell) => {
                // CfC cells don't have explicit volatility adaptation yet
                // Could be implemented with backbone network modulation
                Ok(())
            },
        }
    }
}

/// Main liquid neural network
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LiquidNetwork {
    pub config: LiquidNetworkConfig,
    pub layers: Vec<LiquidLayer>,
    pub output_weights: Vec<Vec<FixedPoint>>,
    pub output_bias: Vec<FixedPoint>,
    pub current_dt: FixedPoint,
    pub current_regime: MarketRegime,
    pub volatility_window: Vec<FixedPoint>,
    pub performance_metrics: PerformanceMetrics,
}

impl LiquidNetwork {
    pub fn new(config: LiquidNetworkConfig) -> Result<Self> {
        if config.input_size == 0 || config.output_size == 0 {
            return Err(LiquidError::InvalidConfiguration(
                "Input and output sizes must be positive".to_owned(),
            ));
        }

        if config.layer_configs.is_empty() {
            return Err(LiquidError::InvalidConfiguration(
                "At least one layer configuration required".to_owned(),
            ));
        }

        // Build layers
        let mut layers = Vec::new();
        let mut current_input_size = config.input_size;

        for layer_config in &config.layer_configs {
            let layer = match layer_config {
                LayerConfig::LTC(ltc_config) => {
                    if ltc_config.input_size != current_input_size {
                        return Err(LiquidError::InvalidConfiguration(format!(
                            "Layer input size {} doesn't match previous layer output size {}",
                            ltc_config.input_size, current_input_size
                        )));
                    }
                    let cell = LTCCell::new(ltc_config.clone())?;
                    current_input_size = ltc_config.hidden_size;
                    LiquidLayer::LTC(cell)
                },
                LayerConfig::CfC(cfc_config) => {
                    if cfc_config.input_size != current_input_size {
                        return Err(LiquidError::InvalidConfiguration(format!(
                            "Layer input size {} doesn't match previous layer output size {}",
                            cfc_config.input_size, current_input_size
                        )));
                    }
                    let cell = CfCCell::new(cfc_config.clone())?;
                    current_input_size = cfc_config.hidden_size;
                    LiquidLayer::CfC(cell)
                },
            };
            layers.push(layer);
        }

        // Initialize output layer weights
        let output_weights: Vec<Vec<FixedPoint>> = (0..config.output_size)
            .map(|i| {
                (0..current_input_size)
                    .map(|j| {
                        let value = ((i * 29 + j * 43) % 1000) as f64 / 1000.0 - 0.5;
                        FixedPoint::from_f64(value / (current_input_size as f64).sqrt())
                    })
                    .collect()
            })
            .collect();

        let output_bias: Vec<FixedPoint> = (0..config.output_size)
            .map(|i| FixedPoint::from_f64((i % 10) as f64 / 100.0))
            .collect();

        let performance_metrics = PerformanceMetrics {
            total_inferences: 0,
            average_inference_time_ns: 0,
            average_inference_time_us: 0.0,
            total_parameters: layers.iter().map(|l| l.parameter_count()).sum::<usize>()
                + output_weights.iter().map(|row| row.len()).sum::<usize>()
                + output_bias.len(),
            current_regime: MarketRegime::Normal,
            regime_switches: 0,
            last_adaptation_time: None,
        };

        Ok(Self {
            current_dt: config.default_dt,
            config,
            layers,
            output_weights,
            output_bias,
            current_regime: MarketRegime::Normal,
            volatility_window: Vec::with_capacity(100),
            performance_metrics,
        })
    }

    /// Forward pass through the entire network
    pub fn forward(&mut self, input: &[FixedPoint]) -> Result<Vec<FixedPoint>> {
        let start_time = Instant::now();

        if input.len() != self.config.input_size {
            return Err(LiquidError::InvalidInput(format!(
                "Expected input size {}, got {}",
                self.config.input_size,
                input.len()
            )));
        }

        // Forward through liquid layers
        let mut current_activations = input.to_vec();
        for layer in &mut self.layers {
            current_activations = layer.forward(&current_activations, self.current_dt)?;
        }

        // Output layer computation
        let mut outputs = Vec::with_capacity(self.config.output_size);
        for i in 0..self.config.output_size {
            let mut sum = self.output_bias[i];

            for (j, activation) in current_activations.iter().enumerate() {
                if j < self.output_weights[i].len() {
                    let contrib = (self.output_weights[i][j] * *activation)?;
                    sum = (sum + contrib)?;
                }
            }

            // Apply output activation if specified
            let output = if let Some(activation_type) = self.config.output_layer.output_activation {
                apply_activation(sum, activation_type)?
            } else if self.config.output_layer.use_linear_output {
                sum
            } else {
                apply_activation(sum, ActivationType::Linear)?
            };

            outputs.push(output);
        }

        // Update performance metrics
        let elapsed = start_time.elapsed();
        self.performance_metrics.total_inferences += 1;
        let total_time_ns = self.performance_metrics.average_inference_time_ns
            * (self.performance_metrics.total_inferences - 1)
            + elapsed.as_nanos() as u64;
        self.performance_metrics.average_inference_time_ns =
            total_time_ns / self.performance_metrics.total_inferences;
        self.performance_metrics.average_inference_time_us =
            self.performance_metrics.average_inference_time_ns as f64 / 1000.0;

        Ok(outputs)
    }

    /// Predict method for compatibility with ML trait
    pub fn predict(&mut self, input: &[f64]) -> MLResult<Vec<f64>> {
        let fixed_input: Vec<FixedPoint> = input.iter().map(|&x| FixedPoint::from_f64(x)).collect();

        let fixed_output = self.forward(&fixed_input).map_err(MLError::from)?;

        let output: Vec<f64> = fixed_output.iter().map(|fp| fp.to_f64()).collect();

        Ok(output)
    }

    /// Update market volatility and adapt network behavior
    pub fn update_market_volatility(&mut self, volatility: FixedPoint) -> Result<()> {
        if !self.config.market_regime_adaptation {
            return Ok(());
        }

        // Add to volatility window (rolling window of size 100)
        if self.volatility_window.len() >= 100 {
            self.volatility_window.remove(0);
        }
        self.volatility_window.push(volatility);

        // Determine market regime based on volatility
        let avg_volatility = if !self.volatility_window.is_empty() {
            let sum: i64 = self.volatility_window.iter().map(|v| v.0).sum();
            FixedPoint(sum / self.volatility_window.len() as i64)
        } else {
            volatility
        };

        let new_regime = if avg_volatility.0 < PRECISION / 5 {
            // < 0.2 - Low volatility maps to Normal
            MarketRegime::Normal
        } else if avg_volatility.0 < PRECISION {
            // < 1.0 - Medium volatility maps to Sideways
            MarketRegime::Sideways
        } else if avg_volatility.0 < 3 * PRECISION {
            // < 3.0 - High volatility with direction maps to Trending
            MarketRegime::Trending
        } else if avg_volatility.0 < 5 * PRECISION {
            // < 5.0 - Very high volatility maps to Bull/Bear (using Bull as default)
            MarketRegime::Bull
        } else {
            // >= 5.0 - Extreme volatility maps to Crisis
            MarketRegime::Crisis
        };

        if new_regime != self.current_regime {
            self.current_regime = new_regime;
            self.performance_metrics.current_regime = new_regime;
            self.performance_metrics.regime_switches += 1;
            self.performance_metrics.last_adaptation_time = Some(
                SystemTime::now()
                    .duration_since(UNIX_EPOCH)
                    .map_err(|e| LiquidError::InferenceError(format!("System time error: {}", e)))?
                    .as_millis() as u64,
            );

            // Adapt time step based on regime
            self.current_dt = match new_regime {
                MarketRegime::Normal => self.config.default_dt,
                MarketRegime::Sideways => FixedPoint(self.config.default_dt.0 / 2),
                MarketRegime::Trending => FixedPoint(self.config.default_dt.0 * 2),
                MarketRegime::Bull => FixedPoint(self.config.default_dt.0 / 4),
                MarketRegime::Bear => FixedPoint(self.config.default_dt.0 / 4),
                MarketRegime::Crisis => FixedPoint(self.config.default_dt.0 / 8),
            };
        }

        // Update all layers with volatility information
        for layer in &mut self.layers {
            layer.update_market_volatility(volatility)?;
        }

        Ok(())
    }

    /// Reset all network states
    pub fn reset_states(&mut self) {
        for layer in &mut self.layers {
            layer.reset_state();
        }
    }

    /// Get network performance metrics
    pub const fn get_performance_metrics(&self) -> &PerformanceMetrics {
        &self.performance_metrics
    }

    /// Get metrics as `HashMap` for compatibility
    pub fn get_metrics(&self) -> HashMap<String, f64> {
        let mut metrics = HashMap::new();
        metrics.insert(
            "total_inferences".to_owned(),
            self.performance_metrics.total_inferences as f64,
        );
        metrics.insert(
            "avg_inference_time_us".to_owned(),
            self.performance_metrics.average_inference_time_us,
        );
        metrics.insert(
            "total_parameters".to_owned(),
            self.performance_metrics.total_parameters as f64,
        );
        metrics.insert(
            "regime_switches".to_owned(),
            self.performance_metrics.regime_switches as f64,
        );
        metrics
    }

    /// Get input size
    pub const fn input_size(&self) -> usize {
        self.config.input_size
    }

    /// Get output size
    pub const fn output_size(&self) -> usize {
        self.config.output_size
    }

    /// Get total number of parameters
    pub const fn parameter_count(&self) -> usize {
        self.performance_metrics.total_parameters
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::liquid::activation::ActivationType;
    use crate::liquid::ode_solvers::SolverType;
    use crate::liquid::PRECISION;
    use common::trading::MarketRegime;


    #[test]
    fn test_liquid_network_creation() -> Result<()> {
        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 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: None,
                dropout_rate: None,
            },
            default_dt: FixedPoint(PRECISION / 100),
            market_regime_adaptation: true,
        };

        let network = LiquidNetwork::new(network_config)?;
        assert_eq!(network.layers.len(), 1);
        assert_eq!(network.output_weights.len(), 2);
        Ok(())
    }

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

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

        let mut network = LiquidNetwork::new(network_config)?;

        let input = vec![
            FixedPoint(PRECISION / 2), // 0.5
            FixedPoint(PRECISION / 4), // 0.25
            FixedPoint(PRECISION / 3), // 0.33
        ];

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

        assert_eq!(output.len(), 1);
        assert!(output[0].is_finite()); // Check for finite values

        // Test multiple forward passes
        for _ in 0..5 {
            let output = network.forward(&input)?;
            assert_eq!(output.len(), 1);
            assert!(output[0].is_finite());
        }
        Ok(())
    }

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

        let network_config = LiquidNetworkConfig {
            network_type: NetworkType::LTC,
            input_size: 2,
            output_size: 1,
            layer_configs: vec![LayerConfig::LTC(ltc_config)],
            output_layer: OutputLayerConfig {
                use_linear_output: true,
                output_activation: None,
                dropout_rate: None,
            },
            default_dt: FixedPoint(PRECISION / 100),
            market_regime_adaptation: true,
        };

        let mut network = LiquidNetwork::new(network_config)?;

        // Update with high volatility
        let high_volatility = FixedPoint(5 * PRECISION); // 5.0
        network.update_market_volatility(high_volatility)?;

        // Check that regime was updated
        let metrics = network.get_performance_metrics();
        assert_ne!(metrics.current_regime, MarketRegime::Normal);
        Ok(())
    }

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

        let network_config = LiquidNetworkConfig {
            network_type: NetworkType::LTC,
            input_size: 2,
            output_size: 1,
            layer_configs: vec![LayerConfig::LTC(ltc_config)],
            output_layer: OutputLayerConfig {
                use_linear_output: true,
                output_activation: None,
                dropout_rate: None,
            },
            default_dt: FixedPoint(PRECISION / 100),
            market_regime_adaptation: false,
        };

        let mut network = LiquidNetwork::new(network_config)?;

        let input = vec![FixedPoint(PRECISION / 2), FixedPoint(PRECISION / 3)];

        // Run multiple inferences
        for _ in 0..10 {
            let _output = network.forward(&input)?;
        }

        let metrics = network.get_performance_metrics();
        assert_eq!(metrics.total_inferences, 10);
        assert!(metrics.average_inference_time_ns > 0);
        assert!(metrics.average_inference_time_us >= 0.0);
        assert!(metrics.total_parameters > 0);
        Ok(())
    }

    #[test]
    fn test_predict_compatibility() -> std::result::Result<(), Box<dyn std::error::Error>> {
        let ltc_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::Tanh,
        };

        let network_config = LiquidNetworkConfig {
            network_type: NetworkType::LTC,
            input_size: 2,
            output_size: 1,
            layer_configs: vec![LayerConfig::LTC(ltc_config)],
            output_layer: OutputLayerConfig {
                use_linear_output: true,
                output_activation: None,
                dropout_rate: None,
            },
            default_dt: FixedPoint(PRECISION / 100),
            market_regime_adaptation: false,
        };

        let mut network = LiquidNetwork::new(network_config)?;

        let input = vec![0.5, 0.25];
        let output = network.predict(&input)?;

        assert_eq!(output.len(), 1);
        assert!(output[0].is_finite());
        Ok(())
    }
}
