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, 8, 33, 33), 'models.resmasking_dropout1', 'resmasking_dropout1', ({(33, 28, 33, 29): '3', (33, 31, 33, 32): '7'}, {}), '(3, 7)', False, 'from models import resmasking_dropout1\n'), ((35, 8, 35, 85), 'torch.load', 'torch.load', ({(35, 19, 35, 84): '"""./saved/checkpoints/Z_resmasking_dropout1_rot30_2019Nov30_13.32"""'}, {}), "('./saved/checkpoints/Z_resmasking_dropout1_rot30_2019Nov30_13.32')", False, 'import torch\n'), ((21, 13, 21, 37), 'torch.squeeze', 'torch.squeeze', ({(21, 27, 21, 33): 'tensor', (21, 35, 21, 36): '0'}, {}), '(tensor, 0)', False, 'import torch\n'), ((22, 13, 22, 34), 'torch.mean', 'torch.mean', ({(22, 24, 22, 30): 'tensor', (22, 32, 22, 33): '0'}, {}), '(tensor, 0)', False, 'import torch\n'), ((24, 13, 24, 34), 'numpy.maximum', 'np.maximum', ({(24, 24, 24, 30): 'tensor', (24, 32, 24, 33): '0'}, {}), '(tensor, 0)', True, 'import numpy as np\n'), ((25, 13, 25, 43), 'cv2.resize', 'cv2.resize', ({(25, 24, 25, 30): 'tensor', (25, 32, 25, 42): '(224, 224)'}, {}), '(tensor, (224, 224))', False, 'import cv2\n'), ((41, 4, 41, 71), 'glob.glob', 'glob.glob', (), '', False, 'import glob\n'), ((43, 17, 43, 45), 'os.path.basename', 'os.path.basename', ({(43, 34, 43, 44): 'image_path'}, {}), '(image_path)', False, 'import os\n'), ((46, 12, 46, 34), 'cv2.imread', 'cv2.imread', ({(46, 23, 46, 33): 'image_path'}, {}), '(image_path)', False, 'import cv2\n'), ((47, 12, 47, 41), 'cv2.resize', 'cv2.resize', ({(47, 23, 47, 28): 'image', (47, 30, 47, 40): '(224, 224)'}, {}), '(image, (224, 224))', False, 'import cv2\n'), ((49, 13, 49, 39), 'torch.unsqueeze', 'torch.unsqueeze', ({(49, 29, 49, 35): 'tensor', (49, 37, 49, 38): '0'}, {}), '(tensor, 0)', False, 'import torch\n'), ((78, 8, 78, 27), 'torch.flatten', 'torch.flatten', ({(78, 22, 78, 23): 'x', (78, 25, 78, 26): '1'}, {}), '(x, 1)', False, 'import torch\n'), ((14, 8, 14, 31), 'torchvision.transforms.transforms.ToPILImage', 'transforms.ToPILImage', ({}, {}), '()', False, 'from torchvision.transforms import transforms\n'), ((15, 8, 15, 29), 'torchvision.transforms.transforms.ToTensor', 'transforms.ToTensor', ({}, {}), '()', False, 'from torchvision.transforms import transforms\n'), ((26, 22, 26, 36), 'numpy.min', 'np.min', ({(26, 29, 26, 35): 'tensor'}, {}), '(tensor)', True, 'import numpy as np\n'), ((27, 22, 27, 36), 'numpy.max', 'np.max', ({(27, 29, 27, 35): 'tensor'}, {}), '(tensor)', True, 'import numpy as np\n'), ((29, 32, 29, 54), 'numpy.uint8', 'np.uint8', ({(29, 41, 29, 53): '255 * tensor'}, {}), '(255 * tensor)', True, 'import numpy as np\n'), ((87, 8, 87, 47), 'numpy.concatenate', 'np.concatenate', (), '', True, 'import numpy as np\n')]