import torch
import pandas as pd
from transformers import (
    AdamW,
    T5ForConditionalGeneration,
    T5TokenizerFast as T5Tokenizer,
    MT5Tokenizer,
    MT5ForConditionalGeneration,
    ByT5Tokenizer,
)
from torch.utils.data import Dataset, DataLoader
import pytorch_lightning as pl
from pytorch_lightning.loggers import MLFlowLogger, WandbLogger
from pytorch_lightning import Trainer
from pytorch_lightning.callbacks.early_stopping import EarlyStopping
from pytorch_lightning import LightningDataModule
from pytorch_lightning import LightningModule
from datasets import load_metric
from tqdm.auto import tqdm

# from dagshub.pytorch_lightning import DAGsHubLogger


torch.cuda.empty_cache()
pl.seed_everything(42)


class DataModule(Dataset):
    """
    Data Module for pytorch
    """

    def __init__(
        self,
        data: pd.DataFrame,
        tokenizer: T5Tokenizer,
        source_max_token_len: int = 512,
        target_max_token_len: int = 512,
    ):
        """
        :param data:
        :param tokenizer:
        :param source_max_token_len:
        :param target_max_token_len:
        """
        self.data = data
        self.target_max_token_len = target_max_token_len
        self.source_max_token_len = source_max_token_len
        self.tokenizer = tokenizer

    def __len__(self):
        return len(self.data)

    def __getitem__(self, index: int):
        data_row = self.data.iloc[index]

        input_encoding = self.tokenizer(
            data_row["input_text"],
            max_length=self.source_max_token_len,
            padding="max_length",
            truncation=True,
            return_attention_mask=True,
            add_special_tokens=True,
            return_tensors="pt",
        )

        output_encoding = self.tokenizer(
            data_row["output_text"],
            max_length=self.target_max_token_len,
            padding="max_length",
            truncation=True,
            return_attention_mask=True,
            add_special_tokens=True,
            return_tensors="pt",
        )

        labels = output_encoding["input_ids"]
        labels[labels == 0] = -100

        return dict(
            keywords=data_row["input_text"],
            text=data_row["output_text"],
            keywords_input_ids=input_encoding["input_ids"].flatten(),
            keywords_attention_mask=input_encoding["attention_mask"].flatten(),
            labels=labels.flatten(),
            labels_attention_mask=output_encoding["attention_mask"].flatten(),
        )


class PLDataModule(LightningDataModule):
    def __init__(
        self,
        train_df: pd.DataFrame,
        test_df: pd.DataFrame,
        tokenizer: T5Tokenizer,
        source_max_token_len: int = 512,
        target_max_token_len: int = 512,
        batch_size: int = 4,
        split: float = 0.1,
        num_workers: int = 2,
    ):
        """
        :param data_df:
        :param tokenizer:
        :param source_max_token_len:
        :param target_max_token_len:
        :param batch_size:
        :param split:
        """
        super().__init__()
        self.train_df = train_df
        self.test_df = test_df
        self.split = split
        self.batch_size = batch_size
        self.target_max_token_len = target_max_token_len
        self.source_max_token_len = source_max_token_len
        self.tokenizer = tokenizer
        self.num_workers = num_workers

    def setup(self, stage=None):
        self.train_dataset = DataModule(
            self.train_df,
            self.tokenizer,
            self.source_max_token_len,
            self.target_max_token_len,
        )
        self.test_dataset = DataModule(
            self.test_df,
            self.tokenizer,
            self.source_max_token_len,
            self.target_max_token_len,
        )

    def train_dataloader(self):
        """training dataloader"""
        return DataLoader(
            self.train_dataset,
            batch_size=self.batch_size,
            shuffle=True,
            num_workers=self.num_workers,
        )

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

    def val_dataloader(self):
        """validation dataloader"""
        return DataLoader(
            self.test_dataset,
            batch_size=self.batch_size,
            shuffle=False,
            num_workers=self.num_workers,
        )


class LightningModel(LightningModule):
    """PyTorch Lightning Model class"""

    def __init__(
        self,
        tokenizer,
        model,
        learning_rate,
        adam_epsilon,
        weight_decay,
        output: str = "outputs",
    ):
        """
        initiates a PyTorch Lightning Model
        Args:
            tokenizer : T5 tokenizer
            model : T5 model
            output (str, optional): output directory to save model checkpoints. Defaults to "outputs".
        """
        super().__init__()
        self.model = model
        self.tokenizer = tokenizer
        self.output = output
        self.learning_rate = learning_rate
        self.adam_epsilon = adam_epsilon
        self.weight_decay = weight_decay

    def forward(self, input_ids, attention_mask, decoder_attention_mask, labels=None):
        """forward step"""
        output = self.model(
            input_ids,
            attention_mask=attention_mask,
            labels=labels,
            decoder_attention_mask=decoder_attention_mask,
        )

        return output.loss, output.logits

    def training_step(self, batch, batch_size):
        """training step"""
        input_ids = batch["keywords_input_ids"]
        attention_mask = batch["keywords_attention_mask"]
        labels = batch["labels"]
        labels_attention_mask = batch["labels_attention_mask"]

        loss, outputs = self(
            input_ids=input_ids,
            attention_mask=attention_mask,
            decoder_attention_mask=labels_attention_mask,
            labels=labels,
        )
        self.log("train_loss", loss, prog_bar=True, logger=True)
        return loss

    def validation_step(self, batch, batch_size):
        """validation step"""
        input_ids = batch["keywords_input_ids"]
        attention_mask = batch["keywords_attention_mask"]
        labels = batch["labels"]
        labels_attention_mask = batch["labels_attention_mask"]

        loss, outputs = self(
            input_ids=input_ids,
            attention_mask=attention_mask,
            decoder_attention_mask=labels_attention_mask,
            labels=labels,
        )
        self.log("val_loss", loss, prog_bar=True, logger=True)
        return loss

    def test_step(self, batch, batch_size):
        """test step"""
        input_ids = batch["keywords_input_ids"]
        attention_mask = batch["keywords_attention_mask"]
        labels = batch["labels"]
        labels_attention_mask = batch["labels_attention_mask"]

        loss, outputs = self(
            input_ids=input_ids,
            attention_mask=attention_mask,
            decoder_attention_mask=labels_attention_mask,
            labels=labels,
        )

        self.log("test_loss", loss, prog_bar=True, logger=True)
        return loss

    def configure_optimizers(self):
        """configure optimizers"""
        model = self.model
        no_decay = ["bias", "LayerNorm.weight"]
        optimizer_grouped_parameters = [
            {
                "params": [
                    p
                    for n, p in model.named_parameters()
                    if not any(nd in n for nd in no_decay)
                ],
                "weight_decay": self.weight_decay,
            },
            {
                "params": [
                    p
                    for n, p in model.named_parameters()
                    if any(nd in n for nd in no_decay)
                ],
                "weight_decay": 0.0,
            },
        ]
        optimizer = AdamW(
            optimizer_grouped_parameters, lr=self.learning_rate, eps=self.adam_epsilon
        )
        self.opt = optimizer
        return [optimizer]


class Summarization:
    """Custom Summarization class"""

    def __init__(self) -> None:
        """initiates Summarization class"""
        pass

    def from_pretrained(self, model_type="t5", model_name="t5-base") -> None:
        """
        loads T5/MT5 Model model for training/finetuning
        Args:
            model_name (str, optional): exact model architecture name, "t5-base" or "t5-large". Defaults to "t5-base".
            :param model_type:
        """
        if model_type == "t5":
            self.tokenizer = T5Tokenizer.from_pretrained(f"{model_name}")
            self.model = T5ForConditionalGeneration.from_pretrained(
                f"{model_name}", return_dict=True
            )
        elif model_type == "mt5":
            self.tokenizer = MT5Tokenizer.from_pretrained(f"{model_name}")
            self.model = MT5ForConditionalGeneration.from_pretrained(
                f"{model_name}", return_dict=True
            )
        elif model_type == "byt5":
            self.tokenizer = ByT5Tokenizer.from_pretrained(f"{model_name}")
            self.model = T5ForConditionalGeneration.from_pretrained(
                f"{model_name}", return_dict=True
            )

    def train(
        self,
        train_df: pd.DataFrame,
        eval_df: pd.DataFrame,
        source_max_token_len: int = 512,
        target_max_token_len: int = 512,
        batch_size: int = 8,
        max_epochs: int = 5,
        use_gpu: bool = True,
        outputdir: str = "models",
        early_stopping_patience_epochs: int = 0,  # 0 to disable early stopping feature
        learning_rate: float = 0.0001,
        adam_epsilon: float = 0.01,
        num_workers: int = 2,
        weight_decay: float = 0.0001,
    ):
        """
        trains T5/MT5 model on custom dataset
        Args:
            train_df (pd.DataFrame): training datarame. Dataframe must have 2 column --> "input_text" and "output_text"
            eval_df ([type], optional): validation datarame. Dataframe must have 2 column --> "input_text" and
            "output_text"
            source_max_token_len (int, optional): max token length of source text. Defaults to 512.
            target_max_token_len (int, optional): max token length of target text. Defaults to 512.
            batch_size (int, optional): batch size. Defaults to 8.
            max_epochs (int, optional): max number of epochs. Defaults to 5.
            use_gpu (bool, optional): if True, model uses gpu for training. Defaults to True.
            outputdir (str, optional): output directory to save model checkpoints. Defaults to "outputs".
            early_stopping_patience_epochs (int, optional): monitors val_loss on epoch end and stops training,
            if val_loss does not improve after the specied number of epochs. set 0 to disable early stopping.
            Defaults to 0 (disabled)
            :param learning_rate:
            :param adam_epsilon:
        """
        self.target_max_token_len = target_max_token_len
        self.data_module = PLDataModule(
            train_df,
            eval_df,
            self.tokenizer,
            batch_size=batch_size,
            source_max_token_len=source_max_token_len,
            target_max_token_len=target_max_token_len,
            num_workers=num_workers,
        )

        self.T5Model = LightningModel(
            tokenizer=self.tokenizer,
            model=self.model,
            output=outputdir,
            learning_rate=learning_rate,
            adam_epsilon=adam_epsilon,
            weight_decay=weight_decay,
        )

        MLlogger = MLFlowLogger(
            experiment_name="Summarization",
            tracking_uri="https://dagshub.com/gagan3012/summarization.mlflow",
        )

        WandLogger = WandbLogger(project="summarization-dagshub")

        # logger = DAGsHubLogger(metrics_path='reports/training_metrics.txt')

        early_stop_callback = (
            [
                EarlyStopping(
                    monitor="val_loss",
                    min_delta=0.00,
                    patience=early_stopping_patience_epochs,
                    verbose=True,
                    mode="min",
                )
            ]
            if early_stopping_patience_epochs > 0
            else None
        )

        gpus = -1 if use_gpu and torch.cuda.is_available() else 0

        trainer = Trainer(
            logger=[WandLogger, MLlogger],
            callbacks=early_stop_callback,
            max_epochs=max_epochs,
            gpus=gpus,
            progress_bar_refresh_rate=5,
        )

        trainer.fit(self.T5Model, self.data_module)

    def load_model(
        self, model_type: str = "t5", model_dir: str = "models", use_gpu: bool = False
    ):
        """
        loads a checkpoint for inferencing/prediction
        Args:
            model_type (str, optional): "t5" or "mt5". Defaults to "t5".
            model_dir (str, optional): path to model directory. Defaults to "outputs".
            use_gpu (bool, optional): if True, model uses gpu for inferencing/prediction. Defaults to True.
        """
        if model_type == "t5":
            self.tokenizer = T5Tokenizer.from_pretrained(f"{model_dir}")
            self.model = T5ForConditionalGeneration.from_pretrained(
                f"{model_dir}", return_dict=True
            )
        elif model_type == "mt5":
            self.tokenizer = MT5Tokenizer.from_pretrained(f"{model_dir}")
            self.model = MT5ForConditionalGeneration.from_pretrained(
                f"{model_dir}", return_dict=True
            )
        elif model_type == "byt5":
            self.tokenizer = ByT5Tokenizer.from_pretrained(f"{model_dir}")
            self.model = T5ForConditionalGeneration.from_pretrained(
                f"{model_dir}", return_dict=True
            )

        if use_gpu:
            if torch.cuda.is_available():
                self.device = torch.device("cuda")
            else:
                raise Exception(
                    "exception ---> no gpu found. set use_gpu=False, to use CPU"
                )
        else:
            self.device = torch.device("cpu")

        self.model = self.model.to(self.device)

    def save_model(self, model_dir="models"):
        """
        Save model to dir
        :param model_dir:
        :return: model is saved
        """
        path = f"{model_dir}"
        self.tokenizer.save_pretrained(path)
        self.model.save_pretrained(path)

    def predict(
        self,
        source_text: str,
        max_length: int = 512,
        num_return_sequences: int = 1,
        num_beams: int = 2,
        top_k: int = 50,
        top_p: float = 0.95,
        do_sample: bool = True,
        repetition_penalty: float = 2.5,
        length_penalty: float = 1.0,
        early_stopping: bool = True,
        skip_special_tokens: bool = True,
        clean_up_tokenization_spaces: bool = True,
    ):
        """
        generates prediction for T5/MT5 model
        Args:
            source_text (str): any text for generating predictions
            max_length (int, optional): max token length of prediction. Defaults to 512.
            num_return_sequences (int, optional): number of predictions to be returned. Defaults to 1.
            num_beams (int, optional): number of beams. Defaults to 2.
            top_k (int, optional): Defaults to 50.
            top_p (float, optional): Defaults to 0.95.
            do_sample (bool, optional): Defaults to True.
            repetition_penalty (float, optional): Defaults to 2.5.
            length_penalty (float, optional): Defaults to 1.0.
            early_stopping (bool, optional): Defaults to True.
            skip_special_tokens (bool, optional): Defaults to True.
            clean_up_tokenization_spaces (bool, optional): Defaults to True.
        Returns:
            list[str]: returns predictions
        """
        input_ids = self.tokenizer.encode(
            source_text, return_tensors="pt", add_special_tokens=True
        )

        input_ids = input_ids.to(self.device)
        generated_ids = self.model.generate(
            input_ids=input_ids,
            num_beams=num_beams,
            max_length=max_length,
            repetition_penalty=repetition_penalty,
            length_penalty=length_penalty,
            early_stopping=early_stopping,
            top_p=top_p,
            top_k=top_k,
            num_return_sequences=num_return_sequences,
        )
        preds = self.tokenizer.decode(
            generated_ids[0],
            skip_special_tokens=skip_special_tokens,
            clean_up_tokenization_spaces=clean_up_tokenization_spaces,
        )
        return preds

    def evaluate(self, test_df: pd.DataFrame, metrics: str = "rouge"):
        metric = load_metric(metrics)
        input_text = test_df["input_text"]
        references = test_df["output_text"]
        references = references.to_list()

        predictions = [self.predict(x) for x in tqdm(input_text)]

        results = metric.compute(predictions=predictions, references=references)

        output = {
            "Rouge_1 Low Precision": results["rouge1"].low.precision,
            "Rouge_1 Low recall": results["rouge1"].low.recall,
            "Rouge_1 Low F1": results["rouge1"].low.fmeasure,
            "Rouge_1 Mid Precision": results["rouge1"].mid.precision,
            "Rouge_1 Mid recall": results["rouge1"].mid.recall,
            "Rouge_1 Mid F1": results["rouge1"].mid.fmeasure,
            "Rouge_1 High Precision": results["rouge1"].high.precision,
            "Rouge_1 High recall": results["rouge1"].high.recall,
            "Rouge_1 High F1": results["rouge1"].high.fmeasure,
            "Rouge_2 Low Precision": results["rouge2"].low.precision,
            "Rouge_2 Low recall": results["rouge2"].low.recall,
            "Rouge_2 Low F1": results["rouge2"].low.fmeasure,
            "Rouge_2 Mid Precision": results["rouge2"].mid.precision,
            "Rouge_2 Mid recall": results["rouge2"].mid.recall,
            "Rouge_2 Mid F1": results["rouge2"].mid.fmeasure,
            "Rouge_2 High Precision": results["rouge2"].high.precision,
            "Rouge_2 High recall": results["rouge2"].high.recall,
            "Rouge_2 High F1": results["rouge2"].high.fmeasure,
            "Rouge_L Low Precision": results["rougeL"].low.precision,
            "Rouge_L Low recall": results["rougeL"].low.recall,
            "Rouge_L Low F1": results["rougeL"].low.fmeasure,
            "Rouge_L Mid Precision": results["rougeL"].mid.precision,
            "Rouge_L Mid recall": results["rougeL"].mid.recall,
            "Rouge_L Mid F1": results["rougeL"].mid.fmeasure,
            "Rouge_L High Precision": results["rougeL"].high.precision,
            "Rouge_L High recall": results["rougeL"].high.recall,
            "Rouge_L High F1": results["rougeL"].high.fmeasure,
            "rougeLsum Low Precision": results["rougeLsum"].low.precision,
            "rougeLsum Low recall": results["rougeLsum"].low.recall,
            "rougeLsum Low F1": results["rougeLsum"].low.fmeasure,
            "rougeLsum Mid Precision": results["rougeLsum"].mid.precision,
            "rougeLsum Mid recall": results["rougeLsum"].mid.recall,
            "rougeLsum Mid F1": results["rougeLsum"].mid.fmeasure,
            "rougeLsum High Precision": results["rougeLsum"].high.precision,
            "rougeLsum High recall": results["rougeLsum"].high.recall,
            "rougeLsum High F1": results["rougeLsum"].high.fmeasure,
        }
        return output