% PURPOSE: Demonstrate Machine Translation with EM

% AUTHOR   : Nando de Freitas.
% DATE     : July 2001


clear;

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%-- SPECIFY SIMULATION PARAMETERS  --%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

EM_tolerance = 1e-4;              % Tolerance.
max_EM_iterations = 20;           % Maximum number of EM steps.


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%--    GENERATE THE TRAINING DATA  --%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%


% I assume the data is given in the following format:
% =================================================================

E = {'b c';'b'};
F = {'x y';'y'};


E = {'the cat sings';'a lovely cat';'the dog dances';'lovely';'big dog'};
%F = {'el gato canta';'un bellisimo gato';'el perro baila';'bellisimo';'perro grande'};
F = {'le chat chante';'un chat beau';'le chien danse';'beau';'grand chien'},

E = {'sun sky sea';
     'sun sky cloud';
     'sun';
    }
F = {'1 1 2 3';
     '1 2 4';
     '2';
    }

E = {'Garcia and groups';
     'the modern groups sell pharmaceuticals';
     'the small groups sell program';
     'the program has been implemented';
     'the groups implemented program';
     'Garcia has implemented program';
     'Garcia do not sell pharmaceuticals and program';
     'the groups are not strong';
}
F = {'Garcia y grupos';
     'los grupos modernos venden medicinas';
     'los grupos pequenos venden programme';
     'los programme a ete mis en application';
     'los grupos mis en application programme';
     'Garcia a mis en application programme';
     'Garcia no venden medicinas y programme';
     'los grupos no son fuertes'; 
}

 
% ==================================================================

% I then transform this format so that e{2}{3}='cat' - i.e. just like in the 
% tech report where we write e{s}{i}. 

[S,arb] = size(E);
L = ones(S,1);
e = cell(S,1);
for s=1:S,
  spaces = findstr(char(E(s)),' ');
  L(s) = length(spaces) + 1;
  sentence = char(E(s)) ;  
  e{s}=cell(L(s),1);
  if L(s)==1,
    word=sentence;
    e{s}{1}=word;
  else
    for i=1:L(s)  
      if i==1,
        word = sentence(1:spaces(1)-1);
      elseif i==L(s),
        word = sentence(spaces(i-1)+1:length(sentence));
      else
        word = sentence(spaces(i-1)+1:spaces(i)-1);
      end;
      e{s}{i} = word;
    end;
  end;
end;
M = ones(S,1);
f = cell(S,1);
a = cell(S,1);
for s=1:S,
  spaces = findstr(char(F(s)),' ');
  M(s) = length(spaces) + 1;
  sentence = char(F(s)) ;  
  f{s}=cell(M(s),1);
  a{s}=cell(M(s),1);   % Intialise the assignments cell.
  if M(s)==1,
    word=sentence;
    f{s}{1}=word;
  else
    for j=1:M(s)  
      if j==1,
        word = sentence(1:spaces(1)-1);
      elseif j==M(s),
        word = sentence(spaces(j-1)+1:length(sentence));
      else
        word = sentence(spaces(j-1)+1:spaces(j)-1);
      end;
      f{s}{j} = word;
    end;
  end;
end;



%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%--        INITIALISATION          --%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

proportions = cell(S,1);        % p(as_j=i)
indicator   = cell(S,1);        % p(as_j=i|f_sj,e_si)
for s=1:S,
  proportions{s} = ones(L(s),M(s));         % Random initialisation.
  for j=1:M(s)
    proportions{s}(:,j) = proportions{s}(:,j)/sum(proportions{s}(:,j)); 
  end;
  indicator{s} = zeros(L(s),M(s));       % Doesn't matter.       
end;  


% Initialising t(f_sj|e_si) is more tricky because we need to differentiate between
% instances and tokens.

L_total = sum(L); % Ignoring repeats 
M_total = sum(M); % Ignoring repeats

e_bag = cell(L_total,1);
L_T = 0;
for s=1:S,
  for i=1:L(s),
    exists_already = sum(strcmp(e_bag,e{s}{i}));
    if exists_already == 0,   % i.e. it's not in the bag.   
      L_T = L_T + 1;
      e_bag{L_T} = e{s}{i};    
    end;
  end;
end;
e_tokens = e_bag(1:L_T);    % All the english tokens.
f_bag = cell(M_total,1);
M_T = 0;
for s=1:S,
  for j=1:M(s),
    exists_already = sum(strcmp(f_bag,f{s}{j}));
    if exists_already == 0,   % i.e. it's not in the bag.   
      M_T = M_T + 1;
      f_bag{M_T} = f{s}{j};    
    end;
  end;
end;
f_tokens = f_bag(1:M_T);    % All the french tokens.

translation = ones(L_T,M_T)/M_T;   % t(f_sj|e_si)



%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%-- PERFORM ML ESTIMATION WITH EM  --%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%


loglikelihood = zeros(max_EM_iterations,1);
iterations = 1;
tol = 1e99; % arbitrary large number;
while ((iterations <= max_EM_iterations) & (tol >= EM_tolerance))
  fprintf('EM iteration = %i / %i ',iterations,max_EM_iterations);

  % ======= 
  % E STEP: 
  % ======= 

  % Compute the posterior association p(as_j=i|f_sj,e_si)
  for s=1:S,
    for j=1:M(s),
      for i=1:L(s),
        m = strmatch(e{s}{i},e_tokens);   % position of token in bag.
        if length(m)>1,
          m = m(find(strcmp(e{s}{i},e_tokens(strmatch(e{s}{i},e_tokens)))==1)); 
        end; 
        n = strmatch(f{s}{j},f_tokens);
        if length(n)>1,
          n = n(find(strcmp(f{s}{j},f_tokens(strmatch(f{s}{j},f_tokens)))==1)); 
        end; 
        indicator{s}(i,j) = proportions{s}(i,j) * translation(m,n);
      end; 
    end;   
    for j=1:M(s),
      indicator{s}(:,j) = indicator{s}(:,j)/sum(indicator{s}(:,j));   % normalise.
    end;
  end;

  % =======
  % M STEP:
  % =======

  % Update mixture proportions p(a_sj=i)

  for m=1:max(M),
    for l=1:max(L),
      correct_size = zeros(S,1);
      the_sum=zeros(l,m);
      for s=1:S,
        if ((l==L(s))&(m==M(s))),
          correct_size(s) = 1;
          the_sum = the_sum + indicator{s}(:,:);
        end;
      end;
      for s=1:S,
        if correct_size(s)==1,
          proportions{s}(:,:) = the_sum;
          for j=1:M(s),
            proportions{s}(:,j)=proportions{s}(:,j)/sum(proportions{s}(:,j));
          end;
%      prop=proportions{s}

        end; 
      end;
   end;
  end;
      

  % Update translation probabilities t(f_sj|e_si)
  for m=1:L_T,
    for n=1:M_T,
      the_sum = 0;
      for s=1:S,
        for j=1:M(s),           
          for i=1:L(s), 
            pair = sum(strcmp(e_tokens(m),e{s}{i}))*sum(strcmp(f_tokens(n),f{s}{j})); 
            if pair
              the_sum = the_sum + pair*indicator{s}(i,j);   
            else
              % don't add.
            end;    
          end;
        end;
      end;  
      translation(m,n) = the_sum;
    end;
    translation(m,:) = translation(m,:)/sum(translation(m,:));
  end;

  % Compute the loglikelihood for monitoring. 
  loglikelihood(iterations)=0;
  
  the_s_sum = 0;
  for s=1:S,
    the_j_sum = 0;
    for j=1:M(s),
      the_i_sum = 0;
      for i=1:L(s),
        m = strmatch(e{s}{i},e_tokens);   % position of token in bag.
        if length(m)>1,
          m = m(find(strcmp(e{s}{i},e_tokens(strmatch(e{s}{i},e_tokens)))==1)); 
        end; 
        n = strmatch(f{s}{j},f_tokens);
        if length(n)>1,
          n = n(find(strcmp(f{s}{j},f_tokens(strmatch(f{s}{j},f_tokens)))==1)); 
        end;   
        the_i_sum = the_i_sum + proportions{s}(i,j) * translation(m,n);
      end;
      the_j_sum =  the_j_sum + log(the_i_sum); 
    end;  
    the_s_sum = the_s_sum + the_j_sum;
  end;
  loglikelihood(iterations) =  loglikelihood(iterations) + the_s_sum; 
  fprintf('loglikelihood = %2.2d \n',loglikelihood(iterations));
  if iterations>1
    tol = loglikelihood(iterations) - loglikelihood(iterations-1);
  end;
  iterations = iterations+1;
end; % End of EM iterations.


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%--        RESULTS AND PLOTS       --%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%tokens = [e_tokens f_tokens]
traslation = translation

figure(1)
clf;
%hold on;
plot(1:iterations-1,loglikelihood(1:iterations-1),'r','LineWidth',2);  
ylabel('Log-likelihood','fontsize',15);
xlabel('EM iterations','fontsize',15);

English = E
French  = F
for m=1:L_T,
  for n=1:M_T,
    e_word = char(e_tokens(m));
    f_word = char(f_tokens(n));
    fprintf('p(F = %10s | E = %s)   =   %1.2f\n',f_word,e_word,translation(m,n)); 
   end;
end;
fprintf('\n');





