=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::GeneTrees::OverallGroupsetQC

=head1 DESCRIPTION

This RunnableDB compares the clusters of this compara database to another one's.

=cut

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

use strict;
use warnings;

use Time::HiRes qw(time gettimeofday tv_interval);

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


sub param_defaults {
    return {
        'unmap_tolerance'       => 0.2,
        'hmm_unmap_tolerance'   => 0.7,
    };
}


=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 @_;

    $self->complete_early('No reuse_db, cannot compare the clusters') unless $self->param('reuse_db');
    $self->param('groupset_tree', $self->compara_dba->get_GeneTreeAdaptor->fetch_all(-tree_type => 'clusterset', -clusterset_id => 'default')->[0]) or $self->die_no_retry("Could not fetch groupset tree");
}


=head2 run

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

=cut

sub run {
    my $self = shift @_;

    $self->overall_groupset_qc();
}



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

sub generate_dbname {
    my ($self, $given_compara_dba) = @_;

    return join('.',
        $given_compara_dba->dbc->host,
        $given_compara_dba->dbc->port,
        $given_compara_dba->dbc->dbname
    );
}


sub overall_groupset_qc {
    my $self     = shift @_;

    my $reuse_compara_dba = $self->get_cached_compara_dba('reuse_db');    # may die if bad parameters

    my $xtb_filename = $self->join_one_pair( $reuse_compara_dba, $self->compara_dba );

    my $xtb_copy_filename = $self->param('cluster_dir') . "/" . "groupset_qc.xtb";
    my $cpcmd = ['cp', $xtb_filename, $xtb_copy_filename];
    $self->run_command($cpcmd, { die_on_failure => 1 });

    my $map_filename = $self->cluster_mapping($xtb_filename, $reuse_compara_dba, $self->compara_dba);

    my $map_copy_filename = $self->param('cluster_dir') . "/" . "groupset_qc.map";
    $cpcmd = ['cp', $map_filename, $map_copy_filename];
    $self->run_command($cpcmd, { die_on_failure => 1 });

    # Cluster mapping can take quite some time,
    # and involves little or no database access.
    $reuse_compara_dba->dbc->disconnect_if_idle;
    $self->compara_dba->dbc->disconnect_if_idle;
    $self->param('groupset_tree')->adaptor->dbc->disconnect_if_idle;

    $self->quantify_mapping($map_filename, $reuse_compara_dba);
}



# ---------------------------------------------------------------------------------------------------------------
#       The following 3 subroutines have been re-written using proper classes in Bio::EnsEMBL::Compara::StableId,
#       We should be using those instead of copy-pasting.
#       Also see Bio::EnsEMBL::Compara::RunnableDB::StableIdMapper on proper manupulating those objects.
# ---------------------------------------------------------------------------------------------------------------


sub fetch_groupset {        # see Bio::EnsEMBL::Compara::StableId::Adaptor::load_compara_ncs
  my $self = shift;
  my $given_compara_dba = shift;

  my $starttime = time();

  my $default_noname = 'NoName';
  my $dataset;

  my $sql = "SELECT gtn.root_id, m2.stable_id FROM gene_tree_node gtn JOIN seq_member m1 USING (seq_member_id) JOIN gene_member m2 USING (gene_member_id)";

  my $sth = $given_compara_dba->dbc->prepare($sql);
  $sth->execute();

  printf("%1.3f secs to fetch entries\n", (time()-$starttime)) if ($self->debug);

  my $counter = 0;

  while(my($cluster_id, $member_stable_id)=$sth->fetchrow()) {
    # print STDERR "ID=$cluster_id NAME=$cluster_name MEM=$member_stable_id\n" if ($self->debug);
    my $cluster_name;
    if (defined($cluster_id)) {
      $cluster_name = 'Node_' . $cluster_id;
    } else {
      $cluster_name = $default_noname; # we need some name here however bogus (for formatting purposes)
    }

    if ($member_stable_id) {
      $dataset->{membership}{$member_stable_id} = $cluster_id;
      $dataset->{clustername}{$cluster_id} = $cluster_name;
    } else {
      $self->throw("Missing stable_id for $cluster_id\n");
    }

    if ($self->debug && ($counter++ % 50000 == 0)) { printf("%10d loaded\n", $counter); }
  }

  return $dataset;
}


sub join_one_pair {         # see Bio::EnsEMBL::Compara::StableId::NamedClusterSetLink::compute_stats
  my ($self, $from_dba, $to_dba) = @_;

  my $from_dbname       = $self->generate_dbname( $from_dba );
  print STDERR "Fetching groupset for $from_dbname\n" if ($self->debug);
  my $from_dataset      = $self->fetch_groupset( $from_dba );
  my $from_membership   = $from_dataset->{membership};
  my $from_clustername  = $from_dataset->{clustername};

  my $to_dbname     = $self->generate_dbname( $to_dba );
  print STDERR "Fetching groupset for $to_dbname\n" if ($self->debug);
  my $to_dataset    = $self->fetch_groupset( $to_dba );
  my $to_membership = $to_dataset->{membership};
  my $to_clustername= $to_dataset->{clustername};

  my %direct     = ();
  my %reverse    = ();
  my %from_size  = ();
  my %to_size    = ();
  my %xto_size   = ();
  my %xfrom_size = ();

  my $total_count  = 0;
  my $common_count = 0;

  while(my ($from_member, $from_family) = each %$from_membership) {
    if(my $to_family = $to_membership->{$from_member}) {
      $direct{$from_family}{$to_family}++;
      $reverse{$to_family}{$from_family}++;
      $from_size{$from_family}++;
      $to_size{$to_family}++;
      $common_count++;
    } else { # strictly disappeared members (disappearing either with or without the families)
      $xfrom_size{$from_family}++;
    }
    $total_count++;
  }
  while(my ($to_member, $to_family) = each %$to_membership) {
    if(not exists $from_membership->{$to_member}) { # strictly new members (going either into existing or new family)
      $xto_size{$to_family}++;
      $total_count++;
    }
  }

  print STDERR "Total number of keys:  $total_count\n" if ($self->debug);
  print STDERR "Number of common keys: $common_count\n" if ($self->debug);

  my $xtb_filename = $self->worker_temp_directory . "/" . $from_dbname . "-" . $to_dbname. '.xtb';
  open(my $out_fh, '>', $xtb_filename) or die "Could not open '$xtb_filename' for writing : $!";
  foreach my $from_id (sort {$a <=> $b} keys %direct) {
    my $from_name = $from_clustername->{$from_id};
    my $subhash = $direct{$from_id};

    foreach my $to_id (sort { $subhash->{$b} <=> $subhash->{$a} } keys %$subhash) {
      my $to_name = $to_clustername->{$to_id};
      my $cnt = $direct{$from_id}{$to_id};

      print $out_fh join("\t", $from_id, $from_name, $from_size{$from_id}, $to_id, $to_name, $to_size{$to_id}, $cnt)."\n";
    }
  }

  foreach my $to_id (sort {$a <=> $b} keys %xto_size) { # iterate through families that contain new members
    next if($reverse{$to_id}); # skip the ones that also have old members (i.e. iterate only through strictly-new families)
    my $to_name = $to_clustername->{$to_id};

    print $out_fh join("\t", 0, '-', 0, $to_id, $to_name, $xto_size{$to_id}, $xto_size{$to_id})."\n";
  }
  foreach my $from_id (sort {$a <=> $b} keys %xfrom_size) { # iterate through families that lost some members
    next if($direct{$from_id}); # skip the families that retained some members (i.e. iterate only through strictly-disappearing families)
    my $from_name = $from_clustername->{$from_id};

    print $out_fh join("\t", $from_id, $from_name, $xfrom_size{$from_id}, 0, '-', 0, $xfrom_size{$from_id})."\n";
  }
  close $out_fh;

  return $xtb_filename;
}


sub cluster_mapping {       # see Bio::EnsEMBL::Compara::StableId::NamedClusterSetLink::maximum_name_reuse
  my ($self, $link_filename, $from_dba, $to_dba) = @_;

  my $premap = ''; # premap will always be empty here
  my $threshold = 0.67;

  my $maj_label =  $threshold ? sprintf("Major_%d", int($threshold*100) ) : 'Majority';
  my @labels = ('Exact', 'Exact_o', $maj_label, $maj_label.'_o', 'NextBest', 'NextBest_o', 'NewName', 'NewName_o', 'NewFam', 'NewFam_o');

  my $default_version = 1;
  my $prefix = 'ENSGT';
  my $to_rel = 2;
  my $stid_counter = 0;

  my $revcontrib;
  my $from2size;
  my $to2size;
  my $xfrom2size;
  my $xto2size;
  my $from2name;
  my $to2name;

  open(my $link_fh, '<', $link_filename) || die "Cannot open '$link_filename' file $@";
  while (my ($from_id, $from_name, $from_size, $to_id, $to_name, $to_size, $contrib) = split(/\s/,<$link_fh>)) {

    next unless($contrib=~/^\d+$/); # skip the header line if present

    if($from_size and $to_size) { # Shared
      $from2name->{$from_id}          = $premap ? ($premap->{$from_id} || die "Premap does not contain mapping for '$from_name' (id=$from_id)") : $from_name;
      $to2name->{$to_id}              = $to_name;

      $revcontrib->{$to_id}{$from_id} = $contrib;
      $from2size->{$from_id}          = $from_size;
      $to2size->{$to_id}              = $to_size;
    } elsif($to_size) { # Newborn
      $to2name->{$to_id}              = $to_name;

      $xto2size->{$to_id}             = $to_size;
    } elsif($from_size) { # Disappearing
      $from2name->{$from_id}          = $premap ? ($premap->{$from_id} || die "Premap does not contain mapping for '$from_name' (id=$from_id)") : $from_name;

      $xfrom2size->{$from_id}         = $from_size;
    }
  }
  close $link_fh;

  # Now we run through the hashes
  my %matchtype_counter = ();
  my %from_taken        = (); # indicates the 'from' name has been taken 
  my %postmap           = (); # the goal of this subroutine is to map between the '$to' and '$given_name'-s

  my $from_dbname   = $self->generate_dbname( $from_dba );
  my $to_dbname     = $self->generate_dbname( $to_dba );

  my $map_filename = $self->worker_temp_directory . "/" . $from_dbname . "-" . $to_dbname. '.map';
  open(my $map_fh, '>', $map_filename) or die "Could not open '$map_filename' for writing : $!";

 TOPAIR: foreach my $topair (sort { $b->[1] <=> $a->[1] } map { [$_,$to2size->{$_}] } keys %$to2size ) {
    my ($to_id, $to_size) = @$topair;

    my $subhash = $revcontrib->{$to_id};

    my $td_counts  = 0;
    my $matchtype  = '';
    my $matchscore = 0;
    my $given_name = ''; # serves both as logical flag and the actual name

  FROMPAIR: foreach my $frompair (sort { $b->[1] <=> $a->[1] } map { [$_,$subhash->{$_}] } keys %$subhash ) {
      my ($from_id, $contrib) = @$frompair;
      my $from_size = $from2size->{$from_id};

      my $from_name = $from2name->{$from_id};
      if (!defined $from_taken{$from_name}) { # means the '$from' name is still unused, so we can reuse it now

        if ($contrib==$from_size and $contrib==$to_size) {
          $matchtype  = 'Exact';
        } elsif ($threshold>0) { # either the majority rule is applicable or we don't bother looking at other possibilities (as they are even smaller)
          if ($contrib/$from_size>=$threshold and $contrib/$to_size>=$threshold) {
            $matchtype  = $maj_label;
          }             # otherwise we have an implicit 'NewName' case
        } else {                # non-threshold mode
          $matchtype = $td_counts ? 'NextBest' : $maj_label;
        }

        if ($matchtype) {
          if ($matchtype eq 'Exact') {
            # $from_name =~ /^(\w+)(?:\.(\d+))?/;
            # $given_name = $1.'.'. (defined($2) ? $2 : $default_version ); # same version (but we may want to make it more obvious)
            $given_name = $from_name;
          } else {
            $from_name =~ /^(\w+)(?:\.(\d+))?/;
            $given_name = $1.'.'. ((defined($2) ? $2 : $default_version)+1); # change the version if the match is not exact (or set it if previously unset)
          }
          $from_taken{$from_name} = 1;
          $matchscore = int(100*$contrib/$to_size);
        }
        last FROMPAIR;

      }                         # if name not taken

      $td_counts++; # counts all attempts, not only the ones where the '$from' name was unused

    }                           # FROMPAIR

    # the following two lines work either if we arrive here from 'last FROMPAIR' after implicit 'NewName'
    # or by exhausting all FROMPAIRS (beacause they were all taken)
    $matchtype  ||= 'NewName';
    $given_name ||= sprintf("%s%04d%010d.%d",$prefix, $to_rel, ++$stid_counter, $default_version);

    print $map_fh (join("\t", $to_id, $to2name->{$to_id}, $given_name, $matchscore)."\n");
    $postmap{$to_id} = $given_name;

    if ($to_size == 1) {
      $matchtype .= '_o';
    }
    $matchtype_counter{$matchtype}++;
  }                             # TOPAIR

  while (my ($to_id, $to_size) = each %$xto2size) {
    my $given_name = sprintf("%s%04d%010d.%d",$prefix, $to_rel, ++$stid_counter, $default_version);
    print $map_fh join("\t", $to_id, $to2name->{$to_id}, $given_name, 0)."\n";
    $postmap{$to_id} = $given_name;

    my $matchtype = ($to_size == 1) ? 'NewFam_o' : 'NewFam';
    $matchtype_counter{$matchtype}++;
  }
  close $map_fh;

  return $map_filename;
}


sub quantify_mapping {
  my ($self, $map_filename, $reuse_compara_dba) = @_;

  my %mapping_stats = ();

  open(my $map_fh, '<', $map_filename) or die "Could not open '$map_filename' for reading : $!";
  my $tag_count = 0;
  while (<$map_fh>) {
    my ($cluster_id, $from_cluster_name, $to_cluster_name, $contribution) = split(" ",$_);
    if ($to_cluster_name =~ /ENSGT/) {
      $mapping_stats{novel}{$cluster_id} = 1;
    } else {
      $to_cluster_name =~ /\_(\d+)/;
      my $reuse_node_id = $1;
      $mapping_stats{mapped}{$cluster_id} = $contribution;
      $mapping_stats{mapped_tagging}{$cluster_id} = $reuse_node_id;
    }
    if ($self->debug && ($tag_count++ % 100 == 0)) { print STDERR "[$tag_count] mapped clusters\n"; }
  }
  close $map_fh;

  my $current_gene_tree_adaptor = $self->compara_dba->get_GeneTreeAdaptor;
  my $reuse_gene_tree_adaptor = $reuse_compara_dba->get_GeneTreeAdaptor;
  foreach my $mapped_cluster_id (keys %{$mapping_stats{mapped_tagging}}) {
    my $reuse_node_id = $mapping_stats{mapped_tagging}{$mapped_cluster_id};
    next unless (defined($reuse_node_id));
    my $reuse_node = $reuse_gene_tree_adaptor->fetch_by_dbID($reuse_node_id);
    next unless (defined($reuse_node));
    next unless $reuse_node->has_tag('aln_runtime');
    my $reuse_aln_runtime_value = $reuse_node->get_all_values_for_tag('aln_runtime')->[0];  # FIXME it seems that there are multiple occurrences of the tag for some trees, so we can't use get_value_for_tag
    my $this_node = $current_gene_tree_adaptor->fetch_by_dbID($mapped_cluster_id);
    next unless (defined($this_node));
    $this_node->store_tag('reuse_node_id',$reuse_node_id);
    $this_node->store_tag('reuse_aln_runtime',$reuse_aln_runtime_value);
    my $contribution = $mapping_stats{mapped}{$mapped_cluster_id};
    next unless (defined($contribution));
    $this_node->store_tag('reuse_contribution',$contribution);
  }

  my $num_novel_clusters = scalar keys %{$mapping_stats{novel}};
  my $num_mapped_clusters = scalar keys %{$mapping_stats{mapped}};

  my $sum_contrib;
  foreach my $mapped_cluster_id (keys %{$mapping_stats{mapped}}) {
    $sum_contrib += $mapping_stats{mapped}{$mapped_cluster_id};
  }
  if ($num_mapped_clusters == 0) {
    $self->warning('No mapped clusters');
    return;
  }

  my $average_mapped_contribution = $sum_contrib / $num_mapped_clusters;

  my $proportion_novel_clusters = $num_novel_clusters/($num_novel_clusters+$num_mapped_clusters);

  print STDERR "# Proportion novel clusters = $proportion_novel_clusters [$num_novel_clusters $num_mapped_clusters]\n";
  print STDERR "# Average contribution mapped clusters = $average_mapped_contribution\n";

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

  $groupset_tree->store_tag('sid_map_novel_cls', $num_novel_clusters);
  $groupset_tree->store_tag('sid_map_mapped_cls', $num_mapped_clusters);
  $groupset_tree->store_tag('sid_map_summary_contrib', $sum_contrib);
  $groupset_tree->store_tag('sid_map_average_contrib', $average_mapped_contribution);
  $groupset_tree->store_tag('sid_prop_novel_cls', $proportion_novel_clusters);

  my $hmm_switch = ($self->_db_has_hmm_annots($self->compara_dba) xor $self->_db_has_hmm_annots($reuse_compara_dba));
  if ($hmm_switch) {
      $self->warning('Clusters created with different clustering methods (HMMs vs blast)');
  }

  my $unmap_tolerance = $self->param_required($hmm_switch ? 'hmm_unmap_tolerance' : 'unmap_tolerance');
  print STDERR "# Unmap tolerance parameter set to $unmap_tolerance\n";

  if ($proportion_novel_clusters > $unmap_tolerance) {
    $self->input_job->transient_error(0);
    die "Quality Check FAILED: Proportion of novel clusters $proportion_novel_clusters > $unmap_tolerance\n";
  }

  return;
}

sub _db_has_hmm_annots {
    my ($self, $dba) = @_;
    my $sql = 'SELECT 1 FROM gene_tree_root_tag WHERE tag = "model_id" LIMIT 1';
    my $sth = $dba->dbc->prepare($sql);
    $sth->execute();
    my ($found_hmm_annot) = $sth->fetchrow_array;
    return $found_hmm_annot;
}

1;
