#!/usr/bin/env python3
'''
Program for solving Poisson equation (Laplace(x) = b) on a rectangular
grid (Nx x Ny).
'''

import numpy as np
import matplotlib.pyplot as plt


def getLinearizedIdx(x,y):
    '''
    Returns index in vector x, which corresponds with point (x,y) on grid.
    '''
    return x+y*Nx

def initMatrix():
    '''
    Creates matrix A and vector b for eqn Ax=b. Three possibilities
    are involved:
    1) Approximation of Laplace operator by finite differences
    x[(i+1,j)]+x[(i-1),j]+x[(i,j+1)]+x[(i,j-1)]-4*x[(i,j)]=h*h*Laplace(x),
    where (i,j) denotes point in grid (=getLinearizedIdx(i,j)).
    2) Dirichlet boundary condition: x[(i,j)] = 0 (more generally x[k]=C,
    where C is constant).
    3) von Neumann boundary condition: d/dx(x[(i,j)]) = 0 or d/dy(x[(i,j)])=0
    approximated by (d/dy(x[(i,j)]) = 0) <=> (x[(i,j)]=x[(i,j+1)]).
    '''
    tmp = 1./(h*h)
    N = Nx*Ny
    A = np.zeros((N,N))
    b = np.zeros(N)
    for i in range(Nx):
        for j in range(Ny):
            k = getLinearizedIdx(i,j)
            # upper plate
            if j==0: #Dirichlet
                A[k,k] = 1.0
                b[k] = 1.0 # nonzero potential
            # lower plate
            elif j==Ny-1: #Dirichlet
                A[k,k] = 1.0
            # circle
            elif (j-Ny/4)**2+(i-Nx/2)**2<=Ny/5: #Dirichlet
                A[k,k] = 1.0
            # sides
            elif i==0: #von Neumann
                A[k,k] = 1.0
                A[k,getLinearizedIdx(i+1,j)] = -1.0
            elif i==Nx-1: #von Neumann
                A[k,k] = 1.0
                A[k,getLinearizedIdx(i-1,j)] = -1.0
            # Laplace operator approximation
            else:
                A[k,k] = -4.0*tmp
                A[k,getLinearizedIdx(i-1,j)] = 1.0*tmp
                A[k,getLinearizedIdx(i+1,j)] = 1.0*tmp
                A[k,getLinearizedIdx(i,j-1)] = 1.0*tmp
                A[k,getLinearizedIdx(i,j+1)] = 1.0*tmp
    return A,b

# parameters
Nx,Ny = 100,100 #grid size
h = 0.1 #size of grid cell in meters (in both dimensions)
# boundaries of area (used only for graph legend)
X0,Y0 = -5.0, 0.0
X1,Y1 = X0+h*Nx, Y0+h*Ny

# actual solving
A,b = initMatrix()
x = np.linalg.solve(A,b)

# displaying of results
plt.imshow(x.reshape((Ny,Nx)), extent=(X0,X1,Y0,Y1))
plt.colorbar().set_label("$\Phi$ [V]")
plt.xlabel("$x$ [m]")
plt.ylabel("$y$ [m]")
plt.show()