sklears-manifold 0.1.2

use scirs2_core::ndarray::Array2;
use std::collections::HashMap;
use std::marker::PhantomData;

pub struct Category<Obj, Mor> {
    objects: Vec<Obj>,
    morphisms: HashMap<(usize, usize), Mor>,
}

impl<Obj, Mor> Default for Category<Obj, Mor>
where
    Obj: Clone,
    Mor: Clone,
{
    fn default() -> Self {
        Self::new()
    }
}

impl<Obj, Mor> Category<Obj, Mor>
where
    Obj: Clone,
    Mor: Clone,
{
    pub fn new() -> Self {
        Self {
            objects: Vec::new(),
            morphisms: HashMap::new(),
        }
    }

    pub fn add_object(&mut self, obj: Obj) -> usize {
        let id = self.objects.len();
        self.objects.push(obj);
        id
    }

    pub fn add_morphism(&mut self, source: usize, target: usize, morphism: Mor) {
        self.morphisms.insert((source, target), morphism);
    }

    pub fn get_object(&self, id: usize) -> Option<&Obj> {
        self.objects.get(id)
    }

    pub fn get_morphism(&self, source: usize, target: usize) -> Option<&Mor> {
        self.morphisms.get(&(source, target))
    }

    pub fn object_count(&self) -> usize {
        self.objects.len()
    }
}

#[allow(dead_code)] // retained for serialization/introspection
pub struct Functor<C1, C2, FObj, FMor>
where
    C1: Clone,
    C2: Clone,
    FObj: Fn(&C1) -> C2,
    FMor: Clone,
{
    object_map: FObj,
    morphism_map: FMor,
    _phantom: PhantomData<(C1, C2)>,
}

impl<C1, C2, FObj, FMor> Functor<C1, C2, FObj, FMor>
where
    C1: Clone,
    C2: Clone,
    FObj: Fn(&C1) -> C2,
    FMor: Clone,
{
    pub fn new(object_map: FObj, morphism_map: FMor) -> Self {
        Self {
            object_map,
            morphism_map,
            _phantom: PhantomData,
        }
    }

    pub fn map_object(&self, obj: &C1) -> C2 {
        (self.object_map)(obj)
    }
}

#[derive(Clone)]
pub struct ManifoldObject {
    /// dimension
    pub dimension: usize,
    /// metric
    pub metric: MetricType,
    /// topology
    pub topology: TopologyType,
}

#[derive(Clone)]
pub enum MetricType {
    /// Euclidean
    Euclidean,
    /// Riemannian
    Riemannian,
    /// Lorentzian
    Lorentzian,
}

#[derive(Clone)]
pub enum TopologyType {
    /// Connected
    Connected,
    /// Disconnected
    Disconnected,
    /// Compact
    Compact,
    /// NonCompact
    NonCompact,
}

#[derive(Clone)]
pub struct EmbeddingMorphism {
    /// source_dim
    pub source_dim: usize,
    /// target_dim
    pub target_dim: usize,
    /// transformation
    pub transformation: TransformationType,
}

#[derive(Clone)]
pub enum TransformationType {
    /// Linear
    Linear(Array2<f64>),
    /// Nonlinear
    Nonlinear(String), // Placeholder for complex transformations
    /// Diffeomorphism
    Diffeomorphism,
}

pub type ManifoldCategory = Category<ManifoldObject, EmbeddingMorphism>;

#[allow(dead_code)] // retained for serialization/introspection
pub struct CategoricalManifoldLearning {
    category: ManifoldCategory,
    current_object_id: Option<usize>,
}

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

impl CategoricalManifoldLearning {
    pub fn new() -> Self {
        Self {
            category: ManifoldCategory::new(),
            current_object_id: None,
        }
    }

    pub fn add_manifold(
        &mut self,
        dimension: usize,
        metric: MetricType,
        topology: TopologyType,
    ) -> usize {
        let manifold = ManifoldObject {
            dimension,
            metric,
            topology,
        };
        self.category.add_object(manifold)
    }

    pub fn add_embedding(
        &mut self,
        source: usize,
        target: usize,
        transformation: TransformationType,
    ) -> Result<(), String> {
        let source_obj = self
            .category
            .get_object(source)
            .ok_or("Source manifold not found")?;
        let target_obj = self
            .category
            .get_object(target)
            .ok_or("Target manifold not found")?;

        let embedding = EmbeddingMorphism {
            source_dim: source_obj.dimension,
            target_dim: target_obj.dimension,
            transformation,
        };

        self.category.add_morphism(source, target, embedding);
        Ok(())
    }

    pub fn compose_embeddings(
        &self,
        first: usize,
        intermediate: usize,
        second: usize,
    ) -> Result<Option<EmbeddingMorphism>, String> {
        let first_embedding = self
            .category
            .get_morphism(first, intermediate)
            .ok_or("First embedding not found")?;
        let second_embedding = self
            .category
            .get_morphism(intermediate, second)
            .ok_or("Second embedding not found")?;

        if first_embedding.target_dim != second_embedding.source_dim {
            return Err("Dimension mismatch in composition".to_string());
        }

        // For simplicity, only handle linear compositions
        match (
            &first_embedding.transformation,
            &second_embedding.transformation,
        ) {
            (TransformationType::Linear(m1), TransformationType::Linear(m2)) => {
                let composed = m2.dot(m1);
                Ok(Some(EmbeddingMorphism {
                    source_dim: first_embedding.source_dim,
                    target_dim: second_embedding.target_dim,
                    transformation: TransformationType::Linear(composed),
                }))
            }
            _ => Ok(None), // Complex compositions not implemented
        }
    }

    pub fn get_manifold_category(&self) -> &ManifoldCategory {
        &self.category
    }
}

#[allow(dead_code)] // retained for serialization/introspection
pub struct FunctorialEmbedding<F> {
    functor: F,
    source_category: ManifoldCategory,
    target_category: ManifoldCategory,
}

impl<F> FunctorialEmbedding<F>
where
    F: Fn(&ManifoldObject) -> ManifoldObject,
{
    pub fn new(functor: F) -> Self {
        Self {
            functor,
            source_category: ManifoldCategory::new(),
            target_category: ManifoldCategory::new(),
        }
    }

    pub fn apply_to_object(&self, obj: &ManifoldObject) -> ManifoldObject {
        (self.functor)(obj)
    }
}

pub fn dimensionality_reduction_functor(
    target_dim: usize,
) -> impl Fn(&ManifoldObject) -> ManifoldObject {
    move |obj: &ManifoldObject| ManifoldObject {
        dimension: target_dim.min(obj.dimension),
        metric: obj.metric.clone(),
        topology: obj.topology.clone(),
    }
}

pub struct ToposStructure {
    /// name
    pub name: String,
    /// presheaves
    pub presheaves: HashMap<String, PresheafData>,
}

#[derive(Clone)]
pub struct PresheafData {
    /// local_sections
    pub local_sections: Array2<f64>,
    /// gluing_data
    pub gluing_data: Vec<GluingMap>,
}

#[derive(Clone)]
pub struct GluingMap {
    /// source_patch
    pub source_patch: usize,
    /// target_patch
    pub target_patch: usize,
    /// transition_function
    pub transition_function: Array2<f64>,
}

impl ToposStructure {
    pub fn new(name: String) -> Self {
        Self {
            name,
            presheaves: HashMap::new(),
        }
    }

    pub fn add_presheaf(&mut self, name: String, data: PresheafData) {
        self.presheaves.insert(name, data);
    }

    pub fn get_presheaf(&self, name: &str) -> Option<&PresheafData> {
        self.presheaves.get(name)
    }
}

pub struct SheafBasedManifoldLearning {
    topos: ToposStructure,
    patch_overlaps: Vec<(usize, usize)>,
}

impl SheafBasedManifoldLearning {
    pub fn new(name: String) -> Self {
        Self {
            topos: ToposStructure::new(name),
            patch_overlaps: Vec::new(),
        }
    }

    pub fn add_local_patch(&mut self, patch_name: String, local_data: Array2<f64>) {
        let presheaf_data = PresheafData {
            local_sections: local_data,
            gluing_data: Vec::new(),
        };
        self.topos.add_presheaf(patch_name, presheaf_data);
    }

    pub fn add_patch_overlap(&mut self, patch1: usize, patch2: usize) {
        self.patch_overlaps.push((patch1, patch2));
    }

    pub fn compute_global_sections(&self) -> Result<Array2<f64>, String> {
        if self.topos.presheaves.is_empty() {
            return Err("No local patches defined".to_string());
        }

        // Simple concatenation for demonstration
        let first_patch = self
            .topos
            .presheaves
            .values()
            .next()
            .expect("operation should succeed");
        let mut global_sections = first_patch.local_sections.clone();

        for (_, presheaf) in self.topos.presheaves.iter().skip(1) {
            if presheaf.local_sections.ncols() == global_sections.ncols() {
                global_sections = scirs2_core::ndarray::concatenate(
                    scirs2_core::ndarray::Axis(0),
                    &[global_sections.view(), presheaf.local_sections.view()],
                )
                .map_err(|e| format!("Concatenation failed: {:?}", e))?;
            }
        }

        Ok(global_sections)
    }
}

pub struct HigherCategoryEmbedding {
    /// level
    pub level: usize,
    /// higher_morphisms
    pub higher_morphisms: HashMap<Vec<usize>, Array2<f64>>,
}

impl HigherCategoryEmbedding {
    pub fn new(level: usize) -> Self {
        Self {
            level,
            higher_morphisms: HashMap::new(),
        }
    }

    pub fn add_higher_morphism(&mut self, path: Vec<usize>, morphism: Array2<f64>) {
        if path.len() == self.level + 1 {
            self.higher_morphisms.insert(path, morphism);
        }
    }

    pub fn get_higher_morphism(&self, path: &[usize]) -> Option<&Array2<f64>> {
        self.higher_morphisms.get(path)
    }

    pub fn compose_higher_morphisms(
        &self,
        path1: &[usize],
        path2: &[usize],
    ) -> Option<Array2<f64>> {
        if let (Some(m1), Some(m2)) = (
            self.higher_morphisms.get(path1),
            self.higher_morphisms.get(path2),
        ) {
            if m1.ncols() == m2.nrows() {
                return Some(m2.dot(m1));
            }
        }
        None
    }
}

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

    #[test]
    fn test_category_creation() {
        let mut category = ManifoldCategory::new();

        let manifold = ManifoldObject {
            dimension: 3,
            metric: MetricType::Euclidean,
            topology: TopologyType::Connected,
        };

        let id = category.add_object(manifold);
        assert_eq!(id, 0);
        assert_eq!(category.object_count(), 1);
    }

    #[test]
    fn test_categorical_manifold_learning() {
        let mut cml = CategoricalManifoldLearning::new();

        let source_id = cml.add_manifold(3, MetricType::Euclidean, TopologyType::Connected);
        let target_id = cml.add_manifold(2, MetricType::Euclidean, TopologyType::Connected);

        let transformation = TransformationType::Linear(Array2::eye(2));
        let result = cml.add_embedding(source_id, target_id, transformation);

        assert!(result.is_ok());
    }

    #[test]
    fn test_functorial_embedding() {
        let functor = dimensionality_reduction_functor(2);
        let fe = FunctorialEmbedding::new(functor);

        let original = ManifoldObject {
            dimension: 5,
            metric: MetricType::Riemannian,
            topology: TopologyType::Compact,
        };

        let reduced = fe.apply_to_object(&original);
        assert_eq!(reduced.dimension, 2);
    }

    #[test]
    fn test_embedding_composition() {
        let mut cml = CategoricalManifoldLearning::new();

        let obj1 = cml.add_manifold(4, MetricType::Euclidean, TopologyType::Connected);
        let obj2 = cml.add_manifold(3, MetricType::Euclidean, TopologyType::Connected);
        let obj3 = cml.add_manifold(2, MetricType::Euclidean, TopologyType::Connected);

        let m1 = Array2::from_shape_vec((3, 4), (0..12).map(|x| x as f64).collect())
            .expect("operation should succeed");
        let m2 = Array2::from_shape_vec((2, 3), (0..6).map(|x| x as f64).collect())
            .expect("operation should succeed");

        cml.add_embedding(obj1, obj2, TransformationType::Linear(m1))
            .expect("operation should succeed");
        cml.add_embedding(obj2, obj3, TransformationType::Linear(m2))
            .expect("operation should succeed");

        let composed = cml
            .compose_embeddings(obj1, obj2, obj3)
            .expect("operation should succeed");
        assert!(composed.is_some());

        let composed_embedding = composed.expect("operation should succeed");
        assert_eq!(composed_embedding.source_dim, 4);
        assert_eq!(composed_embedding.target_dim, 2);
    }

    #[test]
    fn test_topos_structure() {
        let mut topos = ToposStructure::new("TestTopos".to_string());

        let local_data = Array2::from_shape_vec((2, 3), vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0])
            .expect("operation should succeed");
        let presheaf_data = PresheafData {
            local_sections: local_data,
            gluing_data: vec![],
        };

        topos.add_presheaf("patch1".to_string(), presheaf_data);

        let retrieved = topos.get_presheaf("patch1");
        assert!(retrieved.is_some());
        assert_eq!(
            retrieved
                .expect("operation should succeed")
                .local_sections
                .shape(),
            &[2, 3]
        );
    }

    #[test]
    fn test_sheaf_based_manifold_learning() {
        let mut sml = SheafBasedManifoldLearning::new("TestSheaf".to_string());

        let patch1 = Array2::from_shape_vec((2, 3), vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0])
            .expect("operation should succeed");
        let patch2 = Array2::from_shape_vec((2, 3), vec![7.0, 8.0, 9.0, 10.0, 11.0, 12.0])
            .expect("operation should succeed");

        sml.add_local_patch("patch1".to_string(), patch1);
        sml.add_local_patch("patch2".to_string(), patch2);
        sml.add_patch_overlap(0, 1);

        let global = sml.compute_global_sections();
        assert!(global.is_ok());
        assert_eq!(global.expect("operation should succeed").shape(), &[4, 3]);
    }

    #[test]
    fn test_higher_category_embedding() {
        let mut hce = HigherCategoryEmbedding::new(2);

        let morphism = Array2::eye(3);
        hce.add_higher_morphism(vec![0, 1, 2], morphism);

        let retrieved = hce.get_higher_morphism(&[0, 1, 2]);
        assert!(retrieved.is_some());
        assert_eq!(
            retrieved.expect("operation should succeed").shape(),
            &[3, 3]
        );
    }

    #[test]
    fn test_higher_morphism_composition() {
        let mut hce = HigherCategoryEmbedding::new(2);

        let m1 = Array2::from_shape_vec((2, 3), vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0])
            .expect("operation should succeed");
        let m2 = Array2::from_shape_vec((3, 2), vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0])
            .expect("operation should succeed");

        hce.add_higher_morphism(vec![0, 1, 2], m1);
        hce.add_higher_morphism(vec![1, 2, 3], m2);

        let composed = hce.compose_higher_morphisms(&[0, 1, 2], &[1, 2, 3]);
        assert!(composed.is_some());
        assert_eq!(composed.expect("operation should succeed").shape(), &[3, 3]);
    }
}