import random

import torch
from torch import nn
import torch.nn.functional as F
import ffmpeg
import numpy as np
import cv2
from moviepy.editor import VideoFileClip

from .utils import get_frames


class TransNetV2(nn.Module):
    def __init__(self, F=16, L=3, S=2, D=1024):
        super(TransNetV2, self).__init__()
        self.SDDCNN = nn.ModuleList(
            [
                StackedDDCNNV2(
                    in_filters=3, n_blocks=S, filters=F, stochastic_depth_drop_prob=0.0
                )
            ]
            + [
                StackedDDCNNV2(
                    in_filters=(F * 2 ** (i - 1)) * 4, n_blocks=S, filters=F * 2**i
                )
                for i in range(1, L)
            ]
        )

        # 帧相似网络
        self.frame_sim_layer = FrameSimilarity(
            sum([(F * 2**i) * 4 for i in range(L)]),
            lookup_window=101,
            output_dim=128,
            similarity_dim=128,
            use_bias=True,
        )

        # 颜色相似网络
        self.color_hist_layer = ColorHistograms(lookup_window=101, output_dim=128)

        # dropout
        self.dropout = nn.Dropout(0.5)

        output_dim = ((F * 2 ** (L - 1)) * 4) * 3 * 6  #
        output_dim = output_dim + 128  # 使用了帧相似网络, 维度需要加128
        output_dim = output_dim + 128  # 使用了颜色相似网络, 维度需要再加128

        self.fc1 = nn.Linear(output_dim, D)
        self.cls_layer1 = nn.Linear(D, 1)
        self.cls_layer2 = nn.Linear(D, 1)

    def forward(self, inputs):
        # 输入必须为torch.uint8, (h,w)=(27,48)的图片batch样本
        #         assert isinstance(inputs, torch.Tensor) and list(inputs.shape[2:]) == [27, 48, 3] and inputs.dtype == torch.uint8, "incorrect input type and/or shape"

        # uint8 of shape [B, T, H, W, 3] to float of shape [B, 3, T, H, W]
        with torch.autograd.set_detect_anomaly(True):
            x = inputs.permute([0, 4, 1, 2, 3]).float()
            x = x.div_(255.0)

            # 收集每一层的SDDCNN特征图
            block_features = []
            for block in self.SDDCNN:
                x = block(x)
                block_features.append(x)

            x = x.permute(0, 2, 3, 4, 1)  # 把维度从[B, 通道数, T, H, W] 转化为 [B, T, H, W, 通道数]
            x = x.reshape(x.shape[0], x.shape[1], -1)

            x = torch.cat(
                [self.frame_sim_layer(block_features), x], 2
            )  # 在最后一维度cat上block_features输出的特征
            x = torch.cat(
                [self.color_hist_layer(inputs), x], 2
            )  # 在最后一维度cat上color_hist_layer输出的特征

            x = F.relu(self.fc1(x))
            x = self.dropout(x)

            one_hot = self.cls_layer1(x)
            many_hot = self.cls_layer2(x)
            return one_hot, many_hot

    # 预测MP4文件转换帧，并给出对应帧位置
    def predict_video(
        self,
        mp4_file,
        cache_path="",
        c_box=None,
        width=48,
        height=27,
        input_frames=100,
        overlap=30,
        sample_fps=30,
        threshold=0.3,
    ):
        """
        mp4_file: ~/6712566330782010632.mp4
        cache_path: ~/视频单帧数据_h48_w27
        return: [x,x,...] 点位时间
        """
        assert overlap % 2 == 0
        assert input_frames > overlap
        # fps = eval(ffmpeg.probe(mp4_file)['streams'][0]['r_frame_rate']) # 获取视频的视频帧率
        # total_frames = int(ffmpeg.probe(mp4_file)['streams'][0]['nb_frames']) # 获取视频的总帧数
        # duration = float(ffmpeg.probe(mp4_file)['streams'][0]['duration']) # 获取视频的总时长
        video = VideoFileClip(mp4_file)
        # video = video.subclip(0, 60 * 10)
        fps = video.fps
        duration = video.duration
        total_frames = int(duration * fps)
        w, h = video.size
        print(fps, duration, total_frames, w, h)

        if c_box:
            video.crop(*c_box)

        frame_iter = video.iter_frames(fps=sample_fps)
        sample_total_frames = int(sample_fps * duration)
        frame_list = []
        for i in range(sample_total_frames // (input_frames - overlap) + 1):
            # if i==1:
            #     break
            frame_list = frame_list[-overlap:]
            start_frame = i * (input_frames - overlap)
            end_frame = min(start_frame + input_frames, sample_total_frames)
            print("start_frame & end_frame: ", start_frame, end_frame)
            for frame in frame_iter:
                frame = cv2.resize(frame, (width, height))
                frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
                frame_list.append(frame)
                if len(frame_list) == end_frame - start_frame:
                    break
            frames = torch.Tensor(frame_list)  # 获得帧
            if frames.shape[0] < end_frame - start_frame:
                # 原视频的视频时长比音频时长短，体现出来的是原视频最后有声音没画面
                print(
                    "total_frames is wrong: ",
                    total_frames,
                    "-->",
                    start_frame + frames.shape[0],
                )
                # sample_total_frames = start_frame + frames.shape[0]
                # fps = total_frames / duration
            frames = frames.cuda()
            # single_frame_pred和all_frame_pred都是输出window_size长的是否转场概率，
            single_frame_pred, all_frame_pred = self.forward(
                frames.unsqueeze(0)
            )  # 前向推理
            # single_frame_pred = F.softmax(single_frame_pred, dim=-1) # 获得每一帧对应的类别概率
            # single_frame_pred = torch.argmax(single_frame_pred, dim=-1).reshape(-1)
            single_frame_pred = torch.sigmoid(single_frame_pred).reshape(-1)
            all_frame_pred = torch.sigmoid(all_frame_pred).reshape(-1)

            # single_frame_pred = (single_frame_pred>threshold)*1
            if total_frames > end_frame:
                if i == 0:
                    single_frame_pred_label = single_frame_pred[: -overlap // 2]
                    all_frame_pred_label = all_frame_pred[: -overlap // 2]
                else:
                    single_frame_pred_label = torch.cat(
                        (
                            single_frame_pred_label,
                            single_frame_pred[overlap // 2 : -overlap // 2],
                        ),
                        dim=0,
                    )
                    all_frame_pred_label = torch.cat(
                        (
                            all_frame_pred_label,
                            all_frame_pred[overlap // 2 : -overlap // 2],
                        ),
                        dim=0,
                    )
            else:
                if i == 0:
                    single_frame_pred_label = single_frame_pred
                    all_frame_pred_label = all_frame_pred
                else:
                    single_frame_pred_label = torch.cat(
                        (single_frame_pred_label, single_frame_pred[overlap // 2 :]),
                        dim=0,
                    )
                    all_frame_pred_label = torch.cat(
                        (all_frame_pred_label, all_frame_pred[overlap // 2 :]), dim=0
                    )
                break

        single_frame_pred_label = single_frame_pred_label.cpu().numpy()
        all_frame_pred_label = all_frame_pred_label.cpu().numpy()

        return (
            single_frame_pred_label,
            all_frame_pred_label,
            fps,
            total_frames,
            duration,
            h,
            w,
        )

        # transition_index = torch.where(pred_label==1)[0].cpu().numpy() # 转场帧位置
        # transition_index = transition_index.astype(np.float)
        # # 对返回结果做后处理合并相邻帧
        # result_transition = []
        # for i, transition in enumerate(transition_index):
        #     if i == 0:
        #         result_transition.append([transition])
        #     else:
        #         if abs(result_transition[-1][-1]-transition) == 1:
        #             result_transition[-1].append(transition)
        #         else:
        #             result_transition.append([transition])
        #
        # result_transition = [[0]] + [[item[0], item[-1]] if len(item)>1 else [item[0]] for item in result_transition] + [[total_frames]]
        #
        # return result_transition, fps, total_frames, duration, h, w

    def predict_video_2(
        self,
        mp4_file,
        cache_path="",
        c_box=None,
        width=48,
        height=27,
        input_frames=100,
        overlap=30,
        sample_fps=30,
        threshold=0.3,
    ):
        """
        mp4_file: ~/6712566330782010632.mp4
        cache_path: ~/视频单帧数据_h48_w27
        return: [x,x,...] 点位时间
        """
        assert overlap % 2 == 0
        assert input_frames > overlap
        # fps = eval(ffmpeg.probe(mp4_file)['streams'][0]['r_frame_rate']) # 获取视频的视频帧率
        # total_frames = int(ffmpeg.probe(mp4_file)['streams'][0]['nb_frames']) # 获取视频的总帧数
        # duration = float(ffmpeg.probe(mp4_file)['streams'][0]['duration']) # 获取视频的总时长
        video = VideoFileClip(mp4_file)
        # video = video.subclip(0, 60 * 10)
        fps = video.fps
        duration = video.duration
        total_frames = int(duration * fps)
        w, h = video.size
        print(fps, duration, total_frames, w, h)

        if c_box:
            video.crop(*c_box)

        frame_iter = video.iter_frames(fps=sample_fps)
        sample_total_frames = int(sample_fps * duration)
        frame_list = []
        for i in range(sample_total_frames // (input_frames - overlap) + 1):
            # if i==1:
            #     break
            frame_list = frame_list[-overlap:]
            start_frame = i * (input_frames - overlap)
            end_frame = min(start_frame + input_frames, sample_total_frames)
            print("start_frame & end_frame: ", start_frame, end_frame)
            for frame in frame_iter:
                frame = cv2.resize(frame, (width, height))
                frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
                frame_list.append(frame)
                if len(frame_list) == end_frame - start_frame:
                    break
            frames = torch.Tensor(frame_list)  # 获得帧
            if frames.shape[0] < end_frame - start_frame:
                # 原视频的视频时长比音频时长短，体现出来的是原视频最后有声音没画面
                print(
                    "total_frames is wrong: ",
                    total_frames,
                    "-->",
                    start_frame + frames.shape[0],
                )
                # sample_total_frames = start_frame + frames.shape[0]
                # fps = total_frames / duration
            frames = frames.cuda()
            single_frame_pred, all_frame_pred = self.forward(
                frames.unsqueeze(0)
            )  # 前向推理
            # single_frame_pred = F.softmax(single_frame_pred, dim=-1) # 获得每一帧对应的类别概率
            # single_frame_pred = torch.argmax(single_frame_pred, dim=-1).reshape(-1)
            single_frame_pred = torch.sigmoid(single_frame_pred).reshape(-1)
            all_frame_pred = torch.sigmoid(all_frame_pred).reshape(-1)

            # single_frame_pred = (single_frame_pred>threshold)*1
            if total_frames > end_frame:
                if i == 0:
                    single_frame_pred_label = single_frame_pred[: -overlap // 2]
                    all_frame_pred_label = all_frame_pred[: -overlap // 2]
                else:
                    single_frame_pred_label = torch.cat(
                        (
                            single_frame_pred_label,
                            single_frame_pred[overlap // 2 : -overlap // 2],
                        ),
                        dim=0,
                    )
                    all_frame_pred_label = torch.cat(
                        (
                            all_frame_pred_label,
                            all_frame_pred[overlap // 2 : -overlap // 2],
                        ),
                        dim=0,
                    )
            else:
                if i == 0:
                    single_frame_pred_label = single_frame_pred
                    all_frame_pred_label = all_frame_pred
                else:
                    single_frame_pred_label = torch.cat(
                        (single_frame_pred_label, single_frame_pred[overlap // 2 :]),
                        dim=0,
                    )
                    all_frame_pred_label = torch.cat(
                        (all_frame_pred_label, all_frame_pred[overlap // 2 :]), dim=0
                    )
                break

        single_frame_pred_label = single_frame_pred_label.cpu().numpy()
        all_frame_pred_label = all_frame_pred_label.cpu().numpy()

        return (
            single_frame_pred_label,
            all_frame_pred_label,
            fps,
            total_frames,
            duration,
            h,
            w,
        )


class StackedDDCNNV2(nn.Module):
    def __init__(
        self,
        in_filters,
        n_blocks,
        filters,
        shortcut=True,
        pool_type="avg",
        stochastic_depth_drop_prob=0.0,
    ):
        super(StackedDDCNNV2, self).__init__()

        self.shortcut = shortcut
        # 定义DDCNN层
        self.DDCNN = nn.ModuleList(
            [
                DilatedDCNNV2(
                    in_filters if i == 1 else filters * 4,
                    filters,
                    activation=F.relu if i != n_blocks else None,
                )
                for i in range(1, n_blocks + 1)
            ]
        )  # 有n_blocks层数量的DilateDCNNV2模块

        # 定义pool层
        self.pool = (
            nn.MaxPool3d(kernel_size=(1, 2, 2))
            if pool_type == "max"
            else nn.AvgPool3d(kernel_size=(1, 2, 2))
        )
        self.stochastic_depth_drop_prob = stochastic_depth_drop_prob

    def forward(self, inputs):
        x = inputs
        shortcut = None

        # DDCNN层前向传播
        for block in self.DDCNN:
            x = block(x)
            if shortcut is None:  # 记录第一层的结果作为残差连接
                shortcut = x

        x = F.relu(x)
        if self.shortcut is not None:
            if self.stochastic_depth_drop_prob != 0.0:
                if self.training:
                    if random.random() < self.stochastic_depth_drop_prob:
                        x = shortcut
                    else:
                        x = x + shortcut
                else:
                    x = (1 - self.stochastic_depth_drop_prob) * x + shortcut
            else:
                x = x + shortcut

        x = self.pool(x)
        return x


class DilatedDCNNV2(nn.Module):
    def __init__(self, in_filters, filters, batch_norm=True, activation=None):
        super(DilatedDCNNV2, self).__init__()

        self.Conv3D_1 = Conv3DConfigurable(
            in_filters, filters, 1, use_bias=not batch_norm
        )
        self.Conv3D_2 = Conv3DConfigurable(
            in_filters, filters, 2, use_bias=not batch_norm
        )
        self.Conv3D_4 = Conv3DConfigurable(
            in_filters, filters, 4, use_bias=not batch_norm
        )
        self.Conv3D_8 = Conv3DConfigurable(
            in_filters, filters, 8, use_bias=not batch_norm
        )

        self.bn = nn.BatchNorm3d(filters * 4, eps=1e-3) if batch_norm else None
        self.activation = activation  # 激活函数定义

    def forward(self, inputs):
        conv1 = self.Conv3D_1(inputs)
        conv2 = self.Conv3D_2(inputs)
        conv3 = self.Conv3D_4(inputs)
        conv4 = self.Conv3D_8(inputs)

        x = torch.cat([conv1, conv2, conv3, conv4], dim=1)

        if self.bn is not None:
            x = self.bn(x)

        if self.activation is not None:
            x = self.activation(x)

        return x


class Conv3DConfigurable(nn.Module):
    def __init__(
        self, in_filters, filters, dilation_rate, separable=True, use_bias=True
    ):
        super(Conv3DConfigurable, self).__init__()

        if separable:
            # (2+1)D convolution https://arxiv.org/pdf/1711.11248.pdf
            conv1 = nn.Conv3d(
                in_filters,
                2 * filters,
                kernel_size=(1, 3, 3),
                dilation=(1, 1, 1),
                padding=(0, 1, 1),
                bias=False,
            )
            conv2 = nn.Conv3d(
                2 * filters,
                filters,
                kernel_size=(3, 1, 1),
                dilation=(dilation_rate, 1, 1),
                padding=(dilation_rate, 0, 0),
                bias=use_bias,
            )
            self.layers = nn.ModuleList([conv1, conv2])
        else:
            conv = nn.Conv3d(
                in_filters,
                filters,
                kernel_size=3,
                dilation=(dilation_rate, 1, 1),
                padding=(dilation_rate, 1, 1),
                bias=use_bias,
            )
            self.layers = nn.ModuleList([conv])

    def forward(self, inputs):
        x = inputs
        for layer in self.layers:
            x = layer(x)
        return x


# 帧相似网络构建
class FrameSimilarity(nn.Module):
    def __init__(
        self,
        in_filters,
        similarity_dim=128,
        lookup_window=101,
        output_dim=128,
        use_bias=False,
    ):
        super(FrameSimilarity, self).__init__()

        self.projection = nn.Linear(in_filters, similarity_dim, bias=use_bias)
        self.fc = nn.Linear(lookup_window, output_dim)

        self.lookup_window = lookup_window
        assert lookup_window % 2 == 1, "`lookup_window` must be odd integer"

    def forward(self, inputs):
        x = torch.cat([torch.mean(x, dim=[3, 4]) for x in inputs], dim=1)
        x = torch.transpose(x, 1, 2)

        x = self.projection(x)
        x = F.normalize(x, p=2, dim=2)

        batch_size, time_window = x.shape[0], x.shape[1]
        similarities = torch.bmm(
            x, x.transpose(1, 2)
        )  # [batch_size, time_window, time_window]余弦相似度
        similarities_padded = F.pad(
            similarities, [(self.lookup_window - 1) // 2, (self.lookup_window - 1) // 2]
        )

        batch_indices = (
            torch.arange(0, batch_size, device=x.device)
            .view([batch_size, 1, 1])
            .repeat([1, time_window, self.lookup_window])
        )
        time_indices = (
            torch.arange(0, time_window, device=x.device)
            .view([1, time_window, 1])
            .repeat([batch_size, 1, self.lookup_window])
        )
        lookup_indices = (
            torch.arange(0, self.lookup_window, device=x.device)
            .view([1, 1, self.lookup_window])
            .repeat([batch_size, time_window, 1])
            + time_indices
        )

        similarities = similarities_padded[batch_indices, time_indices, lookup_indices]
        return F.relu(self.fc(similarities))


# 颜色相似网络
class ColorHistograms(nn.Module):
    def __init__(self, lookup_window=101, output_dim=None):
        super(ColorHistograms, self).__init__()

        self.fc = (
            nn.Linear(lookup_window, output_dim) if output_dim is not None else None
        )
        self.lookup_window = lookup_window
        assert lookup_window % 2 == 1, "`lookup_window` must be odd integer"

    @staticmethod
    def compute_color_histograms(frames):
        frames = frames.int()

        def get_bin(frames):
            # returns 0 .. 511
            R, G, B = frames[:, :, 0], frames[:, :, 1], frames[:, :, 2]
            R, G, B = R >> 5, G >> 5, B >> 5
            return (R << 6) + (G << 3) + B

        batch_size, time_window, height, width, no_channels = frames.shape
        assert no_channels == 3
        frames_flatten = frames.view(batch_size * time_window, height * width, 3)

        binned_values = get_bin(frames_flatten)
        frame_bin_prefix = (
            torch.arange(0, batch_size * time_window, device=frames.device) << 9
        ).view(-1, 1)
        binned_values = (binned_values + frame_bin_prefix).view(-1)

        histograms = torch.zeros(
            batch_size * time_window * 512, dtype=torch.int32, device=frames.device
        )
        histograms.scatter_add_(
            0,
            binned_values,
            torch.ones(len(binned_values), dtype=torch.int32, device=frames.device),
        )

        histograms = histograms.view(batch_size, time_window, 512).float()
        histograms_normalized = F.normalize(histograms, p=2, dim=2)
        return histograms_normalized

    def forward(self, inputs):
        x = self.compute_color_histograms(inputs)

        batch_size, time_window = x.shape[0], x.shape[1]
        similarities = torch.bmm(
            x, x.transpose(1, 2)
        )  # [batch_size, time_window, time_window]
        similarities_padded = F.pad(
            similarities, [(self.lookup_window - 1) // 2, (self.lookup_window - 1) // 2]
        )

        batch_indices = (
            torch.arange(0, batch_size, device=x.device)
            .view([batch_size, 1, 1])
            .repeat([1, time_window, self.lookup_window])
        )
        time_indices = (
            torch.arange(0, time_window, device=x.device)
            .view([1, time_window, 1])
            .repeat([batch_size, 1, self.lookup_window])
        )
        lookup_indices = (
            torch.arange(0, self.lookup_window, device=x.device)
            .view([1, 1, self.lookup_window])
            .repeat([batch_size, time_window, 1])
            + time_indices
        )

        similarities = similarities_padded[batch_indices, time_indices, lookup_indices]

        if self.fc is not None:
            return F.relu(self.fc(similarities))
        return similarities