#!/usr/bin/env perl

#
# Evaluate a TREC 2010 entity track run.  The measures computed are
# NDCG@R, MAP, Rprec, and P@10.
#
# Version: 1.2.
#
# Usage: eval-entity.pl [-v] [--name] <qrels> <runfile>\n";
#        -v: verbose mode (default: no)
#        --name: give extra credit for a correct NAME for a primary (default: no)
#        <qrels> and <runfile> may be gzipped.
#
# In verbose mode, the runfile is output in an annotated fashion,
# along with the evaluation scores, showing judgments pages
# and names.
#
# nDCG is computed at rank (R), the number of primaries and relevants 
# for the topic.  Relevant pages get gain 1, primary pages gain 3.
# If the --name option is specified, primaries with a correct name
# receive gain 4.
#
# P@10, MAP, and Rprec are computed over primary pages only.
# 
# This version does little if any error checking on the runfile.  It
# is highly recommended that you run check_entity.pl before running
# this eval script, especially for unofficial runs.
#
# This script expects to be able to find GNU gzip and a POSIX-compatible
# sort(1) on the path.
#
# HISTORY
#
# 1.2  adapted to TREC 2010 run and qrels formats
# 1.1  changed NAMEs to case-insensitive match
# 1.0  original release

use strict;
use Getopt::Long;

my $usage = "eval-entity.pl [-v] [--name] <qrels> <runfile>\n";
my $verbose = 0;
my $namecredit = 0;
GetOptions('verbose!' => \$verbose,
	   'name!' => \$namecredit,
    ) or die $usage;

my $qrels_file = shift or die $usage;
my $run_file = shift or die $usage;

my %qrel;
my %class;
my %cref;
my %eval;

open QRELS, $qrels_file or die "Can't open `$qrels_file': $!\n";
while (<QRELS>) {
    chomp;
    my ($topic, $page, $name, $drel, $cat, $nrel) = split;

    $name = lc $name;

    $qrel{$topic}{$page} = $drel;
    $qrel{$topic}{$name} = $nrel;
    $class{$topic}{$page} = $cat;
    $class{$topic}{$name} = $cat;
    $cref{$topic}{$cat} = 1;
}
close QRELS;

# Scan the qrels to build ideal gain vectors for nDCG
for my $topic (keys %qrel) {
    for my $thing (keys %{ $qrel{$topic} }) {
	if ($thing =~ /^clueweb09-en/) {
	    my $rel = $qrel{$topic}{$thing};
	    if ($rel == 1) { $qrel{$topic}{num_rel}++; }
	    if ($rel == 2) { $qrel{$topic}{num_pri}++; }

	    my $gain = $rel;
	    if ($rel == 2) {
		$gain = 3;
		if ($namecredit and
		    exists $cref{$topic}{$class{$topic}{$thing}}) {
		    $gain = 4;
		}
	    }
	    push @{ $qrel{$topic}{list} }, $gain;
	}
    }
    $qrel{$topic}{list} = [ sort {$b<=>$a} @{ $qrel{$topic}{list} } ];
}

my $last_topic = -1;
my $cat = "cat";
$cat = "gzip -dc" if ($run_file =~ /.gz$/);

open RUN, "$cat $run_file | sort -k1,1n -k5,5gr -k3,3r |"
    or die "Can't open `$run_file': $!\n";
while (<RUN>) {
    s/\r//g;
    next if /^$/;
    chomp;

    my ($topic, undef, $page, $rank, $sim, $tag, $name) = split;
    $topic =~ s/^0+//;
    next unless exists $qrel{$topic};

    if ($topic ne $last_topic) {
	if ($last_topic != -1) {
	    if ($verbose) {
		my $i = 1;
		my $str = "";
		for my $gain (@{ $eval{list} }) {
		    $str .= $gain;
		    $str .= " / " unless ($i % 10);
		    $i++;
		}
		print "# Gains: $str\n";
		$i = 1;
		$str = "";
		for my $gain (@{ $qrel{$last_topic}{list} }) {
		    last if $gain == 0;
		    $str .= $gain;
		    $str .= " / " unless ($i % 10);
		    $i++;
		}
		print "# Ideal gains: $str\n";
	    }
	    
	    score(\%eval);
	    %eval = ();
	}
	$eval{topic} = $topic;
	$last_topic = $topic;
    }

    my $rel = 0;
    my $gain = 0;
    $name = lc $name;
    if (exists $qrel{$topic}{$page}) {
	$rel = $qrel{$topic}{$page};
	if ($rel == 1)    { $eval{rel_ret}++;  $gain = 1; }
	elsif ($rel == 2) { $eval{pri_ret}++;  $gain = 3; }
    }

    if ($namecredit and ($rel == 2) and ($qrel{$topic}{$name} == 2)
	and ($class{$topic}{$page} == $class{$topic}{$name})) {
	$gain = 4;
    }

    if ($verbose) {
	printf "#  %-5s $topic Q0 $page $rank $sim $tag $name\n", $gain;
    }

    $eval{num_ret}++;
    push @{ $eval{list} }, $gain;
}

if ($verbose) {
    my $i = 1;
    my $str = "";
    for my $gain (@{ $eval{list} }) {
	$str .= $gain;
	$str .= " / " unless ($i % 10);
	$i++;
    }
    print "# Gains: $str\n";
}
score(\%eval);
means();

my %cum = ();

sub score {
    my ($eval) = @_;

    my $p10 = 0;
    my $dcg = 0;
    my $idcg = 0;
    my $pri_so_far = 0;
    my $sum_prec = 0;
    my $pprec = 0;
    my $topic = $eval->{topic};

    my $rank = 0;
    for my $e (@{ $eval->{list} }) {
	$rank++;
	last if $rank > 100;

	if ($e > 1) {
	    $pri_so_far++;
	    $sum_prec += ($pri_so_far / $rank);
	    if ($rank <= 10) { $p10 += 0.1; }
	}
	
	if ($rank == $qrel{$topic}{num_pri}) {
	    $pprec = $pri_so_far / $rank;
	}
	
	my $ie = $qrel{$topic}{list}[$rank - 1];
	next if $ie == 0;
	$idcg += ($ie / (log($rank + 1)/log(2)));
	$dcg += ($e / (log($rank + 1)/log(2)));
    }

    if ($rank < $qrel{$topic}{num_pri}) {  # run retrieved fewer than R docs
	$pprec = $pri_so_far / $qrel{$topic}{num_pri};
    }
    
    # The run may retrieve a short list, but the ideal gain is up to
    # the number of nonzero gains available in the topic.
    $rank++;
    while ($rank < scalar @{ $qrel{$topic}{list} } &&
	   $qrel{$topic}{list}[$rank - 1] != 0) {
	my $ie = $qrel{$topic}{list}[$rank - 1];
	$idcg += ($ie / (log($rank + 1)/log(2)));
	$rank++;
    }
	
    my $ndcg = $dcg / $idcg;
    my $map = 0;
    if ($qrel{$topic}{num_pri} > 0) {
	$map = $sum_prec / $qrel{$topic}{num_pri};
    }

    printf "num_ret\t$topic\t%6d\n", $eval->{num_ret};
    printf "num_rel\t$topic\t%6d\n", $qrel{$topic}{num_rel};
    printf "rel_ret\t$topic\t%6d\n", $eval->{rel_ret};
    printf "num_pri\t$topic\t%6d\n", $qrel{$topic}{num_pri};
    printf "pri_ret\t$topic\t%6d\n", $eval->{pri_ret};

    printf "P10\t$topic\t%6.4f\n", $p10;
    printf "nDCG_R\t$topic\t%6.4f\n", $ndcg;
    printf "map\t$topic\t%6.4f\n", $map;
    printf "Rprec\t$topic\t%6.4f\n", $pprec;

    $cum{num_ret} += $eval->{num_ret};
    $cum{num_rel} += $qrel{$topic}{num_rel};
    $cum{rel_ret} += $eval->{rel_ret};
    $cum{num_pri} += $qrel{$topic}{num_pri};
    $cum{pri_ret} += $eval->{pri_ret};
    $cum{P10} += $p10;
    $cum{nDCG_R} += $ndcg;
    $cum{map} += $map;
    $cum{Rprec} += $pprec;
}

sub means() {
    my $num_q = scalar keys %qrel;
    printf "num_q\tall\t%6d\n", $num_q;
    printf "num_ret\tall\t%6d\n", $cum{num_ret};
    printf "num_rel\tall\t%6d\n", $cum{num_rel};
    printf "rel_ret\tall\t%6d\n", $cum{rel_ret};
    printf "num_pri\tall\t%6d\n", $cum{num_pri};
    printf "pri_ret\tall\t%6d\n", $cum{pri_ret};
    for my $meas (qw/P10 nDCG_R map Rprec/) {
	my $mean = $cum{$meas} / $num_q;
	printf "$meas\tall\t%6.4f\n", $mean;
    }
}