import tensorflow
import torch
import gradio as gr
from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer

# --- Load the NLP pipeline for text classification ---
classifier = pipeline("text-classification")

# --- Define the function to generate mini-apps based on user input ---
def generate_mini_apps(theme):
    # Use the NLP pipeline to classify the input theme
    classification = classifier(theme)

    # Generate a set of mini-apps based on the classification
    if classification[0]['label'] == 'Productivity':
        mini_apps = [
            'Idea-to-Codebase Generator',
            'Automated GitHub Repo Guardian Angel',
            'AI-Powered IDE'
        ]
    elif classification[0]['label'] == 'Creativity':
        mini_apps = [
            'Brainstorming Assistant',
            'Mood Board Generator',
            'Writing Assistant'
        ]
    elif classification[0]['label'] == 'Well-being':
        mini_apps = [
            'Meditation Guide',
            'Mood Tracker',
            'Sleep Tracker'
        ]
    else:
        mini_apps = ["No matching mini-apps found. Try a different theme."]

    # Return the generated mini-apps
    return mini_apps

# --- Load the model and tokenizer from the provided files ---
model = AutoModelForCausalLM.from_pretrained("./", trust_remote_code=True)  # Load from the current directory
tokenizer = AutoTokenizer.from_pretrained("./")

# --- Define a function to generate text using the model ---
def generate_text(input_text):
    inputs = tokenizer(input_text, return_tensors="pt")
    output = model.generate(**inputs, max_length=50, num_return_sequences=1)
    return tokenizer.decode(output[0], skip_special_tokens=True)

# --- Create the Gradio interface ---
demo = gr.Interface(
    fn=generate_mini_apps,
    inputs=gr.Textbox(label="Enter a theme for your life"),
    outputs=gr.Textbox(label="Generated Mini-Apps"),
    title="AI4ME: Personalized AI Tools",
    description="Enter your hobby/interest/job and we'll generate a set of AI-powered mini-apps tailored to your specific needs."
)

# --- Add a text generation tab ---
with gr.Blocks() as demo_text:
    gr.Markdown("## Text Generation")
    input_text = gr.Textbox(label="Enter your text")
    output_text = gr.Textbox(label="Generated Text")
    input_text.submit(generate_text, inputs=input_text, outputs=output_text)

# --- Launch the Gradio app ---
demo.launch(share=True)  # Share the app publicly
demo_text.launch(share=True)