"""
# 🦠 Cell Instance Segmentation: 🔍 interative data view

For more or future development see https://borda.github.io/kaggle_cell-inst-segm
"""
! ls -l /kaggle/input/sartorius-cell-instance-segmentation
# ! ls -l /kaggle/input/sartorius-cell-instance-segmentation/LIVECell_dataset_2021
"""
Browsing the provided data/images/annotations...
"""
import os
import pandas as pd

PATH_DATASET = "/kaggle/input/sartorius-cell-instance-segmentation"
PATH_IMAGES = os.path.join(PATH_DATASET, "train")
PATH_TRAIN_CSV = os.path.join(PATH_DATASET, "train.csv")

df_train = pd.read_csv(PATH_TRAIN_CSV)
display(df_train.head())
df_train.loc[0, "annotation"]
"""
## Histogram of annotation per image
"""
df_counts = df_train.groupby(['id']).size()
display(df_counts.head())
ax = df_counts.hist(bins=50, grid=True)
ax.set_xlabel("Annotations per image")
ax.set_ylabel("Images with particular annot. count")
"""
## Decode single annotation
"""
import numpy as np
import matplotlib.pyplot as plt

def rle_decode(mask_rle: str, img: np.ndarray = None, img_shape: tuple = None, label: int = 1) -> np.ndarray:
    seq = mask_rle.split()
    starts = np.array(list(map(int, seq[0::2])))
    lengths = np.array(list(map(int, seq[1::2])))
    assert len(starts) == len(lengths)
    ends = starts + lengths
    
    if img is None:
        img = np.zeros((np.product(img_shape), ), dtype=np.uint16)
    else:
        img_shape = img.shape
        img = img.flatten()
    for begin, end in zip(starts, ends):
        img[begin:end] = label
    return img.reshape(img_shape)

mask = rle_decode(df_train.loc[0, "annotation"], img_shape=(df_train.loc[0, "height"], df_train.loc[0, "width"]))
mask = rle_decode(df_train.loc[1, "annotation"], img=mask, label=2)
_= plt.imshow(mask)
"""
## Create complete mask
"""
def create_mask(df_image: pd.DataFrame) -> np.ndarray:
    assert len(df_image["id"].unique()) == 1
    sizes = list(set((row["height"], row["width"]) for _, row in df_image.iterrows()))
    assert len(sizes) == 1
    mask = np.zeros(sizes[0], dtype=np.uint16)
    df_image.reset_index(inplace=True)
    for idx, row in df_image.iterrows():
        mask = rle_decode(row["annotation"], img=mask, label=idx + 1)
    return mask
    

# print(df_train["sample_id"].unique())
mask = create_mask(df_train[df_train["id"] == "0030fd0e6378"])
_= plt.imshow(mask, interpolation='antialiased')
"""
## Interactive view

Note that for interative browsing you need to be in edit mode (it is not supported for saved version)
"""
from ipywidgets import interact, SelectionSlider


def show_image_annot(img_name: str, df_train: pd.DataFrame, img_folder: str):
    print(img_name)
    df_img = df_train[df_train["id"] == img_name]
    path_img = os.path.join(img_folder, f"{img_name}.png")
    img = plt.imread(path_img)
    mask = create_mask(df_img)
    fig, axarr = plt.subplots(ncols=3, figsize=(18, 6))
    axarr[0].imshow(img)
    axarr[1].imshow(img)
    axarr[1].contour(mask, levels=np.unique(mask).tolist(), cmap="inferno", linewidths=0.5)
    axarr[2].imshow(mask, cmap="inferno", interpolation='antialiased')
    return fig


def interactive_show(df_train: pd.DataFrame, img_folder: str):
    uq_images = df_train["id"].unique()
    interact(
        lambda im: plt.show(show_image_annot(im, df_train, img_folder)),
        im=SelectionSlider(
            options=uq_images,
            value="cc40345857dd",  # uq_images[np.random.randint(0, len(uq_images))],
            description='Select image:',
            disabled=False,
            continuous_update=False,
            orientation='horizontal',
            readout=True
        ),
    )

interactive_show(df_train, PATH_IMAGES)