import math 

import torch
import torch.nn as nn
import torch.nn.functional as F


class AMSoftmax(nn.Module):
    def __init__(self, in_features, out_features, s=30.0, m=0.35):
        """
        in_features: размерность входных эмбеддингов
        out_features: количество классов
        s: масштабный множитель (scale)
        m: аддитивный margin
        """
        super(AMSoftmax, self).__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.s = s
        self.m = m
        self.weight = nn.Parameter(torch.FloatTensor(out_features, in_features))
        nn.init.xavier_uniform_(self.weight)
        self.use_labels_when_train = True

    def inference(self, x):
        x_norm = F.normalize(x, p=2, dim=1)
        w_norm = F.normalize(self.weight, p=2, dim=1)
        logits = F.linear(x_norm, w_norm) * self.s
        return logits

    def forward(self, x, labels=None):
        if not self.training or labels is None:
            return self.inference(x)
        # Нормализация входов и весов
        input_norm = F.normalize(x, p=2, dim=1)
        weight_norm = F.normalize(self.weight, p=2, dim=1)

        # Косинус угла между входами и центрами классов
        cosine = F.linear(input_norm, weight_norm)  # [batch_size, num_classes]

        # Скопировать для дальнейшего вычисления
        phi = cosine - self.m

        # Создать one-hot маску
        one_hot = torch.zeros_like(cosine)
        one_hot.scatter_(1, labels.view(-1, 1), 1.0)

        # Применить margin только к целевым логитам
        output = self.s * (one_hot * phi + (1.0 - one_hot) * cosine)
        return output



class AAMSoftmax(nn.Module):
    def __init__(self, in_features, out_features, s=30.0, m=0.50, easy_margin=False):
        """
        in_features: размерность эмбеддинга
        out_features: количество классов
        s: scale (обычно 30)
        m: angular margin (обычно 0.5 радиан)
        easy_margin: использовать "easy margin" трюк или нет
        """
        super(AAMSoftmax, self).__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.s = s
        self.m = m
        self.easy_margin = easy_margin

        self.weight = nn.Parameter(torch.FloatTensor(out_features, in_features))
        nn.init.xavier_uniform_(self.weight)

        self.cos_m = math.cos(m)
        self.sin_m = math.sin(m)
        self.th = math.cos(math.pi - m)
        self.mm = math.sin(math.pi - m) * m
        self.use_labels_when_train = True

    def inference(self, x):
        x_norm = F.normalize(x, p=2, dim=1)
        w_norm = F.normalize(self.weight, p=2, dim=1)
        logits = F.linear(x_norm, w_norm) * self.s
        return logits

    def forward(self, x, labels=None):
        if not self.training or labels is None:
            return self.inference(x)
        # Нормализуем входы и веса
        cosine = F.linear(F.normalize(x), F.normalize(self.weight))  # [B, C]
        sine = torch.sqrt(1.0 - cosine ** 2 + 1e-6)

        # cos(θ + m) = cosθ * cos(m) - sinθ * sin(m)
        phi = cosine * self.cos_m - sine * self.sin_m

        if self.easy_margin:
            # Используем "легкий" трюк, чтобы избежать неустойчивости
            phi = torch.where(cosine > 0, phi, cosine)
        else:
            # Ограничиваем phi снизу
            phi = torch.where(cosine > self.th, phi, cosine - self.mm)

        # One-hot метки
        one_hot = torch.zeros_like(cosine)
        one_hot.scatter_(1, labels.view(-1, 1), 1.0)

        # Вычисляем итоговый логит
        output = self.s * (one_hot * phi + (1.0 - one_hot) * cosine)
        return output


class RAMSoftmax(nn.Module):
    """
    Real Additive Margin Softmax (RAM-Softmax)

    Args:
        in_features: размерность входных эмбеддингов
        out_features: число классов
        s: scale-фактор для логитов
        m: additive margin
        eps: небольшой стабилизатор для sqrt
    """

    def __init__(self, in_features, out_features, s=30.0, m=0.35, eps=1e-6):
        super(RAMSoftmax, self).__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.s = s
        self.m = m
        self.eps = eps
        # веса центров классов
        self.weight = nn.Parameter(torch.FloatTensor(out_features, in_features))
        nn.init.xavier_uniform_(self.weight)

        # Большое отрицательное для «отсечки» легко разделённых классов
        self.register_buffer('large_neg', torch.tensor(-1e9))
        self.use_labels_when_train = True

    def inference(self, x):
        x_norm = F.normalize(x, p=2, dim=1)  # [B, D]
        w_norm = F.normalize(self.weight, p=2, dim=1)  # [C, D]

        # 2) косинус
        cosine = torch.matmul(x_norm, w_norm.t())  # [B, C]

        # 3) масштаб
        logits = cosine * self.s  # [B, C]
        return logits

    def forward(self, x, labels=None):
        if not self.training or labels is None:

            return self.inference(x)
        # 1) нормализуем эмбеддинги и веса
        x_norm = F.normalize(x, p=2, dim=1)  # [B, D]
        w_norm = F.normalize(self.weight, p=2, dim=1)  # [C, D]

        # 2) вычисляем все косинусы
        cosine = F.linear(x_norm, w_norm)  # [B, C]

        # 3) достаём целевой косинус и вычитаем margin
        idx = torch.arange(x.size(0), device=x.device)
        cos_y = cosine[idx, labels]  # [B]
        phi = cos_y - self.m  # [B]

        # 4) заменяем целевой логит на phi, остальные — оставляем как cosine
        logits = cosine.clone()
        logits[idx, labels] = phi

        # 5) маскирование «легко» разделённых: для каждого j≠y
        #    если cos_y - cos_j > m  => отсечь (логит -> large_neg)
        diff = cos_y.unsqueeze(1) - cosine  # [B, C]
        mask = (diff > self.m)  # [B, C]
        mask[idx, labels] = False  # не маскируем целевой
        logits = torch.where(mask, self.large_neg, logits)

        # 6) масштабируем
        logits = logits * self.s

        return logits


class RAAMSoftmax(nn.Module):
    def __init__(self, in_features, out_features, s=30.0, m=0.50, eps=1e-6):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.s = s
        self.m = m
        self.eps = eps

        self.weight = nn.Parameter(torch.FloatTensor(out_features, in_features))
        nn.init.xavier_uniform_(self.weight)

        self.cos_m = math.cos(m)
        self.sin_m = math.sin(m)
        self.large_neg = -1e9  # для masking

    def inference(self, x):
        x_norm = F.normalize(x, p=2, dim=1)
        w_norm = F.normalize(self.weight, p=2, dim=1)
        logits = F.linear(x_norm, w_norm) * self.s
        return logits
    
    def forward(self, x, labels=None):
        if not self.training or labels is None:
            return self.inference(x)
        x = F.normalize(x, dim=1)
        W = F.normalize(self.weight, dim=1)

        cosine = F.linear(x, W)  # [B, C]
        sine = torch.sqrt(1.0 - cosine ** 2 + self.eps)

        cos_theta_y = cosine[torch.arange(x.size(0)), labels]
        phi = cos_theta_y * self.cos_m - sine[torch.arange(x.size(0)), labels] * self.sin_m

        logits = cosine.clone()
        logits[torch.arange(x.size(0)), labels] = phi

        # RAM: отсечка легкоразделённых негативов
        diff = cos_theta_y.unsqueeze(1) - cosine
        mask = (diff > self.m)
        mask[torch.arange(x.size(0)), labels] = False

        logits = torch.where(mask, self.large_neg, logits)

        return self.s * logits


class WeightCrossEntropy(nn.Module):
    def __init__(self, id2name: list, class_distribution: dict):
        super().__init__()
        weight = torch.Tensor([1 / math.sqrt(class_distribution[name]) for name in id2name])
        self.ce = nn.CrossEntropyLoss(weight=weight)
    def forward(self, input, target):
        return self.ce(input, target)


class CrossEntropyLabelSmooth(nn.Module):
    """Cross entropy loss with label smoothing regularizer.

    Reference:
    Szegedy et al. Rethinking the Inception Architecture for Computer Vision. CVPR 2016.
    Equation: y = (1 - epsilon) * y + epsilon / K.

    Args:
        num_classes (int): number of classes.
        epsilon (float): weight.
    """

    def __init__(self, num_classes, epsilon=0.1, id2name=None, class_distribution=None):
        super(CrossEntropyLabelSmooth, self).__init__()
        self.num_classes = num_classes
        self.epsilon = epsilon
        self.logsoftmax = nn.LogSoftmax(dim=1)
        if id2name is not None and class_distribution is not None:
            weights = torch.Tensor([1 / math.sqrt(class_distribution[name]) for name in id2name])
            self.weights = weights.to(torch.float32)
        else:
            self.weights = None

    def forward(self, inputs, targets, use_label_smoothing=True):
        """
        Args:
            inputs: prediction matrix (before softmax) with shape (batch_size, num_classes)
            targets: ground truth labels with shape (b,)
        """
        #targets = torch.zeros(labels.size(0), self.num_classes).to(labels.device)
        #targets.scatter_(1, labels.unsqueeze(1), 1)
        #targets = targets.long()
        log_probs = self.logsoftmax(inputs)
        targets = torch.zeros(log_probs.size()).scatter_(1, targets.unsqueeze(1).data.cpu(), 1).to(targets.device)
        #if self.use_gpu: targets = targets #.to(torch.device('cuda'))
        if use_label_smoothing:
            targets = (1 - self.epsilon) * targets + self.epsilon / self.num_classes
        loss = (- targets * log_probs)
        if self.weights is not None:
            weights = self.weights.to(loss.device)
            loss = loss * weights.unsqueeze(0)  # (batch_size, num_classes)
        loss = loss.sum(dim=1).mean()
        return loss




class AMSoftmaxLoss(nn.Module):
    def __init__(self, in_features, out_features, num_classes, s=30.0, m=0.50, easy_margin=False, epsilon=0.1, id2name=None, class_distribution=None):
        super(AMSoftmaxLoss, self).__init__()
        self.aam = AAMSoftmax(in_features, out_features, s, m, easy_margin)
        self.criterion = CrossEntropyLabelSmooth(num_classes, epsilon, id2name, class_distribution)

    def forward(self, inputs, targets):
        return self.criterion(self.aam(inputs, targets), targets)