
num_concepts = 100;
num_its = 30;

con_item_mp = load('0_continuous_item_mp');

[ num_blobs, num_features ] = size(con_item_mp);
num_images = num_blobs / 10;

dis_item_mp = load('0_discrete_item_sparse_mp');
[ num_images_2, max_num_words_per_image ] = size(dis_item_mp);

num_words = max(max(dis_item_mp));

num_blobs = 10 * num_images;

%
% Initialize
%

%
% Set the random number generator, so we don't get confused by randomness.
%
rand('state',0);

%
% The 10 regions (blobs) per image are stacked, row-wise, with the concepts
% occuring across columns. 
%
p_c_given_b = normalize_rows(rand(num_blobs, num_concepts).^5);

%
% The max_num_words_per_image words are stacked, row-wise, with the concepts
% occuring across columns. Note that we create initializations for the missing
% words (ones that have -99), but we don't care---we just never use it.
%
p_c_given_w = normalize_rows(rand(num_images * max_num_words_per_image, num_concepts).^5);


for it = 1:num_its
    
    % --------------    M Step   -------------------------

    p_c = sum(p_c_given_b,1) + sum(p_c_given_w, 1);
    sum_p_c = sum(p_c);
    p_c = p_c ./ sum_p_c;

    p_w_given_c = zeros(num_concepts, num_words);

    for i = 1:num_images
        for j = 1:max_num_words_per_image
            w = dis_item_mp(i,j);
            if (w > 0)
                p_w_given_c(:,w) = p_w_given_c(:,w) + (p_c_given_w((i-1)*max_num_words_per_image + j,:))';
            end
        end
    end

    p_w_given_c = normalize_rows(p_w_given_c);

    sum_x  = zeros(num_concepts, num_features);
    sum_x2 = zeros(num_concepts, num_features);
    weight = zeros(num_concepts, 1);

    for b = 1:num_blobs
        x = con_item_mp(b, :);

        for c = 1:num_concepts
            sum_x(c,:) = sum_x(c,:) + x * p_c_given_b(b,c);
            sum_x2(c,:) = sum_x2(c,:) + x .* x * p_c_given_b(b,c);
            weight(c) = weight(c) + p_c_given_b(b,c);
        end
    end

    u_of_c = zeros(num_concepts, num_features);
    v_of_c = zeros(num_concepts, num_features);

    for c = 1:num_concepts
        u_of_c(c,:) = sum_x(c,:) / weight(c);
        v_of_c(c,:) = 0.0001 + (sum_x2(c,:) - sum_x(c,:) .* sum_x(c,:) / weight(c)) / weight(c);
        log_root_det(c) = sum(log(v_of_c(c,:))) / 2;
    end

    % -------------------- E step ------------------------------

    ll = 0;
 
    for b=1:num_blobs
        p_b = 0;
 
        log_p_b_and_c = zeros(num_concepts,1);
 
        for c=1:num_concepts
            x = con_item_mp(b,:);
            u = u_of_c(c,:);
            v = v_of_c(c,:);
 
            log_p_b_given_c = -0.5 * sum((x-u).*(x-u) ./ v) - log_root_det(c);
            log_p_b_and_c(c) = log_p_b_given_c + log(p_c(c));
        end
 
        max_log_p_b_and_c = max(log_p_b_and_c);
        log_p_b_and_c = log_p_b_and_c - max_log_p_b_and_c;
        p_b_and_c = exp(log_p_b_and_c);
        p_b = sum(p_b_and_c);
        p_c_given_b(b,:) = p_b_and_c;
 
        ll = ll + log(p_b) + max_log_p_b_and_c;
    end
 
    p_c_given_b = normalize_rows(p_c_given_b);
 
    for i=1:num_images
        p_c_given_image = sum(p_c_given_b((i-1)*10+1:i*10,:),1);
 
        p_c_given_image = p_c_given_image ./ sum(p_c_given_image);
 
        for j=1:max_num_words_per_image
            w = dis_item_mp(i,j);
            if (w > 0)
                p_w_and_c = p_w_given_c(:,w)' .* p_c_given_image;
                sum_p_w_and_c = sum(p_w_and_c);
 
                ll = ll + log(sum_p_w_and_c);
 
                p_c_given_w((i-1) * max_num_words_per_image + j,:) = p_w_and_c / sum_p_w_and_c;
            end
        end
    end
 
    if (it < 2)
        fprintf(1,'%.3e\n', ll);
    else
        fprintf(1,'%.3e      %.3e\n', ll, ll - prev_ll);
    end
 
    prev_ll = ll;
end

save model.mat u_of_c v_of_c log_root_det p_w_given_c p_c


