"""
(C) Steven Byrnes, 2013. This code is released under the MIT license
http://opensource.org/licenses/MIT

This code runs in Python 2.7 or 3.3. It requires imagemagick to be installed;
that's how it assembles images into animated GIFs.

Creates an animation of the electric field vectors of a
surface-plasmon-polariton wave. The metal-dielectric interface is at z=0,
with metal at z<0 and dielectric (such as air) at z>0.
"""
 
# Use Python 3 style division and print
# a/b is real division, a//b is integer division
from __future__ import division, print_function

import numpy as np
from cmath import exp, pi

import matplotlib.pyplot as plt
from matplotlib.path import Path
import matplotlib.patches as patches

import subprocess, os
directory_now = os.path.dirname(os.path.realpath(__file__))

# --- PART 1 OF 2: PHYSICS --- #

def Efield(x_times_kvac, z_times_kvac, t_times_omega, eps_m, eps_d, print_wave_properties=False):
    """"
    Inputs:
      * x_times_kvac and z_times_kvac are the coordinates (x,z) multiplied
          by the vacuum angular wavenumber kvac = omega / c.
      * t_times_nu is the time t multiplied by frequency nu.
      * eps_m and eps_d are the complex permittivities of the metal and dielectric.
    
    Output: The electric field vector (E_x,E_z). It is scaled so that
      E_z(0, 0, tw, e1, e2) = cos(tw) on the dielectric side of the interface.
    """
    # Calculate components of the angular wavevector in the dielectric and
    # metal, as a multiple of kvac = omega / c:
    eps_m = complex(eps_m) #cast to complex, so square roots won't raise errors.
    kx_over_kvac = (eps_m * eps_d / (eps_m + eps_d))**(1/2)
    kzd_over_kvac = (eps_d - kx_over_kvac**2)**(1/2)
    kzm_over_kvac = (eps_m - kx_over_kvac**2)**(1/2)
    
    # Pick the correct square-roots, so that e^(i*k*z) decays away from interface
    if kzd_over_kvac.imag < 0:
        kzd_over_kvac *= -1
    if kzm_over_kvac.imag > 0:
        kzm_over_kvac *= -1
    
    #double-check the boundary condition
    almost_equal = lambda a,b,tolerance : (abs(a-b) / (abs(a) + abs(b))) < tolerance
    if not almost_equal(kzd_over_kvac * eps_m, kzm_over_kvac * eps_d, 1e-10):
        raise ValueError('Something is wrong! Boundary condition fails!')
    
    if print_wave_properties:
        print('kx / kvac = ', kx_over_kvac)
        print('kzd / kvac = ', kzd_over_kvac)
        print('kzm / kvac = ', kzm_over_kvac)
        print('Wavelength / (Vacuum wavelength) = ', 1/(kx_over_kvac.real))
        if kx_over_kvac.imag != 0:
            print('(Decay length) / (Vacuum wavelength) = ', 1/(kx_over_kvac.imag))
        else:
            print('Wave does not decay, it propagates forever.')
        print('(Decay length into dielectric) / (Vacuum wavelength) = ', 1/(kzd_over_kvac.imag))
        print('(Decay length into metal) / (Vacuum wavelength) = ', -1/(kzm_over_kvac.imag))
   
    if z_times_kvac > 0:
        # dielectric
        Ez = exp(1j * kx_over_kvac * x_times_kvac + 1j * kzd_over_kvac * z_times_kvac - 1j * t_times_omega)
        Ex = -Ez * kzd_over_kvac / kx_over_kvac
    else:
        # metal
        Ez = (kzd_over_kvac / kzm_over_kvac) * exp(
                 1j * kx_over_kvac * x_times_kvac + 1j * kzm_over_kvac * z_times_kvac - 1j * t_times_omega)
        Ex = -Ez * kzm_over_kvac / kx_over_kvac
    return (Ex.real, Ez.real)

# --- PART 2 OF 2: DRAWING --- #

def draw_box(x1, x2, y1, y2, color):
    """
    Code to draw a rectangular box in matplotlib ... used to color the metal area.
    See http://matplotlib.org/users/path_tutorial.html
    """
    vertices = [ (x1, y1), (x1, y2), (x2, y2), (x2, y1), (x1, y1)]
    codes = [Path.MOVETO,Path.LINETO,Path.LINETO,Path.LINETO,Path.CLOSEPOLY]
    path = Path(vertices, codes)
    #lw=0 means no outline; zorder=-1 means it shouldn't block the arrows.
    return patches.PathPatch(path, facecolor=color, lw=0, zorder=-1)

def draw_frame(t_times_omega, eps_m, eps_d, aspect_ratio=1.5,
               frac_metal=0.5, fig_width_px=480, x_range_times_kvac=2,
               img_filename=None):
    """
    Draw one frame of the animation.
    
    Inputs:
      * t_times_omega is time multiplied by angular frequency,
          it goes 0 --> 2pi each cycle.
      * eps_m and eps_d are the dielectric constants of the metal and dielectric
      * aspect_ratio is width over height
      * frac_metal is how much of the image is taken up by the metal.
      * fig_width_px is figure width in pixels
      * x_range_times_kvac is the width of the image as a multiple of kvac.
      * "img_filename" is what to save the frame as (or None to not save it).
    """
    # Figure geometry...
    fig_height_px = fig_width_px // aspect_ratio
    dpi = 80 #This number doesn't affect the final animation...
    fig_width_inches = fig_width_px / dpi
    fig_height_inches = fig_height_px / dpi
    
    # Coordinate limits in figure. All are implicitly multiplied by kvac.
    z_range_times_kvac = x_range_times_kvac / aspect_ratio
    xmin = 0
    xmax = x_range_times_kvac
    zmin = -z_range_times_kvac * frac_metal
    zmax = z_range_times_kvac * (1-frac_metal)
    
    # How many arrows to draw?
    num_arrows_x = 15
    num_arrows_z = num_arrows_x // aspect_ratio
    
    # Pick arrow coordinates...
    arrow_x_list, spacing = np.linspace(xmin, xmax, num=num_arrows_x,
                                     endpoint=False, retstep=True)
    arrow_x_list += spacing / 2
    
    arrow_z_list, spacing = np.linspace(zmin, zmax, num=num_arrows_z,
                                     endpoint=False, retstep=True)
    arrow_z_list += spacing / 2

    X,Z = np.meshgrid(arrow_x_list, arrow_z_list)
    
    # Arrow length scale: Larger number = smaller arrows
    arrow_len_scale = 15
    
    # Calculate the length of each arrow
    Ex_func = np.vectorize(lambda x,z : Efield(x,z,t_times_omega,eps_m,eps_d)[0])
    Ex_array = Ex_func(X, Z)
    Ez_func = np.vectorize(lambda x,z : Efield(x,z,t_times_omega,eps_m,eps_d)[1])
    Ez_array = Ez_func(X, Z)
    
    # Open a new figure with correct aspect ratio and pixel count and white background
    fig = plt.figure(figsize = (fig_width_inches,fig_height_inches), dpi=dpi, facecolor='w')
    
    # Draw a new set of axes that fill the entire figure area
    ax=fig.add_axes((0,0,1,1),axisbg='w')
    ax.set_axis_off()
    
    # Color the metal part
    metal_color = '#dddddd' #light gray
    ax.add_patch(draw_box(xmin,xmax,zmin,0,metal_color))
    
    # Draw the arrows
    ax.quiver(X, Z, Ex_array , Ez_array , scale=arrow_len_scale, scale_units='width', pivot='middle')
    
    ax.set_xlim(xmin, xmax)
    ax.set_ylim(zmin, zmax)
    if img_filename is not None:
        fig.savefig(os.path.join(directory_now, img_filename), dpi=dpi)
 
def draw_anim(eps_m, eps_d, anim_filename='anim.gif', frames_in_anim=30,
              total_anim_time_in_sec=2, keep_frame_images=False, **kwargs):
    """
    Create an animated GIF. **kwargs are all the keyword arguments to
    draw_frame()
    
    keep_frame_images=True to save the individual frame image files that make
    up the animation; otherwise they are created and immediately deleted.
    """
    filename_list = ['temp' + str(n) + '.png' for n in range(frames_in_anim)]
    
    for n in range(frames_in_anim):
        draw_frame(2*pi*n/frames_in_anim, eps_m, eps_d,
                   img_filename=filename_list[n], **kwargs)
    
    seconds_per_frame = total_anim_time_in_sec / frames_in_anim
    frame_delay = str(seconds_per_frame * 100)
    command_list = ['convert', '-delay', frame_delay, '-loop', '0'] + filename_list + [anim_filename]
    # Use the "convert" command (part of ImageMagick) to build the animation
    subprocess.call(command_list, cwd=directory_now)
    if keep_frame_images is False:
        for filename in filename_list:
            os.remove(os.path.join(directory_now, filename))

###################################################################

if True:
    # Silver-air interface at 370nm, using data from Palik
    eps_m = -2.62 + 0.626j
    eps_d = 1
    # Print diagnostics
    Efield(0, 0, 0, eps_m, eps_d, print_wave_properties=True)
    # Create animation
    draw_anim(eps_m, eps_d, anim_filename='Silver_370nm_Palik.gif')