from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
import time
import urllib.request

import cv2
import numpy as np


@dataclass
class PoseResult17:
    """Common COCO-17 pose result returned by all backends."""
    xy: np.ndarray          # (17, 2), pixel coordinates
    conf: np.ndarray        # (17,), confidence/visibility
    bbox: np.ndarray | None # [x1, y1, x2, y2]


class PoseBackendBase:
    name = "unknown"

    def __init__(self):
        self.last_inference_ms = 0.0

    def infer(self, frame_bgr: np.ndarray) -> PoseResult17 | None:
        raise NotImplementedError

    def close(self) -> None:
        pass


class UltralyticsPoseBackend(PoseBackendBase):
    name = "Ultralytics"

    def __init__(
        self,
        model_name: str = "yolo26n-pose.pt",
        imgsz: int = 640,
        person_conf: float = 0.30,
    ):
        super().__init__()

        import torch
        from ultralytics import YOLO

        self.model = YOLO(model_name)
        self.imgsz = imgsz
        self.person_conf = person_conf

        self.use_cuda = torch.cuda.is_available()
        self.device = 0 if self.use_cuda else "cpu"

        print("[Pose] backend = Ultralytics")
        print(f"[Pose] model   = {model_name}")
        print(f"[Pose] device  = {self.device}")

        if self.use_cuda:
            print(f"[Pose] GPU     = {torch.cuda.get_device_name(0)}")

    def infer(self, frame_bgr: np.ndarray) -> PoseResult17 | None:
        kwargs = dict(
            source=frame_bgr,
            imgsz=self.imgsz,
            conf=self.person_conf,
            device=self.device,
            verbose=False,
        )

        # Current Ultralytics:
        # quantize=16 replaces deprecated half=True.
        if self.use_cuda:
            kwargs["quantize"] = 16

        t0 = time.perf_counter()

        result = self.model.predict(**kwargs)[0]

        self.last_inference_ms = (
            time.perf_counter() - t0
        ) * 1000.0

        if (
            result.boxes is None
            or result.keypoints is None
            or len(result.boxes) == 0
        ):
            return None

        # Single-subject assumption:
        # choose the largest detected person.
        boxes = (
            result.boxes.xyxy
            .detach()
            .cpu()
            .numpy()
        )

        areas = (
            (boxes[:, 2] - boxes[:, 0])
            * (boxes[:, 3] - boxes[:, 1])
        )

        index = int(np.argmax(areas))

        xy = (
            result.keypoints.xy[index]
            .detach()
            .cpu()
            .numpy()
            .astype(np.float32)
        )

        conf_tensor = result.keypoints.conf

        if conf_tensor is None:
            conf = np.ones(
                17,
                dtype=np.float32
            )
        else:
            conf = (
                conf_tensor[index]
                .detach()
                .cpu()
                .numpy()
                .astype(np.float32)
            )

        bbox = boxes[index].astype(
            np.float32
        )

        return PoseResult17(
            xy=xy,
            conf=conf,
            bbox=bbox
        )


# ============================================================
# MediaPipe Pose Landmarker 33 -> COCO 17
# ============================================================

MEDIAPIPE_TO_COCO17 = [
    0,   # Nose
    2,   # Left eye
    5,   # Right eye
    7,   # Left ear
    8,   # Right ear
    11,  # Left shoulder
    12,  # Right shoulder
    13,  # Left elbow
    14,  # Right elbow
    15,  # Left wrist
    16,  # Right wrist
    23,  # Left hip
    24,  # Right hip
    25,  # Left knee
    26,  # Right knee
    27,  # Left ankle
    28,  # Right ankle
]


MEDIAPIPE_MODELS = {

    "lite": (
        "pose_landmarker_lite.task",

        "https://storage.googleapis.com/"
        "mediapipe-models/pose_landmarker/"
        "pose_landmarker_lite/float16/1/"
        "pose_landmarker_lite.task",
    ),

    "full": (
        "pose_landmarker_full.task",

        "https://storage.googleapis.com/"
        "mediapipe-models/pose_landmarker/"
        "pose_landmarker_full/float16/1/"
        "pose_landmarker_full.task",
    ),

    "heavy": (
        "pose_landmarker_heavy.task",

        "https://storage.googleapis.com/"
        "mediapipe-models/pose_landmarker/"
        "pose_landmarker_heavy/float16/1/"
        "pose_landmarker_heavy.task",
    ),
}


class MediaPipePoseBackend(PoseBackendBase):

    name = "MediaPipe"

    def __init__(
        self,
        model_variant: str = "lite",
        detection_conf: float = 0.40,
        presence_conf: float = 0.40,
        tracking_conf: float = 0.40,
        model_dir: str = "models",
    ):
        super().__init__()

        import mediapipe as mp

        if model_variant not in MEDIAPIPE_MODELS:
            raise ValueError(
                "model_variant must be "
                "lite, full, or heavy"
            )

        self.mp = mp
        self.model_variant = model_variant
        self.last_timestamp_ms = -1

        filename, url = (
            MEDIAPIPE_MODELS[
                model_variant
            ]
        )

        model_dir_path = Path(
            model_dir
        )

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

        model_path = (
            model_dir_path / filename
        )

        # ----------------------------------------------------
        # Automatic model download
        # ----------------------------------------------------

        if not model_path.exists():

            print(
                "[Pose] downloading "
                "MediaPipe model:"
            )

            print(
                f"       {url}"
            )

            urllib.request.urlretrieve(
                url,
                model_path
            )

            print(
                f"[Pose] saved to "
                f"{model_path}"
            )

        # ----------------------------------------------------
        # VIDEO mode:
        # tracking information is reused between frames.
        # ----------------------------------------------------

        options = (
            mp.tasks.vision
            .PoseLandmarkerOptions(

                base_options=(
                    mp.tasks.BaseOptions(
                        model_asset_path=
                        str(model_path)
                    )
                ),

                running_mode=(
                    mp.tasks.vision
                    .RunningMode.VIDEO
                ),

                num_poses=1,

                min_pose_detection_confidence=
                detection_conf,

                min_pose_presence_confidence=
                presence_conf,

                min_tracking_confidence=
                tracking_conf,

                output_segmentation_masks=False,
            )
        )

        self.landmarker = (
            mp.tasks.vision
            .PoseLandmarker
            .create_from_options(
                options
            )
        )

        print(
            "[Pose] backend = MediaPipe"
        )

        print(
            f"[Pose] model   = "
            f"{model_variant}"
        )

        print(
            "[Pose] mode    = VIDEO"
        )

    @staticmethod
    def _landmark_confidence(lm):

        values = []

        visibility = getattr(
            lm,
            "visibility",
            None
        )

        presence = getattr(
            lm,
            "presence",
            None
        )

        if visibility is not None:
            values.append(
                float(visibility)
            )

        if presence is not None:
            values.append(
                float(presence)
            )

        if values:
            return min(values)

        return 1.0

    def infer(
        self,
        frame_bgr: np.ndarray
    ) -> PoseResult17 | None:

        h, w = frame_bgr.shape[:2]

        # BGR -> RGB
        frame_rgb = cv2.cvtColor(
            frame_bgr,
            cv2.COLOR_BGR2RGB
        )

        frame_rgb = np.ascontiguousarray(
            frame_rgb
        )

        mp_image = self.mp.Image(
            image_format=
                self.mp.ImageFormat.SRGB,
            data=frame_rgb,
        )

        # VIDEO mode requires strictly increasing timestamps.
        timestamp_ms = int(
            time.perf_counter()
            * 1000.0
        )

        timestamp_ms = max(
            timestamp_ms,
            self.last_timestamp_ms + 1
        )

        self.last_timestamp_ms = (
            timestamp_ms
        )

        t0 = time.perf_counter()

        result = (
            self.landmarker
            .detect_for_video(
                mp_image,
                timestamp_ms,
            )
        )

        self.last_inference_ms = (
            time.perf_counter() - t0
        ) * 1000.0

        if not result.pose_landmarks:
            return None

        landmarks = (
            result.pose_landmarks[0]
        )

        xy = np.zeros(
            (17, 2),
            dtype=np.float32
        )

        conf = np.zeros(
            17,
            dtype=np.float32
        )

        # ----------------------------------------------------
        # MediaPipe 33 -> COCO17
        # ----------------------------------------------------

        for coco_id, mp_id in enumerate(
            MEDIAPIPE_TO_COCO17
        ):

            lm = landmarks[mp_id]

            xy[coco_id, 0] = (
                float(lm.x) * w
            )

            xy[coco_id, 1] = (
                float(lm.y) * h
            )

            conf[coco_id] = (
                self._landmark_confidence(
                    lm
                )
            )

        # ----------------------------------------------------
        # Approximate bbox from landmarks
        # ----------------------------------------------------

        good = conf >= 0.20

        bbox = None

        if np.count_nonzero(good) >= 4:

            pts = xy[good]

            pad = 15.0

            bbox = np.array(
                [
                    max(
                        0.0,
                        float(
                            pts[:, 0].min()
                            - pad
                        )
                    ),

                    max(
                        0.0,
                        float(
                            pts[:, 1].min()
                            - pad
                        )
                    ),

                    min(
                        float(w - 1),
                        float(
                            pts[:, 0].max()
                            + pad
                        )
                    ),

                    min(
                        float(h - 1),
                        float(
                            pts[:, 1].max()
                            + pad
                        )
                    ),
                ],

                dtype=np.float32,
            )

        return PoseResult17(
            xy=xy,
            conf=conf,
            bbox=bbox
        )

    def close(self):

        self.landmarker.close()


# ============================================================
# Factory
# ============================================================

def create_pose_backend(
    backend: str,
    *,
    ultralytics_model:
        str = "yolo26n-pose.pt",
    ultralytics_imgsz:
        int = 640,
    person_conf:
        float = 0.30,
    mediapipe_model:
        str = "lite",
):

    backend = backend.lower()

    if backend == "ultralytics":

        return UltralyticsPoseBackend(
            model_name=
                ultralytics_model,
            imgsz=
                ultralytics_imgsz,
            person_conf=
                person_conf,
        )

    if backend == "mediapipe":

        return MediaPipePoseBackend(
            model_variant=
                mediapipe_model,
        )

    raise ValueError(
        f"Unknown backend: {backend}"
    )