Cytoscape_table_from_colocation_analysis.py

#!/usr/bin/env python3
# authors: Jakub Barylski, Sophia Bałdysz
# coding: utf-8
"""
Generate Cytoscape-importable edge and node tables that represents network of directed HMM-HMM similarities
generated by analysis of their co-occurrence in protein sequences
"""

from collections import defaultdict, namedtuple
from itertools import chain, product
from pathlib import Path
from typing import Any, List, Union, Dict, Tuple, Collection, DefaultDict, Hashable, Set

import numpy as np
import pandas as pd
from scipy.stats import fisher_exact
from tqdm import tqdm

from utils import checkpoint, BatchParallel, summarize_fasta

PathLike = Union[Path, str]

# dictionary with column numbers for different HMMER3 domtblout file formats
PARSER_DICT = {'hmmsearch': {'prot_id': 0,
                             'hmm_id': 3,
                             'prot_start': 17,
                             'prot_end': 18,
                             'evalue': 6,
                             'ivalue': 12,
                             'score': 13},
               'hmmscan': {'hmm_id': 0,
                           'prot_id': 3,
                           'prot_start': 17,
                           'prot_end': 18,
                           'evalue': 6,
                           'ivalue': 12,
                           'score': 13}}


class Domain:
    """
    Representation of a protein domain (e.g. single contious HMMer3 HMM-sequence alignment)
    :ivar prot_id: identifier of the protein e.g. A0A2K9ATG3
    :ivar hmm_id: identifier of the hidden Markov model e.g. VOG0010, PG_binding_1
    :ivar prot_start: position in the sequence where a domain (alignment) starts
    :ivar score: alignment score (e.g. HMMer3 domain score)
    :ivar overlaps: dictionary cataloging overlaps with other domains/alignments {HMM_1: {overlap_A, overlap_B, (...)}, (...)}
    :ivar merged: is this domain result of merging multiple alignments
    """

    def __init__(self, prot_id: str,
                 hmm_id: str,
                 prot_start: int,
                 prot_end: int,
                 score: float,
                 merged: bool = False):
        """
        Object initialisation (creation of the class instance)
        the functions assumes that has a line from hmmsearch
        and has to be told otherwise if hmmscan was used
        """
        self.prot_id = prot_id
        self.hmm_id = hmm_id
        self.prot_start, self.prot_end = prot_start, prot_end
        self.score = score
        self.overlaps = defaultdict(set)
        self.merged = merged

    def __repr__(self) -> str:
        """
        How an instance should look like in print etc.
        """
        return f'{self.hmm_id} [{self.prot_id} {self.prot_start} -> {self.prot_end}] ({self.score})'

    def __eq__(self, other: Any) -> bool:
        """
        Domains are equivalent if they represent the same part of the protein aligned to the same HMM
        :param other: other object to check for identity
        :return: are compared objects identical?
        """
        if isinstance(other, Domain):
            if self.prot_id == other.prot_id and self.hmm_id == other.hmm_id:
                return self.prot_start == other.prot_start and self.prot_end == other.prot_end
        return False

    def __hash__(self) -> int:
        """
        Domain should be hashed based on its identity (protein locationa and aligned HMM)
        :return: hash number of the domain
        """
        return hash((self.prot_id, self.hmm_id, self.prot_start, self.prot_end))

    def gff(self, id_string: str) -> str:
        """
        Represent a domain as a GFF line
        :param id_string: a short, unique identifier string
        :return: gff-formatted line (no newline (\n) at the end)
        """
        attributes = {'ID': id_string,
                      'HMM': self.hmm_id}
        attribute_string = ';'.join(f'{k}={v}' for k, v in attributes.items())
        return '\t'.join([self.prot_id,
                          'HMMER3',
                          'domain',
                          str(self.prot_start),
                          str(self.prot_end),
                          str(self.score),
                          '.', '.',
                          attribute_string])

    def seq_length(self) -> int:
        """
        Calculate length of domain in AA
        :return: length of domain in AA
        """
        return self.prot_end - self.prot_start + 1

    def does_overlap(self,
                     domain: 'Domain',
                     annotate_domains: bool = False) -> 'Overlap':
        """
        Check if this domain overlaps with the other and return
        overlap object representing the span of the sequence overlap.
        :param domain: other domain to compare 'self'
        :param annotate_domains: should the overlap be noted in 'overlaps' attribute of both partners
        :return: annotation representing the overlap between two other annotations
        """
        assert self.prot_id == domain.prot_id
        if self.prot_end >= domain.prot_start and self.prot_start <= domain.prot_end:
            overlap_start = max(self.prot_start, domain.prot_start)
            overlap_end = min(self.prot_end, domain.prot_end)
            overlap = Overlap(self.prot_id, overlap_start, overlap_end, self, domain)
            if annotate_domains:
                domains_to_annotate = set(overlap.parent_annotations)
                while domains_to_annotate:
                    first_domain = domains_to_annotate.pop()
                    if domains_to_annotate:
                        second_domain = domains_to_annotate.pop()
                        first_domain.overlaps[second_domain.hmm_id].add(overlap)
                        second_domain.overlaps[first_domain.hmm_id].add(overlap)
                    else:
                        first_domain.overlaps[first_domain.hmm_id].add(overlap)
            return overlap
        elif domain.prot_end >= self.prot_start and domain.prot_start <= self.prot_end:
            return domain.does_overlap(self)

    def merge(self, domain: 'Domain') -> 'Domain':
        """
        Create new annotation that is the union of two parent domains
        :param domain: other domain to merge 'self'
        :return: new domain constructed from two original ones
        """
        shared_prot, = {self.prot_id, domain.prot_id}  # check if both domains are placed on the same protein
        shared_hmm, = {self.hmm_id, domain.hmm_id}  # check if both domains refer to the same HMM
        if self.prot_end >= domain.prot_start and self.prot_start <= domain.prot_end:
            start = min(self.prot_start, domain.prot_start)
            end = max(self.prot_end, domain.prot_end)
            score = self.score + domain.score
            merged_domain = Domain(shared_prot, shared_hmm, start, end, score, merged=True)
            return merged_domain
        elif domain.prot_end >= self.prot_start and domain.prot_start <= self.prot_end:
            return domain.merge(self)
        else:
            raise ValueError(f'Illegal merge operation on domains:\n{self}\nand\n{domain}')

    @classmethod
    def from_hmmer(cls, hmmer_line, program) -> 'Domain':
        split_line = hmmer_line.split()
        prot_id = split_line[PARSER_DICT[program]['prot_id']]
        hmm_id = split_line[PARSER_DICT[program]['hmm_id']]
        prot_start, prot_end = int(split_line[PARSER_DICT[program]['prot_start']]), int(split_line[PARSER_DICT[program]['prot_end']])
        score = float(split_line[PARSER_DICT[program]['score']])
        return cls(prot_id, hmm_id, prot_start, prot_end, score)


class AnnotationTrack(set):
    """
    Representation of all annotations in a given sequence (e.g. domains in protein)
    :ivar sequence_lengths: dictionary with lengths of proteins e.g. {A0A2K9ATG3: 193, (...)}
    :ivar n_sequences: total number of sequences in the track
    :ivar n_positions: total number of positions ina all sequences in the track (e.g. sum of all aminoacids in proteins)
    :ivar by_hmm: dictionary representing domains sorted by Markov model they align to
                  {VOG0010: {d1, d23, (...)}, PG_binding_1: {d7, d11, (...)}, (...)}
    :ivar by_sequence: dictionary representing domains sorted by sequence
                      {Q4L555: {d1, d23, (...)}, A0A1F4B2T0: {d7, d11, (...)}, (...)}
    :ivar by_protein_and_hmm: layered dictionary representing domains sorted by protein and HMM
                              {Q4L555: {VOG0010: {d1, (...)}, PG_binding_1: {d7, (...)}}, (...)}
    :ivar hmm_sorted_proteins: dictionary of proteins that align to HMMs
                               {PG_binding_1: {Q4L555, A0A1F4B2T0}, (...)}
    """

    def __init__(self, *args):
        set.__init__(self, *args)
        self.sequence_lengths = {}
        self.n_sequences = 0
        self.n_positions = 0
        self.by_hmm = defaultdict(set)
        self.by_sequence = defaultdict(set)
        self.by_protein_and_hmm = defaultdict(lambda: defaultdict(set))
        self.hmm_sorted_proteins = defaultdict(set)

    def _import_fasta(self, fasta_path: PathLike):
        """
        Read basic sequence stats (number, lengths and total number of positions)
        that will be needed during further analysis of annotations
        :param fasta_path: path to the fasta file
        """
        fasta_path = Path(fasta_path)
        self.sequence_lengths = summarize_fasta(fasta_path)
        self.n_sequences = len(self.sequence_lengths)
        self.n_positions = sum(self.sequence_lengths.values())

    def add(self, domain):
        """
        Add new annotation (domain) to the track
        :param domain: domain to be added
        """
        assert isinstance(domain, Domain)
        self.by_hmm[domain.hmm_id].add(domain)
        self.by_sequence[domain.prot_id].add(domain)
        self.by_protein_and_hmm[domain.prot_id][domain.hmm_id].add(domain)
        self.hmm_sorted_proteins[domain.hmm_id].add(domain.prot_id)
        super().add(domain)

    def remove(self,
               domain: Domain,
               with_protein: bool = False):
        """
        Delete an annotation (domain) from the track
        :param domain: domain to be deleted
        """
        self.by_hmm[domain.hmm_id].remove(domain)
        self.by_sequence[domain.prot_id].remove(domain)
        self.by_protein_and_hmm[domain.prot_id][domain.hmm_id].remove(domain)
        if with_protein:
            self.hmm_sorted_proteins[domain.hmm_id].remove(domain.prot_id)
        super().remove(domain)

    def domains_with_hmms(self,
                          hmms: Collection) -> 'AnnotationTrack':
        """
        Return the child track that contains only hits to specified HMMs
        :param hmms: collection of the HMMs to include in the filtered track
        :return: filtered annotation track
        """
        hmms = set(hmms)
        filtered_domains = [domain for domain in self if domain.hmm_id in hmms]
        filtered_track = AnnotationTrack()
        for d in filtered_domains:
            filtered_track.add(d)
        return filtered_track

    def proteins_with_hmms(self,
                           hmms: Collection) -> 'AnnotationTrack':
        """
         Return the child track that includes only SEQUENCES (proteins) containing HMMs from the list
         :param hmms: collection of the HMMs used to select proteins that will be included in the filtered track
         :return: filtered annotation track
         """
        proteins = set(chain(*[self.hmm_sorted_proteins[h] for h in hmms]))
        filtered_domains = [domain for domain in self if domain.prot_id in proteins]
        filtered_track = AnnotationTrack()
        for d in filtered_domains:
            filtered_domains.append(d)
        return filtered_track

    def merge_redundant(self):
        """
        Marge any overlapping in all studied proteins alignments to the same HMM
        to simplify the analysis of the HMM/HMM co-occurrence
        """

        pre_merge = len(self)

        dereplication_input = list(self.by_protein_and_hmm.values())
        jobs = BatchParallel(dereplicate_domains, n_jobs=25,
                             input_collection=dereplication_input,
                             partition_size=5000)
        all_merged = 0
        for r in jobs.result:
            redundant, merged = r
            [self.remove(r) for r in redundant]
            [self.add(m) for m in merged]
            all_merged += len(merged)

        post_merge = len(self)
        print(f'merged {all_merged} of {pre_merge} redundant domains ({post_merge} remaining)')

    def find_overlaps(self):
        """
        Find annotation overlaps in all imported sequences
        (detected overlaps will be stored in 'overlaps' attribute of overlapping domains)
        """
        with tqdm(total=len(self.by_protein_and_hmm)) as bar:
            bar.set_description('Finding overlaps')
            for protein, domains in self.by_sequence.items():
                sorted_domains = sorted(domains, key=lambda dom: dom.prot_start)
                while sorted_domains:
                    analyzed_domain = sorted_domains.pop(0)
                    [analyzed_domain.does_overlap(compared_domain, annotate_domains=True) for compared_domain in sorted_domains]
                bar.update()

    def find_overlaps_with_hmms(self,
                                selected_hmms: Collection):
        """
        Find annotation overlaps in all imported sequences, but only,
        if at lest one of overlapping domains align to any of specified HMMs
        (detected overlaps will be stored in 'overlaps' attribute of overlapping domains)
        :param selected_hmms: models of infest {'PFAM_PF12671.8', 'VOG00138', (...)}
        """
        with tqdm(total=len(self.by_protein_and_hmm)) as bar:
            bar.set_description(f'Finding overlaps for {len(selected_hmms)} models')
            for protein, domains_for_hmms in self.by_protein_and_hmm.items():
                sorted_domains = sorted(self.by_sequence[protein], key=lambda dom: dom.prot_start)
                domains_of_interest = set(chain(*[domains_for_hmms[hmm] for hmm in selected_hmms if hmm in domains_for_hmms]))
                for analyzed_domain in domains_of_interest:
                    [analyzed_domain.does_overlap(compared_domain, annotate_domains=True) for compared_domain in sorted_domains]
                    sorted_domains.remove(analyzed_domain)
                bar.update()

    def finalize(self):
        """
        Finalize the creation of the track by converting defaulters to built-in dictionaries
        """
        self.by_hmm = standard_dict(self.by_hmm)
        self.by_sequence = standard_dict(self.by_sequence)
        self.by_protein_and_hmm = standard_dict(self.by_protein_and_hmm)
        self.hmm_sorted_proteins = standard_dict(self.hmm_sorted_proteins)

    def summary(self) -> Tuple[
        Dict[str, int],
        Dict[str, List[int]],
        Dict[str, List[int]],
        Dict[str, List[float]]]:
        """
        Generate dictionaries in {'hmm_id': [number_x, number_y, (...)], (...)} format
        summarizing statistics of ech model (HMM) found in analysed sequences.
        :return: total number of proteins aligned to each HMM
                 number of domains aligned to each HMM in different proteins
                 lengths of alignments associated with HMM in different proteins
                 fractions of various proteins covered by alignments associated with HMM
        """
        proteins = {hmm: len(proteins) for hmm, proteins in self.hmm_sorted_proteins.items()}
        counts = defaultdict(list)
        spans = defaultdict(list)
        protein_fractions = defaultdict(list)
        for protein_id, hmm_sorted_domains in self.by_protein_and_hmm.items():
            for hmm, domains in hmm_sorted_domains.items():
                counts[hmm].append(len(domains))
                total_span = sum([d.seq_length() for d in domains])
                spans[hmm].append(total_span)
                protein_fractions[hmm].append(total_span / self.sequence_lengths[protein_id])

        return proteins, standard_dict(counts), standard_dict(spans), standard_dict(protein_fractions)

    def overlap_summary(self) -> Tuple[
        Dict[str, Dict[str, int]],
        Dict[str, Dict[str, List[int]]],
        Dict[str, Dict[str, List[int]]],
        Dict[str, Dict[str, List[float]]]]:
        """
        Generate dictionaries in {hmm1: {hmm2: 133}} or {hmm1: {hmm2: [3, 1, 8 ...]}} format
        summarizing statistics of domain co-occurrence in analysed sequences.
        :return: "shared_proteins" - a dictionary with total numbers of proteins where 2 HMMs are present together (but not necessarily overlap)
                                     {hmm1: {hmm2: 131}}
                 "counts" - a dictionary with list summarizing how many times each domain aligned to one HMM overlaps with the other
                            {hmm1: {hmm2: [3, 1, 8 ...]}}
                 "spans" - a dictionary with list of total length of sequence (AA) in domains aligned to the HMM_A
                           covered by overlaps with the domain aligned to the HMM_B
                           {hmm1: {hmm2: [32, 29, 44 ...]}}
                 "span_fractions" - a dictionary with list of fractions of domains aligned to the HMM_A
                                    covered by overlaps with the domains aligned to the HMM_ B (overlap_length / domain_A_length)
                                    {hmm1: {hmm2: [0.16, 0.145, 0.22 ...]}}
        """
        shared_proteins = defaultdict(lambda: defaultdict(int))
        counts = defaultdict(lambda: defaultdict(list))
        spans = defaultdict(lambda: defaultdict(list))
        span_fractions = defaultdict(lambda: defaultdict(list))

        for domain in self:
            domain_length = domain.seq_length()
            for target_hmm, overlaps in domain.overlaps.items():
                counts[domain.hmm_id][target_hmm].append(len(overlaps))
                total_span = sum([o.seq_length() for o in overlaps])
                spans[domain.hmm_id][target_hmm].append(total_span)
                span_fractions[domain.hmm_id][target_hmm].append(total_span / domain_length)

        for protein, hmms in self.by_protein_and_hmm.items():
            symmetrical_increment(shared_proteins, hmms)

        return standard_dict(shared_proteins), standard_dict(counts), standard_dict(spans), standard_dict(span_fractions)


def dereplicate_domains(hmm_domains: Dict[str, Set[Domain]]):
    """
    Find and marge overlapping alignments to the same HMM in a given protein
    return sets of domains (alignments) to delete and merge
    :param hmm_domains: all domains found in a given proteins
    :return: 'redundant' - domains that have been merged with other and should be deleted
             'merged' - domains produced by merging of overlaps
    """
    redundant = set()
    merged = set()
    for hmm, domains in hmm_domains.items():
        if len(domains) > 1:
            sorted_domains = sorted(domains, key=lambda dom: dom.prot_start)
            analyzed_domain = sorted_domains.pop(0)
            compared_domain = sorted_domains.pop(0)
            while compared_domain:
                if analyzed_domain.does_overlap(compared_domain):
                    if analyzed_domain.merged:
                        merged.remove(analyzed_domain)
                        redundant.add(compared_domain)
                    else:
                        redundant.update({analyzed_domain, compared_domain})
                    merged_domain = analyzed_domain.merge(compared_domain)
                    merged.add(merged_domain)
                    analyzed_domain = merged_domain
                else:
                    analyzed_domain = compared_domain
                compared_domain = sorted_domains.pop(0) if sorted_domains else None
    return redundant, merged


def symmetrical_increment(target_dict: DefaultDict[Hashable, DefaultDict[Hashable, int]],
                          keys: Collection[Hashable],
                          increment: int = 1):
    """
    Increment two corresponding values in a layered {A: {B: X, (...)}, B: {A: X, (...)}, (...)} dictionary
    :param target_dict: a dictionary to update
    :param keys: keys representing an entry pair
    :param increment: number to add to a values
    """
    pairs_to_increment = tuple(product(keys, keys))
    for k0, k1 in pairs_to_increment:
        target_dict[k0][k1] += increment


def standard_dict(target_defaultdict: DefaultDict):
    """
    Convert multi-layered default dictionary to build-in dict.
    :param target_defaultdict: defaultdict to be converted
    :return: a converted dictionary
    """
    if isinstance(target_defaultdict, defaultdict):
        result = {k: standard_dict(v) for k, v in target_defaultdict.items()}
    else:
        result = target_defaultdict
    return result


class Overlap:
    """
    Representation of an overlap between two annotations.
    :ivar prot_id: id of a protein shared by two annotations
    :ivar prot_start: position in protein where overlap starts
    :ivar prot_end: position in protein where overlap ends
    :ivar parent_annotations: annotations that overlapped
    """

    def __init__(self, prot_id,
                 start: int, end: int,
                 annot_0: Domain,
                 annot_1: Domain):
        """
        Object initialisation (creation of the class instance)
        the functions assumes that has a line from hmmsearch
        and has to be told otherwise if hmmscan was used
        """
        self.prot_id = prot_id
        self.prot_start, self.prot_end = start, end
        self.parent_annotations = {annot_0, annot_1}

    def __repr__(self) -> str:
        """
        Overlap is represented as its position in domain and HMMs it align with.
        :return: 'Q9T132 [12-143] (Amidase_2, PGRP)'
        """
        return f'{self.prot_id} [{self.prot_start} - {self.prot_end}] ({self.parent_annotations})'

    def __eq__(self, other) -> bool:
        """
        Overlap is equivalent to another overlap if they came from the same annotation pair
        :param other: object to compare
        :return: are compared objects identical?
        """
        if isinstance(other, Overlap):
            return self.parent_annotations == other.parent_annotations
        return False

    def __hash__(self):
        """
        Overlap should be hashed based on location of the parent annotations
        :return: hash number of the domain
        """
        return hash((tuple(sorted(self.parent_annotations, key=lambda dom: dom.prot_start)), self.prot_start, self.prot_end))

    def seq_length(self) -> int:
        """
        Calculate length of overlap in AA
        :return: length of overlap in AA
        """
        return self.prot_end - self.prot_start + 1


@checkpoint
def read_hmmer(fasta_path: PathLike,
               annotation_path: PathLike,
               program,
               min_score: int = 20,
               pickle_path: Path = None) -> AnnotationTrack:
    """
    Read fasta and raw annotation file from HMMer3
    and return corresponding track object
    :param fasta_path: path to the fasta file with annotated sequences
    :param annotation_path: path to the HMMer3 domtblout file
    :param program: was annotation file generated using hmmscan or hmmsearch
    :param min_score: keep only domains with HMMer3 score above this value
    :param pickle_path: path to a temporary file storing results in case of a crash or re-run with modified parameters
    :return: AnnotationTrack representing parsed annotations
    """
    annotation_path = Path(annotation_path)

    all_domains = AnnotationTrack()
    all_domains._import_fasta(fasta_path)

    print(f'Reading input file at: {annotation_path.as_posix()}')
    with annotation_path.open() as inpt:
        discarded, kept = 0, 0
        for line in inpt:
            if line.strip() and not line.startswith('#'):
                domain = Domain.from_hmmer(line,
                                           program=program)
                if domain.score > min_score:
                    all_domains.add(domain)
                    kept += 1
                else:
                    discarded += 1
            else:
                print(f'Line skipped: {line}')

    print(f'Score filter: {min_score} ({kept} hits kept, {discarded} discarded)')
    all_domains.finalize()
    all_domains.merge_redundant()

    return all_domains


def fisher(common_positions, positions_hmm0, positions_hmm1, all_positions):
    """
    Scores bigrams using Fisher's Exact Test (Pedersen 1996).  Less
    sensitive to small counts than PMI or Chi Sq, but also more expensive
    to compute. Requires scipy. NLTK CORRECTION
    """

    n_ii, n_io, n_oi, n_oo = _contingency(common_positions, positions_hmm0, positions_hmm1, all_positions)

    (odds, pvalue) = fisher_exact([[n_ii, n_io], [n_oi, n_oo]], alternative="greater")
    return pvalue


def _contingency(common, a, b, total):
    """
    Construct NLTK-formatted contingency table based on number of:
    :param common: observations common to both groups (e.g. number of positions/domains/proteins where two HMMs overlap)
    :param a: observations connected with group 'a' (e.g. total number of positions/domains/proteins aligned to HMM1)
    :param b: observations connected with group 'b' (e.g. total number of positions/domains/proteins aligned to HMM2)
    :param total: total common observations (e.g. positions/domains/proteins)
    :return: contingency table suitable for NLTK fisher_exact test
    """
    n_ii = common
    n_xx = total
    n_oi = b - common
    n_io = a - common
    return n_ii, n_oi, n_io, n_xx - n_ii - n_oi - n_io


def bonferroni_holm_with_missing(pvals, n_test) -> np.ndarray:
    """
    Perform Bonferroni–Holm correction for comparisons on the p-value array with missing values.
    Used to save time, inc situation where some test heuristically filered out
    and their p-values were never calculated.
    :param pvals: p-value array with missing values
    :param n_test: real number of tests
    :return: array of p-values corrected for the real number of comparisons
    """
    pv_array = np.array(pvals)
    ranks = pv_array.argsort().argsort()
    corrections = np.full(len(pv_array), n_test) - ranks
    return np.clip(pv_array * corrections, 0, 1)


def jaccard_and_containment(n_common,
                            n_a,
                            n_b) -> Tuple[float, float]:
    """
    Calculate Jaccard similarity coefficient and Jaccard containment for two sets of elements
    e.g. positions/domains/proteins where two HMMs overlap.
    for details see https://www.ncbi.nlm.nih.gov/pmc/articles/PMC6439793/
    :param n_common: number of common elements in two sets (e.g. positions/domains/proteins where two HMMs overlap)
    :param n_a: elements connected with group 'a' (e.g. positions/domains/proteins aligned to HMM1)
    :param n_b: elements connected with group 'b' (e.g. positions/domains/proteins aligned to HMM2)
    :return: value of the Jaccard index and Jaccard containment of 'a' in 'b'
    """
    jaccard_index = n_common / (n_a + n_b - n_common)
    containment_index = n_common / n_a
    return jaccard_index, containment_index


Summary = namedtuple('Summary',
                     ['hmm_proteins',
                      'hmm_counts',
                      'hmm_spans',
                      'hmm_protein_fractions',
                      'shared_proteins',
                      'overlap_counts',
                      'overlap_spans',
                      'overlap_fractions',
                      'n_all_proteins',
                      'n_all_aa',
                      'n_all_domains'])  # Simplified representation of annotation track used to generate overlap network


@checkpoint
def summarize_overlaps(annotation_track: AnnotationTrack,
                       positive_hmms: Set[str],
                       pickle_path: Path = None) -> Summary:
    """
    Annotate overlaps on a given annotation track, then summarize domains and
    :param annotation_track: annotations to analyze
    :param positive_hmms: models of interest that will be the backbone of the network (e.g. models connected to particular enzymatic activity)
    :param pickle_path: path to a temporary file storing results in case of a crash or re-run with modified parameters
    :return: simplified summary of annotation track used to generate overlap network
    """

    annotation_track.find_overlaps_with_hmms(positive_hmms)

    hmm_proteins, hmm_counts, hmm_spans, hmm_protein_fractions = annotation_track.summary()
    shared_proteins, overlap_counts, overlap_spans, overlap_fractions = annotation_track.overlap_summary()

    n_all_proteins, n_all_aa = annotation_track.n_sequences, annotation_track.n_positions
    n_all_domains = len(annotation_track)

    result = Summary(hmm_proteins=hmm_proteins,
                     hmm_counts=hmm_counts,
                     hmm_spans=hmm_spans,
                     hmm_protein_fractions=hmm_protein_fractions,
                     shared_proteins=shared_proteins,
                     overlap_counts=overlap_counts,
                     overlap_spans=overlap_spans,
                     overlap_fractions=overlap_fractions,
                     n_all_proteins=n_all_proteins,
                     n_all_aa=n_all_aa,
                     n_all_domains=n_all_domains)

    return result


def overlap_network(track_summary: Summary,
                    min_count_threshold: int = 3,
                    containment_threshold: float = 0.75) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]:
    """
    Generate edge and node dictionaries that represents network of HMM-HMM similarities
    generated by analysis of their co-occurrence in protein sequences making up the annotation track
    :param track_summary: simplified versio of the annotation track (domains found in multiple sequences)
    :param min_count_threshold: network edge is created only, if two HMMs overlap in at least this number of sequels
    :param containment_threshold: network edge is created only, if at lest this fraction of domains aligned to one HMM
                                  is enveloped in domains aligned to the other
                                  (Jaccard containment of proteins, domains and positions reaches this threshold)
    :return: list of dictionaries representing statistics of edges and nodes
    """

    ts = track_summary

    with tqdm(total=len(ts.hmm_proteins)) as bar:
        edge_records = []
        node_records = []
        for hmm0 in ts.hmm_proteins:
            donor_id = hmm0
            donor = {'ID': donor_id,
                     'proteins.n': ts.hmm_proteins[donor_id],
                     'domains.n': sum(ts.hmm_counts[donor_id]),
                     'domains.median': np.median(ts.hmm_counts[donor_id]),
                     'domains.std': np.std(ts.hmm_counts[donor_id]),
                     'positions.n': sum(ts.hmm_spans[donor_id]),
                     'positions.median': np.median(ts.hmm_spans[donor_id]),
                     'positions.std': np.std(ts.hmm_spans[donor_id]),
                     'protein_fraction.median': np.median(ts.hmm_protein_fractions[donor_id]),
                     'protein_fraction.std': np.std(ts.hmm_protein_fractions[donor_id])}
            donor['domains.fraction'] = donor['domains.n'] / ts.n_all_domains
            donor['positions.fraction'] = donor['positions.n'] / ts.n_all_aa
            node_records.append(donor)
            if donor_id in ts.overlap_counts:
                donor_shared_proteins = ts.shared_proteins[donor_id]
                donor_domains, donor_positions, donor_fractions = ts.overlap_counts[donor_id], ts.overlap_spans[donor_id], ts.overlap_fractions[donor_id]
                for hmm1 in donor_domains:
                    acceptor_id = hmm1
                    pair_shared_prots = donor_shared_proteins[acceptor_id]
                    pair_domains, pair_positions, pair_fractions = donor_domains[acceptor_id], donor_positions[acceptor_id], donor_fractions[acceptor_id]
                    if pair_shared_prots >= min_count_threshold and sum(pair_domains) >= min_count_threshold:
                        acceptor_proteins, acceptor_domains, acceptor_positions = ts.hmm_proteins[acceptor_id], len(ts.hmm_counts[acceptor_id]), sum(ts.hmm_spans[acceptor_id])
                        edge = {'donor': donor_id,
                                'acceptor': acceptor_id,
                                'proteins.k': pair_shared_prots,
                                'domains.k': sum(pair_domains),
                                'domains.uniq_k': len(pair_domains),
                                'domains.median': np.median(pair_domains),
                                'domains.std': np.std(pair_domains),
                                'positions.k': sum(pair_positions),
                                'positions.median': np.median(pair_positions),
                                'positions.std': np.std(pair_positions),
                                'donor_fraction.median': np.median(pair_fractions),
                                'donor_fraction.std': np.std(pair_fractions)}
                        edge['proteins.jaccard'], edge['proteins.containment'] = jaccard_and_containment(edge['proteins.k'], donor['proteins.n'], acceptor_proteins)
                        edge['domains.jaccard'], edge['domains.containment'] = jaccard_and_containment(edge['domains.uniq_k'], donor['domains.n'], acceptor_domains)
                        edge['positions.jaccard'], edge['positions.containment'] = jaccard_and_containment(edge['positions.k'], donor['positions.n'], acceptor_positions)
                        if edge['proteins.containment'] >= containment_threshold and edge['domains.containment'] >= containment_threshold and edge['positions.containment'] >= containment_threshold:
                            edge_records.append(edge)
            bar.set_description(f'Found {len(edge_records)} edges')
            bar.update()

    return edge_records, node_records


if __name__ == '__main__':

    # thresholds

    min_count = 3  # at least this number of proteins with overlap is required to create an edge
    min_containment = 0.75  # smallest containment value required to create an (directed) edge
    score_threshold = 20  # only domains with HMMer3 score above this value are considered during analysis
    """
    https://www.biorxiv.org/content/10.1101/2021.06.24.449764v2.full
    Searching sequence databases for functional homologs using profile HMMs: how to set bit score thresholds?
    (...). Bit scores were used as thresholds rather than E-values since they remain the same irrespective of the size of the database searched. 
    """

    # input and output files

    fasta_file = Path(' (...) /proteins.fasta')
    annotation_file = Path(' (...) /hmmer_output.domtblout')
    hmm_details_table = Path(' (...) /HMM_info.tsv')  # table with names and status of each HMM (e.g. whether HMM was selected as positive) # todo reformat
    output_dir = Path(' (...) /results')

    # parsing inputs

    hmm_details_table = pd.read_table(hmm_details_table, index_col='ID')
    pos_hmms = set(hmm_details_table.index[(hmm_details_table['TYPE'] == 'positive') | (hmm_details_table['TYPE'] == 'minimal')])
    track = read_hmmer(fasta_path=fasta_file,
                       annotation_path=annotation_file,
                       program='hmmsearch',
                       min_score=score_threshold,
                       pickle_path=output_dir.joinpath('AnnotationTrack.tmp'))

    # prepare directory and run overlap search

    output_dir.mkdir(parents=True)

    track_summary = summarize_overlaps(track,
                                       positive_hmms=pos_hmms,
                                       pickle_path=output_dir.joinpath('OverlapSummary.tmp'))

    edges, nodes = overlap_network(track_summary=track_summary,
                                   min_count_threshold=min_count,
                                   containment_threshold=min_containment)

    # write output files

    edge_table = pd.DataFrame.from_records(edges)
    edge_file = output_dir.joinpath('Edge_table.csv')
    edge_table.to_csv(edge_file.as_posix(), index=False)

    node_table = pd.DataFrame.from_records(nodes, index='ID')
    for column in hmm_details_table.columns:
        edge_table[column] = hmm_details_table[column]  # add HMM metadata (model names etc.) to the node table
    node_file = output_dir.joinpath('Node_table.csv')
    node_table.to_csv(node_file.as_posix())