import  torch,logging,imageio 
import mediapipe as mp
import numpy as np
from PIL import Image, ImageEnhance
logging.basicConfig(level=logging.INFO)

import signapse.constants as C
from signapse.create import GENERATOR
from signapse.mp_utils import INTERPLATION, IMAGE_PROCESSING
from signapse.sam2 import get_masks_sam2

mp_face_mesh = mp.solutions.face_mesh
face_mesh =  mp_face_mesh.FaceMesh(max_num_faces=1,
                                    static_image_mode=True,
                                    refine_landmarks=True,
                                    min_detection_confidence=0.8,
                                    min_tracking_confidence=0.8)
face_oval = list(mp_face_mesh.FACEMESH_FACE_OVAL)
face_mesh_mop = [num-176 for num in C.FACE_MESH_MOP]

class SAM_CROPS():
    def __init__(self, args):
        super(SAM_CROPS,self).__init__()
        self.args=args  
        
        
    def get_face_landmarks(self,image):
        image.flags.writeable = False
        results = face_mesh.process(image)
        image.flags.writeable = True    
        if results.multi_face_landmarks:
            data = results.multi_face_landmarks[0] 
            face_landmarks = np.array([[landmark.x, landmark.y] for landmark in data.landmark])
        else:
            face_landmarks = np.zeros((478,2))
        return face_landmarks
    
    def sam_masks(self,img_centre):
        if (self.args.opt.face_crop) or (self.args.opt.hand_crop):
            if (self.args.opt.hand_crop):
                mask = get_masks_sam2(self.args.input_video_path, self.args.stop,tmp_dir = "./inputs/tmp")
                np.save(f'./inputs/SAM_files/{self.args.sign_request_id}.npy',mask)

            if self.args.opt.face_crop:
                mask2=[]
                GENERATOR(self.args).mkdir_safe("./inputs/Face_files")       
                reader = imageio.get_reader(self.args.input_video_path)
                for i, frame in enumerate(reader): 
                    if i < self.args.start:
                        continue     
                    if i% 50==0:
                        logging.info("Frame: {}".format(i))
                    if self.args.stop != -1 and i >= self.args.stop:
                        break        
                    frame = IMAGE_PROCESSING().apply_norm(frame,img_centre,self.args.opt.loadSize)             
                    face_mask = self.get_face_landmarks(frame)    
                    mask2.append(face_mask)
            
                if self.args.opt.smooth_face_crop:
                    mask2 = INTERPLATION().interpolate_All_MP_Tensors(torch.from_numpy(np.array(mask2)),no_zero_start=False)
                    Points = mask2[0,48:]  # compute the point distance based on frame 0
                    mask2[:,face_mesh_mop] = IMAGE_PROCESSING().mouth_centering(mask2[:,face_mesh_mop], sigma=9)
                    # Using the center points of the mouth [13,:] as the center of the face
                    dist_to_centre = np.linalg.norm(Points - Points[13,:], axis=-1)
                    dist_to_centre = dist_to_centre/np.max(dist_to_centre)
                    smoothed_signal = np.array([INTERPLATION().smooth_keypoints(mask2[:,i], i, dist_to_centre[i-48]) for i in range(mask2.shape[1])]) #dist_to_centre[i-48]
                    mask2 = np.transpose(smoothed_signal, (1, 0, 2))
            
                np.save(f'./inputs/Face_files/{self.args.sign_request_id}.npy',mask2)
            

    
def array2pil(frame,enhance=False):
    if enhance:
        image_pil = Image.fromarray(frame.astype(np.uint8)).convert('RGB') 
        image_pil = ImageEnhance.Sharpness(image_pil).enhance(5)
    else:
        image_pil = Image.fromarray(frame.astype(np.uint8)).convert('RGB')      
    return image_pil
 
