# -*- coding: utf-8 -*-

"""This script details the classes for genomic annotations and features
- ExonSet: mapping of exon to genes and other meta-data related
to a paricular alternative splicing event
- TpmSet: gene expression matrix
"""

# Author : zzjfrank
# Date : Sep. 14, 2020

import os
import pandas as pd
import numpy as np
from tqdm import tqdm
import pickle
from collections import Counter
import scipy.stats
try:
    import mygene
    has_mygene = True
except ImportError:
    print("Cannot find mygene package, GeneSet functionality will be limited; you can install it by `pip install mygene`")
    has_mygene = False


class ExonSet:
    def __init__(self, data, darts_id, ens_gene, gene_symbol, event_type):
        self.darts_id = darts_id
        self.ens_gene = ens_gene
        self.gene_symbol = gene_symbol
        self.data = data
        self._event_type = event_type
        # a growing list of attributes for ExonSet...
        self.splice_site_coding_region = None
        self.splice_site_pfam_domain = None
        self.splice_site_eclip = None

    def __add__(self, other):
        assert type(other) is ExonSet
        assert self._event_type  == other._event_type
        df = pd.concat([self.data, other.data], axis=0)
        df = df.loc[~df.index.duplicated(keep='first')]
        exon = ExonSet(
                data=df,
                darts_id = df['darts_id'],
                ens_gene = df['GeneID'],
                gene_symbol = df['geneSymbol'],
                event_type = self._event_type
                )
        return exon

    @property
    def loc(self):
        return self.data.loc

    @property
    def iloc(self):
        return self.data.iloc

    @staticmethod
    def from_rmats(gtf_annot_fp, event_type="SE", strip_version_num=True):
        from .rmats_helper import rmatsID_to_dartsID, annot_dtypes
        darts_id_converter = rmatsID_to_dartsID()[event_type]
        gtf_annot = pd.read_table(gtf_annot_fp)
        for x in annot_dtypes[event_type]:
            gtf_annot[x] = gtf_annot[x].astype(annot_dtypes[event_type][x])
        gtf_annot['darts_id'] = darts_id_converter(gtf_annot)
        gtf_annot.index = gtf_annot['darts_id']
        if strip_version_num is True:
            gtf_annot['GeneID'] = [g.split('.')[0] for g in gtf_annot['GeneID']]
        exon = ExonSet(
                data = gtf_annot,
                darts_id = gtf_annot['darts_id'],
                ens_gene = gtf_annot['GeneID'],
                gene_symbol = gtf_annot['geneSymbol'],
                event_type = event_type
                )
        return exon

    @staticmethod
    def from_suppa(suppa_index_fp, event_type='SE', strip_version_num=True, cache_dir=None):
        from .suppa_helper import suppaID_to_dartsID
        suppa_idx = pd.read_table(suppa_index_fp)
        darts_id_converter = suppaID_to_dartsID()[event_type]
        eids = [darts_id_converter(eid) for eid in suppa_idx['event_id']]
        ens_ids = suppa_idx['gene_id'].tolist()
        if strip_version_num:
            ens_ids = [g.split('.')[0] for g in ens_ids]
        if cache_dir is not None:
            mapper = pickle.load(open(os.path.join(cache_dir, 'jemm.genemap.pkl'), 'rb'))
            print("reloaded 'jemm.genemap.pkl' from cache dir %s" % cache_dir)
            gene_symbols = [mapper[g]['symbol'] if g in mapper else None for g in ens_ids]
        elif has_mygene is True:
            mg = mygene.MyGeneInfo()
            mapper = mg.querymany(ens_ids, scopes='ensemblgene', fields='symbol,entrezgene', species='human')
            mapper = { x['query']: {k:x[k] for k in ['entrezgene', 'symbol'] if k in x} for x in mapper if 'notfound' not in x}
            gene_symbols = [mapper[g]['symbol'] if g in mapper else None for g in ens_ids]
        else:
            gene_symbols = []
        data = pd.DataFrame({'darts_id': eids, 'GeneID': ens_ids, 'geneSymbol': gene_symbols })
        data.index = data['darts_id']
        data = data.loc[~data.index.duplicated(keep='first')]
        exon = ExonSet(
            data = data,
            darts_id = data['darts_id'],
            ens_gene = data['GeneID'],
            gene_symbol = data['geneSymbol'],
            event_type = event_type
            )
        return exon

    def read_rbp_clip(self, peak_fps):
        """reads in a set of files generated by intersecting splice sites with RBP peaks, e.g. clip-seq
        
        Parameters
        ----------
        peak_fps: list, or str

        """
        splice_site_annot = {}
        ANNOT_COL = 10
        if type(peak_fps) is not list:
            peaks_fps = [peak_fps]
        for peak_fp in peak_fps:
            with open(peak_fp, 'r') as f:
                for line in f:
                    ele = line.strip().split()
                    sid = ele[3]
                    if sid in splice_site_annot:
                        splice_site_annot[sid].add(ele[ANNOT_COL])
                    else:
                        splice_site_annot[sid] = set([ele[ANNOT_COL]])
        self.splice_site_rbp_clip = splice_site_annot

    def read_coding_region_from_gtf_by_splice_site(self, gtf_fp, chunk_size=1000000):
        """This method reads in the UTR/Coding annotations for each *splice site* from a GTF file
        """
        splice_site_annot = {}
        for chunk in pd.read_csv(gtf_fp, sep="\t", header=None, comment="#", chunksize=chunk_size):
            for r in range(chunk.shape[0]):
                row = chunk.iloc[r]
                if row[2] != 'CDS' and row[2] != 'UTR':
                    continue
                s5_id = ":".join([ row[0], row[6], str(int(row[3])-1) ])
                s3_id = ":".join([ row[0], row[6], str(row[4]) ])
                if s5_id in splice_site_annot:
                    splice_site_annot[s5_id].add(row[2])
                else:
                    splice_site_annot[s5_id] = set([row[2]])
                if s3_id in splice_site_annot:
                    splice_site_annot[s3_id].add(row[2])
                else:
                    splice_site_annot[s3_id] = set([row[2]])
        for ss in splice_site_annot:
            splice_site_annot[ss] = ",".join(sorted(splice_site_annot[ss]))
        self.splice_site_coding_region = splice_site_annot

    def read_pfam_domain_from_bed_by_splice_site(self, bed_fp, clan_annot_fp=None, is_shorten_domain_name=True):
        bed_df = pd.read_csv(bed_fp, sep="\t", header=None)
        if clan_annot_fp is not None:
            clan_annot = pd.read_table(clan_annot_fp, header=None, index_col=3)[2].to_dict()
            for domain in clan_annot:
                if pd.isna(clan_annot[domain]):
                    clan_annot[domain] = domain

        # There are overlapping annotations between protein domains..
        left = bed_df.apply(lambda x: ":".join([str(x[c]) for c in [0, 5, 1] ]), axis=1)
        right = bed_df.apply(lambda x: ":".join([str(x[c]) for c in [0, 5, 2] ]), axis=1)
        ss_to_domain = {}
        for l, r, domain in zip(*[left, right, bed_df[3]]):
            if is_shorten_domain_name is True:
                domain = domain.split("_exon_")[0]
            if clan_annot_fp is not None:
                domain = clan_annot[domain] if domain in clan_annot else domain
            if l in ss_to_domain:
                ss_to_domain[l] += ';'+domain
            else:
                ss_to_domain[l] = domain
            if r in ss_to_domain:
                ss_to_domain[r] += ';'+domain
            else:
                ss_to_domain[r] = domain
        self.splice_site_pfam_domain = ss_to_domain

    def get_splice_site_annotations(self, eids, annotations):
        """Get a list of annotations for a list of event IDs

        See Also
        --------
        jemm.genomic_annotations.ExonSet.get_splice_site_single_annotation
        """
        dfs = []
        for annotation in annotations:
            df = self.get_splice_site_single_annotation(eids, annotation)
            dfs.append(df)
        dfs = pd.concat(dfs, axis=1)
        return dfs

    def get_splice_site_single_annotation(self, eids, annotation):
        """Get a single splice site annotation for a list of event IDs

        The format of Darts ID is always ordered by from upstream to downstream (regardless of strand). Hence, by selecting the
        first and last coordinates, we can sandwich the property for the alternative exon
        The types are
        specified as:
            - se : upstream_ee - cassette_es - cassette_ee - downstream_ee
            - a5 : depends on strand; short_form - long_form
            - a3 : depends on strand; short_form - long_form
            - ri : upstream_es - retained_es - retained_ee - downstream_ee

        Parameters
        ----------
        eids : list
        annotation : str
            choices from ['coding_region', 'pfam_domain']
        """
        _viable_choices = ['coding_region', 'pfam_domain', 'rbp_clip']
        if annotation == 'coding_region':
            ss_annot = self.splice_site_coding_region
            non_exist = "NonCoding"
            return self._map_flanking_ss(ss_annot=ss_annot, eids=eids, annotation=annotation, non_exist=non_exist)
        elif annotation == 'pfam_domain':
            ss_annot = self.splice_site_pfam_domain
            non_exist = "NoDomain"
            return self._map_all_ss(ss_annot=ss_annot, eids=eids, annotation=annotation, non_exist=non_exist)
        elif annotation == 'rbp_clip':
            ss_annot = self.splice_site_rbp_clip
            non_exist = ''
            return self._map_all_ss(ss_annot=ss_annot, eids=eids, annotation=annotation)
        else:
            raise ValueError("Unknown splice site annotation : %s\n, must be in %s" % (annotation, _viable_choices))

    @staticmethod
    def _map_all_ss(ss_annot, eids, annotation, non_exist=None, separator='|'):
        eid_to_property = {}
        for eid in eids:
            eid_to_property[eid] = []
            chrom, strand, left, cassette_l, cassette_r, right = eid.split(":")
            for coord in [left, cassette_l, cassette_r, right]:
                sid = ":".join([chrom, strand, coord])
                if sid in ss_annot:
                    if type(ss_annot[sid]) is str:
                        eid_to_property[eid].append(ss_annot[sid])
                    else:
                        eid_to_property[eid].extend(list(ss_annot[sid]))
                elif non_exist is not None:
                    eid_to_property[eid].append(non_exist)

            eid_to_property[eid] = separator.join(list(set(eid_to_property[eid])))
        eid_to_property = pd.DataFrame.from_dict(eid_to_property, orient="index", columns=[annotation])
        return eid_to_property

    @staticmethod
    def _map_flanking_ss(ss_annot, eids, annotation, non_exist=None, separator='-'):
        eid_to_property = {}
        for eid in eids:
            chrom, strand, left, _, _, right = eid.split(":")
            left_id = ":".join([chrom, strand, left])
            right_id = ":".join([chrom, strand, right])
            if left_id in ss_annot:
                left_annot = ss_annot[left_id]
            else:
                left_annot = non_exist if non_exist is not None else ''
            if right_id in ss_annot:
                right_annot = ss_annot[right_id]
            else:
                right_annot = non_exist if non_exist is not None else ''
            if strand == '+':
                eid_to_property[eid] = left_annot + separator + right_annot
            else:
                eid_to_property[eid] = right_annot + separator + left_annot
        eid_to_property = pd.DataFrame.from_dict(eid_to_property, orient="index", columns=[annotation])
        return eid_to_property

    @staticmethod
    def enrichment_test(das_eids, metaloader, annot_col, separator='|'):
        """Perform a Fisher's test for the occurence of annotation in foreground vs background

        Parameters
        ----------
        das_eids : dict
            a map of AS types to a list of event IDs
        metaloader : jemm.meta_loader.MetaLoader
            a MetaLoader instance that holds data and regression table
        annot_col : str
        separator : str

        """
        das_cnts = Counter([item for as_type in das_eids for eid_items in metaloader._reg_tables[as_type].loc[das_eids[as_type], annot_col].dropna() for item in eid_items.split(separator)])
        bg_cnts = Counter([item for as_type in das_eids for eid_items in metaloader._reg_tables[as_type][annot_col].dropna() for item in eid_items.split(separator)])
        das_tot_cnt = sum([len(das_eids[x]) for x in das_eids])
        bg_tot_cnt = sum([metaloader._reg_tables[x].shape[0] for x in metaloader.data])
        items = []
        ratios = []
        pvals = []
        obs = []
        expected = []
        for item in bg_cnts:
            das_obs_cnt = das_cnts[item]
            bg_obs_cnt = bg_cnts[item]
            mat = np.array([[das_obs_cnt, bg_obs_cnt], [das_tot_cnt, bg_tot_cnt]])
            res = scipy.stats.fisher_exact(mat)
            items.append(item)
            ratios.append(res[0])
            pvals.append(res[1])
            obs.append(das_obs_cnt)
            expected.append(bg_obs_cnt/bg_tot_cnt*das_tot_cnt)

        test_df = pd.DataFrame({'items': items, 'observed': obs, 'expected': expected, 'ratios': ratios, 'pvals': pvals})
        return test_df


class GeneSet:
    """A class for storing and analyzing gene sets with differential alternative splicing

    Attributes
    ----------
    geneexp : pandas.DataFrame
        A matrix that stores the gene expression in each row; columns are samples
    genemap : dict
        A dictionary that maps gene ID to its synonyms, e.g. ensembl ID, HUGO gene symbol and entrez ID

    Notes
    ------
    For `genemap`, currently expects the keys are Ensembl gene IDs
    """
    def __init__(self, geneexp=None, genemap=None):
        self.geneexp = geneexp
        self.genemap = genemap

        # a growing list of internal attributes
        self.gene_sets = {}      # gene_set_name --> (id_type, [genes])
        self.ranked_genes = {}   # ranked_gene_name --> (gene_ids)

    @staticmethod
    def from_transcript_matrix(fp, t2g_fp, strip_version_num=True):
        """Constructor for GeneSet from a set of transcript expressions
        """
        t2g = pd.read_table(t2g_fp)
        if strip_version_num:
            t2g['ensembl_transcript_id'] = [x.split(".")[0] for x in t2g['ensembl_transcript_id']]
            t2g['ensembl_gene_id'] = [x.split(".")[0] for x in t2g['ensembl_gene_id']]
        gene_index = {k:v for v,k in enumerate(t2g['ensembl_gene_id'].unique().tolist())}
        t2g.index = t2g['ensembl_transcript_id']
        t2g = t2g['ensembl_gene_id'].to_dict()
        tx_exp = pd.read_table(fp, index_col=0)
        n_genes = len(gene_index)
        n_samps = tx_exp.shape[1]
        gene_exp = np.zeros((n_genes, n_samps))
        n_missed = 0
        for tid in tqdm(tx_exp.index):
            if not tid in t2g:
                n_missed +=1
                continue
            tx = tx_exp.loc[tid].to_numpy()
            if len(tx.shape) > 1:
                gene_exp[gene_index[t2g[tid]]] += tx.mean(axis=0)
            else:
                gene_exp[gene_index[t2g[tid]]] += tx
        print("Total n=%i transcripts not found in t2g index"%n_missed)
        gdf = pd.DataFrame(gene_exp, columns=tx_exp.columns, index=[g[0] for g in sorted(gene_index.items(), key=lambda item: item[1])])
        gs = GeneSet(geneexp = gdf)
        return gs

    @staticmethod
    def from_gene_matrix(fp, **kwargs):
        """Contructor for GeneSet from a gene expression matrix

        expects the rows are genes, columns are samples
        """
        gdf = pd.read_table(fp, **kwargs)
        gs = GeneSet(geneexp = gdf)
        return gs

    @staticmethod
    def from_rmats(fp, cache_dir=None, geneset=None):
        """Constructor for GeneSet from a rMATS `fromGTF.${eventType}.txt` file

        Parameters
        ----------
        fp : str
        cache_dir : str
            filepath to cache folder for reloading and storing "jemm.genemap.pkl"
        geneset : None or jemm.genomic_annotation.GeneSet
            if None, starts a new instance; if given an instance of GeneSet, add/update genemap to it
        """
        gtf_annot = pd.read_table(fp)
        gene_id = [g.split('.')[0] for g in gtf_annot['GeneID'].unique()]
        if cache_dir is not None and os.path.isfile(os.path.join(cache_dir, "jemm.genemap.pkl")):
            print("reload from previous cached pickle..")
            mapper = pickle.load(open(os.path.join(cache_dir, "jemm.genemap.pkl"), "rb"))
        elif has_mygene:
            mg = mygene.MyGeneInfo()
            mapper = mg.querymany(gene_id, scopes='ensemblgene', fields='symbol,entrezgene', species='human')
            mapper = { x['query']: {k:x[k] for k in ['entrezgene', 'symbol'] if k in x} for x in mapper if 'notfound' not in x}
            if cache_dir is not None:
                os.makedirs(cache_dir, exist_ok=True)
                pickle.dump(mapper, open(os.path.join(cache_dir, "jemm.genemap.pkl"), "wb"))
        else:
            mapper = {}
        genemap = {g:mapper[g] for g in gene_id if g in mapper}
        if geneset is None:
            geneset = GeneSet(genemap = genemap)
        else:
            geneset.genemap = genemap
        return geneset

    def marginalized_gene_expr_mean(self, covariate, contrast_col, is_log2=True, pseudo_cnt=1):
        """Given a Contrast, get the marginalized gene expression means

        Parameters
        ----------
        covariate : jemm.covariate.Covariate
        """
        covariate.match_index(self.geneexp.columns)
        covariate_sample_df = self.geneexp[covariate.covariate.index]
        contrast_col_values = covariate.covariate[contrast_col].unique()
        marginals = {}
        for val in contrast_col_values:
            key = "%s=%s"%(contrast_col, val)
            marginals[key] = covariate_sample_df[ covariate.covariate.index[(covariate.covariate[contrast_col]==val)] ].apply(np.mean, axis=1)
        df = pd.DataFrame.from_dict(marginals)
        if is_log2:
            df = np.log2(df + pseudo_cnt)
        return df

    def read_geneset_annotations(self, fp, gene_id_type, min_geneset_size=10, max_geneset_size=1000):
        """read in a list of gene set to perform GSEA

        Parameters
        ----------
        fp : str
        gene_id_type : str
        min_geneset_size : int
        max_geneset_size : int

        Raises
        ------
        AssertionError : if `gene_id_type` is not understood
        """
        _viable_id_types = ('entrez', 'ensembl', 'symbol')
        assert gene_id_type in _viable_id_types, "Unknown gene_id type %s; must be in %s" % (gene_id_type, _viable_id_types)
        with open(fp, "r") as f:
            for line in f:
                ele = line.strip().split('\t')
                # plus 2 because the first and second columns are name/annotations
                if len(ele) < min_geneset_size + 2:
                    continue
                if len(ele) > max_geneset_size + 2:
                    continue
                gs_name = "|".join([ele[0], ele[1]])
                self.gene_sets[gs_name] = (gene_id_type, ele[2:])
        #return self.gene_sets

    def parse_ranked_exons(self, regression_table, reg_name, agg_method='best', overwrite=False):
        """Parse a list of ranked exons from regression table, aggregate to gene-level ranking, and
        add them to ranked gene list for downstream analysis on gene set enrichment

        Parameters
        ----------
        regression_table : pandas.DataFrame
            a dataframe of regression table, generated by `jemm.Jemm.save_regression_table`
        reg_name : str
            an identifier for this ranked gene list
        agg_method : str
            method for aggregating exons into genes; supports ('best', 'average')
        overwrite : bool
            overwrite protection for existing `reg_name`; if true, will force overwrite

        Raises
        ------
        AssertionError : if agg_method not understood
        AssertionError : if reg_name exists and overwrite=False
        """
        _viable_agg_methods = ('best', 'average')
        assert agg_method in _viable_agg_methods, "Unknown agg method %s; must be in %s" % (agg_method, _viable_agg_methods)
        if reg_name in self.ranked_genes:
            assert overwrite is True, "ranked_genes with name '%s' exists; abort because overwrite=false" % reg_name
        from collections import defaultdict
        gene_to_rank = defaultdict(list)
        for i in range(regression_table.shape[0]):
            gene = regression_table.iloc[i]['GeneID']
            gene_to_rank[gene].append(i)
        if agg_method == 'best':
            for g in gene_to_rank:
                gene_to_rank[g] = min(gene_to_rank[g])
        elif agg_method == 'average':
                gene_to_rank[g] = np.mean(gene_to_rank[g])
        self.ranked_genes[reg_name] = [_[0] for _ in sorted(gene_to_rank.items(), key=lambda x: x[1])]
        #return self.ranked_genes

    def gsva_ks(self, ranked_gene_name, geneset_name):
        """Perform GeneSet Variation/Enrichment Analysis be K-S test

        Parameters
        ----------
        ranked_gene_name : str
            the name for ranked list
        geneset_name : str
            the name for calculating Enrichment-Score against the gene set

        Raises
        ------
        ValueError : if ranked_gene_name or geneset_name not found

        References
        ----------
        Citation:
            https://bmcbioinformatics.biomedcentral.com/articles/10.1186/1471-2105-14-7

        Based on my old R code::
            # data: 1st column is rank , 2nd column is a factor 'Hit'/'Null', 3rd column is the FDR/delta-psi
            data=data[sample(nrow(data)),]
            data=data[order(data[,1]),]
            hits=which(data[,2]=='Hit')
            null=which(data[,2]=='Null')
            n.x=length(hits)
            n.y=length(null)
            w=c(hits, null)
            ES = cumsum(ifelse(order(w) <= n.x, 1/n.x, -1/n.y))
            hit.pos = which(data[,2]=='Hit')
            nonhit.pos = which(data[,2]!='Hit')

        """
        import scipy.stats
        if ranked_gene_name not in self.ranked_genes:
            raise ValueError('Cannot find %s in ranked_genes' % ranked_gene_name)
        if geneset_name not in self.gene_sets:
            raise ValueError('Cannot find %s in gene_sets' % geneset_name)
        this_rank = self.ranked_genes[ranked_gene_name]
        this_id_type, this_gs = self.gene_sets[geneset_name]
        # convert ID for the ranked list
        if this_id_type == 'entrez':
            this_rank = [self.genemap[g]['entrezgene'] for g in this_rank if g in self.genemap and 'entrezgene' in self.genemap[g]]
        elif this_id_type == 'symbol':
            this_rank = [self.genemap[g]['symbol'] for g in this_rank if g in self.genemap and 'symbol' in self.genemap[g]]
        elif this_id_type == 'ensembl':
            pass
        else:
            raise ValueError('Unsupported id type "%s" at gsea-ks' % this_id_type)
        hits = [i for i in range(len(this_rank)) if this_rank[i] in this_gs]
        nonhits = [i for i in range(len(this_rank)) if this_rank[i] not in this_gs]
        stats, pval = scipy.stats.ks_2samp(hits, nonhits)

        w = np.concatenate([hits, nonhits], axis=0)
        n_x = len(hits)
        n_y = len(nonhits)
        es = np.cumsum([1./n_x if w_ <= n_x else -1./n_y for w_ in np.argsort(w)+1])
        return stats, pval, es, hits, nonhits


