# 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_differences, plot_track, plot_track_in_ts, plot_timesteps_from_desc, plot_tracks_from_desc
import enstools.feature.identification.african_easterly_waves.configuration as cfg
import sys
from enstools.feature.util.data_utils import get_subset_by_description

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

# in_files_all_cv_data = cfg.cv_data_ex


# for climatology
# if len(sys.argv) > 1:
#     proc_summer_of_year = int(sys.argv[1])
#     if len(sys.argv) > 2:
#         proc_month_of_year = int(sys.argv[2])

if len(sys.argv) >= 3 and sys.argv[1] == '-kw':
    # TODO kitweather
    #  only -kw and infer directory from config
    #  also based on latest timestep?
    in_file = sys.argv[2:]
else:
    in_file = cfg.in_files
    out_dir = cfg.out_dir

# init AEWIdentification strategy, can take different parameters
i_strat = AEWIdentification(wt_out_file=False, cv='cv') # , year_summer=proc_summer_of_year, month=proc_month_of_year)
t_strat = AEWTracking()

pipeline.set_identification_strategy(i_strat)
pipeline.set_tracking_strategy(t_strat)

pipeline.set_data_path(in_file)
# execute pipeline
pipeline.execute()

od = pipeline.get_object_desc()

for trackable_set in od.sets:
    # generate graph out of tracked data
    g = DataGraph(trackable_set, t_strat)

    # generate single tracks from tracked data
    # returns list of tracks, also gets added to object description. Also if apply_filter, keep_track can be implemented
    g.generate_tracks(apply_filter=True) # add tracks to OD, applies filtering TODO tracks not in desc.
    tracks = g.set_desc.tracks
    # track = tracks[0]
    # parents of a node: track.get_parents(track.graph.edges[0].parent)
    # childs of a node: track.get_childs(track.graph.edges[0].parent)

    # for track_id, track in enumerate(tracks):
    #    plot_track(track, "track" + str(track_id))

    ds = pipeline.get_data()
    ds_set = get_subset_by_description(ds, trackable_set, '2d')

    plot_differences(g, tracks, cv=ds_set.cv)

# no out data besides plots on kitweather
if sys.argv[1] == '-kw':
    exit()


# out_netcdf_path = data_path + '_streamers.nc'

if len(sys.argv) == 1:
    out_json_path = out_dir + 'aew_desc.json'
    out_dataset_path = out_dir + '05_wt.nc'
elif len(sys.argv) == 2:
    out_json_path = out_dir + 'aew_desc_' + str(proc_summer_of_year) + '.json'
    out_dataset_path = out_dir + '05_wt_' + str(proc_summer_of_year) + '.nc'
else:
    m_str = str(proc_month_of_year).zfill(2)
    out_json_path = out_dir + 'aew_desc_' + str(proc_summer_of_year) + '_' + m_str + '.json'
    out_dataset_path = out_dir + '05_wt_' + str(proc_summer_of_year) + '_' + m_str + '.nc'
    


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

# print("Plot.")
# plot_timesteps_from_desc(od, pipeline.get_data())
# plot_tracks_from_desc(od, None)