import urllib
 
from bs4 import BeautifulSoup
import requests
 
def listFD(url, ext=''):
  page = requests.get(url).text
  soup = BeautifulSoup(page, 'html.parser')
  return [url + '/' + node.get('href') for node in soup.find_all('a') if node.get('href').endswith(ext)]
 
def url_to_image(url):
	# download the image, convert it to a NumPy array, and then read
	# it into OpenCV format
  resp = urllib.request.urlopen(url)
  image = np.asarray(bytearray(resp.read()), dtype="uint8")
  image = cv2.imdecode(image, cv2.IMREAD_COLOR)
 
  # Convert BGR to RGB
  image = image[:, :, [2,1,0]]
 
  return image
 
 
DATA_URL = "http://www.cs.cornell.edu/courses/cs5670/2019sp/projects/pa5/assignment_data/dataset/"
IMAGENET_LABELS_URL = "http://www.cs.cornell.edu/courses/cs5670/2019sp/projects/pa5/assignment_data/imagenet_classes.txt"
DOGS_DIR = "test-dog"
FOOD_DIR = "test-food"
 
food_paths = listFD(DATA_URL + FOOD_DIR, ".jpg")
dog_paths = listFD(DATA_URL + DOGS_DIR, ".jpg")
 
 
#-----------
import os
 
for url in dog_paths:
    name='/home/fffasttime/userres/cv/data/test-dog/'+url.split('/')[-1]
    if os.path.exists(name):
        continue
    urllib.request.urlretrieve(url,name)
#-----------
import os
 
for url in food_paths:
    name='/home/fffasttime/userres/cv/data/test-food/'+url.split('/')[-1]
    if os.path.exists(name):
        continue
    urllib.request.urlretrieve(url,name)

def vis_square(data, title=None):
    """Take a Tensor of shape (n, K, height, width) or (n, K, height, width)
       and visualize each (height, width) thing in a grid of size approx. sqrt(n) by sqrt(n)"""
 
    if data.size(1) > 3:
      data = data.view(-1, 1, data.size(2), data.size(3))
 
    data = to_numpy_image(data)
 
    # normalize data for display
    data = (data - data.min()) / (data.max() - data.min())
 
    # force the number of filters to be square
    n = int(np.ceil(np.sqrt(data.shape[0])))
    padding = (((0, n ** 2 - data.shape[0]),
               (0, 2), (0, 2))                 # add some space between filters
               + ((0, 0),) * (data.ndim - 3))  # don't pad the last dimension (if there is one)
    data = np.pad(data, padding, mode='constant', constant_values=1)  # pad with ones (white)
 
    # tile the filters into an image
    data = data.reshape((n, n) + data.shape[1:]).transpose((0, 2, 1, 3) + tuple(range(4, data.ndim + 1)))
    data = data.reshape((n * data.shape[1], n * data.shape[3]) + data.shape[4:])
 
    data = data.squeeze()
 
    # plot it
    plt.figure(figsize=(8, 8))
    plt.imshow(data)
    plt.axis('off')
    if title: plt.title(title)
 
 
weights = model.features[0].weight.data
vis_square(weights, title="Visualizing filters in conv1")

import functools
 
# Re-define the model to clear any previously registered hooks
model = alexnet(pretrained=True, num_classes=1000)
 
# Define a hook that visualizes a layer output
def show_activations_hook(name, module, input, output):
 
  print(f"Visualizing layer: {name}")
  # For conv/relu layer outputs (BxCxHxW) we plot an image as before
  if output.dim() == 4:
    vis_square(output, f"Activations on: {name}")
 
  # For linear layer outputs, we plot the activations as a line plot
  else:
    feat = output.data.view([-1]).cpu().numpy()
    plt.figure(figsize=(15, 3))
    plt.plot(feat)
    plt.title(f"Activations on: {name}")
 
# Register the hook on the select set of modules
for name, module in model.shortcut_modules():
  hook = functools.partial(show_activations_hook, name)
  module.register_forward_hook(hook)
 
# PyTorch modules work on minibatches and expect the first axis to be the batch axis
# If we run a model on a single image, we must turn this into a batch of size 1
model_input = Variable(example_image[np.newaxis, ...])
 
# Run the forward pass on the model
class_activations = model(model_input)[0]

import torch.nn as nn
import torch.utils.model_zoo as model_zoo
 
 
model_urls = {
    'alexnet': 'https://download.pytorch.org/models/alexnet-owt-4df8aa71.pth',
}
 
INPLACE = False
 
class AlexNet(nn.Module):
 
    def __init__(self, num_classes=1000):
        super(AlexNet, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=11, stride=4, padding=2),
            nn.ReLU(inplace=INPLACE),
            nn.MaxPool2d(kernel_size=3, stride=2),
            nn.Conv2d(64, 192, kernel_size=5, padding=2),
            nn.ReLU(inplace=INPLACE),
            nn.MaxPool2d(kernel_size=3, stride=2),
            nn.Conv2d(192, 384, kernel_size=3, padding=1),
            nn.ReLU(inplace=INPLACE),
            nn.Conv2d(384, 256, kernel_size=3, padding=1),
            nn.ReLU(inplace=INPLACE),
            nn.Conv2d(256, 256, kernel_size=3, padding=1),
            nn.ReLU(inplace=INPLACE),
            nn.MaxPool2d(kernel_size=3, stride=2),
        )
        self.classifier = nn.Sequential(
            nn.Dropout(),
            nn.Linear(256 * 6 * 6, 4096),
            nn.ReLU(inplace=INPLACE),
            nn.Dropout(),
            nn.Linear(4096, 4096),
            nn.ReLU(inplace=INPLACE),
            nn.Linear(4096, num_classes),
        )
 
        self.module_shortcuts = {
            "conv1": self.features[0],
            "relu1": self.features[1],
            "conv2": self.features[3],
            "relu2": self.features[4],
            "conv3": self.features[6],
            "relu3": self.features[7],
            "conv4": self.features[8],
            "relu4": self.features[9],
            "conv5": self.features[10],
            "relu5": self.features[11],
 
            "fc6": self.classifier[2],
            "fc7": self.classifier[5],
            "fc8": self.classifier[6]
        }
 
 
    def __getitem__(self, layer_name):
      if layer_name in self.module_shortcuts:
        return self.module_shortcuts[layer_name]
      return None
 
 
    def shortcut_modules(self):
      for name, mod in self.module_shortcuts.items():
        yield name, mod
 
 
    def forward(self, x):
        x = self.features(x)
        x = x.view(x.size(0), 256 * 6 * 6)
        x = self.classifier(x)
        return x
 
 
def alexnet(pretrained=False, **kwargs):
    r"""AlexNet model architecture from the
    `"One weird trick..." <https://arxiv.org/abs/1404.5997>`_ paper.
    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
    """
    model = AlexNet(**kwargs)
    if pretrained:
        print("Loading pre-trained weights from model zoo")
        model.load_state_dict(model_zoo.load_url(model_urls['alexnet']))
    model.eval()
    return model
 
model = alexnet(pretrained=True, num_classes=1000)

def convert_ilsvrc2012_probs_to_dog_vs_food_probs(probs_ilsvrc):
    """
    Convert from 1000-class ILSVRC probabilities to 2-class "dog vs food"
    incices.  Use the variables "dog_indices" and "food_indices" to map from
    ILSVRC2012 classes to our classes.
    HINT:
    Compute "probs" by first estimating the probability of classes 0 and 1,
    using probs_ilsvrc.  Stack together the two probabilities along axis 1, and
    then normalize (along axis 1).
    :param probs_ilsvrc: shape (N, 1000) probabilities across 1000 ILSVRC classes
    :return probs: shape (N, 2): probabilities of each of the N items as being
        either dog (class 0) or food (class 1).
    """
    # in the ILSVRC2012 dataset, indices 151-268 are dogs and index 924-969 are foods
    dog_indices = range(151, 269)
    food_indices = range(924, 970)
    N, _ = probs_ilsvrc.shape
    probs = np.zeros((N, 2)) # placeholder
    ############################ TODO 2 BEGIN #################################
    probs[:,0] = np.sum(probs_ilsvrc[:,151:269],axis=1)
    probs[:,1] = np.sum(probs_ilsvrc[:,924:970],axis=1)
    norm = np.sum(probs,axis=1)
    probs[:,0] /= norm
    probs[:,1] /= norm
    ############################ TODO 2 END #################################
    return probs

# Re-define the model to clear any previously registered hooks
model = alexnet(pretrained=True, num_classes=1000)
N = len(full_dataset)
 
dataloader = DataLoader(
    full_dataset,
    batch_size=1024,
    shuffle=False,
    num_workers=0,
    pin_memory=False,
    drop_last=False)
 
assert len(dataloader) == 1, "Since batch_size is bigger than the number of examples, we should only have one batch"
 
for batch in dataloader:
  images = batch[0]
  labels = batch[1]
  labels_np = labels.numpy()
 
  class_activations = model(Variable(images))
  ilsvrc_class_probs = F.softmax(class_activations, dim=1)
 
  # Convert from Variable containing FloatTensor to numpy ndarray
  ilsvrc_class_probs_np = ilsvrc_class_probs.data.cpu().numpy()
 
  dogfood_class_probs_np = convert_ilsvrc2012_probs_to_dog_vs_food_probs(ilsvrc_class_probs_np)
 
  assert list(dogfood_class_probs_np.shape) == [N, 2]
 
  np.testing.assert_almost_equal(np.sum(dogfood_class_probs_np, axis=1), np.ones(N), decimal=5)
 
  print("Seems correct!")
 
predicted_class = np.argmax(dogfood_class_probs_np, axis=1)
correct_mask = predicted_class == labels_np
num_correct = np.sum(correct_mask)
accuracy = 100.0 * float(num_correct) / N
 
print(f"Overall accuracy {accuracy} {num_correct} / {N}")
 
for cidx, cname in enumerate(classes):
  cls_mask = labels_np == cidx
  predicted_cls = predicted_class[cls_mask]
  num_correct = np.sum(predicted_cls == cidx)
  cls_acc = 100.0 * float(num_correct) / cls_mask.sum()
 
  print(f"{cname} class accuracy {cls_acc} {num_correct} / {cls_mask.sum()}")

def get_prediction_descending_order_indices(probs, cidx):
    """
    Returns the ordering of probs that would sort it in descending order
    :param probs: (N, 2) probabilities (computed in TODO 2)
    :param cidx: class index (0 or 1)
    :return list of N indices that sorts the array in descending order
    """
    order = range(probs.shape[0]) # placeholder
    ############################ TODO 3 BEGIN #################################
    order = sorted(order, key = lambda x:probs[x][cidx], reverse=True)
    ############################ TODO 3 END #################################
    return order
 
#-----------
 
for cidx, cname in enumerate(classes):
  print(f"Predictions for class: {cname}")
  cls_mask = labels_np == cidx
  predicted_cls = predicted_class[cls_mask]
  cls_probs = dogfood_class_probs_np[cls_mask]
  correct = [p == cidx for p in predicted_cls]
 
  cls_images = [images[i] for i in range(len(images)) if cls_mask[i]]
 
  order = get_prediction_descending_order_indices(cls_probs, cidx)
 
  assert len(order) == cls_mask.sum()
 
  show_images(
    images = [cls_images[i] for i in order],
    correct_list = [correct[i] for i in order],
    titles = [f"Prob: {cls_probs[i, cidx]}" for i in order],
    size=128
  )

def compute_dscore_dimage(scores, image, class_idx):
    """
    Returns the gradient of s_y (the score at index class_idx) with respect to
    the image (data), ds_y / dI.  Note that this is the unnormalized class
    score "s", not the probability "p".
    :param scores: (Variable) shape (1000) the output scores from AlexNet for image
    :param image: (Variable) shape (1, 3, 224, 244) the input image
    :param class_idx: class index in range [0, 999] indicating which class to compute saliency for
    :return grad: (Tensor) shape (3, 224, 224), gradient ds_y / dI
    """
    grad = torch.zeros_like(image) # placeholder
    ############################ TODO 4 BEGIN #################################
    sc = scores[class_idx]
    sc.backward()
    grad = image.grad
    ############################ TODO 4 END #################################
    assert tuple(grad.shape) == (1, 3, 224, 224) # expected shape
    return grad[0]

model = alexnet(pretrained=True)
 
def visualize_saliency(image):
    image_in = Variable(image.unsqueeze(0), requires_grad=True)
    cls_scores = model(image_in)[0]
    max_score, max_idx = torch.max(cls_scores, 0)
 
    grad = compute_dscore_dimage(cls_scores, image_in, max_idx)
 
    vis = grad * grad
    vis, _ = torch.max(vis, 0)
 
    return vis
 
class_datasets = [dog_dataset, food_dataset]
 
num_images = 6
for cidx, cname in enumerate(classes):
 
  print(f"Saliency for class: {cname}")
  in_images = []
  vis_images = []
 
  for i, sample in enumerate(class_datasets[cidx]):
    image_i = sample[0]
    label_i = sample[1]
 
    vis_image = visualize_saliency(image_i)
 
    assert list(vis_image.shape) == [224, 224]
 
    vis_images.append(vis_image.unsqueeze(0))
    in_images.append(image_i)
    if i >= num_images:
      break
 
  row_list = list(zip(in_images, vis_images))
 
  show_image_rows(row_list)
 
  cls_mask = labels_np == cidx
  predicted_cls = predicted_class[cls_mask]
  cls_probs = dogfood_class_probs_np[cls_mask]

def normalized_sgd_with_momentum_update(image, grad, velocity, momentum, learning_rate):
    """
    :param image: (Variable) shape (1, 3, 224, 244) the current solution
    :param grad: (Variable) gradient of the loss with respect to the image
    :param velocity: (Variable) momentum vector "V"
    :param momentum: (float) momentum parameter "mu"
    :param learning_rate: (float) learning rate "alpha"
    :return: (Variable) the updated image and momentum vector (image, velocity)
    """
    ############################ TODO 5a BEGIN #################################
    velocity = momentum * velocity - learning_rate * (grad / np.linalg.norm(grad.data.numpy()))
    image = image + velocity
    ############################ TODO 5a BEGIN #################################
    return image, velocity
 
 
def fooling_image_gradient(target_score, orig_data, image_in, target_class, reg_lambda):
    """
    Compute the gradient for make_fooling_image (dL / dI).
    :param target_score: (Variable) holding the current score assigned to the target class
    :param orig_data: (Variable) shape (1, 3, 224, 224) holding the original image
    :param image_in: (Variable) shape (1, 3, 224, 224) hoding the current solution
    :param target_class: (int) ILSVRC class in range [0, 999]
    :param reg_lambda: (float) weight applied to the regularizer.
    :return grad: (Variable) gradient dL / dI
    """
    grad = torch.zeros_like(image_in) # placeholder
    ############################ TODO 5b BEGIN #################################
    norm = torch.norm(orig_data - image_in)
    loss = -target_score + 0.5 * reg_lambda * norm * norm
    loss.backward()
    grad = image_in.grad
    ############################ TODO 5b END #################################
    assert tuple(grad.shape) == (1, 3, 224, 224) # expected shape
    return grad

model = alexnet(pretrained=True)
 
# Set the model to eval mode
model.eval()
 
def make_fooling_image(image, target_class, learning_rate, regularization,
                       num_iter, momentum, threshold=0.9):
    """
    Fool AlexNet into thinking that any image has a particular class, by perturbing it just a little bit
 
    :param image: starting image CxHxW tensor
    :param target_class: the class that this will become after optimization
    :param learning_rate: either a constant, or a function that returns the learning rate at each iteration
    :param regularization: lambda parameter to multiply the regularizer
    :param num_iter: maximum number of iterations
    :param momentum: amount of momentum to use in the SGD 
    :param threshold: the target score for target_class
    """
    # Create batch dimension and turn into a variable
    image = image[np.newaxis, ...]
 
    print(f"Fooling AlexNet into thinking this is a: {labels_id2word[target_class]}")
 
    # This is the original image (used by the regularizer)
    orig_data = Variable(image.clone())
 
    image_in = Variable(image, requires_grad=True)
    velocity = torch.zeros_like(image_in)
 
    for i in range(num_iter):
        curr_scores = model(image_in)[0]
        curr_probs = F.softmax(curr_scores, 0)
 
        target_prob = curr_probs[target_class]
        target_score = curr_scores[target_class]
 
        # compute the gradient
        grad_wrt_image = fooling_image_gradient(
            target_score, orig_data, image_in, target_class, regularization)
 
        # update the image with the SGD rule
        image_in, velocity = normalized_sgd_with_momentum_update(
            image_in, grad_wrt_image, velocity, momentum, learning_rate)
 
        # Detach the image and velocity so that we don't backprop through
        # multiple iterations of the loop
        image_in = Variable(image_in.data, requires_grad = True)
        velocity = Variable(velocity.data, requires_grad = False)
 
        # Zero the gradients
        model.zero_grad()
 
        # Take the target probability out of the variable (turn it into float)
        target_prob = target_prob.data.item()
 
        # visualize the current state
        print(f"({i+1}/{num_iter}), {target_prob * 100} confidence")
 
        if target_prob > threshold:
            break
 
    delta = (image_in - orig_data).data
    return image_in, delta
 
 
num_images = 2
target_class=113
for cidx, cname in enumerate(classes):
    dataset = class_datasets[cidx]
 
    images_in = []
    fooling_images = []
    deltas = []
 
    for i, input in enumerate(dataset):
      image_in = input[0]
      label = input[1]
 
      fooling_image, delta = make_fooling_image(
          image_in,
          target_class=target_class,
          learning_rate=1e-1,
          regularization=5e-5,
          num_iter=100,
          momentum=0.9)
 
      delta = 0.5 + (5.0/255.0) * delta
 
      images_in.append(image_in)
      fooling_images.append(fooling_image)
      deltas.append(delta)
 
      if i >= num_images:
        break
 
    print ("\nLeft: original, middle: fooling image, right: difference magnified by 5x (gray is 0).\n"
        "AlexNet will classify the middle image in each row as %r with high confidence" % (
        labels_id2word[target_class]))
    show_image_rows(list(zip(images_in, fooling_images, deltas)))

def class_visualization_gradient(target_score, image, target_class, reg_lambda):
    """
    Compute the gradient for make_class_visualization (dL / dI).
    :param target_score: (Variable) holding the current score assigned to the target class
    :param image: (Variable) shape (1, 3, 224, 224) the current solution
    :param target_class: (int) ILSVRC class in range [0, 999]
    :param regularization: (float) weight (lambda) applied to the regularizer.
    :return grad: (Variable) gradient dL / dI
    """
    grad = torch.zeros_like(image) # placeholder
 
    ############################ TODO 6 BEGIN #################################
    norm = torch.norm(image)
    loss = -target_score + 0.5 * reg_lambda * norm * norm
    loss.backward()
    grad = image.grad
    ############################ TODO 6 END #################################
    assert tuple(grad.shape) == (1, 3, 224, 224) # expected shape
    return grad

model = alexnet(pretrained=True)
 
# Set the model to eval mode
model.eval()
 
def make_class_visualization(target_class, learning_rate, regularization, num_iter,
                             max_jitter, blur_sigma, momentum):
    """
    Visualize an ILSVRC2012 class by maximizing the probability of that class, starting from random.
 
    :param target_class: what ILSVRC2012 class to visualize
    :param learning_rate: either a constant, or a function that returns the learning rate at each iteration
    :param regularization: lambda parameter to multiply the regularizer
    :param num_iter: number of iterations
    :param max_jitter: amount of jitter to add for regularization
    :param blur_sigma: blur each iteration by this amount.
    :param momentum: amount of momentum to use in the SGD update
    """
 
    image_np = 100 * np.random.randn(1, 3, 224, 224).astype(np.float32)
    vel_np = np.zeros_like(image_np)
 
    print(f"Generating visualization of: {labels_id2word[target_class]}...")
 
    for i in range(num_iter):
        # Random jitter to regularize
        ox, oy = np.random.randint(-max_jitter, max_jitter+1, 2)
 
        image_np = np.roll(np.roll(image_np, ox, -1), oy, -2)
        vel_np = np.roll(np.roll(vel_np, ox, -1), oy, -2)
 
        image = Variable(torch.from_numpy(image_np), requires_grad=True)
        velocity = Variable(torch.from_numpy(vel_np), requires_grad=False)
 
        # Compute the current class score
        class_scores = model(image)[0]
        class_probs = F.softmax(class_scores, 0)
 
        target_score = class_scores[target_class]
        target_prob = class_probs[target_class]
 
        # Compute the gradient
        grad = class_visualization_gradient(
            target_score, image, target_class, regularization)
 
        # Normalized SGD+Momentum update
        image, velocity = normalized_sgd_with_momentum_update(
            image, grad, velocity, momentum, learning_rate)
 
        # Convert back to numpy
        image_np = image.data.cpu().numpy()
        vel_np = velocity.data.cpu().numpy()
 
        # Undo jitter
        image_np = np.roll(np.roll(image_np, -ox, -1), -oy, -2)
        vel_np = np.roll(np.roll(vel_np, -ox, -1), -oy, -2)
 
        # blur the image every iteration
        for c in range(3):
            image_np[0, c, ...] = gaussian_filter(image_np[0, c, ...], sigma=blur_sigma)
 
        # Convert to float
        target_prob = target_prob.data.item()
        target_score = target_score.data.item()
 
        # Visualize our current result
        print(f"({i+1}/{num_iter}), {target_prob * 100} % confidence, score: {target_score}")
 
    return torch.from_numpy(image_np)[0]