import numpy as np

class ArmController:
    def __init__(self, config, mirror=False):
        self.config = config
        self.mirror = mirror
        self.shoulder_idx = {'left': 11, 'right': 12}
        self.wrist_idx = {'left': 15, 'right': 16}
        self.hip_idx = {'left': 23, 'right': 24}
        self.dead_zone = config.get('dead_zone', 0.2)
        self.debug = config.get('debug', False)
        
    def _get_side_indices(self, side):
        """
        Возвращает (shoulder_idx, wrist_idx) для заданной стороны (left/right).
        При mirror=True интерпретируем сторону как в реальности: левая/правая рука.
        """
        if self.mirror:
            if side == 'left':
                s_idx = 12
                w_idx = 16
            else:
                s_idx = 11
                w_idx = 15
        else:
            if side == 'left':
                s_idx = 11
                w_idx = 15
            else:
                s_idx = 12
                w_idx = 16
        return s_idx, w_idx

    def _get_shoulder_width(self, landmarks):
        """Ширина плеч для нормировки горизонтальных смещений."""
        left = landmarks[11][:2]
        right = landmarks[12][:2]
        width = np.linalg.norm(right - left)
        if width < 50 or width > 300:
            return None
        return width

    def _get_torso_height(self, landmarks):
        """Высота торса для нормировки вертикальных смещений."""
        left_shoulder = landmarks[11][:2]
        right_shoulder = landmarks[12][:2]
        left_hip = landmarks[23][:2]
        right_hip = landmarks[24][:2]
        shoulder_center = (left_shoulder + right_shoulder) / 2
        hip_center = (left_hip + right_hip) / 2
        height = np.linalg.norm(shoulder_center - hip_center)
        if height < 50:
            return None
        return height

    def _horizontal_displacement_rel(self, landmarks, side):
        """
        Нормированное горизонтальное смещение запястья относительно плеча.
        Сторона `side` — это реальная сторона руки
        """
        s_idx, w_idx = self._get_side_indices(side)

        if landmarks[s_idx][3] < 0.5 or landmarks[w_idx][3] < 0.5:
            return 0.0

        shoulder = landmarks[s_idx][:2]
        wrist = landmarks[w_idx][:2]

        shoulder_width = self._get_shoulder_width(landmarks)
        if shoulder_width is None:
            return 0.0

        disp = wrist[0] - shoulder[0]
        return disp / shoulder_width

    def _vertical_displacement_rel(self, landmarks, side):
        """
        Вертикальное смещение: верх/низ запястья относительно плеча.
        Сторона `side` — реальная сторона руки.
        """
        s_idx, w_idx = self._get_side_indices(side)

        if landmarks[s_idx][3] < 0.5 or landmarks[w_idx][3] < 0.5:
            return 0.0

        shoulder = landmarks[s_idx][:2]
        wrist = landmarks[w_idx][:2]

        torso_height = self._get_torso_height(landmarks)
        if torso_height is None:
            return 0.0

        disp = shoulder[1] - wrist[1]
        return disp / torso_height

    def compute_speeds(self, landmarks):
        if (landmarks[11][3] < 0.5 or landmarks[12][3] < 0.5 or
            landmarks[15][3] < 0.5 or landmarks[16][3] < 0.5):
            if self.debug:
                print("Руки не видны")
            return 0.0, 0.0

        linear_side = self.config['linear_arm']
        angular_side = self.config['angular_arm']

        lin_rel = self._horizontal_displacement_rel(landmarks, linear_side)
        ang_rel = self._vertical_displacement_rel(landmarks, angular_side)

        if self.debug:
            print(f"lin_rel={lin_rel:.3f}, ang_rel={ang_rel:.3f}")

        # Линейная скорость (только вперёд)
        if lin_rel < self.dead_zone:
            linear = 0.0
        else:
            linear = min(lin_rel, 1.0) * self.config['max_speed_linear']

        # Угловая скорость
        if abs(ang_rel) < self.dead_zone:
            angular = 0.0
        else:
            ang_rel_clipped = np.clip(ang_rel, -1.0, 1.0)
            angular = ang_rel_clipped * self.config['max_speed_angular']

        return linear, angular