main.py

import asyncio
import io
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image
import micropip
import base64
import sys
from io import StringIO
from pyodide.http import pyfetch

# PyScript specific imports to interact with the browser DOM
import js
from pyscript import display, document

# ==========================================
# 1. GLOBAL STATE & FUNCTION STUBS
# ==========================================
state = {
    'OrigBlue': None, 'OrigGreen': None, 'OrigRed': None,
    'CropBlue': None, 'CropBlue_Mask': None, 'Nuclei': None
}

crop_fn = None
RemovalMethod = None
identify_primary_fn = None

# ==========================================
# 2. ASYNC INITIALIZATION
# ==========================================
async def setup_environment():
    status_text = document.getElementById('loading-status')
    try:
        wheel_url = "https://sspathak.github.io/pyodide-packages/cellprofiler_library-5.0.2-py3-none-any.whl"
        
        print(f"System: Fetching wheel from {wheel_url}...")
        status_text.innerText = "Downloading CellProfiler WASM Library..."
        await micropip.install(wheel_url)
        
        print("System: ✅ Successfully installed cellprofiler_library!")
        status_text.innerText = "Library installed. Compiling dependencies..."
        
        global crop_fn, RemovalMethod, identify_primary_fn
        from cellprofiler_library.modules._crop import crop as crop_fn
        from cellprofiler_library.opts.crop import RemovalMethod
        from cellprofiler_library.modules._identifyprimaryobjects import identifyprimaryobjects as identify_primary_fn
        print("""
        from cellprofiler_library.modules._crop import crop as crop_fn
        from cellprofiler_library.opts.crop import RemovalMethod
        from cellprofiler_library.modules._identifyprimaryobjects import identifyprimaryobjects as identify_primary_fn
        """)
        print("System: ✅ Library modules imported successfully.")
        
        # ----------------------------------------------------
        # THE UI REVEAL SEQUENCE
        # ----------------------------------------------------
        status_text.innerText = "Ready!"
        
        # 1. Fade out the white loading screen
        document.getElementById('loading-overlay').style.opacity = '0'
        
        # 2. Auto-minimize the Developer Dock using our JS function
        # js.toggleDock()
        
        # Wait for the CSS fade transition to finish before hiding overlay completely
        await asyncio.sleep(0.5)
        document.getElementById('loading-overlay').style.display = 'none'
        
        # 3. Reveal the main UI
        document.getElementById('ui-container').style.display = 'block'
        
    except Exception as e:
        print(f"System: ❌ Failed to install or import wheel: {e}")
        status_text.innerText = f"Error: Failed to load environment.\n{e}"
        status_text.style.color = "red"

# ==========================================
# 3. HELPER FUNCTIONS
# ==========================================
async def get_image_from_input(input_id):
    file_list = document.getElementById(input_id).files
    if file_list.length == 0: return None
    
    file = file_list.item(0)
    js_buffer = await file.arrayBuffer()
    js_array = js.Uint8Array.new(js_buffer)
    bytes_data = js_array.to_bytes()
    
    img = Image.open(io.BytesIO(bytes_data)).convert('L')
    return np.array(img, dtype=np.float32) / 255.0

async def get_image_from_url(url_string):
    """Fetches an image from a URL and converts to normalized Numpy grayscale float array."""
    if not url_string.strip(): 
        return None
        
    response = await pyfetch(url_string)
    if not response.ok:
        raise Exception(f"HTTP {response.status}: Failed to fetch {url_string}")
        
    bytes_data = await response.bytes()
    img = Image.open(io.BytesIO(bytes_data)).convert('L')
    return np.array(img, dtype=np.float32) / 255.0

def plot_outputs(img1, img2, title1, title2, target_div_id):
    """Renders Matplotlib figure to a Base64 image string for bulletproof DOM insertion."""
    fig, axes = plt.subplots(1, 2, figsize=(10, 5))
    
    # Plot original grayscale image
    axes[0].imshow(img1, cmap='gray')
    axes[0].set_title(title1)
    axes[0].axis('off')
    
    # Detect if img2 is a segmentation label matrix (integers > 1)
    if img2.dtype in [np.int32, np.int64, np.uint16, np.uint32] and np.max(img2) > 1:
        # CellProfiler Style: Use a discrete colormap but force 0 (background) to black
        import matplotlib.colors as mcolors
        cmap = plt.get_cmap('tab20')
        cmap.set_under('black')
        axes[1].imshow(img2, cmap=cmap, vmin=0.1, interpolation='nearest')
    else:
        # Standard mask or grayscale
        axes[1].imshow(img2, cmap='gray')
        
    axes[1].set_title(title2)
    axes[1].axis('off')
    
    plt.tight_layout()
    
    # ---------------------------------------------------------
    # The Fix: Render to Bytes, encode to Base64, and inject HTML
    # ---------------------------------------------------------
    buf = io.BytesIO()
    fig.savefig(buf, format='png', bbox_inches='tight')
    buf.seek(0)
    
    img_str = base64.b64encode(buf.read()).decode('utf-8')
    img_html = f'<img src="data:image/png;base64,{img_str}" style="max-width: 100%; height: auto; border-radius: 4px;">'
    
    # Insert directly into the DOM
    document.getElementById(target_div_id).innerHTML = img_html
    
    # Now it is completely safe to close the figure and free memory!
    plt.close(fig)

# ==========================================
# 4. UI EVENT HANDLERS
# ==========================================
async def on_load_images(event):
    out_div = document.getElementById("out-uploads")
    out_div.innerText = "⏳ Initializing image load..."
    out_div.style.color = "#ffc107"
    await asyncio.sleep(0.05)
    
    # Check the selected mode
    mode = document.querySelector('input[name="source-type"]:checked').value
    
    try:
        if mode == "upload":
            print("System: Loading 3-channel set from local files...")
            state['OrigBlue'] = await get_image_from_input('upload-blue')
            state['OrigGreen'] = await get_image_from_input('upload-green')
            state['OrigRed'] = await get_image_from_input('upload-red')
        else:
            print("System: Fetching 3-channel set from CellProfiler GitHub...")
            # Note: Restored OrigRed URL mapping
            url_blue = document.getElementById('url-blue').value
            url_green = document.getElementById('url-green').value
            url_red = document.getElementById('url-red').value
            
            state['OrigBlue'] = await get_image_from_url(url_blue)
            state['OrigGreen'] = await get_image_from_url(url_green)
            state['OrigRed'] = await get_image_from_url(url_red)

        # Basic validation: At least Blue and Green are required for the demo pipeline
        if state['OrigBlue'] is not None and state['OrigGreen'] is not None:
            out_div.innerText = f"✅ Success: 3 Channels loaded via {mode.upper()}"
            out_div.style.color = "#28a745"
            print(f"System: Loaded shapes: B:{state['OrigBlue'].shape}, G:{state['OrigGreen'].shape}")
        else:
            out_div.innerText = "⚠️ Load failed: Missing required channels."
            out_div.style.color = "#dc3545"

    except Exception as e:
        print(f"FATAL ERROR during load: {e}")
        out_div.innerText = f"❌ Load Failed: {str(e)}"
        out_div.style.color = "#dc3545"

async def on_crop_b(event):
    status_div = document.getElementById("out-crop-b-status")
    status_div.innerText = "⏳ Processing crop..."
    status_div.style.color = "#ffc107"
    await asyncio.sleep(0.05)
    
    if crop_fn is None or state['OrigBlue'] is None:
        status_div.innerText = "❌ Error: Library or OrigBlue missing."
        status_div.style.color = "#dc3545"
        return
        
    try:
        t = int(document.getElementById("crop-b-top").value)
        b = int(document.getElementById("crop-b-bottom").value)
        l = int(document.getElementById("crop-b-left").value)
        r = int(document.getElementById("crop-b-right").value)
        
        img_shape = state['OrigBlue'].shape
        cropping_mask = np.zeros(img_shape, dtype=bool)
        cropping_mask[t:b, l:r] = True
        
        cropped_pixels, out_mask, _ = crop_fn(
            orig_image_pixels=state['OrigBlue'],
            cropping=cropping_mask, mask=None, orig_image_mask=None,
            removal_method=RemovalMethod.EDGES
        )
        
        state['CropBlue'] = cropped_pixels
        state['CropBlue_Mask'] = out_mask
        
        status_div.innerText = "✅ Crop Complete!"
        status_div.style.color = "#28a745"
        
        # Render and automatically pop open the dropdown
        plot_outputs(state['OrigBlue'], state['CropBlue'], "OrigBlue", "CropBlue", "out-crop-b-viz")
        document.getElementById("details-crop-b").open = True
        
    except Exception as e:
        status_div.innerText = f"❌ Execution Error:\n{e}"
        status_div.style.color = "#dc3545"

async def on_ip(event):
    status_div = document.getElementById("out-ip-status")
    status_div.innerText = "⏳ Identifying nuclei..."
    status_div.style.color = "#ffc107"
    await asyncio.sleep(0.05)
    
    if identify_primary_fn is None or state['CropBlue'] is None:
        status_div.innerText = "❌ Error: Library or CropBlue missing."
        status_div.style.color = "#dc3545"
        return
        
    try:
        min_d = int(document.getElementById("ip-min").value)
        max_d = int(document.getElementById("ip-max").value)
        
        state['Nuclei'] = identify_primary_fn(
            image=state['CropBlue'],
            min_size=min_d, max_size=max_d,
            exclude_size=True, return_cp_output=False
        )
        
        status_div.innerText = "✅ Identify Primary Complete!"
        status_div.style.color = "#28a745"
        
        # Render the segmentation and open the dropdown
        plot_outputs(state['CropBlue'], state['Nuclei'], "CropBlue", "Nuclei Segmentations", "out-ip-viz")
        document.getElementById("details-ip").open = True
        
    except Exception as e:
        status_div.innerText = f"❌ Execution Error:\n{e}"
        status_div.style.color = "#dc3545"

async def on_repl_run(event):
    """Executes custom Python code, captures stdout, and writes directly to the DOM."""
    code_input = document.getElementById("repl-input")
    code = code_input.value
    if not code.strip():
        return
        
    terminal = document.getElementById("boot-terminal")
    
    # 1. Echo the command to the terminal safely
    safe_code = code.strip().replace('<', '&lt;').replace('>', '&gt;')
    terminal.innerHTML += f"<div style='color: #00ff00;'>\n[REPL] >>> {safe_code}</div>"
    
    await asyncio.sleep(0.05)
    
    try:
        # 2. Try to evaluate as an expression first
        try:
            result = eval(code, globals())
            if result is not None:
                safe_res = repr(result).replace('<', '&lt;').replace('>', '&gt;')
                terminal.innerHTML += f"<div style='color: #fff;'>{safe_res}</div>"
                
        # 3. If it's a statement, use exec() and capture ANY internal print() statements
        except SyntaxError:
            old_stdout = sys.stdout
            redirected_output = sys.stdout = StringIO()
            try:
                exec(code, globals())
            finally:
                sys.stdout = old_stdout # Always restore stdout!
                
            # If the user's code printed anything, inject it
            printed_out = redirected_output.getvalue()
            if printed_out:
                safe_out = printed_out.replace('<', '&lt;').replace('>', '&gt;')
                terminal.innerHTML += f"<div style='color: #fff;'>{safe_out}</div>"
                
    except Exception as e:
        terminal.innerHTML += f"<div class='log-error'>[REPL ERROR]: {e}</div>"
        
    # 4. Clear the input and auto-scroll the terminal to the bottom
    code_input.value = ""
    js.eval("var t = document.getElementById('boot-terminal'); t.scrollTop = t.scrollHeight;")
    
# Start the boot sequence
asyncio.ensure_future(setup_environment())