import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModel

tokenizer = AutoTokenizer.from_pretrained('distilbert/distilbert-base-uncased')
model = AutoModel.from_pretrained('distilbert/distilbert-base-uncased')

options = [
    "Visual", "Auditory", "Kinesthetic",
    "Recorded lectures", "Instant messaging (e.g., chat)", "Video calls",
    "Clarify concepts", "Improve grades", "Prepare for exams", "Career guidance",
    "1-2 hours", "3-5 hours", "6+ hours",
    "0 years", " 2 years", "5 years",
    "AIML", "Web dev", "App dev",
    "Slow and steady", "Moderate", "Fast-paced",
    "Explaining with examples", "Interactive discussions", "Problem-solving exercises",
    "Yes", "No",
    "Trouble with specific concepts", "Time management", "Test anxiety",
    "Experienced teacher", "Industry experience", "Should answer question"
]


def calculate_similarity(student_inputs, mentor_inputs):
    """
    Check for values in options that are not in student_inputs.
    Then append a 'not' to such options and place them at same index as in student_inputs.
    Also we need to take care of value shifts after adding 'not' to some options, so use a new list
    """
    # student = []
    # mentor = []

    # for i in range(len(options)):
    #     if options[i] in student_inputs:
    #         student.append(options[i])
    #     else:
    #         student.append(f"NOT {options[i]}")

    # for i in range(len(options)):
    #     if options[i] in mentor_inputs:
    #         mentor.append(options[i])
    #     else:
    #         mentor.append(f"NOT {options[i]}")

    # print(f"Student inputs: {student}")
    # print(f"Mentor inputs: {mentor}")

    student_tokens = tokenizer(student_inputs, return_tensors="pt", padding=True)
    mentor_tokens = tokenizer(mentor_inputs, return_tensors="pt", padding=True)

    with torch.no_grad():
        student_output = model(**student_tokens).last_hidden_state.mean(dim=1)
        mentor_output = model(**mentor_tokens).last_hidden_state.mean(dim=1)

    similarity = torch.nn.functional.cosine_similarity(student_output, mentor_output)
    print(similarity)
    similarity = similarity.mean().item()
    print(similarity)
    normalized_similarity = (similarity - 0.9)/0.1 * 100
    print(normalized_similarity)
    return round(normalized_similarity, 2)


with gr.Blocks() as demo:
    with gr.Column():
        with gr.Row():
            gr.Markdown("Student Options")
            with gr.Column():
                student_inputs = gr.CheckboxGroup(choices=options, label="Student Options")
                mentor_inputs = gr.CheckboxGroup(choices=options, label="Mentor Options")

    with gr.Row():
        btn = gr.Button("Calculate Similarity")

    with gr.Row():
        output = gr.Number()

    btn.click(
        fn=calculate_similarity,
        inputs=[student_inputs, mentor_inputs],
        outputs=output
    )

    demo.launch(debug=True)