import streamlit as st
from transformers import GPT2LMHeadModel, GPT2Tokenizer, AutoModelForCausalLM, AutoTokenizer
from gtts import gTTS

# Load Pretrained Model (GPT-2 for Reminders & Health Tips)
model = GPT2LMHeadModel.from_pretrained("gpt2")
tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
tokenizer.pad_token = tokenizer.eos_token

# Load DialoGPT (for chatbot functionality)
chatbot_tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-medium")
chatbot_model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-medium")

# Function to Generate Medication Reminder
def generate_reminder(prompt):
    inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True)
    outputs = model.generate(
        **inputs, max_length=50, num_return_sequences=1, temperature=0.7, do_sample=True
    )
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

# Function to Generate Health Tip
def generate_health_tip(prompt):
    return generate_reminder(prompt)

# Function to Get Chatbot Response
def chatbot_response(user_input):
    inputs = chatbot_tokenizer.encode(user_input + chatbot_tokenizer.eos_token, return_tensors="pt")
    outputs = chatbot_model.generate(inputs, max_length=1000, pad_token_id=chatbot_tokenizer.eos_token_id)
    response = chatbot_tokenizer.decode(outputs[:, inputs.shape[-1]:][0], skip_special_tokens=True)
    return response

# Function for Text-to-Speech
def text_to_speech(text):
    tts = gTTS(text=text, lang='en')
    tts.save("output.mp3")
    return "output.mp3"

# Streamlit UI
st.title("HealthEase - Medication Reminder and Health Assistant")

# Medication Reminder
st.header("Generate Medication Reminder")
medication_input = st.text_input("Enter medication details (e.g., 'Reminder for insulin at 7 PM:')")
if medication_input:
    reminder = generate_reminder(medication_input)
    st.write("Generated Reminder:", reminder)
    audio_file = text_to_speech(reminder)
    st.audio(audio_file)

# Health Tip
st.header("Generate Health Tip")
health_tip_input = st.text_input("Enter health tip details (e.g., 'Daily health tip for a diabetic patient:')")
if health_tip_input:
    health_tip = generate_health_tip(health_tip_input)
    st.write("Generated Health Tip:", health_tip)
    audio_file = text_to_speech(health_tip)
    st.audio(audio_file)

# Chatbot
st.header("Conversational Chatbot")
user_query = st.text_input("Ask a question (e.g., 'What are the side effects of ibuprofen?')")
if user_query:
    bot_reply = chatbot_response(user_query)
    st.write("Chatbot Reply:", bot_reply)