# Task 1.3
import numpy as np
from typing import Callable
from matplotlib import pyplot as plt

# Config
D = 1e-6            # diffusion coefficient
h = 1               # space domain (max x size)
T = 2e6             # solution end time
nx = 50             # nr of space discretization points
nt = 20000          # nr of time discretization 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 implicit scheme equation matrix
T = (1 + 2 * d) * np.eye(nx) - d * np.eye(nx, k = 1) - d * np.eye(nx, k = -1)
# fix boundary condition equations
T[0, 0] = 1         # Left Dirichlet BC
T[0, 1] = 0
T[-1, -2] = 1       # Right Neumann BC
T[-1, -1] = 0

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

i = 0       # index for plot generation
plt.figure(figsize = (8, 6), dpi = 100)
for t in range(nt):
    # every 400'th time step save a plot
    if t % (nt // 400) == 0:
        plt.clf()
        plt.plot(np.linspace(0, h, nx), C)
        plt.xlim([0, h])
        plt.ylim([0, 1.2])
        plt.savefig(f"plots/task01_3_{i:0>5}.png")
        i += 1
    # update solution using the explicit schema
    C = np.linalg.solve(T, C)
    # fix BC conditions (theoretically, they are set by the update but for
    # stability reasons (numerical) we enforce the correct values)
    C[0] = 1
    C[-1] = C[-2]

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