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


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


def __dir__():
    import watch.tasks.rutgers_material_change_detection.datasets.peo_dataset 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("peo")

    # Visualize one example of dataset.
    root_dir = get_dataset_root_dir("peo")
    task_mode = "ss_arrow_of_time"
    video_slice = generate_video_slice_object(height=300, width=300, n_frames=None, scale=1)
    dset = PassiveEarthObservationDataset(root_dir, "train", video_slice, task_mode, channels="RGB_NIR")
    ex_index = 0
    ex = dset.__getitem__(ex_index)
    dset.visualize_example(
        ex_index, save_path=f"peo_ex_{str(ex_index).zfill(2)}.png", overlay_last_anno=True, num_plot_images=6
    )
    print("TEST 1: Create a visualizion. | PASSED")

    # Test 2: Get a single item from the training set in ss_triplet mode.
    video_slice = generate_video_slice_object(height=150, width=150, n_frames=None, scale=2)
    split = "train"
    task_mode = "ss_triplet"
    dset = PassiveEarthObservationDataset(root_dir, split, video_slice, task_mode, channels="RGB_NIR")
    example = dset.__getitem__(0)
    print("Test 2: Get one example from train split of dataset. | PASSED")

    # Test 3: Get a single item from the validation set in ss_triplet mode.
    video_slice = generate_video_slice_object(height=300, width=300, n_frames=None, scale=1)
    split = "valid"
    task_mode = "ss_triplet"
    dset = PassiveEarthObservationDataset(root_dir, split, video_slice, task_mode, channels="RGB_NIR")
    example = dset.__getitem__(0)
    print("Test 3: Get one example from valid split of dataset. | PASSED")

    # Test 4: Get a single item from the train set in ss_arrow_of_time mode.
    video_slice = generate_video_slice_object(height=300, width=300, n_frames=None, scale=1)
    split = "train"
    task_mode = "ss_arrow_of_time"
    dset = PassiveEarthObservationDataset(root_dir, split, video_slice, task_mode, channels="RGB_NIR")
    example = dset.__getitem__(0)
    print("Test 4: Get one example from valid split of dataset in arrow_of_time mode. | PASSED")

    # Test 5: Visualize an ss_splice_change_index example.
    video_slice = generate_video_slice_object(height=300, width=300, n_frames=10, scale=1)
    split = "train"
    task_mode = "ss_splice_change_index"
    dset = PassiveEarthObservationDataset(root_dir, split, video_slice, task_mode, channels="RGB_NIR")
    ex_index = 52
    example = dset.__getitem__(ex_index)
    dset.visualize_example(
        ex_index,
        save_path=f"peo_ss_splice_change_index_{str(ex_index).zfill(2)}.png",
        overlay_last_anno=False,
        num_plot_images=video_slice.n_frames,
    )
    print("Test 5: Visualize an ss_splice_change_index example.| PASSED")
