# anconda
'''step1: 安装paddlepaddle
conda create -n paddle python=3.6
conda activate paddle
pip install paddlepaddle -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install paddlehub -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install pillow
'''
import paddlehub as hub

class HubUtils(object):
    def __init__(self):
        self.hub = hub
        #初始化时需要加载模型
        self.module_love_words = self.hub.Module(name="ernie_gen_lover_words")          # 情话
        self.module_poetry = self.hub.Module(name="ernie_gen_poetry")                   # 诗歌
        self.module_acrostic_poetry = self.hub.Module(name="ernie_gen_acrostic_poetry") # 藏头诗
        self.module_couplets = self.hub.Module(name="ernie_gen_couplet")                # 对联
    def GetLoveWords(self,text,size=5):
        results = self.module_love_words.generate(texts=[text], use_gpu=False, beam_width=size)
        return results[0]

    def GetAcrostic(self,text,line=4,word=7,size=1):
        #默认返回一首
        self.module_acrostic = self.hub.module = hub.Module(name="ernie_gen_acrostic_poetry", line=line, word=word)
        results = self.module_acrostic.generate(texts=[text], use_gpu=False, beam_width=size)
        return results[0]
    def GetAcrostic_poetry(self,text,line=4, word=7):
        self.module_acrostic_poetry = self.hub.module = hub.Module(name="ernie_gen_acrostic_poetry", line=line, word=word)
        results = self.module_acrostic_poetry.generate(texts=[text], use_gpu=False, beam_width=1)
        return results[0]
    def GetPoetry(self,curx,size=1):
        #输入关键句，词即可生成诗歌，默认一首
        results = self.module_poetry.generate(texts=[curx], use_gpu=False, beam_width=size)
        return results[0]

    def GetCouplet(self,up_couplet,size=1):
        #输入上联给出下联,默认给一个
        results = self.module_couplets.generate(texts=[up_couplet], use_gpu=False, beam_width=size)
        return results[0]
    

if __name__=="__main__":
    Hub = HubUtils()
    while True:
        input_ = input("请选择功能: 1.情话 2.诗歌 3.藏头诗 4.对联 5.退出\n")
        if input_ == "1":
            input_ = input("请输入情话关键词:\n>>>")
            print(Hub.GetLoveWords(input_))
        elif input_ == "2":
            input_ = input("请输入诗歌关键词:\n>>>")
            print(Hub.GetPoetry(input_))
        elif input_ == "3":
            input_ = input("请输入藏头诗关键词:\n>>>")
            print(Hub.GetAcrostic_poetry(input_))
        elif input_ == "4":
            input_ = input("请输入上联:\n>>>")
            print(Hub.GetCouplet(input_))
        elif input_ == "5":
            break