%% Direction solution to the unconstrained optimization problem with the Conjugate Gradient technique
clc;
clear;
format long;

% Define the function (now accepts a 3D array)
f           = @(X) sum(X(:).^4) - 4*sum(X(:).^3) + 6*sum(X(:).^2) + 5;

% Initial Guess: Now we use a 3D array
x_o         = rand(3, 3, 3);  % Random 3D array as the initial guess
fprintf('Initial Objective Function Value: %.6f\n\n', f(x_o));

% Convergence Criteria:
Epsilon     = 1E-5;

% Gradient Calculation (numerical)
J           = compute_gradient(f, x_o);  % Gradient at initial point
S           = -J;  % Initial search direction

% Iteration Counter:
i           = 1;

% Minimization
while norm(S(:)) > Epsilon % Halting Criterion
    % Step Size Optimization (Line Search)
    lambda  = optimize_step_size(f, x_o, S);

    % Update the solution
    x_o_new = x_o + lambda * S;

    % Update the gradient and search direction
    J_old   = J;
    J       = compute_gradient(f, x_o_new);
    S       = update_search_direction(S, J, J_old);

    % Update for next iteration
    x_o     = x_o_new;
    i       = i + 1;
end

% Output
fprintf('Number of Iterations for Convergence: %d\n\n', i);
fprintf('Objective Function Minimum Value: %.6f\n\n', f(x_o));


%% Visualize the optimization

clc;
clear;
format long;

% Define the function (now accepts a 3D array)
f = @(X) sum(X(:).^4) - 4*sum(X(:).^3) + 6*sum(X(:).^2) + 5;

% Initial Guess: Now we use a 3D array
x_o = rand(3, 3, 3);  % Random 3D array as the initial guess
fprintf('Initial Objective Function Value: %.6f\n\n', f(x_o));

% Convergence Criteria:
Epsilon = 1E-5;

% Gradient Calculation (numerical)
J = compute_gradient(f, x_o);  % Gradient at initial point
S = -J;  % Initial search direction

% Iteration Counter:
i = 1;

% Prepare for visualization
fig = figure;
clf
set(gcf,'Position',[50 50 950 750])
set(gca,'FontSize',16,'Box','On','Linewidth',2);
hold on;

% Create a mesh grid for visualization
[X1, X2] = meshgrid(-2:0.1:2, -2:0.1:2); % Adjust domain based on your function
Z = arrayfun(@(x1, x2) f([x1, x2, 0]), X1, X2); % Example 2D slice: f([X1, X2, 0])

% Plot the surface
surf(X1, X2, Z);
xlabel('X1');
ylabel('X2');
zlabel('f(X)');
title('Conjugate Gradient Optimization Process');
colorbar;
scatter3(x_o(1), x_o(2), f(x_o), 'r', 'filled'); % Initial point

% Set the 3D view
view(75, 30);  % Example 3D view: azimuth 45°, elevation 30°

% Minimization with visualization
while norm(S(:)) > Epsilon % Halting Criterion
    % Step Size Optimization (Line Search)
    lambda = optimize_step_size(f, x_o, S);

    % Update the solution
    x_o_new = x_o + lambda * S;

    % Update the gradient and search direction
    J_old = J;
    J = compute_gradient(f, x_o_new);
    S = update_search_direction(S, J, J_old);

    % Update for next iteration
    x_o = x_o_new;
    i = i + 1;
    
    % Visualization Update
    scatter3(x_o(1), x_o(2), f(x_o), 'g', 'filled'); % Current point
    drawnow;  % Update the plot
end

% Output
fprintf('Number of Iterations for Convergence: %d\n\n', i);
fprintf('Objective Function Minimum Value: %.6f\n\n', f(x_o));

%%
% Numerical Gradient Calculation using the finite differences method for 3D
function J  = compute_gradient(f, X)
    epsilon = 1E-5;  % Step size for numerical differentiation
    J       = zeros(size(X));  % Gradient with the same size as the input
    
    % Loop over all elements of the 3D array
    for i = 1:numel(X)
        X1    = X;
        X2    = X;
        % Perturb the element X(i) in both directions
        X1(i) = X1(i) + epsilon;
        X2(i) = X2(i) - epsilon;
        % Central difference formula for gradient
        J(i)  = (f(X1) - f(X2)) / (2 * epsilon);
    end
end

% Line Search (Step Size Optimization)
function lambda = optimize_step_size(f, X, S)
    alpha       = 1;        % Initial step size
    rho         = 0.5;      % Step size reduction factor
    c           = 1E-4;     % Armijo condition constant
    max_iter    = 100;      % Max iterations for backtracking
    
    for k = 1:max_iter
        % Check if the Armijo condition is satisfied
        grad = compute_gradient(f, X);  % Compute gradient
        if f(X + alpha * S) <= f(X) + c * alpha * (S(:)' * grad(:))
            break;
        else
            alpha = rho * alpha;  % Reduce the step size
        end
    end
    lambda = alpha;  % Return the optimized step size
end

% Update Search Direction (Polak-Ribiere method)
function S_new  = update_search_direction(S, J_new, J_old)
    beta        = (norm(J_new(:))^2) / (norm(J_old(:))^2);
    S_new       = -J_new + beta * S;
end