import numpy as np

def solve_heat_equation():
    # Parameters
    nu = 0.05
    x_start, x_end = 0, 2
    y_start, y_end = 0, 2
    t_start, t_end = 0, 0.3777
    nx = 200
    ny = 200
    nt = 2000
    
    dx = (x_end - x_start) / (nx - 1)
    dy = (y_end - y_start) / (ny - 1)
    dt = (t_end - t_start) / (nt - 1)

    # Initialize u
    u = np.ones((nx, ny))
    for i in range(nx):
        for j in range(ny):
            x = x_start + i * dx
            y = y_start + j * dy
            if 0.5 <= x <= 1 and 0.5 <= y <= 1:
                u[i, j] = 2

    # Boundary conditions
    u[0, :] = 1
    u[-1, :] = 1
    u[:, 0] = 1
    u[:, -1] = 1

    # Time loop
    for n in range(nt):
        u_new = u.copy()
        for i in range(1, nx - 1):
            for j in range(1, ny - 1):
                u_new[i, j] = u[i, j] + nu * dt * (
                    (u[i+1, j] - 2*u[i, j] + u[i-1, j]) / dx**2 +
                    (u[i, j+1] - 2*u[i, j] + u[i, j-1]) / dy**2
                )
        u = u_new

        # Boundary conditions (enforce at each time step)
        u[0, :] = 1
        u[-1, :] = 1
        u[:, 0] = 1
        u[:, -1] = 1
    
    # Save the final solution
    np.save('/opt/CFD-Benchmark/PDE_Benchmark/results/prediction/gemini/prompts/u_2D_Diffusion.npy', u)

solve_heat_equation()