%% CFHigham University of Glasgow February 2026

% Code to recreate the results in Experiment 2: Agent Learns Task Order (Figure 4b and 4c)
% NB Results may not exactly replicate results in the paper due to the
% stochastic nature of the Diffusion model.

addpath('Experiment2')

% % To obtain saved results for Figure 4b:
% 
% R4b = load ('abcd0508.mat')
% 
% % saved results for Figure 4b:
% 
% R4c = load ('acbd0508.mat')


addpath('Diffusion')

%% set up DDPM for backward sampling

imageSize=[4 4];
numChannels=1;
numImages=100;

images = randn([imageSize numChannels numImages]);

numNoiseSteps = 100;
betaMin = 1e-4;
betaMax = 0.02;
varianceSchedule = linspace(betaMin,betaMax,numNoiseSteps);

% Compute variance schedule parameters.
alphaBar = cumprod(1 - varianceSchedule);
alphaBarPrev = [1 alphaBar(1:end-1)];
posteriorVariance = varianceSchedule.*(1 - alphaBarPrev)./(1 - alphaBar);

% load DDPM network model
load('net0508.mat')
% load reference data
load('data0508.mat')
load('classdata0508.mat')


%% Fig 4(b) Diffusion prior conditioned on 'ab'


% Reference data - Observations 
img=c1;
img = rescale(img,-1,1);
img=repmat(img,1,1,1,100);
img=dlarray(img,'SSCB');

% Mask
Ms = createMask;
M=Ms(:,:,2); 
M = repmat(M,1,1,1,numImages);


x = randn([imageSize numChannels numImages]);
U=20;    
     for noiseStep = numNoiseSteps:-1:1

         for u=1:U
    
            % Generate random noise.
            targetNoise = randn(size(x),Like=images);

            % Apply noise to the image.
            noisyImage = applyNoiseToImage(img,targetNoise,noiseStep,varianceSchedule);
            x_known = noisyImage;    
    
    if noiseStep ~= 1
        z = randn([imageSize,numChannels,numImages]);
    else
        z = zeros([imageSize,numChannels,numImages]);
    end
        
    % Predict the noise using the network.
    predictedNoise = predict(net,x,noiseStep);
    sqrtOneMinusBeta = sqrt(1 - varianceSchedule(noiseStep));
    addedNoise = sqrt(posteriorVariance(noiseStep))*z;
    predNoise = varianceSchedule(noiseStep)*predictedNoise/sqrt(1 - alphaBar(noiseStep));
    x_unknown = 1/sqrtOneMinusBeta*(x - predNoise) + addedNoise;
    

    x = M.*x_known + (1-M).*x_unknown;
        
         
    if noiseStep ~= 1 & u<U

        z = randn([imageSize,numChannels,numImages]);
        x = sqrtOneMinusBeta*x + sqrt(posteriorVariance(noiseStep))*z;
    else
    end
     

           x2=(extractdata(x(:,:,1,:)));
           x2=reshape(x2,16,numImages);
           sampleData=double(x2)';
           % 
           class = classify(sampleData,trainData,group);
           c=histogram(class);
           p=c.Values/numImages;
           pc=zeros(1,4);
           pc(unique(class))=p;
           
         end
      end      

           
eps=1.0000e-06;
% preference - to be in all 'on' state
p0=[eps eps eps 1];
DK=sum(p0.*(log2(p0)-log2(pc+eps)));
H=-sum(pc./(log2(pc+eps)));
% Expected Free energy
EFE=DK+H;

X = ['p = [' ,num2str(pc),'] and EFE = ',num2str(EFE)];
disp(X)

% Results in paper
%load('abcd0508.mat')


%% Fig 4(b) Diffusion prior conditioned on 'ac'

% Reference Image- Observations 
img=c4;
img = rescale(img,-1,1);
img=repmat(img,1,1,1,100);
img=dlarray(img,'SSCB');

% Mask
Ms = createMask;
M=Ms(:,:,2); 
M = repmat(M,1,1,1,numImages);

x = randn([imageSize numChannels numImages]);
U=20;    
     for noiseStep = numNoiseSteps:-1:1

         for u=1:U
    
            % Generate random noise.
            targetNoise = randn(size(x),Like=images);

            % Apply noise to the image.
            noisyImage = applyNoiseToImage(img,targetNoise,noiseStep,varianceSchedule);
            x_known = noisyImage;    
    
    if noiseStep ~= 1
        z = randn([imageSize,numChannels,numImages]);
    else
        z = zeros([imageSize,numChannels,numImages]);
    end
        
    % Predict the noise using the network.
    predictedNoise = predict(net,x,noiseStep);
    sqrtOneMinusBeta = sqrt(1 - varianceSchedule(noiseStep));
    addedNoise = sqrt(posteriorVariance(noiseStep))*z;
    predNoise = varianceSchedule(noiseStep)*predictedNoise/sqrt(1 - alphaBar(noiseStep));
    x_unknown = 1/sqrtOneMinusBeta*(x - predNoise) + addedNoise;
    

    x = M.*x_known + (1-M).*x_unknown;
        
         
    if noiseStep ~= 1 & u<U

        z = randn([imageSize,numChannels,numImages]);
        x = sqrtOneMinusBeta*x + sqrt(posteriorVariance(noiseStep))*z;
    else
    end
     

           x2=(extractdata(x(:,:,1,:)));
           x2=reshape(x2,16,numImages);
           sampleData=double(x2)';
           % 
           class = classify(sampleData,trainData,group);
           c=histogram(class);
           p=c.Values/numImages;
           pc=zeros(1,4);
           pc(unique(class))=p;
           
         end
      end      

           
eps=1.0000e-06;
% preference - to be in all 'on' state
p0=[eps eps eps 1];
DK=sum(p0.*(log2(p0)-log2(pc+eps)));
H=-sum(pc./(log2(pc+eps)));
% Expected Free energy
EFE=DK+H;


X = ['p = [' ,num2str(pc),'] and EFE = ',num2str(EFE)];
disp(X)

% Results in paper
%load('acbd0508.mat')


  
% Figure to inspect generated images x2
figure
for i=1:100
subplot(10,10,i)
imagesc(reshape(x2(:,i),4,4))
subtitle(num2str(class(i)))
end