import numpy as np
import time
import matplotlib.pyplot as plt

# Question 2:
delta = 0.01 # spatial step size
M = int(1. / delta) # number of nodes - 1
eps = 1e-5 # criterion to stop iteration

def distance(phi, phi_prime):
   """
   Inputs:
   phi: current solution
   phi_prime: new solution
   
   Output:
   relative distance between phi and phi_prime and norm of phi
   """
   return np.sum(np.absolute(phi_prime - phi)), np.sum(np.absolute(phi))

phi_jacobi = np.zeros((M + 1, M + 1))
phi_jacobi[-1,:] = 1. # Each row corresponds to a fixed value of y, and each column to a fixed value of x.
phi_before = np.zeros((M + 1, M + 1))
dist, norm = distance(phi_jacobi, phi_before)
dist /= norm
n = 0
tstart = time.process_time()
while dist > eps:
    phi_before = np.copy(phi_jacobi)
    phi_jacobi[1:M,1:M] = 0.25 * ( phi_before[2:,1:M] + phi_before[:M-1,1:M] + phi_before[1:M,2:] + phi_before[1:M,:M-1] )
    dist, norm = distance(phi_before, phi_jacobi)
    dist /= norm
    n += 1
tend = time.process_time()
print('Number of steps for the Jacobi method=', n)
print('Convergence time for the Jacobi method=', tend - tstart, 's')

plt.figure()        
plt.imshow(phi_jacobi, extent=(0, 1, 0, 1), origin='lower', cmap='rainbow')
plt.xlabel('x')
plt.ylabel('y')
plt.colorbar(label = r'$\phi$')
plt.title('Electrostatic potential between conductors (Jacobi)')
plt.tight_layout()
plt.show()

# Question 3: The Gauss-Seidel method requires a smaller number of iterations, but it cannot be vectorized so it takes longer to converge.
phi_gauss = np.zeros((M + 1, M + 1))
phi_gauss[-1,:] = 1.
phi_before = np.zeros((M + 1, M + 1))
dist, norm = distance(phi_gauss, phi_before)
dist /= norm
n = 0
tstart = time.process_time()
while dist > eps:
    phi_before = np.copy(phi_gauss)
    for j in range(1, M):
        for k in range(1, M):
            phi_gauss[j,k] = 0.25 * ( phi_gauss[j+1,k] + phi_gauss[j-1,k] + phi_gauss[j,k+1] + phi_gauss[j,k-1] )
    dist, norm = distance(phi_before, phi_gauss)
    dist /= norm
    n += 1
tend = time.process_time()
print('Number of steps for the Gauss-Seidel method=', n) 
print('Convergence time for the Gauss-Seidel method=', tend - tstart, 's')

plt.figure()
plt.imshow(phi_gauss, extent=(0, 1, 0, 1), origin='lower', cmap='rainbow')
plt.xlabel('x')
plt.ylabel('y')
plt.colorbar(label = r'$\phi$')
plt.title('Electrostatic potential between conductors (Gauss-Seidel)')
plt.tight_layout()
plt.show()

# Question 4: The overrelaxation method is the method which requires the smallest number of steps (with a well-chosen omega), and therefore it is faster than the Gauss-Seidel method.
w = 1.8
phi_over = np.zeros((M + 1, M + 1))
phi_over[-1,:] = 1.
phi_before = np.zeros((M + 1, M + 1))
dist, norm = distance(phi_over, phi_before)
dist /= norm
n = 0
tstart = time.process_time()
while dist > eps:
    phi_before = np.copy(phi_over)
    for j in range(1, M):
        for k in range(1, M):
            phi_hat = 0.25 * ( phi_over[j+1,k] + phi_over[j-1,k] + phi_over[j,k+1] + phi_over[j,k-1] )
            phi_over[j,k] = ( 1. - w ) * phi_before[j,k] + w * phi_hat
    dist, norm = distance(phi_before, phi_over)
    dist /= norm
    n += 1
tend = time.process_time()
print('Number of steps for the overrelaxation method=', n) 
print('Convergence time for the overrelaxation method=', tend - tstart, 's')

plt.figure()
plt.imshow(phi_over, extent=(0, 1, 0, 1), origin='lower', cmap='rainbow')
plt.xlabel('x')
plt.ylabel('y')
plt.colorbar(label = r'$\phi$')
plt.title('Electrostatic potential between conductors (overrelaxation)')
plt.tight_layout()
plt.show()

# Question 5: The error is larger for the Jacobi method, then for the Gauss-Seidel method, and eventually for the overrelaxation method. The overrelaxation method is a good compromise.
def exact_solution(x, y, N = 100):
    """
    Inputs:
    (x, y): position
    
    Output:
    exact solution of the potential
    """
    phi = 0.
    for i in range(N):
        phi += np.sin((2. * i + 1.) * np.pi * x) * np.sinh((2. * i + 1.) * np.pi * y) / np.sinh((2. * i + 1.) * np.pi) / (2. * i + 1.)
    return phi * 4. / np.pi

xx, yy = np.meshgrid(np.arange(0., 1. + delta, delta), np.arange(0., 1. + delta, delta))
phi_exact = exact_solution(xx, yy)
dist, norm = distance(phi_jacobi, phi_exact)
print('Error with the exact solution for the Jacobi method (in %)=', dist / norm * 100.)
dist, norm = distance(phi_gauss, phi_exact)
print('Error with the exact solution for the Gauss-Seidel method (in %)=', dist / norm * 100.)
dist, norm = distance(phi_over, phi_exact)
print('Error with the exact solution for the overrelaxation method (in %)=', dist / norm * 100.)