#!/usr/bin/env python3

import sys
import re
import argparse
import logging


def get_options():
    parser = argparse.ArgumentParser(description="Parses a domtbl file generated by HMMER, applying its own set of"
                                                 "filters and summarizing those high-quality matches."
                                                 "Optionally, sequences of the high-quality matches can also be written"
                                                 "to a FASTA file.", add_help=False)
    required_args = parser.add_argument_group("Required arguments")
    required_args.add_argument("-i",
                               dest="input",
                               help="Input domain table from HMMER to be parsed",
                               required=True)

    opt_args = parser.add_argument_group("Optional arguments")
    opt_args.add_argument("-f",
                          dest="fasta_in", required=False,
                          help="Path to a FASTA file containing the sequence that were searched against the HMM(s). "
                               "If one isn't provided, the script just prints summary stats.")
    opt_args.add_argument("-o",
                          dest="output", required=False, default="hmm_purified.fasta",
                          help="The name of the FASTA file to write containing sequences of the high-quality hits. "
                               "[default=hmm_purified.fasta]")
    opt_args.add_argument("-p",
                          dest="perc_aligned",
                          type=int,
                          default=90,
                          help="The minimum percentage of the HMM that was covered by the target sequence (ORF) "
                               "for the COG hit to be included [default=90]")
    opt_args.add_argument("-e",
                          dest="min_e",
                          type=float,
                          default=0.0001,
                          help="The largest E-value for the search to be accepted as significant [default=1E-3]")
    opt_args.add_argument("-a",
                          dest="min_acc",
                          type=float,
                          default=0.90,
                          help="The minimum acc threshold of the HMM search for reporting [default=0.90]")
    miscellaneous_opts = parser.add_argument_group("Miscellaneous options")
    miscellaneous_opts.add_argument("-v", "--verbose", action='store_true', default=False,
                                    help='Prints a more verbose runtime log')
    miscellaneous_opts.add_argument("-h", "--help",
                                    action="help",
                                    help="Show this help message and exit")
    args = parser.parse_args()
    return args


def format_hmmer_domtbl_line(line):
    stats = []
    stat = ""
    for c in line:
        if c == ' ':
            if len(stat) > 0:
                stats.append(str(stat))
                stat = ""
            else:
                pass
        else:
            stat += c
    stats.append(str(stat))
    return stats


class HmmSearchStats:
    def __init__(self):
        self.raw_alignments = 0
        self.seqs_identified = 0
        self.fragmented = 0
        self.inverted = 0
        self.glued = 0
        self.dropped = 0
        self.bad = 0
        self.short = 0
        self.multi_alignments = 0  # matches of the same query to a different HMM (>1 lines)

    def num_dropped(self):
        self.dropped = self.inverted + self.bad + self.short
        return self.dropped

    def summarize(self):
        alignment_stat_string = ""
        alignment_stat_string += "\tInitial alignments:\t" + str(self.raw_alignments) + "\n"
        alignment_stat_string += "\tAlignments discarded:\t" + str(self.dropped) + "\n"
        alignment_stat_string += "\tFragmented alignments:\t" + str(self.fragmented) + "\n"
        alignment_stat_string += "\tInversions detected:\t" + str(self.inverted) + "\n"
        alignment_stat_string += "\tAlignments scaffolded:\t" + str(self.glued) + "\n"
        alignment_stat_string += "\tMulti-alignments:\t" + str(self.multi_alignments) + "\n"
        alignment_stat_string += "\tSequences identified:\t" + str(self.seqs_identified) + "\n"
        return alignment_stat_string


class HmmMatch:
    def __init__(self):
        self.genome = ""  # Name of the input file (Metagenome, SAG, MAG, or isolate genome)
        self.target_hmm = ""  # Name of the HMM aligned to
        self.orf = ""  # Name of the ORF, or more generally contig sequence
        self.hmm_len = 0  # Length of the hidden Markov model
        self.start = 0  # Alignment start position on the contig
        self.end = 0  # Alignment end position on the contig
        self.pstart = 0  # Alignment start position on the hmm profile
        self.pend = 0  # Alignment end position on the hmm profile
        self.seq_len = 0  # Length of the query sequence
        self.num = 0
        self.of = 0
        self.desc = ""
        self.acc = 0.0
        self.ieval = 0.0
        self.eval = 0.0
        self.full_score = 0
        self.next_domain = None  # The next domain aligned by hmmsearch

    def get_info(self):
        info_string = "Info for query " + str(self.orf) + ":\n"
        info_string += "\tHMM = " + self.target_hmm + ", length = " + str(self.hmm_len) + "\n"
        info_string += "\tSequence length = " + str(self.seq_len) + "\n"
        info_string += "\tAligned length = " + str(self.end - self.start) + "\n"
        info_string += "\tAlignment start = " + str(self.start) + ", alignment stop = " + str(self.end) + "\n"
        info_string += "\tProfile start = " + str(self.pstart) + ", profile stop = " + str(self.pend) + "\n"
        info_string += "\tNumber " + str(self.num) + " of " + str(self.of) + "\n"
        info_string += "\tcE-value = " + str(self.ieval) + "\n"
        info_string += "\tacc = " + str(self.acc) + "\n"
        info_string += "\tfull score = " + str(self.full_score) + "\n"
        return info_string

    def subsequent_matches(self):
        if not self.next_domain:
            return [self]
        return [self] + self.next_domain.subsequent_matches()

    def collinear(self) -> bool:
        """
        This checks the profile-start and profile-stop positions against the other alignment fragments to ensure
        they are the smallest (derived from the left-most position of the HMM profile).

        :return: Boolean representing whether all domains/alignment fragments are collinear (True) or not (False)
        """
        if not self.next_domain:
            return True
        if self.next_domain.collinear():
            if self.pend < self.next_domain.pend and self.pstart < self.next_domain.pstart:
                return True
        return False

    def contains_duplicate_loci(self, overlap_threshold=0.33) -> bool:
        if not self.next_domain:
            return False
        if self.next_domain.contains_duplicate_loci():
            return True
        for subsequent_match in sorted(self.subsequent_matches(), key=lambda x: x.num):
            if self.num != subsequent_match.num:
                # Check for redundant profile HMM coverage
                p_overlap_len = overlap_length(self.pstart, self.pend,
                                               subsequent_match.pstart, subsequent_match.pend)
                if p_overlap_len > (overlap_threshold*self.hmm_len):
                    return True
            else:
                pass
        return False

    def enumerate(self) -> None:
        i = 1
        next_matches = sorted(self.subsequent_matches(), key=lambda x: x.num)
        of = len(next_matches)
        for hmm_match in next_matches:  # type: HmmMatch
            hmm_match.num = i
            hmm_match.of = of
            i += 1
        return

    def drop_next_alignment(self, index: int) -> None:
        """
        Function for removing an element from the linked list formed by HmmMatch.next_domain

        :param index: The index in the linked list that is to be dropped
        :return: None
        """
        if index == 0:
            logging.error("Unable to drop current HmmMatch from next_domain linked list.\n")
            sys.exit(9)

        fragmented_alignment_data = self.subsequent_matches()
        i = 1
        while i < len(fragmented_alignment_data):
            if i == index:
                try:
                    fragmented_alignment_data[i-1].next_domain = fragmented_alignment_data[i+1]
                except IndexError:
                    fragmented_alignment_data[i - 1].next_domain = None
                return
            i += 1
        logging.warning("Next HmmMatch at index %s was not dropped from linked list.\n" % index)
        return

    def merge_alignment_fragments(self, index: int) -> None:
        # Update the alignments that are being merged into the base/reference alignment
        merging_aln = self.subsequent_matches()[index]
        self.start = min([self.start, merging_aln.start])
        self.end = max([self.end, merging_aln.end])
        self.pstart = min([self.pstart, merging_aln.pstart])
        self.pend = max([self.pend, merging_aln.pend])
        self.ieval = min([self.ieval, merging_aln.ieval])

        self.drop_next_alignment(index)
        return

    def scaffold_domain_alignments(self, seq_length_wobble=1.2) -> None:
        """
        If one or more alignments do not completely redundantly cover the HMM profile,
        overlap or are within a few BPs of each other of the query sequence,
        and do not generate an alignment 120% longer than the HMM profile,
                                                    THEN
        merge the alignment co-ordinates, average the acc, Eval, cEval and make 'num' and 'of' reflect number of alignments
                                                Takes this:
        -------------
                    --------------
                                                                --------------------------------------
                                     ----------------
                                            and converts it to:
        ---------------------------------------------           --------------------------------------
        
        :param seq_length_wobble: The scalar threshold controlling the maximum size of a scaffold.
        Since we're dealing with HMMs, its difficult to estimate the sequence length variance so we're allowing for
        some 'wobble' in how long or short paralogs could be.
        :return: None
        """
        if not self.next_domain:
            return
        self.next_domain.scaffold_domain_alignments()

        if self.contains_inversion():
            return
        i = 1
        while i < len(self.subsequent_matches()):
            next_match = sorted(self.subsequent_matches(), key=lambda x: x.num)[i]  # type: HmmMatch
            if self.num == next_match.num:
                logging.error("Iteration of HmmMatch.subsequent_matches() begins at current instance (%s).\n" % self.num)
                sys.exit(9)
            # Check for sub- or super-sequence orientation on the query sequence
            q_overlap_len = overlap_length(self.start, self.end, next_match.start, next_match.end)
            # Check for redundant profile HMM coverage
            p_overlap_len = overlap_length(self.pstart, self.pend, next_match.pstart, next_match.pend)

            min_profile_covered = min([self.pend - self.pstart,
                                       next_match.pend - next_match.pstart])
            try:
                aln_overlap_proportion = p_overlap_len / min_profile_covered
            except ZeroDivisionError:
                if self.pend - self.pstart < next_match.pend - next_match.pstart:
                    self.drop_next_alignment(0)
                else:
                    self.drop_next_alignment(i)
                continue
            if aln_overlap_proportion > 0.5:
                # They overlap significant regions - are they overlapping sequence or are they repeats?
                if q_overlap_len == next_match.seq_len:
                    # The projected alignment is a subsequence of the base alignment - DROP
                    self.drop_next_alignment(i)
                    i -= 1
                elif q_overlap_len < p_overlap_len:
                    # The two alignments represent repeats of a single profile - KEEP
                    pass
                else:
                    # The two alignments represent something...
                    self.merge_alignment_fragments(i)
                    i -= 1
            else:
                # They do not significantly overlap in the profile, likely part of the same alignment
                a_new_start = min([self.start, next_match.start])
                a_new_end = max([self.end, next_match.end])
                if float(a_new_end - a_new_start) < float(seq_length_wobble * int(self.hmm_len)):
                    # The proximal alignments overlap so they would probably form a homologous sequence - MERGE
                    self.merge_alignment_fragments(i)
                    i -= 1
                else:
                    # The alignments are very far apart from each other and would form a monster sequence - KEEP
                    pass
            i += 1
        return

    def contains_inversion(self) -> bool:
        """
        Asserting this HmmMatch instance is the first alignment (and therefore the left-most alignment on the query)
        this checks for collinearity along the HMM profile and query sequences.
        If they are not collinear and additional operation to look for duplicated HMM profile regions in the query,
        indicating multiple domains, is performed.
        Otherwise the positions have been inverted, or in other words, a section of the gene has been rearranged.

        :return: Failing both tests for collinearity and duplication return True, else False
        """
        if self.of <= 1 or not self.next_domain:
            return False
        if self.collinear():
            return False
        elif self.contains_duplicate_loci():
            return False
        else:
            return True

    def orient_alignments(self) -> dict:
        alignment_relations = dict()
        if not self.next_domain:
            return alignment_relations
        alignment_relations.update(self.next_domain.orient_alignments())

        for next_match in self.subsequent_matches():
            if self.num != next_match.num:
                alignment_relations[(self.num, next_match.num)] = detect_orientation(self.start, self.end,
                                                                                     next_match.start, next_match.end)
        return alignment_relations

    def consolidate_subalignments(self, alignment_relations: dict) -> dict:
        """
        What is the point of this function?
        Identify alignments that can be

        :param alignment_relations:
        :return: Dictionary of HmmMatch instances indexed by their respective query names
        """
        alignments_for_disposal = set()
        scaffolded_alignments = dict()
        for pair in alignment_relations:
            # If there are multiple alignments that span the whole hmm profile, report them both
            base, projected = pair
            if alignment_relations[pair] == "satellite" or alignment_relations[pair] == "overlap":
                pass
            elif alignment_relations[pair] == "supersequence":
                alignments_for_disposal.add(projected)
            elif alignment_relations[pair] == "subsequence":
                alignments_for_disposal.add(base)
            else:
                logging.error("Unexpected alignment comparison state: '" + alignment_relations[pair] + "'\n")
                sys.exit(31)

        for hmm_match in self.subsequent_matches():
            # Filter out the worst of the overlapping alignments that couldn't be scaffolded
            if hmm_match.num not in alignments_for_disposal:
                query_header = ' '.join([hmm_match.orf, hmm_match.desc]) + \
                               '_' + str(hmm_match.num) + '_' + str(hmm_match.of)
                scaffolded_alignments[query_header] = hmm_match

        return scaffolded_alignments


class DomainTableParser(object):

    def __init__(self, dom_tbl):
        self.alignments = {}
        self.i = 0
        self.lines = []
        self.size = 0
        try:
            self.commentPattern = re.compile(r'^#')
            self.src = open(dom_tbl)
        except IOError:
            logging.error("Could not open " + dom_tbl + " or file is not available for reading.\n")
            sys.exit(0)

    def __iter__(self):
        return self

    def read_domtbl_lines(self):
        """
        Function to read the lines in the domain table file,
        skipping those matching the comment pattern
        :return: self.lines is a list populated with the lines
        """
        line = self.src.readline()
        while line:
            comment = self.commentPattern.match(line)
            if not comment:
                self.lines.append(line.strip())
            if not line:
                break
            line = self.src.readline()
        self.size = len(self.lines)

    def next(self):
        """
        Reformat the raw lines of the domain table into
        an easily accessible hmm_domainTable format and perform
        QC to validate the significance of the alignments
        """
        if self.i < self.size:
            hit = format_hmmer_domtbl_line(self.lines[self.i])
            self.prepare_data(hit)
            self.i += 1

            try:
                return self.alignments
            except ValueError:
                return None
        else:
            self.src.close()
            return None

    def prepare_data(self, hit):
        self.alignments['query'] = str(hit[0])
        self.alignments['query_len'] = int(hit[2])
        self.alignments['hmm_name'] = str(hit[3])
        self.alignments['hmm_len'] = int(hit[5])
        self.alignments['Eval'] = float(hit[6])  # Full-sequence E-value (in the case a sequence alignment is split)
        self.alignments['full_score'] = float(hit[7])  # Full-sequence score
        self.alignments['num'] = int(hit[9])  # HMMER is able to detect whether there are multi-hits
        self.alignments['of'] = int(hit[10])  # This is the number of multi-hits for a query
        self.alignments['cEval'] = float(hit[11])  # conditional E-value
        self.alignments['iEval'] = float(hit[12])  # conditional E-value
        self.alignments['pstart'] = int(hit[15])  # First position on HMM profile
        self.alignments['pend'] = int(hit[16])  # Last position on HMM profile
        self.alignments['qstart'] = int(hit[19])  # env coord from
        self.alignments['qend'] = int(hit[20])  # env coord to
        self.alignments['acc'] = float(hit[21])
        self.alignments['desc'] = ' '.join(hit[22:])


def detect_orientation(q_i: int, q_j: int, r_i: int, r_j: int):
    """
    Returns the number of positions the query (base) alignment overlaps with the reference (projected) alignment
    :param q_i: query start position
    :param q_j: query end position
    :param r_i: reference start position
    :param r_j: reference end position
    :return:
    """
    if q_i <= r_i <= q_j:
        if q_i <= r_j <= q_j:
            return "supersequence"
        else:
            return "overlap"
    elif r_i <= q_i <= r_j:
        if r_i <= q_j <= r_j:
            return "subsequence"
        else:
            return "overlap"
    else:
        return "satellite"


def overlap_length(r_i: int, r_j: int, q_i: int, q_j: int) -> int:
    """
    Returns the number of positions the query (base) alignment overlaps with the reference (projected) alignment

    :param q_i: query start position
    :param q_j: query end position
    :param r_i: reference start position
    :param r_j: reference end position
    :return: Number of positions the two alignments overlap
    """
    if r_j < q_i or q_i > r_j:
        # Satellite alignments
        return 0
    else:
        return min(r_j, q_j) - max(r_i, q_i)


def assemble_domain_aligments(first_match: HmmMatch, search_stats: HmmSearchStats):
    distinct_alignments = dict()
    if first_match.next_domain:
        # STEP 1: Scaffold the alignments covering overlapping regions on the query sequence
        frags_pre_scaffolding = len(first_match.subsequent_matches())
        first_match.scaffold_domain_alignments()
        if first_match.contains_inversion():
            search_stats.inverted += 1
            return distinct_alignments
        first_match.enumerate()
        search_stats.glued += (frags_pre_scaffolding - len(first_match.subsequent_matches()))

        # STEP 2: Determine the order and orientation of the alignments
        # alignment_relations = orient_alignments(first_match.subsequent_matches())
        alignment_relations = first_match.orient_alignments()

        # STEP 3: Decide what to do with the fragmented alignments: join or split?
        distinct_alignments = first_match.consolidate_subalignments(alignment_relations)
        first_match.enumerate()
    return distinct_alignments


def format_split_alignments(domain_table, search_stats: HmmSearchStats):
    """
    Handles the alignments where 'of' > 1
    If the alignment covers the whole target HMM or if the distance between the two parts of the alignment
    are very far apart, then the alignment will be divided into two unrelated alignments
    If the alignment parts are near together and/or each part covers a portion of the HMM, then they will be joined
    :param domain_table: DomainTableParser() object
    :param search_stats: HmmSearchStats() instance containing accumulators to track alignment parsing
    :return:
    """
    # Dictionary of single sequence alignments to return
    distinct_alignments = dict()

    # Query-relevant parsing variables
    first_match = previous_match = HmmMatch()
    previous_query_header = ""
    while domain_table.next():
        data = domain_table.alignments
        hmm_match = HmmMatch()
        hmm_match.target_hmm = data['hmm_name']
        hmm_match.hmm_len = data['hmm_len']
        hmm_match.seq_len = data['query_len']
        hmm_match.orf = data['query']
        hmm_match.desc = data['desc']
        hmm_match.start = data['qstart']
        hmm_match.end = data['qend']
        hmm_match.pstart = data['pstart']
        hmm_match.pend = data['pend']
        hmm_match.num = data['num']
        hmm_match.of = data['of']
        hmm_match.acc = data['acc']  # Used for filtering
        hmm_match.ieval = data['iEval']  # Used for filtering
        hmm_match.eval = data['Eval']  # Used for filtering
        hmm_match.full_score = data['full_score']  # Used for filtering

        search_stats.raw_alignments += 1

        query_header = ' '.join([hmm_match.orf, hmm_match.desc])
        # Finish off "old business" (sub-alignments)
        if previous_match.orf != hmm_match.orf and first_match.orf == previous_match.orf:
            distinct_alignments.update(assemble_domain_aligments(first_match, search_stats))
        if hmm_match.target_hmm != previous_match.target_hmm and query_header == previous_query_header:
            # New HMM (target), same ORF (query)
            search_stats.multi_alignments += 1

        # Carry on with this new alignment
        query_header_desc_aln = ' '.join([hmm_match.orf, hmm_match.desc]) + \
                                '_' + str(hmm_match.num) + '_' + str(hmm_match.of)
        if not hmm_match.orf:
            logging.error("Double-line parsing encountered: hmm_match.orf is empty!\n")
            sys.exit(9)

        if data["of"] == 1:
            distinct_alignments[query_header_desc_aln] = hmm_match
        elif hmm_match.num == 1:
            search_stats.fragmented += 1
            first_match = hmm_match
        else:
            search_stats.fragmented += 1
            previous_match.next_domain = hmm_match

        previous_query_header = query_header
        previous_match = hmm_match

    # Check to see if the last alignment was part of multiple alignments, just like before
    if first_match.next_domain and first_match.orf == previous_match.orf:
        distinct_alignments.update(assemble_domain_aligments(first_match, search_stats))

    return distinct_alignments


def filter_poor_hits(args, distinct_alignments, search_stats: HmmSearchStats):
    """
    Filters the homology matches based on their E-values and mean posterior probability of aligned residues from
    the maximum expected accuracy (MEA) calculation.
    Takes into account multiple homology matches of an ORF to a single gene and determines the total length of the
    alignment instead of treating them as individual alignments. This information is used in the next filtering step.
    """
    min_acc = float(args.min_acc)
    min_score = float(args.min_score)
    min_ie = float(args.min_ie)
    min_e = float(args.min_e)

    purified_matches = dict()

    for query_header_desc_aln in sorted(distinct_alignments):
        hmm_match = distinct_alignments[query_header_desc_aln]

        query_header_desc = (hmm_match.orf, hmm_match.desc)
        if query_header_desc not in purified_matches:
            purified_matches[query_header_desc] = list()

        if hmm_match.eval <= min_e and hmm_match.ieval <= min_ie:
            if hmm_match.acc >= min_acc and hmm_match.full_score >= min_score:
                purified_matches[query_header_desc].append(hmm_match)
        else:
            search_stats.dropped += 1
            search_stats.bad += 1

    return purified_matches


def filter_incomplete_hits(args, purified_matches, search_stats: HmmSearchStats):
    complete_gene_hits = list()

    for query in purified_matches:
        for hmm_match in purified_matches[query]:
            ali_len = hmm_match.pend - hmm_match.pstart
            perc_aligned = (float((int(ali_len)*100)/int(hmm_match.hmm_len)))
            if perc_aligned >= args.perc_aligned:
                complete_gene_hits.append(hmm_match)
            else:
                search_stats.dropped += 1
                search_stats.short += 1

    return complete_gene_hits


def renumber_multi_matches(complete_gene_hits: list):
    orf_name_index = dict()
    # Find which names have duplicates
    for match in sorted(complete_gene_hits, key=lambda x: x.orf):  # type: HmmMatch
        if match.orf not in orf_name_index:
            orf_name_index[match.orf] = []
        orf_name_index[match.orf].append(match)
    # Rename the ORFs with more than one alignment
    for orf_name in orf_name_index:
        n = 1
        of = len(orf_name_index[orf_name])
        for match in sorted(orf_name_index[orf_name], key=lambda x: x.num):
            match.num = n
            match.of = of
            n += 1
    return
