import numpy as np
import cv2


def _clip_unit(v):
    return float(np.clip(v, -1.0, 1.0))


def _apply_dead_zone(v, dz):
    return 0.0 if abs(v) < dz else v


def _robust_metrics(landmarks, min_conf=0.5):
    """
    Compute shoulder_center, shoulder_width, torso_height robustly.
    Uses hips if available; otherwise falls back to nose/shoulder geometry.

    Returns:
        shoulder_center (np.array shape (2,))
        shoulder_width (float)
        torso_height (float)
        ok (bool)
    """
    def pt(i):
        return np.array(landmarks[i][:2], dtype=float), float(landmarks[i][3])

    l_sh, c_lsh = pt(11)
    r_sh, c_rsh = pt(12)
    if c_lsh < min_conf or c_rsh < min_conf:
        return None, 0.0, 0.0, False

    shoulder_center = (l_sh + r_sh) / 2.0
    shoulder_width = float(np.linalg.norm(r_sh - l_sh))
    if shoulder_width < 1e-3:
        return shoulder_center, 0.0, 0.0, False

    # Try hips
    l_hip, c_lhip = pt(23)
    r_hip, c_rhip = pt(24)

    if c_lhip >= min_conf and c_rhip >= min_conf:
        hip_center = (l_hip + r_hip) / 2.0
        torso_height = float(np.linalg.norm(hip_center - shoulder_center))
        if torso_height >= 1e-3:
            return shoulder_center, shoulder_width, torso_height, True

    # Fallbacks (upper-body only)
    nose, c_nose = pt(0)
    if c_nose >= min_conf:
        nose_to_shoulder = abs(nose[1] - shoulder_center[1])
        torso_height = max(1.6 * nose_to_shoulder, 0.9 * shoulder_width)
    else:
        torso_height = max(1.2 * shoulder_width, 1.0)

    return shoulder_center, shoulder_width, float(torso_height), True


class ArmControllerMethod1:
    """
    Method 1: Single-hand driving with right wrist.
    - Linear: vertical offset of right wrist from shoulder center (normalized by torso height)
    - Angular: horizontal offset of right wrist from shoulder center (normalized by shoulder width)
    """
    def __init__(self, config, mirror=False):
        self.config = config
        self.mirror = mirror
        self.dead_zone = config.get('dead_zone', 0.1)
        self.debug = config.get('debug', False)
        self.min_conf = config.get('min_conf', 0.5)

    def compute_speeds(self, landmarks, frame_shape=None):
        if landmarks is None:
            return 0.0, 0.0

        # Require: shoulders + right wrist
        need = [11, 12, 16]
        if any(landmarks[i][3] < self.min_conf for i in need):
            return 0.0, 0.0

        shoulder_center, shoulder_width, torso_height, ok = _robust_metrics(landmarks, self.min_conf)
        if not ok or shoulder_width < 1e-3 or torso_height < 1e-3:
            return 0.0, 0.0

        r_wr = landmarks[16][:2]

        # Positive linear when wrist above shoulder center (forward)
        linear = (shoulder_center[1] - r_wr[1]) / torso_height
        # Positive angular when wrist to the right of shoulder center
        angular = (r_wr[0] - shoulder_center[0]) / shoulder_width

        if self.mirror:
            angular = -angular

        linear = _clip_unit(_apply_dead_zone(linear, self.dead_zone))
        angular = _clip_unit(_apply_dead_zone(angular, self.dead_zone))

        if self.debug:
            print(f"[M1] L:{linear:.2f} A:{angular:.2f}")
        return linear, angular

    def draw_overlay(self, frame, landmarks=None):
        if frame is None:
            return frame
        h, w = frame.shape[:2]
        # Draw center cross (screen center approximation)
        cv2.line(frame, (w // 2, 0), (w // 2, h), (0, 0, 0), 1)
        cv2.line(frame, (0, h // 2), (w, h // 2), (0, 0, 0), 1)
        # Draw right wrist
        if landmarks is not None and landmarks[16][3] > 0.5:
            x, y = int(landmarks[16][0]), int(landmarks[16][1])
            cv2.circle(frame, (x, y), 8, (0, 255, 255), -1)
        return frame


class ArmControllerMethod2:
    """
    Method 2: Two-hand blended control.
    - Linear: average vertical offset of both wrists from shoulder center (normalized by torso height)
    - Angular: horizontal balance of wrists around shoulder center (normalized by shoulder width)
    """
    def __init__(self, config, mirror=False):
        self.config = config
        self.mirror = mirror
        self.dead_zone = config.get('dead_zone', 0.1)
        self.debug = config.get('debug', False)
        self.min_conf = config.get('min_conf', 0.5)

    def compute_speeds(self, landmarks, frame_shape=None):
        if landmarks is None:
            return 0.0, 0.0

        # Require: shoulders + both wrists
        need = [11, 12, 15, 16]
        if any(landmarks[i][3] < self.min_conf for i in need):
            return 0.0, 0.0

        shoulder_center, shoulder_width, torso_height, ok = _robust_metrics(landmarks, self.min_conf)
        if not ok or shoulder_width < 1e-3 or torso_height < 1e-3:
            return 0.0, 0.0

        l_wr = landmarks[15][:2]
        r_wr = landmarks[16][:2]

        # Linear: average elevation of both wrists
        lin_l = (shoulder_center[1] - l_wr[1]) / torso_height
        lin_r = (shoulder_center[1] - r_wr[1]) / torso_height
        linear = 0.5 * (lin_l + lin_r)

        # Angular: horizontal balance
        ang = ((r_wr[0] - shoulder_center[0]) - (shoulder_center[0] - l_wr[0])) / shoulder_width
        angular = ang

        if self.mirror:
            angular = -angular

        linear = _clip_unit(_apply_dead_zone(linear, self.dead_zone))
        angular = _clip_unit(_apply_dead_zone(angular, self.dead_zone))

        if self.debug:
            print(f"[M2] L:{linear:.2f} A:{angular:.2f}")
        return linear, angular

    def draw_overlay(self, frame, landmarks=None):
        if frame is None:
            return frame
        if landmarks is not None:
            for idx, color in [(15, (255, 0, 255)), (16, (0, 255, 255))]:
                if landmarks[idx][3] > 0.5:
                    x, y = int(landmarks[idx][0]), int(landmarks[idx][1])
                    cv2.circle(frame, (x, y), 8, color, -1)
        return frame


class ArmControllerMethod3:
    """
    Method 3: Elbow-augmented control.
    - Linear: average vertical offset of elbows (normalized by torso height)
    - Angular: wrist horizontal balance (normalized by shoulder width)
    """
    def __init__(self, config, mirror=False):
        self.config = config
        self.mirror = mirror
        self.dead_zone = config.get('dead_zone', 0.1)
        self.debug = config.get('debug', False)
        self.min_conf = config.get('min_conf', 0.5)

    def compute_speeds(self, landmarks, frame_shape=None):
        if landmarks is None:
            return 0.0, 0.0

        # Require shoulders; prefer elbows for linear; wrists for angular.
        need_base = [11, 12]
        if any(landmarks[i][3] < self.min_conf for i in need_base):
            return 0.0, 0.0

        elbows_ok = (landmarks[13][3] >= self.min_conf and landmarks[14][3] >= self.min_conf)
        wrists_ok = (landmarks[15][3] >= self.min_conf and landmarks[16][3] >= self.min_conf)

        if not elbows_ok and not wrists_ok:
            return 0.0, 0.0

        shoulder_center, shoulder_width, torso_height, ok = _robust_metrics(landmarks, self.min_conf)
        if not ok or shoulder_width < 1e-3 or torso_height < 1e-3:
            return 0.0, 0.0

        # Linear: prefer elbows, fallback to wrists average if elbows missing
        if elbows_ok:
            l_el = landmarks[13][:2]
            r_el = landmarks[14][:2]
            lin_l = (shoulder_center[1] - l_el[1]) / torso_height
            lin_r = (shoulder_center[1] - r_el[1]) / torso_height
            linear = 0.5 * (lin_l + lin_r)
        else:
            l_wr = landmarks[15][:2]
            r_wr = landmarks[16][:2]
            lin_l = (shoulder_center[1] - l_wr[1]) / torso_height
            lin_r = (shoulder_center[1] - r_wr[1]) / torso_height
            linear = 0.5 * (lin_l + lin_r)

        # Angular: use wrists if available, else 0
        if wrists_ok:
            l_wr = landmarks[15][:2]
            r_wr = landmarks[16][:2]
            angular = ((r_wr[0] + l_wr[0]) - 2 * shoulder_center[0]) / shoulder_width
        else:
            angular = 0.0

        if self.mirror:
            angular = -angular

        linear = _clip_unit(_apply_dead_zone(linear, self.dead_zone))
        angular = _clip_unit(_apply_dead_zone(angular, self.dead_zone))

        if self.debug:
            print(f"[M3] L:{linear:.2f} A:{angular:.2f}")
        return linear, angular

    def draw_overlay(self, frame, landmarks=None):
        if frame is None:
            return frame
        if landmarks is not None:
            for idx, color in [(13, (0, 200, 0)), (14, (0, 200, 0)), (15, (0, 255, 255)), (16, (255, 0, 255))]:
                if landmarks[idx][3] > 0.5:
                    x, y = int(landmarks[idx][0]), int(landmarks[idx][1])
                    cv2.circle(frame, (x, y), 6, color, -1)
        return frame


class ArmControllerMethod4:
    """
    Method 4: 3x3 grid based on landmark 19 (right index finger tip).
    Screen split at 2/5 and 3/5 (both axes). Center band = 0.
    Proportional speed away from the center bands.
    """
    def __init__(self, config, mirror=False):
        self.config = config
        self.mirror = mirror
        self.finger_idx = 19  # right index finger tip
        self.debug = config.get('debug', False)

    def compute_speeds(self, landmarks, frame_shape=None):
        linear = 0.0
        angular = 0.0

        if frame_shape is None or landmarks is None:
            return 0.0, 0.0
        if landmarks[self.finger_idx][3] < 0.5:
            return 0.0, 0.0

        h, w = int(frame_shape[0]), int(frame_shape[1])
        x = landmarks[self.finger_idx][0]
        y = landmarks[self.finger_idx][1]

        # Angular (horizontal): center band 2/5..3/5 = 0
        if 2 * w / 5 <= x <= 3 * w / 5:
            angular = 0.0
        elif x > 3 * w / 5:
            angular = (x * 5) / (2 * w) - 1
        else:
            angular = (x - 3 * w / 5) / (2 * w / 5)

        # Linear (vertical): center band 2/5..3/5 = 0
        if 2 * h / 5 <= y <= 3 * h / 5:
            linear = 0.0
        elif y > 3 * h / 5:
            linear = -((y * 5) / (2 * h) - 1)
        else:
            linear = -(y - 3 * h / 5) / (2 * h / 5)

        if self.mirror:
            angular = -angular

        if self.debug:
            print(f"[M4] L:{linear:.2f} A:{angular:.2f}")
        return _clip_unit(linear), _clip_unit(angular)

    def draw_overlay(self, frame, landmarks=None):
        if frame is None:
            return frame

        h, w = frame.shape[:2]
        x1, x2 = int(w * 2 / 5), int(w * 3 / 5)
        y1, y2 = int(h * 2 / 5), int(h * 3 / 5)

        # Grid lines
        cv2.line(frame, (x1, 0), (x1, h), (0, 0, 0), 2)
        cv2.line(frame, (x2, 0), (x2, h), (0, 0, 0), 2)
        cv2.line(frame, (0, y1), (w, y1), (0, 0, 0), 2)
        cv2.line(frame, (0, y2), (w, y2), (0, 0, 0), 2)

        # Highlight active cell + finger
        if landmarks is not None and landmarks[self.finger_idx][3] > 0.5:
            fx, fy = int(landmarks[self.finger_idx][0]), int(landmarks[self.finger_idx][1])
            cx0, cx1 = (0, x1) if fx < x1 else ((x2, w) if fx > x2 else (x1, x2))
            cy0, cy1 = (0, y1) if fy < y1 else ((y2, h) if fy > y2 else (y1, y2))

            overlay = frame.copy()
            cv2.rectangle(overlay, (cx0, cy0), (cx1, cy1), (0, 255, 255), -1)
            frame = cv2.addWeighted(overlay, 0.2, frame, 0.8, 0)

            cv2.circle(frame, (fx, fy), 8, (0, 255, 255), -1)
            cv2.circle(frame, (fx, fy), 12, (0, 120, 120), 2)

        return frame