import logging
import os
import sys
import nltk
from nltk.translate.meteor_score import meteor_score

# Configuración de NLTK
os.environ["NLTK_DATA"] = "/root/nltk_data"
nltk.data.path.append('/root/nltk_data')
logging.getLogger('nltk').setLevel(logging.CRITICAL)

nltk.download('punkt_tab', quiet=True)
nltk.download('wordnet', quiet=True)

from sklearn.feature_extraction.text import CountVectorizer
from sklearn.metrics.pairwise import cosine_similarity
import numpy as np
from sentence_transformers import SentenceTransformer

def calculate_meteor_score(target_output, model_output, language='english', debug=False):
    """
    Calculate the METEOR score between a single target output (reference) and a model output (hypothesis).

    Parameters:
        target_output (str): The reference output (expected response).
        model_output (str): The output generated by the model (hypothesis).
        language (str, optional): The language for which to calculate the METEOR score. Defaults to 'english'.
        debug (bool, optional): If set to True, prints detailed comparison between the target and model output. Defaults to False.

    Returns:
        float: The METEOR score for the given pair of target and model output.

    METEOR, the Metric for Evaluation of Translation with Explicit ORdering, is a semantic evaluation metric 
    that considers word order and synonyms when comparing two pieces of text. It improves over metrics like BLEU 
    by recognizing paraphrases and matching semantically equivalent content, making it particularly valuable in 
    tasks where meaning is more important than exact lexical matches.

    Example usage:
        reference = "This is an example sentence"
        hypothesis = "This is a sample phrase"
        print(calculate_meteor_score(reference, hypothesis, language='english', debug=True))
    """
    # Download necessary NLTK resources for tokenization and synonym matching
    nltk.download('punkt', quiet=True)  # Tokenizer for various languages
    nltk.download('wordnet', quiet=True)  # WordNet for English synonym matching

    # Tokenize and handle empty or null inputs
    if not target_output or not target_output.strip():
        target_tokens = [[]]  # Empty list if the reference string is empty
    else:
        target_tokens = [nltk.word_tokenize(target_output, language=language)]  # Tokenization based on language

    if not model_output or not model_output.strip():
        model_output_tokens = []  # Empty list if the hypothesis string is empty
    else:
        model_output_tokens = nltk.word_tokenize(model_output, language=language)

    # Calculate the METEOR score
    score = meteor_score(target_tokens, model_output_tokens)

    if debug:
        print(f"METEOR Score between '{target_output}' and '{model_output}' (language={language}): {score:.4f}")

    return score


def calculate_cosine_similarity(target_output, model_output, model_type="all-mpnet-base-v2", debug=False):
    """
    Calculate the cosine similarity score between a target output (reference) and a model output (hypothesis).

    Parameters:
        target_output (str): The reference output (expected response).
        model_output (str): The output generated by the model (hypothesis).
        model_type (str, optional): The model type to use for calculating embeddings. 
                                    Options: "multi-qa-MiniLM-L6-cos-v1", "all-mpnet-base-v2", "gtr-t5-xxl". Defaults to "all-mpnet-base-v2".
        debug (bool, optional): If set to True, prints detailed embeddings and the resulting similarity score. Defaults to False.

    Returns:
        float: The cosine similarity score between the target and model output.

    Cosine similarity measures the cosine of the angle between two vectors in a multi-dimensional space,
    providing a numerical value between -1 and 1 that represents the similarity between the target output 
    and the model output.

    **Available model types and their performance:**
    - **"multi-qa-MiniLM-L6-cos-v1"**:
        - Speed: 14,200 sentences/second.
        - Model size: 80 MB.
        - Advantage: Lightweight and fast.
        - Disadvantage: Lower performance on complex semantic tasks.
    - **"all-mpnet-base-v2"** (default):
        - Speed: 2,800 sentences/second.
        - Model size: 420 MB.
        - Advantage: Good balance between accuracy and speed.
        - Disadvantage: Requires more resources than smaller models.
    - **"gtr-t5-xxl"**:
        - Speed: 50 sentences/second.
        - Model size: 9.2 GB.
        - Advantage: High representation capacity for capturing semantic meaning.
        - Disadvantage: Slow and resource-heavy.

    **Example usage:**
        target = "The impact of climate change is significant."
        model_output = "Climate change has a huge effect on the environment."
        score = calculate_cosine_similarity(target, model_output, model_type='all-mpnet-base-v2', debug=True)
        print(f"Cosine Similarity Score: {score:.4f}")
    """
    from torch.nn.functional import cosine_similarity
    import torch

    # Load the specified model for embedding generation
    if model_type == "multi-qa-MiniLM-L6-cos-v1":
        model = SentenceTransformer("multi-qa-MiniLM-L6-cos-v1")
    elif model_type == "all-mpnet-base-v2":
        model = SentenceTransformer("all-mpnet-base-v2")
    elif model_type == "gtr-t5-xxl":
        model = SentenceTransformer("gtr-t5-xxl")
    else:
        raise ValueError("Invalid model_type. Please choose from 'multi-qa-MiniLM-L6-cos-v1', 'all-mpnet-base-v2', or 'gtr-t5-xxl'.")

    # Handle edge cases: if either string is empty or None, return 0.0
    if not target_output or not model_output:
        if debug:
            print(f"Either the target or model output is empty. Returning cosine similarity score: 0.0")
        return 0.0

    # Generate embeddings for the target and model output
    target_embedding = model.encode(target_output, convert_to_tensor=True)
    model_output_embedding = model.encode(model_output, convert_to_tensor=True)

    # Calculate the cosine similarity
    similarity_score = cosine_similarity(target_embedding, model_output_embedding, dim=0).item()

    if debug:
        print(f"Target Output: '{target_output}'")
        print(f"Model Output: '{model_output}'")
        print(f"Cosine Similarity Score: {similarity_score:.4f}")

    return similarity_score



def calculate_metrics_from_logprobs(logprobs, metrics, target_output=None, model_output=None, model_type="all-mpnet-base-v2", debug=False):
    """
    Calculates multiple metrics such as log probability, perplexity, METEOR score, and cosine similarity.

    Args:
        logprobs (list of floats): Log-probabilities for each token in the sequence.
        metrics (list of str): Metrics to calculate ("log_probability", "perplexity", etc.).
        target_output (str): Expected target sequence (for metrics such as METEOR).
        model_output (str): Sequence generated by the model (for comparison metrics).
        model_type (str): Model to use for cosine similarity (optional).
        debug (bool): If True, prints debugging details.

    Returns:
        dict: Results with the calculated metrics.
    """
    import math

    results = {}
    
    # Calculate the average log-probability
    avg_logprob = sum(logprobs) / len(logprobs) if logprobs else float("-inf")

    if "log_probability" in metrics:
        results["log_probability"] = avg_logprob

    if "perplexity" in metrics:
        # Perplexity is computed as the exponential of the negative average log-probability
        perplexity = math.exp(-avg_logprob) if avg_logprob != float("-inf") else float("inf")
        results["perplexity"] = perplexity

    if "meteor_score" in metrics and target_output and model_output:
        # Calculate the METEOR score if both target and model output are provided
        results["meteor_score"] = calculate_meteor_score(target_output, model_output, debug=debug)

    if "cosine_similarity" in metrics and target_output and model_output:
        # Calculate cosine similarity if both target and model output are provided
        results["cosine_similarity"] = calculate_cosine_similarity(target_output, model_output, model_type=model_type, debug=debug)

    
    return results