saorsa_core/adaptive/

learning.rs

// Copyright 2024 Saorsa Labs Limited
//
// This software is dual-licensed under:
// - GNU Affero General Public License v3.0 or later (AGPL-3.0-or-later)
// - Commercial License
//
// For AGPL-3.0 license, see LICENSE-AGPL-3.0
// For commercial licensing, contact: saorsalabs@gmail.com
//
// Unless required by applicable law or agreed to in writing, software
// distributed under these licenses is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

//! Machine learning subsystems for adaptive behavior
//!
//! Includes Thompson Sampling for routing optimization, Q-learning for cache management,
//! and LSTM for churn prediction

use super::beta_distribution::BetaDistribution;
use super::*;
use rand::Rng;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;

/// Thompson Sampling for routing strategy optimization
///
/// Uses Beta distributions to model success rates for each routing strategy
/// per content type, automatically balancing exploration and exploitation
pub struct ThompsonSampling {
    /// Beta distributions for each (content type, strategy) pair
    /// Beta(α, β) where α = successes + 1, β = failures + 1
    arms: Arc<RwLock<HashMap<(ContentType, StrategyChoice), BetaParams>>>,

    /// Minimum number of samples before considering a strategy reliable
    min_samples: u32,

    /// Decay factor for old observations (0.0-1.0)
    decay_factor: f64,

    /// Performance metrics
    metrics: Arc<RwLock<RoutingMetrics>>,
}

/// Beta distribution parameters with proper distribution
#[derive(Debug, Clone)]
struct BetaParams {
    /// Beta distribution instance
    distribution: BetaDistribution,
    /// Total number of trials
    trials: u32,
    /// Last update timestamp
    last_update: std::time::Instant,
}

impl Default for BetaParams {
    fn default() -> Self {
        let distribution = BetaDistribution::new(1.0, 1.0).unwrap_or(BetaDistribution {
            alpha: 1.0,
            beta: 1.0,
        });
        Self {
            distribution,
            trials: 0,
            last_update: std::time::Instant::now(),
        }
    }
}

/// Routing performance metrics
#[derive(Debug, Default, Clone)]
pub struct RoutingMetrics {
    /// Total routing decisions made
    pub total_decisions: u64,
    /// Decisions per content type
    pub decisions_by_type: HashMap<ContentType, u64>,
    /// Success rate per strategy
    pub strategy_success_rates: HashMap<StrategyChoice, f64>,
    /// Average latency per strategy (ms)
    pub strategy_latencies: HashMap<StrategyChoice, f64>,
}

impl Default for ThompsonSampling {
    fn default() -> Self {
        Self::new()
    }
}

impl ThompsonSampling {
    /// Create a new Thompson Sampling instance
    pub fn new() -> Self {
        Self {
            arms: Arc::new(RwLock::new(HashMap::new())),
            min_samples: 10,
            decay_factor: 0.99,
            metrics: Arc::new(RwLock::new(RoutingMetrics::default())),
        }
    }

    /// Select optimal routing strategy for given content type
    pub async fn select_strategy(&self, content_type: ContentType) -> Result<StrategyChoice> {
        let mut arms = self.arms.write().await;
        let mut metrics = self.metrics.write().await;

        metrics.total_decisions += 1;
        *metrics.decisions_by_type.entry(content_type).or_insert(0) += 1;

        let strategies = vec![
            StrategyChoice::Kademlia,
            StrategyChoice::Hyperbolic,
            StrategyChoice::TrustPath,
            StrategyChoice::SOMRegion,
        ];

        let mut best_strategy = StrategyChoice::Kademlia;
        let mut best_sample = 0.0;

        // Sample from each arm's Beta distribution
        for strategy in &strategies {
            let key = (content_type, *strategy);
            let params = arms.entry(key).or_default();

            // Apply decay to old observations
            if params.trials > 0 {
                let elapsed = params.last_update.elapsed().as_secs() as f64;
                let decay = self.decay_factor.powf(elapsed / 3600.0); // Hourly decay
                let alpha = params.distribution.alpha;
                let beta = params.distribution.beta;
                let new_alpha = 1.0 + (alpha - 1.0) * decay;
                let new_beta = 1.0 + (beta - 1.0) * decay;
                params.distribution =
                    BetaDistribution::new(new_alpha, new_beta).unwrap_or(BetaDistribution {
                        alpha: new_alpha.max(f64::MIN_POSITIVE),
                        beta: new_beta.max(f64::MIN_POSITIVE),
                    });
            }

            // Sample from Beta distribution using proper implementation
            let mut rng = rand::thread_rng();
            let sample = params.distribution.sample(&mut rng);

            // Add exploration bonus for under-sampled strategies
            let exploration_bonus = if params.trials < self.min_samples {
                0.1 * (1.0 - (params.trials as f64 / self.min_samples as f64))
            } else {
                0.0
            };

            let adjusted_sample = sample + exploration_bonus;

            if adjusted_sample > best_sample {
                best_sample = adjusted_sample;
                best_strategy = *strategy;
            }
        }

        Ok(best_strategy)
    }

    /// Update strategy performance based on outcome
    pub async fn update(
        &self,
        _content_type: ContentType,
        strategy: StrategyChoice,
        success: bool,
        latency_ms: u64,
    ) -> anyhow::Result<()> {
        let mut arms = self.arms.write().await;
        let mut metrics = self.metrics.write().await;

        let key = (_content_type, strategy);
        let params = arms.entry(key).or_default();

        // Update Beta parameters
        params.distribution.update(success);
        params.trials += 1;
        params.last_update = std::time::Instant::now();

        // Update success rate (exponential moving average)
        let success_rate = params.distribution.mean();
        let current_rate = metrics
            .strategy_success_rates
            .entry(strategy)
            .or_insert(0.5);
        *current_rate = 0.9 * (*current_rate) + 0.1 * success_rate;

        // Update latency (exponential moving average)
        let current_latency = metrics
            .strategy_latencies
            .entry(strategy)
            .or_insert(latency_ms as f64);
        *current_latency = 0.9 * (*current_latency) + 0.1 * (latency_ms as f64);

        Ok(())
    }

    /// Get current performance metrics
    pub async fn get_metrics(&self) -> RoutingMetrics {
        self.metrics.read().await.clone()
    }

    /// Get confidence interval for a strategy's success rate
    pub async fn get_confidence_interval(
        &self,
        _content_type: ContentType,
        strategy: StrategyChoice,
    ) -> (f64, f64) {
        let arms = self.arms.read().await;
        let key = (_content_type, strategy);

        if let Some(params) = arms.get(&key) {
            if params.trials == 0 {
                return (0.0, 1.0);
            }

            // Use the Beta distribution's confidence interval method
            params.distribution.confidence_interval()
        } else {
            (0.0, 1.0)
        }
    }

    /// Reset statistics for a specific strategy
    pub async fn reset_strategy(&self, _content_type: ContentType, strategy: StrategyChoice) {
        let mut arms = self.arms.write().await;
        arms.remove(&(_content_type, strategy));
    }
}

#[async_trait]
impl LearningSystem for ThompsonSampling {
    async fn select_strategy(&self, context: &LearningContext) -> StrategyChoice {
        self.select_strategy(context.content_type)
            .await
            .unwrap_or(StrategyChoice::Kademlia)
    }

    async fn update(
        &mut self,
        context: &LearningContext,
        choice: &StrategyChoice,
        outcome: &Outcome,
    ) {
        let _ = ThompsonSampling::update(
            self,
            context.content_type,
            *choice,
            outcome.success,
            outcome.latency_ms,
        )
        .await;
    }

    async fn metrics(&self) -> LearningMetrics {
        let metrics = self.get_metrics().await;

        LearningMetrics {
            total_decisions: metrics.total_decisions,
            success_rate: metrics.strategy_success_rates.values().sum::<f64>()
                / metrics.strategy_success_rates.len().max(1) as f64,
            avg_latency_ms: metrics.strategy_latencies.values().sum::<f64>()
                / metrics.strategy_latencies.len().max(1) as f64,
            strategy_performance: metrics.strategy_success_rates.clone(),
        }
    }
}

/// Cache statistics
#[derive(Debug, Clone, Default)]
pub struct CacheStats {
    /// Total cache hits
    pub hits: u64,

    /// Total cache misses
    pub misses: u64,

    /// Current cache size in bytes
    pub size_bytes: u64,

    /// Number of items in cache
    pub item_count: u64,

    /// Total evictions
    pub evictions: u64,

    /// Cache hit rate
    pub hit_rate: f64,
}

/// Q-Learning cache manager
pub struct QLearnCacheManager {
    /// Q-table mapping states to action values
    q_table: Arc<tokio::sync::RwLock<HashMap<CacheState, HashMap<CacheAction, f64>>>>,

    /// Learning rate
    learning_rate: f64,

    /// Discount factor
    discount_factor: f64,

    /// Exploration rate (epsilon)
    epsilon: f64,

    /// Cache storage
    cache: Arc<tokio::sync::RwLock<HashMap<ContentHash, CachedContent>>>,

    /// Cache capacity in bytes
    capacity: usize,

    /// Current cache size
    current_size: Arc<std::sync::atomic::AtomicUsize>,

    /// Request statistics for popularity tracking
    request_stats: Arc<tokio::sync::RwLock<HashMap<ContentHash, RequestStats>>>,

    /// Hit/miss statistics
    hit_count: Arc<std::sync::atomic::AtomicU64>,
    miss_count: Arc<std::sync::atomic::AtomicU64>,

    /// Bandwidth tracking
    _bandwidth_used: Arc<std::sync::atomic::AtomicU64>,
}

/// Request statistics for tracking content popularity
#[derive(Debug, Clone)]
pub struct RequestStats {
    /// Total number of requests
    request_count: u64,
    /// Requests in the last hour
    hourly_requests: u64,
    /// Last request timestamp
    last_request: std::time::Instant,
    /// Content size
    content_size: usize,
}

/// Cache state representation
#[derive(Debug, Clone, Hash, PartialEq, Eq)]
pub struct CacheState {
    /// Cache utilization (0-10 buckets)
    utilization_bucket: u8,

    /// Request rate bucket (0-10, bucketed hourly rate)
    request_rate_bucket: u8,

    /// Content popularity score (0-10)
    content_popularity: u8,

    /// Content size bucket (0-10, logarithmic scale)
    size_bucket: u8,
}

/// Actions the cache manager can take
#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq)]
pub enum CacheAction {
    Cache,
    Evict(EvictionPolicy),
    IncreaseReplication,
    DecreaseReplication,
    NoAction,
}

/// Eviction policies
#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq)]
pub enum EvictionPolicy {
    LRU,
    LFU,
    Random,
}

/// Cached content metadata
#[derive(Debug, Clone)]
pub struct CachedContent {
    pub data: Vec<u8>,
    pub access_count: u64,
    pub last_access: std::time::Instant,
    pub insertion_time: std::time::Instant,
}

impl QLearnCacheManager {
    /// Create a new Q-learning cache manager
    pub fn new(capacity: usize) -> Self {
        Self {
            q_table: Arc::new(tokio::sync::RwLock::new(HashMap::new())),
            learning_rate: 0.1,
            discount_factor: 0.9,
            epsilon: 0.1,
            cache: Arc::new(tokio::sync::RwLock::new(HashMap::new())),
            capacity,
            current_size: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
            request_stats: Arc::new(tokio::sync::RwLock::new(HashMap::new())),
            hit_count: Arc::new(std::sync::atomic::AtomicU64::new(0)),
            miss_count: Arc::new(std::sync::atomic::AtomicU64::new(0)),
            _bandwidth_used: Arc::new(std::sync::atomic::AtomicU64::new(0)),
        }
    }

    /// Decide what action to take for a content request
    pub async fn decide_action(&self, content_hash: &ContentHash) -> CacheAction {
        let state = self.get_current_state(content_hash);

        if rand::random::<f64>() < self.epsilon {
            // Explore: random action
            self.random_action()
        } else {
            // Exploit: best known action
            let q_table = self.q_table.read().await;
            q_table
                .get(&state)
                .and_then(|actions| {
                    actions
                        .iter()
                        .max_by(|(_, a), (_, b)| {
                            a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)
                        })
                        .map(|(action, _)| *action)
                })
                .unwrap_or(CacheAction::NoAction)
        }
    }

    /// Update Q-value based on action outcome
    pub async fn update_q_value(
        &self,
        state: CacheState,
        action: CacheAction,
        reward: f64,
        next_state: CacheState,
    ) {
        let mut q_table = self.q_table.write().await;

        let current_q = q_table
            .entry(state.clone())
            .or_insert_with(HashMap::new)
            .get(&action)
            .copied()
            .unwrap_or(0.0);

        let max_next_q = q_table
            .get(&next_state)
            .and_then(|actions| {
                actions
                    .values()
                    .max_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
                    .copied()
            })
            .unwrap_or(0.0);

        let new_q = current_q
            + self.learning_rate * (reward + self.discount_factor * max_next_q - current_q);

        q_table
            .entry(state)
            .or_insert_with(HashMap::new)
            .insert(action, new_q);
    }

    /// Get current cache state
    fn get_current_state(&self, _content_hash: &ContentHash) -> CacheState {
        let utilization = (self.current_size.load(std::sync::atomic::Ordering::Relaxed) * 10
            / self.capacity) as u8;

        // Get request stats synchronously (we'll need to handle this properly)
        // For now, using placeholder values that will be updated in a future method
        let (request_rate_bucket, content_popularity, size_bucket) = (5, 5, 5);

        CacheState {
            utilization_bucket: utilization.min(10),
            request_rate_bucket,
            content_popularity,
            size_bucket,
        }
    }

    /// Get current cache state asynchronously with full stats
    pub async fn get_current_state_async(&self, content_hash: &ContentHash) -> CacheState {
        let utilization = (self.current_size.load(std::sync::atomic::Ordering::Relaxed) * 10
            / self.capacity) as u8;

        let stats = self.request_stats.read().await;
        let (request_rate_bucket, content_popularity, size_bucket) =
            if let Some(stat) = stats.get(content_hash) {
                // Calculate hourly request rate bucket (0-10)
                let hourly_rate = stat.hourly_requests.min(100) / 10;

                // Calculate popularity (0-10) based on total requests
                let popularity = (stat.request_count.min(1000) / 100) as u8;

                // Calculate size bucket (logarithmic scale)
                let size_bucket = match stat.content_size {
                    0..=1_024 => 0,                   // 1KB
                    1_025..=10_240 => 1,              // 10KB
                    10_241..=102_400 => 2,            // 100KB
                    102_401..=1_048_576 => 3,         // 1MB
                    1_048_577..=10_485_760 => 4,      // 10MB
                    10_485_761..=104_857_600 => 5,    // 100MB
                    104_857_601..=1_073_741_824 => 6, // 1GB
                    _ => 7,                           // >1GB
                };

                (hourly_rate as u8, popularity, size_bucket)
            } else {
                (0, 0, 0) // Unknown content
            };

        CacheState {
            utilization_bucket: utilization.min(10),
            request_rate_bucket,
            content_popularity,
            size_bucket,
        }
    }

    /// Get a random action
    fn random_action(&self) -> CacheAction {
        match rand::random::<u8>() % 5 {
            0 => CacheAction::Cache,
            1 => CacheAction::Evict(EvictionPolicy::LRU),
            2 => CacheAction::Evict(EvictionPolicy::LFU),
            3 => CacheAction::IncreaseReplication,
            4 => CacheAction::DecreaseReplication,
            _ => CacheAction::NoAction,
        }
    }

    /// Insert content into cache
    pub async fn insert(&self, hash: ContentHash, data: Vec<u8>) -> bool {
        let size = data.len();

        // Check if we need to evict
        while self.current_size.load(std::sync::atomic::Ordering::Relaxed) + size > self.capacity {
            if !self.evict_one().await {
                return false;
            }
        }

        let mut cache = self.cache.write().await;
        cache.insert(
            hash,
            CachedContent {
                data,
                access_count: 0,
                last_access: std::time::Instant::now(),
                insertion_time: std::time::Instant::now(),
            },
        );

        self.current_size
            .fetch_add(size, std::sync::atomic::Ordering::Relaxed);
        true
    }

    /// Evict one item from cache
    async fn evict_one(&self) -> bool {
        // Simple LRU eviction for now
        let mut cache = self.cache.write().await;
        let oldest = cache
            .iter()
            .min_by_key(|(_, content)| content.last_access)
            .map(|(k, _)| *k);

        if let Some(key) = oldest
            && let Some(value) = cache.remove(&key)
        {
            self.current_size
                .fetch_sub(value.data.len(), std::sync::atomic::Ordering::Relaxed);
            return true;
        }

        false
    }

    /// Get content from cache
    pub async fn get(&self, hash: &ContentHash) -> Option<Vec<u8>> {
        let cache_result = {
            let mut cache = self.cache.write().await;
            cache.get_mut(hash).map(|entry| {
                entry.access_count += 1;
                entry.last_access = std::time::Instant::now();
                entry.data.clone()
            })
        };

        // Update statistics
        if cache_result.is_some() {
            self.hit_count
                .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
        } else {
            self.miss_count
                .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
        }

        // Update request stats
        if let Some(ref data) = cache_result {
            let mut stats = self.request_stats.write().await;
            let stat = stats.entry(*hash).or_insert_with(|| RequestStats {
                request_count: 0,
                hourly_requests: 0,
                last_request: std::time::Instant::now(),
                content_size: data.len(),
            });
            stat.request_count += 1;
            stat.hourly_requests += 1; // In real implementation, this would decay over time
            stat.last_request = std::time::Instant::now();
        }

        cache_result
    }

    /// Calculate reward based on action outcome
    pub fn calculate_reward(&self, action: CacheAction, hit: bool, bandwidth_cost: u64) -> f64 {
        let hits = self.hit_count.load(std::sync::atomic::Ordering::Relaxed) as f64;
        let misses = self.miss_count.load(std::sync::atomic::Ordering::Relaxed) as f64;
        let hit_rate = if hits + misses > 0.0 {
            hits / (hits + misses)
        } else {
            0.0
        };

        // Storage cost (normalized by capacity)
        let storage_cost = self.current_size.load(std::sync::atomic::Ordering::Relaxed) as f64
            / self.capacity as f64;

        // Bandwidth cost (normalized)
        let bandwidth_cost_normalized = bandwidth_cost as f64 / 1_000_000.0; // Per MB

        // Reward function: R = hit_rate - storage_cost - bandwidth_cost
        match action {
            CacheAction::Cache => {
                if hit {
                    hit_rate - storage_cost * 0.1 - bandwidth_cost_normalized * 0.01
                } else {
                    -0.1 - bandwidth_cost_normalized * 0.1 // Penalty for caching unused content
                }
            }
            CacheAction::Evict(_) => {
                if hit {
                    -0.5 // Penalty for evicting needed content
                } else {
                    0.1 - storage_cost * 0.05 // Small reward for freeing space
                }
            }
            CacheAction::IncreaseReplication => {
                if hit {
                    hit_rate * 0.5 - bandwidth_cost_normalized * 0.2
                } else {
                    -0.2 - bandwidth_cost_normalized * 0.2
                }
            }
            CacheAction::DecreaseReplication => {
                if hit {
                    -0.3 // Penalty if content was needed
                } else {
                    0.05 + storage_cost * 0.05 // Small reward for saving resources
                }
            }
            CacheAction::NoAction => {
                hit_rate * 0.1 - storage_cost * 0.01 // Neutral reward
            }
        }
    }

    /// Execute a cache action
    pub async fn execute_action(
        &self,
        hash: &ContentHash,
        action: CacheAction,
        data: Option<Vec<u8>>,
    ) -> Result<()> {
        match action {
            CacheAction::Cache => {
                if let Some(content) = data {
                    self.insert(*hash, content).await;
                }
            }
            CacheAction::Evict(policy) => {
                match policy {
                    EvictionPolicy::LRU => self.evict_lru().await,
                    EvictionPolicy::LFU => self.evict_lfu().await,
                    EvictionPolicy::Random => self.evict_random().await,
                };
            }
            CacheAction::IncreaseReplication => {
                // In a real implementation, this would trigger replication to more nodes
                // For now, just track the decision
            }
            CacheAction::DecreaseReplication => {
                // In a real implementation, this would reduce replication factor
                // For now, just track the decision
            }
            CacheAction::NoAction => {
                // Do nothing
            }
        }
        Ok(())
    }

    /// Evict using LRU policy
    async fn evict_lru(&self) -> bool {
        self.evict_one().await
    }

    /// Evict using LFU policy
    async fn evict_lfu(&self) -> bool {
        let mut cache = self.cache.write().await;
        let least_frequent = cache
            .iter()
            .min_by_key(|(_, content)| content.access_count)
            .map(|(k, _)| *k);

        if let Some(key) = least_frequent
            && let Some(value) = cache.remove(&key)
        {
            self.current_size
                .fetch_sub(value.data.len(), std::sync::atomic::Ordering::Relaxed);
            return true;
        }
        false
    }

    /// Evict random item
    async fn evict_random(&self) -> bool {
        let cache = self.cache.read().await;
        if cache.is_empty() {
            return false;
        }

        let random_idx = rand::random::<usize>() % cache.len();
        let random_key = cache.keys().nth(random_idx).cloned();
        drop(cache);

        if let Some(key) = random_key {
            let mut cache = self.cache.write().await;
            if let Some(value) = cache.remove(&key) {
                self.current_size
                    .fetch_sub(value.data.len(), std::sync::atomic::Ordering::Relaxed);
                return true;
            }
        }
        false
    }

    /// Get cache statistics
    pub fn get_stats(&self) -> CacheStats {
        let hits = self.hit_count.load(std::sync::atomic::Ordering::Relaxed);
        let misses = self.miss_count.load(std::sync::atomic::Ordering::Relaxed);
        let _hit_rate = if hits + misses > 0 {
            hits as f64 / (hits + misses) as f64
        } else {
            0.0
        };

        CacheStats {
            hits,
            misses,
            size_bytes: self.current_size.load(std::sync::atomic::Ordering::Relaxed) as u64,
            item_count: 0, // TODO: Track number of items
            evictions: 0,  // TODO: Track evictions
            hit_rate: if hits + misses > 0 {
                hits as f64 / (hits + misses) as f64
            } else {
                0.0
            },
        }
    }

    /// Decide whether to cache content based on Q-learning
    pub async fn decide_caching(
        &self,
        hash: ContentHash,
        data: Vec<u8>,
        _content_type: ContentType,
    ) -> Result<()> {
        let _state = self.get_current_state_async(&hash).await;
        let action = self.decide_action(&hash).await;

        if matches!(action, CacheAction::Cache) {
            let _ = self.insert(hash, data).await;
        }

        Ok(())
    }

    /// Get cache statistics asynchronously
    pub async fn get_stats_async(&self) -> CacheStats {
        let cache = self.cache.read().await;
        let hit_count = self.hit_count.load(std::sync::atomic::Ordering::Relaxed);
        let miss_count = self.miss_count.load(std::sync::atomic::Ordering::Relaxed);
        let total = hit_count + miss_count;

        CacheStats {
            hits: hit_count,
            misses: miss_count,
            size_bytes: self.current_size.load(std::sync::atomic::Ordering::Relaxed) as u64,
            item_count: cache.len() as u64,
            evictions: 0, // TODO: Track evictions
            hit_rate: if total > 0 {
                hit_count as f64 / total as f64
            } else {
                0.0
            },
        }
    }
}

/// Node behavior features for churn prediction
#[derive(Debug, Clone)]
pub struct NodeFeatures {
    /// Online duration in seconds
    pub online_duration: f64,
    /// Average response time in milliseconds
    pub avg_response_time: f64,
    /// Resource contribution score (0-1)
    pub resource_contribution: f64,
    /// Messages per hour
    pub message_frequency: f64,
    /// Hour of day (0-23)
    pub time_of_day: f64,
    /// Day of week (0-6)
    pub day_of_week: f64,
    /// Historical reliability score (0-1)
    pub historical_reliability: f64,
    /// Number of disconnections in past week
    pub recent_disconnections: f64,
    /// Average session length in hours
    pub avg_session_length: f64,
    /// Connection stability score (0-1)
    pub connection_stability: f64,
}

/// Feature history for pattern analysis
#[derive(Debug, Clone)]
pub struct FeatureHistory {
    /// Node ID
    pub node_id: NodeId,
    /// Feature snapshots over time
    pub snapshots: Vec<(std::time::Instant, NodeFeatures)>,
    /// Session history (start, end)
    pub sessions: Vec<(std::time::Instant, Option<std::time::Instant>)>,
    /// Total uptime in seconds
    pub total_uptime: u64,
    /// Total downtime in seconds
    pub total_downtime: u64,
}

impl Default for FeatureHistory {
    fn default() -> Self {
        Self::new()
    }
}

impl FeatureHistory {
    /// Create a new feature history
    pub fn new() -> Self {
        Self {
            node_id: NodeId { hash: [0u8; 32] },
            snapshots: Vec::new(),
            sessions: Vec::new(),
            total_uptime: 0,
            total_downtime: 0,
        }
    }
}

/// LSTM-based churn predictor
#[derive(Debug)]
pub struct ChurnPredictor {
    /// Prediction cache
    prediction_cache: Arc<tokio::sync::RwLock<HashMap<NodeId, ChurnPrediction>>>,

    /// Feature history for each node
    feature_history: Arc<tokio::sync::RwLock<HashMap<NodeId, FeatureHistory>>>,

    /// Model parameters (simulated LSTM weights)
    model_weights: Arc<tokio::sync::RwLock<ModelWeights>>,

    /// Experience replay buffer for online learning
    experience_buffer: Arc<tokio::sync::RwLock<Vec<TrainingExample>>>,

    /// Maximum buffer size
    max_buffer_size: usize,

    /// Update frequency
    _update_interval: std::time::Duration,
}

/// Simulated LSTM model weights
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct ModelWeights {
    /// Feature importance weights
    pub feature_weights: Vec<f64>,
    /// Time decay factors
    pub time_decay: Vec<f64>,
    /// Pattern weights
    pub pattern_weights: HashMap<String, f64>,
    /// Bias terms
    pub bias: Vec<f64>,
}

impl Default for ModelWeights {
    fn default() -> Self {
        Self {
            // Initialize with reasonable defaults
            feature_weights: vec![
                0.15, // online_duration
                0.20, // avg_response_time
                0.10, // resource_contribution
                0.05, // message_frequency
                0.05, // time_of_day
                0.05, // day_of_week
                0.25, // historical_reliability
                0.10, // recent_disconnections
                0.05, // avg_session_length
                0.00, // connection_stability (will be learned)
            ],
            time_decay: vec![0.9, 0.8, 0.7], // 1h, 6h, 24h
            pattern_weights: HashMap::new(),
            bias: vec![0.1, 0.2, 0.3], // Base probabilities
        }
    }
}

/// Training example for online learning
#[derive(Debug, Clone)]
pub struct TrainingExample {
    pub node_id: NodeId,
    pub features: NodeFeatures,
    pub timestamp: std::time::Instant,
    pub actual_churn_1h: bool,
    pub actual_churn_6h: bool,
    pub actual_churn_24h: bool,
}

/// Churn prediction result
#[derive(Debug, Clone)]
pub struct ChurnPrediction {
    pub probability_1h: f64,
    pub probability_6h: f64,
    pub probability_24h: f64,
    pub confidence: f64,
    pub timestamp: std::time::Instant,
}

impl Default for ChurnPredictor {
    fn default() -> Self {
        Self::new()
    }
}

impl ChurnPredictor {
    /// Create a new churn predictor
    pub fn new() -> Self {
        Self {
            prediction_cache: Arc::new(tokio::sync::RwLock::new(HashMap::new())),
            feature_history: Arc::new(tokio::sync::RwLock::new(HashMap::new())),
            model_weights: Arc::new(tokio::sync::RwLock::new(ModelWeights::default())),
            experience_buffer: Arc::new(tokio::sync::RwLock::new(Vec::new())),
            max_buffer_size: 10000,
            _update_interval: std::time::Duration::from_secs(3600), // 1 hour
        }
    }

    /// Check if content should be replicated based on node churn risk
    pub async fn should_replicate(&self, node_id: &NodeId) -> bool {
        let prediction = self.predict(node_id).await;
        prediction.probability_1h > 0.7
    }

    /// Update node features for prediction
    pub async fn update_node_features(
        &self,
        node_id: &NodeId,
        features: Vec<f64>,
    ) -> anyhow::Result<()> {
        if features.len() != 10 {
            return Err(anyhow::anyhow!(
                "Expected 10 features, got {}",
                features.len()
            ));
        }

        let node_features = NodeFeatures {
            online_duration: features[0],
            avg_response_time: features[1],
            resource_contribution: features[2],
            message_frequency: features[3],
            time_of_day: features[4],
            day_of_week: features[5],
            historical_reliability: features[6],
            recent_disconnections: features[7],
            avg_session_length: features[8],
            connection_stability: features[9],
        };

        // Update feature history
        let mut history = self.feature_history.write().await;
        let entry = history
            .entry(node_id.clone())
            .or_insert(FeatureHistory::new());

        // Add or update session with new features
        if entry.sessions.is_empty() || entry.sessions.last().map(|s| s.1.is_some()).unwrap_or(true)
        {
            // Start a new session if there's no session or the last one has ended
            entry.sessions.push((std::time::Instant::now(), None));
        }

        // Add feature snapshot
        entry
            .snapshots
            .push((std::time::Instant::now(), node_features));

        // Keep only recent snapshots
        while entry.snapshots.len() > 100 {
            entry.snapshots.remove(0);
        }

        Ok(())
    }

    /// Extract features from node behavior
    pub async fn extract_features(&self, node_id: &NodeId) -> Option<NodeFeatures> {
        let history = self.feature_history.read().await;
        let node_history = history.get(node_id)?;

        // Calculate features from history
        let now = std::time::Instant::now();
        let current_session = node_history.sessions.last()?;
        let online_duration = if current_session.1.is_none() {
            now.duration_since(current_session.0).as_secs() as f64
        } else {
            0.0
        };

        // Calculate average session length
        let completed_sessions: Vec<_> = node_history
            .sessions
            .iter()
            .filter_map(|(start, end)| {
                end.as_ref()
                    .map(|e| e.duration_since(*start).as_secs() as f64)
            })
            .collect();
        let avg_session_length = if !completed_sessions.is_empty() {
            completed_sessions.iter().sum::<f64>() / completed_sessions.len() as f64 / 3600.0
        } else {
            1.0 // Default to 1 hour
        };

        // Calculate recent disconnections
        let one_week_ago = now - std::time::Duration::from_secs(7 * 24 * 3600);
        let recent_disconnections = node_history
            .sessions
            .iter()
            .filter(|(start, end)| end.is_some() && *start > one_week_ago)
            .count() as f64;

        // Get latest snapshot for other features
        let latest_snapshot = node_history
            .snapshots
            .last()
            .map(|(_, features)| features.clone())
            .unwrap_or_else(|| NodeFeatures {
                online_duration,
                avg_response_time: 100.0,
                resource_contribution: 0.5,
                message_frequency: 10.0,
                time_of_day: 12.0, // Default to noon
                day_of_week: 3.0,  // Default to Wednesday
                historical_reliability: node_history.total_uptime as f64
                    / (node_history.total_uptime + node_history.total_downtime).max(1) as f64,
                recent_disconnections,
                avg_session_length,
                connection_stability: 1.0 - (recent_disconnections / 7.0).min(1.0),
            });

        Some(NodeFeatures {
            online_duration,
            recent_disconnections,
            avg_session_length,
            historical_reliability: node_history.total_uptime as f64
                / (node_history.total_uptime + node_history.total_downtime).max(1) as f64,
            connection_stability: 1.0 - (recent_disconnections / 7.0).min(1.0),
            ..latest_snapshot
        })
    }

    /// Analyze patterns in node behavior
    async fn analyze_patterns(&self, features: &NodeFeatures) -> HashMap<String, f64> {
        let mut patterns = HashMap::new();

        // Time-based patterns
        let is_night = features.time_of_day < 6.0 || features.time_of_day > 22.0;
        let is_weekend = features.day_of_week == 0.0 || features.day_of_week == 6.0;

        patterns.insert("night_time".to_string(), if is_night { 1.0 } else { 0.0 });
        patterns.insert("weekend".to_string(), if is_weekend { 1.0 } else { 0.0 });

        // Behavior patterns
        patterns.insert(
            "short_session".to_string(),
            if features.online_duration < 1800.0 {
                1.0
            } else {
                0.0
            },
        );
        patterns.insert(
            "unstable".to_string(),
            if features.recent_disconnections > 5.0 {
                1.0
            } else {
                0.0
            },
        );
        patterns.insert(
            "low_contribution".to_string(),
            if features.resource_contribution < 0.3 {
                1.0
            } else {
                0.0
            },
        );
        patterns.insert(
            "slow_response".to_string(),
            if features.avg_response_time > 500.0 {
                1.0
            } else {
                0.0
            },
        );

        // Combined patterns
        let risk_score = (features.recent_disconnections / 10.0).min(1.0) * 0.3
            + (1.0 - features.historical_reliability) * 0.4
            + (1.0 - features.connection_stability) * 0.3;
        patterns.insert(
            "high_risk".to_string(),
            if risk_score > 0.6 { 1.0 } else { 0.0 },
        );

        patterns
    }

    /// Predict churn probability for a node
    pub async fn predict(&self, node_id: &NodeId) -> ChurnPrediction {
        // Check cache first
        {
            let cache = self.prediction_cache.read().await;
            if let Some(cached) = cache.get(node_id)
                && cached.timestamp.elapsed() < std::time::Duration::from_secs(300)
            {
                return cached.clone();
            }
        }

        // Extract features
        let features = match self.extract_features(node_id).await {
            Some(f) => f,
            None => {
                // No history, return low probability
                return ChurnPrediction {
                    probability_1h: 0.1,
                    probability_6h: 0.2,
                    probability_24h: 0.3,
                    confidence: 0.1,
                    timestamp: std::time::Instant::now(),
                };
            }
        };

        // Analyze patterns
        let patterns = self.analyze_patterns(&features).await;

        // Apply model (simulated LSTM)
        let model = self.model_weights.read().await;
        let prediction = self.apply_model(&features, &patterns, &model).await;

        // Cache the prediction
        let mut cache = self.prediction_cache.write().await;
        cache.insert(node_id.clone(), prediction.clone());
        prediction
    }

    /// Apply the model to compute predictions
    async fn apply_model(
        &self,
        features: &NodeFeatures,
        patterns: &HashMap<String, f64>,
        model: &ModelWeights,
    ) -> ChurnPrediction {
        // Convert features to vector
        let feature_vec = [
            features.online_duration / 3600.0,   // Normalize to hours
            features.avg_response_time / 1000.0, // Normalize to seconds
            features.resource_contribution,
            features.message_frequency / 100.0, // Normalize
            features.time_of_day / 24.0,        // Normalize
            features.day_of_week / 7.0,         // Normalize
            features.historical_reliability,
            features.recent_disconnections / 10.0, // Normalize
            features.avg_session_length / 24.0,    // Normalize to days
            features.connection_stability,
        ];

        // Compute base score from features
        let mut base_scores = [0.0; 3]; // 1h, 6h, 24h
        for (i, &weight) in model.feature_weights.iter().enumerate() {
            if i < feature_vec.len() {
                for score in &mut base_scores {
                    *score += weight * feature_vec[i];
                }
            }
        }

        // Apply pattern weights
        let mut pattern_score = 0.0;
        for (pattern, &value) in patterns {
            if let Some(&weight) = model.pattern_weights.get(pattern) {
                pattern_score += weight * value;
            } else {
                // Default weight for unknown patterns
                pattern_score += 0.1 * value;
            }
        }

        // Combine scores with time decay
        let probabilities: Vec<f64> = base_scores
            .iter()
            .zip(&model.time_decay)
            .zip(&model.bias)
            .map(|((base, decay), bias)| {
                let raw_score = base + pattern_score * decay + bias;
                // Sigmoid activation
                1.0 / (1.0 + (-raw_score).exp())
            })
            .collect();

        // Calculate confidence based on feature completeness and history length
        let confidence = 0.8; // Base confidence, would be calculated from history in real implementation

        ChurnPrediction {
            probability_1h: probabilities[0].min(0.99),
            probability_6h: probabilities[1].min(0.99),
            probability_24h: probabilities[2].min(0.99),
            confidence,
            timestamp: std::time::Instant::now(),
        }
    }

    /// Update node behavior tracking
    pub async fn update_node_behavior(
        &self,
        node_id: &NodeId,
        features: NodeFeatures,
    ) -> anyhow::Result<()> {
        let mut history = self.feature_history.write().await;
        let node_history = history
            .entry(node_id.clone())
            .or_insert_with(|| FeatureHistory {
                node_id: node_id.clone(),
                snapshots: Vec::new(),
                sessions: vec![(std::time::Instant::now(), None)],
                total_uptime: 0,
                total_downtime: 0,
            });

        // Add snapshot
        node_history
            .snapshots
            .push((std::time::Instant::now(), features));

        // Keep only recent snapshots (last 24 hours)
        let cutoff = std::time::Instant::now() - std::time::Duration::from_secs(24 * 3600);
        node_history
            .snapshots
            .retain(|(timestamp, _)| *timestamp > cutoff);

        Ok(())
    }

    /// Record node connection event
    pub async fn record_node_event(&self, node_id: &NodeId, event: NodeEvent) -> Result<()> {
        let mut history = self.feature_history.write().await;
        let node_history = history
            .entry(node_id.clone())
            .or_insert_with(|| FeatureHistory {
                node_id: node_id.clone(),
                snapshots: Vec::new(),
                sessions: Vec::new(),
                total_uptime: 0,
                total_downtime: 0,
            });

        match event {
            NodeEvent::Connected => {
                // Start new session
                node_history
                    .sessions
                    .push((std::time::Instant::now(), None));
            }
            NodeEvent::Disconnected => {
                // End current session
                if let Some((start, end)) = node_history.sessions.last_mut()
                    && end.is_none()
                {
                    let now = std::time::Instant::now();
                    *end = Some(now);
                    let session_length = now.duration_since(*start).as_secs();
                    node_history.total_uptime += session_length;
                }
            }
        }

        Ok(())
    }

    /// Add training example for online learning
    pub async fn add_training_example(
        &self,
        node_id: &NodeId,
        features: NodeFeatures,
        actual_churn_1h: bool,
        actual_churn_6h: bool,
        actual_churn_24h: bool,
    ) -> anyhow::Result<()> {
        let example = TrainingExample {
            node_id: node_id.clone(),
            features,
            timestamp: std::time::Instant::now(),
            actual_churn_1h,
            actual_churn_6h,
            actual_churn_24h,
        };

        let mut buffer = self.experience_buffer.write().await;
        buffer.push(example);

        // Maintain buffer size
        if buffer.len() > self.max_buffer_size {
            let drain_count = buffer.len() - self.max_buffer_size;
            buffer.drain(0..drain_count);
        }

        // Trigger model update if enough examples
        if buffer.len() >= 32 && buffer.len() % 32 == 0 {
            self.update_model().await?;
        }

        Ok(())
    }

    /// Update model weights based on experience buffer
    async fn update_model(&self) -> anyhow::Result<()> {
        let buffer = self.experience_buffer.read().await;
        if buffer.is_empty() {
            return Ok(());
        }

        let mut model = self.model_weights.write().await;

        // Simple online learning update (gradient descent simulation)
        let learning_rate = 0.01;
        let batch_size = 32.min(buffer.len());

        // Sample random batch
        let mut rng = rand::thread_rng();
        let batch: Vec<_> = (0..batch_size)
            .map(|_| &buffer[rng.gen_range(0..buffer.len())])
            .collect();

        // Update weights based on prediction errors
        for example in batch {
            // Extract features for this example
            let feature_vec = [
                example.features.online_duration / 3600.0,
                example.features.avg_response_time / 1000.0,
                example.features.resource_contribution,
                example.features.message_frequency / 100.0,
                example.features.time_of_day / 24.0,
                example.features.day_of_week / 7.0,
                example.features.historical_reliability,
                example.features.recent_disconnections / 10.0,
                example.features.avg_session_length / 24.0,
                example.features.connection_stability,
            ];

            // Calculate patterns
            let patterns = self.analyze_patterns(&example.features).await;

            // Get predictions
            let prediction = self.apply_model(&example.features, &patterns, &model).await;

            // Calculate errors
            let errors = [
                if example.actual_churn_1h { 1.0 } else { 0.0 } - prediction.probability_1h,
                if example.actual_churn_6h { 1.0 } else { 0.0 } - prediction.probability_6h,
                if example.actual_churn_24h { 1.0 } else { 0.0 } - prediction.probability_24h,
            ];

            // Update feature weights
            for (i, &feature_value) in feature_vec.iter().enumerate() {
                if i < model.feature_weights.len() {
                    for (j, &error) in errors.iter().enumerate() {
                        model.feature_weights[i] +=
                            learning_rate * error * feature_value * model.time_decay[j];
                    }
                }
            }

            // Update pattern weights
            for (pattern, &value) in &patterns {
                let avg_error = errors.iter().sum::<f64>() / errors.len() as f64;
                model
                    .pattern_weights
                    .entry(pattern.clone())
                    .and_modify(|w| *w += learning_rate * avg_error * value)
                    .or_insert(learning_rate * avg_error * value);
            }
        }

        Ok(())
    }

    /// Save model to disk
    pub async fn save_model(&self, path: &std::path::Path) -> anyhow::Result<()> {
        let model = self.model_weights.read().await;
        let serialized = serde_json::to_string(&*model)?;
        tokio::fs::write(path, serialized).await?;
        Ok(())
    }

    /// Load model from disk
    pub async fn load_model(&self, path: &std::path::Path) -> anyhow::Result<()> {
        let data = tokio::fs::read_to_string(path).await?;
        let loaded_model: ModelWeights = serde_json::from_str(&data)?;
        let mut model = self.model_weights.write().await;
        *model = loaded_model;
        Ok(())
    }
}

/// Node connection event
#[derive(Debug, Clone)]
pub enum NodeEvent {
    Connected,
    Disconnected,
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn test_thompson_sampling_initialization() {
        let ts = ThompsonSampling::new();
        let metrics = ts.get_metrics().await;

        assert_eq!(metrics.total_decisions, 0);
        assert!(metrics.decisions_by_type.is_empty());
        assert!(metrics.strategy_success_rates.is_empty());
    }

    #[tokio::test]
    async fn test_thompson_sampling_selection() -> Result<()> {
        let ts = ThompsonSampling::new();

        // Test selection for different content types
        for content_type in [
            ContentType::DHTLookup,
            ContentType::DataRetrieval,
            ContentType::ComputeRequest,
            ContentType::RealtimeMessage,
        ] {
            let strategy = ts.select_strategy(content_type).await?;
            assert!(matches!(
                strategy,
                StrategyChoice::Kademlia
                    | StrategyChoice::Hyperbolic
                    | StrategyChoice::TrustPath
                    | StrategyChoice::SOMRegion
            ));
        }

        let metrics = ts.get_metrics().await;
        assert_eq!(metrics.total_decisions, 4);
        assert_eq!(metrics.decisions_by_type.len(), 4);
        Ok(())
    }

    #[tokio::test]
    async fn test_thompson_sampling_update() -> Result<()> {
        let ts = ThompsonSampling::new();

        // Heavily reward Hyperbolic strategy for DataRetrieval
        for _ in 0..20 {
            ts.update(
                ContentType::DataRetrieval,
                StrategyChoice::Hyperbolic,
                true,
                50,
            )
            .await?;
        }

        // Penalize Kademlia for DataRetrieval
        for _ in 0..10 {
            ts.update(
                ContentType::DataRetrieval,
                StrategyChoice::Kademlia,
                false,
                200,
            )
            .await?;
        }

        // After training, Hyperbolic should be preferred for DataRetrieval
        let mut hyperbolic_count = 0;
        for _ in 0..100 {
            let strategy = ts.select_strategy(ContentType::DataRetrieval).await?;
            if matches!(strategy, StrategyChoice::Hyperbolic) {
                hyperbolic_count += 1;
            }
        }

        // Should select Hyperbolic significantly more often
        assert!(
            hyperbolic_count > 60,
            "Expected Hyperbolic to be selected more than 60% of the time, got {}%",
            hyperbolic_count
        );
        Ok(())
    }

    #[tokio::test]
    async fn test_confidence_intervals() -> Result<()> {
        let ts = ThompsonSampling::new();

        // Add some successes and failures
        for i in 0..10 {
            ts.update(
                ContentType::DHTLookup,
                StrategyChoice::Kademlia,
                i % 3 != 0, // 70% success rate
                100,
            )
            .await?;
        }

        let (lower, upper) = ts
            .get_confidence_interval(ContentType::DHTLookup, StrategyChoice::Kademlia)
            .await;

        assert!(lower > 0.0 && lower < 1.0);
        assert!(upper > lower && upper <= 1.0);
        assert!(upper - lower < 0.5); // Confidence interval should narrow with data
        Ok(())
    }

    #[tokio::test]
    async fn test_exploration_bonus() -> Result<()> {
        let ts = ThompsonSampling::new();

        // Give one strategy some data
        for _ in 0..15 {
            ts.update(
                ContentType::ComputeRequest,
                StrategyChoice::TrustPath,
                true,
                100,
            )
            .await?;
        }

        // Other strategies should still be explored due to exploration bonus
        let mut strategy_counts = HashMap::new();
        for _ in 0..100 {
            let strategy = ts.select_strategy(ContentType::ComputeRequest).await?;
            *strategy_counts.entry(strategy).or_insert(0) += 1;
        }

        // All strategies should have been tried at least once
        assert!(
            strategy_counts.len() >= 3,
            "Expected at least 3 different strategies to be tried"
        );
        Ok(())
    }

    #[tokio::test]
    async fn test_reset_strategy() -> Result<()> {
        let ts = ThompsonSampling::new();

        // Train a strategy
        for _ in 0..10 {
            ts.update(
                ContentType::RealtimeMessage,
                StrategyChoice::SOMRegion,
                true,
                50,
            )
            .await?;
        }

        // Reset it
        ts.reset_strategy(ContentType::RealtimeMessage, StrategyChoice::SOMRegion)
            .await;

        // Confidence interval should be back to uniform
        let (lower, upper) = ts
            .get_confidence_interval(ContentType::RealtimeMessage, StrategyChoice::SOMRegion)
            .await;

        assert_eq!(lower, 0.0);
        assert_eq!(upper, 1.0);
        Ok(())
    }

    #[tokio::test]
    async fn test_learning_system_trait() {
        let mut ts = ThompsonSampling::new();

        let context = LearningContext {
            content_type: ContentType::DataRetrieval,
            network_conditions: NetworkConditions {
                connected_peers: 100,
                avg_latency_ms: 50.0,
                churn_rate: 0.1,
            },
            historical_performance: vec![0.8, 0.85, 0.9],
        };

        // Test trait methods
        let choice = <ThompsonSampling as LearningSystem>::select_strategy(&ts, &context).await;
        assert!(matches!(
            choice,
            StrategyChoice::Kademlia
                | StrategyChoice::Hyperbolic
                | StrategyChoice::TrustPath
                | StrategyChoice::SOMRegion
        ));

        let outcome = Outcome {
            success: true,
            latency_ms: 45,
            hops: 3,
        };

        <ThompsonSampling as LearningSystem>::update(&mut ts, &context, &choice, &outcome).await;

        let metrics = <ThompsonSampling as LearningSystem>::metrics(&ts).await;
        assert_eq!(metrics.total_decisions, 1);
    }

    #[tokio::test]
    async fn test_cache_manager() {
        let manager = QLearnCacheManager::new(1024);
        let hash = ContentHash([1u8; 32]);

        // Test insertion
        assert!(manager.insert(hash.clone(), vec![0u8; 100]).await);

        // Test retrieval
        assert!(manager.get(&hash).await.is_some());

        // Test Q-learning decision
        let action = manager.decide_action(&hash).await;
        assert!(matches!(
            action,
            CacheAction::Cache
                | CacheAction::Evict(_)
                | CacheAction::IncreaseReplication
                | CacheAction::DecreaseReplication
                | CacheAction::NoAction
        ));
    }

    #[tokio::test]
    async fn test_q_value_update() {
        let manager = QLearnCacheManager::new(1024);

        let state = CacheState {
            utilization_bucket: 5,
            request_rate_bucket: 5,
            content_popularity: 5,
            size_bucket: 3,
        };

        let next_state = CacheState {
            utilization_bucket: 6,
            request_rate_bucket: 5,
            content_popularity: 5,
            size_bucket: 3,
        };

        manager
            .update_q_value(state, CacheAction::Cache, 1.0, next_state)
            .await;

        // Q-value should have been updated
        let q_table = manager.q_table.read().await;
        assert!(q_table.len() > 0);
    }

    #[tokio::test]
    async fn test_churn_predictor() {
        use crate::peer_record::UserId;
        use rand::RngCore;

        let predictor = ChurnPredictor::new();
        let mut hash = [0u8; 32];
        rand::thread_rng().fill_bytes(&mut hash);
        let node_id = UserId::from_bytes(hash);

        let prediction = predictor.predict(&node_id).await;
        assert!(prediction.probability_1h >= 0.0 && prediction.probability_1h <= 1.0);
        assert!(prediction.probability_6h >= 0.0 && prediction.probability_6h <= 1.0);
        assert!(prediction.probability_24h >= 0.0 && prediction.probability_24h <= 1.0);

        // Test caching
        let prediction2 = predictor.predict(&node_id).await;
        assert_eq!(prediction.probability_1h, prediction2.probability_1h);
    }

    #[tokio::test]
    async fn test_cache_eviction_policies() {
        let manager = QLearnCacheManager::new(300); // Small cache for testing

        // Insert multiple items
        let hash1 = ContentHash([1u8; 32]);
        let hash2 = ContentHash([2u8; 32]);
        let hash3 = ContentHash([3u8; 32]);

        manager.insert(hash1.clone(), vec![0u8; 100]).await;
        manager.insert(hash2.clone(), vec![0u8; 100]).await;
        manager.insert(hash3.clone(), vec![0u8; 100]).await;

        // Access hash1 and hash2 to make them more recently used
        manager.get(&hash1).await;
        manager.get(&hash2).await;

        // Force eviction by adding another item
        let hash4 = ContentHash([4u8; 32]);
        manager.insert(hash4.clone(), vec![0u8; 100]).await;

        // hash3 should have been evicted (LRU)
        assert!(manager.get(&hash1).await.is_some());
        assert!(manager.get(&hash2).await.is_some());
        assert!(manager.get(&hash3).await.is_none());
        assert!(manager.get(&hash4).await.is_some());
    }

    #[tokio::test]
    async fn test_reward_calculation() {
        let manager = QLearnCacheManager::new(1024);

        // Insert some content to establish hit rate
        let hash = ContentHash([1u8; 32]);
        manager.insert(hash.clone(), vec![0u8; 100]).await;

        // Generate some hits
        for _ in 0..5 {
            manager.get(&hash).await;
        }

        // Generate some misses
        let miss_hash = ContentHash([2u8; 32]);
        for _ in 0..2 {
            manager.get(&miss_hash).await;
        }

        // Test reward calculation for different actions
        let cache_reward = manager.calculate_reward(CacheAction::Cache, true, 1000);
        assert!(cache_reward > 0.0); // Should be positive for cache hit

        let evict_reward =
            manager.calculate_reward(CacheAction::Evict(EvictionPolicy::LRU), false, 0);
        assert!(evict_reward >= 0.0); // Should be slightly positive for evicting unused content

        let evict_penalty =
            manager.calculate_reward(CacheAction::Evict(EvictionPolicy::LRU), true, 0);
        assert!(evict_penalty < 0.0); // Should be negative for evicting needed content
    }

    #[tokio::test]
    async fn test_cache_statistics() {
        let manager = QLearnCacheManager::new(1024);

        let hash1 = ContentHash([1u8; 32]);
        let hash2 = ContentHash([2u8; 32]);

        // Insert content
        manager.insert(hash1.clone(), vec![0u8; 100]).await;

        // Generate hits and misses
        manager.get(&hash1).await; // Hit
        manager.get(&hash1).await; // Hit
        manager.get(&hash2).await; // Miss

        let stats = manager.get_stats();
        assert_eq!(stats.hits, 2);
        assert_eq!(stats.misses, 1);
        assert!((stats.hit_rate - 0.666).abs() < 0.01); // ~66.6% hit rate
        assert_eq!(stats.size_bytes, 100);
    }

    #[tokio::test]
    async fn test_exploration_vs_exploitation() {
        let manager = QLearnCacheManager::new(1024);

        // Train the Q-table with some states and actions
        let state = CacheState {
            utilization_bucket: 5,
            request_rate_bucket: 7,
            content_popularity: 8,
            size_bucket: 2,
        };

        // Make Cache action very valuable for this state
        for _ in 0..10 {
            manager
                .update_q_value(state.clone(), CacheAction::Cache, 1.0, state.clone())
                .await;
        }

        // Count how often we get Cache action
        let mut cache_count = 0;
        for _ in 0..100 {
            // Temporarily set get_current_state to return our trained state
            // In real test we'd mock this properly
            let action = manager.decide_action(&ContentHash([1u8; 32])).await;
            if matches!(action, CacheAction::Cache) {
                cache_count += 1;
            }
        }

        // With epsilon=0.1, we should get Cache action ~90% of the time
        // Allow some variance
        assert!(cache_count > 80 && cache_count < 95);
    }

    #[tokio::test]
    async fn test_state_representation() {
        let manager = QLearnCacheManager::new(1024);

        // Test state bucketing
        let hash = ContentHash([1u8; 32]);

        // Insert content and track stats
        manager.insert(hash.clone(), vec![0u8; 100]).await;

        // Make some requests to update stats
        for _ in 0..5 {
            manager.get(&hash).await;
        }

        let state = manager.get_current_state_async(&hash).await;

        // Check state bounds
        assert!(state.utilization_bucket <= 10);
        assert!(state.request_rate_bucket <= 10);
        assert!(state.content_popularity <= 10);
        assert!(state.size_bucket <= 10);
    }

    #[tokio::test]
    async fn test_action_execution() -> Result<()> {
        let manager = QLearnCacheManager::new(1024);
        let hash = ContentHash([1u8; 32]);

        // Test Cache action
        manager
            .execute_action(&hash, CacheAction::Cache, Some(vec![0u8; 100]))
            .await?;
        assert!(manager.get(&hash).await.is_some());

        // Test NoAction
        manager
            .execute_action(&hash, CacheAction::NoAction, None)
            .await?;
        assert!(manager.get(&hash).await.is_some()); // Should still be there

        // Test Evict action
        manager
            .execute_action(&hash, CacheAction::Evict(EvictionPolicy::LRU), None)
            .await?;
        // Note: May or may not evict our specific item depending on LRU state

        let stats = manager.get_stats();
        assert!(stats.size_bytes <= 100); // Should be 0 or 100 depending on eviction
        Ok(())
    }

    #[tokio::test]
    async fn test_churn_predictor_initialization() {
        let predictor = ChurnPredictor::new();

        // Test prediction for unknown node
        let node_id = NodeId { hash: [1u8; 32] };
        let prediction = predictor.predict(&node_id).await;

        // Should return low confidence for unknown node
        assert!(prediction.confidence < 0.2);
        assert!(prediction.probability_1h < 0.3);
        assert!(prediction.probability_6h < 0.4);
        assert!(prediction.probability_24h < 0.5);
    }

    #[tokio::test]
    async fn test_churn_predictor_node_events() -> Result<()> {
        let predictor = ChurnPredictor::new();
        let node_id = NodeId { hash: [1u8; 32] };

        // Record connection
        predictor
            .record_node_event(&node_id, NodeEvent::Connected)
            .await?;

        // Record disconnection
        tokio::time::sleep(std::time::Duration::from_millis(100)).await;
        predictor
            .record_node_event(&node_id, NodeEvent::Disconnected)
            .await?;

        // Check that session was recorded
        let features = predictor.extract_features(&node_id).await;
        assert!(features.is_some());
        Ok(())
    }

    #[tokio::test]
    async fn test_churn_predictor_feature_extraction() -> Result<()> {
        let predictor = ChurnPredictor::new();
        let node_id = NodeId { hash: [1u8; 32] };

        // Create node history
        predictor
            .record_node_event(&node_id, NodeEvent::Connected)
            .await?;

        // Update behavior
        let features = NodeFeatures {
            online_duration: 3600.0,
            avg_response_time: 50.0,
            resource_contribution: 0.8,
            message_frequency: 20.0,
            time_of_day: 14.0,
            day_of_week: 2.0,
            historical_reliability: 0.9,
            recent_disconnections: 1.0,
            avg_session_length: 4.0,
            connection_stability: 0.85,
        };

        predictor
            .update_node_behavior(&node_id, features.clone())
            .await?;

        // Extract features
        let extracted = predictor
            .extract_features(&node_id)
            .await
            .ok_or(anyhow::anyhow!("no features extracted"))?;
        assert_eq!(extracted.resource_contribution, 0.8);
        assert_eq!(extracted.avg_response_time, 50.0);
        Ok(())
    }

    #[tokio::test]
    async fn test_churn_predictor_pattern_analysis() {
        let predictor = ChurnPredictor::new();

        // Test night time pattern
        let night_features = NodeFeatures {
            online_duration: 1000.0,
            avg_response_time: 100.0,
            resource_contribution: 0.5,
            message_frequency: 10.0,
            time_of_day: 2.0, // 2 AM
            day_of_week: 3.0,
            historical_reliability: 0.8,
            recent_disconnections: 2.0,
            avg_session_length: 2.0,
            connection_stability: 0.8,
        };

        let patterns = predictor.analyze_patterns(&night_features).await;
        assert_eq!(patterns.get("night_time"), Some(&1.0));
        assert_eq!(patterns.get("weekend"), Some(&0.0));

        // Test unstable pattern
        let unstable_features = NodeFeatures {
            recent_disconnections: 7.0,
            connection_stability: 0.3,
            ..night_features
        };

        let patterns = predictor.analyze_patterns(&unstable_features).await;
        assert_eq!(patterns.get("unstable"), Some(&1.0));
        assert_eq!(patterns.get("high_risk"), Some(&1.0));
    }

    #[tokio::test]
    async fn test_churn_predictor_proactive_replication() -> Result<()> {
        let predictor = ChurnPredictor::new();
        let node_id = NodeId { hash: [1u8; 32] };

        // Create high-risk node history
        predictor
            .record_node_event(&node_id, NodeEvent::Connected)
            .await?;

        let features = NodeFeatures {
            online_duration: 600.0, // Only 10 minutes online
            avg_response_time: 500.0,
            resource_contribution: 0.1,
            message_frequency: 2.0,
            time_of_day: 23.0,
            day_of_week: 5.0,
            historical_reliability: 0.3,
            recent_disconnections: 10.0,
            avg_session_length: 0.5,
            connection_stability: 0.1,
        };

        predictor.update_node_behavior(&node_id, features).await?;

        // Should recommend replication for high-risk node
        let _should_replicate = predictor.should_replicate(&node_id).await;
        // Without full model training, this might not always be true, but test structure is correct
        let prediction = predictor.predict(&node_id).await;
        assert!(prediction.probability_1h > 0.0);
        Ok(())
    }

    #[tokio::test]
    async fn test_churn_predictor_online_learning() -> Result<()> {
        let predictor = ChurnPredictor::new();
        let node_id = NodeId { hash: [1u8; 32] };

        // Add training examples
        for i in 0..40 {
            let features = NodeFeatures {
                online_duration: (i * 1000) as f64,
                avg_response_time: 100.0,
                resource_contribution: 0.5,
                message_frequency: 10.0,
                time_of_day: 12.0,
                day_of_week: 3.0,
                historical_reliability: 0.8,
                recent_disconnections: (i % 5) as f64,
                avg_session_length: 2.0,
                connection_stability: 0.8,
            };

            // Some nodes churn, some don't
            let churned = i % 3 == 0;
            predictor
                .add_training_example(&node_id, features, churned, churned, churned)
                .await?;
        }

        // Model should have been updated after 32 examples
        let prediction = predictor.predict(&node_id).await;
        assert!(prediction.confidence > 0.0);
        Ok(())
    }

    #[tokio::test]
    async fn test_churn_predictor_model_persistence() -> Result<()> {
        let predictor = ChurnPredictor::new();
        let temp_path = std::path::Path::new("/tmp/test_model.json");

        // Save model
        predictor.save_model(temp_path).await?;

        // Load model
        predictor.load_model(temp_path).await?;

        // Clean up
        let _ = std::fs::remove_file(temp_path);
        Ok(())
    }
}