package Mail::Audit::Bayesian;
use Mail::Audit;
use vars q(@VERSION);
$VERSION = 0.001;
1;

package Mail::Audit;

use strict;
use MIME::Words qw(:all);
use DBI;
use Encode;
use Time::HiRes qw(time);


=head2 tokenize([$entity])

Split a message into tokens. This method is automatically called by classify
and add_to_corpus and therefore needs to be rarely invoked directly.

If $entitity is given it is tokenized and the tokens are added to the token 
list of $self. This is used to recursively tokenize nested MIME entities.

=cut

sub tokenize {
    my ($self, $entity) = @_;
    unless (defined $entity) {
	$self->{_bayesian_tokens} = {};
	# if there are too many tokens, the computation will underflow 
	# and we get a probability of 0.5 for every corpus.
	# so we limit the number of tokens to a reasonable number.
	$self->{_bayesian_max_tokens} = 500;

	# bias for low number of samples. We pretend that there
	# are _bayesian_default_strength more messages with each
	# token and that these have a probability of
	# _bayesian_default_belief.
	$self->{_bayesian_default_belief} = 0.5;
	$self->{_bayesian_default_strength} = 1;
	$entity = $self;
    }

    my @h = @{$entity->head()->header()};
    # tokenize header lines.
    # for now we just split each header line into words and 
    # add each each word and the word prepended by the header name
    # to the list of tokens.
    # treating some header lines special may be worthwhile.

    for my $h (@h) {
	my ($k, $v) = split(/:/, $h, 2);
	$self->add_token($k);
	$v = $self->_decode_rfc2047($v);
	my @t = ($v =~ m/(\w+)/g);
	for my $t (@t) {
	    $self->add_token("$k:$t");
	    $self->add_token($t);
	}
    }

    # tokenize the body
    # for MIME entities this depends on the content-type.
    if ($entity->isa('MIME::Entity')) {
	if ($entity->mime_type eq "text/plain") {
	    my $cs = $entity->head->mime_attr("content-type.charset") ;
	    $cs = "us-ascii" unless ($cs && Encode::resolve_alias($cs));
	    my $bh = $entity->bodyhandle();
	    my $io = $bh->open("r");
	    while (defined (my $l = $io->getline())) {
		$l = decode($cs, $l);
		my @t = ($l =~ m/(\w+)/g);
		for my $t (@t) {
		    $self->add_token($t);
		}
	    }       
	    $io->close();
	} else {
	    my @parts = $entity->parts;
	    for (@parts) {
		$self->tokenize($_);
	    }       
	}

    } else {
	for my $l (@{$entity->body}) {
	    my @t = ($l =~ m/(\w+)/g);
	    for my $t (@t) {
		$self->add_token($t);
	    }
	}
    }
}

sub add_token {
    my ($self, $token) = @_;
    if (scalar keys %{$self->{_bayesian_tokens}} < $self->{_bayesian_max_tokens}) {
	$self->{_bayesian_tokens}{$token} = 1;
    }
}

sub tokens {
    my ($self) = @_;
    $self->tokenize unless ($self->{_bayesian_tokens});
    return keys %{$self->{_bayesian_tokens}};
}

sub classify {
    my ($self, @corpora) = @_;

    $self->tokenize unless ($self->{_bayesian_tokens});
    my $tp;
    for my $corpus (@corpora) {
	my $msgcnt = $corpus->message_count;
	next if ($msgcnt == 0);

	my $corpus_name = $corpus->name;
	my $t0 = time;
	my $n = 0;
	my $nt = 0;
	for my $token ($self->tokens) {
	    my $tkcnt = $corpus->token_count($token);
	    my $tkprp = $tkcnt/$msgcnt;
	    $tp->{$token}{$corpus_name} = { p => $tkprp, n => $tkcnt, m => $msgcnt};
	    $n++;
	    $nt += $tkcnt;
	}
	my $t1 = time;
	printf STDERR "%d (%d) tokens  in %f seconds: %f tokens/second\n",
		      $n, $nt, $t1 - $t0, $n/($t1 - $t0)
	    if ($self->{_bayesian_debug});
    }

    # compute belief (f) based on the probabilities and frequencies
    # We remove all completely unknown tokens here - I'm not sure this
    # is correct, but IMHO they don't make any difference (they are unknown
    # in all corpora) and we have less to compute this way (and we don't divide by 0)
    for my $tk (keys %$tp) {
	my $te = $tp->{$tk};
	my $sn = 0;
	my $sm = 0;
	for my $ce (values %$te) {
	    $sn += $ce->{n};
	    $sm += $ce->{m};
	}
	if ($sn > 0) {
	    for my $ce (values %$te) {
		# p = p(in_corpus) / (p(in_corpus) + p(out_of_corpus))
		my $p = $ce->{p} / ($ce->{p} + ($sn - $ce->{n}) / ($sm - $ce->{m}));
		$ce->{f} = ($self->{_bayesian_default_strength} * $self->{_bayesian_default_belief} + $sn * $p)
			    /
			   ($self->{_bayesian_default_strength} + $sn);
	    }
	} else {
	    delete($tp->{$tk});
	}
    }

    # finally, combine beliefs about all tokens, test both null hypotheses
    # for all corpora and return result.
    my @i;
    my $eps = 1E-300;
    for my $corpus (@corpora) {
	my $corpus_name = $corpus->name;
	my $ph = 1;
	my $ps = 1;
	for my $tk (keys %$tp) {
	    my $te = $tp->{$tk};
	    print STDERR "f($tk, $corpus_name) = ", $te->{$corpus_name}->{f}, "\n"
		if ($self->{_bayesian_debug});
	    $ph *= $te->{$corpus_name}->{f};
	    $ps *= (1 - $te->{$corpus_name}->{f});
	}
	# clamp probabilities to positive value to avoid domain error in
	# log function.
	$ph = $ph > $eps ? $ph : $eps;
	$ps = $ps > $eps ? $ps : $eps;
	my $h = $self->ichi2(-2 * log($ph), 2 * scalar(values %$tp));
	my $s = $self->ichi2(-2 * log($ps), 2 * scalar(values %$tp));
	my $i = (1 + $h - $s) / 2;
	print STDERR "ph = $ph, ps = $ps, h = $h, s = $s, i = $i\n"
	    if ($self->{_bayesian_debug});
	push (@i, $i);
    }
    return @i;
}

sub ichi2 {
    my ($self, $chi, $df) = @_;
    my $m = $chi / 2.0;
    my $term = exp(-$m);
    my $sum = $term;
    for my $i (1 .. $df/2) {
        $term *= $m / $i;
        $sum += $term;
    }
    return $sum < 1.0 ? $sum : 1.0;
}


=head2 add_to_corpus

Add this message to the given corpus.

=cut

sub add_to_corpus {
    my ($self, $corpus, $force) = @_;

    $corpus->add($self, $force);
}


=head2 filename [$filename]

Get or set a file name to be associated with this message. The file name
is stored in the database.

=cut

sub filename {
    my ($self, $filename) = @_;

    my $oldfilename = $self->{_bayesian_filename};
    $self->{_bayesian_filename} = $filename if (defined($filename));
    return $oldfilename;
}

=head2 decode_rfc2047

decode a string encoded according to RFC 2047

The returned string is in internal perl string representation and has
the UTF-8 flag set if it contains any non-ascii characters.

=cut

sub _decode_rfc2047 {
    my ($dummy, $enc) = @_;

    my @words = decode_mimewords($enc);

    my $dec = "";
    for (@words) {
	eval {
	    $dec .= $_->[1] ? decode($_->[1], $_->[0]) : $_->[0];
	};
	if ($@) {
	    # if decoding fails for any reason (usually unknown charset)
	    # we just append he encoded word.
	    $dec .= $_->[0];
	}
    }
    return $dec;
}


# vim:sw=4
