# Usage Example
from enstools.feature.pipeline import FeaturePipeline
from enstools.feature.identification.african_easterly_waves import AEWIdentification
from enstools.feature.tracking.african_easterly_waves import AEWTracking
from datetime import timedelta
from enstools.feature.identification._proto_gen import african_easterly_waves_pb2
from os.path import expanduser, join
from enstools.feature.util.graph import DataGraph
from enstools.feature.identification.african_easterly_waves.plotting import plot_track_in_ts, plot_timesteps_from_desc, plot_tracks_from_graph

pipeline = FeaturePipeline(african_easterly_waves_pb2, processing_mode='2d')

data_dir = join(expanduser("~") + '/phd/data/aew/cv/') # 2008cv.nc
# data_dir = join(expanduser("~"), 'phd/data/aew/comp/helene/') # helene_25.nc
in_file = join(data_dir, '2008cv.nc') # 2008cv.nc helene_25.nc
out_traj_dir = join(data_dir, 'trajectories/')

# init AEWIdentification strategy, can take different parameters
i_strat = AEWIdentification(wt_traj_dir=None, cv='cv') # TODO
t_strat = AEWTracking(max_cmp_delta=timedelta(hours=18))

pipeline.set_identification_strategy(i_strat)
pipeline.set_tracking_strategy(t_strat) # None TODO

# data_path = join(expanduser("~") + '/phd/data/aew/cv/2008cv.nc')
# data_path = join(expanduser("~") + '/phd/data/aew/comp/helene/helene_25.nc')
pipeline.set_data_path(in_file)

# execute pipeline
pipeline.execute()

# generate the single tracks
pipeline.generate_tracks()

od = pipeline.get_object_desc() # just the identified objects and their connections in object_desc

# out_netcdf_path = data_path + '_streamers.nc'
out_json_path = in_file[:-3] + 'aew_desc.json'
out_graph_path = in_file[:-3] + 'aew_graph.json'
out_dataset_path = in_file[:-5] + '05_wt.nc'

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

print("Plot.")
#plot_timesteps_from_desc(od, pipeline.get_data())
plot_tracks_from_graph(od, None)
# plot_track_in_ts(od.sets[0].tracks[8])