from rayflare.transfer_matrix_method import tmm_structure
from rayflare.options import default_options
import numpy as np
from solcore import material
import matplotlib.pyplot as plt
from inkstone import Inkstone

glass = material('BK7')()
GaAs = material('GaAs')()
Si = material('Si')()
Air = material("Air")()

# not actually doing a patterned structure, but need to set something for the lattice vectors
lat = ((1, 0), (0, 1))

num_g = 1 # 1 order only because no diffraction in planar structure
wl = 535 # free space wavelength (nm)

glass_nk = glass.n(wl*1e-9) + 1j*glass.k(wl*1e-9)
gaas_nk = GaAs.n(wl*1e-9) + 1j*GaAs.k(wl*1e-9)
si_nk =  Si.n(wl*1e-9) + 1j*Si.k(wl*1e-9)

widths = [np.inf, 50, np.inf] # RayFlare uses inf for incidence and transmission medium

# Incidence medium = glass
layers_oc = np.array([
    glass_nk**2       ,
    gaas_nk**2,
    si_nk**2
])

tmm_stack = [[widths[i1+1], np.array([wl, wl+1]),
              np.array([np.real(np.sqrt(layers_oc)[i1+1])]*2),
              np.array([np.imag(np.sqrt(layers_oc)[i1+1])]*2)] for i1, _ in enumerate(widths[1:-1])]

options = default_options()

options.wavelengths = np.array([wl*1e-9])
options.pol = 's' # this is p-polarization in Inkstone but s-polarization in RayFlare/S4

tmm_str = tmm_structure(tmm_stack, glass, Si)
geom_list = [None, None, None] # no patterned layers

s = Inkstone(lat, num_g)
s.frequency = 1./wl

for i1, _ in enumerate(widths):
    s.AddMaterial(name="layer_" + str(i1 + 1), epsilon=layers_oc[i1])

for i1, wid in enumerate(widths):  # set base layers
    layer_name = "layer_" + str(i1 + 1)

    if wid == float("Inf"):
        s.AddLayer(layer_name, 0, layer_name)

    else:
        s.AddLayer(layer_name, wid, layer_name)


# angle sweep
angle = np.linspace(0, 89.9, 50)

R = np.zeros((len(angle), 2))
T = np.zeros((len(angle), 2))

A_layer = np.zeros((len(angle), len(widths)-2))

A_layer_tmm = np.zeros((len(angle), len(widths)-2))

n_inc = np.real(np.sqrt(layers_oc[0]))

for i1, a in enumerate(angle):
    # inkstone sweep

    # Inkstone calculation
    s.SetExcitation(theta=a, phi=0., s_amplitude=0, p_amplitude=1)
    trapf, trapb = s.GetPowerFlux('layer_' + str(len(widths)))
    incpf, incpb = s.GetPowerFlux('layer_1') # Air power forward and backward

    R[i1, 0] = -incpb/incpf
    T[i1, 0] = trapf/incpf

    A_layer[i1, 0] = (np.sum(s.GetPowerFlux('layer_1')) - np.sum(s.GetPowerFlux('layer_3')))/incpf

    for j1 in np.arange(len(widths)-3) + 3:

        A_layer[i1,j1-2] = (- np.sum(s.GetPowerFlux('layer_' + str(j1+1))) +
                            np.sum(s.GetPowerFlux('layer_' + str(j1))))/incpf

    # TMM (RayFlare) calculation

    options.theta_in = a*np.pi/180
    RAT_tmm = tmm_str.calculate(options)
    R[i1, 1] = RAT_tmm['R']
    T[i1, 1] = RAT_tmm['T']
    A_layer_tmm[i1] = RAT_tmm['A_per_layer']

# Incidence medium = Air

layers_oc = np.array([
    1,
    gaas_nk**2,
    si_nk**2,
])

tmm_stack = [[widths[i1+1], np.array([wl, wl+1]),
              np.array([np.real(np.sqrt(layers_oc)[i1+1])]*2),
              np.array([np.imag(np.sqrt(layers_oc)[i1+1])]*2)] for i1, _ in enumerate(widths[1:-1])]

tmm_str = tmm_structure(tmm_stack, Air, Si)

s = Inkstone(lat, num_g)
s.frequency = 1./wl

for i1, _ in enumerate(widths):
    s.AddMaterial(name="layer_" + str(i1 + 1), epsilon=layers_oc[i1])

for i1, wid in enumerate(widths):  # set base layers
    layer_name = "layer_" + str(i1 + 1)

    if wid == float("Inf"):
        s.AddLayer(layer_name, 0, layer_name)

    else:
        s.AddLayer(layer_name, wid, layer_name)


R_air = np.zeros((len(angle), 2))
T_air = np.zeros((len(angle), 2))

A_layer_air = np.zeros((len(angle), len(widths)-2))

A_layer_tmm_air = np.zeros((len(angle), len(widths)-2))

n_inc = np.real(np.sqrt(layers_oc[0]))

for i1, a in enumerate(angle):
    # inkstone sweep
    options.theta_in = a*np.pi/180

    s.SetExcitation(theta=a, phi=0., s_amplitude=0, p_amplitude=1)
    trapf, trapb = s.GetPowerFlux('layer_' + str(len(widths)))
    incpf, incpb = s.GetPowerFlux('layer_1') # Air power forward and backward

    R_air[i1, 0] = -incpb/incpf
    T_air[i1, 0] = trapf/incpf

    A_layer_air[i1, 0] = (np.sum(s.GetPowerFlux('layer_1')) - np.sum(s.GetPowerFlux('layer_3')))/incpf

    for j1 in np.arange(len(widths)-3) + 3:
        A_layer_air[i1,j1-2] = (- np.sum(s.GetPowerFlux('layer_' + str(j1+1))) +
                            np.sum(s.GetPowerFlux('layer_' + str(j1))))/incpf

    RAT_tmm_air = tmm_str.calculate(options)
    R_air[i1, 1] = RAT_tmm_air['R']
    T_air[i1, 1] = RAT_tmm_air['T']
    A_layer_tmm_air[i1] = RAT_tmm_air['A_per_layer']


plt.figure()
plt.plot(angle, R[:,0], '-k', label='R (Inkstone)')
plt.plot(angle, R[:,1], '-r', label='R (TMM)')
plt.plot(angle, T[:,0], '--k', label='T (Inkstone)')
plt.plot(angle, T[:,1], '--r', label='T (TMM)')
plt.plot(angle, A_layer, '-.k', label='A (Inkstone)')
plt.plot(angle, A_layer_tmm, '-.r', label='A (TMM)')
plt.title(r'$n_{inc}$ = 1.52')
plt.legend()
plt.xlabel('Incidence angle (degrees)')
plt.ylabel('R / A / T (fraction of incident power)')
plt.tight_layout()
plt.show()

plt.figure()
plt.plot(angle, R_air[:,0], '-k', label='R (Inkstone)')
plt.plot(angle, R_air[:,1], '-r', label='R (TMM)')
plt.plot(angle, T_air[:,0], '--k', label='T (Inkstone)')
plt.plot(angle, T_air[:,1], '--r', label='T (TMM)')
plt.plot(angle, A_layer_air, '-.k', label='A (Inkstone)')
plt.plot(angle, A_layer_tmm_air, '-.r', label='A (TMM)')
plt.title(r'$n_{inc}$ = 1')
plt.legend()
plt.xlabel('Incidence angle (degrees)')
plt.ylabel('R / A / T (fraction of incident power)')
plt.tight_layout()
plt.show()


plt.figure()
plt.plot(angle, R[:,0]/R[:,1])
plt.show()