//! Traits for TGNN implementation

use super::types::{TrainingMetrics, ValidationMetrics};
use ml_core::{InferenceResult, MLError, ModelMetadata};
use async_trait::async_trait;
use ndarray::Array2;

/// Core ML model trait for TGNN implementations
#[async_trait]
pub trait MLModel {
    type Config;

    /// Get model metadata
    fn metadata(&self) -> &ModelMetadata;

    /// Check if model is ready for inference
    fn is_ready(&self) -> bool;

    /// Train the model
    async fn train(
        &mut self,
        features: &Array2<f64>,
        targets: &Array2<f64>,
    ) -> Result<TrainingMetrics, MLError>;

    /// Predict using the model
    async fn predict(&self, features: &[f64]) -> Result<InferenceResult, MLError>;

    /// Validate model performance
    async fn validate(
        &self,
        features: &Array2<f64>,
        targets: &Array2<f64>,
    ) -> Result<ValidationMetrics, MLError>;

    /// Update model with new data (online learning)
    async fn update(
        &mut self,
        features: &Array2<f64>,
        targets: &Array2<f64>,
    ) -> Result<(), MLError>;

    /// Save model to file
    async fn save(&self, path: &str) -> Result<(), MLError>;

    /// Load model from file
    async fn load(&mut self, path: &str) -> Result<(), MLError>;

    /// Get model configuration
    fn config(&self) -> Self::Config;

    /// Set model configuration
    fn set_config(&mut self, config: Self::Config) -> Result<(), MLError>;
}

/// Trait for models that can predict from Array2 features (batch prediction)
#[async_trait]
pub trait BatchPredict {
    /// Predict from batch of features
    async fn predict_batch(&self, features: &Array2<f64>) -> Result<Vec<InferenceResult>, MLError>;
}

/// Trait for graph-based models
pub trait GraphModel {
    /// Add node to the model's internal graph
    fn add_node(&mut self, node_id: String, features: Vec<f64>) -> Result<(), MLError>;

    /// Remove node from the model's internal graph
    fn remove_node(&mut self, node_id: &str) -> Result<(), MLError>;

    /// Add edge between nodes
    fn add_edge(&mut self, from: &str, to: &str, weight: f64) -> Result<(), MLError>;

    /// Get neighbors of a node
    fn get_neighbors(&self, node_id: &str) -> Option<Vec<String>>;

    /// Update graph structure
    fn update_graph(&mut self) -> Result<(), MLError>;
}
