import numpy as np

def my_lu(A):
    # The upper triangular matrix U is saved in the upper part of the matrix M (including the diagonal)
    # The lower triangular matrix L is saved in the lower part of the matrix M (not including the diagonal)
    # Do NOT use `scipy.linalg.lu`!
    # You should not use pivoting

    mA = np.copy(A)
    M = np.zeros((len(A), len(A)))
    for i in range(len(A) - 1):
        M[i, i:] = mA[i, i:]
        M[i + 1:, i] = mA[i + 1:, i] / mA[i, i]
        prod = np.einsum('i,j->ij', M[i + 1:, i], M[i, i + 1:])
        mA[i + 1:, i + 1:] -= prod
        M[i + 1:, i + 1:] = mA[i + 1:, i + 1:]
    return M


def my_triangular_solve(M, b):
    # A = LU (L and U are stored in M)
    # A x = b (given A and b, find x)
    # M is a 2D numpy array
    # The upper triangular matrix U is stored in the upper part of the matrix M (including the diagonal)
    # The lower triangular matrix L is stored in the lower part of the matrix M (not including the diagonal)
    # b is a 1D numpy array
    # x is a 1D numpy array
    # Do not use `scipy.linalg.solve_triangular`

    y = np.zeros(len(M))
    for i in range(len(y)):
        currSum = sum([M[i, j] * y[j] for j in range(i)])
        y[i] = b[i] - currSum

    x = np.zeros(len(M))
    for i in range(len(x) - 1, -1, -1):
        currSum = sum([M[i, j] * x[j] for j in range(len(x) - 1, i - 1, -1)])
        x[i] = (y[i] - currSum) / M[i, i]

    return x


def fea_solve(Kpp, Kpf, Kfp, Kff, xp, Ff):
    # Use my_lu and my_triangular_solve
    M = my_lu(Kff)
    xf = my_triangular_solve(M, Ff - np.matmul(Kfp, xp))
    Fp = np.matmul(Kpp, xp) + np.matmul(Kpf, xf)
    return xf, Fp