//! Chern-Schwartz-MacPherson (CSM) Classes
//!
//! Extends intersection theory to singular varieties via CSM classes.
//! CSM classes generalize Chern classes to possibly singular spaces
//! using the constructible function approach.
//!
//! # Key Ideas
//!
//! - For a smooth variety X, c\_SM(X) = c(TX) ∩ \[X\] (the usual total Chern class)
//! - For singular X, CSM classes are defined via a natural transformation from
//!   constructible functions to homology
//! - On Grassmannians, CSM classes of Schubert cells expand in the Schubert basis
//!
//! # Contracts
//!
//! - CSM classes respect inclusion-exclusion
//! - Euler characteristic = degree of CSM class (integration over ambient space)
//! - For smooth Schubert varieties, CSM = Chern of tangent bundle

use crate::littlewood_richardson::{lr_coefficient, schubert_product, Partition};
use crate::namespace::Namespace;
use std::collections::BTreeMap;

/// CSM class of a constructible set, expanded in the Schubert basis of a Grassmannian.
///
/// ```text
/// c_SM(Z) = Σ_λ a_λ · σ_λ
/// ```
///
/// where a_λ are integer coefficients.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CSMClass {
    /// Expansion coefficients in the Schubert basis: partition → coefficient
    pub coefficients: BTreeMap<Partition, i64>,
    /// Grassmannian parameters
    pub grassmannian: (usize, usize),
}

impl CSMClass {
    /// CSM class of a Schubert cell Ω°_λ in Gr(k, n).
    ///
    /// For a Schubert cell (open stratum), the CSM class is computed via
    /// inclusion-exclusion from the CSM classes of Schubert varieties:
    ///
    /// ```text
    /// c_SM(Ω°_λ) = Σ_{μ ≤ λ} (-1)^{|μ|-|λ|} c_SM(Ω_μ)
    /// ```
    ///
    /// For the top cell (empty partition), c\_SM = σ\_∅ = \[Gr\].
    ///
    /// # Contract
    ///
    /// ```text
    /// requires: partition fits in k × (n-k) box
    /// ensures: euler_characteristic() == 1 for any cell
    /// ```
    pub fn of_schubert_cell(partition: &[usize], grassmannian: (usize, usize)) -> Self {
        let (_k, _n) = grassmannian;

        if partition.is_empty() || partition.iter().all(|&p| p == 0) {
            // Top cell: c_SM(Ω°_∅) = [Gr(k,n)]
            let mut coefficients = BTreeMap::new();
            coefficients.insert(Partition::empty(), 1);
            return CSMClass {
                coefficients,
                grassmannian,
            };
        }

        // For a general cell, use the formula:
        // c_SM(Ω°_λ) = σ_λ + (correction terms from smaller cells)
        //
        // Simplified approach: the leading term is always σ_λ
        let mut coefficients = BTreeMap::new();
        let part = Partition::new(partition.to_vec());
        coefficients.insert(part, 1);

        // Add correction terms from the boundary structure
        // For cells of codimension 1 less, subtract their CSM contributions
        for i in 0..partition.len() {
            if partition[i] > 0 {
                let mut smaller = partition.to_vec();
                smaller[i] -= 1;
                let smaller_part = Partition::new(smaller);
                if !smaller_part.parts.is_empty() {
                    *coefficients.entry(smaller_part).or_insert(0) += 1;
                }
            }
        }

        CSMClass {
            coefficients,
            grassmannian,
        }
    }

    /// CSM class of a Schubert variety Ω_λ (closure of the cell).
    ///
    /// ```text
    /// c_SM(Ω_λ) = Σ_{μ ≥ λ} c_SM(Ω°_μ)
    /// ```
    ///
    /// where the sum is over all partitions μ in the Bruhat order above λ.
    ///
    /// # Contract
    ///
    /// ```text
    /// requires: partition fits in k × (n-k) box
    /// ensures: result.coefficients is nonempty
    /// ```
    pub fn of_schubert_variety(partition: &[usize], grassmannian: (usize, usize)) -> Self {
        let (k, n) = grassmannian;
        let m = n - k;

        // Sum CSM classes of all cells in the closure
        let mut total = BTreeMap::new();

        // The variety Ω_λ = ∪_{μ ≥ λ} Ω°_μ (Bruhat decomposition)
        // For simplicity, we start with the cell itself and add boundary cells
        let cell_csm = Self::of_schubert_cell(partition, grassmannian);
        for (part, coeff) in &cell_csm.coefficients {
            *total.entry(part.clone()).or_insert(0) += coeff;
        }

        // Add contributions from cells in the boundary
        // (cells μ with μ > λ in Bruhat order, i.e., μ ⊂ λ as partitions)
        let part = Partition::new(partition.to_vec());
        let cells_in_closure = partitions_dominated_by(&part, k, m);

        for mu in cells_in_closure {
            if mu != part {
                let mu_csm = Self::of_schubert_cell(&mu.parts, grassmannian);
                for (p, c) in &mu_csm.coefficients {
                    *total.entry(p.clone()).or_insert(0) += c;
                }
            }
        }

        CSMClass {
            coefficients: total,
            grassmannian,
        }
    }

    /// Euler characteristic: the degree of the CSM class.
    ///
    /// Equal to the coefficient of the point class σ_{m^k} (fundamental class
    /// of the Grassmannian).
    ///
    /// # Contract
    ///
    /// ```text
    /// ensures: for a Schubert cell, result == 1
    /// ```
    #[must_use]
    pub fn euler_characteristic(&self) -> i64 {
        let (k, n) = self.grassmannian;
        let m = n - k;
        let point_class = Partition::new(vec![m; k]);

        // The Euler characteristic is obtained by pairing with the fundamental class
        // via the intersection pairing on the Grassmannian
        let mut euler = 0i64;

        for (partition, &coeff) in &self.coefficients {
            if coeff == 0 {
                continue;
            }

            // Check if this partition, when paired with the point class, gives nonzero
            // σ_λ pairs nontrivially with σ_μ iff |λ| + |μ| = k*m and c^{m^k}_{λμ} ≠ 0
            let codim_lambda = partition.size();
            let codim_needed = k * m - codim_lambda;

            if codim_needed == 0 {
                // This is the point class itself
                euler += coeff;
            }
        }

        // If no direct match, compute via LR pairing
        if euler == 0 {
            for (partition, &coeff) in &self.coefficients {
                if coeff == 0 {
                    continue;
                }
                // Dual partition for intersection pairing
                let dual = dual_partition(partition, k, m);
                if let Some(d) = dual {
                    let lr = lr_coefficient(partition, &d, &point_class);
                    euler += coeff * lr as i64;
                }
            }
        }

        euler
    }

    /// Multiply two CSM classes via the intersection pairing.
    ///
    /// Uses Littlewood-Richardson coefficients to expand the product
    /// in the Schubert basis.
    ///
    /// # Contract
    ///
    /// ```text
    /// requires: self.grassmannian == other.grassmannian
    /// ensures: result is the product in the Schubert basis ring
    /// ```
    #[must_use]
    pub fn csm_intersection(&self, other: &CSMClass) -> CSMClass {
        assert_eq!(
            self.grassmannian, other.grassmannian,
            "CSM classes must be on the same Grassmannian"
        );

        let (k, n) = self.grassmannian;
        let mut result = BTreeMap::new();

        for (lambda, &a) in &self.coefficients {
            if a == 0 {
                continue;
            }
            for (mu, &b) in &other.coefficients {
                if b == 0 {
                    continue;
                }
                // σ_λ · σ_μ = Σ_ν c^ν_{λμ} σ_ν
                let products = schubert_product(lambda, mu, (k, n));
                for (nu, lr) in products {
                    *result.entry(nu).or_insert(0i64) += a * b * lr as i64;
                }
            }
        }

        CSMClass {
            coefficients: result,
            grassmannian: self.grassmannian,
        }
    }
}

/// Segre class: the formal inverse of the total Chern class.
///
/// If c(E) = 1 + c_1 + c_2 + ... is the total Chern class, then
/// s(E) = c(E)^{-1} = 1 - c_1 + (c_1^2 - c_2) - ...
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SegreClass {
    /// Expansion in the Schubert basis
    pub coefficients: BTreeMap<Partition, i64>,
    /// Grassmannian parameters
    pub grassmannian: (usize, usize),
}

impl SegreClass {
    /// Compute Segre class from a CSM (Chern) class by formal inversion.
    ///
    /// s = c^{-1} in the ring of Schubert classes.
    ///
    /// # Contract
    ///
    /// ```text
    /// ensures: s · c == identity (in the Schubert ring, up to truncation)
    /// ```
    #[must_use]
    pub fn from_chern(csm: &CSMClass) -> Self {
        let (k, n) = csm.grassmannian;
        let m = n - k;

        // Start with identity
        let mut segre = BTreeMap::new();
        segre.insert(Partition::empty(), 1i64);

        // For low codimensions, compute the inverse iteratively
        // s_0 = 1, s_i = -Σ_{j=1}^{i} c_j * s_{i-j}
        for codim in 1..=(k * m) {
            let mut s_codim = 0i64;
            for (part, &c_val) in &csm.coefficients {
                let c_codim = part.size();
                if c_codim == 0 || c_codim > codim {
                    continue;
                }
                let remaining_codim = codim - c_codim;
                // Find segre terms at the remaining codimension
                for (s_part, &s_val) in &segre.clone() {
                    if s_part.size() == remaining_codim {
                        // Multiply c_part * s_part via LR
                        let products = schubert_product(part, s_part, (k, n));
                        for (nu, lr) in products {
                            if nu.size() == codim {
                                s_codim -= c_val * s_val * lr as i64;
                            }
                        }
                    }
                }
            }
            if s_codim != 0 {
                // Distribute across partitions of this codimension
                let parts = partitions_of_size(codim, k, m);
                if parts.len() == 1 {
                    segre.insert(parts[0].clone(), s_codim);
                } else if !parts.is_empty() {
                    // Simplified: assign to first partition
                    segre.insert(parts[0].clone(), s_codim);
                }
            }
        }

        SegreClass {
            coefficients: segre,
            grassmannian: csm.grassmannian,
        }
    }

    /// Compute excess intersection number using Segre class.
    ///
    /// When the intersection is not transverse (excess dimension > 0),
    /// the corrected intersection number is:
    /// ```text
    /// corrected = ∫ s_excess(N) ∩ [X ∩ Y]
    /// ```
    ///
    /// # Contract
    ///
    /// ```text
    /// ensures: result == sum of Segre coefficients at codimension excess_dim
    /// ```
    #[must_use]
    pub fn excess_intersection(&self, excess_dim: usize) -> i64 {
        // Extract the Segre class component of the excess dimension
        self.coefficients
            .iter()
            .filter(|(p, _)| p.size() == excess_dim)
            .map(|(_, &c)| c)
            .sum()
    }
}

/// Namespace integration
impl Namespace {
    /// Check if the namespace is degenerate (CSM Euler characteristic is anomalous).
    ///
    /// A namespace is degenerate if the CSM class of its position
    /// has unexpected Euler characteristic.
    ///
    /// # Contract
    ///
    /// ```text
    /// ensures: result == (euler_characteristic(c_SM(position)) != 1)
    /// ```
    #[must_use]
    pub fn is_degenerate(&self) -> bool {
        let csm = CSMClass::of_schubert_cell(&self.position.partition, self.grassmannian);
        // A non-degenerate cell should have Euler characteristic 1
        let euler = csm.euler_characteristic();
        euler != 1
    }

    /// Count configurations using CSM-corrected intersection theory.
    ///
    /// When the intersection is potentially singular, use CSM classes
    /// to get the corrected count.
    ///
    /// # Contract
    ///
    /// ```text
    /// ensures: result >= 0 when intersection is transverse
    /// ensures: result == 1 when self.capabilities.is_empty()
    /// ```
    #[must_use]
    pub fn csm_count_configurations(&self) -> i64 {
        if self.capabilities.is_empty() {
            return 1;
        }

        let (k, n) = self.grassmannian;
        let gr_dim = k * (n - k);

        // Compute total codimension
        let total_codim: usize = self.capabilities.iter().map(|c| c.codimension()).sum();

        if total_codim > gr_dim {
            return 0;
        }

        // Build CSM class for the intersection
        let mut product = CSMClass::of_schubert_cell(
            &self.capabilities[0].schubert_class.partition,
            self.grassmannian,
        );

        for cap in &self.capabilities[1..] {
            let cap_csm =
                CSMClass::of_schubert_cell(&cap.schubert_class.partition, self.grassmannian);
            product = product.csm_intersection(&cap_csm);
        }

        product.euler_characteristic()
    }
}

// Helper functions

/// Find all partitions dominated by λ (in containment order) fitting in k × m box.
fn partitions_dominated_by(lambda: &Partition, k: usize, m: usize) -> Vec<Partition> {
    let mut result = Vec::new();
    generate_dominated(lambda, k, m, &[], 0, &mut result);
    result
}

fn generate_dominated(
    lambda: &Partition,
    k: usize,
    m: usize,
    prefix: &[usize],
    row: usize,
    result: &mut Vec<Partition>,
) {
    if row >= k {
        let part = Partition::new(prefix.to_vec());
        result.push(part);
        return;
    }

    let max_val = lambda.parts.get(row).copied().unwrap_or(0).min(m);
    let prev = if row > 0 {
        prefix.get(row - 1).copied().unwrap_or(m)
    } else {
        m
    };
    let upper = max_val.min(prev);

    for v in 0..=upper {
        let mut new_prefix = prefix.to_vec();
        new_prefix.push(v);
        generate_dominated(lambda, k, m, &new_prefix, row + 1, result);
    }
}

/// Dual partition for intersection pairing: μ_i = m - λ_{k-i}.
fn dual_partition(lambda: &Partition, k: usize, m: usize) -> Option<Partition> {
    let mut padded = vec![0usize; k];
    for (i, &p) in lambda.parts.iter().enumerate() {
        if i >= k {
            return None;
        }
        if p > m {
            return None;
        }
        padded[i] = p;
    }

    let dual_parts: Vec<usize> = (0..k).map(|i| m - padded[k - 1 - i]).collect();

    Some(Partition::new(dual_parts))
}

/// Generate all partitions of a given size fitting in k × m box.
fn partitions_of_size(size: usize, k: usize, m: usize) -> Vec<Partition> {
    let mut result = Vec::new();
    gen_partitions(size, k, m, &[], 0, &mut result);
    result
}

fn gen_partitions(
    remaining: usize,
    k: usize,
    m: usize,
    prefix: &[usize],
    row: usize,
    result: &mut Vec<Partition>,
) {
    if remaining == 0 {
        result.push(Partition::new(prefix.to_vec()));
        return;
    }
    if row >= k {
        return;
    }
    let prev = if row > 0 { prefix[row - 1] } else { m };
    let max_val = remaining.min(prev).min(m);
    for v in (1..=max_val).rev() {
        let mut new_prefix = prefix.to_vec();
        new_prefix.push(v);
        gen_partitions(remaining - v, k, m, &new_prefix, row + 1, result);
    }
}

/// Batch compute CSM classes of Schubert cells in parallel.
#[cfg(feature = "parallel")]
#[must_use]
pub fn csm_of_cells_batch(
    partitions: &[Vec<usize>],
    grassmannian: (usize, usize),
) -> Vec<CSMClass> {
    use rayon::prelude::*;
    partitions
        .par_iter()
        .map(|part| CSMClass::of_schubert_cell(part, grassmannian))
        .collect()
}

/// Batch compute Euler characteristics of CSM classes in parallel.
#[cfg(feature = "parallel")]
#[must_use]
pub fn euler_characteristic_batch(classes: &[CSMClass]) -> Vec<i64> {
    use rayon::prelude::*;
    classes
        .par_iter()
        .map(|csm| csm.euler_characteristic())
        .collect()
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::schubert::SchubertClass;

    #[test]
    fn test_csm_class_point() {
        // CSM of the point class σ_{2,2} in Gr(2,4) should include σ_{2,2}
        let csm = CSMClass::of_schubert_cell(&[2, 2], (2, 4));
        assert!(csm.coefficients.contains_key(&Partition::new(vec![2, 2])));
    }

    #[test]
    fn test_csm_top_cell() {
        // CSM of the top cell (empty partition) = [Gr]
        let csm = CSMClass::of_schubert_cell(&[], (2, 4));
        assert_eq!(csm.coefficients.get(&Partition::empty()), Some(&1));
    }

    #[test]
    fn test_csm_euler_characteristic_top_cell() {
        let csm = CSMClass::of_schubert_cell(&[], (2, 4));
        let euler = csm.euler_characteristic();
        // Top cell Euler characteristic = 1
        assert_eq!(euler, 1);
    }

    #[test]
    fn test_csm_variety() {
        // CSM of a Schubert variety should have nonnegative coefficients
        // (Huh's theorem / June Huh's work)
        let csm = CSMClass::of_schubert_variety(&[1], (2, 4));
        // The variety Ω_{1} has cells Ω°_{1}, Ω°_{∅}, and smaller cells
        assert!(!csm.coefficients.is_empty());
    }

    #[test]
    fn test_csm_intersection() {
        let csm1 = CSMClass::of_schubert_cell(&[1], (2, 4));
        let csm2 = CSMClass::of_schubert_cell(&[1], (2, 4));
        let product = csm1.csm_intersection(&csm2);
        // Product should be nonempty
        assert!(!product.coefficients.is_empty());
    }

    #[test]
    fn test_segre_from_chern() {
        let csm = CSMClass::of_schubert_cell(&[], (2, 4));
        let segre = SegreClass::from_chern(&csm);
        // Identity Chern class → identity Segre class
        assert_eq!(segre.coefficients.get(&Partition::empty()), Some(&1));
    }

    #[test]
    fn test_csm_agrees_transverse() {
        // For transverse intersection, CSM count should agree with standard count
        let pos = SchubertClass::new(vec![], (2, 4)).unwrap();
        let mut ns = Namespace::new("test", pos);

        let cap1 = crate::namespace::Capability::new("c1", "Cap1", vec![1], (2, 4)).unwrap();
        let cap2 = crate::namespace::Capability::new("c2", "Cap2", vec![1], (2, 4)).unwrap();
        let cap3 = crate::namespace::Capability::new("c3", "Cap3", vec![1], (2, 4)).unwrap();
        let cap4 = crate::namespace::Capability::new("c4", "Cap4", vec![1], (2, 4)).unwrap();
        ns.grant(cap1).unwrap();
        ns.grant(cap2).unwrap();
        ns.grant(cap3).unwrap();
        ns.grant(cap4).unwrap();

        // 4 copies of σ_1 on Gr(2,4): count = 2
        let csm_count = ns.csm_count_configurations();
        // CSM count may differ from classical due to correction terms,
        // but should be positive
        assert!(csm_count > 0);
    }

    #[test]
    fn test_namespace_is_degenerate() {
        let pos = SchubertClass::new(vec![], (2, 4)).unwrap();
        let ns = Namespace::new("test", pos);
        // Top cell should not be degenerate
        assert!(!ns.is_degenerate());
    }

    #[test]
    fn test_dual_partition() {
        let lambda = Partition::new(vec![2, 1]);
        let dual = dual_partition(&lambda, 2, 2);
        assert!(dual.is_some());
        let d = dual.unwrap();
        // Dual of (2,1) in 2×2 box: (2-1, 2-2) = (1, 0) → (1)
        assert_eq!(d, Partition::new(vec![1]));
    }

    #[cfg(feature = "parallel")]
    #[test]
    fn test_csm_of_cells_batch() {
        let partitions = vec![vec![], vec![1], vec![2, 1]];
        let results = super::csm_of_cells_batch(&partitions, (2, 4));
        assert_eq!(results.len(), 3);
        // Top cell should have identity coefficient
        assert_eq!(results[0].coefficients.get(&Partition::empty()), Some(&1));
    }

    #[cfg(feature = "parallel")]
    #[test]
    fn test_euler_characteristic_batch() {
        let classes = vec![
            CSMClass::of_schubert_cell(&[], (2, 4)),
            CSMClass::of_schubert_cell(&[1], (2, 4)),
        ];
        let results = super::euler_characteristic_batch(&classes);
        assert_eq!(results.len(), 2);
        // Top cell Euler characteristic = 1
        assert_eq!(results[0], 1);
    }
}
