#!/usr/bin/env python3

import logging

import torch
import torch.nn as nn
from fairseq.modules import (
    FairseqDropout,
    LayerNorm,
)

logger = logging.getLogger(__name__)


class CTC(nn.Module):
    
    def __init__(self, embed_dim, dictionary_size, dropout, need_layernorm=False):
        super(CTC, self).__init__()

        self.ctc_projection = nn.Linear(embed_dim, dictionary_size, bias=False)

        nn.init.normal_(
            self.ctc_projection.weight, mean=0, std=embed_dim ** -0.5
        )

        self.ctc_dropout_module = FairseqDropout(
            p=dropout, module_name=self.__class__.__name__
        )
        self.softmax = nn.Softmax(dim=-1)
        self.need_layernorm = need_layernorm
        if self.need_layernorm:
            self.LayerNorm = LayerNorm(embed_dim)

    def forward(self, x):
        if self.need_layernorm:
            x = self.LayerNorm(x)

        x = self.ctc_projection(self.ctc_dropout_module(x))
        return x

    @staticmethod
    def softmax(ctc_logit, temperature=1.0):
        return torch.nn.functional.softmax(ctc_logit / temperature, dim=-1)

    @staticmethod
    def log_softmax(ctc_logit, temperature=1.0):
        return torch.nn.functional.log_softmax(ctc_logit / temperature, dim=-1)

    @staticmethod
    def argmax(ctc_logit):
        return torch.argmax(ctc_logit, dim=-1)

