import numpy as np
from scipy.sparse import spdiags
from scipy.linalg import solve_banded
from matplotlib import pyplot as plt

# Config
dx = 0.01           # space step size
h = 1.5             # space domain (max x size)

# plots a solution to file, simply for code size reduction in `solve`
def plot(x, t, C, C0, analytic, plotname):
    plt.clf()
    plt.plot(x, C, "o")
    plt.xlim([0, h])
    plt.title(plotname)
    plt.xlabel("Periodic Space: x")
    plt.ylabel("Wave Form: C(x, t)")
    if analytic is not None:
        plt.plot(x, analytic(x, t, C0))
    plt.savefig(f"plots/{plotname}_{plot.i:0>5}.png")
    plot.i += 1
plot.i = 0

# solver for `C_t + U C_x = 0` given initial state `C`, the current `C0` which
# incorporates `U` as well as the space and time discretization and the number
# of time steps `nt`.
def solve(C, C0, nt, plotname = None, plotskip = 1, analytic = None):
    nx = C.shape[0]
    if plotname is not None:
        plot.i = 0                                      # (re)set plot counter
        x = np.linspace(0, h, nx)                       # space grid points
        fig = plt.figure(figsize = (8, 6), dpi = 100)   # setup plot figure size
        plot(x, 0, C, C0, analytic, plotname)           # plot initial solution
    # space grid point `Indices Minus 1`
    im1 = np.array([nx - 1] + list(range(nx - 1)))
    for t in range(nt):
        # upwind update scheme
        C = (1 - C0) * C + C0 * C[im1]
        if plotname is not None and ((t + 1) % plotskip == 0):
            plot(x, t + 1, C, C0, analytic, plotname)    # plot current solution
    return C

# setup initial conditions
x = np.linspace(0, h, int(h / dx))
gauss = np.exp(-10 * (4 * x - 1)**2)
square = np.array((0.1 < x) * (x < 0.3), dtype = "float")

def gauss_analytic(x, t, C0):
    return np.exp(-10 * (4 * (np.mod(x - t * dx / C0, h + 1e-9)) - 1)**2)
def square_analytic(x, t, C0):
    x = x - t * dx / C0
    x = np.mod(x, h + 1e-9) # hack to avoid shift in the plot
    return np.array((0.1 < x) * (x < 0.3), dtype = "float")

solve(gauss, 1, 200, "task02_gauss_C0_1", analytic = gauss_analytic)
solve(square, 1, 200, "task02_square_C0_1", analytic = square_analytic)
solve(gauss, 0.7, 200, "task02_gauss_C0_0.7", analytic = gauss_analytic)
solve(square, 0.7, 200, "task02_square_C0_0.7", analytic = square_analytic)
solve(gauss, 1.1, 200, "task02_gauss_C0_1.1", analytic = gauss_analytic)
solve(square, 1.1, 200, "task02_square_C0_1.1", analytic = square_analytic)

# To generate videos out of the generated plots
# $> ffmpeg -y -r 60 -i plots/task02_gauss_C0_1_%05d.png -pix_fmt yuv420p task02_gauss_C0_1.mp4
# $> ffmpeg -y -r 60 -i plots/task02_square_C0_1_%05d.png -pix_fmt yuv420p task02_square_C0_1.mp4
# $> ffmpeg -y -r 60 -i plots/task02_gauss_C0_0.7_%05d.png -pix_fmt yuv420p task02_gauss_C0_0.7.mp4
# $> ffmpeg -y -r 60 -i plots/task02_square_C0_0.7_%05d.png -pix_fmt yuv420p task02_square_C0_0.7.mp4
# $> ffmpeg -y -r 60 -i plots/task02_gauss_C0_1.1_%05d.png -pix_fmt yuv420p task02_gauss_C0_1.1.mp4
# $> ffmpeg -y -r 60 -i plots/task02_square_C0_1.1_%05d.png -pix_fmt yuv420p task02_square_C0_1.1.mp4