Neuromorphic Computing

import numpy as np
import matplotlib.pyplot as plt
from typing import Dict, List, Tuple, Optional, Callable
from dataclasses import dataclass
from abc import ABC, abstractmethod
import time
from collections import defaultdict

@dataclass
class Spike:
    """Represents a neural spike event"""
    neuron_id: int
    timestamp: float
    weight: float = 1.0

class SpikingNeuron(ABC):
    """Abstract base class for spiking neuron models"""
    
    def __init__(self, neuron_id: int, config: Dict):
        self.neuron_id = neuron_id
        self.config = config
        
        # State variables
        self.membrane_potential = config.get('v_rest', -70.0)  # mV
        self.threshold = config.get('v_threshold', -50.0)  # mV
        self.reset_potential = config.get('v_reset', -80.0)  # mV
        
        # Timing
        self.last_spike_time = -np.inf
        self.refractory_period = config.get('refractory_period', 2.0)  # ms
        
        # Input tracking
        self.input_spikes = []
        self.synaptic_weights = {}
        
    @abstractmethod
    def update(self, dt: float, current_time: float) -> Optional[Spike]:
        """Update neuron state and return spike if threshold crossed"""
        pass
    
    def add_input_spike(self, spike: Spike, weight: float = 1.0):
        """Add input spike to neuron"""
        self.input_spikes.append((spike, weight))
    
    def is_refractory(self, current_time: float) -> bool:
        """Check if neuron is in refractory period"""
        return (current_time - self.last_spike_time) < self.refractory_period

class LIFNeuron(SpikingNeuron):
    """Leaky Integrate-and-Fire neuron model"""
    
    def __init__(self, neuron_id: int, config: Dict):
        super().__init__(neuron_id, config)
        
        # LIF parameters
        self.tau_membrane = config.get('tau_membrane', 20.0)  # ms
        self.resistance = config.get('resistance', 1.0)  # MOhm
        self.capacitance = config.get('capacitance', 1.0)  # nF
        
    def update(self, dt: float, current_time: float) -> Optional[Spike]:
        """Update LIF neuron dynamics"""
        
        if self.is_refractory(current_time):
            return None
        
        # Calculate input current from spikes
        input_current = self._calculate_input_current(current_time)
        
        # Membrane potential dynamics: C * dV/dt = -(V - V_rest)/R + I
        leak_current = (self.membrane_potential - self.config['v_rest']) / self.resistance
        
        dv_dt = (-leak_current + input_current) / self.capacitance
        self.membrane_potential += dv_dt * dt
        
        # Check for spike
        if self.membrane_potential >= self.threshold:
            # Generate spike
            spike = Spike(self.neuron_id, current_time)
            
            # Reset membrane potential
            self.membrane_potential = self.reset_potential
            self.last_spike_time = current_time
            
            # Clear processed input spikes
            self.input_spikes = []
            
            return spike
        
        return None
    
    def _calculate_input_current(self, current_time: float) -> float:
        """Calculate total input current from synaptic inputs"""
        
        total_current = 0.0
        
        for spike, weight in self.input_spikes:
            # Exponential decay of synaptic current
            tau_syn = self.config.get('tau_synaptic', 5.0)  # ms
            delay = current_time - spike.timestamp
            
            if delay >= 0:
                synaptic_current = weight * np.exp(-delay / tau_syn)
                total_current += synaptic_current
        
        return total_current

class IzhikevichNeuron(SpikingNeuron):
    """Izhikevich neuron model - efficient and versatile"""
    
    def __init__(self, neuron_id: int, config: Dict):
        super().__init__(neuron_id, config)
        
        # Izhikevich parameters
        self.a = config.get('a', 0.02)  # Recovery time constant
        self.b = config.get('b', 0.2)   # Sensitivity of recovery
        self.c = config.get('c', -65.0) # After-spike reset value
        self.d = config.get('d', 8.0)   # After-spike reset of recovery
        
        # State variables
        self.v = config.get('v_rest', -70.0)  # Membrane potential
        self.u = self.b * self.v  # Recovery variable
        
    def update(self, dt: float, current_time: float) -> Optional[Spike]:
        """Update Izhikevich neuron dynamics"""
        
        # Calculate input current
        input_current = self._calculate_input_current(current_time)
        
        # Izhikevich equations
        dv_dt = 0.04 * self.v**2 + 5 * self.v + 140 - self.u + input_current
        du_dt = self.a * (self.b * self.v - self.u)
        
        self.v += dv_dt * dt
        self.u += du_dt * dt
        
        # Check for spike
        if self.v >= 30.0:  # Spike threshold
            # Generate spike
            spike = Spike(self.neuron_id, current_time)
            
            # Reset dynamics
            self.v = self.c
            self.u += self.d
            self.last_spike_time = current_time
            
            # Clear processed input spikes
            self.input_spikes = []
            
            return spike
        
        return None
    
    def _calculate_input_current(self, current_time: float) -> float:
        """Calculate input current for Izhikevich model"""
        
        total_current = 0.0
        
        for spike, weight in self.input_spikes:
            # Alpha function for synaptic current
            tau = self.config.get('tau_synaptic', 5.0)
            delay = current_time - spike.timestamp
            
            if delay >= 0:
                # Alpha function: I(t) = (t/tau) * exp(1 - t/tau)
                alpha_factor = (delay / tau) * np.exp(1 - delay / tau)
                synaptic_current = weight * alpha_factor
                total_current += synaptic_current
        
        return total_current

class STDPSynapse:
    """Spike-Timing Dependent Plasticity synapse"""
    
    def __init__(self, pre_neuron_id: int, post_neuron_id: int, config: Dict):
        self.pre_neuron_id = pre_neuron_id
        self.post_neuron_id = post_neuron_id
        self.config = config
        
        # Synaptic weight
        self.weight = config.get('initial_weight', 0.5)
        self.min_weight = config.get('min_weight', 0.0)
        self.max_weight = config.get('max_weight', 1.0)
        
        # STDP parameters
        self.tau_plus = config.get('tau_plus', 20.0)  # ms
        self.tau_minus = config.get('tau_minus', 20.0)  # ms
        self.a_plus = config.get('a_plus', 0.01)
        self.a_minus = config.get('a_minus', 0.01)
        
        # Spike history
        self.pre_spike_times = []
        self.post_spike_times = []
        
    def update_weight(self, pre_spike_time: Optional[float], 
                     post_spike_time: Optional[float]):
        """Update synaptic weight based on STDP rule"""
        
        # Record spike times
        if pre_spike_time is not None:
            self.pre_spike_times.append(pre_spike_time)
        
        if post_spike_time is not None:
            self.post_spike_times.append(post_spike_time)
        
        # Apply STDP for recent spikes
        if pre_spike_time is not None and post_spike_time is not None:
            dt = post_spike_time - pre_spike_time
            
            if dt > 0:  # Post before pre (potentiation)
                weight_change = self.a_plus * np.exp(-dt / self.tau_plus)
                self.weight += weight_change
            elif dt < 0:  # Pre before post (depression)
                weight_change = -self.a_minus * np.exp(dt / self.tau_minus)
                self.weight += weight_change
            
            # Clip weight to bounds
            self.weight = np.clip(self.weight, self.min_weight, self.max_weight)

class SpikingNeuralNetwork:
    """
    Comprehensive spiking neural network framework
    
    Features:
    - Multiple neuron models (LIF, Izhikevich, etc.)
    - STDP learning and other plasticity rules
    - Event-driven simulation
    - Real-time processing capabilities
    - Encoding/decoding for classical data
    """
    
    def __init__(self, config: Dict):
        self.config = config
        
        # Network components
        self.neurons = {}
        self.synapses = {}
        self.layers = {}
        
        # Simulation state
        self.current_time = 0.0
        self.dt = config.get('dt', 0.1)  # ms
        
        # Event queues
        self.spike_events = []
        self.spike_history = defaultdict(list)
        
        # Monitoring
        self.monitors = {}
        
    def add_neuron(self, neuron_id: int, neuron_type: str, config: Dict):
        """Add neuron to network"""
        
        if neuron_type == 'LIF':
            neuron = LIFNeuron(neuron_id, config)
        elif neuron_type == 'Izhikevich':
            neuron = IzhikevichNeuron(neuron_id, config)
        else:
            raise ValueError(f"Unknown neuron type: {neuron_type}")
        
        self.neurons[neuron_id] = neuron
    
    def add_synapse(self, pre_id: int, post_id: int, weight: float, 
                   plasticity: str = 'static'):
        """Add synapse between neurons"""
        
        synapse_id = (pre_id, post_id)
        
        if plasticity == 'STDP':
            synapse = STDPSynapse(pre_id, post_id, {
                'initial_weight': weight,
                **self.config.get('stdp_config', {})
            })
        else:
            # Static synapse
            synapse = {'weight': weight, 'type': 'static'}
        
        self.synapses[synapse_id] = synapse
    
    def simulate(self, duration: float, input_spikes: List[Spike] = None) -> Dict:
        """Simulate network for given duration"""
        
        if input_spikes is None:
            input_spikes = []
        
        # Sort input spikes by time
        input_spikes.sort(key=lambda s: s.timestamp)
        
        # Simulation loop
        end_time = self.current_time + duration
        input_spike_index = 0
        
        while self.current_time < end_time:
            # Process input spikes at current time
            while (input_spike_index < len(input_spikes) and 
                   input_spikes[input_spike_index].timestamp <= self.current_time):
                
                spike = input_spikes[input_spike_index]
                self._process_input_spike(spike)
                input_spike_index += 1
            
            # Update all neurons
            output_spikes = []
            for neuron_id, neuron in self.neurons.items():
                spike = neuron.update(self.dt, self.current_time)
                
                if spike is not None:
                    output_spikes.append(spike)
                    self.spike_history[neuron_id].append(spike.timestamp)
            
            # Propagate spikes through synapses
            for spike in output_spikes:
                self._propagate_spike(spike)
            
            # Update plasticity
            self._update_plasticity(output_spikes)
            
            # Advance time
            self.current_time += self.dt
        
        return {
            'spike_history': dict(self.spike_history),
            'final_time': self.current_time,
            'total_spikes': sum(len(spikes) for spikes in self.spike_history.values())
        }
    
    def _process_input_spike(self, spike: Spike):
        """Process external input spike"""
        
        if spike.neuron_id in self.neurons:
            neuron = self.neurons[spike.neuron_id]
            neuron.add_input_spike(spike)
    
    def _propagate_spike(self, spike: Spike):
        """Propagate spike through network synapses"""
        
        pre_neuron_id = spike.neuron_id
        
        # Find all synapses from this neuron
        for (pre_id, post_id), synapse in self.synapses.items():
            if pre_id == pre_neuron_id:
                # Get synaptic weight
                if isinstance(synapse, STDPSynapse):
                    weight = synapse.weight
                else:
                    weight = synapse['weight']
                
                # Create weighted spike
                weighted_spike = Spike(
                    neuron_id=pre_neuron_id,
                    timestamp=spike.timestamp,
                    weight=weight
                )
                
                # Send to post-synaptic neuron
                if post_id in self.neurons:
                    post_neuron = self.neurons[post_id]
                    post_neuron.add_input_spike(weighted_spike, weight)
    
    def _update_plasticity(self, spikes: List[Spike]):
        """Update synaptic plasticity"""
        
        for spike in spikes:
            pre_neuron_id = spike.neuron_id
            
            # Update STDP synapses
            for (pre_id, post_id), synapse in self.synapses.items():
                if isinstance(synapse, STDPSynapse):
                    pre_spike_time = spike.timestamp if pre_id == pre_neuron_id else None
                    
                    # Check for post-synaptic spikes
                    post_spike_time = None
                    for post_spike in spikes:
                        if post_spike.neuron_id == post_id:
                            post_spike_time = post_spike.timestamp
                            break
                    
                    synapse.update_weight(pre_spike_time, post_spike_time)

class SpikeEncoder:
    """Encode classical data into spike trains"""
    
    @staticmethod
    def rate_encoding(values: np.ndarray, duration: float, max_rate: float = 100.0) -> List[Spike]:
        """Encode values as Poisson spike trains"""
        
        spikes = []
        
        for neuron_id, value in enumerate(values):
            # Normalize value to rate
            rate = max_rate * max(0, min(1, value))  # Hz
            
            # Generate Poisson spikes
            if rate > 0:
                # Average inter-spike interval
                isi = 1000.0 / rate  # ms
                
                current_time = 0.0
                while current_time < duration:
                    # Exponential distribution for next spike
                    next_isi = np.random.exponential(isi)
                    current_time += next_isi
                    
                    if current_time < duration:
                        spikes.append(Spike(neuron_id, current_time))
        
        return spikes
    
    @staticmethod
    def temporal_encoding(values: np.ndarray, duration: float) -> List[Spike]:
        """Encode values as spike timing"""
        
        spikes = []
        
        for neuron_id, value in enumerate(values):
            # Map value to spike time
            spike_time = duration * max(0, min(1, value))
            spikes.append(Spike(neuron_id, spike_time))
        
        return spikes

class SpikeDecoder:
    """Decode spike trains into classical data"""
    
    @staticmethod
    def rate_decoding(spike_history: Dict[int, List[float]], 
                     window_duration: float) -> np.ndarray:
        """Decode spike trains using firing rates"""
        
        neuron_ids = sorted(spike_history.keys())
        rates = []
        
        for neuron_id in neuron_ids:
            spike_times = spike_history[neuron_id]
            
            # Count spikes in window
            spike_count = len(spike_times)
            
            # Convert to rate (Hz)
            rate = spike_count / (window_duration / 1000.0)
            rates.append(rate)
        
        return np.array(rates)
    
    @staticmethod
    def temporal_decoding(spike_history: Dict[int, List[float]]) -> np.ndarray:
        """Decode using first spike times"""
        
        neuron_ids = sorted(spike_history.keys())
        times = []
        
        for neuron_id in neuron_ids:
            spike_times = spike_history[neuron_id]
            
            if spike_times:
                # Use first spike time
                first_spike = min(spike_times)
                times.append(first_spike)
            else:
                # No spike - use maximum time
                times.append(np.inf)
        
        # Normalize to [0, 1]
        times = np.array(times)
        max_time = np.max(times[times != np.inf])
        
        if max_time > 0:
            times = np.where(times == np.inf, 1.0, times / max_time)
        
        return times

# Example applications
def pattern_recognition_example():
    """Example: Pattern recognition with SNN"""
    
    # Create network
    config = {
        'dt': 0.1,
        'stdp_config': {
            'tau_plus': 20.0,
            'tau_minus': 20.0,
            'a_plus': 0.01,
            'a_minus': 0.01
        }
    }
    
    network = SpikingNeuralNetwork(config)
    
    # Add input layer (4 neurons)
    for i in range(4):
        network.add_neuron(i, 'LIF', {
            'v_rest': -70.0,
            'v_threshold': -50.0,
            'v_reset': -80.0,
            'tau_membrane': 20.0
        })
    
    # Add hidden layer (2 neurons)
    for i in range(4, 6):
        network.add_neuron(i, 'Izhikevich', {
            'a': 0.02,
            'b': 0.2,
            'c': -65.0,
            'd': 8.0
        })
    
    # Add output layer (1 neuron)
    network.add_neuron(6, 'LIF', {
        'v_rest': -70.0,
        'v_threshold': -45.0,  # Lower threshold for output
        'v_reset': -80.0,
        'tau_membrane': 20.0
    })
    
    # Connect layers with STDP synapses
    # Input to hidden
    for i in range(4):
        for j in range(4, 6):
            network.add_synapse(i, j, weight=0.5, plasticity='STDP')
    
    # Hidden to output
    for i in range(4, 6):
        network.add_synapse(i, 6, weight=0.8, plasticity='STDP')
    
    # Training patterns
    patterns = [
        np.array([1.0, 0.0, 1.0, 0.0]),  # Pattern A
        np.array([0.0, 1.0, 0.0, 1.0]),  # Pattern B
    ]
    
    # Training
    for epoch in range(10):
        for pattern_id, pattern in enumerate(patterns):
            # Encode pattern
            input_spikes = SpikeEncoder.rate_encoding(pattern, duration=50.0)
            
            # Simulate
            result = network.simulate(duration=50.0, input_spikes=input_spikes)
            
            print(f"Epoch {epoch}, Pattern {pattern_id}: {len(result['spike_history'].get(6, []))} output spikes")
    
    return network

def real_time_processing_example():
    """Example: Real-time sensor processing"""
    
    print("Real-time SNN processing simulation...")
    
    # Simplified real-time processing loop
    network = SpikingNeuralNetwork({'dt': 1.0})  # 1ms resolution
    
    # Add sensor processing neurons
    for i in range(8):  # 8 sensor inputs
        network.add_neuron(i, 'LIF', {
            'v_rest': -70.0,
            'v_threshold': -50.0,
            'tau_membrane': 10.0  # Fast dynamics
        })
    
    # Simulate real-time processing
    for time_step in range(100):  # 100ms simulation
        # Generate sensor data (simulated)
        sensor_data = np.random.rand(8)
        
        # Encode as spikes
        input_spikes = SpikeEncoder.rate_encoding(sensor_data, duration=1.0, max_rate=200.0)
        
        # Process in real-time
        start_time = time.time()
        result = network.simulate(duration=1.0, input_spikes=input_spikes)
        processing_time = (time.time() - start_time) * 1000  # ms
        
        if time_step % 20 == 0:
            print(f"Time {time_step}ms: {processing_time:.2f}ms processing time")
    
    return "Real-time processing completed"

if __name__ == "__main__":
    print("=== Spiking Neural Network Examples ===")
    
    # Pattern recognition
    print("\n1. Pattern Recognition:")
    pattern_network = pattern_recognition_example()
    
    # Real-time processing
    print("\n2. Real-time Processing:")
    real_time_result = real_time_processing_example()
    print(real_time_result)