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

class GumbelSoftmaxTopK(nn.Module):
    def __init__(self, k, tau):
        super(GumbelSoftmaxTopK, self).__init__()
        self.k = k
        self.tau = tau

    def forward(self, logits):
        logits = torch.log(nn.Softmax(dim=0)(logits).clamp(min=1e-8))

        """
            Adapted from: https://uvadlc-notebooks.readthedocs.io/en/latest/tutorial_notebooks/DL2/sampling/subsets.html#Subset-Sampler-Class
            and https://pytorch.org/docs/stable/_modules/torch/nn/functional.html#gumbel_softmax
        """
            
        m = torch.distributions.gumbel.Gumbel(torch.zeros_like(logits, memory_format=torch.legacy_contiguous_format),
                                              torch.ones_like(logits, memory_format= torch.legacy_contiguous_format))
        gumbels = m.sample()
        y = logits + gumbels
        _EPS = torch.tensor(1e-40).to(logits.device)
        khot = torch.zeros_like(logits, memory_format=torch.legacy_contiguous_format)
        onehot_approx = torch.zeros_like(logits, memory_format=torch.legacy_contiguous_format)

        for i in range(self.k):
            khot_mask = torch.maximum(1.0 - onehot_approx, _EPS)
            y += khot_mask.log()

            onehot_approx = nn.Softmax(dim=0)(y / self.tau)
            khot = torch.add(khot, onehot_approx)

        return khot

class GumbelSoftmax(nn.Module):
    def __init__(self, tau, hard=False):
        super(GumbelSoftmax, self).__init__()
        self.tau = tau
        self.hard = hard

    def forward(self, logits):
        m = torch.distributions.gumbel.Gumbel(torch.zeros_like(logits, memory_format=torch.legacy_contiguous_format),
                                              torch.ones_like(logits, memory_format= torch.legacy_contiguous_format))
        gumbels = m.sample()

        y = logits + gumbels
        y_soft = nn.Softmax(dim=0)(y / self.tau) 

        return y_soft