import torch
import torch.nn as nn
import torch.nn.init as init
import torch.nn.functional as F
import numpy as np

class L0Mask(nn.Module):
    def __init__(self, mask_dim, mask_p = 0.9):
        super().__init__()
        self.mask_setting = 'mask'
        self.mask_scores = nn.Parameter(torch.zeros(mask_dim))
        self.mask_p = mask_p
        self.l, self.r, self.b = -0.1, 1.1, 2 / 3
        self.init_weights()

    def init_weights(self):
        p = (self.mask_p - self.l) / (self.r - self.l)
        init.constant_(self.mask_scores, val=np.log(p / (1 - p)))
        
    def set_temperature(self, temp):
        self.b = temp
        
    def produce_mask(self):
        if self.training:
            u = torch.zeros_like(self.mask_scores).uniform_().clamp(0.0001, 0.9999)
            s = torch.sigmoid((u.log() - (1 - u).log() + self.mask_scores) / self.b)
        else:
            s = torch.sigmoid(self.mask_scores)
        s_bar = s * (self.r - self.l) + self.l
        mask = s_bar.clamp(min=0.0, max=1.0)
        return mask
    
    def regularizer(self):
        return torch.sum(torch.sigmoid(self.mask_scores - self.b * np.log(-self.l / self.r))) / self.mask_scores.numel()

class MaskedLinear(nn.Linear):
    def __init__(self, in_features: int, out_features: int, bias: bool = True, 
                 mask_p: float = 0.9, out_w_per_mask = 1, in_w_per_mask = 1, num_heads = 12):
        super().__init__(in_features=in_features, out_features=out_features, bias=bias)
        self.num_heads = num_heads
        self.out_w_per_mask = out_w_per_mask
        self.in_w_per_mask = in_w_per_mask
        
        assert out_features % out_w_per_mask == 0, "{} % {} not 0".format(out_features, out_w_per_mask)
        assert in_features % in_w_per_mask == 0, "{} % {} not 0".format(in_features, in_w_per_mask)
        mask_dim = (1, out_features // out_w_per_mask, 1, in_features // in_w_per_mask)
        self.mask = L0Mask(mask_dim)
        
        self.cached_activation = None
        self.do_caching = False

    def produce_mask_reshaped(self):
        mask = self.mask.produce_mask()
        mask = mask.repeat(self.out_w_per_mask, 1, self.in_w_per_mask, 1)
        return mask.reshape(self.out_features, self.in_features)

    def produce_mask(self):
        mask = self.mask.produce_mask()
        return mask

    def forward(self, input: torch.tensor):
        # "masked_weight = self.produce_mask_reshaped() * self.weight" is equivalent but slower.
        masked_weight = self.produce_mask() * self.weight.reshape(
            self.out_w_per_mask, self.out_features // self.out_w_per_mask,
            self.in_w_per_mask, self.in_features // self.in_w_per_mask)
        masked_weight = masked_weight.reshape(self.out_features, self.in_features)
        
        act = F.linear(input, masked_weight, self.bias)
        if self.do_caching:
            if self.cached_activation is None:
                self.cached_activation = act.detach()
            else: # only works if subbatched, since maxlen must be constant
                self.cached_activation = torch.cat((
                    self.cached_activation,act.detach()), dim = 0)
        return act
    
    def activate_caching(self, caching = True):
        self.cached_activation = None
        self.do_caching = caching

    @classmethod
    def from_layer(cls, layer, out_w_per_mask, in_w_per_mask):
        assert type(layer) == nn.modules.linear.Linear
        res = cls(in_features = layer.in_features, out_features = layer.out_features, bias = layer.bias is not None, 
                  out_w_per_mask = out_w_per_mask, in_w_per_mask = in_w_per_mask)
        res.weight = layer.weight
        res.bias = layer.bias
        return res # make sure to call cuda