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

Bio::EnsEMBL::Compara::RunnableDB::GeneTrees::FlagUpdateClusters

=head1 DESCRIPTION

1) This module loops through all the genes from the current and previous databases.

2) Identifies and flags the genes that have been updated, deleted or newly added and
stores them as appropriate: "updated_genes_list", "added_genes_list" and "deleted_genes_list"

3) Checks for all the flagged genes (%flagged) in the list of members from all the
root_ids and updated the flag "needs_update" in the gene_tree_root_tag table.

=cut

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

use strict;
use warnings;
use Data::Dumper;

use Bio::EnsEMBL::Compara::Utils::CopyData qw(:insert);

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

sub fetch_input {
    my $self = shift @_;

    #get reuse_compara_dba adaptor
    $self->param( 'reuse_compara_dba', $self->get_cached_compara_dba('reuse_db') );
}

sub run {
    my $self = shift;

    #Get list of genes:
    print "getting prev_hash\n" if ( $self->debug );
    my $prev_hash = _hash_all_sequences_from_db( $self->param('reuse_compara_dba') );
    print "getting curr_hash\n" if ( $self->debug );
    my $curr_hash = _hash_all_sequences_from_db( $self->param('compara_dba') );
    print "getting reused vs current seq_member_ids map\n" if ( $self->debug );
    my $seq_member_id_map = _seq_member_map( $self->param('reuse_compara_dba'), $self->param('compara_dba') );
    $self->param( 'seq_member_id_map', $seq_member_id_map );

    #---------------------------------------------------------------------------------
    #deleted (members that existed in the previous mlss_id but dont exist in the current mlss_id), updated & added arent used by the logic.
    #deleted: is not used at all, but lets keep it here for debugging purpouse.
    #It is just here in case we need to test which genes are different in each case.
    #flag is that hash used my the module.
    #---------------------------------------------------------------------------------
    my ( %flag, %deleted, %updated, %added );
    print "flagging members on the sequence level\n" if ( $self->debug );
    _check_hash_equals( $prev_hash, $curr_hash, \%flag, \%deleted, \%updated, \%added );
    print "DELETED:|" . keys(%deleted) . "|\tUPDATED:|" . keys(%updated) . "|\tADDED:|" . keys(%added) . "|\n" if ( $self->debug );

    print "undef prev_hash\n" if ( $self->debug );
    undef %$prev_hash;

    print "undef curr_hash\n" if ( $self->debug );
    undef %$curr_hash;

    #Get list of current_stable_ids:
    print "getting list of roots\n" if ( $self->debug );
    my ( $current_stable_ids, $reused_stable_ids ) = $self->_get_stable_id_root_id_list( $self->param('compara_dba'), $self->param('reuse_compara_dba') );
    $self->param( 'current_stable_ids', $current_stable_ids );
    $self->param( 'reused_stable_ids',  $reused_stable_ids );

    my %root_ids_2_update;
    $self->param( 'root_ids_2_update', \%root_ids_2_update );

    my %root_ids_2_add;
    $self->param( 'root_ids_2_add', \%root_ids_2_add );

    my %root_ids_2_delete;
    $self->param( 'root_ids_2_delete', \%root_ids_2_delete );

    #get current tree adaptor
    $self->param( 'current_tree_adaptor', $self->param('compara_dba')->get_GeneTreeAdaptor );

    #get reused tree adaptor
    $self->param( 'reused_tree_adaptor', $self->param('reuse_compara_dba')->get_GeneTreeAdaptor );

    print "looping current_stable_ids\n" if ( $self->debug );
    foreach my $current_stable_id ( keys %{$current_stable_ids} ) {

        my $gene_tree_id = $current_stable_ids->{$current_stable_id};

        #get current_gene_tree
        my $current_gene_tree = $self->param('current_tree_adaptor')->fetch_by_dbID($gene_tree_id) or $self->die_no_retry("Could not fetch current_gene_tree with gene_tree_id='$gene_tree_id'");
        $self->throw("no input current_gene_tree") unless $current_gene_tree;
        my @current_members = @{ $current_gene_tree->get_all_Members };
        my $num_of_members  = scalar(@current_members);

        #get reused_gene_tree
        my $reused_gene_tree = $self->param('reused_tree_adaptor')->fetch_by_stable_id($current_stable_id);

        #There is no need for both array and hash structures here. Should change code to use only the hash!
        my @reused_members;
        my %reused_members_hash;
        my %current_members_hash;

        if ($reused_gene_tree) {
            @reused_members = @{ $reused_gene_tree->get_all_Members };

            foreach my $member (@reused_members) {
                my $reused_stable_id = $member->stable_id;
                my $reused_genome_db_id = $member->genome_db_id;
                $reused_members_hash{ $reused_stable_id . " " . $reused_genome_db_id } = 1;
            }
        }

        foreach my $current_member (@current_members) {

            my $current_stable_id_genome_db_id = $current_member->stable_id . " "  . $current_member->genome_db_id;
            $current_members_hash{ $current_stable_id_genome_db_id } = 1;

            if ( exists( $updated{ $current_stable_id_genome_db_id } ) ) {
                $root_ids_2_update{$gene_tree_id}{ $current_stable_id_genome_db_id } = 1;
            }

            if ( exists( $added{ $current_stable_id_genome_db_id } ) ) {
                $root_ids_2_add{$gene_tree_id}{ $current_stable_id_genome_db_id } = 1;
            }

            if ($reused_gene_tree) {
                if ( !exists( $reused_members_hash{ $current_stable_id_genome_db_id } )
                    && ( !exists( $added{ $current_stable_id_genome_db_id } ) )
                    && ( !exists( $updated{ $current_stable_id_genome_db_id } ) )
                ) {
                    $root_ids_2_add{$gene_tree_id}{ $current_stable_id_genome_db_id } = 1;
                }
            }
        }

        if ($reused_gene_tree) {
            foreach my $reused_member (@reused_members) {

                my $reused_stable_id_genome_db_id = $reused_member->stable_id . " " . $reused_member->genome_db_id;
                if ( exists( $deleted{ $reused_stable_id_genome_db_id } ) ) {
                    $root_ids_2_delete{$gene_tree_id}{ $reused_stable_id_genome_db_id } = 1;
                }

                # Removing members that are present in the reused tree and are not in the current tree
                # with the same sequence. This happens if there were diffences in the clustering, since
                # the sequences were the same they were not caught by previous MD5 checksum tests.
                if ( !$current_members_hash{ $reused_stable_id_genome_db_id } ) {
                    $root_ids_2_delete{$gene_tree_id}{ $reused_stable_id_genome_db_id } = 1;
                }

            }
            $reused_gene_tree->release_tree;
        }
        else {

            $current_gene_tree->store_tag( 'new_build', 1 ) || die "Could not store_tag 'new_build'";
        }

        my $members_2_change = scalar( keys( %{ $root_ids_2_add{$gene_tree_id} } ) ) + scalar( keys( %{ $root_ids_2_update{$gene_tree_id} } ) );

        if ( ( $members_2_change/scalar(@current_members) ) >= $self->param_required('update_threshold_trees') ) {
            $current_gene_tree->store_tag( 'new_build', 1 ) || die "Could not store_tag 'new_build'";
        }

        #releasing tree from memory
        $current_gene_tree->release_tree;

    }
}

sub write_output {
    my $self = shift;

    # When a genome is updated it may contain the same sequences and stable ids, but the seq_member_ids
    # will be different since it was re-inserted into the database when loading members.
    # This will cause the copying of the trees from the previous database to fail, since the old seq_member_ids
    # will not be the same as in the current database.
    # We just store the mapping now, it will later be used by Bio::EnsEMBL::Compara::RunnableDB::GeneTrees::CopyTreesFromDB.

    my @mapping_data;
    foreach my $stable_id ( keys %{ $self->param('seq_member_id_map') } ) {
        push(@mapping_data, [ $stable_id, $self->param('seq_member_id_map')->{$stable_id}->{'reused'}, $self->param('seq_member_id_map')->{$stable_id}->{'current'} ]);
    }
    bulk_insert($self->compara_dba->dbc, 'seq_member_id_current_reused_map', \@mapping_data, ['stable_id', 'seq_member_id_reused', 'seq_member_id_current']);
    @mapping_data= (); # To free the memory

    my %flagged;
    foreach my $current_stable_id ( keys %{ $self->param('current_stable_ids') } ) {

        my $gene_tree_id = $self->param('current_stable_ids')->{$current_stable_id};

        my $gene_tree = $self->param('current_tree_adaptor')->fetch_by_dbID($gene_tree_id) or die "Could not fetch gene_tree with gene_tree_id='$gene_tree_id'";

        #root_ids_2_update
        if ( keys( %{ $self->param('root_ids_2_update')->{$gene_tree_id} } ) ) {
            if ( !$flagged{$gene_tree_id} ) {
                $gene_tree->store_tag( 'needs_update',        1 ) || die "Could not store_tag 'needs_update' for $gene_tree_id";
                $gene_tree->store_tag( 'only_needs_deleting', 0 ) || die "Could not store_tag 'only_needs_deleting' for $gene_tree_id";
            }
            $gene_tree->store_tag( 'updated_genes_list', join( ",", keys( %{ $self->param('root_ids_2_update')->{$gene_tree_id} } ) ) ) ||
              die "Could not store_tag 'updated_genes_list' for $gene_tree_id";
            $flagged{$gene_tree_id} = 1;
        }

        #root_ids_2_add
        if ( keys( %{ $self->param('root_ids_2_add')->{$gene_tree_id} } ) ) {
            if ( !$flagged{$gene_tree_id} ) {
                $gene_tree->store_tag( 'needs_update',        1 ) || die "Could not store_tag 'needs_update' for $gene_tree_id";
                $gene_tree->store_tag( 'only_needs_deleting', 0 ) || die "Could not store_tag 'only_needs_deleting' for $gene_tree_id";
            }
            $gene_tree->store_tag( 'added_genes_list', join( ",", keys( %{ $self->param('root_ids_2_add')->{$gene_tree_id} } ) ) ) ||
              die "Could not store_tag 'added_genes_list' for $gene_tree_id";
            $flagged{$gene_tree_id} = 1;
        }

        #root_ids_2_delete
        if ( keys( %{ $self->param('root_ids_2_delete')->{$gene_tree_id} } ) ) {
            if ( ( !$gene_tree->has_tag('only_needs_deleting') ) && ( !$gene_tree->has_tag('needs_update') ) ) {
                $gene_tree->store_tag( 'only_needs_deleting', 1 ) || die "Could not store_tag 'only_needs_deleting' for $gene_tree_id";
            }
            $gene_tree->store_tag( 'deleted_genes_list', join( ",", keys( %{ $self->param('root_ids_2_delete')->{$gene_tree_id} } ) ) ) ||
              die "Could not store_tag 'deleted_genes_list' for $gene_tree_id";
        }

        if ( ( !keys( %{ $self->param('root_ids_2_update')->{$gene_tree_id} } ) ) &&
             ( !keys( %{ $self->param('root_ids_2_add')->{$gene_tree_id} } ) ) &&
             ( !keys( %{ $self->param('root_ids_2_delete')->{$gene_tree_id} } ) ) ) {
            $gene_tree->store_tag( 'identical_copy', 1 ) || die "Could not store_tag 'identical_copy' for $gene_tree_id";
        }
    }

}

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

#This function checks the differences between all the sequences loaded in the current database vs. the re-used database.
#It will not catch more subtles cases where sequences are the same but the were not present in the re-used tree.
sub _check_hash_equals {
    my ( $prev_hash, $curr_hash, $flag, $deleted, $updated, $added ) = @_;

    foreach my $stable_id ( keys %$curr_hash ) {
        if ( !exists( $prev_hash->{$stable_id} ) ) {
            $added->{$stable_id} = 1;
        }
        else {
            if ( $curr_hash->{$stable_id} ne $prev_hash->{$stable_id} ) {
                $updated->{$stable_id} = 1;
            }
        }
    }

    foreach my $stable_id ( keys %$prev_hash ) {
        if ( !exists( $curr_hash->{$stable_id} ) ) {
            $deleted->{$stable_id} = 1;
        }
    }
}

sub _hash_all_sequences_from_db {
    my $compara_dba = shift;

    my %sequence_set;

    my $sql = "SELECT stable_id, MD5(sequence) FROM seq_member JOIN sequence USING (sequence_id)";
    my $sth = $compara_dba->dbc->prepare($sql);
    $sth->execute() || die "Could not execute ($sql)";

    while ( my ( $stable_id, $seq_md5 ) = $sth->fetchrow() ) {
        $sequence_set{$stable_id} = lc $seq_md5;
    }
    $sth->finish();

    return \%sequence_set;
}

sub _get_stable_id_root_id_list {
    my ( $self, $compara_dba, $reuse_compara_dba ) = @_;

    my %current_stable_ids;
    my %reused_stable_ids;

    #current root_ids and stable_ids
    my $sql_current = "SELECT root_id, stable_id FROM gene_tree_root WHERE tree_type = \"tree\" AND clusterset_id=\"default\"";
    my $sth_current = $compara_dba->dbc->prepare($sql_current) || die "Could not prepare query.";
    $sth_current->execute() || die "Could not execute ($sql_current)";

    while ( my ( $root_id, $current_stable_id ) = $sth_current->fetchrow() ) {
        $current_stable_ids{$current_stable_id} = $root_id;
    }
    $sth_current->finish();

    #previous root_ids and stable_ids
    my $sql_reused = "SELECT root_id, stable_id FROM gene_tree_root WHERE tree_type = \"tree\" AND clusterset_id=\"default\" AND stable_id IS NOT NULL";
    my $sth_reused = $reuse_compara_dba->dbc->prepare($sql_reused) || die "Could not prepare query.";
    $sth_reused->execute() || die "Could not execute ($sql_reused)";

    while ( my ( $root_id, $reused_stable_id ) = $sth_reused->fetchrow() ) {
        $reused_stable_ids{$reused_stable_id} = $root_id;
    }
    $sth_reused->finish();

    return ( \%current_stable_ids, \%reused_stable_ids );
}

sub _seq_member_map {
    my $reuse_compara_dba = shift;
    my $compara_dba       = shift;

    my %map;

    my $sql_reuse = "SELECT stable_id, seq_member_id FROM seq_member";
    my $sth_reuse = $reuse_compara_dba->dbc->prepare($sql_reuse);
    $sth_reuse->execute() || die "Could not execute ($sql_reuse)";

    while ( my ( $stable_id_reuse, $seq_member_id_reuse ) = $sth_reuse->fetchrow() ) {
        $map{$stable_id_reuse}{'reused'} = $seq_member_id_reuse;
    }
    $sth_reuse->finish();

    my $sql_current = "SELECT stable_id, seq_member_id FROM seq_member";
    my $sth_current = $compara_dba->dbc->prepare($sql_current);
    $sth_current->execute() || die "Could not execute ($sql_current)";

    while ( my ( $stable_id_current, $seq_member_id_current ) = $sth_current->fetchrow() ) {
        $map{$stable_id_current}{'current'} = $seq_member_id_current;
    }
    $sth_current->finish();

    return \%map;
}

1;
