import numpy as np
import pandas as pd

from skimage.measure import find_contours, regionprops_table

# build a single boolean 3d mask (objects v. bboxes) and return bounding boxes
def build_stack_mask(desc, obmap, stack, filters=None, mask_type='contour', expand_box_by=(0, 0)): # TODO: specify boxes data type
    """

    filters: dict of (min, max) tuples
    expand_box_by: (xy, z) pixelsf
    """

    # validate inputs
    assert len(stack.shape) == 3, stack.shape
    assert mask_type in ('contour', 'box'), mask_type # TODO: replace with call to validator

    for k in filters.keys():
        assert k in ('area', 'solidity')
        vmin, vmax = filters[k]
        assert vmin >= 0

    # build object query
    query_str = 'label > 0'  # always true
    if filters:
        for k in filters.keys():
            assert k in ('area', 'solidity')
            vmin, vmax = filters[k]
            assert vmin >= 0
            query_str = query_str + f' & {k} > {vmin} & {k} < {vmax}'

    # build dataframe of objects, assign z index to each object
    argmax = stack.argmax(axis=0)
    df = (
        pd.DataFrame(
            regionprops_table(
                obmap,
                intensity_image=argmax,
                properties=('label', 'area', 'intensity_mean', 'solidity', 'bbox')
            )
        )
        .query(query_str)
        .rename(
            columns={
                'bbox-0': 'y0',
                'bbox-1': 'x0',
                'bbox-2': 'y1',
                'bbox-3': 'x1',
            }
        )
    )
    df['zi'] = df['intensity_mean'].round().astype('int')

    # make an object map where label is replaced by focus position in stack and background is -1
    lut = np.zeros(obmap.max() + 1) - 1
    lut[df.label] = df.zi

    # convert bounding boxes to slices
    ebxy, ebz = expand_box_by
    nz, h, w = stack.shape

    boxes = []
    for ob in df.itertuples(name='LabeledObject'):
        y0 = max(ob.y0 - ebxy, 0)
        y1 = min(ob.y1 + ebxy, h - 1)
        x0 = max(ob.x0 - ebxy, 0)
        x1 = min(ob.x1 + ebxy, w - 1)
        z0 = max(ob.zi - ebz, 0)
        z1 = min(ob.zi + ebz, nz)

        # relative bounding box positions
        rbb = {
            'y0': ob.y0 - y0,
            'y1': ob.y1 - y0,
            'x0': ob.x0 - x0,
            'x1': ob.x1 - x0,
        }

        sl = np.s_[z0: z1 + 1, y0: y1, x0: x1]

        # compute contours
        obmask = (obmap == ob.label)
        contour = find_contours(obmask)
        mask = obmask[ob.y0: ob.y1, ob.x0: ob.x1]

        boxes.append({
            'info': ob,
            'slice': sl,
            'relative_bounding_box': rbb,
            'contour': contour,
            'mask': mask
        })

    # build mask z-stack
    zi_st = np.zeros(stack.shape, dtype='bool')
    if mask_type == 'contour':
        zi_map = (lut[obmap] + 1.0).astype('int')
        idxs = np.array([zi_map]) - 1
        np.put_along_axis(zi_st, idxs, 1, axis=0)

        # change background level from to 0 in final frame
        zi_st[-1, :, :][obmap == 0] = 0

    elif mask_type == 'box':
        for bb in boxes:
            sl = bb['slice']
            zi_st[sl] = 1

    return zi_st, boxes