import argparse
import logging
import os
import numpy as np
import pandas as pd
import pytorch_lightning as pl
import torch
from pytorch_lightning import loggers as pl_loggers
from torch.utils.data import DataLoader, Dataset
from dataset import KoBARTSummaryDataset
from transformers import BartForConditionalGeneration, PreTrainedTokenizerFast
from transformers.optimization import AdamW, get_cosine_schedule_with_warmup
from kobart import get_pytorch_kobart_model, get_kobart_tokenizer

parser = argparse.ArgumentParser(description='KoBART translation')

parser.add_argument('--checkpoint_path',
                    type=str,
                    help='checkpoint path')

logger = logging.getLogger()
logger.setLevel(logging.INFO)


class ArgsBase():
    @staticmethod
    def add_model_specific_args(parent_parser):
        parser = argparse.ArgumentParser(
            parents=[parent_parser], add_help=False)
        parser.add_argument('--train_file',
                            type=str,
                            default='data/train.tsv',
                            help='train file')

        parser.add_argument('--test_file',
                            type=str,
                            default='data/test.tsv',
                            help='test file')

        parser.add_argument('--batch_size',
                            type=int,
                            default=28,
                            help='')

        parser.add_argument('--max_len',
                            type=int,
                            default=512,
                            help='max seq len')
        return parser

class KobartSummaryModule(pl.LightningDataModule):
    def __init__(self, train_file,
                 test_file, tok,
                 max_len=512,
                 batch_size=8,
                 num_workers=5):
        super().__init__()
        self.batch_size = batch_size
        self.max_len = max_len
        self.train_file_path = train_file
        self.test_file_path = test_file
        if tok is None:
            self.tok = get_kobart_tokenizer()
        else:
            self.tok = tok
        self.num_workers = num_workers

    @staticmethod
    def add_model_specific_args(parent_parser):
        parser = argparse.ArgumentParser(
            parents=[parent_parser], add_help=False)
        parser.add_argument('--num_workers',
                            type=int,
                            default=5,
                            help='num of worker for dataloader')
        return parser

    # OPTIONAL, called for every GPU/machine (assigning state is OK)
    def setup(self, stage):
        # split dataset
        self.train = KoBARTSummaryDataset(self.train_file_path,
                                 self.tok,
                                 self.max_len)
        self.test = KoBARTSummaryDataset(self.test_file_path,
                                self.tok,
                                self.max_len)

    def train_dataloader(self):
        train = DataLoader(self.train,
                           batch_size=self.batch_size,
                           num_workers=self.num_workers, shuffle=True)
        return train

    def val_dataloader(self):
        val = DataLoader(self.test,
                         batch_size=self.batch_size,
                         num_workers=self.num_workers, shuffle=False)
        return val

    def test_dataloader(self):
        test = DataLoader(self.test,
                          batch_size=self.batch_size,
                          num_workers=self.num_workers, shuffle=False)
        return test


class Base(pl.LightningModule):
    def __init__(self, hparams, **kwargs) -> None:
        super(Base, self).__init__()
        self.hparams = hparams

    @staticmethod
    def add_model_specific_args(parent_parser):
        # add model specific args
        parser = argparse.ArgumentParser(
            parents=[parent_parser], add_help=False)

        parser.add_argument('--batch-size',
                            type=int,
                            default=14,
                            help='batch size for training (default: 96)')

        parser.add_argument('--lr',
                            type=float,
                            default=3e-5,
                            help='The initial learning rate')

        parser.add_argument('--warmup_ratio',
                            type=float,
                            default=0.1,
                            help='warmup ratio')

        parser.add_argument('--model_path',
                            type=str,
                            default=None,
                            help='kobart model path')
        return parser

    def configure_optimizers(self):
        # Prepare optimizer
        param_optimizer = list(self.model.named_parameters())
        no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']
        optimizer_grouped_parameters = [
            {'params': [p for n, p in param_optimizer if not any(
                nd in n for nd in no_decay)], 'weight_decay': 0.01},
            {'params': [p for n, p in param_optimizer if any(
                nd in n for nd in no_decay)], 'weight_decay': 0.0}
        ]
        optimizer = AdamW(optimizer_grouped_parameters,
                          lr=self.hparams.lr, correct_bias=False)
        # warm up lr
        num_workers = (self.hparams.gpus if self.hparams.gpus is not None else 1) * (self.hparams.num_nodes if self.hparams.num_nodes is not None else 1)
        data_len = len(self.train_dataloader().dataset)
        logging.info(f'number of workers {num_workers}, data length {data_len}')
        num_train_steps = int(data_len / (self.hparams.batch_size * num_workers) * self.hparams.max_epochs)
        logging.info(f'num_train_steps : {num_train_steps}')
        num_warmup_steps = int(num_train_steps * self.hparams.warmup_ratio)
        logging.info(f'num_warmup_steps : {num_warmup_steps}')
        scheduler = get_cosine_schedule_with_warmup(
            optimizer,
            num_warmup_steps=num_warmup_steps, num_training_steps=num_train_steps)
        lr_scheduler = {'scheduler': scheduler, 
                        'monitor': 'loss', 'interval': 'step',
                        'frequency': 1}
        return [optimizer], [lr_scheduler]


class KoBARTConditionalGeneration(Base):
    def __init__(self, hparams, **kwargs):
        super(KoBARTConditionalGeneration, self).__init__(hparams, **kwargs)
        self.model = BartForConditionalGeneration.from_pretrained(get_pytorch_kobart_model())
        self.model.train()
        self.bos_token = '<s>'
        self.eos_token = '</s>'
        self.pad_token_id = 0
        self.tokenizer = get_kobart_tokenizer()

    def forward(self, inputs):
        attention_mask = inputs['input_ids'].ne(self.pad_token_id).float()
        decoder_attention_mask = inputs['decoder_input_ids'].ne(self.pad_token_id).float()
        
        return self.model(input_ids=inputs['input_ids'],
                          attention_mask=attention_mask,
                          decoder_input_ids=inputs['decoder_input_ids'],
                          decoder_attention_mask=decoder_attention_mask,
                          labels=inputs['labels'], return_dict=True)


    def training_step(self, batch, batch_idx):
        outs = self(batch)
        loss = outs.loss
        self.log('train_loss', loss, prog_bar=True)
        return loss

    def validation_step(self, batch, batch_idx):
        outs = self(batch)
        loss = outs['loss']
        return (loss)

    def validation_epoch_end(self, outputs):
        losses = []
        for loss in outputs:
            losses.append(loss)
        self.log('val_loss', torch.stack(losses).mean(), prog_bar=True)

if __name__ == '__main__':
    parser = Base.add_model_specific_args(parser)
    parser = ArgsBase.add_model_specific_args(parser)
    parser = KobartSummaryModule.add_model_specific_args(parser)
    parser = pl.Trainer.add_argparse_args(parser)
    args = parser.parse_args()
    logging.info(args)

    model = KoBARTConditionalGeneration(args)

    dm = KobartSummaryModule(args.train_file,
                        args.test_file,
                        None,
                        max_len=args.max_len,
                        batch_size=args.batch_size,
                        num_workers=args.num_workers)
    
    checkpoint_callback = pl.callbacks.ModelCheckpoint(monitor='val_loss',
                                                       dirpath=args.default_root_dir,
                                                       filename='model_chp/{epoch:02d}-{val_loss:.3f}',
                                                       verbose=True,
                                                       save_last=True,
                                                       mode='min',
                                                       save_top_k=-1,
                                                       prefix='kobart_translation')
    tb_logger = pl_loggers.TensorBoardLogger(os.path.join(args.default_root_dir, 'tb_logs'))
    lr_logger = pl.callbacks.LearningRateMonitor()
    trainer = pl.Trainer.from_argparse_args(args, logger=tb_logger,
                                            callbacks=[checkpoint_callback, lr_logger])
    trainer.fit(model, dm)

[((15, 9, 15, 66), 'argparse.ArgumentParser', 'argparse.ArgumentParser', (), '', False, 'import argparse\n'), ((21, 9, 21, 28), 'logging.getLogger', 'logging.getLogger', ({}, {}), '()', False, 'import logging\n'), ((210, 13, 210, 49), 'pytorch_lightning.Trainer.add_argparse_args', 'pl.Trainer.add_argparse_args', ({(210, 42, 210, 48): 'parser'}, {}), '(parser)', True, 'import pytorch_lightning as pl\n'), ((212, 4, 212, 22), 'logging.info', 'logging.info', ({(212, 17, 212, 21): 'args'}, {}), '(args)', False, 'import logging\n'), ((223, 26, 230, 83), 'pytorch_lightning.callbacks.ModelCheckpoint', 'pl.callbacks.ModelCheckpoint', (), '', True, 'import pytorch_lightning as pl\n'), ((232, 16, 232, 50), 'pytorch_lightning.callbacks.LearningRateMonitor', 'pl.callbacks.LearningRateMonitor', ({}, {}), '()', True, 'import pytorch_lightning as pl\n'), ((233, 14, 234, 87), 'pytorch_lightning.Trainer.from_argparse_args', 'pl.Trainer.from_argparse_args', (), '', True, 'import pytorch_lightning as pl\n'), ((28, 17, 29, 52), 'argparse.ArgumentParser', 'argparse.ArgumentParser', (), '', False, 'import argparse\n'), ((70, 17, 71, 52), 'argparse.ArgumentParser', 'argparse.ArgumentParser', (), '', False, 'import argparse\n'), ((81, 21, 83, 46), 'dataset.KoBARTSummaryDataset', 'KoBARTSummaryDataset', ({(81, 42, 81, 62): 'self.train_file_path', (82, 33, 82, 41): 'self.tok', (83, 33, 83, 45): 'self.max_len'}, {}), '(self.train_file_path, self.tok, self.max_len)', False, 'from dataset import KoBARTSummaryDataset\n'), ((84, 20, 86, 45), 'dataset.KoBARTSummaryDataset', 'KoBARTSummaryDataset', ({(84, 41, 84, 60): 'self.test_file_path', (85, 32, 85, 40): 'self.tok', (86, 32, 86, 44): 'self.max_len'}, {}), '(self.test_file_path, self.tok, self.max_len)', False, 'from dataset import KoBARTSummaryDataset\n'), ((89, 16, 91, 70), 'torch.utils.data.DataLoader', 'DataLoader', (), '', False, 'from torch.utils.data import DataLoader, Dataset\n'), ((95, 14, 97, 69), 'torch.utils.data.DataLoader', 'DataLoader', (), '', False, 'from torch.utils.data import DataLoader, Dataset\n'), ((101, 15, 103, 70), 'torch.utils.data.DataLoader', 'DataLoader', (), '', False, 'from torch.utils.data import DataLoader, Dataset\n'), ((115, 17, 116, 52), 'argparse.ArgumentParser', 'argparse.ArgumentParser', (), '', False, 'import argparse\n'), ((149, 20, 150, 65), 'transformers.optimization.AdamW', 'AdamW', (), '', False, 'from transformers.optimization import AdamW, get_cosine_schedule_with_warmup\n'), ((154, 8, 154, 80), 'logging.info', 'logging.info', ({(154, 21, 154, 79): 'f"""number of workers {num_workers}, data length {data_len}"""'}, {}), "(f'number of workers {num_workers}, data length {data_len}')", False, 'import logging\n'), ((156, 8, 156, 60), 'logging.info', 'logging.info', ({(156, 21, 156, 59): 'f"""num_train_steps : {num_train_steps}"""'}, {}), "(f'num_train_steps : {num_train_steps}')", False, 'import logging\n'), ((158, 8, 158, 62), 'logging.info', 'logging.info', ({(158, 21, 158, 61): 'f"""num_warmup_steps : {num_warmup_steps}"""'}, {}), "(f'num_warmup_steps : {num_warmup_steps}')", False, 'import logging\n'), ((159, 20, 161, 82), 'transformers.optimization.get_cosine_schedule_with_warmup', 'get_cosine_schedule_with_warmup', (), '', False, 'from transformers.optimization import AdamW, get_cosine_schedule_with_warmup\n'), ((176, 25, 176, 47), 'kobart.get_kobart_tokenizer', 'get_kobart_tokenizer', ({}, {}), '()', False, 'from kobart import get_pytorch_kobart_model, get_kobart_tokenizer\n'), ((231, 45, 231, 91), 'os.path.join', 'os.path.join', ({(231, 58, 231, 79): 'args.default_root_dir', (231, 81, 231, 90): '"""tb_logs"""'}, {}), "(args.default_root_dir, 'tb_logs')", False, 'import os\n'), ((63, 23, 63, 45), 'kobart.get_kobart_tokenizer', 'get_kobart_tokenizer', ({}, {}), '()', False, 'from kobart import get_pytorch_kobart_model, get_kobart_tokenizer\n'), ((171, 66, 171, 92), 'kobart.get_pytorch_kobart_model', 'get_pytorch_kobart_model', ({}, {}), '()', False, 'from kobart import get_pytorch_kobart_model, get_kobart_tokenizer\n'), ((204, 29, 204, 48), 'torch.stack', 'torch.stack', ({(204, 41, 204, 47): 'losses'}, {}), '(losses)', False, 'import torch\n')]