"""Fine-tune BART/T5/Pegasus on CNN/DailyMail or XSum."""

import argparse
import json
import os

import evaluate
import numpy as np
from datasets import load_dataset
from transformers import (
    AutoModelForSeq2SeqLM,
    AutoTokenizer,
    DataCollatorForSeq2Seq,
    Seq2SeqTrainer,
    Seq2SeqTrainingArguments,
    set_seed,
)


DATASETS = {
    "cnn_dailymail": {
        "path": "cnn_dailymail",
        "name": "3.0.0",
        "text": "article",
        "summary": "highlights",
    },
    "xsum": {
        "path": "EdinburghNLP/xsum",
        "name": None,
        "text": "document",
        "summary": "summary",
    },
}


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--dataset", choices=DATASETS, default="cnn_dailymail")
    parser.add_argument("--model", default="facebook/bart-base")
    parser.add_argument("--output-dir", default="outputs/summarization-model")
    parser.add_argument("--epochs", type=float, default=3)
    parser.add_argument("--batch-size", type=int, default=4)
    parser.add_argument("--learning-rate", type=float, default=3e-5)
    parser.add_argument("--max-input-length", type=int, default=1024)
    parser.add_argument("--max-target-length", type=int, default=128)
    parser.add_argument("--train-samples", type=int)
    parser.add_argument("--eval-samples", type=int)
    parser.add_argument("--gradient-accumulation-steps", type=int, default=4)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--fp16", action="store_true")
    parser.add_argument("--push-to-hub", action="store_true")
    parser.add_argument("--hub-model-id")
    return parser.parse_args()


def main():
    args = parse_args()
    set_seed(args.seed)
    config = DATASETS[args.dataset]
    dataset = load_dataset(config["path"], config["name"])
    tokenizer = AutoTokenizer.from_pretrained(args.model)
    model = AutoModelForSeq2SeqLM.from_pretrained(args.model)

    train_data = dataset["train"]
    eval_split = "validation" if "validation" in dataset else "test"
    eval_data = dataset[eval_split]
    if args.train_samples:
        train_data = train_data.select(range(min(args.train_samples, len(train_data))))
    if args.eval_samples:
        eval_data = eval_data.select(range(min(args.eval_samples, len(eval_data))))

    def preprocess(batch):
        model_inputs = tokenizer(
            batch[config["text"]],
            max_length=args.max_input_length,
            truncation=True,
        )
        labels = tokenizer(
            text_target=batch[config["summary"]],
            max_length=args.max_target_length,
            truncation=True,
        )
        model_inputs["labels"] = labels["input_ids"]
        return model_inputs

    remove_columns = train_data.column_names
    train_data = train_data.map(preprocess, batched=True, remove_columns=remove_columns)
    eval_data = eval_data.map(preprocess, batched=True, remove_columns=remove_columns)
    rouge = evaluate.load("rouge")

    def compute_metrics(prediction):
        predictions, labels = prediction
        if isinstance(predictions, tuple):
            predictions = predictions[0]
        labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
        decoded_predictions = tokenizer.batch_decode(predictions, skip_special_tokens=True)
        decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)
        scores = rouge.compute(
            predictions=[text.strip() for text in decoded_predictions],
            references=[text.strip() for text in decoded_labels],
            use_stemmer=True,
        )
        return {key: round(value * 100, 4) for key, value in scores.items()}

    training_args = Seq2SeqTrainingArguments(
        output_dir=args.output_dir,
        num_train_epochs=args.epochs,
        learning_rate=args.learning_rate,
        per_device_train_batch_size=args.batch_size,
        per_device_eval_batch_size=args.batch_size,
        gradient_accumulation_steps=args.gradient_accumulation_steps,
        eval_strategy="epoch",
        save_strategy="epoch",
        logging_steps=50,
        predict_with_generate=True,
        generation_max_length=args.max_target_length,
        load_best_model_at_end=True,
        metric_for_best_model="rougeL",
        greater_is_better=True,
        fp16=args.fp16,
        report_to="none",
        push_to_hub=args.push_to_hub,
        hub_model_id=args.hub_model_id,
    )

    trainer = Seq2SeqTrainer(
        model=model,
        args=training_args,
        train_dataset=train_data,
        eval_dataset=eval_data,
        processing_class=tokenizer,
        data_collator=DataCollatorForSeq2Seq(tokenizer, model=model),
        compute_metrics=compute_metrics,
    )
    trainer.train()
    metrics = trainer.evaluate()
    trainer.save_model(args.output_dir)
    tokenizer.save_pretrained(args.output_dir)
    os.makedirs(args.output_dir, exist_ok=True)
    with open(os.path.join(args.output_dir, "evaluation_metrics.json"), "w", encoding="utf-8") as file:
        json.dump(metrics, file, indent=2)
    print(json.dumps(metrics, indent=2))


if __name__ == "__main__":
    main()