import gradio as gr
from datasets import load_dataset
import random

dataset = load_dataset("hugfaceguy0001/TangshiDalle3Images")['train']
demo = gr.Blocks(theme=gr.themes.Default(text_size="lg"))
N = 149
j = 3 * N + 1
poem_text = dataset[j]['prompt']
poem_info = dataset[j+1]['prompt']
poem_title = poem_info[poem_info.index("《")+1:poem_info.index("》")]
poem_author = poem_info[3:poem_info.index("的")]
poem = f"        {poem_title}    {poem_author}\n{poem_text}"
k = j - 1
poem_image = dataset[k]['image']
def random_poem():
    global N
    global j
    global k
    N = random.randint(0,230)
    j = 3 * N + 1
    poem_text = dataset[j]['prompt']
    poem_info = dataset[j+1]['prompt']
    poem_title = poem_info[poem_info.index("《")+1:poem_info.index("》")]
    poem_author = poem_info[3:poem_info.index("的")]
    poem = f"        {poem_title}    {poem_author}\n{poem_text}"
    k = j + random.randint(-1,1)
    poem_image = dataset[k]['image']
    return poem_image,poem

def change_pic():
    global k
    rand_shift = random.choice([i for i in [-1,0,1] if i != k-j])
    k = j + rand_shift
    return dataset[k]['image']

with demo:
    demo_title = gr.HTML('<h1 align="center">唐诗欣赏</h1>')
    with gr.Row():
        img = gr.Image(poem_image,scale=2)
        poem_textbox = gr.Text(poem,scale=1,label="唐诗",show_copy_button=True)
    with gr.Row():
        btn1 = gr.Button('切换唐诗')
        btn2 = gr.Button('切换插图')
    btn1.click(random_poem,inputs=[],outputs=[img,poem_textbox])
    btn2.click(change_pic,inputs=[],outputs=img)
demo.launch()