# PyTorch implementation of Darknet
# This is a custom, hard-coded version of darknet with 
# YOLOv3 implementation for openimages database. This 
# was written to test viability of implementing YOLO 
# for face detection followed by emotion / sentiment
# analysis.
#
# Configuration, weights and data are hardcoded.
# Additional options include, ability to create
# subset of data with faces exracted for labelling.
#
# Author    : Saikiran Tharimena
# Co-Authors: Kjetil Marinius Sjulsen, Juan Carlos Calvet Lopez
# Project   : Emotion / Sentiment Detection from news images
# Date      : 12 September 2022
# Version   : v0.1
#
# (C) Schibsted ASA

# Libraries
import os
import cv2
import torch
import numpy as np
from utils import *
from darknet import Darknet
from torch.autograd import Variable
from torch.cuda import is_available as check_cuda

# Parameters
batch_size = 1
confidence = 0.25
nms_thresh = 0.30
run_cuda = False

# CFG Files
cwd = os.path.dirname(__file__)
cfg     = cwd + '/cfg/yolov3-openimages.cfg'
data    = cwd + '/cfg/openimages.data'
clsnames= cwd + '/cfg/openimages.names'
weights = cwd + '/cfg/yolov3-openimages.weights'

# Load classes
num_classes = 601
classes = load_classes(clsnames)

# Set up the neural network
print('Load Network')
model = Darknet(cfg)

print('Load Weights')
model.load_weights(weights)

print('Successfully loaded Network')

# Check CUDA
if run_cuda:
    CUDA = check_cuda()
else:
    CUDA = False

# Input dimension
inp_dim = int(model.net_info["height"])

# put the model on GPU
if CUDA:
    model.cuda()

# Set the model in evaluation mode
model.eval()

# face detector
def detect_face(image):
    # Just lazy to update this
    imlist = [image]
    
    loaded_ims = [cv2.imread(x) for x in imlist]

    im_batches = list(map(prep_image, loaded_ims, [inp_dim for x in range(len(imlist))]))
    im_dim_list = [(x.shape[1], x.shape[0]) for x in loaded_ims]
    im_dim_list = torch.FloatTensor(im_dim_list).repeat(1,2)

    leftover = 0
    if (len(im_dim_list) % batch_size):
        leftover = 1

    if batch_size != 1:
        num_batches = len(imlist) // batch_size + leftover            
        im_batches = [torch.cat((im_batches[i*batch_size : min((i +  1)*batch_size,
                            len(im_batches))]))  for i in range(num_batches)]  

    write = 0
    if CUDA:
        im_dim_list = im_dim_list.cuda()
    
    for i, batch in enumerate(im_batches):
        # load the image 
        
        if CUDA:
            batch = batch.cuda()
        with torch.no_grad():
            prediction = model(Variable(batch), CUDA)

        prediction = write_results(prediction, confidence, num_classes, nms_conf = nms_thresh)

        if type(prediction) == int:

            for im_num, image in enumerate(imlist[i*batch_size: min((i +  1)*batch_size, len(imlist))]):
                im_id = i*batch_size + im_num
                
            continue

        prediction[:,0] += i*batch_size    # transform the atribute from index in batch to index in imlist 

        if not write: # If we have't initialised output
            output = prediction  
            write = 1
        else:
            output = torch.cat((output, prediction))

        for im_num, image in enumerate(imlist[i*batch_size: min((i +  1)*batch_size, len(imlist))]):
            im_id = i * batch_size + im_num
            objs = [classes[int(x[-1])] for x in output if int(x[0]) == im_id]

        if CUDA:
            torch.cuda.synchronize()
       
    try:
        output
    except NameError:
        return None

    im_dim_list = torch.index_select(im_dim_list, 0, output[:,0].long())

    scaling_factor = torch.min(608/im_dim_list,1)[0].view(-1,1)

    output[:, [1,3]] -= (inp_dim - scaling_factor*im_dim_list[:,0].view(-1,1))/2
    output[:, [2,4]] -= (inp_dim - scaling_factor*im_dim_list[:,1].view(-1,1))/2

    output[:, 1:5] /= scaling_factor

    for i in range(output.shape[0]):
        output[i, [1,3]] = torch.clamp(output[i, [1,3]], 0.0, im_dim_list[i,0])
        output[i, [2,4]] = torch.clamp(output[i, [2,4]], 0.0, im_dim_list[i,1])
    
    def get_detections(x, results):
        c1 = [int(y) for y in x[1:3]]
        c2 = [int(y) for y in x[3:5]]
        
        det_class = int(x[-1])
        label = "{0}".format(classes[det_class])
        
        return (label, tuple(c1 + c2))
    
    detections = list(map(lambda x: get_detections(x, loaded_ims), output))

    if CUDA:
        torch.cuda.empty_cache()
    
    return loaded_ims[0], detections