"""
Push the trained YOLOv8s honey-bee detector to a Hugging Face model repo.

Earns the "Sharing is Caring" badge. Re-run any time the local
weights/honey_bee_detector.pt is updated.

Usage:
    py scripts/push_yolo_to_hub.py --repo-id maryammeda/apiarist-honey-bee-detector
"""

from __future__ import annotations

import argparse
import os
from pathlib import Path

from dotenv import load_dotenv
from huggingface_hub import HfApi


WEIGHTS_PATH = Path(__file__).parent.parent / "weights" / "honey_bee_detector.pt"


README_TEMPLATE = """---
license: apache-2.0
library_name: ultralytics
tags:
- yolo
- yolov8
- object-detection
- bees
- beekeeping
- computer-vision
pipeline_tag: object-detection
---

# Apiarist Honey-Bee Detector (YOLOv8s)

A custom-trained **YOLOv8s** specialist detector for honeycomb frame
inspection. Built as part of [Apiarist](https://huggingface.co/spaces/build-small-hackathon/Apiarist), a fully-offline AI hive inspector for
backyard beekeepers, made for the
[Build Small Hackathon](https://huggingface.co/build-small-hackathon).

## Classes

Trained on the [hendricks_ricky bee-project](https://universe.roboflow.com/hendricks_ricky-hotmail-de/bee-project) dataset
(3,308 labeled images, including 892 queen-bee examples). Four classes:

- `Queen`, queen bee (larger, elongated abdomen)
- `Worker`, worker / forager bees (the majority class)
- `Drone`, male / drone bees
- `Varroa`, varroa destructor mites visible on bees or comb

## Usage

```python
from ultralytics import YOLO
from PIL import Image

model = YOLO("honey_bee_detector.pt")
results = model(Image.open("hive_frame.jpg"), conf=0.10, device="cpu")
for box in results[0].boxes:
    cls = model.names[int(box.cls.item())]
    conf = float(box.conf.item())
    print(f"{cls}: {conf:.0%}")
```

Recommended per-class confidence thresholds (tuned empirically on real
frame photos):

```python
PER_CLASS_CONF = {
    "Worker": 0.25,
    "Drone": 0.55, # higher, drones false-positive on fingers, shadows
    "Varroa": 0.30, # mites are small
    "Queen": 0.15, # lower, try harder to find her
}
```

## Training

- Base weights: `yolov8s.pt`
- Epochs: 60
- Image size: 640×640
- Batch size: 32
- Hardware: 1× NVIDIA T4 (Modal)
- Cost: ~$0.40 of free hackathon credit

## Limitations

The training set includes ~892 queen images, which is far more than the
typical bee-detection dataset but still limited compared to the
millions of bees in worker class. Queen detection precision is
moderate, when the model labels a queen, verify visually. Use as a
**candidate flag**, not as ground truth.

Varroa mite detection works on close-up frames where mites are visible
on the dorsal side of bees or in cells, but will miss mites hidden
under abdomens.

## License

Apache 2.0. Trained on data released under CC BY 4.0 by the original
dataset authors.

## Citation

If you use this model, please credit:
- This work: Apiarist (Build Small Hackathon entry, 2026)
- Training data: hendricks_ricky/bee-project, Roboflow Universe
"""


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--repo-id",
        required=True,
        help="Target HF repo, e.g. 'maryammeda/apiarist-honey-bee-detector'",
    )
    parser.add_argument(
        "--private",
        action="store_true",
        help="Create repo as private (default public)",
    )
    args = parser.parse_args()

    load_dotenv()
    token = (
        os.environ.get("HF_TOKEN")
        or os.environ.get("HUGGING_FACE_HUB_TOKEN")
    )
    if not token:
        raise SystemExit(
            "Need HF_TOKEN in .env. Get one at "
            "https://huggingface.co/settings/tokens (write scope)."
        )

    if not WEIGHTS_PATH.exists():
        raise SystemExit(f"Weights not found at {WEIGHTS_PATH}")

    print(f"Weights: {WEIGHTS_PATH} ({WEIGHTS_PATH.stat().st_size / 1024 / 1024:.1f} MB)")

    api = HfApi(token=token)
    print(f"Creating model repo {args.repo_id} ...")
    api.create_repo(
        repo_id=args.repo_id,
        repo_type="model",
        private=args.private,
        exist_ok=True,
    )

    # Write a model card README
    print("Uploading README ...")
    api.upload_file(
        path_or_fileobj=README_TEMPLATE.encode("utf-8"),
        path_in_repo="README.md",
        repo_id=args.repo_id,
        repo_type="model",
        commit_message="Add model card",
    )

    print(f"Uploading {WEIGHTS_PATH.name} ...")
    api.upload_file(
        path_or_fileobj=str(WEIGHTS_PATH),
        path_in_repo="honey_bee_detector.pt",
        repo_id=args.repo_id,
        repo_type="model",
        commit_message="Upload YOLOv8s honey-bee detector weights",
    )

    print(
        f"\n[OK] Model live at: "
        f"https://huggingface.co/{args.repo_id}"
    )


if __name__ == "__main__":
    main()