# Task 1.4
import numpy as np
import scipy, scipy.sparse, scipy.linalg
from matplotlib import pyplot as plt

# plotting helper function
def plot(C, x, t, name):
    plt.clf()
    plt.plot(x, C)
    plt.xlim([0, x[-1]])
    plt.ylim([0, 1.2])
    plt.title(f"{name} - time: {t: >12.0f}")
    plt.savefig(f"plots/{name}_{plot.fignr:0>5}.png")
    plot.fignr += 1

# set static variable as function attribute
plot.fignr = 1
# and initialize a figure
plt.figure(figsize = (8, 6), dpi = 100)

# Config
D = 1e-6            # diffusion coefficient
h = 1               # space domain (max x size)
T = 1e5             # solution end time
nx = 50             # nr of space discretization points
nt = 2000           # nr of time points

# derived constants
dx = h / (nx - 1)   # space step size
dt = T / (nt - 1)   # time step size
d = dt * D / dx**2  # stability/stepsize coefficient

# setup matrices for implicit update scheme `A C^n+1 = (2 I - A) C^n = B C^n`
# The matrix `A` is represented as a `3 x (nx - 1)` matrix as expected by
# `scipy.linalg.solve_banded` representing the 3 diagonals of the tridiagonal
# matrix `A`.
A = np.zeros((3, nx - 1))
# upper minor diagonal
A[0, 2:] = -d
# main diagonal
A[1, 0] = 1
A[1, 1:] = 1 + 2 * d
# and lower minor diagonal
A[2, :-2] = -d
A[2, -2] = -2 * d
# The right hand side matrix `B = 2 I - A` which is represented as a sparse
# matrix in diag. layout which allows for fast matrix vector multiplication
B = -A
B[1, ] += 2
B = scipy.sparse.spdiags(B, (1, 0, -1), B.shape[1], B.shape[1])

# set descritized space evaluation points `x` of `C`
x = np.linspace(0, h, nx)

# Set initial solution
C = np.zeros(nx)
C[0] = 1
C[-1] = C[-3] # (0 = 0)

for n in range(nt):
    # generate roughly 400 plots
    if n % (nt // 400) == 0:
        plot(C, x, n * dt, "task01_4")
    # update solution using the implicit schema
    C[:-1] = scipy.linalg.solve_banded((1, 1), A, B @ C[:-1])
    # set/enforce boundary conditions
    C[0] = 1            # left Dirichlet (theoretically not needed)
    C[-1] = C[-3]       # right Neumann condition

# plot final state
plot(C, x, T, "task01_4")

# to convert generated image sequence to video use:
# $> ffmpeg -r 60 -i plots/task01_4_%05d.png -pix_fmt yuv420p video_1_4.mp4