import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from scipy.sparse import diags, csc_matrix
from scipy.sparse.linalg import inv

def fisher_kpp(L, N, T, dt, r, D, res_dpi, save_gif=False, gif_filename="animation.gif"):
    dx = 2 * L / N
    x = np.linspace(-L, L, N)
    u = 0.2*np.where(np.logical_and(x > -L / 10, x < L / 10), 1, 0)
    
    # Build the matrix of the Laplacien
    diagonals = [1, -2, 1]
    offsets = [-1, 0, 1]
    Lap = diags(diagonals, offsets, shape=(N, N)).toarray()
    Lap[0, 1] = 2
    Lap[-1, -2] = 2
    Lap *= D / (dx ** 2)
    
    # print(Lap)

    # Build the linear system
    A = np.eye(N) - dt * Lap
    # print(A)

    # Convert A into a sparce matrix and compute  its inverse
    A_csc = csc_matrix(A)
    A_inv = inv(A_csc)
    
    def update_u(u):
        u_new = A_inv @ (u + dt * r * u * (1 - u))
        return u_new

    fig, ax = plt.subplots(dpi = res_dpi)
    line, = ax.plot(x, u, lw=3)
    ax.axhline(y=1, color='r', linestyle='--')  # Add the horizontal line y=1
    ax.set_xlim(-L, L)
    ax.set_ylim(0, 1.1)
    ax.set_xlabel("x")
    ax.set_ylabel("u")
    ax.set_title("Fisher-KPP equation")
    
    def animate(i):
        nonlocal u
        line.set_ydata(u)
        u = update_u(u)
        print(i*dt)
        return line,

    ani = animation.FuncAnimation(fig, animate, frames=int(T / dt), blit=True, interval=500 * dt, repeat=False)
    
    if save_gif:
        f = r"video_save.gif"
        writergif = animation.PillowWriter(fps=40)
        ani.save(f, writer=writergif)
    else:
        plt.show()

# Parameters
L = 100
N = 500
T = 10
dt = 0.01
r = 10
D = 1
DPI = 100

fisher_kpp(L, N, T, dt, r, D, res_dpi=DPI, save_gif=False, gif_filename="animation.gif")