%% NAFEMS Straight Cantilever GNL Benchmark
% 2D Corotational Beam Element - Geometrically Nonlinear Analysis
%
% Reference: NAFEMS R0065, Section 5.2 (Becker, 1999)
% Formulation: Crisfield's 2D corotational beam (Crisfield 1990, 1997)
%
% Tangent stiffness: K_T = df_int/dd = B^T * K_l * B + K_geo
%   where K_geo = (N/L)*z*z^T + (M1+M2)/L^2 * (r*z^T + z*r^T)

function NAFEMS_Cantilever_GNL()
    close all; clc;

    %% ==================== SECTION 1: PARAMETERS ====================
    E_mod   = 2.1e11;       % Young's modulus [N/m^2]
    nu      = 0.0;          % Poisson's ratio

    L_beam  = 3.2;          % Beam length [m]
    b_side  = 0.1;          % Cross-section side [m]
    h_side  = 0.1;          % Cross-section height [m]

    A_area  = b_side * h_side;              % Area [m^2]
    I_inertia = b_side * h_side^3 / 12.0;   % Second moment of area [m^4]

    Fx_total = -3.844e6;     % Total axial compressive load [N]
    Fy_total = -3.844e3;     % Total transverse perturbation load [N]

    n_elem = 40;   % finer mesh for accuracy under extreme curling

    %% ==================== SECTION 2: MESH GENERATION ====================
    [node_coords, elem_nodes, n_node, n_dof] = generate_mesh(L_beam, n_elem);

    node_dofs = zeros(n_node, 3);
    for i = 1:n_node
        node_dofs(i, :) = [3*i-2, 3*i-1, 3*i];
    end

    fixed_dofs = [1, 2, 3];
    free_dofs  = setdiff(1:n_dof, fixed_dofs);

    L0_vec  = zeros(n_elem, 1);
    c0_vec  = zeros(n_elem, 1);
    s0_vec  = zeros(n_elem, 1);
    for e = 1:n_elem
        n1 = elem_nodes(e, 1); n2 = elem_nodes(e, 2);
        X1 = node_coords(n1, 1); Y1 = node_coords(n1, 2);
        X2 = node_coords(n2, 1); Y2 = node_coords(n2, 2);
        dX = X2 - X1; dY = Y2 - Y1;
        L0_vec(e) = sqrt(dX^2 + dY^2);
        c0_vec(e) = dX / L0_vec(e);
        s0_vec(e) = dY / L0_vec(e);
    end

    elem_dofs = zeros(n_elem, 6);
    for e = 1:n_elem
        n1 = elem_nodes(e, 1); n2 = elem_nodes(e, 2);
        elem_dofs(e, :) = [node_dofs(n1, :), node_dofs(n2, :)];
    end

    tip_node = n_node;

    fprintf('Mesh: %d elements, %d nodes, %d DOFs\n', n_elem, n_node, n_dof);
    fprintf('Cross-section: A = %.6f m^2, I = %.6e m^4\n', A_area, I_inertia);
    fprintf('Euler buckling load (theoretical): P_cr = %.4e N\n', ...
            pi^2 * E_mod * I_inertia / (4 * L_beam^2));

    %% ==================== ANALYSIS A: LINEAR BUCKLING ====================
    fprintf('\n========== ANALYSIS A: LINEAR BUCKLING ==========\n');

    F_ref = zeros(n_dof, 1);
    F_ref(node_dofs(tip_node, 1)) = -1.0;

    K_E = sparse(n_dof, n_dof);
    for e = 1:n_elem
        L0 = L0_vec(e);
        c0 = c0_vec(e); s0 = s0_vec(e);
        edofs = elem_dofs(e, :);
        B0 = compute_B_matrix(c0, s0, L0);
        Kl = local_stiffness_matrix(E_mod, A_area, I_inertia, L0);
        Ke = B0' * Kl * B0;
        K_E(edofs, edofs) = K_E(edofs, edofs) + Ke;
    end

    d_ref = zeros(n_dof, 1);
    d_ref(free_dofs) = K_E(free_dofs, free_dofs) \ F_ref(free_dofs);

    K_G = sparse(n_dof, n_dof);
    for e = 1:n_elem
        L0 = L0_vec(e);
        c0 = c0_vec(e); s0 = s0_vec(e);
        edofs = elem_dofs(e, :);
        d_e = d_ref(edofs);
        B0 = compute_B_matrix(c0, s0, L0);
        d_local = B0 * d_e;
        fl = local_force_vector(E_mod, A_area, I_inertia, L0, d_local(1), d_local(2), d_local(3));
        N_axial = fl(1); M1 = fl(2); M2 = fl(3);
        z0 = [s0; -c0; 0; -s0; c0; 0];
        r0 = [-c0; -s0; 0; c0; s0; 0];
        Kgeo = compute_geometric_stiffness(N_axial, M1, M2, L0, z0, r0);
        K_G(edofs, edofs) = K_G(edofs, edofs) + Kgeo;
    end

    K_E_ff = K_E(free_dofs, free_dofs);
    K_G_ff = K_G(free_dofs, free_dofs);

    A_mat = full(K_E_ff \ K_G_ff);
    [eig_vecs, eig_vals_diag] = eig(A_mat);
    eig_vals = real(diag(eig_vals_diag));

    neg_idx = find(eig_vals < 0);
    eig_vals_neg = eig_vals(neg_idx);
    eig_vecs_neg = eig_vecs(:, neg_idx);

    [~, idx_min_mu] = min(eig_vals_neg);
    mu_cr = eig_vals_neg(idx_min_mu);
    lambda_cr = -1.0 / mu_cr;
    P_cr = lambda_cr * 1.0;

    phi_full = zeros(n_dof, 1);
    phi_full(free_dofs) = real(eig_vecs_neg(:, idx_min_mu));
    phi_full = phi_full / max(abs(phi_full));

    fprintf('First buckling eigenvalue: lambda_1 = %.4e\n', lambda_cr);
    fprintf('Critical buckling load:   P_cr   = %.4e N\n', P_cr);
    fprintf('Theoretical Euler load:   P_euler = %.4e N\n', ...
            pi^2 * E_mod * I_inertia / (4 * L_beam^2));

    %% ==================== ANALYSIS B: GNL ANALYSIS ====================
    fprintf('\n========== ANALYSIS B: GEOMETRIC NONLINEAR ANALYSIS ==========\n');

    % Newton-Raphson parameters
    max_iter    = 30;
    tol_res     = 1.0e-6;      % force residual tolerance
    tol_disp    = 1.0e-8;      % displacement increment tolerance
    max_halvings = 8;           % max step halvings before giving up

    % Load stepping: use gradually increasing steps
    n_steps_total = 500;
    NCL_vec = generate_load_steps(n_steps_total);

    % Storage
    max_store = length(NCL_vec) + 100;
    tip_u_hist = zeros(max_store, 1);
    tip_v_hist = zeros(max_store, 1);
    NCL_hist   = zeros(max_store, 1);

    % For intermediate deformed shapes
    n_intermediate = 20;
    intermediate_shapes = cell(n_intermediate, 1);
    intermediate_NCL    = zeros(n_intermediate, 1);
    shape_stride = max(1, floor(length(NCL_vec) / n_intermediate));

    d_global = zeros(n_dof, 1);
    conv_count = 0;
    ncl_current = 0;
    idx_step = 0;

    while idx_step < length(NCL_vec) && ncl_current < 1.0 - 1e-12
        idx_step = idx_step + 1;
        ncl_target = NCL_vec(idx_step);
        delta_ncl = ncl_target - ncl_current;

        % NR iteration with step halving if needed
        n_halvings = 0;
        converged = false;
        d_save = d_global;

        while n_halvings <= max_halvings
            d_global = d_save;

            % External force at target load level
            ncl_step_target = ncl_current + delta_ncl;
            F_ext = zeros(n_dof, 1);
            F_ext(node_dofs(tip_node, 1)) = Fx_total * ncl_step_target;
            F_ext(node_dofs(tip_node, 2)) = Fy_total * ncl_step_target;

            % Initial residual
            F_int = compute_global_internal_force(d_global, node_coords, elem_nodes, ...
                elem_dofs, n_elem, L0_vec, E_mod, A_area, I_inertia);
            R = F_ext - F_int;
            R(fixed_dofs) = 0;
            norm_R0 = norm(R(free_dofs));

            if norm_R0 < tol_res
                converged = true;
                break;
            end

            for iter = 1:max_iter
                % Assemble tangent stiffness
                K_T = assemble_tangent_stiffness(d_global, node_coords, elem_nodes, ...
                    elem_dofs, n_elem, L0_vec, E_mod, A_area, I_inertia);

                % Solve
                delta_d = zeros(n_dof, 1);
                delta_d(free_dofs) = K_T(free_dofs, free_dofs) \ R(free_dofs);

                % Update
                d_global = d_global + delta_d;

                % Residual at new state
                F_int = compute_global_internal_force(d_global, node_coords, elem_nodes, ...
                    elem_dofs, n_elem, L0_vec, E_mod, A_area, I_inertia);
                R = F_ext - F_int;
                R(fixed_dofs) = 0;
                norm_R = norm(R(free_dofs));
                norm_dd = norm(delta_d(free_dofs));

                if norm_R < tol_res || (norm_dd < tol_disp && norm_R < tol_res * 100)
                    converged = true;
                    break;
                end
            end

            if converged
                break;
            else
                n_halvings = n_halvings + 1;
                delta_ncl = delta_ncl / 2;
            end
        end

        if ~converged
            warning('Failed to converge at NCL=%.6f after %d halvings. Stopping.', ...
                    ncl_current + delta_ncl, n_halvings);
            break;
        end

        ncl_current = ncl_current + delta_ncl;
        d_save = d_global;

        conv_count = conv_count + 1;
        tip_u_hist(conv_count) = d_global(node_dofs(tip_node, 1));
        tip_v_hist(conv_count) = d_global(node_dofs(tip_node, 2));
        NCL_hist(conv_count) = ncl_current;

        if mod(conv_count, 20) == 0 || conv_count == 1
            fprintf('  NCL=%.4f | u=%.6f | v=%.6f | n_halvings=%d\n', ...
                    ncl_current, tip_u_hist(conv_count), tip_v_hist(conv_count), n_halvings);
        end

        % Store intermediate shape
        shape_idx = round(conv_count / shape_stride);
        if shape_idx >= 1 && shape_idx <= n_intermediate && isempty(intermediate_shapes{shape_idx})
            intermediate_shapes{shape_idx} = d_global;
            intermediate_NCL(shape_idx) = ncl_current;
        end
    end

    % Trim
    tip_u_hist = tip_u_hist(1:conv_count);
    tip_v_hist = tip_v_hist(1:conv_count);
    NCL_hist   = NCL_hist(1:conv_count);

    valid_shapes = ~cellfun(@isempty, intermediate_shapes);
    intermediate_shapes = intermediate_shapes(valid_shapes);
    intermediate_NCL = intermediate_NCL(valid_shapes);

    fprintf('\n=== GNL Analysis Complete ===\n');
    fprintf('Final tip displacement at NCL=%.4f:\n', ncl_current);
    fprintf('  u_tip = %.6f m  (reference: ~ -5.05 m)\n', tip_u_hist(end));
    fprintf('  v_tip = %.6f m  (reference: ~ -1.35 m)\n', tip_v_hist(end));

    %% ==================== VISUALIZATION ====================
    fprintf('\n========== GENERATING FIGURES ==========\n');

    % ---- Figure 1: Linear Buckling Mode ----
    figure(1); clf;
    set(gcf, 'Position', [100, 500, 700, 350]);
    X_undef = node_coords(:, 1);
    Y_undef = node_coords(:, 2);
    scale_buckle = 0.3;
    X_buckle = X_undef + scale_buckle * phi_full(node_dofs(:, 1));
    Y_buckle = Y_undef + scale_buckle * phi_full(node_dofs(:, 2));
    plot(X_undef, Y_undef, 'k-', 'LineWidth', 2); hold on;
    plot(X_buckle, Y_buckle, 'r--', 'LineWidth', 2);
    xlabel('X (m)'); ylabel('Y (m)');
    title(sprintf(['Figure 1: 1st Linear Buckling Mode\n' ...
                   'Critical Load Factor \\lambda_1 = %.4e  (P_{cr} = %.4e N)'], ...
                   lambda_cr, P_cr));
    legend('Undeformed', sprintf('1st Buckling Mode (scaled %.1fx)', scale_buckle), ...
           'Location', 'best');
    axis equal; grid on;

    % ---- Figure 2: Load-Displacement Curve ----
    figure(2); clf;
    set(gcf, 'Position', [850, 500, 600, 500]);
    yyaxis left;
    plot(NCL_hist, tip_u_hist, 'b-', 'LineWidth', 1.5);
    ylabel('Horizontal tip displacement u (m)');
    hold on;
    yyaxis right;
    plot(NCL_hist, tip_v_hist, 'r-', 'LineWidth', 1.5);
    ylabel('Vertical tip displacement v (m)');
    xlabel('Normalized Load Factor NCL');
    title('Figure 2: Tip Displacement vs. Normalized Load');
    legend('u (horizontal)', 'v (vertical)', 'Location', 'best');
    grid on;

    % ---- Figure 3: Large Deformation Process ----
    figure(3); clf;
    set(gcf, 'Position', [100, 50, 900, 500]);
    n_shapes = length(intermediate_shapes);
    cmap = jet(n_shapes + 1);
    hold on;
    for s = 1:n_shapes
        d_shape = intermediate_shapes{s};
        X_def = node_coords(:, 1) + d_shape(node_dofs(:, 1));
        Y_def = node_coords(:, 2) + d_shape(node_dofs(:, 2));
        plot(X_def, Y_def, '-', 'Color', [cmap(s, :) 0.4], 'LineWidth', 1.0);
    end
    X_final = node_coords(:, 1) + d_global(node_dofs(:, 1));
    Y_final = node_coords(:, 2) + d_global(node_dofs(:, 2));
    plot(X_final, Y_final, 'k-', 'LineWidth', 2.5);
    plot(node_coords(:, 1), node_coords(:, 2), 'k--', 'LineWidth', 1.0);
    xlabel('X (m)'); ylabel('Y (m)');
    title(sprintf(['Figure 3: Large Deformation Process\n' ...
                   'u_{tip} = %.3f m, v_{tip} = %.3f m at NCL = %.2f'], ...
                   tip_u_hist(end), tip_v_hist(end), ncl_current));
    legend_entries = cell(n_shapes + 2, 1);
    for s = 1:n_shapes
        legend_entries{s} = sprintf('NCL = %.2f', intermediate_NCL(s));
    end
    legend_entries{n_shapes+1} = sprintf('Final (NCL=%.2f)', ncl_current);
    legend_entries{n_shapes+2} = 'Undeformed';
    legend(legend_entries, 'Location', 'bestoutside');
    axis equal; grid on;

    % Save figures to PNG
    saveas(figure(1), 'Figure1_LinearBucklingMode.png');
    saveas(figure(2), 'Figure2_LoadDisplacementCurve.png');
    saveas(figure(3), 'Figure3_LargeDeformationProcess.png');
    fprintf('Figures saved to PNG files.\n');
    fprintf('\n========== PROGRAM COMPLETE ==========\n');
end

%% ==================== LOAD STEP GENERATION ====================

function NCL_vec = generate_load_steps(n_total)
% Generate load step sequence covering full NCL range [0, 1].
% Refined near buckling (NCL ~ 0.1).
    % Allocate more points near the buckling region using tanh spacing
    x = linspace(0, 1, n_total)';
    % Sigmoid to cluster points near NCL=0.1
    center = 0.10;
    width  = 0.05;
    x_transformed = 0.5 * (1 + tanh((x - center) / width));
    % Blend: weighted average of uniform and clustered spacing
    % Use more uniform spacing to ensure we reach NCL=1
    NCL_vec = x;  % uniform spacing, simple and robust
    NCL_vec = NCL_vec(2:end);  % skip NCL=0
end

%% ==================== ELEMENT ROUTINES ====================

function B = compute_B_matrix(c, s, L)
    B = zeros(3, 6);
    B(1, :) = [-c, -s, 0,  c,  s, 0];
    B(2, :) = [-s/L,  c/L, 1,  s/L, -c/L, 0];
    B(3, :) = [-s/L,  c/L, 0,  s/L, -c/L, 1];
end

function Kl = local_stiffness_matrix(E, A, I, L0)
    Kl = zeros(3, 3);
    Kl(1, 1) = E * A / L0;
    Kl(2, 2) = 4 * E * I / L0;
    Kl(2, 3) = 2 * E * I / L0;
    Kl(3, 2) = 2 * E * I / L0;
    Kl(3, 3) = 4 * E * I / L0;
end

function fl = local_force_vector(E, A, I, L0, u_bar, theta1_bar, theta2_bar)
    N  = E * A / L0 * u_bar;
    M1 = E * I / L0 * (4 * theta1_bar + 2 * theta2_bar);
    M2 = E * I / L0 * (2 * theta1_bar + 4 * theta2_bar);
    fl = [N; M1; M2];
end

function Kgeo = compute_geometric_stiffness(N_axial, M1, M2, L, z, r)
    z_col = z(:);
    r_col = r(:);
    Kgeo = (N_axial / L) * (z_col * z_col') ...
         + ((M1 + M2) / L^2) * (r_col * z_col' + z_col * r_col');
end

function F_int = compute_global_internal_force(d_global, node_coords, elem_nodes, ...
    elem_dofs, n_elem, L0_vec, E_mod, A_area, I_inertia)
    n_dof = length(d_global);
    F_int = zeros(n_dof, 1);
    for e = 1:n_elem
        edofs = elem_dofs(e, :);
        d_e = d_global(edofs);
        n1 = elem_nodes(e, 1); n2 = elem_nodes(e, 2);
        L0 = L0_vec(e);
        X1 = node_coords(n1, 1); Y1 = node_coords(n1, 2);
        X2 = node_coords(n2, 1); Y2 = node_coords(n2, 2);
        alpha0 = atan2(Y2 - Y1, X2 - X1);
        x1 = X1 + d_e(1); y1 = Y1 + d_e(2); theta1 = d_e(3);
        x2 = X2 + d_e(4); y2 = Y2 + d_e(5); theta2 = d_e(6);
        dx = x2 - x1; dy = y2 - y1;
        L = sqrt(dx^2 + dy^2);
        c = dx / L; s = dy / L;
        beta = atan2(dy, dx);
        u_bar = L - L0;
        alpha = beta - alpha0;
        theta1_bar = theta1 - alpha;
        theta2_bar = theta2 - alpha;
        B_mat = compute_B_matrix(c, s, L);
        fl = local_force_vector(E_mod, A_area, I_inertia, L0, u_bar, theta1_bar, theta2_bar);
        fg = B_mat' * fl;
        F_int(edofs) = F_int(edofs) + fg;
    end
end

function K_T = assemble_tangent_stiffness(d_global, node_coords, elem_nodes, ...
    elem_dofs, n_elem, L0_vec, E_mod, A_area, I_inertia)
    n_dof = length(d_global);
    K_T = sparse(n_dof, n_dof);
    for e = 1:n_elem
        edofs = elem_dofs(e, :);
        d_e = d_global(edofs);
        n1 = elem_nodes(e, 1); n2 = elem_nodes(e, 2);
        L0 = L0_vec(e);
        X1 = node_coords(n1, 1); Y1 = node_coords(n1, 2);
        X2 = node_coords(n2, 1); Y2 = node_coords(n2, 2);
        alpha0 = atan2(Y2 - Y1, X2 - X1);
        x1 = X1 + d_e(1); y1 = Y1 + d_e(2); theta1 = d_e(3);
        x2 = X2 + d_e(4); y2 = Y2 + d_e(5); theta2 = d_e(6);
        dx = x2 - x1; dy = y2 - y1;
        L = sqrt(dx^2 + dy^2);
        c = dx / L; s = dy / L;
        beta = atan2(dy, dx);
        u_bar = L - L0;
        alpha = beta - alpha0;
        theta1_bar = theta1 - alpha;
        theta2_bar = theta2 - alpha;
        B_mat = compute_B_matrix(c, s, L);
        Kl = local_stiffness_matrix(E_mod, A_area, I_inertia, L0);
        fl = local_force_vector(E_mod, A_area, I_inertia, L0, u_bar, theta1_bar, theta2_bar);
        N_axial = fl(1); M1 = fl(2); M2 = fl(3);
        Ke_mat = B_mat' * Kl * B_mat;
        z_vec = [s; -c; 0; -s; c; 0];
        r_vec = [-c; -s; 0; c; s; 0];
        Ke_geo = compute_geometric_stiffness(N_axial, M1, M2, L, z_vec, r_vec);
        K_T(edofs, edofs) = K_T(edofs, edofs) + Ke_mat + Ke_geo;
    end
end

%% ==================== MESH GENERATION ====================

function [node_coords, elem_nodes, n_node, n_dof] = generate_mesh(L_beam, n_elem)
    n_node = n_elem + 1;
    n_dof  = 3 * n_node;
    node_coords = zeros(n_node, 2);
    for i = 1:n_node
        node_coords(i, 1) = (i - 1) * L_beam / n_elem;
        node_coords(i, 2) = 0.0;
    end
    elem_nodes = zeros(n_elem, 2);
    for e = 1:n_elem
        elem_nodes(e, :) = [e, e + 1];
    end
end
