% CNN Regression: Compare Optimization Methods
% Trains the same CNN architecture with different optimizers on a smooth
% function regression task. Uses custom training loops (dlnetwork) so
% each optimizer can be implemented from scratch.
%
% Methods compared:
%   1. Vanilla GD       (full-batch gradient descent, no momentum)
%   2. SGD              (mini-batch stochastic gradient descent)
%   3. Heavy Ball       (SGD + classical momentum, Polyak 1964)
%   4. Nesterov AGD     (Nesterov accelerated gradient, 1983)
%   5. AdaGrad          (adaptive learning rates, Duchi et al. 2011)
%   6. RMSProp          (Hinton, unpublished lecture notes 2012)
%   7. Adam             (Kingma & Ba, 2015)

function cnn_regression_optimizers()
    %% =====================================================================
    %  1. Generate data
    %  =====================================================================
    rng(0);

    f_true = @(x) sin(2*x) + 0.5*cos(5*x) + 0.3*sin(8*x);

    nData   = 800;
    nTest   = 300;
    xRange  = [-pi, pi];

    xTrain = sort(xRange(1) + diff(xRange) * rand(nData, 1));
    yTrain = f_true(xTrain) + 0.05 * randn(nData, 1);

    xTest = linspace(xRange(1), xRange(2), nTest)';
    yTest = f_true(xTest);

    xPlot = linspace(xRange(1), xRange(2), 500)';
    yPlot = f_true(xPlot);

    %% =====================================================================
    %  2. Encode patches
    %  =====================================================================
    winSize   = 15;
    freqs     = [1, 2, 4, 8];
    nChannels = 2 * length(freqs) + 1;

    XtrainAll = encodePatches(xTrain, freqs, winSize);
    XtestAll  = encodePatches(xTest,  freqs, winSize);
    XplotAll  = encodePatches(xPlot,  freqs, winSize);

    % Convert to dlarray for custom training
    dlXtest = dlarray(single(XtestAll), 'SSCB');
    dlXplot = dlarray(single(XplotAll), 'SSCB');
    dlYtest = single(yTest);

    %% =====================================================================
    %  3. Define optimizers
    %  =====================================================================
    methods = struct('name', {}, 'color', {}, 'marker', {}, 'lr', {});

    methods(1).name = 'Vanilla GD';
    methods(1).color = [0.50 0.50 0.50];
    methods(1).marker = 'v';
    methods(1).lr = 1e-3;

    methods(2).name = 'SGD';
    methods(2).color = [0.85 0.33 0.10];
    methods(2).marker = 'o';
    methods(2).lr = 1e-3;

    methods(3).name = 'Heavy Ball';
    methods(3).color = [0.93 0.69 0.13];
    methods(3).marker = 'd';
    methods(3).lr = 1e-3;

    methods(4).name = 'Nesterov AGD';
    methods(4).color = [0.47 0.67 0.19];
    methods(4).marker = '^';
    methods(4).lr = 1e-3;

    methods(5).name = 'AdaGrad';
    methods(5).color = [0.30 0.75 0.93];
    methods(5).marker = 's';
    methods(5).lr = 1e-2;

    methods(6).name = 'RMSProp';
    methods(6).color = [0.64 0.08 0.18];
    methods(6).marker = 'p';
    methods(6).lr = 1e-3;

    methods(7).name = 'Adam';
    methods(7).color = [0.00 0.45 0.74];
    methods(7).marker = 'h';
    methods(7).lr = 1e-3;

    nMethods = length(methods);

    %% =====================================================================
    %  4. Training parameters
    %  =====================================================================
    nEpochs   = 60;
    miniBatch = 64;
    momentum  = 0.9;       % for Heavy Ball and Nesterov
    beta1     = 0.9;       % Adam first moment
    beta2     = 0.999;     % Adam second moment / RMSProp decay
    epsilon   = 1e-8;      % numerical stability

    %% =====================================================================
    %  5. Set up animated figure
    %  =====================================================================
    fig1 = figure('Name','Optimizer Comparison - Training','Color','w',...
                  'Position',[60 60 1300 850]);

    % Panel 1: all predictions overlaid
    ax1 = subplot(2,2,[1 2],'Parent',fig1);
    hold(ax1,'on'); grid(ax1,'on');
    plot(ax1, xTrain, yTrain, '.', 'Color',[0.82 0.82 0.82], 'MarkerSize',3);
    plot(ax1, xPlot, yPlot, 'k-', 'LineWidth', 2.5);
    hPreds = gobjects(nMethods, 1);
    for m = 1:nMethods
        hPreds(m) = plot(ax1, xPlot, nan(size(xPlot)), '-', ...
                         'Color', methods(m).color, 'LineWidth', 1.8);
    end
    xlim(ax1, xRange); ylim(ax1, [min(yPlot)-0.5, max(yPlot)+0.5]);
    xlabel(ax1,'x'); ylabel(ax1,'y');
    title(ax1,'CNN Predictions by Optimizer','FontSize',14);
    legNames = [{'Train data','True f(x)'}, {methods.name}];
    legend(ax1, legNames, 'Location','southeast','FontSize',8);
    epochText = text(ax1, xRange(1)+0.1, max(yPlot)+0.3, 'Epoch: 0', ...
                     'FontSize',12,'FontWeight','bold');

    % Panel 2: loss curves (log scale)
    ax2 = subplot(2,2,3,'Parent',fig1);
    hold(ax2,'on'); grid(ax2,'on');
    hLoss = gobjects(nMethods, 1);
    for m = 1:nMethods
        hLoss(m) = animatedline(ax2, 'Color', methods(m).color, 'LineWidth', 1.8, ...
                                'Marker', methods(m).marker, 'MarkerSize', 3, ...
                                'MarkerFaceColor', methods(m).color);
    end
    xlabel(ax2,'Epoch'); ylabel(ax2,'Test MSE');
    title(ax2,'Test Loss (log scale)','FontSize',13);
    set(ax2,'YScale','log');
    legend(ax2, {methods.name}, 'Location','northeast','FontSize',7);

    % Panel 3: relative improvement bar chart (updated each epoch)
    ax3 = subplot(2,2,4,'Parent',fig1);
    hold(ax3,'on'); grid(ax3,'on');
    title(ax3,'Current Test MSE by Optimizer','FontSize',13);
    xlabel(ax3,'Optimizer'); ylabel(ax3,'MSE');

    drawnow;

    %% =====================================================================
    %  6. Train all methods
    %  =====================================================================
    % Build a base layer graph (without regression output layer) for dlnetwork
    lgraph = layerGraph([
        imageInputLayer([winSize nChannels 1], 'Name','input','Normalization','none')
        convolution2dLayer([5 nChannels], 32, 'Padding','same', 'Name','conv1')
        batchNormalizationLayer('Name','bn1')
        reluLayer('Name','relu1')
        convolution2dLayer([5 1], 64, 'Padding','same', 'Name','conv2')
        batchNormalizationLayer('Name','bn2')
        reluLayer('Name','relu2')
        convolution2dLayer([3 1], 32, 'Padding','same', 'Name','conv3')
        reluLayer('Name','relu3')
        globalAveragePooling2dLayer('Name','gap')
        fullyConnectedLayer(64, 'Name','fc1')
        reluLayer('Name','relu_fc')
        fullyConnectedLayer(1, 'Name','fc_out')
    ]);

    % Store initial weights so every method starts the same
    rng(42);
    net0 = dlnetwork(lgraph);
    initParams = net0.Learnables;

    % Preallocate results
    allLosses = nan(nMethods, nEpochs);
    allNets   = cell(nMethods, 1);

    % Initialize optimizer states for each method
    optState = cell(nMethods, 1);
    nets     = cell(nMethods, 1);
    for m = 1:nMethods
        nets{m} = net0;  % same initial weights
        switch methods(m).name
            case 'Vanilla GD'
                optState{m} = struct();
            case 'SGD'
                optState{m} = struct();
            case 'Heavy Ball'
                optState{m} = initVelocity(initParams);
            case 'Nesterov AGD'
                optState{m} = initVelocity(initParams);
            case 'AdaGrad'
                optState{m} = initAccumulator(initParams);
            case 'RMSProp'
                optState{m} = initAccumulator(initParams);
            case 'Adam'
                optState{m} = struct('m', {initAccumulator(initParams)}, ...
                                     'v', {initAccumulator(initParams)}, ...
                                     't', 0);
        end
    end

    nTrain = size(XtrainAll, 4);
    dlYtrain = single(yTrain);

    for epoch = 1:nEpochs
        % Shuffle training data
        perm = randperm(nTrain);

        for m = 1:nMethods
            net_m = nets{m};
            lr = methods(m).lr;

            if strcmp(methods(m).name, 'Vanilla GD')
                %% --- Vanilla GD: full-batch ---
                dlX = dlarray(single(XtrainAll), 'SSCB');
                dlY = dlYtrain;
                [loss, grads, state] = dlfeval(@modelGradients, net_m, dlX, dlY);
                net_m.State = state;
                net_m = sgdUpdate(net_m, grads, lr);

            else
                %% --- Mini-batch methods ---
                for bStart = 1:miniBatch:nTrain
                    bEnd = min(bStart + miniBatch - 1, nTrain);
                    idx = perm(bStart:bEnd);
                    dlX = dlarray(single(XtrainAll(:,:,:,idx)), 'SSCB');
                    dlY = dlYtrain(idx);

                    switch methods(m).name
                        case 'SGD'
                            [~, grads, state] = dlfeval(@modelGradients, net_m, dlX, dlY);
                            net_m.State = state;
                            net_m = sgdUpdate(net_m, grads, lr);

                        case 'Heavy Ball'
                            [~, grads, state] = dlfeval(@modelGradients, net_m, dlX, dlY);
                            net_m.State = state;
                            [net_m, optState{m}] = heavyBallUpdate(net_m, grads, lr, momentum, optState{m});

                        case 'Nesterov AGD'
                            % Look-ahead: evaluate gradient at w + mu*v
                            net_look = nesterovLookahead(net_m, optState{m}, momentum);
                            [~, grads, state] = dlfeval(@modelGradients, net_look, dlX, dlY);
                            net_m.State = state;
                            [net_m, optState{m}] = heavyBallUpdate(net_m, grads, lr, momentum, optState{m});

                        case 'AdaGrad'
                            [~, grads, state] = dlfeval(@modelGradients, net_m, dlX, dlY);
                            net_m.State = state;
                            [net_m, optState{m}] = adagradUpdate(net_m, grads, lr, epsilon, optState{m});

                        case 'RMSProp'
                            [~, grads, state] = dlfeval(@modelGradients, net_m, dlX, dlY);
                            net_m.State = state;
                            [net_m, optState{m}] = rmspropUpdate(net_m, grads, lr, beta2, epsilon, optState{m});

                        case 'Adam'
                            [~, grads, state] = dlfeval(@modelGradients, net_m, dlX, dlY);
                            net_m.State = state;
                            [net_m, optState{m}] = adamUpdate(net_m, grads, lr, beta1, beta2, epsilon, optState{m});
                    end
                end
            end

            nets{m} = net_m;

            % Evaluate test MSE
            dlYpred = predict(net_m, dlXtest);
            yPred = extractdata(dlYpred)';
            mse = mean((dlYtest - yPred).^2);
            allLosses(m, epoch) = mse;

            % Update prediction curve
            dlYpredPlot = predict(net_m, dlXplot);
            set(hPreds(m), 'YData', extractdata(dlYpredPlot)');
        end

        % Update epoch text
        set(epochText, 'String', sprintf('Epoch: %d / %d', epoch, nEpochs));

        % Update loss curves
        for m = 1:nMethods
            addpoints(hLoss(m), epoch, allLosses(m, epoch));
        end
        validLosses = allLosses(:, 1:epoch);
        validLosses = validLosses(validLosses > 0);
        if ~isempty(validLosses)
            ylim(ax2, [min(validLosses)*0.3, max(validLosses)*3]);
        end

        % Update bar chart
        cla(ax3);
        currentMSE = allLosses(:, epoch);
        for m = 1:nMethods
            bar(ax3, m, currentMSE(m), 'FaceColor', methods(m).color, 'EdgeColor', 'w');
        end
        set(ax3, 'XTick', 1:nMethods, 'XTickLabel', {methods.name}, ...
                 'XTickLabelRotation', 30, 'FontSize', 7);
        set(ax3, 'YScale', 'log');
        ylabel(ax3, 'MSE (log)');
        title(ax3, sprintf('Test MSE at Epoch %d', epoch), 'FontSize', 13);

        drawnow;
    end

    %% =====================================================================
    %  7. Final comparison figures
    %  =====================================================================

    % --- Figure 2: final predictions side by side ---
    fig2 = figure('Name','Final Predictions','Color','w',...
                  'Position',[80 100 1400 500]);
    nCols = ceil(nMethods / 2);
    for m = 1:nMethods
        ax = subplot(2, nCols, m, 'Parent', fig2);
        hold(ax,'on'); grid(ax,'on');
        plot(ax, xPlot, yPlot, 'k-', 'LineWidth', 1.5);
        dlYp = predict(nets{m}, dlXplot);
        plot(ax, xPlot, extractdata(dlYp)', '-', 'Color', methods(m).color, 'LineWidth', 2);
        xlim(ax, xRange);
        ylim(ax, [min(yPlot)-0.5, max(yPlot)+0.5]);
        title(ax, sprintf('%s\nMSE=%.5f', methods(m).name, allLosses(m,end)), 'FontSize', 10);
        if mod(m-1, nCols) == 0, ylabel(ax,'y'); end
        xlabel(ax,'x');
        set(ax, 'FontSize', 8);
    end
    sgtitle('Final CNN Fits by Optimizer', 'FontSize', 14);

    % --- Figure 3: convergence comparison ---
    fig3 = figure('Name','Convergence Comparison','Color','w',...
                  'Position',[120 80 800 550]);
    hold on; grid on;
    for m = 1:nMethods
        semilogy(1:nEpochs, allLosses(m,:), ['-' methods(m).marker], ...
                 'Color', methods(m).color, 'LineWidth', 1.8, ...
                 'MarkerSize', 4, 'MarkerFaceColor', methods(m).color, ...
                 'MarkerIndices', 1:5:nEpochs);
    end
    set(gca, 'YScale', 'log');
    xlabel('Epoch'); ylabel('Test MSE');
    title('Convergence Comparison (log scale)', 'FontSize', 14);
    legend({methods.name}, 'Location', 'northeast');

    %% Print summary
    fprintf('\n===============================================\n');
    fprintf('  Optimizer Comparison Summary\n');
    fprintf('===============================================\n');
    fprintf('Target: f(x) = sin(2x) + 0.5cos(5x) + 0.3sin(8x)\n');
    fprintf('Train: %d pts, Test: %d pts, Epochs: %d\n', nData, nTest, nEpochs);
    fprintf('-----------------------------------------------\n');
    fprintf('%-15s  %10s  %10s  %10s\n', 'Method', 'Final MSE', 'Best MSE', 'Best Epoch');
    fprintf('-----------------------------------------------\n');
    for m = 1:nMethods
        [bestMSE, bestEp] = min(allLosses(m,:));
        fprintf('%-15s  %10.6f  %10.6f  %10d\n', methods(m).name, ...
                allLosses(m,end), bestMSE, bestEp);
    end
    fprintf('===============================================\n\n');
end

%% =========================================================================
%  Model gradient computation
%  =========================================================================
function [loss, gradients, state] = modelGradients(net, dlX, Y)
    [dlYpred, state] = forward(net, dlX);
    % Output is [1 x 1 x 1 x B]; reshape to [1 x B] to match target
    dlYpred = reshape(dlYpred, 1, []);
    dlYtarget = dlarray(reshape(Y, 1, []), 'CB');
    loss = mean((dlYpred - dlYtarget).^2, 'all');
    gradients = dlgradient(loss, net.Learnables);
end

%% =========================================================================
%  Optimizer update rules
%  =========================================================================

% --- Vanilla SGD / GD ---
function net = sgdUpdate(net, grads, lr)
    for i = 1:size(grads, 1)
        net.Learnables.Value{i} = net.Learnables.Value{i} - lr * grads.Value{i};
    end
end

% --- Heavy Ball (SGD + momentum) ---
function [net, vel] = heavyBallUpdate(net, grads, lr, mu, vel)
    for i = 1:size(grads, 1)
        vel.Value{i} = mu * vel.Value{i} + lr * grads.Value{i};
        net.Learnables.Value{i} = net.Learnables.Value{i} - vel.Value{i};
    end
end

% --- Nesterov look-ahead ---
function netLook = nesterovLookahead(net, vel, mu)
    netLook = net;
    for i = 1:size(net.Learnables, 1)
        netLook.Learnables.Value{i} = net.Learnables.Value{i} - mu * vel.Value{i};
    end
end

% --- AdaGrad ---
function [net, acc] = adagradUpdate(net, grads, lr, eps, acc)
    for i = 1:size(grads, 1)
        acc.Value{i} = acc.Value{i} + grads.Value{i}.^2;
        net.Learnables.Value{i} = net.Learnables.Value{i} ...
            - lr * grads.Value{i} ./ (sqrt(acc.Value{i}) + eps);
    end
end

% --- RMSProp ---
function [net, acc] = rmspropUpdate(net, grads, lr, rho, eps, acc)
    for i = 1:size(grads, 1)
        acc.Value{i} = rho * acc.Value{i} + (1 - rho) * grads.Value{i}.^2;
        net.Learnables.Value{i} = net.Learnables.Value{i} ...
            - lr * grads.Value{i} ./ (sqrt(acc.Value{i}) + eps);
    end
end

% --- Adam ---
function [net, state] = adamUpdate(net, grads, lr, b1, b2, eps, state)
    state.t = state.t + 1;
    for i = 1:size(grads, 1)
        state.m.Value{i} = b1 * state.m.Value{i} + (1 - b1) * grads.Value{i};
        state.v.Value{i} = b2 * state.v.Value{i} + (1 - b2) * grads.Value{i}.^2;
        mHat = state.m.Value{i} / (1 - b1^state.t);
        vHat = state.v.Value{i} / (1 - b2^state.t);
        net.Learnables.Value{i} = net.Learnables.Value{i} ...
            - lr * mHat ./ (sqrt(vHat) + eps);
    end
end

%% =========================================================================
%  State initialization helpers
%  =========================================================================
function vel = initVelocity(learnables)
    vel = learnables;
    for i = 1:size(learnables, 1)
        vel.Value{i} = zeros(size(learnables.Value{i}), 'like', learnables.Value{i});
    end
end

function acc = initAccumulator(learnables)
    acc = learnables;
    for i = 1:size(learnables, 1)
        acc.Value{i} = zeros(size(learnables.Value{i}), 'like', learnables.Value{i});
    end
end

%% =========================================================================
%  Patch encoding (same as cnn_regression.m)
%  =========================================================================
function patches = encodePatches(x, freqs, winSize)
    n = length(x);
    nChannels = 2 * length(freqs) + 1;
    offsets = linspace(-1, 1, winSize);
    patches = zeros(winSize, nChannels, 1, n);
    for i = 1:n
        xLocal = x(i) + offsets;
        patches(:, 1, 1, i) = xLocal(:);
        for f = 1:length(freqs)
            patches(:, 2*f,   1, i) = sin(freqs(f) * xLocal(:));
            patches(:, 2*f+1, 1, i) = cos(freqs(f) * xLocal(:));
        end
    end
end
