"""
**Artemis**: public knowledge graphs enable accessible and scalable drug target discovery.

Github Repository   
-----------------

`GitHub - alethiotx/artemis-paper <https://github.com/alethiotx/artemis-paper>`_
"""

from upsetplot import UpSet, from_contents

def get_all_targets(scores: list):
    """
    Extract all unique target genes from a list of score dictionaries.

    :param scores: A list of dictionaries, where each dictionary contains a ``Target Gene`` key
                   with a value that can be converted to a list using tolist() method
                   (typically a pandas Series or similar object)
    :type scores: list
    :return: A list of all unique target genes found across all dictionaries in the input list
    :rtype: list

    **Example**
    >>> scores = [
    ...     {'Target Gene': pd.Series(['GENE1', 'GENE2'])},
    ...     {'Target Gene': pd.Series(['GENE2', 'GENE3'])}
    ... ]
    >>> get_all_targets(scores)
    ['GENE1', 'GENE2', 'GENE3']
    """
    all_targets = set()

    for d in scores:
        all_targets.update(d['Target Gene'].tolist())

    return(list(all_targets))

def cut_clinical_scores(scores: list, lowest_score = 0):
    """
    Filter clinical scores by removing entries below a threshold.

    This function creates a copy of the input scores list and filters each DataFrame
    within it to retain only rows where the ``Clinical Score`` exceeds the specified
    lowest_score threshold.

    :param scores: List of DataFrames containing clinical score data with a ``Clinical Score`` column
    :type scores: list
    :param lowest_score: Minimum score threshold for filtering, defaults to 0
    :type lowest_score: int or float
    :return: List of filtered DataFrames with scores above the threshold
    :rtype: list

    **Example**
    >>> import pandas as pd
    >>> scores = [pd.DataFrame({'Clinical Score': [0.5, 1.5, 2.0]})]
    >>> filtered = cut_clinical_scores(scores, lowest_score=1.0)
    >>> filtered[0]
        Clinical Score
    1         1.5
    2         2.0
    """
    res = scores.copy()

    for n, d in enumerate(res):
        res[n] = d[d['Clinical Score'] > lowest_score]

    return(res)

def find_overlapping_genes(genes: list, overlap = 1, common_genes = []):
    """
    Find genes that overlap across multiple gene lists.

    :param genes: A list of gene lists to check for overlapping genes.
    :type genes: list
    :param overlap: Minimum number of lists a gene must appear in to be considered overlapping, defaults to 1.
    :type overlap: int
    :param common_genes: Initial list of common genes to include in the result, defaults to [].
    :type common_genes: list
    :return: A list of genes that appear in more than ``overlap`` lists but not in all lists.
    :rtype: list

    **Example**
    >>> genes_list = [['gene1', 'gene2'], ['gene2', 'gene3'], ['gene2', 'gene4']]
    >>> find_overlapping_genes(genes_list, overlap=1)
    ['gene2']

    .. note::

        Genes that appear in all input lists are excluded from the result by default.
        To include them, modify the condition ``d[i] < len(genes)`` to ``d[i] <= len(genes)``.

    """
    d = {}
    overlapping_genes = common_genes.copy()

    for i in [x for y in genes for x in y]:
        if i in d.keys():
            d[i]=d[i]+1
        else:
            d[i]=1
    for i in d.keys():
        if d[i]>overlap and d[i] < len(genes): # the second condition excludes genes present in all lists, if you want to include those, change < to <=, or remove the condition entirely
            overlapping_genes.append(i)
    return overlapping_genes

def uniquify_clinical_scores(scores: list, overlap = 1, common_genes = []):
    """
    Remove overlapping genes from clinical score dataframes to ensure uniqueness.
    This function processes a list of dataframes containing clinical scores and removes
    genes that appear in multiple dataframes above a specified overlap threshold.
    :param scores: List of dataframes, each containing a ``Target Gene`` column with gene identifiers
    :type scores: list
    :param overlap: Minimum number of dataframes a gene must appear in to be considered overlapping, defaults to 1
    :type overlap: int, optional
    :param common_genes: Additional list of genes to always consider as overlapping regardless of frequency, defaults to []
    :type common_genes: list, optional
    :return: List of dataframes with overlapping genes removed from each dataframe
    :rtype: list

    **Example**
    >>> df1 = pd.DataFrame({'Target Gene': ['BRCA1', 'TP53', 'EGFR']})
    >>> df2 = pd.DataFrame({'Target Gene': ['TP53', 'KRAS', 'MYC']})
    >>> result = uniquify_clinical_scores([df1, df2], overlap=2)
    >>> # TP53 will be removed from both dataframes as it appears in 2 or more

    .. note::

        The function uses :func:`find_overlapping_genes` to identify genes that should be removed.
    .. warning::

        This function modifies copies of the input dataframes. The original dataframes remain unchanged.
    """
    genes = []
    for n, d in enumerate(scores):
        genes.append(d['Target Gene'].tolist())

    overlapping_genes = find_overlapping_genes(genes, overlap = overlap, common_genes = common_genes)
    
    res = scores.copy()

    for n, d in enumerate(res):
        res[n] = d[~d['Target Gene'].isin(overlapping_genes)]
    
    return(res)


def uniquify_pathway_genes(genes: list, overlap = 1, common_genes = []):
    """
    Remove overlapping genes from pathway gene lists.
    This function identifies genes that appear in multiple pathways (based on the
    specified overlap threshold) and removes them from each pathway's gene list,
    returning uniquified pathway gene lists.
    :param genes: A list of gene lists, where each inner list represents genes
                  associated with a particular pathway.
    :type genes: list
    :param overlap: The minimum number of pathways a gene must appear in to be
                    considered overlapping and removed. Defaults to 1.
    :type overlap: int, optional
    :param common_genes: A list of genes that should be considered as commonly
                         overlapping regardless of their occurrence count.
    :type common_genes: list, optional
    :return: A copy of the input gene lists with overlapping genes removed from
             each pathway.
    :rtype: list

    **Example**
    >>> genes = [['GENE1', 'GENE2', 'GENE3'], ['GENE2', 'GENE4'], ['GENE3', 'GENE5']]
    >>> uniquify_pathway_genes(genes, overlap=2)
    [['GENE1'], ['GENE4'], ['GENE5']]
    """
    overlapping_genes = find_overlapping_genes(genes, overlap = overlap, common_genes = common_genes)
    
    res = genes.copy()

    for n, y in enumerate(res):
        res[n] = [x for x in y if x not in overlapping_genes]

    return(res)

def prepare_upset(breast, lung, prostate, melanoma, bowel, diabetes, cardiovascular, mode):
    """
    Prepare data for UpSet plot visualization.

    This function processes disease-related data and converts it into a format
    suitable for UpSet plot generation using the `from_contents` function.

    :param breast: Breast cancer data. In ``ct`` mode, should be a DataFrame with
                   ``Target Gene`` column. In ``pg`` mode, should be a list.
    :type breast: pandas.DataFrame or list
    :param lung: Lung cancer data. In ``ct`` mode, should be a DataFrame with
                 ``Target Gene`` column. In ``pg`` mode, should be a list.
    :type lung: pandas.DataFrame or list
    :param prostate: Prostate cancer data. In ``ct`` mode, should be a DataFrame with
                     ``Target Gene`` column. In ``pg`` mode, should be a list.
    :type prostate: pandas.DataFrame or list
    :param melanoma: Melanoma data. In ``ct`` mode, should be a DataFrame with
                     ``Target Gene`` column. In ``pg`` mode, should be a list.
    :type melanoma: pandas.DataFrame or list
    :param bowel: Bowel cancer data. In ``ct`` mode, should be a DataFrame with
                  ``Target Gene`` column. In ``pg`` mode, should be a list.
    :type bowel: pandas.DataFrame or list
    :param diabetes: Diabetes data. In ``ct`` mode, should be a DataFrame with
                     ``Target Gene`` column. In ``pg`` mode, should be a list.
    :type diabetes: pandas.DataFrame or list
    :param cardiovascular: Cardiovascular disease data. In ``ct`` mode, should be a
                           DataFrame with ``Target Gene`` column. In ``pg`` mode,
                           should be a list.
    :type cardiovascular: pandas.DataFrame or list
    :param mode: Processing mode. ``ct`` extracts ``Target Gene`` column from DataFrames,
                 ``pg`` uses the data directly as lists.
    :type mode: str
    :return: Dictionary formatted for UpSet plot with disease names as keys and
             gene lists or data lists as values.
    :rtype: dict
    :raises ValueError: If mode is not ``ct`` or ``pg``.
    """
    if mode == 'ct':
        return from_contents(
            {
                "breast": breast['Target Gene'].tolist(), 
                "lung": lung['Target Gene'].tolist(), 
                "prostate": prostate['Target Gene'].tolist(), 
                "melanoma": melanoma['Target Gene'].tolist(), 
                "bowel": bowel['Target Gene'].tolist(), 
                "diabetes": diabetes['Target Gene'].tolist(), 
                "cardiovascular": cardiovascular['Target Gene'].tolist()
            }
        )
    elif mode == 'pg':
        return from_contents(
            {
                "breast": breast, 
                "lung": lung, 
                "prostate": prostate, 
                "melanoma": melanoma, 
                "bowel": bowel, 
                "diabetes": diabetes, 
                "cardiovascular": cardiovascular
            }
        )
    else:
        raise ValueError("Mode must be either 'ct' or 'pg'")

def create_upset_plot(indications, min_subset_size):
    """
    Create an UpSet plot for visualizing set intersections.

    :param indications: Data structure containing set membership information for creating the UpSet plot
    :type indications: pandas.DataFrame or similar data structure compatible with UpSet
    :param min_subset_size: Minimum size threshold for subsets to be displayed in the plot
    :type min_subset_size: int
    :return: An UpSet plot object configured with the specified parameters
    :rtype: UpSet
    :raises TypeError: If indications is not a compatible data structure
    :raises ValueError: If min_subset_size is negative

    .. note::

        The plot is configured with the following settings:
        
        - subset_size: "count" - shows count of elements in each subset
        - orientation: "vertical" - displays plot in vertical orientation
        - show_counts: "{:d}" - formats counts as integers
        - sort_by: ``cardinality`` - sorts subsets by their size
        - sort_categories_by: ``input`` - maintains input order for categories
        - min_subset_size: filters out subsets smaller than the specified threshold

    .. seealso::

        `UpSet documentation <https://upsetplot.readthedocs.io/>`_
    """
    return UpSet(
        indications, 
        subset_size="count", 
        orientation="vertical", 
        show_counts="{:d}", 
        sort_by='cardinality', 
        sort_categories_by = 'input',
        min_subset_size = min_subset_size,
    )

# run when file is directly executed
if __name__ == '__main__': 
    print("Artemis module. Run tests or examples here.")
    # # 'Breast Cancer', 
    # # 'Lung Cancer',
    # # 'Prostate Cancer',
    # # 'Bowel Cancer',
    # # 'Melanoma',
    # # 'Diabetes Mellitus Type 2',
    # # 'Cardiovascular Disease',
    # df = get_pathway_genes("acute myeloid leukemia")
    # print(df.loc["FLT3", ["gene_count", "rank"]])
    # breast, lung, prostate, melanoma, bowel, diabetes, cardiovascular = load_clinical_scores(date="2025-11-11")
    # print('Clinical scores:\n\n')
    # print([breast, lung, prostate, melanoma, bowel, diabetes, cardiovascular])
    # print('All target genes:\n\n')
    # known_targets = get_all_targets([breast, lung, prostate, melanoma, bowel, diabetes, cardiovascular])
    # print(known_targets)
    # print('Cut clinical scores:\n\n')
    # print(cut_clinical_scores([breast, lung, prostate, melanoma, bowel, diabetes, cardiovascular], lowest_score = 10))
    # breast_pg, lung_pg, prostate_pg, melanoma_pg, bowel_pg, diabetes_pg, cardiovascular_pg = load_pathway_genes(date='2025-09-15', n=50)
    # print('Uniquified clinical scores:\n\n')
    # print(uniquify_clinical_scores([breast, lung, prostate, melanoma, bowel, diabetes, cardiovascular]))
    # print('Uniquified pathway genes:\n\n')
    # print(uniquify_pathway_genes([breast_pg, lung_pg, prostate_pg, melanoma_pg, bowel_pg, diabetes_pg, cardiovascular_pg]))
    # print(prepare_upset(breast, lung, prostate, melanoma, bowel, diabetes, cardiovascular, mode='ct'))

    # # X is always bigger than y!!!
    # X = DataFrame({
    #     'term1' : [0, 1, 2, 3, 4, 5, 4, 3, 2, 1, 0, 1, 2, 3, 4, 5, 0, 0, 1, 1, 2, 2, 0, 1, 2, 3, 4, 5, 4, 3, 2, 1, 0, 1, 2, 3, 4, 5, 0, 0, 1, 1, 2, 2], 
    #     'term2' : [1, 2, 3, 4, 5, 0, 0, 1, 1, 2, 2, 2, 3, 4, 5, 6, 5, 7, 2, 9, 2, 1, 1, 2, 3, 4, 5, 0, 0, 1, 1, 2, 2, 2, 3, 4, 5, 6, 5, 7, 2, 9, 2, 1], 
    #     'term3' : [2, 3, 4, 5, 6, 5, 7, 2, 9, 2, 1, 3, 4, 5, 6, 7, 4, 5, 0, 0, 1, 2, 2, 3, 4, 5, 6, 5, 7, 2, 9, 2, 1, 3, 4, 5, 6, 7, 4, 5, 0, 0, 1, 2], 
    #     'term4' : [3, 4, 5, 6, 7, 4, 5, 0, 0, 1, 2, 9, 2, 1, 3, 4, 5, 6, 7, 5, 7, 2, 3, 4, 5, 6, 7, 4, 5, 0, 0, 1, 2, 9, 2, 1, 3, 4, 5, 6, 7, 5, 7, 2], 
    #     'term5' : [4, 5, 6, 7, 8, 6, 5, 7, 2, 9, 3, 0, 1, 2, 3, 4, 5, 4, 3, 2, 1, 0, 4, 5, 6, 7, 8, 6, 5, 7, 2, 9, 3, 0, 1, 2, 3, 4, 5, 4, 3, 2, 1, 0]
    # })
    # y = DataFrame({
    #     'Target Gene' : ['gene1', 'gene2', 'gene3', 'gene4', 'gene5', 'gene6', 'gene7', 'gene8', 'gene9', 'gene10', 'gene11', 'gene12', 'gene13', 'gene14', 'gene15'], 
    #     'Clinical Score' : [1, 1, 1, 4, 5, 10, 5, 3, 2, 5, 1, 100, 200, 300, 400]
    # })
    # # X is always bigger than y!!!
    # X.index = ['gene1', 'gene2', 'gene3', 'gene4', 'gene5', 'gene6', 'gene7', 'gene8', 'gene9', 'gene10', 'gene11', 'gene12', 'gene13', 'gene14', 'gene15', 'gene16', 'gene17', 'gene18', 'gene19', 'gene20', 'gene21', 'gene22', 'gene23', 'gene24', 'gene25', 'gene26', 'gene27', 'gene28', 'gene29', 'gene30', 'gene31', 'gene32', 'gene33', 'gene34', 'gene35', 'gene36', 'gene37', 'gene38', 'gene39', 'gene40', 'gene41', 'gene42', 'gene43', 'gene44']
    # print('\nInput term matrix:')
    # print(X)
    # print('\nInput clinical scores:')
    # print(y)
    # print('\nCheck binning:')
    # res = pre_model(X, y, bins = 5)
    # print(res['y_encoded'])

    # known_targets = ['gene32', 'gene33']

    # res = pre_model(X, y, known_targets = known_targets, bins = 5)
    # print(res['y_encoded'])

    # print('\nResults of cross validation pipeline:')
    # print(cv_pipeline(X, y, n_iterations = 3))
    # print(cv_pipeline(X, y, y_slot = 'y_encoded', scoring = 'accuracy', n_iterations = 3))
