import os
os.system("pip install -q gradio torch transformers")

import gradio as gr
import torch
import random
from transformers import GPT2LMHeadModel, GPT2Tokenizer

tokenizer = GPT2Tokenizer.from_pretrained('gpt2-medium')
model = GPT2LMHeadModel.from_pretrained('RandomNameAnd6/DharGPT-Small')

def generate_text(prompt):
  input_ids = tokenizer.encode(prompt, return_tensors="pt")
  output = model.generate(input_ids, max_length=48, temperature=0.85, do_sample=True)
  text = tokenizer.decode(output[0], skip_special_tokens=True)
  return text

# Read real titles from file
with open('dhar_mann_titles.txt', 'r') as file:
    dhar_mann_titles = file.readlines()

# Function to generate an AI title (dummy implementation)
def generate_ai_title():
    inputs = tokenizer(["<|startoftext|>"]*1, return_tensors = "pt")
    outputs = model.generate(**inputs, max_new_tokens=50, use_cache=True, temperature=0.85, do_sample=True)
    return (tokenizer.batch_decode(outputs)[0])[15:-13]

# Function to check user's answer and update score
def check_answer(user_choice, real_index, score):
    if (user_choice == "Option 1" and real_index == 0) or (user_choice == "Option 2" and real_index == 1):
        score += 1
        return f"Correct! Your current score is: {score}", score, gr.update(visible=True), gr.update(visible=False)
    else:
        score = 0
        return f"Incorrect. Your score has been reset to: {score}", score, gr.update(visible=False), gr.update(visible=True)

# Function to update options
def update_options():
    real_index = random.choice([0, 1])
    real_title = random.choice(dhar_mann_titles).strip()
    ai_title = generate_ai_title()

    if real_index == 0:
        return real_title, ai_title, real_index
    else:
        return ai_title, real_title, real_index

def create_interface():
    with gr.Blocks() as demo:
        score = gr.State(0)
        real_index_state = gr.State(0)

        score_display = gr.Markdown("## Real or AI - Dhar Mann\n**Current Score: 0**")

        with gr.Row():
            with gr.Column():
                gr.Markdown("### Option 1")
                option1_box = gr.Markdown("")
            with gr.Column():
                gr.Markdown("### Option 2")
                option2_box = gr.Markdown("")

        with gr.Row():
            choice = gr.Radio(["Option 1", "Option 2"], label="Which one do you think is real?")

        submit_button = gr.Button("Submit")
        result_text = gr.Markdown("")
        continue_button = gr.Button("Continue", visible=False)
        restart_button = gr.Button("Restart", visible=False)

        def on_submit(user_choice, option1, option2, real_index, score):
            result, new_score, continue_visibility, restart_visibility = check_answer(user_choice, real_index, score)
            return result, new_score, continue_visibility, restart_visibility

        def on_continue(score):
            option1, option2, real_index = update_options()
            new_score_display = f"## Real or AI - Dhar Mann\n**Current Score: {score}**"
            return option1, option2, real_index, new_score_display, gr.update(value=None), "", gr.update(visible=False), gr.update(visible=False)

        def on_restart():
            return on_continue(0)

        # Initialize options
        option1, option2, real_index = update_options()

        submit_button.click(on_submit, inputs=[choice, option1_box, option2_box, real_index_state, score], outputs=[result_text, score, continue_button, restart_button])
        continue_button.click(on_continue, inputs=score, outputs=[option1_box, option2_box, real_index_state, score_display, choice, result_text, continue_button, restart_button])
        restart_button.click(on_restart, outputs=[option1_box, option2_box, real_index_state, score_display, choice, result_text, continue_button, restart_button])

        # Set initial content for option boxes
        option1_box.value = option1
        option2_box.value = option2
        real_index_state.value = real_index

    return demo

demo = create_interface()
demo.launch()