pub mod arpack;
use ndarray::{s, Array1, Array2, ArrayView1, ArrayView2};
use std::os::raw::{c_char, c_double, c_int};

#[derive(Debug)]
pub enum Error {
    NonSquare,
    IllegalParameters(String),
    Other(i32),
}

impl std::fmt::Display for Error {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        match self {
            Error::NonSquare => f.write_str("Non square matrix."),
            Error::IllegalParameters(s) => write!(f, "Invalid parameters: {}", s),
            Error::Other(i) => write!(f, "Arpack error (code {}", i),
        }
    }
}

impl std::error::Error for Error {}


#[link(name = "arpack_SUN4")]
extern "C" {
    fn dnaupd_(
        ido: *mut c_int,
        bmat: *const c_char,
        n: *const c_int,
        which: *const c_char,
        nev: *const c_int,
        tol: *const c_double,
        resid: *mut c_double,
        ncv: *const c_int,
        v: *mut c_double,
        ldv: *const c_int,
        iparam: *mut c_int,
        ipntr: *mut c_int,
        workd: *mut c_double,
        workl: *mut c_double,
        lworkl: *const c_int,
        info: *mut c_int,
    );

    fn dneupd_(
        rvec: *const c_int,
        howmny: *const c_char,
        select: *const c_int,
        dr: *mut c_double,
        di: *mut c_double,
        z: *mut c_double,
        ldz: *const c_int,
        sigmar: *const c_double,
        sigmai: *const c_double,
        workev: *mut c_double,
        bmat: *const c_char,
        n: *const c_int,
        which: *const c_char,
        nev: *const c_int,
        tol: *const c_double,
        resid: *mut c_double,
        ncv: *const c_int,
        v: *mut c_double,
        ldv: *const c_int,
        iparam: *mut c_int,
        ipntr: *mut c_int,
        workd: *mut c_double,
        workl: *mut c_double,
        lworkl: *const c_int,
        info: *mut c_int,
    );
}

fn laplacian(n: usize) -> Array2<f64> {
    let mut array: Array2<f64> = Array2::zeros((n, n));
    // let dx = 1. / n as f64;
    let dx = 1. as f64;
    let pf = (1. / dx).powi(2);
    array[[0, 0]] = pf * -2.;
    array[[0, 1]] = pf;
    for i in 1..n - 1 {
        array[[i, i]] = pf * -2.;
        array[[i, i - 1]] = pf;
        array[[i, i + 1]] = pf;
    }
    array[[n - 1, n - 1]] = pf * -2.;
    array[[n - 1, n - 2]] = pf;
    array
}

// Implementation of matrix - vector multiplication
// In this example, op is the tridiagonal laplacian
// and v comes from ARPACK
fn av_3diag(op: ArrayView2<f64>, v: ArrayView1<f64>) -> Array1<f64> {
    let n = v.len();
    let mut out = Array1::zeros(n);
    out[0] = op[[0, 0]] * v[0] + op[[0, 1]] * v[1];
    for i in 1..n - 1 {
        out[i] = op[[i, i]] * v[i] + op[[i, i - 1]] * v[i - 1] + op[[i, i + 1]] * v[i + 1];
    }
    out[n - 1] = op[[n - 1, n - 1]] * v[n - 1] + op[[n - 1, n - 2]] * v[n - 2];
    out
}

fn main() -> Result<(), Error> {
    // let n = 10;
    let nx = 1000;
    let maxiter = 2000;
    let array = laplacian(nx);

    // https://www.caam.rice.edu/software/ARPACK/UG/node137.html

    // Reverse communication flag.  IDO must be zero on the first iterate
    let mut ido: c_int = 0;

    // BMAT = 'I' -> standard eigenvalue problem A*x = lambda*x
    // BMAT = 'G' -> generalized eigenvalue problem A*x = lambda*B*x
    let bmat = "I";

    // N       Integer.  (INPUT)
    //         Dimension of the eigenproblem.
    //let n: c_int = (array.dim().0 * array.dim().1) as c_int;
    let n: c_int = array.dim().0 as c_int;

    //  WHICH   Character*2.  (INPUT)
    //          'LM' -> want the NEV eigenvalues of largest magnitude.
    //          'SM' -> want the NEV eigenvalues of smallest magnitude.
    //          'LR' -> want the NEV eigenvalues of largest real part.
    //          'SR' -> want the NEV eigenvalues of smallest real part.
    //          'LI' -> want the NEV eigenvalues of largest imaginary part.
    //          'SI' -> want the NEV eigenvalues of smallest imaginary part.
    let which = "LR";

    //  NEV     Integer.  (INPUT)
    //          Number of eigenvalues of OP to be computed. 0 < NEV < N-1.
    // let nev: c_int = (dim - 10) as c_int;
    let nev: c_int = 6;

    // TOL     Double precision scalar.  (INPUT)
    //         Stopping criterion.
    let tol: c_double = std::f64::EPSILON;

    //  RESID   Double precision array of length N.  (INPUT/OUTPUT)
    //          On INPUT:
    //          If INFO .EQ. 0, a random initial residual vector is used.
    //          If INFO .NE. 0, RESID contains the initial residual vector,
    //                          possibly from a previous run.
    //          On OUTPUT:
    //          RESID contains the final residual vector.
    let mut resid: Array1<c_double> = Array1::zeros(n as usize);

    //  NCV     Integer.  (INPUT)
    //          Number of columns of the matrix V. NCV must satisfy the two
    //          inequalities 2 <= NCV-NEV and NCV <= N.
    //          This will indicate how many Arnoldi vectors are generated
    //          at each iteration.
    let ncv: c_int = std::cmp::max(20, 2 * nev + 1);
    // let ncv: c_int = 20;

    //  V       Double precision array N by NCV.  (OUTPUT)
    //          Contains the final set of Arnoldi basis vectors.
    let mut v: Array2<c_double> = Array2::zeros((n as usize, ncv as usize));

    //  LDV     Integer.  (INPUT)
    //          Leading dimension of V exactly as declared in the calling program.
    let ldv: c_int = n;

    //  IPARAM  Integer array of length 11.  (INPUT/OUTPUT)
    //          IPARAM(1) = ISHIFT: method for selecting the implicit shifts.
    //          The shifts selected at each iteration are used to restart
    //          the Arnoldi iteration in an implicit fashion.
    // https://www.caam.rice.edu/software/ARPACK/UG/node137.html
    let mut iparam = [0 as c_int; 11];
    iparam[0] = 1 as c_int;
    iparam[2] = maxiter as c_int;
    iparam[6] = 1 as c_int;

    //  IPNTR   Integer array of length 14.  (OUTPUT)
    //          Pointer to mark the starting locations in the WORKD and WORKL
    //          arrays for matrices/vectors used by the Arnoldi iteration.
    let mut ipntr = [0 as c_int; 14];

    //  WORKD   Double precision work array of length 3*N.  (REVERSE COMMUNICATION)
    //          Distributed array to be used in the basic Arnoldi iteration
    //          for reverse communication.  The user should not use WORKD
    //          as temporary workspace during the iteration. Upon termination
    //          WORKD(1:N) contains B*RESID(1:N).
    let mut workd: Array1<c_double> = Array1::zeros(3 * n as usize);

    //  LWORKL  Integer.  (INPUT)
    //          LWORKL must be at least 3*NCV**2 + 6*NCV.
    let lworkl: c_int = 3 * ncv.pow(2) + 6 * ncv;

    //  WORKL   Double precision work array of length LWORKL.  (OUTPUT/WORKSPACE)
    //          Private (replicated) array on each PE or array allocated on
    //          the front end.  See Data Distribution Note below.
    let mut workl: Array1<c_double> = Array1::zeros(lworkl as usize);

    //  INFO    Integer.  (INPUT/OUTPUT)
    //          If INFO .EQ. 0, a randomly initial residual vector is used.
    //          If INFO .NE. 0, RESID contains the initial residual vector,
    //                          possibly from a previous run
    let mut info: c_int = 0;

    // dnaupd
    while ido != 99 {
        unsafe {
            dnaupd_(
                &mut ido,
                bmat.as_ptr() as *const i8,
                &n,
                which.as_ptr() as *const i8,
                &nev,
                &tol,
                resid.as_mut_ptr(),
                &ncv,
                v.as_mut_ptr(),
                &ldv,
                iparam.as_mut_ptr(),
                ipntr.as_mut_ptr(),
                workd.as_mut_ptr(),
                workl.as_mut_ptr(),
                &lworkl,
                &mut info,
            );
        }
        // println!("{:?}", ido);
        // println!("{:?}", info);

        if (ido == -1) || (ido == 1) {
            let v = workd
                .slice(s![ipntr[0] as usize - 1..ipntr[0] as usize + n as usize - 1])
                .to_owned();
            let v2 = av_3diag(array.view(), v.view());
            workd
                .slice_mut(s![ipntr[1] as usize - 1..ipntr[1] as usize + n  as usize - 1])
                .assign(&v2);
        }

        // println!("{:?}", ncv - nev);
    }
    println!("INFO {:?}", info);
    
    match info {
        0 | 1 | 2 => {}
        -1 => return Err(Error::IllegalParameters("N must be positive.".to_string())),
        -2 => {
            return Err(Error::IllegalParameters(
                "NEV must be positive.".to_string(),
            ))
        }
        -3 => {
            return Err(Error::IllegalParameters(
                "NCV-NEV >= 2 and less than or equal to N.".to_string(),
            ))
        }
        -4 => {
            return Err(Error::IllegalParameters(
                "Maximum iterations must be greater than 0.".to_string(),
            ))
        }
        -5 => {
            return Err(Error::IllegalParameters(
                "Maximum iterations must be greater than 0.".to_string(),
            ))
        }
        i => return Err(Error::Other(i)),
    }

    // dneupd

    // 1: Return eigenvectors
    let rvec: c_int = true as c_int;

    //  HOWMNY  Character*1  (INPUT)
    //          Specifies the form of the basis for the invariant subspace
    //          corresponding to the converged Ritz values that is to be computed.
    let howmny = "A";

    // SELECT  Logical array of dimension NCV.  (INPUT)
    let select = vec![false as c_int; ncv as usize];

    //  DR      Double precision array of dimension NEV+1. (OUTPUT)
    //          On exit: DR contains the real part of the Ritz approximations
    //          to the eigenvalues of A*z = lambda*B*z.
    let mut dr: Array1<f64> = Array1::zeros((nev + 1) as usize);

    //  DI      Double precision array of dimension NEV+1. (OUTPUT)
    //          On exit: DR contains the imaginary part of the Ritz approximations
    //          to the eigenvalues of A*z = lambda*B*z.
    let mut di: Array1<f64> = Array1::zeros((nev + 1) as usize);

    //  Z       On exit, if RVEC = 1 and HOWMANY = 'A', then the columns of Z represent
    //          approximate eigenvectors (Ritz vectors) corresponding to the NCONV=IPARAM(5)
    //          Ritz values for eigensystem A*z = lambda*B*z.
    let mut z: Array2<f64> = Array2::zeros((n as usize, nev as usize));

    // LDZ      Integer.  (INPUT)
    //          The leading dimension of the array Z.  If Ritz vectors are
    //          desired, then  LDZ >= max( 1, N ).  In any case,  LDZ >= 1.
    let ldz: c_int = n;

    // SIGMAR   If IPARAM(7) = 3 or 4, represents the real part of the shift.
    //          Not referenced if IPARAM(7) = 1 or 2.
    let sigmar: c_double = 0.;
    // SIGMAI   If IPARAM(7) = 3 or 4, represents the imaginary part of the shift.
    //          Not referenced if IPARAM(7) = 1 or 2.
    let sigmai: c_double = 0.;
    // WORKEV   Double precision work array of dimension 3*NCV. (WORKSPACE)
    let mut workev: Array1<f64> = Array1::zeros(3 * ncv as usize);

    unsafe {
        dneupd_(
            &rvec,
            howmny.as_ptr() as *const i8,
            select.as_ptr(),
            dr.as_mut_ptr(),
            di.as_mut_ptr(),
            z.as_mut_ptr(),
            &ldz,
            &sigmar,
            &sigmai,
            workev.as_mut_ptr(),
            // same as dnaupd_
            bmat.as_ptr() as *const i8,
            &n,
            which.as_ptr() as *const i8,
            &nev,
            &tol,
            resid.as_mut_ptr(),
            &ncv,
            v.as_mut_ptr(),
            &ldv,
            iparam.as_mut_ptr(),
            ipntr.as_mut_ptr(),
            workd.as_mut_ptr(),
            workl.as_mut_ptr(),
            &lworkl,
            &mut info,
        );
    }

    //println!("INFO {:?}", info);
    println!("{:?}", dr);
    println!("{:?}", di);



    // match info {
    //     0 | 1 | 2 => {}
    //     -1 => return Err(Error::IllegalParameters("N must be positive.".to_string())),
    //     -2 => {
    //         return Err(Error::IllegalParameters(
    //             "NEV must be positive.".to_string(),
    //         ))
    //     }
    //     -3 => {
    //         return Err(Error::IllegalParameters(
    //             "NCV-NEV >= 2 and less than or equal to N.".to_string(),
    //         ))
    //     }
    //     -4 => {
    //         return Err(Error::IllegalParameters(
    //             "Maximum iterations must be greater than 0.".to_string(),
    //         ))
    //     }
    //     -5 => {
    //         return Err(Error::IllegalParameters(
    //             "Maximum iterations must be greater than 0.".to_string(),
    //         ))
    //     }
    //     i => return Err(Error::Other(i)),
    // }


    // println!("{:?}", ido);
    // println!("{:?}", ii);
    //
    // println!("{:?}", v);
    // println!("");
    // println!("{:?}", v.slice(s![..,0]));
    // println!("");
    // println!("{:?}", v.slice(s![..,ncv-1]));

    println!("Hello, world!");
    Ok(())
}
