import gradio as gr
from transformers import pipeline

# Initialize the model
chatbot = pipeline("text-generation", model="microsoft/DialoGPT-medium")

# Dictionary of scientist information
SCIENTISTS = {
    "Albert Einstein": {
        "intro": "theoretical physicist known for the theory of relativity",
        "style": "I speak with enthusiasm about physics and use thought experiments to explain complex ideas.",
    },
    "Marie Curie": {
        "intro": "physicist and chemist who conducted pioneering research on radioactivity",
        "style": "I am precise in my explanations and passionate about scientific research.",
    },
    "Nikola Tesla": {
        "intro": "inventor and electrical engineer",
        "style": "I share my visions about electricity and the future of technology with great enthusiasm.",
    }
}

def chat_with_scientist(scientist, question):
    # Create a prompt based on the scientist
    info = SCIENTISTS[scientist]
    prompt = f"""You are {scientist}, {info['intro']}. {info['style']}
    
    Question: {question}
    Answer as {scientist}:"""
    
    # Generate response
    response = chatbot(
        prompt,
        max_length=150,
        num_return_sequences=1,
        temperature=0.7,
        do_sample=True
    )
    
    return response[0]['generated_text']

# Create the Gradio interface
demo = gr.Interface(
    fn=chat_with_scientist,
    inputs=[
        gr.Dropdown(
            choices=list(SCIENTISTS.keys()),
            label="Choose a Scientist",
            value="Albert Einstein"
        ),
        gr.Textbox(
            placeholder="Ask your question here...",
            label="Your Question"
        )
    ],
    outputs=gr.Textbox(label="Scientist's Response"),
    title="Chat with Historical Scientists",
    description="""Welcome to the Scientific Time Machine! 
    Choose a scientist and ask them questions about their work, life, or scientific concepts.
    This is an educational tool to help students engage with historical scientists.""",
    examples=[
        ["Albert Einstein", "Can you explain E=mc² in simple terms?"],
        ["Marie Curie", "What inspired you to study radioactivity?"],
        ["Nikola Tesla", "Tell me about your work with electricity."]
    ],
    theme=gr.themes.Soft()
)

# Launch the app
if __name__ == "__main__":
    demo.launch()