import sys
import time
import multiprocessing

"""
This program distributes the compute (decompose the domain)  based on myID and nprocs.
There are issues that are not resolved:
(1) boundary condition was not set correctly
(2) maxdiff1 is only local to each process
(3) the result data in distributed in different processes
"""

# Constants
N = 1024
nprocs = 1
if (len(sys.argv) > 1):
    nprocs = int(sys.argv[1])

if (len(sys.argv) > 2):
    N = int(sys.argv[2])

MAXDIFF = 0.5
iter = 0

# Initialize the two 2D arrays

x = [[0.0 for _ in range(N + 2)] for _ in range(N + 2)]
y = [[0.0 for _ in range(N + 2)] for _ in range(N + 2)]


#distribute computation based on myID and nprocs
def parallelComputeMyBlockOnly(myId):
    global x
    global y
    blockSize = N // nprocs

    print(f"process {myId} running ....")
    
    maxdiff1 = 100000.0
    iteration = 0

    # Iterative computation
    while maxdiff1 > MAXDIFF:
        maxdiff1 = -1.0

        for i in range(myId*blockSize + 1, N+1 if myId == nprocs -1 else (myId+1) * blockSize + 1):
            for j in range (1, N+1):
                y[i][j] = (
                    0.25 * x[i - 1][j] +
                    0.25 * x[i + 1][j] + 0.25 * x[i][j - 1] +
                    0.25 * x[i][j + 1]
                )
                diff = abs(x[i][j] - y[i][j])
                if diff > maxdiff1:
                    maxdiff1 = diff

        for i in range(myId*blockSize + 1, N+1 if myId == nprocs -1 else (myId+1) * blockSize + 1):
            for j in range(1, N + 1):
                x[i][j] = y[i][j]

        
            
        print(f"myID: {myId}, iteration = {iteration}, maxdiff1 = {maxdiff1:.3f}, MAXDIFF = {MAXDIFF:.3f}")
        iteration += 1

    
print(f"{N}x{N} mesh, maxdiff = {MAXDIFF:.3f}")


start_time = time.time()
# Boundary initialization

for i in range(1, N + 2):
    x[i][N + 1] = y[i][N + 1] = 0.1 * i
    x[N + 1][i] = y[N + 1][i] = 0.1 * i
    
#setup processes
processes = []
for i in range(nprocs):
    processes.append(multiprocessing.Process(target=parallelComputeMyBlockOnly, args=(i,)))

for i in range(nprocs):
    processes[i].start()

for i in range(nprocs):
    processes[i].join()

print(f"total time: {time.time() - start_time:.2f} seconds")
    
# write to file
try:
    with open("output.txt", "w") as f:
        # Write x[N//2][j] for j in 0..N+1
        for i in range(N+2): 
            for j in range(N + 2):
                f.write(f"{x[i][j]:.4f}\n")
except IOError as e:
    print("File write error:", e)