from PIL import Image
from face_detection import detect_faces
from heatmap import get_heatmap
from pose import detect_pose, draw_pose

def prepare_image(img_path):
    # detect face
    annotated_image, face_bboxes = detect_faces(img_path)
    # detect pose + bounding_box
    pose_results = detect_pose(img_path)
    pose_result = pose_results[0]
    body_bboxes, body_keypoints = pose_result.boxes, pose_result.keypoints 
    # generate heatmap
    heatmap = get_heatmap(img_path)
    heatmap_pil = Image.fromarray(heatmap).convert('RGB')
    # pose on heatmap
    heatmap_n_pose = draw_pose(heatmap_pil,body_keypoints,0.5) #PIL.Image
    
    # area in bounding_box below head (rectangular)
    face_y1, face_h = face_bboxes[0][1], face_bboxes[0][3] #xywh of first face
    face_y2 = face_y1 + face_h
    body_box = body_bboxes.cpu().numpy().xyxy[0]
    body_box[1] = face_y2 #shift below head
    
    # Calculate the coordinates for cropping
    x1, y1, x2, y2 = map(int, body_box)
    cropped_region = heatmap_n_pose.crop((x1, y1, x2, y2))
    original_image = Image.open(img_path)
    original_image.paste(cropped_region, (x1, y1))
    
    # Return the modified image
    return original_image