#! /usr/local/bin/perl -w

# makeCorpus.pl - script to recreate corpus of ambiguous biomedical
# abbreviations

# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# General Public License for more details.

# You should have received a copy of the GNU General Public License
# along with this program.  If not, see
# <http://www.gnu.org/licenses/>.

use LWP::Simple;
use XML::Simple;


# Some global variables

my $utils = "http://eutils.ncbi.nlm.nih.gov/entrez/eutils";
my $db = "Pubmed";

my $batch_size = 200;

%pubmed_docs = ();


# Check that there is exactly one argument

unless($#ARGV == 1) {
  &usage();
} # unless


# Read through the data file
# Store each line of the data in the %corpus_info hash, using an integer
# as the unique identifier for each example
print STDERR "Reading data file ....";


if($ARGV[1] =~ /(.+)\_data/ && ($ARGV[0] eq "genecorpus" || $ARGV[0] eq "abbrcorpus")) {
  $corpus_label = $1;
} else {
    
      &usage();
} # if/else

my $genecorpus = $ARGV[0] eq "genecorpus" ? 1:0;

%corpus_info = ();
%abbrev_counts = ();

open(IN, "$ARGV[1]") or die "Can't open the input file($ARGV[1]): $!\n";
while(<IN>) {

  if(/^([0-9]+)\.([^\.]+)\.([0-9]+)\.([0-9]+)\.(M?[0-9]+)/) {
    $pmid = $1;
    $abbrev = $2;
    $start_offset = $3;
    $end_offset = $4;
    $sense = $5;
   
    unless(exists $abbrev_counts{$abbrev}) {
      $abbrev_counts{$abbrev} = 1;
    } # unless
    $count = $abbrev_counts{$abbrev}++;

    $corpus_info{$abbrev}{$count}{pmid} = $pmid;
    $corpus_info{$abbrev}{$count}{start} = $start_offset;
    $corpus_info{$abbrev}{$count}{end} = $end_offset;
    $corpus_info{$abbrev}{$count}{sense} = $sense;

  } # if

} # while
close(IN);
print STDERR " done\n";


# Create a directory for the corpus to be written to. Bail out if this directory
# alrady exists.
if(-e $corpus_label){
    rmdir($corpus_label);
}
mkdir($corpus_label) or die "\nNothing to do, directory already exists: $corpus_label/\nDelete this directory and its contents to recreate the corpus\n";


%pubmed_docs = ();
foreach $abbrev (sort {$a cmp $b} keys %corpus_info) {

  print STDERR "Retrieving examples for $abbrev\n";

  # Collect together batched of PMID numbers for abstracts that are contained 
  # in the corpus, retrieve andstore them. 
  %pmid_list = ();
  foreach $count (keys %{$corpus_info{$abbrev}}) {
    $pmid_list{$corpus_info{$abbrev}{$count}{pmid}}++;
  } # foreach 

  $pmids_to_retrieve = "";
  $counter = 0;
  foreach $pmid (keys %pmid_list) {
    $pmids_to_retrieve .= $pmid;
    $pmids_to_retrieve .= " ";
    $counter++;
    if($counter == $batch_size) {
      $pmids_to_retrieve .= " [uid]";
      &get_docs($pmids_to_retrieve);
      $counter = 0;
      $pmids_to_retrieve = "";
    } # if
  } # foreach 
  $pmids_to_retrieve .= " [uid]";  
  &get_docs($pmids_to_retrieve);
  # print "\n";

  # Go through each example in the corpus and print out an adapted verison 
  # of the abstract
  $out_file = $corpus_label."/".$abbrev."_set";
  open(OUT, ">$out_file") or die "Can't create $out_file\n";

  $example_id = 1;

  foreach $count (sort {$a <=> $b} keys %{$corpus_info{$abbrev}}) {
    # Retrieve abstract and adapt it 
    # then print
    $pmid = $corpus_info{$abbrev}{$count}{pmid};
    $start_offset = $corpus_info{$abbrev}{$count}{start};
    $length = $corpus_info{$abbrev}{$count}{end} - $start_offset;
    $sense = $corpus_info{$abbrev}{$count}{sense};

    $formatted_abstract = &printAbstract($pmid, $example_id, $abbrev, $start_offset, $length, $sense);
    unless($formatted_abstract eq "0") {
      print OUT $formatted_abstract;
      $example_id++;
    } # unless

  } # foreach 
  
  close(OUT); 

} # foreach



###############################################################################
#                             SUBROUTINES
###############################################################################


sub get_docs { 

  my ($pmids) = @_;

  # Check how many hits there are
  my $esearch = "$utils/esearch.fcgi?" .
    "tool=abbrev_corpus_generator&db=$db&retmax=1&usehistory=y&term=";

  my $esearch_result = get($esearch . $pmids);

  $esearch_result =~ 
    m|<Count>(\d+)</Count>.*<QueryKey>(\d+)</QueryKey>.*<WebEnv>(\S+)</WebEnv>|s;
  
  my $Count    = $1;
  my $QueryKey = $2;
  my $WebEnv   = $3;
  
  # Retrieve the abstracts from PubMed and store them
  my $efetch = 
    "$utils/efetch.fcgi?" .
      "tool=abbrev_corpus_generator&" .
	"rettype=xml&retmode=text&retstart=0&" .
	  "retmax=$batch_size&" .
	    "db=$db&query_key=$QueryKey&WebEnv=$WebEnv";
  
  my $efetch_result = get($efetch);

  &parse_results($efetch_result);

} # sub get_docs


# sub parse_results
# Reads through XML output from PubMed and stores in a hash

sub parse_results {

  my ($search_result) = @_;

  my $data = XMLin($search_result);
  
  # Go through each article and extract relevant information
  # How this is done depends on how many articles there are in the XML
  # file (since different data structures are returned for each)
   
  # If hash table is returned means that just one article
  if(ref($data->{PubmedArticle}) eq "HASH") {

    $pmid = $$data{PubmedArticle}{MedlineCitation}{PMID};
    $title = $$data{PubmedArticle}{MedlineCitation}{Article}{ArticleTitle};
    if(exists $$data{PubmedArticle}{MedlineCitation}{Article}{Abstract}{AbstractText}) {
      $abstract = 
	$$data{PubmedArticle}{MedlineCitation}{Article}{Abstract}{AbstractText};
    } else {
      $abstract = "";
    } # if/else
    
    @mesh_terms = ();
    if(exists $$data{PubmedArticle}{MedlineCitation}{MeshHeadingList}{MeshHeading}) {
      $type = ref($$data{PubmedArticle}{MedlineCitation}{MeshHeadingList}{MeshHeading});
      # If data type is ARRAY, means there is an array of hashes containing
      # MeSH terms (process each one in turn)
      if($type eq "ARRAY") {
	foreach $mesh_term_hash (@{$$data{PubmedArticle}{MedlineCitation}{MeshHeadingList}{MeshHeading}}) {
	  if(exists $$mesh_term_hash{DescriptorName}{content}) {
	    push(@mesh_terms, $$mesh_term_hash{DescriptorName}{content});
	  } # if
	} # foreach 
      } else {
	# Otherwise the data type is hash and there is just one MeSH term to
	# process
	if(exists $$data{PubmedArticle}{MedlineCitation}{MeshHeadingList}{MeshHeading}{DescriptorName}{content}) {
	  push(@mesh_terms, $$data{PubmedArticle}{MedlineCitation}{MeshHeadingList}{MeshHeading}{DescriptorName}{content});
	} # if
      } # if/else 
    } # if 


    $pubmed_docs{$pmid}{title} = $title;
    $pubmed_docs{$pmid}{abstract} = $abstract;
    foreach $mesh_term (@mesh_terms) {
      $pubmed_docs{$pmid}{mesh}{$mesh_term} = "true";
    } # foreach 
    # $corpus{$abbrev}{$abbrev_count} = &formatAbstract($pmid, $title, $abstract, $abbrev, $abbrev_count, @mesh_terms); 
    
    # If the data structure returned isn't a hash - means there is 
    # an array of hashes. Process each in turn
  } else {

    foreach $article (@{$data->{PubmedArticle}}) {
      
      # Gather the information about this abstract
      $pmid = $$article{MedlineCitation}{PMID};
      $title = $$article{MedlineCitation}{Article}{ArticleTitle};
      if(exists $$article{MedlineCitation}{Article}{Abstract}{AbstractText}) {
	$abstract = 
	  $$article{MedlineCitation}{Article}{Abstract}{AbstractText};
      } else {
	$abstract = "";
      } # if/else
      
      @mesh_terms = ();
      # Need to check how many MeSH terms there are
      if(exists $$article{MedlineCitation}{MeshHeadingList}{MeshHeading}) {
	$type = ref($$article{MedlineCitation}{MeshHeadingList}{MeshHeading});
	# If data type is ARRAY, means there is an array of hashes containing
	# MeSH terms (process each one in turn)
	if($type eq "ARRAY") {
	  
	  foreach $mesh_term_hash (@{$$article{MedlineCitation}{MeshHeadingList}{MeshHeading}}) {
	    if(exists $$mesh_term_hash{DescriptorName}{content}) {
	      push(@mesh_terms, $$mesh_term_hash{DescriptorName}{content});
	    } # if
	  } # foreach 
	  
	} else {
	  # Otherwise the data type is hash and there is just one MeSH term to
	  # process
	  if(exists $$article{MedlineCitation}{MeshHeadingList}{MeshHeading}{DescriptorName}{content}) {
	    push(@mesh_terms, $$article{MedlineCitation}{MeshHeadingList}{MeshHeading}{DescriptorName}{content});
	  } # if
	} # if/else
      } # if
      
      $pubmed_docs{$pmid}{title} = $title;
      $pubmed_docs{$pmid}{abstract} = $abstract;
      foreach $mesh_term (@mesh_terms) {
	$pubmed_docs{$pmid}{mesh}{$mesh_term} = "true";
      } # foreach 

     } # foreach
  } # if/else

} # sub parse_results


# sub printAbstract
# Prints out a single abstract from PubMed in a format similar to the one 
# used by the NLM-WSD corpus. The abstract's text is adapted so that the 
# text containing the abbreviation and expansion (defined by the $start_offset
# and $length arguments) are replaced by the abbreviation ($abbrev argument). 

sub printAbstract {

  my ($pmid, $count, $abbrev, $start_offset, $length, $sense) = @_;

  $title = $pubmed_docs{$pmid}{title};
  $abstract = $pubmed_docs{$pmid}{abstract};
  @mesh_terms = keys %{$pubmed_docs{$pmid}{mesh}};

  # Check to deal with cases where PMID is no longer available from PubMed
  unless(defined $title) {
    return 0;
  } # unless

  # Deal with cases where abstracts don't have sentence breaks at the end 
  # (by adding one)
  if($title =~ /^\[(.+)\]$/) {
    $title = $1;
  } # if
  unless($title =~ /[\.\!\?]$/) {
    $title = $title.".";
  } # unless
  
  $context = "PMID - $pmid\nTI - $title\nAB - $abstract\n";
  # Substitute text containing abbreviation and expansion with just abbreviation
  # Note that the $start_offset value has to be adjusted since this number
  # refers to the XML-formatted version of the abstract
  
  if(length $abstract == 0) {
    if($genecorpus == 0){$start_offset = $start_offset - 26;}
    if($genecorpus == 0){
	substr($title, $start_offset, $length) = $abbrev;
        $context = "PMID - $pmid\nTI - $title\n";
    }
    else{
    	substr($context, $start_offset, $length) = $abbrev;
    }
    $start_context_length = length "PMID - $pmid\nTI - ";
    
    if($genecorpus == 0){$start_offset += $start_context_length;}

  } else {
    if($genecorpus == 0){$start_offset = $start_offset - 30;}
    if($genecorpus == 1){
   	substr($context, $start_offset, $length) =  $abbrev;
    }
    else{
	substr($abstract, $start_offset, $length) =  $abbrev;
        $context = "PMID - $pmid\nTI - $title\nAB - $abstract\n";
    }

    $start_context_length = length "PMID - $pmid\nTI - $title\nAB - ";
    
    if($genecorpus == 0){$start_offset += $start_context_length;}
  } # if/else
 
  $length = length $abbrev;
  # Form the whole abstract to be printed. Find the instance of the ambiguous
  # abbreviation with the abstract and then search around it to identify the 
  # sentence it occurs in. 

  #
  $end_offset = $genecorpus == 1? $start_offset +  length $abbrev : $start_offset + 3 ;

  $context_length = length $context;

  $pre_target = substr($context, 0, $start_offset);
  $post_target = substr($context, $end_offset, $context_length);


  # Clear out the data structures
  $local_pre_target = "";
  $local_post_target = "";
          
  $wide_start_offset = length $pre_target;
  $wide_end_offset = $wide_start_offset + length $abbrev;
  $wide_end_offset--;
    
  # Look for sentence context around target
  # Start with context before target
  
  # Easiest way to do this is to scan from the start of the context 
  # until the start of the current sentence is found
  $local_pre_target = "";
  while($pre_target =~ /\n?(.*?([\.\?\!]|$))/gs) {
    $local_pre_target .= $1;
    
    # Don't let the sentence end as part of a contraction (e.g.)
    # or decimal expression
    if($pre_target =~ /(mg\.\Gkg|St\.\G|e\.\Gg\.|i\.\Ge\.|e\.g\.\G|i\.e\.\G|[0-9]\.\G[0-9])/) {
      # not the end of the sentence -- keep collecting context
    } elsif($pre_target =~ /[\.\?\!]\G/) {
      # real end of sentence. Clear sentence and keep going. 
      $local_pre_target = "";
    } # if/else
  } # while
  
  # Clean up by trimming any leading whitespace
  $local_pre_target =~ s/^[ \n]+//;
  # remove the lines containing "\n", the ambiguious word may contain in the title line,
  # and the sentence detection heuristic will also include the PMID line, so need to get
  # rid of the PMID line, as it is irrelevant.
  if($local_pre_target =~ /\n/){
    $local_pre_target =~ s/(.*?\n)//;
  }
  # Then look for post target
  # Match anything (non-greedily) but don't let last punctuation
  # be part of "i.e." or "e.g."
  $local_post_target = "";
  while($post_target =~ /(.*?[\.\?\!])/gs) {
    $local_post_target .= $1;
    # Don't end the sentence as part of a contraction (e.g.) or 
    # decimal expression
    if($post_target =~ /(mg\.\Gkg|St\.\G|i\.\Ge\.|e\.\Gg\.|i\.e\.\G|e\.g\.\G|[0-9]\.\G[0-9])/) {
      # continue
    } else {
      last;
    } # if/else
  } #while 
  
  $local_start_offset = length $local_pre_target;
  $local_end_offset = $local_start_offset + length $abbrev;
  $local_end_offset--;
  
  $local_context = $local_pre_target.$abbrev.$local_post_target;
  
  # Print out context in reformatted version
  $reformated_abstract = "$count|$pmid.xx.x|$sense\n";
  $reformated_abstract .= "$local_context\n";
  $reformated_abstract .= "$abbrev|$abbrev|xx|xx|$local_start_offset|$local_end_offset|xxxxx|\n";
  $reformated_abstract .= $context;
  $reformated_abstract .= "$abbrev|$abbrev|xx|xx|$wide_start_offset|$wide_end_offset|xxxxx|\n";
  
  # Print out MeSH terms (if there are any) 
  foreach $mesh_term (@mesh_terms) {
    $reformated_abstract .= "MeSH: $mesh_term\n";
  } # foreach  
  $reformated_abstract .= "\n\n";
  
  return $reformated_abstract;
  
} # sub printAbstract


# sub usage
# Stop program and print some usage notes

sub usage {

  die("Usage: $0 [genecorpus|abbrcorpus] *_data\nExample: $0 genecorpus Corpus.100_data\n");

} # usage
