//! Liquid Neural Network Cell Implementations
//!
//! Implements LTC (Liquid Time-constant) and `CfC` (Closed-form Continuous-time)
//! cells with fixed-point arithmetic for ultra-low latency inference.

use std::time::{SystemTime, UNIX_EPOCH};

use super::activation::{self, apply_activation, ActivationType};
use super::ode_solvers::{SolverEnum, SolverFactory, SolverType, VolatilityAwareTimeConstants};
use super::{FixedPoint, LiquidError, Result};
use serde::{Deserialize, Serialize};

/// Configuration for LTC (Liquid Time-constant) cells
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LTCConfig {
    pub input_size: usize,
    pub hidden_size: usize,
    pub tau_min: FixedPoint,
    pub tau_max: FixedPoint,
    pub use_bias: bool,
    pub solver_type: SolverType,
    pub activation: ActivationType,
}

/// Configuration for `CfC` (Closed-form Continuous-time) cells
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CfCConfig {
    pub input_size: usize,
    pub hidden_size: usize,
    pub backbone_layers: Vec<usize>,
    pub mixed_memory: bool,
    pub use_gate: bool,
    pub solver_type: SolverType,
}

/// LTC (Liquid Time-constant) cell implementation
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LTCCell {
    pub config: LTCConfig,
    pub hidden_state: Vec<FixedPoint>,
    pub input_weights: Vec<Vec<FixedPoint>>,
    pub recurrent_weights: Vec<Vec<FixedPoint>>,
    pub bias: Vec<FixedPoint>,
    pub time_constants: Vec<VolatilityAwareTimeConstants>,
    #[serde(skip)]
    solver: Option<SolverEnum>,
    pub inference_count: u64,
    pub last_inference_time: Option<u64>, // Store as timestamp millis instead of Instant
}

impl LTCCell {
    pub fn new(config: LTCConfig) -> Result<Self> {
        if config.input_size == 0 || config.hidden_size == 0 {
            return Err(LiquidError::InvalidConfiguration(
                "Input size and hidden size must be positive".to_owned(),
            ));
        }

        // Initialize weights with small random values (Xavier initialization approximation)
        let weight_scale = FixedPoint::from_f64(1.0 / (config.input_size as f64).sqrt());
        let mut input_weights = Vec::with_capacity(config.hidden_size);
        let mut recurrent_weights = Vec::with_capacity(config.hidden_size);

        for i in 0..config.hidden_size {
            let mut input_row = Vec::with_capacity(config.input_size);
            let mut recurrent_row = Vec::with_capacity(config.hidden_size);

            for j in 0..config.input_size {
                // Simple deterministic initialization based on indices
                let value = ((i * 37 + j * 17) % 1000) as f64 / 1000.0 - 0.5;
                input_row.push(FixedPoint::from_f64(value * weight_scale.to_f64()));
            }

            for j in 0..config.hidden_size {
                let value = ((i * 41 + j * 19) % 1000) as f64 / 1000.0 - 0.5;
                recurrent_row.push(FixedPoint::from_f64(value * weight_scale.to_f64()));
            }

            input_weights.push(input_row);
            recurrent_weights.push(recurrent_row);
        }

        // Initialize bias
        let bias = if config.use_bias {
            (0..config.hidden_size)
                .map(|i| FixedPoint::from_f64((i % 10) as f64 / 100.0)) // Small bias values
                .collect()
        } else {
            vec![FixedPoint::zero(); config.hidden_size]
        };

        // Initialize time constants with volatility awareness
        let base_tau = FixedPoint((config.tau_min.0 + config.tau_max.0) / 2);
        let time_constants = (0..config.hidden_size)
            .map(|_| VolatilityAwareTimeConstants::new(base_tau, config.tau_min, config.tau_max))
            .collect();

        let hidden_state = vec![FixedPoint::zero(); config.hidden_size];
        let solver = Some(SolverFactory::create_solver(config.solver_type));

        Ok(Self {
            config,
            hidden_state,
            input_weights,
            recurrent_weights,
            bias,
            time_constants,
            solver,
            inference_count: 0,
            last_inference_time: None,
        })
    }

    /// Forward pass through the LTC cell
    pub fn forward(&mut self, input: &[FixedPoint], dt: FixedPoint) -> Result<Vec<FixedPoint>> {
        let start_time = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .map_err(|e| LiquidError::InferenceError(format!("System time error: {}", e)))?
            .as_millis() as u64;

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

        let solver = if let Some(s) = &self.solver { s } else {
            // Re-create solver if needed (after deserialization)
            self.solver = Some(SolverFactory::create_solver(self.config.solver_type));
            self.solver.as_ref().ok_or_else(|| {
                LiquidError::InferenceError("Solver not initialized".to_owned())
            })?
        };

        let mut new_hidden_state = Vec::with_capacity(self.config.hidden_size);

        for i in 0..self.config.hidden_size {
            // Compute input contribution
            let mut input_sum = FixedPoint::zero();
            for (j, &inp) in input.iter().enumerate().take(self.config.input_size) {
                let contrib = (self.input_weights[i][j] * inp)?;
                input_sum = (input_sum + contrib)?;
            }

            // Compute recurrent contribution
            let mut recurrent_sum = FixedPoint::zero();
            for j in 0..self.config.hidden_size {
                let contrib = (self.recurrent_weights[i][j] * self.hidden_state[j])?;
                recurrent_sum = (recurrent_sum + contrib)?;
            }

            // Total input
            let total_input = (input_sum + recurrent_sum)? + self.bias[i];

            // Apply activation
            let activated = apply_activation(total_input?, self.config.activation)?;

            // LTC dynamics: dx/dt = (1/tau) * (-x + activated)
            let current_tau = self.time_constants[i].current_tau();
            let current_x = self.hidden_state[i];
            let neg_x = FixedPoint(-current_x.0);
            let diff = (neg_x + activated)?;
            let dx_dt = (diff / current_tau)?;

            // Simple ODE function for this neuron
            let ode_fn = |_x: FixedPoint, _t: FixedPoint| -> FixedPoint { dx_dt };

            // Integrate using ODE solver
            let new_x = solver.step(&ode_fn, current_x, FixedPoint::zero(), dt)?;
            new_hidden_state.push(new_x);
        }

        self.hidden_state = new_hidden_state;
        self.inference_count += 1;
        self.last_inference_time = Some(start_time);

        Ok(self.hidden_state.clone())
    }

    /// Update market volatility for adaptive time constants
    pub fn update_market_volatility(&mut self, volatility: FixedPoint) -> Result<()> {
        for tau in &mut self.time_constants {
            tau.update_volatility(volatility)?;
        }
        Ok(())
    }

    /// Get current time constants
    pub fn get_time_constants(&self) -> Vec<FixedPoint> {
        self.time_constants
            .iter()
            .map(|tau| tau.current_tau())
            .collect()
    }

    /// Reset hidden state
    pub fn reset_state(&mut self) {
        self.hidden_state.fill(FixedPoint::zero());
    }

    /// Get number of parameters
    pub fn parameter_count(&self) -> usize {
        let input_params = self.input_weights.len() * self.input_weights[0].len();
        let recurrent_params = self.recurrent_weights.len() * self.recurrent_weights[0].len();
        let bias_params = if self.config.use_bias {
            self.bias.len()
        } else {
            0
        };
        let tau_params = self.time_constants.len();

        input_params + recurrent_params + bias_params + tau_params
    }
}

/// `CfC` (Closed-form Continuous-time) cell implementation
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CfCCell {
    pub config: CfCConfig,
    pub output_state: Vec<FixedPoint>,
    pub backbone_weights: Vec<Vec<Vec<FixedPoint>>>, // Layer -> Neuron -> Input
    pub backbone_bias: Vec<Vec<FixedPoint>>,
    pub input_weights: Vec<FixedPoint>,
    pub recurrent_weights: Vec<FixedPoint>,
    pub output_weights: Vec<FixedPoint>,
    #[serde(skip)]
    solver: Option<SolverEnum>,
    pub inference_count: u64,
    pub last_inference_time: Option<u64>, // Store as timestamp millis instead of Instant
}

impl CfCCell {
    pub fn new(config: CfCConfig) -> Result<Self> {
        if config.input_size == 0 || config.hidden_size == 0 {
            return Err(LiquidError::InvalidConfiguration(
                "Input size and hidden size must be positive".to_owned(),
            ));
        }

        // Initialize backbone network weights
        let mut backbone_weights = Vec::new();
        let mut backbone_bias = Vec::new();
        let mut prev_size = config.input_size + config.hidden_size; // Input + recurrent

        for &layer_size in &config.backbone_layers {
            let mut layer_weights = Vec::with_capacity(layer_size);
            let mut layer_bias = Vec::with_capacity(layer_size);

            for i in 0..layer_size {
                let mut neuron_weights = Vec::with_capacity(prev_size);
                for j in 0..prev_size {
                    let value = ((i * 23 + j * 31) % 1000) as f64 / 1000.0 - 0.5;
                    neuron_weights.push(FixedPoint::from_f64(value / (prev_size as f64).sqrt()));
                }
                layer_weights.push(neuron_weights);

                let bias_value = (i % 10) as f64 / 100.0;
                layer_bias.push(FixedPoint::from_f64(bias_value));
            }

            backbone_weights.push(layer_weights);
            backbone_bias.push(layer_bias);
            prev_size = layer_size;
        }

        // Initialize final layer weights
        let _backbone_output_size = config
            .backbone_layers
            .last()
            .copied()
            .unwrap_or(config.input_size + config.hidden_size);

        let input_weights: Vec<FixedPoint> = (0..config.hidden_size)
            .map(|i| FixedPoint::from_f64(((i * 13) % 1000) as f64 / 1000.0 - 0.5))
            .collect();

        let recurrent_weights: Vec<FixedPoint> = (0..config.hidden_size)
            .map(|i| FixedPoint::from_f64(((i * 17 + 100) % 1000) as f64 / 1000.0 - 0.5))
            .collect();

        let output_weights: Vec<FixedPoint> = (0..config.hidden_size)
            .map(|i| FixedPoint::from_f64(((i * 19 + 200) % 1000) as f64 / 1000.0 - 0.5))
            .collect();

        let output_state = vec![FixedPoint::zero(); config.hidden_size];
        let solver = Some(SolverFactory::create_solver(config.solver_type));

        Ok(Self {
            config,
            output_state,
            backbone_weights,
            backbone_bias,
            input_weights,
            recurrent_weights,
            output_weights,
            solver,
            inference_count: 0,
            last_inference_time: None,
        })
    }

    /// Forward pass through the `CfC` cell
    pub fn forward(&mut self, input: &[FixedPoint], dt: FixedPoint) -> Result<Vec<FixedPoint>> {
        let start_time = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .map_err(|e| LiquidError::InferenceError(format!("System time error: {}", e)))?
            .as_millis() as u64;

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

        let _solver = if let Some(s) = &self.solver { s } else {
            // Re-create solver if needed (after deserialization)
            self.solver = Some(SolverFactory::create_solver(self.config.solver_type));
            self.solver.as_ref().ok_or_else(|| {
                LiquidError::InferenceError("Solver not initialized".to_owned())
            })?
        };

        // Concatenate input and current state for backbone network
        let mut backbone_input = input.to_vec();
        backbone_input.extend_from_slice(&self.output_state);

        // Forward through backbone network
        let mut current_activations = backbone_input;
        for (layer_idx, layer_weights) in self.backbone_weights.iter().enumerate() {
            let mut next_activations = Vec::with_capacity(layer_weights.len());

            for (neuron_idx, neuron_weights) in layer_weights.iter().enumerate() {
                let mut sum = self.backbone_bias[layer_idx][neuron_idx];

                for (weight_idx, weight) in neuron_weights.iter().enumerate() {
                    if weight_idx < current_activations.len() {
                        let contrib = (*weight * current_activations[weight_idx])?;
                        sum = (sum + contrib)?;
                    }
                }

                // Apply tanh activation for backbone
                let activated = activation::tanh(sum)?;
                next_activations.push(activated);
            }

            current_activations = next_activations;
        }

        // Use backbone output to compute dynamics
        let backbone_output = &current_activations;
        let mut new_state = Vec::with_capacity(self.config.hidden_size);

        for i in 0..self.config.hidden_size {
            let current_x = self.output_state[i];

            // Compute input contribution (simplified)
            let input_contrib = if i < input.len() {
                (self.input_weights[i] * input[i])?
            } else {
                FixedPoint::zero()
            };

            // Compute recurrent contribution
            let recurrent_contrib = (self.recurrent_weights[i] * current_x)?;

            // Use backbone output to modulate dynamics
            let backbone_modulation = if i < backbone_output.len() {
                backbone_output[i]
            } else {
                FixedPoint::zero()
            };

            // CfC closed-form approximation: integrate directly
            let total_input = (input_contrib + recurrent_contrib)? + backbone_modulation;
            let activated = activation::tanh(total_input?)?;

            // Simple integration step (Euler approximation)
            let dx_dt = (activated - current_x)?;
            let step = (dt * dx_dt)?;
            let new_x = (current_x + step)?;

            new_state.push(new_x);
        }

        self.output_state = new_state;
        self.inference_count += 1;
        self.last_inference_time = Some(start_time);

        Ok(self.output_state.clone())
    }

    /// Reset output state
    pub fn reset_state(&mut self) {
        self.output_state.fill(FixedPoint::zero());
    }

    /// Get number of parameters
    pub fn parameter_count(&self) -> usize {
        let backbone_params: usize = self
            .backbone_weights
            .iter()
            .map(|layer| layer.iter().map(|neuron| neuron.len()).sum::<usize>())
            .sum::<usize>();
        let backbone_bias_params: usize = self.backbone_bias.iter().map(|layer| layer.len()).sum();
        let final_params =
            self.input_weights.len() + self.recurrent_weights.len() + self.output_weights.len();

        backbone_params + backbone_bias_params + final_params
    }
}

#[cfg(test)]
#[allow(clippy::needless_range_loop)]
mod tests {
    use super::*;
    use crate::liquid::activation::ActivationType;
    use crate::liquid::ode_solvers::SolverType;
    use crate::liquid::PRECISION;


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

        let cell = LTCCell::new(config)?;
        assert_eq!(cell.hidden_state.len(), 8);
        assert_eq!(cell.input_weights.len(), 8);
        assert_eq!(cell.input_weights[0].len(), 4);
        Ok(())
    }

    #[test]
    fn test_ltc_forward_pass() -> Result<()> {
        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::Tanh,
        };

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

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

        let dt = FixedPoint(PRECISION / 100); // 0.01

        let output = cell.forward(&input, dt)?;

        assert_eq!(output.len(), 3);
        // Check that outputs are finite
        for &out in &output {
            assert!(out.is_finite());
        }
        Ok(())
    }

    #[test]
    fn test_cfc_cell_creation() -> Result<()> {
        let config = CfCConfig {
            input_size: 4,
            hidden_size: 6,
            backbone_layers: vec![8, 8],
            mixed_memory: true,
            use_gate: false,
            solver_type: SolverType::RK4,
        };

        let cell = CfCCell::new(config)?;
        assert_eq!(cell.output_state.len(), 6);
        assert_eq!(cell.backbone_weights.len(), 2); // Two backbone layers
        Ok(())
    }

    #[test]
    fn test_cfc_forward_pass() -> Result<()> {
        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)?;

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

        let dt = FixedPoint(PRECISION / 100); // 0.01

        let output = cell.forward(&input, dt)?;

        assert_eq!(output.len(), 4);
        // Check that outputs are finite
        for &out in &output {
            assert!(out.is_finite());
        }
        Ok(())
    }

    #[test]
    fn test_volatility_adaptation() -> Result<()> {
        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())?;

        let initial_taus = cell.get_time_constants();

        // Apply high volatility
        let high_volatility = FixedPoint(3 * PRECISION); // 3.0
        cell.update_market_volatility(high_volatility)?;

        let adapted_taus = cell.get_time_constants();

        // Time constants should generally decrease with high volatility
        // (though they're clamped to valid ranges)
        for i in 0..initial_taus.len() {
            assert!(adapted_taus[i].0 >= config.tau_min.0);
            assert!(adapted_taus[i].0 <= config.tau_max.0);
        }
        Ok(())
    }
}
