from .face_parsing import BiSeNet
import numpy as np
from .metrics import LPIPS, MS_SSIM, IdScore, ClipHair
import torch.nn as nn
import torch
from torchvision import transforms

class FaceSegmentation(nn.Module):
    def __init__(self, n_classes=19, device='cuda', save_pth='./pretrained_models/79999_iter.pth'):
        super(FaceSegmentation, self).__init__()
        self.net = BiSeNet(n_classes=n_classes).to(device)
        self.net.load_state_dict(torch.load(save_pth))
        self.net.eval()
        self.transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
        ])
        self.device=device



    def get_facemask(self, parsing_anno):
        """
        Returns a binary image of the face.
        """
        # face_attr = {1: 'skin', 2: 'l_brow', 3: 'r_brow', 4: 'l_eye', 5: 'r_eye', 6: 'eye_glass', 7: 'l_ear', 8: 'r_ear', 10: 'nose', 11: 'mouth', 12: 'u_lip', 13: 'l_lip', 14: 'neck'}
        face_attr = torch.tensor([1,2,3,4,5,6,7,8,10,11,12,13,14],device=self.device)
        face_mask = torch.isin(parsing_anno, face_attr)
        return(face_mask.int())


    def get_hairmask(self, parsing_anno):
        """
        Returns a binary image of the hair.
        """ 
        hair_mask = parsing_anno == 17
        return(hair_mask.int())

    def forward(self, img):
        """
        Returns a binary image of the face and hair.
        """
        img = self.transform(img).to(self.device)
        parsing_anno = self.net(img.unsqueeze(0))[0].squeeze(0).argmax(0)
        face_mask = self.get_facemask(parsing_anno).to(self.device)
        hair_mask = self.get_hairmask(parsing_anno).to(self.device)
        return img, face_mask, hair_mask


class FaceMetric(nn.Module):
    def __init__(self, metric_type, eval_face=True, eval_hair=True, device='cuda', seg_save_pth='./pretrained_models/79999_iter.pth'):
        super(FaceMetric, self).__init__()
        if metric_type == 'ms-ssim':
            self.metric =  MS_SSIM()
            self.eval_hair= eval_hair
            self.eval_face= eval_face
        elif metric_type == 'lpips':
            self.metric = LPIPS(device=device)
            self.eval_hair= eval_hair
            self.eval_face= eval_face
        elif metric_type == 'id':
            self.metric = IdScore(device=device)
            self.eval_hair = False
            self.eval_face = eval_face
        elif metric_type == 'cliphair':
            self.metric = ClipHair(device=device)
            self.eval_face = False
            self.eval_hair = eval_hair
        else:
            raise NotImplementedError
        self.parser = FaceSegmentation(device=device, save_pth=seg_save_pth)
        self.device=device


    
    
    def forward(self, x, y):
        face_score, hair_score = None, None
        x_tensor, x_face_seg, x_hair_seg = self.parser(x)
        y_tensor, y_face_seg, y_hair_seg = self.parser(y)
        if self.eval_hair == True:

 
            ## Get union of two hair masks
            #hair_mask = (x_hair_seg + y_hair_seg) > 0

            x_hair = x_tensor * x_hair_seg
            y_hair = y_tensor * y_hair_seg

            hair_score = self.metric(x_hair, y_hair).item()
        if self.eval_face == True:

            ## Get intersection of two face masks
            face_mask = (x_face_seg + y_face_seg) > 1

            x_face = x_tensor * face_mask
            y_face = y_tensor * face_mask

            face_score = self.metric(x_face, y_face).item()

        return face_score, hair_score