from longformer.longformer_encoder_decoder import LongformerEncoderDecoderForConditionalGeneration
from longformer.longformer_encoder_decoder import LongformerEncoderDecoderConfig

from longformer.longformer import LongformerForMaskedLM
from longformer.longformer import LongformerConfig

import torch
from torch.utils.data import DataLoader, Dataset
from pytorch_lightning import Trainer
import pytorch_lightning as pl

seqlen = 1024 * 2
global_size = seqlen // 100
attention_window = 256  # one sided


class CoolDataset(Dataset):
    def __len__(self):
        return 1024  # number of examples

    def __getitem__(self, idx):
        tokne_ids = torch.tensor([5] * seqlen)
        mask = torch.tensor([1] * seqlen)
        mask[:global_size] = 2
        return tokne_ids, mask


class MemoryProfiler(pl.LightningModule):

    def __init__(self, hparams=None):
        super().__init__()
        self.hparams = hparams

        config = LongformerEncoderDecoderConfig.from_pretrained('bart-long-4096')
        # config = LongformerConfig.from_pretrained('roberta-large')
        config.max_position_embeddings = seqlen + 2
        config.gradient_checkpointing = True
        config.attention_mode = 'sliding_chunks'
        # config.attention_mode = 'n2'
        config.attention_window = [attention_window] * config.num_hidden_layers
        config.attention_dilation = [1] * config.num_hidden_layers
        self.model = LongformerEncoderDecoderForConditionalGeneration(config)
        # self.model = LongformerForMaskedLM(config)

    def forward(self, x, y):
        print(seqlen, global_size, attention_window, torch.cuda.max_memory_allocated(x.device) / 1024 ** 3)
        # import ipdb; ipdb.set_trace()
        # return self.model(x, attention_mask=y, decoder_input_ids=x[:, :attention_window * 2], use_cache=False)
        return self.model(x, attention_mask=y)

    def training_step(self, batch, batch_idx):
        # import ipdb; ipdb.set_trace()
        x, y = batch
        y_hat = self(x, y)
        loss = y_hat[0].sum()
        # import ipdb; ipdb.set_trace()
        return {'loss': loss}

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=0.001)

    def train_dataloader(self):
        return DataLoader(CoolDataset(), batch_size=2, num_workers=0)


if __name__ == '__main__':
    model = MemoryProfiler(hparams={})
    trainer = Trainer(gpus=[0], progress_bar_refresh_rate=1, max_epochs=1, amp_level='O2', use_amp=True)
    trainer.fit(model)