"""
Train YOLOv8s on a COMBINED dataset:
- matt-nudi/honey-bee-detection-model-zgjnb v4 (wide inspection shots)
- hendricks_ricky-hotmail-de/bee-project v2 (close-up macro shots with mites + queens)

This gives us generalization across both image styles. The earlier
single-dataset run on hendricks_ricky alone scored mAP 0.99 on queens
but failed on wide-frame photos, because the training distribution
was too narrow.

Class merge plan (both datasets land on these 4 canonical labels):
  Matt Nudi "bee"        + Hendricks "Worker Bee" -> Worker Bee (0)
  Matt Nudi "drone"      + Hendricks "Drone Bee"  -> Drone Bee  (1)
  Matt Nudi "queen"      + Hendricks "Queen Bee"  -> Queen Bee  (2)
  Hendricks "Varroa Mite"                          -> Varroa Mite (3)
  Matt Nudi "pollenbee"                            -> Worker Bee (treated as worker)
"""

import os
from pathlib import Path

import modal


APP_NAME = "apiarist-yolo-combined"
VOLUME_NAME = "apiarist-weights"
EPOCHS = 60
IMG_SIZE = 640
BATCH = 32
BASE_WEIGHTS = "yolov8s.pt"

image = (
    modal.Image.debian_slim(python_version="3.11")
    .pip_install(
        "ultralytics==8.3.81",
        "roboflow==1.1.50",
        "pyyaml",
    )
    .apt_install("libgl1", "libglib2.0-0")
)

vol = modal.Volume.from_name(VOLUME_NAME, create_if_missing=True)

app = modal.App(APP_NAME)


# Class index in the combined dataset
CANONICAL = {
    "Worker Bee": 0,
    "Drone Bee": 1,
    "Queen Bee": 2,
    "Varroa Mite": 3,
}

# Map of source class index -> canonical name, per dataset
HENDRICKS_MAP = {
    0: "Drone Bee",   # Drone Bee
    1: "Queen Bee",   # Queen Bee
    2: "Varroa Mite", # Varroa Mite
    3: "Worker Bee",  # Worker Bee
}
MATT_MAP = {
    0: "Worker Bee",  # bee
    1: "Drone Bee",   # drone
    2: "Worker Bee",  # pollenbee, treat as worker
    3: "Queen Bee",   # queen
}


@app.function(
    image=image,
    gpu="T4",
    volumes={"/weights": vol},
    timeout=3 * 60 * 60,
)
def train(rf_api_key: str) -> str:
    import shutil
    import sys
    from pathlib import Path

    from roboflow import Roboflow
    from ultralytics import YOLO

    print("=" * 60)
    print("Downloading both datasets ...")
    print("=" * 60)
    rf = Roboflow(api_key=rf_api_key)

    # Hendricks
    hend_proj = rf.workspace("hendricks_ricky-hotmail-de").project("bee-project")
    hend_ds = hend_proj.version(2).download("yolov8", location="/tmp/hendricks")

    # Matt Nudi
    matt_proj = rf.workspace("matt-nudi").project("honey-bee-detection-model-zgjnb")
    matt_ds = matt_proj.version(4).download("yolov8", location="/tmp/matt_nudi")

    print(f"  hendricks at {hend_ds.location}")
    print(f"  matt_nudi at {matt_ds.location}")

    # Merge into one dataset folder
    merged = Path("/tmp/combined")
    for split in ("train", "valid", "test"):
        (merged / split / "images").mkdir(parents=True, exist_ok=True)
        (merged / split / "labels").mkdir(parents=True, exist_ok=True)

    def _remap_label_file(src_label: Path, dst_label: Path, src_map: dict):
        """Rewrite a YOLO label file with remapped class indices.
        Drops boxes whose source class doesn't appear in src_map."""
        lines_out = []
        for line in src_label.read_text().splitlines():
            parts = line.strip().split()
            if not parts:
                continue
            try:
                cls_id = int(parts[0])
            except ValueError:
                continue
            canonical_name = src_map.get(cls_id)
            if canonical_name is None:
                continue
            new_id = CANONICAL[canonical_name]
            lines_out.append(" ".join([str(new_id)] + parts[1:]))
        dst_label.write_text("\n".join(lines_out) + ("\n" if lines_out else ""))

    def _absorb(src_root: Path, src_map: dict, prefix: str):
        for split in ("train", "valid", "test"):
            img_dir = src_root / split / "images"
            lbl_dir = src_root / split / "labels"
            if not img_dir.exists():
                continue
            for img in img_dir.iterdir():
                stem = img.stem
                new_img = merged / split / "images" / f"{prefix}_{img.name}"
                new_lbl = merged / split / "labels" / f"{prefix}_{stem}.txt"
                shutil.copy(img, new_img)
                src_lbl = lbl_dir / f"{stem}.txt"
                if src_lbl.exists():
                    _remap_label_file(src_lbl, new_lbl, src_map)

    print("\nAbsorbing Matt Nudi ...")
    _absorb(Path(matt_ds.location), MATT_MAP, "mnudi")
    print("Absorbing Hendricks ...")
    _absorb(Path(hend_ds.location), HENDRICKS_MAP, "hendr")

    # Write the combined data.yaml
    yaml_text = (
        "train: ../train/images\n"
        "val: ../valid/images\n"
        "test: ../test/images\n"
        "nc: 4\n"
        "names: ['Worker Bee', 'Drone Bee', 'Queen Bee', 'Varroa Mite']\n"
    )
    (merged / "data.yaml").write_text(yaml_text)
    print(f"\nCombined dataset ready at {merged}")
    for split in ("train", "valid", "test"):
        n_img = len(list((merged / split / "images").iterdir()))
        n_lbl = len(list((merged / split / "labels").iterdir()))
        print(f"  {split}: {n_img} images, {n_lbl} labels")

    print("\n" + "=" * 60)
    print(f"Training YOLOv8s for {EPOCHS} epochs ...")
    print("=" * 60)
    model = YOLO(BASE_WEIGHTS)
    model.train(
        data=str(merged / "data.yaml"),
        epochs=EPOCHS,
        imgsz=IMG_SIZE,
        batch=BATCH,
        project="/weights",
        name="apiarist_combined",
        exist_ok=True,
        device=0,
        patience=15,
    )

    best_pt = Path("/weights/apiarist_combined/weights/best.pt")
    if not best_pt.exists():
        print("ERROR: best.pt not found after training", file=sys.stderr)
        sys.exit(1)

    size_mb = best_pt.stat().st_size / 1024 / 1024
    print(f"\n[OK] best.pt saved at {best_pt} ({size_mb:.1f} MB)")
    vol.commit()
    return str(best_pt)


@app.local_entrypoint()
def main() -> None:
    from dotenv import load_dotenv

    load_dotenv()
    api_key = os.environ.get("ROBOFLOW_API_KEY")
    if not api_key:
        raise SystemExit(
            "Missing ROBOFLOW_API_KEY in .env. Add it before running."
        )
    print("Kicking off combined Modal training ...")
    weights_path = train.remote(rf_api_key=api_key)
    print("\n" + "=" * 60)
    print(f"DONE. Weights at: {weights_path}")
    print("=" * 60)
    print(
        "\nDownload locally with:\n"
        f"  modal volume get {VOLUME_NAME} /apiarist_combined/weights/best.pt "
        f"weights/honey_bee_detector.pt"
    )