#!/usr/bin/env python3
import imageio
import io
import numpy as np

import matplotlib.pyplot as plt
from matplotlib.figure import Figure
from matplotlib.axes import Axes


def _lim(data: np.ndarray):

    # Turn data flat for a second.
    data = data.flatten()
    data = data[np.isfinite(data)]

    # Find the min and max.
    min_, max_ = np.percentile(data, [0.1, 99.9])

    var_ = (max_ - min_)
    return (min_ - var_ * 0.1, max_ + var_ * 0.1)


def generate_gif(X, U, T, filename: str, xlim=None, ylim=None, label="t={time:0.3f}"):

    dT = T[1] - T[0]

    images = []

    if xlim is None:
        xlim = _lim(X)

    if ylim is None:
        ylim = _lim(U)

    for u, t in zip(U, T):

        fig: Figure = plt.figure()
        ax: Axes = fig.gca()

        # Plot the line.
        plt.plot(X, u, label=label.format(time=t))

        if xlim is not None:
            ax.set_xlim(xlim)

        if ylim is not None:
            ax.set_ylim(ylim)

        # Activate the grid.
        ax.grid(visible=True)

        # Activate the legend.
        ax.legend(loc='upper right')

        with io.BytesIO() as buf:
            fig.savefig(buf, format='png')
            buf.seek(0)
            images.append(imageio.imread(buf))

        plt.close()
    # Save the list of images as a gif
    imageio.mimsave(filename, images, duration=dT, loop=0)


# Snippet 1

import numpy as np


X = np.linspace(0, 1, 11)
dX = X[1] - X[0]

Y = X**2
Y_prime = (Y[1:] - Y[:-1]) / dX

plt.figure()

plt.plot(X, Y, label="f(x)")
plt.plot(X, 2 * X, label="f'(x)")
plt.plot(X[:-1], Y_prime, label="approx f'(x)")
plt.grid()
plt.legend()
plt.savefig("derivative.png")


# Snippet 2

import numpy as np


X = np.linspace(0, 1, 21)
dX = X[1] - X[0]

Y = X**2
Y_prime = (np.roll(Y, -1) - Y) / dX

plt.figure()
plt.plot(X, Y, label="f(x)")
plt.plot(X, 2 * X, label="f'(x)")
plt.plot(X, Y_prime, label="approx f'(x)")
plt.grid()
plt.legend()
plt.savefig("derivative_periodic.png")


# Snippet 3

import numpy as np

def solve_heat(U_0, dT, dX, steps=10, skip=0):

    # If this coefficient is small the simulation is stable.
    r = dT / dX ** 2

    # Initialize first the variables.
    U = np.zeros((steps, X.shape[0]))
    T = np.zeros(steps)

    # Initial conditions.
    U[0] = U_0

    # Start the iterations.
    for i in range(1, steps):

        U_ = U[i - 1]

        for j in range(skip + 1):
            U_ = (1 - 2 * r) * U_ + r * (np.roll(U_, 1) + np.roll(U_, -1))

        U[i] = U_
        T[i] = T[i -1] + dT * (skip + 1)

    return T, U

def initial_condition(x, sigma=0.2):
    return np.exp(-x**2 / (2 * sigma**2)) / np.sqrt(2 * np.pi * sigma ** 2)

# Spatial domain.
X = np.linspace(-1, 1, 123)

# Initial condition.
U_0 = initial_condition(X)

# Solve the equation.
T, U = solve_heat(U_0, dT=0.0001, dX=X[1] - X[0], steps=40, skip=100)

# Generate a GIF.
generate_gif(X, U, T, 'evolution_heat_1d_2.gif')


# Snippet 4

def solve_burgers(U_0, dT, dX, steps=10, skip=0):

    # If this coefficient is small the simulation is stable.
    r = dT / dX

    # Initialize first the variables.
    U = np.zeros((steps, X.shape[0]))
    T = np.zeros(steps)

    # Initial conditions.
    U[0] = U_0

    # Start the iterations.
    for i in range(1, steps):

        U_ = U[i - 1]

        for j in range(skip + 1):
            U_ = U_ * (1 - r * (np.roll(U_, -1) - U_))

        U[i] = U_
        T[i] = T[i -1] + dT * (skip + 1)

    return T, U

def initial_condition(x, sigma=0.2):
    return np.exp(-x**2 / (2 * sigma**2)) / np.sqrt(2 * np.pi * sigma ** 2)

# Spatial domain.
X = np.linspace(-1, 1, 123)

# Initial condition.
U_0 = initial_condition(X)

#
T, U = solve_burgers(U_0, dT=0.0001, dX=X[1] - X[0], steps=20, skip=100)
generate_gif(X, U, T, 'evolution_burgers_1d.gif')


# Snippet 5

X = np.linspace(-1, 1, 1023)
u_0 = initial_condition(X)
T, U = solve_burgers(u_0, dT=0.0000001, dX=X[1] - X[0], steps=20, skip=10000)
generate_gif(X, U, T, 'evolution_burgers_1d_2.gif')


# Snippet 6

def solve_burgers(U_0, dT, dX, mu=0.0, steps=10, skip=0):

    # If these coefficient are too large the simulation is not stable.
    r = dT / dX
    r2 = dT / dX ** 2

    # Initialize first the variables.
    U = np.zeros((steps, X.shape[0]))
    T = np.zeros(steps)

    # Initial conditions.
    U[0] = U_0

    # Start the iterations.
    for i in range(1, steps):

        U_ = U[i - 1]

        for j in range(skip + 1):
            U_ = U_ * (1 - r * (np.roll(U_, -1) - U_) - 2 * r2 * mu) \
                 + mu * r2 * (np.roll(U_, 1) + np.roll(U_, -1))

        U[i] = U_
        T[i] = T[i -1] + dT * (skip + 1)

    return T, U

X = np.linspace(-1, 1, 123)
u_0 = initial_condition(X)
T, U = solve_burgers(u_0, dT=0.00001, mu=0.025, dX=X[1] - X[0], steps=80, skip=4000)
generate_gif(X, U, T, 'evolution_burgers_vis_1d.gif')