import time
import gradio as gr
from transformers import pipeline

# Load the generation model
generator = pipeline("text2text-generation", model="t5-small")

# Define the questions
questions = [
    "Which grade is the child studying?",
    "How old is the child?",
    "What is the gender?",
    "Can you provide the name and location of the child's school?",
    "What are the names of the child's guardians or parents?",
    "What is the chief complaint regarding the child's oral health? If there is none, just say the word 'none', else elaborate only on medication history",
    "Can you provide any relevant medical history for the child? If there is none, just say the word 'none', else elaborate",
    "Does the child take any medications regularly? If there is none, just say the word 'none'. If yes, please specify.",
    "When was the child's previous dental visit? If no visits before, just say the word 'first' or mention the visit number and nothing else",
    "Does the child have any habits such as thumb sucking, tongue thrusting, nail biting, or lip biting? If yes, just list them and don't provide any further details",
    "Does the patient brush their teeth? Just use the words 'once daily', 'twice daily', or 'thrice daily' to answer, nothing else",
    "Does the child experience bleeding gums? Just say 'yes' or 'no' for this and nothing else",
    "Has the child experienced early childhood caries? Just say 'yes' or 'no' and nothing else",
    "Please mention if tooth decay is present with tooth number(s), else just say the word 'none' and nothing else",
    "Have any teeth been fractured? If yes, please mention the tooth number(s), else just say 'none' and nothing else",
    "Is there any pre-shedding mobility of teeth? If yes, please specify, else just say 'none' and nothing else",
    "Does the child have malocclusion? If yes, please provide details, else just say the word 'none' and nothing else",
    "Does the child experience pain, swelling, or abscess? If yes, please provide details, else just say 'none' and nothing else",
    "Are there any other findings you would like to note?",
    "What treatment plan do you recommend? Choose only from Options: (Scaling, Filling, Pulp therapy/RCT, Extraction, Medication, Referral) and nothing else"
]

# Function to generate answers from context
def generate_answers(context):
    start_time = time.time()
    answers = []
    for question in questions:
        # Generate the answer using the model
        answer = generator(f"Context: {context}\nQuestion: {question}", max_length=50)[0]['generated_text']
        answers.append(answer)
    end_time = time.time()
    generation_time = end_time - start_time
    print(f"Time taken to generate answers: {generation_time:.2f} seconds")
    return answers

# Gradio interface
def main(context):
    answers = generate_answers(context)
    return answers

# Gradio app layout
with gr.Blocks() as demo:
    context = gr.Textbox(lines=5, placeholder="Enter the context here", label="Context")
    generate_button = gr.Button("Generate Answers")
    output_table = gr.Dataframe(headers=["Question", "Answer"], row_count=20, col_count=2, visible=False)
    save_button = gr.Button("Save Answers", visible=False)

    def generate_and_show_table(context):
        answers = main(context)
        data = [[question, answer] for question, answer in zip(questions, answers)]
        generate_button.update(interactive=False)
        output_table.update(data, visible=True)
        save_button.update(visible=True)
        return data

    generate_button.click(generate_and_show_table, inputs=context, outputs=[output_table])
    save_button.click(lambda: "Answers displayed in the table", inputs=None, outputs=gr.Textbox(label="Status"))

# Launch the app
demo.launch()