# Task 1.1, 1.2
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

# report stability
if d > 0.5:
    print("NOT Stable")
else:
    print("Stable")

# explicit scheme integration for `u_t - D u_xx = 0` with boundary conditions
# enforced by `set_bounds` and initial conditions `initial`.
def integrate(*, name: str, initial: np.array, set_bounds: Callable[[np.array], None]) -> None:
    C = initial
    # Setup boundary conditions
    set_bounds(C)

    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/{name}_{i:0>5}.png")
            i += 1
        # update solution using the explicit schema
        C[1:-1] += d * (C[2:] - 2 * C[1:-1] + C[:-2])
        # update right Neumann BC
        set_bounds(C)

# Subtask 1 boundary conditions (Dirichlet and Neumann)
def bounds_1(C):
    C[0] = 1
    C[-1] = C[-2]

# Subtask 2 boundary conditions (two Dirichlet)
def bounds_2(C):
    C[0] = 1
    C[-1] = 0

# run simulations
integrate(name = 'task01_1', initial = np.zeros(nx), set_bounds = bounds_1)
integrate(name = 'task01_2', initial = np.zeros(nx), set_bounds = bounds_2)

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