import os
import cv2
import numpy as np
from tqdm import tqdm
import matplotlib.pyplot as plt

def _sort_points(points):
    sort_xy = np.array(sorted(points , key=lambda k: [k[0], k[1]]))
    left_top, left_bottom = (sort_xy[0],sort_xy[1]) if sort_xy[0][1] < sort_xy[1][1] else (sort_xy[1],sort_xy[0])
    right_top, right_bottom = (sort_xy[2],sort_xy[3]) if sort_xy[2][1] < sort_xy[3][1] else (sort_xy[3],sort_xy[2])

    return(np.array([left_top,left_bottom,right_top,right_bottom]))


def _clip_and_mapped_points(corners,x1,y1,x2,y2):
    origin = (x1,y1)
    clip_x = np.clip(corners[:,0],x1,x2)
    clip_y = np.clip(corners[:,1],y1,y2)
    cleft_top,cleft_bottom,cright_top,cright_bottom = np.array(list((zip(clip_x,clip_y))))

    Tleft_top = np.subtract(cleft_top,origin)
    Tleft_bottom = np.subtract(cleft_bottom,origin)
    Tright_top = np.subtract(cright_top,origin)
    Tright_bottom = np.subtract(cright_bottom,origin)

    return np.array([Tleft_top,Tleft_bottom,Tright_top,Tright_bottom],dtype=np.float32)


def perspective_transform(obb_pred, dst_width=300, dst_height=80, plot=False):
    crops = []
    transformed = []
    
    for res in obb_pred[0]:
        x1,y1,x2,y2 = np.array(res.obb.xyxy.cpu().numpy()[0], dtype=np.int32)
        obb_points = res.obb.xyxyxyxy.cpu().numpy()
        
        sorted_corners = _sort_points(obb_points.reshape(4,2))
        mapped_points = _clip_and_mapped_points(sorted_corners,x1,y1,x2,y2)

        org_img = res.orig_img
        crop_lp = org_img[y1:y2,x1:x2,:]
        crops.append(crop_lp)
        
        homography = cv2.getPerspectiveTransform(mapped_points, np.float32([[0,0],[0,dst_height],[dst_width,0],[dst_width,dst_height]]))
        transformed_lp = cv2.warpPerspective(crop_lp, np.float32(homography), (dst_width, dst_height), flags=cv2.INTER_LINEAR)
        transformed.append(transformed_lp)
        
    if plot:
        fig, axes = plt.subplots(len(crops), 2, figsize=(6, 2*len(crops)))
        for i,(crop, transformed) in enumerate(zip(crops, transformed)): 
            axes[i, 0].imshow(crop)
            axes[i, 0].set_title(f'cropped {i + 1}')

            axes[i, 1].imshow(transformed)
            axes[i, 1].set_title('Transformed')

        for ax in axes.flatten():
            ax.axis('off')

        plt.tight_layout()
        plt.show()

    return crops , transformed