#!/usr/bin/env python3

#
# A short python script to demonstrate electromagnetic wave polarization.
#
#
# Xin Tao, 2017/04/13

import numpy as np
from numpy import pi
from matplotlib import pyplot as plt
from matplotlib import animation
from mpl_toolkits.axes_grid.axislines import SubplotZero

def func(i):
    t = i*dt
    phi_y = -w*t+phi_y0
    phi_z = -w*t+phi_z0
    
    ey = [0, ey0*np.cos(phi_y)]
    ez = [0, ez0*np.cos(phi_z)]

    return ey, ez

# initialization function: plot the background of each frame
def init():
    for i in range(0,100,1):
        ey, ez = func(i)
        ax.plot(ey, ez, 'bs', ms=1.5)

    line.set_data([], [])
    return line,

# animation function.  This is called sequentially
def animate(i):
    ey, ez = func(i)

    line.set_data(ey, ez)
    line.set_marker('o')
    line.set_color('k')
    
    return line,

if __name__ == '__main__':

    # Set up two simulation parameters here

    ez0_2_ey0 = 1.0  # ratio of ez/ey
    dphi = pi/2   # dphi = phi_z0 - phi_y0

    # Start simulation  
    ey0 = 1
    ez0 = ey0*ez0_2_ey0

    e0 = np.sqrt(ey0**2+ez0**2)
    e0_max = e0*1.5

    phi_y0 = np.random.random()*2*pi # select a random initial y phase 
    phi_z0 = phi_y0+dphi

    w=2*pi
    dt = 0.01

    # Make animation
    animation_interval=10

    plt.close()
    
    fig = plt.figure(figsize=(6,6))

    ax = SubplotZero(fig, 111)
    fig.add_subplot(ax)

    for direction in ["xzero", "yzero"]:
        ax.axis[direction].set_axisline_style("-|>")
        ax.axis[direction].set_visible(True)

    for direction in ["left", "right", "bottom", "top"]:
        ax.axis[direction].set_visible(False)

    axis_lim = np.array([-1,1])*e0_max   
    ax.set_xlim(axis_lim)
    ax.set_ylim(axis_lim)

    ax.set_xticks([-int(e0_max), int(e0_max)])
    ax.set_yticks([-int(e0_max), int(e0_max)])

    line, = ax.plot([], [], lw=2)

    anim = animation.FuncAnimation(fig, animate, init_func=init,
                                   frames=2000, interval=animation_interval)

    plt.show()