import os
import cv2
import torch
import uuid
import trimesh
import numpy as np
import matplotlib.cm as cm

from mediapipe import solutions
from mediapipe.framework.formats import landmark_pb2

from lib.core.config import cfg
from lib.utils.human_models import mano
from lib.utils.mano_utils import change_flat_hand_mean



class ContactRenderer:
    def __init__(self):
        self.default_mesh_color = [130, 130, 130, 255]
        self.contact_mesh_color = [0, 255, 0, 255]
        # Initialize your hand model here
        with torch.no_grad():
            hand_pose = change_flat_hand_mean(np.zeros((48)), remove=True)
            mano_rest_out = mano.layer['right'](betas=torch.zeros((1, 10)), hand_pose=torch.from_numpy(hand_pose[None, 3:]).float(), global_orient=torch.zeros((1, 3)), transl=torch.zeros((1, 3)))
            self.hand_model_mano = trimesh.Trimesh(mano_rest_out.vertices[0], mano.watertight_face['right'])

    def export_contact_mesh(self, contact_mask, output_path=None):
        vis_contact = contact_mask == 1.0
        self.hand_model_mano.visual.vertex_colors = np.tile(self.default_mesh_color, (self.hand_model_mano.vertices.shape[0], 1))
        self.hand_model_mano.visual.vertex_colors[vis_contact] = self.contact_mesh_color
        if output_path is None:
            output_path = f"/tmp/contact_mesh_{uuid.uuid4().hex}.obj"
        self.hand_model_mano.export(output_path)
        return output_path


# This function is for demo code with mediapipe
MARGIN = 10  # pixels
FONT_SIZE = 1
FONT_THICKNESS = 1
HANDEDNESS_TEXT_COLOR = (88, 205, 54) # vibrant green


def draw_landmarks_on_image(rgb_image, detection_result):
    hand_landmarks_list = detection_result.hand_landmarks
    handedness_list = detection_result.handedness
    annotated_image = np.copy(rgb_image)
    right_hand_bbox = None

    best_score = -1.0
    best_idx = -1

    # Step 1: Find the index of the most confident right hand
    for idx in range(len(hand_landmarks_list)):
        handedness = handedness_list[idx][0]  # get Classification result
        if handedness.category_name == "Right" and handedness.score > best_score:
            best_score = handedness.score
            best_idx = idx

    # Step 2: If a right hand was found, draw it and extract bbox
    if best_idx != -1:
        hand_landmarks = hand_landmarks_list[best_idx]
        handedness = handedness_list[best_idx]

        # Convert landmarks to protobuf for drawing
        hand_landmarks_proto = landmark_pb2.NormalizedLandmarkList()
        hand_landmarks_proto.landmark.extend([
            landmark_pb2.NormalizedLandmark(x=lm.x, y=lm.y, z=lm.z) for lm in hand_landmarks
        ])
        solutions.drawing_utils.draw_landmarks(
            annotated_image,
            hand_landmarks_proto,
            solutions.hands.HAND_CONNECTIONS,
            solutions.drawing_styles.get_default_hand_landmarks_style(),
            solutions.drawing_styles.get_default_hand_connections_style())

        # Compute bbox
        height, width, _ = annotated_image.shape
        x_coords = [lm.x * width for lm in hand_landmarks]
        y_coords = [lm.y * height for lm in hand_landmarks]
        x_min, x_max = int(min(x_coords)), int(max(x_coords))
        y_min, y_max = int(min(y_coords)), int(max(y_coords))
        bb_c_x, bb_c_y = (x_min+x_max)/2, (y_min+y_max)/2
        bb_width, bb_height = x_max-x_min, y_max-y_min

        expand_ratio = cfg.DATASET.ho_big_bbox_expand_ratio

        bb_width_expand, bb_height_expand = expand_ratio * bb_width, expand_ratio * bb_height
        x_min_expand, y_min_expand = bb_c_x - 0.5 * bb_width_expand, bb_c_y - 0.5 * bb_height_expand
        
        right_hand_bbox = [x_min_expand, y_min_expand, bb_width_expand, bb_height_expand]

        # Draw bbox and label
        cv2.rectangle(annotated_image, (x_min, y_min), (x_max, y_max), (0, 255, 0), 2)
        cv2.putText(annotated_image, "Right Hand", (x_min, y_min - MARGIN),
                    cv2.FONT_HERSHEY_DUPLEX, FONT_SIZE, HANDEDNESS_TEXT_COLOR, FONT_THICKNESS, cv2.LINE_AA)

    return annotated_image, right_hand_bbox


def draw_landmarks_on_image_simple(rgb_image, right_hand_bbox):
    # Draw bbox and label
    annotated_image = rgb_image.copy()[..., ::-1]
    annotated_image = np.ascontiguousarray(annotated_image)

    # Expand bbox
    x_min, x_max = int(right_hand_bbox[0]), int(right_hand_bbox[0]+right_hand_bbox[2])
    y_min, y_max = int(right_hand_bbox[1]), int(right_hand_bbox[1]+right_hand_bbox[3])
    bb_c_x, bb_c_y = (x_min+x_max)/2, (y_min+y_max)/2
    bb_width, bb_height = x_max-x_min, y_max-y_min

    expand_ratio = cfg.DATASET.ho_big_bbox_expand_ratio

    bb_width_expand, bb_height_expand = expand_ratio * bb_width, expand_ratio * bb_height
    x_min_expand, y_min_expand = bb_c_x - 0.5 * bb_width_expand, bb_c_y - 0.5 * bb_height_expand
    
    right_hand_bbox = [x_min_expand, y_min_expand, bb_width_expand, bb_height_expand]

    x_min, y_min, x_max, y_max = int(right_hand_bbox[0]), int(right_hand_bbox[1]), int(right_hand_bbox[0]+right_hand_bbox[2]), int(right_hand_bbox[1]+right_hand_bbox[3])
    cv2.rectangle(annotated_image, (x_min, y_min), (x_max, y_max), (0, 255, 0), 2)
    cv2.putText(annotated_image, "Right Hand", (x_min, y_min - MARGIN), cv2.FONT_HERSHEY_DUPLEX, FONT_SIZE, HANDEDNESS_TEXT_COLOR, FONT_THICKNESS, cv2.LINE_AA)

    return annotated_image, right_hand_bbox