%% Wigner Function Transformation
clc
clear

ds = 0.005; %sampling rate at every 5mm
fs = 1/ds; % sampling frequency
L=60;  %sets the number of sample points 
D=0.295;    %sets the physical distance range

c=3e8;
f=3e9; %freq
lambda=c/f;
k=2*pi/lambda;

load('Hy_R1(x1,y1,x2,y2)_36pos_1cm.mat'); % load source CF=R(x1,y1,x2,y2)
CF = permute(Hy_R1,[1,3,2,4]); % CF(x1,x2,y1,y2)

clear Hy_R1
%% 1. Variable change Part 1
for k=1:L 
    for n=1:L
        
        R = CF(:,:,k,n); %R(x1,x2)
        
% variable rotation for each inner matrix

        % variable change on the correlation matrix R(X2,X1) -> R(x,sx)

        for X1 = [1:L]; % nx matrix
            for X2 = [1:L]; % ny matrix
            x = (X1+X2)/2;  % change of variable X1/X2 to x
            sx = (X1-X2);    % change 
            Q(sx+L,2*x-1)=(R(X1,X2)); %Q(sx,x) value
    
            end
        end

        % averaging in between the values already calculated after variable change
        for h = 2:2:2*L-2
            Q(h,L) = (Q(h-1,L)+Q(h,L-1)+Q(h,L+1)+Q(h+1,L))/4; 
        end

        for t=1:L-1
            for m = 1:L-1
    
            Q(m+t,L-m+t) = (Q(m+t-1,L-m+t)+Q(m+t,L-m+t-1)+Q(m+t,L-m+t+1)+Q(m+t+1,L-m+t))/4; 
        
            end
        end
      
% after averaging, arrange the rotated inner matrix into AA(sx,x,x2,y2)

        AA(:,:,k,n) = Q; % AA(sx,x,y2,y1)
        
    end
   
end

clear CF
clear R
clear Q
%% 2. Variable change Part 2

R = AA; % AA(sx,x,y1,y2)

% variable change on the outer correlation matrix R(X1,X2) -> R(x,s)

for y1 = [1:L]; % nx matrix    
    for y2 = [1:L]; % ny matrix
    y = (y1+y2)/2;  % change of variable to y
    sy = (y1-y2);    % change sy=y1-y2 
    q(:,:,sy+L,2*y-1)=(R(:,:,y1,y2)); %D(s,x) value
    end
end


% averaging in between the real value after variable change
for h = 2:2:2*L-2
    q(:,:,h,L) = (q(:,:,h-1,L)+q(:,:,h,L-1)+q(:,:,h,L+1)+q(:,:,h+1,L))/4; 
end

for t=1:L-1
    for m = 1:L-1
    q(:,:,m+t,L-m+t) = (q(:,:,m+t-1,L-m+t)+q(:,:,m+t,L-m+t-1)+q(:,:,m+t,L-m+t+1)+q(:,:,m+t+1,L-m+t))/4; 
    end
end

clear R
clear AA
%% 3. rearrange Q(sx,x,sy,y) -> Q(sy,sx,y,x) and do FFT2 on sx and sy 

CFr=permute(q,[3,1,4,2]); %rearrange Q to be CFr(sy,sx,y,x)

clear q
%%
pmax=lambda/2*fs;
range = 1.5; % new range -1.5<P<1.5 
NFFT=round((2*L-1)*(lambda*fs)/(range*2)); % Defines zero padding value

for y=1:(2*L-1)
    y
    for x=1:(2*L-1)
       
       s=CFr(:,:,y,x);
       S=padarray(s,[(NFFT-(2*L-1))/2 (NFFT-(2*L-1))/2]);
 
       % for range -1.5 to 1.5 with zero padding
       S1=fftshift(fft2((ifftshift(S)))); 
       WF(:,:,y,x)= S1(ceil(-L+(NFFT/2)+1):ceil((L+(NFFT/2)-1)),ceil(-L+(NFFT/2)+1):ceil((L+(NFFT/2)-1))); %WF(py,px,y,x) 
       
    end;
end;

clear CFr
%% 4. Plot WF(py,px,y,x)

figure
x=linspace(-D/2,D/2,2*L-1); % range of x
p=linspace(1.5,-1.5, 2*L-1);
size = 14;

imagesc(x,p,abs(squeeze(WF(60,:,60,:)))); % WF(py,px,y,x)
set(gca,'YDir','normal')
colormap('jet'); % set the colorscheme
colorbar
title('$$WF\,at \,1cm$$','Interpreter','latex', 'FontSize', size)
xlabel('$$x_1 [m]$$','Interpreter','latex', 'FontSize', size);
ylabel('$$Momentum \,[P_x]$$','Interpreter','latex', 'FontSize', size);
set(gca,'fontsize', size)

%% 5. Frobenius Perron Approximation - Wigner Transformed Function (x-domain)

clearvars -except WF
clc

L=60;  %sets the number of sample points 
D=0.295;    %sets the physical distance range
c=3e8;
f=3e9;
lambda=c/f; 
k=2*pi/lambda;

ds = 0.005; %sampling rate
fs = 1/ds; % sampling frequency
pmax=lambda/2*fs;
NFFT=2*L-1;
z=0.09; % propagaion distance

n = 2*L-1; % Resolution
[x1,y1] = meshgrid(linspace(-D/2,D/2,n),linspace(-D/2,D/2,n));

Px = linspace(-1.5,1.5,NFFT);
Py = linspace(-1.5,1.5,NFFT);

for ipy=1:NFFT
    for ipx=1:NFFT
    mag_p(ipy,ipx) = sqrt(Px(ipx)^2+Py(ipy)^2); %mag_p(py,px)
    end
end

Wq1 = zeros(2*L-1,n); % Preallocate Wq

%% 6. FP on WF(py,px,y,x)

for py=1:NFFT; %py
    py
    for px=1:NFFT; %px
    
    Wx = squeeze(WF(py,px,:,:)); % Wx(y,x) from WF3/WF(py,px,y,x)
    
% Transform domain
% Interpolate W with new xq and yq axes

                if abs(mag_p(py,px)) <= 1 % FP propagation
                    xq = x1 - z.*(Px(px)./sqrt(1-abs(mag_p(py,px)).^2));
                    yq = y1 - z.*(Py(py)./sqrt(1-abs(mag_p(py,px)).^2));
                    
                    %Wq1 = interp2(xq,yq,Wx,x1,y1,'cubic',0); %1-D data interpolation 
                    Wq1 = interp2(xq,yq,Wx,x1,y1,'cubic',0); %1-D data interpolation 
                else % Evanescent propagation
                    Wq1 = exp(-2*k*z*sqrt(mag_p(py,px).^2-1))*interp2(x1,y1,Wx,x1,y1,'cubic',0);
                end
        WF2(py,px,:,:) = Wq1; %WF3(py,px,yq,xq)
    end
end

%% 7. Plot WF3(py,px,yq,xq)

figure
px=linspace(1.5,-1.5,NFFT);
imagesc(x1(1,:),px,abs(squeeze(WF2(60,:,60,:)))); %WF4(py,px,yq,xq)
set(gca,'YDir','normal')
colormap('jet'); % set the colorscheme
title('1cm propagated to 15cm')
xlabel('Transformed Position (f(X,Z,P))');
ylabel('Momentum (P)')
colorbar

%% 8. inverse WF 

for y=1:(2*L-1)
    y
    for x=1:(2*L-1)
       
       g=fftshift(ifft2(ifftshift(squeeze(WF2(:,:,y,x)))));
       CFr2(:,:,y,x)= g; % CFr2(sy,sx,y,x)
    end;
                
end;

CFr3 = permute(CFr2,[2,4,1,3]); %CFr3(sx,x,sy,y)

figure
imagesc(abs(squeeze(CFr3(:,:,60,60)))); %CFr2 without zero padding - WF(sx,x,sy,x)
title('CF(sx,x,60,60) - 5cm')
set(gca,'YDir','normal')
colormap('jet'); % set the colorscheme
colorbar

%% 9. zoom in on the sx and sy
CFr4=CFr3(50:70,:,50:70,:);

[sx,sy] = meshgrid(linspace(-D,D,length(CFr4(:,1,1,1))));
[sx2,sy2] = meshgrid(linspace(-D,D,119));

for x=1:2*L-1; %py
    x
    for y=1:2*L-1; %px
   
    Wx = squeeze(CFr4(:,x,:,y)); % Wx(y,x)
    Wq1 = interp2(sx,sy,Wx,sx2,sy2,'cubic',0); %2-D data interpolation 
    
    CFr5(:,x,:,y) = Wq1;
    end
end

figure
imagesc(abs(squeeze(CFr5(:,:,60,60)))); %CFr2 without zero padding - WF(sx,x,sy,x)
title('CF(sx,x,60,60) - 1cm')
set(gca,'YDir','normal')
colormap('jet'); % set the colorscheme
colorbar

%% change of variable CFr2(:,:,sy,y) to CFr2(:,:,y1,y2)
for y1 = [1:L]; % nx matrix    
    for y2 = [1:L]; % ny matrix
    y = (y1+y2)/2;  % change of variable to y1 and y2
    sy = (y1-y2);    % change s=y1-y2 
    Cr(:,:,y1,y2)=CFr5(:,:,sy+L,2*y-1); %Cr(sx,x,y1,y2)
    end
end

%% change of variable F(sx,x) to F(x1,x2)

for k=1:L
    for n=1:L
        
        F = Cr(:,:,k,n);
        
 % variable rotation for each inner matrix

        % variable change on the correlation matrix R(X2,X1) -> R(x,s)

        for X1 = [1:L]; % nx matrix
            for X2 = [1:L]; % ny matrix
            x = (X1+X2)/2;  % change of variable x to X1/X2
            sx = (X1-X2);    % change 
            RR(X1,X2)=F(sx+L,2*x-1); %RR(X1,X2)
    
            end
        end
       CF2(:,:,k,n) = RR; % set CF2(x1,x2,y1,y2)
    end
   
end

%%
load('SF1.mat') % scaling factor

Hy_CF3 = CF2 .* SF1;
%Hy_CF3 = CF2 .* SF1.*(PF_ave(1601)^2);

x=linspace(-D/2,D/2,L); % range of x
y=linspace(-D/2,D/2,L); % range of y
size = 14;

figure
pcolor(x,y,squeeze(abs(Hy_CF3(30,:,30,:)))'); %Hy_CF3(x1,x2,y1,y2)
title('$$CF\,propagated\,to \,10\,cm$$','Interpreter','latex', 'FontSize', size)
set(gca,'YDir','normal')
colormap('jet');
colorbar
shading interp
%caxis([0 14e-4])
xlabel('$$x_1 [m]$$','Interpreter','latex', 'FontSize', size);
ylabel('$$y_1 [m]$$','Interpreter','latex', 'FontSize', size);
set(gca,'fontsize', size)