import cv2
from utils import read_image,get_valid_augs
import torch
import torch.nn.functional as F
from model import ResNet50Model
import numpy as np

CKPT_EMOTION = 'emotion_resnet_model'
CKPT_CELEBRITY = 'celebrity_resnet_model'
FaceInverseTargetMapper = {0: 'Tom Hanks',1: 'Sandra Bullock',2: 'Natalie Portman',3: 'Scarlett Johansson',4: 'Robert Downey Jr',5: 'Nicole Kidman',6: 'Brad Pitt',
 7: 'Hugh Jackman',8: 'Tom Cruise',9: 'Leonardo DiCaprio',10: 'Megan Fox',11: 'Johnny Depp',12: 'Will Smith',13: 'Denzel Washington',14: 'Jennifer Lawrence',15: 'Kate Winslet',16: 'Angelina Jolie'}
EmotionMapper = {0: 'sadness',1: 'contempt',2: 'happiness',3: 'surprise',
 4: 'fear',5: 'anger',6: 'disgust',7: 'neutrality'}
def predict_one_image(path) :
    image = read_image(path)
    image = get_valid_augs()(image=image)['image']
    image = torch.tensor(image,dtype=torch.float)
    image = image.reshape((1,3,224,224))
    emotion_model = ResNet50Model('emotion')
    #loading ckpt
    emotion_model.load_state_dict(torch.load(CKPT_EMOTION,map_location=torch.device('cpu')))

    celebrity_model = ResNet50Model('celebrity')
    #loading ckpt
    celebrity_model.load_state_dict(torch.load(CKPT_CELEBRITY,map_location=torch.device('cpu')))
    
    
    with torch.no_grad() :
        #emotion
        outputs = emotion_model(image)
        outputs = torch.nn.functional.softmax(outputs).cpu().detach().numpy()
        print(outputs.shape)
        emotion = np.argmax(outputs,axis=1)[0]
        emotion_proba = np.max(outputs,axis=1)[0]
        print(emotion_proba)
        #celebrity
        outputs = celebrity_model(image)
        outputs = torch.nn.functional.softmax(outputs).cpu().detach().numpy()
        print(outputs.shape)
        celebrity = np.argmax(outputs,axis=1)[0]
        celebrity_proba = np.max(outputs,axis=1)[0]
        print(celebrity_proba)
    
    if celebrity_proba<0.45 : 
        return f"Unkonwn Person Detected with emotion {EmotionMapper[emotion]} "
    return f"Detected {FaceInverseTargetMapper[celebrity]} with emotion {EmotionMapper[emotion]} "