import torch
import torch.nn as nn
import math
import os
import cv2
import numpy as np
from utils.DATAIMG import DATAIMG
from utils import landmark_processor
import glob
from mathlib.umeyama import umeyama

def op_name(op_name, m):
    m.op_name = op_name
    return m

class Flatten(nn.Module):
    def __init__(self, axis=1):
        super(Flatten, self).__init__()
        self.axis = axis

    def forward(self, x):
        assert self.axis == 1
        x = x.reshape(x.shape[0], -1)
        return x

def flatten(name, axis=1):
    return op_name(name, Flatten(axis))

def conv_relu(name, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1):
    return nn.Sequential(
        op_name(name, nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, True)),
        op_name(name.replace('conv', 'relu'), nn.PReLU(out_channels)),
    )

class BasicBlock(nn.Module):
    def __init__(self, i, j, inplanes):
        super(BasicBlock, self).__init__()

        self.conv1 = conv_relu('conv{}_{}'.format(i, 2 * j), inplanes, inplanes, kernel_size=3, stride=1, padding=1)
        self.conv2 = conv_relu('conv{}_{}'.format(i, 2 * j + 1), inplanes, inplanes, kernel_size=3, stride=1, padding=1)

    def forward(self, x):
        residule = x
        out = self.conv1(x)
        out = self.conv2(out)
        ret = out + residule
        return ret

class ResnetBlock(nn.Module):
    def __init__(self, i, inplanes, outplanes, stride=2, n_blocks=1):
        super(ResnetBlock, self).__init__()

        self.conv = [conv_relu('conv{}_{}'.format(i, 1), inplanes, outplanes, kernel_size=3, stride=stride, padding=1)]
        for m in range(n_blocks):
            self.conv.append(BasicBlock(i, m + 1, outplanes))
        self.conv = nn.Sequential(*self.conv)

    def forward(self, x):
        ret = self.conv(x)
        return ret

class FaceRecognition(nn.Module):
    def __init__(self):
        super(FaceRecognition, self).__init__()
        op_list = []
        ch_num = [3, 32, 64, 128, 128]
        strides = [2, 2, 2, 2]
        n_blocks = [1, 2, 4, 1]
        op_list = []
        op_list += [ResnetBlock(i + 1, ch_num[i], ch_num[i + 1], strides[i], n_blocks[i]) for i in range(len(ch_num) - 1)]
        op_list += [flatten('flatten', 1)]
        op_list += [op_name('fc5', nn.Linear(4608, 256))]
        self.features = nn.Sequential(*op_list)

        model_path, _ = os.path.split(os.path.realpath(__file__))
        weights = torch.load(os.path.join(model_path, 'Face_ResNet.pth'),
                             map_location=lambda storage, loc: storage)
        self.load_state_dict(weights)
        self.eval()

    def forward(self, x):
        features = self.features(x)
        return features

class MomocvFaceRecognition(object):
    def __init__(self, gpu_id=None):
        self.gpu_id = gpu_id
        self.device = torch.device('cuda:{}'.format(gpu_id) if gpu_id is not None else 'cpu')
        self.face_recognition_net = FaceRecognition()
        # self.model_path, _ = os.path.split(os.path.realpath(__file__))
        # weights = torch.load(os.path.join(self.model_path, 'Face_ResNet.pth'),
        #                      map_location=lambda storage, loc: storage)
        # self.face_recognition_net.load_state_dict(weights)
        self.face_recognition_net.to(self.device)
        # self.face_recognition_net.eval()

    def forward(self, images, landmarks):
        dst_size = 90

        with torch.no_grad():
            input_numpy = np.zeros((len(images), 3, dst_size, dst_size), dtype=np.float32)
            assert len(images) == len(landmarks)
            for ix, img in enumerate(images):
                landmark = landmarks[ix]

                mat = landmark_processor.get_transform_mat_for_face_recognition(landmark, dst_size)[0:2]

                tmp = cv2.warpAffine(img, mat, (dst_size, dst_size))

                input_numpy[ix, :, :, :] = (tmp.transpose((2, 0, 1)).astype(np.float32) - 128) / 128

                # cv2.imshow('tmp', tmp)
                # cv2.waitKey()

            in_tensor = torch.from_numpy(input_numpy)
            in_tensor = in_tensor.to(self.device)
            features = self.face_recognition_net(in_tensor)
            features = features.cpu().numpy()
            norm_factor = np.sqrt(np.sum(features ** 2, axis=1))
            norm_factor = norm_factor.reshape(-1, 1)
            features /= norm_factor
            return features

if __name__ == '__main__':
    all_jpegs = glob.glob(r'E:\deepfacelab_data\expression_dst\7201806132018061208311920180612083119\*.jpg')
    for s_filename_path in all_jpegs:
        img = cv2.imread(s_filename_path)

        # cv2.imshow('img', img)
        # cv2.waitKey()

        dflpng = DATAIMG(str(s_filename_path), print_on_no_embedded_data=True)
        if dflpng is None:
            print('ERROR')

        landmarks = dflpng.get_landmarks()

        mmcv = MomocvFaceRecognition()
        features = mmcv.forward([img, img], [landmarks, landmarks])
        # fullyconnected1 = fullyconnected1[0]
        # fullyconnected1 = (np.reshape(fullyconnected1, (2, 87)).transpose((1, 0))).astype(np.int32)
        # occlusion_probe = occlusion_probe[0]
        # for ix, pt in enumerate(fullyconnected1):
        #     cv2.circle(img, (int(pt[0]), int(pt[1])), 1, (0, 255, 0) if occlusion_probe[ix] > 0.1 else (0, 0, 255), 2)
        #     # cv2.putText(img, str(ix), (int(pt[0]), int(pt[1])), cv2.FONT_HERSHEY_SIMPLEX, 0.3, (0, 255, 0), 1)
        # cv2.imshow('img', img)
        # cv2.waitKey()