clear all; close all;

[spmatrix, tokenlist, trainCategory] = readMatrix('MATRIX.TRAIN.1400');

trainMatrix = full(spmatrix);
numTrainDocs = size(trainMatrix, 1);
numTokens = size(trainMatrix, 2);

% trainMatrix is now a (numTrainDocs x numTokens) matrix.
% Each row represents a unique document (email).
% The j-th column of the row $i$ represents the number of times the j-th
% token appeared in email $i$. 

% tokenlist is a long string containing the list of all tokens (words).
% These tokens are easily known by position in the file TOKENS_LIST

% trainCategory is a (1 x numTrainDocs) vector containing the true 
% classifications for the documents just read in. The i-th entry gives the 
% correct class for the i-th email (which corresponds to the i-th row in 
% the document word matrix).

% Spam documents are indicated as class 1, and non-spam as class 0.
% Note that for the SVM, you would want to convert these to +1 and -1.

%------------------------------------------------------------
% Tony Hyun Kim
% CS 229, PS#2, Problem 3(a)
%------------------------------------------------------------

% The model parameters are:
%   p_word(i,k): probability of word k occuring given
%                   i=1 (non-spam) or i=2 (spam) classification
%   p_prior:     class prior probability
p_word = zeros(2,numTokens);

% We are building a Naive Bayes generative model. So, we begin
%   by splitting the training examples by their classes. We'll
%   start with the non-spam examples
sieve = ~logical(full(trainCategory));
for i = 1:2
    class  = trainMatrix(sieve',:);
    Nclass = size(class,1); % Number of examples in class
    
    word_hist = sum(class,1); % Histogram of words in class
    Nwords = sum(word_hist);
    
    % MLE estimate for conditional prob with Laplace smoothing
    p_word(i,:) = (word_hist+1)/(Nwords+numTokens);
    
    sieve = ~sieve; % Select the next class
end

% Compute the prior class probabilities. Note that I use 
%  Laplace smoothing here as well
p_prior = (Nclass+1)/(numTrainDocs+2);
clear sieve i class Nclass word_hist Nwords;

% We base our computation on the log-likelihoods rather than 
%   direct likelihoods (i.e. p_word) from 'nb_train'
ll_word  = log(p_word);
ll_prior = log([1-p_prior p_prior])'; 

% Problem 3(b): Particularly 'spammy' tokens
%------------------------------------------------------------
ll_diff  = diff(ll_word); % Higher values here indicate
                          %     especially spammy words
[~, spammy_index] = sort(ll_diff,'descend');
spammy_index(1:5) % Display the top five spammy tokens