% Autoencoder vs PCA: Curved Decision Boundaries
%
% Demonstrates that an autoencoder can recover nonlinear structure that
% PCA fundamentally cannot. Data lives on a 2D manifold embedded
% nonlinearly in 50 dimensions. Classes are separated by circular
% boundaries in the hidden 2D space -- no linear projection can
% untangle them.
%
% 3 classes defined by concentric rings in latent space:
%   Class 1 (inner):  s^2 + t^2 < r1^2
%   Class 2 (middle): r1^2 < s^2 + t^2 < r2^2
%   Class 3 (outer):  s^2 + t^2 > r2^2
%
% The mapping from 2D latent to 50D observed is nonlinear (sines,
% cosines, products), so PCA cannot recover the circular structure.

function autoencoder_curved_boundary()
    %% =====================================================================
    %  1. Generate data on a nonlinear manifold
    %  =====================================================================
    rng(5);

    nSamples  = 1500;
    obsDim    = 50;      % observed dimensionality
    latentDim = 2;       % autoencoder bottleneck (matches true latent dim)

    % Class boundaries in latent space: concentric circles
    r1 = 1.0;   % inner boundary
    r2 = 1.8;   % outer boundary

    % Generate latent coordinates uniformly in a disk of radius 2.5
    S = zeros(nSamples, 1);
    T = zeros(nSamples, 1);
    Labels = zeros(nSamples, 1);
    classNames = {'Inner', 'Ring', 'Outer'};
    nClasses = 3;

    % Sample roughly equal class sizes by rejection sampling
    count = 0;
    while count < nSamples
        s = 5 * rand - 2.5;
        t = 5 * rand - 2.5;
        r = sqrt(s^2 + t^2);
        if r > 2.8, continue; end  % clip to bounded region
        count = count + 1;
        S(count) = s;
        T(count) = t;
        if r < r1
            Labels(count) = 1;
        elseif r < r2
            Labels(count) = 2;
        else
            Labels(count) = 3;
        end
    end

    % Generate random nonlinear embedding: (s,t) -> nFeatures -> obsDim
    rng(42);
    nFeatures = 20;
    a = randn(nFeatures, 1);
    b = randn(nFeatures, 1);
    c = randn(nFeatures, 1);
    d = randn(nFeatures, 1);
    W = randn(obsDim, nFeatures);  % linear mixing from features to observed

    % Map latent (s,t) -> 50D via nonlinear features + linear mixing
    Data = zeros(nSamples, obsDim);
    for i = 1:nSamples
        si = S(i); ti = T(i);
        features = [sin(a .* si + b .* ti); ...      % sinusoidal [nF x 1]
                    cos(c .* si .* ti); ...            % product    [nF x 1]
                    sin(d .* si.^2 + a .* ti.^2); ... % quadratic  [nF x 1]
                    tanh(b .* ti - c .* si)];          % nonlinear  [nF x 1]
        % features is [4*nFeatures x 1], take first nFeatures via mixing
        phi = features(1:nFeatures);
        Data(i, :) = (W * phi)' + 0.1 * randn(1, obsDim);
    end
    rng(5);  % restore seed for reproducibility of splits

    % 80/20 train/val split per class
    trainIdx = []; valIdx = [];
    for c = 1:nClasses
        cIdx = find(Labels == c);
        cIdx = cIdx(randperm(length(cIdx)));
        nTr = round(0.8 * length(cIdx));
        trainIdx = [trainIdx; cIdx(1:nTr)];     %#ok<AGROW>
        valIdx   = [valIdx;   cIdx(nTr+1:end)];  %#ok<AGROW>
    end
    trainIdx = trainIdx(randperm(length(trainIdx)));
    valIdx   = valIdx(randperm(length(valIdx)));

    DataTrain = Data(trainIdx,:);   LabelsTrain = Labels(trainIdx);
    DataVal   = Data(valIdx,:);     LabelsVal   = Labels(valIdx);
    STrain = S(trainIdx);  TTrain = T(trainIdx);
    SVal   = S(valIdx);    TVal   = T(valIdx);
    nTrain = length(trainIdx);
    nVal   = length(valIdx);

    % Normalize using training stats
    dataMean = mean(DataTrain, 1);
    dataStd  = std(DataTrain, 0, 1) + 1e-8;
    DataTrainNorm = (DataTrain - dataMean) ./ dataStd;
    DataValNorm   = (DataVal   - dataMean) ./ dataStd;

    %% =====================================================================
    %  2. Build autoencoder
    %  =====================================================================
    encSizes = [obsDim, 64, 32, 16, latentDim];
    decSizes = [latentDim, 16, 32, 64, obsDim];

    params = struct;
    for i = 1:length(encSizes)-1
        fi = encSizes(i); fo = encSizes(i+1);
        params.(['enc_W' num2str(i)]) = dlarray(single(randn(fo,fi) * sqrt(2/fi)));
        params.(['enc_b' num2str(i)]) = dlarray(single(zeros(fo,1)));
    end
    for i = 1:length(decSizes)-1
        fi = decSizes(i); fo = decSizes(i+1);
        params.(['dec_W' num2str(i)]) = dlarray(single(randn(fo,fi) * sqrt(2/fi)));
        params.(['dec_b' num2str(i)]) = dlarray(single(zeros(fo,1)));
    end

    %% =====================================================================
    %  3. Training parameters
    %  =====================================================================
    nEpochs   = 300;
    miniBatch = 64;
    lr        = 1e-3;
    beta1 = 0.9; beta2 = 0.999; epsilon = 1e-8;

    adamM = structfun(@(x) dlarray(zeros(size(x),'single')), params, 'UniformOutput', false);
    adamV = structfun(@(x) dlarray(zeros(size(x),'single')), params, 'UniformOutput', false);
    adamT = 0;

    %% =====================================================================
    %  4. Set up figure
    %  =====================================================================
    classColors = [0.00 0.45 0.74;   % inner: blue
                   0.85 0.33 0.10;   % ring: orange
                   0.47 0.67 0.19];  % outer: green

    fig1 = figure('Name','Curved Boundary: AE Training','Color','w',...
                  'Position',[50 40 1400 800]);

    % --- Top-left: true latent space (ground truth) ---
    ax_true = subplot(2,3,1,'Parent',fig1);
    hold(ax_true,'on'); grid(ax_true,'on');
    for c = 1:nClasses
        mask = LabelsTrain == c;
        scatter(ax_true, STrain(mask), TTrain(mask), 15, classColors(c,:), 'filled', ...
                'MarkerFaceAlpha', 0.5);
    end
    % Draw boundary circles
    th = linspace(0, 2*pi, 200);
    plot(ax_true, r1*cos(th), r1*sin(th), 'k--', 'LineWidth', 1.5);
    plot(ax_true, r2*cos(th), r2*sin(th), 'k--', 'LineWidth', 1.5);
    axis(ax_true, 'equal');
    xlabel(ax_true,'s'); ylabel(ax_true,'t');
    title(ax_true,'True Latent Space (hidden)','FontSize',12);
    legend(ax_true, classNames, 'Location','northeast','FontSize',8);

    % --- Top-center: autoencoder latent space ---
    ax_ae = subplot(2,3,2,'Parent',fig1);
    hold(ax_ae,'on'); grid(ax_ae,'on');
    hScatAE = gobjects(nClasses,1);
    for c = 1:nClasses
        hScatAE(c) = scatter(ax_ae, NaN, NaN, 15, classColors(c,:), 'filled', ...
                             'MarkerFaceAlpha', 0.6);
    end
    xlabel(ax_ae,'z_1'); ylabel(ax_ae,'z_2');
    title(ax_ae,'Autoencoder Latent (learned)','FontSize',12);

    % --- Top-right: PCA projection ---
    ax_pca = subplot(2,3,3,'Parent',fig1);
    hold(ax_pca,'on'); grid(ax_pca,'on');
    [coeff, ~, ~, ~, ~, mu_pca] = pca(DataTrainNorm);
    ZpcaVal = (DataValNorm - mu_pca) * coeff(:,1:2);
    for c = 1:nClasses
        mask = LabelsVal == c;
        scatter(ax_pca, ZpcaVal(mask,1), ZpcaVal(mask,2), 15, classColors(c,:), 'filled', ...
                'MarkerFaceAlpha', 0.6);
    end
    xlabel(ax_pca,'PC_1'); ylabel(ax_pca,'PC_2');
    title(ax_pca,'PCA Projection (2D)','FontSize',12);
    legend(ax_pca, classNames, 'Location','best','FontSize',8);

    % --- Bottom-left: loss curves ---
    ax_loss = subplot(2,3,4,'Parent',fig1);
    hold(ax_loss,'on'); grid(ax_loss,'on');
    hLossTr = animatedline(ax_loss,'Color',[0 0.45 0.74],'LineWidth',2,...
                           'Marker','o','MarkerSize',2,'MarkerFaceColor',[0 0.45 0.74]);
    hLossVl = animatedline(ax_loss,'Color',[0.85 0.33 0.1],'LineWidth',2,...
                           'Marker','s','MarkerSize',2,'MarkerFaceColor',[0.85 0.33 0.1]);
    xlabel(ax_loss,'Epoch'); ylabel(ax_loss,'MSE');
    title(ax_loss,'Reconstruction Loss','FontSize',12);
    legend(ax_loss,{'Train','Validation'},'Location','northeast');
    set(ax_loss,'YScale','log');

    % --- Bottom-center: KNN accuracy comparison (bar chart, updated) ---
    ax_knn = subplot(2,3,5,'Parent',fig1);
    hold(ax_knn,'on'); grid(ax_knn,'on');
    title(ax_knn,'KNN Classification Accuracy','FontSize',12);
    ylabel(ax_knn,'Accuracy (%)');

    % --- Bottom-right: latent space with decision boundary attempt ---
    ax_bnd = subplot(2,3,6,'Parent',fig1);
    hold(ax_bnd,'on'); grid(ax_bnd,'on');
    xlabel(ax_bnd,'z_1'); ylabel(ax_bnd,'z_2');
    title(ax_bnd,'AE Latent (Val) with True Labels','FontSize',12);

    epochText = annotation(fig1,'textbox',[0.35 0.005 0.3 0.03],...
        'String','Epoch: 0','FontSize',12,'EdgeColor','none',...
        'HorizontalAlignment','center','FontWeight','bold');

    drawnow;

    %% =====================================================================
    %  5. Training loop
    %  =====================================================================
    dlTrain = dlarray(single(DataTrainNorm)');
    dlVal   = dlarray(single(DataValNorm)');
    trainLosses = []; valLosses = [];

    for epoch = 1:nEpochs
        shuffleIdx = randperm(nTrain);
        epochLoss = 0; nBatches = 0;

        for bStart = 1:miniBatch:nTrain
            bEnd = min(bStart + miniBatch - 1, nTrain);
            bIdx = shuffleIdx(bStart:bEnd);
            dlX = dlTrain(:, bIdx);

            [loss, grads] = dlfeval(@aeGradients, params, dlX, encSizes, decSizes);

            adamT = adamT + 1;
            fnames = fieldnames(grads);
            for fi = 1:length(fnames)
                fn = fnames{fi};
                adamM.(fn) = beta1 * adamM.(fn) + (1-beta1) * grads.(fn);
                adamV.(fn) = beta2 * adamV.(fn) + (1-beta2) * grads.(fn).^2;
                mHat = adamM.(fn) / (1 - beta1^adamT);
                vHat = adamV.(fn) / (1 - beta2^adamT);
                params.(fn) = params.(fn) - lr * mHat ./ (sqrt(vHat) + epsilon);
            end

            epochLoss = epochLoss + extractdata(loss);
            nBatches = nBatches + 1;
        end

        epochLoss = epochLoss / nBatches;
        trainLosses(end+1) = epochLoss; %#ok<AGROW>

        [~, dlValRecon] = aeForward(params, dlVal, encSizes, decSizes);
        valLoss = extractdata(mean((dlVal - dlValRecon).^2, 'all'));
        valLosses(end+1) = valLoss; %#ok<AGROW>

        % --- Update animation every 5 epochs ---
        if mod(epoch, 5) == 0 || epoch == 1 || epoch == nEpochs
            % Autoencoder latent codes
            ZaeTr  = extractdata(aeEncode(params, dlTrain, encSizes))';
            ZaeVal = extractdata(aeEncode(params, dlVal,   encSizes))';

            % Update AE latent scatter (training data)
            for c = 1:nClasses
                mask = LabelsTrain == c;
                set(hScatAE(c), 'XData', ZaeTr(mask,1), 'YData', ZaeTr(mask,2));
            end
            margin = 0.5;
            allZ = [ZaeTr; ZaeVal];
            xlim(ax_ae, [min(allZ(:,1))-margin, max(allZ(:,1))+margin]);
            ylim(ax_ae, [min(allZ(:,2))-margin, max(allZ(:,2))+margin]);

            % Loss curves
            addpoints(hLossTr, epoch, epochLoss);
            addpoints(hLossVl, epoch, valLoss);
            allL = [trainLosses, valLosses];
            ylim(ax_loss, [min(allL)*0.5, max(allL)*2]);

            % KNN accuracy (k=5) on validation set
            knnK = 5;
            % AE latent
            aeAcc = knnAccuracy(ZaeTr, LabelsTrain, ZaeVal, LabelsVal, knnK);
            % PCA
            ZpcaTr = (DataTrainNorm - mu_pca) * coeff(:,1:2);
            pcaAcc = knnAccuracy(ZpcaTr, LabelsTrain, ZpcaVal, LabelsVal, knnK);
            % Raw data
            rawAcc = knnAccuracy(DataTrainNorm, LabelsTrain, DataValNorm, LabelsVal, knnK);

            % Update bar chart
            cla(ax_knn);
            barData = [pcaAcc, aeAcc, rawAcc];
            bh = bar(ax_knn, barData, 0.6);
            bh.FaceColor = 'flat';
            bh.CData = [0.64 0.08 0.18; 0 0.45 0.74; 0.5 0.5 0.5];
            set(ax_knn, 'XTickLabel', {'PCA 2D','AE 2D','Raw 50D'}, 'FontSize', 9);
            ylim(ax_knn, [0 105]);
            ylabel(ax_knn, 'Accuracy (%)');
            title(ax_knn, sprintf('KNN (k=%d) Accuracy', knnK), 'FontSize', 12);
            % Add value labels on bars
            for bi = 1:3
                text(ax_knn, bi, barData(bi)+2, sprintf('%.1f%%', barData(bi)), ...
                     'HorizontalAlignment','center','FontSize',10,'FontWeight','bold');
            end

            % Update bottom-right: AE latent for validation with labels
            cla(ax_bnd);
            for c = 1:nClasses
                mask = LabelsVal == c;
                scatter(ax_bnd, ZaeVal(mask,1), ZaeVal(mask,2), 25, classColors(c,:), ...
                        'filled', 'MarkerFaceAlpha', 0.7);
            end
            xlabel(ax_bnd,'z_1'); ylabel(ax_bnd,'z_2');
            title(ax_bnd,'AE Latent (Validation)','FontSize',12);
            xlim(ax_bnd, [min(allZ(:,1))-margin, max(allZ(:,1))+margin]);
            ylim(ax_bnd, [min(allZ(:,2))-margin, max(allZ(:,2))+margin]);

            set(epochText, 'String', sprintf(...
                'Epoch: %d / %d  |  Train: %.5f  Val: %.5f  |  KNN: PCA=%.1f%%  AE=%.1f%%', ...
                epoch, nEpochs, epochLoss, valLoss, pcaAcc, aeAcc));

            drawnow;
        end
    end

    saveas(fig1, fullfile(fileparts(mfilename('fullpath')), 'ae_curved_training.png'));

    %% =====================================================================
    %  6. Final comparison figure
    %  =====================================================================
    ZaeVal = extractdata(aeEncode(params, dlVal, encSizes))';

    fig2 = figure('Name','Curved Boundary: AE vs PCA','Color','w',...
                  'Position',[80 80 1500 450]);

    % True latent
    ax1 = subplot(1,3,1,'Parent',fig2);
    hold(ax1,'on'); grid(ax1,'on');
    for c = 1:nClasses
        mask = LabelsVal == c;
        scatter(ax1, SVal(mask), TVal(mask), 25, classColors(c,:), 'filled', ...
                'MarkerFaceAlpha', 0.7);
    end
    plot(ax1, r1*cos(th), r1*sin(th), 'k--', 'LineWidth', 2);
    plot(ax1, r2*cos(th), r2*sin(th), 'k--', 'LineWidth', 2);
    axis(ax1,'equal'); xlabel(ax1,'s'); ylabel(ax1,'t');
    title(ax1,'True Latent Space','FontSize',14);
    legend(ax1, classNames, 'Location','northeast');

    % Autoencoder
    ax2 = subplot(1,3,2,'Parent',fig2);
    hold(ax2,'on'); grid(ax2,'on');
    for c = 1:nClasses
        mask = LabelsVal == c;
        scatter(ax2, ZaeVal(mask,1), ZaeVal(mask,2), 25, classColors(c,:), 'filled', ...
                'MarkerFaceAlpha', 0.7);
    end
    xlabel(ax2,'z_1'); ylabel(ax2,'z_2');
    title(ax2,'Autoencoder Latent (2D)','FontSize',14);
    legend(ax2, classNames, 'Location','best');

    % PCA
    ax3 = subplot(1,3,3,'Parent',fig2);
    hold(ax3,'on'); grid(ax3,'on');
    for c = 1:nClasses
        mask = LabelsVal == c;
        scatter(ax3, ZpcaVal(mask,1), ZpcaVal(mask,2), 25, classColors(c,:), 'filled', ...
                'MarkerFaceAlpha', 0.7);
    end
    xlabel(ax3,'PC_1'); ylabel(ax3,'PC_2');
    title(ax3,'PCA Projection (2D)','FontSize',14);
    legend(ax3, classNames, 'Location','best');

    sgtitle('Curved Boundary: Autoencoder Recovers Nonlinear Structure, PCA Cannot', 'FontSize', 15);
    saveas(fig2, fullfile(fileparts(mfilename('fullpath')), 'ae_curved_comparison.png'));

    %% =====================================================================
    %  7. Summary
    %  =====================================================================
    knnK = 5;
    ZaeTr = extractdata(aeEncode(params, dlTrain, encSizes))';
    ZpcaTr = (DataTrainNorm - mu_pca) * coeff(:,1:2);
    aeAcc  = knnAccuracy(ZaeTr, LabelsTrain, ZaeVal, LabelsVal, knnK);
    pcaAcc = knnAccuracy(ZpcaTr, LabelsTrain, ZpcaVal, LabelsVal, knnK);
    rawAcc = knnAccuracy(DataTrainNorm, LabelsTrain, DataValNorm, LabelsVal, knnK);

    fprintf('\n=============================================\n');
    fprintf('  Curved Boundary: AE vs PCA Summary\n');
    fprintf('=============================================\n');
    fprintf('Data: 50D nonlinear embedding of 2D latent space\n');
    fprintf('Classes: 3 concentric rings (circular boundaries)\n');
    fprintf('Samples: %d train / %d val\n', nTrain, nVal);
    fprintf('Epochs: %d\n', nEpochs);
    fprintf('Final train MSE: %.6f\n', trainLosses(end));
    fprintf('Final val   MSE: %.6f\n', valLosses(end));
    fprintf('---------------------------------------------\n');
    fprintf('KNN (k=%d) Classification Accuracy:\n', knnK);
    fprintf('  PCA (2D):          %.1f%%\n', pcaAcc);
    fprintf('  Autoencoder (2D):  %.1f%%\n', aeAcc);
    fprintf('  Raw (50D):         %.1f%%\n', rawAcc);
    fprintf('=============================================\n\n');
end

%% =========================================================================
%  Autoencoder forward / encode / gradients
%  =========================================================================
function [z, xRecon] = aeForward(params, x, encSizes, decSizes)
    z = aeEncode(params, x, encSizes);
    h = z;
    nDec = length(decSizes) - 1;
    for i = 1:nDec
        h = params.(['dec_W' num2str(i)]) * h + params.(['dec_b' num2str(i)]);
        if i < nDec, h = relu(h); end
    end
    xRecon = h;
end

function z = aeEncode(params, x, encSizes)
    h = x;
    nEnc = length(encSizes) - 1;
    for i = 1:nEnc
        h = params.(['enc_W' num2str(i)]) * h + params.(['enc_b' num2str(i)]);
        if i < nEnc, h = relu(h); end
    end
    z = h;
end

function [loss, grads] = aeGradients(params, x, encSizes, decSizes)
    [~, xRecon] = aeForward(params, x, encSizes, decSizes);
    loss = mean((x - xRecon).^2, 'all');
    grads = dlgradient(loss, params);
end

%% =========================================================================
%  KNN classification accuracy
%  =========================================================================
function acc = knnAccuracy(Xtrain, Ytrain, Xtest, Ytest, k)
    nTest = size(Xtest, 1);
    correct = 0;
    for i = 1:nTest
        dists = sum((Xtrain - Xtest(i,:)).^2, 2);
        [~, idx] = sort(dists);
        neighbors = Ytrain(idx(1:k));
        pred = mode(neighbors);
        if pred == Ytest(i)
            correct = correct + 1;
        end
    end
    acc = 100 * correct / nTest;
end
