# Copyright (c) Alibaba, Inc. and its affiliates.

import torch

from .unified_transformer import UnifiedTransformer


class GenUnifiedTransformer(UnifiedTransformer):
    """
    Implement generation unified transformer.
    """

    def __init__(self, model_dir, config, reader, generator):
        super(GenUnifiedTransformer, self).__init__(model_dir, config, reader,
                                                    generator)
        self.understand = config.BPETextField.understand

        if self.use_gpu:
            self.cuda()
        return

    def _forward(self, inputs, is_training, with_label):
        """ Real forward process of model in different mode(train/test). """

        def cat(x, y, dim=1):
            return torch.cat([x, y], dim=dim)

        outputs = {}

        if self.understand or self.policy:
            if self.understand:
                prompt_token = inputs['understand_token']
                prompt_mask = inputs['understand_mask']
                if self.policy:
                    prompt_token = cat(prompt_token, inputs['policy_token'])
                    prompt_mask = cat(prompt_mask, inputs['policy_mask'])
            else:
                prompt_token = inputs['policy_token']
                prompt_mask = inputs['policy_mask']

            enc_embed, dec_embed, prompt_embed = self._encoder_prompt_decoder_network(
                src_token=inputs['src_token'],
                src_mask=inputs['src_mask'],
                tgt_token=inputs['tgt_token'][:, :-1],
                tgt_mask=inputs['tgt_mask'][:, :-1],
                prompt_token=prompt_token,
                prompt_mask=prompt_mask,
                src_pos=inputs['src_pos'],
                src_type=inputs['src_type'],
                src_turn=inputs['src_turn'],
                tgt_pos=inputs['tgt_pos'][:, :-1],
                tgt_type=inputs['tgt_type'][:, :-1],
                tgt_turn=inputs['tgt_turn'][:, :-1])
        else:
            enc_embed, dec_embed = self._encoder_decoder_network(
                src_token=inputs['src_token'],
                src_mask=inputs['src_mask'],
                tgt_token=inputs['tgt_token'][:, :-1],
                tgt_mask=inputs['tgt_mask'][:, :-1],
                src_pos=inputs['src_pos'],
                src_type=inputs['src_type'],
                src_turn=inputs['src_turn'],
                tgt_pos=inputs['tgt_pos'][:, :-1],
                tgt_type=inputs['tgt_type'][:, :-1],
                tgt_turn=inputs['tgt_turn'][:, :-1])

        outputs['dec_probs'] = self._dec_head(dec_embed=dec_embed)
        return outputs

    def _collect_metrics(self, inputs, outputs, with_label, data_file):

        metrics = {}
        loss = 0.

        label = inputs['tgt_token'][:, 1:]
        token_num = torch.sum(torch.sum(inputs['tgt_mask'], dim=1) - 1)
        nll = self.nll_loss(
            torch.log(outputs['dec_probs'] + 1e-12).permute(0, 2, 1), label)
        nll = torch.sum(nll, dim=1)
        token_nll = torch.sum(nll) / token_num
        nll = torch.mean(nll)
        metrics['nll'] = nll
        metrics['token_nll'] = token_nll
        metrics['token_num'] = token_num
        loss = loss + (token_nll if self.token_loss else nll)

        metrics['loss'] = loss
        if self.gpu > 1:
            return nll, token_nll, token_num
        else:
            return metrics

    def _optimize(self, loss, do_update=False, optimizer=None):
        """ Optimize loss function and update model. """
        assert optimizer is not None

        if self.gradient_accumulation_steps > 1:
            loss = loss / self.gradient_accumulation_steps

        loss.backward()

        if self.grad_clip is not None and self.grad_clip > 0:
            torch.nn.utils.clip_grad_norm_(
                parameters=self.parameters(), max_norm=self.grad_clip)

        if do_update:
            optimizer.step()
            optimizer.zero_grad()

        return

    def _init_state(self,
                    src_token,
                    src_mask,
                    src_pos=None,
                    src_type=None,
                    src_turn=None):
        """ Initialize decode state. """
        state = {}
        batch_size = src_token.shape[0]

        src_embed = self.embedder(src_token, src_pos, src_type, src_turn)
        src_embed = self.embed_layer_norm(src_embed)

        mask = self._create_mask(src_mask, append_head=False)

        enc_out = src_embed

        cache = {}
        for _l, layer in enumerate(self.layers):
            cache[f'layer_{_l}'] = {}
            enc_out = layer(enc_out, mask, cache[f'layer_{_l}'])

        state['cache'] = cache
        state['mask'] = mask[:, :1]
        state['batch_size'] = batch_size
        shape = [batch_size, 1, 1]
        state['pred_mask'] = torch.ones(shape, dtype=torch.float32)
        state['pred_pos'] = torch.zeros(shape, dtype=torch.int64)
        state['pred_type'] = torch.zeros(shape, dtype=torch.int64)
        state['pred_turn'] = torch.zeros(shape, dtype=torch.int64)
        if self.use_gpu:
            state['pred_mask'] = state['pred_mask'].cuda()
            state['pred_pos'] = state['pred_pos'].cuda()
            state['pred_type'] = state['pred_type'].cuda()
            state['pred_turn'] = state['pred_turn'].cuda()

        return state

    def _init_prompt_state(self,
                           src_token,
                           src_mask,
                           prompt_token,
                           prompt_mask,
                           src_pos=None,
                           src_type=None,
                           src_turn=None,
                           prompt_pos=None,
                           prompt_type=None,
                           prompt_turn=None):
        """ Initialize decode state. """
        state = {}
        batch_size = src_token.shape[0]

        src_embed = self.embedder(src_token, src_pos, src_type, src_turn)
        prompt_embed = self.embedder(prompt_token, prompt_pos, prompt_type,
                                     prompt_turn)
        embed = torch.cat([src_embed, prompt_embed], dim=1)
        embed = self.embed_layer_norm(embed)
        enc_out = embed

        enc_mask = self._create_mask(src_mask, auto_regressive=False)
        dec_mask = self._create_mask(prompt_mask, auto_regressive=True)
        mask = self._join_mask(enc_mask, dec_mask)

        cache = {}
        for _l, layer in enumerate(self.layers):
            cache[f'layer_{_l}'] = {}
            enc_out = layer(enc_out, mask, cache[f'layer_{_l}'])

        state['cache'] = cache
        state['mask'] = mask[:, -1:]  # state["mask"] = mask[:, :1]
        state['batch_size'] = batch_size
        shape = [batch_size, 1, 1]
        state['pred_mask'] = torch.ones(shape, dtype=torch.float32)
        state['pred_pos'] = torch.zeros(shape, dtype=torch.int64)
        state['pred_type'] = torch.zeros(shape, dtype=torch.int64)
        state['pred_turn'] = torch.zeros(shape, dtype=torch.int64)
        if self.use_gpu:
            state['pred_mask'] = state['pred_mask'].cuda()
            state['pred_pos'] = state['pred_pos'].cuda()
            state['pred_type'] = state['pred_type'].cuda()
            state['pred_turn'] = state['pred_turn'].cuda()

        return state

    def _decode(self, state):
        """ Decoding one time stamp. """

        # shape: [batch_size, 1, seq_len]
        mask = state['mask']

        # shape: [batch_size, 1, 1]
        pred_token = state['pred_token']
        pred_mask = state['pred_mask']
        pred_pos = state['pred_pos']
        pred_type = state['pred_type']
        pred_turn = state['pred_turn']

        # list of shape(len: num_layers): [batch_size, seq_len, hidden_dim]
        cache = state['cache']

        pred_embed = self.embedder(pred_token, pred_pos, pred_type,
                                   pred_turn).squeeze(-2)
        pred_embed = self.embed_layer_norm(pred_embed)

        # shape: [batch_size, 1, seq_len + 1]
        mask = torch.cat([mask, 1 - pred_mask], dim=2)

        # shape: [batch_size, 1, hidden_dim]
        for _l, layer in enumerate(self.layers):
            pred_embed = layer(pred_embed, mask, cache[f'layer_{_l}'])

        # shape: [batch_size, vocab_size]
        pred_probs = self._dec_head(dec_embed=pred_embed[:, 0])
        pred_logits = torch.log(pred_probs)

        state['mask'] = mask
        return pred_logits, state

    def _infer(self,
               inputs,
               start_id=None,
               eos_id=None,
               max_gen_len=None,
               prev_input=None):
        """ Real inference process of model. """

        def cat(x, y, dim=1):
            return torch.cat([x, y], dim=dim)

        # Initial decode state.
        if self.understand or self.policy:
            if self.understand:
                prompt_token = inputs['understand_token']
                prompt_mask = inputs['understand_mask']
                if self.policy:
                    prompt_token = cat(prompt_token, inputs['policy_token'])
                    prompt_mask = cat(prompt_mask, inputs['policy_mask'])
            else:
                prompt_token = inputs['policy_token']
                prompt_mask = inputs['policy_mask']

            state = self._init_prompt_state(
                src_token=inputs['src_token'],
                src_mask=inputs['src_mask'],
                prompt_token=prompt_token,
                prompt_mask=prompt_mask,
                src_pos=inputs['src_pos'],
                src_type=inputs['src_type'],
                src_turn=inputs['src_turn'])
        else:
            state = self._init_state(
                src_token=inputs['src_token'],
                src_mask=inputs['src_mask'],
                src_pos=inputs['src_pos'],
                src_type=inputs['src_type'],
                src_turn=inputs['src_turn'])

        # Generation process.
        gen_results = self.generator(
            step_fn=self._decode,
            state=state,
            start_id=start_id,
            eos_id=eos_id,
            max_gen_len=max_gen_len,
            prev_input=prev_input)

        outputs = gen_results['preds']
        return outputs


GenUnifiedTransformer.register('GenUnifiedTransformer')
