r"""
download_dataset.py
───────────────────
One-time script to download and cache the AG News dataset from HuggingFace Hub.

Run this BEFORE training to ensure the dataset is fully cached locally.
The dataset is ~30 MB and is stored in:
    Windows: C:\Users\<you>\.cache\huggingface\datasets\ag_news\

Usage
─────
    python download_dataset.py
"""
import logging
from datasets import load_dataset

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s  %(levelname)s  %(message)s",
    datefmt="%H:%M:%S",
)
logger = logging.getLogger(__name__)


def main() -> None:
    print("\n" + "─" * 60)
    print("  Downloading AG News dataset from HuggingFace Hub")
    print("  Size: ~30 MB  |  Cached after first download")
    print("─" * 60 + "\n")

    ds = load_dataset("ag_news")

    print("\n  ✓  Download complete!")
    print(f"\n     Split     Count")
    print(f"     ──────    ──────────")
    print(f"     train     {len(ds['train']):>10,}")
    print(f"     test      {len(ds['test']):>10,}")
    print(f"\n     Labels: 0=World  1=Sports  2=Business  3=Sci/Tech\n")

    # Show one sample to confirm it loaded correctly
    sample = ds["train"][42]
    print("  Sample record (index 42):")
    print(f"    text  : {sample['text'][:110]}…")
    print(f"    label : {sample['label']}")
    print()


if __name__ == "__main__":
    main()