%%%%%%%%%%%%%%%%%
% Problem Setup %
%%%%%%%%%%%%%%%%%
clear
close all

% Physical constants
eps0 = 8.8541878e-12;
mu0  = 4e-7 * pi;
c0 = 299792458;

% Cell size
h = 0.0025;

% Waveguide dimensions
Lx = 0.040;
Ly = 0.0225;
Lz = 0.160;

% Number of cells in each direction
Nx = round(Lx / h);
Ny = round(Ly / h);
Nz = round(Lz / h);

% Length of time steps
Dt = h / (c0 * sqrt(3));

% Insignal data
t_max = 16e-9;
Nt    = ceil(t_max / Dt);
t     = (1:Nt)' * Dt;

f_min = 4e9;
f_max = 7e9;
f_mid = (f_max + f_min) / 2;
BWr   = (f_max - f_mid) / f_mid;
f     = ((0:Nt-1)'-floor(Nt/2)) / Nt / Dt;
[s, sq, se]     = gauspuls(t-0.2e-8, f_mid, BWr, -12);

% % Square wave
% s=square(1e9*pi.*t-0.2e-8)
%  s(1:100) = 0;
%  s(500:end) = 0;
% % Sinc
%  s = zeros(length(t),1);
%  s(1:501) = sinc((-250:250)/(1*2*pi));
% plot(s)
 
 
% Allocate field matrices
Ex = zeros(Nx,     Ny + 1, Nz + 1);
Ey = zeros(Nx + 1, Ny,     Nz + 1);
Ez = zeros(Nx + 1, Ny + 1, Nz    );

Hx = zeros(Nx + 1, Ny,     Nz    );
Hy = zeros(Nx    , Ny + 1, Nz    );
Hz = zeros(Nx    , Ny,     Nz + 1);

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Initiation of boundary conditions %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
disp(sprintf('Initiate boundary conditions...'))
NumModesTE = 7; % 01 10 11 20 21 30 31
NumModesTM = 3; % 11 21 31
NumModes   = NumModesTE + NumModesTM;

disp(sprintf('  Compute TE modes'))
[ExTE, EyTE, K2TE] = ComputeTEModes(NumModesTE, Nx, Ny, h, h);
disp(sprintf('  Compute TM modes'))
[ExTM, EyTM, K2TM] = ComputeTMModes(NumModesTM, Nx, Ny, h, h);

ModalEx = [ExTE ExTM];
ModalEy = [EyTE EyTM];
ModalK2 = [K2TE K2TM];
ModalNm = sum(ModalEx.^2) + sum(ModalEy.^2); % Normalizing constants

clear ExTE EyTE ExTM EyTM

% Compute Impulse response for the propagating mode
IR   = zeros(Nt, NumModes);  % Impulse response
s1R  = zeros(Nt, NumModes);  % Reflected signal at z = Dz
s1   = zeros(Nt, NumModes);  % Total signal at z = 0
s2T  = zeros(Nt, NumModes);  % Transmitted signal at z = Lz - Dz
s2   = zeros(Nt, NumModes);  % Total signal at z = Lz

for k = 1:NumModes
  disp(sprintf('  Computing Impulse response for Mode %d', k))
  IR(:,k) = ComputeIR(Dt, h, Nt, ModalK2(k));
end

%%%%%%%%%%%%%%
% Main Loop. %
%%%%%%%%%%%%%%
disp(sprintf('Start time stepping...'))
% Set initial source boundary conditions
sEy = Ey(2:Nx, :, 2);
sEx = Ex(:, 2:Ny, 2);
sEy(:) = s(1) * ModalEy(:,1);
sEx(:) = s(1) * ModalEx(:,1);
Ey(2:Nx, :, 1) = sEy;
Ex(:, 2:Ny, 1) = sEx;
s1(1,1) = s(1);

CH = Dt / (h * mu0);
CE = Dt / (h * eps0);

ks = 400;
for k = 2:Nt

%   if k > 250 && k < 600    
%     figure(99)
%     mesh(0:Nz, 0:Nx, squeeze(0.5*Ey(:,6,:))), axis equal, axis([0 Nz 0 Nx -6 6])
%     xlabel('z');ylabel('x');zlabel('y');
%     caxis([-6 6]), view(145,30)
%     drawnow
%   end

% % Export pretty image.
% if (k == 450) 
%     figure 
%     mesh(0:Nz, 0:Nx, squeeze(0.5*Ey(:,6,:))), axis equal, axis([0 Nz 0 Nx -6 6])
%     xlabel('z');ylabel('x');zlabel('y');
%     caxis([-6 6]), view(145,30)
%     export_fig('TE10-Ey.pdf', '-painters', '-transparent')
%     pause()
% end

  %===================%
  % FDTD update loops %
  %===================%

  % -------------------------------------------------------------
  % | ADD FDTD UPDATE LOOPS                                     |
  % ----------------------------------------------------------  |
  %                                                          |  |
  %                                                         \    /
  %                                                          \  /
  %                                                           \/
  R = c0*Dt/h; % Numerical Dispersion parameter ~0.58 by default. 1 is magic. >1 diverges.
  
%     Ex = zeros(Nx,     Ny + 1, Nz + 1);
%     Ey = zeros(Nx + 1, Ny,     Nz + 1);
%     Ez = zeros(Nx + 1, Ny + 1, Nz    );
% 
%     Hx = zeros(Nx + 1, Ny,     Nz    );
%     Hy = zeros(Nx    , Ny + 1, Nz    );
%     Hz = zeros(Nx    , Ny,     Nz + 1);
  
  
% %% Update Hx
%   for i = 1:Nx+1
%     for j = 1:Ny
%         for k = 1:Nz
%             Hx(i,j,k) = Hx(i,j,k) + CH*...
%                 ((Ey(i,j,k+1)-Ey(i,j,k)) - (Ez(i,j+1,k)-Ez(i,j,k)));
%         end
%     end
%   end
% 
%     %% Update Hy
%   for i = 1:Nx
%     for j = 1:Ny+1
%         for k = 1:Nz
%             Hy(i,j,k) = Hy(i,j,k) + CH*...
%                 ((Ez(i+1,j,k)-Ez(i,j,k)) - (Ex(i,j,k+1)-Ex(i,j,k)));
%         end
%     end
%   end
%   
%     %% Update Hz
%   for i = 1:Nx
%     for j = 1:Ny
%         for k = 1:Nz+1
%             Hz(i,j,k) = Hz(i,j,k) + CH*...
%                 (Ex(i,j+1,k)-Ex(i,j,k)) - (Ey(i+1,j,k)-Ey(i,j,k));
%         end
%     end
%   end
%   
%   
%   %% E-field
%   %% Update Ex except on boundary
%   for i = 1:Nx
%       for j = 2:Ny
%           for k = 2:Nz
%             Ex(i,j,k) = Ex(i,j,k) + CE*...
%                 (Hz(i,j,k)-Hz(i,j-1,k)) - (Hy(i,j,k)-Hy(i,j,k-1));
%           end 
%       end
%   end
%   
%   %% Update Ey
%     for i = 2:Nx
%       for j = 1:Ny
%           for k = 2:Nz
%             Ey(i,j,k) = Ey(i,j,k) + CE*...
%                 (Hx(i,j,k)-Hx(i,j,k-1)) - (Hz(i,j,k)-Hz(i-1,j,k));
%           end 
%       end
%     end
%   
%     %% Update Ez
%     for i = 2:Nx
%       for j = 2:Ny
%           for k = 1:Nz
%             Ez(i,j,k) = Ez(i,j,k) + CE*...
%                 (Hy(i, j, k)-Hy(i-1 ,j, k)) - (Hx(i, j, k)-Hx(i, j-1, k));
%           end 
%       end
%   end
  
  %% Based on Rylander pg.76. 
  Dx = h;
  Dy = h;
  Dz = h;
  
  Hx = Hx + (Dt/mu0) * (diff(Ey,1,3)/Dz - diff(Ez,1,2)/Dy);
  Hy = Hy + (Dt/mu0) * (diff(Ez,1,1)/Dx - diff(Ex,1,3)/Dz);
  Hz = Hz + (Dt/mu0) * (diff(Ex,1,2)/Dy - diff(Ey,1,1)/Dx);
  
  Ex(: ,2:Ny, 2:Nz) = Ex(: ,2:Ny, 2:Nz) + (Dt/eps0) * (diff(Hz(:,:,2:Nz),1,2)/Dy - diff(Hy(:,2:Ny,:),1,3)/Dz);
  Ey(2:Nx ,: ,2:Nz) = Ey(2:Nx ,: ,2:Nz) + (Dt/eps0) * (diff(Hx(2:Nx,:,:),1,3)/Dz - diff(Hz(:,:,2:Nz),1,1)/Dx);
  Ez(2:Nx, 2:Ny, :) = Ez(2:Nx, 2:Ny, :) + (Dt/eps0) * (diff(Hy(:,2:Ny,:),1,1)/Dx - diff(Hx(2:Nx,:,:),1,2)/Dy);
  
  
  %                                                           /\
  %                                                          /  \
  %                                                         /    \
  %                                                          |  |
  % ----------------------------------------------------------  |
  % |                                                           |
  % -------------------------------------------------------------


   
  %==========================%
  % Metal Boundary Conditions %
  %==========================%

  % -------------------------------------------------------------
  % | ADD METAL OBJECTS FOR FILTERING                           |
  % ----------------------------------------------------------  |
  %                                                          |  |
  %                                                         \    /
  %                                                          \  /
  %                                                           \/
  %% Narrow midsection
  
  % ------------    ------------- 
  %            |----|
  %            |----|              x
  % ------------    -------------  |
  %                                 -z
  
%   % Set E and H to 0 within PEC walls.
%   Ex(1:2, :, 25:39) = zeros(2,10,15);
%  Ex(end-1:end, :, 25:39) = zeros(2,10,15);
%   
%   Ey(1:2, :, 25:39) = zeros(2,9,15);
%  Ey(end-1:end, :, 25:39) = zeros(2,9,15);
%   
%  Ez(1:2, :, 25:39) = zeros(2,10,15);
%  Ez(end-1:end, :, 25:39) = zeros(2,10,15);
%  
%    Hx(1:2, :, 25:39) = zeros(2,9,15);
%  Hx(end-1:end, :, 25:39) = zeros(2,9,15);
%   
%   Hy(1:2, :, 25:39) = zeros(2,10,15);
%  Hy(end-1:end, :, 25:39) = zeros(2,10,15);
%   
%  Hz(1:2, :, 25:39) = zeros(2,9,15);
%  Hz(end-1:end, :, 25:39) = zeros(2,9,15);
  
  
  %                                                           /\
  %                                                          /  \
  %                                                         /    \
  %                                                          |  |
  % ----------------------------------------------------------  |
  % |                                                           |
  % -------------------------------------------------------------

  
   
  %==============================%
  % Boundary Conditions at Z = 0 %
  %==============================%  

  % Extract transverse fields at z = dz %
  sEy = Ey(2:Nx, :, 2); 
  sEx = Ex(:, 2:Ny, 2);
  
  % Compute modal voltages %
  s1R(k,:) = (sEy(:)' * ModalEy + sEx(:)' * ModalEx) ./ ModalNm;
  
  % The reflected modal amplitude at z = Dz is the difference between the
  % total modal amplitude and that of the incoming wave. The amplitude of
  % the incoming wave at z = Dz is a convolution of the insignal at z = 0
  % and the impulse response of the wave guide.
  s1R(k,1) = s1R(k,1) - s(1:k-1)' * IR(k-1:-1:1,1);
  
  % Port Amplitude: sum of insignal and convolution of the reflected
  % signal at z = Dz with the impulse response
  for l = 1:NumModes
    s1(k,l) = s1R(1:k-1, l)' * IR(k-1:-1:1, l);
  end
  s1(k,1) = s1(k,1) + s(k);
  
  % Set port 1  boundary conditions %
  sEy(:) = ModalEy * s1(k,:)';
  sEx(:) = ModalEx * s1(k,:)';
  Ey(2:Nx,:,1) = sEy;
  Ex(:,2:Ny,1) = sEx;

  %===============================%
  % Boundary Conditions at Z = Lz %
  %===============================%

  % Extract transverse fields at z = Lz - Dz %
  sEy = Ey(2:Nx, :, Nz); 
  sEx = Ex(:, 2:Ny, Nz);

  % Compute modal voltage at z = Dz
  s2T(k,:) = (sEy(:)' * ModalEy + sEx(:)' * ModalEx) ./ ModalNm;
  
  % Port Amplitude
  for l = 1:NumModes
    s2(k,l) = s2T(1:k-1, l)' * IR(k-1:-1:1, l);
  end
  
  % Set port 2  boundary conditions %
  sEy(:) = ModalEy * s2(k,:)';
  sEx(:) = ModalEx * s2(k,:)';
  Ey(2:Nx,:,Nz+1) = sEy;
  Ex(:,2:Ny,Nz+1) = sEx;

  if(mod(k,100) == 0)
    disp(sprintf('  step %5d of %5d', k, Nt))
  end

end

% remove the incoming signal from the total signal at port 1.
s1(:,1) = s1(:,1) - s; % s1out?

figure(6)
T = t*1e9;
plot(T,s1(:,1),T,s2(:,1))
legend('s_1(t)','s_2(t)')
xlabel('t [ns]');

% -------------------------------------------------------------
% | ADD SCATTERING PARAMETERS AS FUNCTION OF FREQUENCY        |
% ----------------------------------------------------------  |
%                                                          |  |
%                                                         \    /
%                                                          \  /
f = f(f>=0);

S1in   = one_sided_scaled_fft(s);
S1out  = one_sided_scaled_fft(s1R(:,1));
S2out  = one_sided_scaled_fft(s2T(:,1));

S11 = S1out./S1in;
S21 = S2out./S1in;

% Plot transmission coefficient T(w) hedgehog -> sinewave!
figure
plot(f, S21)
grid on
title('T(\omega)')
xlim([3.9e9 6.5e9]);
%ylim([0 2]);
xlabel('f');
ylabel('|T(\omega)|');
hold on

% Plot S2out frequency domain
figure
plot(f, unwrap(angle(S1in)))
hold on
plot(f, unwrap(angle(S2out)))
legend('S1in', 'S2out')
xlim([0 10e9])
grid on


% Plot Analytical Transmission coefficient T(w).
omega = 2*pi.*f;
k = omega.*sqrt(mu0*eps0);
kc = sqrt((1*pi/Lx)^2 + (0*pi/Ly)^2);
beta = sqrt(k.^2 - kc.^2); 
Tanalytical = exp(-1i.*beta.*Lz);
plot(f, abs(Tanalytical))
hold on
fcTE10 = 1/(2*Lx*sqrt(mu0*eps0)); % Cutoff frequency
plot([fcTE10 fcTE10], [-2 10], '--')
ylim([-3 10])
xlim([0 10e9])
legend('T(\omega) Numerical', 'T(\omega) Analytical')
grid on

%% Subplots
figure(2)

theXs = [3.7 7] .* 1e9 ; %6.5 not enough; cutoff frequency 3.74

% Input signal: s
subplot(321)
plot(f, abs(S1in))
title('S1in')
xlim(theXs)
xlabel('f (Hz)')
ylabel('|S(f)|')
legend('S')
grid on
hold on

% S1in phase
subplot(322)
plot(f,unwrap(angle(S1in)))
title('angle S1in')
xlabel('f (Hz)')
xlim(theXs)
ylim([-60 0])
grid on
hold on
% TE10 cutoff frequency (Cheng. eq:10-163)
fcTE10 = 1/(2*Lx*sqrt(mu0*eps0))
plot([fcTE10 fcTE10], [-60 0], '--r')
legend('S1in', 'f_c')

% Transmitted signal: s2T
subplot(323)
plot(f,abs(S2out))
title('S2out')
xlabel('f (Hz)')
xlim(theXs)
legend('S2T')
grid on
hold on

% S1in phase
subplot(324)
plot(f,unwrap(angle(S2out)))
title('angle S2T')
xlabel('f (Hz)')
xlim(theXs)
ylim([-100 10])
grid on
hold on
% TE10 cutoff frequency (Cheng. eq:10-163)
fcTE10 = 1/(2*Lx*sqrt(mu0*eps0))
plot([fcTE10 fcTE10], [-100 10], '--r')
legend('S2T', 'f_c')

% % S11
% subplot(325)
% plot(f,abs(S11))
% title('|S11|')
% xlabel('f (Hz)')
% legend('S11 (R)')
% xlim([0 10e9])
% ylim([0 2])
% grid on
% hold on

% % S11 phase
% subplot(326)
% plot(f,unwrap(angle(S11)))
% title('angle S11')
% xlabel('f (Hz)')
% xlim([0 10e9])
% ylim([-20 20])
% grid on
% hold on
% % TE10 cutoff frequency (Cheng. eq:10-163)
% fcTE10 = 1/(2*Lx*sqrt(mu0*eps0))
% plot([fcTE10 fcTE10], [0 150], '--r')
% legend('S11 (T)', 'f_c')


% S21 mag
subplot(325)
plot(f,abs(S21))
title('S21')
xlabel('f (Hz)')
xlim(theXs)
ylim([0 2])

grid on
hold on
% TE10 cutoff frequency (Cheng. eq:10-163)
fcTE10 = 1/(2*Lx*sqrt(mu0*eps0))
plot([fcTE10 fcTE10], [0 150], '--r')
plot(f,abs(Tanalytical))
legend('S21 (T)', 'f_c','mag Tanalytical')

% S21 phase
subplot(326)
plot(f,unwrap(angle(S21))-12.2) % We offset the phase manually, so that we have zero phase difference at fc. We can then observe phase comparatively in the frequency region of interest.
title('angle S21')
xlabel('f (Hz)')
xlim(theXs)
%ylim([-20 20])
grid on
hold on
% TE10 cutoff frequency (Cheng. eq:10-163)
fcTE10 = 1/(2*Lx*sqrt(mu0*eps0));
plot([fcTE10 fcTE10], [0 150], '--r')
plot(f,unwrap(angle(Tanalytical)))
legend('S21 (T)', 'f_c','phase Tanalytical')

% figure
% plot(f,unwrap(angle(Tanalytical))-345)
% xlim([0 10e9])

%% Compare S1in with S2out (phases)
figure
plot(f, unwrap(angle(S1in)))
hold on
plot(f, unwrap(angle(S2out)))
plot([fcTE10 fcTE10], [0 150], '--r')
legend('S1in', 'S2out')

%% Misc 
Modalfc = sqrt(ModalK2)./(2*pi*sqrt(mu0*eps0));


% for blepz = 1:6
% % Analytical resonant frequency (Rylander eq5.21)
% %f_TE10 = c0/2*( (1/Lx)^2 + (0/Ly)^2 + (blepz/Lz)^2 )^(1/2)
% plot([Modalfc(blepz) Modalfc(blepz)], [0 max(P1)], '--')
% end

%% Plot K vs omega
figure(4)
plot((fcTE10:1000:8e9), sqrt((2*pi*(fcTE10:1000:8e9)).^2*mu0*eps0 - (pi/Lx)^2))
hold on
plot((fcTE10:1000:8e9), sqrt((2*pi*(fcTE10:1000:8e9))./c0))
xlabel('f');ylabel('k_z');
grid on

R = c0*Dt/Dz;

figure
hold on
grid on
k = omega.*sqrt(mu0*eps0);
for R = [1 0.8 c0*Dt/Dz 0.1]
    x = k.*Dz;
    y = R.*sin(x./2);
    
    plot( x,2*asin(y)./Dt.*Dz/c0)
    ylim([0 pi])
    xlim([0 pi])
end

%                                                           /\
%                                                          /  \
%                                                         /    \
%                                                          |  |
% ----------------------------------------------------------  |
% |                                                           |
% -------------------------------------------------------------