from typing import List, Tuple
import algorithm
import numpy as np


def iterative_method(
    A: np.matrix, B: np.matrix, eps: float = 0.0001, max_iterations: int = 100_000
) -> np.matrix:
    if A.shape[0] != A.shape[1] or A.shape[0] != B.shape[0] or B.shape[1] != 1:
        return None

    n = A.shape[0]

    alpha = np.copy(A)
    beta = np.copy(B)

    swaps = []
    for i in range(n):
        if not _swap_rows(alpha, beta):
            return None

        success, swaps_ = _swap_columns(alpha)
        if not success:
            return None

        swaps.extend(swaps_)

    _make_diagonal_dominant(alpha, beta)

    if algorithm.det(alpha) == 0:
        return None

    for i in range(n):
        beta[i, 0] = beta[i, 0] / alpha[i, i]

    for i in range(n):
        for j in range(n):
            if i == j:
                continue

            alpha[i, j] = -(alpha[i, j] / alpha[i, i])

        alpha[i, i] = 0

    i = 0
    X = np.copy(beta)
    while i < max_iterations:
        i += 1

        X_old = np.copy(X)
        X = np.add(np.matmul(alpha, X), beta)

        if np.sum(np.abs(X - X_old)) < eps:
            break

    print(f"iterations: {i}")

    for i, j in swaps:
        X[[i, j], :] = X[[j, i], :]

    return X


# WTF EVEN IS THIS
# EXCEL MOMENT: https://elib.bsu.by/bitstream/123456789/52531/1/42-47.pdf
def _make_diagonal_dominant(a: np.matrix, b: np.matrix):
    n = a.shape[0]
    a_inv = np.linalg.inv(a.T)

    do_changes = [0] * n
    for i in range(n):
        if np.sum(np.abs(a[i, :])) > 2 * np.abs(a[i, i]):
            do_changes[i] = 1

    helper = np.ones((n, n))
    for i in range(n):
        helper[i, i] = 2 * n

    koef = np.zeros((n, n))
    for i in range(n):
        for j in range(n):
            if not do_changes[j]:
                continue
            koef[i, j] = np.sum(a_inv[i, :] * helper[:, j]) * b[i, 0]

    koef_sum = np.sum(koef, axis=0)

    for i in range(n):
        if not do_changes[i]:
            continue
        a[i, :] = np.ones(n)
        a[i, i] = 2 * n
        b[i, 0] = koef_sum[i]


def _swap_rows(a: np.matrix, b: np.matrix) -> bool:
    n = a.shape[0]
    for i in range(n):
        j = algorithm.swap_max_row(a, i)

        if a[i, i] == 0:
            return False

        b[[i, j], :] = b[[j, i], :]

    return True


def _swap_columns(a: np.matrix) -> Tuple[bool, List[Tuple[int, int]] | None]:
    n = a.shape[0]
    swaps = []
    for i in range(n):
        j = algorithm.swap_max_column(a, i)

        if a[i, i] == 0:
            return False, None

        if i != j:
            swaps.append((i, j))

    return True, swaps
