import argparse
from collections import Counter
from pathlib import Path

from datasets import load_dataset
from transformers import AutoTokenizer


def parse_args():
    parser = argparse.ArgumentParser(description="Prepare tokenized IMDB datasets")
    parser.add_argument("--data-dir", default="./data")
    parser.add_argument("--backbone", default="distilbert-base-uncased")
    parser.add_argument("--max-length", type=int, default=128)
    parser.add_argument("--max-train-samples", type=int, default=None)
    parser.add_argument("--max-test-samples", type=int, default=None)
    return parser.parse_args()


def is_valid_text(example):
    text = example["text"]
    return text is not None and len(text.strip()) > 10


def main():
    args = parse_args()
    data_dir = Path(args.data_dir)
    train_path = data_dir / "imdb_train"
    test_path = data_dir / "imdb_test"
    tokenizer_path = data_dir / "tokenizer"
    data_dir.mkdir(parents=True, exist_ok=True)

    dataset = load_dataset("imdb")
    train_ds = dataset["train"]
    test_ds = dataset["test"]

    if args.max_train_samples is not None:
        train_ds = train_ds.select(range(args.max_train_samples))
    if args.max_test_samples is not None:
        test_ds = test_ds.select(range(args.max_test_samples))

    print("raw train label counts:", Counter(train_ds["label"]))
    print("raw test label counts:", Counter(test_ds["label"]))

    tokenizer = AutoTokenizer.from_pretrained(args.backbone)

    def tokenize_batch(examples):
        return tokenizer(
            examples["text"],
            truncation=True,
            max_length=args.max_length,
        )

    train_processed = (
        train_ds.filter(is_valid_text)
        .map(tokenize_batch, batched=True, remove_columns=["text"])
        .rename_column("label", "labels")
    )
    test_processed = (
        test_ds.filter(is_valid_text)
        .map(tokenize_batch, batched=True, remove_columns=["text"])
        .rename_column("label", "labels")
    )

    train_processed.save_to_disk(str(train_path))
    test_processed.save_to_disk(str(test_path))
    tokenizer.save_pretrained(str(tokenizer_path))

    print("saved train dataset:", train_path)
    print("saved test dataset:", test_path)
    print("saved tokenizer:", tokenizer_path)
    print("processed columns:", train_processed.column_names)


if __name__ == "__main__":
    main()