import napari
import numpy as np

from enum import Enum
from functools import partial
from magicgui import magic_factory
from skimage.filters import (
    threshold_isodata,
    threshold_li,
    threshold_otsu,
    threshold_triangle,
    threshold_yen,
)
from skimage.measure import label
from scipy.ndimage import center_of_mass
DATA_PATH = '~/PhD/data/toy_data/easy-no-divide-swaps/Fluo-N2DL-HeLa.tif'

# very simple thresholding based segmentation widget
class Threshold(Enum):
    isodata = partial(threshold_isodata)
    li = partial(threshold_li)
    otsu = partial(threshold_otsu)
    triangle = partial(threshold_triangle)
    yen = partial(threshold_yen)

@magic_factory
def segment_by_threshold(
    img_layer: "napari.layers.Image", threshold: Threshold
) -> "napari.types.LayerDataTuple":

    seg_labels = np.zeros_like(img_layer.data)
    for i in range(len(img_layer.data)):
        frame = img_layer.data[i]
        # need to use threshold.value to get the function from the enum member
        threshold_val = threshold.value(frame)
        binarised_im = frame > threshold_val
        seg_labels[i] = label(binarised_im)

    seg_layer = (seg_labels, {"name": f"{img_layer.name}_seg"}, "labels")

    return seg_layer

viewer = napari.Viewer()
cells_image = viewer.open(DATA_PATH)[0]

widget = segment_by_threshold()
viewer.window.add_dock_widget(widget)

---------------------------------------------------------------------------
FileNotFoundError                         Traceback (most recent call last)
/tmp/ipykernel_4337/68408623.py in ?()
      1 viewer = napari.Viewer()
----> 2 cells_image = viewer.open(DATA_PATH)[0]
      3 
      4 widget = segment_by_threshold()
      5 viewer.window.add_dock_widget(widget)

/opt/hostedtoolcache/Python/3.11.11/x64/lib/python3.11/site-packages/napari/components/viewer_model.py in ?(self, path, stack, plugin, layer_type, **kwargs)
   1224                 # If _path is not a list already, make it a list.
   1225                 _path = [_path] if not isinstance(_path, list) else _path
   1226                 if plugin:
   1227                     added.extend(
-> 1228                         self._add_layers_with_plugins(
   1229                             _path,
   1230                             kwargs=kwargs,
   1231                             plugin=plugin,

/opt/hostedtoolcache/Python/3.11.11/x64/lib/python3.11/site-packages/napari/components/viewer_model.py in ?(self, paths, stack, kwargs, plugin, layer_type)
   1424                 paths, plugin=plugin, stack=stack
   1425             )
   1426         else:
   1427             assert len(paths) == 1
-> 1428             layer_data, hookimpl = read_data_with_plugins(
   1429                 paths, plugin=plugin, stack=stack
   1430             )
   1431 

/opt/hostedtoolcache/Python/3.11.11/x64/lib/python3.11/site-packages/napari/plugins/io.py in ?(paths, plugin, stack)
     74     if not stack:
     75         assert len(paths) == 1
     76     hookimpl: Optional[HookImplementation]
     77 
---> 78     res = _npe2.read(paths, plugin, stack=stack)
     79     if res is not None:
     80         ld_, hookimpl = res
     81         return [] if _is_null_layer_sentinel(ld_) else list(ld_), hookimpl

/opt/hostedtoolcache/Python/3.11.11/x64/lib/python3.11/site-packages/napari/plugins/_npe2.py in ?(paths, plugin, stack)
     56     try:
     57         layer_data, reader = io_utils.read_get_reader(
     58             npe1_path, plugin_name=plugin
     59         )
---> 60     except ValueError as e:
     61         # plugin wasn't passed and no reader was found
     62         if 'No readers returned data' not in str(e):
     63             raise

/opt/hostedtoolcache/Python/3.11.11/x64/lib/python3.11/site-packages/npe2/io_utils.py in ?(path, plugin_name, stack)
     62     if stack is None:
     63         # "npe1" old path
     64         # Napari 0.4.15 and older, hopefully we can drop this and make stack mandatory
     65         new_path, new_stack = v1_to_v2(path)
---> 66         return _read(
     67             new_path, plugin_name=plugin_name, return_reader=True, stack=new_stack
     68         )
     69     else:

/opt/hostedtoolcache/Python/3.11.11/x64/lib/python3.11/site-packages/npe2/io_utils.py in ?(paths, stack, plugin_name, return_reader, _pm)
    166             kwargs={"path": paths, "stack": stack, "_registry": _pm.commands}
    167         )
    168         if read_func is not None:
    169             # if the reader function raises an exception here, we don't try to catch it
--> 170             if layer_data := read_func(paths, stack=stack):
    171                 return (layer_data, rdr) if return_reader else layer_data
    172 
    173     if plugin_name:

/opt/hostedtoolcache/Python/3.11.11/x64/lib/python3.11/site-packages/npe2/manifest/contributions/_readers.py in ?(paths, stack)
     66         @wraps(callable_)
     67         def npe1_compat(paths, *, stack):
     68             path = v2_to_v1(paths, stack)
---> 69             return callable_(path)

/opt/hostedtoolcache/Python/3.11.11/x64/lib/python3.11/site-packages/napari_builtins/io/_read.py in ?(path)
    502 def _magic_imreader(path: str) -> list['LayerData']:
--> 503     return [(magic_imread(path),)]

/opt/hostedtoolcache/Python/3.11.11/x64/lib/python3.11/site-packages/napari_builtins/io/_read.py in ?(filenames, use_dask, stack)
    235             if shape is None:
    236                 shape = zarr_shape
    237         else:
    238             if shape is None:
--> 239                 image = imread(filename)
    240                 shape = image.shape
    241                 dtype = image.dtype
    242             if use_dask:

/opt/hostedtoolcache/Python/3.11.11/x64/lib/python3.11/site-packages/napari_builtins/io/_read.py in ?(filename)
     94     import tifffile
     95 
     96     # Pre-download urls before loading them with tifffile
     97     with file_or_url_context(filename) as filename:
---> 98         return tifffile.imread(str(filename))

/opt/hostedtoolcache/Python/3.11.11/x64/lib/python3.11/site-packages/tifffile/tifffile.py in ?(files, selection, aszarr, key, series, level, squeeze, maxworkers, buffersize, mode, name, offset, size, pattern, axesorder, categories, imread, imreadargs, sort, container, chunkshape, chunkdtype, axestiled, ioworkers, chunkmode, fillvalue, zattrs, multiscales, omexml, out, out_inplace, _multifile, _useframes, **kwargs)
   1217         ):
   1218             files = files[0]
   1219 
   1220         if isinstance(files, str) or not isinstance(files, Sequence):
-> 1221             with TiffFile(
   1222                 files,
   1223                 mode=mode,
   1224                 name=name,

/opt/hostedtoolcache/Python/3.11.11/x64/lib/python3.11/site-packages/tifffile/tifffile.py in ?(self, file, mode, name, offset, size, omexml, _multifile, _useframes, _parent, **is_flags)
   4245                 raise ValueError('invalid OME-XML')
   4246             self._omexml = omexml
   4247             self.is_ome = True
   4248 
-> 4249         fh = FileHandle(file, mode=mode, name=name, offset=offset, size=size)
   4250         self._fh = fh
   4251         self._multifile = True if _multifile is None else bool(_multifile)
   4252         self._files = {fh.name: self}

/opt/hostedtoolcache/Python/3.11.11/x64/lib/python3.11/site-packages/tifffile/tifffile.py in ?(self, file, mode, name, offset, size)
  14622         self._offset = -1 if offset is None else offset
  14623         self._size = -1 if size is None else size
  14624         self._close = True
  14625         self._lock = NullContext()
> 14626         self.open()
  14627         assert self._fh is not None

/opt/hostedtoolcache/Python/3.11.11/x64/lib/python3.11/site-packages/tifffile/tifffile.py in ?(self)
  14641             if self._mode not in {'rb', 'r+b', 'wb', 'xb'}:
  14642                 raise ValueError(f'invalid mode {self._mode}')
  14643             self._file = os.path.realpath(self._file)
  14644             self._dir, self._name = os.path.split(self._file)
> 14645             self._fh = open(self._file, self._mode, encoding=None)
  14646             self._close = True
  14647             self._offset = max(0, self._offset)
  14648         elif isinstance(self._file, FileHandle):

FileNotFoundError: [Errno 2] No such file or directory: '/home/runner/PhD/data/toy_data/easy-no-divide-swaps/Fluo-N2DL-HeLa.tif'

cell_labels = viewer.layers[-1]

def move_point_to_label_center(event):
    points_layer = event.source
    layer_data = event.value
    # round latest point so we can find the label value underneath
    latest_point = np.round(layer_data[-1]).astype(cell_labels.data.dtype)
    label_val_at_point = cell_labels.data[tuple(latest_point)]
    # pass through just the frame we added the point to, so we can get center of mass
    data_slice = latest_point[0]
    label_center = center_of_mass(cells_image.data[data_slice], cell_labels.data[data_slice], label_val_at_point)
    new_point = [data_slice, *label_center]
    # block data event when we're changing the point to avoid an infinite loop
    with points_layer.events.data.blocker():
        points_layer.data[-1] = new_point
    points_layer.refresh()

points_layer = viewer.add_points(ndim=3)
points_layer.events.data.connect(move_point_to_label_center)

# add keybinding on labels layer to add point with same color as the label underneath
@cell_labels.mouse_drag_callbacks.append
def add_point_with_label_color(layer, event):
    if (len(event.modifiers) == 1
            and event.modifiers[0].name == 'Alt'):
        # rounding mouse position to integer coordinates
        data_pos = np.round(event.position).astype(layer.data.dtype)
        label_color = layer.get_color(layer.data[tuple(data_pos)])   
        # add points layer if it doesn't already exist
        if 'cell_points' not in viewer.layers:
            points_layer = viewer.add_points(name='cell_points', ndim=3)
            points_layer.events.data.connect(move_point_to_label_center)
            viewer.layers.selection.active = cell_labels
        # add a point to the (potentially newly created points layer)
        points_layer = viewer.layers['cell_points']
        points_layer.current_face_color = label_color
        points_layer.add(event.position)
        points_layer.selected_data.clear()