def obter_numero(mensagem):
    while True:
        try:
            return int(input(mensagem))
        except ValueError:
            print('\033[1;31mValor INVÁLIDO! Digite apenas números!\033[m')


def verifica_matriz_quadrada(num_linhas, num_colunas):
    if num_linhas != num_colunas:
        print(
            '\033[1;31mErro: Esta implementação suporta apenas matrizes quadradas (mesmo número de linhas e colunas). '
            'Reiniciando o programa...\033[m')
        main()  # Reinicia o programa


def criar_matriz(num_linha, num_coluna):
    matriz = []
    for linha_numero in range(num_linha):
        linha = []
        for coluna_numero in range(num_coluna):
            while True:
                try:
                    valor = int(
                        input(f'\033[1;35mPor Favor, digite o {coluna_numero + 1}º elemento da {linha_numero + 1}º '
                              f'linha: \033[m'))
                    linha.append(valor)
                    break
                except ValueError:
                    print('\033[1;31mValor INVÁLIDO! Digite apenas números!\033[m')
        matriz.append(linha)
    return matriz


def calcular_soma_diagonal_principal(matriz):
    diagonal_principal = []
    soma = 0
    for i in range(len(matriz)):
        diagonal_principal.append(matriz[i][i])
        soma += matriz[i][i]
    print(f'\033[1;31mOs elementos da diagonal principal é {diagonal_principal}.\033[m')
    return soma


def mostrar_matriz(matriz):
    print('\n\033[1;36mMatriz gerada:\033[m')
    for linha in matriz:
        print(f'\033[1;36m{linha}\033[m')


def main():
    print('\033[1;41m*** Inicializando o programa ***\033[m\n')
    numero_linhas = obter_numero('\033[1;34mDigite o número de linhas da matriz: \033[m')
    numero_colunas = obter_numero('\033[1;34mDigite o número de colunas da matriz: \033[m')

    verifica_matriz_quadrada(numero_linhas, numero_colunas)

    matriz = criar_matriz(numero_linhas, numero_colunas)
    print('\n\033[1;34mMatriz inserida:\033[m')
    mostrar_matriz(matriz)

    soma_diagonal = calcular_soma_diagonal_principal(matriz)
    print(f'\n\033[1;34mA soma dos elementos da diagonal principal da matriz é: {soma_diagonal}\n\033[m')


if __name__ == "__main__":
    while True:
        main()
        continuar = input('\033[1;35mDeseja calcular para outra matriz? (s/n): \033[m').lower()
        if continuar != 's':
            print('\033[1;33mFinalizando o programa!!!!\033[m')
            break