# 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.

import torch


class LatencyMetric(object):
    @staticmethod
    def length_from_padding_mask(padding_mask, batch_first: bool = False):
        dim = 1 if batch_first else 0
        return padding_mask.size(dim) - padding_mask.sum(dim=dim, keepdim=True)

    def prepare_latency_metric(
        self,
        delays,
        src_lens,
        target_padding_mask=None,
        batch_first: bool = False,
        start_from_zero: bool = True,
    ):
        assert len(delays.size()) == 2
        assert len(src_lens.size()) == 2

        if start_from_zero:
            delays = delays + 1

        if batch_first:
            # convert to batch_last
            delays = delays.t()
            src_lens = src_lens.t()
            tgt_len, bsz = delays.size()
            _, bsz_1 = src_lens.size()

            if target_padding_mask is not None:
                target_padding_mask = target_padding_mask.t()
                tgt_len_1, bsz_2 = target_padding_mask.size()
                assert tgt_len == tgt_len_1
                assert bsz == bsz_2

        assert bsz == bsz_1

        if target_padding_mask is None:
            tgt_lens = tgt_len * delays.new_ones([1, bsz]).float()
        else:
            # 1, batch_size
            tgt_lens = self.length_from_padding_mask(target_padding_mask, False).float()
            delays = delays.masked_fill(target_padding_mask, 0)

        return delays, src_lens, tgt_lens, target_padding_mask

    def __call__(
        self,
        delays,
        src_lens,
        target_padding_mask=None,
        batch_first: bool = False,
        start_from_zero: bool = True,
    ):
        delays, src_lens, tgt_lens, target_padding_mask = self.prepare_latency_metric(
            delays, src_lens, target_padding_mask, batch_first, start_from_zero
        )
        return self.cal_metric(delays, src_lens, tgt_lens, target_padding_mask)

    @staticmethod
    def cal_metric(delays, src_lens, tgt_lens, target_padding_mask):
        """
        Expected sizes:
        delays: tgt_len, batch_size
        src_lens: 1, batch_size
        target_padding_mask: tgt_len, batch_size
        """
        raise NotImplementedError


class AverageProportion(LatencyMetric):
    """
    Function to calculate Average Proportion from
    Can neural machine translation do simultaneous translation?
    (https://arxiv.org/abs/1606.02012)

    Delays are monotonic steps, range from 1 to src_len.
    Give src x tgt y, AP is calculated as:

    AP = 1 / (|x||y]) sum_i^|Y| deleys_i
    """

    @staticmethod
    def cal_metric(delays, src_lens, tgt_lens, target_padding_mask):
        if target_padding_mask is not None:
            AP = torch.sum(
                delays.masked_fill(target_padding_mask, 0), dim=0, keepdim=True
            )
        else:
            AP = torch.sum(delays, dim=0, keepdim=True)

        AP = AP / (src_lens * tgt_lens)
        return AP


class AverageLagging(LatencyMetric):
    """
    Function to calculate Average Lagging from
    STACL: Simultaneous Translation with Implicit Anticipation
    and Controllable Latency using Prefix-to-Prefix Framework
    (https://arxiv.org/abs/1810.08398)

    Delays are monotonic steps, range from 1 to src_len.
    Give src x tgt y, AP is calculated as:

    AL = 1 / tau sum_i^tau delays_i - (i - 1) / gamma

    Where
    gamma = |y| / |x|
    tau = argmin_i(delays_i = |x|)
    """

    @staticmethod
    def cal_metric(delays, src_lens, tgt_lens, target_padding_mask):
        # tau = argmin_i(delays_i = |x|)
        tgt_len, bsz = delays.size()
        lagging_padding_mask = delays >= src_lens
        lagging_padding_mask = torch.nn.functional.pad(
            lagging_padding_mask.t(), (1, 0)
        ).t()[:-1, :]
        gamma = tgt_lens / src_lens
        lagging = (
            delays
            - torch.arange(delays.size(0))
            .unsqueeze(1)
            .type_as(delays)
            .expand_as(delays)
            / gamma
        )
        lagging.masked_fill_(lagging_padding_mask, 0)
        tau = (1 - lagging_padding_mask.type_as(lagging)).sum(dim=0, keepdim=True)
        AL = lagging.sum(dim=0, keepdim=True) / tau

        return AL


class DifferentiableAverageLagging(LatencyMetric):
    """
    Function to calculate Differentiable Average Lagging from
    Monotonic Infinite Lookback Attention for Simultaneous Machine Translation
    (https://arxiv.org/abs/1906.05218)

    Delays are monotonic steps, range from 0 to src_len-1.
    (In the original paper thery are from 1 to src_len)
    Give src x tgt y, AP is calculated as:

    DAL = 1 / |Y| sum_i^|Y| delays'_i - (i - 1) / gamma

    Where
    delays'_i =
        1. delays_i if i == 1
        2. max(delays_i, delays'_{i-1} + 1 / gamma)

    """

    @staticmethod
    def cal_metric(delays, src_lens, tgt_lens, target_padding_mask):
        tgt_len, bsz = delays.size()

        gamma = tgt_lens / src_lens
        new_delays = torch.zeros_like(delays)

        for i in range(delays.size(0)):
            if i == 0:
                new_delays[i] = delays[i]
            else:
                new_delays[i] = torch.cat(
                    [
                        new_delays[i - 1].unsqueeze(0) + 1 / gamma,
                        delays[i].unsqueeze(0),
                    ],
                    dim=0,
                ).max(dim=0)[0]

        DAL = (
            new_delays
            - torch.arange(delays.size(0))
            .unsqueeze(1)
            .type_as(delays)
            .expand_as(delays)
            / gamma
        )
        if target_padding_mask is not None:
            DAL = DAL.masked_fill(target_padding_mask, 0)

        DAL = DAL.sum(dim=0, keepdim=True) / tgt_lens

        return DAL


class LatencyMetricVariance(LatencyMetric):
    def prepare_latency_metric(
        self,
        delays,
        src_lens,
        target_padding_mask=None,
        batch_first: bool = True,
        start_from_zero: bool = True,
    ):
        assert batch_first
        assert len(delays.size()) == 3
        assert len(src_lens.size()) == 2

        if start_from_zero:
            delays = delays + 1

        # convert to batch_last
        bsz, num_heads_x_layers, tgt_len = delays.size()
        bsz_1, _ = src_lens.size()
        assert bsz == bsz_1

        if target_padding_mask is not None:
            bsz_2, tgt_len_1 = target_padding_mask.size()
            assert tgt_len == tgt_len_1
            assert bsz == bsz_2

        if target_padding_mask is None:
            tgt_lens = tgt_len * delays.new_ones([bsz, tgt_len]).float()
        else:
            # batch_size, 1
            tgt_lens = self.length_from_padding_mask(target_padding_mask, True).float()
            delays = delays.masked_fill(target_padding_mask.unsqueeze(1), 0)

        return delays, src_lens, tgt_lens, target_padding_mask


class VarianceDelay(LatencyMetricVariance):
    @staticmethod
    def cal_metric(delays, src_lens, tgt_lens, target_padding_mask):
        """
        delays : bsz, num_heads_x_layers, tgt_len
        src_lens : bsz, 1
        target_lens : bsz, 1
        target_padding_mask: bsz, tgt_len or None
        """
        if delays.size(1) == 1:
            return delays.new_zeros([1])

        variance_delays = delays.var(dim=1)

        if target_padding_mask is not None:
            variance_delays.masked_fill_(target_padding_mask, 0)

        return variance_delays.sum(dim=1, keepdim=True) / tgt_lens


class LatencyInference(object):
    def __init__(self, start_from_zero=True):
        self.metric_calculator = {
            "differentiable_average_lagging": DifferentiableAverageLagging(),
            "average_lagging": AverageLagging(),
            "average_proportion": AverageProportion(),
        }

        self.start_from_zero = start_from_zero

    def __call__(self, monotonic_step, src_lens):
        """
        monotonic_step range from 0 to src_len. src_len means eos
        delays: bsz, tgt_len
        src_lens: bsz, 1
        """
        if not self.start_from_zero:
            monotonic_step -= 1

        src_lens = src_lens

        delays = monotonic_step.view(
            monotonic_step.size(0), -1, monotonic_step.size(-1)
        ).max(dim=1)[0]

        delays = delays.masked_fill(delays >= src_lens, 0) + (src_lens - 1).expand_as(
            delays
        ).masked_fill(delays < src_lens, 0)
        return_dict = {}
        for key, func in self.metric_calculator.items():
            return_dict[key] = func(
                delays.float(),
                src_lens.float(),
                target_padding_mask=None,
                batch_first=True,
                start_from_zero=True,
            ).t()

        return return_dict


class LatencyTraining(object):
    def __init__(
        self,
        avg_weight,
        var_weight,
        avg_type,
        var_type,
        stay_on_last_token,
        average_method,
    ):
        self.avg_weight = avg_weight
        self.var_weight = var_weight
        self.avg_type = avg_type
        self.var_type = var_type
        self.stay_on_last_token = stay_on_last_token
        self.average_method = average_method

        self.metric_calculator = {
            "differentiable_average_lagging": DifferentiableAverageLagging(),
            "average_lagging": AverageLagging(),
            "average_proportion": AverageProportion(),
        }

        self.variance_calculator = {
            "variance_delay": VarianceDelay(),
        }

    def expected_delays_from_attention(
        self, attention, source_padding_mask=None, target_padding_mask=None
    ):
        if type(attention) == list:
            # bsz, num_heads, tgt_len, src_len
            bsz, num_heads, tgt_len, src_len = attention[0].size()
            attention = torch.cat(attention, dim=1)
            bsz, num_heads_x_layers, tgt_len, src_len = attention.size()
            # bsz * num_heads * num_layers, tgt_len, src_len
            attention = attention.view(-1, tgt_len, src_len)
        else:
            # bsz * num_heads * num_layers, tgt_len, src_len
            bsz, tgt_len, src_len = attention.size()
            num_heads_x_layers = 1
            attention = attention.view(-1, tgt_len, src_len)

        if not self.stay_on_last_token:
            residual_attention = 1 - attention[:, :, :-1].sum(dim=2, keepdim=True)
            attention = torch.cat([attention[:, :, :-1], residual_attention], dim=2)

        # bsz * num_heads_x_num_layers, tgt_len, src_len for MMA
        steps = (
            torch.arange(1, 1 + src_len)
            .unsqueeze(0)
            .unsqueeze(1)
            .expand_as(attention)
            .type_as(attention)
        )

        if source_padding_mask is not None:
            src_offset = (
                source_padding_mask.type_as(attention)
                .sum(dim=1, keepdim=True)
                .expand(bsz, num_heads_x_layers)
                .contiguous()
                .view(-1, 1)
            )
            src_lens = src_len - src_offset
            if source_padding_mask[:, 0].any():
                # Pad left
                src_offset = src_offset.view(-1, 1, 1)
                steps = steps - src_offset
                steps = steps.masked_fill(steps <= 0, 0)
        else:
            src_lens = attention.new_ones([bsz, num_heads_x_layers]) * src_len
            src_lens = src_lens.view(-1, 1)

        # bsz * num_heads_num_layers, tgt_len, src_len
        expected_delays = (
            (steps * attention).sum(dim=2).view(bsz, num_heads_x_layers, tgt_len)
        )

        if target_padding_mask is not None:
            expected_delays.masked_fill_(target_padding_mask.unsqueeze(1), 0)

        return expected_delays, src_lens

    def avg_loss(self, expected_delays, src_lens, target_padding_mask):

        bsz, num_heads_x_layers, tgt_len = expected_delays.size()
        target_padding_mask = (
            target_padding_mask.unsqueeze(1)
            .expand_as(expected_delays)
            .contiguous()
            .view(-1, tgt_len)
        )

        if self.average_method == "average":
            # bsz * tgt_len
            expected_delays = expected_delays.mean(dim=1)
        elif self.average_method == "weighted_average":
            weights = torch.nn.functional.softmax(expected_delays, dim=1)
            expected_delays = torch.sum(expected_delays * weights, dim=1)
        elif self.average_method == "max":
            # bsz * num_heads_x_num_layers, tgt_len
            expected_delays = expected_delays.max(dim=1)[0]
        else:
            raise RuntimeError(f"{self.average_method} is not supported")

        src_lens = src_lens.view(bsz, -1)[:, :1]
        target_padding_mask = target_padding_mask.view(bsz, -1, tgt_len)[:, 0]

        if self.avg_weight > 0.0:
            if self.avg_type in self.metric_calculator:
                average_delays = self.metric_calculator[self.avg_type](
                    expected_delays,
                    src_lens,
                    target_padding_mask,
                    batch_first=True,
                    start_from_zero=False,
                )
            else:
                raise RuntimeError(f"{self.avg_type} is not supported.")

            # bsz * num_heads_x_num_layers, 1
            return self.avg_weight * average_delays.sum()
        else:
            return 0.0

    def var_loss(self, expected_delays, src_lens, target_padding_mask):
        src_lens = src_lens.view(expected_delays.size(0), expected_delays.size(1))[
            :, :1
        ]
        if self.var_weight > 0.0:
            if self.var_type in self.variance_calculator:
                variance_delays = self.variance_calculator[self.var_type](
                    expected_delays,
                    src_lens,
                    target_padding_mask,
                    batch_first=True,
                    start_from_zero=False,
                )
            else:
                raise RuntimeError(f"{self.var_type} is not supported.")

            return self.var_weight * variance_delays.sum()
        else:
            return 0.0

    def loss(self, attention, source_padding_mask=None, target_padding_mask=None):
        expected_delays, src_lens = self.expected_delays_from_attention(
            attention, source_padding_mask, target_padding_mask
        )

        latency_loss = 0

        latency_loss += self.avg_loss(expected_delays, src_lens, target_padding_mask)

        latency_loss += self.var_loss(expected_delays, src_lens, target_padding_mask)

        return latency_loss