import torch
from torch import nn
from torch.nn import functional as F
import math
import numpy as np


class SparseActivations(nn.Module):
    """ Custom Linear layer but mimics a standard linear layer """

    def __init__(self, size_in, size_out, threshold):
        super().__init__()
        self.size_in, self.size_out = size_in, size_out
        self.threshold = threshold
        weights = torch.Tensor(size_out, size_in)
        self.weights = nn.Parameter(weights)  # nn.Parameter is a Tensor that's a module parameter.
        bias = torch.Tensor(size_out)
        self.bias = nn.Parameter(bias)
        self.sigmoid = nn.Sigmoid()
        # initialize weights and biases
        self.init_weights()

    def init_weights(self):
        stdv = 1. / math.sqrt(self.weights.size(1))
        self.weights.data.uniform_(-stdv, stdv)
        if self.bias is not None:
            self.bias.data.uniform_(-stdv, stdv)

    def forward(self, x):
        w_times_x = torch.mm(x, self.weights.t())
        outs = self.sigmoid(torch.add(w_times_x, self.bias))  # w times x + b
        for i in range(0, self.size_out):
            if outs[:, i] > self.threshold:
                outs[:, i] = 1
            else:
                outs[:, i] = 0
        return outs


class ForwardNetwork(nn.Module):

    # Using custom layer
    def __init__(self):
        super().__init__()
        self.l1 = nn.Linear(100, 10000)
        self.custom1 = SparseActivations(10000, 10000, threshold=0.56)
        self.l2 = nn.Linear(10000, 100)
        self.l3 = nn.Linear(100, 10)

    def forward(self, x):
        x = F.relu(self.l1(x))
        x = self.custom1(x)
        return x


def main():
    fr = ForwardNetwork()
    n = torch.tensor(np.random.uniform(0, 1, 100).reshape(1, 100).astype(np.float32))
    results = fr.forward(n).detach().numpy().flatten().astype(np.uint8)
    counter = 0
    on_bit_list = []
    for idx, r in enumerate(results):
        if r == 1:
            counter += 1
            on_bit_list.append(idx)
        else:
            pass
    print(counter)
    print(on_bit_list)


if __name__ == '__main__':
    main()