import torch
import torchvision
import torchvision.transforms as transforms
from tqdm import tqdm_notebook as tqdm
import numpy as np
import matplotlib.pyplot as plt
from sklearn.metrics.pairwise import rbf_kernel
import time


train_loader = torch.utils.data.DataLoader(
    torchvision.datasets.MNIST('../data', train=True, download=True,
                   transform=transforms.Compose([
                       transforms.ToTensor()
                   ])), shuffle=True)
test_loader = torch.utils.data.DataLoader(
    torchvision.datasets.MNIST('../data', train=False, transform=transforms.Compose([
                       transforms.ToTensor()
                   ])), shuffle=True)

train_data = []
train_label = []
test_data = []
test_label = []
for data, target in tqdm(train_loader):
    train_data.append(data.numpy().ravel())
    train_label.append(target)
train_data, train_label = np.array(train_data), np.array(train_label)
for data, target in tqdm(test_loader):
    test_data.append(data.numpy().ravel())
    test_label.append(target)
test_data, test_label = np.array(test_data), np.array(test_label)

gamma = 1. / test_data.shape[1]
train_kernel = rbf_kernel(train_data, gamma=gamma)
test_kernel = rbf_kernel(test_data, train_data, gamma=gamma)

print(clf.score(train_kernel, train_label))
print(clf.score(test_kernel, test_label))

# compute robust error
num_correct = 0
eps_norm = 0.3
for i in tqdm(range(0, 100)):
    score = np.dot(test_kernel[i, :], clf.coef_.T) + clf.intercept_
    if np.argmax(score) != test_label[i]:
        continue
    # now test the robustness of score
    diff = np.dot((clf.coef_ - clf.coef_[test_label[i]]) * 2 * gamma * test_kernel[i, :], test_data[i, :] - train_data)
    if np.any(np.linalg.norm(diff, ord=1, axis=1) * eps_norm > score[test_label[i]] - score) == False:
        num_correct += 1
print(num_correct)