from cropmask.preprocess import PreprocessWorkflow, get_arr_channel_mean, setup_dirs
import time
import dask

param_path = "/home/ryan/work/CropMask_RCNN/cropmask/test_preprocess_config.yaml"
scene_list = [
    "/mnt/cropmaskperm/unpacked_landsat_downloads/LT050280322005012001T2-SC20190818204900", 
    "/mnt/cropmaskperm/unpacked_landsat_downloads/LT050310322005021001T1-SC20190818205059",
    "/mnt/cropmaskperm/unpacked_landsat_downloads/LT050290312005031601T1-SC20190818204935",  
    "/mnt/cropmaskperm/unpacked_landsat_downloads/LT050320312005011601T1-SC20190818205113",
]
labels_path = "/mnt/cropmaskperm/external/nebraska_pivots_projected.geojson"

setup_dirs(param_path)


# this is just to get the train dir path
wflow = PreprocessWorkflow(param_path, 
                             scene_list[0],
                             labels_path)

results = []
for scene_path in scene_list:

    wflow = dask.delayed(PreprocessWorkflow)(param_path, scene_path, labels_path)
    
    band_list = dask.delayed(wflow.yaml_to_band_index)()
        
    product_list = dask.delayed(wflow.get_product_paths)(band_list)
        
    dask.delayed(wflow.load_meta_and_bounds)(product_list)

    dask.delayed(wflow.load_single_scene)(product_list)
        
    dask.delayed(wflow.stack_and_save_bands)()
        
    dask.delayed(wflow.negative_buffer_and_small_filter)(-31, 100)
        
    dask.delayed(wflow.grid_images)()
        
    dask.delayed(wflow.remove_from_gridded)()
        
    dask.delayed(wflow.move_chips_to_folder)()
        
    result = dask.delayed(wflow.connected_components)()

    results.append(result)
    

# https://docs.dask.org/en/stable/delayed-best-practices.html
from dask.distributed import Client, progress

client = Client()  # use dask.distributed by default

x = dask.compute(*results)  # start computation in the background
progress(x)      # watch progress  

tornado.application - ERROR - Exception in callback functools.partial(<bound method IOLoop._discard_future_result of <tornado.platform.asyncio.AsyncIOMainLoop object at 0x7f9d23efb898>>, <Task finished coro=<MultiProgressBar.listen() done, defined at /data/anaconda/envs/cropmask/lib/python3.7/site-packages/distributed/diagnostics/progressbar.py:238> exception=AttributeError("'MultiProgressWidget' object has no attribute 'elapsed_time'")>)
Traceback (most recent call last):
  File "/data/anaconda/envs/cropmask/lib/python3.7/site-packages/tornado/ioloop.py", line 743, in _run_callback
    ret = callback()
  File "/data/anaconda/envs/cropmask/lib/python3.7/site-packages/tornado/ioloop.py", line 767, in _discard_future_result
    future.result()
  File "/data/anaconda/envs/cropmask/lib/python3.7/site-packages/distributed/diagnostics/progressbar.py", line 282, in listen
    self._draw_stop(**response)
  File "/data/anaconda/envs/cropmask/lib/python3.7/site-packages/distributed/diagnostics/progressbar.py", line 381, in _draw_stop
    + "</div>"
AttributeError: 'MultiProgressWidget' object has no attribute 'elapsed_time'