# AUTOGENERATED! DO NOT EDIT! File to edit: notebooks/batch_process_segmentation.ipynb (unless otherwise specified).

__all__ = ['func', 'process_experiment_command', 'process_experiment']

# Cell
# export
import pandas
import dotenv
import os
import matplotlib.pyplot as plt
import numpy

dotenv.load_dotenv()

# Cell

from c7m import process_segmentation, load_tiff_images, focus_stacking
import networkx
from tqdm import tqdm, trange
import click
import multiprocessing
import logging
import cv2
from skimage import img_as_float32


# Cell

def func(i, series_circles, img):
    img = focus_stacking.focus_stack(img)
    img = cv2.bilateralFilter(img, 5, 10., 10.)

    g = process_segmentation.make_overlap_graph(series_circles)
    g_filtered = networkx.subgraph_view(g, filter_node=process_segmentation.lt_mean_filter(g))
    g_filtered = process_segmentation.compute_patch_coordinates(g_filtered)
    g_filtered = networkx.subgraph_view(g_filtered, filter_node=process_segmentation.erronous_segmentations_filter(g_filtered))
    g_filtered = process_segmentation.find_hough_circle_per_node(g_filtered, img, threshold=0.7)
    g_filtered = networkx.subgraph_view(g_filtered, filter_node=process_segmentation.hough_filter(g_filtered))

    data = numpy.empty((len(g_filtered),7), dtype=numpy.float64)
    for j,n in enumerate(g_filtered.nodes()):
        data[j, :4] = list(g_filtered.nodes[n]["patch_coordinates"])
        data[j, 4:] = list(g_filtered.nodes[n]["circle"])

    df = pandas.DataFrame.from_records(data, columns=["0", "1", "2", "3", "cx", "cy", "radius"])
    df["series"] = i

    return df


@click.command()
@click.option("--czi-path", type=click.Path(exists=True, dir_okay=False))
@click.option("--segmentation-csv", type=click.Path(exists=True, dir_okay=False))
@click.option("--output", type=click.Path(dir_okay=False), default=None, help="Path to pickle file to save segmentations.")
@click.option("--processes", type=int, default=6)
def process_experiment_command(*args, **kwargs):
    process_experiment(*args, **kwargs)

def process_experiment(path, segmentation_csv, output=None, processes=1, series_start=None, series_end=None):
    circles, series = process_segmentation.load_hough_segmentation(segmentation_csv)

    if series_end is None and series_start is None:
        r = series
    else:
        r = numpy.arange(series_start, series_end)

    with multiprocessing.Pool(processes) as pool:
        # loop over all series and submit to pool
        promises = []
        for i in r:
            if i not in series:
                print(f"series {i} is not present in the segmentation file")
                continue
            series_circles = circles[circles["series"] == i].reset_index(drop=True)
            img = img_as_float32(load_tiff_images.load_slice(c=1, series=i, z="all", path=path))
            promises.append(pool.apply_async(func, (i, series_circles, img)))

        # get results from the queue
        df_list = []
        for p in tqdm(promises):
            df_list.append(p.get())

    df = pandas.concat(df_list).reset_index(drop=True)
    if output is not None:
        df.to_pickle(output)

    return df