# Jacobi-Iteration Poisson
# n = Dimension der Matrix A (wobei hier n durch Anzahl der Prozesseteilbar sein soll
# mpiexec -np 100 python JacobiParallel.py 1000000
import numpy as np
import sys
import math
from mpi4py import MPI

n = int(sys.argv[1])

comm =  MPI.COMM_WORLD
size = comm.Get_size()
rank = comm.Get_rank()

# lokale Groesse der Arrays (zu bearbeitender Block)
local_n = n/size

if comm.rank < n-size*local_n:
local_n = local_n + 1

# Initialisierung der Arrays
local_A = np.zeros((n+2,local_n+2))
for i in range(0,n+1):
for j in range(0,local_n+1):
    if i == j:
        local_A[i,j] = 4
    elif i==0 and j== 0:
        local_A[i,j] = 0
    elif i-1 == j and i != 0:
        local_A[i,j] = -1
    elif i+1 == j:
        local_A[i,j] = -1
    else:
        local_A[i,j] = 0
local_B = np.zeros((n,local_n))
print ("Fuer den Prozess", rank, "gilt:" , local_A)


iterations = 0
eps = 1e-8 #Konvergenzgrenzwert
error = 2*eps


#Main - Loop
while iterations < 1e4 and error > eps:
   for j in range(1, local_n, 1):
      for i in range(1, n, 1):
          local_B[i,j] = 0.25*(local_A[i+1,j]+local_A[i,j+1]+local_A[i-1,j]+local_A[i,j-1])
          error = error+ abs(local_B[i,j]-local_A[i,j])

 for j in range(1,local_n,1):
     for i in range(1, n, 1):
         local_A[i,j] = local_B[i,j]

if comm.rank % 2 == 1:
    comm.send(local_B[:,1], rank-1)
    comm.recv(local_A[:,0], rank-1)
    if comm.rank < size - 1:
        comm.send(local_B[:,local_n], rank+1)
        comm.recv(local_A[:,local_n+1], rank+1)
else:
    if comm.rank > 0:
        comm.recv(local_A[:,0], rank-1)
        comm.send(local_B[:,1], rank-1)
    if comm.rank < size-1:
        comm.recv(local_A[:,local_n+1], rank+1)
        comm.send(local_B[:,local_n], rank+1)
print (local_B)

iterations = iterations + 1
error = error/n**2


comm.Barrier()
if rank == 0:
A_Gathered = np.zeros((n+2,n+2))
else:
A_Gathered = None
comm.Gather(local_A, A_Gathered, root = 0)
print (" Das ist A_Gathered",A_Gathered)