import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from math import sin, cos, pi

# The heat equation is a partial differential equation (PDE) that describes the distribution of heat (or variation in temperature) in a given region over time. 
# It is an example of a diffusion equation and is given by:

# ∂u/∂t = α * (∂²u/∂x² + ∂²u/∂y²)

# where:
# u(x, y, t) is the temperature distribution as a function of spatial coordinates (x, y) and time t,
# α is the thermal diffusivity of the material,
# ∂u/∂t is the partial derivative of u with respect to time,
# ∂²u/∂x² and ∂²u/∂y² are the second partial derivatives of u with respect to spatial coordinates x and y, respectively.

# In this code, we solve the 2D heat equation using a finite difference method with explicit time stepping.

# Set up the parameters
Lx, Ly = 1.0, 1.0  # Length of the domain in x and y
Nx, Ny = 100, 100  # Number of spatial points in x and y
T = 0.2           # Final time
dt = 0.001        # Time step
Nt = int(T / dt)   # Number of time steps
alpha = 0.01       # Diffusion coefficient

x = np.linspace(0, Lx, Nx)
y = np.linspace(0, Ly, Ny)
dx = Lx / (Nx - 1)
dy = Ly / (Ny - 1)

# Initialize the temperature profile: initial condition (Gaussian)
u0 = np.zeros((Nx, Ny))
# Define the first square
x1_start, x1_end = int(0.2 * Nx), int(0.4 * Nx)
y1_start, y1_end = int(0.2 * Ny), int(0.4 * Ny)
# Define the second square as a "thick border" of a square
x2_start, x2_end = int(0.6 * Nx), int(0.8 * Nx)
y2_start, y2_end = int(0.6 * Ny), int(0.8 * Ny)

# Set the temperature in the "thick border" of the second square
border_thickness = 30  # Define the thickness of the border

# Top border
u0[x2_start:x2_start+border_thickness, y2_start:y2_end] = 1.0
# Bottom border
u0[x2_end-border_thickness:x2_end, y2_start:y2_end] = 1.0
# Left border
u0[x2_start:x2_end, y2_start:y2_start+border_thickness] = 1.0
# Right border
u0[x2_start:x2_end, y2_end-border_thickness:y2_end] = 1.0

# Set the temperature in the first square
u0[x1_start:x1_end, y1_start:y1_end] = 1.0
# Set the temperature in the second square
u0[x2_start:x2_end, y2_start:y2_end] = 1.0

# Function to solve the 2D heat equation with homogeneous boundary conditions
def heat_equation_2d(u, alpha, dt, dx, dy, Nx, Ny):
    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] + alpha * 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
            )
    return u_new

# Set up the figure and axis for the animation
fig, ax = plt.subplots()
cax = ax.imshow(u0, cmap='hot', origin='lower', extent=[0, Lx, 0, Ly])
fig.colorbar(cax)
ax.set_title('2D Heat Equation Simulation')

# Animation function
def animate(frame):
    global u0
    u0 = heat_equation_2d(u0, alpha, dt, dx, dy, Nx, Ny)
    cax.set_array(u0)
    return cax,

# Create animation
anim = FuncAnimation(fig, animate, frames=Nt, interval=4, blit=True)

# Save the animation as a video file
if True:
    anim.save('heat_equation.mp4', writer='ffmpeg', fps=30)
# Display the animation
plt.show()