import sys
import time

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

# 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)]

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
    
maxdiff1 = 100000.0
iteration = 0

# Iterative computation
while maxdiff1 > MAXDIFF:
    maxdiff1 = -1.0
    for i in range(1, N + 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(1, N + 1):
        for j in range(1, N + 1):
            x[i][j] = y[i][j]

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

print(f"total time: {time.time() - start_time:.2f} seconds for {iteration} iterations")
    
# 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)