import warnings
import numpy as np
import pandas as pd
import os
import json
import random
import gradio as gr
import torch
from sklearn.preprocessing import OneHotEncoder
from transformers import AutoModelForSequenceClassification, AutoTokenizer, AutoModelForCausalLM, pipeline
from deap import base, creator, tools, algorithms
import nltk
from nltk.sentiment import SentimentIntensityAnalyzer
from nltk.tokenize import word_tokenize
from nltk.tag import pos_tag
from nltk.chunk import ne_chunk
from textblob import TextBlob
import matplotlib.pyplot as plt
import seaborn as sns
from accelerate import init_empty_weights, load_checkpoint_and_dispatch

warnings.filterwarnings('ignore', category=FutureWarning, module='huggingface_hub.file_download')

# Download necessary NLTK data
nltk.download('vader_lexicon', quiet=True)
nltk.download('punkt', quiet=True)
nltk.download('averaged_perceptron_tagger', quiet=True)
nltk.download('maxent_ne_chunker', quiet=True)
nltk.download('words', quiet=True)

# Initialize Example Dataset (For Emotion Prediction)
data = {
    'context': [
        'I am happy', 'I am sad', 'I am angry', 'I am excited', 'I am calm',
        'I am feeling joyful', 'I am grieving', 'I am feeling peaceful', 'I am frustrated',
        'I am determined', 'I feel resentment', 'I am feeling glorious', 'I am motivated',
        'I am surprised', 'I am fearful', 'I am trusting', 'I feel disgust', 'I am optimistic',
        'I am pessimistic', 'I feel bored', 'I am envious'
    ],
    'emotion': [
        'joy', 'sadness', 'anger', 'joy', 'calmness', 'joy', 'grief', 'calmness', 'anger',
        'determination', 'resentment', 'glory', 'motivation', 'surprise', 'fear', 'trust',
        'disgust', 'optimism', 'pessimism', 'boredom', 'envy'
    ]
}
df = pd.DataFrame(data)

# Encoding the contexts using One-Hot Encoding (memory-efficient)
try:
    encoder = OneHotEncoder(handle_unknown='ignore', sparse_output=True)
except TypeError:
    encoder = OneHotEncoder(handle_unknown='ignore', sparse=True)
contexts_encoded = encoder.fit_transform(df[['context']])

# Encoding emotions
emotions_target = pd.Categorical(df['emotion']).codes
emotion_classes = pd.Categorical(df['emotion']).categories

# Load pre-trained BERT model for emotion prediction
emotion_prediction_model = AutoModelForSequenceClassification.from_pretrained("bhadresh-savani/distilbert-base-uncased-emotion")
emotion_prediction_tokenizer = AutoTokenizer.from_pretrained("bhadresh-savani/distilbert-base-uncased-emotion")

# Load pre-trained large language model and tokenizer for response generation with increased context window
response_model_name = "gpt2-xl"
response_tokenizer = AutoTokenizer.from_pretrained(response_model_name)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
response_model = load_checkpoint_and_dispatch(AutoModelForCausalLM, response_model_name, device_map="auto")

# Set the pad token
response_tokenizer.pad_token = response_tokenizer.eos_token

# Enhanced Emotional States
emotions = {
    'joy': {'percentage': 10, 'motivation': 'positive', 'intensity': 0},
    'sadness': {'percentage': 10, 'motivation': 'negative', 'intensity': 0},
    'anger': {'percentage': 10, 'motivation': 'traumatic or strong', 'intensity': 0},
    'fear': {'percentage': 10, 'motivation': 'defensive', 'intensity': 0},
    'love': {'percentage': 10, 'motivation': 'affectionate', 'intensity': 0},
    'surprise': {'percentage': 10, 'motivation': 'unexpected', 'intensity': 0},
    'neutral': {'percentage': 40, 'motivation': 'balanced', 'intensity': 0},
}

total_percentage = 100
emotion_history_file = 'emotion_history.json'
global conversation_history
conversation_history = []
max_history_length = 1000  # Increase the maximum history length

def load_historical_data(file_path=emotion_history_file):
    if os.path.exists(file_path):
        with open(file_path, 'r') as file:
            return json.load(file)
    return []

def save_historical_data(historical_data, file_path=emotion_history_file):
    with open(file_path, 'w') as file:
        json.dump(historical_data, file)

emotion_history = load_historical_data()


def update_emotion(emotion, percentage, intensity):
    emotions[emotion]['percentage'] += percentage
    emotions[emotion]['intensity'] = intensity

    # Normalize percentages
    total = sum(e['percentage'] for e in emotions.values())
    for e in emotions:
        emotions[e]['percentage'] = (emotions[e]['percentage'] / total) * 100

def normalize_context(context):
    return context.lower().strip()

# Create FitnessMulti and Individual outside of evolve_emotions
creator.create("FitnessMulti", base.Fitness, weights=(-1.0, -0.5, -0.2))
creator.create("Individual", list, fitness=creator.FitnessMulti)

def evaluate(individual):
    emotion_values = individual[:len(emotions)]
    intensities = individual[len(emotions):]
    
    total_diff = abs(100 - sum(emotion_values))
    intensity_range = max(intensities) - min(intensities)
    emotion_balance = max(emotion_values) - min(emotion_values)
    
    return total_diff, intensity_range, emotion_balance

def evolve_emotions():
    toolbox = base.Toolbox()
    toolbox.register("attr_float", random.uniform, 0, 100)
    toolbox.register("attr_intensity", random.uniform, 0, 10)
    toolbox.register("individual", tools.initCycle, creator.Individual,
                     (toolbox.attr_float,) * len(emotions) +
                     (toolbox.attr_intensity,) * len(emotions), n=1)
    toolbox.register("population", tools.initRepeat, list, toolbox.individual)
    toolbox.register("mate", tools.cxTwoPoint)
    toolbox.register("mutate", tools.mutGaussian, mu=0, sigma=1, indpb=0.2)
    toolbox.register("select", tools.selNSGA2)
    toolbox.register("evaluate", evaluate)

    population = toolbox.population(n=100)
    algorithms.eaMuPlusLambda(population, toolbox, mu=50, lambda_=100, cxpb=0.7, mutpb=0.2, ngen=50,
                               stats=None, halloffame=None, verbose=False)

    best_individual = tools.selBest(population, k=1)[0]
    emotion_values = best_individual[:len(emotions)]
    intensities = best_individual[len(emotions):]

    for i, (emotion, data) in enumerate(emotions.items()):
        data['percentage'] = emotion_values[i]
        data['intensity'] = intensities[i]

     # Normalize percentages
    total = sum(e['percentage'] for e in emotions.values())
    for e in emotions:
        emotions[e]['percentage'] = (emotions[e]['percentage'] /total) * 100

def update_emotion_history(emotion, percentage, intensity, context):
    entry = {
        'emotion': emotion,
        'percentage': percentage,
        'intensity': intensity,
        'context': context,
        'timestamp': pd.Timestamp.now().isoformat()
    }
    emotion_history.append(entry)
    save_historical_data(emotion_history)

# Adding 443 features
additional_features = {}
for i in range(443):
    additional_features[f'feature_{i+1}'] = 0

def feature_transformations():
    global additional_features
    for feature in additional_features:
        additional_features[feature] += random.uniform(-1, 1)

def generate_response(input_text, ai_emotion, conversation_history):
    # Prepare a prompt based on the current emotion and input
    prompt = f"You are an AI assistant created by Sephiroth Baptiste, and you are currently feeling {ai_emotion}. Your response should reflect this emotion. Human: {input_text}\nAI:"
    
    # Add conversation history to the prompt
    for entry in conversation_history[-100:]:  # Use last 100 entries for context
        prompt = f"Human: {entry['user']}\nAI: {entry['response']}\n" + prompt
    
    inputs = response_tokenizer(prompt, return_tensors="pt", padding=True, truncation=True, max_length=8192).to(device)
    
    # Adjust generation parameters based on emotion
    temperature = 0.7
    if (ai_emotion == 'anger'):
        temperature = 0.9  # more intense
    elif (ai_emotion == 'calmness'):
        temperature = 0.5  # more composed

    outputs = response_model.generate(
        inputs['input_ids'],
        max_length=500,
        num_return_sequences=1,
        temperature=temperature,
        pad_token_id=response_tokenizer.eos_token_id
    )

    response = response_tokenizer.decode(outputs[0], skip_special_tokens=True)
    response = response.replace(prompt, "").strip()
    conversation_history.append({'user': input_text, 'response': response})
    return response

def process_input(input_text):
    # Predict emotion of the input text
    inputs = emotion_prediction_tokenizer(input_text, return_tensors='pt', truncation=True, padding=True).to(device)
    with torch.no_grad():
        logits = emotion_prediction_model(**inputs).logits

    predicted_class_id = torch.argmax(logits, dim=1).item()
    predicted_emotion = emotion_classes[predicted_class_id]

    # Update emotion percentages and intensities based on predicted emotion
    update_emotion(predicted_emotion, 5, 5)  # Example increment values
    update_emotion_history(predicted_emotion, emotions[predicted_emotion]['percentage'], emotions[predicted_emotion]['intensity'], input_text)

    # Evolve emotions
    evolve_emotions()

    # Generate response
    response = generate_response(input_text, predicted_emotion, conversation_history)
    
    # Feature transformations
    feature_transformations()

    return response

def plot_emotion_distribution():
    emotion_labels = list(emotions.keys())
    emotion_percentages = [emotions[emotion]['percentage'] for emotion in emotion_labels]
    emotion_intensities = [emotions[emotion]['intensity'] for emotion in emotion_labels]

    fig, ax1 = plt.subplots(figsize=(10, 6))

    ax2 = ax1.twinx()
    ax1.bar(emotion_labels, emotion_percentages, color='b', alpha=0.6)
    ax2.plot(emotion_labels, emotion_intensities, color='r', marker='o', linestyle='dashed', linewidth=2)

    ax1.set_xlabel('Emotion')
    ax1.set_ylabel('Percentage', color='b')
    ax2.set_ylabel('Intensity', color='r')

    plt.title('Emotion Distribution and Intensities')
    plt.show()

def clear_conversation_history():
    global conversation_history
    conversation_history = []

# Function to display the history of the 10 most recent conversations
def display_recent_conversations():
    num_conversations = min(len(conversation_history), 10)
    recent_conversations = conversation_history[-num_conversations:]

    conversation_text = ""
    for i, conversation in enumerate(recent_conversations, start=1):
        conversation_text += f"Conversation {i}:\n"
        conversation_text += f"User: {conversation['user']}\n"
        conversation_text += f"AI: {conversation['response']}\n\n"

    return conversation_text.strip()

with gr.Blocks() as chatbot:
    gr.Markdown("# AI Chatbot with Enhanced Emotions")
    
    with gr.Row():
        with gr.Column():
            input_text = gr.Textbox(label="Input Text")
            response_text = gr.Textbox(label="Response", interactive=False)
            send_button = gr.Button("Send")
            clear_button = gr.Button("Clear Conversation History")
    
    with gr.Row():
        recent_conversations = gr.Textbox(label="Recent Conversations", interactive=False)
        update_button = gr.Button("Update Recent Conversations")
    
    with gr.Row():
        emotion_plot = gr.Plot(label="Emotion Distribution and Intensities")
        update_plot_button = gr.Button("Update Emotion Plot")
    
    send_button.click(fn=process_input, inputs=input_text, outputs=response_text)
    clear_button.click(fn=clear_conversation_history)
    update_button.click(fn=display_recent_conversations, outputs=recent_conversations)
    update_plot_button.click(fn=plot_emotion_distribution, outputs=emotion_plot)

chatbot.launch()