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

__all__ = ['load_hough_segmentation', 'make_overlap_graph', 'plot_circles_on_img', 'compute_patch_coordinates',
           'extract_patch_from_image', 'extract_patches_from_image', 'find_hough_circle_per_node', 'hough_filter',
           'erronous_segmentations_filter']

# Cell
# export
import pandas
import dotenv
import os
import matplotlib.pyplot as plt
from importlib import reload
import numpy
import warnings

warnings.filterwarnings('error')

dotenv.load_dotenv()

# Cell
from matplotlib import cm
import math
import itertools
import scipy.spatial
import networkx
from skimage.feature import canny
from skimage.transform import hough_circle, hough_circle_peaks
from skimage.draw import circle_perimeter
from skimage.filters import sobel
from skimage import img_as_float32
import re
import cv2

from c7m import load_tiff_images, denoise, focus_stacking

# Cell
def load_hough_segmentation(csv_file, micron_dims=(189.45, 189.45)):
    circles = pandas.read_csv(csv_file, index_col=False)
    circles = circles.drop(columns=["ID"])
    circles["X frac"] = circles["X (microns)"]/micron_dims[0]
    circles["Y frac"] = circles["Y (microns)"]/micron_dims[1]
    circles["Radius frac"] = circles["Radius (microns)"]/math.sqrt(micron_dims[0]*micron_dims[1])

    # extract series id from image title
    exp = r"^[^#]+#([0-9]+)-.+$"
    circles["series"] = circles["Image title"].apply(lambda a: int(re.match(exp, a).group(1))-1)
    return circles, circles["series"].unique()

# Internal Cell
def compute_overlap(d, r1, r2):
    # https://scipython.com/book/chapter-8-scipy/problems/p84/overlapping-circles/

    area = 0

    # only compute overlap if sum of radii is larger then distance between circle centers
    if r1 + r2 > d + 1e-5:
        r, R = (r1, r2) if r2 >= r1 else (r2, r1)

        try:
            alpha = numpy.arccos((r**2 + d**2 - R**2)/(2*r*d))
        except RuntimeWarning:
            print(r, R, d)
        beta = numpy.arccos((R**2 + d**2 - r**2)/(2*R*d))
        area = beta*R**2 + alpha*r**2 - (1/2)*r**2*numpy.sin(2*alpha) - (1/2)*R**2*numpy.sin(2*beta)

    return area

# Cell
def make_overlap_graph(df):
    g = networkx.Graph()

    # compute distance between all circle centers
    d = scipy.spatial.distance.squareform(scipy.spatial.distance.pdist(df[["X frac", "Y frac"]]))

    # select only those circle pairs whose center to center distance is larger then twice the smallest radius
    radius = df["Radius frac"].max()
    combinations = [c for c in itertools.combinations(df.index, 2) if d[c[0], c[1]] < radius*2]

    for idx1, idx2 in combinations:
        overlap = compute_overlap(d[idx1, idx2], df.iloc[idx1]["Radius frac"], df.iloc[idx2]["Radius frac"])

        g.add_node(idx1, **df.iloc[idx1].to_dict())
        g.add_node(idx2, **df.iloc[idx2].to_dict())

        if overlap > 0:
            g.add_edge(idx1, idx2, weight=overlap)

    return g

# Cell
def plot_circles_on_img(graph, img, highlight=[]):
    fig, ax = plt.subplots(figsize=(10, 10), dpi=100)

    ax.imshow(img, origin="lower", extent=(0, 1, 0, 1))
    ax.set_xlim(0, 1)
    ax.set_ylim(0, 1)

    for idx in graph.nodes:
        color = "orange"
        if idx in highlight:
            color="red"

        circle = graph.nodes[idx]

        c = plt.Circle((circle["X frac"], circle["Y frac"]), circle["Radius frac"], color=color, alpha=0.2)
        ax.add_artist(c)
        ax.annotate(str(idx), (circle["X frac"], circle["Y frac"]), ha="center", va="center")

    return ax

# Internal Cell

def lt_mean_filter(g):
    def func(n):
        if len(list(g.neighbors(n))) == 0:
            return True

        # keep this node if the mean score of its neighbors is lower then its score
        agg_neighbor_score = numpy.mean([g.nodes[k]["Score"] for k in g.neighbors(n)])
        if agg_neighbor_score < g.nodes[n]["Score"]:
            return True

        # check if total overlap with neighbors is less then 50% of its area
        tot_overlap = sum([g.get_edge_data(n, k)["weight"] for k in g.neighbors(n)])
        if tot_overlap < 0.5*numpy.pi*g.nodes[n]["Radius frac"]**2:
            return True

        return False
    return func

# Internal Cell

def create_circular_mask(h, w, center=None, radius=None):
    # https://stackoverflow.com/a/44874588/2977466

    if center is None: # use the middle of the image
        center = (int(w/2), int(h/2))
    if radius is None: # use the smallest distance between the center and image walls
        radius = min(center[0], center[1], w-center[0], h-center[1])

    Y, X = numpy.ogrid[:h, :w]
    dist_from_center = numpy.sqrt((X - center[0])**2 + (Y-center[1])**2)

    mask = dist_from_center <= radius
    return mask

# Cell

def compute_patch_coordinates(graph):
    for index in graph.nodes:
        circle = graph.nodes[index]

        x = circle["X frac"]
        y = circle["Y frac"]
        radius = circle["Radius frac"]

        circle["patch_coordinates"] = (max(0, y-radius), y+radius, max(0, x-radius), x+radius)
    return graph

def extract_patch_from_image(coords, img, masked=None, fill=0):
    size = img.shape[0]
    patch = img[
        int(coords[0]*size):int(coords[1]*size),
        int(coords[2]*size):int(coords[3]*size)
    ]
    if masked is not None:
        mask = create_circular_mask(patch.shape[0], patch.shape[1], (masked[0], masked[1]), masked[2])
        return numpy.where(mask, patch, fill)
    return patch

def extract_patches_from_image(graph, img, masked=False):
    patches = []

    for index in graph.nodes:
        coords = graph.nodes[index]["patch_coordinates"]

        m = None
        if masked:
            m = graph.nodes[index]["circle"]

        patch = extract_patch_from_image(coords, img, m)
        patches.append(patch)

    return patches

# Cell
def find_hough_circle_per_node(graph, img, threshold=0.7):
    for index in graph.nodes:
        coords = graph.nodes[index]["patch_coordinates"]
        size = img.shape[0]
        patch = img[
            int(coords[0]*size):int(coords[1]*size),
            int(coords[2]*size):int(coords[3]*size)
        ]

        edges = sobel(patch)
        bin_edges = numpy.where(edges > numpy.mean(edges), 1, 0)
        try_radii = numpy.arange(min(patch.shape)*0.3, min(patch.shape)*0.5)
        result = hough_circle(bin_edges, try_radii)
        accums, cx, cy, radii = hough_circle_peaks(result, try_radii, total_num_peaks=1, threshold=threshold)

        if len(accums) > 0:
            max_idx = numpy.argmax(accums)
            graph.nodes[index]["circle"] = cx[max_idx], cy[max_idx], radii[max_idx]
        else:
            graph.nodes[index]["circle"] = None

    return graph

# Cell
def hough_filter(graph):
    def func(index):
        if graph.nodes[index]["circle"] is not None:
            return True

    return func

# Cell
def erronous_segmentations_filter(graph):
    def func(index):
        coords = graph.nodes[index]["patch_coordinates"]

        # check if width and height aren't equal
        if numpy.abs(coords[0]-coords[1]) < 1e-3:
            return False
        if numpy.abs(coords[2]-coords[3]) < 1e-3:
            return False

        return True
    return func