#!/usr/bin/python3
import numpy as np
from scipy.linalg import expm
import matplotlib.pyplot as pl
## =========================================================
def convmat(A,M,N):

    # DETERMINE SIZE OF A
    Nx,Ny = np.shape(A)

    # COMPUTE INDICES OF SPATIAL HARMONICS
    p = np.arange(-M,M+1)
    q = np.arange(-N,N+1)
    P,Q = 2*M+1, 2*N+1

    # COMPUTE FOURIER COEFFICENTS OF A
    A = np.fft.fftshift(np.fft.fft2(A))/(Nx*Ny)
    
    # COMPUTE ARRAY INDICES OF CENTER HARMONIC
    p0 = np.floor(Nx/2)
    q0 = np.floor(Ny/2)
    
    row = (Q-1)*P + (P-1)
    col = (Q-1)*P + (P-1)

    C = np.zeros((row+1,col+1),dtype=np.complex128)

    # COMPUTE CONVOLUTION MATRICES
    for qrow in range(Q):
        for prow in range(P):
            
            row = qrow*P + prow
           
            for qcol in range(Q):
                for pcol in range(P):
                    
                    col = qcol*P + pcol
                    
                    pfft = p[prow] - p[pcol]
                    qfft = q[qrow] - q[qcol]

                    C[row,col] = A[int(p0+pfft),int(q0+qfft)]

    return C
# ========================================================
def rect(Nx,Lx,Ny,Ly,wx,wy,alpha=0):
    # Nx: number of pixels along x
    # Ny: number of pixels along y
    # Lx: Unit cell length along x
    # Ly: Unit cell length along y
    # wx: rectangle width along x
    # wy: rectangle width along y
    # alpha: rotation angle in degrees

    # create grid
    x = np.linspace(-Lx/2, Lx/2, Nx)
    y = np.linspace(-Ly/2, Ly/2, Ny)
    X, Y = np.meshgrid(x,y)

    # rotate rectangle
    ct, st = np.cos(np.deg2rad(alpha)), np.sin(np.deg2rad(alpha))
    Xr =  ct * X + st * Y
    Yr = -st * X + ct * Y

    # create rotated rectangle
    A = ((np.abs(Xr) <= wx/2) & (np.abs(Yr) <= wy/2)).astype(float)

    return A
## =========================================================
##                      DASHBOARD
## =========================================================

er1 = 1         # epsilon in reflection medium
mr1 = 1         # mu in reflection medium
er2 = 1         # epsilon in transmission medium
mr2 = 1         # mu in transmission medium

pol_inc = 'TM'  # input polarizations: TM, TE

wavelength = 1000e-9  # wavelength in [m]

theta = 0     # incidence angle in rad
phi   = 0     # incidence angle in rad

Lx = 800e-9    # physical dimension of the unit cell along x in [m]
Ly = 800e-9    # physical dimension of the unit cell along y in [m]
Lz = [500e-9]  # list containing the thickness of each layer 

M = 6          # Number of DO along x
N = 6          # Number of DO along y

# DEFINE LAYER(S)
r = 250e-9      # circle radius in [m]

Nx = 1024       # number of cells along x
Ny = Nx         # number of cells along y

mr = np.ones((Nx,Ny))

# GRID ARRAYS
dx = Lx/Nx
xa = np.arange(Nx)*dx
xa = xa - np.mean(xa)

dy = Ly/Ny
ya = np.arange(Ny)*dy
ya = ya - np.mean(ya)

# CREATE CIRCLE
Y,X = np.meshgrid(ya,xa)
A = (X**2 + Y**2) <= r**2

er = 12*A + (1-A)

# COMPUTE CONVOLUTION MATRICES
ERc = [convmat(er,M,N)]
MRc = [convmat(mr,M,N)]

## =========================================================
## START COMPUTATION
## =========================================================
def redheffer(SA11,SA12,SA21,SA22,SB11,SB12,SB21,SB22):

    I = np.identity(len(SA11))

    D = SA12.dot(np.linalg.inv(I - SB11.dot(SA22)))
    F = SB21.dot(np.linalg.inv(I - SA22.dot(SB11)))

    SAB11 = SA11 + D.dot(SB11.dot(SA21))
    SAB12 = D.dot(SB12)
    SAB21 = F.dot(SA21)
    SAB22 = SB22 + F.dot(SA22.dot(SB12))

    return SAB11,SAB12,SAB21,SAB22
# ========================================================
# incident normalized wave vector
kinc = np.array([np.sin(theta)*np.cos(phi),\
np.sin(theta)*np.sin(phi), np.cos(theta)])

# vector normal to medium
nvec = np.array([0,0,1])

# definition of polarization vector
if theta == 0:
    ate = np.array([0,1,0])
else:
    ate = np.cross(kinc,nvec)/np.linalg.norm(np.cross(kinc,nvec))

atm  = np.cross(ate,kinc)/np.linalg.norm(np.cross(ate,kinc))

if theta == 0:
    atm0, ate0 = atm, ate
    atm = np.cos(phi)*atm0 - np.sin(phi)*ate0
    ate = np.sin(phi)*atm0 + np.cos(phi)*ate0

if pol_inc == 'TM':
    Pol = atm
elif pol_inc == 'TE':
    Pol = ate

# Create Dirac vector
MN = (2*M+1)*(2*N+1)
Delta = np.zeros((MN,1))
middle = int(np.floor(MN/2))
Delta[middle] = 1

# Compute Source Field
Esrc = np.concatenate((Pol[0]*Delta, Pol[1]*Delta))

# Initialize scattering matrices
One, ID = np.identity(MN), np.identity(MN*2)
SR_0 = np.zeros(((MN)*2,(MN)*2))
ST_0 = ID

# ====================================================
# Number of spatial harmonics
m = np.arange(-M,M+1)
n = np.arange(-N,N+1)

# Generate k-vectors
mr1er1 = mr1*er1
mr2er2 = mr2*er2

kinc  = np.emath.sqrt(mr1er1)*kinc

kx    = kinc[0]-m*wavelength/Lx
ky    = kinc[1]-n*wavelength/Ly

ky,kx = np.meshgrid(ky,kx)

kzRef = np.conj(np.emath.sqrt(np.conj(mr1er1)-kx**2-ky**2))
kzTrn = np.conj(np.emath.sqrt(np.conj(mr2er2)-kx**2-ky**2))
kzGap = np.conj(np.emath.sqrt(1-kx**2-ky**2))

Kx    = np.diag(kx.flatten('F'))
Ky    = np.diag(ky.flatten('F'))
KzRef = np.diag(kzRef.flatten('F'))
KzTrn = np.diag(kzTrn.flatten('F'))
KzGap = np.diag(kzGap.flatten('F'))

KxKy, KyKx, KyKy, KxKx = Kx.dot(Ky), Ky.dot(Kx), Ky.dot(Ky), Kx.dot(Kx)

# Reset scattering matrices
S11d, S22d, S12d, S21d = SR_0, SR_0, ST_0, ST_0

# Compute eigen-modes of gap medium
Zer = np.zeros((MN,MN))
Q0 = np.concatenate((np.concatenate((KxKy, One-KxKx),axis=1),\
                     np.concatenate((KyKy-One, -KyKx),axis=1)))

L0 = np.concatenate((np.concatenate((1j*KzGap, Zer),axis=1),\
                     np.concatenate((Zer, 1j*KzGap),axis=1)))

V0 = Q0.dot(np.linalg.inv(L0))
iV0 = np.linalg.inv(V0)

# Layer loop
for n in range(len(ERc)):
    
    ER, MR = ERc[n], MRc[n]
    iER, iMR = np.linalg.inv(ER), np.linalg.inv(MR)
    KxiER, KyiER, KxiMR, KyiMR = Kx.dot(iER), Ky.dot(iER), Kx.dot(iMR), Ky.dot(iMR)

    Pm = np.concatenate((np.concatenate((KxiER.dot(Ky), MR-KxiER.dot(Kx)),axis=1),\
                         np.concatenate((KyiER.dot(Ky)-MR, -KyiER.dot(Kx)),axis=1)))

    Qm = np.concatenate((np.concatenate((KxiMR.dot(Ky), ER-KxiMR.dot(Kx)),axis=1),\
                         np.concatenate((KyiMR.dot(Ky)-ER, -KyiMR.dot(Kx)),axis=1)))

    omega = Pm.dot(Qm)
    
    eV, W = np.linalg.eig(omega)
    eV = np.diag(np.emath.sqrt(eV))
    V = Qm.dot(W).dot(np.linalg.inv(eV))
    iW, iV = np.linalg.inv(W), np.linalg.inv(V).dot(V0)

    A = iW + iV
    B = iW - iV

    X = expm(-eV*Lz[n]*2*np.pi/wavelength)

    iA = np.linalg.inv(A)
    Num = A - X.dot(B).dot(iA).dot(X).dot(B)

    S11 = np.linalg.solve(Num, X.dot(B).dot(iA).dot(X).dot(A) - B)
    S12 = np.linalg.solve(Num, X.dot(A - B.dot(iA).dot(B)))
    S21 = S12
    S22 = S11

    # Compute device scattering matrices
    S11d,S12d,S21d,S22d = redheffer(S11d,S12d,S21d,S22d,S11,S12,S21,S22)

# ================================================== 
# Compute reflection and transmission sides connection S-matrices
# Reflection side
mr1er1O = mr1er1*One
Qref = 1/mr1*np.concatenate((np.concatenate((KxKy, mr1er1O-KxKx),axis=1),\
                             np.concatenate((KyKy-mr1er1O, -KyKx),axis=1)))

Lref = np.concatenate((np.concatenate((1j*KzRef, Zer),axis=1),\
                       np.concatenate((Zer, 1j*KzRef),axis=1)))

Vref = Qref.dot(np.linalg.inv(Lref))
iV0Vref = iV0.dot(Vref)

A = ID + iV0Vref
B = ID - iV0Vref

iA = np.linalg.inv(A)
iAB = iA.dot(B)

S11ref = -iAB 
S12ref = 2*iA
S21ref = 0.5*(A - B.dot(iAB))
S22ref = B.dot(iA)

# Transmission side
mr2er2O = mr2er2*One
Qtrn = 1/mr2*np.concatenate((np.concatenate((KxKy, mr2er2O-KxKx),axis=1),\
                             np.concatenate((KyKy-mr2er2O, -KyKx),axis=1)))

Ltrn = np.concatenate((np.concatenate((1j*KzTrn, Zer),axis=1),\
                       np.concatenate((Zer, 1j*KzTrn),axis=1)))

Vtrn = Qtrn.dot(np.linalg.inv(Ltrn))
iV0Vtrn = iV0.dot(Vtrn)

A = ID + iV0Vtrn
B = ID - iV0Vtrn

iA  = np.linalg.inv(A)
iAB = iA.dot(B)

S11trn = B.dot(iA)
S12trn = 0.5*(A - B.dot(iAB))
S21trn = 2*iA
S22trn = -iAB

# COMPUTE GLOBAL SCATTERING MATRIX
S11,S12,S21,S22 = redheffer(S11ref,S12ref,S21ref,S22ref,S11d,S12d,S21d,S22d)
S11,S12,S21,S22 = redheffer(S11,S12,S21,S22,S11trn,S12trn,S21trn,S22trn)

# ================================================== 
# COMPUTE REFLECTED AND TRANSMITTED FIELDS
# Compute Refl/Trans Fields
T = S21.dot(Esrc)
R = S11.dot(Esrc)

Rx = R[0:MN]
Ry = R[MN:]
Rz = -np.linalg.inv(KzRef).dot(Kx.dot(Rx) + Ky.dot(Ry))

Tx = T[0:MN]
Ty = T[MN:]
Tz = -np.linalg.inv(KzTrn).dot(Kx.dot(Tx) + Ky.dot(Ty))

# transmission power coefficient
Tcoef = np.real(mr1/mr2*KzTrn/kinc[2])

# getting the zeroth order scattering parameters
r0 = np.array([Rx[middle,0], Ry[middle,0], Rz[middle,0]])
t0 = np.array([Tx[middle,0], Ty[middle,0], Tz[middle,0]])

FA = np.emath.sqrt(Tcoef[middle,middle]) # transmission field amplitude

r0tm, t0tm = atm.dot(r0), FA*atm.dot(t0)
r0te, t0te = ate.dot(r0), FA*ate.dot(t0)

# ================================================== 
# COMPUTE DIFFRACTION EFFICIENCES
# For reflected field
reff = np.abs(Rx)**2 + np.abs(Ry)**2 + np.abs(Rz)**2
R = np.real(KzRef/kinc[2]).dot(reff)
Rtot = np.sum(R)

# For transmitted field
teff = np.abs(Tx)**2 + np.abs(Ty)**2 + np.abs(Tz)**2
T = Tcoef.dot(teff)
Ttot = np.sum(T)

# ================================================== 

print(Rtot,Ttot)