import cv2
import numpy as np
import mediapipe as mp

mp_holistic = mp.solutions.holistic  # Holistic model
mp_drawing = mp.solutions.drawing_utils  # Drawing utilities


def mediapipe_detection(image, model):
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)  # Color conversion from BGR to RGB
    image.flags.writeable = False                   # Image is no longer writeable
    results = model.process(image)                  # Make prediction
    image.flags.writeable = True                    # Image is no longer writeable
    image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)  # Color conversion RGB to BGR
    return image, results


def draw_styled_landmarks(image,results):

    # Draw pose connection
    mp_drawing.draw_landmarks(image, results.pose_landmarks, mp_holistic.POSE_CONNECTIONS,
                              mp_drawing.DrawingSpec(color=(0, 0, 255), thickness=1, circle_radius=1),
                              mp_drawing.DrawingSpec(color=(80, 110, 10), thickness=1, circle_radius=1)
                              )
    # Draw left hand connection
    mp_drawing.draw_landmarks(image, results.left_hand_landmarks, mp_holistic.HAND_CONNECTIONS,
                              mp_drawing.DrawingSpec(color=(0, 0, 255), thickness=1, circle_radius=2),
                              mp_drawing.DrawingSpec(color=(80, 110, 10), thickness=1, circle_radius=1)
                              )
    # Draw right hand connection
    mp_drawing.draw_landmarks(image, results.right_hand_landmarks, mp_holistic.HAND_CONNECTIONS,
                              mp_drawing.DrawingSpec(color=(0, 0, 255), thickness=1, circle_radius=2),
                              mp_drawing.DrawingSpec(color=(80, 110, 10), thickness=1, circle_radius=1)
                              )


def extract_keypoints(results):
    pose = np.array([[res.x, res.y, res.z, res.visibility] for res in results.pose_landmarks.landmark]).flatten() if results.pose_landmarks else np.zeros(33*4)
    lh = np.array([[res.x, res.y, res.z] for res in results.left_hand_landmarks.landmark]).flatten() if results.left_hand_landmarks else np.zeros(21*3)
    rh = np.array([[res.x, res.y, res.z] for res in results.right_hand_landmarks.landmark]).flatten() if results.right_hand_landmarks else np.zeros(21*3)

    return np.concatenate([pose, lh, rh])


# --- STGCN Helpers ---
def get_adjacency_matrix():
    A = np.eye(75)
    for conn in mp_holistic.POSE_CONNECTIONS:
        A[conn[0], conn[1]] = 1;
        A[conn[1], conn[0]] = 1
    for conn in mp_holistic.HAND_CONNECTIONS:
        A[conn[0] + 33, conn[1] + 33] = 1;
        A[conn[1] + 33, conn[0] + 33] = 1
        A[conn[0] + 54, conn[1] + 54] = 1;
        A[conn[1] + 54, conn[0] + 54] = 1
    return A


def reshape_for_stgcn(X):
    N, T, _ = X.shape
    X_new = np.zeros((N, T, 75, 3))
    for i in range(N):
        for t in range(T):
            frame = X[i, t]
            pose = frame[0:132].reshape(33, 4)[:, :3]
            lh = frame[132:195].reshape(21, 3)
            rh = frame[195:258].reshape(21, 3)
            X_new[i, t] = np.concatenate([pose, lh, rh], axis=0)
    return X_new.transpose(0, 3, 1, 2)