/* This file is part of the Amalthea library.
 *
 * Copyright (C) 2019, 2021 Eugene 'Vindex' Stulin
 *
 * Distributed under the Boost Software License 1.0 or (at your option)
 * the GNU Lesser General Public License 3.0 or later.
 */

module amalthea.matrixmath;

import std.exception;
import std.conv;


/*******************************************************************************
 * The function gets minor of matrix by i,j.
 */
auto calcMinor(T)(in T[][] mat, in size_t i, in size_t j) {
    auto m = mat.length > i+1 ? (mat[0 .. i] ~ mat[i+1 .. $]).dup
                              : mat[0 .. i].dup;
    if (m[0].length > j+1) {
        foreach(rn; 0 .. m.length) {
            m[rn] = m[rn][0 .. j] ~ m[rn][j+1 .. $];
        }
    } else {
        foreach(rn; 0 .. m.length) {
            m[rn] = m[rn][0 .. j];
        }
    }
    return m.dup;
}


/*******************************************************************************
 * The function gets determinant of matrix.
 */
T calcDeterminant(T)(in T[][] m) {
    enforce(m.length == m[0].length, "Matrix is non-square.");
    T d = 0;
    size_t i = 0;
    if (m.length == 1 && m[0].length == 1) {
        return m[0][0];
    }
    T minusOne = -1;
    foreach(j; 0 .. m.length) {
        d += m[i][j] * minusOne^^(i+j) * calcDeterminant(calcMinor(m, i, j));
    }
    return d;
}
alias det = calcDeterminant;
unittest {
    auto m = [[-8, 16], [4, 0]];
    assert(det(m) == -64);

    m = [
        [-5, -3, -1],
        [10, 11, -5],
        [ 4,  8, 16]
    ];
    assert(det(m) == -576);

    m = [
        [-5, -3, -1,  1],
        [ 4,  8, 16, 12],
        [10, 11, -5, 10],
        [ 9,  8,  7,  6]
    ];
    assert(m.det == 1040);
}

/*******************************************************************************
 * The function gets the sum of two matrices.
 */
auto summarize(T)(in T[][] m1, in T[][] m2) {
    auto rowsNumber1 = m1.length;
    auto columnsNumber1 = m1[0].length;
    auto rowsNumber2 = m2.length;
    auto columnsNumber2 = m2[0].length;
    enforce(
        columnsNumber1 == columnsNumber2
        && rowsNumber1 == rowsNumber2,
        "Wrong summation."
    );
    auto m3 = new T[][](rowsNumber1, columnsNumber1);
    foreach(i; 0 .. rowsNumber1) {
        foreach(j;  0 .. columnsNumber1) {
            m3[i][j] = m1[i][j] + m2[i][j];
        }
    }
    return m3;
}
unittest {
    int[][] mat1 = [[-3, 9], [7, 16]];
    int[][] mat2 = [[7, -4], [1, -15]];
    assert(summarize(mat1, mat2) == [[4, 5], [8, 1]]);
}


/*******************************************************************************
 * The function multiplies two matrices.
 */
auto multiply(T1, T2)(in T1[][] m1, in T2[][] m2) {
    auto rowsNumber1 = m1.length;
    auto columnsNumber1 = m1[0].length;
    auto rowsNumber2 = m2.length;
    auto columnsNumber2 = m2[0].length;
    enforce(columnsNumber1 == rowsNumber2, "Wrong multiplication.");
    auto m3 = new typeof(m1[0][0] + m2[0][0] + 0)[][]
        (rowsNumber1, columnsNumber2);
    foreach(i; 0 .. rowsNumber1) {
        foreach(j; 0 .. columnsNumber2) {
            m3[i][j] = 0;
            foreach(r; 0 .. columnsNumber1) {
                m3[i][j] += m1[i][r] * m2[r][j];
            }
        }
    }
    return m3;
}
unittest {
    int[][] mat1 = [[5, 1, 9]];
    int[][] mat2 = [[4], [7], [-1]];
    auto res = multiply(mat1, mat2);
    assert(res == [[18]]);
    int[][] arr1 = [[1, 2], [3, 5]];
    int[][] arr2 = [[4, 0], [8, 16]];
    res = multiply(arr1, arr2);
    assert(res == [[20, 32], [52, 80]]);
}


/*******************************************************************************
 * The function multiplies the matrix by a number.
 */
auto multiply(T, E)(in T[][] m, in E value)
if (is(typeof(m[0][0] != value) == bool)) {
    auto res = to!(typeof(m[0][0] + value + 0)[][])(m.dup);
    foreach(i; 0 .. m.length) {
        foreach(j; 0 .. m[0].length) {
            res[i][j] = m[i][j] * value;
        }
    }
    return res;
}
auto multiply(T, E)(in E value, in T[][] m)
if (is(typeof(m[0][0] != value) == bool)) {
    return multiply(m, value);
}


/*******************************************************************************
 * Matrix transpose.
 */
T[][] transpose(T)(in T[][] m) {
    auto res = to!(typeof(m[0][0] + 0)[][])(m);
    foreach(i; 0 .. m.length) {
        foreach(j; 0 .. m[0].length) {
            res[i][j] = m[j][i];
        }
    }
    return res;
}
unittest {
    auto m = [
        [1, 2, 3],
        [4, 5, 6],
        [7, 8, 9]
    ];
    auto expected = [
        [1, 4, 7],
        [2, 5, 8],
        [3, 6, 9]
    ];
    assert(transpose(m) == expected);
}

private {
    enum DimCount(T: U[], U) = is(T: U[], U) ? 1 + DimCount!U : 0;
    enum DimCount(T) = 0;
}
unittest {
    assert(DimCount!(int[]) == 1);
    assert(DimCount!(int[][]) == 2);
    assert(DimCount!(int[4][3]) == 2);
    assert(DimCount!(int[][2][]) == 3);
    assert(DimCount!(string) == 1);
}


/*******************************************************************************
 * Number of array dimensions.
 */
size_t ndim(T)(T matrix) {
    return DimCount!T;
}
///
unittest {
    int[] arr1d;
    int[][] arr2d;
    int[3][][1] arr3d;
    int[][][][][][][] arr7d;
    assert(arr1d.ndim == 1);
    assert(arr2d.ndim == 2);
    assert(arr3d.ndim == 3);
    assert(arr7d.ndim == 7);
}


/*******************************************************************************
 * Lengths of the corresponding array dimensions.
 */
size_t[] shape(A)(A arr) {
    import std.range : ElementType;
    static if (DimCount!A == 1) {
        return [arr.length];
    } else {
        bool arrIsEmpty = arr.length == 0;
        if (!arrIsEmpty) {
            return shape!(ElementType!A)(arr[0]) ~ [arr.length];
        } else {
            return shape!(ElementType!A)((ElementType!A).init) ~ [arr.length];
        }
    }
}
///
unittest {
    int[4][5][6] arr456;
    assert(arr456.shape == [4, 5, 6]);
    int[][][] arr000;
    assert(arr000.shape == [0, 0, 0]);
    int[3][][7] arr307;
    assert(arr307.shape == [3, 0, 7]);
    int[][][7] arr007;
    assert(arr007.shape == [0, 0, 7]);
    int[5][][] arr500;
    assert(arr500.shape == [5, 0, 0]);
    int[] dynArr = new int[](8);
    assert(dynArr.shape == [8]);
}


/*******************************************************************************
 * The template function returns a two-dimensional array with ones
 * on the main diagonal and zeros elsewhere.
 */
T[][] eye(T = double)(size_t matrixDimension) {
    T[][] arr = new T[][](matrixDimension, matrixDimension);
    foreach(i; 0 .. matrixDimension) {
        arr[i][0 .. $] = 0;
        arr[i][i] = 1;
    }
    return arr;
}