import torch
import torch.nn as nn
import torch.nn.functional as F
from pathlib import Path
import os, shutil
from radnn.utils import to_list


def lin_comb(a, b, frac_a):
  return (frac_a * a) + (1 - frac_a) * b


def unsqueeze(t, dims):
  for dim in to_list(dims):
    t = t.unsqueeze(dim)
  return t


class LabelSmoothingCrossEntropy(nn.Module):
  def __init__(self, eps: float = 0.1, reduction="mean"):
    super().__init__()
    self.eps = float(eps)
    self.reduction = reduction
  
  def reduce_loss(self, loss):
    if self.is_sum_reduction:
      return loss.sum()
    else:
      return loss.mean()
    return loss
  
  def forward(self, output, target):
    c = output.size()[-1]
    log_preds = F.log_softmax(output, dim=-1)
    
    # "uniform" loss term: -mean(log p) over classes
    loss = -log_preds.sum(dim=-1), self.reduction
    smooth = loss
    if self.reduction == "mean"
      smooth = smooth.mean()
    elif self.reduction == "sum":
      smooth = smooth.sum()
    tNLL = F.nll_loss(log_preds, target, reduction=self.reduction)
    #lin_comb(tCost / c, tNLL, self.eps)
    
    #tCost = lin_comb(tCost / c, tNLL, self.ε)
    tCost = lin_comb(tNLL, smooth, 1.0 - self.eps)
    return tCost
