import gradio as gr
from cijiang import CiJiangRhymer
from cijiang.utils import print_results
from colorama import Fore, Style, init
import spaces


rhymer = CiJiangRhymer(strict=False, tone=True)

@spaces.GPU  # Add this decorator
def get_results(input_text, target_rhyme, beam_width, num_candidates, top_results):
    out = rhymer.get_rhymes(text_with_placeholder=input_text, target_rhyme=target_rhyme, beam_width=beam_width, num_candidates=num_candidates)
    mask_count = input_text.count("[M]")
    context = input_text.split('[M]')[0]

    all_outputs = f"======= 韵脚: |{target_rhyme}|\n"
    for i, (seq, log_prob) in enumerate(out[:top_results]): 
        rhymes = seq[-mask_count:].split()
        colored_rhymes = [part if idx < mask_count else part for idx, part in enumerate(rhymes)]
        colored_rhymes = ''.join(colored_rhymes)  # Join the parts back together

        all_outputs += f"{i+1}. {context}{colored_rhymes} (score: {log_prob:.3f})\n"
    all_outputs += "=" * 40
    return all_outputs


demo = gr.Interface(
    fn=get_results,
    inputs=[
        gr.Textbox(
            label="Input Text",
            placeholder="Enter your text with [M] markers for rhyme positions",
            lines=3,
            value="[M][M][M][M]"
        ),
        gr.Textbox(
            label="Target Rhyme",
            placeholder="Enter target rhyme pattern (Chinese characters or pinyin)",
            value="乱比高低"
        ),
        gr.Slider(
            minimum=1,
            maximum=100,
            value=20,
            step=1,
            label="Beam Width"
        ),
        gr.Slider(
            minimum=100,
            maximum=10000,
            value=4000,
            step=100,
            label="Number of Candidates"
        ),
        gr.Slider(
            minimum=1,
            maximum=50,
            value=20,
            step=1,
            label="Top Results to Show"
        )
    ],
    outputs=gr.Textbox(
        label="Generated Rhymes",
        lines=15,
        max_lines=20
    ),
    title="CiJiang Chinese Poetry Rhymer",
    description="Generate Chinese poetry rhymes using the CiJiang library. Use [M] markers to indicate positions where rhymes should be generated. Supports both Chinese characters and pinyin as target rhymes.",
    examples=[
        ["[M][M][M][M]", "乱比高低", 20, 4000, 10],
        ["[M][M][M][M]", "luan4 bi3 gao1 di1", 20, 4000, 10],
        ["我总是天亮了[M][M]", "排队", 20, 4000, 10],
        ["我们两个人的verse价值一辆[M][M][M][M]", "韩国低级", 20, 4000, 10],
        ["我还是觉得音乐的[M][M][M][M][M]", "背靠背钻研", 100, 4000, 10],
    ]
)
demo.launch(share=True)