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)

[((498, 555), 'argparse.ArgumentParser', 'argparse.ArgumentParser', ([], {'description': '"""KoBART translation"""'}), "(description='KoBART translation')\n", (521, 555), False, 'import argparse\n'), ((682, 701), 'logging.getLogger', 'logging.getLogger', ([], {}), '()\n', (699, 701), False, 'import logging\n'), ((7987, 8023), 'pytorch_lightning.Trainer.add_argparse_args', 'pl.Trainer.add_argparse_args', (['parser'], {}), '(parser)\n', (8015, 8023), True, 'import pytorch_lightning as pl\n'), ((8059, 8077), 'logging.info', 'logging.info', (['args'], {}), '(args)\n', (8071, 8077), False, 'import logging\n'), ((8425, 8653), 'pytorch_lightning.callbacks.ModelCheckpoint', '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"""'}), "(monitor='val_loss', dirpath=args.\n    default_root_dir, filename='model_chp/{epoch:02d}-{val_loss:.3f}',\n    verbose=True, save_last=True, mode='min', save_top_k=-1, prefix=\n    'kobart_translation')\n", (8453, 8653), True, 'import pytorch_lightning as pl\n'), ((9134, 9168), 'pytorch_lightning.callbacks.LearningRateMonitor', 'pl.callbacks.LearningRateMonitor', ([], {}), '()\n', (9166, 9168), True, 'import pytorch_lightning as pl\n'), ((9183, 9285), 'pytorch_lightning.Trainer.from_argparse_args', 'pl.Trainer.from_argparse_args', (['args'], {'logger': 'tb_logger', 'callbacks': '[checkpoint_callback, lr_logger]'}), '(args, logger=tb_logger, callbacks=[\n    checkpoint_callback, lr_logger])\n', (9212, 9285), True, 'import pytorch_lightning as pl\n'), ((835, 899), 'argparse.ArgumentParser', 'argparse.ArgumentParser', ([], {'parents': '[parent_parser]', 'add_help': '(False)'}), '(parents=[parent_parser], add_help=False)\n', (858, 899), False, 'import argparse\n'), ((2253, 2317), 'argparse.ArgumentParser', 'argparse.ArgumentParser', ([], {'parents': '[parent_parser]', 'add_help': '(False)'}), '(parents=[parent_parser], add_help=False)\n', (2276, 2317), False, 'import argparse\n'), ((2683, 2749), 'dataset.KoBARTSummaryDataset', 'KoBARTSummaryDataset', (['self.train_file_path', 'self.tok', 'self.max_len'], {}), '(self.train_file_path, self.tok, self.max_len)\n', (2703, 2749), False, 'from dataset import KoBARTSummaryDataset\n'), ((2836, 2901), 'dataset.KoBARTSummaryDataset', 'KoBARTSummaryDataset', (['self.test_file_path', 'self.tok', 'self.max_len'], {}), '(self.test_file_path, self.tok, self.max_len)\n', (2856, 2901), False, 'from dataset import KoBARTSummaryDataset\n'), ((3015, 3114), 'torch.utils.data.DataLoader', 'DataLoader', (['self.train'], {'batch_size': 'self.batch_size', 'num_workers': 'self.num_workers', 'shuffle': '(True)'}), '(self.train, batch_size=self.batch_size, num_workers=self.\n    num_workers, shuffle=True)\n', (3025, 3114), False, 'from torch.utils.data import DataLoader, Dataset\n'), ((3230, 3329), 'torch.utils.data.DataLoader', 'DataLoader', (['self.test'], {'batch_size': 'self.batch_size', 'num_workers': 'self.num_workers', 'shuffle': '(False)'}), '(self.test, batch_size=self.batch_size, num_workers=self.\n    num_workers, shuffle=False)\n', (3240, 3329), False, 'from torch.utils.data import DataLoader, Dataset\n'), ((3441, 3540), 'torch.utils.data.DataLoader', 'DataLoader', (['self.test'], {'batch_size': 'self.batch_size', 'num_workers': 'self.num_workers', 'shuffle': '(False)'}), '(self.test, batch_size=self.batch_size, num_workers=self.\n    num_workers, shuffle=False)\n', (3451, 3540), False, 'from torch.utils.data import DataLoader, Dataset\n'), ((3879, 3943), 'argparse.ArgumentParser', 'argparse.ArgumentParser', ([], {'parents': '[parent_parser]', 'add_help': '(False)'}), '(parents=[parent_parser], add_help=False)\n', (3902, 3943), False, 'import argparse\n'), ((5241, 5316), 'transformers.optimization.AdamW', 'AdamW', (['optimizer_grouped_parameters'], {'lr': 'self.hparams.lr', 'correct_bias': '(False)'}), '(optimizer_grouped_parameters, lr=self.hparams.lr, correct_bias=False)\n', (5246, 5316), False, 'from transformers.optimization import AdamW, get_cosine_schedule_with_warmup\n'), ((5582, 5654), 'logging.info', 'logging.info', (['f"""number of workers {num_workers}, data length {data_len}"""'], {}), "(f'number of workers {num_workers}, data length {data_len}')\n", (5594, 5654), False, 'import logging\n'), ((5771, 5823), 'logging.info', 'logging.info', (['f"""num_train_steps : {num_train_steps}"""'], {}), "(f'num_train_steps : {num_train_steps}')\n", (5783, 5823), False, 'import logging\n'), ((5908, 5962), 'logging.info', 'logging.info', (['f"""num_warmup_steps : {num_warmup_steps}"""'], {}), "(f'num_warmup_steps : {num_warmup_steps}')\n", (5920, 5962), False, 'import logging\n'), ((5983, 6101), 'transformers.optimization.get_cosine_schedule_with_warmup', 'get_cosine_schedule_with_warmup', (['optimizer'], {'num_warmup_steps': 'num_warmup_steps', 'num_training_steps': 'num_train_steps'}), '(optimizer, num_warmup_steps=\n    num_warmup_steps, num_training_steps=num_train_steps)\n', (6014, 6101), False, 'from transformers.optimization import AdamW, get_cosine_schedule_with_warmup\n'), ((6719, 6741), 'kobart.get_kobart_tokenizer', 'get_kobart_tokenizer', ([], {}), '()\n', (6739, 6741), False, 'from kobart import get_pytorch_kobart_model, get_kobart_tokenizer\n'), ((9070, 9116), 'os.path.join', 'os.path.join', (['args.default_root_dir', '"""tb_logs"""'], {}), "(args.default_root_dir, 'tb_logs')\n", (9082, 9116), False, 'import os\n'), ((2066, 2088), 'kobart.get_kobart_tokenizer', 'get_kobart_tokenizer', ([], {}), '()\n', (2086, 2088), False, 'from kobart import get_pytorch_kobart_model, get_kobart_tokenizer\n'), ((6546, 6572), 'kobart.get_pytorch_kobart_model', 'get_pytorch_kobart_model', ([], {}), '()\n', (6570, 6572), False, 'from kobart import get_pytorch_kobart_model, get_kobart_tokenizer\n'), ((7734, 7753), 'torch.stack', 'torch.stack', (['losses'], {}), '(losses)\n', (7745, 7753), False, 'import torch\n')]