# -*- encoding: utf-8 -*-
"""
@File    :   cuda2d_model.py
@Time    :   2021/10/02 01:36:32
@Author  :   Ming Ding
@Contact :   dm18@mails.tsinghua.edu.cn
"""

# here put the import lib
import math
import torch
import torch.nn.functional as F

from helm.common.optional_dependencies import handle_module_not_found_error

try:
    from SwissArmyTransformer.model.base_model import BaseModel, BaseMixin
    from SwissArmyTransformer.mpu.utils import sqrt
    from SwissArmyTransformer.mpu import ColumnParallelLinear, RowParallelLinear
    from SwissArmyTransformer.model.transformer import unscaled_init_method, split_tensor_along_last_dim
    from SwissArmyTransformer.ops.local_attention_function import f_similar, f_weighting
except ModuleNotFoundError as e:
    handle_module_not_found_error(e, ["heim"])


class PositionEmbeddingMixin(BaseMixin):
    def __init__(
        self, additional_sequence_length, hidden_size, init_method_std=0.02, reinit_slice=slice(512, 512 + 400)
    ):
        super(PositionEmbeddingMixin, self).__init__()
        self.reinit_slice = reinit_slice
        self.position_embeddings = torch.nn.Embedding(additional_sequence_length, hidden_size)
        torch.nn.init.normal_(self.position_embeddings.weight, mean=0.0, std=init_method_std)

    def reinit(self, parent_model=None):
        old_weights = self.transformer.position_embeddings.weight.data[self.reinit_slice]
        old_len, hidden_size = old_weights.shape
        assert hidden_size == self.position_embeddings.weight.shape[-1]
        old_edge, new_edge = sqrt(old_len), sqrt(self.position_embeddings.weight.shape[-2])
        assert new_edge % old_edge == 0
        self.position_embeddings.weight.data.view(
            new_edge // old_edge, old_edge, new_edge // old_edge, old_edge, hidden_size
        ).copy_(old_weights.view(1, old_edge, 1, old_edge, hidden_size))
        # self.position_embeddings.weight.data.view(-1, old_len, hidden_size).copy_(old_weights)


class AttentionMixin(BaseMixin):
    def __init__(
        self,
        num_layers,
        hidden_size,
        init_method=unscaled_init_method(0.02),
        output_layer_init_method=unscaled_init_method(0.02),
    ):
        super(AttentionMixin, self).__init__()
        self.num_layers = num_layers  # replace attention in the LAST n layers
        self.query_key_value = torch.nn.ModuleList(
            [
                ColumnParallelLinear(
                    hidden_size, 3 * hidden_size, stride=3, gather_output=False, init_method=init_method
                )
                for layer_id in range(num_layers)
            ]
        )
        self.dense = torch.nn.ModuleList(
            [
                RowParallelLinear(
                    hidden_size, hidden_size, input_is_parallel=True, init_method=output_layer_init_method
                )
                for layer_id in range(num_layers)
            ]
        )

    def reinit(self, parent_model=None):
        start_layer = len(self.transformer.layers) - self.num_layers
        assert start_layer >= 0
        for layer_id in range(self.num_layers):
            old_attention = self.transformer.layers[start_layer + layer_id].attention
            self.query_key_value[layer_id].weight.data.copy_(old_attention.query_key_value.weight.data)  # type: ignore
            self.query_key_value[layer_id].bias.data.copy_(old_attention.query_key_value.bias.data)  # type: ignore
            self.dense[layer_id].weight.data.copy_(old_attention.dense.weight.data)  # type: ignore
            self.dense[layer_id].bias.data.copy_(old_attention.dense.bias.data)  # type: ignore


class DsrModel(BaseModel):
    def __init__(self, args, transformer=None):
        super().__init__(args, transformer=transformer)
        self.original_sequence_length = args.max_sequence_length
        additional_seqlen = args.new_sequence_length - args.max_sequence_length
        self.add_mixin("extra_position_embedding", PositionEmbeddingMixin(additional_seqlen, args.hidden_size))
        self.add_mixin("attention_plus", AttentionMixin(num_layers=args.num_layers, hidden_size=args.hidden_size))
        self.layout = args.layout
        # [PAD]... [ROI1] text ... [BOI1] {layout[0]} 1024 {layout[1]} [EOI1] 4095 {layout[2]}
        self.kernel_size = args.kernel_size
        self.kernel_size2 = args.kernel_size2
        self.log_attention_weights = None

    def position_embedding_forward(self, position_ids, **kw_args):
        position = position_ids[..., : self.layout[1]]
        position_plus = position_ids[..., self.layout[1] :] - self.original_sequence_length
        position_embeddings = torch.cat(
            (
                self.transformer.position_embeddings(position),
                self.get_mixin("extra_position_embedding").position_embeddings(position_plus),
            ),
            dim=-2,
        )
        return position_embeddings

    def attention_forward(self, hidden_states, mask, layer_id=None, log_attention_weights=None, **kw_args):
        attn_module = self.transformer.layers[layer_id].attention
        # attention_plus on all layers
        query_key_value_plus = self.get_mixin("attention_plus").query_key_value[layer_id]
        dense_plus = self.get_mixin("attention_plus").dense[layer_id]
        # split two parts
        hidden_states_plus = hidden_states[:, self.layout[1] :]
        hidden_states = hidden_states[:, : self.layout[1]]
        # base model qkv
        mixed_raw_layer = attn_module.query_key_value(hidden_states)
        q0, k0, v0 = split_tensor_along_last_dim(mixed_raw_layer, 3)
        # cuda2d model qkv
        mixed_raw_layer = query_key_value_plus(hidden_states_plus)
        q1, k1, v1 = split_tensor_along_last_dim(mixed_raw_layer, 3)

        dropout_fn = attn_module.attention_dropout if self.training else None

        # cuda2d attention
        context_layer0, context_layer1 = sparse_attention_2d_light(
            q0,
            k0,
            v0,
            q1,
            k1,
            v1,
            mask,
            n_head=attn_module.num_attention_heads_per_partition,
            text_len=self.layout[0],
            kernel_size=self.kernel_size,
            kernel_size2=self.kernel_size2,
            attention_dropout=dropout_fn,
            log_attention_weights=log_attention_weights,
            add_scalar=(kw_args["add_scalar"] if "add_scalar" in kw_args else 0),
        )

        output_0 = attn_module.dense(context_layer0)
        output_1 = dense_plus(context_layer1)
        output = torch.cat((output_0, output_1), dim=1)

        return output

    def final_forward(self, logits, **kwargs):
        logits_parallel = logits
        logits_parallel = torch.nn.functional.linear(
            logits_parallel.float(), self.transformer.word_embeddings.weight[:20000].float()
        )
        # logits_parallel = torch.nn.functional.linear(logits_parallel, self.transformer.word_embeddings.weight[:20000])
        return logits_parallel

    def disable_untrainable_params(self):
        self.transformer.requires_grad_(False)

    @classmethod
    def add_model_specific_args(cls, parser):
        group = parser.add_argument_group("Cuda2dModel", "cuda2d model configurations")
        group.add_argument("--kernel-size", type=int, default=5)
        group.add_argument("--kernel-size2", type=int, default=5)
        group.add_argument("--layout", type=str, default="96,496,4096")
        group.add_argument("--new-sequence-length", type=int, default=4096)
        return parser


def sparse_attention_2d_light(
    q0,
    k0,
    v0,
    q1,
    k1,
    v1,
    attention_mask,
    n_head,
    text_len,
    kernel_size=9,
    kernel_size2=7,
    attention_dropout=None,
    log_attention_weights=None,
    add_scalar=0,
    **kwargs,
):
    """
    q0, k0, v0: [batch_size, 1088, hidden_size]
    q1, k1, v1: [batch_size, 4096, h2]
    n_head: int
    attention_mask: [batch_size, 1088, 1088]
    """
    b, s0, h0 = q0.shape
    b, s1, h1 = q1.shape
    h, l0, l1 = h0 // n_head, sqrt(s0 - text_len), sqrt(s1)

    q0 = q0.reshape(b, s0, n_head, h).permute(0, 2, 1, 3)
    v0 = v0.reshape(b, s0, n_head, h).permute(0, 2, 1, 3)
    k0T = k0.reshape(b, s0, n_head, h).permute(0, 2, 3, 1)

    # standard attention for level 0
    attention_scores = torch.matmul(q0 / math.sqrt(q0.shape[-1]), k0T)

    if log_attention_weights is not None:
        attention_scores += log_attention_weights
    attention_scores = torch.mul(attention_scores, attention_mask) - 10000.0 * (1.0 - attention_mask)

    attention_probs0 = F.softmax(attention_scores, dim=-1)

    # local attention for level 1
    q1 = (
        (q1.view(b, s1, n_head, h1 // n_head).permute(0, 2, 3, 1) / math.sqrt(h1 // n_head))
        .contiguous()
        .view(b * n_head, h1 // n_head, l1, l1)
    )
    k1 = k1.view(b, s1, n_head, h1 // n_head).permute(0, 2, 3, 1).contiguous().view(b * n_head, h1 // n_head, l1, l1)
    v1 = v1.view(b, s1, n_head, h1 // n_head).permute(0, 2, 3, 1).contiguous().view(b * n_head, h1 // n_head, l1, l1)
    # scores_1_to_1 = f_similar(q1, k1, kernel_size*2-1, kernel_size, True)
    scores_1_to_1 = f_similar(q1, k1, kernel_size * 2 - 1, kernel_size, False)

    # cross attention
    k0T = k0T[..., -(l0**2) :].reshape(b * n_head, h, l0, l0).contiguous()
    scores_1_to_0 = f_similar(q1, k0T, kernel_size2, kernel_size2, False)  # [b*n_head, l1, l1, field]
    scores_1 = torch.cat(
        (
            scores_1_to_0.view(b * n_head, -1, scores_1_to_0.shape[3]) + add_scalar,
            scores_1_to_1.view(b * n_head, -1, scores_1_to_1.shape[3]),
        ),
        dim=-1,
    )
    attention_probs1 = F.softmax(scores_1, dim=-1)

    if attention_dropout is not None:
        # with get_cuda_rng_tracker().fork():
        attention_probs0 = attention_dropout(attention_probs0)
        attention_probs1 = attention_dropout(attention_probs1)

    # weighting for level 0
    context0 = torch.matmul(attention_probs0, v0)  # [b, n_head, s0, h]
    # weighting for level 1
    probs_1_to_1 = attention_probs1[:, :, -scores_1_to_1.shape[3] :].view_as(scores_1_to_1)
    # context1_to_1 = f_weighting(v1, probs_1_to_1.contiguous(), kernel_size*2-1, kernel_size, True)
    context1_to_1 = f_weighting(v1, probs_1_to_1.contiguous(), kernel_size * 2 - 1, kernel_size, False)

    context1 = context1_to_1.view(b, n_head * h, l1**2)
    # weighting for cross attention
    probs_1_to_0 = attention_probs1[:, :, : scores_1_to_0.shape[3]].view_as(scores_1_to_0)
    v0_part = v0[:, :, -(l0**2) :].transpose(-1, -2).contiguous().view(b * n_head, h, l0, l0)
    context1_to_0 = f_weighting(v0_part, probs_1_to_0.contiguous(), kernel_size2, kernel_size2, False)
    context1_to_0 = context1_to_0.view(b, n_head * h, l1**2)
    context1 = context1 + context1_to_0
    return context0.transpose(1, 2).reshape(b, s0, h0), context1.transpose(-1, -2)
