from enstools.feature.pipeline import FeaturePipeline
from enstools.feature.identification.threshold import DoubleThresholdIdentification,SingleThresholdIdentification
from enstools.feature.tracking.overlap_double_threshold_tracking import OverlapDoubleThresholdTracking
from enstools.feature.identification._proto_gen import threshold_pb2
from os.path import expanduser

# set the pb_reference to the compiled pb2.py file (see proto_gen directory)
pipeline = FeaturePipeline(threshold_pb2, processing_mode='3d')

# change this to an identification technique that actually does something: existing one or implement your own

path = expanduser("~") + '/phd/data/enstools-feature/pv_error_soeren.nc' #  ERA5/vietnam/t0_glob.nc' #  set data path(s) here
pipeline.set_data_path(path)

i_strat = DoubleThresholdIdentification("E",5,10,">",processing_mode='3d') # set the Identification strategy
t_strat = OverlapDoubleThresholdTracking("inner_thresh_mask","outer_thresh_mask","outer_first") # set the tracking strategy

pipeline.set_identification_strategy(i_strat)
pipeline.set_tracking_strategy(t_strat) # or None as argument if no tracking

# execute pipeline
pipeline.execute()

# generate the single tracks
pipeline.generate_tracks()

out_json_path = path[:-3] + '_desc.json'
out_graph_path = path[:-3] + '_graph.json'
print(out_json_path)
od = pipeline.get_object_desc() # identified objects, and graph-structure if tracking done

pipeline.save_result(description_type='json', description_path=out_json_path) # dataset_path=...