# Autogenerated via:
# python ~/code/watch/dev/maintain/mirror_package_geowatch.py
from watch.tasks.rutgers_material_seg.experiments.cluster_labeling import print, dset, sys


def __getattr__(key):
    import watch.tasks.rutgers_material_seg.experiments.cluster_labeling as mirror
    return getattr(mirror, key)


def __dir__():
    import watch.tasks.rutgers_material_seg.experiments.cluster_labeling as mirror
    return dir(mirror)


if __name__ == '__main__':

    app = QApplication(sys.argv)
    # save_kwcoco_path = "/home/native/core534_data/datasets/smart_watch/processed/drop0_aligned_v2.1/material_labels.kwcoco.json"
    resume = ""
    # resume = "/media/native/data/data/smart_watch_dvc/drop0_aligned_msi/material_labels2.kwcoco.json"
    save_kwcoco_path = "/media/native/data/data/smart_watch_dvc/drop0_aligned_msi/material_labels2.kwcoco.json"

    coco_fpath = ub.expandpath('/media/native/data/data/smart_watch_dvc/drop0_aligned_msi/data_fielded.kwcoco.json')
    # coco_fpath = ub.expandpath('/home/native/core534_data/datasets/smart_watch/processed/drop0_aligned_v2.1/data_fielded_filtered.kwcoco.json')

    # "cirrus", "coastal", "costal", "green", "lwir11", "lwir12", "nir", "pan", "red", "swir16", "swir22"]
    # channels = 'red|green|blue|nir|lwir11|lwir12|swir16|swir22'#|coastal|costal'
    channels = 'red|green|blue|nir|swir16|swir22|cirrus'
    # channels = 'red|green|blue|nir|swir22|cirrus'

    expected_channels = channels.split('|')
    print(expected_channels)
    # expected_channels = [ "red"  ,"green", "blue", "nir", "lwir11", "lwir12", "swir16", "swir22", "cirrus"]
    dset = kwcoco.CocoDataset(coco_fpath)

    # Only select images with the correct channels
    gids_to_remove = []
    for gid, img in dset.index.imgs.items():
        try:
            # print(img['auxiliary'][0]['channels'])
            image_auxiliary_list = img['auxiliary']
            bands_available = [x['channels'] for x in image_auxiliary_list]
            # print(bands_available)
            if not all(elem in bands_available for elem in expected_channels):
                gids_to_remove.append(gid)
        except Exception:
            gids_to_remove.append(gid)
            continue

    dset.remove_images(gids_to_remove)

    sampler = ndsampler.CocoSampler(dset)

    number_of_timestamps, h, w = 3, 512, 512
    window_dims = (number_of_timestamps, h, w)  # [t,h,w]
    input_dims = (h, w)

    dataset = SequenceDataset(sampler, window_dims, input_dims, channels)
    loader = dataset.make_loader(batch_size=1)

    window = Window(dataset, dset, resume, save_path=save_kwcoco_path)

    sys.exit(app.exec_())
