import matplotlib.pyplot as plt
import numpy as np
N=70  #2N=140, hence the spatial discretisation step is 1/140
e=0.01  #e is larger than the spatial discretization step
a= np.full((2*N, 2*N), -1/(e*2*N)) #Full quadratic matrix (2N,2N) for the integral term. We need 2N of such matrices for every x_i of the discretisation
for k in range(2*N):
    a[k][k]=np.sign(k-N+0.5)*(-2*N+2*k+1)+(1/e)-(1/(e*2*N)) #Diagonal of the (2N,2N) matrix
A = np.zeros((4*N*N, 4*N*N))   #Big matrix for the linear system of equation
for k in range(2*N):
    A[k*2*N:(k+1)*2*N,k*2*N:(k+1)*2*N]=a    #On the diagaonal there are 2N of the (2N,2N) matrices constructed before
for k in range(N):
    for m in range(1,2*N):
        A[2*N*(m-1)+k][2*N*m+k]=1+2*k-2*N  #Terms for the spatial derivative when n<0
for k in range(N,2*N):
    for m in range(1,2*N):
        A[2*N*m+k][2*N*(m-1)+k]=-1-2*k+2*N  #Terms for the spatial derivative when n>0
for k in range(N):
    A[2*N*(2*N-1)+k][2*N*(2*N-1)+k]=4*N-4*k-2+(1/e)-(1/(e*2*N))   #Diagonal terms for x=1-1/(4N)
for k in range(N,2*N):
    A[k][k]=-4*N+4*k+2+(1/e)-(1/(e*2*N))  #Diagonal terms for x=1/(4N)
B= np.zeros(4*N*N)
for j in range(N,2*N):
    B[j]=(2+4*j-4*N)*(1+2*j-2*N)/(2*N)  #The vector B contains the boundary condition
x = np.linalg.solve(A, B)   #Solve the linear system of equations
b = x.reshape(2*N, 2*N)    #each row of b represents the value of J at a given spatial point for all directions
c= b.transpose()            #each row of c represents the value of J at a given direction for all spatial points
D= np.zeros(2*N)
for i in range(2*N):
    D[i]=(2*i+1)/(4*N)        #The spatial discretization
fig = plt.figure(figsize=(12, 6))
plt.plot(D,c[20], 'tab:blue')
plt.plot(D, c[55], 'tab:orange')
plt.plot(D,c[80], 'tab:green')
plt.plot(D, c[130], 'tab:red')
plt.plot([0.04, 0.04],[0, 0.8],  color='tab:gray', linestyle='dashed')
plt.xlabel("x")
plt.ylabel("J(x,n)")
plt.show()