import plotly
plotly.__version__

import numpy as np
import plotly.graph_objs as go

"""
Explanation: Electromagnetic wave animation
An electromagnetic wave is a combination of electric and magnetic fields that vibrate together in space and time, in a synchronous way. Suppose that the end points of the  electric field vectors describe, during the field oscillation, the wave $E_z=A\sin(y/\lambda-\omega t), x=0$, while the end points of the magnetic field vectors run the wave $E_x=A\sin(y/\lambda-\omega t), z=0$. 
To simulate/animate the electromagnetic wave propagation we plot at each mpmemt t in  the time interval of simulation,  the electric field and the magnetic field vectors.
End of explanation
"""


def rot_matrix(theta):
    return np.array([[np.cos(theta), -np.sin(theta)], 
                     [np.sin(theta), np.cos(theta)]])

def get_arrows(start, end, arrow_angle,  plane=2, fract=0.1):
    """
    this i  function defines 3d vectors/quivers
    - start -numpy array of x, y, z-coordinates of the arrow starting points; shape (3, m)
      start[0,:] contains x-coordinates, etc
    - end - numpy array with the same shape as start; contains on rows the x,  y and z-coords 
      of ending points of the arrow  
    - the arrowhead is an isosceles triangle with the equal sides forming an angle of 2*arrow_angle radians 
    - plane=0 or 2 depending on the plane where the vectors are drawn (plane=0 i.e. x=0, plane=2, z=0)
    """
    start = np.asarray(start)
    end = np.asarray(end)
    m = start[0,:].shape[0]
    arr_dir = start-end
    arr_dir_norm = np.linalg.norm(arr_dir, axis=0)
    arr_dir = fract*arr_dir/arr_dir_norm[None,:] # the arrowhead is a fraction fract from the unit vector
    if plane == 2: 
        v = np.einsum('ji, im -> jm', rot_matrix(arrow_angle), arr_dir[:plane,:]) # Einstein summation                                                                               # rotation to all vectors at a time
        w = np.einsum('ji, im -> jm', rot_matrix(-arrow_angle), arr_dir[:plane, :])
        v = np.append(v, [[0]*m], axis=0) 
        w = np.append(w, [[0]*m], axis=0) 
    elif plane == 0:
        v = np.einsum('ji, im -> jm',  rot_matrix(arrow_angle), arr_dir[1:,:])                         
        w = np.einsum('ji, im -> jm', rot_matrix(-arrow_angle), arr_dir[1:, :])
        v = np.append([[0]*m], v, axis=0)
        w = np.append([[0]*m], w, axis=0)
    else: raise ValueError('the key plane must be 0 or 2')    
    p = end+v
    q = end+w

    suppx = np.stack((start[0,:], end[0,:], np.nan*np.ones(m ))) #supp is the line segment as support for arrow 
    suppy = np.stack((start[1,:], end[1,:], np.nan*np.ones(m )))
    suppz = np.stack((start[2,:], end[2,:], np.nan*np.ones(m )))
    x = suppx.flatten('F')#Fortran type flattening
    y = suppy.flatten('F')
    z = suppz.flatten('F')
    x = list(map(lambda u: None if np.isnan(u) else u, x))
    y = list(map(lambda u: None if np.isnan(u) else u, y))
    z = list(map(lambda u: None if np.isnan(u) else u, z))
    
    #headx, heady, headz are the x, y, z coordinates of the triangle vertices
    headx = np.stack((end[0,:], p[0,:], q[0,:], end[0,:], np.nan*np.ones(m)))
    heady = np.stack((end[1,:], p[1,:], q[1,:], end[1,:], np.nan*np.ones(m)))
    headz = np.stack((end[2,:], p[2,:], q[2,:], end[2,:], np.nan*np.ones(m)))               
    xx = headx.flatten('F')
    yy = heady.flatten('F')
    zz = headz.flatten('F')               
    xx = list(map(lambda u: None if np.isnan(u) else u, xx))
    yy = list(map(lambda u: None if np.isnan(u) else u, yy))
    zz = list(map(lambda u: None if np.isnan(u) else u, zz))               
    
    return x, y, z, xx, yy, zz
    

"""
Explanation: Define a function that returns data to draw  vectors of a vector field with values in the plane yOz or xOy (i.e.  electric or magnetic field vectors, in our case).
End of explanation
"""


a = 2
b = 5
 
xblue = [-a, a, a , -a, -a]
yblue = [-b, -b, b, b, -b]
zblue = [0]*5

xred = [0]*5+[None, 0, 0, 0]
yred = [-b, -b, b, b, -b, None, -b, b]
zred = [a, -a, -a, a, a, None, 0, 0]

x_Oy = [0, 0]
y_Oy = [-b, b]
z_Oy = [0, 0]


"""
Explanation: Define data for fixed traces in each frame (tr0, tr1, tr2, described below), representing the two orthogonal planes of the electric and magnetic field, and their common line:
End of explanation
"""


A = 1 # wave amplitude
lambdA = 0.5 # wavelength
omega = 1 # angular frequency
t = np.arange(0., 10., 0.2)# interval of simulation
Y = np.arange(-b, b, 0.2) # a grid of an interval on Oy, where the vector fields are evaluated
X = np.zeros(Y.shape[0])
ZZe = np.zeros(Y.shape[0])

nr_frames = t.shape[0]
theta = np.pi/13 # the characteristic angle of each arrow
start = np.stack((X, Y, np.zeros(X.shape))) # the numpy array of starting points of the the two classes of vectors
start.shape

"""
Explanation: Set the wave parameters and the interval of simulation:
End of explanation
"""


Ze = A*np.sin(Y/lambdA-omega*t[0])
end1 = np.stack((X, Y, Ze))
x1, y1, z1, xx1, yy1, zz1 = get_arrows(start, end1, theta, plane=0)
XXe = A*np.sin(Y/lambdA-omega*t[0])
end2 = np.stack((XXe, Y, ZZe))
x2, y2, z2, xx2, yy2, zz2 = get_arrows(start, end2, theta, plane=2)

tr0 = dict(type='scatter3d',  # a rectangle in xOy  
         x=xblue,
         y=yblue,
         z=zblue, 
         mode='lines',
         line=dict(width=1.5, color='blue'))
tr1 = dict(type='scatter3d',# a rectangle in yOz
         x=xred,
         y=yred,
         z=zred, 
         mode='lines',
         line=dict(width=1.5, color='red'))
tr2 = dict(type='scatter3d',#line of direction Oy
         x=x_Oy,
         y=y_Oy,
         z=z_Oy, 
         mode='lines',
         line=dict(width=1.5, color='rgb(140,140,140)'))      

"""
Explanation: Define data representing the vectors of the two vector fields:
End of explanation
"""


tr3 = dict(
         type='scatter3d',
         x=x1,
         y=y1,
         z=z1,
         mode='lines', 
         line=dict(color='red', width=1.5),
         name=''
        )
                                   
tr4 = dict(
         type='scatter3d',
         x=xx1,
         y=yy1,
         z=zz1,      
         mode='lines', 
         line=dict(color='red', width=2), 
         name='' 
        )
tr5 = dict(
         type='scatter3d',
         x=x2,
         y=y2,
         z=z2,
         mode='lines', 
         line=dict(color='blue', width=1.5),
         name=''
        )
tr6 = dict(
         type='scatter3d',
         x=xx2,
         y=yy2,
         z=zz2,  
         mode='lines', 
         line=dict(color='blue', width=2), 
         name=''
        )

"""
Explanation: The following four traces are the base traces updated by the animation frames:
End of explanation
"""


data = [tr0, tr1, tr2, tr3, tr4, tr5, tr6]

frames=[]
for k in range(nr_frames):
    Ze = A*np.sin(Y/lambdA-omega*t[k])
    end1 = np.stack((X, Y, Ze))
    x1, y1, z1, xx1, yy1, zz1 = get_arrows(start, end1, theta, plane=0)
    XXe = A*np.sin(Y/lambdA-omega*t[k])
    end2 = np.stack((XXe, Y, ZZe))
    x2, y2, z2, xx2, yy2, zz2 = get_arrows(start, end2, theta, plane=2)
    frames += [dict(data=[dict(type='scatter3d', 
                               x=x1,
                               y=y1,
                               z=z1),
                            
                          dict(type='scatter3d',
                              x=xx1,
                              y=yy1,
                              z=zz1),
                          dict(type='scatter3d',
                               x=x2, 
                               y=y2, 
                               z=z2),
                          dict(type='scatter3d',
                               x=xx2,
                               y=yy2,
                               z=zz2)],
                  traces=[3, 4, 5, 6]
                 )]
 

"""
Explanation: Define data to be plotted in each animation frame:
End of explanation
"""


title='Electromagnetic wave propagating in the positive Oy direction<br>'+\
'The  electric field vectors (red) are included in the yz-plane,<br> and  the magnetic field vectors (blue), in xy'
layout = dict(title=title,
              font=dict(family='Balto'),
              autosize=False,
              width=700,
              height=700,
              showlegend=False,
              scene=dict(camera = dict(eye=dict(x=1.22, y=0.55, z=0.3)),
                         aspectratio=dict(x=1, y=1, z=0.65),
                         xaxis_visible=False,
                         yaxis_visible=False, 
                         zaxis_visible=False,
                        ),
              updatemenus=[dict(type='buttons', showactive=False,
                                y=0.75,
                                x=1.05,
                                xanchor='left',
                                yanchor='top',
                                pad=dict(t=0, l=10),
                                buttons=[dict(label='Play',
                                              method='animate',
                                              args=[None, 
                                                    dict(frame=dict(duration=80, 
                                                                    redraw=True),
                                                         transition=dict(duration=0),
                                                         fromcurrent=True,
                                                         mode='immediate'
                                                        )
                                                   ]
                                             )
                                        ]
                               )
                          ]     
          )


import chart_studio.plotly as py
fig = go.Figure(data=data, layout=layout, frames=frames)
py.iplot(fig, filename='anim-electromagwave')

from IPython.core.display import HTML
def  css_styling():
    styles = open("./custom.css", "r").read()
    return HTML(styles)
css_styling()

"""
Explanation: Set the plot layout:
End of explanation
"""


