classdef eGreedy < handle
    %EGREEDY implements Epsilon Greedy Strategy for Multi-Armed Bandit
    
    properties
        n_arms                                  % number of arms
        epsilon                                 % ratio of exploration
        trials                                  % counter for trials
        rewards                                 % accumulates rewards earned
        scores                                  % track expected rewards
    end
    
    methods
        function self = eGreedy(n_arms,epsilon) % initialize object
            self.n_arms = n_arms;               % number of arms
            self.trials = zeros(1,n_arms);      % initialize empty vector
            self.rewards = zeros(1,n_arms);     % initialize with 0
            self.scores = ones(1,n_arms);       % initialize with 1
            if nargin == 1                      % epsilon not given
                self.epsilon = 0.1;             % use default 0.1
            else                                % otherwise
                self.epsilon = epsilon;         % use user input
            end
        end
        
        function s = score(self, arm)           % compute expected reward
            s = self.scores(end, :);            % carry over current scores
            s(arm) = (self.rewards(arm) + 1) ...% update the score
                / (self.trials(arm) + 1);        
        end
        
        function choice = choose(self)          % choose an arm
            if rand <= self.epsilon             % (epsilon) times
                choice = randi(self.n_arms);    % explore all arms
            else                                % (1 - epsilon) times
                [~, choice] = ...               % exploit the best arm
                    max(self.scores(end,:));
            end       
        end
        
        function update(self, arm, reward)      % update data
            self.trials(arm) =  ...             % increment counter
                self.trials(arm) + 1; 
            self.rewards(arm) = ...             % increment reward
                self.rewards(arm) + reward;
            self.scores(end + 1,:) = ...        % append scores
                self.score(arm);
        end
        
        function plot(self,title_str,options)   % plot scores
            n = size(self.scores,1);            % number of trials
            figure                              % new figure
            plot(1:n, self.scores)              % plot
            xlim([0 n])                         % x-axis limits  
            title(title_str)                    % add title
            xlabel('Trials')                    % add x-axis label
            ylabel('Expected Reward')           % add y-axis label
            legend(options)                     % add legend      
        end
        
    end
    
end

