close all; clear all; clc;

                               
Grid.hr = 0.001;  
Grid.ht = 0.005; 
Grid.r = Grid.hr/2:Grid.hr:1-Grid.hr/2;            % Mid-point grid
Grid.theta = 0:Grid.ht:2*pi; 
[rs, thetas] = meshgrid(Grid.r, Grid.theta); 
Grid.rs = rs; 
Grid.thetas = thetas; 

c = 12;

load('PSWFdata/Concentrations.mat'); 

%% Visualizations 
close all; 
set(0,'DefaultLineMarkerSize',  6);
set(0,'DefaultTextFontSize',    14);
set(0,'DefaultAxesFontSize',    12);
set(0,'DefaultLineLineWidth',   1.4);
set(0,'defaultfigurecolor',     'w');

figure(1); 
plotEigenvalues(cseq); 


figure(2);
cseqset = [cseq; cseq(2:end, :)]; 
plot(sort(cseqset(:), 'descend'));  grid on; 
xlabel('(N,n)');
ylabel('$\lambda_{N,n}$', 'Interpreter','latex');


figure(3);
plotPSWF(1, 2, Grid); 
fprintf('Eigen-energy Ratio = %.3f\n', (sum(cseq(1,:)) + 2*sum(cseq(2:end, :), 'all'))/(c^2/4));

fprintf('End of this script.\n'); 




%% Check the alpha eigenvalues. 

% load N1.mat; 
% 
% g = diskNumericalTransform(RN(1,:), 1, c, Grid); 
% 
% figure;
% subplot(1, 2, 1);
% plotInDisk(g, Grid);
% 
% subplot(1, 2, 2);
% f = (TNp.')*RN(1,:); 
% plotInDisk(f, Grid);
% 
% alpha_est = sum((g.*conj(f)).*r, 'all')*hr*ht;
% a = alpha_p(1);





%% Utils
function f = diskNumericalTransform(radiusPart, m, c, Grid)
% Transform the function defined in separate form by radiusPart and
% angularPart, together with parameter c. 
% m: exp(im theta) in the angular part. 
    hr = Grid.hr; 
    r = Grid.r;
    theta = Grid.theta; 
    
    Lr = length(r);

    R = zeros(1, Lr); 
    for idx_r = 1:Lr
        R(idx_r) = 2*pi*sum(radiusPart.*r.*besselj(m,c*r(idx_r)*r))*hr; 
    end
    
    f = ((1i)^m)*exp(1i*m*theta.')*R;

end

function plotInDisk(f, Grid)
    rs = Grid.rs;
    ts = Grid.thetas; 
    xs = rs.*cos(ts);
    ys = rs.*sin(ts);

    spr = 1:10:length(Grid.r); 
    spt = 1:10:length(Grid.theta); 
    
    zs = real(f); 
    mesh(xs, ys, zs); axis equal; grid off; box on; 
    xlabel('x_1'); ylabel('x_2'); zlabel('Real(f)'); 
    view(0, 90); 
    colorbar; 
    colormap("jet");
end


function plotPSWF(angularOrder, radialOrder, Grid)

    if angularOrder == 0
        load('N0.mat'); 
        plot(r, R0(radialOrder+1,:)); hold on;
        legend({sprintf("$R_{%d, %d},\\quad c=%.3f$", angularOrder, radialOrder, c)}, 'Interpreter','latex');
        xlabel('r'); ylabel('R(r)'); 

        spr = 1:10:length(Grid.r); 
        spt = 1:10:length(Grid.theta); 
        

        [rs, ts] = meshgrid(Grid.r(spr), Grid.theta(spt)); 
        xs = rs.*cos(ts);
        ys = rs.*sin(ts);
        
        figure;
        zs = real(kron(T(spt).', R0(radialOrder+1,spr))); 
        surf(xs, ys, zs); axis equal; grid off; box on; 
        xlabel('x_1'); ylabel('x_2'); zlabel('PSWF val'); 
        view(0, 90); 
        colorbar; 
        colormap("jet");
        
    else
        load(sprintf('PSWFdata/N%d.mat', angularOrder)); 
    
        plot(r, RN(radialOrder+1,:)); hold on;
        legend({sprintf("$R_{%d, %d},\\quad c=%.3f$", angularOrder, radialOrder, c)}, 'Interpreter','latex');
        xlabel('r'); ylabel('R(r)'); 

        spr = 1:4:length(Grid.r); 
        spt = 1:4:length(Grid.theta); 
        
        [rs, ts] = meshgrid(Grid.r(spr), Grid.theta(spt)); 
        xs = rs.*cos(ts);
        ys = rs.*sin(ts);
        
        figure;
        colormap("hsv"); 
        zs = (kron(TNp(spt).', RN(radialOrder+1,spr))); 
%         cs = zeros([size(zs), 3]); 
%         
%         for idx_t = 1:length(spt)
%             for idx_r = 1:length(spr)
%                 ii = angle()
%                 cs(idx_t, idx_r, :) = cmap(ii, :); 
%             end
%         end

        cs = angle(zs); 

        surf(xs, ys, abs(zs), cs, 'FaceAlpha',0.7); axis equal; grid off; box on; 
        xlabel('x_1'); ylabel('x_2'); zlabel('PSWF val'); 
        caxis([-pi, pi]); 
        view(68, 40); 

        colorbar; 
        shading flat; 
    end

end

function plotEigenvalues(cseq)
    [Nmax, nmax] = size(cseq); 
    Nmax = Nmax - 1;
    
    cseq(cseq >= 1) = 1;
    cseq(cseq <= 0) = 0; 

    axis equal; 
    xlim([-Nmax, Nmax]); 
    ylim([-1/2, 10]);
    
    for idx = 1:nmax
        rectangle('Position',[-1/2,-1/2+(idx-1),1,1], 'FaceColor', rgba([0,0,0], 1-cseq(1, idx)), 'LineStyle', 'none');
    end

    for N = 1:Nmax
        for idx = 1:nmax
            rectangle('Position',[-1/2-N,-1/2+(idx-1),1,1], 'FaceColor', rgba([0,0,0], 1-cseq(N+1, idx)), 'LineStyle', 'none');
        end

        for idx = 1:nmax
            rectangle('Position',[-1/2+N,-1/2+(idx-1),1,1], 'FaceColor', rgba([0,0,0], 1-cseq(N+1, idx)), 'LineStyle', 'none');
        end
    end

    xlabel('Angular Order N'); 
    ylabel('Radial Order n');
    
    colormap(flipud(colormap("gray")));
    colorbar();

end

function rgb = rgba(c, a)
    rgb = a*[1,1,1] + (1-a)*c;
end


