from pathlib import Path

import pandas as pd

from extensions.chaeo.workflows import export_patches_from_multichannel_zstack
from extensions.ilastik.models import IlastikPixelClassifierModel
from model_server.accessors import InMemoryDataAccessor, write_accessor_data_to_file

if __name__ == '__main__':
    where_czi = Path(
        'z:/rhodes/projects/proj0004-marine-photoactivation/data/exp0038/AutoMic/20230906-163415/Selection'
    )
    where_output = Path(
        'c:/Users/rhodes/projects/proj0011-plankton-seg/exp0009/batch_output'
    )
    csv_args = {'mode': 'w', 'header': True} # when creating file
    px_ilp = Path.home() / 'model-server' / 'ilastik' / 'AF405-bodies_boundaries.ilp'
    px_model = IlastikPixelClassifierModel(
        params={'project_file': px_ilp}
    )

    for ff in where_czi.iterdir():
        print(ff)
        if not ff.suffix.upper() == '.CZI':
            continue

        result = export_patches_from_multichannel_zstack(
            input_zstack_path=where_czi/ff,
            px_model=px_model,
            pxmap_threshold=0.25,
            pixel_class=0,
            zmask_channel=0,
            patches_channel=4,
            where_output=where_output,
            mask_type='boxes',
            zmask_filters={'area': (1e3, 1e5)},
            zmask_expand_box_by=(128, 3),
        )

        # parse and record results
        df = result['dataframe']
        df['filename'] = ff.name
        df.to_csv(where_output / 'df_objects.csv', **csv_args)
        pd.DataFrame(result['timer_results'], index=[0]).to_csv(where_output / 'timer_results.csv', **csv_args)
        csv_args = {'mode': 'a', 'header': False} # append to CSV from here on

        # export intermediate data if flagged
        for k in result['interm'].keys():
            write_accessor_data_to_file(
                where_output / k / (ff.stem + '.tif'),
                InMemoryDataAccessor(result['interm'][k])
                )