from typing import Any, Callable, Dict

from torchvision import transforms as T

from rikai.types.vision import Image

"""
Adapted from https://github.com/pytorch/pytorch.github.io/blob/site/assets/hub/pytorch_vision_resnet.ipynb
"""  # noqa E501


def pre_processing(options: Dict[str, Any]) -> Callable:
    """
    All pre-trained models expect input images normalized in the same way, i.e.
    mini-batches of 3-channel RGB images of shape (3 x H x W), where H and W
    are expected to be at least 224. The images have to be loaded in to a range
    of [0, 1] and then normalized using mean = [0.485, 0.456, 0.406] and std =
    [0.229, 0.224, 0.225].
    """
    return T.Compose(
        [
            T.Resize(256),
            T.CenterCrop(224),
            T.ToTensor(),
            T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        ]
    )


def post_processing(options: Dict[str, Any]) -> Callable:
    def post_process_func(batch):
        results = []
        for result in batch:
            results.append(result.detach().cpu().tolist())
        return results

    return post_process_func


OUTPUT_SCHEMA = "array<float>"
