# Copyright (c) 2021 Graphcore Ltd. All rights reserved.
import collections
import math

import numpy as np
import pytest
import torch
import torch.nn.functional as F

NEG_INF = -float("inf")


@pytest.mark.parametrize(
    "max_time,batch_size,num_classes,blank,beam_width,top_paths",
    [
        (10, 5, 15, 0, 10, 1),
        (10, 5, 15, 14, 10, 1),
        (10, 5, 15, 14, 10, 10),
    ],
)
def test_ctc_beam_search_decoder(
    op_tester, max_time, batch_size, num_classes, blank, beam_width, top_paths
):
    """Test that the CTC beam search decoder is interfaced correctly.

    Rigorous tests for correctness are done lower in the stack, here only the
    interface to the decoder is tested.
    """
    torch.manual_seed(42)

    # Input data.
    logits_data = torch.rand(max_time, batch_size, num_classes)
    log_probs_data = F.log_softmax(logits_data, 2, dtype=torch.float32).numpy()
    data_lengths_data = np.full((batch_size,), max_time, np.uint32)

    def init_builder(builder):
        log_probs = builder.addInputTensor(log_probs_data)
        data_lengths = builder.addInputTensor(data_lengths_data)

        (
            label_probs,
            label_lengths,
            decoded_labels,
        ) = builder.aiGraphcore.ctcbeamsearchdecoder(
            [log_probs, data_lengths], blank, beam_width, top_paths
        )

        builder.addOutputTensor(label_probs)  # batchSize, topPaths
        builder.addOutputTensor(label_lengths)  # batchSize, topPaths
        builder.addOutputTensor(decoded_labels)  # batchSize, topPaths, maxTime

        return [label_probs, label_lengths, decoded_labels]

    def reference(ref_data):
        # Initialise empty arrays to store reference data.
        target_label_probs = np.empty((batch_size, top_paths), dtype=np.float32)
        target_label_lengths = np.empty((batch_size, top_paths), dtype=np.uint32)
        target_decoded_labels = np.empty(
            (batch_size, top_paths, max_time), dtype=np.uint32
        )

        # Generate the reference data.
        for b in range(batch_size):
            (
                target_label_probs[b, :],
                target_label_lengths[b, :],
                target_decoded_labels[b, :, :],
            ) = reference_decode(log_probs_data[:, b, :], beam_width, top_paths, blank)

        # The data in decoded_labels after the indices specified in
        # label_lengths is arbitrary, and should not be used for testing. Here
        # we substitute this data in the target_decoded_labels with the one that
        # was generated by our CTC decoder.
        decoded_labels_mask = np.zeros((batch_size, top_paths, max_time), dtype=np.bool)
        for p in range(top_paths):
            for b in range(batch_size):
                decoded_labels_mask[b, p, target_label_lengths[b, p] :] = 1
        # Copy the unwanted decoded_labels from our CTC decoder to the reference
        # data.
        decoded_labels = ref_data.getOutputTensor(2)
        target_decoded_labels[decoded_labels_mask] = decoded_labels[decoded_labels_mask]

        return [target_label_probs, target_label_lengths, target_decoded_labels]

    op_tester.setPatterns([], enableRuntimeAsserts=False)
    op_tester.run(init_builder, reference, step_type="infer")


def make_new_beam():
    """Utility function used in reference_decode()."""
    fn = lambda: (NEG_INF, NEG_INF)
    return collections.defaultdict(fn)


def logsumexp(*args):
    """Stable log sum exp used in reference_decode()."""
    if all(a == NEG_INF for a in args):
        return NEG_INF
    a_max = max(args)
    lsp = math.log(sum(math.exp(a - a_max) for a in args))
    return a_max + lsp


def reference_decode(log_probs, beam_width, top_paths, blank):
    """Performs CTC for the given log probabilities using prefix beam search.

    This implementation was adapted from
    [this gist](https://gist.github.com/awni/56369a90d03953e370f3964c826ed4b0)
    and cited on distil [here](https://distill.pub/2017/ctc/).


    Args:
        log_probs (np.array): The output log probabilities (e.g. post-logsoftmax) for each
            time step. Should be an array of shape (time x output dim).
        beam_width (int): Size of the beam to use during inference.
        top_paths (int): Number of top paths to return.
        blank (int): Index of the CTC blank label.

    Returns:
        tuple: Returns the output label sequence and the corresponding length
            and log-likelihood estimated by the decoder.
    """

    T, S = log_probs.shape

    beam = [(tuple(), (0.0, NEG_INF))]

    for t in range(T):
        next_beam = make_new_beam()

        for s in range(S):
            p = log_probs[t, s]

            for prefix, (p_b, p_nb) in beam:
                if s == blank:
                    n_p_b, n_p_nb = next_beam[prefix]
                    n_p_b = logsumexp(n_p_b, p_b + p, p_nb + p)
                    next_beam[prefix] = (n_p_b, n_p_nb)
                    continue

                end_t = prefix[-1] if prefix else None
                n_prefix = prefix + (s,)
                n_p_b, n_p_nb = next_beam[n_prefix]
                if s != end_t:
                    n_p_nb = logsumexp(n_p_nb, p_b + p, p_nb + p)
                else:
                    n_p_nb = logsumexp(n_p_nb, p_b + p)

                next_beam[n_prefix] = (n_p_b, n_p_nb)

                if s == end_t:
                    n_p_b, n_p_nb = next_beam[prefix]
                    n_p_nb = logsumexp(n_p_nb, p_nb + p)
                    next_beam[prefix] = (n_p_b, n_p_nb)

        beam = sorted(next_beam.items(), key=lambda x: logsumexp(*x[1]), reverse=True)
        beam = beam[:beam_width]

    max_time = log_probs.shape[0]
    label_probs = [logsumexp(*beam[i][1]) for i in range(top_paths)]
    label_lengths = [len(beam[i][0]) for i in range(top_paths)]
    decoded_labels = [
        beam[i][0] + tuple([0] * (max_time - len(beam[i][0]))) for i in range(top_paths)
    ]

    return label_probs, label_lengths, decoded_labels