"""
Chunk Impact Test
"""

import numpy as np
from numpy import dot
from numpy.linalg import norm

from .test_utils import embedding_generator


def chunk_impact_test(contexts, response=None, threshold=0.5):
    """
    Provides the Impact score for each chunk

    Args:
        context (list) : The list of Responses given to the model
        response (str, optional): The response generated by the model.
        threshold (float): The threshold for Chunk Test.

    Returns:
        dict: A dictionary containing overall score based on contexts and response
              pass and fail based on threshold
    """
    if type(contexts) is not list:
        print("Ensure to send context as list of contexts for each prompt")
        raise ValueError("Context not in correct format")

    similarity_threshold = 0.5

    context_embeddings = [embedding_generator(c)[0] for c in contexts]
    response_embedding = embedding_generator([response])[0]
    res = []
    for i, embedding in enumerate(context_embeddings):
        score = dot(embedding, response_embedding) / (
            norm(embedding) * norm(response_embedding)
        )
        if score > similarity_threshold:
            res.append(
                {
                    "score": score,
                    "context": contexts[i],
                }
            )
    # sort the res dict by scores in descending order

    res = sorted(res, key=lambda x: x["score"], reverse=True)

    # score is the mean of res['scores']
    if res:
        relevant_score = np.mean([item["score"] for item in res])
        relevant_contexts = [item["context"] for item in res]
    else:
        relevant_score = 0
        relevant_contexts = []

    is_passed = bool(relevant_score > threshold)

    reason = ""
    for r in res:
        reason += f"Score: {r['score']:.3f} -> {r['context'][0]}\n\n"

    result = {
        "response": response,
        "score": relevant_score,
        "is_passed": is_passed,
        "context": relevant_contexts,
        "threshold": threshold,
        "reason": reason,
    }
    return result
