"""
Minimal loader for the STP release on HuggingFace.

    import torch
    from loader import STPEncoder

    enc = STPEncoder.from_pretrained('yangjiange/STP')
    feats = enc.encode(images)      # (N, 3, 224, 224) -> (N, 768)

Or straight from the hub without this file:

    import torch, models_stp
    from huggingface_hub import hf_hub_download

    path = hf_hub_download('yangjiange/STP', 'stp_vitb.pth')
    model = models_stp.mae_vit_base_patch16()
    from util.ckpt_io import load_model_state
    state, _ = load_model_state(path)
    model.load_state_dict(state, strict=True)
    model.eval()
"""
import os

import numpy as np
import torch
from PIL import Image

IMAGENET_MEAN = (0.485, 0.456, 0.406)
IMAGENET_STD = (0.229, 0.224, 0.225)


class STPEncoder:
    """Frozen STP encoder: images in, [CLS] representations out.

    Only the encoder is used. The spatial and temporal decoders exist for
    pre-training and are irrelevant once you are training a policy.
    """

    def __init__(self, model, device='cpu', image_size=224):
        self.model = model
        self.device = device
        self.image_size = image_size

    # ------------------------------------------------------------------
    @classmethod
    def from_pretrained(cls, repo_id_or_path='yangjiange/STP', filename='stp_vitb.pth',
                        device=None, repo_root=None):
        import models_stp

        device = device or ('cuda' if torch.cuda.is_available() else 'cpu')

        if os.path.isfile(repo_id_or_path):
            path = repo_id_or_path
        else:
            from huggingface_hub import hf_hub_download
            path = hf_hub_download(repo_id=repo_id_or_path, filename=filename,
                                   cache_dir=None)

        import sys as _sys, os as _os
        _root = repo_root or _os.path.dirname(_os.path.dirname(_os.path.abspath(__file__)))
        if _root not in _sys.path:
            _sys.path.insert(0, _root)
        from util.ckpt_io import load_model_state
        state, _ = load_model_state(path)

        if repo_root and repo_root not in os.sys.path:
            os.sys.path.insert(0, repo_root)

        model = models_stp.mae_vit_base_patch16()
        model.load_state_dict(state, strict=True)
        model.eval().to(device)
        return cls(model, device)

    # ------------------------------------------------------------------
    def preprocess(self, img):
        if not isinstance(img, Image.Image):
            img = Image.fromarray(np.asarray(img))
        if img.size != (self.image_size, self.image_size):
            img = img.resize((self.image_size, self.image_size), Image.BICUBIC)
        x = torch.from_numpy(np.asarray(img.convert('RGB'), dtype=np.float32) / 255.0)
        x = x.permute(2, 0, 1)
        mean = torch.tensor(IMAGENET_MEAN).view(3, 1, 1)
        std = torch.tensor(IMAGENET_STD).view(3, 1, 1)
        return (x - mean) / std

    @torch.no_grad()
    def encode(self, images, batch_size=32):
        """PIL images / numpy arrays / a tensor -> (N, 768) numpy array."""
        if isinstance(images, torch.Tensor):
            out = []
            for i in range(0, len(images), batch_size):
                out.append(self.model.forward_features(
                    images[i:i + batch_size].to(self.device)).float().cpu())
            return torch.cat(out, 0).numpy()

        if isinstance(images, Image.Image):
            images = [images]
        feats = []
        for i in range(0, len(images), batch_size):
            batch = torch.stack([self.preprocess(im) for im in images[i:i + batch_size]])
            feats.append(self.model.forward_features(batch.to(self.device)).float().cpu())
        return torch.cat(feats, 0).numpy()