import gradio as gr
from gradio import Interface, inputs, outputs

import spacy
from transformers import AutoModelForCausalLM, AutoTokenizer

nlp = spacy.load('es_core_news_sm')
tokenizer = AutoTokenizer.from_pretrained("OpenAssistant/oasst-sft-4-pythia-12b-epoch-3.5")
model = AutoModelForCausalLM.from_pretrained("OpenAssistant/oasst-sft-4-pythia-12b-epoch-3.5")

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

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

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

def generate_hint(sentence: str, user_word: str, user_pos: str):
    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: str, user_word: str, user_pos: str):
    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=[
                      inputs.Textbox(lines=2, placeholder='Introduce una frase aquí...'),
                      inputs.Textbox(lines=1, placeholder='Introduce una palabra aquí...'),
                      inputs.Dropdown(choices=['Selecciona una función gramatical...', 'sustantivo', 'verbo', 'adjetivo', 'artículo'])
                  ], 
                  outputs=[outputs.Textbox(), outputs.Textbox()])
iface.launch()