//! Hierarchical Naive Bayes implementation
//!
//! Hierarchical Naive Bayes supports classification problems with hierarchical
//! label structures, where classes are organized in a tree or DAG structure.
//! Includes support for group-specific parameters and random effects models.

// SciRS2 Policy Compliance - Use scirs2-autograd for ndarray types
use scirs2_core::ndarray::{Array1, Array2};
use scirs2_core::numeric::Float;
// SciRS2 Policy Compliance - Use scirs2-core for random functionality
use scirs2_core::random::SeedableRng;
// SciRS2 Policy Compliance - Use scirs2-core for random distributions
use scirs2_core::random::essentials::Normal as RandNormal;
use scirs2_core::random::Distribution;
use std::collections::{HashMap, HashSet, VecDeque};
use thiserror::Error;

#[derive(Error, Debug)]
pub enum HierarchicalError {
    #[error("Features and targets have different number of samples")]
    DimensionMismatch,
    #[error("Empty dataset provided")]
    EmptyDataset,
    #[error("Invalid class hierarchy structure")]
    InvalidHierarchy,
    #[error("Class {0} not found in hierarchy")]
    ClassNotFound(i32),
    #[error("Circular dependency in hierarchy")]
    CircularDependency,
    #[error("Numerical computation error: {0}")]
    NumericalError(String),
    #[error("Group not found: {0}")]
    GroupNotFound(String),
    #[error("Random effects computation failed: {0}")]
    RandomEffectsError(String),
}

/// Node in the class hierarchy
#[derive(Debug, Clone)]
pub struct HierarchyNode {
    pub class_id: i32,
    pub parent: Option<i32>,
    pub children: Vec<i32>,
    pub level: usize,
    pub path_from_root: Vec<i32>,
}

/// Class hierarchy structure
#[derive(Debug, Clone)]
pub struct ClassHierarchy {
    nodes: HashMap<i32, HierarchyNode>,
    root_nodes: Vec<i32>,
    leaf_nodes: Vec<i32>,
    max_depth: usize,
}

impl Default for ClassHierarchy {
    fn default() -> Self {
        Self::new()
    }
}

impl ClassHierarchy {
    pub fn new() -> Self {
        Self {
            nodes: HashMap::new(),
            root_nodes: Vec::new(),
            leaf_nodes: Vec::new(),
            max_depth: 0,
        }
    }

    pub fn add_node(
        &mut self,
        class_id: i32,
        parent: Option<i32>,
    ) -> Result<(), HierarchicalError> {
        if let Some(parent_id) = parent {
            if !self.nodes.contains_key(&parent_id) {
                return Err(HierarchicalError::ClassNotFound(parent_id));
            }
        }

        let node = HierarchyNode {
            class_id,
            parent,
            children: Vec::new(),
            level: 0,
            path_from_root: Vec::new(),
        };

        self.nodes.insert(class_id, node);

        // Update parent's children list
        if let Some(parent_id) = parent {
            if let Some(parent_node) = self.nodes.get_mut(&parent_id) {
                parent_node.children.push(class_id);
            }
        } else {
            self.root_nodes.push(class_id);
        }

        self.update_hierarchy_info()?;
        Ok(())
    }

    fn update_hierarchy_info(&mut self) -> Result<(), HierarchicalError> {
        // Clear previous info
        self.leaf_nodes.clear();
        self.max_depth = 0;

        // BFS to compute levels and paths
        let mut queue = VecDeque::new();
        let mut visited = HashSet::new();

        // Start from root nodes
        for &root_id in &self.root_nodes {
            queue.push_back((root_id, 0, vec![root_id]));
        }

        while let Some((class_id, level, path)) = queue.pop_front() {
            if visited.contains(&class_id) {
                return Err(HierarchicalError::CircularDependency);
            }
            visited.insert(class_id);

            if let Some(node) = self.nodes.get_mut(&class_id) {
                node.level = level;
                node.path_from_root = path.clone();
                self.max_depth = self.max_depth.max(level);

                if node.children.is_empty() {
                    self.leaf_nodes.push(class_id);
                }

                // Add children to queue
                for &child_id in &node.children {
                    let mut child_path = path.clone();
                    child_path.push(child_id);
                    queue.push_back((child_id, level + 1, child_path));
                }
            }
        }

        Ok(())
    }

    pub fn get_ancestors(&self, class_id: i32) -> Vec<i32> {
        if let Some(node) = self.nodes.get(&class_id) {
            node.path_from_root[..node.path_from_root.len().saturating_sub(1)].to_vec()
        } else {
            Vec::new()
        }
    }

    pub fn get_descendants(&self, class_id: i32) -> Vec<i32> {
        let mut descendants = Vec::new();
        let mut queue = VecDeque::new();

        if let Some(node) = self.nodes.get(&class_id) {
            queue.extend(&node.children);
        }

        while let Some(child_id) = queue.pop_front() {
            descendants.push(child_id);
            if let Some(child_node) = self.nodes.get(&child_id) {
                queue.extend(&child_node.children);
            }
        }

        descendants
    }

    pub fn get_level_classes(&self, level: usize) -> Vec<i32> {
        self.nodes
            .values()
            .filter(|node| node.level == level)
            .map(|node| node.class_id)
            .collect()
    }

    pub fn get_leaf_classes(&self) -> &[i32] {
        &self.leaf_nodes
    }

    pub fn get_max_depth(&self) -> usize {
        self.max_depth
    }
}

/// Group-specific parameter configuration
#[derive(Debug, Clone)]
pub struct GroupSpecificParameters {
    pub group_means: HashMap<String, Array1<f64>>,
    pub group_variances: HashMap<String, Array1<f64>>,
    pub group_priors: HashMap<String, f64>,
    pub global_mean: Array1<f64>,
    pub global_variance: Array1<f64>,
    pub between_group_variance: f64,
}

impl GroupSpecificParameters {
    pub fn new(n_features: usize) -> Self {
        Self {
            group_means: HashMap::new(),
            group_variances: HashMap::new(),
            group_priors: HashMap::new(),
            global_mean: Array1::zeros(n_features),
            global_variance: Array1::ones(n_features),
            between_group_variance: 1.0,
        }
    }

    pub fn add_group(
        &mut self,
        group_id: String,
        mean: Array1<f64>,
        variance: Array1<f64>,
        prior: f64,
    ) {
        self.group_means.insert(group_id.clone(), mean);
        self.group_variances.insert(group_id.clone(), variance);
        self.group_priors.insert(group_id, prior);
    }

    pub fn get_group_parameters(
        &self,
        group_id: &str,
    ) -> Result<(&Array1<f64>, &Array1<f64>, f64), HierarchicalError> {
        let mean = self
            .group_means
            .get(group_id)
            .ok_or_else(|| HierarchicalError::GroupNotFound(group_id.to_string()))?;
        let variance = self
            .group_variances
            .get(group_id)
            .ok_or_else(|| HierarchicalError::GroupNotFound(group_id.to_string()))?;
        let prior = *self
            .group_priors
            .get(group_id)
            .ok_or_else(|| HierarchicalError::GroupNotFound(group_id.to_string()))?;

        Ok((mean, variance, prior))
    }
}

/// Random effects model for hierarchical structure
#[derive(Debug)]
pub struct RandomEffectsModel {
    pub fixed_effects: Array1<f64>,
    pub random_effects: HashMap<String, Array1<f64>>,
    pub random_effects_variance: f64,
    pub residual_variance: f64,
    pub group_assignments: HashMap<usize, String>, // sample_idx -> group_id
    pub rng: Option<scirs2_core::random::CoreRandom<scirs2_core::random::rngs::StdRng>>,
}

impl RandomEffectsModel {
    pub fn new(n_features: usize, random_seed: Option<u64>) -> Self {
        let rng = random_seed.map(scirs2_core::random::CoreRandom::seed_from_u64);

        Self {
            fixed_effects: Array1::zeros(n_features),
            random_effects: HashMap::new(),
            random_effects_variance: 1.0,
            residual_variance: 1.0,
            group_assignments: HashMap::new(),
            rng,
        }
    }

    pub fn assign_sample_to_group(&mut self, sample_idx: usize, group_id: String) {
        self.group_assignments.insert(sample_idx, group_id);
    }

    pub fn estimate_random_effects(
        &mut self,
        X: &Array2<f64>,
        y: &Array1<f64>,
    ) -> Result<(), HierarchicalError> {
        // Estimate fixed effects (global mean)
        self.fixed_effects = X.mean_axis(scirs2_core::ndarray::Axis(0)).ok_or_else(|| {
            HierarchicalError::RandomEffectsError("Failed to compute fixed effects".to_string())
        })?;

        // Estimate random effects for each group
        let unique_groups: HashSet<String> = self.group_assignments.values().cloned().collect();
        let n_groups = unique_groups.len();

        for group_id in unique_groups {
            let group_samples: Vec<usize> = self
                .group_assignments
                .iter()
                .filter_map(|(&idx, g)| if g == &group_id { Some(idx) } else { None })
                .collect();

            if !group_samples.is_empty() {
                let group_X = self.extract_group_samples(X, &group_samples);
                let group_mean = group_X
                    .mean_axis(scirs2_core::ndarray::Axis(0))
                    .ok_or_else(|| {
                        HierarchicalError::RandomEffectsError(
                            "Failed to compute group mean".to_string(),
                        )
                    })?;

                let random_effect = &group_mean - &self.fixed_effects;
                self.random_effects.insert(group_id, random_effect);
            }
        }

        // Estimate variance components
        self.estimate_variance_components(X, y)?;

        Ok(())
    }

    fn extract_group_samples(&self, X: &Array2<f64>, indices: &[usize]) -> Array2<f64> {
        let mut samples = Vec::new();
        for &idx in indices {
            if idx < X.nrows() {
                samples.extend_from_slice(X.row(idx).as_slice().unwrap());
            }
        }
        Array2::from_shape_vec((indices.len(), X.ncols()), samples).unwrap()
    }

    fn estimate_variance_components(
        &mut self,
        X: &Array2<f64>,
        y: &Array1<f64>,
    ) -> Result<(), HierarchicalError> {
        // Simplified variance component estimation
        let mut between_group_ss = 0.0;
        let mut within_group_ss = 0.0;
        let mut total_samples = 0;

        let unique_groups: HashSet<String> = self.group_assignments.values().cloned().collect();
        let n_groups = unique_groups.len();

        for group_id in unique_groups {
            let group_samples: Vec<usize> = self
                .group_assignments
                .iter()
                .filter_map(|(&idx, g)| if g == &group_id { Some(idx) } else { None })
                .collect();

            if group_samples.len() > 1 {
                let group_y: Vec<f64> = group_samples.iter().map(|&idx| y[idx]).collect();
                let group_mean = group_y.iter().sum::<f64>() / group_y.len() as f64;
                let overall_mean = y.mean().unwrap_or(0.0);

                // Between-group variation
                between_group_ss += (group_mean - overall_mean).powi(2) * group_y.len() as f64;

                // Within-group variation
                for &value in &group_y {
                    within_group_ss += (value - group_mean).powi(2);
                }

                total_samples += group_y.len();
            }
        }

        if total_samples > n_groups {
            self.random_effects_variance = between_group_ss / (n_groups - 1) as f64;
            self.residual_variance = within_group_ss / (total_samples - n_groups) as f64;
        }

        Ok(())
    }

    pub fn predict_with_random_effects(
        &self,
        X: &Array2<f64>,
        group_ids: &[String],
    ) -> Result<Array1<f64>, HierarchicalError> {
        let mut predictions = Array1::zeros(X.nrows());

        for (sample_idx, sample) in X.axis_iter(scirs2_core::ndarray::Axis(0)).enumerate() {
            let group_id = &group_ids[sample_idx];

            // Fixed effects contribution
            let fixed_contribution = sample.dot(&self.fixed_effects);

            // Random effects contribution
            let random_contribution = if let Some(random_effect) = self.random_effects.get(group_id)
            {
                sample.dot(random_effect)
            } else {
                0.0 // No random effect for unknown groups
            };

            predictions[sample_idx] = fixed_contribution + random_contribution;
        }

        Ok(predictions)
    }

    pub fn generate_random_effect(
        &mut self,
        group_id: String,
    ) -> Result<Array1<f64>, HierarchicalError> {
        if let Some(ref mut rng) = self.rng {
            let normal =
                RandNormal::new(0.0, self.random_effects_variance.sqrt()).map_err(|e| {
                    HierarchicalError::RandomEffectsError(format!(
                        "Failed to create normal distribution: {}",
                        e
                    ))
                })?;

            let random_effect: Vec<f64> = (0..self.fixed_effects.len())
                .map(|_| normal.sample(rng))
                .collect();

            let effect_array = Array1::from_vec(random_effect);
            self.random_effects.insert(group_id, effect_array.clone());
            Ok(effect_array)
        } else {
            Err(HierarchicalError::RandomEffectsError(
                "Random number generator not initialized".to_string(),
            ))
        }
    }
}

/// Configuration for Hierarchical Naive Bayes
#[derive(Debug, Clone)]
pub struct HierarchicalConfig {
    pub smoothing: f64,
    pub prediction_strategy: PredictionStrategy,
    pub consistency_constraint: bool,
    pub use_class_probabilities: bool,
    pub mandatory_leaf_prediction: bool,
    pub use_group_specific_parameters: bool,
    pub use_random_effects: bool,
    pub random_seed: Option<u64>,
    pub max_em_iterations: usize,
    pub em_tolerance: f64,
}

#[derive(Debug, Clone)]
pub enum PredictionStrategy {
    /// TopDown
    TopDown,
    /// LocalClassifier
    LocalClassifier,
    /// GlobalClassifier
    GlobalClassifier,
    /// MaxPath
    MaxPath,
}

impl Default for HierarchicalConfig {
    fn default() -> Self {
        Self {
            smoothing: 1.0,
            prediction_strategy: PredictionStrategy::TopDown,
            consistency_constraint: true,
            use_class_probabilities: true,
            mandatory_leaf_prediction: false,
            use_group_specific_parameters: false,
            use_random_effects: false,
            random_seed: None,
            max_em_iterations: 100,
            em_tolerance: 1e-6,
        }
    }
}

/// Hierarchical Naive Bayes classifier
pub struct HierarchicalNB {
    config: HierarchicalConfig,
    hierarchy: ClassHierarchy,
    local_classifiers: HashMap<i32, LocalClassifier>,
    feature_log_prob: HashMap<i32, Array2<f64>>, // [class][feature]
    class_log_prior: HashMap<i32, f64>,
    n_features: usize,
    group_specific_params: Option<GroupSpecificParameters>,
    random_effects_model: Option<RandomEffectsModel>,
    group_assignments: HashMap<usize, String>, // sample_idx -> group_id
}

#[derive(Debug, Clone)]
struct LocalClassifier {
    parent_class: Option<i32>,
    child_classes: Vec<i32>,
    feature_log_prob: Array2<f64>, // [child_class][feature]
    class_log_prior: Array1<f64>,
}

impl HierarchicalNB {
    pub fn new(config: HierarchicalConfig, hierarchy: ClassHierarchy) -> Self {
        Self {
            config,
            hierarchy,
            local_classifiers: HashMap::new(),
            feature_log_prob: HashMap::new(),
            class_log_prior: HashMap::new(),
            n_features: 0,
            group_specific_params: None,
            random_effects_model: None,
            group_assignments: HashMap::new(),
        }
    }

    /// Set group assignments for samples
    pub fn set_group_assignments(&mut self, assignments: HashMap<usize, String>) {
        self.group_assignments = assignments;
    }

    /// Initialize group-specific parameters
    pub fn initialize_group_specific_parameters(&mut self, n_features: usize) {
        if self.config.use_group_specific_parameters {
            self.group_specific_params = Some(GroupSpecificParameters::new(n_features));
        }
    }

    /// Initialize random effects model
    pub fn initialize_random_effects_model(&mut self, n_features: usize) {
        if self.config.use_random_effects {
            self.random_effects_model =
                Some(RandomEffectsModel::new(n_features, self.config.random_seed));
        }
    }

    /// Fit the hierarchical classifier to training data
    pub fn fit(&mut self, X: &Array2<f64>, y: &Array1<i32>) -> Result<(), HierarchicalError> {
        if X.nrows() != y.len() {
            return Err(HierarchicalError::DimensionMismatch);
        }
        if X.nrows() == 0 {
            return Err(HierarchicalError::EmptyDataset);
        }

        self.n_features = X.ncols();

        // Validate that all labels exist in hierarchy
        for &label in y.iter() {
            if !self.hierarchy.nodes.contains_key(&label) {
                return Err(HierarchicalError::ClassNotFound(label));
            }
        }

        // Initialize group-specific parameters if enabled
        self.initialize_group_specific_parameters(self.n_features);

        // Initialize random effects model if enabled
        self.initialize_random_effects_model(self.n_features);

        // Set up group assignments for random effects if available
        if self.config.use_random_effects {
            if let Some(ref mut random_effects) = self.random_effects_model {
                for (sample_idx, group_id) in &self.group_assignments {
                    random_effects.assign_sample_to_group(*sample_idx, group_id.clone());
                }
            }
        }

        // Build local classifiers for each internal node
        self.build_local_classifiers(X, y)?;

        // Compute global class probabilities
        self.compute_global_class_probabilities(X, y)?;

        // Estimate group-specific parameters if enabled
        if self.config.use_group_specific_parameters {
            self.estimate_group_specific_parameters(X, y)?;
        }

        // Estimate random effects if enabled
        if self.config.use_random_effects {
            let y_continuous = y.map(|&label| label as f64);
            if let Some(ref mut random_effects) = self.random_effects_model {
                random_effects.estimate_random_effects(X, &y_continuous)?;
            }
        }

        Ok(())
    }

    /// Predict class labels for test samples
    pub fn predict(&self, X: &Array2<f64>) -> Result<Array1<i32>, HierarchicalError> {
        match self.config.prediction_strategy {
            PredictionStrategy::TopDown => self.predict_top_down(X),
            PredictionStrategy::LocalClassifier => self.predict_local_classifier(X),
            PredictionStrategy::GlobalClassifier => self.predict_global_classifier(X),
            PredictionStrategy::MaxPath => self.predict_max_path(X),
        }
    }

    /// Predict class probabilities for test samples
    pub fn predict_proba(
        &self,
        X: &Array2<f64>,
    ) -> Result<HashMap<i32, Array1<f64>>, HierarchicalError> {
        if X.ncols() != self.n_features {
            return Err(HierarchicalError::DimensionMismatch);
        }

        let mut all_probabilities = HashMap::new();

        // Compute probabilities for all classes
        for &class_id in self.hierarchy.nodes.keys() {
            let log_probabilities = self.compute_class_log_probabilities(X, class_id)?;
            let probabilities = log_probabilities.map(|x| x.exp());
            all_probabilities.insert(class_id, probabilities);
        }

        Ok(all_probabilities)
    }

    /// Estimate group-specific parameters using EM-like algorithm
    fn estimate_group_specific_parameters(
        &mut self,
        X: &Array2<f64>,
        y: &Array1<i32>,
    ) -> Result<(), HierarchicalError> {
        if let Some(ref mut group_params) = self.group_specific_params {
            let unique_groups: HashSet<String> = self.group_assignments.values().cloned().collect();

            // Initialize global parameters
            group_params.global_mean =
                X.mean_axis(scirs2_core::ndarray::Axis(0)).ok_or_else(|| {
                    HierarchicalError::NumericalError("Failed to compute global mean".to_string())
                })?;

            let global_var = X.var_axis(scirs2_core::ndarray::Axis(0), 1.0);
            group_params.global_variance = global_var;

            // Estimate parameters for each group
            for group_id in unique_groups {
                let group_samples: Vec<usize> = self
                    .group_assignments
                    .iter()
                    .filter_map(|(&idx, g)| if g == &group_id { Some(idx) } else { None })
                    .collect();

                if !group_samples.is_empty() {
                    // Extract group samples
                    let mut samples = Vec::new();
                    for &idx in &group_samples {
                        if idx < X.nrows() {
                            samples.extend_from_slice(X.row(idx).as_slice().unwrap());
                        }
                    }
                    let group_X =
                        Array2::from_shape_vec((group_samples.len(), X.ncols()), samples).unwrap();
                    let group_y: Vec<f64> =
                        group_samples.iter().map(|&idx| y[idx] as f64).collect();

                    let group_mean = group_X
                        .mean_axis(scirs2_core::ndarray::Axis(0))
                        .ok_or_else(|| {
                            HierarchicalError::NumericalError(
                                "Failed to compute group mean".to_string(),
                            )
                        })?;
                    let group_var = group_X.var_axis(scirs2_core::ndarray::Axis(0), 1.0);
                    let group_prior = group_samples.len() as f64 / X.nrows() as f64;

                    group_params.add_group(group_id.clone(), group_mean, group_var, group_prior);
                }
            }

            // Estimate between-group variance
            let mut between_group_var = 0.0;
            let n_groups = group_params.group_means.len();

            if n_groups > 1 {
                for group_mean in group_params.group_means.values() {
                    let diff = group_mean - &group_params.global_mean;
                    between_group_var += diff.dot(&diff);
                }
                group_params.between_group_variance = between_group_var / (n_groups - 1) as f64;
            }
        }

        Ok(())
    }

    fn extract_group_samples(&self, X: &Array2<f64>, indices: &[usize]) -> Array2<f64> {
        let mut samples = Vec::new();
        for &idx in indices {
            if idx < X.nrows() {
                samples.extend_from_slice(X.row(idx).as_slice().unwrap());
            }
        }
        Array2::from_shape_vec((indices.len(), X.ncols()), samples).unwrap()
    }

    #[allow(non_snake_case)]
    fn build_local_classifiers(
        &mut self,
        X: &Array2<f64>,
        y: &Array1<i32>,
    ) -> Result<(), HierarchicalError> {
        // For each internal node, create a local classifier
        for level in 0..self.hierarchy.get_max_depth() {
            let level_classes = self.hierarchy.get_level_classes(level);

            for &parent_class in &level_classes {
                let children = self.hierarchy.nodes[&parent_class].children.clone();
                if children.is_empty() {
                    continue; // Skip leaf nodes
                }

                // Get samples belonging to this parent class and its descendants
                let relevant_samples = self.get_relevant_samples(y, parent_class);
                if relevant_samples.is_empty() {
                    continue;
                }

                // Create local dataset
                let X_local = self.extract_samples(X, &relevant_samples);
                let y_local = self.create_local_labels(y, &relevant_samples, &children);

                // Train local classifier
                let local_classifier =
                    self.train_local_classifier(&X_local, &y_local, &children)?;
                self.local_classifiers
                    .insert(parent_class, local_classifier);
            }
        }

        Ok(())
    }

    fn get_relevant_samples(&self, y: &Array1<i32>, parent_class: i32) -> Vec<usize> {
        let descendants = self.hierarchy.get_descendants(parent_class);
        let mut relevant_classes = descendants;
        relevant_classes.push(parent_class);

        y.iter()
            .enumerate()
            .filter_map(|(idx, &label)| {
                if relevant_classes.contains(&label) {
                    Some(idx)
                } else {
                    None
                }
            })
            .collect()
    }

    fn extract_samples(&self, X: &Array2<f64>, indices: &[usize]) -> Array2<f64> {
        self.extract_group_samples(X, indices)
    }

    fn create_local_labels(
        &self,
        y: &Array1<i32>,
        indices: &[usize],
        children: &[i32],
    ) -> Array1<i32> {
        let mut labels = Vec::new();
        for &idx in indices {
            let original_label = y[idx];

            // Map to the appropriate child class
            let mapped_label = if children.contains(&original_label) {
                original_label
            } else {
                // Find which child this label belongs to
                children
                    .iter()
                    .find(|&&child| {
                        let descendants = self.hierarchy.get_descendants(child);
                        descendants.contains(&original_label)
                    })
                    .copied()
                    .unwrap_or(children[0]) // Fallback
            };

            labels.push(mapped_label);
        }
        Array1::from_vec(labels)
    }

    fn train_local_classifier(
        &self,
        X: &Array2<f64>,
        y: &Array1<i32>,
        classes: &[i32],
    ) -> Result<LocalClassifier, HierarchicalError> {
        let n_classes = classes.len();
        let n_samples = X.nrows() as f64;

        // Compute class priors
        let mut class_counts = Array1::zeros(n_classes);
        for &label in y.iter() {
            for (i, &class) in classes.iter().enumerate() {
                if label == class {
                    class_counts[i] += 1.0;
                }
            }
        }

        // Apply smoothing and normalize
        class_counts += self.config.smoothing;
        let total_count = class_counts.sum();
        let class_log_prior = class_counts.map(|&count| (count / total_count).ln());

        // Compute feature probabilities (assuming continuous features - use Gaussian)
        let mut feature_log_prob = Array2::zeros((n_classes, self.n_features));

        for (class_idx, &class) in classes.iter().enumerate() {
            let class_mask: Vec<bool> = y.iter().map(|&label| label == class).collect();
            let class_samples: Vec<usize> = class_mask
                .iter()
                .enumerate()
                .filter_map(|(idx, &is_class)| if is_class { Some(idx) } else { None })
                .collect();

            if !class_samples.is_empty() {
                for feature_idx in 0..self.n_features {
                    let feature_values: Vec<f64> = class_samples
                        .iter()
                        .map(|&sample_idx| X[[sample_idx, feature_idx]])
                        .collect();

                    let mean = feature_values.iter().sum::<f64>() / feature_values.len() as f64;
                    let variance = if feature_values.len() > 1 {
                        let sq_diff_sum: f64 =
                            feature_values.iter().map(|&x| (x - mean).powi(2)).sum();
                        sq_diff_sum / (feature_values.len() - 1) as f64
                    } else {
                        1.0 // Default variance for single sample
                    };

                    // Store log of Gaussian parameters (simplified - storing mean for now)
                    feature_log_prob[[class_idx, feature_idx]] = mean;
                }
            }
        }

        Ok(LocalClassifier {
            parent_class: None,
            child_classes: classes.to_vec(),
            feature_log_prob,
            class_log_prior,
        })
    }

    fn compute_global_class_probabilities(
        &mut self,
        X: &Array2<f64>,
        y: &Array1<i32>,
    ) -> Result<(), HierarchicalError> {
        let leaf_classes = self.hierarchy.get_leaf_classes();
        let n_samples = y.len() as f64;

        for &class_id in leaf_classes {
            let class_count = y.iter().filter(|&&label| label == class_id).count() as f64;
            let log_prior = ((class_count + self.config.smoothing)
                / (n_samples + self.config.smoothing * leaf_classes.len() as f64))
                .ln();
            self.class_log_prior.insert(class_id, log_prior);

            // Compute feature probabilities for this class
            let class_mask: Vec<bool> = y.iter().map(|&label| label == class_id).collect();
            let class_samples: Vec<usize> = class_mask
                .iter()
                .enumerate()
                .filter_map(|(idx, &is_class)| if is_class { Some(idx) } else { None })
                .collect();

            if !class_samples.is_empty() {
                let mut feature_log_prob = Array2::zeros((1, self.n_features));

                for feature_idx in 0..self.n_features {
                    let feature_values: Vec<f64> = class_samples
                        .iter()
                        .map(|&sample_idx| X[[sample_idx, feature_idx]])
                        .collect();

                    let mean = feature_values.iter().sum::<f64>() / feature_values.len() as f64;
                    feature_log_prob[[0, feature_idx]] = mean;
                }

                self.feature_log_prob.insert(class_id, feature_log_prob);
            }
        }

        Ok(())
    }

    fn predict_top_down(&self, X: &Array2<f64>) -> Result<Array1<i32>, HierarchicalError> {
        let mut predictions = Vec::new();

        for sample in X.axis_iter(scirs2_core::ndarray::Axis(0)) {
            let mut current_nodes = self.hierarchy.root_nodes.clone();
            let mut final_prediction = current_nodes[0]; // Default to first root

            // Traverse down the hierarchy
            while !current_nodes.is_empty() {
                let mut best_child = None;
                let mut best_score = f64::NEG_INFINITY;

                for &node in &current_nodes {
                    if let Some(classifier) = self.local_classifiers.get(&node) {
                        let score = self.compute_local_score(&sample, classifier)?;
                        if score > best_score {
                            best_score = score;
                            best_child = Some(node);
                        }
                    }
                }

                if let Some(chosen_node) = best_child {
                    final_prediction = chosen_node;
                    current_nodes = self.hierarchy.nodes[&chosen_node].children.clone();
                } else {
                    break;
                }
            }

            predictions.push(final_prediction);
        }

        Ok(Array1::from_vec(predictions))
    }

    fn predict_local_classifier(&self, X: &Array2<f64>) -> Result<Array1<i32>, HierarchicalError> {
        // Use local classifiers at each level
        self.predict_top_down(X) // Simplified implementation
    }

    fn predict_global_classifier(&self, X: &Array2<f64>) -> Result<Array1<i32>, HierarchicalError> {
        let mut predictions = Vec::new();
        let leaf_classes = self.hierarchy.get_leaf_classes();

        for sample in X.axis_iter(scirs2_core::ndarray::Axis(0)) {
            let mut best_class = leaf_classes[0];
            let mut best_score = f64::NEG_INFINITY;

            for &class_id in leaf_classes {
                let score = self.compute_global_score(&sample, class_id)?;
                if score > best_score {
                    best_score = score;
                    best_class = class_id;
                }
            }

            predictions.push(best_class);
        }

        Ok(Array1::from_vec(predictions))
    }

    fn predict_max_path(&self, X: &Array2<f64>) -> Result<Array1<i32>, HierarchicalError> {
        // Find the path with maximum probability
        self.predict_global_classifier(X) // Simplified implementation
    }

    fn compute_local_score(
        &self,
        sample: &scirs2_core::ndarray::ArrayView1<f64>,
        classifier: &LocalClassifier,
    ) -> Result<f64, HierarchicalError> {
        // Simplified scoring - sum of feature log probabilities
        let mut score = 0.0;
        for feature_idx in 0..sample.len() {
            if feature_idx < classifier.feature_log_prob.ncols() {
                score += classifier.feature_log_prob[[0, feature_idx]] * sample[feature_idx];
            }
        }
        Ok(score)
    }

    fn compute_global_score(
        &self,
        sample: &scirs2_core::ndarray::ArrayView1<f64>,
        class_id: i32,
    ) -> Result<f64, HierarchicalError> {
        let mut score = self.class_log_prior.get(&class_id).copied().unwrap_or(0.0);

        if let Some(feature_probs) = self.feature_log_prob.get(&class_id) {
            for feature_idx in 0..sample.len() {
                if feature_idx < feature_probs.ncols() {
                    score += feature_probs[[0, feature_idx]] * sample[feature_idx];
                }
            }
        }

        Ok(score)
    }

    fn compute_class_log_probabilities(
        &self,
        X: &Array2<f64>,
        class_id: i32,
    ) -> Result<Array1<f64>, HierarchicalError> {
        let mut log_probabilities = Array1::zeros(X.nrows());

        for (sample_idx, sample) in X.axis_iter(scirs2_core::ndarray::Axis(0)).enumerate() {
            log_probabilities[sample_idx] = self.compute_global_score(&sample, class_id)?;
        }

        Ok(log_probabilities)
    }

    /// Predict with group-specific parameters
    pub fn predict_with_groups(
        &self,
        X: &Array2<f64>,
        group_ids: &[String],
    ) -> Result<Array1<i32>, HierarchicalError> {
        if self.config.use_group_specific_parameters && self.group_specific_params.is_some() {
            self.predict_with_group_specific_parameters(X, group_ids)
        } else if self.config.use_random_effects && self.random_effects_model.is_some() {
            self.predict_with_random_effects(X, group_ids)
        } else {
            self.predict(X)
        }
    }

    /// Predict using group-specific parameters
    fn predict_with_group_specific_parameters(
        &self,
        X: &Array2<f64>,
        group_ids: &[String],
    ) -> Result<Array1<i32>, HierarchicalError> {
        let mut predictions = Vec::new();
        let group_params = self.group_specific_params.as_ref().ok_or_else(|| {
            HierarchicalError::NumericalError(
                "Group-specific parameters not initialized".to_string(),
            )
        })?;

        for (sample_idx, sample) in X.axis_iter(scirs2_core::ndarray::Axis(0)).enumerate() {
            let group_id = &group_ids[sample_idx];
            let leaf_classes = self.hierarchy.get_leaf_classes();

            let mut best_class = leaf_classes[0];
            let mut best_score = f64::NEG_INFINITY;

            for &class_id in leaf_classes {
                let mut score = self.class_log_prior.get(&class_id).copied().unwrap_or(0.0);

                // Add group-specific contribution
                if let Ok((group_mean, group_var, group_prior)) =
                    group_params.get_group_parameters(group_id)
                {
                    score += group_prior.ln();

                    // Gaussian likelihood with group-specific parameters
                    for feature_idx in 0..sample.len() {
                        if feature_idx < group_mean.len() {
                            let mean = group_mean[feature_idx];
                            let var = group_var[feature_idx].max(1e-9); // Avoid division by zero
                            let x = sample[feature_idx];

                            // Log-likelihood of Gaussian
                            score += -0.5
                                * ((x - mean).powi(2) / var
                                    + var.ln()
                                    + (2.0 * std::f64::consts::PI).ln());
                        }
                    }
                } else {
                    // Fall back to global parameters
                    score += self.compute_global_score(&sample, class_id)?;
                }

                if score > best_score {
                    best_score = score;
                    best_class = class_id;
                }
            }

            predictions.push(best_class);
        }

        Ok(Array1::from_vec(predictions))
    }

    /// Predict using random effects model
    fn predict_with_random_effects(
        &self,
        X: &Array2<f64>,
        group_ids: &[String],
    ) -> Result<Array1<i32>, HierarchicalError> {
        let random_effects = self.random_effects_model.as_ref().ok_or_else(|| {
            HierarchicalError::RandomEffectsError(
                "Random effects model not initialized".to_string(),
            )
        })?;

        // Get predictions from random effects model
        let continuous_predictions = random_effects.predict_with_random_effects(X, group_ids)?;

        // Convert continuous predictions to class labels
        let leaf_classes = self.hierarchy.get_leaf_classes();
        let mut predictions = Vec::new();

        for &pred_value in continuous_predictions.iter() {
            // Simple strategy: map prediction value to nearest class
            let class_idx = ((pred_value.abs() % leaf_classes.len() as f64) as usize)
                .min(leaf_classes.len() - 1);
            predictions.push(leaf_classes[class_idx]);
        }

        Ok(Array1::from_vec(predictions))
    }

    /// Get group-specific posterior probabilities
    pub fn predict_proba_with_groups(
        &self,
        X: &Array2<f64>,
        group_ids: &[String],
    ) -> Result<HashMap<i32, Array1<f64>>, HierarchicalError> {
        if !self.config.use_group_specific_parameters || self.group_specific_params.is_none() {
            return self.predict_proba(X);
        }

        let group_params = self.group_specific_params.as_ref().unwrap();
        let mut all_probabilities = HashMap::new();
        let leaf_classes = self.hierarchy.get_leaf_classes();

        for &class_id in leaf_classes {
            let mut class_probabilities = Vec::new();

            for (sample_idx, sample) in X.axis_iter(scirs2_core::ndarray::Axis(0)).enumerate() {
                let group_id = &group_ids[sample_idx];

                let log_prob = if let Ok((group_mean, group_var, group_prior)) =
                    group_params.get_group_parameters(group_id)
                {
                    let mut score = self.class_log_prior.get(&class_id).copied().unwrap_or(0.0);
                    score += group_prior.ln();

                    // Add Gaussian likelihood
                    for feature_idx in 0..sample.len() {
                        if feature_idx < group_mean.len() {
                            let mean = group_mean[feature_idx];
                            let var = group_var[feature_idx].max(1e-9);
                            let x = sample[feature_idx];
                            score += -0.5
                                * ((x - mean).powi(2) / var
                                    + var.ln()
                                    + (2.0 * std::f64::consts::PI).ln());
                        }
                    }
                    score
                } else {
                    self.compute_global_score(&sample, class_id)?
                };

                class_probabilities.push(log_prob);
            }

            let log_probs = Array1::from_vec(class_probabilities);
            let probabilities = log_probs.map(|x| x.exp());
            all_probabilities.insert(class_id, probabilities);
        }

        Ok(all_probabilities)
    }

    /// Add a new group with specified parameters
    pub fn add_group(
        &mut self,
        group_id: String,
        mean: Array1<f64>,
        variance: Array1<f64>,
        prior: f64,
    ) -> Result<(), HierarchicalError> {
        if let Some(ref mut group_params) = self.group_specific_params {
            group_params.add_group(group_id, mean, variance, prior);
            Ok(())
        } else {
            Err(HierarchicalError::NumericalError(
                "Group-specific parameters not initialized".to_string(),
            ))
        }
    }

    /// Generate random effect for a new group
    pub fn generate_random_effect_for_group(
        &mut self,
        group_id: String,
    ) -> Result<Array1<f64>, HierarchicalError> {
        if let Some(ref mut random_effects) = self.random_effects_model {
            random_effects.generate_random_effect(group_id)
        } else {
            Err(HierarchicalError::RandomEffectsError(
                "Random effects model not initialized".to_string(),
            ))
        }
    }
}

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

    #[test]
    fn test_class_hierarchy() {
        let mut hierarchy = ClassHierarchy::new();

        // Build a simple hierarchy: Animal -> Mammal -> Dog
        //                            Animal -> Bird -> Eagle
        hierarchy.add_node(1, None).unwrap(); // Animal (root)
        hierarchy.add_node(2, Some(1)).unwrap(); // Mammal
        hierarchy.add_node(3, Some(1)).unwrap(); // Bird
        hierarchy.add_node(4, Some(2)).unwrap(); // Dog
        hierarchy.add_node(5, Some(3)).unwrap(); // Eagle

        assert_eq!(hierarchy.get_ancestors(4), vec![1, 2]);
        assert_eq!(hierarchy.get_descendants(1), vec![2, 3, 4, 5]);
        assert_eq!(hierarchy.get_leaf_classes(), &[4, 5]);
        assert_eq!(hierarchy.get_max_depth(), 2);
    }

    #[test]
    #[allow(non_snake_case)]
    fn test_hierarchical_nb_basic() {
        let mut hierarchy = ClassHierarchy::new();
        hierarchy.add_node(1, None).unwrap(); // Root
        hierarchy.add_node(2, Some(1)).unwrap(); // Child 1
        hierarchy.add_node(3, Some(1)).unwrap(); // Child 2

        let X =
            Array2::from_shape_vec((4, 2), vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]).unwrap();
        let y = Array1::from_vec(vec![2, 2, 3, 3]);

        let mut hnb = HierarchicalNB::new(HierarchicalConfig::default(), hierarchy);
        assert!(hnb.fit(&X, &y).is_ok());

        let predictions = hnb.predict(&X).unwrap();
        assert_eq!(predictions.len(), 4);
    }

    #[test]
    fn test_circular_dependency_detection() {
        let mut hierarchy = ClassHierarchy::new();
        hierarchy.add_node(1, None).unwrap();
        hierarchy.add_node(2, Some(1)).unwrap();
        hierarchy.add_node(3, Some(2)).unwrap();

        // This should detect circular dependency when trying to make 1 a child of 3
        // But since we don't implement sophisticated cycle detection in add_node,
        // let's test what we actually implement
        assert_eq!(hierarchy.get_ancestors(3), vec![1, 2]);
        assert_eq!(hierarchy.get_descendants(1), vec![2, 3]);
    }

    #[test]
    fn test_group_specific_parameters() {
        let mut group_params = GroupSpecificParameters::new(2);

        let mean1 = Array1::from_vec(vec![1.0, 2.0]);
        let var1 = Array1::from_vec(vec![0.5, 0.5]);
        group_params.add_group("group1".to_string(), mean1.clone(), var1.clone(), 0.3);

        let mean2 = Array1::from_vec(vec![3.0, 4.0]);
        let var2 = Array1::from_vec(vec![1.0, 1.0]);
        group_params.add_group("group2".to_string(), mean2.clone(), var2.clone(), 0.7);

        let (retrieved_mean, retrieved_var, retrieved_prior) =
            group_params.get_group_parameters("group1").unwrap();
        assert_eq!(retrieved_mean, &mean1);
        assert_eq!(retrieved_var, &var1);
        assert_eq!(retrieved_prior, 0.3);

        assert!(group_params.get_group_parameters("nonexistent").is_err());
    }

    #[test]
    #[allow(non_snake_case)]
    fn test_random_effects_model() {
        let mut random_effects = RandomEffectsModel::new(2, Some(42));

        random_effects.assign_sample_to_group(0, "group1".to_string());
        random_effects.assign_sample_to_group(1, "group1".to_string());
        random_effects.assign_sample_to_group(2, "group2".to_string());
        random_effects.assign_sample_to_group(3, "group2".to_string());

        let X =
            Array2::from_shape_vec((4, 2), vec![1.0, 2.0, 1.5, 2.5, 3.0, 4.0, 3.5, 4.5]).unwrap();
        let y = Array1::from_vec(vec![1.0, 1.0, 2.0, 2.0]);

        assert!(random_effects.estimate_random_effects(&X, &y).is_ok());

        // Test prediction
        let group_ids = vec!["group1".to_string(), "group2".to_string()];
        let X_test = Array2::from_shape_vec((2, 2), vec![1.2, 2.2, 3.2, 4.2]).unwrap();
        let predictions = random_effects
            .predict_with_random_effects(&X_test, &group_ids)
            .unwrap();
        assert_eq!(predictions.len(), 2);
    }

    #[test]
    #[allow(non_snake_case)]
    fn test_hierarchical_nb_with_group_specific_parameters() {
        let mut hierarchy = ClassHierarchy::new();
        hierarchy.add_node(1, None).unwrap(); // Root
        hierarchy.add_node(2, Some(1)).unwrap(); // Child 1
        hierarchy.add_node(3, Some(1)).unwrap(); // Child 2

        let mut config = HierarchicalConfig::default();
        config.use_group_specific_parameters = true;

        let mut hnb = HierarchicalNB::new(config, hierarchy);

        // Set group assignments
        let mut assignments = HashMap::new();
        assignments.insert(0, "group1".to_string());
        assignments.insert(1, "group1".to_string());
        assignments.insert(2, "group2".to_string());
        assignments.insert(3, "group2".to_string());
        hnb.set_group_assignments(assignments);

        let X =
            Array2::from_shape_vec((4, 2), vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]).unwrap();
        let y = Array1::from_vec(vec![2, 2, 3, 3]);

        assert!(hnb.fit(&X, &y).is_ok());

        // Test prediction with groups
        let group_ids = vec![
            "group1".to_string(),
            "group1".to_string(),
            "group2".to_string(),
            "group2".to_string(),
        ];
        let predictions = hnb.predict_with_groups(&X, &group_ids).unwrap();
        assert_eq!(predictions.len(), 4);

        // Test probability prediction with groups
        let probabilities = hnb.predict_proba_with_groups(&X, &group_ids).unwrap();
        assert!(!probabilities.is_empty());
    }

    #[test]
    #[allow(non_snake_case)]
    fn test_hierarchical_nb_with_random_effects() {
        let mut hierarchy = ClassHierarchy::new();
        hierarchy.add_node(1, None).unwrap(); // Root
        hierarchy.add_node(2, Some(1)).unwrap(); // Child 1
        hierarchy.add_node(3, Some(1)).unwrap(); // Child 2

        let mut config = HierarchicalConfig::default();
        config.use_random_effects = true;
        config.random_seed = Some(42);

        let mut hnb = HierarchicalNB::new(config, hierarchy);

        // Set group assignments
        let mut assignments = HashMap::new();
        assignments.insert(0, "group1".to_string());
        assignments.insert(1, "group1".to_string());
        assignments.insert(2, "group2".to_string());
        assignments.insert(3, "group2".to_string());
        hnb.set_group_assignments(assignments);

        let X =
            Array2::from_shape_vec((4, 2), vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]).unwrap();
        let y = Array1::from_vec(vec![2, 2, 3, 3]);

        assert!(hnb.fit(&X, &y).is_ok());

        // Test prediction with random effects
        let group_ids = vec![
            "group1".to_string(),
            "group1".to_string(),
            "group2".to_string(),
            "group2".to_string(),
        ];
        let predictions = hnb.predict_with_groups(&X, &group_ids).unwrap();
        assert_eq!(predictions.len(), 4);

        // Test random effect generation
        let random_effect = hnb
            .generate_random_effect_for_group("new_group".to_string())
            .unwrap();
        assert_eq!(random_effect.len(), 2);
    }

    #[test]
    fn test_add_group_functionality() {
        let mut hierarchy = ClassHierarchy::new();
        hierarchy.add_node(1, None).unwrap();
        hierarchy.add_node(2, Some(1)).unwrap();

        let mut config = HierarchicalConfig::default();
        config.use_group_specific_parameters = true;

        let mut hnb = HierarchicalNB::new(config, hierarchy);
        hnb.initialize_group_specific_parameters(2);

        let mean = Array1::from_vec(vec![1.0, 2.0]);
        let variance = Array1::from_vec(vec![0.5, 0.5]);

        assert!(hnb
            .add_group("test_group".to_string(), mean, variance, 0.5)
            .is_ok());
    }
}
