"""
<h1>MNISTSuperpixels Interpretation</h1>
"""
!pip install pytorch-lightning
!pip install torch-scatter -f https://pytorch-geometric.com/whl/torch-1.7.0+cu102.html
!pip install torch-sparse -f https://pytorch-geometric.com/whl/torch-1.7.0+cu102.html
!pip install torch-cluster -f https://pytorch-geometric.com/whl/torch-1.7.0+cu102.html
!pip install torch-spline-conv -f https://pytorch-geometric.com/whl/torch-1.7.0+cu102.html
!pip install torch-geometric
!pip install captum
!pip install torch_summary
import pytorch_lightning as pl
import torchvision.transforms
import pandas as pd
import torch
from sklearn.model_selection import train_test_split
from torch.utils.data import DataLoader
from torch.nn import Linear, ReLU
from torch.optim import Adam
from torch.optim.lr_scheduler import StepLR
import torch.nn.functional as F
from torch_geometric.utils import normalized_cut
from torch_geometric.nn import SplineConv
from torch_geometric.nn import (SplineConv, graclus, max_pool, max_pool_x, global_max_pool,
                                global_mean_pool, GCNConv)
import torch.nn.functional as F
from torch_geometric.data import Batch,Data
from torch_geometric.datasets import MNISTSuperpixels
import torch_geometric.transforms as T
from torch_geometric.data import DataLoader

def normalized_cut_2d(edge_index, pos):
    row, col = edge_index
    edge_attr = torch.norm(pos[row] - pos[col], p=2, dim=1)
    return normalized_cut(edge_index, edge_attr, num_nodes=pos.size(0))


class MNISTSplineNet(pl.LightningModule):
    def __init__(self,in_features,out_features):
        super(MNISTSplineNet, self).__init__()
        self.conv1 = SplineConv(in_features, 32, dim=2, kernel_size=5)
        self.conv2 = SplineConv(32, 64, dim=2, kernel_size=5)
        self.relu1 = ReLU()
        self.relu2 = ReLU()
        self.relu3 = ReLU()
        
        
        self.fc1 = torch.nn.Linear(64, 128)
        self.fc2 = torch.nn.Linear(128, out_features)
        
    def forward(self,data):
        print(data)
        data.x = self.relu1(self.conv1(data.x, data.edge_index, data.edge_attr))
        print(data)
        weight = normalized_cut_2d(data.edge_index, data.pos)
        cluster = graclus(data.edge_index, weight, data.x.size(0))
#         data.edge_attr = None
        print(data)
        data = max_pool(cluster, data, transform=self.transform)
        data.x = F.elu(self.conv2(data.x, data.edge_index, data.edge_attr))
        print(data)
        weight = normalized_cut_2d(data.edge_index, data.pos)
        cluster = graclus(data.edge_index, weight, data.x.size(0))
        x, batch = max_pool_x(cluster, data.x, data.batch)
        print(data)
        x = global_mean_pool(x, batch)
        x = F.elu(self.fc1(x))
        x = F.dropout(x, training=self.training)
        return F.log_softmax(self.fc2(x), dim=1)
    
    def prepare_data(self):
        self.transform = T.Cartesian(cat=False)

        MNISTSuperpixels.url = ('https://graphics.cs.tu-dortmund.de/fileadmin/ls7-www/misc/cvpr/'
            'mnist_superpixels.tar.gz')

        train_dataset = MNISTSuperpixels('mnist',True, transform=self.transform)
        test_dataset = MNISTSuperpixels('mnist',False, transform=self.transform)
        train_set, val_set = torch.utils.data.random_split(train_dataset, [50000, 10000])
        self.train_dataset = train_set
        self.val_dataset = val_set
        self.test_dataset = test_dataset
        
            
    def train_dataloader(self):
        return DataLoader(self.train_dataset,batch_size=32, shuffle=True, num_workers=16)
    
    def val_dataloader(self):
        return DataLoader(self.val_dataset,batch_size=32, shuffle=False, num_workers=16)
    
    def test_dataloader(self):
        return DataLoader(self.test_dataset,batch_size=32, shuffle=False, num_workers=16)
    
    def configure_optimizers(self):
        optimizer = Adam(self.parameters(), lr=0.01)
        return optimizer
    
    def training_step(self,batch,batch_idx):
        return F.nll_loss(self.forward(batch), batch.y)
    
    def validation_step(self, batch,batch_idx):
        y_hat = self.forward(batch)
        loss = F.nll_loss(y_hat,batch.y)
        pred = y_hat.max(1)[1]
        return {'val_loss':loss}
#         this.correct += pred.eq(data.y).sum().item()
#         return {'val_loss':loss, 'correct':correct}
    
    def validation_epoch_end(self,outputs):
        avg_loss = torch.stack([x['val_loss'] for x in outputs]).mean()
        tensorboard_logs = {'val_loss':avg_loss}
        return {'avg_val_loss' : avg_loss, 'log':tensorboard_logs}
        


model = MNISTSplineNet(in_features=1, out_features=10)
model.prepare_data()
train_dataset  = model.train_dataset

print("Number of graphs: %i" % (len(train_dataset)))

graph = train_dataset[0]
print("Number of nodes: %i" % (graph.num_nodes))
print("Number of features in node: %i" % (graph.num_node_features))

print("There are 60000 homogenous graphs, each node respresents digit and has features represented as extracted features.")
print()

print("Number of edge feature: %i" %(graph.num_edge_features))
print("Edges number of edges: %i" %(graph.edge_index.size()[1]))
print("Example edge feature (pos): " + str(graph.edge_index[:,1000])) 


print("Edges have 2 features that are the indexes of nodes which are interconnected (from_node_idx,to_node_idx).")
print()

classes = [data.y.item() for data in train_dataset]
print("Example y classes: " + str(classes[0:10]))
print("Minimal class index: " + str(min(classes)))
print("Maximal class index: " + str(max(classes)))
output_classes = max(classes)-min(classes)+1
print("Number of output classes: " + str(output_classes))
print("Output variable represents 10 distinct product categories.")
do_train = False


if do_train:
#     trainer = pl.Trainer(max_epochs=8)
    trainer = pl.Trainer(gpus=-1, max_epochs=8)
    trainer.fit(model)
else:
    #Load weights
    pass
#Node Feature explainations

import torch_geometric
from captum.attr import Saliency,IntegratedGradients,GradientShap, FeatureAblation, DeepLift

def attributation_for_feature_forward(x, data):
    data.x=x
    return model(data)


data_sample = next(iter(DataLoader(model.train_dataset,batch_size=1, shuffle=True, num_workers=8)))

ig = Saliency(attributation_for_feature_forward)
attr = ig.attribute(data_sample.x, additional_forward_args=((data_sample,)), target=data_sample.y)

print("Node features size: " + str(list(data_sample.x.size())))
print("Attributes size: " + str(list(attr.size())))
print()
print("Minimal value: " + str(torch.min(attr).item()))
print("Maximal value: " + str(torch.max(attr).item()))
print()
# print("Top 5 feature indexes: " + str(torch.topk(torch.mean(attr,axis=0),5)))
# print()
# print("Top 5 node indexes: " + str(torch.topk(torch.mean(attr,axis=1),5)))
# print()
import torch_geometric
from torch.nn import Module

# https://github.com/rusty1s/pytorch_geometric/issues/1730

class ExplainNet(Module):
    def __init__(self,model, data):
        super(ExplainNet, self).__init__()
        self.model = model
        self.data = data

    def forward(self,x, edge_index):
        self.model.forward(self.data)
        
explainer = torch_geometric.nn.GNNExplainer(ExplainNet(model, data_sample))

node_idx = 1
# node_feat_mask, edge_mask = explainer.explain_node(node_idx,data_sample.x, data_sample.edge_index)
explainer.explain_graph(data=data_sample)

maximal_feature_importance = torch.max(node_feat_mask).item()
maximal_edge_importance = torch.max(edge_mask).item()

print("Maximal feature importance: %d " % (maximal_feature_importance))
print("Maximal edge importance: %d " % (maximal_edge_importance,))
print()
#Node Feature explainations

import torch_geometric
from captum.attr import Saliency,IntegratedGradients,GradientShap, FeatureAblation, DeepLift

def attributation_for_feature_forward(x, data):
    data.x=x
    return model(data)


data_sample = next(iter(DataLoader(model.train_dataset,batch_size=1, shuffle=True, num_workers=8)))


ig = IntegratedGradients(attributation_for_feature_forward)
attr = ig.attribute(data_sample.x, additional_forward_args=((data_sample,)), target=data_sample.y, baselines=torch.zeros_like(data_sample.x), internal_batch_size=1)

print("Node features size: " + str(list(data_sample.x.size())))
print("Attributes size: " + str(list(attr.size())))
print()
print("Minimal value: " + str(torch.min(attr).item()))
print("Maximal value: " + str(torch.max(attr).item()))
# print()
# print("Top 5 feature indexes: " + str(torch.topk(torch.mean(attr,axis=0),5)))
# print()
# print("Top 5 node indexes: " + str(torch.topk(torch.mean(attr,axis=1),5)))
# print()
# #Node Feature explainations

# import torch_geometric
# from captum.attr import Saliency,IntegratedGradients,GradientShap, FeatureAblation, DeepLift, GuidedBackprop

# def attributation_for_feature_forward(x, data):
# #     print("--")
# #     print(data)
# #     print(data.x.size())
# #     print("x")
# #     print(x.size())
#     data.x=x
#     return model(data)


# data_sample = next(iter(DataLoader(model.train_dataset,batch_size=1, shuffle=True, num_workers=8)))


# ig = GuidedBackprop(model.relu1)
# attr = ig.attribute(data_sample.x, additional_forward_args=((data_sample,)), target=data_sample.y, baselines=torch.zeros_like(data_sample.x))

# print("Node features size: " + str(list(data_sample.x.size())))
# print("Attributes size: " + str(list(attr.size())))
# print()
# print("Minimal value: " + str(torch.min(attr).item()))
# print("Maximal value: " + str(torch.max(attr).item()))
# # print()
# # print("Top 5 feature indexes: " + str(torch.topk(torch.mean(attr,axis=0),5)))
# # print()
# # print("Top 5 node indexes: " + str(torch.topk(torch.mean(attr,axis=1),5)))
# # print()
#Node Feature explainations

import torch_geometric
from captum.attr import Saliency,IntegratedGradients,GradientShap, FeatureAblation, DeepLift, GuidedBackprop

def attributation_for_feature_forward(x, data):
    data.x=x
    return model(data)


data_sample = next(iter(DataLoader(model.train_dataset,batch_size=1, shuffle=True, num_workers=8)))
print(data_sample)
print(len(data_sample.to_data_list()))
data = data_sample.to_data_list()[0]
data.x = data_sample.x
print(data_sample.x.size())
data = torch_geometric.data.Data(x=data_sample.x,edge_index=data_sample.edge_index,y=data_sample.y, pos=data_sample.pos, edge_attr=data_sample.edge_attr)
data.batch = None
print(data)
ig = GradientShap(attributation_for_feature_forward)
attr = ig.attribute(data.x, additional_forward_args=((data,)), target=data.y, baselines=torch.zeros_like(data.x))

print("Node features size: " + str(list(data_sample.x.size())))
print("Attributes size: " + str(list(attr.size())))
print()
print("Minimal value: " + str(torch.min(attr).item()))
print("Maximal value: " + str(torch.max(attr).item()))
# print()
# print("Top 5 feature indexes: " + str(torch.topk(torch.mean(attr,axis=0),5)))
# print()
# print("Top 5 node indexes: " + str(torch.topk(torch.mean(attr,axis=1),5)))
# print()
#Node Feature explainations

import torch_geometric
from captum.attr import Saliency,IntegratedGradients,GradientShap, FeatureAblation, DeepLift

def attributation_for_feature_forward(x, data):
    return model(data)


data_sample = next(iter(DataLoader(model.train_dataset,batch_size=1, shuffle=True, num_workers=8)))


feature_mask = torch.zeros_like(data_sample.x)

ig = FeatureAblation(attributation_for_feature_forward)
attr = ig.attribute(data_sample.x, additional_forward_args=((data_sample,)), target=data_sample.y, baselines=torch.zeros_like(data_sample.x))

print("Node features size: " + str(list(data_sample.x.size())))
print("Attributes size: " + str(list(attr.size())))
print()
print("Minimal value: " + str(torch.min(attr).item()))
print("Maximal value: " + str(torch.max(attr).item()))
# print()
# print("Top 5 feature indexes: " + str(torch.topk(torch.mean(attr,axis=0),5)))
# print()
# print("Top 5 node indexes: " + str(torch.topk(torch.mean(attr,axis=1),5)))
# print()
if do_train:
    torch.save(model.state_dict(), "model.pth")