% Autoencoder for Dimensionality Reduction on Synthetic 3D Point Clouds
%
% Trains a fully-connected autoencoder to compress 3D point clouds (150-dim)
% into a 2D latent space, then compares the learned embedding to PCA.
%
% 5 shape classes: Sphere, Torus, Helix, Swiss Roll, Cube
% Each sample = 50 points x 3 coordinates = 150-dimensional vector
%
% Architecture:
%   Encoder: 150 -> 128 -> 64 -> 32 -> 2   (latent)
%   Decoder: 2   -> 32  -> 64 -> 128 -> 150 (reconstruction)

function autoencoder_pointcloud()
    %% =====================================================================
    %  1. Generate synthetic 3D point cloud dataset
    %  =====================================================================
    rng(7);

    nPerClass  = 200;
    nPtsCloud  = 50;     % points per cloud
    inputDim   = nPtsCloud * 3;  % flattened: 150
    classNames = {'Sphere','Torus','Helix','Swiss Roll','Cube'};
    nClasses   = length(classNames);
    nTotal     = nPerClass * nClasses;

    Data   = zeros(nTotal, inputDim);
    Labels = zeros(nTotal, 1);

    for c = 1:nClasses
        for s = 1:nPerClass
            idx = (c-1)*nPerClass + s;
            Labels(idx) = c;
            pts = generateShape(classNames{c}, nPtsCloud);
            % No rotation — fixed orientation so shape differences are easier to learn
            % Small noise for sample variety
            pts = pts + 0.05 * randn(size(pts));
            Data(idx, :) = pts(:)';  % flatten to row vector
        end
    end

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

    DataTrain   = Data(trainIdx, :);
    LabelsTrain = Labels(trainIdx);
    DataVal     = Data(valIdx, :);
    LabelsVal   = Labels(valIdx);
    nTrain      = size(DataTrain, 1);
    nVal        = size(DataVal, 1);

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

    %% =====================================================================
    %  2. Build autoencoder as dlnetwork
    %  =====================================================================
    % Encoder and decoder in one network: input -> bottleneck -> output
    % We implement as a custom forward function using learnable parameters

    latentDim = 2;
    encSizes = [inputDim, 128, 64, 32, latentDim];
    decSizes = [latentDim, 32, 64, 128, inputDim];

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

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

    % Adam state
    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. Pick example samples for reconstruction display (one per class)
    %  =====================================================================
    % Pick one example per class from validation set for reconstruction display
    exampleIdx = zeros(1, nClasses);
    for c = 1:nClasses
        candidates = find(LabelsVal == c);
        exampleIdx(c) = candidates(1);
    end

    %% =====================================================================
    %  5. Set up animated figure
    %  =====================================================================
    classColors = [0.85 0.33 0.10;   % orange
                   0.00 0.45 0.74;   % blue
                   0.47 0.67 0.19;   % green
                   0.64 0.08 0.18;   % crimson
                   0.56 0.27 0.68];  % purple

    fig1 = figure('Name','Autoencoder Training','Color','w',...
                  'Position',[60 40 1300 850]);

    % --- Top-left: 2D latent space ---
    ax_lat = subplot(2,2,1,'Parent',fig1);
    hold(ax_lat,'on'); grid(ax_lat,'on');
    hScatter = gobjects(nClasses, 1);
    for c = 1:nClasses
        hScatter(c) = scatter(ax_lat, NaN, NaN, 20, classColors(c,:), 'filled', ...
                              'MarkerFaceAlpha', 0.6);
    end
    xlabel(ax_lat,'z_1'); ylabel(ax_lat,'z_2');
    title(ax_lat,'Autoencoder 2D Latent Space','FontSize',13);
    legend(ax_lat, classNames, 'Location','best','FontSize',8);

    % --- Top-right: loss curves (train + val) ---
    ax_loss = subplot(2,2,2,'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',3,'MarkerFaceColor',[0 0.45 0.74]);
    hLossVl = animatedline(ax_loss, 'Color',[0.85 0.33 0.1], 'LineWidth',2, ...
                           'Marker','s','MarkerSize',3,'MarkerFaceColor',[0.85 0.33 0.1]);
    xlabel(ax_loss,'Epoch'); ylabel(ax_loss,'MSE Loss');
    title(ax_loss,'Reconstruction Loss','FontSize',13);
    legend(ax_loss, {'Train','Validation'}, 'Location','northeast');
    set(ax_loss,'YScale','log');

    % --- Bottom: example reconstructions (5 pairs: original + recon) ---
    axRecon = gobjects(nClasses, 2);
    for c = 1:nClasses
        % Original
        axRecon(c,1) = subplot(2, 2*nClasses, 2*nClasses + 2*(c-1) + 1, 'Parent', fig1);
        hold(axRecon(c,1),'on');
        title(axRecon(c,1), classNames{c}, 'FontSize', 8);
        view(axRecon(c,1), 30, 25);
        axis(axRecon(c,1), 'equal');
        set(axRecon(c,1), 'XTick',[],'YTick',[],'ZTick',[]);
        % Reconstructed
        axRecon(c,2) = subplot(2, 2*nClasses, 2*nClasses + 2*(c-1) + 2, 'Parent', fig1);
        hold(axRecon(c,2),'on');
        title(axRecon(c,2), 'Recon', 'FontSize', 8);
        view(axRecon(c,2), 30, 25);
        axis(axRecon(c,2), 'equal');
        set(axRecon(c,2), 'XTick',[],'YTick',[],'ZTick',[]);
    end

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

    drawnow;

    %% =====================================================================
    %  6. Training loop
    %  =====================================================================
    dlDataTrain = dlarray(single(DataTrainNorm)');  % [150 x nTrain]
    dlDataVal   = dlarray(single(DataValNorm)');    % [150 x nVal]
    trainLosses = [];
    valLosses   = [];

    for epoch = 1:nEpochs
        % Shuffle training data
        shuffleIdx = randperm(nTrain);
        epochLoss = 0;
        nBatches = 0;

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

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

            % Adam update
            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>

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

        % --- Update animation every 2 epochs ---
        if mod(epoch, 2) == 0 || epoch == 1 || epoch == nEpochs
            % Encode all data (train + val) for latent space plot
            Ztr  = extractdata(aeEncode(params, dlDataTrain, encSizes))';
            Zval = extractdata(aeEncode(params, dlDataVal,   encSizes))';

            % Update latent scatter (show val points as filled, train as faint)
            for c = 1:nClasses
                maskTr  = LabelsTrain == c;
                maskVal = LabelsVal == c;
                % Combine for display
                Zall = [Ztr(maskTr,:); Zval(maskVal,:)];
                set(hScatter(c), 'XData', Zall(:,1), 'YData', Zall(:,2));
            end
            % Auto-scale axes
            Zall = [Ztr; Zval];
            margin = 0.5;
            xlim(ax_lat, [min(Zall(:,1))-margin, max(Zall(:,1))+margin]);
            ylim(ax_lat, [min(Zall(:,2))-margin, max(Zall(:,2))+margin]);

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

            % Update reconstruction examples (from validation set)
            for c = 1:nClasses
                ei = exampleIdx(c);
                dlXi = dlDataVal(:, ei);
                [~, dlRecon] = aeForward(params, dlXi, encSizes, decSizes);
                reconVec = double(extractdata(dlRecon))';
                reconVec = reconVec .* dataStd + dataMean;
                reconPts = reshape(reconVec, [], 3);

                origPts = reshape(DataVal(ei,:), [], 3);

                % Original
                cla(axRecon(c,1));
                scatter3(axRecon(c,1), origPts(:,1), origPts(:,2), origPts(:,3), ...
                         15, classColors(c,:), 'filled');
                view(axRecon(c,1), 30, 25);
                axis(axRecon(c,1), 'equal');
                title(axRecon(c,1), classNames{c}, 'FontSize', 8);
                set(axRecon(c,1), 'XTick',[],'YTick',[],'ZTick',[]);

                % Reconstructed
                cla(axRecon(c,2));
                scatter3(axRecon(c,2), reconPts(:,1), reconPts(:,2), reconPts(:,3), ...
                         15, classColors(c,:), 'filled', 'MarkerFaceAlpha', 0.6);
                view(axRecon(c,2), 30, 25);
                axis(axRecon(c,2), 'equal');
                title(axRecon(c,2), 'Recon', 'FontSize', 8, 'Color', [0.4 0.4 0.4]);
                set(axRecon(c,2), 'XTick',[],'YTick',[],'ZTick',[]);
            end

            set(epochText, 'String', sprintf('Epoch: %d / %d  |  Train: %.5f  Val: %.5f', ...
                epoch, nEpochs, epochLoss, valLoss));

            drawnow;
        end
    end

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

    %% =====================================================================
    %  7. Final: Autoencoder vs PCA comparison (validation set)
    %  =====================================================================
    % Autoencoder latent codes on validation set
    ZaeVal = extractdata(aeEncode(params, dlDataVal, encSizes))';  % [nVal x 2]

    % PCA: fit on training data, project validation data
    [coeff, ~, ~, ~, ~, mu_pca] = pca(DataTrainNorm);
    ZpcaVal = (DataValNorm - mu_pca) * coeff(:, 1:2);

    fig2 = figure('Name','Autoencoder vs PCA (Validation)','Color','w',...
                  'Position',[100 100 1100 500]);

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

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

    sgtitle('Nonlinear Autoencoder vs Linear PCA  (Validation Set)','FontSize',15);
    saveas(fig2, fullfile(fileparts(mfilename('fullpath')), 'ae_vs_pca.png'));

    %% =====================================================================
    %  8. Final: Reconstruction gallery
    %  =====================================================================
    fig3 = figure('Name','Reconstruction Gallery','Color','w',...
                  'Position',[120 80 1300 500]);

    for c = 1:nClasses
        ei = exampleIdx(c);
        dlXi = dlDataVal(:, ei);
        [~, dlRecon] = aeForward(params, dlXi, encSizes, decSizes);
        reconVec = double(extractdata(dlRecon))';
        reconVec = reconVec .* dataStd + dataMean;
        reconPts = reshape(reconVec, [], 3);
        origPts  = reshape(DataVal(ei,:), [], 3);

        % Original
        ax = subplot(2, nClasses, c, 'Parent', fig3);
        scatter3(ax, origPts(:,1), origPts(:,2), origPts(:,3), ...
                 25, classColors(c,:), 'filled');
        view(ax, 30, 25); axis(ax, 'equal');
        title(ax, sprintf('%s (original)', classNames{c}), 'FontSize', 10);
        set(ax, 'XTick',[],'YTick',[],'ZTick',[]);

        % Reconstructed
        ax = subplot(2, nClasses, nClasses + c, 'Parent', fig3);
        scatter3(ax, reconPts(:,1), reconPts(:,2), reconPts(:,3), ...
                 25, classColors(c,:), 'filled', 'MarkerFaceAlpha', 0.7);
        view(ax, 30, 25); axis(ax, 'equal');
        title(ax, 'Reconstructed', 'FontSize', 10, 'Color', [0.4 0.4 0.4]);
        set(ax, 'XTick',[],'YTick',[],'ZTick',[]);
    end
    sgtitle('Original vs Reconstructed Point Clouds (Validation)', 'FontSize', 14);
    saveas(fig3, fullfile(fileparts(mfilename('fullpath')), 'ae_recon_gallery.png'));

    %% =====================================================================
    %  9. Print summary
    %  =====================================================================
    fprintf('\n=============================================\n');
    fprintf('  Autoencoder Point Cloud Summary\n');
    fprintf('=============================================\n');
    fprintf('Shapes: %s\n', strjoin(classNames, ', '));
    fprintf('Samples: %d per class (%d train / %d val)\n', nPerClass, nTrain, nVal);
    fprintf('Input dim: %d (50 pts x 3 coords)\n', inputDim);
    fprintf('Latent dim: 2\n');
    fprintf('Epochs: %d\n', nEpochs);
    fprintf('Final train MSE: %.6f\n', trainLosses(end));
    fprintf('Final val   MSE: %.6f\n', valLosses(end));
    fprintf('---------------------------------------------\n');
    fprintf('Per-class reconstruction MSE (validation):\n');
    for c = 1:nClasses
        mask = LabelsVal == c;
        dlXc = dlDataVal(:, mask);
        [~, dlRc] = aeForward(params, dlXc, encSizes, decSizes);
        classMSE = mean((dlXc - dlRc).^2, 'all');
        fprintf('  %-12s: %.6f\n', classNames{c}, extractdata(classMSE));
    end
    fprintf('=============================================\n\n');
end

%% =========================================================================
%  Autoencoder forward pass
%  =========================================================================
function [z, xRecon] = aeForward(params, x, encSizes, decSizes)
    % Encode
    z = aeEncode(params, x, encSizes);
    % Decode
    h = z;
    nDecLayers = length(decSizes) - 1;
    for i = 1:nDecLayers
        W = params.(['dec_W' num2str(i)]);
        b = params.(['dec_b' num2str(i)]);
        h = W * h + b;
        if i < nDecLayers
            h = relu(h);  % no activation on output layer
        end
    end
    xRecon = h;
end

%% =========================================================================
%  Encoder only (for extracting latent codes)
%  =========================================================================
function z = aeEncode(params, x, encSizes)
    h = x;
    nEncLayers = length(encSizes) - 1;
    for i = 1:nEncLayers
        W = params.(['enc_W' num2str(i)]);
        b = params.(['enc_b' num2str(i)]);
        h = W * h + b;
        if i < nEncLayers
            h = relu(h);  % no activation on bottleneck (latent is unbounded)
        end
    end
    z = h;
end

%% =========================================================================
%  Gradient computation
%  =========================================================================
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

%% =========================================================================
%  Shape generators: return [nPts x 3]
%  =========================================================================
function pts = generateShape(shapeName, nPts)
    switch shapeName
        case 'Sphere'
            % Evenly-spaced points on unit sphere (golden spiral)
            indices = (0:nPts-1)' + 0.5;
            phi   = acos(1 - 2 * indices / nPts);
            theta = pi * (1 + sqrt(5)) * indices;
            r = 0.8 + 0.2 * rand;  % slight radius variation
            pts = r * [sin(phi).*cos(theta), sin(phi).*sin(theta), cos(phi)];

        case 'Torus'
            R = 1.0;  % major radius
            rr = 0.3 + 0.1 * rand;  % minor radius with variation
            % Evenly-spaced parametric grid
            phi   = linspace(0, 2*pi, nPts)' + 0.1*rand;
            theta = linspace(0, 4*pi, nPts)' + 0.1*rand;
            pts = [(R + rr*cos(theta)).*cos(phi), ...
                   (R + rr*cos(theta)).*sin(phi), ...
                   rr*sin(theta)];

        case 'Helix'
            nTurns = 2 + rand;  % 2-3 turns
            t = linspace(0, nTurns * 2 * pi, nPts)';
            r = 0.5 + 0.15 * rand;
            pts = [r*cos(t), r*sin(t), t / (2*pi) * 0.5];

        case 'Swiss Roll'
            % Evenly-spaced parameter then sorted
            t = linspace(1.5*pi, 4.5*pi, nPts)';
            h = linspace(0, 2, nPts)' .* (0.8 + 0.4*rand);  % height variation
            pts = [t.*cos(t) / (3*pi), h, t.*sin(t) / (3*pi)];

        case 'Cube'
            % Distribute points evenly across 6 faces, sorted by face
            pts = zeros(nPts, 3);
            ptsPerFace = ceil(nPts / 6);
            idx = 0;
            for face = 1:6
                nF = min(ptsPerFace, nPts - idx);
                side = ceil(sqrt(nF));
                [uu, vv] = meshgrid(linspace(-0.5,0.5,side), linspace(-0.5,0.5,side));
                uu = uu(1:nF); vv = vv(1:nF);
                switch face
                    case 1, pts(idx+1:idx+nF,:) = [ 0.5*ones(nF,1), uu(:), vv(:)];
                    case 2, pts(idx+1:idx+nF,:) = [-0.5*ones(nF,1), uu(:), vv(:)];
                    case 3, pts(idx+1:idx+nF,:) = [uu(:),  0.5*ones(nF,1), vv(:)];
                    case 4, pts(idx+1:idx+nF,:) = [uu(:), -0.5*ones(nF,1), vv(:)];
                    case 5, pts(idx+1:idx+nF,:) = [uu(:), vv(:),  0.5*ones(nF,1)];
                    case 6, pts(idx+1:idx+nF,:) = [uu(:), vv(:), -0.5*ones(nF,1)];
                end
                idx = idx + nF;
            end
            pts = pts(1:nPts, :);
    end
end

%% =========================================================================
%  Random 3D rotation matrix
%  =========================================================================
function R = randomRotation()
    % Random rotation via QR decomposition of random matrix
    [Q, Rr] = qr(randn(3));
    R = Q * diag(sign(diag(Rr)));  % ensure proper rotation (det = +1)
    if det(R) < 0
        R(:,1) = -R(:,1);
    end
end
