"""
This module contains functions that use LLMs to compare two cell type labels.
"""
from concurrent.futures import ThreadPoolExecutor, as_completed

import pandas as pd

from pandas import DataFrame
from anndata import AnnData

from anndict.llm import call_llm
from anndict.utils import normalize_string
from anndict.annotate.cells.benchmarking import create_label_df


def ai_compare_cell_types_binary(
    label1: str,
    label2: str
) -> str:
    """
    Compares two cell type labels using AI.

    Parameters
    -----------
    label1
        The first item to compare.

    label2
        The second item to compare.

    Returns
    --------
    The comparison result generated by AI.
    """

    # Check if normalized labels are the same--Shortcut for exact match
    if normalize_string(label1) == normalize_string(label2):
        return 'yes'

    # Prepare the prompt
    gpt_prompt = f"1) {label1} 2) {label2} -> "

    # Prepare the messages for the Chat Completions API
    messages = [
        {"role": "system", "content": "You are a terse molecular biologist who determines if two labels refer to the same cell type. Respond only with 'yes' or 'no'. Examples: 1) CD8-positive T cell 2) T cell -> yes; 1) epithelial cell 2) intrahepatic cholangiocyte -> yes; 1) Macrophage 2) Endothelial Cell -> no; 1) B cell 2) Plasma Cell -> yes"},
        {"role": "user", "content": gpt_prompt}
    ]

    # Call the LLM using the call_llm function
    comparison_result = call_llm(
        messages=messages,
        max_tokens=20,
        temperature=0
    )

    return comparison_result


def ai_compare_cell_types_categorical(
    label1: str,
    label2: str
) -> str:
    """
    Compares two cell type labels using an LLM.

    Parameters
    -----------
    label1
        The first item to compare.

    label2
        The second item to compare.

    Returns
    --------
    The comparison result generated by an LLM.
    """
    # Check if normalized labels are the same--Shortcut for exact match
    if normalize_string(label1) == normalize_string(label2):
        return 'perfect match'

    # Prepare the prompt
    gpt_prompt = f"1) {label1} 2) {label2} -> "

    # Prepare the messages for the Chat Completions API
    messages = [
        {"role": "system", "content": "You are a terse molecular biologist who assesses the degree to which two labels refer to the same cell type. Respond only with 'perfect match', 'partial match', or 'no match'. Examples: 1) CD8-positive T cell 2) T cell -> partial match; 1) epithelial cell 2) intrahepatic cholangiocyte -> partial match; 1) Macrophage 2) Endothelial Cell -> no match; 1) macrophage 2) Macrophage. -> perfect match; 1) B cell 2) Plasma Cell -> partial match"},
        {"role": "user", "content": gpt_prompt}
    ]

    # Call the LLM using the call_llm function
    comparison_result = call_llm(
        messages=messages,
        max_tokens=30,
        temperature=0
    )

    return comparison_result

def ai_compare_cell_type_labels_pairwise(
    adata: AnnData,
    cols1: list[str],
    cols2: list[str],
    new_col_prefix: str = 'agreement',
    comparison_level: str = 'binary'
) -> dict[tuple[str,str], DataFrame]:
    """
    Compare cell type labels by finding unique combinations between labels in ``cols1`` and ``cols2``,
    applying the comparison, and mapping the results back to ``adata.obs``.

    Parameters
    -----------
    adata
        an :class:`AnnData`.

    cols1:
        :class:`List` of columns to compare against cols2.

    cols2:
        :class:`List` of columns to compare with cols1.

    new_col_prefix:
        The base name for the new comparison result columns.

    comparison_level:
        Determines which comparison function to use. Can be either ``'binary'`` or ``'categorical'``. If ``'binary'``, cell type labels are assessed for matching. If ``'categorical'``, the quality of the label match is assessed as one of three levels: ``'perfect'``, ``'partial'``, or ``'not-matching'``.

    Returns
    --------
    Dictionary with keys as tuples of (col1, col2) and values as DataFrames with the comparison results.
    """

    # Call create_label_df to generate the label combinations DataFrame
    label_combinations = create_label_df(adata, cols1, cols2)

    # Define the comparison and cleaning functions based on the level
    def binary_compare(row):
        return ai_compare_cell_types_binary(row['col1'], row['col2'])
    def binary_clean(x):
        return 1 if x.lower() == 'yes' else 0 if x.lower() == 'no' else None
    def categorical_compare(row):
        return ai_compare_cell_types_categorical(row['col1'], row['col2'])
    def categorical_clean(x):
        return 0 if x.lower() == 'no match' else 1 if x.lower() == 'partial match' else 2 if x.lower() == 'perfect match' else None

    if comparison_level == 'binary':
        comparison_func, cleaning_func = binary_compare, binary_clean
    elif comparison_level == 'categorical':
        comparison_func, cleaning_func = categorical_compare, categorical_clean
    else:
        raise ValueError("comparison_level must be either 'binary' or 'categorical'.")

    # Convert the label_combinations DataFrame into a list of dictionaries for parallel processing
    label_records = label_combinations.to_dict('records')

    # Use ThreadPoolExecutor to apply the comparison function using threads for API calls
    with ThreadPoolExecutor() as executor:
        # Submit comparison tasks in parallel
        future_to_row = {executor.submit(comparison_func, row): row for row in label_records}

        # Process the results as they are completed
        for future in as_completed(future_to_row):
            row = future_to_row[future]
            row['raw_agreement'] = future.result()

    # Re-create the DataFrame from the processed records
    label_combinations = pd.DataFrame(label_records)

    # Apply the cleaning function to the 'agreement' column
    label_combinations['agreement'] = label_combinations['raw_agreement'].apply(cleaning_func)

    # Initialize a dictionary to store results
    results = {}

    # Iterate over each pair of cols1 and cols2 and map the comparison results back
    for col1 in cols1:
        for col2 in cols2:
            if col1 == col2:
                continue

            # Define the new column name in adata.obs for this comparison
            new_col_name = f"{new_col_prefix}_{col1}_{col2}"

            # Prepare a temporary DataFrame from adata.obs
            temp_obs = adata.obs[[col1, col2]].copy()
            temp_obs.reset_index(inplace=True)  # Reset index to preserve cell indices
            temp_obs.columns = ['index', 'col1_label', 'col2_label']

            # Merge temp_obs with label_combinations on the label columns
            merged_df = pd.merge(
                temp_obs,
                label_combinations[['col1', 'col2', 'agreement']],
                left_on=['col1_label', 'col2_label'],
                right_on=['col1', 'col2'],
                how='left'
            )

            # Assign the 'agreement' values back to adata.obs
            adata.obs[new_col_name] = merged_df['agreement'].values

            # Store the relevant portion of label_combinations in the results dictionary
            results[(col1, col2)] = label_combinations[
                label_combinations['col1'].isin(adata.obs[col1]) &
                label_combinations['col2'].isin(adata.obs[col2])
            ]

    return results
