oxirs-embed 0.2.4

//! Real-time Fine-tuning System
//!
//! This module implements real-time fine-tuning capabilities for embedding models
//! with incremental learning, online adaptation, and dynamic model updates.

use crate::{EmbeddingModel, ModelConfig, TrainingStats, Triple, Vector};
use anyhow::{anyhow, Result};
use async_trait::async_trait;
use chrono::{DateTime, Utc};
use scirs2_core::ndarray_ext::{Array1, Array2};
use scirs2_core::random::{Random, RngExt};
use serde::{Deserialize, Serialize};
use std::collections::{HashMap, VecDeque};
use uuid::Uuid;

/// Configuration for real-time fine-tuning
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RealTimeFinetuningConfig {
    pub base_config: ModelConfig,
    /// Learning rate for online updates
    pub online_learning_rate: f32,
    /// Buffer size for experience replay
    pub replay_buffer_size: usize,
    /// Batch size for online updates
    pub online_batch_size: usize,
    /// Adaptation threshold for triggering updates
    pub adaptation_threshold: f32,
    /// Memory decay factor
    pub memory_decay: f32,
    /// Update frequency (every N examples)
    pub update_frequency: usize,
    /// Catastrophic forgetting prevention
    pub forgetting_prevention: ForgettingPreventionConfig,
    /// Online evaluation settings
    pub online_evaluation: OnlineEvaluationConfig,
}

impl Default for RealTimeFinetuningConfig {
    fn default() -> Self {
        Self {
            base_config: ModelConfig::default(),
            online_learning_rate: 1e-4,
            replay_buffer_size: 10000,
            online_batch_size: 32,
            adaptation_threshold: 0.1,
            memory_decay: 0.99,
            update_frequency: 10,
            forgetting_prevention: ForgettingPreventionConfig::default(),
            online_evaluation: OnlineEvaluationConfig::default(),
        }
    }
}

/// Catastrophic forgetting prevention configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ForgettingPreventionConfig {
    /// Use elastic weight consolidation
    pub use_ewc: bool,
    /// EWC regularization strength
    pub ewc_lambda: f32,
    /// Use progressive neural networks
    pub use_progressive_nets: bool,
    /// Use memory replay
    pub use_memory_replay: bool,
    /// Memory replay ratio
    pub replay_ratio: f32,
}

impl Default for ForgettingPreventionConfig {
    fn default() -> Self {
        Self {
            use_ewc: true,
            ewc_lambda: 0.4,
            use_progressive_nets: false,
            use_memory_replay: true,
            replay_ratio: 0.3,
        }
    }
}

/// Online evaluation configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct OnlineEvaluationConfig {
    /// Sliding window size for evaluation
    pub window_size: usize,
    /// Evaluation frequency
    pub eval_frequency: usize,
    /// Performance metrics to track
    pub metrics: Vec<OnlineMetric>,
    /// Early stopping criteria
    pub early_stopping: EarlyStoppingConfig,
}

impl Default for OnlineEvaluationConfig {
    fn default() -> Self {
        Self {
            window_size: 1000,
            eval_frequency: 100,
            metrics: vec![
                OnlineMetric::Loss,
                OnlineMetric::Accuracy,
                OnlineMetric::Drift,
                OnlineMetric::Forgetting,
            ],
            early_stopping: EarlyStoppingConfig::default(),
        }
    }
}

/// Online metrics to track
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum OnlineMetric {
    Loss,
    Accuracy,
    Drift,
    Forgetting,
    Plasticity,
    Stability,
}

/// Early stopping configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EarlyStoppingConfig {
    /// Patience (number of evaluations without improvement)
    pub patience: usize,
    /// Minimum improvement threshold
    pub min_improvement: f32,
    /// Metric to monitor
    pub monitor_metric: OnlineMetric,
}

impl Default for EarlyStoppingConfig {
    fn default() -> Self {
        Self {
            patience: 10,
            min_improvement: 1e-4,
            monitor_metric: OnlineMetric::Loss,
        }
    }
}

/// Experience replay buffer entry
#[derive(Debug, Clone)]
pub struct ExperienceEntry {
    pub input: Array1<f32>,
    pub target: Array1<f32>,
    pub timestamp: DateTime<Utc>,
    pub importance: f32,
    pub task_id: Option<String>,
}

/// Online performance tracking
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct OnlinePerformanceTracker {
    pub recent_losses: VecDeque<f32>,
    pub recent_accuracies: VecDeque<f32>,
    pub drift_scores: VecDeque<f32>,
    pub forgetting_scores: VecDeque<f32>,
    pub update_count: usize,
    pub last_evaluation: DateTime<Utc>,
}

impl OnlinePerformanceTracker {
    pub fn new(window_size: usize) -> Self {
        Self {
            recent_losses: VecDeque::with_capacity(window_size),
            recent_accuracies: VecDeque::with_capacity(window_size),
            drift_scores: VecDeque::with_capacity(window_size),
            forgetting_scores: VecDeque::with_capacity(window_size),
            update_count: 0,
            last_evaluation: Utc::now(),
        }
    }

    pub fn update_metrics(&mut self, loss: f32, accuracy: f32, drift: f32, forgetting: f32) {
        self.recent_losses.push_back(loss);
        self.recent_accuracies.push_back(accuracy);
        self.drift_scores.push_back(drift);
        self.forgetting_scores.push_back(forgetting);

        // Maintain window size
        if self.recent_losses.len() > self.recent_losses.capacity() {
            self.recent_losses.pop_front();
        }
        if self.recent_accuracies.len() > self.recent_accuracies.capacity() {
            self.recent_accuracies.pop_front();
        }
        if self.drift_scores.len() > self.drift_scores.capacity() {
            self.drift_scores.pop_front();
        }
        if self.forgetting_scores.len() > self.forgetting_scores.capacity() {
            self.forgetting_scores.pop_front();
        }

        self.update_count += 1;
        self.last_evaluation = Utc::now();
    }

    pub fn get_average_loss(&self) -> f32 {
        if self.recent_losses.is_empty() {
            0.0
        } else {
            self.recent_losses.iter().sum::<f32>() / self.recent_losses.len() as f32
        }
    }

    pub fn get_average_accuracy(&self) -> f32 {
        if self.recent_accuracies.is_empty() {
            0.0
        } else {
            self.recent_accuracies.iter().sum::<f32>() / self.recent_accuracies.len() as f32
        }
    }

    pub fn get_drift_score(&self) -> f32 {
        if self.drift_scores.is_empty() {
            0.0
        } else {
            self.drift_scores.iter().sum::<f32>() / self.drift_scores.len() as f32
        }
    }

    pub fn get_forgetting_score(&self) -> f32 {
        if self.forgetting_scores.is_empty() {
            0.0
        } else {
            self.forgetting_scores.iter().sum::<f32>() / self.forgetting_scores.len() as f32
        }
    }
}

/// Real-time fine-tuning model
#[derive(Debug)]
pub struct RealTimeFinetuningModel {
    pub config: RealTimeFinetuningConfig,
    pub model_id: Uuid,

    /// Core model parameters
    pub embeddings: Array2<f32>,
    pub fisher_information: Array2<f32>, // For EWC
    pub optimal_parameters: Array2<f32>, // For EWC

    /// Experience replay buffer
    pub replay_buffer: VecDeque<ExperienceEntry>,

    /// Online performance tracking
    pub performance_tracker: OnlinePerformanceTracker,

    /// Entity and relation mappings
    pub entities: HashMap<String, usize>,
    pub relations: HashMap<String, usize>,

    /// Training state
    pub examples_seen: usize,
    pub last_update: DateTime<Utc>,
    pub is_adapting: bool,

    /// Task-specific memory
    pub task_memory: HashMap<String, Array2<f32>>,
    pub current_task: Option<String>,

    /// Statistics
    pub training_stats: Option<TrainingStats>,
    pub is_trained: bool,
}

impl RealTimeFinetuningModel {
    /// Create new real-time fine-tuning model
    pub fn new(config: RealTimeFinetuningConfig) -> Self {
        let model_id = Uuid::new_v4();
        let dimensions = config.base_config.dimensions;

        Self {
            config: config.clone(),
            model_id,
            embeddings: Array2::zeros((0, dimensions)),
            fisher_information: Array2::zeros((0, dimensions)),
            optimal_parameters: Array2::zeros((0, dimensions)),
            replay_buffer: VecDeque::with_capacity(config.replay_buffer_size),
            performance_tracker: OnlinePerformanceTracker::new(
                config.online_evaluation.window_size,
            ),
            entities: HashMap::new(),
            relations: HashMap::new(),
            examples_seen: 0,
            last_update: Utc::now(),
            is_adapting: false,
            task_memory: HashMap::new(),
            current_task: None,
            training_stats: None,
            is_trained: false,
        }
    }

    /// Add new example for online learning
    pub async fn add_example(
        &mut self,
        input: Array1<f32>,
        target: Array1<f32>,
        task_id: Option<String>,
    ) -> Result<()> {
        // Initialize network if needed
        if self.embeddings.nrows() == 0 {
            let input_dim = input.len();
            let output_dim = target.len();
            self.embeddings = Array2::from_shape_fn((output_dim, input_dim), |(_, _)| {
                let mut random = Random::default();
                (random.random::<f32>() - 0.5) * 0.1
            });
            self.fisher_information = Array2::zeros((output_dim, input_dim));
            self.optimal_parameters = Array2::zeros((output_dim, input_dim));
        }

        // Add to replay buffer
        let entry = ExperienceEntry {
            input: input.clone(),
            target: target.clone(),
            timestamp: Utc::now(),
            importance: 1.0, // Can be computed based on novelty/difficulty
            task_id: task_id.clone(),
        };

        self.replay_buffer.push_back(entry);
        if self.replay_buffer.len() > self.config.replay_buffer_size {
            self.replay_buffer.pop_front();
        }

        self.examples_seen += 1;

        // Trigger adaptation if threshold met
        if self.should_adapt() {
            self.adapt_online().await?;
        }

        Ok(())
    }

    /// Check if model should adapt
    fn should_adapt(&self) -> bool {
        // Adapt every N examples or if performance drops
        if self.examples_seen % self.config.update_frequency == 0 {
            return true;
        }

        // Adapt if performance drops below threshold
        let current_loss = self.performance_tracker.get_average_loss();
        if current_loss > self.config.adaptation_threshold {
            return true;
        }

        false
    }

    /// Perform online adaptation
    pub async fn adapt_online(&mut self) -> Result<()> {
        if self.replay_buffer.is_empty() {
            return Ok(());
        }

        self.is_adapting = true;

        // Sample batch from replay buffer
        let batch = self.sample_replay_batch();

        // Compute gradients
        let gradients = self.compute_gradients(&batch)?;

        // Apply EWC regularization if enabled
        let regularized_gradients = if self.config.forgetting_prevention.use_ewc {
            self.apply_ewc_regularization(gradients)?
        } else {
            gradients
        };

        // Update parameters
        self.update_parameters(regularized_gradients)?;

        // Update Fisher information for EWC
        if self.config.forgetting_prevention.use_ewc {
            self.update_fisher_information(&batch)?;
        }

        // Evaluate performance
        self.evaluate_online_performance().await?;

        self.last_update = Utc::now();
        self.is_adapting = false;

        Ok(())
    }

    /// Sample batch from replay buffer
    fn sample_replay_batch(&self) -> Vec<ExperienceEntry> {
        let batch_size = self.config.online_batch_size.min(self.replay_buffer.len());
        let mut batch = Vec::with_capacity(batch_size);

        // Sample with importance-based probability
        for _ in 0..batch_size {
            let mut random = Random::default();
            let idx = random.random_range(0..self.replay_buffer.len());
            batch.push(self.replay_buffer[idx].clone());
        }

        batch
    }

    /// Compute gradients for batch
    fn compute_gradients(&self, batch: &[ExperienceEntry]) -> Result<Array2<f32>> {
        let dimensions = self.config.base_config.dimensions;
        let mut gradients = Array2::zeros((batch.len(), dimensions));

        for (i, entry) in batch.iter().enumerate() {
            // Simplified gradient computation
            // In practice, this would involve backpropagation through the model
            let prediction = self.forward_pass(&entry.input)?;
            let error = &entry.target - &prediction;

            // Simple gradient: error * input
            let gradient = &error * &entry.input;
            gradients.row_mut(i).assign(&gradient);
        }

        Ok(gradients)
    }

    /// Apply EWC regularization to gradients
    fn apply_ewc_regularization(&self, gradients: Array2<f32>) -> Result<Array2<f32>> {
        let lambda = self.config.forgetting_prevention.ewc_lambda;

        // EWC penalty: λ * F * (θ - θ*)
        let ewc_penalty =
            &self.fisher_information * (&self.embeddings - &self.optimal_parameters) * lambda;

        // Regularized gradients
        let mut regularized = gradients;
        for i in 0..regularized.nrows().min(ewc_penalty.nrows()) {
            for j in 0..regularized.ncols().min(ewc_penalty.ncols()) {
                regularized[[i, j]] -= ewc_penalty[[i, j]];
            }
        }

        Ok(regularized)
    }

    /// Update model parameters
    fn update_parameters(&mut self, gradients: Array2<f32>) -> Result<()> {
        let learning_rate = self.config.online_learning_rate;

        // Apply gradients with learning rate
        let update = &gradients * learning_rate;

        // Ensure embeddings matrix has the right shape
        if self.embeddings.nrows() < gradients.nrows() {
            let dimensions = self.config.base_config.dimensions;
            let new_rows = gradients.nrows();
            self.embeddings = Array2::from_shape_fn((new_rows, dimensions), |_| {
                let mut random = Random::default();
                random.random::<f32>() * 0.1
            });
        }

        // Update embeddings
        let rows_to_update = update.nrows().min(self.embeddings.nrows());
        let cols_to_update = update.ncols().min(self.embeddings.ncols());

        for i in 0..rows_to_update {
            for j in 0..cols_to_update {
                self.embeddings[[i, j]] += update[[i, j]];
            }
        }

        Ok(())
    }

    /// Update Fisher Information Matrix for EWC
    fn update_fisher_information(&mut self, batch: &[ExperienceEntry]) -> Result<()> {
        let dimensions = self.config.base_config.dimensions;
        let mut fisher_update = Array2::zeros((batch.len(), dimensions));

        for (i, entry) in batch.iter().enumerate() {
            // Compute second-order derivatives (simplified)
            let prediction = self.forward_pass(&entry.input)?;
            let second_derivative = prediction.mapv(|x| x * (1.0 - x)); // Sigmoid derivative approximation
            fisher_update.row_mut(i).assign(&second_derivative);
        }

        // Update Fisher information with exponential moving average
        let decay = self.config.memory_decay;

        // Resize Fisher information if needed
        if self.fisher_information.nrows() < fisher_update.nrows() {
            self.fisher_information = Array2::zeros((fisher_update.nrows(), dimensions));
        }

        let rows_to_update = fisher_update.nrows().min(self.fisher_information.nrows());
        let cols_to_update = fisher_update.ncols().min(self.fisher_information.ncols());

        for i in 0..rows_to_update {
            for j in 0..cols_to_update {
                self.fisher_information[[i, j]] =
                    decay * self.fisher_information[[i, j]] + (1.0 - decay) * fisher_update[[i, j]];
            }
        }

        Ok(())
    }

    /// Forward pass through the model
    fn forward_pass(&self, input: &Array1<f32>) -> Result<Array1<f32>> {
        if self.embeddings.is_empty() {
            return Ok(Array1::zeros(input.len()));
        }

        // Simple linear transformation
        let input_len = input.len().min(self.embeddings.ncols());
        let output_len = self.embeddings.nrows();
        let mut output = Array1::zeros(output_len);

        for i in 0..output_len {
            let mut sum = 0.0;
            for j in 0..input_len {
                sum += self.embeddings[[i, j]] * input[j];
            }
            output[i] = sum.tanh(); // Apply activation
        }

        Ok(output)
    }

    /// Evaluate online performance
    async fn evaluate_online_performance(&mut self) -> Result<()> {
        if self.replay_buffer.is_empty() {
            return Ok(());
        }

        let mut total_loss = 0.0;
        let mut total_accuracy = 0.0;
        let mut total_drift = 0.0;
        let mut total_forgetting = 0.0;
        let sample_size = self
            .config
            .online_evaluation
            .window_size
            .min(self.replay_buffer.len());

        for i in 0..sample_size {
            let idx = self.replay_buffer.len() - 1 - i; // Recent examples
            let entry = &self.replay_buffer[idx];

            let prediction = self.forward_pass(&entry.input)?;

            // Compute loss (MSE)
            let diff = &entry.target - &prediction;
            let loss = diff.dot(&diff) / diff.len() as f32;
            total_loss += loss;

            // Compute accuracy (simplified)
            let accuracy = 1.0 / (1.0 + loss);
            total_accuracy += accuracy;

            // Compute drift (change in prediction distribution)
            let drift = self.compute_drift_score(&prediction)?;
            total_drift += drift;

            // Compute forgetting (performance on old tasks)
            let forgetting = self.compute_forgetting_score(&entry.input, &entry.target)?;
            total_forgetting += forgetting;
        }

        let avg_loss = total_loss / sample_size as f32;
        let avg_accuracy = total_accuracy / sample_size as f32;
        let avg_drift = total_drift / sample_size as f32;
        let avg_forgetting = total_forgetting / sample_size as f32;

        self.performance_tracker
            .update_metrics(avg_loss, avg_accuracy, avg_drift, avg_forgetting);

        Ok(())
    }

    /// Compute drift score
    fn compute_drift_score(&self, prediction: &Array1<f32>) -> Result<f32> {
        // Simplified drift detection based on prediction distribution
        let mean = prediction.mean().unwrap_or(0.0);
        let variance = prediction.var(0.0);
        let drift_score = (mean.abs() + variance).min(1.0);
        Ok(drift_score)
    }

    /// Compute forgetting score
    fn compute_forgetting_score(&self, input: &Array1<f32>, target: &Array1<f32>) -> Result<f32> {
        let prediction = self.forward_pass(input)?;
        let diff = target - &prediction;
        let forgetting_score = diff.dot(&diff).sqrt() / target.len() as f32;
        Ok(forgetting_score.min(1.0))
    }

    /// Set current task context
    pub fn set_current_task(&mut self, task_id: Option<String>) {
        self.current_task = task_id;
    }

    /// Save task-specific parameters
    pub fn save_task_parameters(&mut self, task_id: String) -> Result<()> {
        self.task_memory.insert(task_id, self.embeddings.clone());
        Ok(())
    }

    /// Load task-specific parameters
    pub fn load_task_parameters(&mut self, task_id: &str) -> Result<()> {
        if let Some(task_params) = self.task_memory.get(task_id) {
            self.embeddings = task_params.clone();
        }
        Ok(())
    }

    /// Get online performance statistics
    pub fn get_online_stats(&self) -> HashMap<String, f32> {
        let mut stats = HashMap::new();

        stats.insert(
            "average_loss".to_string(),
            self.performance_tracker.get_average_loss(),
        );
        stats.insert(
            "average_accuracy".to_string(),
            self.performance_tracker.get_average_accuracy(),
        );
        stats.insert(
            "drift_score".to_string(),
            self.performance_tracker.get_drift_score(),
        );
        stats.insert(
            "forgetting_score".to_string(),
            self.performance_tracker.get_forgetting_score(),
        );
        stats.insert("examples_seen".to_string(), self.examples_seen as f32);
        stats.insert(
            "update_count".to_string(),
            self.performance_tracker.update_count as f32,
        );
        stats.insert(
            "replay_buffer_size".to_string(),
            self.replay_buffer.len() as f32,
        );

        stats
    }
}

#[async_trait]
impl EmbeddingModel for RealTimeFinetuningModel {
    fn config(&self) -> &ModelConfig {
        &self.config.base_config
    }

    fn model_id(&self) -> &Uuid {
        &self.model_id
    }

    fn model_type(&self) -> &'static str {
        "RealTimeFinetuningModel"
    }

    fn add_triple(&mut self, triple: Triple) -> Result<()> {
        let subject_str = triple.subject.iri.clone();
        let predicate_str = triple.predicate.iri.clone();
        let object_str = triple.object.iri.clone();

        // Add entities
        let next_entity_id = self.entities.len();
        self.entities.entry(subject_str).or_insert(next_entity_id);
        let next_entity_id = self.entities.len();
        self.entities.entry(object_str).or_insert(next_entity_id);

        // Add relation
        let next_relation_id = self.relations.len();
        self.relations
            .entry(predicate_str)
            .or_insert(next_relation_id);

        Ok(())
    }

    async fn train(&mut self, epochs: Option<usize>) -> Result<TrainingStats> {
        let epochs = epochs.unwrap_or(self.config.base_config.max_epochs);
        let start_time = std::time::Instant::now();

        let mut loss_history = Vec::new();

        for epoch in 0..epochs {
            // Simulate training with online adaptation
            let epoch_loss = {
                let mut random = Random::default();
                0.1 * random.random::<f64>()
            };
            loss_history.push(epoch_loss);

            // Simulate adding examples and adapting
            if epoch % 10 == 0 && !self.replay_buffer.is_empty() {
                self.adapt_online().await?;
            }

            if epoch > 10 && epoch_loss < 1e-6 {
                break;
            }
        }

        let training_time = start_time.elapsed().as_secs_f64();
        let final_loss = loss_history.last().copied().unwrap_or(0.0);

        let stats = TrainingStats {
            epochs_completed: loss_history.len(),
            final_loss,
            training_time_seconds: training_time,
            convergence_achieved: final_loss < 1e-4,
            loss_history,
        };

        self.training_stats = Some(stats.clone());
        self.is_trained = true;

        Ok(stats)
    }

    fn get_entity_embedding(&self, entity: &str) -> Result<Vector> {
        if let Some(&entity_id) = self.entities.get(entity) {
            if entity_id < self.embeddings.nrows() {
                let embedding = self.embeddings.row(entity_id);
                return Ok(Vector::new(embedding.to_vec()));
            }
        }
        Err(anyhow!("Entity not found: {}", entity))
    }

    fn get_relation_embedding(&self, relation: &str) -> Result<Vector> {
        if let Some(&relation_id) = self.relations.get(relation) {
            if relation_id < self.embeddings.nrows() {
                let embedding = self.embeddings.row(relation_id);
                return Ok(Vector::new(embedding.to_vec()));
            }
        }
        Err(anyhow!("Relation not found: {}", relation))
    }

    fn score_triple(&self, subject: &str, predicate: &str, object: &str) -> Result<f64> {
        let subject_emb = self.get_entity_embedding(subject)?;
        let predicate_emb = self.get_relation_embedding(predicate)?;
        let object_emb = self.get_entity_embedding(object)?;

        // Simple TransE-style scoring
        let subject_arr = Array1::from_vec(subject_emb.values);
        let predicate_arr = Array1::from_vec(predicate_emb.values);
        let object_arr = Array1::from_vec(object_emb.values);

        let predicted = &subject_arr + &predicate_arr;
        let diff = &predicted - &object_arr;
        let distance = diff.dot(&diff).sqrt();

        Ok(-distance as f64)
    }

    fn predict_objects(
        &self,
        subject: &str,
        predicate: &str,
        k: usize,
    ) -> Result<Vec<(String, f64)>> {
        let mut scores = Vec::new();

        for entity in self.entities.keys() {
            if entity != subject {
                let score = self.score_triple(subject, predicate, entity)?;
                scores.push((entity.clone(), score));
            }
        }

        scores.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
        scores.truncate(k);

        Ok(scores)
    }

    fn predict_subjects(
        &self,
        predicate: &str,
        object: &str,
        k: usize,
    ) -> Result<Vec<(String, f64)>> {
        let mut scores = Vec::new();

        for entity in self.entities.keys() {
            if entity != object {
                let score = self.score_triple(entity, predicate, object)?;
                scores.push((entity.clone(), score));
            }
        }

        scores.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
        scores.truncate(k);

        Ok(scores)
    }

    fn predict_relations(
        &self,
        subject: &str,
        object: &str,
        k: usize,
    ) -> Result<Vec<(String, f64)>> {
        let mut scores = Vec::new();

        for relation in self.relations.keys() {
            let score = self.score_triple(subject, relation, object)?;
            scores.push((relation.clone(), score));
        }

        scores.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
        scores.truncate(k);

        Ok(scores)
    }

    fn get_entities(&self) -> Vec<String> {
        self.entities.keys().cloned().collect()
    }

    fn get_relations(&self) -> Vec<String> {
        self.relations.keys().cloned().collect()
    }

    fn get_stats(&self) -> crate::ModelStats {
        crate::ModelStats {
            num_entities: self.entities.len(),
            num_relations: self.relations.len(),
            num_triples: 0,
            dimensions: self.config.base_config.dimensions,
            is_trained: self.is_trained,
            model_type: self.model_type().to_string(),
            creation_time: Utc::now(),
            last_training_time: if self.is_trained {
                Some(Utc::now())
            } else {
                None
            },
        }
    }

    fn save(&self, _path: &str) -> Result<()> {
        Ok(())
    }

    fn load(&mut self, _path: &str) -> Result<()> {
        Ok(())
    }

    fn clear(&mut self) {
        self.entities.clear();
        self.relations.clear();
        self.embeddings = Array2::zeros((0, self.config.base_config.dimensions));
        self.replay_buffer.clear();
        self.performance_tracker =
            OnlinePerformanceTracker::new(self.config.online_evaluation.window_size);
        self.examples_seen = 0;
        self.is_trained = false;
        self.training_stats = None;
    }

    fn is_trained(&self) -> bool {
        self.is_trained
    }

    async fn encode(&self, texts: &[String]) -> Result<Vec<Vec<f32>>> {
        let mut results = Vec::new();

        for text in texts {
            // Simple text encoding
            let mut embedding = vec![0.0f32; self.config.base_config.dimensions];
            for (i, c) in text.chars().enumerate() {
                if i >= self.config.base_config.dimensions {
                    break;
                }
                embedding[i] = (c as u8 as f32) / 255.0;
            }
            results.push(embedding);
        }

        Ok(results)
    }
}

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

    #[test]
    fn test_real_time_finetuning_config_default() {
        let config = RealTimeFinetuningConfig::default();
        assert_eq!(config.online_learning_rate, 1e-4);
        assert_eq!(config.replay_buffer_size, 10000);
        assert_eq!(config.online_batch_size, 32);
    }

    #[test]
    fn test_experience_entry_creation() {
        let entry = ExperienceEntry {
            input: Array1::from_vec(vec![1.0, 2.0, 3.0]),
            target: Array1::from_vec(vec![4.0, 5.0, 6.0]),
            timestamp: Utc::now(),
            importance: 1.0,
            task_id: Some("task1".to_string()),
        };

        assert_eq!(entry.input.len(), 3);
        assert_eq!(entry.target.len(), 3);
        assert!(entry.importance > 0.0);
    }

    #[test]
    fn test_online_performance_tracker() {
        let mut tracker = OnlinePerformanceTracker::new(10);
        tracker.update_metrics(0.5, 0.8, 0.1, 0.2);

        assert_eq!(tracker.get_average_loss(), 0.5);
        assert_eq!(tracker.get_average_accuracy(), 0.8);
        assert_eq!(tracker.update_count, 1);
    }

    #[test]
    fn test_real_time_finetuning_model_creation() {
        let config = RealTimeFinetuningConfig::default();
        let model = RealTimeFinetuningModel::new(config);

        assert_eq!(model.entities.len(), 0);
        assert_eq!(model.examples_seen, 0);
        assert!(!model.is_adapting);
    }

    #[tokio::test]
    async fn test_add_example_and_adaptation() {
        let config = RealTimeFinetuningConfig {
            base_config: ModelConfig {
                dimensions: 3, // Match array size
                ..Default::default()
            },
            update_frequency: 1, // Adapt on every example
            ..Default::default()
        };
        let mut model = RealTimeFinetuningModel::new(config);

        let input = Array1::from_vec(vec![1.0, 2.0, 3.0]);
        let target = Array1::from_vec(vec![4.0, 5.0, 6.0]);

        model
            .add_example(input, target, Some("task1".to_string()))
            .await
            .expect("should succeed");

        assert_eq!(model.examples_seen, 1);
        assert_eq!(model.replay_buffer.len(), 1);
    }

    #[tokio::test]
    async fn test_task_memory_management() {
        let config = RealTimeFinetuningConfig::default();
        let mut model = RealTimeFinetuningModel::new(config);

        // Initialize embeddings
        model.embeddings = Array2::from_shape_fn((5, 10), |_| {
            let mut random = Random::default();
            random.random::<f32>()
        });

        // Save task parameters
        model
            .save_task_parameters("task1".to_string())
            .expect("should succeed");

        // Modify embeddings
        model.embeddings *= 2.0;

        // Load task parameters
        model.load_task_parameters("task1").expect("should succeed");

        assert!(model.task_memory.contains_key("task1"));
    }

    #[test]
    fn test_online_stats() {
        let mut config = RealTimeFinetuningConfig::default();
        config.online_evaluation.window_size = 5;
        let model = RealTimeFinetuningModel::new(config);

        let stats = model.get_online_stats();

        assert!(stats.contains_key("average_loss"));
        assert!(stats.contains_key("examples_seen"));
        assert!(stats.contains_key("replay_buffer_size"));
        assert_eq!(stats["examples_seen"], 0.0);
    }

    #[tokio::test]
    async fn test_real_time_training() {
        let config = RealTimeFinetuningConfig {
            base_config: ModelConfig {
                dimensions: 3, // Match array size
                ..Default::default()
            },
            ..Default::default()
        };
        let mut model = RealTimeFinetuningModel::new(config);

        // Add some examples to replay buffer
        let input = Array1::from_vec(vec![1.0, 2.0, 3.0]);
        let target = Array1::from_vec(vec![4.0, 5.0, 6.0]);
        model
            .add_example(input, target, None)
            .await
            .expect("should succeed");

        let stats = model.train(Some(5)).await.expect("should succeed");
        assert_eq!(stats.epochs_completed, 5);
        assert!(model.is_trained());
    }
}