import gradio as gr
import nltk
import numpy as np
import tflearn
import random
import json
import pickle
import torch
from nltk.tokenize import word_tokenize
from nltk.stem.lancaster import LancasterStemmer
from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline
import requests
from bs4 import BeautifulSoup
import pandas as pd
import geocoder  # Use geocoder to get latitude/longitude from city

# Ensure necessary NLTK resources are downloaded
nltk.download('punkt')

# Initialize the stemmer
stemmer = LancasterStemmer()

# Load intents.json
try:
    with open("intents.json") as file:
        data = json.load(file)
except FileNotFoundError:
    raise FileNotFoundError("Error: 'intents.json' file not found. Ensure it exists in the current directory.")

# Load preprocessed data from pickle
try:
    with open("data.pickle", "rb") as f:
        words, labels, training, output = pickle.load(f)
except FileNotFoundError:
    raise FileNotFoundError("Error: 'data.pickle' file not found. Ensure it exists and matches the model.")

# Build the model structure
net = tflearn.input_data(shape=[None, len(training[0])])
net = tflearn.fully_connected(net, 8)
net = tflearn.fully_connected(net, 8)
net = tflearn.fully_connected(net, len(output[0]), activation="softmax")
net = tflearn.regression(net)

# Load the trained model
model = tflearn.DNN(net)
try:
    model.load("MentalHealthChatBotmodel.tflearn")
except FileNotFoundError:
    raise FileNotFoundError("Error: Trained model file 'MentalHealthChatBotmodel.tflearn' not found.")

# Function to process user input into a bag-of-words format
def bag_of_words(s, words):
    bag = [0 for _ in range(len(words))]
    s_words = word_tokenize(s)
    s_words = [stemmer.stem(word.lower()) for word in s_words if word.lower() in words]
    for se in s_words:
        for i, w in enumerate(words):
            if w == se:
                bag[i] = 1
    return np.array(bag)

# Chat function
def chat(message, history, state):
    history = history or []
    message = message.lower()
    try:
        # Predict the tag
        results = model.predict([bag_of_words(message, words)])
        results_index = np.argmax(results)
        tag = labels[results_index]
        
        # Match tag with intent and choose a random response
        for tg in data["intents"]:
            if tg['tag'] == tag:
                responses = tg['responses']
                response = random.choice(responses)
                break
        else:
            response = "I'm sorry, I didn't understand that. Could you please rephrase?"

        # Add emoticons to the response
        emoticon_dict = {
            "joy": "😊",
            "anger": "😡",
            "fear": "😨",
            "sadness": "😔",
            "surprise": "😲",
            "neutral": "😐"
        }

        # Add the emotion-related emoticon to the response
        for tg in data["intents"]:
            if tg['tag'] == tag:
                emotion = tg.get('emotion', 'neutral')  # Default to neutral if no emotion is defined
                response = f"{response} {emoticon_dict.get(emotion, '😐')}"
                break
        
        history.append((message, response))

        # Transition to the next feature (sentiment analysis)
        state['step'] = 2  # Move to sentiment analysis
    except Exception as e:
        response = f"An error occurred: {str(e)}"

    return history, history, state

# Load pre-trained model and tokenizer for sentiment analysis
tokenizer = AutoTokenizer.from_pretrained("cardiffnlp/twitter-roberta-base-sentiment")
sentiment_model = AutoModelForSequenceClassification.from_pretrained("cardiffnlp/twitter-roberta-base-sentiment")

# Function for sentiment analysis
def analyze_sentiment(text, state):
    inputs = tokenizer(text, return_tensors="pt")
    with torch.no_grad():
        outputs = sentiment_model(**inputs)
    predicted_class = torch.argmax(outputs.logits, dim=1).item()
    sentiment = ["Negative", "Neutral", "Positive"][predicted_class]
    
    # Add emoticon to sentiment
    sentiment_emojis = {
        "Negative": "😞",
        "Neutral": "😐",
        "Positive": "😊"
    }
    sentiment_with_emoji = f"{sentiment} {sentiment_emojis.get(sentiment, '😐')}"
    
    # Transition to emotion detection
    state['step'] = 3  # Move to emotion detection and suggestions
    return sentiment_with_emoji, state

# Load pre-trained model and tokenizer for emotion detection
emotion_tokenizer = AutoTokenizer.from_pretrained("j-hartmann/emotion-english-distilroberta-base")
emotion_model = AutoModelForSequenceClassification.from_pretrained("j-hartmann/emotion-english-distilroberta-base")

# Function for emotion detection and suggestions
def detect_emotion(text, state):
    pipe = pipeline("text-classification", model=emotion_model, tokenizer=emotion_tokenizer)
    result = pipe(text)
    emotion = result[0]['label']
    
    # Provide suggestions based on detected emotion
    suggestions = provide_suggestions(emotion)
    
    # Transition to wellness professional search
    state['step'] = 4  # Move to wellness professional search
    return emotion, suggestions, state

# Suggestions based on detected emotion
def provide_suggestions(emotion):
    resources = {
        'joy': {
            'message': "You're feeling happy! Keep up the great mood! 😊",
            'articles': [
                "[Relaxation Techniques](https://www.helpguide.org/mental-health/meditation/mindful-breathing-meditation)",
                "[Dealing with Stress](https://www.helpguide.org/mental-health/anxiety/tips-for-dealing-with-anxiety)"
            ],
            'videos': "[Watch Relaxation Video](https://youtu.be/m1vaUGtyo-A)"
        },
        'anger': {
            'message': "You're feeling angry. It's okay to feel this way. Let's try to calm down. 😡",
            'articles': [
                "[Emotional Wellness Toolkit](https://www.nih.gov/health-information/emotional-wellness-toolkit)",
                "[Stress Management Tips](https://www.health.harvard.edu/health-a-to-z)"
            ],
            'videos': "[Watch Anger Management Video](https://youtu.be/MIc299Flibs)"
        },
        'fear': {
            'message': "You're feeling fearful. Take a moment to breathe and relax. 😨",
            'articles': [
                "[Mindfulness Practices](https://www.helpguide.org/mental-health/meditation/mindful-breathing-meditation)",
                "[Coping with Anxiety](https://www.helpguide.org/mental-health/anxiety/tips-for-dealing-with-anxiety)"
            ],
            'videos': "[Watch Coping Video](https://youtu.be/yGKKz185M5o)"
        },
        'sadness': {
            'message': "You're feeling sad. It's okay to take a break. 😔",
            'articles': [
                "[Emotional Wellness Toolkit](https://www.nih.gov/health-information/emotional-wellness-toolkit)",
                "[Dealing with Anxiety](https://www.helpguide.org/mental-health/anxiety/tips-for-dealing-with-anxiety)"
            ],
            'videos': "[Watch Sadness Relief Video](https://youtu.be/-e-4Kx5px_I)"
        },
        'surprise': {
            'message': "You're feeling surprised. It's okay to feel neutral! 😲",
            'articles': [
                "[Managing Stress](https://www.health.harvard.edu/health-a-to-z)",
                "[Coping Strategies](https://www.helpguide.org/mental-health/anxiety/tips-for-dealing-with-anxiety)"
            ],
            'videos': "[Watch Stress Relief Video](https://youtu.be/m1vaUGtyo-A)"
        }
    }
    return resources.get(emotion, {'message': "Stay calm. 🙂", 'articles': [], 'videos': []})

# Function to find wellness professionals
def find_wellness_professionals(location, state):
    # Geocode the location to get latitude and longitude
    g = geocoder.osm(location)  # Using OpenStreetMap's geocoding service
    if g.ok:
        location_coords = f"{g.lat},{g.lng}"
    else:
        return "Sorry, could not retrieve coordinates for the location. Please try again.", state
    
    query = "therapist OR counselor OR mental health professional OR marriage and family therapist OR psychotherapist OR psychiatrist OR psychologist in " + location
    api_key = "GOOGLE_API_KEY"  # Replace with your own API key
    radius = 50000  # 50 km radius
    
    google_places_data = get_all_places(query, location_coords, radius, api_key)
    if google_places_data:
        response = "Wellness professionals near you:\n"
        for place in google_places_data:
            response += f"- {place['name']} at {place['formatted_address']}\n"
    else:
        response = "Sorry, no wellness professionals found in your area. Please try another location."

    return response, state

# Call Google Places API
def get_all_places(query, location, radius, api_key):
    search_url = f"https://maps.googleapis.com/maps/api/place/textsearch/json?query={query}&location={location}&radius={radius}&key={api_key}"
    response = requests.get(search_url).json()
    
    if 'results' in response:
        return response['results']
    return []

# Gradio UI components
def create_ui():
    with gr.Blocks() as demo:
        state = gr.State()
        chatbot = gr.Chatbot(elem_id="chatbot", label="Mental Health Chatbot")
        message_input = gr.Textbox(placeholder="Ask me something...", label="Enter your message")
        sentiment_output = gr.Textbox(placeholder="Sentiment result", label="Sentiment")
        emotion_output = gr.Textbox(placeholder="Detected emotion", label="Emotion")
        wellness_output = gr.Textbox(placeholder="Wellness professional list", label="Wellness Professionals")
        
        message_input.submit(chat, [message_input, chatbot, state], [chatbot, chatbot, state])
        message_input.submit(analyze_sentiment, [message_input, state], [sentiment_output, state])
        sentiment_output.submit(detect_emotion, [sentiment_output, state], [emotion_output, wellness_output, state])
        return demo

# Launch the Gradio interface
demo = create_ui()
demo.launch(debug=True)