"""
FiftyOne models.

| Copyright 2017-2022, Voxel51, Inc.
| `voxel51.com <https://voxel51.com/>`_
|
"""
import contextlib
import inspect
import logging

import numpy as np

import eta.core.image as etai
import eta.core.frameutils as etaf
import eta.core.learning as etal
import eta.core.models as etam
import eta.core.utils as etau
import eta.core.video as etav
import eta.core.web as etaw

import fiftyone as fo
import fiftyone.core.labels as fol
import fiftyone.core.media as fom
import fiftyone.core.utils as fou
import fiftyone.core.validation as fov

tud = fou.lazy_import("torch.utils.data")

foue = fou.lazy_import("fiftyone.utils.eta")
fouf = fou.lazy_import("fiftyone.utils.flash")
foup = fou.lazy_import("fiftyone.utils.patches")
fout = fou.lazy_import("fiftyone.utils.torch")


logger = logging.getLogger(__name__)


_ALLOWED_PATCH_TYPES = (
    fol.Detection,
    fol.Detections,
    fol.Polyline,
    fol.Polylines,
)


def apply_model(
    samples,
    model,
    label_field="predictions",
    confidence_thresh=None,
    store_logits=False,
    batch_size=None,
    num_workers=None,
    skip_failures=True,
    **kwargs,
):
    """Applies the :class:`FiftyOne model <Model>` or
    :class:`Lightning Flash model <flash:flash.core.model.Task>` to the samples
    in the collection.

    This method supports all of the following cases:

    -   Applying an image :class:`Model` to an image collection
    -   Applying an image :class:`Model` to the frames of a video collection
    -   Applying a video :class:`Model` to a video collection
    -   Applying a :class:`flash:flash.core.model.Task` to an image or video
        collection

    Args:
        samples: a :class:`fiftyone.core.collections.SampleCollection`
        model: a :class:`Model` or :class:`flash:flash.core.model.Task`
        label_field ("predictions"): the name of the field in which to store
            the model predictions. When performing inference on video frames,
            the "frames." prefix is optional
        confidence_thresh (None): an optional confidence threshold to apply to
            any applicable labels generated by the model
        store_logits (False): whether to store logits for the model
            predictions. This is only supported when the provided ``model`` has
            logits, ``model.has_logits == True``
        batch_size (None): an optional batch size to use, if the model supports
            batching
        num_workers (None): the number of workers to use when loading images.
            Only applicable for Torch-based models
        skip_failures (True): whether to gracefully continue without raising an
            error if predictions cannot be generated for a sample. Only
            applicable to :class:`Model` instances
        **kwargs: optional model-specific keyword arguments passed through
            to the underlying inference implementation
    """
    if _is_flash_model(model):
        return fouf.apply_flash_model(
            samples,
            model,
            label_field=label_field,
            confidence_thresh=confidence_thresh,
            store_logits=store_logits,
            batch_size=batch_size,
            num_workers=num_workers,
            **kwargs,
        )

    if not isinstance(model, Model):
        raise ValueError(
            "Model must be a %s or %s; found %s"
            % (Model, _BASE_FLASH_TYPE, type(model))
        )

    if samples.media_type == fom.IMAGE:
        fov.validate_image_collection(samples)
    elif samples.media_type == fom.GROUP:
        raise fom.SelectGroupSlicesError((fom.IMAGE, fom.VIDEO))
    elif samples.media_type != fom.VIDEO:
        raise fom.MediaTypeError(
            "Unsupported media type '%s'" % samples.media_type
        )

    if model.media_type == "video" and samples.media_type != fom.VIDEO:
        raise ValueError(
            "Video models can only be applied to video collections"
        )

    if model.media_type not in ("image", "video"):
        raise ValueError(
            "Unsupported model `media_type=%s`. Supported values are "
            "('image', 'video')" % model.media_type
        )

    if store_logits and not model.has_logits:
        raise ValueError(
            "The provided model does not expose logits "
            "(model.has_logits = %s)" % model.has_logits
        )

    use_data_loader = (
        isinstance(model, TorchModelMixin) and samples.media_type == fom.IMAGE
    )

    if num_workers is not None and not use_data_loader:
        logger.warning(
            "Ignoring `num_workers` parameter; only supported for Torch models"
        )

    with contextlib.ExitStack() as context:
        try:
            if confidence_thresh is not None:
                cthresh = confidence_thresh
            else:
                cthresh = model.config.confidence_thresh

            # pylint: disable=no-member
            context.enter_context(
                fou.SetAttributes(model.config, confidence_thresh=cthresh)
            )
        except:
            pass

        if store_logits:
            # pylint: disable=no-member
            context.enter_context(fou.SetAttributes(model, store_logits=True))

        if use_data_loader:
            # pylint: disable=no-member
            context.enter_context(fou.SetAttributes(model, preprocess=False))

        # pylint: disable=no-member
        context.enter_context(model)

        if samples.media_type == fom.VIDEO and model.media_type == "video":
            return _apply_video_model(
                samples, model, label_field, confidence_thresh, skip_failures
            )

        batch_size = _parse_batch_size(batch_size, model, use_data_loader)

        if samples.media_type == fom.VIDEO and model.media_type == "image":
            label_field, _ = samples._handle_frame_field(label_field)

            if batch_size is not None:
                return _apply_image_model_to_frames_batch(
                    samples,
                    model,
                    label_field,
                    confidence_thresh,
                    batch_size,
                    skip_failures,
                )

            return _apply_image_model_to_frames_single(
                samples, model, label_field, confidence_thresh, skip_failures
            )

        if use_data_loader:
            return _apply_image_model_data_loader(
                samples,
                model,
                label_field,
                confidence_thresh,
                batch_size,
                num_workers,
                skip_failures,
            )

        if batch_size is not None:
            return _apply_image_model_batch(
                samples,
                model,
                label_field,
                confidence_thresh,
                batch_size,
                skip_failures,
            )

        return _apply_image_model_single(
            samples, model, label_field, confidence_thresh, skip_failures
        )


_BASE_FLASH_TYPE = "flash.core.model.Task"


def _is_flash_model(model):
    for cls in inspect.getmro(type(model)):
        if etau.get_class_name(cls) == _BASE_FLASH_TYPE:
            return True

    return False


def _apply_image_model_single(
    samples, model, label_field, confidence_thresh, skip_failures
):
    samples = samples.select_fields()

    with fou.ProgressBar() as pb:
        for sample in pb(samples):
            try:
                img = etai.read(sample.filepath)
                labels = model.predict(img)

                sample.add_labels(
                    labels, label_field, confidence_thresh=confidence_thresh
                )
            except Exception as e:
                if not skip_failures:
                    raise e

                logger.warning("Sample: %s\nError: %s\n", sample.id, e)


def _apply_image_model_batch(
    samples, model, label_field, confidence_thresh, batch_size, skip_failures
):
    samples = samples.select_fields()
    samples_loader = fou.iter_batches(samples, batch_size)

    with fou.ProgressBar(samples) as pb:
        for sample_batch in samples_loader:
            try:
                imgs = [etai.read(sample.filepath) for sample in sample_batch]
                labels_batch = model.predict_all(imgs)

                for sample, labels in zip(sample_batch, labels_batch):
                    sample.add_labels(
                        labels,
                        label_field,
                        confidence_thresh=confidence_thresh,
                    )

            except Exception as e:
                if not skip_failures:
                    raise e

                logger.warning(
                    "Batch: %s - %s\nError: %s\n",
                    sample_batch[0].id,
                    sample_batch[-1].id,
                    e,
                )

            pb.update(len(sample_batch))


def _apply_image_model_data_loader(
    samples,
    model,
    label_field,
    confidence_thresh,
    batch_size,
    num_workers,
    skip_failures,
):
    samples = samples.select_fields()
    samples_loader = fou.iter_batches(samples, batch_size)
    data_loader = _make_data_loader(
        samples, model, batch_size, num_workers, skip_failures
    )

    with fou.ProgressBar(samples) as pb:
        for sample_batch, imgs in zip(samples_loader, data_loader):
            try:
                if isinstance(imgs, Exception):
                    raise imgs

                labels_batch = model.predict_all(imgs)

                for sample, labels in zip(sample_batch, labels_batch):
                    sample.add_labels(
                        labels,
                        label_field,
                        confidence_thresh=confidence_thresh,
                    )

            except Exception as e:
                if not skip_failures:
                    raise e

                logger.warning(
                    "Batch: %s - %s\nError: %s\n",
                    sample_batch[0].id,
                    sample_batch[-1].id,
                    e,
                )

            pb.update(len(sample_batch))


def _apply_image_model_to_frames_single(
    samples, model, label_field, confidence_thresh, skip_failures
):
    samples = samples.select_fields()
    frame_counts, total_frame_count = _get_frame_counts(samples)
    is_clips = samples._dataset._is_clips

    with fou.ProgressBar(total=total_frame_count) as pb:
        for idx, sample in enumerate(samples):
            if is_clips:
                frames = etaf.FrameRange(*sample.support)
            else:
                frames = None

            try:
                with etav.FFmpegVideoReader(
                    sample.filepath, frames=frames
                ) as video_reader:
                    for img in video_reader:
                        labels = model.predict(img)

                        sample.add_labels(
                            {video_reader.frame_number: labels},
                            label_field,
                            confidence_thresh=confidence_thresh,
                        )

                        pb.update()

            except Exception as e:
                if not skip_failures:
                    raise e

                logger.warning("Sample: %s\nError: %s\n", sample.id, e)

            # Explicitly set in case actual # frames differed from expected #
            pb.set_iteration(frame_counts[idx])


def _apply_image_model_to_frames_batch(
    samples, model, label_field, confidence_thresh, batch_size, skip_failures
):
    samples = samples.select_fields()
    frame_counts, total_frame_count = _get_frame_counts(samples)
    is_clips = samples._dataset._is_clips

    with fou.ProgressBar(total=total_frame_count) as pb:
        for idx, sample in enumerate(samples):
            if is_clips:
                frames = etaf.FrameRange(*sample.support)
            else:
                frames = None

            try:
                with etav.FFmpegVideoReader(
                    sample.filepath, frames=frames
                ) as video_reader:
                    for fns, imgs in _iter_batches(video_reader, batch_size):
                        labels_batch = model.predict_all(imgs)

                        sample.add_labels(
                            {
                                fn: labels
                                for fn, labels in zip(fns, labels_batch)
                            },
                            label_field,
                            confidence_thresh=confidence_thresh,
                        )

                        pb.update(len(imgs))

            except Exception as e:
                if not skip_failures:
                    raise e

                logger.warning("Sample: %s\nError: %s\n", sample.id, e)

            # Explicitly set in case actual # frames differed from expected #
            pb.set_iteration(frame_counts[idx])


def _apply_video_model(
    samples, model, label_field, confidence_thresh, skip_failures
):
    samples = samples.select_fields()
    is_clips = samples._dataset._is_clips

    with fou.ProgressBar() as pb:
        for sample in pb(samples):
            if is_clips:
                frames = etaf.FrameRange(*sample.support)
            else:
                frames = None

            try:
                with etav.FFmpegVideoReader(
                    sample.filepath, frames=frames
                ) as video_reader:
                    labels = model.predict(video_reader)

                sample.add_labels(
                    labels, label_field, confidence_thresh=confidence_thresh
                )
            except Exception as e:
                if not skip_failures:
                    raise e

                logger.warning("Sample: %s\nError: %s\n", sample.id, e)


def _get_frame_counts(samples):
    if samples._dataset._is_clips:
        expr = fo.ViewField("support")[1] - fo.ViewField("support")[0] + 1
        frame_counts = samples.values(expr)
    else:
        samples.compute_metadata()
        frame_counts = samples.values("metadata.total_frame_count")

    frame_counts = np.cumsum([(fc or 0) for fc in frame_counts])
    total_frame_count = frame_counts[-1] if frame_counts.size > 0 else 0

    return frame_counts, total_frame_count


def _iter_batches(video_reader, batch_size):
    frame_numbers = []
    imgs = []
    for img in video_reader:
        imgs.append(img)
        frame_numbers.append(video_reader.frame_number)
        if len(imgs) >= batch_size:
            yield frame_numbers, imgs
            frame_numbers = []
            imgs = []

    if frame_numbers:
        yield frame_numbers, imgs


def _make_data_loader(samples, model, batch_size, num_workers, skip_failures):
    # This function supports DataLoaders that emit numpy arrays that can
    # therefore be used for non-Torch models; but we do not currenly use this
    # functionality
    use_numpy = not isinstance(model, TorchModelMixin)

    if num_workers is None:
        num_workers = fout.recommend_num_workers()

    dataset = fout.TorchImageDataset(
        samples=samples,
        transform=model.transforms,
        use_numpy=use_numpy,
        force_rgb=True,
        skip_failures=skip_failures,
    )

    def handle_errors(batch):
        errors = [b for b in batch if isinstance(b, Exception)]

        if not errors:
            return None

        if len(errors) == 1:
            return errors[0]

        return Exception("\n" + "\n".join([str(e) for e in errors]))

    if model.ragged_batches:

        def collate_fn(batch):
            error = handle_errors(batch)
            if error is not None:
                return error

            return batch  # return list

    elif use_numpy:

        def collate_fn(batch):
            error = handle_errors(batch)
            if error is not None:
                return error

            try:
                return np.stack(batch)
            except Exception as e:
                if not skip_failures:
                    raise e

                return e

    else:

        def collate_fn(batch):
            error = handle_errors(batch)
            if error is not None:
                return error

            try:
                return tud.dataloader.default_collate(batch)
            except Exception as e:
                if not skip_failures:
                    raise e

                return e

    if batch_size is None:
        batch_size = 1

    return tud.DataLoader(
        dataset,
        batch_size=batch_size,
        num_workers=num_workers,
        collate_fn=collate_fn,
    )


def compute_embeddings(
    samples,
    model,
    embeddings_field=None,
    batch_size=None,
    num_workers=None,
    skip_failures=True,
    **kwargs,
):
    """Computes embeddings for the samples in the collection using the given
    :class:`FiftyOne model <Model>` or
    :class:`Lightning Flash model <flash:flash.core.model.Task>`.

    This method supports all the following cases:

    -   Using an image :class:`Model` to compute embeddings for an image
        collection
    -   Using an image :class:`Model` to compute frame embeddings for a video
        collection
    -   Using a video :class:`Model` to compute embeddings for a video
        collection
    -   Using an :ref:`ImageEmbedder <flash:image_embedder>` to compute
        embeddings for an image collection

    When using a :class:`FiftyOne model <Model>`, the model must expose
    embeddings, i.e., :meth:`Model.has_embeddings` must return ``True``.

    If an ``embeddings_field`` is provided, the embeddings are saved to the
    samples; otherwise, the embeddings are returned in-memory.

    Args:
        samples: a :class:`fiftyone.core.collections.SampleCollection`
        model: a :class:`Model` or :class:`flash:flash.core.model.Task>`
        embeddings_field (None): the name of a field in which to store the
            embeddings. When computing video frame embeddings, the "frames."
            prefix is optional
        batch_size (None): an optional batch size to use, if the model supports
            batching
        num_workers (None): the number of workers to use when loading images.
            Only applicable for Torch-based models
        skip_failures (True): whether to gracefully continue without raising an
            error if embeddings cannot be generated for a sample. Only
            applicable to :class:`Model` instances
        **kwargs: optional model-specific keyword arguments passed through
            to the underlying inference implementation

    Returns:
        one of the following:

        -   ``None``, if an ``embeddings_field`` is provided
        -   a ``num_samples x num_dim`` array of embeddings, when computing
            embeddings for image/video collections with image/video models,
            respectively, and no ``embeddings_field`` is provided. If
            ``skip_failures`` is ``True`` and any errors are detected, a list
            of length ``num_samples`` is returned instead containing all
            successfully computed embedding vectors along with ``None`` entries
            for samples for which embeddings could not be computed
        -   a dictionary mapping sample IDs to ``num_frames x num_dim`` arrays
            of embeddings, when computing frame embeddings for video
            collections using an image model. If ``skip_failures`` is ``True``
            and any errors are detected, the values of this dictionary will
            contain arrays of embeddings for all frames 1, 2, ... until the
            error occurred, or ``None`` if no embeddings were computed at all
    """
    if _is_flash_model(model):
        return fouf.compute_flash_embeddings(
            samples,
            model,
            embeddings_field=embeddings_field,
            batch_size=batch_size,
            num_workers=num_workers,
            **kwargs,
        )

    if not isinstance(model, Model):
        raise ValueError(
            "Model must be a %s or %s; found %s"
            % (Model, _BASE_FLASH_TYPE, type(model))
        )

    if not model.has_embeddings:
        raise ValueError(
            "Model must expose embeddings; found model.has_embeddings = %s"
            % model.has_embeddings
        )

    if samples.media_type == fom.IMAGE:
        fov.validate_image_collection(samples)
    elif samples.media_type == fom.GROUP:
        raise fom.SelectGroupSlicesError((fom.IMAGE, fom.VIDEO))
    elif samples.media_type != fom.VIDEO:
        raise fom.MediaTypeError(
            "Unsupported media type '%s'" % samples.media_type
        )

    if model.media_type == "video" and samples.media_type != fom.VIDEO:
        raise ValueError(
            "Video models can only be applied to video collections"
        )

    if model.media_type not in ("image", "video"):
        raise ValueError(
            "Unsupported model `media_type=%s`. Supported values are "
            "('image', 'video')" % model.media_type
        )

    use_data_loader = (
        isinstance(model, TorchModelMixin) and samples.media_type == fom.IMAGE
    )

    if num_workers is not None and not use_data_loader:
        logger.warning(
            "Ignoring `num_workers` parameter; only supported for Torch models"
        )

    with contextlib.ExitStack() as context:
        if use_data_loader:
            # pylint: disable=no-member
            context.enter_context(fou.SetAttributes(model, preprocess=False))

        # pylint: disable=no-member
        context.enter_context(model)

        if samples.media_type == fom.VIDEO and model.media_type == "video":
            return _compute_video_embeddings(
                samples, model, embeddings_field, skip_failures
            )

        batch_size = _parse_batch_size(batch_size, model, use_data_loader)

        if samples.media_type == fom.VIDEO and model.media_type == "image":
            if embeddings_field is not None:
                embeddings_field, _ = samples._handle_frame_field(
                    embeddings_field
                )

            if batch_size is not None:
                return _compute_frame_embeddings_batch(
                    samples, model, embeddings_field, batch_size, skip_failures
                )

            return _compute_frame_embeddings_single(
                samples, model, embeddings_field, skip_failures
            )

        if use_data_loader:
            return _compute_image_embeddings_data_loader(
                samples,
                model,
                embeddings_field,
                batch_size,
                num_workers,
                skip_failures,
            )

        if batch_size is not None:
            return _compute_image_embeddings_batch(
                samples, model, embeddings_field, batch_size, skip_failures
            )

        return _compute_image_embeddings_single(
            samples, model, embeddings_field, skip_failures
        )


def _compute_image_embeddings_single(
    samples, model, embeddings_field, skip_failures
):
    samples = samples.select_fields()
    embeddings = []

    errors = False

    with fou.ProgressBar() as pb:
        for sample in pb(samples):
            embedding = None

            try:
                img = etai.read(sample.filepath)
                embedding = model.embed(img)[0]
            except Exception as e:
                if not skip_failures:
                    raise e

                errors = True
                logger.warning("Sample: %s\nError: %s\n", sample.id, e)

            if embeddings_field:
                sample[embeddings_field] = embedding
                sample.save()
            else:
                embeddings.append(embedding)

    if embeddings_field:
        return None

    if errors:
        return embeddings  # may contain None, must return as list

    return np.stack(embeddings)


def _compute_image_embeddings_batch(
    samples, model, embeddings_field, batch_size, skip_failures
):
    samples = samples.select_fields()
    samples_loader = fou.iter_batches(samples, batch_size)

    embeddings = []
    errors = False

    with fou.ProgressBar(samples) as pb:
        for sample_batch in samples_loader:
            embeddings_batch = [None] * len(sample_batch)

            try:
                imgs = [etai.read(sample.filepath) for sample in sample_batch]
                embeddings_batch = list(model.embed_all(imgs))  # list of 1D
            except Exception as e:
                if not skip_failures:
                    raise e

                errors = True
                logger.warning(
                    "Batch: %s - %s\nError: %s\n",
                    sample_batch[0].id,
                    sample_batch[-1].id,
                    e,
                )

            if embeddings_field:
                for sample, embedding in zip(sample_batch, embeddings_batch):
                    sample[embeddings_field] = embedding
                    sample.save()
            else:
                embeddings.extend(embeddings_batch)

            pb.update(len(sample_batch))

    if embeddings_field:
        return None

    if errors:
        return embeddings  # may contain None, must return as list

    return np.stack(embeddings)


def _compute_image_embeddings_data_loader(
    samples, model, embeddings_field, batch_size, num_workers, skip_failures
):
    samples = samples.select_fields()
    samples_loader = fou.iter_batches(samples, batch_size)
    data_loader = _make_data_loader(
        samples, model, batch_size, num_workers, skip_failures
    )

    embeddings = []
    errors = False

    with fou.ProgressBar(samples) as pb:
        for sample_batch, imgs in zip(samples_loader, data_loader):
            embeddings_batch = [None] * len(sample_batch)

            try:
                if isinstance(imgs, Exception):
                    raise imgs

                embeddings_batch = list(model.embed_all(imgs))  # list of 1D
            except Exception as e:
                if not skip_failures:
                    raise e

                errors = True
                logger.warning(
                    "Batch: %s - %s\nError: %s\n",
                    sample_batch[0].id,
                    sample_batch[-1].id,
                    e,
                )

            if embeddings_field:
                for sample, embedding in zip(sample_batch, embeddings_batch):
                    sample[embeddings_field] = embedding
                    sample.save()
            else:
                embeddings.extend(embeddings_batch)

            pb.update(len(sample_batch))

    if embeddings_field:
        return None

    if errors:
        return embeddings  # may contain None, must return as list

    return np.stack(embeddings)


def _compute_frame_embeddings_single(
    samples, model, embeddings_field, skip_failures
):
    samples = samples.select_fields()
    frame_counts, total_frame_count = _get_frame_counts(samples)
    is_clips = samples._dataset._is_clips

    embeddings_dict = {}

    with fou.ProgressBar(total=total_frame_count) as pb:
        for idx, sample in enumerate(samples):
            embeddings = []

            if is_clips:
                frames = etaf.FrameRange(*sample.support)
            else:
                frames = None

            try:
                with etav.FFmpegVideoReader(
                    sample.filepath, frames=frames
                ) as video_reader:
                    for img in video_reader:
                        embedding = model.embed(img)[0]

                        if embeddings_field is not None:
                            sample.add_labels(
                                {video_reader.frame_number: embedding},
                                embeddings_field,
                            )
                        else:
                            embeddings.append(embedding)

                        pb.update()

            except Exception as e:
                if not skip_failures:
                    raise e

                logger.warning("Sample: %s\nError: %s\n", sample.id, e)

            if embeddings_field is None:
                if embeddings:
                    embeddings = np.stack(embeddings)
                else:
                    embeddings = None

                embeddings_dict[sample.id] = embeddings

            # Explicitly set in case actual # frames differed from expected #
            pb.set_iteration(frame_counts[idx])

    if embeddings_field:
        return None

    return embeddings_dict


def _compute_frame_embeddings_batch(
    samples, model, embeddings_field, batch_size, skip_failures
):
    samples = samples.select_fields()
    frame_counts, total_frame_count = _get_frame_counts(samples)
    is_clips = samples._dataset._is_clips

    embeddings_dict = {}

    with fou.ProgressBar(total=total_frame_count) as pb:
        for idx, sample in enumerate(samples):
            embeddings = []

            if is_clips:
                frames = etaf.FrameRange(*sample.support)
            else:
                frames = None

            try:
                with etav.FFmpegVideoReader(
                    sample.filepath, frames=frames
                ) as video_reader:
                    for fns, imgs in _iter_batches(video_reader, batch_size):
                        embeddings_batch = list(model.embed_all(imgs))

                        if embeddings_field is not None:
                            sample.add_labels(
                                {
                                    fn: embedding
                                    for fn, embedding in zip(
                                        fns, embeddings_batch
                                    )
                                },
                                embeddings_field,
                            )
                        else:
                            embeddings.extend(embeddings_batch)

                        pb.update(len(imgs))

            except Exception as e:
                if not skip_failures:
                    raise e

                logger.warning("Sample: %s\nError: %s\n", sample.id, e)

            if embeddings_field is None:
                if embeddings:
                    embeddings = np.stack(embeddings)
                else:
                    embeddings = None

                embeddings_dict[sample.id] = embeddings

            # Explicitly set in case actual # frames differed from expected #
            pb.set_iteration(frame_counts[idx])

    if embeddings_field:
        return None

    return embeddings_dict


def _compute_video_embeddings(samples, model, embeddings_field, skip_failures):
    samples = samples.select_fields()
    is_clips = samples._dataset._is_clips

    embeddings = []
    errors = False

    with fou.ProgressBar() as pb:
        for sample in pb(samples):
            if is_clips:
                frames = etaf.FrameRange(*sample.support)
            else:
                frames = None

            try:
                with etav.FFmpegVideoReader(
                    sample.filepath, frames=frames
                ) as video_reader:
                    embedding = model.embed(video_reader)[0]

            except Exception as e:
                if not skip_failures:
                    raise e

                errors = True
                logger.warning("Sample: %s\nError: %s\n", sample.id, e)

            if embeddings_field:
                sample[embeddings_field] = embedding
                sample.save()
            else:
                embeddings.append(embedding)

    if embeddings_field:
        return None

    if errors:
        return embeddings  # may contain None, must return as list

    return np.stack(embeddings)


def compute_patch_embeddings(
    samples,
    model,
    patches_field,
    embeddings_field=None,
    force_square=False,
    alpha=None,
    handle_missing="skip",
    batch_size=None,
    num_workers=None,
    skip_failures=True,
):
    """Computes embeddings for the image patches defined by ``patches_field``
    of the samples in the collection using the given :class:`Model`.

    This method supports all the following cases:

    -   Using an image model to compute patch embeddings for an image
        collection
    -   Using an image model to compute frame patch embeddings for a video
        collection

    The ``model`` must expose embeddings, i.e., :meth:`Model.has_embeddings`
    must return ``True``.

    If an ``embeddings_field`` is provided, the embeddings are saved to the
    samples; otherwise, the embeddings are returned in-memory.

    Args:
        samples: a :class:`fiftyone.core.collections.SampleCollection`
        model: a :class:`Model`
        patches_field: the name of the field defining the image patches in each
            sample to embed. Must be of type
            :class:`fiftyone.core.labels.Detection`,
            :class:`fiftyone.core.labels.Detections`,
            :class:`fiftyone.core.labels.Polyline`, or
            :class:`fiftyone.core.labels.Polylines`. When computing video frame
            embeddings, the "frames." prefix is optional
        embeddings_field (None): the name of a field in which to store the
            embeddings. When computing video frame embeddings, the "frames."
            prefix is optional
        force_square (False): whether to minimally manipulate the patch
            bounding boxes into squares prior to extraction
        alpha (None): an optional expansion/contraction to apply to the patches
            before extracting them, in ``[-1, inf)``. If provided, the length
            and width of the box are expanded (or contracted, when
            ``alpha < 0``) by ``(100 * alpha)%``. For example, set
            ``alpha = 1.1`` to expand the boxes by 10%, and set ``alpha = 0.9``
            to contract the boxes by 10%
        handle_missing ("skip"): how to handle images with no patches.
            Supported values are:

            -   "skip": skip the image and assign its embedding as ``None``
            -   "image": use the whole image as a single patch
            -   "error": raise an error

        batch_size (None): an optional batch size to use, if the model supports
            batching
        num_workers (None): the number of workers to use when loading images.
            Only applicable for Torch models
        skip_failures (True): whether to gracefully continue without raising an
            error if embeddings cannot be generated for a sample

    Returns:
        one of the following:

        -   ``None``, if an ``embeddings_field`` is provided
        -   a dict mapping sample IDs to ``num_patches x num_dim`` arrays of
            patch embeddings, when computing patch embeddings for image
            collections and no ``embeddings_field`` is provided. If
            ``skip_failures`` is ``True`` and any errors are detected, this
            dictionary will contain ``None`` values for any samples for which
            embeddings could not be computed
        -   a dict of dicts mapping sample IDs to frame numbers to
            ``num_patches x num_dim`` arrays of patch embeddings, when
            computing patch embeddings for the frames of video collections and
            no ``embeddings_field`` is provided. If ``skip_failures`` is
            ``True`` and any errors are detected, this nested dict will contain
            missing or ``None`` values to indicate uncomputable embeddings
    """
    if not isinstance(model, Model):
        raise ValueError(
            "Model must be a %s instance; found %s" % (Model, type(model))
        )

    if not model.has_embeddings:
        raise ValueError(
            "Model must expose embeddings; found model.has_embeddings = %s"
            % model.has_embeddings
        )

    if model.media_type != "image":
        raise ValueError(
            "This method only supports image models; found "
            "model.media_type = %s" % model.media_type
        )

    _handle_missing_supported = {"skip", "image", "error"}
    if handle_missing not in _handle_missing_supported:
        raise ValueError(
            "Unsupported handle_missing = '%s'; supported values are %s"
            % (handle_missing, _handle_missing_supported)
        )

    use_data_loader = (
        isinstance(model, TorchModelMixin) and samples.media_type == fom.IMAGE
    )

    if num_workers is not None and not use_data_loader:
        logger.warning(
            "Ignoring `num_workers` parameter; only supported for Torch models"
        )

    if samples.media_type == fom.IMAGE:
        fov.validate_image_collection(samples)
        fov.validate_collection_label_fields(
            samples, patches_field, _ALLOWED_PATCH_TYPES
        )
    elif samples.media_type == fom.VIDEO:
        patches_field, _ = samples._handle_frame_field(patches_field)
        if embeddings_field is not None:
            embeddings_field, _ = samples._handle_frame_field(embeddings_field)

        fov.validate_collection_label_fields(
            samples,
            samples._FRAMES_PREFIX + patches_field,
            _ALLOWED_PATCH_TYPES,
        )
    elif samples.media_type == fom.GROUP:
        raise fom.SelectGroupSlicesError((fom.IMAGE, fom.VIDEO))
    else:
        raise fom.MediaTypeError(
            "Unsupported media type '%s'" % samples.media_type
        )

    batch_size = _parse_batch_size(batch_size, model, use_data_loader)

    with contextlib.ExitStack() as context:
        if use_data_loader:
            # pylint: disable=no-member
            context.enter_context(fou.SetAttributes(model, preprocess=False))

        # pylint: disable=no-member
        context.enter_context(model)

        if samples.media_type == fom.VIDEO:
            return _embed_frame_patches(
                samples,
                model,
                patches_field,
                embeddings_field,
                force_square,
                alpha,
                handle_missing,
                batch_size,
                skip_failures,
            )

        if use_data_loader:
            return _embed_patches_data_loader(
                samples,
                model,
                patches_field,
                embeddings_field,
                force_square,
                alpha,
                handle_missing,
                batch_size,
                num_workers,
                skip_failures,
            )

        return _embed_patches(
            samples,
            model,
            patches_field,
            embeddings_field,
            force_square,
            alpha,
            handle_missing,
            batch_size,
            skip_failures,
        )


def _embed_patches(
    samples,
    model,
    patches_field,
    embeddings_field,
    force_square,
    alpha,
    handle_missing,
    batch_size,
    skip_failures,
):
    samples = samples.select_fields(patches_field)

    embeddings_dict = {}

    with fou.ProgressBar() as pb:
        for sample in pb(samples):
            embeddings = None

            try:
                patches = foup.parse_patches(
                    sample, patches_field, handle_missing=handle_missing
                )

                if patches is not None:
                    img = etai.read(sample.filepath)

                    if batch_size is None:
                        embeddings = _embed_patches_single(
                            model, img, patches, force_square, alpha
                        )
                    else:
                        embeddings = _embed_patches_batch(
                            model,
                            img,
                            patches,
                            force_square,
                            alpha,
                            batch_size,
                        )

            except Exception as e:
                if not skip_failures:
                    raise e

                logger.warning("Sample: %s\nError: %s\n", sample.id, e)

            if embeddings_field:
                sample[embeddings_field] = embeddings
                sample.save()
            else:
                embeddings_dict[sample.id] = embeddings

    if embeddings_field:
        return None

    return embeddings_dict


def _embed_patches_single(model, img, detections, force_square, alpha):
    embeddings = []
    for detection in detections.detections:
        patch = foup.extract_patch(
            img, detection, force_square=force_square, alpha=alpha
        )
        embedding = model.embed(patch)
        embeddings.append(embedding)

    return np.concatenate(embeddings)


def _embed_patches_batch(
    model, img, detections, force_square, alpha, batch_size
):
    embeddings = []
    for detection_batch in fou.iter_batches(detections.detections, batch_size):
        patches = [
            foup.extract_patch(
                img, detection, force_square=force_square, alpha=alpha
            )
            for detection in detection_batch
        ]
        embeddings_batch = model.embed_all(patches)
        embeddings.append(embeddings_batch)

    return np.concatenate(embeddings)


def _embed_patches_data_loader(
    samples,
    model,
    patches_field,
    embeddings_field,
    force_square,
    alpha,
    handle_missing,
    batch_size,
    num_workers,
    skip_failures,
):
    samples = samples.select_fields(patches_field)
    data_loader = _make_patch_data_loader(
        samples,
        model,
        patches_field,
        force_square,
        alpha,
        handle_missing,
        num_workers,
        skip_failures,
    )

    embeddings_dict = {}

    with fou.ProgressBar(samples) as pb:
        for sample, patches in pb(zip(samples, data_loader)):
            embeddings = None

            try:
                if isinstance(patches, Exception):
                    raise patches

                if patches is not None:
                    embeddings = []
                    for patches_batch in fou.iter_slices(patches, batch_size):
                        embeddings_batch = model.embed_all(patches_batch)
                        embeddings.append(embeddings_batch)

                    embeddings = np.concatenate(embeddings)

            except Exception as e:
                if not skip_failures:
                    raise e

                logger.warning("Sample: %s\nError: %s\n", sample.id, e)

            if embeddings_field:
                sample[embeddings_field] = embeddings
                sample.save()
            else:
                embeddings_dict[sample.id] = embeddings

    if embeddings_field:
        return None

    return embeddings_dict


def _embed_frame_patches(
    samples,
    model,
    patches_field,
    embeddings_field,
    force_square,
    alpha,
    handle_missing,
    batch_size,
    skip_failures,
):
    samples = samples.select_fields(samples._FRAMES_PREFIX + patches_field)
    frame_counts, total_frame_count = _get_frame_counts(samples)
    is_clips = samples._dataset._is_clips

    embeddings_dict = {}

    with fou.ProgressBar(total=total_frame_count) as pb:
        for idx, sample in enumerate(samples):
            if is_clips:
                frames = etaf.FrameRange(*sample.support)
            else:
                frames = None

            frame_embeddings_dict = {}

            try:
                with etav.FFmpegVideoReader(
                    sample.filepath, frames=frames
                ) as video_reader:
                    for img in video_reader:
                        frame_number = video_reader.frame_number
                        frame = sample.frames[frame_number]

                        patches = foup.parse_patches(
                            frame, patches_field, handle_missing=handle_missing
                        )

                        if patches is not None:
                            if batch_size is None:
                                embeddings = _embed_patches_single(
                                    model, img, patches, force_square, alpha
                                )
                            else:
                                embeddings = _embed_patches_batch(
                                    model,
                                    img,
                                    patches,
                                    force_square,
                                    alpha,
                                    batch_size,
                                )

                        if embeddings_field is not None:
                            sample.add_labels(
                                {frame_number: embeddings},
                                embeddings_field,
                            )
                        else:
                            frame_embeddings_dict[frame_number] = embeddings

                        pb.update()

            except Exception as e:
                if not skip_failures:
                    raise e

                logger.warning("Sample: %s\nError: %s\n", sample.id, e)

            if embeddings_field is None:
                embeddings_dict[sample.id] = frame_embeddings_dict

            # Explicitly set in case actual # frames differed from expected #
            pb.set_iteration(frame_counts[idx])

    if embeddings_field:
        return None

    return embeddings_dict


def _make_patch_data_loader(
    samples,
    model,
    patches_field,
    force_square,
    alpha,
    handle_missing,
    num_workers,
    skip_failures,
):
    # This function supports DataLoaders that emit numpy arrays that can
    # therefore be used for non-Torch models; but we do not currenly use this
    # functionality
    use_numpy = not isinstance(model, TorchModelMixin)

    if num_workers is None:
        num_workers = fout.recommend_num_workers()

    dataset = fout.TorchImagePatchesDataset(
        samples=samples,
        patches_field=patches_field,
        handle_missing=handle_missing,
        transform=model.transforms,
        ragged_batches=model.ragged_batches,
        use_numpy=use_numpy,
        force_rgb=True,
        force_square=force_square,
        alpha=alpha,
        skip_failures=skip_failures,
    )

    return tud.DataLoader(
        dataset,
        batch_size=1,
        num_workers=num_workers,
        collate_fn=lambda batch: batch[0],  # return patches directly
    )


def _parse_batch_size(batch_size, model, use_data_loader):
    if batch_size is None:
        batch_size = fo.config.default_batch_size

    if batch_size is not None and batch_size > 1 and model.ragged_batches:
        logger.warning("Model does not support batching")
        batch_size = None

    if use_data_loader and batch_size is None:
        batch_size = 1

    return batch_size


def load_model(model_config_dict, model_path=None, **kwargs):
    """Loads the model specified by the given :class:`ModelConfig` dict.

    Args:
        model_config_dict: a :class:`ModelConfig` dict
        model_path (None): an optional model path to inject into the
            ``model_path`` field of the model's ``Config`` instance, which must
            implement the ``eta.core.learning.HasPublishedModel`` interface.
            This is useful when working with a model whose weights are stored
            locally and do not need to be downloaded
        **kwargs: optional keyword arguments to inject into the model's
            ``Config`` instance

    Returns:
        a :class:`Model` instance
    """
    # Inject config args
    if kwargs:
        if model_config_dict["type"] == etau.get_class_name(foue.ETAModel):
            _merge_config(model_config_dict["config"]["config"], kwargs)
        else:
            _merge_config(model_config_dict["config"], kwargs)

    # Load model config
    config = ModelConfig.from_dict(model_config_dict)

    #
    # Inject model path
    #
    # Models must be implemented in one of the following ways in order for
    # us to know how to inject ``model_path``:
    #
    # (1) Their config implements ``eta.core.learning.HasPublishedModel``
    #
    # (2) Their config is an ``fiftyone.utils.eta.ETAModelConfig`` whose
    #     embedded config implements ``eta.core.learning.HasPublishedModel``
    #
    if model_path:
        if isinstance(config.config, etal.HasPublishedModel):
            config.config.model_name = None
            config.config.model_path = model_path
        elif isinstance(config.config, foue.ETAModelConfig) and isinstance(
            config.config.config, etal.HasPublishedModel
        ):
            config.config.config.model_name = None
            config.config.config.model_path = model_path
        else:
            raise ValueError(
                "Model config must implement the %s interface"
                % etal.HasPublishedModel
            )

    # Build model
    return config.build()


def _merge_config(d, kwargs):
    for k, v in kwargs.items():
        if k in d and isinstance(d[k], dict):
            d[k].update(v)
        else:
            d[k] = v


class ModelConfig(etal.ModelConfig):
    """Base configuration class that encapsulates the name of a :class:`Model`
    and an instance of its associated Config class.

    Args:
        type: the fully-qualified class name of the :class:`Model` subclass
        config: an instance of the Config class associated with the model
    """

    pass


class Model(etal.Model):
    """Abstract base class for models.

    This class declares the following conventions:

    (a)     :meth:`Model.__init__` should take a single ``config`` argument
            that is an instance of ``<Model>Config``

    (b)     Models implement the context manager interface. This means that
            models can optionally use context to perform any necessary setup
            and teardown, and so any code that builds a model should use the
            ``with`` syntax
    """

    def __enter__(self):
        return self

    def __exit__(self, *args):
        pass

    @property
    def media_type(self):
        """The media type processed by the model.

        Supported values are "image" and "video".
        """
        raise NotImplementedError("subclasses must implement media_type")

    @property
    def has_logits(self):
        """Whether this instance can generate logits for its predictions.

        This method returns ``False`` by default. Methods that can generate
        logits will override this via implementing the
        :class:`LogitsMixin` interface.
        """
        return False

    @property
    def has_embeddings(self):
        """Whether this instance can generate embeddings.

        This method returns ``False`` by default. Methods that can generate
        embeddings will override this via implementing the
        :class:`EmbeddingsMixin` interface.
        """
        return False

    @property
    def ragged_batches(self):
        """True/False whether :meth:`transforms` may return tensors of
        different sizes. If True, then passing ragged lists of data to
        :meth:`predict_all` is not allowed.
        """
        raise NotImplementedError("subclasses must implement ragged_batches")

    @property
    def transforms(self):
        """The preprocessing function that will/must be applied to each input
        before prediction, or ``None`` if no preprocessing is performed.
        """
        raise NotImplementedError("subclasses must implement transforms")

    @property
    def preprocess(self):
        """Whether to apply :meth:`transforms` during inference (True) or to
        assume that they have already been applied (False).
        """
        raise NotImplementedError("subclasses must implement preprocess")

    @preprocess.setter
    def preprocess(self, value):
        raise NotImplementedError("subclasses must implement preprocess")

    def predict(self, arg):
        """Peforms prediction on the given data.

        Image models should support, at minimum, processing ``arg`` values that
        are uint8 numpy arrays (HWC).

        Video models should support, at minimum, processing ``arg`` values that
        are ``eta.core.video.VideoReader`` instances.

        Args:
            arg: the data

        Returns:
            a :class:`fiftyone.core.labels.Label` instance or dict of
            :class:`fiftyone.core.labels.Label` instances containing the
            predictions
        """
        raise NotImplementedError("subclasses must implement predict()")

    def predict_all(self, args):
        """Performs prediction on the given iterable of data.

        Image models should support, at minimum, processing ``args`` values
        that are either lists of uint8 numpy arrays (HWC) or numpy array
        tensors (NHWC).

        Video models should support, at minimum, processing ``args`` values
        that are lists of ``eta.core.video.VideoReader`` instances.

        Subclasses can override this method to increase efficiency, but, by
        default, this method simply iterates over the data and applies
        :meth:`predict` to each.

        Args:
            args: an iterable of data

        Returns:
            a list of :class:`fiftyone.core.labels.Label` instances or a list
            of dicts of :class:`fiftyone.core.labels.Label` instances
            containing the predictions
        """
        return [self.predict(arg) for arg in args]


class LogitsMixin(object):
    """Mixin for :class:`Model` classes that can generate logits for their
    predictions.

    This mixin allows for the possibility that only some instances of a class
    are capable of generating logits, per the value of the
    :meth:`has_logits` property.
    """

    def __init__(self):
        self._store_logits = False

    @property
    def store_logits(self):
        """Whether the model should store logits in its predictions."""
        return self._store_logits

    @store_logits.setter
    def store_logits(self, flag):
        if flag and not self.has_logits:
            raise ValueError("This model cannot generate logits to store")

        self._store_logits = flag

    @property
    def has_logits(self):
        """Whether this instance can generate logits."""
        raise NotImplementedError("subclasses must implement has_logits")


class EmbeddingsMixin(object):
    """Mixin for :class:`Model` classes that can generate embeddings for
    their predictions.

    This mixin allows for the possibility that only some instances of a class
    are capable of generating embeddings, per the value of the
    :meth:`has_embeddings` property.
    """

    @property
    def has_embeddings(self):
        """Whether this instance has embeddings."""
        raise NotImplementedError("subclasses must implement has_embeddings")

    def get_embeddings(self):
        """Returns the embeddings generated by the last forward pass of the
        model.

        By convention, this method should always return an array whose first
        axis represents batch size (which will always be 1 when :meth:`predict`
        was last used).

        Returns:
            a numpy array containing the embedding(s)
        """
        raise NotImplementedError("subclasses must implement get_embeddings()")

    def embed(self, arg):
        """Generates an embedding for the given data.

        Subclasses can override this method to increase efficiency, but, by
        default, this method simply calls :meth:`predict` and then returns
        :meth:`get_embeddings`.

        Args:
            arg: the data. See :meth:`predict` for details

        Returns:
            a numpy array containing the embedding
        """
        # pylint: disable=no-member
        self.predict(arg)
        return self.get_embeddings()

    def embed_all(self, args):
        """Generates embeddings for the given iterable of data.

        Subclasses can override this method to increase efficiency, but, by
        default, this method simply iterates over the data and applies
        :meth:`embed` to each.

        Args:
            args: an iterable of data. See :meth:`predict_all` for details

        Returns:
            a numpy array containing the embeddings stacked along axis 0
        """
        return np.stack([self.embed(arg) for arg in args])


class TorchModelMixin(object):
    """Mixin for :class:`Model` classes that support feeding data for inference
    via a :class:`torch:torch.utils.data.DataLoader`.

    Models implementing this mixin must expose via their
    :meth:`Model.transforms` property the
    :mod:`torchvision:torchvision.transforms` function that will/must be
    applied to each input before prediction.
    """

    pass


class ModelManagerConfig(etam.ModelManagerConfig):
    """Config settings for a :class:`ModelManager`.

    Args:
        url (None): the URL of the file
        google_drive_id (None): the ID of the file in Google Drive
        extract_archive (None): whether to extract the downloaded model, which
            is assumed to be an archive
        delete_archive (None): whether to delete the archive after extracting
            it, if applicable
    """

    def __init__(self, d):
        super().__init__(d)

        self.url = self.parse_string(d, "url", default=None)
        self.google_drive_id = self.parse_string(
            d, "google_drive_id", default=None
        )


class ModelManager(etam.ModelManager):
    """Class for downloading FiftyOne models from the web."""

    @staticmethod
    def upload_model(model_path, *args, **kwargs):
        raise NotImplementedError("Uploading models via API is not supported")

    def _download_model(self, model_path):
        if self.config.google_drive_id:
            gid = self.config.google_drive_id
            logger.info("Downloading model from Google Drive ID '%s'...", gid)
            etaw.download_google_drive_file(gid, path=model_path)
        elif self.config.url:
            url = self.config.url
            logger.info("Downloading model from '%s'...", url)
            etaw.download_file(url, path=model_path)
        else:
            raise ValueError(
                "Invalid ModelManagerConfig '%s'" % str(self.config)
            )

    def delete_model(self):
        raise NotImplementedError("Deleting models via API is not supported")
