import logging
import os
import json
from pathlib import Path
from typing import List

import torch
import numpy

from towhee import register
from towhee.operator.base import NNOperator
from towhee.types.video_frame import VideoFrame
from towhee.models.utils.video_transforms import get_configs, transform_video
from towhee.models.uniformer.uniformer import create_model

log = logging.getLogger()

@register(output_schema=['labels', 'scores', 'features'])
class Uniformer(NNOperator):
    """
    Generate a list of class labels given a video input data.
    Default labels are from [Kinetics400 Dataset](https://deepmind.com/research/open-source/kinetics).
    Args:
        model_name (`str`):
            Supported model names:
            - uniformer_k400_s8
            - uniformer_k400_s16
            - uniformer_k400_b8
            - uniformer_k400_b16
        skip_preprocess (`str`):
            Flag to skip video transforms.
        predict (`bool`):
            Flag to control whether predict labels. If False, then return video embedding.
        classmap (`dict=None`):
            The dictionary maps classes to integers.
        topk (`int=5`):
            The number of classification labels to be returned (ordered by possibility from high to low).
    """
    def __init__(self,
                 model_name: str = 'uniformer_k400_s8',
                 framework: str = 'pytorch',
                 skip_preprocess: bool = False,
                 classmap: dict = None,
                 topk: int = 5,
                 ):
        super().__init__(framework=framework)
        self.model_name = model_name
        self.skip_preprocess = skip_preprocess
        self.topk = topk
        if 'k400' in model_name:
            self.dataset_name = 'kinetics_400'
        if classmap is None:
            class_file = os.path.join(str(Path(__file__).parent), self.dataset_name+'.txt')
            self.classmap = {}
            with open(class_file, "r") as f:
                for line in f.readlines():
                    line = line.split('\t')
                    k = str(line[0])
                    v = int(line[1])
                    self.classmap[v] = str(k).replace('"', '')
        else:
            self.classmap = classmap
        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
        
        self.input_mean = [0.45, 0.45, 0.45]
        self.input_std = [0.225, 0.225, 0.225]

        model_path = str(Path(__file__).parent)
        model_path = {
            'uniformer_k400_s8': os.path.join(model_path, 'uniformer_small_k400_8x8.pth'),
            'uniformer_k400_s16': os.path.join(model_path, 'uniformer_small_k400_16x4.pth'),
            'uniformer_k400_b8': os.path.join(model_path, 'uniformer_base_k400_8x8.pth'),
            'uniformer_k400_b16': os.path.join(model_path, 'uniformer_base_k400_16x4.pth'),
        }
        self.weights_path = model_path[model_name]
        self.model = create_model(model_name=model_name, pretrained=True, weights_path=self.weights_path, device=self.device)
        if '8' in model_name:
            self.transform_cfgs = get_configs(
                    side_size=256,
                    crop_size=224,
                    num_frames=8,
                    mean=self.input_mean,
                    std=self.input_std,
                    )
        elif '16' in model_name:
            self.transform_cfgs = get_configs(
                    side_size=256,
                    crop_size=224,
                    num_frames=16,
                    mean=self.input_mean,
                    std=self.input_std,
                    )
        self.model.eval()

    def __call__(self, video: List[VideoFrame]):
        """
        Args:
            video (`List[VideoFrame]`):
                Video path in string.

        Returns:
            (labels, scores)
                A tuple of lists (labels, scores).
            OR emb
                Video embedding.
        """
        # Convert list of towhee.types.Image to numpy.ndarray in float32
        video = numpy.stack([img.astype(numpy.float32)/255. for img in video], axis=0)
        assert len(video.shape) == 4
        video = video.transpose(3, 0, 1, 2)  # twhc -> ctwh

        # Transform video data given configs
        if self.skip_preprocess:
            self.transform_cfgs.update(num_frames=None)

        data = transform_video(
                    video=video,
                    **self.transform_cfgs
                    )
        inputs = data.to(self.device)[None, ...]

        feats = self.model.forward_features(inputs)
        feats = feats.flatten(2).mean(-1)
        features = feats.to('cpu').squeeze(0).detach().numpy()

        outs = self.model.head(feats)
        post_act = torch.nn.Softmax(dim=1)
        preds = post_act(outs)
        pred_scores, pred_classes = preds.topk(k=self.topk)
        labels = [self.classmap[int(i)] for i in pred_classes[0]]
        scores = [round(float(x), 5) for x in pred_scores[0]]

        return labels, scores, features