import numpy as np
import cv2
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import make_axes_locatable

def getDenseMask(landmarks, h, w):
    RL, LL, H = landmarks[:44], landmarks[44:94], landmarks[94:]
    img = np.zeros([h, w], dtype='uint8')
    RL = RL.reshape(-1, 1, 2).astype('int')
    LL = LL.reshape(-1, 1, 2).astype('int')
    H = H.reshape(-1, 1, 2).astype('int')
    img = cv2.drawContours(img, [RL], -1, 1, -1)
    img = cv2.drawContours(img, [LL], -1, 1, -1)
    img = cv2.drawContours(img, [H], -1, 2, -1)
    return img

def drawOnTop(img, landmarks, original_shape):
    h, w = original_shape
    output = getDenseMask(landmarks, h, w)
    image = np.zeros([h,w,3])
    image[:,:,0] = img + 0.3*(output==1).astype('float') - 0.1*(output==2).astype('float')
    image[:,:,1] = img + 0.3*(output==2).astype('float') - 0.1*(output==1).astype('float')
    image[:,:,2] = img - 0.1*(output==1).astype('float') - 0.2*(output==2).astype('float')
    image = np.clip(image,0,1)
    RL, LL, H = landmarks[:44], landmarks[44:94], landmarks[94:]
    for l in RL: image = cv2.circle(image,(int(l[0]),int(l[1])),5,(1,0,1),-1)
    for l in LL: image = cv2.circle(image,(int(l[0]),int(l[1])),5,(1,0,1),-1)
    for l in H: image = cv2.circle(image,(int(l[0]),int(l[1])),5,(1,1,0),-1)
    return image

def create_overlay(img, landmarks):
    h, w = img.shape[:2]
    dense_mask = getDenseMask(landmarks, h, w)
    overlay = np.zeros([h, w, 3])
    
    overlay[:,:,0] = img + 0.3 * (dense_mask == 1).astype('float') - 0.1 * (dense_mask == 2).astype('float')
    overlay[:,:,1] = img + 0.3 * (dense_mask == 2).astype('float') - 0.1 * (dense_mask == 1).astype('float')
    overlay[:,:,2] = img - 0.1 * (dense_mask == 1).astype('float') - 0.2 * (dense_mask == 2).astype('float')
    overlay = np.clip(overlay, 0, 1)
    
    return overlay

def plot_side_by_side_comparison(img_orig, means_orig, uncertainty_orig, img_corr, means_corr, uncertainty_corr):

    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 7))
    
    fig.set_constrained_layout(True)

    vmax = max(np.max(np.mean(uncertainty_orig, axis=1)), np.max(np.mean(uncertainty_corr, axis=1)))

    # --- Original ---
    overlay_orig = create_overlay(img_orig, means_orig)
    ax1.imshow(overlay_orig)
    scatter1 = ax1.scatter(
        means_orig[:, 0], means_orig[:, 1],
        c=np.mean(uncertainty_orig, axis=1),
        cmap='hot', s=50, vmin=0, vmax=vmax
    )
    ax1.set_title("Original", fontsize=16, pad=10)
    ax1.axis('off')

    # --- Corrupted ---
    overlay_corr = create_overlay(img_corr, means_corr)
    ax2.imshow(overlay_corr)
    scatter2 = ax2.scatter(
        means_corr[:, 0], means_corr[:, 1],
        c=np.mean(uncertainty_corr, axis=1),
        cmap='hot', s=50, vmin=0, vmax=vmax
    )
    ax2.set_title("Corrupted", fontsize=16, pad=10)
    ax2.axis('off')

    # Shared colorbar
    cbar = fig.colorbar(scatter2, ax=[ax1, ax2], fraction=0.046, pad=0.01, shrink=0.85)
    cbar.ax.tick_params(labelsize=10)

    return fig