from pathlib import Path
from typing import Dict

from extensions.ilastik.models import IlastikPixelClassifierModel
from extensions.chaeo.annotators import draw_boxes_on_3d_image
from extensions.chaeo.products import export_patches_from_zstack
from extensions.chaeo.zmask import build_zmask_from_object_mask
from model_server.accessors import generate_file_accessor, InMemoryDataAccessor, write_accessor_data_to_file
from model_server.workflows import Timer

def export_patches_from_multichannel_zstack(
        input_zstack_path: str,
        ilastik_project_file: str,
        pxmap_threshold: float,
        pixel_class: int,
        zmask_channel: int,
        patches_channel: int,
        where_output: str,
        mask_type: str = 'boxes',
        zmask_filters: Dict = None,
        zmask_expand_box_by: int = None,
) -> Dict:

    ti = Timer()
    stack = generate_file_accessor(Path(input_zstack_path))
    fstem = Path(input_zstack_path).stem
    ti.click('file_input')
    assert stack.nz > 1, 'Expecting z-stack'

    # MIP and classify pixels
    mip = InMemoryDataAccessor(
        stack.get_one_channel_data(channel=0).data.max(axis=-1, keepdims=True)
    )
    px_model = IlastikPixelClassifierModel(
        params={'project_file': Path(ilastik_project_file)}
    )
    pxmap, _ = px_model.infer(mip)
    ti.click('infer_pixel_probability')

    write_accessor_data_to_file(
        Path(where_output) / 'pixel_probabilities' / (fstem + '.tif'),
        pxmap
    )
    ti.click('export_pixel_probability')

    obmask = InMemoryDataAccessor(
        pxmap.data > pxmap_threshold
    )
    ti.click('threshold_pixel_mask')

    # make zmask
    zmask, zmask_meta, df, interm = build_zmask_from_object_mask(
        obmask.get_one_channel_data(pixel_class),
        stack.get_one_channel_data(zmask_channel),
        mask_type=mask_type,
        filters=zmask_filters,
        expand_box_by=zmask_expand_box_by,
    )
    zmask_acc = InMemoryDataAccessor(zmask)
    ti.click('generate_zmasks')

    # export patches
    files = export_patches_from_zstack(
        Path(where_output) / '2d_patches',
        stack.get_one_channel_data(patches_channel),
        zmask_meta,
        prefix=fstem,
        draw_bounding_box=True,
        rescale_clip=0.0,
    )
    ti.click('export_patches')

    # export annotated zstack
    annotated = InMemoryDataAccessor(
        draw_boxes_on_3d_image(
            stack.get_one_channel_data(patches_channel).data,
            zmask_meta
        )
    )
    write_accessor_data_to_file(
        Path(where_output) / 'annotated_zstacks' / (fstem + '.tif'),
        annotated
    )
    ti.click('export_annotated_zstack')

    # from extensions.chaeo.zmask import build_image_flattening_zmask_from_points
    #
    # dff = df[df['keeper'] == True]
    # build_image_flattening_zmask_from_points(
    #     dff['centroid-0'],
    #     dff['centroid-1'],
    #     dff['zi'],
    #     stack.get_one_channel_data(patches_channel).data,
    # )

    return {
        'pixel_model_id': px_model.model_id,
        'input_filepath': input_zstack_path,
        'number_of_objects': len(zmask_meta),
        'success': True,
        'timer_results': ti.events,
        'dataframe': df,
        'interm': interm,
    }