'''Routines for solving a linear system of equations.'''

import numpy as np
from numpy.typing import NDArray

def lu_factorization(
        mm: NDArray[np.float_],
        tolerance: float = 1e-6
) -> tuple[NDArray[np.float_], NDArray[np.float_], NDArray[np.float_]]:
    '''Computes the LUP factorization of a matrix

    Args:
        mm: Matrix to factorize
        tolerance: Greatest absolute value considered as 0

    Returns:
        ll: Lower triangular matrix
        uu: Upper triangular matrix
        pp: Permutation matrix
    '''
    nn = mm.shape[0]

    uu = mm.copy()
    ll = np.zeros((nn, nn), dtype=np.float_)
    pp = np.eye(nn, dtype=np.float_)

    i = 0
    j = 0

    while i < nn and j < nn:
        pivot_row = i + int(np.argmax(np.abs(uu[i:, j])))

        if np.abs(uu[pivot_row, j]) < tolerance:
            j = j + 1
            continue

        if i != pivot_row:
            # Swap rows
            uu[[i, pivot_row]] = uu[[pivot_row, i]]
            ll[[i, pivot_row]] = ll[[pivot_row, i]]
            pp[[i, pivot_row]] = pp[[pivot_row, i]]

        for k in range(i + 1, nn):
            q = uu[k, j] / uu[i, j]
            uu[k, :] = uu[k, :] - q * uu[i, :]
            ll[k, j] = q

        i = i + 1
        j = j + 1

    return ll + np.eye(nn), uu, pp

def forward_substitution(
        aa: NDArray[np.float_],
        bb: NDArray[np.float_]
) -> NDArray[np.float_]:
    '''Solves (T x = b), where T is a lower triangular matrix

    Args:
        aa: lower triangular matrix
        bb: vector

    Returns:
        x: solution of the system of equations
    '''
    nn = aa.shape[0]
    xx = np.zeros((nn, 1), dtype=np.float_)
    for i in range(nn):
        xx[i] = (bb[i] - np.dot(aa[i, 0:i], xx[0:i])) / aa[i, i]
    return xx

def back_substitution(
        aa: NDArray[np.float_],
        bb: NDArray[np.float_]
) -> NDArray[np.float_]:
    '''Solves (T x = b), where T is a upper triangular matrix

    Args:
        aa: upper triangular matrix
        bb: vector

    Returns:
        x: solution of the system of equations
    '''
    nn = aa.shape[0]
    xx = np.zeros((nn, 1), dtype=np.float_)
    for i in range(nn - 1, -1, -1):
        xx[i] = (bb[i] - np.dot(aa[i, i:], xx[i:nn])) / aa[i, i]
    return xx

def gaussian_eliminate(
        aa: NDArray[np.float_],
        bb: NDArray[np.float_],
        tolerance: float = 1e-6
) -> NDArray[np.float_] | None:
    '''Solves a linear system of equations (A x = b) by LUP factorization.

    Args:
        aa: upper triangular matrix
        bb: vector
        tolerance: Greatest absolute value considered as 0

    Returns:
        x: solution of the system of equations
    '''
    ll, uu, pp = lu_factorization(aa)

    nn = uu.shape[0]
    # Check if rank of matrix is lower than nn
    if np.abs(uu[nn - 1, nn - 1]) < tolerance:
        return None

    # L y = P @ b
    y = forward_substitution(ll, pp @ bb)
    # U @ x = y
    x = back_substitution(uu, y)
    return x[:, 0]