"""

We build our architecture on top of the ANs proposed in

@InProceedings{huang2018and,

  title={Unsupervised Deep Learning by Neighbourhood Discovery},

  author={Jiabo Huang, Qi Dong, Shaogang Gong and Xiatian Zhu},

  booktitle={Proceedings of the International Conference on machine learning (ICML)},

  year={2019},

}

The code is available online under https://github.com/Raymond-sci/AND

"""

from torch.autograd import Function

import os

import torchvision.models as models

import math

import torch

import torch.nn.functional as F

import torch.nn as nn

def resnet18(pretrained=False):

    model = models.resnet18(pretrained=pretrained)

    model.fc = Identity()

    model.avgpool = Identity()

    return model

def resnet34(pretrained=False):

    model = models.resnet34(pretrained=pretrained)

    model.fc = Identity()

    model.avgpool = Identity()

    return model

def resnet50(pretrained=False):

    model = models.resnet50(pretrained=pretrained)

    model.fc = Identity()

    model.avgpool = Identity()

    return model

class Identity(nn.Module):

    def __init__(self):

        super(Identity, self).__init__()

    def forward(self, x):

        return x

class Backbone(nn.Module):

    def __init__(self, name='resnet18', pretrained=True, freeze_all=False):

        super(Backbone, self).__init__()

        self.name = name

        self.freeze_all = freeze_all

        self.pretrained = pretrained

        if name == 'resnet18':

            self.backbone = resnet18(pretrained=self.pretrained)

        if name == 'resnet34':

            self.backbone = resnet34(pretrained=self.pretrained)  

        elif name == 'resnet50':

            self.backbone = resnet50(pretrained=self.pretrained)

        if self.freeze_all:

            # List all layers (even inside sequential module)

            layers = [module for module in self.backbone.modules() if type(module) != nn.Sequential]

            for layer in layers:

                if hasattr(layer, 'requires_grad_'):

                    layer.requires_grad_(False)

    def forward(self, x):

        return self.backbone(x)

class SimpleDecoder(nn.Module):

    def __init__(self, hidden_dimension=512):

        super(SimpleDecoder, self).__init__()

        self.conv_up_5 = nn.Conv2d(hidden_dimension, hidden_dimension//2, 3, padding=1)

        self.conv_up_4 = nn.Conv2d(hidden_dimension//2, hidden_dimension//4, 3, padding=1)

        self.conv_up_3 = nn.Conv2d(hidden_dimension//4, hidden_dimension//8, 3, padding=1)

        self.conv_up_2 = nn.Conv2d(hidden_dimension//8, hidden_dimension//16, 3, padding=1)

        self.conv_up_1 = nn.Conv2d(hidden_dimension//16, hidden_dimension//32, 5, padding=2)

        self.decoder = nn.Conv2d(hidden_dimension//32, 1, 5, padding=2)

    def forward(self, z):

        h = nn.ReLU()(self.conv_up_5(z))

        h = nn.Upsample(scale_factor=2, mode='bicubic', align_corners=False)(h)

        h = nn.ReLU()(self.conv_up_4(h))

        h = nn.Upsample(scale_factor=2, mode='bicubic', align_corners=False)(h)

        h = nn.ReLU()(self.conv_up_3(h))

        h = nn.Upsample(scale_factor=2, mode='bicubic', align_corners=False)(h)

        h = nn.ReLU()(self.conv_up_2(h))

        h = nn.Upsample(scale_factor=2, mode='bicubic', align_corners=False)(h)

        h = nn.ReLU()(self.conv_up_1(h))

        h = nn.Upsample(scale_factor=2, mode='bicubic', align_corners=False)(h)

        x_hat = nn.Sigmoid()(self.decoder(h))

        return x_hat

class CAE_DNR(nn.Module):

    def __init__(self, pretrained=True, n_channels=3, hidden_dimension=512, name = 'resnet18',npc_dimension = 256):

        super(CAE_DNR, self).__init__()

        self.n_channels = n_channels

        self.name = name

        self.encoder = Backbone(name= self.name, pretrained=pretrained, freeze_all=False)

        if self.n_channels != self.encoder.backbone.conv1.in_channels:

            conv1 = nn.Conv2d(2, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)

            data = self.encoder.backbone.conv1.weight.data[:, :2, :, :]  # Better than nothing ... ?

            self.encoder.backbone.conv1 = conv1

            self.encoder.backbone.conv1.weight.data = data

            self.fc = nn.Linear(hidden_dimension,npc_dimension)

        self.relu = nn.ReLU(inplace=True)

        self.hidden_dimension = hidden_dimension

        self.decoder = SimpleDecoder(hidden_dimension=hidden_dimension)

    def restore_model(self, paths):

        for attr, path in paths.items():

            self._load(attr=attr, path=path)

        return self

    def _load(self, attr, path):

        if not os.path.exists(path):

            print('Unknown path: {}'.format(path))

        if not hasattr(self, attr):

            print('No attribute: {}'.format(attr))

        self.__getattr__(attr).load_state_dict(torch.load(path), strict=True)

        return self

    def forward(self, x, decode=False):

        z = self.encode(x, pool=False)

        zb = nn.AvgPool2d(2, 2)(z).squeeze(dim=3).squeeze(dim=2)

        zp = self.fc(zb)

        zp = self.relu(zp)

        if decode:

            x_hat = self.decoder(z)

        else:

            x_hat = None

        return x_hat, zp, zb

    def encode(self, x, pool=False):

        h = self.encoder(x)

        h = h.view((x.shape[0], -1, 2, 2))

        if pool:

            return nn.AvgPool2d(2, 2)(h).squeeze(dim=3).squeeze(dim=2)

        else:

            return h

    def latent_variable(self, x_in, projectionHead):

        _, zp, zb = self.forward(x_in, decode=True)

        if projectionHead:

            return zp

        else:

            return zb

class NonParametricClassifierOP(Function):

    @staticmethod

    def forward(self, x, y, memory, params):

        T = params[0].item()

        # inner product

        out = torch.mm(x.data, memory.t())

        out.div_(T)  # batchSize * N

        self.save_for_backward(x, memory, y, params)

        return out

    @staticmethod

    def backward(self, gradOutput):

        x, memory, y, params = self.saved_tensors

        T = params[0].item()

        momentum = params[1].item()

        # add temperature

        gradOutput.data.div_(T)

        # gradient of linear

        gradInput = torch.mm(gradOutput.data, memory)

        gradInput.resize_as_(x)

        # update the memory

        weight_pos = memory.index_select(0, y.data.view(-1)).resize_as_(x)

        weight_pos.mul_(momentum)

        weight_pos.add_(torch.mul(x.data, 1 - momentum))

        w_norm = weight_pos.pow(2).sum(1, keepdim=True).pow(0.5)

        updated_weight = weight_pos.div(w_norm)

        memory.index_copy_(0, y, updated_weight)

        return gradInput, None, None, None, None

class NonParametricClassifier(nn.Module):

    """Non-parametric Classifier

    Non-parametric Classifier from

    "Unsupervised Feature Learning via Non-Parametric Instance Discrimination"

    Extends:

        nn.Module

    """

    def __init__(self, inputSize, outputSize, T=0.05, momentum=0.5):

        """Non-parametric Classifier initial functin

        Initial function for non-parametric classifier

        Arguments:

            inputSize {int} -- in-channels dims

            outputSize {int} -- out-channels dims

        Keyword Arguments:

            T {int} -- distribution temperate (default: {0.05})

            momentum {int} -- memory update momentum (default: {0.5})

        """

        super(NonParametricClassifier, self).__init__()

        self.nLem = outputSize

        self.register_buffer('params', torch.tensor([T, momentum]))

        stdv = 1. / math.sqrt(inputSize / 3)

        self.register_buffer('memory', torch.rand(outputSize, inputSize).mul_(2 * stdv).add_(-stdv))

    def forward(self, x, y):

        out = NonParametricClassifierOP.apply(x, y, self.memory, self.params)

        return out

class ANsDiscovery(nn.Module):

    """Discovery ANs

    Discovery ANs according to current round, select_rate and most importantly,

    all sample's corresponding entropy

    """

    def __init__(self, nsamples):

        """Object used to discovery ANs

        Discovery ANs according to the total amount of samples, ANs selection

        rate, ANs size

        Arguments:

            nsamples {int} -- total number of sampels

            select_rate {float} -- ANs selection rate

            ans_size {int} -- ANs size

        Keyword Arguments:

            device {str} -- [description] (default: {'cpu'})

        """

        super(ANsDiscovery, self).__init__()

        # not going to use ``register_buffer'' as

        # they are determined by configs

        self.select_rate = 0.25

        self.ANs_size = 1

        # number of samples

        self.register_buffer('samples_num', torch.tensor(nsamples))

        # indexes list of anchor samples

        self.register_buffer('anchor_indexes', torch.LongTensor(nsamples//2))

        # indexes list of instance samples

        self.register_buffer('instance_indexes', torch.arange(nsamples//2).long())

        # anchor samples' and instance samples' position

        self.register_buffer('position', -1 * torch.arange(nsamples).long() - 1)

        # anchor samples' neighbours

        self.register_buffer('neighbours', torch.LongTensor(nsamples//2, 1))

        # each sample's entropy

        self.register_buffer('entropy', torch.FloatTensor(nsamples))

        # consistency

        self.register_buffer('consistency', torch.tensor(0.))

    def get_ANs_num(self, round):

        """Get number of ANs

        Get number of ANs at target round according to the select rate

        Arguments:

            round {int} -- target round

        Returns:

            int -- number of ANs

        """

        return int(self.samples_num.float() * self.select_rate * round)

    def update(self, round, npc, cheat_labels=None):

        """Update ANs

        Discovery new ANs and update `anchor_indexes`, `instance_indexes` and

        `neighbours`

        Arguments:

            round {int} -- target round

            npc {Module} -- non-parametric classifier

            cheat_labels {list} -- used to compute consistency of chosen ANs only

        Returns:

            number -- [updated consistency]

        """

        with torch.no_grad():

            batch_size = 100

            ANs_num = self.get_ANs_num(round)

            features = npc.memory

            for start in range(0, self.samples_num, batch_size):

                end = start + batch_size

                end = min(end, self.samples_num)

                preds = F.softmax(npc(features[start:end], None), 1)

                self.entropy[start:end] = -(preds * preds.log()).sum(1)

            # get the anchor list and instance list according to the computed

            # entropy

            self.anchor_indexes = self.entropy.topk(ANs_num, largest=False)[1]

            self.instance_indexes = (torch.ones_like(self.position)

                                     .scatter_(0, self.anchor_indexes, 0)

                                     .nonzero().view(-1))

            anchor_entropy = self.entropy.index_select(0, self.anchor_indexes)

            instance_entropy = self.entropy.index_select(0, self.instance_indexes)

            # get position

            # if the anchor sample x whose index is i while position is j, then

            # sample x_i is the j-th anchor sample at current round

            # if the instance sample x whose index is i while position is j, then

            # sample x_i is the (-j-1)-th instance sample at current round

            instance_cnt = 0

            for i in range(self.samples_num):

                # for anchor samples

                if (i == self.anchor_indexes).any():

                    self.position[i] = (self.anchor_indexes == i).max(0)[1]

                    continue

                # for instance samples

                instance_cnt -= 1

                self.position[i] = instance_cnt

            anchor_features = features.index_select(0, self.anchor_indexes)

            self.neighbours = (torch.LongTensor(ANs_num, self.ANs_size)

                               .to('cuda'))

            for start in range(0, ANs_num, batch_size):

                end = start + batch_size

                end = min(end, ANs_num)

                sims = torch.mm(anchor_features[start:end], features.t())

                sims.scatter_(1, self.anchor_indexes[start:end].view(-1, 1), -1.)

                _, self.neighbours[start:end] = (

                    sims.topk(self.ANs_size, largest=True, dim=1))

            # if cheat labels is provided, then compute consistency

            if cheat_labels is None:

                return 0.

            anchor_label = cheat_labels.index_select(0, self.anchor_indexes)

            neighbour_label = cheat_labels.index_select(0,

                                                        self.neighbours.view(-1)).view_as(self.neighbours)

            self.consistency = ((anchor_label.view(-1, 1) == neighbour_label)

                                .float().mean())

            return self.consistency

class Criterion(nn.Module):

    def __init__(self):

        super(Criterion, self).__init__()

    def calculate_loss(self, x, y, ANs):

        batch_size, _ = x.shape

        # split anchor and instance list

        anchor_indexes, instance_indexes = self._split(y[:batch_size//2], ANs)

        preds = F.softmax(x, 1)

        l_ans = torch.tensor(0).cuda()

        if anchor_indexes.size(0) > 0:

            # compute loss for anchor samples

            y_ans = y.index_select(0, anchor_indexes)

            y_ans_p = y.index_select(0, anchor_indexes + batch_size//2)

            y_ans_neighbour = ANs.position.index_select(0, y_ans)

            neighbours = ANs.neighbours.index_select(0, y_ans_neighbour)

            # p_i = \sum_{j \in \Omega_i} p_{i,j}

            x_ans = preds.index_select(0, anchor_indexes)

            x_ans_p = preds.index_select(0, anchor_indexes + batch_size//2)

            x_ans_neighbour = x_ans.gather(1, neighbours).sum(1)

            x_ans_p = x_ans_p.gather(1, y_ans_p.view(-1, 1)).view(-1)

            x_ans = x_ans.gather(1, y_ans.view(-1, 1)).view(-1)

            # sum all terms : self + sim + neighbors

            # NLL: l = -log(p_i)

            l_ans = -1 * torch.log(x_ans + x_ans_p + x_ans_neighbour).sum(0)

        l_inst = torch.tensor(0).cuda()

        if instance_indexes.size(0) > 0:

            # compute loss for instance samples

            y_inst = y.index_select(0, instance_indexes)

            y_inst_p = y.index_select(0, instance_indexes + batch_size//2)

            x_inst = preds.index_select(0, instance_indexes)

            x_inst_p = preds.index_select(0, instance_indexes + batch_size//2)

            # p_i = p_{i, i}

            x_inst = x_inst.gather(1, y_inst.view(-1, 1))

            x_inst_p = x_inst_p.gather(1, y_inst_p.view(-1, 1))

            # NLL: l = -log(p_i)

            l_inst = -1 * torch.log(x_inst + x_inst_p).sum(0)

        return l_inst / batch_size, l_ans / batch_size

    def _split(self, y, ANs):

        pos = ANs.position.index_select(0, y.view(-1))

        return (pos >= 0).nonzero().view(-1), (pos < 0).nonzero().view(-1)

    def forward(self, x_out, index, npc, ANs_discovery, x_hat, zp):

        z_n = torch.div(zp, torch.norm(zp+1e-12, p=2, dim=1, keepdim=True))

        outputs = npc(z_n, index)  # For each image get similarity with neighbour

        loss_inst, loss_ans = self.calculate_loss(outputs, index, ANs_discovery)

        loss = loss_inst + loss_ans

        l_loss = {'loss': loss, 'loss_inst': loss_inst, 'loss_ans': loss_ans}

        if x_hat is not None:

            loss_mse = nn.MSELoss()(x_hat, x_out)

            loss = loss + loss_mse

            l_loss['loss_mse'] = loss_mse

            l_loss['loss'] = loss

        return l_loss