# AUTOGENERATED! DO NOT EDIT! File to edit: notebooks/03_CTC_loss_banded.ipynb (unless otherwise specified).

__all__ = ['device', 'window_start_positions', 'loss_masked', 'prepare_inputs', 'loss_basic', 'fwd_bwd', 'loss_py',
           'loss_cupy', 'cupy_funcs', 'soft_alignments', 'viterbi_alignments']

# Cell
from functools import partial
import torch
import numpy as np
import cupy as cp

from .core import semiring, Log, Max
from .utils import *
from seqdist import ctc

device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')

# Cell
def window_start_positions(base_alignments, input_lengths, width, num_states):
    t = torch.arange(base_alignments.size(0), device=base_alignments.device)
    base_alignments = torch.where(t[:, None] < input_lengths, base_alignments, num_states - 1)
    return torch.min(torch.max(base_alignments-width//2, torch.zeros_like(base_alignments)), (num_states-width))

def loss_masked(logits, targets, input_lengths, target_lengths, base_alignments, width):
    state_scores, repeat_mask, final_states, input_lengths = ctc.prepare_inputs(logits.log_softmax(2), targets, input_lengths, target_lengths)
    T, N, Lp = state_scores.shape
    window_starts = window_start_positions(base_alignments, input_lengths, width, target_lengths*2 + 1).unsqueeze(-1)
    states = torch.arange(Lp, device=logits.device)
    state_scores = torch.where((window_starts <= states) & (states < window_starts + width), state_scores, logits.new_full((1,), Log.zero))
    logz = ctc._logz_fwd(state_scores, repeat_mask, final_states, input_lengths)
    return -(logz / target_lengths).mean()

# Cell
def prepare_inputs(scores, targets, input_lengths, target_lengths, base_alignments, width:int):
    (T, N, C), device = scores.shape, scores.device
    num_states = target_lengths*2 + 1
    states = ctc.interleave_blanks(targets, blank_idx=0)
    repeat_mask = torch.nn.functional.pad(states[:, 2:] == states[:, :-2], (2, 2), value=0.)
    window_starts = window_start_positions(base_alignments, input_lengths, width, num_states)
    windows = torch.arange(width, device=device) + window_starts.unsqueeze(-1)

    state_scores = torch.gather(scores, 2, states.expand(T, -1, -1).gather(2, windows))
    final_states = (torch.stack([num_states - 2, num_states - 1], 1) - window_starts[input_lengths - 1, torch.arange(N), None])

    return state_scores, repeat_mask, final_states, input_lengths, window_starts

def _logz_fwd(state_scores, repeat_mask, final_states, input_lengths, window_starts, S:semiring=Log):
    T, N, W = state_scores.shape
    zeros = state_scores.new_full((N, W + 2), S.zero)
    alpha = state_scores.new_full((T + 1, N, W), S.zero)
    alpha[0, :, 0] = S.one
    prev_pos = window_starts.new_zeros(N, 1)
    range_W = torch.arange(W + 2, device=state_scores.device)
    for t in range(0, T):
        a = torch.cat([zeros[:, :2], alpha[t], zeros[:, :2]], dim=1)
        a[:, 2:] = S.sum(torch.stack([
                            a[:, 2:], a[:, 1:-1], torch.where(repeat_mask.gather(1, prev_pos + range_W), zeros, a[:, :-2])]), dim=0)
        pos = window_starts[t, :, None]
        alpha[t+1] = S.mul(state_scores[t], a.gather(1, (pos-prev_pos) + range_W[2:]))
        prev_pos = pos
    return S.sum(alpha[input_lengths, torch.arange(N)].gather(1, final_states), dim=1)

def loss_basic(logits, targets, input_lengths, target_lengths, base_alignments, width):
    log_probs = logits.log_softmax(2)
    logz = _logz_fwd(*prepare_inputs(log_probs, targets, input_lengths, target_lengths, base_alignments, width))
    return -(logz / target_lengths).mean()

# Cell
def fwd_bwd(state_scores, repeat_mask, final_states, input_lengths, window_starts, fwd_bwd_impl, S:semiring=Log):
    T, N, W = state_scores.shape
    alpha, beta = [state_scores.new_full((T+1, N, W), S.zero) for _ in range(2)]
    alpha[0, :, 0] = S.one
    beta[input_lengths, torch.arange(N)] = state_scores.new_full((N, W), S.zero).scatter_(1, final_states, S.one)
    window_starts = torch.cat([window_starts.new_zeros(1, N), window_starts], dim=0)
    alpha_T = fwd_bwd_impl(alpha, beta, state_scores, repeat_mask, input_lengths, window_starts, S)
    logz = S.sum(alpha_T.gather(1, final_states), dim=1)
    return alpha, beta, logz

def _fwd_bwd_py(alpha, beta, state_scores, repeat_mask, input_lengths, window_starts, S:semiring=Log):
    T, N, W = state_scores.shape
    zeros = state_scores.new_full((N, W + 2), S.zero)
    range_W = torch.arange(W + 2, device=state_scores.device)
    #fwd
    prev_pos = window_starts[0, :, None]
    for t in range(0, T):
        a = torch.cat([zeros[:, :2], alpha[t], zeros[:, :2]], dim=1)
        a[:, 2:] = S.sum(torch.stack([
                            a[:, 2:], a[:, 1:-1], torch.where(repeat_mask.gather(1, prev_pos + range_W), zeros, a[:, :-2])]), dim=0)
        pos = window_starts[t + 1, :, None]
        alpha[t+1] = S.mul(state_scores[t], a.gather(1, (pos-prev_pos) + range_W[2:]))
        prev_pos = pos
    #bwd
    for t in range(T, 0, -1):
        b = torch.cat([zeros[:, :2], S.mul(beta[t], state_scores[t-1]), zeros[:, :2]], dim=1)
        pos = window_starts[t - 1, :, None]
        b[:, :-2] = S.sum(torch.stack([b[:, :-2], b[:, 1:-1], torch.where(repeat_mask.gather(1, prev_pos + range_W), zeros, b[:, 2:])]), dim=0)
        beta[t-1, t <= input_lengths] = b.gather(1, (pos-prev_pos) + range_W[2:])[t <= input_lengths]
        prev_pos = pos

    return alpha[input_lengths, torch.arange(N)]

class _Logz(torch.autograd.Function):
    @staticmethod
    def forward(ctx, state_scores, repeat_mask, final_states, input_lengths, window_starts, fwd_bwd_impl):
        alpha, beta, logz = fwd_bwd(state_scores, repeat_mask, final_states, input_lengths, window_starts, fwd_bwd_impl, Log)
        ctx.save_for_backward(alpha, beta, input_lengths)
        return logz

    @staticmethod
    def backward(ctx, grad):
        alpha, beta, input_lengths = ctx.saved_tensors
        g = torch.softmax(alpha[1:] + beta[1:], dim=2) * ctc.masked_grad(grad.expand(alpha.size(0)-1, -1), input_lengths)
        return g, None, None, None, None, None

def loss_py(logits, targets, input_lengths, target_lengths, base_alignments, width):
    logz = _Logz.apply(*prepare_inputs(logits.log_softmax(2), targets, input_lengths, target_lengths, base_alignments, width), _fwd_bwd_py)
    return - (logz / target_lengths).mean()

# Cell
cupy_funcs = {
    (torch.float32, Log): load_cupy_func('cuda/ctc_banded.cu', 'fwd_bwd_banded', FLOAT='float',  SUM='logsumexp3', MUL='add', ZERO='{:E}'.format(Log.zero)),
    (torch.float64, Log): load_cupy_func('cuda/ctc_banded.cu', 'fwd_bwd_banded', FLOAT='double', SUM='logsumexp3', MUL='add', ZERO='{:E}'.format(Log.zero)),
}

def _fwd_bwd_cupy(alpha, beta, state_scores, repeat_mask, input_lengths, window_starts, S:ctc.semiring):
    T, N, W = state_scores.shape
    _, L = repeat_mask.shape #Note this L is Lp + 2
    alpha_T = torch.empty_like(alpha[0])
    with cp.cuda.Device(state_scores.device.index):
        cupy_funcs[(state_scores.dtype, S)](grid=(N, 2, 1), block=(W, 1, 1), shared_mem=2*8*W,
               args=(alpha_T.data_ptr(), alpha.data_ptr(), beta.data_ptr(), state_scores.data_ptr(), repeat_mask.data_ptr(),
                     input_lengths.data_ptr(), window_starts.data_ptr(), N, L, W))
    return alpha_T

def loss_cupy(logits, targets, input_lengths, target_lengths, base_alignments, width):
    logz = _Logz.apply(*prepare_inputs(logits.log_softmax(2), targets, input_lengths, target_lengths, base_alignments, width), _fwd_bwd_cupy)
    return - (logz / target_lengths).mean()

# Cell
cupy_funcs[(torch.float32, Max)] = load_cupy_func('cuda/ctc_banded.cu', 'fwd_bwd_banded', FLOAT='float',  SUM='max3', MUL='add', ZERO='{:E}'.format(Max.zero))
cupy_funcs[(torch.float64, Max)] = load_cupy_func('cuda/ctc_banded.cu', 'fwd_bwd_banded', FLOAT='double', SUM='max3', MUL='add', ZERO='{:E}'.format(Max.zero))

class _LogzViterbi(torch.autograd.Function):
    @staticmethod
    def forward(ctx, state_scores, repeat_mask, final_states, input_lengths, window_starts, fwd_bwd_impl):
        alpha, beta, logz = fwd_bwd(state_scores, repeat_mask, final_states, input_lengths, window_starts, fwd_bwd_impl, Max)
        ctx.save_for_backward(alpha, beta, input_lengths)
        return logz

    @staticmethod
    def backward(ctx, grad):
        alpha, beta, input_lengths = ctx.saved_tensors
        g = Max.dsum(alpha[1:] + beta[1:], dim=2) * ctc.masked_grad(grad.expand(alpha.size(0) - 1, -1), input_lengths)
        return g, None, None, None, None, None

def soft_alignments(logits, targets, input_lengths, target_lengths, base_alignments, width):
    state_scores, *args = prepare_inputs(logits.log_softmax(2), targets, input_lengths, target_lengths, base_alignments, width)
    _Logz.apply(state_scores.detach_().requires_grad_(), *args, _fwd_bwd_cupy).sum().backward()
    return state_scores.grad

def viterbi_alignments(logits, targets, input_lengths, target_lengths, base_alignments, width):
    state_scores, *args = prepare_inputs(logits.log_softmax(2), targets, input_lengths, target_lengths, base_alignments, width)
    _LogzViterbi.apply(state_scores.detach_().requires_grad_(), *args, _fwd_bwd_cupy).sum().backward()
    return state_scores.grad