import numpy as np
from mpi4py import MPI

from matrix_mpi import MatrixMPI
from vector_mpi import VectorMPI

np.random.seed(0)
###############
# Setting up MPI
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
size = comm.Get_size()
for i0 in range(size):
    if i0 == rank:
        print(f"Hello from rank {i0} of size {size}!")
    comm.Barrier()
# Testing the vector class
comm.Barrier()
if rank == 0:
    print("\n\nTesting the vector class --------------(output only from rank 0)---------------------\n\n")
###############

### 1a Initialization
if rank == 0:
    print("Start 1a initialization")
# from list
n_x1 = 10
x1_list = [i0 for i0 in
           range(n_x1)]  # must work for 4 ranks, too, so even if entries are not evenly distributable between ranks
vec_x1 = VectorMPI(x1_list)
# from numpy array
n_x2 = 10000
x2_list = np.random.uniform(0, 1, n_x2)
vec_x2 = VectorMPI(x2_list)
if rank == 0:
    print("End 1a\n")
comm.Barrier()

### 1b __str__ function, string representation
if rank == 0:
    print("Start 1b string representation")
vec_x1_str = str(vec_x1)
vec_x2_str = str(vec_x2)
if rank == 0:
    print(f"vec_x1 values: {vec_x1_str}")
    print(f"vec_x2 values: {vec_x2_str}")
    print("End 1b\n")
comm.Barrier()

### 1c shape and transpose
if rank == 0:
    print("Start 1c shape and transpose")
# shape property of the vector
vec_x1_shape = vec_x1.shape()
vec_x2_T_shape = vec_x2.T().shape()
vec_x2_T_T_shape = vec_x2.T().T().shape()
if rank == 0:
    print(f"vec_x1 has shape {vec_x1_shape} | must be ({n_x1},1)")  # Hint: numpy.reshape
    # transposition (property) of the vector, only "cosmetic" change
    print(f"vec_x2.T() has shape {vec_x2_T_shape} | must be (1,{n_x2})")
    print(f"vec_x2.T().T() has shape {vec_x2_T_T_shape} | must be ({n_x2},1)")
    print("End 1c\n")
comm.Barrier()

### 1d addition and subtraction    
if rank == 0:
    print("Start 1d addition and subtraction    ")
# initialization
n_a = 10
n_b = 10
a_list = [i0 for i0 in range(n_a)]
b_list = [1 / (i0 + 1) for i0 in range(n_b)]
a = VectorMPI(a_list)
b = VectorMPI(b_list)
# computation
x = a + b
y = a - b
numpy_x_compare_norm = (x - VectorMPI(np.array(a_list) + np.array(b_list))).norm()
numpy_y_compare_norm = (y - VectorMPI(np.array(a_list) - np.array(b_list))).norm()
if rank == 0:
    print(f"norm(a + b - numpy) = {numpy_x_compare_norm} | must be < 1e-8")
    print(f"norm(a - b - numpy) = {numpy_x_compare_norm} | must be < 1e-8")
try:
    x_false = a + b.T()
except ValueError as e:
    if rank == 0:
        print("The correct result is this error: " + str(e))
else:
    if rank == 0:
        print("ERROR: It is required to raise a system error, e. g., ValueError, since dimensions mismatch!")
try:
    x_false = a - b.T()
except ValueError as e:
    if rank == 0:
        print("The correct result is this error: " + str(e))
else:
    if rank == 0:
        print("ERROR:It is required to raise a system error, e. g., ValueError, since dimensions mismatch!")
if rank == 0:
    print("End 1d\n")
comm.Barrier()

### 1e multiplication
if rank == 0:
    print("Start 1e multiplication")
# initialization
n_a = 10
n_b = 10
a_list = [i0 for i0 in range(n_a)]
b_list = [1 / (i0 + 1) for i0 in range(n_b)]
a = VectorMPI(a_list)
b = VectorMPI(b_list)
# computation with vectors
x = a.T() * b  # scalar
y = a * b  # vector
numpy_x_compare_norm = np.linalg.norm(x - np.sum(np.array(a_list) * np.array(b_list)))
numpy_y_compare_norm = (y - VectorMPI(np.array(a_list) * np.array(b_list))).norm()
# numpy_z_compare_norm = (z - MatrixMPI(comm,np.outer(np.array(a_list),np.array(b_list)))).norm()
if rank == 0:
    print(f"norm(a.T() * b - numpy) = {numpy_x_compare_norm} | must be < 1e-8")
    print(f"norm(a * b - numpy) = {numpy_y_compare_norm} | must be < 1e-8")
# print(f"norm(b * a.T() - numpy) = \n{numpy_z_compare_norm} | must be 1e-8")
# computation with scalars
x = a * 5
y = 0.1 * b.T()
numpy_x_compare_norm = (x - VectorMPI(np.array(a_list) * 5)).norm()
numpy_y_compare_norm = (y - VectorMPI(np.array(b_list) * 0.1).T()).norm()
if rank == 0:
    print(f"norm(a * 5 - numpy) = {numpy_x_compare_norm} | must be < 1e-8")
    print(f"norm(0.1 * b.T() - numpy) = {numpy_y_compare_norm} | must be < 1e-8")
    print("End 1e\n")
comm.Barrier()

### 1f division
if rank == 0:
    print("Start 1f division")
# initialization
n_a = 10
n_b = 10
a_list = [i0 for i0 in range(n_a)]
b_list = [1 / (i0 + 1) for i0 in range(n_b)]
a = VectorMPI(a_list)
b = VectorMPI(b_list)
# computation with vectors
x = a / b
y = a / 5
numpy_x_compare_norm = (x - VectorMPI(np.array(a_list) / np.array(b_list))).norm()
numpy_y_compare_norm = (y - VectorMPI(np.array(a_list) / 5)).norm()
if rank == 0:
    print(f"norm(a / b - numpy) = {numpy_x_compare_norm} | must be < 1e-8")
    print(f"norm(a / 5 - numpy) = {numpy_y_compare_norm} | must be < 1e-8")
    print("End 1f\n")
comm.Barrier()

### 1g norm
if rank == 0:
    print("Start 1g norm")
# initialization
a_list = [1 / (i0 + 1) for i0 in range(10)]
a = VectorMPI(a_list)
# computation
a_norm = a.norm()
a_normalized = a.normalize()
a_normalized_str = str(a_normalized)
numpy_comparison_norm = (a_normalized - VectorMPI(np.array(a_list) / np.linalg.norm(a_list))).norm()
if rank == 0:
    print(f"a_norm = {a_norm} | must be {np.linalg.norm(a_list)}")
    print(f"norm(a_normalize-np.a_normalize) = {numpy_comparison_norm} | must be < 1e-8")
    print("End 1g\n")
comm.Barrier()

### 1h negation
if rank == 0:
    print("Start 1h negation")
# initialization
a_list = [1 / (i0 + 1) for i0 in range(10)]
a = VectorMPI(a_list)
# computation
x = -a
x_str = str(x)
if rank == 0:
    print(f"-a = {x_str} | must be {-np.array(a_list)}")
    print("End 1h\n")
comm.Barrier()

### 1i manipulation
if rank == 0:
    print("Start 1i manipulation")
# initialization
n_a = 10
a_list = [1 / (i0 + 1) for i0 in range(n_a)]
a = VectorMPI(a_list)
a_idx = [1, 2, 9, 7, 8]
a_values = a[a_idx]
if rank == 0:
    print(
        f"a[{str(a_idx)}] = {str(a_values)} | must be {np.array(a_list).reshape(n_a, 1)[np.array(a_idx)].reshape(len(a_idx), )}")
a[a_idx] = [-1, -1, -1, -1, -1]
a_str = str(a)
np_a = np.array(a_list)
np_a[a_idx] = [-1, -1, -1, -1, -1]
if rank == 0:
    print(f"a = {a_str} | must be {np_a}")
    print("End 1i\n")
comm.Barrier()

###############
# Testing the matrix class
comm.Barrier()
if rank == 0:
    print("\n\nTesting the matrix class -----------------------------------\n\n")
###############

### 2a Initialization
if rank == 0:
    print("Start 2a initialization")
n_a1 = 10
n_a2 = 5
a_list = np.array([[(i0 + 1) * (i1 + 1) for i0 in range(n_a1)] for i1 in range(n_a2)])
A = MatrixMPI(a_list)
B = MatrixMPI(structure="tridiagonal", data=[-1, 2, -1], n=12)
c_list = [i0 for i0 in range(n_a1 * n_a1)]
C = MatrixMPI(c_list, shape=(n_a1, n_a1))
D = MatrixMPI(structure="tridiagonal", data=[-1, 2, -1], n=50)
# D = MatrixMPI(model="sheet1ex1", n=50)
if rank == 0:
    print("End 2a\n")
comm.Barrier()

### 2b __str__ function, string representation
if rank == 0:
    print("Start 2b string representation")
# print(B.__str__(full = True))
A_str = str(A)
B_str = str(B)
C_str = str(C)
D_str = str(D)
if rank == 0:
    print(f"Matrix A (numbers):\n{A_str}")
    print(f"Matrix B (tridiagonal):\n{B_str}")
    print(f"Matrix C (list of numbers):\n{C_str}")
    print(f"Matrix D (sheet1ex1):\n{D_str}")
    print("End 2b\n")
comm.Barrier()

### 2c shape and transpose
if rank == 0:
    print("Start 2c shape and transpose")
# Initialization
A_shape = A.shape()
A_T_shape = A.T().shape()
A_T_T_shape = A.T().T().shape()
if rank == 0:
    print(f"A has shape {A_shape} | must be {np.array(a_list).shape}")
    print(f"A.T() has shape {A_T_shape} | must be {np.array(a_list).T.shape}")
    print(f"A.T().T() has shape {A_T_T_shape} | must be {np.array(a_list).T.T.shape}")
    print("End 2c\n")
comm.Barrier()

# ### DEBUG norms ###
# mat = np.array([[(i0+1)*(i1+1) for i0 in range(10)] for i1 in range(5)])
# local_max_each_row = np.amax(np.abs(mat),axis=1)
# if rank == 0:
#     print(f"Matrix\n{str(mat)}\nlocal max each row {str(local_max_each_row)}")
# # frobenius norm
# A = MatrixMPI(comm,mat)
# a_norm_fro = A.norm("frobenius")
# if rank == 0:
#     print(f"A.norm('frobenius') = {a_norm_fro} | must be {np.linalg.norm(np.array([[(i0+1)*(i1+1) for i0 in range(10)] for i1 in range(5)]),'fro')}")
# # row sum norm
# a_norm_row = A.norm("row sum")
# if rank == 0:
#     print(f"A.norm('row sum') = {a_norm_row} | must be {np.max(local_max_each_row)}")
# # col sum norm


### 2d addition and subtraction    
if rank == 0:
    print("Start 2d addition and subtraction    ")
# Initialization 
n = 10
A = MatrixMPI(structure="diagonal", data=[3], offset=0, n=n)
A21 = MatrixMPI(structure="diagonal", data=[-1], offset=-1, n=n)
A12 = MatrixMPI(structure="diagonal", data=[-1], offset=+1, n=n)
B = MatrixMPI(structure="diagonal", data=[1], offset=0, n=n)
# computation
C = A + A21 + A12 - B
D = C - MatrixMPI(structure='tridiagonal', data=[-1, 2, -1], n=n)
d_norm = D.norm()
A_str = str(5 + A - 3)
if rank == 0:
    print(f"norm(A + A21 + A12 - B - tridiag) = {d_norm} | must be < 1e-8")
    print(f"5+A-3 = \n{A_str}")
    print("End 2d\n")
comm.Barrier()

### 2e multiplication
if rank == 0:
    print("Start 2e multiplication")
# initialization
n = 10
a_mat = np.array([[(i0 + 1) / (i1 + 1) for i1 in range(8)] for i0 in range(n)])
b_mat = np.array([[(i0 + 1) / (i1 + 1) for i1 in range(n)] for i0 in range(8)])
c_mat = np.array([[(i0 + 1) / (i1 + 1) for i1 in range(n)] for i0 in range(n)])
d_mat = np.array([[(i0 + 1) / (i1 + 1) for i1 in range(17)] for i0 in range(n)])
A = MatrixMPI(a_mat)
B = MatrixMPI(b_mat)
C = MatrixMPI(c_mat)
D = MatrixMPI(d_mat)
x_vec = np.array([i0 + 1 for i0 in range(n)])
y_vec = np.array([n * (i0 + 1) for i0 in range(n)])
x = VectorMPI(x_vec)
y = VectorMPI(y_vec)
# computation matrix scalar
norm5 = (5 * A - MatrixMPI(5 * np.array(a_mat))).norm()
# computation matrix vector
norm6 = (C * x - VectorMPI(np.array(c_mat) @ np.array(x_vec))).norm()
norm7 = (D.T() * x - VectorMPI(np.array(d_mat).T @ np.array(x_vec))).norm()
y_shape = (D.T() * x).shape()
if rank == 0:
    print(f"Norm of (5*A - np.(5*A)) is {norm5} | must be < 1e-8")
    print(f"Norm of (C*x - np.(C*x)) is {norm6} | must be < 1e-8")
    print(f"Norm of (D.T*x - np.(D.T*x)) is {norm7} | must be < 1e-8 | shape(D.T*x) is {y_shape} | must be (17,1)")
# computation matrix multiplication
A_str = str(A)
B_str = str(B)
if rank == 0:
    print(f"DEBUG: A\n{A_str}\nDEBUG: B\n{B_str}")
norm1 = (A * B - MatrixMPI(np.array(a_mat) @ np.array(b_mat))).norm()
norm2 = (A.T() * A - MatrixMPI(np.array(a_mat).T @ np.array(a_mat))).norm()
norm3 = (A * A.T() - MatrixMPI(np.array(a_mat) @ np.array(a_mat).T)).norm()
norm4 = (B.T() * A.T() - MatrixMPI(np.array(b_mat).T @ np.array(a_mat).T)).norm()
if rank == 0:
    print(f"Norm of (A*B - np.(A*B)) is {norm1} | must be < 1e-8")
    print(f"Norm of (A.T()*A - np(A.T()*A)) is {norm2} | must be < 1e-8")
    print(f"Norm of (A*A.T() - np(A*A.T())) is {norm3} | must be < 1e-8")
    print(f"Norm of (B.T()*A.T() - np.(B.T()*A.T())) is {norm4} | must be < 1e-8")
    print("End 2e\n")
comm.Barrier()

### 2f division
if rank == 0:
    print("Start 2f division")
# initialization
A = MatrixMPI(a_mat)
# computation
print(f"Norm of (A/5 - np.(A/5)) is {(A / 5 - MatrixMPI(np.array(a_mat) / 5)).norm()} | must be < 1e-8")
if rank == 0:
    print("End 2f\n")
comm.Barrier()

### 2g norm
if rank == 0:
    print("Start 2g norm")
A = MatrixMPI(structure="tridiagonal", n=50, data=[-1, 2, -1])
print(f"Frobenius norm of tridiagonal matrix: {A.norm('frobenius')} | must be 17.263")
print(f"Row sum norm of tridiagonal matrix: {A.norm('row sum')} | must be 4")
print(f"Col sum norm of tridiagonal matrix: {A.norm('col sum')} | must be 4")
if rank == 0:
    print("End 2g\n")
comm.Barrier()

### 2h negation
if rank == 0:
    print("Start 2h negation")
A = MatrixMPI(structure="tridiagonal", n=50, data=[-1, 2, 1])
print(f"Norm of (A + (-A)) is {(A + (-A)).norm('frobenius')} | must be < 1e-8")
if rank == 0:
    print("End 2h\n")
comm.Barrier()

### 2i manipulation
if rank == 0:
    print("Start 2i manipulation")
A = MatrixMPI(structure="tridiagonal", n=10, data=[-1, 2, 1])
A[1, 1] = 4
A[[1, 2, 3], 2] = [-5, -10, 100]
print(str(A))
if rank == 0:
    print("End 2i\n")
comm.Barrier()