#!/usr/bin/perl

# Open output files
open WPFILE,">wp.mat" || die("no openings\n");
open SPFILE,">sp.mat" || die("no openings\n");

# Read sentences (one per line)
while (<stdin>) {
  chomp;
  $sen[$num_sen]=[split / /];
  foreach $i (@{$sen[$num_sen++]}) {
    $words{$i}=++$num_words if ! $words{$i};
  }
}
print "$num_words different words\n";
print "$num_sen sentences\n";

######################################################
#Init the model
######################################################
for $cur_sense (0..1) {
  $tempsum=0.0;
  # random noise for probabilities for words in context
  for ($cur_word=1; $cur_word<=$num_words; $cur_word++) {
    $wp[$cur_sense][$cur_word]=rand()+1000;
    $tempsum+=$wp[$cur_sense][$cur_word];
  }
  # normalise wordprobs to 1
  for ($cur_word=1; $cur_word<=$num_words; $cur_word++) {
    $wp[$cur_sense][$cur_word]/=$tempsum;
  }
  # Equal probabilities for the priors
  $sp[$cur_sense]=0.5;
}


for ($iter=0;$iter<100;$iter++) {
############################################################
# Expectation step: update sentence probabilities
###########################################################
  for ($s=0; $s<$num_sen; $s++) {
    $sum=0.0;
    for $cur_sense (0..1) {
      $tempprob=1.0;
      foreach $i (@{$sen[$s]}) {
	$tempprob*=$wp[$cur_sense][$words{$i}]*$sp[$cur_sense];
      }
      $senprob[$cur_sense][$s]=$tempprob;
      $sum+=$senprob[$cur_sense][$s];
    }
    # normalize
    for $cur_sense (0..1) {
      $senprob[$cur_sense][$s]/=$sum;
    }
  }

#############################################################
# Maximization step: update context probabilities
#############################################################
  $spsum=0.0;
  for $cur_sense (0..1) {
    $wsum=0.0;
    for $cur_word (keys %words) {
      $tempsum=0.0;
      for ($s=0; $s<$num_sen; $s++) {
	$tempsum+= $senprob[$cur_sense][$s] if (find_word($cur_word,$sen[$s]));
      }
      $wp[$cur_sense][$words{$cur_word}]=$tempsum;
      $wsum+=$wp[$cur_sense][$words{$cur_word}];
    }
    # normalize
    for $cur_word (keys %words) {
      $wp[$cur_sense][$words{$cur_word}]/=$wsum;
    }

    # priors
    $sptemp=0.0;
    for ($s=0; $s<$num_sen; $s++) {
      $sptemp+=$senprob[$cur_sense][$s];
    }
    $sp[$cur_sense]=$sptemp;
    $spsum+=$sptemp;
  }
  for $cur_sense (0..1) {
    $sp[$cur_sense]/=$spsum;
    if ($iter<15) {    $sp[$cur_sense]=0.5; }
  }

#############################################################
# Print some running stats
############################################################
  $stat_senprob=0.0;
  for ($s=0; $s<$num_sen; $s++) {
    for $cur_sense (0..1) {
      $stat_senprob+=$senprob[$cur_sense][$s];
    }
  }
  if ($iter % 1 ==0) {
    print "Iter $iter, average sentence probability ".$stat_senprob/$num_sen."\n";
    printf "priors: %.2f %.2f\n", $sp[0],$sp[1];
    print_sentence_stats();
    print_word_probabilities();
  }
  if (($iter <20 && $iter%2==0) ||  $iter %10 ==0) {
    write_matlab_plot_data();
  }
}
print_sentences();

sub find_word {
  my $word=shift;
  my $wlist=shift;
  foreach (@$wlist) {
    if ( $_ eq $word) {
      return 1;
    }
  }
  return 0;
}

sub print_word_probabilities {
  print "Word probabilities:\n";
  print "\tsense0\t\tsense1\n";
  for $w (keys %words) {
    print "$w";
    for $cur_sense (0..1) {
      printf "\t%.2f",$wp[$cur_sense][$words{$w}];
    }
    print "\n";
  }
}

sub print_sentence_stats {
  print "x:\tsense1\tsense2\n";
  for ($s=0; $s<$num_sen; $s++) {
    print $s+1 .":";
    for $cur_sense (0..1) {
      printf "\t%.2f",$senprob[$cur_sense][$s];
    }
    print "\n";
  }
}

sub print_sentences {
  for ($s=0;$s<$num_sen;$s++) {
    print $s+1 .":";
    foreach $i (@{$sen[$s]}) {
      print " $i";
    }
    print "\n";
  }
}

sub write_matlab_plot_data {
  # Print word probs
  printf WPFILE "%d",$iter;
  for (keys %words) {
    print "!!!$_\n";
    printf WPFILE " %.3f", 
    ($wp[0][$words{$_}]/( $wp[0][$words{$_}]+$wp[1][$words{$_}]));
  }
  printf WPFILE "\n";

  # Print sentence probs
  printf SPFILE "%d",$iter;
  for ($i=0;$i<$num_sen;$i++) {
    printf SPFILE " %.3f",$senprob[0][$i];
  }
  printf SPFILE "\n";
}
