oxirs-embed 0.2.4

//! Vision-Language-Graph Multi-Modal Integration
//!
//! This module implements advanced multi-modal integration for vision, language, and knowledge graphs
//! with features including:
//! - Multi-modal transformers with cross-attention
//! - Joint representation learning
//! - Zero-shot and few-shot transfer learning
//! - Meta-learning for adaptation
//! - Vision-text-graph unified embedding spaces

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

pub mod config;
pub mod encoders;
pub mod meta_learner;
pub mod transformer;

pub use config::*;
pub use encoders::*;
pub use meta_learner::*;
pub use transformer::*;

/// Vision-Language-Graph embedding model
#[derive(Debug)]
pub struct VisionLanguageGraphModel {
    pub config: VisionLanguageGraphConfig,
    pub model_id: Uuid,
    /// Vision encoder
    pub vision_encoder: VisionEncoder,
    /// Language encoder
    pub language_encoder: LanguageEncoder,
    /// Graph encoder
    pub graph_encoder: GraphEncoder,
    /// Multi-modal transformer
    pub multimodal_transformer: MultiModalTransformer,
    /// Meta-learner for adaptation
    pub meta_learner: MetaLearner,
    /// Cached embeddings
    pub vision_embeddings: HashMap<String, Array1<f32>>,
    pub language_embeddings: HashMap<String, Array1<f32>>,
    pub graph_embeddings: HashMap<String, Array1<f32>>,
    pub unified_embeddings: HashMap<String, Array1<f32>>,
    /// Training state
    pub training_stats: Option<TrainingStats>,
    pub is_trained: bool,
}

/// Vision encoder
impl VisionLanguageGraphModel {
    /// Create new vision-language-graph model
    pub fn new(config: VisionLanguageGraphConfig) -> Self {
        let model_id = Uuid::new_v4();

        let vision_encoder = VisionEncoder::new(config.vision_config.clone());
        let language_encoder = LanguageEncoder::new(config.language_config.clone());
        let graph_encoder = GraphEncoder::new(config.graph_config.clone());
        let multimodal_transformer = MultiModalTransformer::new(config.transformer_config.clone());
        let meta_learner = MetaLearner::new(config.meta_learning_config.clone());

        Self {
            config,
            model_id,
            vision_encoder,
            language_encoder,
            graph_encoder,
            multimodal_transformer,
            meta_learner,
            vision_embeddings: HashMap::new(),
            language_embeddings: HashMap::new(),
            graph_embeddings: HashMap::new(),
            unified_embeddings: HashMap::new(),
            training_stats: None,
            is_trained: false,
        }
    }

    /// Generate unified multi-modal embedding
    pub async fn generate_unified_embedding(
        &mut self,
        image: Option<&Array3<f32>>,
        text: Option<&str>,
        graph_data: Option<(&Array2<f32>, &Array2<f32>, &Array2<f32>)>,
    ) -> Result<Array1<f32>> {
        let mut embeddings = Vec::new();

        // Vision embedding
        let vision_emb = if let Some(img) = image {
            let emb = self.vision_encoder.encode_image(img)?;
            self.vision_embeddings
                .insert("current_image".to_string(), emb.clone());
            emb
        } else {
            Array1::zeros(self.config.vision_config.vision_dim)
        };
        embeddings.push(vision_emb.clone());

        // Language embedding
        let language_emb = if let Some(txt) = text {
            let emb = self.language_encoder.encode_text(txt)?;
            self.language_embeddings
                .insert("current_text".to_string(), emb.clone());
            emb
        } else {
            Array1::zeros(self.config.language_config.language_dim)
        };
        embeddings.push(language_emb.clone());

        // Graph embedding
        let graph_emb = if let Some((nodes, edges, adj)) = graph_data {
            let emb = self.graph_encoder.encode_graph(nodes, edges, adj)?;
            self.graph_embeddings
                .insert("current_graph".to_string(), emb.clone());
            emb
        } else {
            Array1::zeros(self.config.graph_config.graph_dim)
        };
        embeddings.push(graph_emb.clone());

        // Fuse embeddings
        let unified_emb =
            self.multimodal_transformer
                .fuse_embeddings(&vision_emb, &language_emb, &graph_emb)?;

        self.unified_embeddings
            .insert("current_unified".to_string(), unified_emb.clone());

        Ok(unified_emb)
    }

    /// Zero-shot prediction
    pub fn zero_shot_predict(
        &self,
        query_embedding: &Array1<f32>,
        class_prototypes: &HashMap<String, Array1<f32>>,
    ) -> Result<String> {
        let mut best_class = String::new();
        let mut best_score = f32::NEG_INFINITY;

        for (class_name, prototype) in class_prototypes {
            let score = self.cosine_similarity(query_embedding, prototype);
            if score > best_score {
                best_score = score;
                best_class = class_name.clone();
            }
        }

        Ok(best_class)
    }

    /// Few-shot adaptation
    pub fn few_shot_adapt(
        &mut self,
        support_examples: &[(Array1<f32>, String)],
        query_examples: &[Array1<f32>],
    ) -> Result<Vec<String>> {
        // Convert support examples to meta-learning format
        let support_set: Vec<(Array1<f32>, Array1<f32>)> = support_examples
            .iter()
            .map(|(emb, label)| {
                let label_emb = Array1::from_vec(vec![label.len() as f32]); // Simplified label encoding
                (emb.clone(), label_emb)
            })
            .collect();

        let query_set: Vec<(Array1<f32>, Array1<f32>)> = query_examples
            .iter()
            .map(|emb| (emb.clone(), Array1::zeros(1)))
            .collect();

        // Adapt meta-learner
        let _adapted_params = self.meta_learner.adapt_to_task(&support_set, &query_set)?;

        // Make predictions on query set
        let mut predictions = Vec::new();

        for query_emb in query_examples {
            // Find nearest support example
            let mut best_label = String::new();
            let mut best_distance = f32::INFINITY;

            for (support_emb, label) in support_examples {
                let distance = self.euclidean_distance(query_emb, support_emb);
                if distance < best_distance {
                    best_distance = distance;
                    best_label = label.clone();
                }
            }

            predictions.push(best_label);
        }

        Ok(predictions)
    }

    /// Cosine similarity
    fn cosine_similarity(&self, a: &Array1<f32>, b: &Array1<f32>) -> f32 {
        let dot_product = a.dot(b);
        let norm_a = a.dot(a).sqrt();
        let norm_b = b.dot(b).sqrt();

        if norm_a > 0.0 && norm_b > 0.0 {
            dot_product / (norm_a * norm_b)
        } else {
            0.0
        }
    }

    /// Euclidean distance
    fn euclidean_distance(&self, a: &Array1<f32>, b: &Array1<f32>) -> f32 {
        let diff = a - b;
        diff.dot(&diff).sqrt()
    }
}

/// Multi-modal statistics
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct VisionLanguageGraphStats {
    pub num_vision_samples: usize,
    pub num_language_samples: usize,
    pub num_graph_samples: usize,
    pub num_unified_embeddings: usize,
    pub vision_dim: usize,
    pub language_dim: usize,
    pub graph_dim: usize,
    pub unified_dim: usize,
    pub zero_shot_accuracy: f32,
    pub few_shot_accuracy: f32,
    pub cross_modal_alignment_score: f32,
}

impl Default for VisionLanguageGraphStats {
    fn default() -> Self {
        Self {
            num_vision_samples: 0,
            num_language_samples: 0,
            num_graph_samples: 0,
            num_unified_embeddings: 0,
            vision_dim: 768,
            language_dim: 768,
            graph_dim: 512,
            unified_dim: 768,
            zero_shot_accuracy: 0.0,
            few_shot_accuracy: 0.0,
            cross_modal_alignment_score: 0.0,
        }
    }
}

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

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

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

    fn add_triple(&mut self, _triple: Triple) -> Result<()> {
        // Implementation would process triples for graph structure
        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 multi-modal training
            let epoch_loss = self.train_epoch().await?;
            loss_history.push(epoch_loss);

            if epoch > 10 && epoch_loss < 1e-4 {
                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(embedding) = self.unified_embeddings.get(entity) {
            Ok(Vector::new(embedding.to_vec()))
        } else {
            Err(anyhow!("Entity not found: {}", entity))
        }
    }

    fn get_relation_embedding(&self, relation: &str) -> Result<Vector> {
        if let Some(embedding) = self.unified_embeddings.get(relation) {
            Ok(Vector::new(embedding.to_vec()))
        } else {
            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.unified_embeddings.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.unified_embeddings.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.unified_embeddings.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.unified_embeddings.keys().cloned().collect()
    }

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

    fn get_stats(&self) -> ModelStats {
        ModelStats {
            num_entities: self.unified_embeddings.len(),
            num_relations: self.unified_embeddings.len(),
            num_triples: 0,
            dimensions: self.config.transformer_config.unified_dim,
            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.vision_embeddings.clear();
        self.language_embeddings.clear();
        self.graph_embeddings.clear();
        self.unified_embeddings.clear();
        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 {
            let embedding = self.language_encoder.encode_text(text)?;
            results.push(embedding.to_vec());
        }

        Ok(results)
    }
}

impl VisionLanguageGraphModel {
    /// Training epoch for multi-modal model
    async fn train_epoch(&mut self) -> Result<f64> {
        // Simulate multi-modal training loss
        let mut random = Random::default();
        let vision_loss = 0.1 * random.random::<f64>();
        let language_loss = 0.1 * random.random::<f64>();
        let graph_loss = 0.1 * random.random::<f64>();
        let fusion_loss = 0.1 * random.random::<f64>();

        let total_loss = vision_loss + language_loss + graph_loss + fusion_loss;

        Ok(total_loss)
    }
}

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

    /// Create a small VisionLanguageGraphConfig suitable for fast tests in debug builds.
    /// All dimensions are reduced to avoid multi-second matrix allocations.
    fn small_test_config() -> VisionLanguageGraphConfig {
        VisionLanguageGraphConfig {
            vision_config: VisionEncoderConfig {
                image_size: (32, 32),
                channels: 3,
                patch_size: (8, 8),
                vision_dim: 32,
                cnn_config: CNNConfig {
                    num_layers: 2,
                    filter_sizes: vec![8, 16],
                    stride_sizes: vec![2, 2],
                    ..CNNConfig::default()
                },
                vit_config: ViTConfig {
                    num_layers: 2,
                    num_heads: 2,
                    mlp_dim: 64,
                    ..ViTConfig::default()
                },
                ..VisionEncoderConfig::default()
            },
            language_config: LanguageEncoderConfig {
                vocab_size: 256,
                language_dim: 32,
                max_seq_length: 16,
                transformer_config: LanguageTransformerConfig {
                    num_layers: 2,
                    num_heads: 2,
                    hidden_dim: 32,
                    intermediate_dim: 64,
                    ..LanguageTransformerConfig::default()
                },
                ..LanguageEncoderConfig::default()
            },
            graph_config: GraphEncoderConfig {
                node_dim: 16,
                edge_dim: 8,
                graph_dim: 32,
                num_layers: 2,
                ..GraphEncoderConfig::default()
            },
            transformer_config: MultiModalTransformerConfig {
                unified_dim: 32,
                num_fusion_layers: 2,
                cross_attention_config: CrossAttentionConfig {
                    num_heads: 2,
                    head_dim: 16,
                    ..CrossAttentionConfig::default()
                },
                ..MultiModalTransformerConfig::default()
            },
            ..VisionLanguageGraphConfig::default()
        }
    }

    #[test]
    fn test_vision_language_graph_config_default() {
        let config = VisionLanguageGraphConfig::default();
        assert_eq!(config.vision_config.vision_dim, 768);
        assert_eq!(config.language_config.language_dim, 768);
        assert_eq!(config.graph_config.graph_dim, 768); // Updated to match unified_dim
    }

    #[test]
    fn test_vision_encoder_creation() {
        let config = VisionEncoderConfig::default();
        let encoder = VisionEncoder::new(config);
        assert!(!encoder.cnn_parameters.is_empty());
        assert!(!encoder.vit_parameters.is_empty());
    }

    #[test]
    fn test_language_encoder_creation() {
        // Use small dimensions to avoid timeout in debug builds
        let config = LanguageEncoderConfig {
            vocab_size: 256,
            language_dim: 32,
            max_seq_length: 16,
            transformer_config: LanguageTransformerConfig {
                num_layers: 2,
                num_heads: 2,
                hidden_dim: 32,
                intermediate_dim: 64,
                ..LanguageTransformerConfig::default()
            },
            ..LanguageEncoderConfig::default()
        };
        let encoder = LanguageEncoder::new(config);
        assert_eq!(encoder.token_embeddings.nrows(), 256);
        assert_eq!(encoder.position_embeddings.nrows(), 16);
    }

    #[test]
    fn test_graph_encoder_creation() {
        let config = GraphEncoderConfig::default();
        let encoder = GraphEncoder::new(config);
        assert!(!encoder.node_parameters.is_empty());
        assert!(!encoder.edge_parameters.is_empty());
    }

    #[test]
    fn test_multimodal_transformer_creation() {
        let config = MultiModalTransformerConfig::default();
        let transformer = MultiModalTransformer::new(config);
        assert!(!transformer.cross_attention_params.is_empty());
        assert!(!transformer.fusion_params.is_empty());
    }

    #[test]
    #[cfg_attr(debug_assertions, ignore = "Model initialization slow in debug builds")]
    fn test_vision_language_graph_model_creation() {
        let config = VisionLanguageGraphConfig::default();
        let model = VisionLanguageGraphModel::new(config);
        assert!(!model.is_trained);
        assert_eq!(model.unified_embeddings.len(), 0);
    }

    #[test]
    fn test_vision_encoder_image_encoding() {
        let config = VisionEncoderConfig::default();
        let encoder = VisionEncoder::new(config);

        let mut random = Random::default();
        let image = Array3::from_shape_fn((224, 224, 3), |_| random.random::<f32>());
        let embedding = encoder.encode_image(&image).expect("should succeed");

        assert_eq!(embedding.len(), encoder.config.vision_dim);
    }

    #[test]
    fn test_language_encoder_text_encoding() {
        // Use small dimensions to avoid timeout in debug builds
        let config = LanguageEncoderConfig {
            vocab_size: 256,
            language_dim: 32,
            max_seq_length: 16,
            transformer_config: LanguageTransformerConfig {
                num_layers: 2,
                num_heads: 2,
                hidden_dim: 32,
                intermediate_dim: 64,
                ..LanguageTransformerConfig::default()
            },
            ..LanguageEncoderConfig::default()
        };
        let encoder = LanguageEncoder::new(config);

        let text = "Hello world, this is a test";
        let embedding = encoder
            .encode_text(text)
            .expect("encode_text should succeed");

        assert_eq!(embedding.len(), encoder.config.language_dim);
    }

    #[test]
    fn test_graph_encoder_graph_encoding() {
        let config = GraphEncoderConfig::default();
        let node_dim = config.node_dim;
        let edge_dim = config.edge_dim;
        let encoder = GraphEncoder::new(config);

        let mut random = Random::default();
        let node_features = Array2::from_shape_fn((5, node_dim), |_| random.random::<f32>());
        let edge_features = Array2::from_shape_fn((10, edge_dim), |_| random.random::<f32>());
        let adjacency = Array2::eye(5);

        let embedding = encoder
            .encode_graph(&node_features, &edge_features, &adjacency)
            .expect("should succeed");

        assert_eq!(embedding.len(), encoder.config.graph_dim);
    }

    #[tokio::test]
    #[cfg_attr(debug_assertions, ignore = "Embedding tests require release builds")]
    async fn test_unified_embedding_generation() {
        let config = VisionLanguageGraphConfig::default();
        let mut model = VisionLanguageGraphModel::new(config);

        let mut random = Random::default();
        let image = Array3::from_shape_fn((224, 224, 3), |_| random.random::<f32>());
        let text = "A beautiful landscape with mountains";
        let node_features = Array2::from_shape_fn((3, 256), |_| random.random::<f32>());
        let edge_features = Array2::from_shape_fn((6, 128), |_| random.random::<f32>());
        let adjacency = Array2::eye(3);

        let unified_embedding = model
            .generate_unified_embedding(
                Some(&image),
                Some(text),
                Some((&node_features, &edge_features, &adjacency)),
            )
            .await
            .expect("should succeed");

        assert!(!unified_embedding.is_empty());
        assert_eq!(model.vision_embeddings.len(), 1);
        assert_eq!(model.language_embeddings.len(), 1);
        assert_eq!(model.graph_embeddings.len(), 1);
        assert_eq!(model.unified_embeddings.len(), 1);
    }

    #[test]
    fn test_zero_shot_prediction() {
        let config = small_test_config();
        let model = VisionLanguageGraphModel::new(config);

        let mut random = Random::default();
        let query = Array1::from_shape_fn(32, |_| random.random::<f32>());

        let mut prototypes = HashMap::new();
        let mut random = Random::default();
        prototypes.insert(
            "class1".to_string(),
            Array1::from_shape_fn(32, |_| random.random::<f32>()),
        );
        let mut random = Random::default();
        prototypes.insert(
            "class2".to_string(),
            Array1::from_shape_fn(32, |_| random.random::<f32>()),
        );

        let prediction = model
            .zero_shot_predict(&query, &prototypes)
            .expect("zero_shot_predict should succeed");
        assert!(prototypes.contains_key(&prediction));
    }

    #[test]
    #[cfg_attr(debug_assertions, ignore = "Embedding tests require release builds")]
    fn test_few_shot_adaptation() {
        let config = VisionLanguageGraphConfig::default();
        let mut model = VisionLanguageGraphModel::new(config);

        let mut random = Random::default();
        let support_examples = vec![
            (
                Array1::from_shape_fn(512, |_| random.random::<f32>()),
                "cat".to_string(),
            ),
            (
                Array1::from_shape_fn(512, |_| random.random::<f32>()),
                "dog".to_string(),
            ),
        ];

        let mut random = Random::default();
        let query_examples = vec![
            Array1::from_shape_fn(512, |_| random.random::<f32>()),
            Array1::from_shape_fn(512, |_| random.random::<f32>()),
        ];

        let predictions = model
            .few_shot_adapt(&support_examples, &query_examples)
            .expect("should succeed");
        assert_eq!(predictions.len(), 2);
    }

    #[test]
    fn test_meta_learner_adaptation() {
        let config = MetaLearningConfig::default();
        let mut meta_learner = MetaLearner::new(config);

        let mut random = Random::default();
        let support_set = vec![
            (
                Array1::from_shape_fn(512, |_| random.random::<f32>()),
                Array1::from_vec(vec![1.0]),
            ),
            (
                Array1::from_shape_fn(512, |_| random.random::<f32>()),
                Array1::from_vec(vec![0.0]),
            ),
        ];

        let query_set = vec![];

        let adapted_params = meta_learner
            .adapt_to_task(&support_set, &query_set)
            .expect("should succeed");
        assert!(!adapted_params.is_empty());
    }

    #[tokio::test]
    async fn test_vision_language_graph_training() {
        let config = small_test_config();
        let mut model = VisionLanguageGraphModel::new(config);

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

    #[tokio::test]
    #[cfg_attr(debug_assertions, ignore = "Embedding tests require release builds")]
    async fn test_vision_language_graph_encoding() {
        let config = VisionLanguageGraphConfig::default();
        let expected_dim = config.language_config.language_dim;
        let model = VisionLanguageGraphModel::new(config);

        let texts = vec!["hello world".to_string(), "test encoding".to_string()];
        let embeddings = model.encode(&texts).await.expect("should succeed");

        assert_eq!(embeddings.len(), 2);
        assert_eq!(embeddings[0].len(), expected_dim);
    }
}