import pytorch_lightning as pl
import torch
from torch.optim.lr_scheduler import ExponentialLR


class AbstractLightningModel(pl.LightningModule):
    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=self.hparams.lr)
        scheduler = ExponentialLR(optimizer,
                                  gamma=0.99,
                                  last_epoch=-1,
                                  verbose=False)
        sch = {
            "scheduler": scheduler,
            "interval": "step",
        }  # called after each training step
        return [optimizer], [sch]


class AbstractPrototypeModel(AbstractLightningModel):
    @property
    def prototypes(self):
        return self.proto_layer.components.detach().cpu()