import gradio as gr
from bioemu.sample import main as sample
import os
import shutil

def run_bioemu(seq, num_samples, outdir):
    if not outdir.strip():
        outdir = "outputs"   # 默认目录
    os.makedirs(outdir, exist_ok=True)

    # 运行 BioEmu 采样
    sample(
        sequence=seq,
        num_samples=int(num_samples),
        output_dir=outdir
    )

    # 打包结果为 zip
    zip_path = f"{outdir}.zip"
    if os.path.exists(zip_path):
        os.remove(zip_path)  # 避免覆盖错误
    shutil.make_archive(outdir, 'zip', outdir)

    return f"运行完成！结果保存在 {outdir} 文件夹下。", zip_path

with gr.Blocks() as demo:
    gr.Markdown("## 🧬 BioEmu Demo: 生成蛋白质结构集合")

    seq = gr.Textbox(label="输入蛋白质序列", value="GYDPETGTWG")
    num_samples = gr.Slider(1, 20, value=5, step=1, label="采样数量")
    outdir = gr.Textbox(label="输出目录", value="outputs", placeholder="比如：results/chignolin")

    run_btn = gr.Button("运行 BioEmu")
    output_info = gr.Textbox(label="运行信息")
    output_file = gr.File(label="下载结果（zip 格式）")

    run_btn.click(fn=run_bioemu, inputs=[seq, num_samples, outdir], outputs=[output_info, output_file])

demo.launch()