#!/usr/bin/env perl

#
# Evaluate a TREC 2009 entity track run.  The measures computed are
# NDCG and P@10.
#
# Version: 1.1.
#
# Usage: eval-entity.pl [-v] [--wiki] [--name] <qrels> <runfile>\n";
#        -v: verbose mode (default: no)
#        --wiki: give credit for primary WP entries (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 for HPs and WPs,
# and primary NAMEs.
#
# 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.
#
# The NDCG that is computed is NDCG at R, where R is the number of primary
# and relevant HP (and WP, in WP mode) pages.
#
# HISTORY
#
# 1.1  changed NAMEs to case-insensitive match
# 1.0  original release

use strict;
use Getopt::Long;

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

# allow qrels and runfiles to be compressed with gzip
@ARGV = map { /.gz$/ ? "gzip -dc $_ |" : $_ } @ARGV;

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 ($lead, $thing, $rel, $cat) = split;
    my ($topic, $type) = split /-/, $lead;
    $type =~ s/HP[123]/HP/;

    next if $type eq "WP" and $wpcredit == 0;
    next if $type eq "NAME" and $namecredit == 0;

    if ($type eq "NAME") {
	$thing = lc $thing;
    }

    $qrel{$topic}{$thing} = $rel;
    $class{$topic}{$thing} = $cat;
    $cref{$topic}{$cat}{$type} = 1;

    if ($type eq "HP" or $type eq "WP") {
	if ($rel == 1) {
	    $qrel{$topic}{num_rel}++;
	}
	if ($rel == 2) {
	    $qrel{$topic}{num_pri}++;
	}
    }
}
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 == 2 and 
		exists $cref{$topic}{$class{$topic}{$thing}}{NAME}) {
		$rel = 3;
	    }
	    push @{ $qrel{$topic}{list} }, $rel;
	}
    }
    $qrel{$topic}{list} = [ sort {$b<=>$a} @{ $qrel{$topic}{list} } ];
}

my $last_topic = -1;
my $rec_eval = -1;
my %rec_class = ();
open RUN, $run_file or die "Can't open `$run_file': $!\n";
while (<RUN>) {
    s/\r//g;
    next if /^$/;
    next if /SUPPORT/;

    chomp;
    my ($lead, undef, $thing, $rank, $sim, $tag) = split;
    my ($topic, $type) = split /-/, $lead;
    $topic =~ s/^0+//;

    if ($type eq "HP1" and $rec_eval != -1) {
	print "# -------------------- Record gain: $rec_eval\n" if $verbose;
	push @{ $eval{list} }, $rec_eval;
	$eval{num_ret}++;
	$rec_eval = 0;
	%rec_class = ();
    }
    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;
    }

    $type =~ s/HP[123]/HP/;
    my $rel = 0;
    if ($type eq "NAME") {
	$thing = lc $thing;
    }
    if (exists $qrel{$topic}{$thing}) {
	$rel = $qrel{$topic}{$thing};
	if ($type eq "HP" or $type eq "WP") {
	    if ($rel == 1) {
		$eval{rel_ret}++;
	    }
	    if ($rel == 2) {
		$eval{pri_ret}++;
	    }
	}
    }

    if ($type eq "HP") {
	if ($rel == 2) {
	    if ($thing !~ /clueweb09-enwp/) {
		$rec_class{$class{$topic}{$thing}} = 1;
	    }
	}

	if ($rec_eval < $rel) {
	    $rec_eval = $rel;
	}	 
    } elsif ($type eq "WP") {
	if ($rel == 2) {
	    $rec_class{$class{$topic}{$thing}} = 1;
	}
	if ($rec_eval < $rel) {
	    $rec_eval = $rel;
	}
    } elsif ($type eq "NAME") {
	if ($rel == 2 and exists $rec_class{$class{$topic}{$thing}}) {
	    $rec_eval = 3;
	}
    }

    if ($verbose) {
	if ($type eq "WP" and $wpcredit == 0) { $rel = "-"; }
	if ($type eq "NAME" and $namecredit == 0) { $rel = "-"; }
	printf "#  %-5s $lead Q0 $thing $rank $sim $tag\n", $rel;
    }
}

if ($rec_eval != -1) {
    print "# -------------------- Record gain: $rec_eval\n" if $verbose;
    $eval{num_ret}++;
    push @{ $eval{list} }, $rec_eval;
}

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 $topic = $eval->{topic};

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

	if ($e > 1 && $rank <= 10) {
	    $p10 += 0.1;
	}

	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)));
    }
    
    # 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;
    
    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;

    $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;
}

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/) {
	my $mean = $cum{$meas} / $num_q;
	printf "$meas\tall\t%6.4f\n", $mean;
    }
}

