oxirs-embed 0.3.1

//! Neural-Symbolic Integration — model engine
//!
//! Defines [`NeuralSymbolicModel`], its inherent inference/learning methods, and
//! its [`EmbeddingModel`] trait implementation, combining a neural forward pass
//! with symbolic rule application and constraint satisfaction.

use crate::neural_symbolic_integration_types::{
    ActivationFunction, KnowledgeRule, LogicalFormula, NeuralSymbolicConfig,
};
use crate::{EmbeddingModel, ModelConfig, 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 std::collections::HashMap;
use uuid::Uuid;

/// Neural-symbolic integration model
#[derive(Debug)]
pub struct NeuralSymbolicModel {
    pub config: NeuralSymbolicConfig,
    pub model_id: Uuid,

    /// Neural components
    pub neural_layers: Vec<Array2<f32>>,
    pub attention_weights: Array3<f32>,

    /// Symbolic components
    pub knowledge_base: Vec<KnowledgeRule>,
    pub logical_formulas: Vec<LogicalFormula>,
    pub symbol_embeddings: HashMap<String, Array1<f32>>,

    /// Integration layers
    pub neural_to_symbolic: Array2<f32>,
    pub symbolic_to_neural: Array2<f32>,
    pub fusion_weights: Array2<f32>,

    /// Constraint satisfaction
    pub constraints: Vec<LogicalFormula>,
    pub constraint_weights: Array1<f32>,

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

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

impl NeuralSymbolicModel {
    /// Create new neural-symbolic model
    pub fn new(config: NeuralSymbolicConfig) -> Self {
        let model_id = Uuid::new_v4();
        let dimensions = config.base_config.dimensions;

        // Initialize neural layers with proper dimensions
        let mut neural_layers = Vec::new();
        let layer_configs = &config.architecture_config.neural_config.layers;

        for (i, layer_config) in layer_configs.iter().enumerate() {
            let input_size = if i == 0 {
                dimensions // First layer takes configured input dimension
            } else {
                layer_configs[i - 1].size // Subsequent layers take previous layer's output
            };

            let output_size = if i == layer_configs.len() - 1 {
                dimensions // Last layer outputs configured dimension
            } else {
                layer_config.size // Middle layers use configured size
            };

            neural_layers.push(Array2::from_shape_fn((output_size, input_size), |_| {
                let mut random = Random::default();
                random.random::<f32>() * 0.1
            }));
        }

        Self {
            config,
            model_id,
            neural_layers,
            attention_weights: Array3::from_shape_fn((8, dimensions, dimensions), |_| {
                let mut random = Random::default();
                random.random::<f32>() * 0.1
            }),
            knowledge_base: Vec::new(),
            logical_formulas: Vec::new(),
            symbol_embeddings: HashMap::new(),
            neural_to_symbolic: Array2::from_shape_fn((dimensions, dimensions), |_| {
                let mut random = Random::default();
                random.random::<f32>() * 0.1
            }),
            symbolic_to_neural: Array2::from_shape_fn((dimensions, dimensions), |_| {
                let mut random = Random::default();
                random.random::<f32>() * 0.1
            }),
            fusion_weights: Array2::from_shape_fn((dimensions, dimensions * 2), |_| {
                let mut random = Random::default();
                random.random::<f32>() * 0.1
            }),
            constraints: Vec::new(),
            constraint_weights: Array1::from_shape_fn(10, |_| 1.0),
            entities: HashMap::new(),
            relations: HashMap::new(),
            training_stats: None,
            is_trained: false,
        }
    }

    /// Add knowledge rule
    pub fn add_knowledge_rule(&mut self, rule: KnowledgeRule) {
        self.knowledge_base.push(rule);
    }

    /// Add logical constraint
    pub fn add_constraint(&mut self, constraint: LogicalFormula) {
        self.constraints.push(constraint);
    }

    /// Forward pass through neural component
    fn neural_forward(&self, input: &Array1<f32>) -> Result<Array1<f32>> {
        let mut activation = input.clone();

        for (i, layer) in self.neural_layers.iter().enumerate() {
            // Linear transformation
            activation = layer.dot(&activation);

            // Apply activation function
            let activation_fn = &self.config.architecture_config.neural_config.activations[i];
            activation = match activation_fn {
                ActivationFunction::ReLU => activation.mapv(|x| x.max(0.0)),
                ActivationFunction::Sigmoid => activation.mapv(|x| 1.0 / (1.0 + (-x).exp())),
                ActivationFunction::Tanh => activation.mapv(|x| x.tanh()),
                ActivationFunction::GELU => {
                    activation.mapv(|x| x * 0.5 * (1.0 + (x * 0.797_884_6).tanh()))
                }
                ActivationFunction::Swish => activation.mapv(|x| x * (1.0 / (1.0 + (-x).exp()))),
                ActivationFunction::LogicActivation => activation.mapv(|x| (x.tanh() + 1.0) / 2.0), // Maps to [0,1]
                _ => activation.mapv(|x| x.max(0.0)),
            };
        }

        Ok(activation)
    }

    /// Forward pass through symbolic component
    fn symbolic_forward(&self, input: &Array1<f32>) -> Result<Array1<f32>> {
        let mut symbolic_state = HashMap::new();

        // Ground neural input to symbolic facts
        for (i, &value) in input.iter().enumerate() {
            let symbol = format!("input_{i}");
            symbolic_state.insert(symbol, value);
        }

        // Apply knowledge rules
        let mut inferred_facts = symbolic_state.clone();

        for _ in 0..self.config.symbolic_config.rule_based.max_depth {
            let mut new_facts = inferred_facts.clone();
            let mut facts_added = false;

            for rule in &self.knowledge_base {
                if let Some((predicate, value)) = rule.apply(&inferred_facts) {
                    if !new_facts.contains_key(&predicate) || new_facts[&predicate] < value {
                        new_facts.insert(predicate, value);
                        facts_added = true;
                    }
                }
            }

            if !facts_added {
                break;
            }

            inferred_facts = new_facts;
        }

        // Convert back to vector
        let mut output = Array1::zeros(input.len());
        for (i, symbol) in (0..input.len()).map(|i| format!("output_{i}")).enumerate() {
            if let Some(&value) = inferred_facts.get(&symbol) {
                output[i] = value;
            }
        }

        Ok(output)
    }

    /// Integrate neural and symbolic components
    pub fn integrated_forward(&self, input: &Array1<f32>) -> Result<Array1<f32>> {
        // Neural forward pass
        let neural_output = self.neural_forward(input)?;

        // Map neural output to symbolic space
        let symbolic_input = self.neural_to_symbolic.dot(&neural_output);

        // Symbolic forward pass
        let symbolic_output = self.symbolic_forward(&symbolic_input)?;

        // Map symbolic output back to neural space
        let neural_symbolic_output = self.symbolic_to_neural.dot(&symbolic_output);

        // Fuse neural and neural-symbolic outputs
        let fused_input = Array1::from_iter(
            neural_output
                .iter()
                .chain(neural_symbolic_output.iter())
                .cloned(),
        );

        let fused_output = self.fusion_weights.dot(&fused_input);

        // Apply constraints
        let constrained_output = self.apply_constraints(fused_output)?;

        Ok(constrained_output)
    }

    /// Apply logical constraints
    fn apply_constraints(&self, mut output: Array1<f32>) -> Result<Array1<f32>> {
        if self.constraints.is_empty() {
            return Ok(output);
        }

        // Convert output to symbolic facts
        let mut facts = HashMap::new();
        for (i, &value) in output.iter().enumerate() {
            facts.insert(format!("output_{i}"), value);
        }

        // Evaluate constraints and adjust output
        for (constraint, &weight) in self.constraints.iter().zip(self.constraint_weights.iter()) {
            let constraint_satisfaction = constraint.evaluate(&facts);

            // If constraint is not satisfied, adjust output
            if constraint_satisfaction < 0.8 {
                let adjustment_factor = (0.8 - constraint_satisfaction) * weight * 0.1;
                output *= 1.0 - adjustment_factor;
            }
        }

        Ok(output)
    }

    /// Learn symbolic rules from examples
    pub fn learn_symbolic_rules(&mut self, examples: &[(Array1<f32>, Array1<f32>)]) -> Result<()> {
        // Simple rule learning algorithm
        let mut candidate_rules = Vec::new();

        for (input, output) in examples.iter() {
            // Create candidate rules based on input-output patterns
            for j in 0..input.len() {
                for k in 0..output.len() {
                    if input[j] > 0.5 && output[k] > 0.5 {
                        // Create rule: input_j -> output_k
                        let antecedent = LogicalFormula::new_atom(format!("input_{j}"));
                        let consequent = LogicalFormula::new_atom(format!("output_{k}"));
                        let rule =
                            KnowledgeRule::new(format!("rule_{j}_{k}"), antecedent, consequent);
                        candidate_rules.push(rule);
                    }
                }
            }
        }

        // Evaluate and filter rules based on support and confidence
        for rule in candidate_rules {
            let mut support = 0;
            let mut confidence_sum = 0.0;

            for (input, output) in examples {
                let mut facts = HashMap::new();
                for (i, &value) in input.iter().enumerate() {
                    facts.insert(format!("input_{i}"), value);
                }

                if let Some((predicate, predicted_value)) = rule.apply(&facts) {
                    if let Some(index) = predicate
                        .strip_prefix("output_")
                        .and_then(|s| s.parse::<usize>().ok())
                    {
                        if index < output.len() {
                            let actual_value = output[index];
                            let error = (predicted_value - actual_value).abs();
                            if error < 0.2 {
                                support += 1;
                                confidence_sum += 1.0 - error;
                            }
                        }
                    }
                }
            }

            if support >= 3 && confidence_sum / support as f32 > 0.7 {
                self.add_knowledge_rule(rule);
            }
        }

        Ok(())
    }

    /// Compute semantic loss
    pub fn compute_semantic_loss(
        &self,
        predictions: &Array1<f32>,
        targets: &Array1<f32>,
    ) -> Result<f32> {
        // Standard MSE loss
        let mse_loss = {
            let diff = predictions - targets;
            diff.dot(&diff) / predictions.len() as f32
        };

        // Constraint violation loss
        let constraint_loss = {
            let mut facts = HashMap::new();
            for (i, &value) in predictions.iter().enumerate() {
                facts.insert(format!("output_{i}"), value);
            }

            let mut total_violation = 0.0;
            for constraint in &self.constraints {
                let satisfaction = constraint.evaluate(&facts);
                if satisfaction < 1.0 {
                    total_violation += (1.0 - satisfaction).powi(2);
                }
            }
            total_violation / self.constraints.len().max(1) as f32
        };

        // Rule consistency loss
        let rule_loss = {
            let mut facts = HashMap::new();
            for (i, &value) in predictions.iter().enumerate() {
                facts.insert(format!("input_{i}"), value);
            }

            let mut total_inconsistency = 0.0;
            for rule in &self.knowledge_base {
                if let Some((predicate, predicted_value)) = rule.apply(&facts) {
                    if let Some(index) = predicate
                        .strip_prefix("output_")
                        .and_then(|s| s.parse::<usize>().ok())
                    {
                        if index < predictions.len() {
                            let actual_value = predictions[index];
                            let inconsistency = (predicted_value - actual_value).powi(2);
                            total_inconsistency += inconsistency * rule.weight;
                        }
                    }
                }
            }
            total_inconsistency / self.knowledge_base.len().max(1) as f32
        };

        // Combine losses
        let total_loss = mse_loss + 0.1 * constraint_loss + 0.1 * rule_loss;

        Ok(total_loss)
    }

    /// Explain prediction using symbolic reasoning
    pub fn explain_prediction(
        &self,
        input: &Array1<f32>,
        prediction: &Array1<f32>,
    ) -> Result<String> {
        let mut explanation = String::new();
        explanation.push_str("Prediction Explanation:\n");

        // Ground input to facts
        let mut facts = HashMap::new();
        for (i, &value) in input.iter().enumerate() {
            facts.insert(format!("input_{i}"), value);
        }

        // Find activated rules
        let mut activated_rules = Vec::new();
        for rule in &self.knowledge_base {
            let antecedent_value = rule.antecedent.evaluate(&facts);
            if antecedent_value > 0.5 {
                activated_rules.push((rule, antecedent_value));
            }
        }

        if !activated_rules.is_empty() {
            explanation.push_str("\nActivated Rules:\n");
            for (rule, activation) in activated_rules {
                explanation.push_str(&format!(
                    "- Rule {}: {} (activation: {:.2})\n",
                    rule.id, rule.id, activation
                ));
            }
        }

        // Check constraint satisfaction
        let mut constraint_violations = Vec::new();
        let mut prediction_facts = HashMap::new();
        for (i, &value) in prediction.iter().enumerate() {
            prediction_facts.insert(format!("output_{i}"), value);
        }

        for constraint in &self.constraints {
            let satisfaction = constraint.evaluate(&prediction_facts);
            if satisfaction < 0.8 {
                constraint_violations.push(satisfaction);
            }
        }

        if !constraint_violations.is_empty() {
            explanation.push_str("\nConstraint Violations:\n");
            for (i, violation) in constraint_violations.iter().enumerate() {
                explanation.push_str(&format!(
                    "- Constraint {i}: satisfaction = {violation:.2}\n"
                ));
            }
        }

        Ok(explanation)
    }
}

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

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

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

    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.clone())
            .or_insert(next_entity_id);
        let next_entity_id = self.entities.len();
        self.entities
            .entry(object_str.clone())
            .or_insert(next_entity_id);

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

        // Create symbolic representation
        let rule_id = format!("{subject_str}_{predicate_str}");
        let antecedent = LogicalFormula::new_atom(subject_str);
        let consequent = LogicalFormula::new_atom(object_str);
        let rule = KnowledgeRule::new(rule_id, antecedent, consequent);
        self.add_knowledge_rule(rule);

        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 neural-symbolic training
            let epoch_loss = {
                let mut random = Random::default();
                0.1 * random.random::<f64>()
            };
            loss_history.push(epoch_loss);

            // Learn symbolic rules periodically
            if epoch % 10 == 0 && epoch > 0 {
                // Simulate learning from examples
                let examples = vec![
                    (
                        Array1::from_vec(vec![1.0, 0.0, 1.0]),
                        Array1::from_vec(vec![1.0, 1.0]),
                    ),
                    (
                        Array1::from_vec(vec![0.0, 1.0, 0.0]),
                        Array1::from_vec(vec![0.0, 1.0]),
                    ),
                ];
                self.learn_symbolic_rules(&examples)?;
            }

            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) {
            // Generate embedding from neural-symbolic integration
            let input = Array1::from_shape_fn(self.config.base_config.dimensions, |i| {
                if i == entity_id % self.config.base_config.dimensions {
                    1.0
                } else {
                    0.0
                }
            });

            if let Ok(embedding) = self.integrated_forward(&input) {
                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) {
            // Generate embedding from neural-symbolic integration
            let input = Array1::from_shape_fn(self.config.base_config.dimensions, |i| {
                if i == relation_id % self.config.base_config.dimensions {
                    1.0
                } else {
                    0.0
                }
            });

            if let Ok(embedding) = self.integrated_forward(&input) {
                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> {
        // Use symbolic reasoning for scoring
        let mut facts = HashMap::new();
        facts.insert(subject.to_string(), 1.0);
        facts.insert(predicate.to_string(), 1.0);

        // Check if any rules support this triple
        let mut max_score: f32 = 0.0;
        for rule in &self.knowledge_base {
            let antecedent_value = rule.antecedent.evaluate(&facts);
            let consequent_value = rule.consequent.evaluate(&facts);
            let rule_score = antecedent_value * consequent_value * rule.confidence;
            max_score = max_score.max(rule_score);
        }

        Ok(max_score 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.knowledge_base.clear();
        self.logical_formulas.clear();
        self.symbol_embeddings.clear();
        self.constraints.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 {
            // Use neural-symbolic integration for encoding
            let input = Array1::from_shape_fn(self.config.base_config.dimensions, |i| {
                if i < text.len() {
                    (text
                        .chars()
                        .nth(i)
                        .expect("index should be within text length") as u8
                        as f32)
                        / 255.0
                } else {
                    0.0
                }
            });

            match self.integrated_forward(&input) {
                Ok(embedding) => {
                    results.push(embedding.to_vec());
                }
                _ => {
                    results.push(vec![0.0; self.config.base_config.dimensions]);
                }
            }
        }

        Ok(results)
    }
}