#!/usr/bin/perl

# Written by Jason Rennie (jr6b@andrew.cmu.edu)

# Format of data files is as follows:
# line 1: target value names each separated by a space
# line 2: # of total words for all documents in each target value each
#         separated by a space
# line 3: number of documents for each category
# line 4: a vocab word, followed by pairs of category numbers and the word
#         frequency for that category

# When a vocab word is refiled it is assumed that this is the desired
#      target value (or folder, in this case).

# tv = target value

# This program is written in an attempt to emulate the MH refile command
#   while adding the benefit of allowing it to learn user preference with
#   a bayesian-style learning algorithm.  
# The only options which it currently emulates are '-src +folder' and '-file
#   file'.  It supports refiling multiple messages, but does not support
#   filing messages into multiple folders.

open(TMP, ">> /tmp/refile.info")
    || warn "Could not open /tmp/refile.info: $!\n";

print TMP "irefile called ";
 print TMP "(@ARGV)\n";

$home_dir = $ENV{"HOME"}."/";
 
print TMP "Home dir = $home_dir\n";

&initialize_file_locks;

open(MH, "< $home_dir.mh_profile") 
    || die "No .mh_profile found: $!\n";

while (<MH>) {
    if (m/Path:\s*(\w+)/) {
	$mail_path = $1;
	$mail_path .= "/";
    }
    if (m/draft-folder:\s*(\w+)/) {
	$draft_folder = $1;
    }
    if (m/context:\s*(\S+)/i) {
	$context_file = $1;
    }
    if (m/Inbox:\s*(\S+)/i) {
	$inbox_folder = $1;
    }
}

if (!$context) {
    $context_file = $home_dir.$mail_path."context";
}
if (!$inbox_folder) {
    $inbox_folder = "inbox";
}

print TMP "Mail Path = $mail_path\n";
print TMP "Draft Fodler = $draft_folder\n";
print TMP "Context File = $context_file\n";
print TMP "Inbox Folder = $inbox_folder\n";

&write_idata_lock;

&read_idata_file;

$queue_file = $home_dir.".idata.queue";
    
if (open(QUEUE, "< $queue_file"))
{
    while (!eof(QUEUE))
    {
	&read_next_queue_message;
	&update_global_stats;
	&clear_queue_variables;
	# words are aged with the following variable during writing of .idata
	$messages_processed++;
    }
    &die_if_rw_lock_is_set($rw_lock);
    &read_write_idata_lock;
    &write_idata_file;
    system "rm ".$queue_file;
    &unlock_idata;
}
else
{
    print TMP "Could not open ${queue_file}\n";
    print STDERR "Could not open ${queue_file}\n";
}

print TMP "\n";
close(TMP);

exit 0;

################

# If the PID *doesn't* exist, then something has died.  It *should* be safe
# to proceed, unless the 'total lock' ($rw_lock) is turned on.  In that 
# case, something very bad has happened, and we should warn the user, and just
# plain abort.
sub die_if_rw_lock_is_set 
{
    my($rw_lock) = @_;
    if (-e $rw_lock) {
	$oldpid = `cat $rw_lock`;
	chop $oldpid;
	open(PROCS, "ls /proc | grep $oldpid|");
	while (<PROCS>) {
	    if (/^0*$oldpid\s*/) {
# Another writer is running...
		&ouch("Another writer running as PID $oldpid\n");
      }
	}
	# Other process not found - BAD problem.
	&ouch("Another writer was running as PID $oldpid, but seems to have died.\n\nMake sure $home_dir.idata looks OK, remove $rw_lock, and try again.\n");
    }
    
}

sub die_if_other_writer_running 
{
    my($w_lock) =@_;
    if (-e $w_lock) 
    {
	$oldpid = `cat $w_lock`;
	chop $oldpid;
	open(PROCS, "ls /proc | grep $oldpid|");
	while (<PROCS>) {
	    if (/^0*$oldpid\s*/) {
		# Another writer is running...
		&ouch("Another writer already running as PID $oldpid\n");
	    }
	}
# Other process not found - minor problem.
	print TMP "Another writer was running as PID $oldpid, but seems to have died.\n";
	print TMP "Will overwrite $w_lock file.\n";
	print "Another writer was running as PID $oldpid, but seems to have died.\n";
	print "Will overwrite $w_lock file.\n";
    }
}

sub initialize_file_locks
{
    # Define locks...
    $w_lock = $home_dir . ".idata.write.lock";
    $rw_lock = $home_dir . ".idata.lock";

    &die_if_rw_lock_is_set($rw_lock);
    &die_if_other_writer_running($w_lock);
}

# completely restricts access to .idata
sub read_write_idata_lock
{
    open(RWLOCK, ">$rw_lock");
    print RWLOCK $$;  # Process ID
    close(RWLOCK);
}

# restricts writing access to .idata
sub write_idata_lock
{
    open(WLOCK, ">$w_lock");
    print WLOCK $$;  # Write Process ID out to lock file
    close(WLOCK);
}

sub unlock_idata
{
    unlink $w_lock, $rw_lock;
}

sub ouch {
    local($string) = @_;
    print TMP "$string: \n";
    die "$string: $!\n";
}

sub print_usage {
    print "Usage: irefile [msgs] [-src +folder] [-file file] [MH options] +folder\n";
    print "       irefile -help\n\n";
    print " msgs         => text files to be refiled into +'folder'\n";
    print " -src +folder => looks for 'msgs' in specified folder\n";
    print " -file file   => disregards previous two options; uses 'file' as\n";
    print "                 source message\n";
    print " +folder      => folder to direct messages into\n";
    print " MH options   => all command line options will be passed to the\n";
    print "                 MH refile, so all MH refile options are valid\n\n";
    exit(0);
}


sub read_idata_file
{
    $tot_words = 0;
    
    # target values and 'n' both start with index 1
    $tv[0] = 0;
    
    $data_file = $home_dir.".idata";
    
    if (open(DATA, "< $data_file")) 
    {
	# reads in category names
	$line = <DATA>;
	while ($line =~ s/(\S+)\s+//) 
	{
	    $tv[$#tv+1] = $1;
	}
	
	# reads in total number of words for each category
	$line = <DATA>;
	$i = 1;
	while ($line =~ s/(\d+)\s+//) 
	{
	    $n{$tv[$i++]} = $1;
	}
	if ($i <= $#tv)
	{
	    &ouch("Bad data file format - line #2");
	}
	    
	# reads in number of documents for each category
	$line = <DATA>;
	$i = 1;
	while ($line =~ s/(\d+)\s+//)
	{
	    $docs{$tv[$i++]} = $1;
	    $tot_docs += $1;
	}
	if ($i <= $#tv)
	{
	    &ouch("Bad data file format - line #3");
	}
	
	# reads in actual word frequencies
	$word_array[0] = "";
	
	while (($line = <DATA>) && (!($line =~ m/^\s*$/)))
	{
	    $tot_words++;
	    
	    # reads word
	    $line =~ s/^(\S+)\s+// || &ouch("Couldn't read vocab word");
	    $word_array[$#word_array+1] = $1;
	    $index{$1} = $#word_array;
	    
	    # reads age of word
	    $line =~ s/^(\d+)\s+// || &ouch("Couldn't read age of word");
	    $age{$word_array[$#word_array]} = $1;
	    
	    # reads in frequency per category   cat_num:freq
	    while ($line =~ s/(\d+):(\d+)\s+//) 
	    {
		$nk{$word_array[$#word_array].",".$tv[$1]} = $2;
		$sum_nk{$word_array[$#word_array]} += $2;
	    }
	    
	    if (!($line =~ m/^\s*$/))
	    {
		&ouch("Vocab line # $tot_words has bad format");
	    }
	}
	close(DATA);
	print TMP "Read data file\n";
    }
}


sub clear_queue_variables
{
    @msg_word_array = 0;
    @freq_array = 0;
    @add_folder = 0;
    @delete_folder = 0;
    undef %msg_word_index;
}

# Queue file format is as follows.  For Queue file lexing, a word is a sequence
#   of space separated characters.
# A message is begun by any line which contains as its first character
#   something other than a whitespace character.
# The first line of the file should always begin a message.
# A message is ended by two consecutive feedline characters or by the end
#   of the file.  It is not necessary that the last line of the file be
#   ended with a feedline.
# Lines which are part of a message may be formatted in any fashion as long
#   as they could not be interpreted as the end of a message.
# A word/frequency line can be identified by a line which begins with a
#   sequence of non-space characters, followed by one or more spaces, followed
#   by an integer.
# Each message should be started with a line which indicates the database
#   action and the folder name(s) to perform the actions on.  Valid actions
#   include 'add' and 'delete'.  The first word on the line to match
#   an action name will be considered the action for the message.  All other
#   words will be considered folder names on which to perform the action.
#   More than one action line may be included in a message.  To move the
#   stats for a message from folder A to folder B, one would include the
#   lines "add B" and "delete A".
# Each message is represented as a sequence of consecutive lines, each of
#   which begins with a word, followed by a space followed by an integer.

# Reads in information about next message in QUEUE file.

sub read_next_queue_message
{
    $msg_word_array[0] = "nil";
    $freq_array[0] = "nil";

    $line = <QUEUE>;
    while ($line =~ /^\s/) { $line = <QUEUE>; }
    
    # while the first character isn't a space...
    while ($line =~ /^\S/)
    {
	$action = "";
	if ($line =~ /^(\S+)\s+(\d+)/)
	{
	    $word = $1;
	    $freq = $2;
	    
	    # print "Found word/frequency pair: ${word}/${freq}\n";

	    # if this is the first time we've seen the word...
	    if ($msg_word_index{$word} == 0)
	    {
		$msg_word_array[@msg_word_array] = $word;
		$msg_word_index{$word} = $#msg_word_array;
		$freq_array[$msg_word_index{$word}] = $freq;
	    }
	    else
	    {
		$freq_array[$msg_word_index{$word}] += $freq;
	    }
	}
	else
	{
	    @action_words = split(/\s+/, $line);
	    foreach $word (@action_words)
	    {
		if ($word eq "add" && $action eq "")
		{
		    $action = "add";
		}
		elsif ($word eq "delete" && $action eq "")
		{
		    $action = "delete";
		}
	    }
	    foreach $word (@action_words)
	    {
		if ($action eq "add" && $word ne "add")
		{
		    $add_folder[@add_folder] = $word;
		    # print "Add folder: $word\n";
		}
		elsif ($action eq "delete" && $word ne "delete")
		{
		    $delete_folder[@delete_folder] = $word;
		    # print "Delete folder: $word\n";
		}
	    }
	}
	$line = <QUEUE>;
    }
}


# Uses the information obtained from reading the queue to
#   update data read in from .idata
sub update_global_stats
{
    # the nesting here represents the format of .idata and may allow more
    #   efficient updating if the in memory representation of .idata is changed
    for ($i = 1; $i < @msg_word_array; $i++)
    {
	$this_word = $msg_word_array[$i];
	$this_freq = $freq_array[$i];

	# check for existence
	if ($index{$this_word} == 0)
	{
	    $word_array[@word_array] = $this_word;
	    $index{$this_word} = $#word_array;
	}

	foreach $folder (@add_folder)
	{
	    $nk{$this_word.",".$folder} += $this_freq;
	    $n{$folder} += $this_freq;
	}

	foreach $folder (@delete_folder)
	{
	    $nk{$this_word.",".$folder} -= $this_freq;
	    $n{$folder} -= $this_freq;
	}
    }
    foreach $folder (@add_folder)
    {
	$docs{$folder}++;
    }
    
    foreach $folder (@delete_folder)
    {
	$docs{$folder}--;
	if ($docs{$folder} < 0) { $docs{$folder} = 0; }
    }
}


sub write_idata_file
{
    # reopen the data file and write the learned information

    open(DATA, "> $home_dir.idata")
	|| &ouch("Could not open ${home_dir}.idata\n");
    
    open(FOLDERS, "< ${home_dir}${mail_path}.folders")
	|| &ouch("Couldn't open ${home_dir}${mail_path}.folders");
    
    # prints category names to data file (uses .folders file to determine folders)
    $i = 0;
    while ($folder = <FOLDERS>) 
    {
	chop $folder;
	print DATA $folder." ";
	$f_tv[++$i] = $folder;
    }
    
    close(FOLDER);
    
    print DATA "\n";
    
    $tv[0] = "nil";
    $n{"nil"} = 0;
    
    # prints total word frequency values
    for ($i=1; $i <= $#f_tv; $i++) 
    {
	if ($n{$f_tv[$i]} > 0) {
	    print DATA $n{$f_tv[$i]}." ";
	} else {
	    print DATA "0 ";
	}
    }
    print DATA "\n";
    
    # prints number of documents for each category
    for ($i=1; $i <= $#f_tv; $i++) 
    {
	if ($n{$f_tv[$i]} > 0) {
	    print DATA $docs{$f_tv[$i]}." ";
	} else {
	    print DATA "0 ";
	}
    }
    print DATA "\n";
    
    for ($i=1; $i <= $#word_array; $i++) 
    {
	$this_word = $word_array[$i];
	$this_age = $age{$this_word} + $messages_processed;
	
	# if the word is infrequent, discard it (don't print it to DATA)
	if ($sum_nk{$this_word} > (log($this_age)/log(2) - 1))
	{
	    print DATA $this_word." ".$this_age." ";
	    
	    for ($j=1; $j <= $#f_tv; $j++)
	    {
		$this_nk = $nk{$this_word.",".$f_tv[$j]};
		if ($this_nk > 0) {
		    print DATA $j.":".$this_nk." ";
		}
	    }
	    print DATA "\n";
	}
    }
    
    close(DATA);
}
