import numpy as np
import pytest

import aesara
import aesara.tensor as aet
from aesara.tensor.nnet.ctc import (
    ConnectionistTemporalClassification,
    ctc,
    ctc_available,
)
from tests import unittest_tools as utt


def setup_torch_case():
    # Test obtained from Torch tutorial at:
    # https://github.com/baidu-research/warp-ctc/blob/master/torch_binding/TUTORIAL.md

    # Layout, from slowest to fastest changing dimension, is (time, batchSize, inputLayerSize)
    activations = np.asarray(
        [
            [[0, 0, 0, 0, 0], [1, 2, 3, 4, 5], [-5, -4, -3, -2, -1]],
            [[0, 0, 0, 0, 0], [6, 7, 8, 9, 10], [-10, -9, -8, -7, -6]],
            [[0, 0, 0, 0, 0], [11, 12, 13, 14, 15], [-15, -14, -13, -12, -11]],
        ],
        dtype=np.float32,
    )
    # Duration of each sequence
    activation_times = np.asarray([1, 3, 3], dtype=np.int32)
    # Labels for each sequence
    labels = np.asarray([[1, -1], [3, 3], [2, 3]], dtype=np.int32)

    expected_costs = np.asarray(
        [1.609437943, 7.355742931, 4.938849926], dtype=np.float32
    )

    grads = [
        [
            [0.2, -0.8, 0.2, 0.2, 0.2],
            [0.01165623125, 0.03168492019, 0.08612854034, -0.7658783197, 0.636408627],
            [-0.02115798369, 0.03168492019, -0.8810571432, 0.2341216654, 0.636408627],
        ],
        [
            [0, 0, 0, 0, 0],
            [-0.9883437753, 0.03168492019, 0.08612854034, 0.2341216654, 0.636408627],
            [-0.02115798369, 0.03168492019, -0.1891518533, -0.4577836394, 0.636408627],
        ],
        [
            [0, 0, 0, 0, 0],
            [0.01165623125, 0.03168492019, 0.08612854034, -0.7658783197, 0.636408627],
            [-0.02115798369, 0.03168492019, 0.08612854034, -0.7330639958, 0.636408627],
        ],
    ]
    expected_gradients = np.asarray(grads, dtype=np.float32)

    return [activations, labels, activation_times, expected_costs, expected_gradients]


def setup_ctc_case():
    activations = np.asarray(
        [
            [[0.1, 0.6, 0.1, 0.1, 0.1], [0.1, 0.1, 0.6, 0.1, 0.1]],
            [[0.6, 0.1, 0.1, 0.1, 0.1], [0.1, 0.1, 0.5, 0.2, 0.1]],
        ],
        dtype=np.float32,
    )

    activation_times = np.asarray([2, 2], dtype=np.int32)

    labels = np.asarray([[1, 2], [1, 2]], dtype=np.int32)

    expected_costs = np.asarray([2.962858438, 3.053659201], dtype=np.float32)

    grads = [
        [
            [0.177031219, -0.7081246376, 0.177031219, 0.177031219, 0.177031219],
            [0.177031219, -0.8229685426, 0.291875124, 0.177031219, 0.177031219],
        ],
        [
            [0.291875124, 0.177031219, -0.8229685426, 0.177031219, 0.177031219],
            [0.1786672771, 0.1786672771, -0.7334594727, 0.1974578798, 0.1786672771],
        ],
    ]

    expected_gradients = np.asarray(grads, dtype=np.float32)

    return [activations, labels, activation_times, expected_costs, expected_gradients]


def setup_grad_case():
    activations = np.asarray(
        [
            [[0.1, 0.6, 0.1, 0.1, 0.1], [0.1, 0.1, 0.6, 0.1, 0.1]],
            [[0.6, 0.1, 0.1, 0.1, 0.1], [0.1, 0.1, 0.5, 0.2, 0.1]],
        ],
        dtype=np.float32,
    )

    activation_times = np.asarray([2, 2], dtype=np.int32)

    labels = np.asarray([[1, 2], [1, 2]], dtype=np.int32)

    return [activations, labels, activation_times]


@pytest.mark.skipif(
    not ctc_available(), reason="Optional library warp-ctc not available"
)
@pytest.mark.skipif(
    aesara.config.mode == "FAST_COMPILE" or aesara.config.cxx == "",
    reason="We need a c compiler",
)
class TestCTC:
    """
    Test Baidu CTC wrapper implementation.

    Expected values for costs and gradients are obtained through an external
    C implementation, that uses the library directly.
    """

    def run_ctc(
        self, activations, labels, input_length, expected_costs, expected_grads
    ):
        # Create symbolic variables
        t_activations = aesara.shared(activations, name="activations")
        t_activation_times = aesara.shared(input_length, name="activation_times")
        t_labels = aesara.shared(labels, name="labels")

        t_cost = ctc(t_activations, t_labels, t_activation_times)
        # Symbolic gradient of CTC cost
        t_grad = aet.grad(aet.mean(t_cost), t_activations)
        # Compile symbolic functions
        train = aesara.function([], [t_cost, t_grad])

        cost, grad = train()

        utt.assert_allclose(expected_grads / cost.shape[0], grad)
        utt.assert_allclose(expected_costs, cost)

        self.check_grads_disabled(t_activations, t_labels, t_activation_times)

    def check_grads_disabled(self, activations, labels, input_length):
        """
        Check if optimization to disable gradients is working
        """
        ctc_cost = ctc(activations, labels, input_length)
        ctc_function = aesara.function([], [ctc_cost])
        for node in ctc_function.maker.fgraph.apply_nodes:
            if isinstance(node.op, ConnectionistTemporalClassification):
                assert node.op.compute_grad is False

    def test_torch_case(self):
        (
            activations,
            labels,
            input_length,
            expected_costs,
            expected_grads,
        ) = setup_torch_case()
        self.run_ctc(activations, labels, input_length, expected_costs, expected_grads)

    def test_ctc(self):
        (
            activations,
            labels,
            input_length,
            expected_costs,
            expected_grads,
        ) = setup_ctc_case()
        self.run_ctc(activations, labels, input_length, expected_costs, expected_grads)

    def test_verify_grad(self):
        def ctc_op_functor(labels, in_lengths):
            def wrapper(acts):
                # Create auxiliary symbolic variables
                t_activation_times = aesara.shared(in_lengths, name="activation_times")
                t_labels = aesara.shared(labels, name="labels")
                return ctc(acts, t_labels, t_activation_times)

            return wrapper

        activations, labels, activation_times = setup_grad_case()

        ctc_op = ctc_op_functor(labels, activation_times)

        utt.verify_grad(ctc_op, [activations])
