###########################################################################################################################################
#||||- - - |8.19.2025| - - -                               ||   LIQUID BAYES   ||                               - - - |1990two| - - -|||| #
###########################################################################################################################################
"""
Mathematical Foundation & Conceptual Documentation
-------------------------------------------------

CORE PRINCIPLE:
Combines liquid state machines (continuous neural dynamics) with Bayesian inference
to create adaptive neural systems where probabilistic confidence modulates the
evolution of continuous dynamical states, enabling exploration-exploitation balance.

MATHEMATICAL FOUNDATION:
=======================

1. LIQUID STATE MACHINE DYNAMICS:
   dx/dt = -x/τ + W_rec·σ(x) + W_in·u + b
   
   Where:
   - x: liquid state vector (membrane potentials)
   - τ: time constant (liquid viscosity)
   - W_rec: recurrent connection matrix
   - W_in: input weight matrix
   - σ: nonlinear activation function
   - u: external input signal
   - b: bias terms

2. CONFIDENCE-MODULATED EVOLUTION:
   dx/dt = c·f(x,u) + (1-c)·ε·η
   
   Where:
   - c: Bayesian confidence score ∈ [0,1]
   - f(x,u): standard liquid dynamics
   - ε: exploration rate parameter
   - η: Gaussian noise for exploration
   
   High confidence → smooth, deterministic evolution
   Low confidence → exploratory, stochastic behavior

3. BAYESIAN CONFIDENCE ESTIMATION:
   P(θ|D) ∝ P(D|θ)·P(θ)
   
   Confidence = 1 - H(P(θ|D))
   
   Where:
   - H: entropy function
   - Low entropy → high confidence
   - High entropy → low confidence

4. BELIEF PROPAGATION:
   For Bayesian network with variables X₁...Xₙ:
   
   P(Xi|parents(Xi)) = normalize(evidence(Xi) · ∏ messages)
   
   Iterative message passing for approximate inference.

5. TEMPORAL INTEGRATION:
   x(t+dt) = x(t) + dt·dx/dt
   
   Euler integration for continuous-time dynamics.

CONCEPTUAL REASONING:
====================

WHY LIQUID + BAYESIAN?
- Traditional neural networks lack temporal dynamics
- Liquid state machines provide rich temporal processing
- Bayesian inference quantifies uncertainty in decisions
- Confidence feedback enables adaptive exploration

KEY INNOVATIONS:
1. **Confidence-Modulated Dynamics**: Uncertainty controls exploration
2. **Temporal Bayesian Networks**: Dynamic probabilistic reasoning
3. **Adaptive Time Constants**: Liquid viscosity adapts to confidence
4. **Hierarchical Uncertainty**: Multiple levels of uncertainty quantification
5. **Exploration-Exploitation Balance**: Automatic trade-off via confidence

APPLICATIONS:
- Adaptive control systems with uncertainty quantification
- Time-series prediction with confidence bounds
- Reinforcement learning with liquid state representations
- Robust decision-making under uncertainty
- Continuous learning systems

COMPLEXITY ANALYSIS:
- Liquid Evolution: O(d²) where d = state dimension
- Bayesian Inference: O(n·k²) where n = variables, k = states per variable
- Chain Execution: O(T·(d² + n·k²)) where T = chain steps
- Memory: O(d² + n²·k²) for connection matrices

BIOLOGICAL INSPIRATION:
- Membrane potential dynamics in neural circuits
- Confidence-based neuromodulation (dopamine, norepinephrine)
- Bayesian brain hypothesis for uncertainty processing
- Liquid computing in cortical microcircuits
"""

from __future__ import annotations
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import math
from collections import defaultdict
from typing import List, Dict, Tuple, Optional

SAFE_MIN = -1e6
SAFE_MAX = 1e6
EPS = 1e-8

#||||- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - 𓅸 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -||||#

def make_safe(tensor, min_val=SAFE_MIN, max_val=SAFE_MAX):
    tensor = torch.where(torch.isnan(tensor), torch.tensor(0.0, device=tensor.device, dtype=tensor.dtype), tensor)
    tensor = torch.where(torch.isinf(tensor), torch.tensor(max_val, device=tensor.device, dtype=tensor.dtype), tensor)
    return torch.clamp(tensor, min_val, max_val)

def safe_softmax(x, dim=-1, temperature=1.0):
    x = x.to(dtype=torch.float32)
    x = make_safe(x, min_val=-50, max_val=50)
    # Guard temperature and improve numerical stability
    if isinstance(temperature, torch.Tensor):
        temperature = float(temperature.detach().cpu().item())
    temperature = max(float(temperature), EPS)
    x = x / temperature
    x = x - x.amax(dim=dim, keepdim=True)
    return F.softmax(x, dim=dim)

###########################################################################################################################################
#################################################- - -   LIQUID DYNAMICS CORE   - - -######################################################

class LiquidDynamicsCore(nn.Module):
    """Continuous neural dynamics with confidence-modulated evolution.
    
    Implements liquid state machine dynamics where the evolution of continuous
    neural states is modulated by Bayesian confidence estimates, enabling
    adaptive exploration-exploitation behavior.
    
    Mathematical Details:
    - Standard dynamics: dx/dt = -x/τ + W_rec·σ(x) + W_in·u + b
    - Confidence modulation: dx/dt = c·standard + (1-c)·exploration
    - Euler integration: x(t+dt) = x(t) + dt·dx/dt
    
    The liquid state represents membrane potentials of a continuous neural
    circuit with recurrent connections and external inputs.
    """
    def __init__(self, state_dim, input_dim, liquid_time_constant=1.0):
        super().__init__()
        self.state_dim = state_dim
        self.input_dim = input_dim
        self.liquid_time_constant = nn.Parameter(torch.tensor(liquid_time_constant))
        
        # Liquid dynamics parameters
        self.W_rec = nn.Parameter(torch.randn(state_dim, state_dim) * 0.1)  # Recurrent weights
        self.W_in = nn.Parameter(torch.randn(state_dim, input_dim) * 0.1)   # Input weights
        self.bias = nn.Parameter(torch.zeros(state_dim))
        
        # Nonlinear activation function
        self.activation = nn.Tanh()
        
        # Membrane potential (liquid state)
        self.register_buffer('liquid_state', torch.zeros(1, state_dim))
        
        # Uncertainty injection parameters
        self.noise_scale = nn.Parameter(torch.tensor(0.1))
        self.exploration_rate = nn.Parameter(torch.tensor(0.05))
        
    def reset_state(self, batch_size=1):
        """Reset the liquid state (preserve buffer & device)."""
        with torch.no_grad():
            if self.liquid_state.shape[0] != batch_size:
                self.liquid_state = torch.zeros(
                    batch_size, self.state_dim,
                    device=self.liquid_state.device,
                    dtype=self.liquid_state.dtype,
                )
            else:
                self.liquid_state.zero_()
    
    def evolve_liquid(self, input_signal, confidence_weight=1.0, dt=0.1):
        """Evolve liquid state with confidence-modulated dynamics.
        
        Implements the core liquid state machine evolution with Bayesian
        confidence modulation. High confidence leads to smooth, deterministic
        evolution while low confidence enables exploration through noise injection.
        
        Mathematical Details:
        - Decay term: -x/τ
        - Recurrent term: W_rec·tanh(x)
        - Input term: W_in·u
        - Confidence modulation: c·dynamics + (1-c)·exploration
        
        Args:
            input_signal: External input to liquid [batch_size, input_dim]
            confidence_weight: Confidence score(s) ∈ [0,1] modulating evolution
            dt: Integration time step
            
        Returns:
            Updated liquid state [batch_size, state_dim]
        """
        batch_size = input_signal.shape[0]
        
        if self.liquid_state.shape[0] != batch_size:
            self.reset_state(batch_size)
        
        # Compute liquid dynamics: dx/dt = -x/τ + W_rec*σ(x) + W_in*u + b
        tau = torch.clamp(self.liquid_time_constant, 0.1, 10.0)
        
        # Recurrent dynamics
        recurrent_input = torch.matmul(self.activation(self.liquid_state), self.W_rec.T)
        
        # External input
        external_input = torch.matmul(input_signal, self.W_in.T)
        
        # Combined dynamics
        dynamics = (-self.liquid_state / tau + recurrent_input + external_input + self.bias)
        
        # Confidence-modulated evolution
        if isinstance(confidence_weight, torch.Tensor):
            if confidence_weight.dim() == 1:
                confidence_weight = confidence_weight.unsqueeze(-1)
            confidence_weight = confidence_weight.to(self.liquid_state.dtype)
        else:
            confidence_weight = torch.tensor(confidence_weight, device=self.liquid_state.device, dtype=self.liquid_state.dtype)
        
        # High confidence → smooth evolution, Low confidence → exploration
        exploration_noise = torch.randn_like(self.liquid_state) * self.noise_scale
        exploration_strength = (1.0 - confidence_weight) * self.exploration_rate
        
        modulated_dynamics = confidence_weight * dynamics + exploration_strength * exploration_noise
        
        # Euler integration (keep buffer identity)
        self.liquid_state.add_(dt * make_safe(modulated_dynamics))
        
        return self.liquid_state.clone()
    
    def get_liquid_features(self):
        """Extract multiple feature representations from liquid state.
        
        Returns:
            Dictionary containing:
            - raw_state: Raw membrane potentials
            - activated_state: Nonlinearly activated state
            - state_energy: L2 energy of state vector
            - state_entropy: Entropy of state distribution
        """
        return {
            'raw_state': self.liquid_state.clone(),
            'activated_state': self.activation(self.liquid_state),
            'state_energy': torch.sum(self.liquid_state ** 2, dim=-1, keepdim=True),
            'state_entropy': self._compute_state_entropy()
        }
    
    def _compute_state_entropy(self):
        """Compute entropy of liquid state distribution.
        
        Treats liquid state as a probability distribution (after softmax)
        and computes Shannon entropy: H = -Σ p(x)·log(p(x))
        
        Returns:
            Entropy values [batch_size, 1]
        """
        state_probs = safe_softmax(self.liquid_state, dim=-1, temperature=1.0)
        entropy = -torch.sum(state_probs * torch.log(state_probs + EPS), dim=-1, keepdim=True)
        return entropy

###########################################################################################################################################
############################################- - -   BAYESIAN CONFIDENCE NETWORK   - - -###################################################

class BayesianConfidenceNetwork(nn.Module):
    """Bayesian network for confidence estimation from liquid states.
    
    Implements a simplified Bayesian network that performs probabilistic
    inference over discrete variables extracted from continuous liquid states.
    Uses belief propagation for approximate inference and estimates confidence
    based on posterior entropy.
    
    Mathematical Framework:
    - Variables: X₁, X₂, ..., Xₙ with discrete states
    - Evidence: E(Xi) from liquid state features
    - Conditional probabilities: P(Xi|parents(Xi))
    - Posterior beliefs: P(Xi|evidence)
    - Confidence: 1 - H(P(Xi|evidence))
    """
    def __init__(self, state_dim, num_variables=5, num_states_per_var=3):
        super().__init__()
        self.state_dim = state_dim
        self.num_variables = num_variables
        self.num_states_per_var = num_states_per_var
        
        # Feature extraction from liquid state
        self.feature_extractor = nn.Sequential(
            nn.Linear(state_dim, state_dim * 2),
            nn.LayerNorm(state_dim * 2),
            nn.ReLU(),
            nn.Linear(state_dim * 2, num_variables * num_states_per_var)
        )
        
        # Bayesian network structure (simplified - fully connected for now)
        # Each variable can influence every other variable
        self.conditional_prob_tables = nn.ParameterList([
            nn.Parameter(torch.randn(num_states_per_var, num_states_per_var * (num_variables - 1)) * 0.1)
            for _ in range(num_variables)
        ])
        
        # Prior probabilities for each variable
        self.priors = nn.Parameter(torch.ones(num_variables, num_states_per_var))
        
        # Confidence calibration network
        self.confidence_net = nn.Sequential(
            nn.Linear(num_variables, num_variables * 2),
            nn.ReLU(),
            nn.Linear(num_variables * 2, 1),
            nn.Sigmoid()
        )
        
        # Uncertainty quantification
        self.uncertainty_estimator = nn.Sequential(
            nn.Linear(state_dim, state_dim),
            nn.ReLU(),
            nn.Linear(state_dim, 1),
            nn.Sigmoid()
        )
        
    def extract_variable_beliefs(self, liquid_features):
        """Extract discrete variable beliefs from continuous liquid state.
        
        Maps high-dimensional continuous liquid state to evidence for
        discrete variables in the Bayesian network.
        
        Args:
            liquid_features: Dictionary containing liquid state features
            
        Returns:
            Variable beliefs [batch_size, num_variables, num_states_per_var]
        """
        # Get features from liquid state
        liquid_state = liquid_features['activated_state']
        
        # Extract variable evidence
        evidence = self.feature_extractor(liquid_state)
        evidence = evidence.view(-1, self.num_variables, self.num_states_per_var)
        
        # Convert to probability distributions
        variable_beliefs = safe_softmax(evidence, dim=-1)
        
        return variable_beliefs
    
    def bayesian_inference(self, variable_beliefs):
        """Perform approximate Bayesian inference via belief propagation.
        
        Implements simplified belief propagation algorithm to compute
        posterior beliefs over variables given evidence.
        
        Mathematical Details:
        - Initialize with priors: P(Xi)
        - Iterate belief updates: P(Xi) ← normalize(evidence(Xi) · ∏ messages)
        - Messages based on conditional probability tables
        
        Args:
            variable_beliefs: Evidence for variables [batch_size, num_vars, num_states]
            
        Returns:
            Posterior beliefs [batch_size, num_variables, num_states_per_var]
        """
        batch_size = variable_beliefs.shape[0]
        device = variable_beliefs.device
        
        # Initialize with priors
        current_beliefs = safe_softmax(self.priors.unsqueeze(0).expand(batch_size, -1, -1), dim=-1)
        
        # Iterative belief propagation (simplified)
        for iteration in range(3):  # Few iterations for efficiency
            new_beliefs = current_beliefs.clone()
            
            for var_idx in range(self.num_variables):
                # Get evidence for this variable
                evidence = variable_beliefs[:, var_idx, :]
                
                # Get conditional probabilities from other variables
                if self.num_variables > 1:
                    other_var_beliefs = torch.cat([
                        current_beliefs[:, :var_idx].flatten(1),
                        current_beliefs[:, var_idx+1:].flatten(1)
                    ], dim=1)
                else:
                    other_var_beliefs = torch.zeros(batch_size, 0, device=device)
                
                # Compute conditional probabilities
                if other_var_beliefs.shape[1] > 0:
                    cond_probs = torch.matmul(other_var_beliefs, self.conditional_prob_tables[var_idx].T)
                    cond_probs = safe_softmax(cond_probs, dim=-1)
                else:
                    cond_probs = torch.ones_like(evidence) / self.num_states_per_var
                
                # Combine evidence with conditional probabilities
                combined = evidence * cond_probs
                new_beliefs[:, var_idx, :] = safe_softmax(combined, dim=-1)
            
            current_beliefs = new_beliefs
        
        return current_beliefs
    
    def compute_confidence(self, beliefs, liquid_features):
        """Compute confidence score from Bayesian beliefs and liquid features.
        
        Combines multiple sources of confidence information:
        1. Belief sharpness (low entropy = high confidence)
        2. Neural confidence estimation
        3. Liquid state uncertainty
        
        Mathematical Details:
        - Entropy confidence: 1 - H(beliefs)/H_max
        - Combined confidence: weighted average of sources
        
        Args:
            beliefs: Posterior beliefs [batch_size, num_vars, num_states]
            liquid_features: Dictionary of liquid state features
            
        Returns:
            Confidence scores [batch_size, 1]
        """
        # Confidence based on belief sharpness (low entropy = high confidence)
        belief_entropy = -torch.sum(beliefs * torch.log(beliefs + EPS), dim=-1)
        avg_entropy = belief_entropy.mean(dim=-1, keepdim=True)
        
        # Normalize entropy to confidence (low entropy = high confidence)
        max_entropy = math.log(self.num_states_per_var)
        entropy_confidence = 1.0 - (avg_entropy / max_entropy)
        
        # Additional confidence from neural network
        nn_confidence = self.confidence_net(belief_entropy)
        
        # Uncertainty from liquid state
        liquid_uncertainty = self.uncertainty_estimator(liquid_features['raw_state'])
        state_confidence = 1.0 - liquid_uncertainty
        
        # Combine confidence sources
        total_confidence = 0.4 * entropy_confidence + 0.3 * nn_confidence + 0.3 * state_confidence
        
        return torch.clamp(total_confidence, 0.0, 1.0)
    
    def forward(self, liquid_features):
        """Complete forward pass: feature extraction → inference → confidence.
        
        Args:
            liquid_features: Dictionary containing liquid state features
            
        Returns:
            Dictionary containing:
            - beliefs: Posterior beliefs over variables
            - confidence: Overall confidence score
            - variable_beliefs: Raw variable evidence
        """
        # Extract variable beliefs from liquid state
        variable_beliefs = self.extract_variable_beliefs(liquid_features)
        
        # Perform Bayesian inference
        posterior_beliefs = self.bayesian_inference(variable_beliefs)
        
        # Compute confidence
        confidence = self.compute_confidence(posterior_beliefs, liquid_features)
        
        return {
            'beliefs': posterior_beliefs,
            'confidence': confidence,
            'variable_beliefs': variable_beliefs
        }

###########################################################################################################################################
############################################- - -   LIQUID BAYES CHAIN   - - -############################################################

class LiquidBayesChain(nn.Module):
    """Complete Liquid-Bayes system with iterative refinement chain.
    
    Implements the full Liquid-Bayes architecture where liquid dynamics
    and Bayesian confidence estimation form a feedback loop over multiple
    chain steps, enabling progressive refinement of predictions.
    
    Architecture:
    1. Liquid evolution (confidence-modulated)
    2. Bayesian confidence estimation
    3. Feedback to liquid dynamics
    4. Repeat for multiple chain steps
    5. Final prediction with uncertainty quantification
    
    The chain allows the system to iteratively improve its predictions
    by using confidence estimates to guide further exploration or exploitation.
    """
    def __init__(self, input_dim, state_dim, output_dim, num_chain_steps=3):
        super().__init__()
        self.input_dim = input_dim
        self.state_dim = state_dim
        self.output_dim = output_dim
        self.num_chain_steps = num_chain_steps
        
        # Core components
        self.liquid_core = LiquidDynamicsCore(state_dim, input_dim)
        self.bayesian_confidence = BayesianConfidenceNetwork(state_dim)
        
        # Final prediction network
        self.final_predictor = nn.Sequential(
            nn.Linear(state_dim, state_dim * 2),
            nn.LayerNorm(state_dim * 2),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(state_dim * 2, output_dim)
        )
        
        # Final Bayesian uncertainty estimation
        self.final_bayesian = BayesianConfidenceNetwork(output_dim, num_variables=3, num_states_per_var=4)
        
        # Chain step weights (learnable importance of each step)
        self.step_weights = nn.Parameter(torch.ones(num_chain_steps))
        
    def single_chain_step(self, input_signal, step_idx=0):
        """Execute one step of the Liquid-Bayes feedback chain.
        
        Each chain step consists of:
        1. Evolve liquid dynamics (with confidence modulation if not first step)
        2. Extract features from liquid state
        3. Perform Bayesian confidence assessment
        
        Args:
            input_signal: External input to liquid [batch_size, input_dim]
            step_idx: Current step index in chain
            
        Returns:
            Dictionary containing step outputs and intermediate states
        """
        # Step 1: Evolve liquid state
        if step_idx == 0:
            # First step - no confidence modulation
            liquid_state = self.liquid_core.evolve_liquid(input_signal, confidence_weight=1.0)
        else:
            # Get confidence from previous step
            liquid_features = self.liquid_core.get_liquid_features()
            bayes_output = self.bayesian_confidence(liquid_features)
            confidence = bayes_output['confidence']
            
            # Step 2: Confidence-modulated liquid evolution
            liquid_state = self.liquid_core.evolve_liquid(input_signal, confidence_weight=confidence)
        
        # Get updated liquid features
        liquid_features = self.liquid_core.get_liquid_features()
        
        # Step 3: Bayesian confidence assessment
        bayes_output = self.bayesian_confidence(liquid_features)
        
        return {
            'liquid_state': liquid_state,
            'liquid_features': liquid_features,
            'bayes_output': bayes_output,
            'confidence': bayes_output['confidence']
        }
    
    def forward(self, input_signal, return_chain_states=False):
        """Execute complete Liquid-Bayes chain with iterative refinement.
        
        Args:
            input_signal: Input to process [batch_size, input_dim]
            return_chain_states: Whether to return intermediate chain states
            
        Returns:
            Dictionary containing:
            - prediction: Final output predictions
            - final_confidence: Weighted confidence across chain
            - final_beliefs: Final Bayesian beliefs
            - prediction_uncertainty: Uncertainty in predictions
            - chain_states: Intermediate states (if requested)
        """
        batch_size = input_signal.shape[0]
        
        # Reset liquid state
        self.liquid_core.reset_state(batch_size)
        
        # Store chain states for analysis
        chain_states = []
        
        # Execute chain steps
        for step in range(self.num_chain_steps):
            step_output = self.single_chain_step(input_signal, step_idx=step)
            step_output['step_idx'] = step
            chain_states.append(step_output)
        
        # Final prediction from last liquid state
        final_liquid_state = chain_states[-1]['liquid_features']['activated_state']
        prediction_logits = self.final_predictor(final_liquid_state)
        
        # Final Bayesian uncertainty quantification
        prediction_features = {
            'raw_state': prediction_logits,
            'activated_state': torch.tanh(prediction_logits)
        }
        final_bayes = self.final_bayesian(prediction_features)
        
        # Weighted combination of confidence scores across chain
        step_weights = safe_softmax(self.step_weights, dim=0)
        weighted_confidence = sum(
            step_weights[i] * chain_states[i]['confidence'] 
            for i in range(self.num_chain_steps)
        )
        
        output = {
            'prediction': prediction_logits,
            'final_confidence': weighted_confidence,
            'final_beliefs': final_bayes['beliefs'],
            'prediction_uncertainty': 1.0 - final_bayes['confidence']
        }
        
        if return_chain_states:
            output['chain_states'] = chain_states
        
        return output
    
    def predict_with_uncertainty(self, input_signal):
        """Make predictions with comprehensive uncertainty quantification.
        
        Provides detailed uncertainty analysis including:
        - Final prediction confidence
        - Chain-step progression
        - Liquid state entropy evolution
        
        Args:
            input_signal: Input to process [batch_size, input_dim]
            
        Returns:
            Dictionary with comprehensive uncertainty information
        """
        output = self.forward(input_signal, return_chain_states=True)
        
        # Extract uncertainty information
        uncertainty_info = {
            'prediction': output['prediction'],
            'confidence': output['final_confidence'],
            'prediction_uncertainty': output['prediction_uncertainty'],
            'chain_confidences': [state['confidence'] for state in output['chain_states']],
            'liquid_entropies': [state['liquid_features']['state_entropy'] for state in output['chain_states']]
        }
        
        return uncertainty_info

###########################################################################################################################################
##################################################- - -   DEMO AND TESTING   - - -#########################################################

def test_liquid_bayes_chain():
    print(" Testing Liquid Bayes Chain - Probabilistic Control of Continuous Dynamics")
    print("=" * 80)
    
    # Create Liquid-Bayes chain
    input_dim = 32
    state_dim = 64
    output_dim = 10
    
    model = LiquidBayesChain(
        input_dim=input_dim,
        state_dim=state_dim,
        output_dim=output_dim,
        num_chain_steps=4
    )
    
    print(f"Created Liquid-Bayes Chain:")
    print(f"  - Input dimension: {input_dim}")
    print(f"  - Liquid state dimension: {state_dim}")
    print(f"  - Output dimension: {output_dim}")
    print(f"  - Chain steps: {model.num_chain_steps}")
    
    # Generate test data
    batch_size = 8
    test_input = torch.randn(batch_size, input_dim)
    
    print(f"\nTesting with batch size: {batch_size}")
    
    # Test forward pass
    print("\nExecuting Liquid-Bayes chain...")
    output = model(test_input, return_chain_states=True)
    
    print("Chain execution results:")
    print(f"  - Final prediction shape: {output['prediction'].shape}")
    print(f"  - Average confidence: {output['final_confidence'].mean():.3f}")
    print(f"  - Average uncertainty: {output['prediction_uncertainty'].mean():.3f}")
    
    # Analyze chain progression
    print("\nChain step analysis:")
    for i, state in enumerate(output['chain_states']):
        conf = state['confidence'].mean().item()
        entropy = state['liquid_features']['state_entropy'].mean().item()
        print(f"  Step {i+1}: Confidence={conf:.3f}, Liquid Entropy={entropy:.3f}")
    
    # Test uncertainty prediction
    print("\nTesting uncertainty quantification...")
    uncertainty_output = model.predict_with_uncertainty(test_input[:3])
    
    print("Uncertainty analysis:")
    for i in range(3):
        conf = uncertainty_output['confidence'][i].item()
        pred_unc = uncertainty_output['prediction_uncertainty'][i].item()
        print(f"  Sample {i+1}: Confidence={conf:.3f}, Prediction Uncertainty={pred_unc:.3f}")
    
    # Test adaptive behavior
    print("\nTesting adaptive behavior with different inputs...")
    
    # High-confidence input (structured pattern)
    structured_input = torch.ones(1, input_dim) * 0.5
    struct_output = model(structured_input)
    struct_conf = struct_output['final_confidence'].item()
    
    # Low-confidence input (random noise)
    noisy_input = torch.randn(1, input_dim) * 2.0
    noisy_output = model(noisy_input)
    noisy_conf = noisy_output['final_confidence'].item()
    
    print(f"  Structured input confidence: {struct_conf:.3f}")
    print(f"  Noisy input confidence: {noisy_conf:.3f}")
    print(f"  Confidence difference: {abs(struct_conf - noisy_conf):.3f}")
    
    print("\nLiquid-Bayes Chain test completed!")
    print("✓ Liquid dynamics evolve with Bayesian confidence modulation")
    print("✓ Probabilistic feedback loop controls exploration vs exploitation")
    print("✓ Full uncertainty quantification throughout the chain")
    print("✓ Adaptive behavior based on input characteristics")
    
    return True

def confidence_modulation_demo():
    """Demonstrate how Bayesian confidence modulates liquid evolution."""
    print("\n" + "="*60)
    print(" CONFIDENCE MODULATION DEMO")
    print("="*60)
    
    # Create simplified model
    model = LiquidBayesChain(input_dim=16, state_dim=32, output_dim=5, num_chain_steps=3)
    
    # Test with controlled confidence scenarios
    scenarios = [
        ("High Confidence Input", torch.ones(1, 16) * 0.3),  # Structured
        ("Medium Confidence Input", torch.randn(1, 16) * 0.5),  # Moderate noise
        ("Low Confidence Input", torch.randn(1, 16) * 2.0),  # High noise
    ]
    
    print("Testing confidence-driven adaptation:")
    
    for name, test_input in scenarios:
        output = model(test_input, return_chain_states=True)
        
        # Extract confidence progression
        confidences = [state['confidence'].item() for state in output['chain_states']]
        entropies = [state['liquid_features']['state_entropy'].item() for state in output['chain_states']]
        
        print(f"\n{name}:")
        print(f"  Chain confidences: {[f'{c:.3f}' for c in confidences]}")
        print(f"  Liquid entropies: {[f'{e:.3f}' for e in entropies]}")
        print(f"  Final confidence: {output['final_confidence'].item():.3f}")
    
    print("\n Demo shows how liquid dynamics adapt based on Bayesian confidence!")
    print("   High confidence → stable evolution, Low confidence → exploration")

if __name__ == "__main__":
    test_liquid_bayes_chain()
    confidence_modulation_demo()

###########################################################################################################################################
###########################################################################################################################################