#  Tencent is pleased to support the open source community by making GNES available.
#
#  Copyright (C) 2019 THL A29 Limited, a Tencent company. All rights reserved.
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#  http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.

from typing import List

import numpy as np
import random
from .base import BaseVideoPreprocessor
from ..helper import get_video_frames, phash_descriptor
from ...proto import gnes_pb2, array2blob


class FFmpegPreprocessor(BaseVideoPreprocessor):

    def __init__(self,
                 duplicate_rm: bool = True,
                 use_phash_weight: bool = False,
                 phash_thresh: int = 5,
                 *args,
                 **kwargs):
        super().__init__(*args, **kwargs)
        self.phash_thresh = phash_thresh
        self.duplicate_rm = duplicate_rm
        self.use_phash_weight = use_phash_weight

        self._ffmpeg_kwargs = kwargs

    def apply(self, doc: 'gnes_pb2.Document') -> None:
        super().apply(doc)

        # video could't be processed from ndarray!
        # only bytes can be passed into ffmpeg pipeline
        if doc.raw_bytes:
            frames = get_video_frames(doc.raw_bytes, **self._ffmpeg_kwargs)

            # remove dupliated key frames by phash value
            if self.duplicate_rm:
                frames = self.duplicate_rm_hash(frames)

            if self.use_phash_weight:
                weight = FFmpegPreprocessor.pic_weight(frames)
            else:
                weight = [1 / len(frames)] * len(frames)

            for ci, chunk in enumerate(frames):
                c = doc.chunks.add()
                c.doc_id = doc.doc_id
                c.blob.CopyFrom(array2blob(chunk))
                c.offset_1d = ci
                c.weight = weight[ci]

        else:
            self.logger.error('bad document: "raw_bytes" is empty!')

    @staticmethod
    def pic_weight(images: List['np.ndarray']) -> List[float]:
        import cv2
        weight = np.zeros([len(images)])
        # n_channel is usually 3 for RGB images
        n_channel = images[0].shape[-1]
        for i, image in enumerate(images):
            weight[i] = sum([
                cv2.calcHist([image], [_], None, [256], [0, 256]).var()
                for _ in range(n_channel)
            ])
        weight = weight / weight.sum()

        # normalized result
        weight = np.exp(-weight * 10)
        return weight / weight.sum()

    def duplicate_rm_hash(self,
                          images: List['np.ndarray']) -> List['np.ndarray']:
        hash_list = [phash_descriptor(_) for _ in images]
        ret = []
        for i, h in enumerate(hash_list):
            flag = 1
            if len(ret) >= 1:
                # only keep images with high phash diff
                # comparing only last kept 9 pics
                for j in range(1, min(len(ret) + 1, 9)):
                    dist = abs(ret[-j][1] - h)
                    if dist < self.phash_thresh:
                        flag = 0
                        break

            if flag:
                ret.append((i, h))

        return [images[_[0]] for _ in ret]


class FFmpegVideoSegmentor(BaseVideoPreprocessor):
    def __init__(self,
                 segment_method: str = 'cut_by_frame',
                 segment_interval: int = -1,
                 segment_num: int = 3,
                 max_frames_per_doc: int = -1,
                 audio_interval: int = 30,
                 sample_rate: int = 16000,
                 *args,
                 **kwargs):
        super().__init__(*args, **kwargs)
        self.segment_method = segment_method
        self.segment_interval = segment_interval
        self.segment_num = segment_num
        self.max_frames_per_doc = max_frames_per_doc
        self.audio_interval = audio_interval
        self.sample_rate = sample_rate
        self._ffmpeg_kwargs = kwargs

    def apply(self, doc: 'gnes_pb2.Document') -> None:
        super().apply(doc)
        if doc.raw_bytes:
            frames = get_video_frames(doc.raw_bytes, **self._ffmpeg_kwargs)
            if self.max_frames_per_doc > 0:
                random_id = random.sample(range(len(frames)),
                                          k=min(self.max_frames_per_doc, len(frames)))
                frames = [frames[i] for i in sorted(random_id)]

            sub_videos = []
            if len(frames) >= 1:
                # cut by frame: should specify how many frames to cut
                if self.segment_method == 'cut_by_frame':
                    if self.segment_interval == -1:
                        sub_videos = [frames]
                    else:
                        sub_videos = [frames[_: _ + self.segment_interval]
                                      for _ in range(0, len(frames), self.segment_interval)]
                # cut by num: should specify how many chunks for each doc
                elif self.segment_method == 'cut_by_num':
                    if self.segment_num >= 2:
                        _interval = int(len(frames) / self.segment_num)
                        sub_videos = [frames[_: _ + _interval]
                                      for _ in range(0, len(frames), _interval)]
                    else:
                        sub_videos = [frames]

                for ci, chunk in enumerate(sub_videos):
                    c = doc.chunks.add()
                    c.doc_id = doc.doc_id
                    c.blob.CopyFrom(array2blob(np.array(chunk, dtype=np.uint8)))
                    c.offset_1d = ci
                    c.weight = 1 / len(sub_videos)

            else:
                self.logger.info('bad document: no key frames extracted')
        else:
            self.logger.error('bad document: "raw_bytes" is empty!')
