# Source: Lightning Bolts by William Falcon and Kyunghyun Cho
# https://github.com/PyTorchLightning/lightning-bolts/blob/master/pl_bolts/models/self_supervised/evaluator.py

from torch import nn


class SSLEvaluator(nn.Module):

    def __init__(self, n_input, n_classes, n_hidden=512, p=0.1):
        super().__init__()
        self.n_input = n_input
        self.n_classes = n_classes
        self.n_hidden = n_hidden
        if n_hidden is None:
            # use linear classifier
            self.block_forward = nn.Sequential(
                # Flatten(),
                nn.Dropout(p=p),
                nn.Linear(n_input, n_classes, bias=True),
            )
        else:
            # use simple MLP classifier
            self.block_forward = nn.Sequential(
                # Flatten(),
                nn.Dropout(p=p),
                nn.Linear(n_input, n_hidden, bias=False),
                nn.BatchNorm1d(n_hidden),
                nn.ReLU(inplace=True),
                nn.Dropout(p=p),
                nn.Linear(n_hidden, n_classes, bias=True),
            )

    def forward(self, x):
        not_logits = self.block_forward(x)
        return not_logits


class Flatten(nn.Module):

    def __init__(self):
        super(Flatten, self).__init__()

    def forward(self, input_tensor):
        return input_tensor.view(input_tensor.size(0), -1)