app1.py

# app.py
"""
MedSketch AI: Advanced Clinical Diagram Generator

A Streamlit application leveraging AI models (GPT-4o, potentially Stable Diffusion)
to generate medical diagrams based on user prompts, with options for styling,
metadata association, and annotations.
"""

import os
import json
import logging
from io import BytesIO
from typing import List, Dict, Any, Optional, Tuple

import streamlit as st
from streamlit_drawable_canvas import st_canvas
from PIL import Image
import openai
from openai import OpenAI, OpenAIError  # Use modern OpenAI client and error types

# ─── Constants ───────────────────────────────────────────────────────────────

APP_TITLE = "MedSketch AI – Advanced Clinical Diagram Generator"
DEFAULT_MODEL = "GPT-4o (Vision)" # Updated model name
STABLE_DIFFUSION_MODEL = "Stable Diffusion LoRA" # Placeholder name
MODEL_OPTIONS = [DEFAULT_MODEL, STABLE_DIFFUSION_MODEL]
STYLE_PRESETS = ["Anatomical Diagram", "H&E Histology", "IHC Pathology", "Custom"]
DEFAULT_STYLE = "Anatomical Diagram"
DEFAULT_STRENGTH = 0.7
IMAGE_SIZE = "1024x1024"
CANVAS_SIZE = 512
ANNOTATION_COLOR = "rgba(255, 0, 0, 0.3)" # Red with transparency
ANNOTATION_STROKE_WIDTH = 2
SESSION_STATE_ANNOTATIONS = "medsketch_annotations"
SESSION_STATE_HISTORY = "medsketch_history" # Store generated images too

# ─── Setup & Configuration ────────────────────────────────────────────────────

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

st.set_page_config(
    page_title=APP_TITLE,
    layout="wide",
    initial_sidebar_state="expanded",
    menu_items={
        'About': f"{APP_TITLE} - Generates medical diagrams using AI.",
        'Get Help': None, # Add a link if you have one
        'Report a bug': None # Add a link if you have one
    }
)

# Initialize OpenAI Client (Best Practice)
# Use st.secrets for deployment, fallback to env var for local dev
api_key = st.secrets.get("OPENAI_API_KEY", os.getenv("OPENAI_API_KEY"))
if not api_key:
    st.error("🚨 OpenAI API Key not found! Please set it in Streamlit secrets or environment variables.", icon="🚨")
    st.stop() # Halt execution if no key

try:
    client = OpenAI(api_key=api_key)
    logger.info("OpenAI client initialized successfully.")
except Exception as e:
    st.error(f"🚨 Failed to initialize OpenAI client: {e}", icon="🚨")
    logger.exception("OpenAI client initialization failed.")
    st.stop()

# ─── Helper Functions ─────────────────────────────────────────────────────────

def generate_openai_image(prompt: str, style: str, strength: float) -> Image.Image:
    """
    Generates an image using the OpenAI API (GPT-4o).

    Args:
        prompt: The user's text prompt.
        style: The selected style preset.
        strength: The stylization strength (conceptually used in prompt).

    Returns:
        A PIL Image object.

    Raises:
        OpenAIError: If the API call fails.
        IOError: If the image data cannot be processed.
    """
    logger.info(f"Requesting OpenAI image generation for prompt: '{prompt}' with style '{style}'")
    full_prompt = f"Style: [{style}], Strength: [{strength:.2f}] - Generate the following medical illustration: {prompt}"
    try:
        response = client.images.generate(
            model="dall-e-3", # Or "gpt-4o" if/when available via this endpoint. DALL-E 3 is current standard.
            prompt=full_prompt,
            size=IMAGE_SIZE,
            quality="standard", # or "hd"
            n=1,
            response_format="url" # Or "b64_json" to avoid a second request
        )
        image_url = response.data[0].url
        logger.info(f"Image generated successfully, URL: {image_url}")

        # Fetch the image data from the URL
        # Note: Using response_format="b64_json" would avoid this extra step
        import requests # Need to import requests library
        image_response = requests.get(image_url, timeout=30) # Add timeout
        image_response.raise_for_status() # Check for HTTP errors

        img_data = BytesIO(image_response.content)
        img = Image.open(img_data)
        return img

    except OpenAIError as e:
        logger.error(f"OpenAI API error: {e}")
        st.error(f"❌ OpenAI API Error: {e}", icon="❌")
        raise
    except requests.exceptions.RequestException as e:
        logger.error(f"Failed to download image from URL {image_url}: {e}")
        st.error(f"❌ Network Error: Failed to download image. {e}", icon="❌")
        raise IOError(f"Failed to download image: {e}") from e
    except Exception as e:
        logger.exception(f"An unexpected error occurred during OpenAI image generation: {e}")
        st.error(f"❌ An unexpected error occurred: {e}", icon="❌")
        raise


def generate_sd_image(prompt: str, style: str, strength: float) -> Image.Image:
    """
    Placeholder for generating an image using a Stable Diffusion LoRA model.
    Replace this with your actual implementation.

    Args:
        prompt: The user's text prompt.
        style: The selected style preset.
        strength: The stylization strength.

    Returns:
        A PIL Image object (dummy implementation).

    Raises:
        NotImplementedError: As this is a placeholder.
    """
    logger.warning("Stable Diffusion LoRA model generation is not implemented. Returning placeholder.")
    st.warning("🚧 Stable Diffusion LoRA generation is not yet implemented. Using placeholder.", icon="🚧")

    # --- Placeholder Implementation ---
    # Replace this with actual SD model call
    # For now, create a simple dummy image with text
    img = Image.new('RGB', (CANVAS_SIZE, CANVAS_SIZE), color = (210, 210, 210))
    from PIL import ImageDraw
    d = ImageDraw.Draw(img)
    d.text((10,10), f"Stable Diffusion Placeholder\nStyle: {style}\nPrompt: {prompt[:50]}...", fill=(0,0,0))
    # --- End Placeholder ---

    # Simulate some processing time
    import time
    time.sleep(1)
    return img
    # raise NotImplementedError("Stable Diffusion LoRA generation is not yet available.")


def display_result(image: Image.Image, prompt: str, index: int, total: int) -> Optional[List[Dict[str, Any]]]:
    """
    Displays a generated image, download button, and annotation canvas.

    Args:
        image: The PIL Image to display.
        prompt: The prompt used to generate the image.
        index: The index of the current image in a batch.
        total: The total number of images in the batch.

    Returns:
        Annotation data (list of dicts) if annotations were made, otherwise None.
    """
    st.image(image, caption=f"Result {index + 1}/{total}: {prompt}", use_container_width=True)

    # Prepare image for download
    buf = BytesIO()
    image.save(buf, format="PNG")
    buf.seek(0)

    st.download_button(
        label="⬇️ Download PNG",
        data=buf,
        file_name=f"medsketch_{index+1}_{prompt[:20].replace(' ', '_')}.png",
        mime="image/png",
        key=f"download_{index}"
    )

    # Annotation Canvas
    st.markdown("**✏️ Annotate:**")
    # Resize image for canvas if needed, maintaining aspect ratio (optional)
    # For simplicity, we assume the canvas size matches desired annotation size
    canvas_image = image.copy()
    canvas_image.thumbnail((CANVAS_SIZE, CANVAS_SIZE))

    canvas_result = st_canvas(
        fill_color=ANNOTATION_COLOR,
        stroke_width=ANNOTATION_STROKE_WIDTH,
        background_image=canvas_image,
        update_streamlit=True, # Update in real-time
        height=canvas_image.height,
        width=canvas_image.width,
        drawing_mode="freedraw", # Or choose other modes like "line", "rect", etc.
        key=f"canvas_{index}"
    )

    if canvas_result.json_data and canvas_result.json_data.get("objects"):
        return canvas_result.json_data["objects"]
    return None

# ─── Initialize Session State ───────────────────────────────────────────────

if SESSION_STATE_ANNOTATIONS not in st.session_state:
    st.session_state[SESSION_STATE_ANNOTATIONS] = {} # Dict[prompt, List[annotation_objects]]
if SESSION_STATE_HISTORY not in st.session_state:
    st.session_state[SESSION_STATE_HISTORY] = [] # List[Dict[str, Any]] storing generation results

# ─── Sidebar: Settings & Metadata ───────────────────────────────────────────

with st.sidebar:
    st.header("⚙️ Generation Settings")
    model_choice = st.selectbox(
        "Select Model",
        options=MODEL_OPTIONS,
        index=MODEL_OPTIONS.index(DEFAULT_MODEL),
        help="Choose the AI model for image generation."
    )

    style_preset = st.radio(
        "Select Preset Style",
        options=STYLE_PRESETS,
        index=STYLE_PRESETS.index(DEFAULT_STYLE),
        horizontal=True, # More compact layout
        help="Apply a predefined visual style to the generation."
    )
    # Allow custom style input only if "Custom" is selected
    custom_style_input = ""
    if style_preset == "Custom":
        custom_style_input = st.text_input("Enter Custom Style Description:", key="custom_style")
    final_style = custom_style_input if style_preset == "Custom" else style_preset


    strength = st.slider(
        "Stylization Strength",
        min_value=0.1,
        max_value=1.0,
        value=DEFAULT_STRENGTH,
        step=0.05,
        help="Controls how strongly the chosen style influences the result (conceptual)."
    )

    st.markdown("---")
    st.header("📋 Optional Metadata")
    patient_id = st.text_input("Patient / Case ID", key="patient_id", help="Associate with a specific patient or case.")
    roi = st.text_input("Region of Interest (ROI)", key="roi", help="Specify the anatomical region shown.")
    umls_code = st.text_input("UMLS / SNOMED CT Code", key="umls_code", help="Link to relevant medical ontology codes.")

    # Add a clear history button
    st.markdown("---")
    if st.button("⚠️ Clear History & Annotations", help="Removes all generated images and annotations from this session."):
        st.session_state[SESSION_STATE_ANNOTATIONS] = {}
        st.session_state[SESSION_STATE_HISTORY] = []
        st.rerun() # Refresh the page to reflect cleared state

# ─── Main Application Area ───────────────────────────────────────────────────

st.title(APP_TITLE)
st.markdown("Generate medical illustrations from text descriptions using AI. Annotate and export your results.")

# --- Prompt Input Area ---
prompt_input_area = st.container()
with prompt_input_area:
    st.subheader("📝 Enter Prompt(s)")
    st.caption("Enter one prompt per line to generate multiple images in a batch.")
    raw_prompts = st.text_area(
        "Describe the medical diagram(s) you need:",
        placeholder=(
            "Example 1: A sagittal view of the human knee joint, labeling the ACL, PCL, meniscus, femur, and tibia.\n"
            "Example 2: High-power field H&E stain of lung adenocarcinoma showing glandular formation.\n"
            "Example 3: Immunohistochemistry (IHC) stain for PD-L1 in tonsil tissue, showing positive staining on immune cells."
        ),
        height=150, # Slightly larger height
        label_visibility="collapsed"
    )
    prompts: List[str] = [p.strip() for p in raw_prompts.splitlines() if p.strip()]

    # --- Generation Trigger ---
    generate_button = st.button(
        f"🚀 Generate Diagram{'s' if len(prompts) > 1 else ''}",
        type="primary",
        disabled=not prompts, # Disable if no prompts
        use_container_width=True
    )

# --- Generation and Display Area ---
results_area = st.container()
if generate_button:
    if not prompts:
        st.warning("⚠️ Please enter at least one prompt description.", icon="⚠️")
    else:
        logger.info(f"Starting generation for {len(prompts)} prompts using model '{model_choice}'.")
        num_prompts = len(prompts)
        max_cols = 3 # Adjust number of columns based on screen width or preference
        cols = st.columns(min(max_cols, num_prompts))

        # Use a progress bar for batch generation
        progress_bar = st.progress(0, text=f"Initializing generation...")

        for i, prompt in enumerate(prompts):
            col_index = i % max_cols
            with cols[col_index]:
                st.markdown(f"--- \n**Processing: {i+1}/{num_prompts}**")
                spinner_msg = f"Generating image {i+1}/{num_prompts} for prompt: \"{prompt[:50]}...\""
                with st.spinner(spinner_msg):
                    try:
                        # Select generation function based on model choice
                        if model_choice == DEFAULT_MODEL:
                            generated_image = generate_openai_image(prompt, final_style, strength)
                        elif model_choice == STABLE_DIFFUSION_MODEL:
                            generated_image = generate_sd_image(prompt, final_style, strength)
                        else:
                            st.error(f"Unknown model selected: {model_choice}", icon="❌")
                            continue # Skip to next prompt

                        # Display result and get annotations
                        annotations = display_result(generated_image, prompt, i, num_prompts)

                        # Store results and annotations in session state
                        result_data = {
                            "prompt": prompt,
                            "model": model_choice,
                            "style": final_style,
                            "strength": strength,
                            "metadata": {
                                "patient_id": patient_id,
                                "roi": roi,
                                "umls_code": umls_code,
                            },
                            # Store image data efficiently (e.g., as base64 or keep PIL object if memory allows)
                            # For simplicity here, we might just store prompt and annotations.
                            # Storing images in session state can consume a lot of memory.
                            # Let's store the prompt reference and annotations.
                            "image_ref_index": i # Reference to this generation instance
                        }
                        st.session_state[SESSION_STATE_HISTORY].append(result_data)

                        if annotations:
                            st.session_state[SESSION_STATE_ANNOTATIONS][prompt] = annotations
                            st.success(f"Annotations saved for prompt {i+1}.", icon="✅")

                    except (OpenAIError, IOError, NotImplementedError, Exception) as e:
                        # Errors are logged and displayed by the generation functions
                        st.error(f"Failed to generate image for prompt: '{prompt}'. Error: {e}", icon="🔥")
                        # Optionally add failed attempts to history?
                        st.session_state[SESSION_STATE_HISTORY].append({
                            "prompt": prompt, "status": "failed", "error": str(e)
                        })

            # Update progress bar
            progress_val = (i + 1) / num_prompts
            progress_bar.progress(progress_val, text=f"Generated {i+1}/{num_prompts} images...")

        progress_bar.progress(1.0, text="Batch generation complete!")
        st.toast(f"Finished generating {num_prompts} image(s)!", icon="🎉")
        # Explicitly clear the progress bar after completion
        # (Streamlit often handles this, but explicit removal can be cleaner)
        # Consider removing or hiding the progress bar element if needed after completion.


# ─── History & Exports Section ───────────────────────────────────────────────

history_area = st.container()
with history_area:
    # Use session state history which is more robust
    if st.session_state[SESSION_STATE_HISTORY]:
        st.markdown("---")
        st.subheader("📚 Session History & Annotations")
        st.caption("Review generated images (if stored) and their annotations from this session.")

        # Display stored history (simplified view focusing on annotations)
        for idx, item in enumerate(st.session_state[SESSION_STATE_HISTORY]):
            if item.get("status") == "failed":
                 st.warning(f"**Prompt {idx+1} (Failed):** {item['prompt']} \n *Error: {item['error']}*", icon="⚠️")
            else:
                prompt_key = item["prompt"]
                st.markdown(f"**Prompt {idx+1}:** `{prompt_key}`")
                st.markdown(f"*Model: {item['model']}, Style: {item['style']}*")
                # Display metadata if present
                meta = item.get('metadata', {})
                if any(meta.values()):
                    meta_str = ", ".join([f"{k}: {v}" for k, v in meta.items() if v])
                    st.markdown(f"*Metadata: {meta_str}*")

                # Check for annotations for this prompt
                annotations = st.session_state[SESSION_STATE_ANNOTATIONS].get(prompt_key)
                if annotations:
                    with st.expander(f"View {len(annotations)} Annotation(s)"):
                        st.json(annotations)
                else:
                    st.caption("_(No annotations made for this item yet)_")
            st.markdown("---") # Separator between history items


        # --- Export Annotations ---
        if st.session_state[SESSION_STATE_ANNOTATIONS]:
            st.markdown("---")
            st.subheader("⬇️ Export Annotations")
            try:
                # Prepare data with metadata included per annotation set
                export_data = {}
                # Find corresponding history item to enrich annotation export
                history_map = {item['prompt']: item for item in st.session_state[SESSION_STATE_HISTORY] if item.get('status') != 'failed'}

                for prompt, ann_objs in st.session_state[SESSION_STATE_ANNOTATIONS].items():
                     history_item = history_map.get(prompt)
                     export_data[prompt] = {
                         "annotations": ann_objs,
                         "generation_details": {
                             "model": history_item.get('model'),
                             "style": history_item.get('style'),
                             "strength": history_item.get('strength'),
                         } if history_item else None,
                         "metadata": history_item.get('metadata') if history_item else None
                     }

                json_data = json.dumps(export_data, indent=2)
                st.download_button(
                    label="⬇️ Export All Annotations (JSON)",
                    data=json_data,
                    file_name="medsketch_session_annotations.json",
                    mime="application/json",
                    help="Download all annotations made during this session, including associated metadata."
                )
            except Exception as e:
                st.error(f"Failed to prepare annotations for download: {e}")
                logger.error(f"Error preparing JSON export: {e}")

    elif generate_button: # If generate was clicked but history is empty (e.g., all failed)
        st.info("No successful generations or annotations in the current session yet.")

# Add a footer (optional)
st.markdown("---")
st.caption("MedSketch AI - Powered by Streamlit and OpenAI")