cylindra.core

def collect_projects(
    files: PathLike | Iterable[PathLike],
    *,
    skip_exc: bool = False,
) -> ProjectSequence:
    """Collect project files into a ProjectSequence object.

    Examples
    --------
    ``` python
    collect_projects("path/to/dir/*.json")
    collect_projects(["path/to/dir-0/*.zip", "path/to/dir-1/*.zip"])
    ```

    Parameters
    ----------
    files : path-like or iterable of path-like
        Project file paths or glob pattern(s).
    """
    from cylindra.project import ProjectSequence
    from cylindra.utils import unwrap_wildcard

    _files = unwrap_wildcard(files)
    if len(_files) == 0:
        raise FileNotFoundError(f"No project files found from the input {files!r}.")
    return ProjectSequence.from_paths(_files, check_scale=False, skip_exc=skip_exc)

def instance(create=False):
    """Get the current CylindraMainWidget instance."""
    ins = _CURRENT_INSTANCE
    if ins is None and create:
        ins = start()
    return ins

def read_batch_project(file: PathLike) -> CylindraBatchProject:
    """Read the Cylindra batch project file."""
    from cylindra.project import CylindraBatchProject

    return CylindraBatchProject.from_file(file)

def read_molecules(
    file: PathLike,
    pos_cols: Sequence[str] = ("z", "y", "x"),
    rot_cols: Sequence[str] = ("zvec", "yvec", "xvec"),
) -> Molecules:
    """Read a molecules CSV or parquet file.

    Parameters
    ----------
    file : PathLike
        File path.
    pos_cols : sequence of str, default ("z", "y", "x")
        Column names for the molecule positions.
    rot_cols : sequence of str, default ("zvec", "yvec", "xvec")
        Column names for the molecule rotation vectors.

    Returns
    -------
    Molecules
        Molecules object.
    """
    from acryo import Molecules

    path = Path(file)
    return Molecules.from_file(path, pos_cols=list(pos_cols), rot_cols=list(rot_cols))

def read_project(file: PathLike) -> CylindraProject:
    """Read the Cylindra project file."""
    from cylindra.project import CylindraProject

    return CylindraProject.from_file(file)

def read_spline(file: PathLike) -> CylSpline:
    """Read a spline file.

    Parameters
    ----------
    file : PathLike
        File path.

    Returns
    -------
    CylSpline
        CylSpline object.
    """
    from cylindra.components import CylSpline

    return CylSpline.from_json(file)

def start(
    project_file: str | None = None,
    viewer: napari.Viewer | None = None,
    *,
    log_level: int | str = "INFO",
    headless: bool = False,
    add_main_widget: bool = True,
    run: bool = True,
) -> CylindraMainWidget:
    """Start napari viewer and dock cylindra widget as a dock widget.

    Parameters
    ----------
    project_file : path-like, optional
        If given, load the project file.
    viewer : napari.Viewer
        Give a viewer object and this viewer will be used as the parent.
    log_level : int or str, default "INFO"
        Log level. See `logging` module for details.
    headless : bool, default False
        If True, do not show the viewer.
    """
    from cylindra.widgets import CylindraMainWidget  # noqa: I001
    import impy as ip
    import matplotlib.pyplot as plt
    import napari
    import numpy as np
    import polars as pl
    from IPython import get_ipython
    from magicclass import logging

    from cylindra._config import init_config

    global _CURRENT_INSTANCE

    if viewer is None:
        viewer = napari.Viewer(show=not headless)
    elif not isinstance(viewer, napari.Viewer):
        raise TypeError(f"viewer must be a napari.Viewer object, got {type(viewer)}")

    init_config()
    ui = CylindraMainWidget()
    ui.macro.options.max_undo = 16
    ACTIVE_WIDGETS.add(ui)

    # set logger
    logger = logging.getLogger("cylindra")
    formatter = logging.Formatter(fmt="%(levelname)s || %(message)s")
    logger.widget.setFormatter(formatter)
    logger.widget.min_height = 150

    # set log level
    if isinstance(log_level, str):
        log_level = log_level.upper()
        if log_level in ("DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"):
            log_level = getattr(logging, log_level)
        else:
            raise ValueError(f"Invalid log level: {log_level}")
    logger.setLevel(log_level)

    # set polars display options
    pl.Config().set_tbl_width_chars(120)

    if add_main_widget:
        dock = viewer.window.add_dock_widget(
            ui, area="right", allowed_areas=["right"], name="cylindra"
        )
        dock.setMinimumHeight(300)
    viewer.window.add_dock_widget(logger.widget, name="Log")
    ui.macro.options.syntax_highlight = False

    if project_file is not None:
        ui.load_project(project_file)
    _CURRENT_INSTANCE = ui

    with suppress(Exception):
        # update console namespace
        viewer.window._qt_viewer.console.push(
            {
                ".ui": ui,  # only available from namespace dict
                "ui": ui,
                "np": np,
                "ip": ip,
                "pl": pl,
                "plt": plt,
                "Path": Path,
            }
        )

    with suppress(Exception):  # This block uses private API.
        # napari viewer does not disconnect layer events when the viewer is closed,
        # so we need to do it manually
        @viewer.window._qt_window.destroyed.connect
        def _on_destroy():
            viewer.layers.events.removing.disconnect()
            viewer.layers.events.removed.disconnect()

        # napari-console disables calltips by default. It's better to enable it.
        viewer.window._qt_viewer.console.enable_calltips = True

    viewer.camera.orientation = ("away", "down", "right")

    @viewer.bind_key("J", overwrite=True)
    def _focus_up(v: napari.Viewer):
        v.dims.set_current_step(0, v.dims.current_step[0] - 4)

    @viewer.bind_key("K", overwrite=True)
    def _focus_down(v: napari.Viewer):
        v.dims.set_current_step(0, v.dims.current_step[0] + 4)

    @viewer.bind_key("D", overwrite=True)
    def _draw(v: napari.Viewer):
        v.layers.selection = {ui._reserved_layers.work}
        ui._reserved_layers.work.mode = "add"

    ui.show(run=run)
    if add_main_widget:
        try:  # Just in case
            # avoid accidentally closing/hiding the dock widget
            dock.title.close_button.clicked.disconnect()
            dock.title.hide_button.clicked.disconnect()
        except Exception:  # pragma: no cover
            print("Failed to disconnect the close/hide button of the dock widget.")

    # Programmatically run `%matplotlib inline` magic
    if ipy := get_ipython():
        ipy.run_line_magic("matplotlib", "inline")
        with suppress(Exception):  # This function uses private API of IPython.
            _allow_completion_on_cylindra_objects(ipy)

    return ui

def start_as_plugin(run: bool = True):
    """Start Cylindra as a napari plugin"""
    import napari
    from magicclass import logging

    ui = start(
        viewer=napari.current_viewer(),
        add_main_widget=False,
        run=run,
    )
    logger = logging.getLogger("cylindra")
    logger.widget.height = 160
    return ui

def view_project(project_file: PathLike, show: bool = True):
    """View the Cylindra project file."""
    from cylindra.project import CylindraProject

    widget = CylindraProject.from_file(project_file).make_project_viewer()
    if show:
        widget.show(run=False)
    ACTIVE_WIDGETS.add(widget)
    return widget