import pandas as pd
from typing import Optional, List, Union
import numpy as np
import bilbystats as bs
import asyncio
from sacrebleu.metrics import BLEU
import logging

from bert_score import score as bertyscorey

import torch

#import nltk

def translate(text: str, model_name: str = "google/gemini-2.5-flash-lite", languageout: str = "English", languagein: str = "Chinese") -> str:
    """
    Translate input text into a specified language using a language model.

    This function sends the input text to a language model API with an 
    instruction prompt that requests translation into the desired output 
    language. The translation result is returned as plain text.

    ---------------------------------------------------------------------------
    ARGUMENTS:
    text : str
        The input text to be translated.
    model_name : str, optional (default="gpt-4o")
        The name of the language model to be used for translation.
    languageout : str, optional (default="English")
        The target language for translation.

    ---------------------------------------------------------------------------
    OUTPUT:
    translation : str
        The translated text in the target language.
    ---------------------------------------------------------------------------
    AUTHORS: Samuel Davenport
    ---------------------------------------------------------------------------
    """
    if type(text) == list:
        translated_texts = []
        for i in range(len(text)):
            translated_texts.append(bs.translate(text[i], model_name, languageout, languagein = languagein))
        return translated_texts
        
    if languagein == "Chinese":
        addon = "(no Chinese characters)"
    else: 
        addon = ""

    instructions = '''Translate the provided ''' + languagein + ''' text into English with these requirements:
1. Translate all content completely and accurately—do not omit, summarize, or skip any details
2. Output must contain only English text ''' + addon + '''
3. Maintain the original tone, style, and nuance while ensuring the English reads naturally and idiomatically
4. Preserve formatting, paragraph breaks, and structure from the original text
Only return the English translation. Do not add anything else.'''
    translation = bs.llm_api(text, instructions, model_name=model_name)

    return translation

async def batch_translate(texts: str, model_name: str = "google/gemini-2.5-flash-lite", languageout: str = "English", languagein: str = "Chinese") -> str:
    """
    Translate input text into a specified language using a language model.

    This function sends the input text to a language model API with an 
    instruction prompt that requests translation into the desired output 
    language. The translation result is returned as plain text.

    ---------------------------------------------------------------------------
    ARGUMENTS:
    text : str
        The input text to be translated.
    model_name : str, optional (default="gpt-4o")
        The name of the language model to be used for translation.
    languageout : str, optional (default="English")
        The target language for translation.

    ---------------------------------------------------------------------------
    OUTPUT:
    translation : str
        The translated text in the target language.
    ---------------------------------------------------------------------------
    AUTHORS: Samuel Davenport
    ---------------------------------------------------------------------------
    """
        
    if languagein == "Chinese":
        addon = "(no Chinese characters)"
    else: 
        addon = ""

    instructions = '''Translate the provided ''' + languagein + ''' text into English with these requirements:
1. Translate all content completely and accurately—do not omit, summarize, or skip any details
2. Output must contain only English text ''' + addon + '''
3. Maintain the original tone, style, and nuance while ensuring the English reads naturally and idiomatically
4. Preserve formatting, paragraph breaks, and structure from the original text
Only return the English translation. Do not add anything else.'''
    translation = await bs.batch_llm_api_async(texts, instructions, model_name=model_name)

    return translation

def translate_orig(text: str, model_name: str = "google/gemini-2.5-flash-lite", languageout: str = "English", languagein: str = "Chinese") -> str:
    """
    Translate input text into a specified language using a language model.

    This function sends the input text to a language model API with an 
    instruction prompt that requests translation into the desired output 
    language. The translation result is returned as plain text.

    ---------------------------------------------------------------------------
    ARGUMENTS:
    text : str
        The input text to be translated.
    model_name : str, optional (default="gpt-4o")
        The name of the language model to be used for translation.
    languageout : str, optional (default="English")
        The target language for translation.

    ---------------------------------------------------------------------------
    OUTPUT:
    translation : str
        The translated text in the target language.
    ---------------------------------------------------------------------------
    AUTHORS: Samuel Davenport
    ---------------------------------------------------------------------------
    """
    if type(text) == list:
        translated_texts = []
        for i in range(len(text)):
            translated_texts.append(bs.translate(text[i], model_name, languageout))
        return translated_texts
        
    instructions = "You are a translator. Only return the " + \
        languageout+" translation. Do not add anything else."
    translation = bs.llm_api(text, instructions, model_name=model_name)
    if model_name == 'gemini':
        translation = translation.replace('\n', '')

    return translation

async def translate_chinese(texts,  model_name: str = "gpt-4o", max_concurrent: int = 75, 
                            max_retries: int = 3, base_delay: float = 1.0, temperature = 0, timeout = 5):
    instructions ='''Translate the provided Chinese text into English with these requirements:
1. Translate all content completely and accurately—do not omit, summarize, or skip any details
2. Output must contain only English text (no Chinese characters)
3. Maintain the original tone, style, and nuance while ensuring the English reads naturally and idiomatically
4. Preserve formatting, paragraph breaks, and structure from the original text
Only return the English translation. Do not add anything else.'''

    translations = await bs.batch_llm_api_async(
        texts=texts,
        instructions=instructions,
        model_name=model_name,
        max_concurrent=max_concurrent,
        max_retries=max_retries,
        base_delay=base_delay,
        temperature = temperature,
        timeout = timeout
    )
    
    return translations

async def translate_batch_async(texts: List[str], model_name: str = "gpt-4o", 
                                languageout: str = "English", max_concurrent: int = 75, 
                                max_retries: int = 3, 
                                base_delay: float = 1.0, temperature = 0, timeout = 5) -> List[str]:
    """
    Translate multiple texts into a specified language using async batch processing.

    This function takes a list of text strings and translates them all to the specified
    target language using the batch_llm_api_async function for efficient concurrent processing.
    It maintains the same order as the input list and handles errors gracefully.

    ---------------------------------------------------------------------------
    ARGUMENTS:
    texts : List[str]
        List of text strings to be translated.
    model_name : str, optional (default="gpt-4o")
        The name of the language model to be used for translation.
    languageout : str, optional (default="English")
        The target language for translation.
    max_concurrent : int, optional (default=5)
        Maximum number of concurrent requests to make to the API.
    batch_size : int, optional (default=10)
        Number of requests to process in each batch.
    max_retries : int, optional (default=3)
        Maximum number of retry attempts for failed requests.
    base_delay : float, optional (default=1.0)
        Base delay in seconds for exponential backoff on retries.

    ---------------------------------------------------------------------------
    OUTPUT:
    translations : List[str]
        List of translated texts in the same order as the input, with empty strings
        for any texts that failed to translate.

    ---------------------------------------------------------------------------
    AUTHORS: Samuel Davenport
    ---------------------------------------------------------------------------
    """
    # Create the translation instruction
    instructions = f"You are a translator. Only return the {languageout} translation. Do not add anything else."
    
    # Use the existing batch_llm_api_async function
    translations = await bs.batch_llm_api_async(
        texts=texts,
        instructions=instructions,
        model_name=model_name,
        max_concurrent=max_concurrent,
        max_retries=max_retries,
        base_delay=base_delay,
        temperature = temperature,
        timeout = timeout
    )
    
    # Post-process translations if using Gemini (remove newlines as in original function)
    if model_name == 'gemini':
        translations = [translation.replace('\n', '') for translation in translations]
    
    return translations


def translate_df(df: pd.DataFrame, covariate: str, languageout: str, model_name: str = "gpt-4o") -> pd.DataFrame:
    """
    Translate a column in a DataFrame to a specified language and add as a new column.

    This function takes a DataFrame column containing text, translates each entry
    to the specified output language using a language model, and adds the translations
    as a new column with the naming convention '{original_column}_{language}'.

    ---------------------------------------------------------------------------
    ARGUMENTS:
    df : pd.DataFrame
        The input DataFrame containing the column to be translated.
    covariate : str
        The name of the column in the DataFrame to translate.
    languageout : str
        The target language for translation (e.g., "Chinese", "Spanish", "French").
    model_name : str, optional (default="gpt-4o")
        The name of the language model to use for translation.

    ---------------------------------------------------------------------------
    OUTPUT:
    df : pd.DataFrame
        The original DataFrame with an additional column containing the translations,
        named '{covariate}_{languageout}'.

    ---------------------------------------------------------------------------
    RAISES:
    ValueError:
        If the specified covariate column does not exist in the DataFrame.

    ---------------------------------------------------------------------------
    AUTHORS: Samuel Davenport
    ---------------------------------------------------------------------------
    """
    # Check if the covariate column exists
    if covariate not in df.columns:
        raise ValueError(f"Column '{covariate}' not found in DataFrame")

    # Create the new column name
    new_column_name = f"{covariate}_{languageout}"

    # Initialize the new column
    df[new_column_name] = None

    print(f"Starting translation of '{covariate}' column to {languageout}...")
    print(f"Processing {len(df)} entries using {model_name}")

    # Translate each entry
    for i in range(len(df)):
        # Show progress using bilbystats loader function
        bs.loader(i, len(df))

        try:
            # Get the original text
            original_text = df.loc[i, covariate]

            # Skip if the text is null or empty
            if pd.isna(original_text) or original_text == "":
                df.loc[i, new_column_name] = None
                continue

            # Translate using bilbystats translate function
            translated_text = bs.translate(
                original_text, model_name=model_name, languageout=languageout)

            # Store the translation
            df.loc[i, new_column_name] = translated_text

        except Exception as e:
            print(f"\nError translating entry {i}: {e}")
            df.loc[i, new_column_name] = None

    # Final progress update
    bs.loader(len(df), len(df))
    print(
        f"\nTranslation complete! New column '{new_column_name}' added to DataFrame.")

    return df


# Example usage:
# Load your dataframe
# df = pd.read_csv('golddatasetsinhakhandait.csv')

# Translate the News column to Chinese
# df = translate_df(df, covariate='News', languageout='Chinese', model_name='gpt-4o')

# Translate to other languages
# df = translate_df(df, covariate='News', languageout='Spanish', model_name='gpt-4o')
# df = translate_df(df, covariate='News', languageout='French', model_name='claude')

# Display some examples
# print("\nFirst 3 translations:")
# for i in range(min(3, len(df))):
#     print(f"\nEntry {i}:")
#     print(f"English: {df.loc[i, 'News']}")
#     print(f"Chinese: {df.loc[i, 'News_Chinese']}")

# Save the updated dataframe
# df.to_csv('translated_dataset.csv', index=False)

def bleu(cands: list[str], refs: list[str]) -> float:
    """
    Compute the BLEU score for a set of candidate translations against references.

    This function calculates the BLEU (Bilingual Evaluation Understudy) score 
    using the `sacrebleu.BLEU` implementation. The references are wrapped into 
    a list of lists as required by the scoring function.

    ---------------------------------------------------------------------------
    ARGUMENTS:
    cands : list of str
        Candidate translations to be evaluated.
    refs : list of str
        Reference translations for comparison.

    ---------------------------------------------------------------------------
    OUTPUT:
    score : float
        The computed BLEU score.
    ---------------------------------------------------------------------------
    AUTHORS: Samuel Davenport
    ---------------------------------------------------------------------------
    """
    origbleu = BLEU()

    refs = [refs]

    # Compute BLEU
    score = origbleu.corpus_score(cands, refs)
    score = score.score/100
    score = np.array(score)
    
    return score

def calculate_meteor(candidates, references):
    """
    Calculate METEOR scores for multiple candidate and reference texts.
    
    Args:
        candidates (list[str]): Candidate translations
        references (list[str]): Reference translations (can be multiple per candidate)
    
    Returns:
        list[float]: METEOR scores for each candidate
    """
    scores = []
    from nltk.tokenize import word_tokenize
    
    for candidate in candidates:

        candidate_tokens = word_tokenize(candidate.lower())
        
        # Tokenize all references
        reference_tokens = [word_tokenize(ref.lower()) for ref in references]
        
        # meteor_score can handle multiple references for one candidate
        from nltk.translate.meteor_score import meteor_score
        score = meteor_score(reference_tokens, candidate_tokens)
        scores.append(score)
    
    return scores

def bert_score(candidate, reference, lang="en"):
    P, R, F1 = bertyscorey(candidate, reference, lang=lang, rescale_with_baseline=True, verbose=False)

    F1 = np.array(F1)
    
    return F1

def load_comet():
    from comet import download_model, load_from_checkpoint

    model_path = download_model("Unbabel/wmt22-comet-da")
    model = load_from_checkpoint(model_path)
    return model

def comet_score(original, candidate, reference, model = None, batch_size  = 8):
    data = [{ "src": original,
          "mt": candidate,
          "ref": reference}]
    
    if not model:
        model = load_comet()
        
    # Temporarily disable MPS detection
    original_mps_available = torch.backends.mps.is_available

    try:
        torch.backends.mps.is_available = lambda: False
        model_output = model.predict(data, batch_size=batch_size, gpus=0, num_workers=0)
    finally:
        torch.backends.mps.is_available = original_mps_available
    
    score = model_output['scores'][0]
    
    return score

def translation_metric(candidates, references, score_type = 'bleu', original_texts = None, model = None):
    import statistics
    if score_type == 'bleu':
        score = bs.bleu(candidates, references)
    elif score_type == 'averagebleu':
        listofbleu = []
        for i in range(len(candidates)):
            current_bleu = bs.bleu([candidates[i]], [references[i]])
            listofbleu.append(current_bleu)
        listofbleu = np.array(listofbleu)
        score = np.mean(listofbleu)
        
        return score, statistics.stdev(listofbleu)

    elif score_type == 'comet':
        score = bs.comet_score(original_texts, candidates, references, model = model)
        if not original_texts:
            raise ValueError("For comet scoring you also need to provide the original texts.")
    elif score_type == 'meteor':
        scores = calculate_meteor(candidates, references)
        score = sum(scores) / len(scores)
    elif score_type == 'bert_score':
        scores = bs.bert_score(candidates, references)
        score = np.mean(scores)
    else:
        raise ValueError("Not an available score type")   
        
    return score
    