"""
This module handles annotating a set of genes sets by considering each gene set in the context of the other sets of genes.
"""

from anndata import AnnData

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: list[list[str]],
    cell_types: list[str] = None,
    tissues: list[str] = None,
    subtype: bool = False
) -> list[str]:
    """
    Returns cell type labels for multiple lists of marker genes as determined by an LLM.

    Parameters
    -------------
    gene_lists
        A list containing multiple lists of genes to be described.

    cell_types
        A list of cell types to provide context to the LLM. 
        Each entry corresponds to a gene list in ``gene_lists``.

    tissues
        A list of tissues of origin to provide context to the LLM. 
        Each entry corresponds to a gene list in ``gene_lists``.

    subtype
        Whether to try to annotate cell type (if ``False``) or subtype (if ``True``).

    Returns
    --------
    The cell type labels generated by an LLM for each gene list.
    """
    if not gene_lists:
        raise ValueError("Empty list of gene lists passed to ai_cell_type_by_comparison")

    # 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: AnnData,
    groupby: str,
    n_top_genes: int,
    *,
    new_label_column: str ='ai_cell_type_by_comparison',
    cell_type_of_origin_col: str = None,
    tissue_of_origin_col: str = None,
    **kwargs
) -> AnnData:
    """
    Annotate cell types by comparison using LLMs.

    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
        Column name in adata.obs for grouping cells.

    n_top_genes
        Number of top genes to consider for annotation.

    new_label_column
        Name of the column to store the AI-generated cell type labels.

    cell_type_of_origin_col
        Name of a column in ``adata.obs`` that contains the cell type 
        of orgin. Used for context to the LLM.

    tissue_of_origin_col
        Name of a column in ``adata.obs`` that contains the tissue of orgin. 
        Used to provide context to the LLM.

    adt_key
        Used by :func:`adata_dict_fapply` or :func:`adata_dict_fapply_return` 
        when passing this function to them.

    Returns
    -------
    :class:`AnnData` with LLM-generated cell type labels stored in ``adata.obs[label_column]``.
    """
    return ai_annotate_by_comparison(func=ai_cell_types_by_comparison, adata=adata, groupby=groupby, n_top_genes=n_top_genes, new_label_column=new_label_column, cell_type_of_origin_col=cell_type_of_origin_col, tissue_of_origin_col=tissue_of_origin_col, **kwargs)
