import numpy as np
import matplotlib.pyplot as plt

def solve_poisson_gauss_seidel(u, bc_mask, bc_values, tol=1e-6, max_iter=50000):
    """
    Solve Laplace equation on a 2D grid using Gauss-Seidel iteration.
    u         : initial guess (will be overwritten)
    bc_mask   : boolean array, True where BCs are fixed
    bc_values : fixed values at bc_mask positions
    tol       : stopping criterion (max change < tol)
    max_iter  : maximum number of iterations
    returns   : u, iteration count
    """
    ny, nx = u.shape
    # Apply initial BCs
    u[bc_mask] = bc_values[bc_mask]
    
    iter_count = 0
    for _ in range(max_iter):
        max_diff = 0.0
        # Gauss-Seidel sweep (red-black ordering would be faster, but simple is fine)
        for i in range(1, ny-1):
            for j in range(1, nx-1):
                if bc_mask[i, j]:
                    continue
                new_val = 0.25 * (u[i-1, j] + u[i+1, j] + u[i, j-1] + u[i, j+1])
                diff = abs(new_val - u[i, j])
                if diff > max_diff:
                    max_diff = diff
                u[i, j] = new_val
        iter_count += 1
        if max_diff < tol:
            break
    return u, iter_count

# Grid parameters
nx, ny = 100, 100
x = np.linspace(0, 1, nx)
y = np.linspace(0, 1, ny)
X, Y = np.meshgrid(x, y)

# Boundary condition mask (True where BC is fixed)
bc_mask = np.zeros((ny, nx), dtype=bool)
bc_mask[0, :] = True   # top
bc_mask[-1, :] = True  # bottom
bc_mask[:, 0] = True   # left
bc_mask[:, -1] = True  # right

# ---- Initial boundary conditions ----
u_initial = np.zeros((ny, nx))
# Left: 100, Right: 0, Top/Bottom: linear from 100 to 0
u_initial[0, :] = 100 * (1 - x)      # top
u_initial[-1, :] = 100 * (1 - x)     # bottom
u_initial[:, 0] = 100                # left
u_initial[:, -1] = 0                 # right

# Solve for initial BC (this is our "base state")
u_base = u_initial.copy()
u_base, iter_base = solve_poisson_gauss_seidel(u_base, bc_mask, u_initial, tol=1e-6)
print(f"Initial BC solution converged in {iter_base} iterations")

# ---- New boundary conditions (changed left wall) ----
u_new_bc = u_initial.copy()
u_new_bc[:, 0] = 50   # left wall changed from 100 to 50

# Cold start: from zero initial guess
u_cold = np.zeros_like(u_new_bc)
u_cold, iter_cold = solve_poisson_gauss_seidel(u_cold, bc_mask, u_new_bc, tol=1e-6)
print(f"Cold start (zero guess) converged in {iter_cold} iterations")

# Warm start: use previous solution as initial guess
u_warm = u_base.copy()   # <--- this is the "held near-finished state"
# But we must enforce new BCs on the initial guess
u_warm[bc_mask] = u_new_bc[bc_mask]
u_warm, iter_warm = solve_poisson_gauss_seidel(u_warm, bc_mask, u_new_bc, tol=1e-6)
print(f"Warm start (previous solution) converged in {iter_warm} iterations")

# Verify that both solutions are essentially identical
error = np.max(np.abs(u_cold - u_warm))
print(f"Max difference between cold and warm solutions: {error:.2e} (effectively zero)")

# Speedup
if iter_cold > 0:
    print(f"\nSpeedup factor: {iter_cold / iter_warm:.1f}x fewer iterations")
    print(f"Energy saving (approx): {100 * (1 - iter_warm/iter_cold):.1f}%")

# Optional visualisation (shows that both solutions are identical)
if nx <= 100:  # avoid huge plots if resolution high, but still works
    fig, axes = plt.subplots(1, 3, figsize=(12, 4))
    im0 = axes[0].contourf(X, Y, u_base, levels=20, cmap='hot')
    axes[0].set_title("Base solution (old BC)")
    im1 = axes[1].contourf(X, Y, u_cold, levels=20, cmap='hot')
    axes[1].set_title("New BC: cold start")
    im2 = axes[2].contourf(X, Y, u_warm, levels=20, cmap='hot')
    axes[2].set_title("New BC: warm start")
    plt.colorbar(im2, ax=axes, orientation='horizontal', pad=0.05)
    plt.tight_layout()
    plt.show()