import numpy as np
from mpi4py import MPI

from matrix import Matrix
from vector import Vector

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 = Vector(comm, x1_list)
vec_x1 = Vector(x1_list)
# from numpy array
n_x2 = 10000
x2_list = np.random.uniform(0, 1, n_x2)
#vec_x2 = Vector(comm, x2_list)
vec_x2 = Vector(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 substraction
if rank == 0:
    print("Start 1d addition and substraction")
# intitialization
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 = Vector(comm, a_list)
a = Vector(a_list)
#b = Vector(comm, b_list)
b = Vector(b_list)
# computation
x = a + b
y = a - b
#numpy_x_compare_norm = (x - Vector(comm, np.array(a_list) + np.array(b_list))).norm()
numpy_x_compare_norm = (x - Vector(np.array(a_list) + np.array(b_list))).norm()
#numpy_y_compare_norm = (y - Vector(comm, np.array(a_list) - np.array(b_list))).norm()
numpy_y_compare_norm = (y - Vector(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")
# intitialization
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 = Vector(comm, a_list)
a = Vector(a_list)
#b = Vector(comm, b_list)
b = Vector(b_list)
# computation with vectors
x = a.T() * b  # scalar
y = a * b  # vector
# z = b * a.T() # matrix
numpy_x_compare_norm = np.linalg.norm(x - np.sum(np.array(a_list) * np.array(b_list)))
#numpy_y_compare_norm = (y - Vector(comm, np.array(a_list) * np.array(b_list))).norm()
numpy_y_compare_norm = (y - Vector(np.array(a_list) * np.array(b_list))).norm()
# numpy_z_compare_norm = (z - Matrix(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 - Vector(comm, np.array(a_list) * 5)).norm()
numpy_x_compare_norm = (x - Vector(np.array(a_list) * 5)).norm()
#numpy_y_compare_norm = (y - Vector(comm, np.array(b_list) * 0.1, transpose=True)).norm()
numpy_y_compare_norm = (y - Vector(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 divison
if rank == 0:
    print("Start 1f divison")
# intitialization
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 = Vector(comm, a_list)
a = Vector(a_list)
#b = Vector(comm, b_list)
b = Vector(b_list)
# computation with vectors
x = a / b
y = a / 5
#numpy_x_compare_norm = (x - Vector(comm, np.array(a_list) / np.array(b_list))).norm()
numpy_x_compare_norm = (x - Vector(np.array(a_list) / np.array(b_list))).norm()
#numpy_y_compare_norm = (y - Vector(comm, np.array(a_list) / 5)).norm()
numpy_y_compare_norm = (y - Vector(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")
# intitialization
a_list = [1 / (i0 + 1) for i0 in range(10)]
#a = Vector(comm, a_list)
a = Vector(a_list)
# computation
a_norm = a.norm()
a_normalized = a.normalize()
a_normalized_str = str(a_normalized)
#numpy_comparison_norm = (a_normalized - Vector(comm, np.array(a_list) / np.linalg.norm(a_list))).norm()
numpy_comparison_norm = (a_normalized - Vector(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")
# intitialization
a_list = [1 / (i0 + 1) for i0 in range(10)]
#a = Vector(comm, a_list)
a = Vector(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")
# intitialization
n_a = 10
a_list = [1 / (i0 + 1) for i0 in range(n_a)]
#a = Vector(comm, a_list)
a = Vector(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 = Matrix(comm, a_list)
B = Matrix(comm, structure="tridiagonal", size=12)
c_list = [i0 for i0 in range(n_a1 * n_a1)]
C = Matrix(comm, c_list, shape=(n_a1, n_a1))
D = Matrix(comm, model="sheet1ex1", size=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 = Matrix(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 substraction
if rank == 0:
    print("Start 2d addition and substraction")
# Initialization 
n = 10
A = Matrix(comm, structure="diagonal", diag_values=[3], offset=0, size=n)
A21 = Matrix(comm, structure="diagonal", diag_values=[-1], offset=-1, size=n)
A12 = Matrix(comm, structure="diagonal", diag_values=[-1], offset=+1, size=n)
B = Matrix(comm, structure="diagonal", diag_values=[1], offset=0, size=n)
# computation
C = A + A21 + A12 - B
D = C - Matrix(comm, structure='tridiagonal', diag_values=[-1, 2, -1], size=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 = Matrix(comm, a_mat)
B = Matrix(comm, b_mat)
C = Matrix(comm, c_mat)
D = Matrix(comm, 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 = Vector(comm, x_vec)
y = Vector(comm, y_vec)
# computation matrix scalar
norm5 = (5 * A - Matrix(comm, 5 * np.array(a_mat))).norm()
# computation matrix vector
norm6 = (C * x - Vector(comm, np.array(c_mat) @ np.array(x_vec))).norm()
norm7 = (D.T() * x - Vector(comm, 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 - Matrix(comm, np.array(a_mat) @ np.array(b_mat))).norm()
norm2 = (A.T() * A - Matrix(np.array(a_mat).T @ np.array(a_mat))).norm()
norm3 = (A * A.T() - Matrix(np.array(a_mat) @ np.array(a_mat).T)).norm()
norm4 = (B.T() * A.T() - Matrix(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 divison 
if rank == 0:
	print("Start 2f divison")
# initialization
A = Matrix(a_mat)
# computation
print(f"Norm of (A/5 - np.(A/5)) is {(A/5 - Matrix(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 = Matrix(structure="tridiagonal",given_size=50,given_values=[-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 2")
print(f"Col sum norm of tridiagonal matrix: {A.norm('col sum')} | must be 2")
if rank == 0:
	print("End 2g\n")
comm.Barrier()

### 2h negation
if rank == 0:
	print("Start 2h negation")
A = Matrix(structure="tridiagonal", given_size=50, given_values=[-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 = Matrix(structure="tridiagonal", given_size=10, given_values=[-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()

'''
"""