=head1 LICENSE

See the NOTICE file distributed with this work for additional information
regarding copyright ownership.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

     http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.

=head1 NAME

Bio::EnsEMBL::Compara::RunnableDB::ncRNAtrees::NCRecoverEPO

=head1 DESCRIPTION

This runnable recovers the genomic alignments from the GenomicAlignBlocks and
GenomicAlignTrees as stored by EPO_EXTENDED multiple whole genome alignment methods

=cut

=head1 SYNOPSIS

my $db           = Bio::EnsEMBL::Compara::DBAdaptor->new($locator);
my $nc_recover_epo = Bio::EnsEMBL::Compara::RunnableDB::ncRNAtrees::NCRecoverEPO->new
  (
   -db         => $db,
   -input_id   => $input_id,
   -analysis   => $analysis
  );
$nc_recover_epo->fetch_input(); #reads from DB
$nc_recover_epo->run();
$nc_recover_epo->write_output(); #writes to DB

=cut


package Bio::EnsEMBL::Compara::RunnableDB::ncRNAtrees::NCRecoverEPO;

use strict;
use warnings;

use Bio::EnsEMBL::Registry;

use base ('Bio::EnsEMBL::Compara::RunnableDB::BaseRunnable');

sub fetch_input {
  my $self = shift @_;

  my $mlss_id    = $self->param_required('mlss_id');
  my $nc_tree_id = $self->param_required('gene_tree_id');

  $self->param('nc_tree', $self->compara_dba->get_GeneTreeAdaptor->fetch_by_dbID($nc_tree_id));
  $self->param('nc_tree') || die "Could not fetch the tree for root_id=$nc_tree_id";

  if ($self->param('max_members') and scalar(@{$self->param('nc_tree')->get_all_Members}) > $self->param('max_members')) {
    $self->input_job->autoflow(0);
    $self->dataflow_output_id(undef, -1);
    $self->complete_early('Too many members, going to _himem');
  }

  $self->param('gene_member_adaptor', $self->compara_dba->get_GeneMemberAdaptor);
  $self->param('treenode_adaptor', $self->compara_dba->get_GeneTreeNodeAdaptor);

  my $epo_dba = $self->get_cached_compara_dba('epo_db');
  $self->param('epo_gab_adaptor', $epo_dba->get_GenomicAlignBlockAdaptor);
  $self->param('epo_mlss_adaptor', $epo_dba->get_MethodLinkSpeciesSetAdaptor);
}

sub run {
  my $self = shift @_;

  $self->iterate_over_lowcov_mlsss;
}

sub write_output {
  my $self = shift @_;

  $self->remove_low_cov_predictions;
}


##########################################
#
# internal methods
#
##########################################

# This is currently not called
sub run_ncrecoverepo {
  my $self = shift;

  my $root_id = $self->param('nc_tree')->root_id;

  my %present_gdbs     = ();
  my %absent_gdbs      = ();
  my %present_epo_gdbs = ();
  
  # NOTE: 'epo_gdb' used to be a hash { genome_db_id => 1 } containing all the low-coverage species
  # Find absent gdbs
  foreach my $leaf (@{$self->param('nc_tree')->get_all_leaves}) {
      $present_gdbs{$leaf->genome_db_id}++;
  }
  foreach my $present_gdb (keys %present_gdbs) {
    if (defined($self->param('epo_gdb')->{$present_gdb})) {
      $present_epo_gdbs{$present_gdb} = 1;
    }
  }
  foreach my $epo_gdb (keys %{$self->param('epo_gdb')}) {
    if (!defined($present_gdbs{$epo_gdb})) {
      $absent_gdbs{$epo_gdb} = 1;
    }
  }

  my %nc_tree_gene_ids = ();
  my $seq_length = 0;

  my $leaves = $self->param('nc_tree')->get_all_leaves;
  foreach my $leaf (@$leaves) {
    my $description = $leaf->description; $description =~ /Gene\:(\S+)/; my $gene_id = $1;
    $nc_tree_gene_ids{$gene_id} = 1;
    $seq_length += length($leaf->sequence);
  }
  
  my $avg_seq_length = $seq_length/(scalar @$leaves);

  my %pecan_restricted_gab_hash = ();

  if (keys %absent_gdbs) {
    my $pecan_mlss = $self->param('epo_mlss_adaptor')->fetch_all_by_method_link_type('PECAN')->[0];

    foreach my $leaf (@{$self->param('nc_tree')->get_all_leaves}) {
      my $gdb_name = $leaf->genome_db->name;
      print STDERR "# PECAN $gdb_name\n" if ($self->debug);
      next unless(defined($present_epo_gdbs{$leaf->genome_db_id})); # Only for the ones in genomic alignments
      my $slice = $leaf->genome_db->db_adaptor->get_SliceAdaptor->fetch_by_transcript_stable_id($leaf->stable_id);
      next unless (defined($slice));
      my $genomic_align_blocks = $self->param('epo_gab_adaptor')->fetch_all_by_MethodLinkSpeciesSet_Slice($pecan_mlss,$slice);
      next unless(0 < scalar(@$genomic_align_blocks));
      foreach my $genomic_align_block (@$genomic_align_blocks) {
        my $pecan_restricted_gab = $genomic_align_block->restrict_between_reference_positions($slice->start,$slice->end);
        next unless (defined($pecan_restricted_gab));
        my $gab_start = $pecan_restricted_gab->{restricted_aln_start};
        my $gab_end   = $genomic_align_block->length - $pecan_restricted_gab->{restricted_aln_end};
        my $boundary = 10;
        if (defined($pecan_restricted_gab_hash{$genomic_align_block->dbID})) {
          if (
              abs($pecan_restricted_gab_hash{$genomic_align_block->dbID}{start} - $gab_start) < $boundary &&
              abs($pecan_restricted_gab_hash{$genomic_align_block->dbID}{end}   - $gab_end) < $boundary &&
              abs($pecan_restricted_gab_hash{$genomic_align_block->dbID}{slice_length} - $slice->length) < $boundary
             ) {
            # same genomic alignment region, dont need to go through it again
            print STDERR "#   same genomic alignment region, dont need to go through it again\n" if ($self->debug);
            next;
          }
        }
        $pecan_restricted_gab_hash{$genomic_align_block->dbID}{start}          = $gab_start;
        $pecan_restricted_gab_hash{$genomic_align_block->dbID}{end}            = $gab_end;
        $pecan_restricted_gab_hash{$genomic_align_block->dbID}{slice_length}   = $slice->length;
        foreach my $genomic_align (@{$pecan_restricted_gab->get_all_GenomicAligns}) {
          my $ga_gdb = $genomic_align->genome_db;
          next if ($ga_gdb->dbID == $leaf->genome_db_id);
          my $core_adaptor = Bio::EnsEMBL::Registry->get_DBAdaptor($ga_gdb->name, "core");
          $self->throw("Cannot access core db") unless(defined($core_adaptor));
          $core_adaptor->dbc->disconnect_when_inactive(0);
          $genomic_align->dnafrag->genome_db->db_adaptor($core_adaptor);
          my $other_slice = $genomic_align->get_Slice;
          my $other_genome_db_id = $genomic_align->genome_db->dbID;
          next unless (defined ($other_slice));
          my $genes = $other_slice->get_all_Genes;
          my $found_prediction; my $validated_prediction = 0;
          print STDERR "#   Other genome: ", $genomic_align->genome_db->name, "\n" if ($self->debug);
          foreach my $gene (@$genes) {
            my $gene_member = $self->param('gene_member_adaptor')->fetch_by_Gene($gene);
            if (defined($nc_tree_gene_ids{$gene_member->stable_id})) {
              $found_prediction->{$gene_member->dbID} = 1;
              $validated_prediction = 1;
            } elsif ($gene->biotype !~ /coding/) {
              $found_prediction->{$gene_member->dbID} = 1;
              print STDERR "#     " . $gene_member->stable_id . ", biotype:", $gene->biotype, "\n" if ($self->debug);
            }
          }
          if (defined($found_prediction) && 0 == $validated_prediction) {
            foreach my $found_gene_member_id (keys %$found_prediction) {
              # Store it in the table
              my $gene_member = $self->param('gene_member_adaptor')->fetch_by_dbID($found_gene_member_id);
              my $sth = $self->compara_dba->dbc->prepare
                ("INSERT IGNORE INTO recovered_member 
                           (node_id,
                            stable_id,
                            genome_db_id) VALUES (?,?,?)");
              $sth->execute($root_id,
                            $gene_member->stable_id,
                            $other_genome_db_id);
              $sth->finish;
              # See if we can match the RFAM name or RFAM id
              next unless (defined($gene_member));
              # FIXME: this code cannot work because nctree_adaptor is not defined !
              my $other_tree = $self->param('nctree_adaptor')->fetch_by_Member_root_id($gene_member);
              if (defined($other_tree)) {
                my $other_tree_id = $other_tree->node_id;
                print STDERR  "#     found_description and gene_member, but already in tree $other_tree_id [$root_id]\n" if ($self->debug);
                next;
              }
              my $description = $gene_member->description;
              $description =~ /Acc:(\w+)/;
              my $acc_description = $1 || '';
              my $model_id = $self->param('nc_tree')->get_value_for_tag('model_id');
              if ($acc_description eq $model_id) {
                $self->param('predictions_to_add')->{$found_gene_member_id} = 1;
              } else {
                print STDERR "#     found_prediction but Acc not mapped: $acc_description [$model_id]\n" if ($self->debug);
              }
            }
          }
          # We don't have a gene prediction here, so we try to predict one
          if (!defined($found_prediction)) {
            my $start   = $other_slice->start;
            my $end     = $other_slice->end;
            my $seqname = $other_slice->seq_region_name;
            my $sequence = $other_slice->seq; $sequence =~ s/N//g;
            my $length = length($sequence);
            next if (0 == $length);
            next if (($avg_seq_length/$length) > 1.2 ||
                     ($avg_seq_length/$length) < 0.8);
            my $found_gene_stable_id = "$seqname:$start-$end";
            my $sth = $self->compara_dba->dbc->prepare
              ("INSERT IGNORE INTO recovered_member 
                           (node_id,
                            stable_id,
                            genome_db_id) VALUES (?,?,?)");
            $sth->execute($root_id,
                          $found_gene_stable_id,
                          $other_genome_db_id);
            $sth->finish;
            print STDERR "#     no_prediction\n" if ($self->debug);
            # Use the RFAM model to see if the sequence is good
          }
        }
      }
    }
  }

  return 1;
}

# This one is called
sub iterate_over_lowcov_mlsss {
    my $self = shift @_;
    my $epolow_mlsss = $self->param('epo_mlss_adaptor')->fetch_all_by_method_link_type('EPO_EXTENDED');
    unless (scalar(@$epolow_mlsss)) {
        die "Could not find an 'EPO_EXTENDED' MLSS in ".$self->param('epo_db')."\n";
    }
    my @gab_ids;
    $self->param('low_cov_leaves_to_delete', []);
    foreach my $epo_low_mlss (@$epolow_mlsss) {
        my $epo_hc_mlss = $epo_low_mlss->get_linked_mlss_by_tag('base_mlss_id')
            || die "Could not find the matching 'EPO' MLSS in ".$self->param('epo_db')."\n";
        my %hc_gdb_id = (map {$_->dbID => 1} @{$epo_hc_mlss->species_set->genome_dbs});
        my @lowcov_gdbs = grep {not exists $hc_gdb_id{$_->dbID}} @{$epo_low_mlss->species_set->genome_dbs};
        my %low_gdb_id = (map {$_->dbID => 1} @lowcov_gdbs);
        my $gab_id = $self->run_low_coverage_best_in_alignment($epo_low_mlss, \%hc_gdb_id, \%low_gdb_id);
        push @gab_ids, $gab_id if $gab_id;
    }
    $self->param('nc_tree')->store_tag('ncrecoverepo_filter_gab_id', \@gab_ids) if scalar(@gab_ids);
}

# This one too
sub run_low_coverage_best_in_alignment {
  my $self = shift;
  my $epo_low_mlss = shift;
  my $hc_gdb_id = shift;
  my $lc_gdb_id = shift;

  my %gdb_per_dbID = map {$_->dbID => $_} @{ $self->compara_dba->get_GenomeDBAdaptor->fetch_all() };
  # now we can disconnect
  $self->compara_dba->dbc->disconnect_if_idle();

  my %members_per_genome_db_id;
  foreach my $member (@{ $self->param('nc_tree')->get_all_Members }) {
      push @{ $members_per_genome_db_id{$member->genome_db_id} }, $member;
  }

  my %all_seen_gab_ids = ();

  # First round to get the candidate GenomicAlignTrees
  # We first iterate over the high-coverage genome_dbs
  # This way, we group the queries to the same core database
  foreach my $gdb_id (keys %{$hc_gdb_id}) {

   my $genome_db = $gdb_per_dbID{$gdb_id};
   if (! defined $genome_db){
       $self->warning("genome_db_id: $gdb_id is not found in the current DB. Probably an old id was being used.");
       next;
   }

   my $gdb_name = $genome_db->name;
   print STDERR "doing $gdb_name\n" if $self->debug;

   my $leaves = $members_per_genome_db_id{$gdb_id};
   my $core_db_adaptor = $genome_db->db_adaptor;
   $core_db_adaptor->dbc->prevent_disconnect( sub {
   foreach my $leaf (@$leaves) {
    my $slice = $core_db_adaptor->get_SliceAdaptor->fetch_by_transcript_stable_id($leaf->stable_id);
    my $genomic_align_blocks = $self->param('epo_gab_adaptor')->fetch_all_by_MethodLinkSpeciesSet_Slice($epo_low_mlss, $slice);
    print STDERR scalar(@$genomic_align_blocks), " blocks for ", $leaf->stable_id, "\n" if $self->debug;
    foreach my $genomic_align_block (@$genomic_align_blocks) {
        if (not defined $genomic_align_block->dbID) {
            # Happens when the block has been restricted
            next;
        }
      $all_seen_gab_ids{$genomic_align_block->dbID}++;
    }
   }
   } );
  }

  # This selects the GAB with the highest number of species
  my $max = 0; my $max_gabID;
  foreach my $gabID (keys %all_seen_gab_ids) {
    my $count = $all_seen_gab_ids{$gabID};
    if ($count > $max) {$max = $count; $max_gabID = $gabID};
  }
  return undef unless $max;
  print STDERR "BEST_GAB: $max_gabID ($max species)\n" if $self->debug;

  # Second round to get the low-covs on the max_gabID
  # We apply the same trick as above
  foreach my $gdb_id (keys %{$lc_gdb_id}) {

   my $genome_db = $gdb_per_dbID{$gdb_id};
   if (! defined $genome_db){
       $self->warning("genome_db_id: $gdb_id is not found in the current DB. Probably an old id was being used.");
       next;
   }
   my $gdb_name = $genome_db->name;
   print STDERR "working on $gdb_name\n" if $self->debug;
   my $leaves = $members_per_genome_db_id{$gdb_id};
   my $core_db_adaptor = $genome_db->db_adaptor;
   $core_db_adaptor->dbc->prevent_disconnect( sub {
   foreach my $leaf (@$leaves) {
    my $slice = $core_db_adaptor->get_SliceAdaptor->fetch_by_transcript_stable_id($leaf->stable_id);
    my $low_cov_genomic_align_blocks = $self->param('epo_gab_adaptor')->fetch_all_by_MethodLinkSpeciesSet_Slice($epo_low_mlss,$slice);
    unless (0 < scalar(@$low_cov_genomic_align_blocks)) {
      push @{ $self->param('low_cov_leaves_to_delete') }, $leaf;
      print STDERR $leaf->stable_id, " has no alignments -> will be removed\n" if $self->debug;
      next;
    }
    print STDERR scalar(@$low_cov_genomic_align_blocks), " blocks for ", $leaf->stable_id, "\n" if $self->debug;
    my $deleted = 0;
    foreach my $low_cov_genomic_align_block (@$low_cov_genomic_align_blocks) {
        die unless $max_gabID;
      if ($low_cov_genomic_align_block->original_dbID != $max_gabID) {
        # We delete this leaf because it's a low_cov slice that is not in the epo_low_cov, so it's the best in alignment
        push @{ $self->param('low_cov_leaves_to_delete') }, $leaf;
        print STDERR $leaf->stable_id, " is not in GAB $max_gabID -> will be removed\n" if $self->debug;
        $deleted = 1;
        last;
      }
    }
    unless ($deleted) {
      print STDERR $leaf->stable_id, " is in GAB $max_gabID -> will be kept \n" if $self->debug;
    }
   }
   } );
  }
  # We don't need the connection to the EPO database any more;
  $self->param('epo_gab_adaptor')->dbc->disconnect_if_idle();

  return $max_gabID;
  $self->param('nc_tree')->store_tag('ncrecoverepo_best_gab_id', $max_gabID);
}

sub remove_low_cov_predictions {
  my $self = shift;
  my $nc_tree = $self->param('nc_tree');
  my $root_id = $nc_tree->root_id;

  # Remove low cov members that are not best in alignment
  foreach my $leaf_to_delete (@{ $self->param('low_cov_leaves_to_delete') }) {
      print STDERR "removing low_cov prediction ", $leaf_to_delete->stable_id, "\n" if($self->debug);
      $self->call_within_transaction(sub {
        $self->param('treenode_adaptor')->remove_seq_member($leaf_to_delete);
      });
  }
  #calc residue count total
  my $leafcount = scalar(@{$nc_tree->get_all_leaves});

  ## Remove the tree if now it is too small
  if ($leafcount < 2) {
      my $gene_tree_adaptor = $self->compara_dba->get_GeneTreeAdaptor;
      $gene_tree_adaptor->delete_tree($nc_tree);

      ## TODO: Dying here prevents storing the tag for a disappeared tree.
      ## The problem is that it also prevents adding new members by add_matching_predictions below.
      ## For now, this is not a problem, since add_matching_predictions is not finished, but we may need to
      ## make sure that it is running properly once it is finished.
      $self->input_job->autoflow(0);
      $self->complete_early("$root_id tree has become too short ($leafcount leaf/ves)\n");
  }
  $nc_tree->store_tag('gene_count', $leafcount);

  return 1;
}

# This is currently not called
sub add_matching_predictions {
  my $self = shift;

  # Insert the members that are found new and have matching Acc
  foreach my $gene_member_id_to_add (keys %{$self->param('predictions_to_add')}) {
    my $gene_member = $self->param('gene_member_adaptor')->fetch_by_dbID($gene_member_id_to_add);
    # Incorporate the canonical seq_member into the cluster
    my $node = new Bio::EnsEMBL::Compara::GeneTreeMember;
    $node->seq_member_id($gene_member->get_canonical_SeqMember->seq_member_id);
    $self->param('nc_tree')->root->add_child($node);

    #the building method uses seq_member_id's to reference unique nodes
    #which are stored in the node_id value, copy to seq_member_id
    # We won't do the store until the end, otherwise it will affect the main loop
    print STDERR "adding matching prediction $gene_member_id_to_add\n" if($self->debug);
    $self->param('treenode_adaptor')->store_node($node);
  }

  #calc residue count total
  my $leafcount = scalar(@{$self->param('nc_tree')->get_all_leaves});
  $self->param('nc_tree')->store_tag('gene_count', $leafcount);

  return 1;
}

1;
