import pandas as pd
import pysam
from celescope.tools import utils
from celescope.tools.step import Step, s_common
from collections import defaultdict
from celescope.__init__ import HELP_DICT


def count_pathseq(args):
    with Count_pathseq(args, display_title="Mapping to Microbe") as runner:
        runner.run()


class Count_pathseq(Step):
    def __init__(self, args, display_title=None):
        Step.__init__(self, args, display_title=display_title)
        # match
        if args.match_dir:
            self.match_dict = utils.parse_match_dir(args.match_dir)
            self.match_barcode = self.match_dict["match_barcode"]
        else:
            self.match_barcode, _ = utils.get_barcode_from_matrix_dir(args.matrix_dir)

        # out
        self.raw_matrix_file = f"{self.outdir}/{self.sample}_raw_UMI_matrix.tsv.gz"
        self.filtered_matrix_file = (
            f"{self.outdir}/{self.sample}_filtered_UMI_matrix.tsv.gz"
        )
        self.outs = [self.raw_matrix_file, self.filtered_matrix_file]

        # metrics
        self.total_reads = 0
        self.cell_reads = 0
        self.cell_microbe_reads = 0
        self.cell_microbe_umi = 0
        self.n_unambiguous_umi = 0

    @utils.add_log
    def pathseq_score_to_dict(self):
        """
        Returns: {taxID: genus_name}
        {
        '1703962': 'Rhizobium',
        '1720346': 'Rhizobium',
        }
        """
        pathseq_score = open(self.args.pathseq_score_file, "r")
        genus_dict = {}
        set_for_genera = set()
        for each_line in pathseq_score:
            each_line = each_line.rstrip("\n")
            each_line_list = each_line.split("\t")
            level = each_line_list[2]
            tax = each_line_list[3]
            if level == "genus":
                set_for_genera.add(tax)
            if "|" in each_line_list[1]:
                name_string_list = each_line_list[1].split("|")
                for n in range(len(name_string_list)):
                    pointer = -n - 1  ## -1,-2,-3
                    if "_" not in name_string_list[pointer]:
                        name = name_string_list[pointer]
                        break
                    if "unclassified" in name_string_list[pointer]:
                        name = name_string_list[pointer]
                        break
                id = each_line_list[0]
                genus_dict[id] = name
        print("len(dict_for_genus) = ", len(genus_dict))
        return genus_dict

    @utils.add_log
    def parse_unmap_bam(self) -> dict[str, dict[str, str]]:
        """
        returns: {read_name: {CB: CB, UB: UB}}
        """
        unmap_bam = pysam.AlignmentFile(self.args.unmap_bam_file, threads=self.thread)
        unmap_read_dict = {}
        for segment in unmap_bam:
            self.total_reads += 1
            if segment.has_tag("CB") and segment.has_tag("UB"):
                read_name, CB, UB = (
                    segment.query_name,
                    segment.get_tag("CB"),
                    segment.get_tag("UB"),
                )
                unmap_read_dict[read_name] = {
                    "CB": CB,
                    "UB": UB,
                }
        self.cell_reads = len(unmap_read_dict)
        self.add_metric(
            name="Number of Reads in Cells",
            value=self.cell_reads,
            total=self.total_reads,
            help_info="Number of reads with cell barcode match with scRNA-Seq library.",
        )
        return unmap_read_dict

    @utils.add_log
    def parse_pathseq_bam(self, unmap_read_dict: dict[str, dict[str, str]]):
        """
        YP: taxIDs
        AS: Alignment score generated by pathseq

        returns: {read_name: {YP: taxIDs, AS: alignment_score}}
        """
        pathseq_bam = pysam.AlignmentFile(
            self.args.pathseq_bam_file, "rb", threads=self.thread
        )
        read_dict = {}
        for segment in pathseq_bam:
            if segment.has_tag("YP"):
                read_name, YP, AS = (
                    segment.query_name,
                    segment.get_tag("YP"),
                    int(segment.get_tag("AS")),
                )
                if read_name not in unmap_read_dict:
                    continue
                self.cell_microbe_reads += 1
                read_dict[read_name] = {
                    "YP": YP,
                    "AS": AS,
                }
        self.add_metric(
            name="Number of Cell Reads mapping to Microbe",
            value=self.cell_microbe_reads,
            total=self.cell_reads,
            help_info="Number of reads in cells with YP tag in the BAM file. The YP read tag lists the NCBI taxonomy IDs of any aligned species meeting the alignment identity criteria.",
        )
        return read_dict

    @utils.add_log
    def select_highest_AS(self, read_dict, unmap_read_dict: dict[str, dict[str, str]]):
        """
        collapse read with same (CB,UB), select the one with highest AS
        returns: {CB, UB: {YP: YP, AS: AS}
        """
        read_highest_AS = {}
        for read_name in read_dict:
            CB, UB = unmap_read_dict[read_name]["CB"], unmap_read_dict[read_name]["UB"]
            if (CB, UB) not in read_highest_AS or read_dict[read_name][
                "AS"
            ] > read_highest_AS[(CB, UB)]["AS"]:
                read_highest_AS[(CB, UB)] = read_dict[read_name]
        self.cell_microbe_umi = len(read_highest_AS)
        self.add_metric(
            name="Number of Cell UMIs Mapping to Microbe",
            value=self.cell_microbe_umi,
            help_info="Referring to the approach of InvadeSeq, merge the reads with the same barcode and UMI, and select the read with the highest mapping quality as the representative.",
        )
        return read_highest_AS

    @utils.add_log
    def select_unique_YP(self, read_highest_AS, genus_dict):
        cb_genus_umi = defaultdict(lambda: defaultdict(int))
        for CB, UB in read_highest_AS:
            YP = read_highest_AS[(CB, UB)]["YP"]
            genus = set([genus_dict[id] for id in YP.split(",")])
            if len(genus) == 1:
                genus = genus.pop()
                cb_genus_umi[CB][genus] += 1
                self.n_unambiguous_umi += 1
        self.add_metric(
            name="Number of Unambiguous UMIs",
            value=self.n_unambiguous_umi,
            total=self.cell_microbe_umi,
            help_info="If all YP tags of a UMI belong to the same genus, this UMI is defined as unambiguous UMIs. These unambiguous UMIs are used to create the expression matrix.",
        )
        self.saturation = 1 - self.cell_microbe_umi / self.cell_microbe_reads
        self.add_metric(
            name="Estimated Saturation",
            value=self.saturation,
            value_type="fraction",
            help_info="saturation = 1 - n_cell_microbe_UMIs / n_cell_microbe_reads",
        )
        return cb_genus_umi

    @utils.add_log
    def write_matrix(self, cb_genus_umi):
        """
        rows are genus and columns are cb
        """
        df_raw = pd.DataFrame.from_dict(cb_genus_umi, orient="index").transpose()
        df_raw = df_raw.fillna(0).astype(int)
        df_raw.to_csv(self.raw_matrix_file, sep="\t")

        df_filtered = df_raw.reindex(columns=self.match_barcode, fill_value=0)
        df_filtered = df_filtered.astype(int)
        df_filtered = df_filtered.loc[(df_filtered != 0).any(axis=1)]
        df_filtered.to_csv(self.filtered_matrix_file, sep="\t")

    def run(self):
        genus_dict = self.pathseq_score_to_dict()
        unmap_read_dict = self.parse_unmap_bam()
        read_dict = self.parse_pathseq_bam(unmap_read_dict)
        read_highest_AS = self.select_highest_AS(read_dict, unmap_read_dict)
        cb_genus_umi = self.select_unique_YP(read_highest_AS, genus_dict)
        self.write_matrix(cb_genus_umi)


def get_opts_count_pathseq(parser, sub_program):
    if sub_program:
        parser.add_argument(
            "--pathseq_bam_file", help="pathseq bam file", required=True
        )
        parser.add_argument(
            "--pathseq_score_file", help="pathseq score file", required=True
        )
        parser.add_argument("--unmap_bam_file", help="unmap bam file", required=True)
        parser.add_argument("--match_dir", help=HELP_DICT["match_dir"])
        parser.add_argument("--matrix_dir")
        s_common(parser)
    return parser
