# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/03_model.ipynb (unless otherwise specified).

__all__ = ['sift_frame_sim', 'simclr_frame_sim', 'SimCLRDataset', 'imagenet_normalize_transform',
           'get_train_transforms', 'get_val_transforms', 'NTXEntCriterion', 'SimCLRModel']

# Cell
import cv2
import torch
import torchvision

import numpy as np
import torch.nn as nn

from pytorch_lightning import Trainer
from pytorch_lightning.core.lightning import LightningModule
from .features import *
from torch.autograd import Variable
from torch.utils.data import DataLoader, Dataset
from torchvision import datasets, models, transforms
from sklearn.metrics.pairwise import cosine_similarity

# Cell
def sift_frame_sim(features_i, features_j, codebook, df, vw):
    vws_i = np.expand_dims(codebook.predict(features_i), axis=0)
    bowv_i = np.expand_dims(np.histogram(vws_i, bins = range(vw + 1))[0], axis=0)
    tf_idf_i = calc_tf_idf(bowv_i, df)

    vws_j = np.expand_dims(codebook.predict(features_j), axis=0)
    bowv_j = np.expand_dims(np.histogram(vws_j, bins = range(vw + 1))[0], axis=0)
    tf_idf_j = calc_tf_idf(bowv_j, df)

    sim = np.dot(tf_idf_i, tf_idf_j) / (np.linalg.norm(tf_idf_i) * np.linalg.norm(tf_idf_j))
    return sim

# Cell
def simclr_frame_sim(features_i, features_j, codebook, df, vw):
    sim = cosine_similarity(features_i, features_j)[0][0]
    return sim

# Cell

"""
All code from: https://github.com/dthiagarajan/simclr_pytorch
Only included directory for easy loading of checkpoints
"""
class SimCLRDataset(Dataset):
    def __init__(self, dataset):
        """Initialize a wrapper of a generic image classification dataset for SimCLR training.

        Args:
            dataset (torch.utils.data.Dataset): an image PyTorch dataset - when iterating over it
                it should return something of the form (image) or (image, label).
        """
        self.dataset = dataset

    def __getitem__(self, index):
        dataset_item = self.dataset[index]
        if type(dataset_item) is tuple:
            image = dataset_item[0]
        else:
            image = dataset_item
        return image, image

    def __len__(self):
        return len(self.dataset)

    @staticmethod
    def mixup(x, alpha=0.4):
        batch_size = x.size()[0] // 2
        if alpha > 0:
            lam = np.random.beta(alpha, alpha, batch_size)
            lam = np.concatenate(
                [lam[:, None], 1 - lam[:, None]], 1
            ).max(1)[:, None, None, None]
            lam = torch.from_numpy(lam).float()
            if torch.cuda.is_available():
                lam = lam.cuda()
        else:
            lam = 1.
        # This is SimCLR specific - we want to use the same mixing for the augmented pairs
        lam = torch.cat([lam, lam])
        index = torch.randperm(batch_size)
        # This is SimCLR specific - we want to use the same permutation on the augmented pairs
        index = torch.cat([index, batch_size + index])
        if torch.cuda.is_available():
            index = index.cuda()
        mixed_x = lam * x + (1 - lam) * x[index, :]

        return mixed_x, lam


def imagenet_normalize_transform():
    return transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])


def get_train_transforms(size=224, color_jitter_prob=0.8, grayscale_prob=0.2):
    color_jitter = transforms.ColorJitter(0.8, 0.8, 0.8, 0.2)
    return transforms.Compose([
        transforms.RandomResizedCrop(size=(size, size)),
        transforms.RandomHorizontalFlip(),
        transforms.RandomApply([color_jitter], p=color_jitter_prob),
        transforms.RandomGrayscale(p=grayscale_prob),
        transforms.ToTensor(),
        imagenet_normalize_transform()
    ])


def get_val_transforms(size=224):
    return transforms.Compose([
        transforms.Resize(size=(size, size)),
        transforms.ToTensor(),
        imagenet_normalize_transform()
    ])


import torch.nn as nn


class NTXEntCriterion(nn.Module):
    """Normalized, temperature-scaled cross-entropy criterion, as suggested in the SimCLR paper.

    Parameters:
        temperature (float, optional): temperature to scale the confidences. Defaults to 0.5.
    """
    criterion = nn.CrossEntropyLoss(reduction="sum")
    similarity = nn.CosineSimilarity(dim=2)

    def __init__(self, temperature=0.5):
        super(NTXEntCriterion, self).__init__()
        self.temperature = temperature
        self.batch_size = None
        self.mask = None

    def mask_correlated_samples(self, batch_size):
        """Masks examples in a batch and it's augmented pair for computing the valid summands for
            the criterion.

        Args:
            batch_size (int): batch size of the individual batch (not including it's augmented pair)

        Returns:
            torch.Tensor: a mask (tensor of 0s and 1s), where 1s indicates a pair of examples in a
                batch that will contribute to the overall batch loss
        """
        mask = torch.ones((batch_size * 2, batch_size * 2), dtype=bool)
        mask = mask.fill_diagonal_(0)
        for i in range(batch_size):
            mask[i, batch_size + i] = 0
            mask[batch_size + i, i] = 0
        return mask

    def compute_similarities(self, z_i, z_j, temperature):
        """Computes the similarities between two projections `z_i` and `z_j`, scaling based on
            `temperature`.

        Args:
            z_i (torch.Tensor): projection of a batch
            z_j (torch.Tensor): projection of the augmented pair for the batch
            temperature (float): temperature to scale the similarity by

        Returns:
            torch.Tensor: tensor of similarities for the positive and negative pairs
        """
        batch_size = len(z_i)
        mask = self.mask_correlated_samples(batch_size)

        p1 = torch.cat((z_i, z_j), dim=0)
        sim = self.similarity(p1.unsqueeze(1), p1.unsqueeze(0)) / temperature

        sim_i_j = torch.diag(sim, batch_size)
        sim_j_i = torch.diag(sim, -batch_size)

        positive_samples = torch.cat((sim_i_j, sim_j_i), dim=0).reshape(
            batch_size * 2, 1
        )
        negative_samples = sim[mask].reshape(batch_size * 2, -1)

        logits = torch.cat((positive_samples, negative_samples), dim=1)
        return logits

    def forward(self, z):
        """Computes the loss for a batch and its augmented pair.

        Args:
            z (torch.Tensor): tensor of a batch and it's augmented pair, concatenated

        Returns:
            torch.Tensor: loss for the given batch
        """
        double_batch_size = len(z)
        batch_size = double_batch_size // 2
        z_i, z_j = z[:double_batch_size // 2], z[double_batch_size // 2:]
        if self.batch_size is None or batch_size != self.batch_size:
            self.batch_size = batch_size
            self.mask = None

        if self.mask is None:
            self.mask = self.mask_correlated_samples(self.batch_size)

        logits = self.compute_similarities(z_i, z_j, self.temperature)
        labels = torch.zeros(self.batch_size * 2).long()
        logits, labels = logits.to(z.device), labels.to(z.device)
        loss = self.criterion(logits, labels)
        loss /= 2 * self.batch_size
        return loss

class SimCLRModel(LightningModule):
    """SimCLR training network for a generic torchvision model (restricted to `allowed_models`). """

    allowed_models = ['resnet18', 'resnet34', 'resnet50', 'resnet101', 'resnet152']
    allowed_datasets = ['CIFAR10', 'CIFAR100', 'STL10', 'SVHN']

    def __init__(
        self, model_name='resnet18', pretrained=True, projection_dim=64, temperature=0.5,
        batch_size=128, image_size=224, save_hparams=True
    ):
        super().__init__()
        layers = list(getattr(torchvision.models, model_name)(pretrained=pretrained).children())
        self.model = nn.Sequential(*layers[:-1])
        self.projection_head = nn.Linear(layers[-1].in_features, projection_dim)
        self.loss = NTXEntCriterion(temperature=temperature)
        self.batch_size = batch_size
        self.image_size = image_size
        self.prepare_data()

    def forward(self, x):
        out = self.model(x)
        out = out.view(x.size(0), -1)
        out = self.projection_head(out)
        return out

    def training_step(self, batch, batch_idx):
        projections = self(batch)
        loss = self.loss(projections)
        tensorboard_logs = {'train_loss': loss}
        self.logger.scalar('loss', loss)
        return {'loss': loss, 'log': tensorboard_logs}

    def training_epoch_end(self, outputs):
        loss_mean = torch.stack([x['loss'] for x in outputs]).mean()
        return {'train_loss': loss_mean}

    def configure_optimizers(self):
        return torch.optim.Adam([
            {'params': self.model.parameters(), 'lr': 0.00001},
            {'params': self.projection_head.parameters(), 'lr': 0.001}
        ])

    def prepare_data(self):
        pass

    def collate_fn(self, batch):
        return torch.cat([torch.stack([b[0] for b in batch]), torch.stack([b[1] for b in batch])])

    def train_dataloader(self):
        return DataLoader(
            self.train_dataset, batch_size=self.batch_size, num_workers=64, shuffle=True,
            collate_fn=self.collate_fn
        )