# Copyright (c) 2019-2020, NVIDIA CORPORATION. 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, Optional
import functools

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


class OptionalParameterList(nn.ParameterList):
    def extra_repr(self):
        child_lines = []
        for k, p in self._parameters.items():
            if p is not None:
                size_str = 'x'.join(str(size) for size in p.size())
                device_str = '' if not p.is_cuda else ' (GPU {})'.format(p.get_device())
                parastr = 'Parameter containing: [{} of size {}{}]'.format(
                    torch.typename(p), size_str, device_str)
                child_lines.append('  (' + str(k) + '): ' + parastr)
        tmpstr = '\n'.join(child_lines)
        return tmpstr


class ProjectedAdaptiveLogSoftmax(nn.Module):
    def __init__(self, n_token, d_embed, d_proj, cutoffs, div_val=1,
                 tie_projs=None, out_layers_weights=None, out_projs=None,
                 keep_order=False,
                 bias_scale=0.0,
                 dropout=0.0,
                 ):
        super().__init__()

        self.n_token = n_token
        self.d_embed = d_embed
        self.d_proj = d_proj

        self.cutoffs = list(cutoffs) + [n_token]
        self.cutoff_ends = [0] + self.cutoffs
        self.div_val = div_val

        self.shortlist_size = self.cutoffs[0]
        self.n_clusters = len(self.cutoffs) - 1
        self.head_size = self.shortlist_size + self.n_clusters

        # [21-09-15 AG]: bake the first False into the definition, just as [0] is built into the cutoffs
        if tie_projs is None: tie_projs = []
        elif isinstance(tie_projs, bool): tie_projs = [tie_projs] * len(cutoffs)
        else: tie_projs = list(tie_projs)
        tie_projs = [False] + tie_projs
        self.tie_projs = tie_projs

        if self.n_clusters > 0:
            self.cluster_weight = nn.Parameter(torch.zeros(self.n_clusters, self.d_embed))
            self.cluster_bias = nn.Parameter(torch.zeros(self.n_clusters))

        if not out_layers_weights:
            self.out_layers_weights = nn.ParameterList()
        else:
            self.out_layers_weights = out_layers_weights

        self.out_layers_biases = nn.ParameterList()

        self.shared_out_projs = out_projs
        self.out_projs = OptionalParameterList()

        self.dropout = dropout
        self.drop = nn.Dropout(dropout)

        if div_val == 1:
            if d_proj != d_embed:
                for i in range(len(self.cutoffs)):
                    if tie_projs[i]:
                        self.out_projs.append(None)
                    else:
                        self.out_projs.append(
                            nn.Parameter(torch.zeros(d_proj, d_embed))
                        )
            else:
                # self.out_projs = [None] * len(self.cutoffs)
                self.out_projs.append(None)

            self.out_layers_biases.append(
                nn.Parameter(torch.zeros(n_token))
                )

            if not out_layers_weights:
                self.out_layers_weights.append(
                    nn.Parameter(torch.zeros(n_token, d_embed))
                    )
        else:
            for i in range(len(self.cutoffs)):
                l_idx, r_idx = self.cutoff_ends[i], self.cutoff_ends[i+1]
                d_emb_i = d_embed // (div_val ** i)

                if tie_projs[i]:
                    self.out_projs.append(None)
                else:
                    self.out_projs.append(
                        nn.Parameter(torch.zeros(d_proj, d_emb_i))
                    )

                self.out_layers_biases.append(
                    nn.Parameter(torch.zeros(r_idx - l_idx))
                    )
                if not out_layers_weights:
                    self.out_layers_weights.append(
                        nn.Parameter(torch.zeros(r_idx - l_idx, d_emb_i))
                        )
        for bias in self.out_layers_biases:
            bound = bias_scale * d_proj ** -.5
            nn.init.uniform_(bias, -bound, bound)


        self.keep_order = keep_order

    def _compute_logit(self, hidden, weight, bias, proj):
        if proj is None:
            logit = F.linear(hidden, weight, bias=bias)
        else:
            if self.dropout > 0.0:
                logit = hidden @ proj
                logit = self.drop(logit)
                logit = logit @ weight.t()
            else:
                logit = torch.einsum('bd,de,ev->bv', (hidden, proj, weight.t()))
            if bias is not None:
                logit = logit + bias
        return logit

    def get_out_proj(self, i):
        if self.tie_projs[i]:
            if len(self.shared_out_projs) == 0:
                return None
            elif len(self.shared_out_projs) == 1:
                return self.shared_out_projs[0]
            else:
                return self.shared_out_projs[i]
        else:
            return self.out_projs[i]

    def forward(self, hidden, target, keep_order=False, key_padding_mask=None, *args, **kwargs):
        # [21-09-15 AG]: TODO may need to handle key_padding_mask
        '''
            hidden :: [len*bsz x d_proj]
            target :: [len*bsz]
        '''

        hidden = hidden.reshape(-1, hidden.size(-1))
        target = target.reshape(-1)
        if hidden.size(0) != target.size(0):
            print(hidden.shape, target.shape)
            raise RuntimeError('Input and target should have the same size '
                               'in the batch dimension.')

        if self.n_clusters == 0:
            logit = self._compute_logit(hidden, self.out_layers_weights[0],
                                        self.out_layers_biases[0], self.get_out_proj(0))
            nll = -F.log_softmax(logit, dim=-1) \
                    .gather(1, target.unsqueeze(1)).squeeze(1)
        else:
            # construct weights and biases
            weights, biases = [], []
            for i in range(len(self.cutoffs)):
                if self.div_val == 1:
                    l_idx, r_idx = self.cutoff_ends[i], self.cutoff_ends[i + 1]
                    weight_i = self.out_layers_weights[0][l_idx:r_idx]
                    bias_i = self.out_layers_biases[0][l_idx:r_idx]
                else:
                    weight_i = self.out_layers_weights[i]
                    bias_i = self.out_layers_biases[i]

                if i == 0:
                    weight_i = torch.cat(
                        [weight_i, self.cluster_weight], dim=0)
                    bias_i = torch.cat(
                        [bias_i, self.cluster_bias], dim=0)

                weights.append(weight_i)
                biases.append(bias_i)

            head_weight, head_bias, head_proj = weights[0], biases[0], self.get_out_proj(0)

            head_logit = self._compute_logit(hidden, head_weight, head_bias, head_proj)
            head_logprob = F.log_softmax(head_logit, dim=1)

            nll = torch.zeros_like(target, dtype=hidden.dtype, device=hidden.device)

            offset = 0
            cutoff_values = [0] + self.cutoffs
            for i in range(len(cutoff_values) - 1):
                l_idx, r_idx = cutoff_values[i], cutoff_values[i + 1]

                mask_i = (target >= l_idx) & (target < r_idx)
                indices_i = mask_i.nonzero(as_tuple=False).squeeze()

                if indices_i.numel() == 0:
                    continue

                target_i = target.index_select(0, indices_i) - l_idx
                head_logprob_i = head_logprob.index_select(0, indices_i)

                if i == 0:
                    logprob_i = head_logprob_i.gather(1, target_i[:, None]).squeeze(1)
                else:
                    weight_i, bias_i, proj_i = weights[i], biases[i], self.get_out_proj(i)

                    hidden_i = hidden.index_select(0, indices_i)

                    tail_logit_i = self._compute_logit(hidden_i, weight_i, bias_i, proj_i)
                    tail_logprob_i = F.log_softmax(tail_logit_i, dim=1)

                    logprob_i = head_logprob_i[:, -i] \
                        + tail_logprob_i.gather(1, target_i[:, None]).squeeze(1)

                if self.keep_order or keep_order:
                    nll.index_copy_(0, indices_i, -logprob_i)
                else:
                    nll[offset:offset+logprob_i.size(0)].copy_(-logprob_i)

                offset += logprob_i.size(0)

        return nll.mean() # TODO maybe cases for length or padding_mask


class AdaptiveEmbedding(nn.Module):
    """ Copy of transformers.AdaptiveEmbedding that works with fp16 by replacing the index_put_ operation

    Initialization has been fixed for the case when d_proj = d_embed
    """
    def __init__(self, n_token, d_embed, d_proj, cutoffs : List[int], div_val=1, init_scale=1.0, sample_softmax=False, dropout=0.0):
        super().__init__()

        self.n_token = n_token
        self.d_embed = d_embed

        self.cutoffs = list(cutoffs) + [n_token]
        self.div_val = div_val
        self.d_proj = d_proj
        self.drop = nn.Dropout(dropout) if dropout > 0.0 else nn.Identity()

        self.emb_scale = d_proj ** 0.5

        self.cutoff_ends = [0] + self.cutoffs

        self.emb_layers = nn.ModuleList()
        self.emb_projs = nn.ParameterList()
        if div_val == 1:
            self.emb_layers.append(nn.Embedding(n_token, d_embed, sparse=sample_softmax > 0))
            _init_embed(self.emb_layers[-1].weight, d_embed, init_scale)
            # torch.nn.init.normal_(self.emb_layers[-1].weight, mean=0, std=init_scale * d_embed ** -.5)
            if d_proj != d_embed: # TODO
                # self.emb_projs.append(nn.Parameter(torch.FloatTensor(d_proj, d_embed)))
                self.emb_projs.append(nn.Parameter(torch.FloatTensor(d_proj, d_embed)))
                # torch.nn.init.normal_(self.emb_projs[-1], mean=0, std=init_scale * 1./self.emb_scale)
                _init_proj(self.emb_projs[-1], d_proj, init_scale)
        else:
            for i in range(len(self.cutoffs)):
                l_idx, r_idx = self.cutoff_ends[i], self.cutoff_ends[i + 1]
                d_emb_i = d_embed // (div_val ** i)
                self.emb_layers.append(nn.Embedding(r_idx - l_idx, d_emb_i))
                # torch.nn.init.normal_(self.emb_layers[-1].weight, mean=0, std=init_scale * d_emb_i ** -.5)
                _init_embed(self.emb_layers[-1].weight, d_emb_i, init_scale)
                self.emb_projs.append(nn.Parameter(torch.FloatTensor(d_proj, d_emb_i)))
                # torch.nn.init.normal_(self.emb_projs[-1], mean=0, std=init_scale * 1./self.emb_scale)
                _init_proj(self.emb_projs[-1], d_proj, init_scale)

    def forward(self, inp, *args, **kwargs):
        if self.div_val == 1:
            embed = self.emb_layers[0](inp)
            embed = self.drop(embed)
            if self.d_proj != self.d_embed:
                embed = F.linear(embed, self.emb_projs[0])
        else:
            param = next(self.parameters())
            inp_flat = inp.view(-1)

            # Changes
            # emb_flat = torch.zeros([inp_flat.size(0), self.d_proj], dtype=param.dtype, device=param.device)
            embeddings = []
            indices = torch.zeros_like(inp_flat) # empty should work as long as cutoffs[-1] > max token
            _total_tokens = 0

            # emb_flat = inp.new_zeros(inp_flat.size(0), self.d_proj)
            for i in range(len(self.cutoffs)):
                l_idx, r_idx = self.cutoff_ends[i], self.cutoff_ends[i + 1]

                mask_i = (inp_flat >= l_idx) & (inp_flat < r_idx)
                indices_i = mask_i.nonzero().squeeze(-1) # shape (_tokens,)

                _tokens = indices_i.numel()
                if _tokens == 0:
                    continue

                inp_i = inp_flat.index_select(0, indices_i) - l_idx
                emb_i = self.emb_layers[i](inp_i)
                emb_i = self.drop(emb_i)
                emb_i = F.linear(emb_i, self.emb_projs[i])

                # Changes
                embeddings.append(emb_i)
                indices.index_put_(
                    (indices_i,),
                    torch.arange(_tokens, device=inp.device) + _total_tokens
                )
                _total_tokens += _tokens

                # emb_flat.index_copy_(0, indices_i, emb_i)
            embeddings = torch.cat(embeddings, dim=0)
            emb_flat = embeddings[indices]

            embed_shape = inp.size() + (self.d_proj,)
            embed = emb_flat.view(embed_shape)

        embed.mul_(self.emb_scale)
        # embed.div_(self.emb_scale)

        return embed


def _init_weight(weight, d : int, init_scale : Optional[float], default=None):
    assert init_scale or default
    if init_scale is None:
        std = default
    else:
        std = init_scale * (d ** -0.5)
    nn.init.normal_(weight, mean=0, std=std)

_init_embed = functools.partial(_init_weight, default=0.02)
_init_proj = functools.partial(_init_weight, default=0.01)

### Just for this codebase, we need to squeeze the last dimension because inputs are always given as (B, L, D) instead of (B, L)
import src.models.nn.utils as U
# AdaptiveEmbedding = U.Squeeze(AdaptiveEmbedding)

[((341, 14, 341, 59), 'functools.partial', 'functools.partial', (), '', False, 'import functools\n'), ((342, 13, 342, 58), 'functools.partial', 'functools.partial', (), '', False, 'import functools\n'), ((339, 4, 339, 44), 'torch.nn.init.normal_', 'nn.init.normal_', (), '', True, 'import torch.nn as nn\n'), ((74, 33, 74, 51), 'torch.nn.ParameterList', 'nn.ParameterList', ({}, {}), '()', True, 'import torch.nn as nn\n'), ((80, 20, 80, 39), 'torch.nn.Dropout', 'nn.Dropout', ({(80, 31, 80, 38): 'dropout'}, {}), '(dropout)', True, 'import torch.nn as nn\n'), ((258, 26, 258, 41), 'torch.nn.ModuleList', 'nn.ModuleList', ({}, {}), '()', True, 'import torch.nn as nn\n'), ((259, 25, 259, 43), 'torch.nn.ParameterList', 'nn.ParameterList', ({}, {}), '()', True, 'import torch.nn as nn\n'), ((70, 38, 70, 56), 'torch.nn.ParameterList', 'nn.ParameterList', ({}, {}), '()', True, 'import torch.nn as nn\n'), ((124, 12, 124, 49), 'torch.nn.init.uniform_', 'nn.init.uniform_', ({(124, 29, 124, 33): 'bias', (124, 35, 124, 41): '(-bound)', (124, 43, 124, 48): 'bound'}, {}), '(bias, -bound, bound)', True, 'import torch.nn as nn\n'), ((131, 20, 131, 55), 'torch.nn.functional.linear', 'F.linear', (), '', True, 'import torch.nn.functional as F\n'), ((197, 27, 197, 59), 'torch.nn.functional.log_softmax', 'F.log_softmax', (), '', True, 'import torch.nn.functional as F\n'), ((199, 18, 199, 84), 'torch.zeros_like', 'torch.zeros_like', (), '', False, 'import torch\n'), ((252, 20, 252, 39), 'torch.nn.Dropout', 'nn.Dropout', ({(252, 31, 252, 38): 'dropout'}, {}), '(dropout)', True, 'import torch.nn as nn\n'), ((252, 62, 252, 75), 'torch.nn.Identity', 'nn.Identity', ({}, {}), '()', True, 'import torch.nn as nn\n'), ((293, 22, 293, 48), 'torch.zeros_like', 'torch.zeros_like', ({(293, 39, 293, 47): 'inp_flat'}, {}), '(inp_flat)', False, 'import torch\n'), ((321, 25, 321, 53), 'torch.cat', 'torch.cat', (), '', False, 'import torch\n'), ((66, 47, 66, 89), 'torch.zeros', 'torch.zeros', ({(66, 59, 66, 74): 'self.n_clusters', (66, 76, 66, 88): 'self.d_embed'}, {}), '(self.n_clusters, self.d_embed)', False, 'import torch\n'), ((67, 45, 67, 73), 'torch.zeros', 'torch.zeros', ({(67, 57, 67, 72): 'self.n_clusters'}, {}), '(self.n_clusters)', False, 'import torch\n'), ((261, 35, 261, 92), 'torch.nn.Embedding', 'nn.Embedding', (), '', True, 'import torch.nn as nn\n'), ((285, 24, 285, 58), 'torch.nn.functional.linear', 'F.linear', ({(285, 33, 285, 38): 'embed', (285, 40, 285, 57): 'self.emb_projs[0]'}, {}), '(embed, self.emb_projs[0])', True, 'import torch.nn.functional as F\n'), ((310, 24, 310, 58), 'torch.nn.functional.linear', 'F.linear', ({(310, 33, 310, 38): 'emb_i', (310, 40, 310, 57): 'self.emb_projs[i]'}, {}), '(emb_i, self.emb_projs[i])', True, 'import torch.nn.functional as F\n'), ((31, 20, 31, 37), 'torch.typename', 'torch.typename', ({(31, 35, 31, 36): 'p'}, {}), '(p)', False, 'import torch\n'), ((96, 29, 96, 49), 'torch.zeros', 'torch.zeros', ({(96, 41, 96, 48): 'n_token'}, {}), '(n_token)', False, 'import torch\n'), ((186, 31, 187, 63), 'torch.cat', 'torch.cat', (), '', False, 'import torch\n'), ((188, 29, 189, 59), 'torch.cat', 'torch.cat', (), '', False, 'import torch\n'), ((223, 37, 223, 71), 'torch.nn.functional.log_softmax', 'F.log_softmax', (), '', True, 'import torch.nn.functional as F\n'), ((273, 39, 273, 75), 'torch.nn.Embedding', 'nn.Embedding', ({(273, 52, 273, 65): '(r_idx - l_idx)', (273, 67, 273, 74): 'd_emb_i'}, {}), '(r_idx - l_idx, d_emb_i)', True, 'import torch.nn as nn\n'), ((101, 33, 101, 62), 'torch.zeros', 'torch.zeros', ({(101, 45, 101, 52): 'n_token', (101, 54, 101, 61): 'd_embed'}, {}), '(n_token, d_embed)', False, 'import torch\n'), ((116, 33, 116, 59), 'torch.zeros', 'torch.zeros', ({(116, 45, 116, 58): '(r_idx - l_idx)'}, {}), '(r_idx - l_idx)', False, 'import torch\n'), ((266, 51, 266, 85), 'torch.FloatTensor', 'torch.FloatTensor', ({(266, 69, 266, 75): 'd_proj', (266, 77, 266, 84): 'd_embed'}, {}), '(d_proj, d_embed)', False, 'import torch\n'), ((276, 51, 276, 85), 'torch.FloatTensor', 'torch.FloatTensor', ({(276, 69, 276, 75): 'd_proj', (276, 77, 276, 84): 'd_emb_i'}, {}), '(d_proj, d_emb_i)', False, 'import torch\n'), ((316, 20, 316, 60), 'torch.arange', 'torch.arange', (), '', False, 'import torch\n'), ((112, 37, 112, 65), 'torch.zeros', 'torch.zeros', ({(112, 49, 112, 55): 'd_proj', (112, 57, 112, 64): 'd_emb_i'}, {}), '(d_proj, d_emb_i)', False, 'import torch\n'), ((120, 37, 120, 72), 'torch.zeros', 'torch.zeros', ({(120, 49, 120, 62): '(r_idx - l_idx)', (120, 64, 120, 71): 'd_emb_i'}, {}), '(r_idx - l_idx, d_emb_i)', False, 'import torch\n'), ((89, 41, 89, 69), 'torch.zeros', 'torch.zeros', ({(89, 53, 89, 59): 'd_proj', (89, 61, 89, 68): 'd_embed'}, {}), '(d_proj, d_embed)', False, 'import torch\n'), ((171, 19, 171, 47), 'torch.nn.functional.log_softmax', 'F.log_softmax', (), '', True, 'import torch.nn.functional as F\n')]