#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2022/5/8 10:04 AM
u"""
Fetch CDS information from a gtf then mapping to genomic position
"""
from collections import defaultdict, namedtuple
from typing import Optional

from loguru import logger
from pysam import asGTF
from pysam.libctabix import Tabixfile

from sashimi.base.CoordinateMap import CoordinateMapper
from sashimi.base.GenomicLoci import GenomicLoci
from sashimi.base.Transcript import Transcript
from sashimi.base.pyUniprot import Uniprot


class CdsProtein(GenomicLoci):
    u"""
    Re-collect cds information
    """

    def __init__(
            self,
            chromosome: str,
            start: int,
            end: int,
            strand: str,
            cds: dict,
            gene_id: set,
            transcript_id: set,
            proxy: Optional[str] = None,
            timeout: int = 10
    ):
        u"""
        Re-iter the gtf and collection all CDS information from gtf, not truncated genomic region at here
        :param chromosome: the chromosome id of the  given genes
        :param start: the minima start site of the given genes
        :param end: the maxima start site of the given genes
        :param strand: strand of gene
        :param cds: a list of Transcript object
        :param gene_id: a set of gene_id
        :param transcript_id: a set of transcript_id
        """
        super().__init__(
            chromosome=chromosome,
            start=start,
            end=end,
            strand=strand)
        self.proxy = proxy
        self.timeout = timeout
        self.cds = cds
        self.pep = self.__init_protein__() if cds is not None else None
        self.gene_id = gene_id
        self.transcript_id = transcript_id

    def __len__(self):
        return len(self.pep)

    def __add__(self, other):
        assert isinstance(other, CdsProtein), "only CdsProtein and CdsProtein could be added"
        self.gene_id.update(other.gene_id)
        self.transcript_id.update(other.transcript_id)
        self.cds.update(other.cds)
        self.pep.update(other.pep)
        self.region = GenomicLoci(
            chromosome=self.chromosome,
            start=min(self.start, other.start),
            end=max(self.end, other.start),
            strand="."
        )

    # @property
    # def truncated_pep(self):
    #
    #     return

    @classmethod
    def __re_iter_gtf__(
            cls,
            gtf_tabix: Tabixfile,
            chromosome: str,
            transcript_id: set,
            gene_id: set,
            strand: str = "*",
            proxy: Optional[str] = None,
            timeout: int = 10
    ):
        u"""
        Generate the CdsProtein object
        :param gtf_tabix: a Tabix file object generated by pysam
        :param chromosome: the chromosome id of the given genes
        :param transcript_id: the set of transcript_id to fetch CDS region
        :param gene_id: the set of gene_id to fetch CDS region, and also used to stop the iteration
        :param strand: the strand of gene
        :return:
        """
        relevant_exon_iterator = gtf_tabix.fetch(reference=chromosome, parser=asGTF())

        boundary_for_iter = []

        cds_dict = defaultdict(list)

        for line in relevant_exon_iterator:

            if strand != "*" and line.strand != strand:
                continue
            if "transcript_id" not in line.keys():
                continue

            try:
                if len(boundary_for_iter) == len(gene_id) and \
                        line.end > max(boundary_for_iter):
                    break

                # update the lower boundary to stop the iter.
                if line.gene_id in gene_id:
                    boundary_for_iter.append(line.end)

                if line.transcript_id in transcript_id and \
                        line.feature == "CDS":
                    cds_dict[line.transcript_id].append(line)

            except KeyError as err:
                logger.warning(err)

        cds_contain = {}
        start_site = float('Inf')
        end_site = float('-Inf')

        if len(cds_dict) != 0:

            for current_trans_id in cds_dict.keys():
                current_exon_list = cds_dict[current_trans_id]
                current_start = min(map(lambda x: x.start + 1, current_exon_list))
                current_end = max(map(lambda x: x.end, current_exon_list))

                cds_contain[current_trans_id] = Transcript(
                    chromosome=chromosome,
                    start=current_start,
                    end=current_end,
                    strand=current_exon_list[0].strand,
                    exons=current_exon_list,
                    gene=current_exon_list[0].gene_name \
                        if "gene_name" in current_exon_list[0].keys() else "",
                    gene_id=current_exon_list[0].gene_id \
                        if "gene_id" in current_exon_list[0].keys() else "",
                    transcript=current_exon_list[0].transcript_name \
                        if "transcript_name" in current_exon_list[0].keys() else "",
                    transcript_id=current_exon_list[0].transcript_id \
                        if "transcript_id" in current_exon_list[0].keys() else ""
                )
                start_site = current_start if current_start < start_site else start_site
                end_site = current_end if current_end > end_site else end_site

            return cls(
                chromosome=chromosome,
                start=start_site,
                end=end_site,
                strand=strand,
                cds=cds_contain,
                gene_id=gene_id,
                transcript_id=transcript_id,
                proxy=proxy,
                timeout=timeout
            )

        return None

    def __init_protein__(self):
        u"""
        Initiate the CdsProtein object
        :return:
        """

        protein_info = defaultdict(list)
        Domain_region = namedtuple(
            "DomainGenomicRegion",
            ["category", "type", "description", "unique_id", "start", "end"]
        )

        for current_transcript_id in self.cds.keys():
            current_pep = Uniprot(
                uniprot_id=current_transcript_id,
                cds_len=len(self.cds[current_transcript_id]),
                proxy=self.proxy, timeout=self.timeout
            )

            if not current_pep.domain:
                protein_info[current_transcript_id] = []
                continue

            coord_mapper = CoordinateMapper(
                coordinates_list=self.cds[current_transcript_id].exon_list,
                strand=self.cds[current_transcript_id].strand
            )

            domain_res = defaultdict(list)
            for domain in current_pep.domain:
                current_unique_id = domain.unique_id
                pep_genomic_coord = coord_mapper.pep_to_cds(
                    int(domain.begin),
                    int(domain.end)
                )

                # keep the protein genomic coordinate into a list
                current_domain_res = []
                for sub_domain_start, sub_domain_end in pep_genomic_coord.se:
                    current_domain_res.append(
                        Domain_region._make(
                            [domain.category,
                             domain.type,
                             domain.description,
                             current_unique_id,
                             sub_domain_start,
                             sub_domain_end
                             ]
                        )
                    )
                domain_res[current_unique_id].extend([current_domain_res])

            for domain_unique_id, domain_list in domain_res.items():
                start_site = min([
                    min(map(lambda x: x.start, i)) for i in domain_list
                ])

                end_site = max([max(map(lambda x: x.end, i)) for i in domain_list])

                protein_info[current_transcript_id].append(
                    Transcript(
                        chromosome=self.chromosome,
                        start=start_site,
                        end=end_site,
                        strand=self.cds[current_transcript_id].strand,
                        exons=domain_list,
                        gene=domain_list[0][0].unique_id,
                        domain_type=domain_list[0][0].type,
                        domain_description=domain_list[0][0].description,
                        domain_category=domain_list[0][0].category,
                        category="protein"
                    )
                )

        return protein_info


if __name__ == '__main__':
    pass
