###album catalog: mycatalog

from album.runner.api import get_args, setup

env_file = """
channels:
  - conda-forge
  - defaults
dependencies:
  - python=3.11
  - pip
  - zarr
  - ome-zarr
  - numpy<2
  - scipy
  - scikit-image
  - trimesh
  - pip:
    - album
    - "copick[all]>=0.5.2"
"""

args = [
    {
        "name": "copick_config_path",
        "description": "Path to the copick config file",
        "type": "string",
        "required": True,
    },
    {
        "name": "run_names",
        "description": "List of comma-separated run names to process",
        "type": "string",
        "required": False,
        "default": "",
    },
    {
        "name": "voxel_spacing",
        "description": "Voxel spacing for the tomograms",
        "type": "float",
        "required": False,
        "default": 10.0,
    },
    {
        "name": "tomo_type",
        "description": "Type of tomogram",
        "type": "string",
        "required": False,
        "default": "wbp",
    },
    {
        "name": "num_points",
        "description": "Number of random points to generate",
        "type": "integer",
        "required": False,
        "default": 10,
    },
    {
        "name": "out_object",
        "description": "Name of the output pickable object.",
        "type": "string",
        "required": False,
        "default": "random-points",
    },
    {
        "name": "out_user",
        "description": "User/Tool name for output points.",
        "type": "string",
        "required": False,
        "default": "solution-01",
    },
    {
        "name": "out_session",
        "description": "Output session, indicating this set was generated by a tool.",
        "type": "string",
        "required": False,
        "default": "0",
    },
]


def run():
    # Imports
    from typing import List, Sequence

    import copick
    import numpy as np
    import zarr
    from copick.models import CopickLocation, CopickPoint

    # Parse arguments
    args = get_args()
    copick_config_path = args.copick_config_path
    run_names = args.run_names.split(",")
    voxel_spacing = args.voxel_spacing
    tomo_type = args.tomo_type
    num_points = args.num_points
    out_object = args.out_object
    out_user = args.out_user
    out_session = args.out_session

    # Function definitions
    def generate_random_points(npoints: int, mdim: Sequence[float]) -> List[CopickPoint]:
        """Generate a set of random points."""
        points = []
        for _i in range(npoints):
            point = CopickPoint(
                location=CopickLocation(
                    x=np.random.rand(1) * mdim[0],
                    y=np.random.rand(1) * mdim[1],
                    z=np.random.rand(1) * mdim[2],
                ),
            )
            points.append(point)
        return points

    # Load copick project root
    root = copick.from_file(copick_config_path)

    # If no run names are provided, process all runs
    if run_names == [""]:
        run_names = [r.name for r in root.runs]

    # Process runs
    for run_name in run_names:
        print(f"Processing run {run_name}")
        run = root.get_run(run_name)

        # Get the physical tomogram dimensions
        vs = run.get_voxel_spacing(voxel_spacing)
        tomo = vs.get_tomogram(tomo_type)
        pixel_max_dim = zarr.open(tomo.zarr())["0"].shape[::-1]
        max_dim = np.array([d * voxel_spacing for d in pixel_max_dim])

        # If picks of the same type already exist, we will get and overwrite them
        picks = run.get_picks(object_name=out_object, user_id=out_user, session_id=out_session)

        # If picks do not exist, we will generate new picks and add them to the run
        if len(picks) == 0:
            picks = run.new_picks(object_name=out_object, user_id=out_user, session_id=out_session)
        else:
            picks = picks[0]

        points = generate_random_points(num_points, max_dim)
        picks.points = points
        picks.store()

    print("Processing complete.")


setup(
    group="copick",
    name="random_points",
    version="0.1.0",
    title="Random Points",
    description="This solution generates a set of random points for each run in a copick project.",
    solution_creators=["Alice", "Bob"],
    tags=["copick", "points "],
    license="MIT",
    album_api_version="0.5.1",
    args=args,
    run=run,
    dependencies={"environment_file": env_file},
)
