# backend/ml/train_pose_model.py

import argparse
import json
import shutil
from datetime import datetime, timezone
from pathlib import Path

from ultralytics import YOLO
from ultralytics.models.yolo.pose.train import PoseTrainer

BASE_MODEL_NAME = "yolo26n-pose.pt"


def emit_progress(stage, progress, message="", **extra):
    out = {
        "type": "progress",
        "stage": stage,
        "progress": max(0.0, min(1.0, float(progress))),
        "message": message,
    }
    out.update(extra)
    print(json.dumps(out, ensure_ascii=False), flush=True)


def count_yolo_samples(dataset_root: Path, split: str) -> int:
    images_dir = dataset_root / "images" / split
    labels_dir = dataset_root / "labels" / split

    if not images_dir.exists() or not labels_dir.exists():
        return 0

    image_exts = {".jpg", ".jpeg", ".png", ".webp"}
    count = 0

    for image_path in images_dir.iterdir():
        if not image_path.is_file():
            continue

        if image_path.suffix.lower() not in image_exts:
            continue

        label_path = labels_dir / f"{image_path.stem}.txt"
        if label_path.exists() and label_path.is_file():
            count += 1

    return count


def safe_int(value, fallback):
    try:
        return int(value)
    except Exception:
        return fallback


def existing_file(path):
    try:
        p = Path(path).expanduser().resolve()
        if p.exists() and p.is_file() and p.stat().st_size > 0:
            return p
    except Exception:
        pass
    return None


def base_model_candidates(root):
    script_dir = Path(__file__).resolve().parent

    return [
        script_dir / BASE_MODEL_NAME,
        Path.cwd() / BASE_MODEL_NAME,
        root / BASE_MODEL_NAME,
        root.parent / BASE_MODEL_NAME,
        root.parent.parent / BASE_MODEL_NAME,
    ]


def resolve_base_model(root, requested):
    requested = str(requested or "").strip()

    if requested and requested != BASE_MODEL_NAME:
        p = existing_file(requested)
        if p:
            return str(p)
        return requested

    for candidate in base_model_candidates(root):
        p = existing_file(candidate)
        if p:
            return str(p)

    return requested or BASE_MODEL_NAME


def batch_sample_ids(batch):
    if not isinstance(batch, dict):
        return []

    image_paths = batch.get("im_file")
    if not image_paths:
        return []

    if isinstance(image_paths, (str, Path)):
        image_paths = [image_paths]

    out = []
    seen = set()

    for image_path in image_paths:
        stem = Path(str(image_path)).stem.strip()
        if stem and stem not in seen:
            seen.add(stem)
            out.append(stem)

    return out


def progress_pose_trainer(train_count, val_count, train_device, fallback_epochs):
    class ProgressPoseTrainer(PoseTrainer):
        def __init__(self, *args, **kwargs):
            super().__init__(*args, **kwargs)
            self._preview_epoch = 0
            self._preview_batch = 0

        def preprocess_batch(self, batch):
            epoch = int(getattr(self, "epoch", 0)) + 1
            total_epochs = int(getattr(self, "epochs", fallback_epochs) or fallback_epochs)

            if epoch != self._preview_epoch:
                self._preview_epoch = epoch
                self._preview_batch = 0

            self._preview_batch += 1

            sample_ids = batch_sample_ids(batch)
            if sample_ids:
                total_batches = max(1, len(self.train_loader))
                completed = (epoch - 1) + min(1.0, self._preview_batch / total_batches)
                progress = 0.04 + 0.90 * (completed / max(1, total_epochs))

                emit_progress(
                    "pose",
                    progress,
                    "Pose Detector trainiert…",
                    epoch=epoch,
                    epochs=total_epochs,
                    sampleId=sample_ids[0],
                    trainSamples=train_count,
                    valSamples=val_count,
                    device=str(train_device),
                )

            return super().preprocess_batch(batch)

    return ProgressPoseTrainer


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--root", required=True)
    parser.add_argument("--base", default="yolo26n-pose.pt")
    parser.add_argument("--epochs", default="80")
    parser.add_argument("--imgsz", default="640")
    parser.add_argument("--device", default="auto")
    parser.add_argument("--workers", default="2")
    parser.add_argument("--batch", default="0")
    parser.add_argument("--threads", default="0")
    parser.add_argument("--patience", default="20")
    args = parser.parse_args()

    import torch

    root = Path(args.root).resolve()
    dataset_root = root / "pose" / "dataset"
    yaml_path = dataset_root / "dataset.yaml"
    runs_dir = root / "pose" / "runs"
    out_dir = root / "pose" / "model"

    epochs = max(1, safe_int(args.epochs, 80))
    imgsz = max(64, safe_int(args.imgsz, 640))
    workers = max(0, safe_int(args.workers, 2))
    batch_size = max(0, safe_int(args.batch, 0))
    threads = max(0, safe_int(args.threads, 0))
    patience = max(0, safe_int(args.patience, 20))

    if threads > 0:
        torch.set_num_threads(threads)
        try:
            torch.set_num_interop_threads(max(1, min(threads, 4)))
        except Exception:
            pass

    if str(args.device).lower() == "auto":
        train_device = 0 if torch.cuda.is_available() else "cpu"
    else:
        train_device = args.device

    if not yaml_path.exists():
        raise SystemExit(f"dataset.yaml not found: {yaml_path}")

    train_count = count_yolo_samples(dataset_root, "train")
    val_count = count_yolo_samples(dataset_root, "val")

    emit_progress(
        "pose",
        0.01,
        "YOLO-Pose-Dataset wird geprüft…",
        trainSamples=train_count,
        valSamples=val_count,
        epochs=epochs,
        imgsz=imgsz,
        device=str(train_device),
    )

    if train_count <= 0:
        raise SystemExit("no YOLO pose train samples found")

    if val_count <= 0:
        raise SystemExit("no YOLO pose val samples found")

    emit_progress(
        "pose",
        0.03,
        "YOLO-Pose-Basismodell wird geladen…",
        base=resolve_base_model(root, args.base),
        device=str(train_device),
    )

    base_model = resolve_base_model(root, args.base)
    model = YOLO(base_model)

    best_epoch = 0
    last_epoch = 0
    best_map50 = 0.0
    best_map5095 = 0.0

    def on_train_epoch_start(trainer):
        epoch = int(getattr(trainer, "epoch", 0)) + 1
        total = int(getattr(trainer, "epochs", epochs) or epochs)

        emit_progress(
            "pose",
            0.04 + 0.90 * ((epoch - 1) / max(1, total)),
            "Pose Detector trainiert…",
            epoch=epoch,
            epochs=total,
            trainSamples=train_count,
            valSamples=val_count,
            device=str(train_device),
        )

    def on_train_epoch_end(trainer):
        nonlocal last_epoch

        epoch = int(getattr(trainer, "epoch", 0)) + 1
        total = int(getattr(trainer, "epochs", epochs) or epochs)
        last_epoch = max(last_epoch, epoch)

        emit_progress(
            "pose",
            0.04 + 0.90 * (epoch / max(1, total)),
            "Pose Detector trainiert…",
            epoch=epoch,
            epochs=total,
            trainSamples=train_count,
            valSamples=val_count,
            device=str(train_device),
        )

    def on_fit_epoch_end(trainer):
        nonlocal best_epoch, best_map50, best_map5095

        epoch = int(getattr(trainer, "epoch", 0)) + 1
        total = int(getattr(trainer, "epochs", epochs) or epochs)
        metrics = getattr(trainer, "metrics", None) or {}

        map50 = (
            metrics.get("metrics/mAP50(B)")
            or metrics.get("metrics/mAP50")
            or metrics.get("mAP50")
        )
        map5095 = (
            metrics.get("metrics/mAP50-95(B)")
            or metrics.get("metrics/mAP50-95")
            or metrics.get("mAP50-95")
        )

        if map50 is not None or map5095 is not None:
            best_epoch = epoch

            m50 = float(map50) if map50 is not None else 0.0
            if m50 >= best_map50:
                best_map50 = m50
                if map5095 is not None:
                    best_map5095 = float(map5095)

        emit_progress(
            "pose",
            0.04 + 0.90 * (epoch / max(1, total)),
            "Pose Detector validiert…",
            epoch=epoch,
            epochs=total,
            mAP50=map50,
            mAP5095=map5095,
            device=str(train_device),
        )

    model.add_callback("on_train_epoch_start", on_train_epoch_start)
    model.add_callback("on_train_epoch_end", on_train_epoch_end)
    model.add_callback("on_fit_epoch_end", on_fit_epoch_end)

    emit_progress(
        "pose",
        0.05,
        "Pose Detector Training startet…",
        trainSamples=train_count,
        valSamples=val_count,
        epochs=epochs,
        imgsz=imgsz,
        device=str(train_device),
    )

    train_kwargs = {
        "trainer": progress_pose_trainer(
            train_count,
            val_count,
            train_device,
            epochs,
        ),
        "data": str(yaml_path),
        "epochs": epochs,
        "imgsz": imgsz,
        "project": str(runs_dir),
        "name": "pose",
        "exist_ok": True,
        "device": train_device,
        "workers": workers,
        "patience": patience,
    }
    if batch_size > 0:
        train_kwargs["batch"] = batch_size

    result = model.train(**train_kwargs)

    emit_progress(
        "pose",
        0.96,
        "Bestes YOLO-Pose-Modell wird übernommen…",
        lastEpoch=last_epoch,
        bestEpoch=best_epoch,
        device=str(train_device),
    )

    best = runs_dir / "pose" / "weights" / "best.pt"
    last = runs_dir / "pose" / "weights" / "last.pt"

    if not best.exists():
        if last.exists():
            best = last
        else:
            raise SystemExit(f"best.pt not found after training: {best}")

    out_dir.mkdir(parents=True, exist_ok=True)

    final_model = out_dir / "best.pt"
    shutil.copy2(best, final_model)

    result_path = runs_dir / "pose"
    status = {
        "ok": True,
        "model": str(final_model),
        "sourceModel": str(best),
        "runs": str(result_path),
        "trainedAt": datetime.now(timezone.utc).isoformat(),
        "trainSamples": train_count,
        "valSamples": val_count,
        "epochs": epochs,
        "imgsz": imgsz,
        "batchSize": batch_size,
        "threads": threads,
        "workers": workers,
        "device": str(train_device),
        "mAP50": round(best_map50, 4),
        "mAP5095": round(best_map5095, 4),
        "bestEpoch": best_epoch,
    }

    with (out_dir / "status.json").open("w", encoding="utf-8") as f:
        json.dump(status, f, ensure_ascii=False, indent=2)

    emit_progress(
        "pose",
        1.0,
        "Pose Detector fertig.",
        **status,
    )

    print(json.dumps(status, ensure_ascii=False), flush=True)


if __name__ == "__main__":
    main()