% CNN Classification: Train a 2D Convolutional Neural Network on image-like
% patches for a multi-class classification problem.
%
% Problem: classify 2D points (x1, x2) into one of 4 classes defined by
% a nonlinear boundary (concentric rings + angular split).
%
% Each point is encoded as a small multi-channel image patch, then fed to a
% CNN. Animation shows the decision boundary evolving during training.

function cnn_classification()
    %% =====================================================================
    %  1. Generate 2D classification data
    %  =====================================================================
    rng(1);

    nPerClass = 300;
    nClasses  = 4;
    nTotal    = nPerClass * nClasses;

    [xTrain, yTrain] = generateSpiralData(nPerClass, nClasses, 0.15);
    [xVal,   yVal]   = generateSpiralData(round(nPerClass*0.3), nClasses, 0.15);

    yTrainCat = categorical(yTrain);
    yValCat   = categorical(yVal);

    %% =====================================================================
    %  2. Encode 2D points as multi-channel image patches
    %  =====================================================================
    patchSize = 9;   % spatial size of the patch (patchSize x patchSize)
    freqs = [1, 2, 4];
    nChannels = 2 * length(freqs) * 2 + 2;  % sin/cos for each freq for x1,x2 + raw x1,x2

    XtrainDL = encode2DPatches(xTrain, freqs, patchSize);
    XvalDL   = encode2DPatches(xVal,   freqs, patchSize);

    %% =====================================================================
    %  3. Define CNN architecture for classification
    %  =====================================================================
    layers = [
        imageInputLayer([patchSize patchSize nChannels], 'Name','input', ...
                        'Normalization','none')

        convolution2dLayer(3, 32, 'Padding','same', 'Name','conv1')
        batchNormalizationLayer('Name','bn1')
        reluLayer('Name','relu1')

        convolution2dLayer(3, 64, 'Padding','same', 'Name','conv2')
        batchNormalizationLayer('Name','bn2')
        reluLayer('Name','relu2')

        convolution2dLayer(3, 64, 'Padding','same', 'Name','conv3')
        batchNormalizationLayer('Name','bn3')
        reluLayer('Name','relu3')

        globalAveragePooling2dLayer('Name','gap')

        fullyConnectedLayer(64, 'Name','fc1')
        reluLayer('Name','relu_fc')
        dropoutLayer(0.3, 'Name','drop')
        fullyConnectedLayer(nClasses, 'Name','fc_out')

        softmaxLayer('Name','softmax')
        classificationLayer('Name','output')
    ];

    %% =====================================================================
    %  4. Prepare decision boundary grid
    %  =====================================================================
    gridRes = 120;
    gx = linspace(-2.2, 2.2, gridRes);
    gy = linspace(-2.2, 2.2, gridRes);
    [GX, GY] = meshgrid(gx, gy);
    gridPts = [GX(:), GY(:)];
    XgridDL = encode2DPatches(gridPts, freqs, patchSize);

    %% =====================================================================
    %  5. Set up animated figure
    %  =====================================================================
    classColors = [0.85 0.33 0.1;    % class 1: orange
                   0 0.45 0.74;      % class 2: blue
                   0.47 0.67 0.19;   % class 3: green
                   0.64 0.08 0.18];  % class 4: crimson

    fig = figure('Name','CNN Classification Training','Color','w',...
                 'Position',[80 40 1200 850]);

    % --- Panel 1: decision boundary + data (large) ---
    ax1 = subplot(2,2,[1 2],'Parent',fig);
    hold(ax1,'on');
    hImg = imagesc(ax1, gx, gy, zeros(gridRes, gridRes, 3));
    set(hImg, 'AlphaData', 0.35);
    % Plot training data by class
    for c = 1:nClasses
        mask = yTrain == c;
        plot(ax1, xTrain(mask,1), xTrain(mask,2), 'o', ...
             'MarkerSize', 4, 'MarkerFaceColor', classColors(c,:), ...
             'MarkerEdgeColor', 'w', 'LineWidth', 0.3);
    end
    xlim(ax1, [-2.2 2.2]); ylim(ax1, [-2.2 2.2]);
    axis(ax1, 'equal'); axis(ax1, 'tight');
    xlabel(ax1,'x_1'); ylabel(ax1,'x_2');
    title(ax1,'Decision Boundary Evolution','FontSize',14);
    epochText = text(ax1, -2.1, 2.0, 'Epoch: 0', 'FontSize',12, ...
                     'FontWeight','bold', 'Color','w', ...
                     'BackgroundColor',[0.2 0.2 0.2 0.7]);

    % --- Panel 2: loss curves ---
    ax2 = subplot(2,2,3,'Parent',fig);
    hold(ax2,'on'); grid(ax2,'on');
    hLossTrain = animatedline(ax2, 'Color',[0 0.45 0.74], 'LineWidth',2, ...
                              'Marker','o','MarkerSize',3,'MarkerFaceColor',[0 0.45 0.74]);
    hLossVal   = animatedline(ax2, 'Color',[0.85 0.33 0.1], 'LineWidth',2, ...
                              'Marker','s','MarkerSize',3,'MarkerFaceColor',[0.85 0.33 0.1]);
    xlabel(ax2,'Epoch'); ylabel(ax2,'Cross-Entropy Loss');
    title(ax2,'Training & Validation Loss','FontSize',13);
    legend(ax2,{'Train','Validation'},'Location','northeast');

    % --- Panel 3: accuracy curves ---
    ax3 = subplot(2,2,4,'Parent',fig);
    hold(ax3,'on'); grid(ax3,'on');
    hAccTrain = animatedline(ax3, 'Color',[0 0.45 0.74], 'LineWidth',2, ...
                             'Marker','o','MarkerSize',3,'MarkerFaceColor',[0 0.45 0.74]);
    hAccVal   = animatedline(ax3, 'Color',[0.85 0.33 0.1], 'LineWidth',2, ...
                             'Marker','s','MarkerSize',3,'MarkerFaceColor',[0.85 0.33 0.1]);
    plot(ax3, [0 100], [100 100], 'k--', 'LineWidth', 0.8);
    xlabel(ax3,'Epoch'); ylabel(ax3,'Accuracy (%)');
    title(ax3,'Training & Validation Accuracy','FontSize',13);
    legend(ax3,{'Train','Validation','100%'},'Location','southeast');
    ylim(ax3,[0 105]);

    drawnow;

    %% =====================================================================
    %  6. Train epoch by epoch with animation
    %  =====================================================================
    nEpochs   = 80;
    miniBatch = 64;
    learnRate = 1e-3;

    options = trainingOptions('adam', ...
        'MaxEpochs', 1, ...
        'MiniBatchSize', miniBatch, ...
        'InitialLearnRate', learnRate, ...
        'Shuffle', 'every-epoch', ...
        'Verbose', false, ...
        'Plots', 'none');

    net = [];
    trainLosses = []; valLosses = [];
    trainAccs   = []; valAccs   = [];
    snapshots = struct('epoch',{},'boundary',{},'trainAcc',{},'valAcc',{});

    for epoch = 1:nEpochs
        % Train one epoch
        if epoch == 1
            [net, info] = trainNetwork(XtrainDL, yTrainCat, layers, options);
        else
            [net, info] = trainNetwork(XtrainDL, yTrainCat, net.Layers, options);
        end

        % Predict on grid for decision boundary
        gridScores = predict(net, XgridDL, 'MiniBatchSize', 512);
        [~, gridLabels] = max(gridScores, [], 2);
        boundaryImg = reshape(gridLabels, gridRes, gridRes);

        % Color the decision boundary image
        rgbImg = zeros(gridRes, gridRes, 3);
        for c = 1:nClasses
            mask = boundaryImg == c;
            for ch = 1:3
                rgbImg(:,:,ch) = rgbImg(:,:,ch) + classColors(c,ch) * double(mask);
            end
        end
        set(hImg, 'CData', rgbImg);

        % Compute training accuracy and loss
        predTrain = classify(net, XtrainDL, 'MiniBatchSize', 256);
        trainAcc  = 100 * mean(predTrain == yTrainCat);
        scoreTrain = predict(net, XtrainDL, 'MiniBatchSize', 256);
        trainLoss  = crossEntropyLoss(scoreTrain, yTrain, nClasses);

        % Compute validation accuracy and loss
        predVal = classify(net, XvalDL, 'MiniBatchSize', 256);
        valAcc  = 100 * mean(predVal == yValCat);
        scoreVal = predict(net, XvalDL, 'MiniBatchSize', 256);
        valLoss  = crossEntropyLoss(scoreVal, yVal, nClasses);

        trainLosses(end+1) = trainLoss; %#ok<AGROW>
        valLosses(end+1)   = valLoss;   %#ok<AGROW>
        trainAccs(end+1)   = trainAcc;  %#ok<AGROW>
        valAccs(end+1)     = valAcc;    %#ok<AGROW>

        % Update plots
        set(epochText, 'String', sprintf('Epoch: %d / %d  |  Val Acc: %.1f%%', ...
            epoch, nEpochs, valAcc));

        addpoints(hLossTrain, epoch, trainLoss);
        addpoints(hLossVal,   epoch, valLoss);
        ylim(ax2, [0, max([trainLosses, valLosses])*1.1 + 0.01]);

        addpoints(hAccTrain, epoch, trainAcc);
        addpoints(hAccVal,   epoch, valAcc);

        drawnow;

        % Store snapshot
        snapshots(end+1) = struct('epoch',epoch,'boundary',rgbImg,...
                                  'trainAcc',trainAcc,'valAcc',valAcc); %#ok<AGROW>
    end

    %% =====================================================================
    %  7. Final: confusion matrix + boundary progression
    %  =====================================================================

    % --- Confusion matrix ---
    figure('Name','Confusion Matrix','Color','w','Position',[200 200 500 450]);
    cm = confusionmat(double(yValCat), double(predVal));
    confusionchart(cm, arrayfun(@(c) sprintf('Class %d',c), 1:nClasses, 'Uni', false));
    title(sprintf('Validation Confusion Matrix (Acc: %.1f%%)', valAcc));

    % --- Boundary progression ---
    figure('Name','Decision Boundary Progression','Color','w',...
           'Position',[120 80 1300 350]);
    snapIdx = round(linspace(1, length(snapshots), 6));
    for i = 1:6
        s = snapshots(snapIdx(i));
        ax = subplot(1,6,i);
        hold(ax,'on');
        imagesc(ax, gx, gy, s.boundary);
        set(ax.Children(end), 'AlphaData', 0.4);
        for c = 1:nClasses
            mask = yTrain == c;
            plot(ax, xTrain(mask,1), xTrain(mask,2), '.', ...
                 'Color', classColors(c,:), 'MarkerSize', 3);
        end
        axis(ax,'equal'); axis(ax,'tight');
        xlim(ax,[-2.2 2.2]); ylim(ax,[-2.2 2.2]);
        title(ax, sprintf('Epoch %d\nAcc=%.1f%%', s.epoch, s.valAcc), 'FontSize',10);
        if i == 1, ylabel(ax,'x_2'); end
        xlabel(ax,'x_1');
        set(ax,'FontSize',8,'XTick',[],'YTick',[]);
    end
    sgtitle('Decision Boundary Progression: from random to learned','FontSize',13);

    %% Save output figure
    saveas(fig, 'cnn_classification.png');

    %% Print summary
    fprintf('\n=== CNN Classification Summary ===\n');
    fprintf('Data: 4-class spiral, %d train / %d val points\n', nTotal, size(xVal,1));
    fprintf('Patch: %dx%d, Channels: %d\n', patchSize, patchSize, nChannels);
    fprintf('Epochs: %d\n', nEpochs);
    fprintf('Final train acc: %.1f%%,  val acc: %.1f%%\n', trainAccs(end), valAccs(end));
    fprintf('Final train loss: %.4f,  val loss: %.4f\n', trainLosses(end), valLosses(end));
    fprintf('==================================\n');
end

%% =========================================================================
%  Helper: generate spiral dataset with nClasses arms
%  =========================================================================
function [X, Y] = generateSpiralData(nPerClass, nClasses, noise)
    X = zeros(nPerClass * nClasses, 2);
    Y = zeros(nPerClass * nClasses, 1);
    for c = 1:nClasses
        idx = (c-1)*nPerClass + (1:nPerClass);
        t = linspace(0.2, 2.5, nPerClass)' + noise * randn(nPerClass, 1);
        angle = t + (c-1) * 2 * pi / nClasses;
        r = t * 0.7;
        X(idx, 1) = r .* cos(angle);
        X(idx, 2) = r .* sin(angle);
        Y(idx) = c;
    end
    % Shuffle
    perm = randperm(size(X,1));
    X = X(perm, :);
    Y = Y(perm);
end

%% =========================================================================
%  Helper: encode 2D points as multi-channel image patches
%  Output: [patchSize x patchSize x nChannels x n]
%  =========================================================================
function patches = encode2DPatches(X, freqs, patchSize)
    n = size(X, 1);
    nFreqs = length(freqs);
    nChannels = 2 * nFreqs * 2 + 2;  % (sin+cos) * nFreqs * 2dims + raw x1,x2

    offsets = linspace(-1, 1, patchSize);
    [OX, OY] = meshgrid(offsets, offsets);  % patchSize x patchSize grids

    patches = zeros(patchSize, patchSize, nChannels, n);
    for i = 1:n
        x1 = X(i, 1);
        x2 = X(i, 2);
        localX1 = x1 + OX;  % patchSize x patchSize
        localX2 = x2 + OY;

        ch = 1;
        patches(:,:,ch,i) = localX1; ch = ch + 1;
        patches(:,:,ch,i) = localX2; ch = ch + 1;
        for f = 1:nFreqs
            patches(:,:,ch,i) = sin(freqs(f) * localX1); ch = ch + 1;
            patches(:,:,ch,i) = cos(freqs(f) * localX1); ch = ch + 1;
            patches(:,:,ch,i) = sin(freqs(f) * localX2); ch = ch + 1;
            patches(:,:,ch,i) = cos(freqs(f) * localX2); ch = ch + 1;
        end
    end
end

%% =========================================================================
%  Helper: cross-entropy loss from softmax scores and integer labels
%  =========================================================================
function loss = crossEntropyLoss(scores, labels, nClasses)
    n = length(labels);
    scores = max(scores, 1e-12);  % avoid log(0)
    idx = sub2ind([n, nClasses], (1:n)', labels(:));
    loss = -mean(log(scores(idx)));
end
