#!/usr/bin/env python
# coding: utf-8

# In[1]:


import numpy as np
import matplotlib.pyplot as plt
from matplotlib import animation


# ### **Task 1**
# Calculate transfer matrix \\(\hat{M}\\) of the medium.

# In[2]:


def transfermatrix(thickness, epsilon, polarisation, wavelength, kz):
    '''Computes the transfer matrix of a given stratified medium.
    
    Parameters
    ----------
    thickness : 1d-array
        Thicknesses of the layers in µm.
    epsilon : 1d-array
        Relative dielectric permittivity of the layers.
    polarisation : str
        Polarisation of the computed field, either 'TE' or 'TM'.
    wavelength : float
        The wavelength of the incident light in µm.
    kz : float
        Transverse wavevector in 1/µm.

    Returns
    -------
    M : 2d-array
        The transfer matrix of the stratified media.
    '''
    # initialize system transfer matrix to identity matrix
    M = np.eye(2, dtype=np.complex128)
    # determine q, depending on the polarization
    if polarisation == 'TE':
        q = np.ones_like(epsilon)
    elif polarisation == 'TM':
        q = 1.0/epsilon
    else:
        raise ValueError('Invalid input: '
                         'polarisation must either be "TE" or "TM"')
    # vacuum wavenumber
    k0 = 2.0*np.pi/wavelength
    # wave vector component normal to layer stack for each layer
    kx = np.sqrt(k0**2*epsilon - kz**2, dtype=np.complex128)

    # iterate over layers
    for kxi, di, qi in zip(kx, thickness, q):
        c = np.cos(kxi*di)
        s = np.sin(kxi*di)
        ka = kxi*qi
        # matrix for ith layer
        m = np.array([[    c, s/ka],
                      [-ka*s,    c]])
        # prepend new layer transfer matrix to system transfer matrix
        M = np.matmul(m,M)

    return M


# ### **Task 2**
# Determine reflection and transmission coefficients as a function of wavelength.

# In[3]:


def spectrum(thickness, epsilon, polarisation, wavelength, angle_inc,
             n_in, n_out):
    '''Computes the reflection and transmission of a stratified medium.

    Parameters
    ----------
    thickness : 1d-array
        Thicknesses of the layers in µm.
    epsilon : 1d-array
        Relative dielectric permittivity of the layers.
    polarisation : str
        Polarisation of the computed field, either 'TE' or 'TM'.
    wavelength : 1d-array
        The wavelength of the incident light in µm.
    angle_inc : float
        The angle of incidence in degree (not radian!).
    n_in, n_out : float
        The refractive indices of the input and output layers.

    Returns
    -------
    T : 1d-array
        Transmitted amplitude
    R : 1d-array
        Reflected amplitude
    tau : 1d-array
        Transmitted energy
    rho : 1d-array
        Reflected energy
    '''

    # Definition of the paramters of the input and output layers
    epsilon_in = n_in**2
    epsilon_out = n_out**2

    if polarisation == 'TE':
        q_in = 1
        q_out = 1
    elif polarisation == 'TM':
        q_in = 1/epsilon_in
        q_out = 1/epsilon_out
    else:
        raise ValueError('Invalid input: '
                          'polarisation must either be "TE" or "TM"')

    k0 = 2.0*np.pi/wavelength # vacuum wavenumber
    kz = k0*n_in*np.sin(np.deg2rad(angle_inc))
    kx_in = np.sqrt(epsilon_in*k0**2 - kz**2, dtype=np.complex128)
    kx_out = np.sqrt(epsilon_out*k0**2 - kz**2, dtype=np.complex128)

    R = np.zeros(wavelength.shape, dtype=np.complex128)
    N = np.zeros(wavelength.shape, dtype=np.complex128)

    # iterate over wavelengths
    for i, (lami, kzi, kx_outi, kx_ini) in enumerate(zip(wavelength, kz, kx_out, kx_in)):
        M = transfermatrix(thickness, epsilon, polarisation, lami, kzi)
        N[i] = (q_in*kx_ini*M[1,1] + q_out*kx_outi*M[0,0]
              + 1j*(M[1,0] - q_in*kx_ini*q_out*kx_outi*M[0,1]))
        R[i] = (q_in*kx_ini*M[1,1] - q_out*kx_outi*M[0,0]
              - 1j*(M[1,0] + q_in*kx_ini*q_out*kx_outi*M[0,1]))
    R /= N

    # calculate remaining coefficients
    T = 2.0*q_in*kx_in/N
    rho = np.real(R*np.conj(R))
    tau = np.real(q_out*kx_out)/np.real(q_in*kx_in)*np.real(T*np.conj(T))

    return T, R, tau, rho


# ### **Task 3**
# Computing fields everywhere in space.

# In[4]:


def field(thickness, epsilon, polarisation, wavelength, kz, n_in, n_out,
          Nx, l_in, l_out):
    '''Computes the field inside a stratified medium.

    The medium starts at x = 0 on the entrance side. The transmitted field
    has a magnitude of unity.

    Parameters
    ----------
    thickness : 1d-array
        Thicknesses of the layers in µm.
    epsilon : 1d-array
        Relative dielectric permittivity of the layers.
    polarisation : str
        Polarisation of the computed field, either 'TE' or 'TM'.
    wavelength : float
        The wavelength of the incident light in µm.
    kz : float
        Transverse wavevector in 1/µm.
    n_in, n_out : float
        The refractive indices of the input and output layers.
    Nx : int
        Number of points where the field will be computed.
    l_in, l_out : float
        Additional thickness of the input and output layers where the field
        will be computed.

    Returns
    -------
    f : 1d-array
        Field structure
    index : 1d-array
        Refractive index distribution
    x : 1d-array
        Spatial coordinates
    '''
    # Input layer for x < 0; and output layer for x > 0; illumination from the input side
    epsilon_in = n_in**2
    epsilon_out = n_out**2

    # extension of the vectors epsilon and thickness to take the input
    #  and output layers into account
    thickness = np.concatenate(([l_in], thickness, [l_out]))
    epsilon = np.concatenate(([epsilon_in], epsilon, [epsilon_out]))

    # flip layers (calculation proceeds backwars from the transmitted field)
    epsilon = epsilon[::-1]
    thickness = thickness[::-1]

    # determine q, depending on the polarization
    if polarisation == 'TE':
        q = np.ones_like(epsilon)
    elif polarisation == 'TM':
        q = 1.0/epsilon
    else:
        raise ValueError('Invalid input: '
                         'polarisation must either be "TE" or "TM"')

    # vacuum wavenumber
    k0 = 2.0*np.pi/wavelength
    # wave vector component normal to layer stack
    kx = np.sqrt(epsilon*k0**2 - kz**2, dtype=np.complex128)

    # output layer parameters
    q_out = q[0]
    kx_out = kx[0]

    # further computation starts from the transmitted field because the fields
    # are calculated from the back
    incident_vec = np.array([[1.0], [-1.0j*q_out*kx_out]])

    # definition of output positions
    x = np.linspace(0, np.sum(thickness), Nx)

    curr_layer = 0 #current ith layer
    pos_in_layer = 0.0
    thickness_below = 0.0 
    M = np.eye(2, dtype=np.complex128)
    f = np.zeros(x.shape, dtype=np.complex128)
    index = np.zeros(x.shape, dtype=epsilon.dtype)

    for i, xi in enumerate(x):
        # get postion within current layer
        pos_in_layer = xi - thickness_below

        # check if a layer interface has been crossed
        if pos_in_layer > thickness[curr_layer]:
            # propagate until layer interface
            c = np.cos(kx[curr_layer]*thickness[curr_layer])
            s = np.sin(kx[curr_layer]*thickness[curr_layer])
            ka = kx[curr_layer]*q[curr_layer]
            m = np.array([[    c, s/ka],
                          [-ka*s,    c]])
            # update transfer matrix
            M = m@M
            # update state variables
            pos_in_layer = pos_in_layer - thickness[curr_layer]
            thickness_below = thickness_below + thickness[curr_layer]
            curr_layer += 1

        # propagate within layer to current position
        c = np.cos(kx[curr_layer]*pos_in_layer)
        s = np.sin(kx[curr_layer]*pos_in_layer)
        ka = kx[curr_layer]*q[curr_layer]
        m = np.array([[    c, s/ka],
                      [-ka*s,    c]])
        out_vector = m@M@incident_vec
        f[i] = out_vector[0,0]
        index[i] = np.sqrt(epsilon[curr_layer])

    # at the end, the fields have to be flipped
    f = f[::-1]
    index = index[::-1]

    return f, index, x


# In[35]:


def timeanimation(x, f, index, steps, periods):
    ''' Animation of a quasi-stationary field.

    Parameters
    ----------
    x : 1d-array
        Spatial coordinates
    f : 1d-array
        Field
    index : 1d-array
        Refractive index
    steps : int
        Total number of time points
    periods : int
        Number of the oscillation periods.

    Returns
    -------
    ani : matplotlib.animation.FuncAnimation
        The time animation of the field
    '''
    # based on https://matplotlib.org/gallery/animation/simple_anim.html

    freq = periods/(steps - 1)
    max_f = np.abs(f).max()
    max_index = index.real.max()

    # helper function to calculate field at step n
    def field_at_step(n):
        return np.real(f*np.exp(-2.0j*np.pi*freq*n))/max_f*max_index

    # set up initial plot
    fig = plt.figure()
    line_f, line_eps = plt.plot(x, field_at_step(0),
                                x, index.real)
    plt.xlabel('x [µm]')
    plt.ylabel('normalized field (real part)')
    plt.legend(['EM field', 'refr. index'], loc='lower right', frameon=False)
    plt.xlim(x[[0,-1]])
    plt.ylim(np.array([-1.1, 1.1])*max_index)

    # function that updates plot data during animation
    def animate(i):
        line_f.set_ydata(field_at_step(i))
        return (line_f,)

    # function that inits plot data for animation (clean state for blitting)
#    def init():
#        line_f.set_ydata(x*np.nan)
#        return (line_f,)

    ani = animation.FuncAnimation(fig, animate, #init_func=init,  blit=True,
                                  save_count=steps, interval=100)

    plt.show()
    return ani


# In[5]:


def bragg(n1, n2, d1, d2, N):
    '''Generates the stack parameters of a Bragg mirror
    The Bragg mirror starts at the incidence side with layer 1 and is
    terminated by layer 2

    Parameters
    ----------
    n1, n2 : float or complex
        Refractive indices of the layers of one period
    d1, d2 : float
        Thicknesses of layers of one period
    N  : in
        Number of periods

    Returns
    -------
    epsilon : 1d-array
        Vector containing the permittivities
    thickness : 1d-array
        Vector containing the thicknesses
    '''

    # find suitable type for epsilon
    dt = np.common_type(np.array([n1]), np.array([n2]))
    epsilon = np.zeros((2*N,), dtype = dt)
    epsilon[::2] = n1**2
    epsilon[1::2] = n2**2
    thickness = np.zeros((2*N,))
    thickness[::2] = d1
    thickness[1::2] = d2
    return epsilon, thickness


# In[6]:


# script for using the above functions
plt.rcParams.update({
        'figure.figsize': (12/2.54, 9/2.54),
        'figure.subplot.bottom': 0.15,
        'figure.subplot.left': 0.15,
        'figure.subplot.right': 0.95,
        'figure.subplot.top': 0.9,
        'axes.grid': True,
})
plt.close('all')
save_figures = False

# %% task 1: transfer matrix %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
n1 = np.sqrt(2.25)
n2 = np.sqrt(15.21)
d1 = 0.13
d2 = 0.05
N  = 5
epsilon, thickness = bragg(n1, n2,d1, d2, N)

wavelength = 0.78
kz = 0.0
polarisation = 'TE'
M = transfermatrix(thickness, epsilon, polarisation, wavelength, kz)
print('M = {0}'.format(M))
print('det(M) = {0}'.format(np.linalg.det(M)))
print('eig(M) = {0}'.format(np.linalg.eig(M)[0]))

wavelength = 1.2
kz = 0.0
polarisation = 'TE'
M = transfermatrix(thickness, epsilon, polarisation, wavelength, kz)
print('M = {0}'.format(M))
print('det(M) = {0}'.format(np.linalg.det(M)))
print('eig(M) = {0}'.format(np.linalg.eig(M)[0]))


# %% task 2: reflection and transmission spectrum %%%%%%%%%%%%%%%%%%%%%%%%%%%%%

# Bragg mirror with 5 periods
epsilon, thickness = bragg(n1, n2, d1, d2, N)
n_in = 1
n_out = 1.5
angle_inc = 0
wavelength_vector = np.linspace(0.5, 1.5, 1001)
T, R, tau, rho = spectrum(thickness, epsilon, polarisation,
                          wavelength_vector, angle_inc, n_in, n_out)

## plot reflectance and transmittance
plt.figure()
plt.plot(wavelength_vector, tau,
         wavelength_vector, rho)
plt.xlim(wavelength_vector[[0,-1]])
plt.ylim([0, 1])
plt.xlabel('wavelength [µm]')
plt.ylabel('reflectance, transmittance')
plt.legend(['transmittance', 'reflectance'], loc='center', frameon=False)
#if save_figures:
#    plt.savefig('spectrum_Bragg_mirror.pdf', dpi=300)
#
## plot transmittance on log scale
plt.figure()
plt.semilogy(wavelength_vector, tau)
plt.xlim(wavelength_vector[[0,-1]])
plt.xlabel('wavelength [µm]')
plt.ylabel('transmittance')
plt.legend(['transmittance'], loc='lower right', frameon=False)
plt.show()
#if save_figures:
#    plt.savefig('spectrum_Bragg_mirror_log.pdf', dpi=300)
#
#print('Maximum absorption: {0}'.format(np.abs(1 - tau - rho).max()))


# %% task 3: field distribution %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
# bragg mirror
epsilon, thickness = bragg(n1, n2, d1, d2, N)
l_in = 1.0
l_out = 1.0
Nx = 1000
wavelength = 0.78
kz = 0
f, index, x = field(thickness, epsilon, polarisation, wavelength, kz,
                    n_in ,n_out, Nx, l_in, l_out)

## plot magnitude of field
plt.figure()
plt.plot(x, np.abs(f)/np.abs(f).max()*index.real.max(),
         x, index.real)
plt.xlim(x[[0,-1]])
plt.ylim([0, 1.1*index.real.max()])
plt.xlabel('x [µm]')
plt.ylabel('normalized field (magnitude)')
plt.legend(['EM field', 'refr. index'], loc='lower right', frameon=False)
#if save_figures:
#    plt.savefig('field_Bragg_mirror_magnitude.pdf', dpi=300)
#
## plot real part of field
plt.figure()
plt.plot(x, f.real/np.abs(f).max()*index.real.max(),
         x, index.real)
plt.xlim(x[[0,-1]])
plt.ylim(np.array([-1.1, 1.1])*index.real.max())
plt.xlabel('x [µm]')
plt.ylabel('normalized field (real part)')
plt.legend(['EM field', 'refr. index'], loc='lower right', frameon=False)
plt.show()
#if save_figures:
#    plt.savefig('field_Bragg_mirror_real_part.pdf', dpi=300)

# %% task 4: time animation of field %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
steps = 200
periods = 10
ani = timeanimation(x, f, index, steps, periods)

# requires Ffmpeg
#if save_figures:
#ani.save("field_FP_resonator.mp4")


# In[ ]: