mdarray-linalg-lapack 0.2.0

use std::{ptr, ptr::null_mut};

use mdarray::{Dim, Layout, Shape, Slice};
use mdarray_linalg::{
    eig::{EigError, SchurError},
    utils::transpose_in_place,
};
use num_complex::ComplexFloat;

use super::scalar::{LapackScalar, NeedsRwork};

pub(super) fn geig<
    La: Layout,
    Ler: Layout,
    Lei: Layout,
    Lvl: Layout,
    Lvr: Layout,
    T: ComplexFloat + Default + LapackScalar + NeedsRwork<Elem = T>,
    D0: Dim,
    D1: Dim,
>(
    a: &mut Slice<T, (D0, D1), La>,
    eigenvalues_real: &mut Slice<T, (D0,), Ler>,
    eigenvalues_imag: &mut Slice<T, (D0,), Lei>,
    mut left_eigenvectors: Option<&mut Slice<T, (D0, D1), Lvl>>,
    mut right_eigenvectors: Option<&mut Slice<T, (D0, D1), Lvr>>,
) -> Result<(), EigError>
where
    T::Real: Into<T>,
{
    let ash = *a.shape();
    let (m, n) = (ash.dim(0), ash.dim(1));

    if m != n {
        return Err(EigError::NotSquareMatrix);
    }

    let jobvl = if left_eigenvectors.is_some() {
        'V'
    } else {
        'N'
    };
    let jobvr = if right_eigenvectors.is_some() {
        'V'
    } else {
        'N'
    };

    // Validate dimensions if eigenvectors are requested
    if let Some(ref vl) = left_eigenvectors {
        let vlsh = *vl.shape();
        let (mvl, nvl) = (vlsh.dim(0), vlsh.dim(1));

        assert_eq!(
            mvl, n,
            "Left eigenvectors must have same number of rows as A"
        );
        assert_eq!(nvl, n, "Left eigenvectors must be square (n × n)");
    }

    if let Some(ref vr) = right_eigenvectors {
        let vrsh = *vr.shape();
        let (mvr, nvr) = (vrsh.dim(0), vrsh.dim(1));

        assert_eq!(
            mvr, n,
            "Right eigenvectors must have same number of rows as A"
        );
        assert_eq!(nvr, n, "Right eigenvectors must be square (n × n)");
    }

    // Validate eigenvalue vectors dimensions
    let ersh = *eigenvalues_real.shape();
    let ner = ersh.dim(0);

    let eish = *eigenvalues_imag.shape();
    let nei = eish.dim(0);

    assert_eq!(ner, n, "Real eigenvalues must have n elements");
    assert_eq!(nei, n, "Imaginary eigenvalues must have n elements");

    let vl_ptr: *mut T = left_eigenvectors
        .as_mut()
        .map_or(null_mut(), |x| x.as_mut_ptr());
    let vr_ptr: *mut T = right_eigenvectors
        .as_mut()
        .map_or(null_mut(), |x| x.as_mut_ptr());

    let info = call_geev(
        a,
        n.try_into().unwrap(),
        eigenvalues_real.as_mut_ptr(),
        eigenvalues_imag.as_mut_ptr(),
        vl_ptr,
        vr_ptr,
        jobvl,
        jobvr,
    );

    if info < 0 {
        panic!(
            "Invalid argument to EIG: the {}-th parameter had an illegal value.",
            -info
        );
    } else if info > 0 {
        Err(EigError::BackendDidNotConverge { iterations: info })
    } else {
        // Transpose eigenvectors if computed (LAPACK returns column-major)
        if let Some(ref mut vl) = left_eigenvectors {
            transpose_in_place(vl);
        }
        if let Some(ref mut vr) = right_eigenvectors {
            transpose_in_place(vr);
        }
        Ok(())
    }
}

pub(super) fn geigh<
    La: Layout,
    Lw: Layout,
    T: ComplexFloat + Default + LapackScalar + NeedsRwork<Elem = T>,
    D0: Dim,
    D1: Dim,
>(
    a: &mut Slice<T, (D0, D1), La>,
    eigenvalues: &mut Slice<T::Real, (D0,), Lw>,
) -> Result<(), EigError>
where
    T::Real: Into<T>,
{
    let ash = *a.shape();
    let (m, n) = (ash.dim(0), ash.dim(1));

    if m != n {
        return Err(EigError::NotSquareMatrix);
    }

    assert_eq!(
        eigenvalues.shape().dim(0),
        n,
        "Eigenvalues must have n elements"
    );

    let info = call_syev(
        a,
        n.try_into().unwrap(),
        eigenvalues.as_mut_ptr(),
        'V', // Always compute eigenvectors for geigh
        'U', // Use upper triangle
    );

    if info < 0 {
        panic!(
            "Invalid argument to EIGH: the {}-th parameter had an illegal value.",
            -info
        );
    } else if info > 0 {
        Err(EigError::BackendDidNotConverge { iterations: info })
    } else {
        Ok(())
    }
}

#[allow(clippy::too_many_arguments)]
fn call_geev<
    T: ComplexFloat + Default + LapackScalar + NeedsRwork<Elem = T>,
    La: Layout,
    D0: Dim,
    D1: Dim,
>(
    a: &mut Slice<T, (D0, D1), La>,
    n: i32,
    wr_ptr: *mut T,
    wi_ptr: *mut T,
    vl_ptr: *mut T,
    vr_ptr: *mut T,
    jobvl: char,
    jobvr: char,
) -> i32
where
    T::Real: Into<T>,
{
    let mut work = T::allocate(1);
    let lwork = -1i32;
    let mut info = 0;

    let row_major = a.stride(1) == 1;
    assert!(
        row_major || a.stride(0) == 1,
        "a must be contiguous in one dimension"
    );

    if row_major {
        transpose_in_place(a);
    }

    let mut rwork = vec![0.0; T::rwork_len_geev(n)];

    // First call to query optimal workspace size
    unsafe {
        T::lapack_geev(
            jobvl as i8,
            jobvr as i8,
            n,
            a.as_mut_ptr(),
            n,
            wr_ptr,
            wi_ptr,
            vl_ptr,
            n,
            vr_ptr,
            n,
            work.as_mut_ptr(),
            lwork,
            rwork.as_mut_ptr() as *mut _,
            &mut info,
        );
    }

    let lwork = T::lwork_from_query(&work[0]);
    let mut work = T::allocate(lwork);

    // Second call with optimal workspace
    unsafe {
        T::lapack_geev(
            jobvl as i8,
            jobvr as i8,
            n,
            a.as_mut_ptr(),
            n,
            wr_ptr,
            wi_ptr,
            vl_ptr,
            n,
            vr_ptr,
            n,
            work.as_mut_ptr(),
            lwork,
            rwork.as_mut_ptr() as *mut _,
            &mut info,
        );
    }

    info
}

fn call_syev<
    T: ComplexFloat + Default + LapackScalar + NeedsRwork<Elem = T>,
    La: Layout,
    D0: Dim,
    D1: Dim,
>(
    a: &mut Slice<T, (D0, D1), La>,
    n: i32,
    w_ptr: *mut T::Real,
    jobz: char,
    uplo: char,
) -> i32
where
    T::Real: Into<T>,
{
    let mut work = T::allocate(1);
    let lwork = -1i32;
    let mut info = 0;

    let row_major = a.stride(1) == 1;
    assert!(
        row_major || a.stride(0) == 1,
        "a must be contiguous in one dimension"
    );

    if row_major {
        transpose_in_place(a);
    }

    let mut rwork = vec![0.0; T::rwork_len_syev(n)];

    // First call to query optimal workspace size
    unsafe {
        T::lapack_syev(
            jobz as i8,
            uplo as i8,
            n,
            a.as_mut_ptr(),
            n,
            w_ptr.cast(),
            work.as_mut_ptr(),
            lwork,
            rwork.as_mut_ptr() as *mut _,
            &mut info,
        );
    }

    let lwork = T::lwork_from_query(&work[0]);
    let mut work = T::allocate(lwork);

    // Second call with optimal workspace
    unsafe {
        T::lapack_syev(
            jobz as i8,
            uplo as i8,
            n,
            a.as_mut_ptr(),
            n,
            w_ptr.cast(),
            work.as_mut_ptr(),
            lwork,
            rwork.as_mut_ptr() as *mut _,
            &mut info,
        );
    }

    info
}

pub(super) fn gees<
    La: Layout,
    Lwr: Layout,
    Lwi: Layout,
    Lvs: Layout,
    T: ComplexFloat + Default + LapackScalar + NeedsRwork<Elem = T>,
    D0: Dim,
    D1: Dim,
>(
    a: &mut Slice<T, (D0, D1), La>,
    eigenvalues_real: &mut Slice<T, (D0,), Lwr>,
    eigenvalues_imag: &mut Slice<T, (D0,), Lwi>,
    schur_vectors: &mut Slice<T, (D0, D1), Lvs>,
) -> Result<(), SchurError>
where
    T::Real: Into<T>,
{
    let ash = *a.shape();
    let (m, n) = (ash.dim(0) as i32, ash.dim(1) as i32);

    if m != n {
        return Err(SchurError::NotSquareMatrix);
    }

    let jobvs = b'V';

    let svsh = *schur_vectors.shape();
    let (mvs, nvs) = (svsh.dim(0) as i32, svsh.dim(1) as i32);

    assert_eq!(mvs, n, "Schur vectors must have same number of rows as A");
    assert_eq!(nvs, n, "Schur vectors must be square (n × n)");

    let ersh = *eigenvalues_real.shape();
    let eish = *eigenvalues_imag.shape();

    let nwr = ersh.dim(0) as i32;
    let nwi = eish.dim(0) as i32;

    assert_eq!(nwr, n, "Real eigenvalues must have n elements");
    assert_eq!(nwi, n, "Imaginary eigenvalues must have n elements");

    transpose_in_place(a);

    // Allocate workspace
    let mut sdim = 0i32;
    let lwork = {
        let mut work = vec![T::default(); 1];
        let mut info = 0;
        let rwork_len = T::rwork_len_gees(n);
        let mut rwork = vec![T::default(); rwork_len];
        let mut bwork_storage = if n > 0 {
            Some(vec![0i32; n as usize])
        } else {
            None
        }; // keep the vec alive

        let bwork: *mut i32 = bwork_storage
            .as_mut()
            .map_or(ptr::null_mut(), |v| v.as_mut_ptr());

        unsafe {
            T::lapack_gees(
                jobvs.try_into().unwrap(),
                b'N'.try_into().unwrap(), // No sorting
                ptr::null_mut(),          // Null select as *mut c_void
                n,
                a.as_mut_ptr(),
                n,
                &mut sdim,
                eigenvalues_real.as_mut_ptr(),
                eigenvalues_imag.as_mut_ptr(),
                schur_vectors.as_mut_ptr(),
                n,
                work.as_mut_ptr(),
                -1, // Query workspace size
                rwork.as_mut_ptr(),
                bwork,
                &mut info,
            );
        }

        if info != 0 {
            panic!("Error during workspace query: info = {info}");
        }
        T::lwork_from_query(&work[0])
    };

    let mut work = T::allocate(lwork);
    let rwork_len = T::rwork_len_gees(n);
    let mut rwork = vec![T::default(); rwork_len];
    let mut bwork_storage = if n > 0 {
        Some(vec![0i32; n as usize])
    } else {
        None
    }; // keep the vec alive

    let bwork: *mut i32 = bwork_storage
        .as_mut()
        .map_or(ptr::null_mut(), |v| v.as_mut_ptr());

    let mut info = 0;
    unsafe {
        T::lapack_gees(
            jobvs.try_into().unwrap(),
            b'N'.try_into().unwrap(), // No sorting
            ptr::null_mut(),          // Null select as *mut c_void
            n,
            a.as_mut_ptr(),
            n,
            &mut sdim,
            eigenvalues_real.as_mut_ptr(),
            eigenvalues_imag.as_mut_ptr(),
            schur_vectors.as_mut_ptr(),
            n,
            work.as_mut_ptr(),
            lwork,
            rwork.as_mut_ptr(),
            bwork,
            &mut info,
        );
    }

    if info < 0 {
        return Err(SchurError::BackendError(-info));
    } else if info > 0 {
        return Err(SchurError::BackendDidNotConverge { iterations: info });
    }
    Ok(())
}
pub(super) fn gees_complex<
    La: Layout,
    Lw: Layout,
    Lvs: Layout,
    T: ComplexFloat + Default + LapackScalar + NeedsRwork<Elem = T>,
    D0: Dim,
    D1: Dim,
>(
    a: &mut Slice<T, (D0, D1), La>,
    eigenvalues: &mut Slice<T, (D0,), Lw>,       // shape (1, n)
    schur_vectors: &mut Slice<T, (D0, D1), Lvs>, // shape (n, n)
) -> Result<(), SchurError> {
    let ash = *a.shape();
    let (m, n) = (ash.dim(0), ash.dim(1));

    if m != n {
        return Err(SchurError::NotSquareMatrix);
    }

    let jobvs = b'V';

    // shape checks

    let svsh = *schur_vectors.shape();
    let (mvs, nvs) = (svsh.dim(0), svsh.dim(1));
    assert_eq!(mvs, n, "Schur vectors must have same number of rows as A");
    assert_eq!(nvs, n, "Schur vectors must be square (n × n)");

    transpose_in_place(a);

    // --- workspace query (lwork) ---
    let mut sdim: i32 = 0;

    let lwork = {
        let mut work = vec![T::default(); 1];
        let mut info = 0i32;

        // rwork is real workspace
        let mut rwork = vec![0.; 1];

        // keep bwork alive
        let mut bwork_storage = if n > 0 { Some(vec![0i32; n]) } else { None };
        let bwork: *mut i32 = bwork_storage
            .as_mut()
            .map_or(ptr::null_mut(), |v| v.as_mut_ptr());

        unsafe {
            T::lapack_gees(
                jobvs.try_into().unwrap(),
                b'N'.try_into().unwrap(), // no sorting
                ptr::null_mut(),          // select (unused)
                n.try_into().unwrap(),
                a.as_mut_ptr(),
                n.try_into().unwrap(),
                &mut sdim,
                eigenvalues.as_mut_ptr(),     // wr
                ptr::null_mut(),              // _wi (unused for complex)
                schur_vectors.as_mut_ptr(),   // vs
                n.try_into().unwrap(),        // ldvs
                work.as_mut_ptr(),            // work (query)
                -1,                           // lwork = -1 -> query
                rwork.as_mut_ptr() as *mut _, // rwork (cast to match expected pointer)
                bwork,
                &mut info,
            );
        }

        if info != 0 {
            panic!("Error during workspace query (complex gees): info = {info}");
        }
        T::lwork_from_query(&work[0])
    };

    // allocate real work + bwork + rwork
    let mut work = T::allocate(lwork);
    let rwork_len = T::rwork_len_gees(n.try_into().unwrap());
    // let mut rwork = T::allocate(rwork_len as i32);
    let mut rwork = vec![0.0; rwork_len];
    let mut bwork_storage = if n > 0 { Some(vec![0i32; n]) } else { None };
    let bwork: *mut i32 = bwork_storage
        .as_mut()
        .map_or(ptr::null_mut(), |v| v.as_mut_ptr());

    let mut info = 0i32;
    unsafe {
        T::lapack_gees(
            jobvs.try_into().unwrap(),
            b'N'.try_into().unwrap(), // no sorting
            ptr::null_mut(),          // select
            n.try_into().unwrap(),
            a.as_mut_ptr(),
            n.try_into().unwrap(),
            &mut sdim,
            eigenvalues.as_mut_ptr(),   // wr
            ptr::null_mut(),            // _wi (unused)
            schur_vectors.as_mut_ptr(), // vs
            n.try_into().unwrap(),      // ldvs
            work.as_mut_ptr(),
            lwork,
            rwork.as_mut_ptr() as *mut _,
            bwork,
            &mut info,
        );
    }

    if info < 0 {
        return Err(SchurError::BackendError(-info));
    } else if info > 0 {
        return Err(SchurError::BackendDidNotConverge { iterations: info });
    }

    Ok(())
}