import gradio as gr, os
from bootstrap_upstream import fetch_upstream
from pipeline.rules import zone_lus_score, ZoneFeatures
from pipeline import rtlucassen_wrapper as rtl
from pipeline import medivlad_wrapper as mvl

def setup_check():
    results = fetch_upstream()
    folders = {
        "B-line_detection": os.path.isdir("external/B-line_detection"),
        "LUS-Segmentation-RT": os.path.isdir("external/LUS-Segmentation-RT"),
        "MeDiVLAD": os.path.isdir("external/MeDiVLAD"),
        "weights_rtlucassen": os.path.isdir("models/rtlucassen"),
        "weights_medivlad": os.path.isdir("models/medivlad"),
    }
    return {"fetched": results, "folders": folders}

def score_demo_zone(_clip):
    try:
        b_counts, _ = rtl.run_bline_detector(_clip)
        median_b = sorted(b_counts)[len(b_counts)//2] if b_counts else 0
        z = ZoneFeatures(b_count=median_b, b_area_frac=0.0, has_consolidation=False)
        return {"zone_score": zone_lus_score(z), "median_b": median_b}
    except Exception as e:
        return {"error": str(e)}

def run_medivlad(_clip):
    try:
        return mvl.run_video_severity(_clip)
    except Exception as e:
        return {"error": str(e)}

with gr.Blocks() as demo:
    gr.Markdown("## NeoLUS — Original Repos (private research)\n**Note:** Upload authorized weights to `/models/*`. Not for clinical use.")
    with gr.Accordion("Setup / Status", open=True):
        gr.Button("Run setup check").click(setup_check, None, gr.JSON())

    with gr.Tab("Zone Scoring"):
        clip = gr.Video(label="Zone clip")
        go = gr.Button("Analyze zone")
        out = gr.JSON()
        go.click(score_demo_zone, clip, out)

    with gr.Tab("Video Severity (MeDiVLAD)"):
        clip2 = gr.Video()
        go2 = gr.Button("Run")
        out2 = gr.JSON()
        go2.click(run_medivlad, clip2, out2)

if __name__ == "__main__":
    demo.launch()