# Autogenerated via:
# python ~/code/watch/dev/maintain/mirror_package_geowatch.py
from watch.tasks.rutgers_material_change_detection.datasets.iarpa_kwdataset import print, print, dset, print, dset, print


def __getattr__(key):
    import watch.tasks.rutgers_material_change_detection.datasets.iarpa_kwdataset as mirror
    return getattr(mirror, key)


def __dir__():
    import watch.tasks.rutgers_material_change_detection.datasets.iarpa_kwdataset as mirror
    return dir(mirror)


if __name__ == '__main__':
    from watch.tasks.rutgers_material_change_detection.utils.util_paths import get_dataset_root_dir
    from watch.tasks.rutgers_material_change_detection.utils.util_misc import generate_video_slice_object

    root_dir = get_dataset_root_dir("iarpa_drop1")
    # root_dir = get_dataset_root_dir('iarpa_drop1_ta1')
    task_mode = "total_bin_change"

    # Create dataset and load one example.
    split = "train"
    video_slice = generate_video_slice_object(height=150, width=150, n_frames=None, scale=2)
    dset = IARPA_KWDATASET(root_dir, split, video_slice, task_mode, channels="RGB_NIR", normalize_mode="local")
    example = dset.__getitem__(4)
    print("TEST 1: Load one example from dataset. | PASSED")

    # Check that validation set of dataset works.
    split = "valid"
    video_slice = generate_video_slice_object(height=400, width=400, n_frames=5, scale=None)
    dset = IARPA_KWDATASET(root_dir, split, video_slice, task_mode, channels="RGB_NIR")
    example = dset.__getitem__(0)
    print("TEST 2: Get an example from validation set. | PASSED")

    # Visualize one example of dataset.
    split = "valid"
    video_slice = generate_video_slice_object(height=300, width=300, n_frames=None, scale=None)
    dset = IARPA_KWDATASET(root_dir, split, video_slice, task_mode, channels="RGB_NIR")
    ex_index = 1
    dset.visualize_example(ex_index, save_path=f"iarpa_data_ex_{str(ex_index).zfill(2)}.png", overlay_last_anno=True)
    print("TEST 3: Create a visualizion. | PASSED")

    # Visualize datetimes of one example.
    # split = 'train'
    # video_slice = generate_video_slice_object(height=100, width=100, n_frames=None, scale=None)
    # dset = IARPA_KWDATASET(root_dir, split, video_slice, task_mode, channels='RGB_NIR')
    # ex_index = 10
    # dset.visualize_dates(ex_index, save_path=f'iarpa_data_dates_ex_{str(ex_index).zfill(2)}.png')
    # print('TEST 4: Create a dates visualizion. | PASSED')

    # Visualize image annotations over time.
    split = "valid"
    video_slice = generate_video_slice_object(height=800, width=800, n_frames=None, scale=None)
    dset = IARPA_KWDATASET(root_dir, split, video_slice, task_mode, channels="RGB_NIR")
    ex_index = 0
    dset.generate_propagated_annotation_gif(ex_index)
    print("TEST 5: Generate a video of annotations over a region. | PASSED")
