from typing import List, Union

from pydantic import BaseModel


class PatchParams(BaseModel):
    draw_bounding_box: bool = False
    draw_contour: bool = False
    draw_mask: bool = False
    rescale_clip: float = 0.001
    focus_metric: str = 'max_sobel'
    rgb_overlay_channels: List[Union[int, None]] = [None, None, None]
    rgb_overlay_weights: List[float] = [1.0, 1.0, 1.0]
    pad_to: int = 256


class AnnotatedZStackParams(BaseModel):
    draw_label: bool = False

class RoiFilterRange(BaseModel):
    min: float
    max: float

class RoiFilter(BaseModel):
    area: Union[RoiFilterRange, None] = None
    solidity: Union[RoiFilterRange, None] = None


class RoiSetMetaParams(BaseModel):
    mask_type: str = 'boxes'
    filters: RoiFilter = {}
    expand_box_by: List[int] = [128, 0]


class RoiSetExportParams(BaseModel):
    pixel_probabilities: bool = False
    patches_3d: Union[PatchParams, None] = None
    patches_2d_for_annotation: Union[PatchParams, None] = None
    patches_2d_for_training: Union[PatchParams, None] = None
    patch_masks: bool = False
    annotated_zstacks: Union[AnnotatedZStackParams, None] = None
    object_classes: bool = False
    dataframe: bool = False