gbrt_rs/io/

model_io.rs

#![allow(dead_code)]

//! Model Persistence and Management for Gradient Boosted Models
//! 
//! This module provides comprehensive functionality for saving, loading, and managing
//! trained [`GradientBooster`] models. It supports multiple serialization formats,
//! model versioning, integrity verification, and a model registry for organizing
//! multiple models.
//! 
//! # Features
//! 
//! - **Multiple Formats**: Save models as binary (compact) or JSON (human-readable)
//! - **Metadata Tracking**: Store model configuration, training parameters, and custom metadata
//! - **Integrity Verification**: Validate model integrity during loading
//! - **Version Management**: Track model versions and detect compatibility issues
//! - **Model Registry**: Centralized management of multiple models with name-based lookup
//! - **Feature Importance**: Optionally export feature importance scores
//! - **Compression**: Optional gzip compression for reduced disk usage

use crate::boosting::{GradientBooster, GBRTConfig};  // Changed from core to boosting
use crate::utils::{ModelSerializer, ModelMetadata, SerializationFormat, SerializationError};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::fs::{File, create_dir_all, read_dir};
use std::io::{BufReader, BufWriter};
use std::path::{Path, PathBuf};
use thiserror::Error;


/// Errors that can occur during model I/O operations.
///
/// This error type covers all failure modes: file system errors, serialization
/// failures, version mismatches, integrity violations, and missing models.
#[derive(Error, Debug)]
pub enum ModelIOError {
    /// File system I/O error.
    #[error("IO error: {0}")]
    IoError(#[from] std::io::Error),

    /// Model serialization/deserialization error.
    #[error("Serialization error: {0}")]
    SerializationError(#[from] crate::utils::SerializationError),

    /// Parquet file reading error.
    #[error("Parquet error: {0}")]
    ParquetError(#[from] parquet::errors::ParquetError),

    /// General model error (e.g., untrained model).
    #[error("Model error: {0}")]
    ModelError(String),

    /// Model file is invalid or corrupted.
    #[error("Invalid model file: {0}")]
    InvalidModelFile(String),

    /// Version incompatibility between saved model and current library.
    #[error("Version mismatch: expected {expected}, got {actual}")]
    VersionMismatch { expected: String, actual: String },

    /// Unsupported file format.
    #[error("Unsupported format: {format}")]
    UnsupportedFormat { format: String },

    /// Model file not found at specified path.
    #[error("Model not found: {0}")]
    ModelNotFound(String),
}

pub type ModelIOResult<T> = std::result::Result<T, ModelIOError>;

/// Supported model serialization formats.
/// 
/// - **Binary**: Compact, fast serialization using bincode (default)
/// - **JSON**: Human-readable text format for inspection and debugging
#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
pub enum ModelFormat {
    /// Binary format for compact storage and fast loading.
    Binary,
    /// JSON format for human-readability and interoperability.
    Json,
}

impl std::str::FromStr for ModelFormat {
    type Err = ModelIOError;

    /// Parses ModelFormat from string.
    /// 
    /// # Supported Values
    /// - "binary", "bin" → `ModelFormat::Binary`
    /// - "json" → `ModelFormat::Json`
    /// 
    /// # Errors
    /// - `ModelIOError::UnsupportedFormat` if string doesn't match any known format
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s.to_lowercase().as_str() {
            "binary" | "bin" => Ok(ModelFormat::Binary),
            "json" => Ok(ModelFormat::Json),
            _ => Err(ModelIOError::UnsupportedFormat {
                format: s.to_string(),
            }),
        }
    }
}

impl std::fmt::Display for ModelFormat {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            ModelFormat::Binary => write!(f, "binary"),
            ModelFormat::Json => write!(f, "json"),
        }
    }
}

/// Configuration options for saving models.
/// 
/// Controls serialization format, compression, and what additional data to include.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SaveOptions {
    /// Serialization format to use (Binary or JSON).
    pub format: ModelFormat,
    /// Whether to compress the model file.
    pub compression: bool,
    /// Whether to include training data in the save (not yet implemented).
    pub include_training_data: bool,
    /// Whether to include feature importance scores.
    pub include_feature_importance: bool,
    /// Additional metadata to store with the model.
    pub metadata: HashMap<String, String>,
}

impl Default for SaveOptions {
    fn default() -> Self {
        Self {
            format: ModelFormat::Json,
            compression: false,
            include_training_data: false,
            include_feature_importance: true,
            metadata: HashMap::new(),
        }
    }
}

impl SaveOptions {
    /// Creates default SaveOptions.
    pub fn new() -> Self {
        Self::default()
    }

    /// Sets the serialization format.
    pub fn with_format(mut self, format: ModelFormat) -> Self {
        self.format = format;
        self
    }

    /// Enables or disables compression.
    pub fn with_compression(mut self, compression: bool) -> Self {
        self.compression = compression;
        self
    }

    /// Sets whether to include training data (future feature).
    pub fn with_training_data(mut self, include: bool) -> Self {
        self.include_training_data = include;
        self
    }

    /// Sets whether to include feature importance scores.
    pub fn with_feature_importance(mut self, include: bool) -> Self {
        self.include_feature_importance = include;
        self
    }

    /// Replaces all metadata with provided map.
    pub fn with_metadata(mut self, metadata: HashMap<String, String>) -> Self {
        self.metadata = metadata;
        self
    }

    /// Adds a single metadata key-value pair.
    pub fn add_metadata(mut self, key: &str, value: &str) -> Self {
        self.metadata.insert(key.to_string(), value.to_string());
        self
    }
}

/// Configuration options for loading models.
/// 
/// Controls validation and integrity checks during model loading.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LoadOptions {
    /// Whether to verify model integrity (default: true).
    pub verify_integrity: bool,
    /// Whether to enforce strict version checking (default: false).
    pub strict_version_check: bool,
}

impl Default for LoadOptions {
    fn default() -> Self {
        Self {
            verify_integrity: true,
            strict_version_check: false,
        }
    }
}

impl LoadOptions {
    /// Creates default LoadOptions.
    pub fn new() -> Self {
        Self::default()
    }

    /// Enables or disables integrity verification.
    pub fn with_verify_integrity(mut self, verify: bool) -> Self {
        self.verify_integrity = verify;
        self
    }

    /// Enables or disables strict version checking.
    pub fn with_strict_version_check(mut self, strict: bool) -> Self {
        self.strict_version_check = strict;
        self
    }
}

/// Main model I/O handler for saving and loading gradient boosted models.
/// 
/// `ModelIO` provides a high-level interface for model persistence operations.
/// It can be configured with custom save/load options and handles all aspects
/// of model serialization, metadata management, and validation.
pub struct ModelIO {
    save_options: SaveOptions,
    load_options: LoadOptions,
}

impl ModelIO {
    /// Creates a new ModelIO with default options.
    /// 
    /// # Returns
    /// `Ok(ModelIO)` with default save and load options.
    pub fn new() -> ModelIOResult<Self> {
        Ok(Self {
            save_options: SaveOptions::default(),
            load_options: LoadOptions::default(),
        })
    }

    /// Creates a ModelIO with custom save options.
    pub fn with_save_options(save_options: SaveOptions) -> Self {
        Self {
            save_options,
            load_options: LoadOptions::default(),
        }
    }

    /// Creates a ModelIO with custom load options.
    pub fn with_load_options(load_options: LoadOptions) -> Self {
        Self {
            save_options: SaveOptions::default(),
            load_options,
        }
    }

    /// Saves a gradient boosting model to file.
    /// 
    /// The model is saved as a directory containing:
    /// - `{model_name}.json`: The main model file
    /// - `{model_name}_metadata.json`: Model metadata and hyperparameters
    /// - `{model_name}_feature_importance.json`: Feature importance scores (if enabled)
    /// - Additional files for training data (future feature)
    /// 
    /// # Parameters
    /// - `model`: Trained gradient boosting model
    /// - `path`: Directory path where model will be saved
    /// - `model_name`: Base name for model files
    /// 
    /// # Errors
    /// - `ModelIOError::ModelError` if model is untrained
    /// - `ModelIOError::IoError` if directory creation or file writing fails
    /// - `ModelIOError::SerializationError` if serialization fails
    pub fn save_model<P: AsRef<Path>>(
        &self,
        model: &GradientBooster,  
        path: P,
        model_name: &str,
    ) -> ModelIOResult<()> {
        save_model(model, path, model_name, &self.save_options)
    }

    /// Loads a gradient boosting model from file or directory.
    /// 
    /// # Path Resolution
    /// 
    /// - If `path` is a file: Loads that file directly
    /// - If `path` is a directory: Looks for `gbrt_model.json` inside the directory
    /// 
    /// # Parameters
    /// - `path`: Path to model file or directory
    /// 
    /// # Errors
    /// - `ModelIOError::ModelNotFound` if model file doesn't exist
    /// - `ModelIOError::SerializationError` if deserialization fails
    /// - `ModelIOError::InvalidModelFile` if integrity check fails
    /// - `ModelIOError::VersionMismatch` if version check fails and strict mode enabled
    pub fn load_model<P: AsRef<Path>>(&self, path: P) -> ModelIOResult<GradientBooster> {
        let path = path.as_ref();

        // ✅ If path is a directory, look for gbrt_model.json inside
        let model_file = if path.is_dir() {
            path.join("gbrt_model.json")
        } else {
            path.to_path_buf()
        };

        if !model_file.exists() {
            return Err(ModelIOError::ModelNotFound(
                format!("Model file not found: {}", model_file.display())
            ));
        }

        let serializer = ModelSerializer::new();
        let (model, metadata) = serializer.load_model(&model_file)?;

        // ✅ Verify integrity if requested
        if self.load_options.verify_integrity {
            verify_model_integrity(&model, &metadata)?;
        }

        // ✅ Version check
        if self.load_options.strict_version_check {
            metadata.validate_version()
                .map_err(|e| {
                    match e {
                        SerializationError::VersionMismatch { expected, actual } => {
                            ModelIOError::VersionMismatch {
                                expected: expected.clone(),
                                actual: actual.clone(),
                            }
                        }
                        other => ModelIOError::SerializationError(other),
                    }
                })?;
        }

        Ok(model)
    }

    /// Exports a model to a different format.
    /// 
    /// This allows converting between Binary and JSON formats without retraining.
    /// 
    /// # Parameters
    /// - `model`: Model to export
    /// - `path`: Output file path
    /// - `format`: Target format (Binary or JSON)
    /// 
    /// # Errors
    /// - `ModelIOError::SerializationError` if conversion fails
    /// - `ModelIOError::IoError` if file writing fails
    pub fn export_model<P: AsRef<Path>>(
        &self,
        model: &GradientBooster,  // Changed from GBRT to GradientBooster
        path: P,
        format: ModelFormat,
    ) -> ModelIOResult<()> {
        export_model(model, path, format)
    }

    /// Imports a model from a specific format.
    /// 
    /// # Parameters
    /// - `path`: Path to model file
    /// - `format`: Format of the model file
    /// 
    /// # Returns
    /// Loaded model
    /// 
    /// # Errors
    /// - `ModelIOError::SerializationError` if deserialization fails
    /// - `ModelIOError::IoError` if file reading fails
    pub fn import_model<P: AsRef<Path>>(&self, path: P, format: ModelFormat) -> ModelIOResult<GradientBooster> {  // Changed return type
        import_model(path, format)
    }

    /// Lists all saved model files in a directory.
    /// 
    /// Searches for files with extensions: `.bin`, `.json`, `.gz`
    /// 
    /// # Parameters
    /// - `directory`: Directory to search
    /// 
    /// # Returns
    /// Sorted vector of model file paths
    /// 
    /// # Errors
    /// - `ModelIOError::IoError` if directory cannot be read
    pub fn list_saved_models<P: AsRef<Path>>(&self, directory: P) -> ModelIOResult<Vec<PathBuf>> {
        list_saved_models(directory)
    }

    /// Retrieves metadata for a saved model.
    /// 
    /// # Parameters
    /// - `path`: Path to model file or directory
    /// 
    /// # Returns
    /// Model metadata including configuration, version, and parameters
    /// 
    /// # Errors
    /// - `ModelIOError::InvalidModelFile` if metadata file is missing or corrupted
    /// - `ModelIOError::SerializationError` if metadata cannot be parsed
    pub fn get_model_info<P: AsRef<Path>>(&self, path: P) -> ModelIOResult<ModelMetadata> {
        get_model_info(path)
    }

    /// Validates a model file for integrity and compatibility.
    /// 
    /// Performs basic checks:
    /// - Model is trained
    /// - Number of trees matches metadata
    /// - Feature dimensions are consistent
    /// - Model type is correct
    /// 
    /// # Parameters
    /// - `path`: Path to model file
    /// 
    /// # Returns
    /// `Ok(true)` if model is valid, `Err` with details otherwise
    /// 
    /// # Errors
    /// - `ModelIOError::InvalidModelFile` if any validation check fails
    pub fn validate_model_file<P: AsRef<Path>>(&self, path: P) -> ModelIOResult<bool> {
        validate_model_file(path)
    }
}

impl Default for ModelIO {
    fn default() -> Self {
        Self::new().unwrap()
    }
}

// ============================================================================
// Standalone Model I/O Functions
// ============================================================================

/// Saves a gradient boosting model to a directory.
/// 
/// Creates a directory structure with:
/// - `{model_name}.json`: Main model file with tree ensembles
/// - `{model_name}_metadata.json`: Hyperparameters and configuration
/// - `{model_name}_feature_importance.json`: Per-feature importance scores (if enabled)
/// 
/// # Parameters
/// - `model`: Trained gradient boosting model
/// - `path`: Directory path for saving
/// - `model_name`: Base name for output files
/// - `options`: Save configuration (format, compression, metadata, etc.)
/// 
    /// # Errors
/// - `ModelIOError::ModelError` if model is untrained
/// - `ModelIOError::IoError` if directory or file creation fails
/// - `ModelIOError::SerializationError` if serialization fails
pub fn save_model<P: AsRef<Path>>(
    model: &GradientBooster,  // Changed from GBRT to GradientBooster
    path: P,
    model_name: &str,
    options: &SaveOptions,
) -> ModelIOResult<()> {
    if !model.is_trained() {
        return Err(ModelIOError::ModelError("Cannot save untrained model".to_string()));
    }

    let path = path.as_ref();
    create_dir_all(path.parent().unwrap_or(Path::new(".")))?;

    // Create metadata
    let mut metadata = ModelMetadata::new(
        "GBRT",
        model.feature_importance().len(),
        model.n_trees(),
    );

    // Add configuration to metadata
    let config = model.config();
    metadata = metadata
        .with_parameter("n_estimators", &config.n_estimators.to_string())
        .with_parameter("learning_rate", &config.learning_rate.to_string())
        .with_parameter("loss", &config.loss.to_string())
        .with_parameter("subsample", &config.subsample.to_string());

    // Add custom metadata
    for (key, value) in &options.metadata {
        metadata = metadata.with_parameter(key, value);
    }

    // Choose serialization format
    let serialization_format = match options.format {
        ModelFormat::Binary => SerializationFormat::Bincode,
        ModelFormat::Json => SerializationFormat::Json,
    };

    let serializer = ModelSerializer::new()
        .with_format(serialization_format)
        .with_compression(options.compression);

    // Create file path
    let file_path = path.join(format!("{}.json", model_name));

    serializer.save_model(model, &metadata, &file_path)?;

    // Save additional information if requested
    if options.include_feature_importance {
        save_feature_importance(model, path, model_name)?;
    }


    // Save metadata
    save_metadata(&metadata, path, model_name)?;  // Add this line

    Ok(())
}

/// Loads a gradient boosting model from file.
/// 
/// # Parameters
/// - `path`: Path to model file
/// - `options`: Load configuration (verification, version checking)
/// 
/// # Returns
/// Loaded model with validated metadata
/// 
/// # Errors
/// - `ModelIOError::ModelNotFound` if file doesn't exist
/// - `ModelIOError::SerializationError` if deserialization fails
/// - `ModelIOError::InvalidModelFile` if integrity check fails
/// - `ModelIOError::VersionMismatch` if strict version check fails
pub fn load_model<P: AsRef<Path>>(
    path: P,
    options: &LoadOptions,
) -> ModelIOResult<GradientBooster> {  // Changed return type
    let path = path.as_ref();

    if !path.exists() {
        return Err(ModelIOError::ModelNotFound(path.to_string_lossy().to_string()));
    }

    let serializer = ModelSerializer::new();
    let (model, metadata) = serializer.load_model(path)?;

    // Verify integrity if requested
    if options.verify_integrity {
        verify_model_integrity(&model, &metadata)?;
    }

    // Version check
    if options.strict_version_check {
        metadata.validate_version()
            .map_err(|e| {
                match e {
                    SerializationError::VersionMismatch { expected, actual } => {
                        ModelIOError::VersionMismatch {
                            expected: expected.clone(),
                            actual: actual.clone(),
                        }
                    }
                    other => ModelIOError::SerializationError(other),
                }
            })?;
    }

    Ok(model)
}


/// Exports a model to a different format.
/// 
/// Allows format conversion without retraining. Useful for sharing models
/// or switching between compact binary and readable JSON formats.
/// 
/// # Parameters
/// - `model`: Model to export
/// - `path`: Output file path
/// - `format`: Target format
/// 
/// # Errors
/// - `ModelIOError::SerializationError` if conversion fails
/// - `ModelIOError::IoError` if file cannot be written
pub fn export_model<P: AsRef<Path>>(
    model: &GradientBooster,  // Changed from GBRT to GradientBooster
    path: P,
    format: ModelFormat,
) -> ModelIOResult<()> {
    match format {
        ModelFormat::Json => {
            // Export as human-readable JSON
            let json = serde_json::to_string_pretty(model)
                .map_err(|e| ModelIOError::SerializationError(e.into()))?;
            std::fs::write(path, json)?;
            Ok(())
        }
        ModelFormat::Binary => {
            // Use standard save with binary format
            let options = SaveOptions::new().with_format(ModelFormat::Binary);
            save_model(model, path, "export", &options)
        }
    }
}

/// Imports a model from a specific format.
/// 
/// # Parameters
/// - `path`: Path to model file
/// - `format`: Expected format of the file
/// 
/// # Returns
/// Loaded model
/// 
/// # Errors
/// - `ModelIOError::SerializationError` if deserialization fails
/// - `ModelIOError::IoError` if file cannot be read
/// - `ModelIOError::InvalidModelFile` if format doesn't match content
pub fn import_model<P: AsRef<Path>>(
    path: P,
    format: ModelFormat,
) -> ModelIOResult<GradientBooster> {  // Changed return type
    match format {
        ModelFormat::Json => {
            let file = File::open(path)?;
            let reader = BufReader::new(file);
            let model: GradientBooster = serde_json::from_reader(reader)  // Changed type
                .map_err(|e| ModelIOError::SerializationError(e.into()))?;
            Ok(model)
        }
        ModelFormat::Binary => {
            load_model(path, &LoadOptions::default())
        }
    }
}

/// Lists all model files in a directory.
/// 
/// Searches for files with extensions: `.bin`, `.json`, `.gz`
/// 
/// # Parameters
/// - `directory`: Directory to search
/// 
/// # Returns
/// Sorted vector of absolute paths to model files
/// 
/// # Errors
/// - `ModelIOError::IoError` if directory cannot be read
pub fn list_saved_models<P: AsRef<Path>>(directory: P) -> ModelIOResult<Vec<PathBuf>> {
    let directory = directory.as_ref();
    let mut model_files = Vec::new();

    if !directory.exists() {
        return Ok(model_files);
    }

    for entry in read_dir(directory)? {
        let entry = entry?;
        let path = entry.path();

        if path.is_file() {
            if let Some(ext) = path.extension() {
                if ext == "bin" || ext == "json" || ext == "gz" {
                    model_files.push(path);
                }
            }
        }
    }

    model_files.sort();
    Ok(model_files)
}

/// Retrieves metadata for a saved model.
/// 
/// Metadata includes model type, version, hyperparameters, and custom parameters.
/// 
/// # Parameters
/// - `path`: Path to model file or directory
/// 
/// # Returns
/// Deserialized ModelMetadata
/// 
/// # Errors
/// - `ModelIOError::InvalidModelFile` if metadata file is missing
/// - `ModelIOError::SerializationError` if JSON parsing fails
/// 
/// # Path Resolution
/// - If `path` is a file like `model.json`, looks for `model_metadata.json`
/// - If `path` is a directory, looks for `gbrt_model_metadata.json`
pub fn get_model_info<P: AsRef<Path>>(path: P) -> ModelIOResult<ModelMetadata> {
    let path = path.as_ref();

    // Determine metadata file path
    let metadata_path = if path.is_file() {
        // If given a file like "gbrt_model.json", look for "gbrt_model_metadata.json"
        if let Some(stem) = path.file_stem().and_then(|s| s.to_str()) {
            path.with_file_name(format!("{}_metadata.json", stem))
        } else {
            return Err(ModelIOError::InvalidModelFile("Invalid file name".to_string()));
        }
    } else {
        // If given a directory, look for "gbrt_model_metadata.json"
        path.join("gbrt_model_metadata.json")
    };

    if metadata_path.exists() {
        let file = File::open(metadata_path)?;
        let reader = BufReader::new(file);
        let metadata: ModelMetadata = serde_json::from_reader(reader)
            .map_err(|e| ModelIOError::SerializationError(e.into()))?;
        Ok(metadata)
    } else {
        Err(ModelIOError::InvalidModelFile(
            format!("Metadata file not found: {}", metadata_path.display())
        ))
    }
}

/// Validates a model file by checking integrity and consistency.
/// 
/// Performs comprehensive validation:
/// - Model is trained
/// - Tree count matches metadata
/// - Feature dimensions match metadata
/// - Model type is correct
/// 
/// # Parameters
/// - `path`: Path to model file
/// 
/// # Returns
/// `Ok(true)` if all checks pass
/// 
/// # Errors
/// `Err` with specific violation details if any check fails
pub fn validate_model_file<P: AsRef<Path>>(path: P) -> ModelIOResult<bool> {
    let metadata = get_model_info(path)?;
    
    // Basic validation checks
    if metadata.num_features == 0 {
        return Err(ModelIOError::InvalidModelFile("Invalid number of features".to_string()));
    }

    if metadata.model_type != "GBRT" {
        return Err(ModelIOError::InvalidModelFile(format!("Invalid model type: {}", metadata.model_type)));
    }

    Ok(true)
}

// ============================================================================
// Helper Functions
// ============================================================================

/// Saves feature importance scores to a JSON file.
/// 
/// Creates a sorted list of features by importance score.
fn save_feature_importance<P: AsRef<Path>>(
    model: &GradientBooster,  // Changed from GBRT to GradientBooster
    path: P,
    model_name: &str,
) -> ModelIOResult<()> {
    let importance = model.feature_importance();
    let feature_names: Vec<String> = (0..importance.len())
        .map(|i| format!("feature_{}", i))
        .collect();

    let mut importance_data = Vec::new();
    for (i, (&importance, name)) in importance.iter().zip(feature_names.iter()).enumerate() {
        importance_data.push(FeatureImportance {
            feature_index: i,
            feature_name: name.clone(),
            importance,
        });
    }

    // Sort by importance (descending)
    importance_data.sort_by(|a, b| b.importance.partial_cmp(&a.importance).unwrap());

    let file_path = path.as_ref().join(format!("{}_feature_importance.json", model_name));
    let file = File::create(file_path)?;
    let writer = BufWriter::new(file);

    serde_json::to_writer_pretty(writer, &importance_data)
        .map_err(|e| ModelIOError::SerializationError(e.into()))?;

    Ok(())
}

/// Saves model metadata to a JSON file.
fn save_metadata<P: AsRef<Path>>(
    metadata: &ModelMetadata,
    path: P,
    model_name: &str,
) -> ModelIOResult<()> {
    let file_path = path.as_ref().join(format!("{}_metadata.json", model_name));
    let file = File::create(file_path)?;
    let writer = BufWriter::new(file);

    serde_json::to_writer_pretty(writer, metadata)
        .map_err(|e| ModelIOError::SerializationError(e.into()))?;

    Ok(())
}

/// Verifies model integrity by comparing against metadata.
///
/// Checks:
/// - Model is trained
/// - Tree count matches metadata
/// - Feature dimension matches metadata
fn verify_model_integrity(model: &GradientBooster, metadata: &ModelMetadata) -> ModelIOResult<()> {  // Changed parameter type
    // Check if model is trained
    if !model.is_trained() {
        return Err(ModelIOError::InvalidModelFile("Model appears to be untrained".to_string()));
    }

    // Check number of trees
    if model.n_trees() != metadata.num_trees {
        return Err(ModelIOError::InvalidModelFile(
            format!("Tree count mismatch: expected {}, got {}", metadata.num_trees, model.n_trees())
        ));
    }

    // Check feature dimensions
    if model.feature_importance().len() != metadata.num_features {
        return Err(ModelIOError::InvalidModelFile(
            format!("Feature dimension mismatch: expected {}, got {}", metadata.num_features, model.feature_importance().len())
        ));
    }

    Ok(())
}

/// Internal representation of feature importance for JSON export.
#[derive(Serialize, Deserialize)]
struct FeatureImportance {
    feature_index: usize,
    feature_name: String,
    importance: f64,
}


// ============================================================================
// Model Registry for Managing Multiple Models
// ============================================================================

/// Centralized registry for managing multiple models by name.
/// 
/// The registry stores a JSON map of model names to file paths, enabling
/// easy lookup and organization of models.
/// 
/// # Registry Structure
/// 
/// The registry is stored as `models.json` in the specified directory:
/// ```json
/// {
///   "iris_classifier": "models/iris.json",
///   "housing_predictor": "models/housing.json"
/// }
/// ```
pub struct ModelRegistry {
    registry_path: PathBuf,
}

impl ModelRegistry {
    /// Creates a new model registry in the specified directory.
    ///
    /// # Parameters
    /// - `path`: Directory where registry file (`models.json`) will be stored
    pub fn new<P: AsRef<Path>>(path: P) -> Self {
        Self {
            registry_path: path.as_ref().to_path_buf(),
        }
    }

    /// Registers a model in the registry.
    /// 
    /// Creates or updates the `models.json` file with the model name and path.
    /// 
    /// # Parameters
    /// - `model_name`: Unique identifier for the model
    /// - `model_path`: Path to the saved model file
    /// 
    /// # Errors
    /// - `ModelIOError::SerializationError` if registry JSON cannot be written
    /// - `ModelIOError::IoError` if registry directory cannot be created
    pub fn register_model(&self, model_name: &str, model_path: &Path) -> ModelIOResult<()> {
        let registry_file = self.registry_path.join("models.json");
        let mut registry: HashMap<String, String> = if registry_file.exists() {
            let file = File::open(&registry_file)?;
            let reader = BufReader::new(file);
            serde_json::from_reader(reader).unwrap_or_default()
        } else {
            HashMap::new()
        };

        registry.insert(model_name.to_string(), model_path.to_string_lossy().to_string());

        let file = File::create(registry_file)?;
        let writer = BufWriter::new(file);
        serde_json::to_writer_pretty(writer, &registry)
            .map_err(|e| ModelIOError::SerializationError(e.into()))?;

        Ok(())
    }

   /// Retrieves the file path for a registered model.
    /// 
    /// # Parameters
    /// - `model_name`: Name of the registered model
    /// 
    /// # Returns
    /// Path to the model file
    /// 
    /// # Errors
    /// - `ModelIOError::ModelNotFound` if model name is not in registry
    /// - `ModelIOError::SerializationError` if registry JSON cannot be read
    pub fn get_model_path(&self, model_name: &str) -> ModelIOResult<PathBuf> {
        let registry_file = self.registry_path.join("models.json");
        if !registry_file.exists() {
            return Err(ModelIOError::ModelNotFound(format!("Model registry not found: {}", model_name)));
        }

        let file = File::open(registry_file)?;
        let reader = BufReader::new(file);
        let registry: HashMap<String, String> = serde_json::from_reader(reader)
            .map_err(|e| ModelIOError::SerializationError(e.into()))?;

        registry
            .get(model_name)
            .map(|path| PathBuf::from(path))
            .ok_or_else(|| ModelIOError::ModelNotFound(model_name.to_string()))
    }

    /// Lists all registered model names.
    /// 
    /// # Returns
    /// Vector of model names sorted alphabetically
    /// 
    /// # Errors
    /// - `ModelIOError::SerializationError` if registry file cannot be read
    pub fn list_registered_models(&self) -> ModelIOResult<Vec<String>> {
        let registry_file = self.registry_path.join("models.json");
        if !registry_file.exists() {
            return Ok(Vec::new());
        }

        let file = File::open(registry_file)?;
        let reader = BufReader::new(file);
        let registry: HashMap<String, String> = serde_json::from_reader(reader)
            .map_err(|e| ModelIOError::SerializationError(e.into()))?;

        Ok(registry.keys().cloned().collect())
    }
}