import torch
import custom_ctc_cpp

Tensor = torch.Tensor

class CustomCTCLossFunction(torch.autograd.Function):
    @staticmethod
    def forward(
        ctx,
        log_probs: Tensor,
        targets: Tensor,
        realval: Tensor,
        targets_realval: Tensor,
        input_lengths: Tensor,
        target_lengths: Tensor,
        sigma: float = 1,
        blank: int = 0,
        blank1: int = 0,
        reduction: str = "mean",
        zero_infinity: bool = False
        ):
        assert reduction in ['none', 'mean']
        if isinstance(input_lengths, list) :
            input_lengths = Tensor(input_lengths).long().to(log_probs.device)
        if isinstance(target_lengths, list) :
            target_lengths = Tensor(target_lengths).long().to(log_probs.device)
        neg_log_likelihood, log_alpha = custom_ctc_cpp.forward(log_probs, targets, realval, targets_realval, input_lengths, target_lengths, sigma, blank, blank1, zero_infinity)
        ctx.save_for_backward(neg_log_likelihood, log_alpha, log_probs, targets, realval, targets_realval, input_lengths, target_lengths)
        ctx.blank = blank
        ctx.blank1 = blank1
        ctx.zero_infinity = zero_infinity
        ctx.sigma = sigma
        ctx.reduction = reduction
        if reduction == 'mean' :
            return (neg_log_likelihood / target_lengths.clamp_min(1)).mean()
        return neg_log_likelihood

    @staticmethod
    def backward(ctx, grad_out):
        neg_log_likelihood, log_alpha, log_probs, targets, realval, targets_realval, input_lengths, target_lengths = ctx.saved_tensors
        if ctx.reduction == 'mean' :
            if grad_out.numel() == 0 :
                grad_out = torch.ones_like(neg_log_likelihood)
            else :
                grad_out = grad_out.view(1).tile(neg_log_likelihood.size(0))
            grad_out /= target_lengths.clamp_min(1)
            grad_out /= log_probs.size(0)
        outputs_cls, outputs_realval = custom_ctc_cpp.backward(grad_out, log_probs, targets, realval, targets_realval, input_lengths, target_lengths, neg_log_likelihood, log_alpha, ctx.sigma, ctx.blank, ctx.blank1, ctx.zero_infinity)
        return outputs_cls, None, outputs_realval, None, None, None, None, None, None, None, None