import gradio as gr
from gradio import Interface
from gradio.components import Textbox, Dropdown

import spacy
from transformers import GPT2LMHeadModel, GPT2Tokenizer

nlp = spacy.load('es_core_news_sm')
tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
model = GPT2LMHeadModel.from_pretrained("gpt2")

pos_map = {
    'sustantivo': 'NOUN',
    'verbo': 'VERB',
    'adjetivo': 'ADJ',
    'artículo': 'DET'
}

def identify_pos(sentence):
    doc = nlp(sentence)
    pos_tags = [(token.text, token.pos_) for token in doc]
    return pos_tags

def game_logic(sentence, user_word, user_pos):
    correct_answers = identify_pos(sentence)
    user_pos = pos_map[user_pos.lower()]
    for word, pos in correct_answers:
        if word == user_word:
            if pos.lower() == user_pos.lower():
                return True, f'¡Correcto! "{user_word}" es un {user_pos}.'
            else:
                return False, f'Incorrecto. "{user_word}" no es un {user_pos}, es un {pos}.'
    return False, f'La palabra "{user_word}" no se encuentra en la frase.'

def generate_hint(sentence, user_word, user_pos):
    inputs = tokenizer.encode(f'Give a hint about the {user_pos} in the sentence "{sentence}"', return_tensors='pt')
    outputs = model.generate(inputs, max_length=150, num_return_sequences=1, no_repeat_ngram_size=2)
    hint = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return hint

def main(sentence, user_word, user_pos):
    if sentence and user_word and user_pos and user_pos != 'Selecciona una función gramatical...':
        correct, message = game_logic(sentence, user_word, user_pos)
        hint = generate_hint(sentence, user_word, user_pos)
        return message, hint
    else:
        return 'Por favor, introduce una frase, una palabra y selecciona una función gramatical válida (sustantivo, verbo, adjetivo, artículo).', ''

iface = Interface(fn=main, 
                  inputs=[
                      Textbox(lines=2, placeholder='Introduce una frase aquí...'),
                      Textbox(lines=1, placeholder='Introduce una palabra aquí...'),
                      Dropdown(choices=['Selecciona una función gramatical...', 'sustantivo', 'verbo', 'adjetivo', 'artículo'])
                  ], 
                  outputs=[Textbox(), Textbox()])
iface.launch()