#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
A collection of functions for studying sentiment
"""
import torch
from transformers import pipeline
import bilbystats as bs
from importlib import resources
from dotenv import load_dotenv
import os

# Use importlib to get a temporary path to the installed .env file
env_path = resources.files(
    "bilbystats.defaults").joinpath("local_defaults.env")

# Load environment variables from the .env file
load_dotenv(dotenv_path=env_path)


def sentiment_xAI(sentence: str, label: str, model_name: str = 'llama',
                  counterlabel: str = None) -> str:
    """
    Generate an explanatory rationale for a sentiment label assigned to a Chinese sentence.

    This function constructs a natural language prompt to query a language model 
    for an explanation of why a given sentiment label (e.g., "positive", 
    "neutral") is appropriate for a Chinese-language sentence, especially 
    from contexts like news reports or government websites. Optionally, it can 
    also elicit a contrastive explanation comparing the primary label against 
    an alternative counterlabel.

    ---------------------------------------------------------------------------
    PARAMETERS:
    sentence : str
        The Chinese-language sentence to be analyzed.
    label : str
        The sentiment label assigned to the sentence (e.g., "positive", "neutral").
    model_name : str, optional (default='llama')
        The name of the language model to be queried for generating the explanation.
    counterlabel : str, optional (default=None)
        An alternative label to contrast against the primary label, 
        prompting an explanation of why it is less appropriate.

    ---------------------------------------------------------------------------
    RETURNS:
    out : str
        The explanation(s) generated by the language model for the sentiment classification.
    ---------------------------------------------------------------------------
    AUTHORS: Samuel Davenport
    ---------------------------------------------------------------------------
    """
    instructions = 'You are a sentiment analysis expert for Chinese language text, specializing in analyzing sentences from news reports and government websites.'

    content = f'''The following phrase has been classified (out of negative, slightly negative, neutral, slightly positive and positive) based on **sentiment the author intended to express**, not how a reader might react.

Here is the phrase: {sentence}

Here is the label: {label}

I'd like you to provide a 1-2 sentence explanation of why this label applies'''

    if counterlabel:
        content = content + \
            f'and a secondary explanation of why the label is not {counterlabel}.\n'

    content = content + \
        '''Return your answer as\nExplanation: <brief explanation (1–2 sentences) of why this label applies>\n'''

    if counterlabel:
        content = content + \
            f'''Second Explanation: The label "{label}" is more appropriate than "{counterlabel}" <brief explanation (1–2 sentences) of why>'''

    out = bs.llm_api(content, instructions, model_name)
    return out


def sentiment_explanation(sentence: str, label: str,
                          model_name: str = 'llama') -> str:
    """
    Generate a brief explanation for a sentiment label assigned to a Chinese sentence.

    This function sends a prompt to a language model to produce a short explanation
    (1–2 sentences) justifying a given sentiment label applied to a Chinese sentence.
    The explanation is framed from the perspective of the author's intended sentiment
    rather than the reader's interpretation.

    ---------------------------------------------------------------------------
    ARGUMENTS:
    sentence : str
        The Chinese sentence for which sentiment has been analyzed.
    label : str
        The sentiment label assigned to the sentence. One of:
        'negative', 'slightly negative', 'neutral', 'slightly positive', 'positive'.
    model_name : str, optional (default='llama')
        The name of the language model used to generate the explanation.

    ---------------------------------------------------------------------------
    OUTPUT:
    out : str
        A 1–2 sentence explanation of why the sentiment label applies to the given sentence.
    ---------------------------------------------------------------------------
    AUTHORS: Samuel Davenport
    ---------------------------------------------------------------------------
    """
    instructions = f'''You are a sentiment analysis expert, specializing in analyzing sentences from news reports and government websites. 
You will be given a phrase which has been classified (out of negative, slightly negative, neutral, slightly positive and positive) based 
on **sentiment the author intended to express**, not how a reader might react.
Give your best explanation possible even if the label doesn't clearly apply.'''

    content = f'''
Here is the phrase: {sentence}

Here is the assigned label: {label}

I'd like you to provide a 1-2 sentence explanation of why this label applies.

Return your answer as
<brief explanation (1–2 sentences) of why this label applies>'''

    out = bs.llm_api(content, instructions, model_name)

    return out


def chatgpt_sentiment(content):
    instructions = '''You are a sentiment analysis expert, specializing in analyzing the sentiment of text from news reports and government websites. 
Classify each input sentence based on the **sentiment the author intended to express**, not how a reader might react.

- Output ONLY the JSON object, nothing else
- Do NOT use markdown code blocks or backticks
- Start your response directly with the opening brace {

Expected output format:
{
  "label": "<one of: negative, slightly negative, neutral, slightly positive, positive>",
  "explanation": "<brief explanation in english (1–2 sentences) of why this label applies>"
}
'''
    output = bs.openai_api(instructions, content)
    return output


def get_sentiment_score(text):
    """
    Analyze the sentiment of a given text using FinBERT for financial sentiment analysis.

    This function uses the pre-trained FinBERT model to evaluate the sentiment of the input 
    text. The sentiment is returned as a continuous score with the following adjustments:
    - Negative sentiment: returns the negative of the model's score.
    - Neutral sentiment: returns the model's score shifted by +10.
    - Positive sentiment: returns the raw model score.

    The function automatically detects GPU availability and utilizes it for faster processing 
    when possible.

    ---------------------------------------------------------------------------
    ARGUMENTS:
    text : str
        The text for which the sentiment score is to be calculated.

    ---------------------------------------------------------------------------
    OUTPUT:
    out : float
        A continuous sentiment score representing the text's sentiment:
        - Negative sentiment: negative of the raw score.
        - Neutral sentiment: raw score + 10.
        - Positive sentiment: raw score.

    ---------------------------------------------------------------------------
    Copyright (C) - 2025 - Samuel Davenport
    ---------------------------------------------------------------------------
    """
    # Check if a GPU is available, otherwise default to CPU
    device = 0 if torch.cuda.is_available() else -1

    # Load FinBERT for financial sentiment analysis, using GPU if available
    sentiment_analyzer = pipeline(
        'sentiment-analysis',
        model='yiyanghkust/finbert-tone',
        device=device
    )

    result = sentiment_analyzer(text)
    # Return the continuous sentiment score
    if result[0]['label'] == 'Negative':
        out = -result[0]['score']
    elif result[0]['label'] == 'Neutral':
        out = result[0]['score'] + 10
    else:
        out = result[0]['score']
    return out


def predict_sentence_sentiment(sentences: list, model_dict: dict = None, padding = False):
    """
    Predict sentiment labels for a list of sentences using a pre-trained model.

    This function uses a pre-trained sentence-level sentiment analysis model to
    generate sentiment predictions for the input sentences. If no model is 
    provided, a default pre-trained model is loaded. The predicted sentiment
    labels are adjusted by subtracting 2 from the raw predictions.

    ---------------------------------------------------------------------------
    ARGUMENTS:
    sentences : list of str
        A list of sentences for which sentiment predictions are to be made.
    model_dict : dict, optional
        A dictionary containing the pre-trained model and any required metadata. 
        If None, a default model is loaded using `bs.load_pretrained_model`.

    ---------------------------------------------------------------------------
    OUTPUT:
    out : dict
        A dictionary containing sentiment prediction results. The key 
        'pred_labels' contains adjusted sentiment labels.
    ---------------------------------------------------------------------------
    AUTHORS: Samuel Davenport
    ---------------------------------------------------------------------------
    """
    if model_dict == None:
        model_dict = bs.load_pretrained_model('sentence_sentiment')
        
    out = bs.predict(sentences, model_dict, padding = padding)
    out['pred_labels'] -= 2
    return out


def sentiment_ts(text: str, model_dict: dict = None, padding = False):
    """
    Generate a sentiment time series from a Chinese text.

    This function splits a Chinese text into sentences, applies a sentence-level 
    sentiment model to each sentence, and returns a time series of sentiment labels.
    If no model is provided, a default pre-trained model is loaded.

    ---------------------------------------------------------------------------
    ARGUMENTS:
    text : str
        A string of Chinese text to be analyzed for sentiment.
    model_dict : dict, optional
        A dictionary containing the pre-trained model and related metadata.
        If None, a default model is loaded using `bs.load_pretrained_model`.

    ---------------------------------------------------------------------------
    OUTPUT:
    ts : np.ndarray
        An array of sentiment labels corresponding to the sentences in the input text.
    ---------------------------------------------------------------------------
    AUTHORS: Samuel Davenport
    ---------------------------------------------------------------------------
    """
    if model_dict == None:
        model_dict = bs.load_pretrained_model('sentence_sentiment')
        
    sentences = bs.get_sentences(text, language = 'chinese')
    predictions = predict_sentence_sentiment(sentences, model_dict, padding = padding)
    ts = predictions['pred_labels']
    logits = predictions['logits']
    return ts, logits, sentences

def sentiment_mapping(df, orig_covariate, new_covariate, mapping = 'coarse_sentiment', way = 0):
    """
    Maps sentiment categories from the original column to numerical values
    and saves them in a new column.

    Parameters:
    - df (pd.DataFrame): The input DataFrame.
    - orig_covariate (str): The name of the column containing sentiment labels.
    - new_covariate (str): The name of the new column to store mapped values.

    Returns:
    - pd.DataFrame: The DataFrame with the new mapped column added.
    """
    #if way == 0:
    if mapping == 'fine_sentiment':
        direction_mapping = {
            'negative': 0,
            'slightly negative': 1,
            'neutral': 2,
            'slightly positive': 3,
            'positive': 4
            }
    if mapping == 'very_fine_sentiment':
        direction_mapping = {
            'very negative': 0,
            'negative': 1,
            'slightly negative': 2,
            'neutral': 3,
            'slightly positive': 4,
            'positive': 5,
            'very positive': 6
            }
    elif mapping == 'coarse_sentiment':
        direction_mapping = {
            'negative': 0,
            'neutral': 1,
            'positive': 2
            }
    
    if way == 1:
        direction_mapping = {v: k for k, v in direction_mapping.items()}

    df[new_covariate] = df[orig_covariate].map(direction_mapping)
    
    return df

def df_map(df, orig_covariate, new_covariate, orig_labels, new_labels):
    mapping = dict(zip(orig_labels, new_labels))
    df[new_covariate] = df[orig_covariate].map(mapping)
    return df