import argparse
import logging
import math
import re

import numpy as np
import torch
from accelerate import Accelerator
from accelerate.utils import set_seed
from torch.utils.data import DataLoader
from tqdm.auto import tqdm
from transformers import get_scheduler, AutoTokenizer, AdamW, SchedulerType, AutoModelForSeq2SeqLM, \
    DataCollatorWithPadding

from datasets import load_dataset

logger = logging.getLogger(__name__)


def get_parser():
    parser = argparse.ArgumentParser(description="Train ELI5 seq2seq answer generation model")
    parser.add_argument(
        "--dataset_name",
        type=str,
        default="vblagoje/lfqa",
        help="The name of the dataset to use (via the datasets library).",
    )

    parser.add_argument(
        "--per_device_train_batch_size",
        type=int,
        default=4,
    )

    parser.add_argument(
        "--per_device_eval_batch_size",
        type=int,
        default=4,
        help="Batch size (per device) for the evaluation dataloader.",
    )

    parser.add_argument(
        "--pretrained_model_name",
        type=str,
        default="facebook/bart-large",
    )

    parser.add_argument(
        "--model_save_name",
        type=str,
        default="eli5_bart_model",
    )

    parser.add_argument(
        "--learning_rate",
        type=float,
        default=2e-4,
    )

    parser.add_argument(
        "--weight_decay",
        type=float,
        default=0.0,
        help="Weight decay to use."
    )

    parser.add_argument(
        "--log_freq",
        type=int,
        default=100,
        help="Log train/validation loss every log_freq update steps"
    )

    parser.add_argument(
        "--ignore_pad_token_for_loss",
        type=bool,
        default=True,
        help="Whether to ignore the tokens corresponding to " "padded labels in the loss computation or not.",
    )

    parser.add_argument(
        "--num_train_epochs",
        type=int,
        default=3,
    )

    parser.add_argument(
        "--max_train_steps",
        type=int,
        default=None,
        help="Total number of training steps to perform. If provided, overrides num_train_epochs.",
    )

    parser.add_argument(
        "--gradient_accumulation_steps",
        type=int,
        default=16,
        help="Number of updates steps to accumulate before performing a backward/update pass.",
    )

    parser.add_argument(
        "--pad_to_max_length",
        action="store_true",
        help="If passed, pad all samples to `max_length`. Otherwise, dynamic padding is used.",
    )

    parser.add_argument(
        "--overwrite_cache", type=bool, default=None, help="Overwrite the cached training and evaluation sets"
    )

    parser.add_argument(
        "--max_source_length",
        type=int,
        default=1024,
        help="The maximum total input sequence length after "
             "tokenization.Sequences longer than this will be truncated, sequences shorter will be padded.",
    )

    parser.add_argument(
        "--max_target_length",
        type=int,
        default=360,
        help="The maximum total sequence length for target text after "
             "tokenization. Sequences longer than this will be truncated, sequences shorter will be padded."
    )

    parser.add_argument(
        "--lr_scheduler_type",
        type=SchedulerType,
        default="linear",  # this is linear with warmup
        help="The scheduler type to use.",
        choices=["linear", "cosine", "cosine_with_restarts", "polynomial", "constant", "constant_with_warmup"],
    )

    parser.add_argument(
        "--num_warmup_steps",
        type=int,
        default=None,
        help="Number of steps for the warmup in the lr scheduler."
    )

    parser.add_argument(
        "--warmup_percentage",
        type=float,
        default=0.08,
        help="Number of steps for the warmup in the lr scheduler."
    )
    return parser


def cleanup_references(text):
    # URL reference where we need to remove both the link text and URL
    # ...and this letter is used by most biographers as the cornerstone of Lee's personal
    # views on slavery ([1](_URL_2_ & pg=PA173), [2](_URL_1_), [3](_URL_5_)).
    # ...and this letter is used by most biographers as the cornerstone of Lee's personal views on slavery.
    result = re.sub(r"[\(\s]*\[\d+\]\([^)]+\)[,)]*", "", text, 0, re.MULTILINE)

    # URL reference where we need to preserve link text but remove URL
    # At the outbreak of the Civil War, [Leyburn left his church](_URL_19_) and joined the South.
    # At the outbreak of the Civil War, Leyburn left his church and joined the South.
    result = re.sub(r"\[([^]]+)\]\([^)]+\)", "\\1", result, 0, re.MULTILINE)

    # lastly remove just dangling _URL_[0-9]_ URL references
    result = re.sub(r"_URL_\d_", "", result, 0, re.MULTILINE)
    return result


def clean_answer(text):
    result = cleanup_references(text)
    result = result.replace("\n", " ")
    result = re.sub(r"\s\s+", " ", result)
    result = re.sub(r"BULLET::::-", "", result)
    return result.strip()


def clean_question(text):
    result = cleanup_references(text)
    result = result.replace("\n", " ")
    result = re.sub(r"\s\s+", " ", result)
    result = result.replace("[deleted]", "")
    return result.lower().strip()


def prepare_support_docs(example):
    provenances = example["output"][-1]["provenance"]
    context = "<P> " + " <P> ".join([p["text"] for p in provenances])
    return {"context": context}


def preprocess_eli5(examples, **fn_kwargs):
    document_cache = fn_kwargs["document_cache"]
    training = fn_kwargs.get("training", True)
    extra_answer_threshold = fn_kwargs.get("extra_answer_threshold", 3)
    include_selftext = fn_kwargs.get("include_selftext", False)
    exclude_answer_patterns = fn_kwargs.get("exclude_answer_patterns", [])

    questions, contexts, answers = [], [], []
    for q_id, question, selftext, answer in zip(examples["q_id"], examples["title"], examples["selftext"],
                                                examples["answers"]):
        accepted_answer_idx = []
        if training:
            accepted_answer_idx = [idx for idx, score in enumerate(answer["score"]) if
                                   score > extra_answer_threshold]
        if not training or not accepted_answer_idx:
            accepted_answer_idx = [0]
        document = document_cache[q_id]
        for idx in accepted_answer_idx:
            skip_answer = any([p.search(answer["text"][idx]) for p in exclude_answer_patterns])
            if skip_answer:
                continue
            if include_selftext:
                questions.append(clean_question(f"{question} {selftext}"))
            else:
                questions.append(clean_question(question))
            contexts.append(document.lower().strip())
            answers.append(clean_answer(answer["text"][idx]))

    return {"question": questions, "context": contexts, "answer": answers}


def eval_qa_s2s_epoch(model, dataloader, accelerator, args):
    model.eval()
    num_eval_steps = math.ceil(len(dataloader))
    progress_bar = tqdm(range(num_eval_steps), disable=not accelerator.is_local_main_process)
    total_loss = 0.
    with torch.no_grad():
        for step, batch in enumerate(dataloader):
            outputs = model(**batch)
            loss = outputs.loss
            total_loss += loss.item()
            progress_bar.update(1)
            progress_bar.set_postfix(loss=round((total_loss / (step + 1)), 3))
        return total_loss / (step + 1)


def train(config):
    set_seed(42)
    args = config["args"]
    eli5 = load_dataset(args.dataset_name)

    support_docs = load_dataset("vblagoje/lfqa_support_docs")

    # Initialize the accelerator. We will let the accelerator handle device placement for us in this example.
    accelerator = Accelerator()
    # Make one log on every process with the configuration for debugging.
    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        level=logging.INFO,
    )
    logger.setLevel(logging.INFO if accelerator.is_local_main_process else logging.ERROR)
    logger.info(accelerator.state)

    tokenizer = AutoTokenizer.from_pretrained(args.pretrained_model_name)
    model = AutoModelForSeq2SeqLM.from_pretrained(args.pretrained_model_name)

    # Optimizer
    # Split weights in two groups, one with weight decay and the other not.
    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": args.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=args.learning_rate, weight_decay=args.weight_decay)

    processed_datasets = {}
    support_docs_prepared = {}
    with accelerator.main_process_first():
        for split in ["train", "validation"]:
            support_docs_prepared[split] = support_docs[split].map(prepare_support_docs,
                                                                   batched=False,
                                                                   cache_file_name=f"./support_docs_{split}.arrow",
                                                                   load_from_cache_file=not args.overwrite_cache,
                                                                   desc="Preparing support docs",
                                                                   )
        column_names = eli5["train"].column_names
        for split in ["train", "validation"]:
            d_cache = dict([(e["id"], e["context"]) for e in tqdm(support_docs_prepared[split],
                                                                  desc=f"Adding support docs to LFQA {split}")])
            processed_datasets[split] = eli5[split].map(preprocess_eli5,
                                                        batched=True,
                                                        remove_columns=column_names,
                                                        cache_file_name=f"./processed_datasets_{split}.arrow",
                                                        load_from_cache_file=not args.overwrite_cache,
                                                        desc="Preparing dataset for tokenization",
                                                        fn_kwargs={"document_cache": d_cache,
                                                                   "training": split == "train",
                                                                   "exclude_answer_patterns": [re.compile("not sure what you"),
                                                                                               re.compile("\n\n >")]}
                                                        )

    padding = "max_length" if args.pad_to_max_length else False
    # Temporarily set max_target_length for training.
    max_target_length = args.max_target_length

    label_pad_token_id = -100 if args.ignore_pad_token_for_loss else tokenizer.pad_token_id

    def tokenize_dataset(examples):
        inputs = ["question: {} context: {}".format(q, c) for q, c in zip(examples["question"], examples["context"])]
        targets = examples["answer"]
        model_inputs = tokenizer(inputs, max_length=args.max_source_length, padding=padding, truncation=True)

        # Setup the tokenizer for targets
        with tokenizer.as_target_tokenizer():
            labels = tokenizer(targets, max_length=max_target_length, padding=True, truncation=True,
                               return_tensors="np")

        model_inputs["decoder_input_ids"] = labels["input_ids"][:, :-1].tolist()
        # replace pad_token_id with label_pad_token_id to avoid loss calculation on those tokens
        labels["input_ids"] = np.where(labels["input_ids"] == tokenizer.pad_token_id,
                                       label_pad_token_id, labels["input_ids"])

        model_inputs["labels"] = labels["input_ids"][:, 1:].tolist()
        return model_inputs

    tokenized_datasets = {}
    with accelerator.main_process_first():
        for split, dataset in processed_datasets.items():
            tokenized_datasets[split] = dataset.map(
                tokenize_dataset,
                batched=True,
                cache_file_name=f"./tokenized_dataset_{split}.arrow",
                remove_columns=dataset.column_names,
                load_from_cache_file=not args.overwrite_cache,
                desc="Running tokenizer on dataset"
            )

    train_dataset = tokenized_datasets["train"]
    eval_dataset = tokenized_datasets["validation"]
    train_dataset.set_format(type='torch')
    eval_dataset.set_format(type='torch')

    data_collator = DataCollatorWithPadding(tokenizer, "max_length")

    # first epoch we don't shuffle
    train_dataloader = DataLoader(train_dataset, shuffle=False, batch_size=args.per_device_train_batch_size,
                                  collate_fn=data_collator)
    eval_dataloader = DataLoader(eval_dataset, batch_size=args.per_device_eval_batch_size, collate_fn=data_collator)

    # train the model
    model, optimizer, train_dataloader, eval_dataloader = accelerator.prepare(model, optimizer, train_dataloader,
                                                                              eval_dataloader)
    # Scheduler and math around the number of training steps.
    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
    if args.max_train_steps is None:
        args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
    else:
        args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)

    num_warmup_steps = args.num_warmup_steps if args.num_warmup_steps else math.ceil(args.max_train_steps *
                                                                                     args.warmup_percentage)
    scheduler = get_scheduler(
        name=args.lr_scheduler_type,
        optimizer=optimizer,
        num_warmup_steps=num_warmup_steps,
        num_training_steps=args.max_train_steps,
    )
    # Train!
    total_batch_size = args.per_device_train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps

    logger.info("***** Running training *****")
    logger.info(f"  Num examples = {len(train_dataset)}")
    logger.info(f"  Num eval examples = {len(eval_dataset)}")
    logger.info(f"  Num Epochs = {args.num_train_epochs}")
    logger.info(f"  Instantaneous batch size per device = {args.per_device_train_batch_size}")
    logger.info(f"  Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}")
    logger.info(f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}")
    logger.info(f"  Total optimization steps = {args.max_train_steps}")
    logger.info(f"  Warmup steps = {num_warmup_steps}")
    logger.info(f"  Logging training progress every {args.log_freq} optimization steps")

    # Only show the progress bar once on each machine.
    progress_bar = tqdm(range(args.max_train_steps), disable=not accelerator.is_local_main_process)
    completed_steps = 0
    switched_train_dataloader = False
    for epoch in range(args.num_train_epochs):
        model.train()
        if epoch > 0 and not switched_train_dataloader:
            train_dataloader = DataLoader(train_dataset, batch_size=args.per_device_train_batch_size,
                                          shuffle=True, collate_fn=data_collator)
            train_dataloader = accelerator.prepare(train_dataloader)
            switched_train_dataloader = True

        for step, batch in enumerate(train_dataloader):
            outputs = model(**batch)
            loss = torch.mean(outputs.loss)
            accelerator.backward(loss)
            if ((step + 1) % args.gradient_accumulation_steps == 0) or (step + 1 == len(train_dataloader)):
                optimizer.step()
                scheduler.step()
                optimizer.zero_grad()
                progress_bar.update(1)
                progress_bar.set_postfix(loss=round(loss.item(), 3))
                completed_steps += 1

            if completed_steps >= args.max_train_steps:
                break

            if step % (args.log_freq * args.gradient_accumulation_steps) == 0:
                validation_loss = eval_qa_s2s_epoch(model, eval_dataloader, accelerator, args)
                model.train()
                logger.info(f"Train loss {loss.item()} , validation loss {validation_loss}")
                if args.wandb and accelerator.is_local_main_process:
                    import wandb
                    wandb.log({"loss": loss.item(),
                               "lr": scheduler.get_last_lr()[0],
                               "validation_loss": validation_loss,
                               "completed_steps": completed_steps})

        logger.info("Saving model {}".format(args.model_save_name))
        accelerator.wait_for_everyone()
        unwrapped_model = accelerator.unwrap_model(model)
        accelerator.save(unwrapped_model.state_dict(), "{}_{}.bin".format(args.model_save_name, epoch))

        # Calculating the validation loss over epoch
        validation_loss = eval_qa_s2s_epoch(model, eval_dataloader, accelerator, args)

        logger.info("Epoch: {}".format(epoch))
        logger.info("Validation loss: {}".format(validation_loss))


def main():
    parser = get_parser()
    parser.add_argument(
        "--wandb",
        action="store_true",
        help="If true, use W&B logging",
    )
    main_args, _ = parser.parse_known_args()
    config = {"args": main_args}
    if main_args.wandb:
        import wandb
        wandb.init(project="Bart_ELI5")
    train(config=config)


main()