import streamlit as st
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras.preprocessing.sequence import pad_sequences
from tensorflow.keras.preprocessing.text import Tokenizer
import numpy as np
from transformers import pipeline, GPT2LMHeadModel, GPT2Tokenizer
import random
import os
from PIL import Image

# -----------------------------
# App Settings
# -----------------------------
st.title("🧠 AI Therapist (Multimodal)")

MAX_SEQUENCE_LENGTH = 100
BREATHING_FOLDER = "breathing_sounds"
SLEEP_AUDIO = os.path.join(BREATHING_FOLDER, "Sleep-meditation-music-free.mp3")

# -----------------------------
# Load Models
# -----------------------------
# Mental Health Text Classifier
mental_model = keras.models.load_model("mental_health_text_model.h5")
tokenizer = Tokenizer(num_words=20000, oov_token="<OOV>")
# Load tokenizer if saved separately
# tokenizer = keras.preprocessing.text.tokenizer_from_json(open("tokenizer.json").read())

# GPT-2 for empathetic replies
gpt2_model_name = "gpt2-medium"
gpt2_tokenizer = GPT2Tokenizer.from_pretrained(gpt2_model_name)
gpt2_model = GPT2LMHeadModel.from_pretrained(gpt2_model_name)

# Speech Emotion Recognition (Hugging Face pre-trained)
speech_classifier = pipeline("audio-classification", model="superb/wav2vec2-large-superb-er")  # replace with your own if needed

# Facial Emotion Detection
face_emotion_model = keras.models.load_model("face_emotion_model.h5")
face_classes = ['Angry','Disgust','Fear','Happy','Neutral','Sad','Surprise']

# -----------------------------
# Helper Functions
# -----------------------------
def predict_mental_health(text):
    seq = tokenizer.texts_to_sequences([text])
    padded = pad_sequences(seq, maxlen=MAX_SEQUENCE_LENGTH, padding='post')
    pred = mental_model.predict(padded)
    classes = ['Anxiety','Bipolar','Depression','Normal','Personality disorder','Stress','Suicidal']
    return classes[np.argmax(pred)]

def generate_gpt2_reply(user_input):
    input_ids = gpt2_tokenizer.encode(
        f"As a compassionate Japanese therapist, respond empathetically: {user_input}", return_tensors="pt"
    )
    output_ids = gpt2_model.generate(
        input_ids,
        max_length=150,
        num_return_sequences=1,
        no_repeat_ngram_size=2,
        do_sample=True,
        top_k=50,
        top_p=0.95,
        temperature=0.7
    )
    reply = gpt2_tokenizer.decode(output_ids[0], skip_special_tokens=True)
    return reply.replace(f"As a compassionate Japanese therapist, respond empathetically: {user_input}", "").strip()

def random_breathing_sound():
    files = os.listdir(BREATHING_FOLDER)
    breathing_files = [f for f in files if f.endswith(".mp3") and "Sleep-meditation" not in f]
    return os.path.join(BREATHING_FOLDER, random.choice(breathing_files))

def predict_face_emotion(image):
    image = image.resize((48,48)).convert("L")
    img_array = np.array(image)/255.0
    img_array = np.expand_dims(img_array, axis=0)
    img_array = np.expand_dims(img_array, axis=-1)
    pred = face_emotion_model.predict(img_array)
    return face_classes[np.argmax(pred)]

# -----------------------------
# Streamlit UI
# -----------------------------
st.header("Describe your feelings or problems:")
user_text = st.text_area("Write here...", height=150)

st.header("Upload your selfie for facial emotion detection:")
user_image = st.file_uploader("Choose an image...", type=["jpg","jpeg","png"])

st.header("Upload your voice for speech emotion detection:")
user_audio = st.file_uploader("Choose audio...", type=["wav","mp3"])

if st.button("Get Advice"):
    if user_text.strip() != "":
        mh_class = predict_mental_health(user_text)
        st.subheader(f"Detected Mental State: {mh_class}")

        therapist_reply = generate_gpt2_reply(user_text)
        st.markdown(f"**AI Therapist:** {therapist_reply}")

    if user_image is not None:
        image = Image.open(user_image)
        face_emotion = predict_face_emotion(image)
        st.subheader(f"Detected Facial Emotion: {face_emotion}")

    if user_audio is not None:
        speech_emotion = speech_classifier(user_audio.name)[0]
        st.subheader(f"Detected Speech Emotion: {speech_emotion['label']} ({speech_emotion['score']:.2f})")

    # Breathing Exercise
    st.audio(random_breathing_sound(), format="audio/mp3", start_time=0)
    st.markdown("💨 *Try this breathing exercise to relax*")

    # Sleep Meditation
    st.audio(SLEEP_AUDIO, format="audio/mp3", start_time=0)
    st.markdown("🛌 *Optional sleep meditation music*")