import random

import fiftyone as fo
import fiftyone.utils.selection as fous
import fiftyone.zoo as foz


NUM_SAMPLES_TO_SELECT = 3
MAX_OBJECTS_PER_SAMPLE_TO_SELECT = 5


def count_detections(sample_collection, label_field):
    num_objects = 0
    for sample in sample_collection:
        num_objects += len(sample[label_field].detections)

    return num_objects


dataset = foz.load_zoo_dataset("quickstart")


# Generate some random object selections
selected_objects = []
for sample in dataset.take(NUM_SAMPLES_TO_SELECT):
    detections = sample.ground_truth.detections
    num_objects = random.randint(
        1, min(len(detections), MAX_OBJECTS_PER_SAMPLE_TO_SELECT)
    )
    for detection in random.sample(detections, num_objects):
        selected_objects.append(
            {
                "sample_id": sample.id,
                "field": "ground_truth",
                "object_id": detection.id,
            }
        )

print("Selected objects:\n%s\n" % fo.pformat(selected_objects))

# Get only selected objects
selected_view = fous.select_objects(dataset, selected_objects)

# Exclude selected objects
excluded_view = fous.exclude_objects(dataset, selected_objects)

total_objects = count_detections(dataset, "ground_truth")
num_selected_objects = len(selected_objects)
num_objects_in_selected_view = count_detections(selected_view, "ground_truth")
num_objects_in_excluded_view = count_detections(excluded_view, "ground_truth")
num_objects_excluded = total_objects - num_objects_in_excluded_view

print("Total objects: %d" % total_objects)
print("Selected objects: %d" % num_selected_objects)
print("Objects in selected view: %d" % num_objects_in_selected_view)
print("Objects in excluded view: %d (%d missing)" % (
    num_objects_in_excluded_view, num_objects_excluded
))