%% CFHigham University of Glasgow February 2026

% Code to estimate the time per policy (perpol) 
% Example state space dimension N^2=4^2 (N=4) and number of repeats T=10

addpath(genpath('Experiment3'))

%% A simple environment: Grid-world
% Code for planning in a simple environment is a MATLAB version based on
% python code from Tutorial 1: Active inference from scratch by Conor Heins
% https://pymdp-rtd.readthedocs.io/en/latest/notebooks/active_inference_from_scratch.html

% The A matrix or $P(o\mid s)$
N=4;
n_states = N^2;
n_observations = N^2;

A = zeros(n_states, n_observations);

% Create an umambiguous or 'noise-less' mapping between hidden states and observations
A=diag(ones(1,n_states));


%% The B matrix or $P(s_{t}\mid s_{t-1}, u_{t-1})$

l1=1:N;
l1a=repmat(l1,N,1);
l1b=repmat(l1,1,N);
l1a=reshape(l1a,N^2,1);
grid_locations=[l1a,l1b'];

actions = ["UP", "DOWN", "LEFT", "RIGHT", "STAY"];

%def create_B_matrix():
  B = zeros(size(grid_locations,1), size(grid_locations,1), length(actions));

  for i=1:length(actions)

    for j=1:size(grid_locations,1)

      %y, x = grid_location
       y=grid_locations(j,1);
       x=grid_locations(j,2);
       curr_state=grid_locations(:,1)==y&grid_locations(:,2)==x;
       if actions(i)=="UP" 
        next_x = x;
       if y > 1
        next_y = y - 1; 
       else 
        next_y = y; 
       end
       end
       
       if actions(i)=="DOWN" 
        next_x = x;
       if y < 3
        next_y = y + 1; 
       else 
        next_y = y;
       end
       end
       
       if actions(i) == "LEFT" 
        next_y = y;
       if x > 1 
        next_x = x - 1;  
       else 
        next_x = x; 
       end 
       end
       
             
       if actions(i) == "RIGHT" 
        next_y = y;
       if x < 3
        next_x = x + 1;
       else 
        next_x = x ;
       end 
       end
       
       if actions(i) == "STAY"
        next_x = x;
        next_y = y;
        end
     
        
       next_state=grid_locations(:,1)==next_y&grid_locations(:,2)==next_x;
       B(next_state,curr_state,i)=1;
    end
  
  end

%% The prior over observations: the $\mathbf{C}$ vector or $\tilde{P}(o)$

% Create an empty vector to store the preferences over observations
C = zeros(1,n_observations);

%""" Choose an observation index to be the 'desired' rewarding index, and fill out the C vector accordingly """
desired_location = [N,N]; % choose a desired location
%desired_location_index = grid_locations.index(desired_location) % get the linear index of the grid location, in terms of 0 through 8
desired_location_index = find(grid_locations(:,1)==desired_location(1)&grid_locations(:,2)==desired_location(2));
C(desired_location_index) = 1.0; % set the preference for that location to be 100%, i.e. 1.0

bar(C)
title("Preferences over observations")

%%  Initialize prior, first observation, and policies
  
%  prior = D # initial prior should be the D vector

starting_state = [1,1];

D = zeros(n_observations,1);
prior_belief_location = [1,1]; % choose a desired location
prior_belief_index = find(grid_locations(:,1)==prior_belief_location(1)&grid_locations(:,2)==prior_belief_location(2));
D(prior_belief_index) = 1.0;

prior = D;
obs=starting_state;

%% create policies
pol=zeros(4,2^4);
ct=0;
for i=[2,4]
    for i2=[2,4]
        for i3=[2,4]
            for i4=[2,4]
                ct=ct+1;
                pol(:,ct)=[i,i2,i3,i4];
            end
        end
    end
end
policies=pol;

%% Timings

T = 10;
epsilon=1e-16;
clear tocj

for j=1:10
  prior = D;
  obs=starting_state;
  Qt=zeros(size(policies,2),T);
  Pt=zeros(length(actions),T);
  CA=zeros(1,T);
  obst=zeros(2,T);
  tic
  for t = 1:T

    obs_idx = find(grid_locations(:,1)==obs(1)&grid_locations(:,2)==obs(2));
    obst(:,t)=obs;
    % perform inference over hidden states
    
    qs_current = infer_states(obs_idx, A, prior, epsilon);
    bar(qs_current)
    title(strcat('Beliefs about location at time t=',num2str(t-1)))

    % calculate expected free energy of actions
    
    G = calculate_G_policies(A, B, C, qs_current, policies, epsilon);

    % to get action posterior, we marginalize P(u|pi) with the probabilities of each policy Q(pi), given by \sigma(-G)
    Q_pi = softmax(-G);
    Qt(:,t)=Q_pi; 
    % compute the probability of each action
    P_u = compute_prob_actions(actions, policies, Q_pi);
    Pt(:,t)=P_u;
    % sample action from probability distribution over actions
    chosen_action = datasample(1:5,1,'Weights',P_u);
    CA(t)=chosen_action;

    % compute prior for next timestep of inference
    prior = B(:,:,chosen_action)*qs_current;

    % step the generative process and get new observation
    
    % update generative process
    action_label = actions(chosen_action);
    
    [~, b]=max(qs_current);
    current_state = grid_locations(b,:);
    obs = get_env_N(current_state,action_label,N);
  
    end
    tocj(j)=toc;
  end

  %% saved results
  % N=4;T=10;
  Tj = load('tocj_dp_4_10_10.mat');
  %perpol = mean(Tj.tocj(2:end));

  % N=8;T=18
  Tj = load('tocj_dp_8_18_10.mat');
  %perpol = mean(Tj.tocj(2:end));

  % N=32;T=68
  Tj = load('tocj_dp_32_68.mat');
  %perpol = mean(Tj.tocj(2:end));
  
  
  % N=64;T=5
  Tj = load('tocj_dp_64_5.mat');
  %perpol = mean(Tj.tocj(2:end));

  % N=128;T=5
  Tj = load('tocj_dp_128_5_4.mat');
  %perpol = mean(Tj.tocj(2:end));

