sci-form 0.15.2

//! Self-Consistent Field (SCF) solver with DIIS acceleration.
//!
//! Solves the Roothaan-Hall equations iteratively:
//! $$\mathbf{FC} = \mathbf{SC}\boldsymbol{\varepsilon}$$
//!
//! Uses Löwdin orthogonalization ($S^{-1/2}$) to transform the generalized
//! eigenvalue problem into a standard one, then applies DIIS for convergence.

use super::fock::{build_fock, electronic_energy};
use nalgebra::{DMatrix, DVector, SymmetricEigen};

#[cfg(feature = "experimental-gpu")]
fn matrix_to_row_major(matrix: &DMatrix<f64>) -> Vec<f64> {
    let nrows = matrix.nrows();
    let ncols = matrix.ncols();
    let mut data = Vec::with_capacity(nrows * ncols);
    for row in 0..nrows {
        for col in 0..ncols {
            data.push(matrix[(row, col)]);
        }
    }
    data
}

#[cfg(feature = "experimental-gpu")]
fn expand_packed_eris(eris: &[f64], n_basis: usize) -> Vec<f64> {
    let mut full = vec![0.0; n_basis * n_basis * n_basis * n_basis];
    for mu in 0..n_basis {
        for nu in 0..n_basis {
            for lam in 0..n_basis {
                for sig in 0..n_basis {
                    let packed = super::integrals::get_eri(eris, mu, nu, lam, sig, n_basis);
                    let idx = mu * n_basis * n_basis * n_basis
                        + nu * n_basis * n_basis
                        + lam * n_basis
                        + sig;
                    full[idx] = packed;
                }
            }
        }
    }
    full
}

/// SCF convergence result.
///
/// Contains converged matrices (density, MO coefficients) and derived
/// scalars. The `DMatrix<f64>` fields prevent `Serialize`/`Deserialize`;
/// use `Hf3cResult` for the serializable public API.
#[derive(Debug, Clone)]
pub struct ScfResult {
    /// Converged electronic energy (Hartree).
    pub energy: f64,
    /// Orbital energies (eigenvalues).
    pub orbital_energies: Vec<f64>,
    /// MO coefficient matrix C (columns are MOs).
    pub coefficients: DMatrix<f64>,
    /// Converged density matrix P.
    pub density: DMatrix<f64>,
    /// Number of SCF iterations.
    pub iterations: usize,
    /// Whether SCF converged.
    pub converged: bool,
}

/// SCF solver configuration.
pub struct ScfConfig {
    pub max_iter: usize,
    pub energy_threshold: f64,
    pub density_threshold: f64,
    pub diis_size: usize,
    /// Level shift for virtual orbitals (Hartree). 0.0 = disabled.
    /// Typical values: 0.1-0.5 Eh for difficult convergence.
    pub level_shift: f64,
}

impl Default for ScfConfig {
    fn default() -> Self {
        ScfConfig {
            max_iter: 100,
            energy_threshold: 1e-8,
            density_threshold: 1e-6,
            diis_size: 6,
            level_shift: 0.0,
        }
    }
}

/// Run the SCF procedure.
pub fn solve_scf(
    h_core: &DMatrix<f64>,
    s_mat: &DMatrix<f64>,
    eris: &[f64],
    _gpu_eris_full: Option<&[f64]>,
    n_electrons: usize,
    config: &ScfConfig,
) -> ScfResult {
    let n = h_core.nrows();
    let n_occ = n_electrons / 2;

    #[cfg(feature = "experimental-gpu")]
    let gpu_ctx = if n >= 4 {
        crate::gpu::context::GpuContext::try_create().ok()
    } else {
        None
    };

    #[cfg(feature = "experimental-gpu")]
    let h_core_row_major = matrix_to_row_major(h_core);

    #[cfg(feature = "experimental-gpu")]
    let eris_full = _gpu_eris_full
        .map(|tensor| tensor.to_vec())
        .or_else(|| gpu_ctx.as_ref().map(|_| expand_packed_eris(eris, n)));

    // Löwdin orthogonalization: S^{-1/2}
    let s_half_inv = lowdin_orthogonalization(s_mat);

    // Initial guess: diagonalize H_core
    let (mut energies, mut coeffs) = diagonalize_fock(h_core, &s_half_inv);
    let mut density = build_density(&coeffs, n_occ);

    let mut prev_energy = 0.0;
    let mut converged = false;
    let mut iterations = 0;

    // DIIS storage
    let mut diis_focks: Vec<DMatrix<f64>> = Vec::new();
    let mut diis_errors: Vec<DMatrix<f64>> = Vec::new();

    let mut min_error_norm = f64::MAX;
    let mut consecutive_increases = 0u32;

    for iter in 0..config.max_iter {
        iterations = iter + 1;

        let fock = {
            #[cfg(feature = "experimental-gpu")]
            {
                if let (Some(ctx), Some(eri_tensor)) = (gpu_ctx.as_ref(), eris_full.as_ref()) {
                    let density_row_major = matrix_to_row_major(&density);
                    if let Ok(fock_flat) = crate::gpu::fock_build_gpu::build_fock_gpu(
                        ctx,
                        &h_core_row_major,
                        &density_row_major,
                        eri_tensor,
                        n,
                    ) {
                        DMatrix::from_row_slice(n, n, &fock_flat)
                    } else {
                        build_fock(h_core, &density, eris, n)
                    }
                } else {
                    build_fock(h_core, &density, eris, n)
                }
            }

            #[cfg(not(feature = "experimental-gpu"))]
            {
                build_fock(h_core, &density, eris, n)
            }
        };
        let energy = electronic_energy(&density, h_core, &fock);

        // DIIS error: FPS - SPF
        let error = &fock * &density * s_mat - s_mat * &density * &fock;
        let error_norm = error.iter().map(|v| v * v).sum::<f64>().sqrt();

        // Robust DIIS reset: require 3 consecutive error increases before resetting,
        // rather than a single large jump (which can occur during normal oscillations).
        if error_norm > min_error_norm * 1.5 {
            consecutive_increases += 1;
        } else {
            consecutive_increases = 0;
        }
        if consecutive_increases >= 3 && diis_focks.len() > 2 {
            diis_focks.clear();
            diis_errors.clear();
            consecutive_increases = 0;
        }
        if error_norm < min_error_norm {
            min_error_norm = error_norm;
        }

        // DIIS extrapolation
        diis_focks.push(fock.clone());
        diis_errors.push(error);
        if diis_focks.len() > config.diis_size {
            diis_focks.remove(0);
            diis_errors.remove(0);
        }

        let fock_diis = if diis_focks.len() >= 2 {
            diis_extrapolate(&diis_focks, &diis_errors)
        } else {
            fock
        };

        // Apply level shifting to virtual orbitals if configured
        let fock_shifted = if config.level_shift > 0.0 && n_occ < n {
            // Level shifting: add shift to virtual orbital block of Fock matrix
            // in the orthogonal basis. This helps convergence for metallic systems.
            let f_orth = s_half_inv.transpose() * &fock_diis * &s_half_inv;
            let eigen = SymmetricEigen::new(f_orth);
            let mut shifted_evals = eigen.eigenvalues.clone();
            let mut idx_sorted: Vec<usize> = (0..n).collect();
            idx_sorted.sort_by(|&a, &b| shifted_evals[a].partial_cmp(&shifted_evals[b]).unwrap());
            for &idx in idx_sorted.iter().skip(n_occ) {
                shifted_evals[idx] += config.level_shift;
            }
            let v = &eigen.eigenvectors;
            let d = DMatrix::from_diagonal(&shifted_evals);
            let f_shifted = v * d * v.transpose();
            // Transform back to AO basis: F_AO = S^{1/2} * F_orth * S^{1/2}^T
            // But for diagonalize_fock, it transforms again, so keep in AO
            &s_half_inv * f_shifted * s_half_inv.transpose()
        } else {
            fock_diis
        };

        let (new_energies, new_coeffs) = diagonalize_fock(&fock_shifted, &s_half_inv);
        let new_density = build_density(&new_coeffs, n_occ);

        let de = (energy - prev_energy).abs();

        if de < config.energy_threshold && error_norm < config.density_threshold {
            converged = true;
            energies = new_energies;
            coeffs = new_coeffs;
            density = new_density;
            break;
        }

        prev_energy = energy;
        energies = new_energies;
        coeffs = new_coeffs;
        density = new_density;
    }

    let final_energy = electronic_energy(&density, h_core, &build_fock(h_core, &density, eris, n));

    ScfResult {
        energy: final_energy,
        orbital_energies: energies.as_slice().to_vec(),
        coefficients: coeffs,
        density,
        iterations,
        converged,
    }
}

fn lowdin_orthogonalization(s: &DMatrix<f64>) -> DMatrix<f64> {
    let eigen = SymmetricEigen::new(s.clone());
    let n = s.nrows();
    let mut s_inv_half = DMatrix::zeros(n, n);

    let max_val = eigen
        .eigenvalues
        .iter()
        .copied()
        .fold(0.0f64, |a, b| a.max(b));
    let threshold = max_val * 1e-10;

    for i in 0..n {
        let val = eigen.eigenvalues[i];
        if val > threshold {
            let factor = 1.0 / val.sqrt();
            let col = eigen.eigenvectors.column(i);
            s_inv_half += factor * col * col.transpose();
        }
    }
    s_inv_half
}

fn diagonalize_fock(
    fock: &DMatrix<f64>,
    s_half_inv: &DMatrix<f64>,
) -> (DVector<f64>, DMatrix<f64>) {
    let f_prime = s_half_inv.transpose() * fock * s_half_inv;
    let eigen = SymmetricEigen::new(f_prime);

    // Sort eigenvalues/vectors by energy
    let n = eigen.eigenvalues.len();
    let mut indices: Vec<usize> = (0..n).collect();
    indices.sort_by(|&a, &b| {
        eigen.eigenvalues[a]
            .partial_cmp(&eigen.eigenvalues[b])
            .unwrap()
    });

    let sorted_energies = DVector::from_fn(n, |i, _| eigen.eigenvalues[indices[i]]);
    let sorted_vecs = DMatrix::from_fn(n, n, |r, c| eigen.eigenvectors[(r, indices[c])]);

    // Back-transform: C = S^{-1/2} C'
    let coeffs = s_half_inv * sorted_vecs;
    (sorted_energies, coeffs)
}

fn build_density(coeffs: &DMatrix<f64>, n_occ: usize) -> DMatrix<f64> {
    let n = coeffs.nrows();
    let mut density = DMatrix::zeros(n, n);
    for i in 0..n_occ {
        let col = coeffs.column(i);
        density += 2.0 * &col * col.transpose();
    }
    density
}

fn diis_extrapolate(focks: &[DMatrix<f64>], errors: &[DMatrix<f64>]) -> DMatrix<f64> {
    let m = errors.len();
    let mut b = DMatrix::zeros(m + 1, m + 1);

    for i in 0..m {
        for j in 0..=i {
            let bij: f64 = errors[i]
                .iter()
                .zip(errors[j].iter())
                .map(|(a, b)| a * b)
                .sum();
            b[(i, j)] = bij;
            b[(j, i)] = bij;
        }
    }
    for i in 0..m {
        b[(m, i)] = -1.0;
        b[(i, m)] = -1.0;
    }

    let mut rhs = DVector::zeros(m + 1);
    rhs[m] = -1.0;

    // Solve B·c = rhs using pseudo-inverse
    let svd = b.svd(true, true);
    let c = match svd.solve(&rhs, 1e-10) {
        Ok(c) => c,
        Err(_) => {
            // Fallback: use latest Fock
            return focks.last().unwrap().clone();
        }
    };

    // Guard against ill-conditioned DIIS: if any coefficient is too large,
    // the extrapolation is unreliable (undershoot/overshoot). Fall back.
    let max_coeff = c.iter().take(m).map(|v| v.abs()).fold(0.0f64, f64::max);
    if max_coeff > 10.0 {
        return focks.last().unwrap().clone();
    }

    let mut f_diis = DMatrix::zeros(focks[0].nrows(), focks[0].ncols());
    for i in 0..m {
        f_diis += c[i] * &focks[i];
    }
    f_diis
}

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

    #[test]
    fn test_lowdin_identity() {
        let s = DMatrix::identity(3, 3);
        let s_inv = lowdin_orthogonalization(&s);
        for i in 0..3 {
            for j in 0..3 {
                let expected = if i == j { 1.0 } else { 0.0 };
                assert!(
                    (s_inv[(i, j)] - expected).abs() < 1e-10,
                    "S^{{-1/2}} of identity should be identity"
                );
            }
        }
    }
}