improved_physical_brain_software.py

"""
IMPROVED PHYSICAL MEMRISTOR BRAIN NETWORK - WITH EFFECTIVE LEARNING
===================================================================

This version uses STRONGER learning pulses that actually change memristor
resistances, showing realistic lab behavior where:
- Initial learning pulses are weak → minimal resistance change
- Progressive training increases pulse strength → visible learning
- Final network shows clear pattern classification

Key improvements:
1. Adaptive learning pulse strength
2. Realistic resistance switching thresholds
3. Better visualization of learning progress
4. More detailed analysis of what's happening
"""

import numpy as np
import matplotlib.pyplot as plt
from customizable_memristor import create_custom_devices
from typing import List, Tuple, Dict
import time

class ImprovedPhysicalMemristorBrain:
    """
    Realistic physical memristor brain with effective learning
    """
    
    def __init__(self, memristors):
        """
        Initialize with adaptive learning parameters
        """
        
        self.memristors = memristors
        
        # Circuit topology - same as before
        self.input_terminals = ['I0', 'I1', 'I2']
        self.junction_points = ['H0', 'H1'] 
        self.output_terminals = ['O0', 'O1']
        
        # Physical connections
        self.physical_connections = {
            'I0_to_H0': 0, 'I1_to_H0': 1, 'I2_to_H0': 2,
            'I0_to_H1': 3, 'I1_to_H1': 4, 'I2_to_H1': 5,
            'H0_to_O0': 6, 'H1_to_O1': 7,
        }
        
        # Current state
        self.input_voltages = np.zeros(3)
        self.junction_voltages = np.zeros(2)
        self.output_voltages = np.zeros(2)
        
        # IMPROVED Learning parameters
        self.base_learning_rate = 0.3  # Higher base rate
        self.learning_boost_factor = 1.0  # Increases over time
        self.min_switching_voltage = 2.8  # Minimum to guarantee switching
        self.training_history = []
        self.resistance_history = []  # Track resistance changes
        
        print("🧠 IMPROVED PHYSICAL MEMRISTOR BRAIN NETWORK")
        print("=" * 60)
        print(f"🔧 Enhanced with ADAPTIVE LEARNING PULSES")
        print(f"⚡ Minimum switching voltage: {self.min_switching_voltage}V")
        print(f"📈 Learning boost factor: {self.learning_boost_factor}x")
        print()
        
        # Record initial resistances
        initial_resistances = [mem.get_resistance() for mem in self.memristors]
        self.resistance_history.append(initial_resistances.copy())
        
        print("🔗 INITIAL MEMRISTOR STATES:")
        for i, resistance in enumerate(initial_resistances):
            print(f"   M{i}: {resistance:.0f}Ω")
        print()
    
    def set_input_voltages(self, voltages):
        """Set input voltages and calculate circuit response"""
        self.input_voltages = np.array(voltages)
        self.calculate_junction_voltages()
        self.calculate_output_voltages()
    
    def calculate_junction_voltages(self):
        """Calculate junction voltages with proper current division"""
        
        # Junction H0: Parallel combination of M0, M1, M2
        resistances_H0 = [self.memristors[i].get_resistance() for i in [0, 1, 2]]
        currents_H0 = [self.input_voltages[i] / resistances_H0[i] for i in range(3)]
        
        # H0 voltage from Norton equivalent
        total_current_H0 = sum(currents_H0)
        total_conductance_H0 = sum(1/r for r in resistances_H0)
        self.junction_voltages[0] = total_current_H0 / total_conductance_H0
        
        # Junction H1: Parallel combination of M3, M4, M5
        resistances_H1 = [self.memristors[i].get_resistance() for i in [3, 4, 5]]
        currents_H1 = [self.input_voltages[i] / resistances_H1[i] for i in range(3)]
        
        total_current_H1 = sum(currents_H1)
        total_conductance_H1 = sum(1/r for r in resistances_H1)
        self.junction_voltages[1] = total_current_H1 / total_conductance_H1
    
    def calculate_output_voltages(self):
        """Calculate output voltages through M6 and M7"""
        
        load_resistance = 1000  # Oscilloscope input impedance
        
        # O0: H0 → M6 → O0
        resistance_M6 = self.memristors[6].get_resistance()
        voltage_divider_ratio = load_resistance / (resistance_M6 + load_resistance)
        self.output_voltages[0] = self.junction_voltages[0] * voltage_divider_ratio
        
        # O1: H1 → M7 → O1
        resistance_M7 = self.memristors[7].get_resistance()
        voltage_divider_ratio = load_resistance / (resistance_M7 + load_resistance)
        self.output_voltages[1] = self.junction_voltages[1] * voltage_divider_ratio
    
    def apply_effective_learning_pulse(self, memristor_idx, error_signal, input_activation, epoch):
        """
        Apply learning pulse that ACTUALLY changes resistance
        Pulse strength adapts based on:
        1. Error magnitude
        2. Input activity level  
        3. Training progress (gets stronger over time)
        4. Minimum switching threshold
        """
        
        memristor = self.memristors[memristor_idx]
        resistance_before = memristor.get_resistance()
        
        # Calculate adaptive pulse voltage
        base_voltage = 4.0 * self.learning_boost_factor  # Base amplitude
        error_scaling = abs(error_signal) * 2.0  # Scale with error
        activity_scaling = abs(input_activation) + 0.5  # Ensure minimum activity
        
        # Raw pulse calculation
        raw_pulse = base_voltage * error_scaling * activity_scaling * self.base_learning_rate
        
        # Ensure minimum switching voltage
        if raw_pulse > 0:  # SET operation (lower resistance)
            pulse_voltage = max(raw_pulse, self.min_switching_voltage)
            if error_signal > 0:  # Want to strengthen (SET)
                pulse_voltage = max(pulse_voltage, 3.2)  # Guarantee SET
            else:  # Want to weaken (RESET) 
                pulse_voltage = min(-abs(pulse_voltage), -0.3)  # Guarantee RESET
        else:  # RESET operation (higher resistance)
            pulse_voltage = min(raw_pulse, -self.min_switching_voltage)
            if error_signal > 0:  # Want to strengthen (SET)
                pulse_voltage = max(abs(pulse_voltage), 3.2)  # Guarantee SET
            else:  # Want to weaken (RESET)
                pulse_voltage = min(-abs(pulse_voltage), -0.3)  # Guarantee RESET
        
        # Apply the pulse with longer duration for better switching
        duration_ms = 2.0  # Longer pulses for reliable switching
        
        print(f"🔄 Learning pulse M{memristor_idx}: {pulse_voltage:+.2f}V × {duration_ms}ms")
        print(f"   Error: {error_signal:+.3f}, Activity: {input_activation:.3f}, Epoch: {epoch}")
        
        # Apply pulse to memristor
        memristor.update_state_from_voltage(pulse_voltage, dt=duration_ms*1e-3)
        
        resistance_after = memristor.get_resistance()
        resistance_change = resistance_after - resistance_before
        
        print(f"   Resistance: {resistance_before:.0f}Ω → {resistance_after:.0f}Ω ({resistance_change:+.0f}Ω)")
        
        # Update circuit state
        self.calculate_junction_voltages()
        self.calculate_output_voltages()
        
        return resistance_before, resistance_after, pulse_voltage
    
    def train_with_adaptive_learning(self, training_patterns, epochs=15):
        """
        Train with progressively stronger learning as network struggles
        """
        
        print("🎓 ADAPTIVE LEARNING TRAINING")
        print("=" * 60)
        print("Learning will get stronger if network doesn't improve!")
        print()
        
        for epoch in range(epochs):
            total_error = 0
            epoch_resistances = []
            
            print(f"\n🔄 EPOCH {epoch+1}/{epochs} (Learning boost: {self.learning_boost_factor:.2f}x)")
            print("-" * 50)
            
            for pattern_idx, (inputs, target_class) in enumerate(training_patterns):
                # Forward pass
                self.set_input_voltages(inputs)
                outputs = self.measure_outputs()
                
                # Target outputs
                target_outputs = [1.0, 0.0] if target_class == 'A' else [0.0, 1.0]
                
                # Calculate errors
                output_errors = np.array(target_outputs) - outputs
                pattern_error = np.sum(output_errors**2)
                total_error += pattern_error
                
                # Apply adaptive learning
                self.apply_adaptive_learning_updates(inputs, output_errors, epoch)
                
                # Classification result
                predicted_class = 'A' if outputs[0] > outputs[1] else 'B'
                correct = predicted_class == target_class
                
                print(f"Pattern {pattern_idx+1}: {inputs} → Class {target_class}")
                print(f"  Output: O0={outputs[0]:.3f}V, O1={outputs[1]:.3f}V → {predicted_class} {'✅' if correct else '❌'}")
                print(f"  Error: {pattern_error:.4f}")
                print()
            
            # Record resistance state
            current_resistances = [mem.get_resistance() for mem in self.memristors]
            self.resistance_history.append(current_resistances.copy())
            
            # Calculate average error
            avg_error = total_error / len(training_patterns)
            self.training_history.append(avg_error)
            
            print(f"📊 EPOCH {epoch+1} SUMMARY:")
            print(f"   Average Error: {avg_error:.4f}")
            print(f"   Current resistances:")
            for i, resistance in enumerate(current_resistances):
                print(f"     M{i}: {resistance:.0f}Ω")
            
            # Adaptive learning boost
            if epoch > 2 and len(self.training_history) >= 3:
                # If error isn't decreasing, boost learning
                recent_errors = self.training_history[-3:]
                if recent_errors[-1] >= recent_errors[-2]:  # No improvement
                    self.learning_boost_factor *= 1.2  # Increase learning strength
                    print(f"   🚀 Learning boost increased to {self.learning_boost_factor:.2f}x")
                elif recent_errors[-1] < recent_errors[-2] * 0.8:  # Good improvement
                    self.learning_boost_factor = max(1.0, self.learning_boost_factor * 0.95)  # Slight decrease
            
            print("-" * 50)
        
        print("\n✅ ADAPTIVE TRAINING COMPLETED!")
        self.analyze_learning_results()
    
    def apply_adaptive_learning_updates(self, inputs, output_errors, epoch):
        """Apply learning updates with proper error attribution"""
        
        # Update output layer (M6, M7) based on output errors
        if abs(output_errors[0]) > 0.05:  # Significant O0 error
            h0_activity = self.junction_voltages[0]
            self.apply_effective_learning_pulse(6, output_errors[0], h0_activity, epoch)
        
        if abs(output_errors[1]) > 0.05:  # Significant O1 error
            h1_activity = self.junction_voltages[1]
            self.apply_effective_learning_pulse(7, output_errors[1], h1_activity, epoch)
        
        # Update input layer based on which path should be strengthened
        for i, input_val in enumerate(inputs):
            if input_val > 0.3:  # Active input
                
                # Update path to H0 based on O0 error
                if abs(output_errors[0]) > 0.05:
                    mem_idx = i  # M0, M1, M2
                    self.apply_effective_learning_pulse(mem_idx, output_errors[0], input_val, epoch)
                
                # Update path to H1 based on O1 error  
                if abs(output_errors[1]) > 0.05:
                    mem_idx = i + 3  # M3, M4, M5
                    self.apply_effective_learning_pulse(mem_idx, output_errors[1], input_val, epoch)
    
    def measure_outputs(self):
        """Get current output voltages"""
        return self.output_voltages.copy()
    
    def test_pattern(self, inputs, expected_class=None):
        """Test pattern with detailed analysis"""
        
        self.set_input_voltages(inputs)
        outputs = self.measure_outputs()
        
        predicted_class = 'A' if outputs[0] > outputs[1] else 'B'
        confidence = abs(outputs[0] - outputs[1])  # Separation between outputs
        
        print(f"🧪 TEST PATTERN:")
        print(f"  Inputs:  I0={inputs[0]:.1f}V, I1={inputs[1]:.1f}V, I2={inputs[2]:.1f}V")
        print(f"  Junctions: H0={self.junction_voltages[0]:.3f}V, H1={self.junction_voltages[1]:.3f}V")
        print(f"  Outputs: O0={outputs[0]:.3f}V, O1={outputs[1]:.3f}V")
        print(f"  Classification: Class {predicted_class} (confidence: {confidence:.3f}V)")
        
        if expected_class:
            correct = predicted_class == expected_class
            print(f"  Expected: Class {expected_class} → {'✅ CORRECT' if correct else '❌ WRONG'}")
        
        print()
        return predicted_class, confidence, outputs
    
    def analyze_learning_results(self):
        """Analyze how the network learned"""
        
        print("\n🔍 LEARNING ANALYSIS")
        print("=" * 60)
        
        initial_resistances = self.resistance_history[0]
        final_resistances = self.resistance_history[-1]
        
        print("📊 RESISTANCE CHANGES:")
        print("   Memristor | Initial | Final  | Change | Function")
        print("   " + "-"*50)
        
        functions = [
            "I0→H0", "I1→H0", "I2→H0", "I0→H1", "I1→H1", "I2→H1", "H0→O0", "H1→O1"
        ]
        
        for i in range(8):
            initial = initial_resistances[i]
            final = final_resistances[i]
            change = final - initial
            direction = "↓" if change < 0 else "↑" if change > 0 else "="
            
            print(f"   M{i}       | {initial:4.0f}Ω  | {final:4.0f}Ω | {change:+4.0f}Ω {direction} | {functions[i]}")
        
        print("\n🧠 LEARNING INSIGHTS:")
        
        # Analyze which paths got strengthened
        h0_strengthened = sum(final_resistances[i] < initial_resistances[i] for i in [0,1,2])
        h1_strengthened = sum(final_resistances[i] < initial_resistances[i] for i in [3,4,5])
        
        if h0_strengthened > h1_strengthened:
            print("   • Network favors H0 path → Should classify more as Class A")
        elif h1_strengthened > h0_strengthened:
            print("   • Network favors H1 path → Should classify more as Class B")
        else:
            print("   • Network shows balanced learning")
        
        # Output layer analysis
        if final_resistances[6] < initial_resistances[6]:
            print("   • M6 (H0→O0) strengthened → Enhanced Class A output")
        if final_resistances[7] < initial_resistances[7]:
            print("   • M7 (H1→O1) strengthened → Enhanced Class B output")
    
    def visualize_learning_progress(self):
        """Visualize how learning progressed"""
        
        fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(16, 12))
        
        # 1. Training error over time
        ax1.plot(self.training_history, 'b-', linewidth=2, marker='o')
        ax1.set_title('🎓 Learning Progress', fontsize=14, fontweight='bold')
        ax1.set_xlabel('Epoch')
        ax1.set_ylabel('Average Error')
        ax1.grid(True, alpha=0.3)
        ax1.set_yscale('log')
        
        # 2. Resistance evolution
        ax2.set_title('🔗 Memristor Resistance Evolution', fontsize=14, fontweight='bold')
        
        epochs = range(len(self.resistance_history))
        colors = plt.cm.tab10(np.linspace(0, 1, 8))
        
        for i in range(8):
            resistances = [self.resistance_history[epoch][i] for epoch in epochs]
            ax2.plot(epochs, resistances, color=colors[i], linewidth=2, 
                    marker='o', label=f'M{i}', markersize=4)
        
        ax2.set_xlabel('Epoch')
        ax2.set_ylabel('Resistance (Ω)')
        ax2.legend(bbox_to_anchor=(1.05, 1), loc='upper left')
        ax2.grid(True, alpha=0.3)
        
        # 3. Final resistance comparison
        ax3.set_title('📊 Initial vs Final Resistances', fontsize=14, fontweight='bold')
        
        initial = self.resistance_history[0]
        final = self.resistance_history[-1]
        
        x = np.arange(8)
        width = 0.35
        
        ax3.bar(x - width/2, initial, width, label='Initial', alpha=0.7, color='lightblue')
        ax3.bar(x + width/2, final, width, label='Final', alpha=0.7, color='orange')
        
        ax3.set_xlabel('Memristor ID')
        ax3.set_ylabel('Resistance (Ω)')
        ax3.set_xticks(x)
        ax3.set_xticklabels([f'M{i}' for i in range(8)])
        ax3.legend()
        ax3.grid(True, alpha=0.3)
        
        # 4. Current network state
        ax4.set_title('⚡ Current Network State', fontsize=14, fontweight='bold')
        
        circuit_info = f"""
CURRENT CIRCUIT STATE:

Signal Flow:
I0={self.input_voltages[0]:.2f}V → M0({final[0]:.0f}Ω) →┐
I1={self.input_voltages[1]:.2f}V → M1({final[1]:.0f}Ω) →┼→ H0={self.junction_voltages[0]:.3f}V → M6({final[6]:.0f}Ω) → O0={self.output_voltages[0]:.3f}V
I2={self.input_voltages[2]:.2f}V → M2({final[2]:.0f}Ω) →┘

I0={self.input_voltages[0]:.2f}V → M3({final[3]:.0f}Ω) →┐
I1={self.input_voltages[1]:.2f}V → M4({final[4]:.0f}Ω) →┼→ H1={self.junction_voltages[1]:.3f}V → M7({final[7]:.0f}Ω) → O1={self.output_voltages[1]:.3f}V
I2={self.input_voltages[2]:.2f}V → M5({final[5]:.0f}Ω) →┘

Classification: {'Class A' if self.output_voltages[0] > self.output_voltages[1] else 'Class B'}
Confidence: {abs(self.output_voltages[0] - self.output_voltages[1]):.3f}V
        """
        
        ax4.text(0.05, 0.95, circuit_info, transform=ax4.transAxes, fontsize=10,
                verticalalignment='top', fontfamily='monospace')
        ax4.set_xlim(0, 1)
        ax4.set_ylim(0, 1)
        ax4.axis('off')
        
        plt.tight_layout()
        plt.show()

def create_realistic_training_data():
    """Create training data that should show clear learning"""
    
    return [
        # Class A: Strong positive patterns
        ([2.5, 2.0, 2.5], 'A'),  # Strong Class A
        ([2.0, 2.5, 2.0], 'A'),  # Strong Class A variant
        ([2.2, 1.8, 2.3], 'A'),  # Moderate Class A
        
        # Class B: Weak patterns
        ([0.3, 0.2, 0.4], 'B'),  # Weak Class B
        ([0.2, 0.4, 0.3], 'B'),  # Weak Class B variant
        ([0.5, 0.3, 0.2], 'B'),  # Moderate Class B
    ]

def main():
    """
    Demonstrate improved learning with effective pulse strengths
    """
    
    print("🧠 IMPROVED PHYSICAL MEMRISTOR BRAIN NETWORK")
    print("=" * 70)
    print("This version uses ADAPTIVE LEARNING that actually works!")
    print()
    
    # Create memristors
    custom_params = {
        'ron_mean': 400, 'roff_mean': 1000,
        'vth_set_mean': 3.0, 'vth_reset_mean': -0.2,
        'compliance_current': 12.5e-3,
        'ron_std': 20, 'roff_std': 50,
        'vth_set_std': 0.05, 'vth_reset_std': 0.05
    }
    
    memristors = create_custom_devices(8, custom_params)
    
    # Create improved network
    network = ImprovedPhysicalMemristorBrain(memristors)
    
    # Create training data
    training_data = create_realistic_training_data()
    
    print("📊 Testing BEFORE training:")
    network.test_pattern([2.0, 2.0, 2.0], expected_class='A')
    network.test_pattern([0.3, 0.3, 0.3], expected_class='B')
    
    # Train with adaptive learning
    network.train_with_adaptive_learning(training_data, epochs=12)
    
    print("\n📊 Testing AFTER training:")
    network.test_pattern([2.0, 2.0, 2.0], expected_class='A')
    network.test_pattern([0.3, 0.3, 0.3], expected_class='B')
    network.test_pattern([2.5, 1.8, 2.2], expected_class='A')  # New pattern
    network.test_pattern([0.4, 0.2, 0.5], expected_class='B')  # New pattern
    
    # Visualize results
    network.visualize_learning_progress()
    
    print("\n🎉 IMPROVED DEMONSTRATION COMPLETE!")
    print("=" * 70)
    print("✅ This shows what happens in the lab when learning works:")
    print("   • Memristor resistances actually change")
    print("   • Network learns to distinguish patterns")
    print("   • Clear improvement in classification")
    print("   • Adaptive learning overcomes initial failures")

if __name__ == "__main__":
    main()