//! Implementation of symmetry analysis for multi-determinantal wavefunctions.

use std::collections::HashSet;
use std::fmt;
use std::ops::Mul;

use anyhow::{self, ensure, format_err, Context};
use approx;
use derive_builder::Builder;
use itertools::Itertools;
use log;
use ndarray::{s, Array1, Array2, Array3, Axis, Ix2};
use ndarray_linalg::{
    eig::Eig,
    eigh::Eigh,
    types::{Lapack, Scalar},
    UPLO,
};
use num_complex::{Complex, ComplexFloat};
use num_traits::{Float, ToPrimitive, Zero};

use crate::analysis::{
    fn_calc_xmat_complex, fn_calc_xmat_real, EigenvalueComparisonMode, Orbit, OrbitIterator,
    Overlap, RepAnalysis,
};
use crate::auxiliary::misc::complex_modified_gram_schmidt;
use crate::chartab::chartab_group::CharacterProperties;
use crate::chartab::{DecompositionError, SubspaceDecomposable};
use crate::group::GroupType;
use crate::io::format::{log_subtitle, qsym2_output, QSym2Output};
use crate::symmetry::symmetry_element::symmetry_operation::SpecialSymmetryTransformation;
use crate::symmetry::symmetry_group::SymmetryGroupProperties;
use crate::symmetry::symmetry_transformation::{SymmetryTransformable, SymmetryTransformationKind};
use crate::target::determinant::SlaterDeterminant;
use crate::target::noci::basis::{Basis, OrbitBasis};
use crate::target::noci::multideterminant::MultiDeterminant;

// =======
// Overlap
// =======

impl<'a, T, B> Overlap<T, Ix2> for MultiDeterminant<'a, T, B>
where
    T: Lapack
        + ComplexFloat<Real = <T as Scalar>::Real>
        + fmt::Debug
        + Mul<<T as ComplexFloat>::Real, Output = T>,
    <T as ComplexFloat>::Real: fmt::Debug
        + approx::RelativeEq<<T as ComplexFloat>::Real>
        + approx::AbsDiffEq<Epsilon = <T as Scalar>::Real>,
    B: Basis<SlaterDeterminant<'a, T>> + Clone,
{
    fn complex_symmetric(&self) -> bool {
        self.complex_symmetric
    }

    /// Computes the overlap between two multi-determinantal wavefunctions.
    ///
    /// Determinants with fractional electrons are currently not supported.
    ///
    /// When one or both of the multi-determinantal wavefunctions have been acted on by an
    /// antiunitary operation, the correct Hermitian or complex-symmetric metric will be chosen in
    /// the evalulation of the overlap.
    ///
    /// # Arguments
    ///
    /// * `metric` - The atomic-orbital overlap matrix with respect to the conventional sesquilinear
    /// inner product.
    /// * `metric_h` - The atomic-orbital overlap matrix with respect to the bilinear inner product.
    ///
    /// # Panics
    ///
    /// Panics if `self` and `other` have mismatched spin constraints or numbers of coefficient
    /// matrices, or if fractional occupation numbers are detected.
    fn overlap(
        &self,
        other: &Self,
        metric: Option<&Array2<T>>,
        metric_h: Option<&Array2<T>>,
    ) -> Result<T, anyhow::Error> {
        ensure!(
            self.spin_constraint() == other.spin_constraint(),
            "Inconsistent spin constraints between `self` and `other`."
        );
        ensure!(
            self.coefficients.len() == other.coefficients.len(),
            "Inconsistent numbers of coefficients between `self` and `other`."
        );

        let s_dets = self.basis.iter().collect::<Result<Vec<_>, _>>()?;
        let o_dets = other.basis.iter().collect::<Result<Vec<_>, _>>()?;
        let swx_vec = s_dets
            .iter()
            .cartesian_product(o_dets.iter())
            .map(|(w, x)| w.overlap(x, metric, metric_h))
            .collect::<Result<Vec<_>, _>>()?;
        let d = self.coefficients.len();
        let swx = Array2::from_shape_vec((d, d), swx_vec)?;

        if self.complex_symmetric {
            Ok(self.coefficients.t().dot(&swx).dot(&other.coefficients))
        } else {
            Ok(self
                .coefficients
                .t()
                .mapv(|x| x.conj())
                .dot(&swx)
                .dot(&other.coefficients))
        }
    }

    /// Returns the mathematical definition of the overlap between two multi-determinantal
    /// wavefunctions.
    fn overlap_definition(&self) -> String {
        let k = if self.complex_symmetric() { "κ " } else { "" };
        format!("⟨{k}Ψ_1|Ψ_2⟩ = ∫ [{k}Ψ_1(x^Ne)]* Ψ_2(x^Ne) dx^Ne")
    }
}

// =============================
// MultiDeterminantSymmetryOrbit
// =============================

// -----------------
// Struct definition
// -----------------

/// Structure to manage symmetry orbits (*i.e.* orbits generated by symmetry groups) of
/// multi-determinantal wavefunctions.
#[derive(Builder, Clone)]
pub struct MultiDeterminantSymmetryOrbit<'a, 'g, G, T, B>
where
    G: SymmetryGroupProperties,
    T: ComplexFloat + fmt::Debug + Lapack,
    B: 'a + Basis<SlaterDeterminant<'a, T>> + Clone,
    MultiDeterminant<'a, T, B>: SymmetryTransformable,
{
    /// The generating symmetry group.
    group: &'g G,

    /// The origin multi-determinantal wavefunction of the orbit.
    origin: &'a MultiDeterminant<'a, T, B>,

    /// The threshold for determining zero eigenvalues in the orbit overlap matrix.
    pub(crate) linear_independence_threshold: <T as ComplexFloat>::Real,

    /// The threshold for determining if calculated multiplicities in representation analysis are
    /// integral.
    integrality_threshold: <T as ComplexFloat>::Real,

    /// The kind of transformation determining the way the symmetry operations in `group` act on
    /// [`Self::origin`].
    symmetry_transformation_kind: SymmetryTransformationKind,

    /// The overlap matrix between the symmetry-equivalent multi-eterminantal wavefunctions in the
    /// orbit.
    #[builder(setter(skip), default = "None")]
    smat: Option<Array2<T>>,

    /// The eigenvalues of the overlap matrix between the symmetry-equivalent multi-determinantal
    /// wavefunctions in the orbit.
    #[builder(setter(skip), default = "None")]
    pub(crate) smat_eigvals: Option<Array1<T>>,

    /// The $`\mathbf{X}`$ matrix for the overlap matrix between the symmetry-equivalent
    /// multi-determinantal wavefunctions in the orbit.
    ///
    /// See [`RepAnalysis::xmat`] for further information.
    #[builder(setter(skip), default = "None")]
    xmat: Option<Array2<T>>,

    /// An enumerated type specifying the comparison mode for filtering out orbit overlap
    /// eigenvalues.
    eigenvalue_comparison_mode: EigenvalueComparisonMode,
}

// ----------------------------
// Struct method implementation
// ----------------------------

impl<'a, 'g, G, T, B> MultiDeterminantSymmetryOrbit<'a, 'g, G, T, B>
where
    G: SymmetryGroupProperties + Clone,
    T: ComplexFloat + fmt::Debug + Lapack,
    B: 'a + Basis<SlaterDeterminant<'a, T>> + Clone,
    MultiDeterminant<'a, T, B>: SymmetryTransformable,
{
    /// Returns a builder for constructing a new multi-determinantal wavefunction symmetry orbit.
    pub fn builder() -> MultiDeterminantSymmetryOrbitBuilder<'a, 'g, G, T, B> {
        MultiDeterminantSymmetryOrbitBuilder::default()
    }

    /// Returns the origin of the multi-determinantal wavefunction symmetry orbit.
    pub fn origin(&self) -> &MultiDeterminant<'a, T, B> {
        self.origin
    }
}

impl<'a, 'g, G, B> MultiDeterminantSymmetryOrbit<'a, 'g, G, f64, B>
where
    G: SymmetryGroupProperties,
    B: 'a + Basis<SlaterDeterminant<'a, f64>> + Clone,
    MultiDeterminant<'a, f64, B>: SymmetryTransformable,
{
    fn_calc_xmat_real!(
        /// Calculates the $`\mathbf{X}`$ matrix for real and symmetric overlap matrix
        /// $`\mathbf{S}`$ between the symmetry-equivalent Slater determinants in the orbit.
        ///
        /// The resulting $`\mathbf{X}`$ is stored in the orbit.
        ///
        /// # Arguments
        ///
        /// * `preserves_full_rank` - If `true`, when $`\mathbf{S}`$ is already of full rank, then
        /// $`\mathbf{X}`$ is set to be the identity matrix to avoid mixing the orbit determinants.
        /// If `false`, $`\mathbf{X}`$ also orthogonalises $`\mathbf{S}`$ even when it is already of
        /// full rank.
        pub calc_xmat
    );
}

impl<'a, 'g, G, T, B> MultiDeterminantSymmetryOrbit<'a, 'g, G, Complex<T>, B>
where
    G: SymmetryGroupProperties,
    T: Float + Scalar<Complex = Complex<T>>,
    Complex<T>: ComplexFloat<Real = T> + Scalar<Real = T, Complex = Complex<T>> + Lapack,
    B: 'a + Basis<SlaterDeterminant<'a, Complex<T>>> + Clone,
    MultiDeterminant<'a, Complex<T>, B>: SymmetryTransformable + Overlap<Complex<T>, Ix2>,
{
    fn_calc_xmat_complex!(
        /// Calculates the $`\mathbf{X}`$ matrix for complex and symmetric or Hermitian overlap
        /// matrix $`\mathbf{S}`$ between the symmetry-equivalent Slater determinants in the orbit.
        ///
        /// The resulting $`\mathbf{X}`$ is stored in the orbit.
        ///
        /// # Arguments
        ///
        /// * `preserves_full_rank` - If `true`, when $`\mathbf{S}`$ is already of full rank, then
        /// $`\mathbf{X}`$ is set to be the identity matrix to avoid mixing the orbit determinants.
        /// If `false`, $`\mathbf{X}`$ also orthogonalises $`\mathbf{S}`$ even when it is already of
        /// full rank.
        pub calc_xmat
    );
}

// ---------------------
// Trait implementations
// ---------------------

// ~~~~~
// Orbit
// ~~~~~

impl<'a, 'g, G, T, B> Orbit<G, MultiDeterminant<'a, T, B>>
    for MultiDeterminantSymmetryOrbit<'a, 'g, G, T, B>
where
    G: SymmetryGroupProperties,
    T: ComplexFloat + fmt::Debug + Lapack,
    B: 'a + Basis<SlaterDeterminant<'a, T>> + Clone,
    MultiDeterminant<'a, T, B>: SymmetryTransformable,
{
    type OrbitIter = OrbitIterator<'a, G, MultiDeterminant<'a, T, B>>;

    fn group(&self) -> &G {
        self.group
    }

    fn origin(&self) -> &MultiDeterminant<'a, T, B> {
        self.origin
    }

    fn iter(&self) -> Self::OrbitIter {
        OrbitIterator::new(
            self.group,
            self.origin,
            match self.symmetry_transformation_kind {
                SymmetryTransformationKind::Spatial => |op, multidet| {
                    multidet.sym_transform_spatial(op).with_context(|| {
                        format!("Unable to apply `{op}` spatially on the origin multi-determinantal wavefunction")
                    })
                },
                SymmetryTransformationKind::SpatialWithSpinTimeReversal => |op, multidet| {
                    multidet.sym_transform_spatial_with_spintimerev(op).with_context(|| {
                        format!("Unable to apply `{op}` spatially (with spin-including time reversal) on the origin multi-determinantal wavefunction")
                    })
                },
                SymmetryTransformationKind::Spin => |op, multidet| {
                    multidet.sym_transform_spin(op).with_context(|| {
                        format!("Unable to apply `{op}` spin-wise on the origin multi-determinantal wavefunction")
                    })
                },
                SymmetryTransformationKind::SpinSpatial => |op, multidet| {
                    multidet.sym_transform_spin_spatial(op).with_context(|| {
                        format!("Unable to apply `{op}` spin-spatially on the origin multi-determinantal wavefunction")
                    })
                },
            },
        )
    }
}

// ~~~~~~~~~~~
// RepAnalysis
// ~~~~~~~~~~~

impl<'a, 'g, G, T, B> RepAnalysis<G, MultiDeterminant<'a, T, B>, T, Ix2>
    for MultiDeterminantSymmetryOrbit<'a, 'g, G, T, B>
where
    G: SymmetryGroupProperties,
    G::CharTab: SubspaceDecomposable<T>,
    T: Lapack
        + ComplexFloat<Real = <T as Scalar>::Real>
        + fmt::Debug
        + Mul<<T as ComplexFloat>::Real, Output = T>,
    <T as ComplexFloat>::Real: fmt::Debug
        + Zero
        + From<u16>
        + ToPrimitive
        + approx::RelativeEq<<T as ComplexFloat>::Real>
        + approx::AbsDiffEq<Epsilon = <T as Scalar>::Real>,
    B: 'a + Basis<SlaterDeterminant<'a, T>> + Clone,
    MultiDeterminant<'a, T, B>: SymmetryTransformable,
{
    fn set_smat(&mut self, smat: Array2<T>) {
        self.smat = Some(smat)
    }

    fn smat(&self) -> Option<&Array2<T>> {
        self.smat.as_ref()
    }

    fn xmat(&self) -> &Array2<T> {
        self.xmat
            .as_ref()
            .expect("Orbit overlap orthogonalisation matrix not found.")
    }

    fn norm_preserving_scalar_map(&self, i: usize) -> Result<fn(T) -> T, anyhow::Error> {
        if self.origin.complex_symmetric {
            Err(format_err!("`norm_preserving_scalar_map` is currently not implemented for complex symmetric overlaps."))
        } else {
            if self
                .group
                .get_index(i)
                .unwrap_or_else(|| panic!("Group operation index `{i}` not found."))
                .contains_time_reversal()
            {
                Ok(ComplexFloat::conj)
            } else {
                Ok(|x| x)
            }
        }
    }

    fn integrality_threshold(&self) -> <T as ComplexFloat>::Real {
        self.integrality_threshold
    }

    fn eigenvalue_comparison_mode(&self) -> &EigenvalueComparisonMode {
        &self.eigenvalue_comparison_mode
    }

    /// Reduces the representation or corepresentation spanned by the multi-determinantal
    /// wavefunctions in the orbit to a direct sum of the irreducible representations or
    /// corepresentations of the generating symmetry group.
    ///
    /// # Returns
    ///
    /// The decomposed result.
    ///
    /// # Errors
    ///
    /// Errors if the decomposition fails, *e.g.* because one or more calculated multiplicities
    /// are non-integral, or also because the combination of group type, transformation type, and
    /// oddity of the number of electrons would not give sensible symmetry results. In particular,
    /// spin or spin-spatial symmetry analysis of odd-electron systems in unitary-represented
    /// magnetic groups is not valid.
    fn analyse_rep(
        &self,
    ) -> Result<
        <<G as CharacterProperties>::CharTab as SubspaceDecomposable<T>>::Decomposition,
        DecompositionError,
    > {
        log::debug!("Analysing representation symmetry for a multi-determinantal wavefunction...");
        let mut nelectrons_set = self
            .origin()
            .basis
            .iter()
            .map(|det_res| det_res.and_then(|det| {
                let det_nelectrons_float = det.nelectrons();
                if approx::relative_eq!(
                    det_nelectrons_float.round(),
                    det_nelectrons_float,
                    epsilon = self.integrality_threshold,
                    max_relative = self.integrality_threshold
                ) {
                    det_nelectrons_float.round().to_usize().ok_or(format_err!("Unable to convert the number of electrons `{det_nelectrons_float:.7}` to `usize`."))
                } else {
                    Err(format_err!("Fractional number of electrons encountered: `{det_nelectrons_float:.7}`"))
                }
            }))
            .collect::<Result<HashSet<_>, _>>()
            .map_err(|err| DecompositionError(err.to_string()))?;
        if nelectrons_set.len() != 1 {
            Err(DecompositionError("Symmetry analysis for multi-determinantal wavefunctions with multiple numbers of electrons is not yet supported.".to_string()))
        } else {
            let nelectrons = nelectrons_set.drain().next().ok_or(DecompositionError(
                "Unable to obtain the number of electrons.".to_string(),
            ))?;
            let (valid_symmetry, err_str) = if nelectrons.rem_euclid(2) == 0 {
                // Even number of electrons; always valid
                (true, String::new())
            } else {
                // Odd number of electrons; validity depends on group and orbit type
                match self.symmetry_transformation_kind {
                    SymmetryTransformationKind::Spatial => (true, String::new()),
                    SymmetryTransformationKind::SpatialWithSpinTimeReversal
                        | SymmetryTransformationKind::Spin
                        | SymmetryTransformationKind::SpinSpatial => {
                        match self.group().group_type() {
                            GroupType::Ordinary(_) => (true, String::new()),
                            GroupType::MagneticGrey(_) | GroupType::MagneticBlackWhite(_) => {
                                (!self.group().unitary_represented(),
                                "Unitary-represented magnetic groups cannot be used for symmetry analysis of odd-electron systems where spin is treated explicitly.".to_string())
                            }
                        }
                    }
                }
            };

            if valid_symmetry {
                let chis = self
                    .calc_characters()
                    .map_err(|err| DecompositionError(err.to_string()))?;
                log::debug!("Characters calculated.");

                log_subtitle("Multi-determinantal wavefunction orbit characters");
                qsym2_output!("");
                self.characters_to_string(&chis, self.integrality_threshold)
                    .log_output_display();
                qsym2_output!("");

                let res = self.group().character_table().reduce_characters(
                    &chis.iter().map(|(cc, chi)| (cc, *chi)).collect::<Vec<_>>(),
                    self.integrality_threshold(),
                );
                log::debug!("Characters reduced.");
                log::debug!("Analysing representation symmetry for a multi-determinantal wavefunction... Done.");
                res
            } else {
                Err(DecompositionError(err_str))
            }
        }
    }
}

// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
// Optimised implementation for multi-determinantal wavefunctions constructed from orbits
// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
impl<'a, 'go, 'g, G, T>
    MultiDeterminantSymmetryOrbit<'a, 'go, G, T, OrbitBasis<'g, G, SlaterDeterminant<'a, T>>>
where
    G: SymmetryGroupProperties + Clone,
    G::CharTab: SubspaceDecomposable<T>,
    T: Lapack
        + ComplexFloat<Real = <T as Scalar>::Real>
        + fmt::Debug
        + Mul<<T as ComplexFloat>::Real, Output = T>,
    <T as ComplexFloat>::Real: fmt::Debug
        + Zero
        + From<u16>
        + ToPrimitive
        + approx::RelativeEq<<T as ComplexFloat>::Real>
        + approx::AbsDiffEq<Epsilon = <T as Scalar>::Real>,
    MultiDeterminant<'a, T, OrbitBasis<'g, G, SlaterDeterminant<'a, T>>>: SymmetryTransformable,
{
    /// Calculates and stores the overlap matrix between multi-determinantal wavefunctions in the
    /// orbit, with respect to a metric of the basis in which the constituting Slater determinants
    /// are expressed.
    ///
    /// This function is particularly optimised for multi-determinantal wavefunctions constructed
    /// from orbits of origin Slater determinants such that the multi-determinantal wavefunctions
    /// are never explicitly transformed by group operations.
    ///
    /// # Arguments
    ///
    /// * `metric` - The metric of the basis in which the orbit items are expressed.
    /// * `metric_h` - The complex-symmetric metric of the basis in which the orbit items are
    /// expressed. This is required if antiunitary operations are involved.
    /// * `use_cayley_table` - A boolean indicating if the Cayley table of the group should be used
    /// to speed up the computation of the overlap matrix. If `false`, this will revert back to the
    /// non-optimised overlap matrix calculation.
    pub(crate) fn calc_smat_optimised(
        &mut self,
        metric: Option<&Array2<T>>,
        metric_h: Option<&Array2<T>>,
        use_cayley_table: bool,
    ) -> Result<&mut Self, anyhow::Error> {
        ensure!(
            self.group.name() == self.origin.basis.group().name(),
            "Multi-determinantal wavefunction orbit-generating group does not match symmetry analysis group."
        );

        if let (Some(ctb), true) = (self.group().cayley_table(), use_cayley_table) {
            log::debug!("Cayley table available. Group closure will be used to speed up overlap matrix computation.");

            let order = self.group.order();
            let mut smat = Array2::<T>::zeros((order, order));
            let multidet_0 = self.origin();
            let det_origins = multidet_0.basis.origins();
            let n_det_origins = det_origins.len();

            let detov_kpwx_vec = multidet_0
                .basis
                .iter()
                .collect::<Result<Vec<_>, _>>()?
                .into_iter()
                .cartesian_product(det_origins.iter())
                .map(|(w, x)| w.overlap(x, metric, metric_h))
                .collect::<Result<Vec<_>, _>>()?;
            let detov_kpwx =
                Array3::from_shape_vec((order, n_det_origins, n_det_origins), detov_kpwx_vec)?;
            let ovs = (0..order)
                .map(|k| {
                    [
                        (0..order),
                        (0..order),
                        (0..n_det_origins),
                        (0..n_det_origins),
                    ]
                    .into_iter()
                    .multi_cartesian_product()
                    .try_fold(T::zero(), |acc, v| {
                        let ip = v[0];
                        let jp = v[1];
                        let w = v[2];
                        let x = v[3];
                        let aipw = multidet_0.coefficients[ip * n_det_origins + w];
                        let ajpx = multidet_0.coefficients[jp * n_det_origins + x];

                        let jpinv = ctb.slice(s![.., jp]).iter().position(|&x| x == 0).ok_or(
                            format_err!("Unable to find the inverse of group element `{jp}`."),
                        )?;
                        let kp = ctb[(jpinv, ctb[(k, ip)])];
                        let ukipwjpx = self.norm_preserving_scalar_map(jp)?(detov_kpwx[(kp, w, x)]);

                        Ok::<_, anyhow::Error>(
                            acc + self.norm_preserving_scalar_map(k)?(aipw.conj()) * ukipwjpx * ajpx,
                        )
                    })
                })
                .collect::<Result<Vec<_>, _>>()?;
            for (i, j) in (0..order).cartesian_product(0..order) {
                let jinv = ctb
                    .slice(s![.., j])
                    .iter()
                    .position(|&x| x == 0)
                    .ok_or(format_err!(
                        "Unable to find the inverse of group element `{j}`."
                    ))?;
                let jinv_i = ctb[(jinv, i)];
                smat[(i, j)] = self.norm_preserving_scalar_map(jinv)?(ovs[jinv_i]);
            }
            if self.origin().complex_symmetric() {
                self.set_smat(
                    (smat.clone() + smat.t().to_owned()).mapv(|x| x / (T::one() + T::one())),
                )
            } else {
                self.set_smat(
                    (smat.clone() + smat.t().to_owned().mapv(|x| x.conj()))
                        .mapv(|x| x / (T::one() + T::one())),
                )
            }
            Ok(self)
        } else {
            self.calc_smat(metric, metric_h, use_cayley_table)
        }
    }
}