import streamlit as st
from transformers import pipeline, AutoModelForCTC, AutoProcessor, Wav2Vec2ForCTC
import pyttsx3
import random
import torchaudio
import torch  # Add this import statement to fix the NameError
from gtts import gTTS
import os

# Set up the app
st.set_page_config(page_title="ChiliDron", layout="wide", page_icon="🎨")

# Load Hugging Face models using cache
@st.cache_resource
def load_storytelling_model():
    return pipeline("text-generation", model="bookbot/gpt2-indo-medium-kids-stories")

@st.cache_resource
def load_phoneme_model():
    model = Wav2Vec2ForCTC.from_pretrained("mirfan899/kids_phoneme_sm_model")
    processor = AutoProcessor.from_pretrained("mirfan899/kids_phoneme_sm_model")
    return model, processor

# Load models
story_generator = load_storytelling_model()
phoneme_model, phoneme_processor = load_phoneme_model()

# Initialize TTS engine
def initialize_tts():
    try:
        engine = pyttsx3.init()
        return engine
    except Exception as e:
        st.error(f"Error initializing text-to-speech engine: {e}")
        return None

# App title and description
st.title("ChiliDron")
st.subheader("A fun, interactive, and educational app for kids!")

# Sidebar navigation
st.sidebar.title("Navigate")
options = st.sidebar.radio("Go to", ["Interactive Storyteller", "Joke and Fun Facts Generator", "Phoneme Practice", "Guided Meditation"])

# Interactive Storyteller
if options == "Interactive Storyteller":
    st.header("📝 Interactive Storyteller")
    st.write("Create your own story by selecting characters and plot points!")

    character = st.selectbox("Choose a character:", ["A Brave Knight", "A Clever Princess", "A Funny Alien", "A Curious Scientist"])
    setting = st.selectbox("Choose a setting:", ["In a magical forest", "On a distant planet", "In an ancient castle", "At a science lab"])
    plot = st.selectbox("Choose a plot point:", ["Finds a mysterious map", "Discovers a secret passage", "Meets a talking animal", "Invents a cool gadget"])

    if st.button("Create Story"):
        with st.spinner("Generating your story..."):
            prompt = f"Once upon a time, {character} was {setting} when they {plot}. And then,"
            story = story_generator(prompt, max_length=500, num_return_sequences=1)
            st.success(story[0]['generated_text'])

# Joke and Fun Facts Generator
elif options == "Joke and Fun Facts Generator":
    st.header("🤣 Joke and Fun Facts Generator")
    st.write("Get ready to laugh and learn with random jokes and fun facts!")

    jokes = [
        "Why did the scarecrow win an award? Because he was outstanding in his field!",
        "Why don’t scientists trust atoms? Because they make up everything!",
        "What do you get if you cross a vampire with a snowman? Frostbite!"
    ]

    fun_facts = [
        "Did you know that honey never spoils? Archaeologists have found pots of honey in ancient Egyptian tombs that are over 3,000 years old and still perfectly edible.",
        "A day on Venus is longer than a year on Venus! It takes about 243 Earth days for Venus to rotate once, but only 225 Earth days to orbit the Sun.",
        "Octopuses have three hearts and blue blood!"
    ]

    if st.button("Tell me a joke!"):
        st.write(random.choice(jokes))

    if st.button("Tell me a fun fact!"):
        st.write(random.choice(fun_facts))

# Phoneme Practice
elif options == "Phoneme Practice":
    st.header("🔤 Phoneme Practice")
    st.write("Learn how to pronounce words correctly!")

    uploaded_audio = st.file_uploader("Upload an audio file to analyze phonemes", type=["wav", "mp3"])

    if uploaded_audio:
        with st.spinner("Analyzing phonemes..."):
            waveform, sample_rate = torchaudio.load(uploaded_audio)

            # Ensure mono audio by averaging channels if stereo
            if waveform.shape[0] > 1:
                waveform = waveform.mean(dim=0, keepdim=True)
            
            # Resample if needed
            if sample_rate != 16000:
                resampler = torchaudio.transforms.Resample(orig_freq=sample_rate, new_freq=16000)
                waveform = resampler(waveform)

            input_values = phoneme_processor(waveform.squeeze(), return_tensors="pt", sampling_rate=16000).input_values
            logits = phoneme_model(input_values).logits
            predicted_ids = torch.argmax(logits, dim=-1)
            transcription = phoneme_processor.batch_decode(predicted_ids)
            st.write(f"Transcription: {transcription[0]}")

# Guided Meditation for Kids
if options == "Guided Meditation":
    st.header("🧘‍♂️ Guided Meditation for Kids")
    st.write("Relax and listen to a calming meditation.")

    meditation_text = "Close your eyes and take a deep breath. Imagine you are in a peaceful forest, with birds singing and leaves rustling in the wind."

    if st.button("Play Meditation"):
        with st.spinner("Generating meditation audio..."):
            try:
                tts = gTTS(text=meditation_text, lang='en')
                tts.save('meditation.mp3')
                st.audio('meditation.mp3', format='audio/mp3')
            except Exception as e:
                st.error(f"Error generating or playing the audio: {e}")

# Footer
st.markdown("<footer style='text-align: center; padding-top: 20px;'>Made with ❤️ for Kids by ChiliDron Team</footer>", unsafe_allow_html=True)