"""Evaluate a summarization checkpoint with ROUGE and corpus BLEU."""

import argparse
import json

import evaluate
from datasets import load_dataset
from transformers import pipeline

from train import DATASETS


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", default="sshleifer/distilbart-cnn-12-6")
    parser.add_argument("--dataset", choices=DATASETS, default="cnn_dailymail")
    parser.add_argument("--split", default="test")
    parser.add_argument("--samples", type=int, default=100)
    parser.add_argument("--batch-size", type=int, default=4)
    parser.add_argument("--output", default="evaluation_results.json")
    return parser.parse_args()


def main():
    args = parse_args()
    config = DATASETS[args.dataset]
    data = load_dataset(config["path"], config["name"], split=args.split)
    data = data.select(range(min(args.samples, len(data))))
    summarizer = pipeline("summarization", model=args.model, device_map="auto")
    predictions = []

    for start in range(0, len(data), args.batch_size):
        texts = data[start : start + args.batch_size][config["text"]]
        outputs = summarizer(
            texts,
            truncation=True,
            max_length=128,
            min_length=20,
            num_beams=4,
        )
        predictions.extend(item["summary_text"] for item in outputs)

    references = list(data[config["summary"]])
    rouge = evaluate.load("rouge").compute(
        predictions=predictions,
        references=references,
        use_stemmer=True,
    )
    bleu = evaluate.load("bleu").compute(
        predictions=predictions,
        references=[[reference] for reference in references],
    )
    results = {
        **{key: round(value * 100, 4) for key, value in rouge.items()},
        "bleu": round(bleu["bleu"] * 100, 4),
        "model": args.model,
        "dataset": args.dataset,
        "samples": len(data),
    }
    with open(args.output, "w", encoding="utf-8") as file:
        json.dump(results, file, indent=2)
    print(json.dumps(results, indent=2))


if __name__ == "__main__":
    main()