#![allow(non_snake_case)]
use ndarray::prelude::*;
use ndarray::Zip;

fn main() {
    println!("Example 1:");
    let A: Array2<f64> = arr2(&[
        [1.0, 3.0, 5.0],
        [2.0, 4.0, 7.0],
        [1.0, 1.0, 0.0],
    ]);
    println!("A \n {}", A);
    let (L, U, P) = lu_decomp(&A);
    println!("L \n {}", L);
    println!("U \n {}", U);
    println!("P \n {}", P);

    println!("\nExample 2:");
    let A: Array2<f64> = arr2(&[
        [11.0, 9.0, 24.0, 2.0],
        [1.0, 5.0, 2.0, 6.0],
        [3.0, 17.0, 18.0, 1.0],
        [2.0, 5.0, 7.0, 1.0],
    ]);
    let A:Array2<f64> = array![
            [4.3552 , 6.25851, 4.12662, 1.93708, 0.21272, 3.25683, 6.53326],
            [4.24746, 1.84137, 6.71904, 0.59754, 3.5806 , 3.63597, 5.347  ],
            [2.30479, 1.70591, 3.05354, 1.82188, 5.27839, 7.9166 , 2.04607],
            [2.40158, 6.38524, 7.90296, 4.69683, 6.63801, 7.32958, 1.45936],
            [0.42456, 6.47456, 1.55398, 8.28979, 4.20987, 0.90401, 4.94587],
            [5.78903, 1.92032, 6.20261, 5.78543, 1.94331, 8.25178, 7.47273],
            [1.44797, 7.41157, 7.69495, 8.90113, 3.05983, 0.41582, 6.42932]];

    println!("A \n {}", A);
    let (mut L, mut U, mut P) = lu_decomp(&A);
    println!("L \n {}", L);
    println!("U \n {}", U);
    println!("P \n {}", P);

    let mut _q = lu_decomp(&A);
    for i in 0 .. 1000000 {
        _q = lu_decomp(&A);
    }
}

fn lu_decomp<T: NdFloat>(A: &Array2<T>) -> (Array2<T>, Array2<T>, Array2<T>) {
    fn pivot<T: NdFloat>(A: &Array2<T>) -> Array2<T> {
        fn swap<T: NdFloat>(A: &mut Array2<T>, ir1: usize, ir2: usize) {
            let (.., mut rest) = A.view_mut().split_at(Axis(0), ir1);
            let (r0, mut rest) = rest.view_mut().split_at(Axis(0), 1);
            let (.., mut rest) = rest.view_mut().split_at(Axis(0), ir2 - ir1 - 1);
            let (r1, ..) = rest.view_mut().split_at(Axis(0), 1);

            Zip::from(r0).and(r1).for_each(std::mem::swap);
        }

        let n = A.raw_dim()[0];
        let mut P: Array2<T> = Array::eye(n);
        for (idx, col) in A.axis_iter(Axis(1)).enumerate() {
            // find idx of maximum value in column i
            let mut mp = idx;
            for i in idx .. n {
                if col[mp].abs() < col[i].abs() {
                    mp = i;
                }
            }
            // swap rows if necessary
            if mp != idx {
                swap(&mut P, idx, mp);
            }
        }

        P
    }

    let d = A.raw_dim();
    let n = d[0];
    assert_eq!(n, d[1], "LU decomposition must take a square matrix.");

    let P = pivot(&A);
    let pA = P.dot(A);

    let mut L: Array2<T> = Array::eye(n);
    let mut U: Array2<T> = Array::zeros((n, n));

    for c in 0 .. n {
        for r in 0 .. n {
            let pAs = pA[[r, c]] - U.slice(s![0..r, c]).dot(&L.slice(s![r, 0..r]));

            if r < c + 1 { // U
                U[[r, c]] = pAs;
            } else { // L
                L[[r, c]] = (pAs) / U[[c, c]];
            }
        }
    }

    (L, U, P)
}