% TMM implementation; refer to lecture notes for notations and derivations
% input: arrays of N interfaces, and N coordinates (last interface is
% virtual); amplitude of the mode(s) incicent on the stack from the firs
% (bottom-most) layer

% output: array of modal amplitudes within each layer

global c0const mu0const; c0const=3e8; 

e1= 2.25; 

lam0=1;  
omg0=2*pi*c0const./lam0;
ang=50; 


zj=[0 0.1]; %last interface is virtual
eps1j=[e1 1+0.00i]; 
% eps2j=[e1 1+0.001i]; 
% muj=[1 1]; 

% nprd=10; 
% zj=zeros(1,nprd*2+2); 
% eps1j=ones(1,nprd*2+2); 
% e1=9; e2=4; 
% % d1=1.5; d2=1.; 
% d1=0.33/4; d2=0.125; 
% for np=1:nprd
%     zj(2*np)=zj(2*np-1)+d1; 
%     zj(2*np+1)=zj(2*np)+d2; 
%     eps1j(2*np)=e1; 
%     eps1j(2*np+1)=e2; 
% end 
% zj(end)=zj(nprd*2+1); 
 
muj=ones(1,length(eps1j));
eps2j=eps1j; 

kx0=omg0*sqrt(eps1j(1))*sind(ang)/c0const; 
x=(-40:0.1:40); 
fx=exp(-(x-5).^2/5).*exp(1i*kx0*x)/sqrt(length(x));

fx=fx.';

[kxArr,fk]=dftI(x,fx,kx0); 

figure(2)
clf
subplot(1,2,1)
plot(x,abs(fx))
subplot(1,2,2)
plot(kxArr,abs(fk))

%% compute all

xVec=(0:0.1:10); 
zVec=(-5:0.01:5)+0; 

[xArr,zArr]=meshgrid(xVec,zVec); 
EFldTot=0*xArr; HFldTot=0*xArr; 

for ikx=1:length(kxArr)
    
    if (abs(fk(ikx))>max(abs(fk))*1e-5)
        
        kx0=kxArr(ikx);
        ainc=[1;1]*fk(ikx); 

        [Tcur,Rcur,aPlArr,aMinArr,kzPlArr]=tmmFun(zj,eps1j,eps2j,muj,omg0,kx0,ainc);

        if isnan(Rcur(1))
            aPlArr(1,:)=0; 
            aMinArr(1,:)=0; 
        end 
        if isnan(Rcur(2))
            aPlArr(2,:)=0; 
            aMinArr(2,:)=0; 
        end 
        [~,EFld,~,~,HFld,~]=tmmFieldFun(zj,eps1j,eps2j,muj,omg0,kx0,...
            aPlArr,aMinArr,xArr,zArr); 

        EFldTot=EFldTot+EFld; 
        HFldTot=HFldTot+HFld; 
    end %if 
end %for

%% plot the fields
figure(1)
subplot(1,2,1)
surf(xArr/lam0,zArr/lam0,real(EFldTot),'EdgeColor','none')
set(gca,'FontSize',14)
xlabel('x/\lambda_0')
ylabel('z/\lambda_0')
daspect([1 1 1])
title('E')
colorbar 
view(2)
colormap jet

subplot(1,2,2)
surf(xArr/lam0,zArr/lam0,real(HFldTot),'EdgeColor','none')
set(gca,'FontSize',14)
xlabel('x/\lambda_0')
ylabel('z/\lambda_0')
daspect([1 1 1])
title('H')
colorbar 
view(2)
view(2)
colormap jet

disp(['T1=',num2str(Tcur(1)),'; R1=',num2str(Rcur(1))])
disp(['T2=',num2str(Tcur(2)),'; R2=',num2str(Rcur(2))])