#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

from typing import Any, Dict

import torch
import torch.nn as nn
from classy_vision.generic.util import is_pos_int, is_pos_int_list

from . import register_model
from .classy_model import ClassyModel, ClassyModelEvaluationMode
from .resnext3d_stage import ResStage
from .resnext3d_stem import R2Plus1DStem, ResNeXt3DStem


model_stems = {
    "r2plus1d_stem": R2Plus1DStem,
    "resnext3d_stem": ResNeXt3DStem,
    # For more types of model stem, add them below
}


class ResNeXt3DBase(ClassyModel):
    def __init__(
        self,
        input_key,
        input_planes,
        clip_crop_size,
        frames_per_clip,
        num_blocks,
        stem_name,
        stem_planes,
        stem_temporal_kernel,
        stem_spatial_kernel,
        stem_maxpool,
    ):
        """
        ResNeXt3DBase implements everything in ResNeXt3D model except the
        construction of 4 stages. See more details in ResNeXt3D.
        """
        super(ResNeXt3DBase, self).__init__()

        self._input_key = input_key
        self.input_planes = input_planes
        self.clip_crop_size = clip_crop_size
        self.frames_per_clip = frames_per_clip
        self.num_blocks = num_blocks

        assert stem_name in model_stems, "unknown stem: %s" % stem_name
        self.stem = model_stems[stem_name](
            stem_temporal_kernel,
            stem_spatial_kernel,
            input_planes,
            stem_planes,
            stem_maxpool,
        )

    @classmethod
    def _parse_config(cls, config):
        ret_config = {}
        required_args = [
            "input_planes",
            "clip_crop_size",
            "skip_transformation_type",
            "residual_transformation_type",
            "frames_per_clip",
            "num_blocks",
        ]
        for arg in required_args:
            assert arg in config, "resnext3d model requires argument %s" % arg
            ret_config[arg] = config[arg]

        # Default setting for model stem, which is considered as stage 0. Stage
        # index starts from 0 as implemented in ResStageBase._block_name() method.
        #   stem_planes: No. of output channles of conv op in stem
        #   stem_temporal_kernel: temporal size of conv op in stem
        #   stem_spatial_kernel: spatial size of conv op in stem
        #   stem_maxpool: by default, spatial maxpool op is disabled in stem
        ret_config.update(
            {
                "input_key": config.get("input_key", None),
                "stem_name": config.get("stem_name", "resnext3d_stem"),
                "stem_planes": config.get("stem_planes", 64),
                "stem_temporal_kernel": config.get("stem_temporal_kernel", 3),
                "stem_spatial_kernel": config.get("stem_spatial_kernel", 7),
                "stem_maxpool": config.get("stem_maxpool", False),
            }
        )
        # Default setting for model stages 1, 2, 3 and 4
        #   stage_planes: No. of output channel of 1st conv op in stage 1
        #   stage_temporal_kernel_basis: Basis of temporal kernel sizes for each of
        #       the stage.
        #   temporal_conv_1x1: if True, do temporal convolution in the fist
        #      1x1 Conv3d. Otherwise, do it in the second 3x3 Conv3d (default settting)
        #   stage_temporal_stride: temporal stride for each stage
        #   stage_spatial_stride: spatial stride for each stage
        #   num_groups: No. of groups in 2nd (group) conv in the residual transformation
        #   width_per_group: No. of channels per group in 2nd (group) conv in the
        #     residual transformation
        ret_config.update(
            {
                "stage_planes": config.get("stage_planes", 256),
                "stage_temporal_kernel_basis": config.get(
                    "stage_temporal_kernel_basis", [[3], [3], [3], [3]]
                ),
                "temporal_conv_1x1": config.get(
                    "temporal_conv_1x1", [False, False, False, False]
                ),
                "stage_temporal_stride": config.get(
                    "stage_temporal_stride", [1, 2, 2, 2]
                ),
                "stage_spatial_stride": config.get(
                    "stage_spatial_stride", [1, 2, 2, 2]
                ),
                "num_groups": config.get("num_groups", 1),
                "width_per_group": config.get("width_per_group", 64),
            }
        )
        # Default setting for model parameter initialization
        ret_config.update(
            {
                "zero_init_residual_transform": config.get(
                    "zero_init_residual_transform", False
                )
            }
        )
        assert is_pos_int_list(ret_config["num_blocks"])
        assert is_pos_int(ret_config["stem_planes"])
        assert is_pos_int(ret_config["stem_temporal_kernel"])
        assert is_pos_int(ret_config["stem_spatial_kernel"])
        assert type(ret_config["stem_maxpool"]) == bool
        assert is_pos_int(ret_config["stage_planes"])
        assert type(ret_config["stage_temporal_kernel_basis"]) == list
        assert all(
            is_pos_int_list(l) for l in ret_config["stage_temporal_kernel_basis"]
        )
        assert type(ret_config["temporal_conv_1x1"]) == list
        assert is_pos_int_list(ret_config["stage_temporal_stride"])
        assert is_pos_int_list(ret_config["stage_spatial_stride"])
        assert is_pos_int(ret_config["num_groups"])
        assert is_pos_int(ret_config["width_per_group"])
        return ret_config

    def _init_parameter(self, zero_init_residual_transform):
        for m in self.modules():
            if isinstance(m, nn.Conv3d):
                if (
                    hasattr(m, "final_transform_op")
                    and m.final_transform_op
                    and zero_init_residual_transform
                ):
                    nn.init.constant_(m.weight, 0)
                else:
                    nn.init.kaiming_normal_(
                        m.weight, mode="fan_out", nonlinearity="relu"
                    )
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.BatchNorm3d) and m.affine:
                if (
                    hasattr(m, "final_transform_op")
                    and m.final_transform_op
                    and zero_init_residual_transform
                ):
                    batchnorm_weight = 0.0
                else:
                    batchnorm_weight = 1.0
                nn.init.constant_(m.weight, batchnorm_weight)
                nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.Linear):
                nn.init.normal_(m.weight, mean=0.0, std=0.01)
                nn.init.constant_(m.bias, 0)

    def set_classy_state(self, state):
        # We need to support both regular checkpoint loading and 2D conv weight
        # inflation into 3D conv weight in this function.
        self.load_head_states(state)
        current_state = self.state_dict()
        for name, weight_src in state["model"]["trunk"].items():
            assert name in current_state, (
                "weight %s is not found in ResNeXt3D model" % name
            )
            weight_tgt = current_state[name]
            assert (
                weight_src.dim() == weight_tgt.dim()
            ), "weight of source- and target 3D convolution should have same dimension"
            if (
                weight_src.dim() == 5
                and weight_src.shape[2] == 1
                and weight_tgt.shape[2] > 1
            ):
                # Find a source weight tensor where temporal dimension is 1. If the
                # temporal dimension of the current weight tensor with the same name
                # is larger than 1, we inflate the source weight tensor before
                # loading it. Such parameter inflation was first introduced in
                # the paper (https://arxiv.org/abs/1705.07750). It can achieve a
                # better initialization compared to random initialization.
                assert (
                    weight_src.shape[-2:] == weight_tgt.shape[-2:]
                    and weight_src.shape[:2] == weight_tgt.shape[:2]
                ), "weight shapes of source- and target 3D convolution mismatch"
                weight_src_inflated = (
                    weight_src.repeat(1, 1, weight_tgt.shape[2], 1, 1)
                    / weight_tgt.shape[2]
                )
                weight_src = weight_src_inflated
            else:
                assert all(
                    weight_src.size(d) == weight_tgt.size(d)
                    for d in range(weight_src.dim())
                ), (
                    "the shapes of source and target weight mismatch: %s Vs %s"
                    % (str(weight_src.size()), str(weight_tgt.size()))
                )

            current_state[name] = weight_src.clone()
        super().load_state_dict(current_state)

    def forward(self, x):
        """
        Args:
            x (dict or torch.Tensor): video input.
                When its type is dict, the dataset is a video dataset, and its
                content is like {"video": torch.tensor, "audio": torch.tensor}.
                When its type is torch.Tensor, the dataset is an image dataset.
        """
        assert isinstance(x, dict) or isinstance(
            x, torch.Tensor
        ), "x must be either a dictionary or a torch.Tensor"
        if isinstance(x, dict):
            assert self._input_key is not None and self._input_key in x, (
                "input key (%s) not in the input" % self._input_key
            )
            x = x[self._input_key]
        else:
            assert (
                self._input_key is None
            ), "when input of forward pass is a tensor, input key should not be set"
            assert x.dim() == 4 or x.dim() == 5, "tensor x must be 4D/5D tensor"
            if x.dim() == 4:
                # x is a 4D tensor of size N x C x H x W and is prepared from an
                # image dataset. We insert a temporal axis make it 5D of size
                # N x C x T x H x W
                x = torch.unsqueeze(x, 2)

        out = self.stem([x])
        out = self.stages(out)

        head_outputs = self.execute_heads()
        if len(head_outputs) == 0:
            raise Exception("Expecting at least one head that generates output")
        elif len(head_outputs) == 1:
            return list(head_outputs.values())[0]
        else:
            return head_outputs

    @property
    def input_shape(self):
        """
        Shape of video model input can vary in the following cases
        - At training stage, input are video frame croppings of fixed size.
        - At test stage, input are original video frames to support Fully Convolutional
            evaluation and its size can vary video by video
        """
        # Input shape is used by tensorboard hook. We put the input shape at
        # training stage for profiling and visualization purpose.
        return (
            self.input_planes,
            self.frames_per_clip,
            self.clip_crop_size,
            self.clip_crop_size,
        )

    @property
    def output_shape(self):
        return (1, None)

    @property
    def model_depth(self):
        return sum(self.num_blocks)

    @property
    def evaluation_mode(self):
        return ClassyModelEvaluationMode.VIDEO_CLIP_AVERAGING

    @property
    def input_key(self):
        return self._input_key


@register_model("resnext3d")
class ResNeXt3D(ResNeXt3DBase):
    """
    Implementation of:
        1. Conventional `post-activated 3D ResNe(X)t <https://arxiv.org/
        abs/1812.03982>`_.

        2. `Pre-activated 3D ResNe(X)t <https://arxiv.org/abs/1811.12814>`_.
        The model consists of one stem, a number of stages, and one or multiple
        heads that are attached to different blocks in the stage.
    """

    def __init__(
        self,
        input_key,
        input_planes,
        clip_crop_size,
        skip_transformation_type,
        residual_transformation_type,
        frames_per_clip,
        num_blocks,
        stem_name,
        stem_planes,
        stem_temporal_kernel,
        stem_spatial_kernel,
        stem_maxpool,
        stage_planes,
        stage_temporal_kernel_basis,
        temporal_conv_1x1,
        stage_temporal_stride,
        stage_spatial_stride,
        num_groups,
        width_per_group,
        zero_init_residual_transform,
    ):
        """
        Args:
            input_key (str): a key that can index into model input that is
                of dict type.
            input_planes (int): the channel dimension of the input. Normally 3 is used
                for rgb input.
            clip_crop_size (int): spatial cropping size of video clip at train time.
            skip_transformation_type (str): the type of skip transformation.
            residual_transformation_type (str): the type of residual transformation.
            frames_per_clip (int): Number of frames in a video clip.
            num_blocks (list): list of the number of blocks in stages.
            stem_name (str): name of model stem.
            stem_planes (int): the output dimension of the convolution in the model
                stem.
            stem_temporal_kernel (int): the temporal kernel size of the convolution
                in the model stem.
            stem_spatial_kernel (int): the spatial kernel size of the convolution
                in the model stem.
            stem_maxpool (bool): If true, perform max pooling.
            stage_planes (int): the output channel dimension of the 1st residual stage
            stage_temporal_kernel_basis (list): Basis of temporal kernel sizes for
                each of the stage.
            temporal_conv_1x1 (bool): Only useful for BottleneckTransformation.
                In a pathaway, if True, do temporal convolution in the first 1x1
                Conv3d. Otherwise, do it in the second 3x3 Conv3d.
            stage_temporal_stride (int): the temporal stride of the residual
                transformation.
            stage_spatial_stride (int): the spatial stride of the the residual
                transformation.
            num_groups (int): number of groups for the convolution.
                num_groups = 1 is for standard ResNet like networks, and
                num_groups > 1 is for ResNeXt like networks.
            width_per_group (int): Number of channels per group in 2nd (group)
                conv in the residual transformation in the first stage
            zero_init_residual_transform (bool): if true, the weight of last
                operation, which could be either BatchNorm3D in post-activated
                transformation or Conv3D in pre-activated transformation, in the
                residual transformation is initialized to zero
        """
        super(ResNeXt3D, self).__init__(
            input_key,
            input_planes,
            clip_crop_size,
            frames_per_clip,
            num_blocks,
            stem_name,
            stem_planes,
            stem_temporal_kernel,
            stem_spatial_kernel,
            stem_maxpool,
        )

        num_stages = len(num_blocks)
        out_planes = [stage_planes * 2 ** i for i in range(num_stages)]
        in_planes = [stem_planes] + out_planes[:-1]
        inner_planes = [
            num_groups * width_per_group * 2 ** i for i in range(num_stages)
        ]

        stages = []
        for s in range(num_stages):
            stage = ResStage(
                s + 1,  # stem is viewed as stage 0, and following stages start from 1
                [in_planes[s]],
                [out_planes[s]],
                [inner_planes[s]],
                [stage_temporal_kernel_basis[s]],
                [temporal_conv_1x1[s]],
                [stage_temporal_stride[s]],
                [stage_spatial_stride[s]],
                [num_blocks[s]],
                [num_groups],
                skip_transformation_type,
                residual_transformation_type,
                block_callback=self.build_attachable_block,
                disable_pre_activation=(s == 0),
                final_stage=(s == (num_stages - 1)),
            )
            stages.append(stage)

        self.stages = nn.Sequential(*stages)
        self._init_parameter(zero_init_residual_transform)

    @classmethod
    def from_config(cls, config: Dict[str, Any]) -> "ResNeXt3D":
        """Instantiates a ResNeXt3D from a configuration.

        Args:
            config: A configuration for a ResNeXt3D.
                See :func:`__init__` for parameters expected in the config.

        Returns:
            A ResNeXt3D instance.
        """
        ret_config = ResNeXt3D._parse_config(config)
        return cls(**ret_config)
