A = 1;
lambda = 1;
k = 2*pi/lambda;

x_range = linspace(0,15,100);
y_range = linspace(-5,5,100);
z_range = linspace(-5,5,100);

[X, Y, Z] = meshgrid(x_range, y_range, z_range);

R = @(X, Y, Z, x0, y0, z0) sqrt((X - x0).^2 + (Y - y0).^2 + (Z - z0).^2);

E = @(X, Y, Z, x0, y0, z0) (A./R(X, Y, Z, x0, y0, z0)) .* exp(-1j * k * R(X, Y, Z, x0, y0, z0));

%% Plot a single phasor in 3D space

E1 = E(X, Y, Z, 0, 0, 0);

c = 0.2;

figure;

% Crests
p1 = patch(isosurface(x_range, y_range, z_range, real(E1),  c));
isonormals(x_range, y_range, z_range, real(E1), p1);
p1.FaceColor = 'cyan'; p1.EdgeColor = 'none';

% Troughs
p2 = patch(isosurface(x_range, y_range, z_range, real(E1), -c));
isonormals(x_range, y_range, z_range, real(E1), p2);
p2.FaceColor = 'red'; p2.EdgeColor = 'none';
camlight; lighting gouraud;

xlabel('x (\lambda)'); ylabel('y (\lambda)'); zlabel('z (\lambda)')
title('3D Wave Propagation from a Single Source')
legend([p1 p2], 'Crest', 'Trough')
axis tight; view(3)

%% Plots the 3D Wave Propagation from two sources

E1 = E(X, Y, Z, 0, 0, 0);
E2 = E(X, Y, Z, 0, 1, 0);

E_total = E1 + E2;

c = 0.2;

figure;

% Crests
p1 = patch(isosurface(x_range, y_range, z_range, real(E_total),  c));
isonormals(x_range, y_range, z_range, real(E_total), p1);
p1.FaceColor = 'cyan'; p1.EdgeColor = 'none';

% Troughs
p2 = patch(isosurface(x_range, y_range, z_range, real(E_total), -c));
isonormals(x_range, y_range, z_range, real(E_total), p2);
p2.FaceColor = 'red'; p2.EdgeColor = 'none';
camlight; lighting gouraud;

xlabel('x (\lambda)'); ylabel('y (\lambda)'); zlabel('z (\lambda)')
title('3D Wave Propagation from Two Sources')
legend([p1 p2], 'Crest', 'Trough')
axis tight; view([45 45])

%% Plots the 3D Wave Propagation from 2x2 Array

E1 = E(X, Y, Z, 0, 0, 0);
E2 = E(X, Y, Z, 0, 1, 0);
E3 = E(X, Y, Z, 0, 0, 1);
E4 = E(X, Y, Z, 0, 1, 1);

E_total = E1 + E2 + E3 + E4;

c = 0.2;

figure;

% Crests
p1 = patch(isosurface(x_range, y_range, z_range, real(E_total),  c));
isonormals(x_range, y_range, z_range, real(E_total), p1);
p1.FaceColor = 'cyan'; p1.EdgeColor = 'none';

% Troughs
p2 = patch(isosurface(x_range, y_range, z_range, real(E_total), -c));
isonormals(x_range, y_range, z_range, real(E_total), p2);
p2.FaceColor = 'red'; p2.EdgeColor = 'none';
camlight; lighting gouraud;

xlabel('x (\lambda)'); ylabel('y (\lambda)'); zlabel('z (\lambda)')
title('3D Wave Propagation from 2×2 Array')
legend([p1 p2], 'Crest', 'Trough')
axis tight; view([50 20])

%% Plots the 3D Wave Propagation from NxN Array

x_range = linspace(0,30,100);
y_range = linspace(-10,10,100);
z_range = linspace(-10,10,100);

[X, Y, Z] = meshgrid(x_range, y_range, z_range);

function E_total = generateNxNArray(Ny, Nz, X, Y, Z, E, spacing)
    E_total = zeros(size(X));

    y_positions = ((0:Ny-1) - (Ny-1)/2) * spacing;
    z_positions = ((0:Nz-1) - (Nz-1)/2) * spacing;

    for i = 1:Ny
        for j = 1:Nz
            E_total = E_total + E(X, Y, Z, 0, y_positions(i), z_positions(j));
        end
    end
end

Ny = 4;
Nz = 4;
spacing = .5;
E_total = generateNxNArray(Ny, Nz, X, Y, Z, E, spacing);

c = 0.2;

figure;

% Crests
p1 = patch(isosurface(x_range, y_range, z_range, real(E_total),  c));
isonormals(x_range, y_range, z_range, real(E_total), p1);
p1.FaceColor = 'cyan'; p1.EdgeColor = 'none';

% Troughs
p2 = patch(isosurface(x_range, y_range, z_range, real(E_total), -c));
isonormals(x_range, y_range, z_range, real(E_total), p2);
p2.FaceColor = 'red'; p2.EdgeColor = 'none';
camlight; lighting gouraud;

xlabel('x (\lambda)'); ylabel('y (\lambda)'); zlabel('z (\lambda)')
title(sprintf('3D Wave Propagation from %d×%d Array', Ny, Nz))
legend([p1 p2], 'Crest', 'Trough')
axis tight; view([50 20])

%% Plots the 3D Wave Propagation from NxN Steered Array

psi_n = @(y_off, z_off, theta, phi) -k * (y_off * sin(theta) + z_off * sin(phi));

function E_total = generateNxNSteeredArray(Ny, Nz, X, Y, Z, E, psi_n, spacing, theta, phi)
    E_total = zeros(size(X));

    y_positions = ((0:Ny-1) - (Ny-1)/2) * spacing;
    z_positions = ((0:Nz-1) - (Nz-1)/2) * spacing;

    for i = 1:Ny
        for j = 1:Nz
            E_total = E_total + (E(X, Y, Z, 0, y_positions(i), z_positions(j)) .* exp(1j * psi_n(y_positions(i), z_positions(j), theta, phi)));
        end
    end
end

x_range = linspace(0,30,100);
y_range = linspace(-30,30,100);
z_range = linspace(-30,30,100);

[X, Y, Z] = meshgrid(x_range, y_range, z_range);

theta = pi/6;   % Elevation
phi = pi/6;     % Azimuth
Ny = 4;
Nz = 4;
spacing = .5;

E_total = generateNxNSteeredArray(Ny, Nz, X, Y, Z, E, psi_n, spacing, theta, phi);

c = 0.2;

figure;

% Crests
p1 = patch(isosurface(x_range, y_range, z_range, real(E_total),  c));
isonormals(x_range, y_range, z_range, real(E_total), p1);
p1.FaceColor = 'cyan'; p1.EdgeColor = 'none';

% Troughs
p2 = patch(isosurface(x_range, y_range, z_range, real(E_total), -c));
isonormals(x_range, y_range, z_range, real(E_total), p2);
p2.FaceColor = 'red'; p2.EdgeColor = 'none';
camlight; lighting gouraud;

xlabel('x (\lambda)'); ylabel('y (\lambda)'); zlabel('z (\lambda)')
title(sprintf('3D Wave Propagation from %d×%d Steered Array', Ny, Nz))
legend([p1 p2], 'Crest', 'Trough')
axis tight; view([80 5])

%% Array Factor Plot of NxN Sources

phi_geo = @(y_off, z_off, theta, phi) k * y_off * sin(theta) + k * z_off * sin(phi);
psi_n = @(y_off, z_off, theta, phi) -k * (y_off * sin(theta) + z_off * sin(phi));

function AF = generateAF(Ny, Nz, phi_geo, psi_n, spacing, theta0, phi0, THETA, PHI)
    y_positions = ((0:Ny-1) - (Ny-1)/2) * spacing;
    z_positions = ((0:Nz-1) - (Nz-1)/2) * spacing;

    AF = zeros(size(THETA));

    for i = 1:length(y_positions)
        y_n = y_positions(i);
        for j = 1:length(z_positions)
            z_n = z_positions(j);
            AF = AF + exp(1j * (phi_geo(y_n, z_n, THETA, PHI) + psi_n(y_n, z_n, theta0, phi0)));
        end
    end
end

theta_vec = linspace(-pi/2, pi/2, 1000);
phi_vec = linspace(-pi/2, pi/2, 1000);
[THETA, PHI] = meshgrid(theta_vec, phi_vec);

theta0 = -pi/4;
phi0 = pi/6;
Ny = 8;
Nz = 4;
spacing = 0.5;

AF = generateAF(Ny, Nz, phi_geo, psi_n, spacing, theta0, phi0, THETA, PHI);

P = abs(AF).^2;

figure;
imagesc(theta_vec*180/pi, phi_vec*180/pi, P);
colorbar;
colormap parula;
xlabel('\theta (degrees)');
ylabel('\phi (degrees)');
title('Radiation Pattern P(\theta, \phi)');
axis xy;

P_norm = P / max(P(:));

X_surf = P_norm .* cos(PHI) .* cos(THETA);
Y_surf = P_norm .* cos(PHI) .* sin(THETA);
Z_surf = P_norm .* sin(PHI);

figure;
surf(X_surf, Y_surf, Z_surf, P_norm);
shading interp;
colorbar;
colormap parula;
xlabel('x'); ylabel('y'); zlabel('z');
title('3D Polar Radiation Pattern');
axis equal;

%% Gain Function

theta_vec = linspace(-pi/2, pi/2, 1000);
phi_vec = linspace(-pi/2, pi/2, 1000);
[THETA, PHI] = meshgrid(theta_vec, phi_vec);

theta0 = -pi/4;
phi0 = pi/6;
Ny = 8;
Nz = 4;
spacing = 0.5;

AF = generateAF(Ny, Nz, phi_geo, psi_n, spacing, theta0, phi0, THETA, PHI);

P = abs(AF).^2;

P_total = trapz(phi_vec, trapz(theta_vec, P));

[~, idx_theta] = min(abs(theta_vec - theta0));
[~, idx_phi]   = min(abs(phi_vec - phi0));
AF0 = AF(idx_phi, idx_theta);

G = abs(AF0)^2 / P_total;

Ny = 8;
Nz = 5;
spacing = 0.5;
theta0 = pi/3;
phi0 = -pi/8;
alpha = 0.01;
max_iterations = 10000;
grad_tol = 1e-3;

theta_vec = linspace(-pi/2, pi/2, 100);
phi_vec   = linspace(-pi/2, pi/2, 100);
[THETA, PHI] = meshgrid(theta_vec, phi_vec);

y_positions = ((0:Ny-1) - (Ny-1)/2) * spacing;
z_positions = ((0:Nz-1) - (Nz-1)/2) * spacing;

num_elements = Ny * Nz;
y_flat = zeros(1, num_elements);
z_flat = zeros(1, num_elements);
n = 1;
for i = 1:Ny
    for j = 1:Nz
        y_flat(n) = y_positions(i);
        z_flat(n) = z_positions(j);
        n = n + 1;
    end
end

% Precompute geometric phases — shape: (Nphi, Ntheta, num_elements)
GEO_PHASE = zeros(length(phi_vec), length(theta_vec), num_elements);
for n = 1:num_elements
    GEO_PHASE(:,:,n) = phi_geo(y_flat(n), z_flat(n), THETA, PHI);
end

% Precompute target indices
[~, idx_theta] = min(abs(theta_vec - theta0));
[~, idx_phi] = min(abs(phi_vec   - phi0));

% Precompute geometric phase at target point — shape: (1, num_elements)
GEO_PHASE_target = zeros(1, num_elements);
for n = 1:num_elements
    GEO_PHASE_target(n) = phi_geo(y_flat(n), z_flat(n), theta0, phi0);
end

psi_current = 0.01 * randn(1, num_elements);
G_history   = [];
psi_history = [];

for iter = 1:max_iterations
    % Step 1: Compute all contributions and AF
    % CONTRIB shape: (Nphi, Ntheta, num_elements)
    CONTRIB = exp(1j * (GEO_PHASE + reshape(psi_current, 1, 1, num_elements)));
    AF = sum(CONTRIB, 3);

    % Step 2: Compute P, P_total, AF0, G
    P = abs(AF).^2;
    P_total = trapz(phi_vec, trapz(theta_vec, P, 2));
    AF0 = AF(idx_phi, idx_theta);
    G = abs(AF0)^2 / P_total;
    G_history(end+1) = G;

    % Step 3: Compute gradient (fully vectorized)
    % Contribution of each element at target — shape: (1, num_elements)
    contrib_at_target = exp(1j * (GEO_PHASE_target + psi_current));

    % dP_total/dpsi for each element — shape: (num_elements, 1)
    dPtotal_dpsi = squeeze(trapz(phi_vec, trapz(theta_vec, ...
        2 * real(1j * conj(AF) .* CONTRIB), 2)));

    grad = (2 * real(conj(AF0) * 1j * contrib_at_target) * P_total ...
        - abs(AF0)^2 * dPtotal_dpsi') / P_total^2;

    % Check convergence
    grad_norm = norm(grad);
    if grad_norm < grad_tol
        fprintf('Converged at iteration %d (grad norm = %.2e)\n', iter, grad_norm);
        break;
    end

    % Step 4: Update psi
    psi_current = psi_current + alpha * grad;
    psi_history(end+1, :) = psi_current;
end

% Plot G convergence
figure;
plot(1:length(G_history), G_history, 'LineWidth', 1.5)
xlabel('Iteration')
ylabel('G')
title(sprintf('Gain vs iteration, %d×%d array, θ₀=%.2f, φ₀=%.2f', Ny, Nz, theta0, phi0))
grid on

% Plot psi history
figure;
plot(1:size(psi_history,1), psi_history, 'LineWidth', 1.5)
xlabel('Iteration')
ylabel('ψ (radians)')
title(sprintf('ψ convergence, %d×%d array, θ₀=%.2f, φ₀=%.2f', Ny, Nz, theta0, phi0))
% legend(arrayfun(@(n) sprintf('element %d', n), 1:num_elements, 'UniformOutput', false))
grid on


% Compare to analytical solution
psi_analytical = zeros(1, num_elements);
for n = 1:num_elements
    psi_analytical(n) = -k * (y_flat(n) * sin(theta0) + z_flat(n) * sin(phi0));
end

disp('Gradient ascent psi:')
disp(psi_current)
disp('Analytical psi:')
disp(psi_analytical)
disp('Difference (should be close to 0 or integer):')
disp((psi_current - psi_analytical) / (2*pi))

% Find meaningful stages from G_history
G_start  = 1;                                          % before anything happens
G_10  = round(0.1 * size(psi_history, 1));          % early — beam just starting to form
G_half = round(interp1(G_history, 1:length(G_history), mean([G_history(1) G_history(end)]))); % halfway up the gain curve
G_end = size(psi_history, 1);                       % converged

iters = [G_start, G_10, G_half, G_end];

titles = {'Initial (random ψ)', 'Early formation', 'Mid convergence', 'Converged'};

figure;
for i = 1:4
    psi_plot = psi_history(iters(i), :);
    CONTRIB_plot = exp(1j * (GEO_PHASE + reshape(psi_plot, 1, 1, num_elements)));
    AF_plot = sum(CONTRIB_plot, 3);
    P_plot = abs(AF_plot).^2;

    subplot(2, 2, i);
    imagesc(theta_vec*180/pi, phi_vec*180/pi, P_plot);
    colormap parula; 
    % colorbar;
    xlabel('\theta (degrees)'); ylabel('\phi (degrees)');
    % title(sprintf('%s\n(iter %d, G=%.2f)', titles{i}, iters(i), G_history(iters(i))));
    axis xy;
end
sgtitle(sprintf('%d×%d Array Beam Evolution, θ₀=%.0f°, φ₀=%.0f°', ...
    Ny, Nz, theta0*180/pi, phi0*180/pi));

% Plot AF progression over psi_history as 2D heatmaps
iters = round(linspace(1, size(psi_history, 1), 20));
figure;
for i = 1:20
    psi_plot = psi_history(iters(i), :);
    CONTRIB_plot = exp(1j * (GEO_PHASE + reshape(psi_plot, 1, 1, num_elements)));
    AF_plot = sum(CONTRIB_plot, 3);
    P_plot = abs(AF_plot).^2;
        imagesc(theta_vec*180/pi, phi_vec*180/pi, P_plot);
    colorbar; colormap parula;
    xlabel('\theta (degrees)'); ylabel('\phi (degrees)');
    title(sprintf('Iteration %d / %d', iters(i), size(psi_history,1)));
    axis xy; drawnow;
    pause(0.3);
end