import gradio as gr
from PIL import Image

# Single image analysis function (your existing logic)
def analyze_image(image, min_size, circularity, do_necrosis=True):
    import Organoid_Analyzer_AI2_HF as analyzer
    processed_img, picname, excelname = analyzer.main([image, min_size, circularity,do_necrosis])
    return Image.fromarray(processed_img), picname, excelname

# Z-stack analysis function (adapt with your own logic)
def analyze_zstack(images, min_size, circularity, distance_cutoff=30, overlap_cutoff=0.4, do_necrosis=True):
    import Organoid_Analyzer_AI_zstack2_HF as analyzer
    images = [Image.open(f.name) for f in images]
    processed_img, picname, excelname = analyzer.main(
        [images, min_size, circularity, distance_cutoff, overlap_cutoff, do_necrosis]
    )
    return Image.fromarray(processed_img), picname, excelname

with gr.Blocks() as demo:
    gr.Markdown("# AI Organoid Analyzer\nUpload an image (or Z-Stack) to run organoid analysis.")

    z_stack_checkbox = gr.Checkbox(label="Enable Z-Stack", value=False)
    do_necrosis_checkbox = gr.Checkbox(label="Enable necrosis detection", value=False)
    image_input_single = gr.Image(type="pil", label="Upload Image", visible=True)
    image_input_multi = gr.File(file_count="multiple", type="filepath", label="Upload Z-Stack Images", visible=False)
    min_size_input = gr.Number(label="Minimum Organoid Size (pixels)", value=1000)
    circularity_input = gr.Number(label="Minimum Circularity", value=0.25)
    zstack_distance_cutoff_input = gr.Number(label="Z-Stack Centroid Distance Cutoff (pixels)", value=30, visible=False)
    zstack_overlap_cutoff_input = gr.Number(label="Z-Stack Colony Overlap Cutoff", value=0.4, visible=False)
    output_image = gr.Image(type="pil", label="Analyzed Image")
    output_file_img = gr.File(label="Download Image")
    output_file_excel = gr.File(label="Download results (Excel)")
    process_btn = gr.Button("Process")

    def toggle_inputs(z_stack_enabled):
        return (
            gr.update(visible=not z_stack_enabled),  # single input
            gr.update(visible=z_stack_enabled),      # multi input
            gr.update(visible=z_stack_enabled),      # zstack_distance_cutoff_input
            gr.update(visible=z_stack_enabled)       # zstack_overlap_cutoff_input
        )

    z_stack_checkbox.change(
    toggle_inputs,
    inputs=z_stack_checkbox,
    outputs=[
        image_input_single, 
        image_input_multi, 
        zstack_distance_cutoff_input, 
        zstack_overlap_cutoff_input])

    def conditional_analyze(z_stack, single_image, multi_images, min_size, circularity, do_necrosis, distance_cutoff, overlap_cutoff):
        if z_stack:
            return analyze_zstack(multi_images, min_size, circularity, distance_cutoff, overlap_cutoff, do_necrosis)
        else:
            return analyze_image(single_image, min_size, circularity, do_necrosis)

    process_btn.click(
    conditional_analyze,
    inputs=[
        z_stack_checkbox,
        image_input_single,
        image_input_multi,
        min_size_input,
        circularity_input,
        do_necrosis_checkbox,
        zstack_distance_cutoff_input,
        zstack_overlap_cutoff_input
    ],
    outputs=[output_image, output_file_img, output_file_excel])

demo.launch()