import numpy as np
import cv2


class ArmController:
    def __init__(self, config, mirror=False):
        self.config = config
        self.mirror = mirror

        # arms assignment
        self.linear_arm = config.get('linear_arm', 'right')
        self.angular_arm = config.get('angular_arm', 'left')

        self.wrist_idx = {'left': 15, 'right': 16}
        self.shoulder_idx = {'left': 11, 'right': 12}
        self.hip_idx = {'left': 23, 'right': 24}

        self.dead_zone = config.get('dead_zone', 0.2)
        self.debug = config.get('debug', False)

        # separate max speeds
        self.max_speed_linear = config.get('max_speed_linear', 1.0)
        self.max_speed_angular = config.get('max_speed_angular', 1.0)

    def _get_side_indices(self, side):
        """Returns the wrist landmark index for a logical side, accounting for mirror."""
        if self.mirror:
            w_idx = self.wrist_idx['right'] if side == 'left' else self.wrist_idx['left']
        else:
            w_idx = self.wrist_idx[side]
        return w_idx

    def _get_wrist_norm(self, landmarks, side):
        """
        Returns normalized wrist position in the range -1..1 relative to the body.
        Uses shoulders as horizontal reference and shoulder/hip as vertical reference.
        Returns None if the wrist is not visible.
        landmarks: array [N, 4] -> [x, y, z, visibility], x/y are normalized 0..1 (MediaPipe).
        """
        w_idx = self._get_side_indices(side)

        if landmarks[w_idx][3] < 0.5:
            return None

        wx = float(landmarks[w_idx][0])
        wy = float(landmarks[w_idx][1])

        # Reference points (body center)
        ls_x, ls_y = landmarks[self.shoulder_idx['left']][0], landmarks[self.shoulder_idx['left']][1]
        rs_x, rs_y = landmarks[self.shoulder_idx['right']][0], landmarks[self.shoulder_idx['right']][1]
        lh_y = landmarks[self.hip_idx['left']][1]
        rh_y = landmarks[self.hip_idx['right']][1]

        center_x = (ls_x + rs_x) / 2.0
        shoulder_y = (ls_y + rs_y) / 2.0
        hip_y = (lh_y + rh_y) / 2.0

        shoulder_width = abs(ls_x - rs_x)
        torso_height = abs(hip_y - shoulder_y)

        # avoid division by zero
        shoulder_width = max(shoulder_width, 1e-3)
        torso_height = max(torso_height, 1e-3)

        # normalized offset from body center, scaled by body size
        nx = (wx - center_x) / shoulder_width
        # vertical: up = positive; measured from shoulder line
        ny = (shoulder_y - wy) / torso_height

        nx = float(np.clip(nx, -1.0, 1.0))
        ny = float(np.clip(ny, -1.0, 1.0))

        return nx, ny

    def _apply_dead_zone(self, value):
        """Removes dead zone and rescales -1..1 -> -1..1."""
        if abs(value) < self.dead_zone:
            return 0.0
        sign = 1.0 if value > 0 else -1.0
        scaled = (abs(value) - self.dead_zone) / (1.0 - self.dead_zone)
        return sign * float(np.clip(scaled, 0.0, 1.0))

    def compute_speeds(self, landmarks):
        """
        Main API for the ROS node.
        Returns (linear, angular), both in the range -1..1.

        - linear   : controlled by the vertical position of the linear_arm wrist (up = forward)
        - angular  : controlled by the horizontal position of the angular_arm wrist
        """
        if landmarks is None:
            return 0.0, 0.0

        # ---- Linear speed from linear_arm (vertical axis) ----
        linear = 0.0
        lin_pos = self._get_wrist_norm(landmarks, self.linear_arm)
        if lin_pos is not None:
            _, ny = lin_pos
            linear = self._apply_dead_zone(ny) * self.max_speed_linear

        # ---- Angular speed from angular_arm (horizontal axis) ----
        angular = 0.0
        ang_pos = self._get_wrist_norm(landmarks, self.angular_arm)
        if ang_pos is not None:
            nx, _ = ang_pos
            # invert so that hand-to-the-left = turn left (adjust sign if needed)
            angular = -self._apply_dead_zone(nx) * self.max_speed_angular

        linear = float(np.clip(linear, -1.0, 1.0))
        angular = float(np.clip(angular, -1.0, 1.0))

        if self.debug:
            print(f"[ArmController] linear={linear:.2f} angular={angular:.2f}")

        return linear, angular