# -*- coding: utf-8 -*-
"""Copy of my_first_few_shot_classifier.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1HkLGEN5BILDlpF186Ufh-B729ex8Xnpc

# Your own few-shot classification model ready in 15mn with PyTorch

[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/sicara/easy-few-shot-learning/blob/master/notebooks/my_first_few_shot_classifier.ipynb)

I have been working on few-shot classification for a while now. The more I talk about it, the more the people around me seem to feel that it's some kind of dark magic. Even sadder: I noticed that very few actually used it on their projects. I think that's too bad, so I decided to make a tutorial so you'll have no excuse to deprive yourself of the power of few-shot learning methods.

In 15 minutes and just a few lines of code, we are going to implement
the [Prototypical Networks](https://arxiv.org/abs/1703.05175). It's the favorite method of
many few-shot learning researchers (~2000 citations in 3 years), because 1) it works well,
and 2) it's incredibly easy to grasp and to implement.

## Discovering Prototypical Networks
First, let's install the [tutorial GitHub repository](https://github.com/sicara/easy-few-shot-learning) and import some packages. If you're on Colab right now, you should also check that you're using a GPU (Edit > Notebook settings).
"""


import torch
from torch import nn, optim
from torch.utils.data import DataLoader
from torchvision import transforms
from torchvision.datasets import Flowers102
from torchvision.models import resnet18
from tqdm import tqdm
import matplotlib.pyplot as plt
import numpy as np
from torch.optim.lr_scheduler import MultiStepLR

from easyfsl.samplers import TaskSampler
from easyfsl.utils import plot_images, sliding_average
from torch import nn, Tensor
from easyfsl.methods import FewShotClassifier
from easyfsl.modules.predesigned_modules import (
    default_matching_networks_support_encoder,
    default_matching_networks_query_encoder,
)
import random
from easyfsl.methods import FewShotClassifier
from easyfsl.modules.predesigned_modules import (
    default_matching_networks_support_encoder,
    default_matching_networks_query_encoder,
)


from torch.utils.data import DataLoader
from torchvision.datasets import Flowers102
from torchvision import transforms

batch_size = 64
n_workers = 12
image_size = 224
train_set = Flowers102(
    root="./data",
    split = 'train',
    transform=transforms.Compose(
        [
            transforms.RandomResizedCrop(image_size),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
        ]
    ),
    download=True,
)

val_set = Flowers102(
    root="./data",
    split = 'val',
    transform=transforms.Compose(
        [
            transforms.RandomResizedCrop(image_size),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
        ]
    ),
    download=True,
)



test_set = Flowers102(
    root="./data",
    split = 'test',
    transform=transforms.Compose(
        [
            transforms.RandomResizedCrop(image_size),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
        ]
    ),
    download=True,
)

''''
Modified prototypical network to train a siamese version of the embedder network
'''

class PrototypicalNetworks(nn.Module):
    def __init__(self, backbone: nn.Module):
        super(PrototypicalNetworks, self).__init__()
        self.backbone = backbone
 
    def return_a_p_n( self,
        support_images: torch.Tensor,
        support_labels: torch.Tensor,
        query_images: torch.Tensor,
        query_labels: torch.Tensor):
        '''
        This function outputs the embedding for the anchor, positive and negative images.
        The function picks an anchor label at random from the support set. The positive and negative images
        are picked up from the query set.
        The function is called during model.fit for triplet loss calculation
        '''
      
        z_support = self.backbone.forward(support_images)
        z_query = self.backbone.forward(query_images)
        label = random.choice(support_labels)
      
        anchor = torch.cat([z_support[torch.nonzero(support_labels == label)]])
        anchor = torch.reshape(anchor, (anchor.shape[0],anchor.shape[-1]))
        
        positive = torch.cat([z_query[torch.nonzero(query_labels == label)]])
        positive = torch.reshape(positive, (positive.shape[0],positive.shape[-1]))
        negative = torch.cat([z_query[torch.nonzero(query_labels != label)]])
 
        negative = torch.reshape(negative, (negative.shape[0],negative.shape[-1]))
        #print(negative.shape)
 
        return anchor,positive,negative
 
    def forward(
        self,
        support_images: torch.Tensor,
        support_labels: torch.Tensor,
        query_images: torch.Tensor,
    ) -> torch.Tensor:
        """
        Predict query labels using labeled support images.
        """
        # Extract the features of support and query images
        z_support = self.backbone.forward(support_images)
        z_query = self.backbone.forward(query_images)
 
 
 
        # Infer the number of different classes from the labels of the support set
        n_way = len(torch.unique(support_labels))
        # Prototype i is the mean of all instances of features corresponding to labels == i
        z_proto = torch.cat(
            [
                z_support[torch.nonzero(support_labels == label)].mean(0)
                for label in range(n_way)
            ]
        )
        # Compute the euclidean distance from queries to prototypes
        dists = torch.cdist(z_query, z_proto)
 
        # And here is the super complicated operation to transform those distances into classification scores!
        scores = -dists
        return scores
convolutional_network = resnet18(pretrained=True)
convolutional_network.fc =  nn.Sequential(
    nn.Flatten(),nn.Linear(512,512)).to('cuda')
 
print(convolutional_network)
 
model = PrototypicalNetworks(convolutional_network).cuda()


N_WAY = 5 # Number of classes in a task
N_SHOT = 4  # Number of images per class in the support set
N_QUERY = 5  # Number of images per class in the query set
N_EVALUATION_TASKS = 500

# The sampler needs a dataset with a "get_labels" method. Check the code if you have any doubt!
test_set.get_labels = lambda: [instance[1] for instance in test_set]
test_sampler = TaskSampler(
    test_set, n_way=N_WAY, n_shot=N_SHOT, n_query=N_QUERY, n_tasks=N_EVALUATION_TASKS
)

test_loader = DataLoader(
    test_set,
    batch_sampler=test_sampler,
    num_workers=12,
    pin_memory=True,
    collate_fn=test_sampler.episodic_collate_fn,
)

"""We created a dataloader that will feed us with 5-way 5-shot tasks (the most common setting in the litterature).
Now, as every data scientist should do before launching opaque training scripts,
let's take a look at our dataset.
"""

(
    example_support_images,
    example_support_labels,
    example_query_images,
    example_query_labels,
    example_class_ids,
) = next(iter(test_loader))

#plot_images(example_support_images, "support images", images_per_row=N_SHOT)
#plot_images(example_query_images, "query images", images_per_row=N_QUERY)

"""For both support and query set, you should have one line for each class.

How does our model perform on this task?
"""

model.eval()
example_scores = model(
    example_support_images.cuda(),
    example_support_labels.cuda(),
    example_query_images.cuda()
).detach()

_, example_predicted_labels = torch.max(example_scores.data, 1)

# print("Ground Truth / Predicted")
# for i in range(len(example_query_labels)):
#     print(
#         f"{example_scores} / {example_predicted_labels}"
#     )

"""This doesn't look bad: keep in mind that the model was trained on very different images, and has only seen 5 examples for each class!

Now that we have a first idea, let's see more precisely how good our model is.
"""

def evaluate_on_one_task(
    support_images: torch.Tensor,
    support_labels: torch.Tensor,
    query_images: torch.Tensor,
    query_labels: torch.Tensor,
):
    """
    Returns the number of correct predictions of query labels, and the total number of predictions.
    """
    model.eval()
    return (
        torch.max(
            model(support_images.cuda(), support_labels.cuda(),query_images.cuda())
            .detach()
            .data,
            1,
        )[1]
        == query_labels.cuda()
    ).sum().item(), len(query_labels)


def evaluate(data_loader: DataLoader):
    # We'll count everything and compute the ratio at the end
    total_predictions = 0
    correct_predictions = 0

    # eval mode affects the behaviour of some layers (such as batch normalization or dropout)
    # no_grad() tells torch not to keep in memory the whole computational graph (it's more lightweight this way)
    model.eval()
    with torch.no_grad():
        for episode_index, (
            support_images,
            support_labels,
            query_images,
            query_labels,
            class_ids,
        ) in tqdm(enumerate(data_loader), total=len(data_loader)):

            correct, total = evaluate_on_one_task(
                support_images, support_labels, query_images, query_labels
            )

            total_predictions += total
            correct_predictions += correct

    print(
        f"Model tested on {len(data_loader)} tasks. Accuracy: {(100 * correct_predictions/total_predictions):.2f}%"
    )
    return correct_predictions / total_predictions

evaluate(test_loader)

N_TRAINING_EPISODES_PER_EPOCH = 100
N_VALIDATION_TASKS = 100

train_set.get_labels = lambda: [instance[1] for instance in train_set]
train_sampler = TaskSampler(
    train_set, n_way=2, n_shot=5, n_query=5, n_tasks=N_TRAINING_EPISODES_PER_EPOCH
)
train_loader = DataLoader(
    train_set,
    batch_sampler=train_sampler,
    num_workers=12,
    pin_memory=True,
    collate_fn=train_sampler.episodic_collate_fn,
)

val_set.get_labels = lambda: [instance[1] for instance in val_set]
val_sampler = TaskSampler(
    val_set, n_way=2, n_shot=5, n_query=5, n_tasks=N_VALIDATION_TASKS
)
val_loader = DataLoader(
    val_set,
    batch_sampler=val_sampler,
    num_workers=12,
    pin_memory=True,
    collate_fn=val_sampler.episodic_collate_fn,
)

criterion = nn.TripletMarginLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
n_epochs =40
scheduler_milestones = [120, 160]
scheduler_gamma = 0.1
learning_rate = 1e-2
train_scheduler = MultiStepLR(
    optimizer,
    milestones=scheduler_milestones,
    gamma=scheduler_gamma,
)

def fit(
    support_images: torch.Tensor,
    support_labels: torch.Tensor,
    query_images: torch.Tensor,
    query_labels: torch.Tensor,
) -> float:
    optimizer.zero_grad()
    # model.process_support_set()
    a,p,n = model.return_a_p_n(
        support_images.cuda(), support_labels.cuda(),
        query_images.cuda(), query_labels.cuda()
    )
 
    loss = criterion(a.cuda(),p.cuda(),n.cuda())
    loss.backward()
    optimizer.step()
 
    return loss.item()

"""To train the model, we are just going to iterate over a large number of randomly generated few-shot classification tasks,
and let the `fit` method update our model after each task. This is called **episodic training**.
"""

# Train the model yourself with this cell
def training_epoch(
    model: FewShotClassifier, data_loader: DataLoader):
    log_update_frequency = 10

    all_loss = []
    model.train()
    with tqdm(enumerate(data_loader), total=len(data_loader), desc="Training") as tqdm_train:
        for episode_index, (
                support_images,
                support_labels,
                query_images,
                query_labels,
                _,
            ) in tqdm_train:
            loss_value = fit(support_images, support_labels, query_images, query_labels)
            #loss_value is a float
            all_loss.append(loss_value)

            if episode_index % log_update_frequency == 0:
                tqdm_train.set_postfix(loss=sliding_average(all_loss, log_update_frequency))
    return np.mean(all_loss)

def validate_epoch(
    model: FewShotClassifier, data_loader: DataLoader):
    val_loss=0.0  
    val_correct=0.0  
    val_loss_history=[]  
    val_correct_history=[]  
    model.eval()                                              
    
    with torch.no_grad(): 
        with tqdm(enumerate(data_loader), total=len(data_loader), desc="Validating") as tqdm_train:
            for episode_index, (
                    support_images,
                    support_labels,
                    query_images,
                    query_labels,
                    _,
                ) in tqdm_train:
                classification_scores = model(support_images.cuda(), support_labels.cuda(),query_images.cuda()
                )

                a,p,n = model.return_a_p_n(
                    support_images.cuda(), support_labels.cuda(),
                    query_images.cuda(), query_labels.cuda()
                )
            
                loss_value = criterion(a.cuda(),n.cuda(),p.cuda())
                
                _,val_preds=torch.max(classification_scores ,1)  
                val_loss+=loss_value.item()  
                val_correct+=torch.sum(val_preds==query_labels.cuda().data) 
                val_epoch_loss=val_loss/len(data_loader)  
                val_epoch_acc=val_correct.float()/len(data_loader)  
                val_loss_history.append(val_epoch_loss)  
                val_correct_history.append(val_epoch_acc)  

    return val_epoch_loss


"""testing and Validation!"""
val_results = []
train_loss_results = []
val_loss_results = []
train_results = []

best_state = model.state_dict()
best_validation_accuracy = 0.0
for epoch in range(n_epochs):
    print(f"Epoch {epoch}")
    average_loss = training_epoch(model, train_loader)
    train_loss_results.append(average_loss)
    average_val_loss = validate_epoch(model, val_loader)
    val_loss_results.append(average_val_loss)

    validation_accuracy = evaluate(val_loader)
    val_results.append(validation_accuracy)
    train_accuracy = evaluate(train_loader)
    train_results.append(train_accuracy)

    if validation_accuracy > best_validation_accuracy:
        best_validation_accuracy = validation_accuracy
        best_state = model.state_dict()
        print(f"We found a new best model! with val accuracy {best_validation_accuracy}")

    # Warn the scheduler that we did an epoch
    # so it knows when to decrease the learning rate
    train_scheduler.step()


"""Testing"""

evaluate(test_loader)
