import os
import glob
import cv2
import numpy as np
import torch
from torchvision.transforms import transforms
from natsort import natsorted
from models import resmasking_dropout1
from utils.datasets.fer2013dataset import EMOTION_DICT
from barez import show

transform = transforms.Compose(
    [
        transforms.ToPILImage(),
        transforms.ToTensor(),
    ]
)


def activations_mask(tensor):
    tensor = torch.squeeze(tensor, 0)
    tensor = torch.mean(tensor, 0)
    tensor = tensor.detach().cpu().numpy()
    tensor = np.maximum(tensor, 0)
    tensor = cv2.resize(tensor, (224, 224))
    tensor = tensor - np.min(tensor)
    tensor = tensor / np.max(tensor)

    heatmap = cv2.applyColorMap(np.uint8(255 * tensor), cv2.COLORMAP_JET)
    return heatmap


model = resmasking_dropout1(3, 7)
# state = torch.load('./saved/checkpoints/resmasking_dropout1_rot30_2019Nov17_14.33')
state = torch.load("./saved/checkpoints/Z_resmasking_dropout1_rot30_2019Nov30_13.32")
model.load_state_dict(state["net"])
model.cuda()
model.eval()

for image_path in natsorted(
    glob.glob("/home/z/research/bkemo/images/**/*.png", recursive=True)
):
    image_name = os.path.basename(image_path)
    print(image_name)
    # image_path = '/home/z/research/bkemo/images/disgust/0.0_dc10a3_1976_0.png'
    image = cv2.imread(image_path)
    image = cv2.resize(image, (224, 224))
    tensor = transform(image)
    tensor = torch.unsqueeze(tensor, 0)
    tensor = tensor.cuda()

    # output = model(tensor)

    x = model.conv1(tensor)  # 112
    x = model.bn1(x)
    x = model.relu(x)
    x = model.maxpool(x)  # 56

    x = model.layer1(x)  # 56
    m = model.mask1(x)
    x = x * (1 + m)

    x = model.layer2(x)  # 28
    m = model.mask2(x)
    x = x * (1 + m)

    x = model.layer3(x)  # 14
    heat_1 = activations_mask(x)
    m = model.mask3(x)
    x = x * (1 + m)
    # heat_2 = activations_mask(m)

    x = model.layer4(x)  # 7
    m = model.mask4(x)
    x = x * (1 + m)

    x = model.avgpool(x)
    x = torch.flatten(x, 1)

    output = model.fc(x)

    # print(np.sum(heat_1 - heat_2))

    # show(np.concatenate((image, heat_1, heat_2), axis=1))
    cv2.imwrite(
        "./masking_provements/{}".format(image_name),
        np.concatenate((image, heat_1), axis=1),
    )
    # np.concatenate((image, heat_1, heat_2), axis=1))

    # output = output.cpu().numpy()
    # print(EMOTION_DICT[torch.argmax(output, 1).item()])

[(33, 'models.resmasking_dropout1', 'resmasking_dropout1', False, 'from models import resmasking_dropout1\n'), (35, 'torch.load', 'torch.load', False, 'import torch\n'), (21, 'torch.squeeze', 'torch.squeeze', False, 'import torch\n'), (22, 'torch.mean', 'torch.mean', False, 'import torch\n'), (24, 'numpy.maximum', 'np.maximum', True, 'import numpy as np\n'), (25, 'cv2.resize', 'cv2.resize', False, 'import cv2\n'), (41, 'glob.glob', 'glob.glob', False, 'import glob\n'), (43, 'os.path.basename', 'os.path.basename', False, 'import os\n'), (46, 'cv2.imread', 'cv2.imread', False, 'import cv2\n'), (47, 'cv2.resize', 'cv2.resize', False, 'import cv2\n'), (49, 'torch.unsqueeze', 'torch.unsqueeze', False, 'import torch\n'), (78, 'torch.flatten', 'torch.flatten', False, 'import torch\n'), (14, 'torchvision.transforms.transforms.ToPILImage', 'transforms.ToPILImage', False, 'from torchvision.transforms import transforms\n'), (15, 'torchvision.transforms.transforms.ToTensor', 'transforms.ToTensor', False, 'from torchvision.transforms import transforms\n'), (26, 'numpy.min', 'np.min', True, 'import numpy as np\n'), (27, 'numpy.max', 'np.max', True, 'import numpy as np\n'), (29, 'numpy.uint8', 'np.uint8', True, 'import numpy as np\n'), (87, 'numpy.concatenate', 'np.concatenate', True, 'import numpy as np\n')]