"""Pytorch-lightning module for causal language modeling.
"""

__all__ = ("GPT2LitModel",)

import pytorch_lightning as pl
import torch


class GPT2LitModel(pl.LightningModule):
    """Lightning module for autoregressive (causal) transformer language modeling.
    Successfully tested on HuggingFace `GPT2LMHeadModel`.
    """

    def __init__(self, transformer, batch_size: int, learning_rate: float,
                 final_learning_rate: float, weight_decay: float, adam_eps: float,
                 adam_betas: tuple, scheduler_T_max: int,
                 save_model_every: int = 10_000, checkpoint: str = ""):
        super().__init__()
        self.save_hyperparameters(ignore=("transformer", "save_model_every",
                                          "checkpoints"))
        self.transformer = transformer
        self.save_model_every = save_model_every
        self.checkpoint = checkpoint or "./gpt2litmodel-logs"

    def forward(self, *args, **kwargs):
        return self.transformer(*args, **kwargs)

    def training_step(self, batch, batch_idx):
        outputs = self(**batch)

        if self.save_model_every > 0 and batch_idx % self.save_model_every == 0:
            self.transformer.save_pretrained(self.checkpoint)

        return {'loss': outputs['loss']}

    def training_epoch_end(self, outputs):
        if self.save_model_every > 0:
            self.transformer.save_pretrained(self.checkpoint)

        losses = [step_output["loss"] for step_output in outputs]
        mean_loss = torch.tensor(losses).mean()
        ppl = torch.exp(mean_loss)

        self.log("ppl", ppl, on_step=False, on_epoch=True, prog_bar=True)

    def configure_optimizers(self):
        parameters = self.named_parameters()
        no_decay = ["bias", "LayerNorm.weight"]
        grouped_parameters = [
            {"params": [p for n, p in parameters
                        if not any(nd in n for nd in no_decay)],
             "weight_decay": self.hparams.weight_decay},
            {"params": [p for n, p in parameters
                        if any(nd in n for nd in no_decay)],
             "weight_decay": 0.0}]
        optimizer = torch.optim.Adam(
            grouped_parameters, lr=self.hparams.learning_rate,
            weight_decay=self.hparams.weight_decay,
            eps=self.hparams.adam_eps, betas=self.hparams.adam_betas)

        lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
            optimizer, self.hparams.scheduler_T_max,
            eta_min=self.hparams.final_learning_rate)

        return {'optimizer': optimizer,
                'lr_scheduler': {'scheduler': lr_scheduler,
                                 'interval': 'step', 'frequency': 1}}