trustformers-models 0.1.1

use trustformers_core::{
    errors::{Result, TrustformersError},
    layers::{LayerNorm, Linear},
    tensor::{DType, Tensor},
    traits::Layer,
};

use super::{
    config::{BiologicalConfig, NeuronModel, PlasticityType},
    model::BiologicalModelOutput,
};

/// Neuron state for spiking neural networks
#[derive(Debug, Clone)]
pub struct NeuronState {
    /// Membrane potential
    pub v_mem: Tensor,
    /// Recovery variable (for Izhikevich model)
    pub u_recovery: Option<Tensor>,
    /// Adaptation current (for AdExp model)
    pub adaptation: Option<Tensor>,
    /// Refractory time remaining
    pub refractory_time: Tensor,
    /// Spike output
    pub spikes: Tensor,
}

/// Synaptic state for plasticity
#[derive(Debug, Clone)]
pub struct SynapticState {
    /// Synaptic weights
    pub weights: Tensor,
    /// Presynaptic traces
    pub pre_traces: Tensor,
    /// Postsynaptic traces
    pub post_traces: Tensor,
    /// Eligibility traces
    pub eligibility: Tensor,
}

/// Spiking neural network layer
#[derive(Debug)]
pub struct SpikingLayer {
    /// Configuration
    pub config: BiologicalConfig,
    /// Input projection
    pub input_projection: Linear,
    /// Recurrent projection
    pub recurrent_projection: Linear,
    /// Neuron states
    pub neuron_states: Option<NeuronState>,
    /// Synaptic states
    pub synaptic_states: Option<SynapticState>,
    /// Layer normalization
    pub layer_norm: LayerNorm,
}

impl SpikingLayer {
    /// Create a new spiking layer
    pub fn new(config: &BiologicalConfig) -> Result<Self> {
        let input_projection =
            Linear::new(config.d_model, config.neurons_per_layer, config.use_bias);
        let recurrent_projection =
            Linear::new(config.neurons_per_layer, config.neurons_per_layer, false);
        let layer_norm = LayerNorm::new(vec![config.neurons_per_layer], 1e-12)?;

        Ok(Self {
            config: config.clone(),
            input_projection,
            recurrent_projection,
            neuron_states: None,
            synaptic_states: None,
            layer_norm,
        })
    }

    /// Initialize neuron states
    pub fn init_states(&mut self, batch_size: usize) -> Result<()> {
        let neurons_per_layer = self.config.neurons_per_layer;

        // Initialize membrane potential
        let v_mem = Tensor::zeros(&[batch_size, neurons_per_layer])?;

        // Initialize recovery variable for Izhikevich model
        let u_recovery = if matches!(self.config.neuron_model, NeuronModel::Izhikevich) {
            Some(Tensor::zeros(&[batch_size, neurons_per_layer])?)
        } else {
            None
        };

        // Initialize adaptation current for AdExp model
        let adaptation = if matches!(self.config.neuron_model, NeuronModel::AdaptiveExponentialIF) {
            Some(Tensor::zeros(&[batch_size, neurons_per_layer])?)
        } else {
            None
        };

        let refractory_time = Tensor::zeros(&[batch_size, neurons_per_layer])?;
        let spikes = Tensor::zeros(&[batch_size, neurons_per_layer])?;

        self.neuron_states = Some(NeuronState {
            v_mem,
            u_recovery,
            adaptation,
            refractory_time,
            spikes,
        });

        // Initialize synaptic states
        let weights = Tensor::randn(&[neurons_per_layer, neurons_per_layer])?
            .scalar_mul(self.config.initializer_range)?;
        let pre_traces = Tensor::zeros(&[batch_size, neurons_per_layer])?;
        let post_traces = Tensor::zeros(&[batch_size, neurons_per_layer])?;
        let eligibility = Tensor::zeros(&[neurons_per_layer, neurons_per_layer])?;

        self.synaptic_states = Some(SynapticState {
            weights,
            pre_traces,
            post_traces,
            eligibility,
        });

        Ok(())
    }

    /// Forward pass through spiking layer
    pub fn forward(&mut self, input: &Tensor) -> Result<Tensor> {
        let batch_size = input.shape()[0];
        let seq_len = input.shape()[1];

        // Initialize states if not present
        if self.neuron_states.is_none() {
            self.init_states(batch_size)?;
        }

        let mut outputs = Vec::new();

        // Process each time step
        for t in 0..seq_len {
            let input_t = input.slice(1, t, t + 1)?.squeeze(1)?;
            let output_t = self.forward_timestep(&input_t)?;
            outputs.push(output_t);
        }

        // Stack outputs
        let mut output = outputs[0].clone();
        for i in 1..outputs.len() {
            output = Tensor::concat(&[output, outputs[i].clone()], 1)?;
        }

        Ok(output)
    }

    /// Forward pass for a single timestep
    fn forward_timestep(&mut self, input: &Tensor) -> Result<Tensor> {
        // Project input
        let input_current = self.input_projection.forward(input.clone())?;

        // Recurrent current
        let recurrent_current = {
            let neuron_states = self.neuron_states.as_ref().ok_or_else(|| {
                TrustformersError::runtime_error(
                    "Neuron states not initialized in forward_timestep".to_string(),
                )
            })?;
            self.recurrent_projection.forward(neuron_states.spikes.clone())?
        };

        // Total input current
        let total_current = input_current.add(&recurrent_current)?;

        // Update neuron dynamics based on model type (simplified for compilation)
        {
            let neuron_states = self.neuron_states.as_mut().ok_or_else(|| {
                TrustformersError::runtime_error(
                    "Neuron states not initialized in forward_timestep".to_string(),
                )
            })?;
            // Simple leaky integrate and fire dynamics for all models for now
            let dt = 0.001; // timestep
            let tau = 0.02; // membrane time constant

            // Update membrane potential: dv/dt = (I - v) / tau
            let decay = neuron_states.v_mem.mul_scalar(1.0 - dt / tau)?;
            let input_term = total_current.mul_scalar(dt / tau)?;
            neuron_states.v_mem = decay.add(&input_term)?;

            // Check for spikes (threshold = 1.0) - create spike mask manually
            let threshold = Tensor::full(1.0, neuron_states.v_mem.shape())?;
            let spike_mask = neuron_states.v_mem.greater(&threshold)?;
            neuron_states.spikes = spike_mask.clone();

            // Reset spiked neurons
            let reset_mask = spike_mask.mul_scalar(-1.0)?; // Reset to 0
            neuron_states.v_mem = neuron_states.v_mem.add(&reset_mask)?;
        }

        // Update synaptic plasticity (simplified)
        {
            let neuron_states = self.neuron_states.as_mut().ok_or_else(|| {
                TrustformersError::runtime_error(
                    "Neuron states not initialized in forward_timestep".to_string(),
                )
            })?;
            let synaptic_states = self.synaptic_states.as_mut().ok_or_else(|| {
                TrustformersError::runtime_error(
                    "Synaptic states not initialized in forward_timestep".to_string(),
                )
            })?;

            // Simple STDP-like plasticity update
            let learning_rate = 0.001;
            let pre_spike_trace = neuron_states.spikes.mul_scalar(0.1)?;
            let post_spike_trace = neuron_states.spikes.mul_scalar(0.1)?;

            // Update synaptic weights based on spike timing
            let weight_update = pre_spike_trace
                .matmul(&post_spike_trace.transpose(0, 1)?)?
                .mul_scalar(learning_rate)?;
            synaptic_states.weights = synaptic_states.weights.add(&weight_update)?;

            // Decay plasticity traces
            synaptic_states.pre_traces = synaptic_states.pre_traces.mul_scalar(0.99)?;
            synaptic_states.post_traces = synaptic_states.post_traces.mul_scalar(0.99)?;
        }

        // Apply layer normalization to spikes
        let output = {
            let neuron_states = self.neuron_states.as_ref().ok_or_else(|| {
                TrustformersError::runtime_error(
                    "Neuron states not initialized in forward_timestep".to_string(),
                )
            })?;
            self.layer_norm.forward(neuron_states.spikes.clone())?
        };

        Ok(output)
    }

    /// Update Leaky Integrate-and-Fire dynamics
    #[allow(dead_code)]
    fn update_lif_dynamics(&self, current: &Tensor, states: &mut NeuronState) -> Result<()> {
        let dt = self.config.dt;
        let tau_mem = self.config.tau_mem;
        let v_threshold = self.config.v_threshold;
        let v_reset = self.config.v_reset;

        // Membrane potential decay
        let decay = (-dt / tau_mem).exp();
        states.v_mem = states.v_mem.mul_scalar(decay)?.add(&current.mul_scalar(dt / tau_mem)?)?;

        // Add noise
        let noise = Tensor::randn_like(&states.v_mem)?.scalar_mul(self.config.noise_variance)?;
        states.v_mem = states.v_mem.add(&noise)?;

        // Check for spikes
        let threshold_tensor = Tensor::full(v_threshold, states.v_mem.shape())?;
        let spike_mask = states.v_mem.greater(&threshold_tensor)?;
        states.spikes = spike_mask.to_dtype(DType::F32)?;

        // Reset membrane potential for spiking neurons
        let reset_tensor = Tensor::full(v_reset, states.v_mem.shape())?;
        // Simple conditional reset: multiply by (1 - spikes) and add reset_value * spikes
        let one_tensor = Tensor::ones_like(&states.spikes)?;
        let inverse_spikes = one_tensor.sub(&states.spikes)?;
        states.v_mem =
            states.v_mem.mul(&inverse_spikes)?.add(&reset_tensor.mul(&states.spikes)?)?;

        // Update refractory period
        states.refractory_time = states.refractory_time.sub_scalar(dt)?;
        states.refractory_time = states.refractory_time.clamp(0.0, f32::INFINITY)?;

        // Set refractory time for new spikes
        let new_refractory = states.spikes.mul_scalar(self.config.refractory_period)?;
        states.refractory_time = states.refractory_time.add(&new_refractory)?;

        Ok(())
    }

    /// Update Izhikevich model dynamics
    #[allow(dead_code)]
    fn update_izhikevich_dynamics(&self, current: &Tensor, states: &mut NeuronState) -> Result<()> {
        let dt = self.config.dt;
        let v_threshold = self.config.v_threshold;
        let _v_reset = self.config.v_reset;

        // Izhikevich parameters
        let a = 0.02; // recovery time constant
        let b = 0.2; // recovery coupling
        let c = -65.0; // reset value
        let d = 2.0; // after-spike reset of u

        let u_recovery = states.u_recovery.as_mut().ok_or_else(|| {
            TrustformersError::runtime_error(
                "Recovery variable not initialized for Izhikevich model".to_string(),
            )
        })?;

        // Membrane potential dynamics
        let v_squared = states.v_mem.pow_scalar(2.0)?;
        let dv_dt = v_squared
            .mul_scalar(0.04)?
            .add(&states.v_mem.mul_scalar(5.0)?)?
            .add_scalar(140.0)?
            .sub(u_recovery)?
            .add(current)?;
        states.v_mem = states.v_mem.add(&dv_dt.mul_scalar(dt)?)?;

        // Recovery variable dynamics
        let du_dt = states.v_mem.mul_scalar(b)?.sub(u_recovery)?.mul_scalar(a)?;
        *u_recovery = u_recovery.add(&du_dt.mul_scalar(dt)?)?;

        // Check for spikes
        let threshold_tensor = Tensor::full(v_threshold, states.v_mem.shape())?;
        let spike_mask = states.v_mem.greater(&threshold_tensor)?;
        states.spikes = spike_mask.to_dtype(DType::F32)?;

        // Reset for spiking neurons using conditional arithmetic
        let reset_tensor = Tensor::full(c, states.v_mem.shape())?;
        let one_tensor = Tensor::ones_like(&states.spikes)?;
        let inverse_spikes = one_tensor.sub(&states.spikes)?;
        states.v_mem =
            states.v_mem.mul(&inverse_spikes)?.add(&reset_tensor.mul(&states.spikes)?)?;

        // Update u_recovery for spiking neurons
        let u_increment = u_recovery.add_scalar(d)?;
        *u_recovery = u_recovery.mul(&inverse_spikes)?.add(&u_increment.mul(&states.spikes)?)?;

        Ok(())
    }

    /// Update Hodgkin-Huxley dynamics (simplified)
    #[allow(dead_code)]
    fn update_hh_dynamics(&self, current: &Tensor, states: &mut NeuronState) -> Result<()> {
        // Simplified HH model - just use LIF for now
        self.update_lif_dynamics(current, states)
    }

    /// Update Adaptive Exponential Integrate-and-Fire dynamics
    #[allow(dead_code)]
    fn update_adexp_dynamics(&self, current: &Tensor, states: &mut NeuronState) -> Result<()> {
        let dt = self.config.dt;
        let tau_mem = self.config.tau_mem;
        let v_threshold = self.config.v_threshold;
        let v_reset = self.config.v_reset;

        // AdExp parameters
        let delta_t = 2.0; // slope factor
        let v_t = -50.0; // threshold slope
        let tau_w = 30.0; // adaptation time constant
        let a = 2.0; // adaptation coupling
        let b = 0.0; // adaptation increment

        let adaptation = states.adaptation.as_mut().ok_or_else(|| {
            TrustformersError::runtime_error(
                "Adaptation current not initialized for AdExp model".to_string(),
            )
        })?;

        // Exponential term
        let exp_term = states.v_mem.sub_scalar(v_t)?.div_scalar(delta_t)?.exp()?;
        let exp_current = exp_term.mul_scalar(delta_t)?;

        // Membrane potential dynamics
        let dv_dt = states
            .v_mem
            .mul_scalar(-1.0 / tau_mem)?
            .add(&exp_current.mul_scalar(1.0 / tau_mem)?)?
            .sub(&adaptation.mul_scalar(1.0 / tau_mem)?)?
            .add(&current.mul_scalar(1.0 / tau_mem)?)?;
        states.v_mem = states.v_mem.add(&dv_dt.mul_scalar(dt)?)?;

        // Adaptation current dynamics
        let dw_dt = states.v_mem.mul_scalar(a)?.sub(adaptation)?.mul_scalar(1.0 / tau_w)?;
        *adaptation = adaptation.add(&dw_dt.mul_scalar(dt)?)?;

        // Check for spikes
        let threshold_tensor = Tensor::full(v_threshold, states.v_mem.shape())?;
        let spike_mask = states.v_mem.greater(&threshold_tensor)?;
        states.spikes = spike_mask.to_dtype(DType::F32)?;

        // Reset for spiking neurons using conditional arithmetic
        let one_tensor = Tensor::ones_like(&states.spikes)?;
        let inverse_spikes = one_tensor.sub(&states.spikes)?;
        let reset_tensor = Tensor::full(v_reset, states.v_mem.shape())?;
        states.v_mem =
            states.v_mem.mul(&inverse_spikes)?.add(&reset_tensor.mul(&states.spikes)?)?;

        // Update adaptation for spiking neurons
        let adaptation_increment = adaptation.add_scalar(b)?;
        *adaptation = adaptation
            .mul(&inverse_spikes)?
            .add(&adaptation_increment.mul(&states.spikes)?)?;

        Ok(())
    }

    /// Update Spike Response Model dynamics
    #[allow(dead_code)]
    fn update_srm_dynamics(&self, current: &Tensor, states: &mut NeuronState) -> Result<()> {
        // Simplified SRM - use LIF for now
        self.update_lif_dynamics(current, states)
    }

    /// Update synaptic plasticity
    #[allow(dead_code)]
    fn update_plasticity(
        &self,
        neuron_states: &mut NeuronState,
        synaptic_states: &mut SynapticState,
    ) -> Result<()> {
        let dt = self.config.dt;
        let learning_rate = self.config.learning_rate;
        let tau_trace = self.config.tau_syn;

        // Update traces
        let trace_decay = (-dt / tau_trace).exp();
        synaptic_states.pre_traces = synaptic_states.pre_traces.mul_scalar(trace_decay)?;
        synaptic_states.post_traces = synaptic_states.post_traces.mul_scalar(trace_decay)?;

        // Add spikes to traces
        synaptic_states.pre_traces = synaptic_states.pre_traces.add(&neuron_states.spikes)?;
        synaptic_states.post_traces = synaptic_states.post_traces.add(&neuron_states.spikes)?;

        // Update weights based on plasticity type
        match self.config.plasticity_type {
            PlasticityType::STDP => {
                self.update_stdp_weights(neuron_states, synaptic_states, learning_rate)?;
            },
            PlasticityType::Hebbian => {
                self.update_hebbian_weights(neuron_states, synaptic_states, learning_rate)?;
            },
            PlasticityType::AntiHebbian => {
                self.update_anti_hebbian_weights(neuron_states, synaptic_states, learning_rate)?;
            },
            PlasticityType::Homeostatic => {
                self.update_homeostatic_weights(neuron_states, synaptic_states, learning_rate)?;
            },
            PlasticityType::Metaplasticity => {
                self.update_metaplasticity_weights(neuron_states, synaptic_states, learning_rate)?;
            },
        }

        Ok(())
    }

    /// Update weights using STDP
    #[allow(dead_code)]
    fn update_stdp_weights(
        &self,
        neuron_states: &NeuronState,
        synaptic_states: &mut SynapticState,
        lr: f32,
    ) -> Result<()> {
        // Pre-before-post: LTP
        let ltp = neuron_states
            .spikes
            .unsqueeze(1)?
            .matmul(&synaptic_states.pre_traces.unsqueeze(0)?)?;

        // Post-before-pre: LTD
        let ltd = synaptic_states
            .post_traces
            .unsqueeze(1)?
            .matmul(&neuron_states.spikes.unsqueeze(0)?)?;

        // Weight update
        let weight_update = ltp.sub(&ltd)?.mul_scalar(lr)?;
        synaptic_states.weights = synaptic_states.weights.add(&weight_update)?;

        Ok(())
    }

    /// Update weights using Hebbian learning
    #[allow(dead_code)]
    fn update_hebbian_weights(
        &self,
        neuron_states: &NeuronState,
        synaptic_states: &mut SynapticState,
        lr: f32,
    ) -> Result<()> {
        let hebbian_update =
            neuron_states.spikes.unsqueeze(1)?.matmul(&neuron_states.spikes.unsqueeze(0)?)?;
        let weight_update = hebbian_update.mul_scalar(lr)?;
        synaptic_states.weights = synaptic_states.weights.add(&weight_update)?;

        Ok(())
    }

    /// Update weights using Anti-Hebbian learning
    #[allow(dead_code)]
    fn update_anti_hebbian_weights(
        &self,
        neuron_states: &NeuronState,
        synaptic_states: &mut SynapticState,
        lr: f32,
    ) -> Result<()> {
        let anti_hebbian_update =
            neuron_states.spikes.unsqueeze(1)?.matmul(&neuron_states.spikes.unsqueeze(0)?)?;
        let weight_update = anti_hebbian_update.mul_scalar(-lr)?;
        synaptic_states.weights = synaptic_states.weights.add(&weight_update)?;

        Ok(())
    }

    /// Update weights using homeostatic plasticity
    #[allow(dead_code)]
    fn update_homeostatic_weights(
        &self,
        neuron_states: &NeuronState,
        synaptic_states: &mut SynapticState,
        lr: f32,
    ) -> Result<()> {
        let target_rate = self.config.target_rate;
        let current_rate = neuron_states.spikes.mean()?;
        let rate_error = current_rate.sub_scalar(target_rate)?;

        let homeostatic_update = rate_error.mul_scalar(-lr)?;
        let homeostatic_scalar = homeostatic_update.to_scalar()?;
        synaptic_states.weights = synaptic_states.weights.add_scalar(homeostatic_scalar)?;

        Ok(())
    }

    /// Update weights using metaplasticity
    #[allow(dead_code)]
    fn update_metaplasticity_weights(
        &self,
        neuron_states: &NeuronState,
        synaptic_states: &mut SynapticState,
        lr: f32,
    ) -> Result<()> {
        // Combine STDP with homeostatic regulation
        self.update_stdp_weights(neuron_states, synaptic_states, lr * 0.8)?;
        self.update_homeostatic_weights(neuron_states, synaptic_states, lr * 0.2)?;

        Ok(())
    }

    /// Reset neuron states
    pub fn reset_states(&mut self) -> Result<()> {
        if let Some(states) = &mut self.neuron_states {
            states.v_mem = Tensor::zeros_like(&states.v_mem)?;
            states.refractory_time = Tensor::zeros_like(&states.refractory_time)?;
            states.spikes = Tensor::zeros_like(&states.spikes)?;

            if let Some(u_recovery) = &mut states.u_recovery {
                *u_recovery = Tensor::zeros_like(u_recovery)?;
            }

            if let Some(adaptation) = &mut states.adaptation {
                *adaptation = Tensor::zeros_like(adaptation)?;
            }
        }

        if let Some(synaptic_states) = &mut self.synaptic_states {
            synaptic_states.pre_traces = Tensor::zeros_like(&synaptic_states.pre_traces)?;
            synaptic_states.post_traces = Tensor::zeros_like(&synaptic_states.post_traces)?;
            synaptic_states.eligibility = Tensor::zeros_like(&synaptic_states.eligibility)?;
        }

        Ok(())
    }

    /// Get parameter count
    pub fn parameter_count(&self) -> usize {
        self.input_projection.parameter_count()
            + self.recurrent_projection.parameter_count()
            + self.layer_norm.parameter_count()
    }

    /// Get memory usage in MB
    pub fn memory_usage(&self) -> f32 {
        let param_memory = self.parameter_count() as f32 * 4.0 / 1_000_000.0; // 4 bytes per f32
        let state_memory = if self.neuron_states.is_some() {
            self.config.neurons_per_layer as f32 * 4.0 * 4.0 / 1_000_000.0 // 4 tensors per neuron
        } else {
            0.0
        };
        param_memory + state_memory
    }
}

/// Spiking neural network model
#[derive(Debug)]
pub struct SpikingNeuralNetwork {
    /// Configuration
    pub config: BiologicalConfig,
    /// Spiking layers
    pub layers: Vec<SpikingLayer>,
    /// Output projection
    pub output_projection: Linear,
}

impl SpikingNeuralNetwork {
    /// Create a new spiking neural network
    pub fn new(config: &BiologicalConfig) -> Result<Self> {
        let mut layers = Vec::new();
        for _ in 0..config.n_layer {
            layers.push(SpikingLayer::new(config)?);
        }

        let output_projection =
            Linear::new(config.neurons_per_layer, config.d_model, config.use_bias);

        Ok(Self {
            config: config.clone(),
            layers,
            output_projection,
        })
    }

    /// Forward pass through the network
    pub fn forward(&mut self, input: &Tensor) -> Result<BiologicalModelOutput> {
        let mut hidden_states = input.clone();
        let mut all_spike_trains = Vec::new();

        // Pass through all layers
        for layer in &mut self.layers {
            hidden_states = layer.forward(&hidden_states)?;

            // Collect spike trains
            if let Some(states) = &layer.neuron_states {
                all_spike_trains.push(states.spikes.clone());
            }
        }

        // Project to output dimension
        let output = self.output_projection.forward(hidden_states)?;

        // Stack spike trains
        let spike_trains = if !all_spike_trains.is_empty() {
            let mut stacked = all_spike_trains[0].clone();
            for i in 1..all_spike_trains.len() {
                stacked = Tensor::concat(&[stacked, all_spike_trains[i].clone()], 2)?;
            }
            Some(stacked)
        } else {
            None
        };

        Ok(BiologicalModelOutput {
            hidden_states: output,
            spike_trains,
            memory_states: None,
            attention_weights: None,
            capsule_outputs: None,
            dendritic_activations: None,
            plasticity_traces: None,
        })
    }

    /// Update plasticity for all layers
    pub fn update_plasticity(&mut self, _targets: &Tensor) -> Result<()> {
        // Plasticity is updated during forward pass
        Ok(())
    }

    /// Reset states for all layers
    pub fn reset_states(&mut self) -> Result<()> {
        for layer in &mut self.layers {
            layer.reset_states()?;
        }
        Ok(())
    }

    /// Get parameter count
    pub fn parameter_count(&self) -> usize {
        self.layers.iter().map(|l| l.parameter_count()).sum::<usize>()
            + self.output_projection.parameter_count()
    }

    /// Get memory usage in MB
    pub fn memory_usage(&self) -> f32 {
        self.layers.iter().map(|l| l.memory_usage()).sum::<f32>()
            + (self.output_projection.parameter_count() as f32 * 4.0 / 1_000_000.0)
    }
}