#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
build_training_set.py — assemble a YOLO training set from external datasets,
mapped onto Alami's 7 material buckets, plus background hard-negatives.

GPU-FREE. Runs on a laptop / Colab CPU / CI. Produces the dataset that
train_yolov8_seg.py (GPU) then consumes.

Sources it understands:
  * TACO (COCO json)  — recommended download: Kaggle mirror kneroma/tacotrashdataset
                        (single batch, no Flickr rate limits; see docs/TRAINING_DATA.md)
  * any COCO-format dataset (category names are mapped via ml/serving/labels.to_bucket)
  * background images (no trash) -> emitted with EMPTY label files = hard negatives,
    the key lever against false positives (docs/TRAINING_DATA.md §4)

Class mapping is deterministic and auditable: every source category name goes
through the SAME bucket mapper the server uses (ml/serving/labels.py), so the
training labels match what the API emits at inference.

Usage:
  # convert a COCO dataset (e.g. downloaded TACO, UAVVaste) into our 7-class YOLO format
  python ml/scripts/build_training_set.py coco \
      --ann /data/taco/annotations.json --images /data/taco/images \
      --out ml/datasets/merged

  # ingest a folder-per-class CLASSIFICATION dataset (TrashNet, RealWaste,
  # e-waste/organic sets): each image gets ONE near-full-frame box of its
  # class — a weak label that works because these datasets show a single,
  # centered object. This is how we fill the dead organic/ewaste classes.
  python ml/scripts/build_training_set.py classification \
      --images /data/realwaste --out ml/datasets/merged

  # add background hard-negatives (empty labels) at ~1:3 ratio
  python ml/scripts/build_training_set.py negatives \
      --images /data/backgrounds --out ml/datasets/merged --split 0.85

  # ingest an ALREADY-YOLO-format dataset (e.g. MRS Trash Detection) by
  # remapping its class ids onto our 7 buckets (works for bbox AND seg lines)
  python ml/scripts/build_training_set.py yolo \
      --images /data/mrs/images/train --labels /data/mrs/labels/train \
      --names /data/mrs/classes.json --out ml/datasets/merged --prefix mrs
"""
from __future__ import annotations

import argparse
import json
import shutil
import sys
from collections import Counter
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

REPO = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(REPO))
from ml.serving.labels import to_bucket, BUCKETS  # noqa: E402

# fixed class order for the model — index == YOLO class id
CLASS_ORDER = list(BUCKETS)
CLASS_ID = {name: i for i, name in enumerate(CLASS_ORDER)}


def _clamp01(v: float) -> float:
    return min(max(v, 0.0), 1.0)


def _fmt_line(cid: int, coords: List[float]) -> str:
    return str(cid) + " " + " ".join(f"{c:.6f}" for c in coords)


def bbox_polygon(cx: float, cy: float, w: float, h: float) -> List[float]:
    """Normalized center-box -> 4-corner polygon (YOLO-seg segment format)."""
    x1, y1 = _clamp01(cx - w / 2.0), _clamp01(cy - h / 2.0)
    x2, y2 = _clamp01(cx + w / 2.0), _clamp01(cy + h / 2.0)
    return [x1, y1, x2, y1, x2, y2, x1, y2]


def _largest_polygon(segmentation: Any) -> Optional[List[float]]:
    """Pick the largest polygon (shoelace area) from a COCO segmentation.
    Returns None for RLE/crowd/malformed segmentations (caller falls back to bbox)."""
    if not isinstance(segmentation, list) or not segmentation:
        return None
    best, best_area = None, -1.0
    for poly in segmentation:
        if not isinstance(poly, list) or len(poly) < 6 or len(poly) % 2 != 0:
            continue
        try:
            xs, ys = poly[0::2], poly[1::2]
            area = abs(sum(xs[i] * ys[(i + 1) % len(ys)] - xs[(i + 1) % len(xs)] * ys[i]
                           for i in range(len(xs)))) / 2.0
        except TypeError:
            continue
        if area > best_area:
            best, best_area = poly, area
    return best


def coco_to_yolo_records(coco: Dict[str, Any]) -> Tuple[Dict[int, List[str]], Counter]:
    """Pure transform: COCO dict -> {image_id: [yolo-SEG label lines]} + bucket stats.

    Each COCO category name is mapped to one of Alami's 7 buckets via to_bucket,
    so the label ids match the server's output classes. Emits SEGMENT polygons
    (required to train yolov8-seg): the real COCO segmentation polygon when
    present (TACO, MTA, Roboflow exports), otherwise the bbox as a 4-corner
    polygon — a valid degenerate mask.
    """
    images = {img["id"]: img for img in coco.get("images", [])}
    cats = {c["id"]: c.get("name", "") for c in coco.get("categories", [])}
    labels: Dict[int, List[str]] = {}
    stats: Counter = Counter()

    for ann in coco.get("annotations", []):
        img = images.get(ann.get("image_id"))
        if not img:
            continue
        w0, h0 = float(img.get("width", 0)), float(img.get("height", 0))
        if w0 <= 0 or h0 <= 0:
            continue
        bbox = ann.get("bbox")
        if not bbox or len(bbox) < 4:
            continue
        x, y, bw, bh = [float(v) for v in bbox[:4]]
        if bw <= 0 or bh <= 0:
            continue
        bucket = to_bucket(cats.get(ann.get("category_id"), ""))
        cid = CLASS_ID[bucket]
        poly = _largest_polygon(ann.get("segmentation"))
        if poly is not None:
            coords = [_clamp01(v / (w0 if i % 2 == 0 else h0)) for i, v in enumerate(poly)]
        else:
            cx = _clamp01((x + bw / 2.0) / w0)
            cy = _clamp01((y + bh / 2.0) / h0)
            coords = bbox_polygon(cx, cy, _clamp01(bw / w0), _clamp01(bh / h0))
        labels.setdefault(ann["image_id"], []).append(_fmt_line(cid, coords))
        stats[bucket] += 1
    return labels, stats


def _write_dataset_yaml(out: Path) -> None:
    import yaml
    ds = {
        "path": str(out.resolve()),
        "train": "images/train", "val": "images/val", "test": "images/val",
        "nc": len(CLASS_ORDER), "names": CLASS_ORDER,
    }
    (out / "dataset.yaml").write_text(yaml.safe_dump(ds, sort_keys=False), encoding="utf-8")
    (out / "names.json").write_text(json.dumps(CLASS_ORDER, ensure_ascii=False, indent=2),
                                    encoding="utf-8")


def _ensure_dirs(out: Path) -> None:
    for sub in ["images/train", "images/val", "labels/train", "labels/val"]:
        (out / sub).mkdir(parents=True, exist_ok=True)


def cmd_coco(args) -> int:
    coco = json.loads(Path(args.ann).read_text(encoding="utf-8"))
    labels, stats = coco_to_yolo_records(coco)
    images = {img["id"]: img for img in coco.get("images", [])}
    out = Path(args.out)
    _ensure_dirs(out)
    src_images = Path(args.images)

    ids = sorted(labels.keys())
    cut = int(len(ids) * args.split)
    written = 0
    for i, img_id in enumerate(ids):
        split = "train" if i < cut else "val"
        img = images[img_id]
        fname = img.get("file_name")
        if not fname:
            continue
        src = src_images / fname
        if not src.exists():
            continue
        stem = Path(fname).name
        shutil.copy2(src, out / f"images/{split}" / stem)
        (out / f"labels/{split}" / (Path(stem).stem + ".txt")).write_text(
            "\n".join(labels[img_id]) + "\n", encoding="utf-8")
        written += 1

    _write_dataset_yaml(out)
    print(f"COCO->YOLO: {written} labelled images written to {out}")
    print(f"bucket distribution: {dict(stats)}")
    missing = [b for b in CLASS_ORDER if stats.get(b, 0) == 0]
    if missing:
        print(f"NOTE: no boxes for {missing} in this source — add a targeted dataset "
              f"(see docs/TRAINING_DATA.md §3.2).")
    return 0


IMG_EXTS = (".jpg", ".jpeg", ".png", ".webp")


def cmd_classification(args) -> int:
    """Ingest a folder-per-class classification dataset as weak detection labels.

    Layout expected:  <images>/<class_name>/*.jpg  (e.g. RealWaste, TrashNet,
    Kaggle e-waste). Each image gets ONE near-full-frame box (margin trims the
    border) labelled with to_bucket(<class_name>). Class folders that map to
    'other' can be skipped with --skip-other to avoid diluting the signal.
    """
    src = Path(args.images)
    out = Path(args.out)
    _ensure_dirs(out)
    margin = float(args.margin)
    assert 0.0 <= margin < 0.5, "--margin must be in [0, 0.5)"

    stats: Counter = Counter()
    skipped_other = 0
    written = 0
    class_dirs = sorted(d for d in src.iterdir() if d.is_dir())
    if not class_dirs:
        print(f"No class subfolders found under {src} — expected <class>/<images>.")
        return 1

    for cdir in class_dirs:
        bucket = to_bucket(cdir.name)
        if bucket == "other" and args.skip_other:
            skipped_other += 1
            continue
        cid = CLASS_ID[bucket]
        imgs = sorted(p for p in cdir.rglob("*") if p.suffix.lower() in IMG_EXTS)
        cut = int(len(imgs) * args.split)
        for i, p in enumerate(imgs):
            split = "train" if i < cut else "val"
            stem = f"cls_{cdir.name}_{p.stem}{p.suffix.lower()}"
            shutil.copy2(p, out / f"images/{split}" / stem)
            w = h = 1.0 - 2.0 * margin
            # weak near-full-frame box as 4-corner polygon (yolov8-seg format)
            (out / f"labels/{split}" / (Path(stem).stem + ".txt")).write_text(
                _fmt_line(cid, bbox_polygon(0.5, 0.5, w, h)) + "\n", encoding="utf-8")
            stats[bucket] += 1
            written += 1

    if not (out / "names.json").exists():
        _write_dataset_yaml(out)
    print(f"classification->YOLO: {written} images written to {out} "
          f"(weak full-frame boxes, margin={margin})")
    print(f"bucket distribution: {dict(stats)}")
    if skipped_other:
        print(f"skipped {skipped_other} class folder(s) mapping to 'other' (--skip-other)")
    return 0


def load_class_names(path: Path) -> List[str]:
    """Load source class names from classes.json (list), dataset.yaml
    (names: list|{id: name}) or a plain txt (one name per line)."""
    text = path.read_text(encoding="utf-8")
    if path.suffix.lower() in (".yaml", ".yml"):
        import yaml
        names = (yaml.safe_load(text) or {}).get("names")
        if isinstance(names, dict):
            return [str(names[k]) for k in sorted(names, key=int)]
        return [str(n) for n in (names or [])]
    if path.suffix.lower() == ".json":
        data = json.loads(text)
        if isinstance(data, dict):  # {"0": "Plastic", ...} or {"names": [...]}
            if "names" in data:
                data = data["names"]
            else:
                return [str(data[k]) for k in sorted(data, key=lambda x: int(x))]
        return [str(n) for n in data]
    return [ln.strip() for ln in text.splitlines() if ln.strip()]


def remap_yolo_label_lines(lines: List[str], names: List[str],
                           skip_other: bool = False) -> Tuple[List[str], Counter]:
    """Pure transform: YOLO label lines with SOURCE class ids -> yolo-SEG lines
    with Alami bucket ids. Polygon lines pass through with remapped id; bbox
    lines (cid cx cy w h) are converted to 4-corner polygons, because training
    yolov8-seg requires segment labels for every object."""
    out_lines: List[str] = []
    stats: Counter = Counter()
    for ln in lines:
        parts = ln.split()
        if len(parts) < 5:  # need at least cid + 4 coords
            continue
        try:
            src_id = int(float(parts[0]))
            coords = [float(v) for v in parts[1:]]
        except ValueError:
            continue
        if not (0 <= src_id < len(names)):
            continue
        bucket = to_bucket(names[src_id])
        if bucket == "other" and skip_other:
            continue
        cid = CLASS_ID[bucket]
        if len(coords) == 4:  # bbox -> degenerate rectangle polygon
            cx, cy, w, h = coords
            out_lines.append(_fmt_line(cid, bbox_polygon(cx, cy, w, h)))
        else:                 # already a polygon
            out_lines.append(_fmt_line(cid, coords))
        stats[bucket] += 1
    return out_lines, stats


def cmd_yolo(args) -> int:
    """Ingest a YOLO-format dataset (images + labels + class names), remapping
    source class ids onto our 7 buckets. Supports bbox and segmentation label
    lines. Images whose labels all drop out (invalid/skipped) are not copied."""
    out = Path(args.out)
    _ensure_dirs(out)
    names = load_class_names(Path(args.names))
    if not names:
        print(f"No class names loaded from {args.names}")
        return 1
    img_dir, lbl_dir = Path(args.images), Path(args.labels)
    imgs = sorted(p for p in img_dir.rglob("*") if p.suffix.lower() in IMG_EXTS)
    prefix = args.prefix or img_dir.parent.name
    stats: Counter = Counter()
    written = skipped = 0
    cut = int(len(imgs) * args.split)
    for i, p in enumerate(imgs):
        lbl = lbl_dir / (p.stem + ".txt")
        if not lbl.exists():
            skipped += 1
            continue
        lines, s = remap_yolo_label_lines(
            lbl.read_text(encoding="utf-8").splitlines(), names, args.skip_other)
        if not lines:
            skipped += 1
            continue
        split = "train" if i < cut else "val"
        stem = f"yolo_{prefix}_{p.stem}{p.suffix.lower()}"
        shutil.copy2(p, out / f"images/{split}" / stem)
        (out / f"labels/{split}" / (Path(stem).stem + ".txt")).write_text(
            "\n".join(lines) + "\n", encoding="utf-8")
        stats.update(s)
        written += 1
    if not (out / "names.json").exists():
        _write_dataset_yaml(out)
    print(f"YOLO->YOLO: {written} images remapped to {out} ({skipped} skipped: "
          "no/empty/unmappable labels)")
    print(f"bucket distribution: {dict(stats)}")
    return 0


def cmd_negatives(args) -> int:
    """Add background images as hard negatives (empty label files)."""
    out = Path(args.out)
    _ensure_dirs(out)
    src = Path(args.images)
    imgs = [p for p in src.rglob("*") if p.suffix.lower() in (".jpg", ".jpeg", ".png", ".webp")]
    cut = int(len(imgs) * args.split)
    for i, p in enumerate(imgs):
        split = "train" if i < cut else "val"
        stem = f"neg__{p.stem}{p.suffix}"
        shutil.copy2(p, out / f"images/{split}" / stem)
        # EMPTY label file == "no object here" == hard negative
        (out / f"labels/{split}" / (Path(stem).stem + ".txt")).write_text("", encoding="utf-8")
    if not (out / "names.json").exists():
        _write_dataset_yaml(out)
    print(f"hard-negatives: {len(imgs)} background images added (empty labels) to {out}")
    return 0


def main(argv=None) -> int:
    ap = argparse.ArgumentParser(description="Build a 7-class YOLO training set from external data.")
    sub = ap.add_subparsers(dest="cmd", required=True)

    p_coco = sub.add_parser("coco", help="convert a COCO dataset (e.g. TACO)")
    p_coco.add_argument("--ann", required=True, help="COCO annotations.json")
    p_coco.add_argument("--images", required=True, help="image directory")
    p_coco.add_argument("--out", default="ml/datasets/merged")
    p_coco.add_argument("--split", type=float, default=0.85)
    p_coco.set_defaults(func=cmd_coco)

    p_cls = sub.add_parser("classification",
                           help="ingest folder-per-class dataset (weak full-frame boxes)")
    p_cls.add_argument("--images", required=True, help="root dir with <class>/<images> layout")
    p_cls.add_argument("--out", default="ml/datasets/merged")
    p_cls.add_argument("--split", type=float, default=0.85)
    p_cls.add_argument("--margin", type=float, default=0.02,
                       help="border fraction trimmed off the full-frame box (default 0.02)")
    p_cls.add_argument("--skip-other", action="store_true",
                       help="skip class folders that map to the 'other' bucket")
    p_cls.set_defaults(func=cmd_classification)

    p_yolo = sub.add_parser("yolo", help="ingest a YOLO-format dataset with class-id remap")
    p_yolo.add_argument("--images", required=True, help="image directory (searched recursively)")
    p_yolo.add_argument("--labels", required=True, help="label directory (matching <stem>.txt)")
    p_yolo.add_argument("--names", required=True,
                        help="source class names: classes.json | dataset.yaml | names.txt")
    p_yolo.add_argument("--out", default="ml/datasets/merged")
    p_yolo.add_argument("--split", type=float, default=0.85)
    p_yolo.add_argument("--prefix", default="", help="filename prefix to avoid collisions")
    p_yolo.add_argument("--skip-other", action="store_true",
                        help="drop boxes that map to the 'other' bucket")
    p_yolo.set_defaults(func=cmd_yolo)

    p_neg = sub.add_parser("negatives", help="add background hard-negatives")
    p_neg.add_argument("--images", required=True, help="directory of background (no-trash) images")
    p_neg.add_argument("--out", default="ml/datasets/merged")
    p_neg.add_argument("--split", type=float, default=0.85)
    p_neg.set_defaults(func=cmd_negatives)

    args = ap.parse_args(argv)
    return args.func(args)


if __name__ == "__main__":
    sys.exit(main())