moyo 0.10.0

use std::collections::HashSet;

use nalgebra::{Dyn, Matrix3, OMatrix, U3, Vector3};
use serde::Serialize;

use super::point_group::{iter_trans_mat_basis, iter_unimodular_trans_mat};
use super::rotation_type::identify_rotation_type;
use super::space_group::match_origin_shift;
use crate::base::{
    AngleTolerance, Lattice, MoyoError, Operation, Operations, Rotation, Translation,
    UnimodularTransformation, project_rotations,
};
use crate::math::SNF;
use crate::search::search_bravais_group;

///
/// Return integral normalizer of the point group representative up to its centralizer.
///
/// Integral normalizers for all arithmetic point groups up to their centralizers.
/// Generate integral normalizer of the given space group up to its centralizer.
/// Because the factor group of the integral normalizer by the centralizer is isomorphic to a finite permutation group, the output is guaranteed to be finite.
///
/// Be careful that the input primitive operations should be in a reduced basis.
/// This function relies on the bounded search in `iter_unimodular_trans_mat`.
/// For a non-reduced basis, the output is only best effort and may not be exhaustive.
pub fn integral_normalizer(
    prim_operations: &Operations,
    prim_generators: &Operations,
    epsilon: f64,
) -> Vec<UnimodularTransformation> {
    let prim_rotations = project_rotations(prim_operations);
    let prim_rotation_generators = project_rotations(prim_generators);

    let rotation_types = prim_rotations
        .iter()
        .map(identify_rotation_type)
        .collect::<Vec<_>>();

    let mut conjugators = vec![];
    for trans_mat_basis in
        iter_trans_mat_basis(prim_rotations, rotation_types, prim_rotation_generators)
    {
        for prim_trans_mat in iter_unimodular_trans_mat(trans_mat_basis) {
            if let Some(origin_shift) =
                match_origin_shift(prim_operations, &prim_trans_mat, prim_generators, epsilon)
            {
                conjugators.push(UnimodularTransformation::new(prim_trans_mat, origin_shift));
                break;
            }
        }
    }
    conjugators
}

/// Euclidean normalizer of a space group.
///
/// The Euclidean normalizer N_E(G) is the set of affine transformations (P, p)
/// that map G to itself while preserving the metric of the lattice:
///
///   N_E(G) = { (P, p) | (P, p)^{-1} G (P, p) = G,  P^T M P = M }
///
/// Where M is the metric tensor of the primitive lattice. Unlike the affine
/// normalizer, N_E(G) is finite (modulo continuous translations along polar
/// directions) because the metric constraint restricts the linear part to the
/// Bravais group of the lattice.
#[derive(Debug, Clone, Serialize)]
pub struct Normalizer {
    /// Basis vectors for the discrete part of the translation subgroup of the
    /// normalizer modulo Z^3 (in the input primitive basis).
    ///
    /// The full translation subgroup is generated by these vectors together
    /// with the integer lattice Z^3.
    pub translations: Vec<Translation>,
    /// Coset representatives (P, p) of the normalizer modulo its translation
    /// subgroup, in the input primitive basis. The linear part `P` lies in the
    /// Bravais group of the lattice, so `det(P) = +/- 1` (proper when
    /// `det = +1`, improper otherwise).
    pub coset_representatives: Vec<UnimodularTransformation>,
    /// Directions of continuous translations of the normalizer (non-empty for
    /// pyroelectric / polar groups), in the input primitive basis.
    pub continuous_translation_directions: Vec<Vector3<f64>>,
}

impl Normalizer {
    /// Compute the Euclidean normalizer of a space group given in a primitive
    /// basis paired with its primitive lattice.
    ///
    /// `prim_lattice` MUST be the primitive lattice; `prim_operations` and
    /// `prim_generators` MUST be expressed in that primitive basis. Passing a
    /// conventional (non-primitive) lattice is a usage error.
    ///
    /// `symprec` and `angle_tolerance` control the tolerance for the
    /// Bravais-group enumeration and origin-shift matching.
    /// `preserve_chirality = true` restricts the normalizer to the
    /// chirality-preserving subgroup N_E^+(G) (det(P) = +1 only).
    pub fn from_lattice(
        prim_lattice: &Lattice,
        prim_operations: &Operations,
        prim_generators: &Operations,
        symprec: f64,
        angle_tolerance: AngleTolerance,
        preserve_chirality: bool,
    ) -> Result<Self, MoyoError> {
        // Reduce to a Minkowski basis (precondition for `search_bravais_group`)
        // and transform operations into that basis. Results are mapped back to
        // the input basis at the end.
        let (reduced_lattice, reduced_trans_mat) = prim_lattice.minkowski_reduce()?;
        let to_reduced = UnimodularTransformation::from_linear(reduced_trans_mat);
        let reduced_operations = to_reduced.transform_operations(prim_operations);
        let reduced_generators = to_reduced.transform_operations(prim_generators);
        let reduced_rotations = project_rotations(&reduced_operations);
        let reduced_rotation_set: HashSet<Rotation> = reduced_rotations.iter().copied().collect();

        // 1. Bravais group of the reduced primitive lattice (reuses the
        //    enumeration from `crate::search`).
        let bravais = search_bravais_group(&reduced_lattice, symprec, angle_tolerance)?;
        let epsilon = symprec / reduced_lattice.volume().powf(1.0 / 3.0);

        // 2. Filter linear parts that conjugate the point group to itself, then
        //    solve for the matching translation via match_origin_shift.
        let mut coset_representatives_reduced = Vec::new();
        for p in &bravais {
            let p_f = p.map(|e| e as f64);
            if preserve_chirality && p_f.determinant().round() as i32 != 1 {
                continue;
            }
            let p_inv = p_f.try_inverse().unwrap().map(|e| e.round() as i32);
            let preserves = reduced_rotations.iter().all(|w| {
                let conjugated = p_inv * w * p;
                reduced_rotation_set.contains(&conjugated)
            });
            if !preserves {
                continue;
            }
            if let Some(origin_shift) =
                match_origin_shift(&reduced_operations, p, &reduced_generators, epsilon)
            {
                coset_representatives_reduced.push(UnimodularTransformation::new(*p, origin_shift));
            }
        }

        // 3. Translation subgroup (discrete part) and continuous directions:
        //    both come from the SNF of the stacked `(W - I)` matrix.
        let (translations_reduced, continuous_reduced) =
            normalizer_translations(&reduced_rotations);

        // 4. Map results back to the input primitive basis: under the basis
        //    change matrix T = reduced_trans_mat, a normalizer element
        //    (P_red, p_red) becomes (T P_red T^{-1}, T p_red).
        let to_input = to_reduced.inverse();
        let coset_representatives = coset_representatives_reduced
            .iter()
            .map(|cr| {
                let op = Operation::new(cr.linear, cr.origin_shift);
                let mapped = to_input.transform_operation(&op);
                UnimodularTransformation::new(
                    mapped.rotation,
                    mapped.translation.map(|e| e.rem_euclid(1.0)),
                )
            })
            .collect();

        let trans_mat_f = to_reduced.linear_as_f64();
        let translations = translations_reduced
            .iter()
            .map(|t| (trans_mat_f * t).map(|e| e.rem_euclid(1.0)))
            .collect();
        let continuous_translation_directions =
            continuous_reduced.iter().map(|d| trans_mat_f * d).collect();

        Ok(Self {
            translations,
            coset_representatives,
            continuous_translation_directions,
        })
    }
}

/// Solve `(W - I) p ≡ 0` for all rotations W via SNF of the stacked
/// `(W - I)` matrix. Returns:
///
/// * **Discrete generators (mod Z^3)**: each diagonal entry `d_i > 1` gives
///   the non-trivial fractional translation `R[:, i] / d_i mod 1`. Excludes
///   the trivial Z^3 generators.
/// * **Continuous directions (over R)**: each diagonal entry `d_i = 0`
///   corresponds to a free coordinate in the SNF, so the column `R[:, i]`
///   spans a polar (continuous-translation) direction of the normalizer.
fn normalizer_translations(prim_rotations: &[Rotation]) -> (Vec<Translation>, Vec<Vector3<f64>>) {
    if prim_rotations.is_empty() {
        return (Vec::new(), Vec::new());
    }
    let n = prim_rotations.len();
    let mut a = OMatrix::<i32, Dyn, U3>::zeros(3 * n);
    for (k, w) in prim_rotations.iter().enumerate() {
        let wmi = w - Matrix3::<i32>::identity();
        for i in 0..3 {
            for j in 0..3 {
                a[(3 * k + i, j)] = wmi[(i, j)];
            }
        }
    }

    let snf = SNF::new(&a);
    let mut discrete = Vec::new();
    let mut continuous = Vec::new();
    for i in 0..3 {
        let column = Vector3::new(snf.r[(0, i)], snf.r[(1, i)], snf.r[(2, i)]);
        match snf.d[(i, i)] {
            0 => continuous.push(column.map(|e| e as f64)),
            d if d > 1 => {
                let t = column.map(|e| (e as f64) / (d as f64));
                discrete.push(t.map(|e| e.rem_euclid(1.0)));
            }
            _ => {}
        }
    }
    (discrete, continuous)
}

#[cfg(test)]
mod tests {
    use std::collections::HashSet;

    use nalgebra::matrix;

    use super::*;
    use crate::base::{AngleTolerance, Lattice, Operation, UnimodularLinear, traverse};
    use crate::data::HallSymbol;

    const TEST_SYMPREC: f64 = 1e-4;
    const TEST_ANGLE: AngleTolerance = AngleTolerance::Default;

    fn cubic_lattice(a: f64) -> Lattice {
        Lattice::new(matrix![
            a, 0.0, 0.0;
            0.0, a, 0.0;
            0.0, 0.0, a;
        ])
    }

    fn hexagonal_lattice(a: f64, c: f64) -> Lattice {
        let s = (3.0_f64).sqrt() / 2.0;
        Lattice::new(matrix![
            a, 0.0, 0.0;
            -0.5 * a, s * a, 0.0;
            0.0, 0.0, c;
        ])
    }

    fn orthorhombic_lattice(a: f64, b: f64, c: f64) -> Lattice {
        Lattice::new(matrix![
            a, 0.0, 0.0;
            0.0, b, 0.0;
            0.0, 0.0, c;
        ])
    }

    /// Verify each coset rep conjugates G to itself (rotation-set check).
    fn assert_normalizer_conjugates_g(
        prim_lattice: &Lattice,
        prim_operations: &Operations,
        prim_generators: &Operations,
        preserve_chirality: bool,
    ) -> Normalizer {
        let normalizer = Normalizer::from_lattice(
            prim_lattice,
            prim_operations,
            prim_generators,
            TEST_SYMPREC,
            TEST_ANGLE,
            preserve_chirality,
        )
        .unwrap();
        let rotations: HashSet<Rotation> = prim_operations.iter().map(|op| op.rotation).collect();
        for cr in &normalizer.coset_representatives {
            let p = cr.linear;
            let p_f = p.map(|e| e as f64);
            let p_inv = p_f.try_inverse().unwrap().map(|e| e.round() as i32);
            let conjugated: HashSet<Rotation> = rotations.iter().map(|w| p_inv * w * p).collect();
            assert_eq!(
                rotations, conjugated,
                "coset rep {:?} did not conjugate G to itself",
                cr
            );
        }
        // Identity is always a coset rep.
        assert!(
            normalizer.coset_representatives.iter().any(|cr| {
                cr.linear == UnimodularLinear::identity()
                    && cr
                        .origin_shift
                        .iter()
                        .all(|&v| (v - v.round()).abs() < 1e-6)
            }),
            "identity not present in coset representatives"
        );
        normalizer
    }

    #[test]
    fn test_normalizer_p1() {
        // P1 (Hall #1): on a cubic metric the linear normalizer is the cubic
        // Bravais group (m-3m, 48 elements); P1 is pyroelectric in all 3 axes.
        let hall_symbol = HallSymbol::from_hall_number(1).unwrap();
        let prim_operations = hall_symbol.primitive_traverse();
        let prim_generators = hall_symbol.primitive_generators();
        let normalizer = Normalizer::from_lattice(
            &cubic_lattice(1.0),
            &prim_operations,
            &prim_generators,
            TEST_SYMPREC,
            TEST_ANGLE,
            false,
        )
        .unwrap();
        assert_eq!(normalizer.coset_representatives.len(), 48);
        assert_eq!(normalizer.continuous_translation_directions.len(), 3);
    }

    #[test]
    fn test_normalizer_p_minus_1() {
        // P-1 (Hall #2): centrosymmetric, not pyroelectric.
        let hall_symbol = HallSymbol::from_hall_number(2).unwrap();
        let prim_operations = hall_symbol.primitive_traverse();
        let prim_generators = hall_symbol.primitive_generators();
        let normalizer = assert_normalizer_conjugates_g(
            &cubic_lattice(1.0),
            &prim_operations,
            &prim_generators,
            false,
        );
        assert_eq!(normalizer.continuous_translation_directions.len(), 0);
        assert_eq!(normalizer.coset_representatives.len(), 48);
    }

    #[test]
    fn test_normalizer_pmm2() {
        // Pmm2: polar along z. Build operations directly to avoid relying on
        // a specific Hall-number ordering.
        use nalgebra::matrix;
        let identity = Operation::new(Rotation::identity(), Translation::zeros());
        let two_z = Operation::new(matrix![-1, 0, 0; 0, -1, 0; 0, 0, 1], Translation::zeros());
        let m_x = Operation::new(matrix![-1, 0, 0; 0, 1, 0; 0, 0, 1], Translation::zeros());
        let m_y = Operation::new(matrix![1, 0, 0; 0, -1, 0; 0, 0, 1], Translation::zeros());
        let prim_operations = vec![identity, two_z, m_x.clone(), m_y];
        let prim_generators = vec![m_x];
        let normalizer = assert_normalizer_conjugates_g(
            &orthorhombic_lattice(1.0, 1.3, 1.7),
            &prim_operations,
            &prim_generators,
            false,
        );
        assert_eq!(normalizer.continuous_translation_directions.len(), 1);
    }

    #[test]
    fn test_normalizer_pmmm() {
        // Pmmm: centrosymmetric orthorhombic, built directly.
        use nalgebra::matrix;
        let identity = Rotation::identity();
        let two_z = matrix![-1, 0, 0; 0, -1, 0; 0, 0, 1];
        let two_x = matrix![1, 0, 0; 0, -1, 0; 0, 0, -1];
        let two_y = matrix![-1, 0, 0; 0, 1, 0; 0, 0, -1];
        let inv = -identity;
        let m_x = -two_x;
        let m_y = -two_y;
        let m_z = -two_z;
        let zero = Translation::zeros();
        let prim_operations: Operations = [identity, two_x, two_y, two_z, inv, m_x, m_y, m_z]
            .iter()
            .map(|r| Operation::new(*r, zero))
            .collect();
        let prim_generators = vec![
            Operation::new(two_x, zero),
            Operation::new(two_y, zero),
            Operation::new(inv, zero),
        ];
        let normalizer = assert_normalizer_conjugates_g(
            &orthorhombic_lattice(1.0, 1.3, 1.7),
            &prim_operations,
            &prim_generators,
            false,
        );
        assert_eq!(normalizer.continuous_translation_directions.len(), 0);
    }

    #[test]
    fn test_chirality_preserving_for_sohncke() {
        // P432 (Hall #517 / SG #207): Sohncke (chiral). The full Euclidean
        // normalizer includes improper operations; the chirality-preserving
        // variant should drop them.
        let hall_symbol = HallSymbol::from_hall_number(517).unwrap();
        let prim_operations = hall_symbol.primitive_traverse();
        let prim_generators = hall_symbol.primitive_generators();
        let full = Normalizer::from_lattice(
            &cubic_lattice(1.0),
            &prim_operations,
            &prim_generators,
            TEST_SYMPREC,
            TEST_ANGLE,
            false,
        )
        .unwrap();
        let preserving = Normalizer::from_lattice(
            &cubic_lattice(1.0),
            &prim_operations,
            &prim_generators,
            TEST_SYMPREC,
            TEST_ANGLE,
            true,
        )
        .unwrap();
        assert!(preserving.coset_representatives.len() < full.coset_representatives.len());
        // All preserving reps must have det = +1.
        for cr in &preserving.coset_representatives {
            let det = cr.linear.map(|e| e as f64).determinant().round() as i32;
            assert_eq!(det, 1);
        }
    }

    #[test]
    fn test_normalizer_p43n() {
        // P-43n (Hall #515 / SG #218). Verify the algorithm runs and the
        // identity is included.
        let hall_symbol = HallSymbol::from_hall_number(515).unwrap();
        let prim_operations = hall_symbol.primitive_traverse();
        let prim_generators = hall_symbol.primitive_generators();
        let normalizer = assert_normalizer_conjugates_g(
            &cubic_lattice(1.0),
            &prim_operations,
            &prim_generators,
            false,
        );
        // For P-43n on cubic metric the coset count should divide |B(L)| = 48.
        assert!(48 % normalizer.coset_representatives.len() == 0);
    }

    #[test]
    fn test_normalizer_translation_subgroup_contains_z3() {
        // For each non-trivial translation t, applying (I, t) to G must
        // preserve G modulo Z^3.
        let hall_symbol = HallSymbol::from_hall_number(515).unwrap();
        let prim_operations = hall_symbol.primitive_traverse();
        let prim_generators = hall_symbol.primitive_generators();
        let normalizer = Normalizer::from_lattice(
            &cubic_lattice(1.0),
            &prim_operations,
            &prim_generators,
            TEST_SYMPREC,
            TEST_ANGLE,
            false,
        )
        .unwrap();
        for t in &normalizer.translations {
            let shifted: Operations = prim_operations
                .iter()
                .map(|op| {
                    // (I, t)^{-1} (W, w) (I, t) = (W, w + (W - I) t).
                    let new_t = (op.translation
                        + (op.rotation - Matrix3::<i32>::identity()).map(|e| e as f64) * t)
                        .map(|e| e.rem_euclid(1.0));
                    Operation::new(op.rotation, new_t)
                })
                .collect();
            let original: HashSet<(Rotation, [i64; 3])> = prim_operations
                .iter()
                .map(|op| (op.rotation, quantize_translation(&op.translation)))
                .collect();
            let new_set: HashSet<(Rotation, [i64; 3])> = shifted
                .iter()
                .map(|op| (op.rotation, quantize_translation(&op.translation)))
                .collect();
            assert_eq!(original, new_set);
        }
    }

    fn quantize_translation(t: &Vector3<f64>) -> [i64; 3] {
        let denom: i64 = 24;
        let denom_f = denom as f64;
        [
            ((t[0] * denom_f).round() as i64).rem_euclid(denom),
            ((t[1] * denom_f).round() as i64).rem_euclid(denom),
            ((t[2] * denom_f).round() as i64).rem_euclid(denom),
        ]
    }

    #[test]
    fn test_normalizer_hexagonal_p6mm() {
        // P6mm (polar along c): lattice has hexagonal Bravais group, and the
        // continuous translation direction must be along the c axis.
        use nalgebra::matrix;
        let six_z = matrix![1, -1, 0; 1, 0, 0; 0, 0, 1];
        let m_x = matrix![1, -1, 0; 0, -1, 0; 0, 0, 1];
        let prim_operations: Operations = traverse(&vec![six_z, m_x])
            .iter()
            .map(|r| Operation::new(*r, Translation::zeros()))
            .collect();
        assert_eq!(prim_operations.len(), 12);
        let prim_generators = vec![
            Operation::new(six_z, Translation::zeros()),
            Operation::new(m_x, Translation::zeros()),
        ];
        let normalizer = assert_normalizer_conjugates_g(
            &hexagonal_lattice(1.0, 1.5),
            &prim_operations,
            &prim_generators,
            false,
        );
        assert_eq!(normalizer.continuous_translation_directions.len(), 1);
    }
}