class FocalLoss(nn.Module):
    def __init__(self, alpha=0.25, gamma=2, reduction='mean'):
        super(FocalLoss, self).__init__()
        self.alpha = alpha
        self.gamma = gamma
        self.reduction = reduction

    def forward(self, output, target):
        if output.dim != target.dim:
            # convert target to onehot
            target = torch.zeros_like(output).scatter_(1, target.unsqueeze(1), 1)
        # convert output to presudo probability
        p = output.sigmoid()
        a = target * self.alpha + (1 - target) * (1 - self.alpha)
        pt = 1 - (target * p + (1 - target) * (1 - p))
        focal_weight = a * pt.pow(self.gamma)

        focal_loss = F.binary_cross_entropy(p, target, weight=focal_weight, reduction=self.reduction)
        return focal_loss

def lgb_focal_loss(y_pred, dtrain, alpha, gamma):
    targets = dtrain.label.reshape(-1)
    n_classes = len(np.unique(targets))
    # encode targets to one hot
    t = np.eye(n_classes)[targets]

    def focal_loss(x):
        # sigmoid => pseudo probability
        p = 1 / (1 + np.exp(-x))
        a = t * alpha + (1 - t) * (1 - alpha)
        g_coef = (1 - (t * p + (1 - t) * (1 - p))) ** gamma
        bce = t * np.log(p) + (1 - t) * np.log(1 - p)
        return - (a * g_coef * bce).sum(axis=1)

    grad = derivative(focal_loss, y_pred, n=1, dx=1e-6)
    hess = derivative(focal_loss, y_pred, n=2, dx=1e-6)
    return grad, hess