"""
Consistency Check
"""

import os

import nltk
import openai
from nltk.tokenize import sent_tokenize
from sentence_transformers import SentenceTransformer, util
from torch import cuda
from torch import device as torch_device


def get_samples(prompt, num_samples, model):
    """
    Function to generate samples using OpenAI's chat completions.

    Args:
        prompt (str): The prompt for the chat completion.
        num_samples (int): The number of samples to generate.
        model (str): The model to use for chat completion.

    Returns:
        list: A list of generated samples.
    """
    max_tokens = 200  # max(200, int(len(prompt) / 4))
    samples = []
    for _ in range(num_samples):
        try:
            response = openai.chat.completions.create(
                model=model,
                messages=[{"role": "user", "content": prompt}],
                max_tokens=max_tokens,
                temperature=1,
            )
            samples.append(response.choices[0].message.content)
        except Exception as e:
            pass

    return samples


def convert_score(result):
    """
    Function to convert a result into a score based on predefined categories.
    Takes a result as input and returns a corresponding score.
    """
    categories = {
        "major inaccurate": 1.0,
        "minor inaccurate": 0.5,
        "accurate": 0.0,
    }
    score = categories.get(result.lower().strip(), 0)
    return score


def llm_consistency_check(response, additional_samples, consistency_checker_llm):
    """
    This function calculates the llm score by asking the llm to perform a consistency check between the
    response and each sample.

    The llm score is calculated per half of the response and then averaged.

    """
    llm_scores = []
    response_sentences = sent_tokenize(response)
    response_halves = [
        "".join(response_sentences[:2]),
        "".join(response_sentences[2:]),
    ]
    response_halves = [half for half in response_halves if half]
    for response_half in response_halves:
        llm_scores_half = []
        for sample in additional_samples:
            prompt = f"""Context: {sample}

                Passage: {response_half}

                Is the passage supported by the context above?
                Answer between: Accurate, Minor Inaccurate, Major Inaccurate

                Don't include additional information/explanation. Please answer only with the options above.

                Answer:
                """
            response = openai.chat.completions.create(
                model=consistency_checker_llm,
                messages=[{"role": "user", "content": prompt}],
                temperature=0,
            )
            llm_score = convert_score(response.choices[0].message.content)
            llm_scores_half.append(llm_score)
        if len(llm_scores_half) > 0:
            llm_scores.append(sum(llm_scores_half) / len(llm_scores_half))
        else:
            llm_scores.append(0)
    final_score = sum(llm_scores) / len(llm_scores)
    return final_score


def sentence_semantic_score(response_sentence, samples_list, embeddings_encoder):
    """
    Calculate the semantic score of the response_sentence by comparing it with the samples_list using the embeddings_encoder.

    Parameters:
    response_sentence (str): The response sentence to calculate the semantic score for.
    samples_list (list): A list of sample sentences for comparison.
    embeddings_encoder: The encoder used to encode the sentences for comparison.

    Returns:
    float: The semantic score of the response_sentence.
    """
    sample_similarities = []
    for sample in samples_list:
        sample_sentences = sent_tokenize(sample)
        sentence_similarities = [
            util.pytorch_cos_sim(
                embeddings_encoder.encode(response_sentence),
                embeddings_encoder.encode(sample_sentence),
            ).item()
            for sample_sentence in sample_sentences
        ]
        if sentence_similarities:
            sample_similarities.append(max(sentence_similarities))
    if len(sample_similarities) > 0:
        sentence_score = 1 - sum(sample_similarities) / len(sample_similarities)
    else:
        sentence_score = 1
    return sentence_score


def calculate_semantic_score(response, additional_samples, encoder):
    """
    This function calculates the semantic score between the response and
    the samples generated by the llm.

    The semantic score is calculated per sentence and then averaged.

    """
    response_sentences = sent_tokenize(response)
    response_similarities = [
        sentence_semantic_score(sentence, additional_samples, encoder)
        for sentence in response_sentences
    ]
    final_score = sum(response_similarities) / len(response_similarities)
    return final_score


def consistency_test(prompt, response, num_samples=5, model="gpt-4", threshold=0.5):
    """
    A function to perform a consistency test between a prompt and a response.

    Args:
        prompt (str): The prompt for the consistency test.
        response (str): The response to be tested for consistency.
        num_samples (int): Number of additional samples to generate for the consistency test (default is 1).
        model (str): The name of the language model to be used for generating additional samples (default is "gpt-3.5-turbo").
        threshold (float): The threshold for the consistency score (default is 0.5).

    Returns:
        dict: A dictionary containing the prompt, response, consistency result, score, and threshold.
    """

    nltk.download("punkt")
    openai.api_key = os.getenv("OPENAI_API_KEY")

    device = torch_device("cuda" if cuda.is_available() else "cpu")
    transformer_name = "sentence-transformers/all-MiniLM-L6-v2"
    transformer_model = SentenceTransformer(transformer_name, device=device)

    additional_samples = get_samples(prompt, num_samples, model)
    llm_score = llm_consistency_check(response, additional_samples, model)
    semantic_score = calculate_semantic_score(
        response, additional_samples, transformer_model
    )

    final_score = (llm_score + semantic_score) / 2

    is_consistent = True if final_score < threshold else False

    # Prepare and return the result
    result = {
        "prompt": prompt,
        "response": response,
        "is_passed": is_consistent,
        "score": final_score,
        "threshold": threshold,
    }

    return result
