# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
#
# Modifications copyright (c) 2021 DocYard Authors. All Rights Reserve.
#
# 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 __future__ import absolute_import, division, print_function

from collections import OrderedDict

import torch
from torch import nn
from torch.nn import functional as F

from .self_attention import WrapEncoder, WrapEncoderForFeature

gradient_clip = 10

from functools import partial


class PVAM(nn.Module):
    def __init__(
        self,
        in_channels,
        char_num,
        max_text_length,
        num_heads,
        num_encoder_tus,
        hidden_dims,
    ):
        super(PVAM, self).__init__()
        self.char_num = char_num
        self.max_length = max_text_length
        self.num_heads = num_heads
        self.num_encoder_TUs = num_encoder_tus
        self.hidden_dims = hidden_dims
        # Transformer encoder
        t = 256
        self.wrap_encoder_for_feature = WrapEncoderForFeature(
            src_vocab_size=1,
            max_length=t,
            n_layer=self.num_encoder_TUs,
            n_head=self.num_heads,
            d_key=int(self.hidden_dims / self.num_heads),
            d_value=int(self.hidden_dims / self.num_heads),
            d_model=self.hidden_dims,
            d_inner_hid=self.hidden_dims,
            prepostprocess_dropout=0.1,
            attention_dropout=0.1,
            relu_dropout=0.1,
            preprocess_cmd="n",
            postprocess_cmd="da",
            weight_sharing=True,
        )

        # PVAM
        self.flatten0 = nn.Flatten(start_dim=0, end_dim=1)
        self.fc0 = nn.Linear(
            in_features=in_channels,
            out_features=in_channels,
        )
        self.emb = nn.Embedding(
            num_embeddings=self.max_length, embedding_dim=in_channels
        )
        self.flatten1 = nn.Flatten(start_dim=0, end_dim=2)
        self.fc1 = nn.Linear(
            in_features=in_channels, out_features=1, bias=False
        )

    def forward(self, inputs, encoder_word_pos, gsrm_word_pos):
        b, c, h, w = inputs.shape
        conv_features = torch.reshape(inputs, shape=(-1, c, h * w))
        conv_features = conv_features.permute(0, 2, 1)
        # transformer encoder
        b, t, c = conv_features.shape

        enc_inputs = [conv_features, encoder_word_pos, None]
        word_features = self.wrap_encoder_for_feature(enc_inputs)

        # pvam
        b, t, c = word_features.shape
        word_features = self.fc0(word_features)
        word_features_ = torch.reshape(word_features, (-1, 1, t, c))
        word_features_ = torch.tile(word_features_, (1, self.max_length, 1, 1))
        word_pos_feature = self.emb(gsrm_word_pos)
        word_pos_feature_ = torch.reshape(
            word_pos_feature, (-1, self.max_length, 1, c)
        )
        word_pos_feature_ = torch.tile(word_pos_feature_, (1, 1, t, 1))
        y = word_pos_feature_ + word_features_
        y = F.tanh(y)
        attention_weight = self.fc1(y)
        attention_weight = torch.reshape(
            attention_weight, shape=(-1, self.max_length, t)
        )
        attention_weight = F.softmax(attention_weight, dim=-1)
        pvam_features = torch.bmm(
            attention_weight, word_features
        )  # [b, max_length, c]
        return pvam_features


class GSRM(nn.Module):
    def __init__(
        self,
        in_channels,
        char_num,
        max_text_length,
        num_heads,
        num_encoder_tus,
        num_decoder_tus,
        hidden_dims,
    ):
        super(GSRM, self).__init__()
        self.char_num = char_num
        self.max_length = max_text_length
        self.num_heads = num_heads
        self.num_encoder_TUs = num_encoder_tus
        self.num_decoder_TUs = num_decoder_tus
        self.hidden_dims = hidden_dims

        self.fc0 = nn.Linear(
            in_features=in_channels, out_features=self.char_num
        )
        self.wrap_encoder0 = WrapEncoder(
            src_vocab_size=self.char_num + 1,
            max_length=self.max_length,
            n_layer=self.num_decoder_TUs,
            n_head=self.num_heads,
            d_key=int(self.hidden_dims / self.num_heads),
            d_value=int(self.hidden_dims / self.num_heads),
            d_model=self.hidden_dims,
            d_inner_hid=self.hidden_dims,
            prepostprocess_dropout=0.1,
            attention_dropout=0.1,
            relu_dropout=0.1,
            preprocess_cmd="n",
            postprocess_cmd="da",
            weight_sharing=True,
        )

        self.wrap_encoder1 = WrapEncoder(
            src_vocab_size=self.char_num + 1,
            max_length=self.max_length,
            n_layer=self.num_decoder_TUs,
            n_head=self.num_heads,
            d_key=int(self.hidden_dims / self.num_heads),
            d_value=int(self.hidden_dims / self.num_heads),
            d_model=self.hidden_dims,
            d_inner_hid=self.hidden_dims,
            prepostprocess_dropout=0.1,
            attention_dropout=0.1,
            relu_dropout=0.1,
            preprocess_cmd="n",
            postprocess_cmd="da",
            weight_sharing=True,
        )

        # self.mul = lambda x: torch.matmul(x,
        #                                    (self.wrap_encoder0.prepare_decoder.emb0.weightk).transpose(-2, -1))  # ! This is an error here, weightk is wrong correct it by visualizing torch model_dict

        self.mul = partial(
            self.f, self.wrap_encoder0.prepare_decoder.emb0.weight
        )

    @staticmethod
    def f(w, x):
        return torch.matmul(x, w.transpose(-2, -1))

    def forward(
        self, inputs, gsrm_word_pos, gsrm_slf_attn_bias1, gsrm_slf_attn_bias2
    ):
        # ===== GSRM Visual-to-semantic embedding block =====
        b, t, c = inputs.shape
        pvam_features = torch.reshape(inputs, (-1, c))
        word_out = self.fc0(pvam_features)
        word_ids = torch.argmax(F.softmax(word_out), dim=1)
        word_ids = torch.reshape(word_ids, shape=(-1, t, 1))

        # ===== GSRM Semantic reasoning block =====
        """
        This module is achieved through bi-transformers,
        ngram_feature1 is the froward one, ngram_fetaure2 is the backward one
        """
        pad_idx = self.char_num

        word1 = word_ids.float()
        pad1 = nn.ConstantPad1d((1, 0), value=1.0 * pad_idx)
        word1 = pad1(word1.permute(0, 2, 1)).permute(0, 2, 1)
        word1 = word1.long()
        word1 = word1[:, :-1, :]
        word2 = word_ids

        enc_inputs_1 = [word1, gsrm_word_pos, gsrm_slf_attn_bias1]
        enc_inputs_2 = [word2, gsrm_word_pos, gsrm_slf_attn_bias2]

        gsrm_feature1 = self.wrap_encoder0(enc_inputs_1)
        gsrm_feature2 = self.wrap_encoder1(enc_inputs_2)

        pad = nn.ConstantPad1d((0, 1), value=0.0)
        gsrm_feature2 = pad(gsrm_feature2.permute(0, 2, 1)).permute(0, 2, 1)
        gsrm_feature2 = gsrm_feature2[
            :,
            1:,
        ]
        gsrm_features = gsrm_feature1 + gsrm_feature2

        gsrm_out = self.mul(gsrm_features)

        b, t, c = gsrm_out.shape
        gsrm_out = torch.reshape(gsrm_out, (-1, c))

        return gsrm_features, word_out, gsrm_out


class VSFD(nn.Module):
    def __init__(self, in_channels=512, pvam_ch=512, char_num=38):
        super(VSFD, self).__init__()
        self.char_num = char_num
        self.fc0 = nn.Linear(in_features=in_channels * 2, out_features=pvam_ch)
        self.fc1 = nn.Linear(in_features=pvam_ch, out_features=self.char_num)

    def forward(self, pvam_feature, gsrm_feature):
        b, t, c1 = pvam_feature.shape
        b, t, c2 = gsrm_feature.shape
        combine_feature_ = torch.cat([pvam_feature, gsrm_feature], dim=2)
        img_comb_feature_ = torch.reshape(
            combine_feature_, shape=(-1, c1 + c2)
        )
        img_comb_feature_map = self.fc0(img_comb_feature_)
        img_comb_feature_map = torch.sigmoid(img_comb_feature_map)
        img_comb_feature_map = torch.reshape(
            img_comb_feature_map, shape=(-1, t, c1)
        )
        combine_feature = (
            img_comb_feature_map * pvam_feature
            + (1.0 - img_comb_feature_map) * gsrm_feature
        )
        img_comb_feature = torch.reshape(combine_feature, shape=(-1, c1))

        out = self.fc1(img_comb_feature)
        return out


class SRNHead(nn.Module):
    def __init__(
        self,
        in_channels,
        out_channels,
        max_text_length,
        num_heads,
        num_encoder_TUs,
        num_decoder_TUs,
        hidden_dims,
        **kwargs
    ):
        super(SRNHead, self).__init__()
        self.char_num = out_channels
        self.max_length = max_text_length
        self.num_heads = num_heads
        self.num_encoder_TUs = num_encoder_TUs
        self.num_decoder_TUs = num_decoder_TUs
        self.hidden_dims = hidden_dims

        self.pvam = PVAM(
            in_channels=in_channels,
            char_num=self.char_num,
            max_text_length=self.max_length,
            num_heads=self.num_heads,
            num_encoder_tus=self.num_encoder_TUs,
            hidden_dims=self.hidden_dims,
        )

        self.gsrm = GSRM(
            in_channels=in_channels,
            char_num=self.char_num,
            max_text_length=self.max_length,
            num_heads=self.num_heads,
            num_encoder_tus=self.num_encoder_TUs,
            num_decoder_tus=self.num_decoder_TUs,
            hidden_dims=self.hidden_dims,
        )
        self.vsfd = VSFD(in_channels=in_channels, char_num=self.char_num)

        self.gsrm.wrap_encoder1.prepare_decoder.emb0 = (
            self.gsrm.wrap_encoder0.prepare_decoder.emb0
        )

    def forward(self, inputs, others):
        encoder_word_pos = others[0]
        gsrm_word_pos = others[1]
        gsrm_slf_attn_bias1 = others[2]
        gsrm_slf_attn_bias2 = others[3]

        pvam_feature = self.pvam(inputs, encoder_word_pos, gsrm_word_pos)

        gsrm_feature, word_out, gsrm_out = self.gsrm(
            pvam_feature,
            gsrm_word_pos,
            gsrm_slf_attn_bias1,
            gsrm_slf_attn_bias2,
        )

        final_out = self.vsfd(pvam_feature, gsrm_feature)
        if not self.training:
            final_out = F.softmax(final_out, dim=1)

        _, decoded_out = torch.topk(final_out, k=1)

        predicts = OrderedDict(
            [
                ("predict", final_out),
                ("pvam_feature", pvam_feature),
                ("decoded_out", decoded_out),
                ("word_out", word_out),
                ("gsrm_out", gsrm_out),
            ]
        )

        return predicts
