import numpy as np
import torch
from torch import nn

def create_test_data(samples, vocaburary, max_seqlen, lengths, sos, eos, pad):
    tokens = []
    for length in lengths:
        token = np.random.randint(low=0, high=vocaburary-4, size=length).tolist()
        token = [sos] + token + [eos]
        if len(token) < max_seqlen:
            token += [pad] * (max_seqlen - len(token))
        assert len(token) == max_seqlen
        tokens.append(token)
    tokens = np.array(tokens)

    logits = np.random.rand(samples, max_seqlen, vocaburary)

    # Convert to tensor.
    tokens = torch.from_numpy(tokens)
    logits = torch.from_numpy(logits)
    return logits, tokens

def compute_loss(loss_fn, logits, tokens):
    total_loss = 0
    seqlen = tokens.shape[-1]
    for t_index in range(1, seqlen):
        # Skip sos.
        logit = logits[:, t_index-1, :]
        token = tokens[:, t_index]
        loss = loss_fn(logit, token)
        if torch.isnan(loss).any():
            message = f"t_index:{t_index},\n logit:{logit},\n token:{token}"
            raise ValueError(message)
        total_loss += loss
    return total_loss

samples = 5
vocaburary = 13
sos = vocaburary - 3
eos = vocaburary - 2
pad = vocaburary - 1
max_seqlen = 10

lengths = np.random.randint(low=1, high=max_seqlen-2, size=samples)

lengths[lengths >= max_seqlen-2] = max_seqlen - 3
print(f"lengths of each sample:{lengths}")
# Create tokens.
logits, tokens = create_test_data(samples, vocaburary, max_seqlen, lengths,
                                  sos, eos, pad)
print(f"Shape of logits:{logits.shape}")
print(f"Tokens:{tokens}")

# Loss using torch's loss naively.
loss_fn = nn.CrossEntropyLoss(reduction="mean", ignore_index=pad)
try:
    total_loss = compute_loss(loss_fn, logits, tokens)
except ValueError as inst:
    print(type(inst))
    print(inst)

lengths[-1] = max_seqlen - 2
print(f"lengths of each sample:{lengths}")
# Create tokens.
logits, tokens = create_test_data(samples, vocaburary, max_seqlen, lengths,
                                  sos, eos, pad)
print(f"Shape of logits:{logits.shape}")
print(f"Tokens:{tokens}")

loss_fn = nn.CrossEntropyLoss(reduction="mean", ignore_index=pad)
try:
    total_loss = compute_loss(loss_fn, logits, tokens)
except ValueError as inst:
    print(inst)
print(f"Total loss (Naive):{total_loss}")

print(f"Total loss (Normalized by max_seqlen):{total_loss / max_seqlen}")

loss_fn = nn.CrossEntropyLoss(reduction="none", ignore_index=pad)
try:
    total_loss = compute_loss(loss_fn, logits, tokens)
except ValueError as inst:
    print(inst)
print(f"Total loss (Raw):{total_loss}")

total_loss /= (lengths + 1)  # Add <eos>.
print(f"Total loss (Normalized by samples lengths):{total_loss}")

total_loss = total_loss.mean()
print(f"Total loss (Averaged after normalized):{total_loss}")

class LabelSmoothingCrossEntropyLoss(nn.Module):
    """Cross-entropy loss with label smoothing.

    For the detail, please refer
    "Rethinking the Inception Architecture for Computer Vision"
    https://arxiv.org/abs/1512.00567
    """
    def __init__(self, weight=None, ignore_indices=None, reduction="none",
                 label_smoothing=0.0):
        super().__init__()
        self.weight = weight
        if isinstance(ignore_indices, int):
            self.ignore_indices = [ignore_indices]
        else:
            self.ignore_indices = ignore_indices
        assert reduction in ["none",
                             "mean_batch_prior", "mean_temporal_prior",
                             "sum"]
        self.reduction = reduction
        assert label_smoothing >= 0.0
        assert label_smoothing <= 1.0
        self.label_smoothing = label_smoothing

    def _isnotin_ignore(self, target):
        ignore = torch.tensor(self.ignore_indices, dtype=target.dtype,
                              device=target.device)
        isin = (target[..., None] == ignore).any(-1)
        return isin.bitwise_not()

    def _calc_loss(self, logit_t, target_t):
        logit_mask = torch.ones(logit_t.shape[-1],
                                dtype=logit_t.dtype,
                                device=logit_t.device)
        target_mask = torch.ones(target_t.shape,
                                 dtype=logit_t.dtype,
                                 device=logit_t.device)
        if self.ignore_indices is not None:
            logit_mask[self.ignore_indices] = 0
            target_mask = self._isnotin_ignore(target_t).float()
        if self.weight is None:
            weight = torch.ones(logit_t.shape[-1],
                                dtype=logit_t.dtype,
                                device=logit_t.device)
        else:
            weight = self.weight.to(dtype=logit_t.dtype, device=logit_t.device)
        # Calculate CE.
        logprobs = F.log_softmax(logit_t, dim=-1)
        logprobs_m = logprobs * weight * logit_mask
        nll_loss = -logprobs_m.gather(dim=-1, index=target_t.unsqueeze(1))
        nll_loss = nll_loss.squeeze(1)
        smooth_loss = -logprobs_m.sum(dim=-1) / logit_mask.sum()
        smooth_loss *= target_mask
        loss = (1 - self.label_smoothing) * nll_loss + self.label_smoothing * smooth_loss
        return loss

    def forward(self, logit, target):
        """Perform forward computation.

        # Args:
          - logit: `[N, C]` or `[N, C, T]`
          - target: `[N]` or [N, T]
        """
        # Check format.
        if len(logit.shape) == 2:
            logit = logit.unsqueeze(-1)
        if len(target.shape) == 1:
            target = target.unsqueeze(-1)
        assert len(logit.shape) == 3, f"{logit.shape}"
        assert len(target.shape) == 2, f"{target.shape}"
        assert logit.shape[0] == target.shape[0], f"{logit.shape, target.shape}"
        assert logit.shape[-1] == target.shape[-1], f"{logit.shape, target.shape}"

        loss = 0
        for t in range(target.shape[-1]):
            _loss = self._calc_loss(logit[:, :, t], target[:, t])
            # Reduction should be conducted in a loop when reduction is
            # mean_batch_prior.
            if self.reduction == "mean_batch_prior":
                if self.ignore_indices is not None:
                    denom = len([t for t in target[:, t]
                                 if t not in self.ignore_indices])
                else:
                    denom = logit.shape[0]
                _loss /= max(denom, 1)
            loss += _loss

        # Reduction.
        if self.reduction == "sum":
            loss = loss.sum()
        # Temporal Normalization.
        if self.reduction == "mean_batch_prior":
            loss = loss.sum() / target.shape[-1]
        if self.reduction == "mean_temporal_prior":
            target_lengths = self._isnotin_ignore(target).sum(dim=-1)
            loss /= torch.clamp(target_lengths, min=1)
            loss = loss.mean()
        return loss

loss_fn = LabelSmoothingCrossEntropyLoss(
    ignore_indices=pad, reduction="mean_temporal_prior")
# `[N, T, C] -> [N, C, T]`
logits = torch.permute(logits, [0, 2, 1])
total_loss = loss_fn(logits[:, :, :-1], tokens[:, 1:])
print(f"Total loss (LabelSmoothingCrossEntropyLoss): {total_loss}")