import math
import torch
from torch.nn.parameter import Parameter
from torch.nn import functional as F
from torch.nn import Module
#
# class CosineLinear(Module):
#     def __init__(self, in_features, out_features, sigma=True):
#         super(CosineLinear, self).__init__()
#         self.in_features = in_features
#         self.out_features = out_features
#         self.radius = Parameter(torch.Tensor(1))
#         self.weight = Parameter(torch.Tensor(out_features, in_features))
#
#         if sigma:
#             self.sigma = Parameter(torch.Tensor(1))
#         else:
#             self.register_parameter('sigma', None)
#         self.reset_parameters()
#
#     def reset_parameters(self):
#         stdv = 1. / math.sqrt(self.weight.size(1))
#         self.weight.data.uniform_(-stdv, stdv)
#         if self.sigma is not None:
#             self.sigma.data.fill_(1)
#         self.radius.data.fill_(0)
#
#     def forward(self, input):
#         out = F.linear(F.normalize(input, p=2,dim=1), F.normalize(self.weight, p=2, dim=1))
#         # import pdb;pdb.set_trace()
#         if self.sigma is not None:
#             out = self.sigma * out
#         return out
#
# class SplitCosineLinear(Module):
#     def __init__(self, in_features, out_features1, out_features2, sigma=True):
#         super(SplitCosineLinear, self).__init__()
#         self.in_features = in_features
#         self.out_features = out_features1 + out_features2
#         self.fc1 = CosineLinear(in_features, out_features1, False)
#         self.fc2 = CosineLinear(in_features, out_features2, False)
#
#         if sigma:
#             self.sigma = Parameter(torch.Tensor(1))
#             self.sigma.data.fill_(1)
#         else:
#             self.register_parameter('sigma', None)
#
#     def forward(self, x):
#         out1 = self.fc1(x)
#         out2 = self.fc2(x)
#         out = torch.cat((out1, out2), dim=1)
#         if self.sigma is not None:
#             out = self.sigma * out
#         return out



class CosineLinear(Module):
    def __init__(self, in_features, out_features, sigma=True):
        super(CosineLinear, self).__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.weight = Parameter(torch.Tensor(out_features, in_features))

        if sigma:
            self.sigma = Parameter(torch.Tensor(1))
        else:
            self.register_parameter('sigma', None)
        self.reset_parameters()

    def reset_parameters(self):
        stdv = 1. / math.sqrt(self.weight.size(1))
        self.weight.data.uniform_(-stdv, stdv)
        if self.sigma is not None:
            self.sigma.data.fill_(1)

    def forward(self, input):
        out = F.linear(input, self.weight)
        # import pdb;pdb.set_trace()
        if self.sigma is not None:
            out = self.sigma * out
        return out

class SplitCosineLinear(Module):
    def __init__(self, in_features, out_features1, out_features2, sigma=True):
        super(SplitCosineLinear, self).__init__()
        self.in_features = in_features
        self.out_features = out_features1 + out_features2
        self.fc1 = CosineLinear(in_features, out_features1, False)
        self.fc2 = CosineLinear(in_features, out_features2, False)

        if sigma:
            self.sigma = Parameter(torch.Tensor(1))
            self.sigma.data.fill_(1)
        else:
            self.register_parameter('sigma', None)

    def forward(self, x):
        out1 = self.fc1(x)
        out2 = self.fc2(x)
        out = torch.cat((out1, out2), dim=1)
        if self.sigma is not None:
            out = self.sigma * out
        return out
