% CNN Regression: Train a 1D Convolutional Neural Network to fit a smooth function
% Demonstrates how a CNN learns a nonlinear mapping from x to y = f(x)
% using local receptive fields (convolutions) on windowed input patches.
%
% Architecture: Input patch -> Conv1D -> ReLU -> Conv1D -> ReLU -> FC -> output
% Animation shows the predicted curve evolving during training.

function cnn_regression()
    %% =====================================================================
    %  1. Generate data: smooth target function
    %  =====================================================================
    rng(0);

    % Target function: mixture of sinusoids
    f_true = @(x) sin(2*x) + 0.5*cos(5*x) + 0.3*sin(8*x);

    nData   = 800;           % training points
    nTest   = 300;           % test points
    xRange  = [-pi, pi];

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

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

    % Dense grid for plotting the true function
    xPlot = linspace(xRange(1), xRange(2), 500)';
    yPlot = f_true(xPlot);

    %% =====================================================================
    %  2. Create windowed input patches
    %  =====================================================================
    % We turn the scalar regression x->y into a patch-based problem:
    % for each x_i, the input is a local window of neighboring x values
    % encoded as a short 1D signal. This gives the CNN spatial structure.

    winSize = 15;  % patch width (odd)
    halfWin = floor(winSize / 2);

    % Encode each x as a Gabor-like feature vector (gives local structure)
    freqs = [1, 2, 4, 8];
    nChannels = 2 * length(freqs) + 1;  % sin + cos for each freq, plus x itself

    % Encode as [winSize x nChannels x 1 x nSamples] for imageInputLayer
    XtrainDL = encodePatches(xTrain, freqs, winSize);
    XtestDL  = encodePatches(xTest,  freqs, winSize);
    XplotDL  = encodePatches(xPlot,  freqs, winSize);

    %% =====================================================================
    %  3. Define CNN architecture
    %  =====================================================================
    layers = [
        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')

        regressionLayer('Name','output')
    ];

    %% =====================================================================
    %  4. Training with custom output function for animation
    %  =====================================================================
    % Store snapshots during training
    snapshots = struct('epoch',{},'yPred',{},'loss',{});
    epochLosses = [];

    % --- Set up figure ---
    fig = figure('Name','CNN Regression Training','Color','w',...
                 'Position',[100 60 1100 800]);

    % Panel 1: fit evolution
    ax1 = subplot(2,2,[1 2],'Parent',fig);
    hold(ax1,'on'); grid(ax1,'on');
    plot(ax1, xTrain, yTrain, '.', 'Color', [0.7 0.7 0.7], 'MarkerSize', 4);
    hTrue = plot(ax1, xPlot, yPlot, 'k-', 'LineWidth', 2);
    hPred = plot(ax1, xPlot, zeros(size(xPlot)), '-', 'Color', [0.85 0.33 0.1], ...
                 'LineWidth', 2.5);
    xlim(ax1, xRange); ylim(ax1, [min(yPlot)-0.5, max(yPlot)+0.5]);
    xlabel(ax1,'x'); ylabel(ax1,'y');
    title(ax1,'CNN Prediction vs True Function','FontSize',14);
    legend(ax1,{'Train data','True f(x)','CNN prediction'},'Location','southeast');
    epochText = text(ax1, xRange(1)+0.1, max(yPlot)+0.3, 'Epoch: 0', ...
                     'FontSize',12,'FontWeight','bold');

    % Panel 2: loss curve
    ax2 = subplot(2,2,3,'Parent',fig);
    hold(ax2,'on'); grid(ax2,'on');
    hLoss = animatedline(ax2, 'Color', [0 0.45 0.74], 'LineWidth', 2, ...
                          'Marker','o','MarkerSize',4,'MarkerFaceColor',[0 0.45 0.74]);
    xlabel(ax2,'Epoch'); ylabel(ax2,'MSE Loss');
    title(ax2,'Training Loss','FontSize',13);
    set(ax2,'YScale','log');

    % Panel 3: residuals
    ax3 = subplot(2,2,4,'Parent',fig);
    hold(ax3,'on'); grid(ax3,'on');
    hResid = plot(ax3, xTest, zeros(size(xTest)), '-s', 'Color', [0.47 0.67 0.19], ...
                  'LineWidth', 1, 'MarkerSize', 3, 'MarkerFaceColor', [0.47 0.67 0.19]);
    plot(ax3, xRange, [0 0], 'k--', 'LineWidth', 1);
    xlim(ax3, xRange);
    xlabel(ax3,'x'); ylabel(ax3,'Residual');
    title(ax3,'Prediction Residuals on Test Set','FontSize',13);

    drawnow;

    %% =====================================================================
    %  5. Train in epochs manually for animation control
    %  =====================================================================
    nEpochs     = 60;
    miniBatch   = 64;
    learnRate   = 1e-3;

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

    net = [];

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

        % Predict on plot grid and test set
        yPredPlot = predict(net, XplotDL, 'MiniBatchSize', 256);
        yPredTest = predict(net, XtestDL, 'MiniBatchSize', 256);
        residuals = yTest - yPredTest;
        mse = mean(residuals.^2);
        epochLosses(end+1) = mse; %#ok<AGROW>

        % Update animation
        set(hPred, 'YData', yPredPlot);
        set(epochText, 'String', sprintf('Epoch: %d / %d   MSE: %.5f', epoch, nEpochs, mse));

        addpoints(hLoss, epoch, mse);
        ylim(ax2, [min(epochLosses)*0.5, max(epochLosses)*2]);

        set(hResid, 'YData', residuals);
        maxRes = max(abs(residuals)) * 1.2 + 0.01;
        ylim(ax3, [-maxRes, maxRes]);

        drawnow;

        % Store snapshot
        snapshots(end+1) = struct('epoch', epoch, 'yPred', yPredPlot, 'loss', mse); %#ok<AGROW>
    end

    %% =====================================================================
    %  6. Final summary figure: snapshots of learning progression
    %  =====================================================================
    figure('Name','Learning Progression','Color','w',...
           'Position',[150 100 1200 400]);

    snapIdx = round(linspace(1, length(snapshots), 6));
    for i = 1:6
        s = snapshots(snapIdx(i));
        ax = subplot(1,6,i);
        hold(ax,'on'); grid(ax,'on');
        plot(ax, xPlot, yPlot, 'k-', 'LineWidth', 1.5);
        plot(ax, xPlot, s.yPred, '-', 'Color', [0.85 0.33 0.1], 'LineWidth', 1.5);
        xlim(ax, xRange);
        ylim(ax, [min(yPlot)-0.5, max(yPlot)+0.5]);
        title(ax, sprintf('Epoch %d\nMSE=%.4f', s.epoch, s.loss), 'FontSize', 10);
        if i == 1
            ylabel(ax, 'y');
        end
        xlabel(ax, 'x');
        set(ax, 'FontSize', 8);
    end
    sgtitle('CNN Learning Progression: from random to fitted', 'FontSize', 13);

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

    %% Print summary
    fprintf('\n=== CNN Regression Summary ===\n');
    fprintf('Target: f(x) = sin(2x) + 0.5*cos(5x) + 0.3*sin(8x)\n');
    fprintf('Training points: %d,  Test points: %d\n', nData, nTest);
    fprintf('Patch size: %d,  Input channels: %d\n', winSize, nChannels);
    fprintf('Epochs: %d,  Final test MSE: %.6f\n', nEpochs, epochLosses(end));
    fprintf('Final test RMSE: %.6f\n', sqrt(epochLosses(end)));
    fprintf('==============================\n');
end

%% =========================================================================
%  Helper: encode scalar x values into windowed multi-channel patches
%  =========================================================================
function patches = encodePatches(x, freqs, winSize)
    % For each x_i, create a [winSize x nChannels] patch where
    % columns are: x offsets, sin(freq*x_local), cos(freq*x_local)
    % Output: [winSize x nChannels x 1 x n] ready for imageInputLayer
    n = length(x);
    nChannels = 2 * length(freqs) + 1;

    offsets = linspace(-1, 1, winSize);  % normalized local window

    % Build as [winSize x nChannels x 1 x n]
    patches = zeros(winSize, nChannels, 1, n);
    for i = 1:n
        xLocal = x(i) + offsets;                   % 1 x winSize
        patches(:, 1, 1, i) = xLocal(:);           % position channel
        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
