agentic_tabs.py

"""
Agentic Tabs for NeuroSAM3.

Defines the new Gradio tabs:
1. Agent Chat — conversational interface with LLM reasoning
2. Clinical Pipeline — one-click automated analysis
3. Research Pipeline — batch analysis with statistics
4. Model Comparison — side-by-side model evaluation

These tabs are prepended before the existing manual tools.
"""

import gradio as gr
from typing import Optional
from config import LLM_MODELS, DEFAULT_LLM_PROVIDER, BIOMEDCLIP_DEFAULT_LABELS
from logger_config import logger


def create_agent_chat_tab():
    """Create the Agent Chat tab — conversational neuroimaging interface."""
    with gr.Tab("Agent Chat", id="agent_chat"):
        gr.Markdown("""
        ### AI-Powered Neuroimaging Assistant
        Upload a medical image and describe what you need in natural language.
        The agent will classify, segment, measure, and report — automatically choosing the best tools.

        **Examples:**
        - "Segment the brain in this CT scan"
        - "Is there a tumor in this MRI? Measure it."
        - "Analyze this image and generate a clinical report"
        - "Classify this image and tell me the modality"
        """)

        with gr.Row():
            with gr.Column(scale=1):
                agent_image = gr.File(
                    label="Upload Medical Image (DICOM, PNG, JPG)",
                    file_types=[".dcm", ".png", ".jpg", ".jpeg"],
                    type="filepath",
                )
                agent_message = gr.Textbox(
                    label="Your Request",
                    placeholder="e.g., 'Segment the tumor and generate a report'",
                    lines=3,
                )
                with gr.Row():
                    agent_provider = gr.Dropdown(
                        choices=list(LLM_MODELS.keys()),
                        value=DEFAULT_LLM_PROVIDER,
                        label="LLM Provider",
                    )
                    agent_mode = gr.Dropdown(
                        choices=["auto", "clinical", "research"],
                        value="auto",
                        label="Mode",
                    )
                agent_submit = gr.Button("Send to Agent", variant="primary", size="lg")

            with gr.Column(scale=2):
                agent_output_image = gr.Image(label="Segmentation Output", type="filepath")
                agent_response = gr.Markdown(label="Agent Response")
                agent_steps = gr.JSON(label="Agent Steps (tool calls)", visible=False)

        def handle_agent_request(image_path, message, provider, mode):
            if not message:
                return None, "Please enter a request.", []
            if not image_path:
                return None, "Please upload a medical image.", []

            try:
                from PIL import Image as PILImage
                from agent_orchestrator import orchestrator

                orchestrator.set_provider(provider)
                pil_image = PILImage.open(image_path).convert("RGB")

                result = orchestrator.process_request(
                    message=message,
                    image=pil_image,
                    image_path=image_path,
                    mode=mode,
                )

                output_img = result["outputs"][0] if result.get("outputs") else None
                response_text = result.get("response", "No response generated.")
                steps = result.get("steps", [])

                return output_img, response_text, steps

            except Exception as e:
                logger.error(f"Agent request error: {e}", exc_info=True)
                return None, f"Error: {str(e)}", []

        agent_submit.click(
            fn=handle_agent_request,
            inputs=[agent_image, agent_message, agent_provider, agent_mode],
            outputs=[agent_output_image, agent_response, agent_steps],
        )


def create_clinical_pipeline_tab():
    """Create the Clinical Pipeline tab — one-click automated analysis."""
    with gr.Tab("Clinical Pipeline", id="clinical_pipeline"):
        gr.Markdown("""
        ### One-Click Clinical Analysis
        Upload an image and get a complete automated analysis:
        **Classify** -> **Segment** -> **Measure** -> **Report**

        No prompts needed — the AI determines the best approach automatically.
        """)

        with gr.Row():
            with gr.Column():
                clinical_image = gr.File(
                    label="Upload Medical Image",
                    file_types=[".dcm", ".png", ".jpg", ".jpeg"],
                    type="filepath",
                )
                clinical_context = gr.Textbox(
                    label="Clinical Context (optional)",
                    placeholder="e.g., '45yo male, headache for 2 weeks, rule out mass'",
                    lines=2,
                )
                with gr.Row():
                    clinical_modality = gr.Dropdown(
                        choices=["auto", "CT", "MRI"],
                        value="auto",
                        label="Modality",
                    )
                    clinical_style = gr.Dropdown(
                        choices=["radiology", "neurosurgery", "research"],
                        value="radiology",
                        label="Report Style",
                    )
                clinical_run = gr.Button("Run Clinical Pipeline", variant="primary", size="lg")

            with gr.Column():
                clinical_output_image = gr.Image(label="Segmentation Result", type="filepath")
                clinical_status = gr.Markdown(label="Pipeline Status")

        clinical_report_output = gr.Markdown(label="Clinical Report")
        clinical_stats = gr.JSON(label="Measurements", visible=True)

        def handle_clinical_pipeline(image_path, context, modality, style):
            if not image_path:
                return None, "Please upload an image.", "", {}

            try:
                from PIL import Image as PILImage
                from clinical_pipeline import run_clinical_pipeline

                pil_image = PILImage.open(image_path).convert("RGB")

                result = run_clinical_pipeline(
                    image_path=image_path,
                    pil_image=pil_image,
                    clinical_context=context,
                    modality_hint=modality,
                    report_style=style,
                )

                seg_img = result.get("segmentation_image")
                status = f"**Status:** {result['status']}\n\n"
                status += f"**Steps completed:** {', '.join(result['steps_completed'])}\n"
                if result["errors"]:
                    status += f"\n**Warnings:** {'; '.join(result['errors'])}"

                report = result.get("report", "Report generation pending.")
                stats = result.get("statistics", {})

                return seg_img, status, report, stats

            except Exception as e:
                logger.error(f"Clinical pipeline error: {e}", exc_info=True)
                return None, f"Error: {str(e)}", "", {}

        clinical_run.click(
            fn=handle_clinical_pipeline,
            inputs=[clinical_image, clinical_context, clinical_modality, clinical_style],
            outputs=[clinical_output_image, clinical_status, clinical_report_output, clinical_stats],
        )


def create_research_pipeline_tab():
    """Create the Research Pipeline tab — batch analysis."""
    with gr.Tab("Research Pipeline", id="research_pipeline"):
        gr.Markdown("""
        ### Batch Research Analysis
        Upload multiple images for consistent batch processing with cross-subject statistics.

        **Outputs:** Per-image statistics CSV, aggregate summary, ZIP export bundle.
        """)

        with gr.Row():
            with gr.Column():
                research_images = gr.File(
                    label="Upload Multiple Images",
                    file_types=[".dcm", ".png", ".jpg", ".jpeg"],
                    type="filepath",
                    file_count="multiple",
                )
                research_prompt = gr.Textbox(
                    label="Segmentation Prompt (consistent across all images)",
                    value="brain",
                    placeholder="e.g., 'ventricles', 'tumor', 'brain'",
                )
                with gr.Row():
                    research_modality = gr.Dropdown(
                        choices=["CT", "MRI"],
                        value="MRI",
                        label="Modality",
                    )
                    research_group = gr.Checkbox(
                        label="Group by Subject",
                        value=True,
                    )
                research_run = gr.Button("Run Research Pipeline", variant="primary", size="lg")

            with gr.Column():
                research_status = gr.Markdown(label="Pipeline Status")
                research_aggregate = gr.JSON(label="Aggregate Statistics")

        research_summary = gr.Textbox(label="Research Summary", lines=15, interactive=False)
        research_download = gr.File(label="Download Export Bundle (ZIP)")

        def handle_research_pipeline(image_files, prompt, modality, group_by):
            if not image_files:
                return "No images uploaded.", {}, "", None

            try:
                from research_pipeline import run_research_pipeline

                # Handle file paths from Gradio
                if isinstance(image_files, str):
                    file_paths = [image_files]
                else:
                    file_paths = image_files if isinstance(image_files, list) else [image_files]

                result = run_research_pipeline(
                    image_files=file_paths,
                    prompt=prompt,
                    modality=modality,
                    group_by_subject=group_by,
                )

                status = f"**Status:** {result['status']}\n"
                status += f"**Processed:** {result['processed']}/{result['total_images']}\n"
                if result['failed'] > 0:
                    status += f"**Failed:** {result['failed']}"

                agg = result.get("aggregate_stats", {})
                summary = result.get("summary", "No summary available.")
                zip_path = result.get("zip_path")

                return status, agg, summary, zip_path

            except Exception as e:
                logger.error(f"Research pipeline error: {e}", exc_info=True)
                return f"Error: {str(e)}", {}, "", None

        research_run.click(
            fn=handle_research_pipeline,
            inputs=[research_images, research_prompt, research_modality, research_group],
            outputs=[research_status, research_aggregate, research_summary, research_download],
        )


def create_model_comparison_tab():
    """Create Model Comparison tab — side-by-side model evaluation."""
    with gr.Tab("Model Comparison", id="model_comparison"):
        gr.Markdown("""
        ### Model Comparison
        Run the same image through multiple segmentation models and compare results side-by-side.

        **Available models:** SAM3 (text-prompt), MedSAM (bounding-box), U-Net (placeholder)
        """)

        with gr.Row():
            with gr.Column():
                compare_image = gr.File(
                    label="Upload Medical Image",
                    file_types=[".dcm", ".png", ".jpg", ".jpeg"],
                    type="filepath",
                )
                compare_prompt = gr.Textbox(
                    label="Segmentation Prompt",
                    value="brain",
                )
                compare_modality = gr.Dropdown(
                    choices=["CT", "MRI"],
                    value="MRI",
                    label="Modality",
                )
                compare_run = gr.Button("Compare Models", variant="primary")

        with gr.Row():
            compare_sam3_output = gr.Image(label="SAM3 Result", type="filepath")
            compare_medsam_output = gr.Image(label="MedSAM Result", type="filepath")

        compare_report = gr.Markdown(label="Comparison Report")

        def handle_model_comparison(image_path, prompt, modality):
            if not image_path:
                return None, None, "Please upload an image."

            try:
                from PIL import Image as PILImage
                from models import run_sam3_inference, is_model_loaded
                from medsam_inference import segment_with_medsam, is_medsam_loaded
                from utils import create_output_image, combine_masks
                from report_generator import generate_comparison_report
                import numpy as np

                pil_image = PILImage.open(image_path).convert("RGB")
                w, h = pil_image.size

                results_list = []
                models_used = []

                # SAM3
                sam3_output = None
                if is_model_loaded():
                    sam3_results = run_sam3_inference(pil_image, prompt, threshold=0.1, mask_threshold=0.0)
                    if sam3_results and "masks" in sam3_results and sam3_results["masks"]:
                        mask = combine_masks(sam3_results["masks"])
                        sam3_output = create_output_image(pil_image, mask, f"SAM3: {prompt}")
                        results_list.append({"mask": mask, "score": 0.9})
                    else:
                        results_list.append(None)
                else:
                    results_list.append(None)
                models_used.append("SAM3")

                # MedSAM (use center bounding box as default)
                medsam_output = None
                bbox = [int(w*0.1), int(h*0.1), int(w*0.9), int(h*0.9)]
                medsam_result = segment_with_medsam(pil_image, bbox, modality)
                if medsam_result and medsam_result.get("mask") is not None:
                    medsam_output = create_output_image(
                        pil_image, medsam_result["mask"], "MedSAM"
                    )
                    results_list.append(medsam_result)
                else:
                    results_list.append(None)
                    medsam_output = None
                models_used.append("MedSAM")

                # Generate comparison report
                report = generate_comparison_report(results_list, models_used)

                return sam3_output, medsam_output, report

            except Exception as e:
                logger.error(f"Model comparison error: {e}", exc_info=True)
                return None, None, f"Error: {str(e)}"

        compare_run.click(
            fn=handle_model_comparison,
            inputs=[compare_image, compare_prompt, compare_modality],
            outputs=[compare_sam3_output, compare_medsam_output, compare_report],
        )


def add_agentic_tabs_to_app(demo_blocks):
    """
    Add all agentic tabs to an existing Gradio Blocks instance.
    Call this within the gr.Blocks() context before existing tabs.
    """
    create_agent_chat_tab()
    create_clinical_pipeline_tab()
    create_research_pipeline_tab()
    create_model_comparison_tab()