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 = 10000
) -> 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_)

    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

    if algorithm.norm(alpha) >= 1.0:
        return None

    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

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

    return X


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
