import numpy as np # linear algebra
import pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)
import os
filename = None
for dirname, _, filenames in os.walk('/kaggle/input'):
    
    for filename in filenames:
        if filename == 'train.pt':
            dataset_dir = dirname
        
        print(os.path.join(dirname, filename))
print("\nInput directory: ", dirname)
"""
### Our approach uses Pytorch Geometric. It is not currently included in Kaggle's environment, so we must install it via  `pip`. 

Unfortunately, building these wheel files for CUDA=11.0 is *reaaalllllly slow* (>1 hour) for some reason, and that is the version available for the Kaggle environment.
"""
!pip install torch-scatter torch-sparse torch-cluster torch-spline-conv torch-geometric -f https://data.pyg.org/whl/torch-1.9.0+cu110.html
import torch
from torch_geometric.data import DataLoader,Dataset,Data
"""
#### Now that this is out of the way, let us *read the Pytorch Geometric version of our dataset* (created on a previous kernel: https://www.kaggle.com/isonettv/converting-kddbr-2021-data-to-pytorch-files/notebook). 

There are 2 files, **'train.pt'** and **'test.pt'**. Both of them are lists, with each element of them being an instance of the TSP problem, represented as a ``pytorch_geometric.data.Data`` object. The attributes of each element are:
1. **x**: tensor with shape [200,2]. The coordinates for up to 200 vertices. Padded out with zeros in case the number is < 200.
2. **m**: tensor with shape [200]. A mask such that ``x[m]`` returns the coordinates of the actual vertices.
3. **y_index**: tensor with shape [?,2]. Each row corresponds to some edge $(i,j)$ that is part of the solution. Only present in `train.pt`.
"""
import torch
from torch_geometric.data import DataLoader,Dataset,Data
import pandas as pd
import numpy as np
import gc
import torch.nn.functional as F
import os.path as osp

device = 'cuda'

L = torch.load(osp.join(dataset_dir,'train.pt'),
               map_location=device)
L_train = [L[i] for i in range(len(L)) if i % 10 > 0]
L_val = [L[i] for i in range(len(L)) if i % 10 == 0]
import random
random.Random(4).shuffle(L_val)

L_test = torch.load(osp.join(dataset_dir,'test.pt'),
                    map_location=device)

print("Training example: ", L_train[0])
print("Test example: ", L_train[0])
"""
### Next, we implement the ``Dataset`` interface from Pytorch Geometric. When we index the dataset (`dataset['train'][i]`, for example), we look up the corresponding TSP problem and apply some small transformations only at runtime, saving up on memory:

1. We dynamically compute the KNN graph (k=20 is a hyperparameter)
2. We transform the target edge list (**y_index**) into a dense adjacency matrix, which we use for the cost function.
"""
import torch
from torch_geometric.data import DataLoader,Dataset,Data
import pandas as pd
import numpy as np
import gc
import torch.nn.functional as F
from torch_geometric.utils import to_dense_adj
from torch_geometric.transforms import KNNGraph,Distance

from torch_geometric.nn import knn_graph



class KDD_Dataset(Dataset):
    """
    Base class representing a dataset for the competition
    """
    
    distance_transf = Distance()
    knn_transf = KNNGraph(k=20)
    def __init__(self, L,mode='train'):
        """
        Args:
            csv_file (string): Path to the csv file with annotations.
            root_dir (string): Directory with all the .pt files.
            transform (callable, optional): Optional transform to be applied
                on a sample.
        """
        self.L = L
        
        
    def __len__(self):
        return len(self.L)
    def __getitem__(self, idx):
        sample = self.L[idx]
        num_nodes = torch.where(sample.m==True)[0][-1] + 1
        
        num_nodes_vector = num_nodes * torch.ones(200,device=sample.x.device)
        num_nodes_vector[sample.m.view(-1)==False] = 1
        
        m = sample.m.view(-1,1).bool()
            
        
        data = Data(m=m, pos=sample.x[0:num_nodes,0:num_nodes],
                         num_nodes=num_nodes,
                         num_nodes_vector=num_nodes_vector)
        if hasattr(sample,'y_index'):
            data.y = to_dense_adj(sample.y_index.transpose(0,1).long(),max_num_nodes=200)
        
        data = self.knn_transf(data)
        data = self.distance_transf(data)
        
        data.num_nodes=200
        data.pos = sample.x
        data.x = sample.x
        
        data.n_edges = data.edge_index.shape[1]

        return data   

datasets = {"train":KDD_Dataset(L_train,mode='train'),
            "val":KDD_Dataset(L_val,mode='val'),
            "test":KDD_Dataset(L_test,mode='test'),
           }

batch_size = 32
dataloaders = {"train":DataLoader(datasets['train'],batch_size=batch_size,shuffle=True),
            "val":DataLoader(datasets['val'],batch_size=batch_size,shuffle=True),
            "test":DataLoader(datasets['test'],batch_size=batch_size),
           }
for k in dataloaders.keys():
    print(f"Length of {k} dataset : {len(datasets[k])}")
"""
### The next step is to define our model. 

#### We feed our input data to the **edge convolutional operator** from the **“Dynamic Graph CNN for Learning on Point Clouds”** paper :https://arxiv.org/abs/1801.07829. After a few layers of EdgeConv, we obtain our final prediction using a **Multilayer Perceptron**. LeakyRELUs serve as the activation function.

#### We have not performed any sort of extensive hyperparameter tuning, though this configuration seemed to perform best from our (limited) experiments
"""
import torch
import pandas as pd
from torch_geometric.nn import MessagePassing
from torch_geometric.utils import add_self_loops, degree

from torch.nn import Sequential, Linear, BatchNorm1d, ReLU, LeakyReLU, Sigmoid
from torch_geometric.nn.conv import GATConv,TransformerConv,GCNConv,GINConv,EdgeConv
import torch.nn.functional as F
from torch_scatter import scatter
class ChallengeModel(torch.nn.Module):

    
    def __init__(self,num_iterations=5,H=64,**kwargs):
        super(ChallengeModel, self).__init__(**kwargs)
        self.num_iterations = num_iterations
        
        
        self.H = H
        self.convs = []
        
        for i in range(self.num_iterations):
            _in = 2 if i == 0 else H
            _out = H
            mlp = Sequential(Linear(2 * _in, 128),
                       LeakyReLU(),
                       Linear(128, _out))
            self.convs.append(EdgeConv(mlp).cuda())
        
        setattr(self,'convs',torch.nn.ModuleList(getattr(self,'convs')))
        
        
        self.edgeMLP = Sequential(Linear(2 * H, 512),
                       LeakyReLU(),
                       Linear(512, 256),
                       LeakyReLU(),
                       Linear(256, 1))
    
    def forward(self,data):
        edge_index = data.edge_index.clone().long()
        x = sample.pos.float().clone()
        
        for i in range(len(self.convs)):
            x = F.leaky_relu(self.convs[i](x,edge_index))
        
        """
            To create the input for our MLP edge classifier, we simply concatenate the outputs for
            the source and target vertices.
        """
        X_source = x[edge_index[0,:],:]
        X_target = x[edge_index[1,:],:]
        X = torch.cat([X_source,X_target],axis=1).clone()
        def act(x):
            return F.leaky_relu(x)
        
        X = self.edgeMLP(X)

        L = []
        num_batches = sample.batch.max()+1

        batch = data.batch.clone().to('cuda')
        
        """
            We use this helper function to convert the prediction ´´X´´ to an adjacency matrix format.
            This yields B 200x200 matrices in total. 
            
            Note that any edge that is not present to the KNN graph will have its entry corresponding 
        to 0 in the dense format.
        
        """
        mat = to_dense_adj(edge_index,edge_attr=X
                               ,max_num_nodes=200,
                         batch=batch).squeeze(-1)
        del batch

        return mat
               
"""
### With our model defined, we will need to add code for the training loop. 


### Our custom loss function uses the same kind of weight balancing as previous approaches :https://arxiv.org/abs/1906.01227https://arxiv.org/abs/1906.01227. 
$$w_0 = \frac{n^2}{c(n^2-2n)} $$
$$w_1 = \frac{n^2}{c\times 2n} $$
In our case, $c=2$ because we perform binary classification. 

---

The final classification probabilities have to be converted to binary labels, maximizing the F1 score. One possibility is to select the best threshold hyperparameter using the validation data. However, currently we simply set a fixed hyperparameter ``THRESH``.  
"""
from tqdm import tqdm,trange
import sklearn.metrics
from sklearn.metrics import f1_score

def _xent(preds,actuals):
    """ Here we assume both preds and actuals are square matrices with the relevant nodes only.
        preds sums up to 1, actuals sums up to 2.
    """
    num_nodes = preds.shape[0]
    preds = preds.reshape(-1)
    actuals = actuals.reshape(-1)
    
    w0 =  (num_nodes**2) / (2*(num_nodes**2 - 2*num_nodes))
    w1 = (num_nodes**2) / (2*(2*num_nodes))
    
    weights = w0 * torch.ones(num_nodes**2,device=preds.device)
    weights[actuals > 0] = w1

    
    return F.binary_cross_entropy_with_logits(preds,actuals,
                                              weight=weights,
                                              reduction='mean')

def custom_XENT(preds,actuals,edge_index,m,num_nodes_vector):
    """ edge_index[1,:] now points from 0 to """
    nb = actuals.shape[0]
    m = m.clone().reshape(nb,200)
    loss = 0
    for i in range(nb):
        num_nodes = torch.where(m[i,:]==True)[0][-1].item() + 1
        loss += _xent(preds[i,0:num_nodes,0:num_nodes],
                      actuals[i,0:num_nodes,0:num_nodes])

    
    loss /= nb
    return loss

THRESH = 0.56
def _metrics(preds,actuals):
    preds = np.array(preds > THRESH,dtype=np.int32)
    #print(preds,actuals)
    return sklearn.metrics.f1_score(preds,actuals)

def custom_metrics(preds,actuals,edge_index,m,num_nodes_vector):
    """ edge_index[1,:] now points from 0 to """
    nb = actuals.shape[0]
    m = m.clone().reshape(nb,200)
    loss = 0
    preds = torch.sigmoid(preds).clone().detach().cpu().numpy()
    actuals = actuals.clone().detach().cpu().numpy()
    
    for i in range(nb):
        num_nodes = torch.where(m[i,:]==True)[0][-1].item() + 1

        
        loss += _metrics(preds[i,0:num_nodes,0:num_nodes].reshape((-1,)),
                      actuals[i,0:num_nodes,0:num_nodes].reshape((-1,)))

    
    loss /= nb
    return loss

import numpy as np
best_f1 = -1
model = ChallengeModel().cuda()

from torch.utils.tensorboard import SummaryWriter
writer = SummaryWriter()
num_epochs = 10
opt = torch.optim.Adam(lr=1e-3,params=model.parameters())
step = 0

#model.load_state_dict('./model/model_3650.pt')
torch.manual_seed(420)
        
        
for epoch in range(num_epochs):
    
        
    for mode in ['train','val']:
        if mode == 'train':
                model.train()
        else:
                model.eval()
        stats = {'loss':[],
                 'f1_score':[],
                }
                
        running_loss = []
        out_mean = []
        occup_mean = []
        actual_out_mean = []
        cnt = 0
        total=len(dataloaders[mode])//10 if mode == 'train' else len(dataloaders[mode])//10
        #mode_maybe_shuffle = 'train' if mode == 'train' else mode
        for i,sample in tqdm(enumerate(dataloaders[mode]),
                             total=total):
                
                
                with torch.set_grad_enabled(mode == 'train'):
                    if i == total:
                        break
                    if mode == 'train':
                        opt.zero_grad()
                    cnt += 1
                    
                    out = model(sample)
                    
                    
                    #sklearn.metrics.f1_score()
                    
                    if mode == 'train':
                        loss = custom_XENT(out,sample.y,sample.edge_index.clone().long(),
                                          sample.m,sample.num_nodes_vector)
                        loss_np = loss.cpu().item()
                        loss.backward()
                        opt.step()
                    
                    f1_score = custom_metrics(out,sample.y,sample.edge_index.clone().long(),
                                          sample.m,sample.num_nodes_vector)
                    _stats = {'loss':loss,
                              'f1_score':torch.as_tensor(f1_score*100.0,device=loss.device)
                             }
                    for k in _stats.keys():
                        stats[k].append(_stats[k].cpu().item())
                    del _stats
               

        print("======================================")
        print(f"Epoch {epoch} - Avg stats ({mode})")
        for k in stats.keys():
            print(f'Mean {k}: {np.array(stats[k]).mean()}')
            epoch_f1 = np.array(stats['f1_score']).mean()
            if best_f1 < epoch_f1:
                best_f1 = epoch_f1
                torch.save(model.state_dict(),f'./best_model.pt')
                print("New best")
    print("Flushed")
    step += 1
    torch.cuda.empty_cache()

    
"""
### Lastly, we put our model into evaluation mode, and  save the predictions to a CSV file.
"""
model.eval()
model.load_state_dict(model.state_dict(),f'best_model.pt')

mode='test'


def get_pred(preds,actuals,edge_index,m,num_nodes_vector,upload_file,cnt):
    """ edge_index[1,:] now points from 0 to """
    nb = preds.shape[0]
    m = m.clone().reshape(nb,200)
    loss = 0
    preds = torch.sigmoid(preds).clone().detach().cpu().numpy()
    for i in range(nb):
        num_nodes = torch.where(m[i,:]==True)[0][-1].item() + 1
        for i1 in range(num_nodes):
            for j1 in range(num_nodes):
                if i1 == j1:
                    continue
                else:
                    upload_file.write(f'{cnt+i}:{i1}-{j1},{1 if preds[i,i1,j1]>THRESH else 0}\n')
import os 
cnt = 0
with open('submission.csv', "w") as upload_file:
    
    upload_file.write('Id,Predicted\n')
    
    for i,sample in tqdm(enumerate(dataloaders[mode]),
                                 total=len(dataloaders[mode])):
        out = model(sample)
        with torch.set_grad_enabled(False):
            pred = get_pred(out,sample.y,sample.edge_index.clone().long(),
                                              sample.m,sample.num_nodes_vector,
                              upload_file,cnt)
            
            
            cnt += sample.batch.max().item() + 1