# -*- coding: utf-8 -*-
"""
Created on Wed Apr  1 08:55:40 2020

@author: sugis
"""
import glob
import os
import os.path as osp
import random
import torch
import torch.nn as nn
import torch.optim as optim
from tqdm import tqdm
from torch.optim import lr_scheduler

from torchvision import models, transforms
from PIL import Image
import matplotlib.pyplot as plt
import torch.utils.data as data
import torchvision
import numpy as np
import cv2
from torch.autograd import Function
Image.LOAD_TRUNCATED_IMAGES = True

import math as math

# In[ ]:

# 学習済みのVGG-16モデルをロード
# VGG-16モデルのインスタンスを生成
use_pretrained = True  # 学習済みのパラメータを使用
net = models.vgg16(pretrained=use_pretrained)

# VGG16の最後の出力層の出力ユニットをアリとハチの2つに付け替える
#net.classifier[6] = nn.Linear(in_features=4096, out_features=3)
net.classifier[6] = nn.Linear(in_features=4096, out_features=2)

# PyTorchのネットワークパラメータのロード
load_path = './weights_fine_tuning.pth'


if torch.cuda.is_available():
    load_weights = torch.load(load_path)
    net.load_state_dict(load_weights)
    
else:
    # GPU上で保存された重みをCPU上でロードする場合
    load_weights = torch.load(load_path, map_location={'cuda:0': 'cpu'})
    net.load_state_dict(load_weights)


# In[ ]: gradcam

class FeatureExtractor():
    """ Class for extracting activations and 
    registering gradients from targetted intermediate layers """

    def __init__(self, model, target_layers):
        self.model = model
        self.target_layers = target_layers
        self.gradients = []

    def save_gradient(self, grad):
        self.gradients.append(grad)

    def __call__(self, x):
        outputs = []
        self.gradients = []
        for name, module in self.model._modules.items():
            x = module(x)
            if name in self.target_layers:
                x.register_hook(self.save_gradient)
                outputs += [x]
        return outputs, x


class ModelOutputs():
    """ Class for making a forward pass, and getting:
    1. The network output.
    2. Activations from intermeddiate targetted layers.
    3. Gradients from intermeddiate targetted layers. """

    def __init__(self, model, target_layers):
        self.model = model
        self.feature_extractor = FeatureExtractor(self.model.features, target_layers)

    def get_gradients(self):
        return self.feature_extractor.gradients

    def __call__(self, x):
        target_activations, output = self.feature_extractor(x)
        output = output.view(output.size(0), -1)
        output = self.model.classifier(output)
        return target_activations, output

def preprocess_image(img):
    means = [0.485, 0.456, 0.406]
    stds = [0.229, 0.224, 0.225]

    preprocessed_img = img.copy()[:, :, ::-1]
    for i in range(3):
        preprocessed_img[:, :, i] = preprocessed_img[:, :, i] - means[i]
        preprocessed_img[:, :, i] = preprocessed_img[:, :, i] / stds[i]
    preprocessed_img = \
        np.ascontiguousarray(np.transpose(preprocessed_img, (2, 0, 1)))
    preprocessed_img = torch.from_numpy(preprocessed_img)
    preprocessed_img.unsqueeze_(0)
    input = preprocessed_img.requires_grad_(True)
    return input

def show_cam_on_image(img, mask):
    heatmap = cv2.applyColorMap(np.uint8(255 * mask), cv2.COLORMAP_JET)
    heatmap = np.float32(heatmap) / 255
    cam = heatmap + np.float32(img)
    cam = cam / np.max(cam)
    plt.imshow(np.uint8(255 * cam))
    plt.show()
    cv2.imwrite("cam.jpg", np.uint8(255 * cam))
   
class GradCam:
    def __init__(self, model, target_layer_names, use_cuda):
        self.model = model
        self.model.eval()
        self.cuda = use_cuda
        if self.cuda:
            self.model = model.cuda()

        self.extractor = ModelOutputs(self.model, target_layer_names)

    def forward(self, input):
        return self.model(input)

    def __call__(self, input, index=None):
        if self.cuda:
            features, output = self.extractor(input.cuda())
        else:
            features, output = self.extractor(input)

        if index == None:
            index = np.argmax(output.cpu().data.numpy())

        one_hot = np.zeros((1, output.size()[-1]), dtype=np.float32)
        one_hot[0][index] = 1
        one_hot = torch.from_numpy(one_hot).requires_grad_(True)
        if self.cuda:
            one_hot = torch.sum(one_hot.cuda() * output)
        else:
            one_hot = torch.sum(one_hot * output)

        self.model.features.zero_grad()
        self.model.classifier.zero_grad()
        one_hot.backward(retain_graph=True)

        grads_val = self.extractor.get_gradients()[-1].cpu().data.numpy()

        target = features[-1]
        target = target.cpu().data.numpy()[0, :]

        weights = np.mean(grads_val, axis=(2, 3))[0, :]
        cam = np.zeros(target.shape[1:], dtype=np.float32)

        for i, w in enumerate(weights):
            cam += w * target[i, :, :]

        cam = np.maximum(cam, 0)
        cam = cv2.resize(cam, (224, 224))
        cam = cam - np.min(cam)
        cam = cam / np.max(cam)
        return cam

class GuidedBackpropReLU(Function):

    @staticmethod
    def forward(self, input):
        positive_mask = (input > 0).type_as(input)
        output = torch.addcmul(torch.zeros(input.size()).type_as(input), input, positive_mask)
        self.save_for_backward(input, output)
        return output

    @staticmethod
    def backward(self, grad_output):
        input, output = self.saved_tensors
        grad_input = None

        positive_mask_1 = (input > 0).type_as(grad_output)
        positive_mask_2 = (grad_output > 0).type_as(grad_output)
        grad_input = torch.addcmul(torch.zeros(input.size()).type_as(input),
                                   torch.addcmul(torch.zeros(input.size()).type_as(input), grad_output,
                                                 positive_mask_1), positive_mask_2)

        return grad_input


class GuidedBackpropReLUModel:
    def __init__(self, model, use_cuda):
        self.model = model
        self.model.eval()
        self.cuda = use_cuda
        if self.cuda:
            self.model = model.cuda()

        # replace ReLU with GuidedBackpropReLU
        for idx, module in self.model.features._modules.items():
            if module.__class__.__name__ == 'ReLU':
                self.model.features._modules[idx] = GuidedBackpropReLU.apply

    def forward(self, input):
        return self.model(input)

    def __call__(self, input, index=None):
        if self.cuda:
            output = self.forward(input.cuda())
        else:
            output = self.forward(input)

        if index == None:
            index = np.argmax(output.cpu().data.numpy())

        one_hot = np.zeros((1, output.size()[-1]), dtype=np.float32)
        one_hot[0][index] = 1
        one_hot = torch.from_numpy(one_hot).requires_grad_(True)
        if self.cuda:
            one_hot = torch.sum(one_hot.cuda() * output)
        else:
            one_hot = torch.sum(one_hot * output)

        # self.model.features.zero_grad()
        # self.model.classifier.zero_grad()
        one_hot.backward(retain_graph=True)

        output = input.grad.cpu().data.numpy()
        output = output[0, :, :, :]

        return output

def deprocess_image(img):
    """ see https://github.com/jacobgil/keras-grad-cam/blob/master/grad-cam.py#L65 """
    img = img - np.mean(img)
    img = img / (np.std(img) + 1e-5)
    img = img * 0.1
    img = img + 0.5
    img = np.clip(img, 0, 1)
    return np.uint8(img*255)


def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--use-cuda', action='store_true', default=False,
                        help='Use NVIDIA GPU acceleration')
    parser.add_argument('--image-path', type=str, default='./examples/both.png',
                        help='Input image path')
    args = parser.parse_args()
    args.use_cuda = args.use_cuda and torch.cuda.is_available()
    if args.use_cuda:
        print("Using GPU for acceleration")
    else:
        print("Using CPU for computation")

    return args
#args.use_cuda
    

# In[ ]: predict

 # In[ ]:
class ImageTransform():
    
    def __init__(self,resize,mean,std):
        self.data_transform={
            'train':transforms.Compose([
                transforms.RandomResizedCrop(
                    resize,scale=(0.5,1.0)),
                transforms.RandomHorizontalFlip(),
                transforms.ToTensor(),
                transforms.Normalize(mean,std)
            ]),
            'val':transforms.Compose([
                transforms.Resize(resize),
                transforms.CenterCrop(resize),  # 画像中央をresize×resizeで切り取り
                transforms.ToTensor(),
                transforms.Normalize(mean,std)
            ])
    }

    def __call__(self, img,phase='train'):
        
        return self.data_transform[phase](img)

size = 224
mean = (0.485, 0.456, 0.406)
std = (0.229, 0.224, 0.225)

#transform_test=ImageTransform(size, mean, std)
#img_transformed_test = transform_test(img, phase='train')
# 乱数のシードを設定
torch.manual_seed(1234)
np.random.seed(1234)
random.seed(1234)

torch.backends.cudnn.deterministic =True
torch.backends.cudnn.benchmark=False

# 入力画像の前処理をするクラス
# 訓練時と推論時で処理が異なる


# In[ ]:
def make_datapath_list(phase='train'):
    
    rootpath="./bridgedata/"
    target_path=osp.join(rootpath+phase+'/**/*.jpg')
    
    path_list=[]
    
    for path in glob.glob(target_path):
        path_list.append(path)
    
    return path_list

# 実行
train_list = make_datapath_list(phase="train")
val_list = make_datapath_list(phase="val")
test_list = make_datapath_list(phase="test")

print(len(train_list))

# In[ ]:

class MyDatasets(data.Dataset):
    
    def __init__(self, file_list,transform = None,phase='train'):
        self.file_list=file_list
        self.transform = transform
        self.phase = phase

    def __len__(self):
        return len(self.file_list)

    def __getitem__(self, index):
        
        img_path=self.file_list[index]
        img=Image.open(img_path)
        
        img_transformed=self.transform(
                img,self.phase)
        
        fp=osp.basename(osp.dirname(img_path))
        
        if self.phase=='train':
            label=fp
        elif self.phase=='val':
            label=fp
            
        if label=='girder':
            label=0
        elif label=='nogirder':
            label=1

        
        return img_transformed,label


train_test = MyDatasets(
    file_list=test_list, transform=ImageTransform(size, mean, std), phase='val')

size = 224

imsize = size
 
loader = transforms.Compose([
    transforms.Resize(imsize),  # scale imported image
    transforms.ToTensor()])  # transform it into a torch tensor
 
def image_loader(image_name):
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    image = Image.open(image_name)
    # fake batch dimension required to fit network's input dimensions
    image = loader(image).unsqueeze(0)
    return image.to(device, torch.float)
 
classes=('girder','nogirder')

#def show(img):
#    npimg = img.numpy()
#    plt.imshow(np.transpose(npimg, (1,2,0)), interpolation='nearest')

#show(torchvision.utils.make_grid(train_dataset1.__getitem__(index)[0], padding=1))
num=0
row=3
col=3


#path1="./testjpg"
#image_path1 = os.listdir(path1)


grad_cam = GradCam(model=net,target_layer_names=["26"], use_cuda=True)

#for i in test_list:
for i in range(9):

    num+=1
    index = np.random.randint(len(test_list))
    
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    
    img_g=Image.open(test_list[index])
    img_g2 =img_g.resize((224,224),resample=0)
    img_g2 = np.float32(img_g2)/255
    input = preprocess_image(img_g2)

    target_index = None
    mask = grad_cam(input, target_index)
    
    plt.figure(1)
    plt.subplot(row,col,num)
    show_cam_on_image(img_g2, mask)
    
    transform_test=ImageTransform(size, mean, std)
    img_transformed_test = transform_test(img_g, phase='val')
    image1=img_transformed_test
    npimg1 = image1.numpy()
    
    image1=image1.to(device, torch.float)
    image1 = image1.unsqueeze(0)
    
    output1=net(image1)
    _,predicted = torch.max(output1, 1)
    preds = torch.max(output1, 1)
    kakuritu=nn.functional.softmax(output1)
    kakuritum=np.str(np.round(kakuritu[0][preds.indices].item(),2))
    print(kakuritu[0][preds.indices])
    print(classes[preds.indices])
    #print(classes[train_test.__getitem__(index)[1]])
    
    plt.figure(2)
    plt.subplot(row,col,num)
    plt.imshow(np.transpose(npimg1, (1,2,0)), interpolation='nearest')
    plt.axis('off')
    plt.title(classes[preds.indices]+kakuritum)
    plt.show()