import gradio as gr
import matplotlib.pyplot as plt
from PIL import Image
import io

# 数式を描画して画像として保存する関数
def render_math_expression(expression):
    # Matplotlib を使用して数式を画像に描画
    plt.figure(figsize=(10, 2))
    plt.text(0.5, 0.5, f"${expression}$", fontsize=20, ha='center', va='center')
    plt.axis('off')  # 軸を非表示
    buf = io.BytesIO()
    plt.savefig(buf, format='png', bbox_inches='tight', dpi=300)
    buf.seek(0)
    plt.close()
    # Pillow Image に変換して返す
    return Image.open(buf)

# 初期値として与える数式
default_math = r"\sum_{t=1}^T \E_{y_t \sim {\tilde P(y_t| y_0)}} \left\| \frac{y_t - \sqrt{\bar{\alpha}_t}y_0}{\sqrt{1-\bar{\alpha}_t}} - \epsilon_\theta(y_t, t)\right\|^2."

# Gradio インターフェースの作成
with gr.Blocks() as demo:
    gr.Markdown("## 数式を画像として描画")
    with gr.Row():
        math_input = gr.Textbox(
            value=default_math,
            label="LaTeX 数式入力",
            lines=3,
            max_lines=5,
            placeholder="LaTeX の形式で数式を入力してください",
        )
    image_output = gr.Image(label="レンダリング結果", type="pil")

    # ボタンで描画実行
    render_button = gr.Button("レンダリング")
    render_button.click(render_math_expression, inputs=[math_input], outputs=[image_output])

# アプリを起動
demo.launch()