import gradio as gr

import random

endings = ['ый', 'ий', 'ая', 'яя', 'ое', 'ее', 'ые', 'ие', 'ому', 'ему', 'ой', 'ей', 'ых', 'их', 'ым', 'им', 'ого', 'его', 'ую', 'юю', 'ого', 'его']

def gen(text):
    a = text
    count = random.randint(5, 15)
    _dataset = open("dataset.txt", "r").readlines()
    _a = a.split()
    __dataset = {el: {"content": i} for el, i in enumerate(_dataset)}
    dataset = ' '.join([i['content'].replace('\n', '') for i in __dataset.values() if any(ii in i['content'] for ii in _a)])
    
    if len(dataset.strip()) == 0:
        return random.choice(_dataset).strip()
  
    words = {}
    text = dataset.split()
    random.shuffle(text)
  
    for el, i in enumerate(text):
        words[el] = {
            "word": i,
            "vp": len(i) > 3 and (i[-3:] in endings or i[-2:] in endings)
        }
  
    result = ""
    previous_value = None
    _from = 0
    _to = count

  
    all_vp = False if str(words).count("'vp': True") == 0 else True
    all_not_vp = False if str(words).count("'vp': False") == 0 else True
    
    if all_vp or all_not_vp:
        for i in words.values():
            if _from == 3:
                break
            _from += 1
            result += f"{i['word']} "
    else:
        for i in words.values():
            if _from == _to:
                break
            _from += 1
            current_value = str(i['vp'])
            if current_value != previous_value:
                result += f"{i['word']} "
                previous_value = current_value
  
    return result.lower()

iface = gr.Interface(fn=gen, inputs="text", outputs="text")
iface.launch()