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

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

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

# Task 1.1 equation system matrix `C^n+1 = A1 C^n`
def A1(d):
    a1 = np.zeros([3, nx])
    # superdiagonal (above the main diagonal)
    a1[0, 1] = 0
    a1[0, 2:] = d
    # main diagonal
    a1[1, 0] = 1
    a1[1, 1:] = 1 - 2 * d
    # subdiagonal (below the main diagonal)
    a1[2, :-2] = d
    a1[2, -2] = 2 * d
    # convert to sparse tridiagonal matrix
    return spdiags(a1, [1, 0, -1], nx, nx)

# Task 1.2 equation system matrix `C^n+1 = A2 C^n`
def A2(d):
    a2 = A1(d).data
    # addapt boundary condition coefficients, the rest is identical
    a2[0, -1] = a2[1, -1] = a2[2, -2] = 0
    # convert to sparse tridiag
    return spdiags(a2, [1, 0, -1], nx, nx)

# Task 1.3 equation system matrix for update rule `A3 C^n+1 = C^n`
def A3(d):
    a3 = np.zeros([3, nx])
    # superdiagonal
    a3[0, 1] = 0
    a3[0, 2:] = -d
    # main diagonal
    a3[1, 0] = 1
    a3[1, 1:] = 1 + 2 * d
    # supdiagonal
    a3[2, :-2] = -d
    a3[2, -2] = -2 * d
    # as tridiag
    return spdiags(a3, [1, 0, -1], nx, nx)

# Task 1.4 equation system matrix `A4`` in `(I + A4) C^n+1 = (I - A4) C^n`
def A4(d):
    a4 = np.zeros([3, nx])
    # superdiagonal
    a4[0, 2:] = -d / 2
    # main diagonal
    a4[1, 1:] = d
    # supdiagonal
    a4[2, :-2] = -d / 2
    a4[2, -2] = -d
    return spdiags(a4, [1, 0, -1], nx, nx)
# setup `I + A4` and `I - A4`
def IpA4(d):
    ipA4 = A4(d).data
    ipA4[1, ] += 1
    return spdiags(ipA4, [1, 0, -1], nx, nx)
def ImA4(d):
    imA4 = -A4(d).data
    imA4[1, ] += 1
    return spdiags(imA4, [1, 0, -1], nx, nx)

# plots a solution to file, simply for code size reduction in `solve`
def plot(x, C, name):
    plt.clf()
    plt.plot(x, C)
    plt.xlim([0, h])
    plt.ylim([0, 1.2])
    plt.title(name)
    plt.xlabel("Distance from Source: x")
    plt.ylabel("Concentration: C(x, t)")
    plt.savefig(f"plots/{name}_{plot.i:0>5}.png")
    plot.i += 1
plot.i = 0

# solver for `C_t - D C_xx = 0` given discretized equation system matrices for
# a time update rule `L C^n+1 = R C^n` starting from an initial solution `C`.
# The update is iterated for `nt` time steps corresponding to a solution at
# time `T = dt nt`
def solve(C, L, R, nt, plotname = None, plotskip = 5):
    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, C, plotname)                            # plot initial solution
    for t in range(nt):
        rhs = C if R is None else R @ C
        C = rhs if L is None else solve_banded((1, 1), L.data, rhs)
        if plotname is not None and ((t + 1) % plotskip == 0):
            plot(x, C, plotname)                        # plot current solution
    return C

# setup initial solution
C0 = np.zeros(nx)
C0[0] = 1

# run simulation and generate plots for each subtask
C1 = solve(C0, None, A1(d), nt, "task01_1")
Cx = solve(C0, None, A1(0.55), 200, "task01_x", 1)  # run with unstable `d > 0.5`
C2 = solve(C0, None, A2(d), nt, "task01_2")
C3 = solve(C0, A3(d), None, nt, "task01_3")
C4 = solve(C0, IpA4(d), ImA4(d), nt, "task01_4")

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

# compute till no time change, e.g. `t -> inf`
C1inf = solve(C0, None, A1(d), 20000)
C2inf = solve(C0, None, A2(d), 20000)

x = np.linspace(0, h, nx)                       # space grid points
fig = plt.figure(figsize = (8, 6), dpi = 100)   # setup plot figure size
plot.i = 0                                      # reset plot counter
plot(x, C1inf, "Cinf")
plot(x, C2inf, "Cinf")


# # All plots in one
C1 = solve(C0, None, A1(d), nt)
C2 = solve(C0, None, A2(d), nt)
C3 = solve(C0, A3(d), None, nt)
C4 = solve(C0, IpA4(d), ImA4(d), nt)

x = np.linspace(0, h, nx)                       # space grid points
plt.figure(figsize = (8, 6), dpi = 100)   # setup plot figure size
plt.plot(x, C1, label = "1.1")
plt.plot(x, C2, label = "1.2")
plt.plot(x, C3, label = "1.3")
plt.plot(x, C4, label = "1.4")
plt.xlim([0, h])
plt.ylim([0, 1.2])
plt.title("All")
plt.legend()
plt.xlabel("Distance from Source: x")
plt.ylabel("Concentration: C(x, t)")
plt.savefig(f"plots/task01_all.png")