# Import necessary libraries
import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM

# Load the model and tokenizer
model_name = "meta-llama/Llama-3.2-1B-Instruct"  # or "meta-llama/Llama-3.2-3B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

def answer_question(question, max_tokens=100):
    """Generate an answer to a given question about photography."""
    if not question.strip():
        return "Please enter a question."

    inputs = tokenizer(question, return_tensors="pt")
    outputs = model.generate(**inputs, max_length=max_tokens, pad_token_id=tokenizer.eos_token_id, temperature=0.7, top_p=0.9)
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

def generate_practice(subject, max_length=400, temperature=0.7, top_p=0.9):
    """Generate a concise photography exercise for a given subject."""
    if not subject.strip():
        return "Please select a photography subject."

    prompt = (f"Create a concise photography exercise for {subject}. "
              f"The exercise should include: "
              f"1. Objective: One sentence about what students should learn. "
              f"2. Materials: List essential equipment. "
              f"3. Steps: Three to four concise instructions. "
              f"4. Expected outcomes: One sentence on what students should achieve.")

    inputs = tokenizer(prompt, return_tensors="pt")
    outputs = model.generate(**inputs, pad_token_id=tokenizer.eos_token_id, max_length=max_length, temperature=temperature, top_p=top_p)
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

# Define the Gradio interface using Blocks
with gr.Blocks() as demo:
    # Title and Description
    gr.Markdown("# 📸 Photography Learning Assistant")
    gr.Markdown("Welcome to the **Photography Learning Assistant**! Use the Q&A section to ask questions or generate exercises.")

    # Q&A Section
    gr.Markdown("### 📝 Q&A")
    question_input = gr.Textbox(label="Photography Question", placeholder="Enter a question (e.g., What is the rule of thirds?)", lines=2)
    max_tokens_slider = gr.Slider(minimum=50, maximum=500, step=50, value=100, label="Max Tokens")
    answer_button = gr.Button("Get Answer")
    answer_output = gr.Textbox(label="Answer", lines=10)
    answer_button.click(fn=answer_question, inputs=[question_input, max_tokens_slider], outputs=answer_output)

    gr.Markdown("#### 💡 Sample Questions")
    gr.Markdown("""
    - What are different types of photography?
    - Explain the exposure triangle like you would explain to a 5-year-old.
    """)


    # Generate Practice Exercise Section
    gr.Markdown("### 🎯 Generate Practice Exercise")
    subject_dropdown = gr.Radio(choices=["Composition", "Lighting", "Camera Settings", "Exposure", "Post-Processing"], label="Photography Subject")
    max_length_slider = gr.Slider(minimum=100, maximum=800, step=50, value=400, label="Max Length")
    temperature_slider = gr.Slider(minimum=0.1, maximum=1.0, step=0.1, value=0.7, label="Temperature")
    top_p_slider = gr.Slider(minimum=0.1, maximum=1.0, step=0.1, value=0.9, label="Top P")
    generate_button = gr.Button("Generate Exercise")
    practice_output = gr.Textbox(label="Generated Practice Exercise", lines=15)
    generate_button.click(fn=generate_practice, inputs=[subject_dropdown, max_length_slider, temperature_slider, top_p_slider], outputs=practice_output)

# Launch the Gradio app
demo.launch(share=True)