import torch
from torch import nn
from torch.optim import Adam

from ml_lib.modules.blocks import ConvModule
from ml_lib.modules.utils import LightningBaseModule


class BinaryClassifier(LightningBaseModule):
    @classmethod
    def name(cls):
        return cls.__name__

    def configure_optimizers(self):
        return Adam(lr=self.hparams.train.lr)

    def training_step(self, batch_xy, batch_nb, *args, **kwargs):
        batch_x, batch_y = batch_xy
        y = self(batch_y)
        loss = self.criterion(y, batch_y)
        return dict(loss=loss)

    def validation_step(self, batch_xy, **kwargs):
        batch_x, batch_y = batch_xy
        y = self(batch_y)
        val_loss = self.criterion(y, batch_y)
        return dict(val_loss=val_loss)

    def validation_epoch_end(self, outputs):
        over_all_val_loss = torch.mean(torch.stack([output['val_loss'] for output in outputs]))

    def __init__(self, hparams):
        super(BinaryClassifier, self).__init__(hparams)

        self.criterion = nn.BCELoss()

        # Additional parameters
        self.in_shape = ()

        #

        # Model Modules
        self.conv_1 = ConvModule(self.in_shape, 32, 5, )
        self.conv_2 = ConvModule(64)
        self.conv_3 = ConvModule(128)

    def forward(self, batch, **kwargs):
        return batch