#
# Test that S4 can reproduce the Fresnel coefficients
#
# all lengths normalized in terms of a
# frequency normalized in terms of 2*pi*c/a

# this test file plots a frequency spectrum and field plot for the python extension

import S4
# https://web.stanford.edu/group/fan/S4/ # website
# https://github.com/victorliu/S4 # Original author
# https://github.com/phoebe-p/S4 # installed version
import numpy as np
import matplotlib.pyplot as plt

from sympy.physics.optics import Medium, fresnel_coefficients
# https://docs.sympy.org/latest/modules/physics/optics/utils.html
# https://en.wikipedia.org/wiki/Fresnel_equations#Complex_amplitude_reflection_and_transmission_coefficients
# https://en.wikipedia.org/wiki/Fresnel_equations#Alternative_forms
from math import pi, cos, asin, sin

def Tintsymp(t,angle_of_incidence,n1,n2):
    angle_of_refraction = asin(n1*sin(angle_of_incidence)/n2)
    return (n2*cos(angle_of_refraction))/(n1*cos(angle_of_incidence)) *abs(t)**2


# SymPy commands
n_glass = 1.5
n_vac = 1
AirMed = Medium('Air',n=n_vac)
GlassMed = Medium('Glass',n=n_glass)

# S4 simulation setup
s = S4.New(Lattice=1,NumBasis=100)
s.SetMaterial(Name='glass',Epsilon=(n_glass**2))
s.SetMaterial(Name='vac',Epsilon=(n_vac**2))
s.AddLayer(Name='air',Thickness=0,Material='vac')
s.AddLayer(Name='gla',Thickness=0,Material='glass')
s.SetFrequency(1./0.5)


# angle sweep
I = []
R = []
T = []
Rs = []
Ts = []
angle = np.linspace(0, 90, 91)


for a in angle:
    # S4 sweep
    s.SetExcitationPlanewave(IncidenceAngles=(a,0),sAmplitude=1,pAmplitude=0,Order=0)
    glapf, glapb = s.GetPowerFlux('gla')
    airpf, airpb = s.GetPowerFlux('air') # Air power forward and backward

    I.append(airpf)
    R.append(abs(airpb/airpf))
    T.append(abs(glapf/airpf))

    # SymPy sweep
    a_rad = a*pi/180.
    fres = fresnel_coefficients(a_rad,AirMed,GlassMed)
    Rs.append(abs(fres[1])**2)
    Ts.append(Tintsymp(fres[3],a_rad,n_vac,n_glass))


#%% plotting
plt.figure()
plt.plot(angle, R, label='R (RCWA)')
plt.plot(angle, T, label='T (RCWA)')
plt.plot(angle, Rs, label='R (Fresnel)')
plt.plot(angle, Ts, label = 'T (Fresnel)')
plt.xlabel('Angle (degrees)')
plt.ylabel('R/T')
plt.title('R and T, s/TE, via RCWA and via the Fresnel equations')
plt.legend()
plt.savefig('Fresnel_S4_TE.png')