import cv2
import numpy as np
import matplotlib.pyplot as plt
import ultralytics
from ultralytics import YOLO
print(ultralytics.__version__)
from ultralytics.yolo.utils.ops import scale_image
model = YOLO('yolov8n-face.pt')
model_seg = YOLO('yolov8s-seg.pt')



def crop_passport_size(img_arr,size):
    if size==3545:
        passport_size = (310, 410)
    elif size==34:
        passport_size = (260, 346)
    # Load the input image
    # original_image = cv2.imread(input_image_path)
    original_image = img_arr
    # Define the dimensions for the passport-size photo (35mm x 45mm)
    # passport_size = (310, 410)

    # Resize the image to the passport-size dimensions
    resized_image = cv2.resize(original_image, (passport_size[0], passport_size[1]))

    # Save the cropped passport-size image
#     w_output = add_border(output, 4, 4, 4, 4, (0,0,0))
#     w_output = add_border(w_output, 15, 15, 15, 15, (255,255,255))
#     cv2.imwrite(output_image_path, resized_image)
    return resized_image

def add_border(image, top, bottom, left, right, color):
    return cv2.copyMakeBorder(image, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)

def horizontal_merge(number,crop_img,size):
    if size==3545:
        white_image = np.ones((450, 350, 3), dtype=np.uint8) * 255
        img_per_row = 6
    elif size==34:
        white_image = np.ones((386, 300, 3), dtype=np.uint8) * 255
        img_per_row = 7
    # white_image = np.ones((450, 350, 3), dtype=np.uint8) * 255
    list_img = []
    # for i in range(5):
    for i in range(number):
        list_img.append(crop_img)
    for i in range(img_per_row-number):
        list_img.append(white_image)
    merge_img = cv2.hconcat(list_img)
    return merge_img

def merge_all(number,input_path,size):
    
    input_image_path = input_path  # Replace with the path to your input image
    # output_image_path = output_path  # Replace with the desired output path
    output=crop_passport_size(input_image_path,size)

    w_output = add_border(output, 4, 4, 4, 4, (0,0,0))
    w_output = add_border(w_output, 15, 15, 15, 15, (255,255,255))
    w_output = add_border(w_output, 1, 1, 1, 1, (0,0,0))

    if size==3545:
        img_per_row = 6
        bottom_part = np.ones((270, 2100, 3), dtype=np.uint8) * 255
    elif size==34:
        img_per_row = 7
        bottom_part = np.ones((268, 2100, 3), dtype=np.uint8) * 255

    img_2 = []
    num_ = number
    if num_%img_per_row == 0:
        for i in range(int(num_/img_per_row)):
            img1 = horizontal_merge(img_per_row,w_output,size)
            img_2.append(img1)
    #         print('fill')
        for i in range(img_per_row-int(num_/img_per_row)):
            img1_1 = horizontal_merge(0,w_output,size)
            img_2.append(img1_1)
    #         print('blank')
    elif num_%img_per_row != 0:
        for i in range(int(num_//img_per_row)):
            img1 = horizontal_merge(img_per_row,w_output,size)
            img_2.append(img1)
    #         print('fill')
        img1 = horizontal_merge(num_%img_per_row,w_output,size)
        img_2.append(img1)
    #     print('some')
        for i in range(img_per_row-int(num_//img_per_row)-1):
            img1 = horizontal_merge(0,w_output,size)
            img_2.append(img1)
    #         print('blank')

    merge_img_final = cv2.vconcat(img_2)
    # print(merge_img_final.shape)
    # bottom_part = np.ones((270, 2100, 3), dtype=np.uint8) * 255
    merge_img_final = cv2.vconcat([merge_img_final,bottom_part])
    # print(merge_img_final.shape)
    merge_img_final = add_border(merge_img_final, 52, 32, 30, 30, (255,255,255))
    # print(merge_img_final.shape)
    # cv2.imwrite(output_path, merge_img_final)
    # print(merge_img_final.shape)
    return merge_img_final

def align_crop_image(img):
#     img = cv2.imread(img)
    img = cv2.copyMakeBorder(img, 50, 0, 0, 0, cv2.BORDER_CONSTANT, value=(255,255,255))
    results = model(img)
    if len(results[0].boxes.data)==1:
        b_box_up = int(results[0].boxes.data[0][0]),int(results[0].boxes.data[0][1])
        b_box_down = int(results[0].boxes.data[0][2]),int(results[0].boxes.data[0][3])
        left_eye=int(results[0].keypoints.data[0][0][0]),int(results[0].keypoints.data[0][0][1])
        right_eye=int(results[0].keypoints.data[0][1][0]),int(results[0].keypoints.data[0][1][1])
        nose=int(results[0].keypoints.data[0][2][0]),int(results[0].keypoints.data[0][2][1])
        left_lip=int(results[0].keypoints.data[0][3][0]),int(results[0].keypoints.data[0][3][1])
        right_lip=int(results[0].keypoints.data[0][4][0]),int(results[0].keypoints.data[0][4][1])
        left_eye_center = left_eye
        left_eye_x = left_eye_center[0] 
        left_eye_y = left_eye_center[1]

        right_eye_center = right_eye
        right_eye_x = right_eye_center[0]
        right_eye_y = right_eye_center[1]
        
        if left_eye_y > right_eye_y:
            A = (right_eye_x, left_eye_y)
            # Integer -1 indicates that the image will rotate in the clockwise direction
            direction = -1 
        else:
            A = (left_eye_x, right_eye_y)
            # Integer 1 indicates that image will rotate in the counter clockwise  
            # direction
            direction = 1 
        delta_x = right_eye_x - left_eye_x
        delta_y = right_eye_y - left_eye_y
        angle=np.arctan(delta_y/delta_x)
        angle = (angle * 180) / np.pi
        h, w = img.shape[:2]
        # Calculating a center point of the image
        # Integer division "//"" ensures that we receive whole numbers
        center = (w // 2, h // 2)
        # Defining a matrix M and calling
        # cv2.getRotationMatrix2D method
        M = cv2.getRotationMatrix2D(center, (angle), 1.0)
        # Applying the rotation to our image using the
        # cv2.warpAffine method
        rotated = cv2.warpAffine(img, M, (w, h))
        results = model(rotated)
        b_box_up = int(results[0].boxes.data[0][0]),int(results[0].boxes.data[0][1])
        b_box_down = int(results[0].boxes.data[0][2]),int(results[0].boxes.data[0][3])
        box_height = b_box_down[0]-b_box_up[0]
        box_width = b_box_down[1]-b_box_up[1]
        left_eye=int(results[0].keypoints.data[0][0][0]),int(results[0].keypoints.data[0][0][1])
        right_eye=int(results[0].keypoints.data[0][1][0]),int(results[0].keypoints.data[0][1][1])
        nose=int(results[0].keypoints.data[0][2][0]),int(results[0].keypoints.data[0][2][1])
        left_lip=int(results[0].keypoints.data[0][3][0]),int(results[0].keypoints.data[0][3][1])
        right_lip=int(results[0].keypoints.data[0][4][0]),int(results[0].keypoints.data[0][4][1])
        left_eye_center = left_eye
        left_eye_x = left_eye_center[0] 
        left_eye_y = left_eye_center[1]

        right_eye_center = right_eye
        right_eye_x = right_eye_center[0]
        right_eye_y = right_eye_center[1]

        up_crop = b_box_up[0]-int(box_height/2),b_box_up[1]-int(box_width/3)
        down_crop = b_box_down[0]+int(box_height/2),right_lip[1]+box_width
        final_image = rotated[up_crop[1]:down_crop[1],up_crop[0]:down_crop[0]]
        resized_image =cv2.resize(final_image, (210, 297))
        return resized_image
    else:
        return None
    
    
def predict_on_image(img, conf):
    result = model_seg(img, conf=conf)[0]

    # detection
    # result.boxes.xyxy   # box with xyxy format, (N, 4)
    cls = result.boxes.cls.cpu().numpy()    # cls, (N, 1)
    probs = result.boxes.conf.cpu().numpy()  # confidence score, (N, 1)
    boxes = result.boxes.xyxy.cpu().numpy()   # box with xyxy format, (N, 4)

    # segmentation
    if result.masks:
        masks = result.masks.data.cpu().numpy()     # masks, (N, H, W)
        masks = np.moveaxis(masks, 0, -1) # masks, (H, W, N)
        print(masks.shape)
        # rescale masks to original image
        masks = scale_image(masks, result.masks.orig_shape)
        masks = np.moveaxis(masks, -1, 0) # masks, (N, H, W)

        return boxes, masks, cls, probs
    else:
        return None


def overlay(image, mask, color, alpha, resize=None):
    """Combines image and its segmentation mask into a single image.
    https://www.kaggle.com/code/purplejester/showing-samples-with-segmentation-mask-overlay

    Params:
        image: Training image. np.ndarray,
        mask: Segmentation mask. np.ndarray,
        color: Color for segmentation mask rendering.  tuple[int, int, int] = (255, 0, 0)
        alpha: Segmentation mask's transparency. float = 0.5,
        resize: If provided, both image and its mask are resized before blending them together.
        tuple[int, int] = (1024, 1024))

    Returns:
        image_combined: The combined image. np.ndarray

    """
    for i in range(len(mask)):
        for j in range(len(mask[i])):
    #         print(masks[0][i][j])
            mask[i][j] = abs(mask[i][j] - 1)
    color = color[::-1]
    colored_mask = np.expand_dims(mask, 0).repeat(3, axis=0)
    colored_mask = np.moveaxis(colored_mask, 0, -1)
    masked = np.ma.MaskedArray(image, mask=colored_mask, fill_value=color)
    image_overlay = masked.filled()
    if resize is not None:
        image = cv2.resize(image.transpose(1, 2, 0), resize)
        image_overlay = cv2.resize(image_overlay.transpose(1, 2, 0), resize)

    image_combined = cv2.addWeighted(image, 1 - alpha, image_overlay, alpha, 0)

    return image_combined

def add_name_dob(img,name,dob):
    img = cv2.resize(img,(420,594))
    img = cv2.rectangle(img,(0,520),(420,594),(255,255,255),-1)
    dob = 'D.O.B: '+dob
    font = 1#cv2.FONT_HERSHEY_SIMPLEX
    font_scale = 2
    font_thickness = 2

    # Get the size of the text
    text_size = cv2.getTextSize(name, font, font_scale, font_thickness)[0]
    dob_size = cv2.getTextSize(dob, font, font_scale, font_thickness)[0]

    # Calculate the position to center the text
    text_x = (img.shape[1] - text_size[0]) // 2
    text_y = (img.shape[0] + text_size[1]) // 2
    dob_x = (img.shape[1] - dob_size[0]) // 2
    dob_y = (img.shape[0] + dob_size[1]) // 2

    # Put the text on the image
    cv2.putText(img, name, (text_x, 547), font, font_scale, (0, 0, 0), font_thickness)
    cv2.putText(img, dob, (dob_x, 583), font, font_scale, (0, 0, 0), font_thickness)

    return img