=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.

=cut


=head1 CONTACT

  Please email comments or questions to the public Ensembl
  developers list at <http://lists.ensembl.org/mailman/listinfo/dev>.

  Questions may also be sent to the Ensembl help desk at
  <http://www.ensembl.org/Help/Contact>.

=head1 NAME

Bio::EnsEMBL::Compara::RunnableDB::GeneTrees::OrthoTree

=head1 DESCRIPTION

This Analysis/RunnableDB is designed to take GeneTree as input

This must already have a rooted tree with duplication/sepeciation tags
on the nodes.

It analyzes that tree structure to pick Orthologues and Paralogs for
each genepair.

input_id/parameters format eg: "{'tree_id'=>1234}"
    tree_id : use 'id' to fetch a cluster from the GeneTree

=head1 SYNOPSIS

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

=head1 AUTHORSHIP

Ensembl Team. Individual contributions can be found in the GIT log.

=head1 APPENDIX

The rest of the documentation details each of the object methods.
Internal methods are usually preceded with an underscore (_)

=cut

package Bio::EnsEMBL::Compara::RunnableDB::GeneTrees::OrthoTree;

use strict;
use warnings;

use feature qw(switch);

use IO::File;
use File::Basename;
use List::Util qw(max);
use Scalar::Util qw(looks_like_number);
use File::Path qw(make_path);

use Bio::EnsEMBL::Compara::Homology;
use Bio::EnsEMBL::Compara::MethodLinkSpeciesSet;
use Bio::EnsEMBL::Compara::Graph::Link;
use Bio::EnsEMBL::Compara::Graph::Node;
use Bio::EnsEMBL::Compara::Graph::NewickParser;
use Bio::EnsEMBL::Compara::Utils::Preloader;
use Bio::EnsEMBL::Compara::Utils::FlatFile qw(map_row_to_header check_column_integrity check_line_counts);

use Bio::EnsEMBL::Hive::Utils 'stringify';  # import 'stringify()'

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


sub param_defaults {
    my $self = shift;
    return {
            %{ $self->SUPER::param_defaults() },
            'tree_scale'            => 1,
            'store_homologies'      => 1,
            'no_between'            => 0.25, # dont store all possible_orthologs
            '_readonly'             => 0,
            'tag_split_genes'       => 0,
            'input_clusterset_id'   => undef,
    };
}


=head2 fetch_input

    Title   :   fetch_input
    Usage   :   $self->fetch_input
    Function:   Fetches input data from the database
    Returns :   none
    Args    :   none

=cut


sub fetch_input {
    my $self = shift @_;

    unless ( $self->param('output_flatfile') ) {
        $self->param('homologyDBA', $self->compara_dba->get_HomologyAdaptor);
    }

    $self->param('seen_seq_member_ids', {});
    $self->param('seen_gene_member_ids', {});

    my $tree_id = $self->param_required('gene_tree_id');
    my $gene_tree = $self->compara_dba->get_GeneTreeAdaptor->fetch_by_root_id($tree_id) or $self->die_no_retry("Could not fetch gene_tree with tree_id='$tree_id'");
    if ($self->param('input_clusterset_id') and $self->param('input_clusterset_id') ne 'default') {
        $gene_tree = $gene_tree->alternative_trees->{$self->param('input_clusterset_id')};
        die sprintf('Cannot find a "%s" tree for tree_id=%d', $self->param('input_clusterset_id'), $self->param('gene_tree_id')) unless $gene_tree;
    }

    Bio::EnsEMBL::Compara::Utils::Preloader::load_all_sequences($self->compara_dba->get_SequenceAdaptor, undef, $gene_tree);
    $self->param('gene_tree', $gene_tree->root);

    if($self->debug) {
        $gene_tree->print_tree($self->param('tree_scale'));
        printf("%d genes in tree\n", scalar(@{$gene_tree->root->get_all_leaves}));
    }

    $self->param('homology_consistency', {});
    $self->param('has_match', {});
    $self->param('orthotree_homology_counts', {});
    $self->param('n_stored_homologies', 0);
}


=head2 run

    Title   :   run
    Usage   :   $self->run
    Function:   runs OrthoTree
    Returns :   none
    Args    :   none

=cut

sub run {
    my $self = shift @_;

    $self->prepare_analysis();
}


=head2 write_output

    Title   :   write_output
    Usage   :   $self->write_output

    Function: parse clustalw output and update homology and
              homology_member tables
    Returns : none 
    Args    : none 

=cut

sub write_output {
    my $self = shift @_;

    $self->delete_old_homologies unless $self->param('_readonly');
    $self->_create_flatfile if $self->param('output_flatfile');
    $self->run_analysis;
    $self->print_summary;
    $self->compara_dba->dbc->disconnect_if_idle;
    $self->_hc_flatfile if $self->param('output_flatfile');
}


sub post_cleanup {
  my $self = shift;

  if($self->param('gene_tree')) {
    printf("OrthoTree::post_cleanup  releasing gene_tree\n") if($self->debug);
    $self->param('gene_tree')->release_tree;
    $self->param('gene_tree', undef);
  }

  $self->SUPER::post_cleanup if $self->can("SUPER::post_cleanup");
}


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

sub prepare_analysis {
    my $self = shift;

    my $gene_tree = $self->param('gene_tree');

    # precalculate the ancestor species_hash
    printf("Calculating ancestor species hash\n") if ($self->debug);
    $self->get_ancestor_species_hash($gene_tree);
}


sub run_analysis {
  my $self = shift;

  my $gene_tree = $self->param('gene_tree');

  #compare every gene in the tree with every other each gene/gene
  #pairing is a potential ortholog/paralog and thus we need to analyze
  #every possibility
  printf("%d genes in tree\n", scalar(@{$gene_tree->get_all_leaves})) if $self->debug;

  foreach my $ancestor (reverse @{$gene_tree->get_all_nodes}) {
    next unless scalar(@{$ancestor->children});
    my ($child1, $child2) = @{$ancestor->children};
    my $leaves1 = $child1->get_all_leaves;
    my $leaves2 = $child2->get_all_leaves;
    my @pair_group;
    foreach my $gene1 (@$leaves1) {
     foreach my $gene2 (@$leaves2) {
      my $genepairlink = new Bio::EnsEMBL::Compara::Graph::Link($gene1, $gene2);
      $genepairlink->add_tag("ancestor", $ancestor);
      $genepairlink->add_tag("subtree1", $child1);
      $genepairlink->add_tag("subtree2", $child2);

      my $node_type = $ancestor->get_value_for_tag('node_type');
      if ($ancestor->is_speciation) {
          $self->tag_genepairlink($genepairlink, $self->tag_orthologues($genepairlink), 1);
      } elsif ($node_type eq 'dubious') {
          $self->tag_genepairlink($genepairlink, $self->tag_orthologues($genepairlink), 0);
      } elsif ($node_type eq 'gene_split') {
          $self->tag_genepairlink($genepairlink, 'gene_split', 1);
      } elsif ($node_type eq 'duplication') {
          push @pair_group, $genepairlink;
      } else {
          die sprintf("Unknown node type '%s' for node_id %d\n", $ancestor->get_value_for_tag('node_type'), $ancestor->node_id);
      }

     }
    }

      my @good_ones = ();
      foreach my $genepairlink (@pair_group) {
          my ($pep1, $pep2) = $genepairlink->get_nodes;
          if ($pep1->genome_db_id == $pep2->genome_db_id) {
              push @good_ones, [$genepairlink, 'within_species_paralog', 1];
          } elsif (($genepairlink->get_value_for_tag('ancestor')->duplication_confidence_score < $self->param('no_between')) and $self->is_closest_homologue($genepairlink)) {
              push @good_ones, [$genepairlink, $self->tag_orthologues($genepairlink), 0];
          }
      }
      foreach my $par (@good_ones) {
          $self->tag_genepairlink(@$par);
      }
  }

    close $self->param('output_filehandle') if $self->param('output_filehandle');
}

sub print_summary {
  my $self = shift;

  #display summary stats of analysis 
  if($self->debug) {
    printf("orthotree homologies\n");
    foreach my $type (keys(%{$self->param('orthotree_homology_counts')})) {
      printf ( "  %13s : %d\n", $type, $self->param('orthotree_homology_counts')->{$type} );
    }
  }

    if ($self->param('output_flatfile')) {
        printf("written to %s\n", $self->param('output_flatfile'));
    } else {
        printf("written to %s\n", $self->compara_dba->url);
    }

  $self->check_homology_consistency;

  my $counts_str = stringify($self->param('orthotree_homology_counts'));
  $self->param('gene_tree')->tree->store_tag('OrthoTree_types_hashstr', $counts_str) unless ($self->param('_readonly'));
}


sub display_link_analysis
{
  my $self = shift;
  my $genepairlink = shift;

  #display raw feature analysis
  my ($gene1, $gene2) = $genepairlink->get_nodes;
  my $ancestor = $genepairlink->get_value_for_tag('ancestor');
  printf("%21s(%7d) - %21s(%7d) : %10.3f dist : ",
    $gene1->gene_member->stable_id, $gene1->gene_member_id,
    $gene2->gene_member->stable_id, $gene2->gene_member_id,
    $genepairlink->distance_between);

  printf("%5s ", "");
  printf("%5s ", "");

  print("ancestor:(");
  my $node_type = $ancestor->get_value_for_tag('node_type', '');
  if ($node_type eq 'duplication') {
    print "DUP ";
  } elsif ($node_type eq 'dubious') {
    print "DD  ";
  } elsif ($node_type eq 'gene_split') {
    print "SPL ";
  } else {
    print "    ";
  }
  printf("%9s)", $ancestor->node_id);
  if ($ancestor->has_tag('duplication_confidence_score')) {
    printf(" %.4f ", $ancestor->get_value_for_tag('duplication_confidence_score'));
  } else {
    print " N/A    ";
  }

  printf(" %s %d %s\n",
         $genepairlink->get_value_for_tag('orthotree_type'),
         $genepairlink->get_value_for_tag('is_tree_compliant'),
         $ancestor->taxonomy_level() // 'N/A',
        );
}


sub get_ancestor_species_hash
{
    my $self = shift;
    my $node = shift;

    my $species_hash = $node->get_value_for_tag('species_hash');
    return $species_hash if($species_hash);

    $species_hash = {};

    if($node->isa('Bio::EnsEMBL::Compara::GeneTreeMember')) {
        my $node_genome_db_id = $node->genome_db_id;
        $species_hash->{$node_genome_db_id} = 1;
        $node->add_tag('species_hash', $species_hash);
        return $species_hash;
    }

    foreach my $child (@{$node->children}) {
        my $t_species_hash = $self->get_ancestor_species_hash($child);
        foreach my $genome_db_id (keys(%$t_species_hash)) {
            unless(defined($species_hash->{$genome_db_id})) {
                $species_hash->{$genome_db_id} = $t_species_hash->{$genome_db_id};
            } else {
                $species_hash->{$genome_db_id} += $t_species_hash->{$genome_db_id};
            }
        }
    }

    $node->add_tag("species_hash", $species_hash);
    return $species_hash;
}


sub delete_old_homologies {
    my $self = shift;

    my $outfile = $self->param('output_flatfile');
    if ( defined $outfile && -e $outfile ) {
        unlink $outfile;
        return;
    }

    my $tree_node_id = $self->param('gene_tree_id');

    # New method all in one go -- requires key on tree_node_id
    print "deleting old homologies\n" if ($self->debug);

    # Delete first the members
    my $sql1 = 'DELETE homology_member FROM homology JOIN homology_member USING (homology_id) WHERE gene_tree_root_id = ?';
    my $sth1 = $self->compara_dba->dbc->prepare($sql1);
    $sth1->execute($tree_node_id);
    $sth1->finish;

    # And then the homologies
    my $sql2 = 'DELETE FROM homology WHERE gene_tree_root_id = ?';
    my $sth2 = $self->compara_dba->dbc->prepare($sql2);
    $sth2->execute($tree_node_id);
    $sth2->finish;
}



########################################################
#
# Classification analysis
#
########################################################


sub tag_genepairlink
{
    my $self = shift;
    my $genepairlink = shift;
    my $orthotree_type = shift;
    my $is_tree_compliant = shift;

    $genepairlink->add_tag('orthotree_type', $orthotree_type);
    $genepairlink->add_tag('is_tree_compliant', $is_tree_compliant);

    if ($orthotree_type =~ /ortholog/) {
        my ($pep1, $pep2) = $genepairlink->get_nodes;
        my $has_match = $self->param('has_match');
        $has_match->{$pep1->seq_member_id}->{$pep2->genome_db_id} = $genepairlink;
        $has_match->{$pep2->seq_member_id}->{$pep1->genome_db_id} = $genepairlink;
    }

    $self->param('orthotree_homology_counts')->{$orthotree_type}++;
    my $n = $self->param('n_stored_homologies') + 1;
    $self->param('n_stored_homologies', $n);
    print STDERR "$n homologies\n" unless $n % 1000;

    $self->store_gene_link_as_homology($genepairlink) if $self->param('store_homologies');

}


sub tag_orthologues
{
    my $self = shift;
    my $genepairlink = shift;

    my ($pep1, $pep2) = $genepairlink->get_nodes;
    my $ancestor = $genepairlink->get_value_for_tag('ancestor');
    my $species_hash = $self->get_ancestor_species_hash($ancestor);
    my $count1 = $species_hash->{$pep1->genome_db_id};
    my $count2 = $species_hash->{$pep2->genome_db_id};

    if ($count1 == 1 and $count2 == 1) {
        return 'ortholog_one2one';
    } elsif ($count1 == 1 or $count2 == 1) {
        return 'ortholog_one2many';
    } else {
        return 'ortholog_many2many';
    }
}




sub is_closest_homologue
{
    my $self = shift;
    my $genepairlink = shift;

    my $has_match = $self->param('has_match');
    my ($pep1, $pep2) = $genepairlink->get_nodes;
    my $s1 = $self->get_ancestor_species_hash($genepairlink->get_value_for_tag('subtree1'));
    my $s2 = $self->get_ancestor_species_hash($genepairlink->get_value_for_tag('subtree2'));

    return 0 if $has_match->{$pep1->seq_member_id}->{$pep2->genome_db_id} or $has_match->{$pep2->seq_member_id}->{$pep1->genome_db_id};
    return 0 if exists $s1->{$pep2->genome_db_id};
    return 0 if exists $s2->{$pep1->genome_db_id};
    return 1;
}



########################################################
#
# Tree input/output section
#
########################################################

sub store_gene_link_as_homology {
  my $self = shift;
  my $genepairlink  = shift;

  $self->display_link_analysis($genepairlink) if($self->debug>2);
  my $type = $genepairlink->get_value_for_tag('orthotree_type');
  return unless($type);
  my $is_tree_compliant = $genepairlink->get_value_for_tag('is_tree_compliant');
  my $ancestor = $genepairlink->get_value_for_tag('ancestor');

  my ($gene1, $gene2) = $genepairlink->get_nodes;

  # get the mlss from the database
  my $mlss_type;
  my $gdbs;
  my $gdb1 = $gene1->genome_db;
  $gdb1 = $gdb1->principal_genome_db if $gdb1->genome_component;
  # Here, we need to be smart about choosing the mlss and the homology type
  if ($type =~ /^ortholog/) {
      my $gdb2 = $gene2->genome_db;
      $gdb2 = $gdb2->principal_genome_db if $gdb2->genome_component;
      if ($gdb1->is_polyploid and $gdb2->is_polyploid and ($gdb1->dbID == $gdb2->dbID)) {
          $mlss_type = 'ENSEMBL_HOMOEOLOGUES';
          $type      =~ s/ortholog/homoeolog/;
          $gdbs      = [$gdb1];
          #### temp fix triticum_aestivum
          if ( (($gdb1->name =~ m/^(triticum_aestivum|triticum_turgidum).*$/) and ($gene1->genome_db->genome_component eq 'U')) or
               (($gdb2->name =~ m/^(triticum_aestivum|triticum_turgidum).*$/) and ($gene2->genome_db->genome_component eq 'U')) ) {
              $mlss_type = 'ENSEMBL_PARALOGUES';
              $type      = 'within_species_paralog';
          }
          ####
      } else {
          $mlss_type = 'ENSEMBL_ORTHOLOGUES';
          $gdbs      = [$gdb1, $gdb2];
      }

  } elsif ($type eq 'alt_allele') {
      $mlss_type = 'ENSEMBL_PROJECTIONS';
      $gdbs      = [$gdb1];
  } else {
      $mlss_type = 'ENSEMBL_PARALOGUES';
      $gdbs      = [$gdb1];
  }

  my $mlss = $self->compara_dba->get_MethodLinkSpeciesSetAdaptor->fetch_by_method_link_type_GenomeDBs($mlss_type, $gdbs);

  # create an Homology object
  my $homology = new Bio::EnsEMBL::Compara::Homology;
  $homology->description($type);
  $homology->is_tree_compliant($is_tree_compliant);
  $homology->gene_tree_node($ancestor) if $ancestor;
  $homology->method_link_species_set($mlss);
  $homology->_species_tree_node_id($ancestor->get_value_for_tag('species_tree_node_id')) if $ancestor;
  
  $homology->add_Member($gene1->Bio::EnsEMBL::Compara::AlignedMember::copy);
  $homology->add_Member($gene2->Bio::EnsEMBL::Compara::AlignedMember::copy);
  $homology->update_alignment_stats;

  my $key = $mlss->dbID . "_" . $gene1->dbID;
  $self->param('homology_consistency')->{$key}{$type} = 1;

  # at this stage, gene_split have been retrieved from the node types
  if ($self->param('tag_split_genes')) {
    # Potential split genes: within_species_paralog that do not overlap at all
    if (($type eq 'within_species_paralog') and ($homology->get_all_Members->[0]->perc_cov == 0) and ($homology->get_all_Members->[1]->perc_cov == 0)) {
        $self->param('orthotree_homology_counts')->{'within_species_paralog'}--;
        $homology->description('gene_split');
        $homology->is_tree_compliant(0);
        $self->param('orthotree_homology_counts')->{'gene_split'}++;
    }
  }

    my $output_fh = $self->param('output_filehandle');
    if ( $output_fh ) {
        $homology->adaptor($self->compara_dba->get_HomologyAdaptor);
        # create homology_id from seq_member_ids as this will be unique
        $homology->dbID( join('_', sort map {$_->seq_member_id} @{ $homology->get_all_Members }) );
        foreach my $member (@{ $homology->get_all_Members }) {
            $self->param('seen_seq_member_ids')->{$member->seq_member_id} = 1;
            $self->param('seen_gene_member_ids')->{$member->gene_member_id} = 1;
        }
        print $output_fh join("\t", @{$homology->object_summary}) . "\n";
    } else {
        $self->param('homologyDBA')->store($homology) unless $self->param('_readonly');
    }

  return $homology;
}


sub check_homology_consistency {
    my $self = shift;

    print "checking homology consistency\n" if ($self->debug);
    my $bad_key = undef;

    foreach my $mlss_member_id ( keys %{$self->param('homology_consistency')} ) {
        my $count = scalar(keys %{$self->param('homology_consistency')->{$mlss_member_id}});

        next if $count == 1;
        next if $count == 2 and exists $self->param('homology_consistency')->{$mlss_member_id}->{gene_split} and exists $self->param('homology_consistency')->{$mlss_member_id}->{within_species_paralog};

        my ($mlss, $seq_member_id) = split("_", $mlss_member_id);
        next if $count > 1 and grep {$_->is_polyploid} @{$self->compara_dba->get_MethodLinkSpeciesSetAdaptor->fetch_by_dbID($mlss)->species_set->genome_dbs};

        $bad_key = "mlss seq_member_id : $mlss $seq_member_id";
        print "$bad_key\n" if ($self->debug);
    }
    $self->throw("Inconsistent homologies: $bad_key") if defined $bad_key;
}

sub _create_flatfile {
    my $self = shift;

    my $outfile = $self->param('output_flatfile');

    # create directory
    make_path(dirname($outfile));

    # open file handle and print header
    open( my $outfh, '>', $outfile ) or die "Cannot open $outfile for writing";
    print $outfh join("\t", @{ $Bio::EnsEMBL::Compara::Homology::object_summary_headers }) . "\n";
    $self->param('output_filehandle', $outfh);
    return $outfh;
}

sub _hc_flatfile {
    my $self = shift;

    my $file = $self->param('output_flatfile');
    print "running HCs on homologies:\n" if $self->debug;

    # first, check the integrity of the file
    $self->_check_file_integrity($file);
    print "\t- file integrity ok\n" if $self->debug;

    # then HC the homologies themselves
    my $seen_seq_member_ids  = $self->param('seen_seq_member_ids');
    my $seen_gene_member_ids = $self->param('seen_gene_member_ids');

    open( my $fh, '<', $file ) or die "Cannot open $file for reading";
    my $header_line = <$fh>;
    my @header_cols = split(/\s+/, $header_line);

    my $c = 1;
    my %seen_homologies;
    while ( my $line = <$fh> ) {
        my $row = map_row_to_header($line, \@header_cols);
        my ( $st_id, $hom_st_id ) = ($row->{stable_id}, $row->{homology_stable_id});

        # A pair of gene can only appear in 1 homology at most
        if (my $prev_line = ($seen_homologies{"$st_id - $hom_st_id"} || $seen_homologies{"$hom_st_id - $st_id"}) ) {
            die sprintf(
                "Pair (%s - %s) seen more than once (lines %d and %d)",
                ( $st_id, $hom_st_id, $prev_line, $c )
            );
        }
        $seen_homologies{"$st_id - $hom_st_id"} = $c;
        $seen_homologies{"$hom_st_id - $st_id"} = $c;

        # Checks that the member_ids are valid
        unless ($seen_seq_member_ids->{$row->{'seq_member_id'}}) {
            die sprintf("seq_member_id=%d does not exist in the tree (line %d)", $row->{'seq_member_id'}, $c);
        }
        unless ($seen_gene_member_ids->{$row->{'gene_member_id'}}) {
            die sprintf("gene_member_id=%d does not exist in the tree (line %d)", $row->{'gene_member_id'}, $c);
        }
        unless ($seen_seq_member_ids->{$row->{'homology_seq_member_id'}}) {
            die sprintf("seq_member_id=%d does not exist in the tree (line %d)", $row->{'homology_seq_member_id'}, $c);
        }
        unless ($seen_gene_member_ids->{$row->{'homology_gene_member_id'}}) {
            die sprintf("gene_member_id=%d does not exist in the tree (line %d)", $row->{'homology_gene_member_id'}, $c);
        }

        $c++;
    }
    print "\t- homologies ok\n\n" if $self->debug;
}

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

    # check line counts
    my $exp_lines = $self->param('n_stored_homologies') + 1; # incl header line
    check_line_counts($file, $exp_lines);

    # check column integrity
    check_column_integrity($file);
}

1;
