# -*- coding:utf-8 -*-
"""
@author: hbs
@date: 2018-10-19
description:
    The output result of MACS and MeTPeak shows the location of a peak. The peak need to be annotated by the record in
    GTF format file.
    element_priority is the priority of the element which we focus on it. Take RNA as a case in point, RNA consists of
    stop codon, 3'utr, CDS, 5'utr and each element was granted a priority value. The less the value is, the higher
    priority it represents.
    Finally, each peak got only one record, it may not matches with RNA elements, or it matches one or a few elements.
    The match result with highest priority remains at last.
"""
import pandas as pd
from peak_reader import PeakReader
from gtf_handler import GtfReader
from interval_tree import IntervalTree
import warnings
import argparse
import sys
# avoid throwing out RecursiveError
sys.setrecursionlimit(1000000)


class PeakMapper:
    """
    mapping the peaks to the genome element.
    """
    def __init__(self, gtf_file, element_priority, element_distance, peak_file):
        self.gr = GtfReader(gtf_file, element_priority)
        self.pr = PeakReader(peak_file)
        self.it = IntervalTree()
        self.priority = element_priority
        self.distance = element_distance

    def build_gtf_tree(self):
        self.gr.load_gtf()
        return self.gr.gtf_tree

    def load_peak_data(self):
        self.pr.peak_info()
        return self.pr.peaks

    def peak_mapping(self, gtf_tree, peaks):
        """
        match one peak to the best transcript element
        :param gtf_tree: interval tree object which generated by gtf_file, type IntervalTree object
        :param peaks: peak records load from bed file, type pandas DataFrame
        :return:
        """
        peaks["element_name"], peaks["location"], peaks["element_start"], peaks["element_end"] = None, None, None, None
        res = peaks.apply(self.peak_location, axis=1, args=(gtf_tree, )).values
        s1, s2, s3, s4 = pd.Series([i[0] for i in res]), pd.Series([i[1] for i in res]), pd.Series([i[2] for i in res]), pd.Series([i[3] for i in res])
        peaks["element_name"], peaks["location"], peaks["element_start"], peaks["element_end"] = s4, s1, s2, s3
        peaks.dropna(inplace=True)
        return peaks

    def peak_location(self, peak_record, gtf_tree):
        chr_num = peak_record["chr"]
        strand = peak_record["strand"]
        peak_center = peak_record["peak_center"]
        try:
            interval = gtf_tree[chr_num][strand]["interval"]
            if peak_center not in range(int(interval[0]), interval[1]+1):
                return None, None, None, None
            gene_tree = gtf_tree[chr_num][strand]["gene_tree"]
            gene_tree_search = self.it.search(gene_tree.root, peak_center)

            # the place where peak center locates is not a gene
            if not gene_tree_search:
                return None, None, None, None
            # if the gene tree search result is not False, gene_tree_search is a list contains GeneNode object
            transcript_tree_search = []
            for g in gene_tree_search:
                transcript_tree = g.compose
                transcript_tree_search += self.it.search(transcript_tree.root, peak_center)
            # the place where peak center locates has no transcript
            if not transcript_tree_search:
                return None, None, None, None
            # if the tracript tree search result is not False, transcript_tree_search is a list
            # which contains TranscriptNode object
            element_tree_search = []
            for t in transcript_tree_search:
                element_tree = t.compose
                element_tree_search += self.it.search(element_tree.root, peak_center)
            # the place where peak center locates matchs no element
            if not element_tree_search:
                return None, None, None, None
            best_priority = float("inf")
            best_ratio = -1 * float("inf")
            element_start, element_end = 0, 0
            locate_element = None
            for e in element_tree_search:
                element_name = e.name
                element_priority = self.priority[element_name]
                if best_priority >= element_priority:
                    best_priority = element_priority
                    if strand == "+":
                        distance = peak_center - e.interval_start + 1
                    elif strand == "-":
                        distance = e.interval_end - peak_center + 1
                    ratio = distance / (e.interval_end - e.interval_start) * 100 + self.distance[element_name]
                    if best_ratio < ratio:
                        element_start = e.interval_start
                        element_end = e.interval_end
                        locate_element = element_name
                        best_ratio = ratio
            return best_ratio, element_start, element_end, locate_element

        except KeyError:
            return None, None, None


if __name__ == "__main__":
    warnings.filterwarnings("ignore")
    # parser = argparse.ArgumentParser(prog="peak_mapping", description="map the merip-seq, medip-seq peaks on "
    #                                                                   "genome elements")
    # parser.add_argument("-g", "--gtf", action="store", type=str, required=True, help="file path of GTF format file")
    # parser.add_argument("-p", "--priority", action="store", type=dict, required=True,
    #                     help="element priorities, type python dict. Key is element, value represent priority.")
    # parser.add_argument("-d", "--distance", action="store", type=dict, required=True, help="for future figuring")
    # parser.add_argument("-pk", "--peak",  action="store", type=str, required=True,
    #                     help="file path of BED format peak file")
    # parser.add_argument("-o", "--output", action="store", type=str, required=True, help="file path of output file")
    # args = parser.parse_args()
    # gtf = args.gtf
    # priority = args.priority
    # distance = args.distance
    # peak = args.peak
    # output = args.output
    #
    # pm = PeakMapper(gtf, priority, distance, peak)
    # tree = pm.build_gtf_tree()
    # peak = pm.load_peak_data()
    # peak = pm.peak_mapping(tree, peak)
    # peak.to_csv(output, sep="\t", index=False, mode="w", encoding="utf-8")

    # just test code
    pm = PeakMapper("/data/nanopore/merip_seq_data/metpeak_calling_res/peak_location/reference_files/simple.gtf",
                    {"stop_codon": 1, "three_prime_utr": 2, "CDS": 3, "five_prime_utr": 4},
                    {"stop_codon": 200, "three_prime_utr": 200, "CDS": 100, "five_prime_utr": 0},
                    "/data/nanopore/merip_seq_data/metpeak_calling_res/peak_location/reference_files/peak.bed")
    tree = pm.build_gtf_tree()
    peak = pm.load_peak_data()
    peak = pm.peak_mapping(tree, peak)
    peak.to_csv("/data/nanopore/merip_seq_data/metpeak_calling_res/peak_location/reference_files/peak_loc.tsv",
                sep="\t", index=False, mode="w", encoding="utf-8")



