"""
Kaggle competition: https://www.kaggle.com/c/digit-recognizer/
"""
# This Python 3 environment comes with many helpful analytics libraries installed
# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python
# For example, here's several helpful packages to load

import numpy as np # linear algebra
import pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)

# Input data files are available in the read-only "../input/" directory
# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory

import os
for dirname, _, filenames in os.walk('/kaggle/input'):
    for filename in filenames:
        print(os.path.join(dirname, filename))

# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using "Save & Run All" 
# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session
"""
We need to import Torch's libraries
"""
import torchvision
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import transforms
from torch.utils.data import DataLoader, Dataset
import wandb

from kaggle_secrets import UserSecretsClient
user_secrets = UserSecretsClient()
secret_value_0 = user_secrets.get_secret("wand_key")

wandb.login(key=secret_value_0)
wandb.init(project='MNIST_MLP', save_code=True)
"""
## Data preparation

A custom dataset which uses the CSV from Kaggle, avoid downloading the dataset from internet
"""
class MyMNISTDataset(Dataset):
    
    def __init__(self, file_path, transform = transforms.Compose([transforms.ToPILImage()]), test_data=False, use_gpu=torch.cuda.is_available()):
        # read the data
        df = pd.read_csv(file_path)
        # for test data we don't have any target
        # MNIST images are 28 by 28, grey colors
        if test_data:
            self.X = df.values.reshape((-1,28,28)).astype(np.uint8)[:,:,:,None]
            self.y = None
        else:
            self.X = df.iloc[:,1:].values.reshape((-1,28,28)).astype(np.uint8)[:,:,:,None]
            self.y = torch.from_numpy(df.iloc[:,0].values)
        self.transform = transform
        self.use_gpu = use_gpu
    
    def __len__(self):
        return len(self.X)

    def __getitem__(self, idx):
        data = self.transform(self.X[idx])
        if self.y is not None:
            target = self.y[idx]
            if self.use_gpu:
                data = data.cuda()
                target = target.cuda()
            return data, target
        else:
            if self.use_gpu:
                data = data.cuda()
            return data
transformations=transforms.Compose([transforms.ToPILImage(), 
                                    transforms.ToTensor(), 
                                    transforms.Normalize(mean=(0.5,), std=(0.5,))])

train_dataset = MyMNISTDataset('/kaggle/input/digit-recognizer/train.csv', transform=transformations, test_data=False)
test_dataset = MyMNISTDataset('/kaggle/input/digit-recognizer/test.csv', transform=transformations, test_data=True)

# TODO split train between validation and training set
training_size = int(0.7 * len(train_dataset))
validation_size = len(train_dataset) - training_size
train, val = torch.utils.data.random_split(train_dataset, [training_size, validation_size], generator=torch.Generator().manual_seed(0))

# create data loader for train and test set
data_loader_train = torch.utils.data.DataLoader(train, batch_size=32, shuffle=True)
data_loader_val = torch.utils.data.DataLoader(val, batch_size=32, shuffle=True)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=32, shuffle=False)
"""
## MLP

### Define model architecture
You need to reach at least 70% accuracy on the test set
"""
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Sequential(         
            nn.Conv2d(
                in_channels=1,              
                out_channels=16,            
                kernel_size=5,              
                stride=1,                   
                padding=2,                  
            ),                              
            nn.ReLU(),                      
            nn.MaxPool2d(kernel_size=2),    
        )
        self.conv2 = nn.Sequential(         
            nn.Conv2d(16, 32, 5, 1, 2),     
            nn.ReLU(),                      
            nn.MaxPool2d(2),                
        )
        # fully connected layer, output 10 classes
        self.out = nn.Linear(32 * 7 * 7, 10)
    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        # flatten the output of conv2 to (batch_size, 32 * 7 * 7)
        x = x.view(x.size(0), -1)       
        output = self.out(x)
        return output    

# TODO

#class Net(nn.Module):    
 #   def __init__(self, input_size=784, output_size=10, layers=[120,84]):
  #      super(Net, self).__init__()
   #     #TODO
    #    self.d1 = nn.Linear(input_size,layers[0]) #layer1
     #   self.d2 = nn.Linear(layers[0],layers[1]) #layer2
      #  self.d3 = nn.Linear(layers[1],output_size) #output layer

 #   def forward(self, x):
  #     x = F.relu(self.d1(x))
   #     x = F.relu(self.d2(x))
    #    x = self.d3(x)
     #   return F.log_softmax(x, dim=1)
28*28
"""
### Init the model and put it on GPU
"""
# TODO
model = Net()
model.cuda()
print(model)
"""
### Training loop
Log the accuracy and the loss to wandb
"""
# TODO
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
num_epochs = 100
train_losses = []
val_losses = []
train_correct = []
val_correct = []

for epoch in range(num_epochs):
    trn_corr = 0
    val_corr = 0
    #TODO TRAINING LOOP
    for i, y in data_loader_train:
        output = model(i) # Forward pass 
        loss = criterion(output, y)
        
        trn_predicted = torch.argmax(output.data,1)[1].data.squeeze()
        batch_corr = (trn_predicted == y).sum().item() / float(y.size(0))
        trn_corr += batch_corr
        
        optimizer.zero_grad() # please don't forget!
        loss.backward() # remember: You need to tell wrt to what the gradient is computed
        optimizer.step() # do a step in the gradient direction
    
    train_losses.append(loss)
    train_correct.append(trn_corr)
    
    with torch.no_grad():
        #TODO VALIDATION LOOP
        for z, j in data_loader_val:
            y_val = model(z)
            loss2 = criterion(y_val,j)
            
            val_predicted = torch.argmax(y_val.data,1)[1].data.squeeze()
            val_corr += (val_predicted == j).sum().item() / float(j.size(0))
    
    val_losses.append(loss2)
    val_correct.append(val_corr)
    
    print(f'epoch {epoch + 1}, training_loss {loss.item()}, validation_loss {loss2.item()}, Train acccuracy {train_correct[-1]}, Validation acccuracy {val_correct[-1]}')
    wandb.log({'num_epochs': num_epochs, 'training_loss': loss,'Train acccuracy': train_correct, 'validation_loss': loss2, 'Validation acccuracy': val_correct})

    # SAVE THE MODEL
    torch.save(model.state_dict(), 'my_model.pt')
"""
### Make prediction
And submit to Kaggle for grading
"""
with torch.no_grad():
    model.eval()
    test_pred = torch.LongTensor()
    for i, data in enumerate(test_loader):
        output = model(data)
        _, predicted = torch.max(output.data, 1)
        predicted = predicted.cpu()
        test_pred = torch.cat((test_pred, predicted), dim=0)
    out_df = pd.DataFrame(np.c_[np.arange(1, len(test_dataset)+1)[:,None], test_pred.numpy()], columns=['ImageId', 'Label'])
    out_df.to_csv('submission.csv', index=False)
"""
### 10 most misclafficied image
"""
import matplotlib.pyplot as plt


unshuffle_train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
                                           batch_size=32, shuffle=False)
with torch.no_grad():
    model.eval()
    missclasified = torch.DoubleTensor()
    for batch_idx, (data, target) in enumerate(unshuffle_train_loader):
        output = model(data)
        prob, predicted = torch.max(output.data, 1)
        predicted = predicted.cpu()
        target = target.cpu()
        prob = prob.cpu().double()
        missclassified_prob = torch.where(predicted == target, 0., prob)
        missclasified = torch.cat((missclasified, missclassified_prob), dim=0)
    most_misclassified = torch.argsort(missclasified, descending=True)
    top_ten_misclassified = most_misclassified[:10]
for misclassified in top_ten_misclassified:
    plt.imshow(train_dataset[misclassified][0].cpu().reshape(28,28))
    with torch.no_grad():
        data, target = train_dataset[misclassified]
        data = data.reshape(1, 1, 28,28)
        output = model(data)
        _, predicted = torch.max(output.data, 1)
        plt.title(f'Predicted: {predicted.item()}, Ground truth: {target}')
    plt.show()