"""
This module handles annotating a set of genes sets by considering each gene set in the context of the other sets of genes.
"""
from anndict.utils import enforce_semantic_list
from anndict.llm import retry_call_llm
from .base import ai_annotate_by_comparison


def ai_cell_types_by_comparison(gene_lists, cell_types=None, tissues=None, subtype=False):
    """
    Returns cell type labels for multiple lists of marker genes as determined by AI.
    Args:
    gene_lists (list of lists): A list containing multiple lists of genes to be described.
    cell_type (str, optional): The cell type to provide context for the AI.
    tissue (str, optional): The tissue of origin to provide context for the AI.
    Returns:
    list of str: The cell type labels generated by AI for each gene list.
    """
    # Enforce semantic_list for each gene list
    for gene_list in gene_lists:
        enforce_semantic_list(gene_list)

    # Prepare the system prompt
    system_prompt = (
        "You are a terse molecular biologist. You respond in a few words and without restating any part of the question. "
        f"Compare and contrast gene sets to identify the most likely cell {'sub' if subtype else ''}type based on marker genes."
    )

    # Prepare the initial user prompt for contrasting all gene lists
    # initial_prompt = f"Tissue: {tissues}, " if tissue else ""
    # initial_prompt += f"Cell Type: {cell_type}, " if cell_type else ""
    initial_prompt = "Briefly compare and contrast the following gene sets:\n"
    for i, gene_list in enumerate(gene_lists, 1):
        tissue_str = " " + ', '.join(tissues[i]) if tissues and tissues[i] else ""
        cell_type_str = " " + ', '.join(cell_types[i]) if cell_types and cell_types[i] else ""

        initial_prompt += f"{i}){tissue_str}{cell_type_str} {('    '.join(gene_list))}\n"

    # Initialize the conversation
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": initial_prompt}
    ]

    # Get the initial contrast response
    contrast_response = retry_call_llm(
        messages=messages,
        process_response=lambda x: x,
        failure_handler=lambda: "Failed to contrast gene sets",
        call_llm_kwargs={'max_tokens': 300, 'temperature': 0},
        max_attempts=1
    )

    # Append the contrast response to the conversation
    messages.append({"role": "assistant", "content": contrast_response})

    messages.append({"role": "user", "content": "Provide only the new label. "})

    # Process each gene list
    cell_subtype_labels = []
    for i, gene_list in enumerate(gene_lists, 1):
        tissue_str = " " + ', '.join(tissues[i]) if tissues and tissues[i] else ""
        cell_type_str = " " + ', '.join(cell_types[i]) if cell_types and cell_types[i] else ""

        gene_set_prompt = f"What is the cell{tissue_str}{cell_type_str} {'sub' if subtype else ''}type label for the gene set: {('    '.join(gene_list))}?"
        messages.append({"role": "user", "content": gene_set_prompt})

        # Get the subtype label
        subtype_label = retry_call_llm(
            messages=messages,
            process_response=lambda x: x.strip(),
            failure_handler=lambda cell_type_str=cell_type_str, i=i: cell_type_str if cell_types and cell_types[i] else "Unknown",
            call_llm_kwargs={'max_tokens': 50, 'temperature': 0},
            max_attempts=1
        )

        cell_subtype_labels.append(subtype_label)
        messages.append({"role": "assistant", "content": subtype_label})

    # print(f"{messages}")
    return cell_subtype_labels


def ai_annotate_cell_type_by_comparison(adata, groupby, n_top_genes, label_column='ai_cell_type_by_comparison', cell_type_of_origin_col=None, tissue_of_origin_col=None, adt_key=None, **kwargs):
    """
    Annotate cell types by comparison using AI.

    This function wraps the ai_annotate function to perform cell type annotation
    based on comparison of gene expression profiles.

    Parameters:
    adata : AnnData Annotated data matrix.
    groupby : str Column name in adata.obs for grouping cells.
    n_top_genes : int Number of top genes to consider for annotation.
    label_column : str, optional Name of the column to store the AI-generated cell type labels (default: 'ai_cell_type_by_comparison').

    Returns:
    AnnData Annotated data with AI-generated cell type labels.
    """
    # print(f"number of unique categories: {len(adata.obs[groupby].unique())}")
    # if tissue_of_origin_col:
    #     tissue = adata.obs[tissue_of_origin_col].unique()
    #     if len(tissue == 1):
    #         tissue = tissue[0]
    #     else:
    #         raise ValueError(f"Multiple tissues of_origin found in adata.obs[{tissue_of_origin_col}]. Currently must have only one tissue of origin per cell type. Pick a different tissue of origin column or set tissue_of_origin_col=None")
    # else:
    #     tissue = None
    return ai_annotate_by_comparison(func=ai_cell_types_by_comparison, adata=adata, groupby=groupby, n_top_genes=n_top_genes, label_column=label_column, cell_type=adt_key, cell_type_of_origin_col=cell_type_of_origin_col, tissue_of_origin_col=tissue_of_origin_col, **kwargs)
