"""
Illustration of the heat equation

Solve the heat equation using finite differences and Forward Euler.

Based on: https://commons.wikimedia.org/wiki/File:Heat_eqn.gif
"""

from __future__ import division, print_function
import numpy as np
from mayavi import mlab
import subprocess

path_to_convert = "C:\Program Files\ImageMagick-6.9.3\convert.exe"

def step_function(N, scale, X, Y, shape="crescent"):
    """Function that is 1 on a set and 0 outside of it"""
    shapes = ["crescent", "cylinder", "hexagon", "superquadric", "smiley"]
    
    if shape not in shapes:
        shape = "crescent"

    if shape == "cylinder":
        Z = np.ones_like(X)
        Z[X**2 + Y**2 < 0.5] = 0
        Z[X**2 + Y**2 > 2] = 0

    if shape == "superquadric":
        Z = np.ones_like(X)
        Z[np.abs(X)**0.5 + np.abs(Y)**0.5 > 1.5] = 0

    if shape == "hexagon":
        Z = np.ones_like(X)
        hexa = 2*np.abs(X) + np.abs(X - Y*np.sqrt(3)) +\
            np.abs(X + Y*np.sqrt(3))
        Z[hexa > 6] = 0

    if shape == "crescent":
        c = 2
        d = -1
        e = 1
        f = 0.5
        k = 1.2
        shift = 10        
        Z = (c**2 - (X/e - d)**2 - (Y/f)**2)**2 + k*(c + d - X/e)**3 - shift
        Z = 1 - np.maximum(np.sign(Z), 0)
        
    if shape == "smiley":
        Z = np.ones_like(X)
        fac = 1.2
        x_eye = 0.5
        y_eye = 0.4
        bicorn = fac**2*(Y + 0.3)**2*(1 - fac**2*X**2) -\
                (fac**2*X**2 - 2*fac*(Y + 0.3) - 1)**2
        left_eye = (X + x_eye)**2/0.1 + (Y - y_eye)**2/0.4 - 1
        right_eye = (X - x_eye)**2/0.1 + (Y - y_eye)**2/0.4 - 1
        Z[X**2 + Y**2 > 2] = 0
        Z[bicorn > 0] = 0
        Z[left_eye < 0] = 0
        Z[right_eye < 0] = 0

    Z = scale * Z
    return Z

def data_gen(num):
    # Solve the heat equation with zero boundary conditions
    for cont in range(ntime_anim):
        Z[1:N-1, 1:N-1] = Z[1:N-1, 1:N-1] + dt*(Z[2:N, 1:N-1] +
                             Z[0:N-2, 1:N-1] + Z[1:N-1, 0:N-2] +
                             Z[1:N-1, 2:N] - 4*Z[1:N-1, 1:N-1])/dx**2

    surf = mlab.surf(X, Y, Z, colormap='autumn', warp_scale=1)
    # Change the visualization parameters.
    surf.actor.property.interpolation = 'phong'
    surf.actor.property.specular = 0.3
    surf.actor.property.specular_power = 20
    surf.module_manager.scalar_lut_manager.reverse_lut = True
    surf.module_manager.scalar_lut_manager.data_range = np.array([ 0.,  scale])

    return surf

N = 500  # Grid points
L = 2.5  # Box size
X, Y = np.mgrid[-L:L:N*1j, -L:L:N*1j]
scale = 2
Z = step_function(N, scale, X, Y, shape="crescent")
CFL = 0.125
dx = X[1, 0] - X[0, 0]
dy = dx
dt = CFL*dx**2
end_time = 0.05
time = np.arange(0, end_time, dt)
nframes = 50
ntime = time.shape[0]
ntime_anim = int(ntime/nframes)

#%% Plot frames
fname = "heat"
bgcolor = (1, 1, 1)
fig = mlab.figure(size=(1200, 1000), bgcolor=bgcolor)
fig.scene.camera.azimuth(180)
mlab.get_engine()
engine = mlab.get_engine()
scene = engine.scenes[0]
for cont in range(nframes):
    mlab.clf()
    surf = data_gen(cont)
    scene.scene.camera.position = [-8, -8,  7]
    scene.scene.camera.clipping_range = [7, 22]
    scene.scene.camera.focal_point = [0, 0, 1]
    print(cont)
    mlab.savefig("{}_{n:02d}.png".format(fname, n=cont))

#%% Generate video
args = [path_to_convert, "-delay", "10", "-loop" , "0", fname + "_*.png",
        fname + ".gif"]
subprocess.call(args, shell=True)
subprocess.call(["del", "/Q", fname + "*.png"], shell=True)
print("Done!")