bubcount/cli.py

"""
bubcount — microbubble sizing CLI
"""
from pathlib import Path
from typing import Optional

import numpy as np
import typer
from rich.columns import Columns
from rich.console import Console
from rich.panel import Panel
from rich.progress import (
    BarColumn,
    MofNCompleteColumn,
    Progress,
    SpinnerColumn,
    TaskProgressColumn,
    TextColumn,
    TimeElapsedColumn,
)
from rich.table import Table
from rich.text import Text
from skimage import io

from .analyzer import BubbleAnalyzer
from .data import AnalysisResults
from .exporter import ResultsExporter
from .params import AnalysisParameters

console = Console()
app = typer.Typer(
    name="bubcount",
    help="Microbubble sizing from optical microscopy using Cellpose.",
    add_completion=False,
)


# ---------------------------------------------------------------------------
# Main command
# ---------------------------------------------------------------------------

@app.command()
def run(
    image_dir: Path = typer.Argument(..., help="Directory containing microscopy images"),
    output: Optional[Path] = typer.Option(None, "--output", "-o", help="Output directory (default: <image_dir>/results)"),
    model: str = typer.Option("cpsam", "--model", "-m", help="Cellpose pretrained model path or name"),
    scale: float = typer.Option(0.0825, "--scale", help="Micrometres per pixel"),
    volume: float = typer.Option(0.00089, "--volume", help="Sample volume per frame (μL)"),
    min_diam: float = typer.Option(0.5, "--min-diam", help="Minimum bubble diameter (μm)"),
    max_diam: float = typer.Option(50.0, "--max-diam", help="Maximum bubble diameter (μm)"),
    no_gpu: bool = typer.Option(False, "--no-gpu", help="Disable GPU"),
    min_circularity: float = typer.Option(0.5, "--min-circularity", help="Minimum circularity (0–1)"),
    max_aspect_ratio: float = typer.Option(2.0, "--max-aspect-ratio", help="Maximum aspect ratio"),
):
    _print_header()

    image_dir = image_dir.expanduser().resolve()
    if not image_dir.exists():
        console.print(f"[red]Error:[/red] directory not found: {image_dir}")
        raise typer.Exit(1)

    output_dir = (output or image_dir / "results").expanduser().resolve()

    params = AnalysisParameters(
        scale_um_per_pixel=scale,
        sample_volume_per_frame_uL=volume,
        min_diameter_um=min_diam,
        max_diameter_um=max_diam,
        pretrained_model=model,
        gpu=not no_gpu,
        min_circularity=min_circularity,
        max_aspect_ratio=max_aspect_ratio,
    )

    _print_params(image_dir, output_dir, params)

    # ---- load model --------------------------------------------------------
    with console.status("[bold cyan]Loading segmentation model…[/bold cyan]"):
        try:
            analyzer = BubbleAnalyzer(params=params)
        except Exception as exc:
            console.print(f"[red]Failed to load model:[/red] {exc}")
            raise typer.Exit(1)


    # ---- discover images ---------------------------------------------------
    image_files = analyzer.list_images(image_dir)
    if not image_files:
        console.print(f"[red]No images found in {image_dir}[/red]")
        raise typer.Exit(1)

    # ---- analyse with live progress ----------------------------------------
    from .data import AnalysisResults

    results = AnalysisResults(
        sample_name=image_dir.name,
        parameters=params,
    )

    with Progress(
        SpinnerColumn(),
        TextColumn("[progress.description]{task.description}"),
        BarColumn(),
        MofNCompleteColumn(),
        TaskProgressColumn(),
        TimeElapsedColumn(),
        console=console,
        transient=False,
    ) as progress:
        task = progress.add_task(
            f"[cyan]Analysing {image_dir.name}[/cyan]", total=len(image_files)
        )

        for img_path in image_files:
            progress.update(task, description=f"[cyan]{img_path.name}[/cyan]")
            image = io.imread(str(img_path))
            frame = analyzer.analyze_image(image, img_path.stem, img_path)
            results.frames.append(frame)
            results.all_bubbles.extend(frame.bubbles)
            progress.advance(task)

    # ---- export ------------------------------------------------------------
    with console.status("[bold cyan]Exporting results…[/bold cyan]"):
        exporter = ResultsExporter(results, output_dir)
        exporter.export_all()

    # ---- summary -----------------------------------------------------------
    _print_summary(results, output_dir)


# ---------------------------------------------------------------------------
# Rich helpers
# ---------------------------------------------------------------------------

def _print_header():
    title = Text("bubcount", style="bold white")
    subtitle = Text("microbubble sizing · Cellpose", style="dim")
    console.print()
    console.print(
        Panel(
            f"[bold white]bubcount[/bold white]  [dim]microbubble sizing · Cellpose[/dim]",
            expand=False,
            border_style="bright_cyan",
            padding=(0, 2),
        )
    )
    console.print()


def _print_params(image_dir: Path, output_dir: Path, p: AnalysisParameters):
    t = Table.grid(padding=(0, 2))
    t.add_column(style="dim")
    t.add_column()
    t.add_row("Input",   str(image_dir))
    t.add_row("Output",  str(output_dir))
    t.add_row("Model",   str(p.pretrained_model))
    t.add_row("Scale",   f"{p.scale_um_per_pixel} μm/pixel")
    t.add_row("Volume",  f"{p.sample_volume_per_frame_uL} μL/frame")
    t.add_row("Range",   f"{p.min_diameter_um}–{p.max_diameter_um} μm")
    console.print(Panel(t, title="[bold]Parameters[/bold]", border_style="cyan", expand=False))
    console.print()


def _print_summary(results: AnalysisResults, output_dir: Path):
    d = results.diameters
    p = results.parameters
    total_vol = p.sample_volume_per_frame_uL * results.num_frames
    conc = results.total_bubbles / total_vol if total_vol > 0 else 0

    # Stats table
    stats = Table(show_header=False, box=None, padding=(0, 2))
    stats.add_column(style="dim", no_wrap=True)
    stats.add_column(justify="right")

    stats.add_row("Frames analysed",    str(results.num_frames))
    stats.add_row("Bubbles accepted",   f"[bold green]{results.total_bubbles}[/bold green]")
    stats.add_row("Bubbles rejected",   f"[yellow]{results.total_rejected}[/yellow]")

    if len(d) > 0:
        stats.add_row("", "")
        stats.add_row("Mean diameter",  f"{np.mean(d):.2f} ± {np.std(d):.2f} μm")
        stats.add_row("Median diameter", f"{np.median(d):.2f} μm")
        stats.add_row("Range",          f"{np.min(d):.2f}–{np.max(d):.2f} μm")
        stats.add_row("", "")
        stats.add_row("Concentration",  f"{conc:.3e} bubbles/μL")

    console.print()
    console.print(
        Panel(stats, title="[bold green]Results[/bold green]", border_style="green", expand=False)
    )

    # Output files
    files = Table(show_header=False, box=None, padding=(0, 1))
    files.add_column(style="dim cyan", no_wrap=True)
    files.add_column(style="dim")
    for f in sorted(output_dir.iterdir()):
        files.add_row(f.name, "")

    console.print()
    console.print(
        Panel(
            files,
            title=f"[bold]Output[/bold]  [dim]{output_dir}[/dim]",
            border_style="cyan",
            expand=False,
        )
    )
    console.print()