# -*- coding: utf-8 -*-
from __future__ import print_function  # for Python 3 compatibility
#import matplotlib        # If the code exits without having shown any plots, remove
#matplotlib.use("TkAgg")  # the comment characters at the beginning of these two lines.
import poisson2d
import timeit
import scipy.linalg
import numpy as np
import matplotlib.pyplot as plt

N = 50
# Generate 25 charge distributions with side length N
rhos = poisson2d.getRhos(N)

# Plot the charge distributions
plt.figure("Charge distribution")
poisson2d.plotImages(rhos)

# Get the matrix to approximate the ODE
A = poisson2d.getMatrix(N)

# Approximate the ODE for the different charge distributions
phis=[]
for rho in rhos:
    # The rho "image" must be flattened into a vector
    rhoflat = rho.reshape((N * N))
    # We get out the (flattened) potential
    phiflat = scipy.linalg.solve(A, rhoflat)
    # Reshape the potential
    phi = phiflat.reshape((N, N))
    phis.append(phi)

# Plot the potentials
plt.figure("Potentials")
poisson2d.plotImages(phis)


# Time a single run of the Gauss elimination
# Get a single rho
rho = rhos[0]
rhoflat = rho.reshape((N*N))
# Time gauss elimination
def f():
    scipy.linalg.solve(A, rhoflat)
t = timeit.timeit(f, number=1)

print("{} seconds".format(t))

plt.show()