"""
WS-DAN models
Hu et al.,
"See Better Before Looking Closer: Weakly Supervised Data Augmentation Network for Fine-Grained Visual Classification",
arXiv:1901.09891
"""
import logging
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

import models.resnet as resnet
from models.inception import inception_v3, BasicConv2d
import models.coatnet as coatnet

import random

__all__ = ['WSDAN_CAL']
EPSILON = 1e-6


def weights_init_classifier(m):
    classname = m.__class__.__name__
    if classname.find('Linear') != -1:
        nn.init.normal_(m.weight, std=0.001)
        if m.bias:
            nn.init.constant_(m.bias, 0.0)

def weights_init_kaiming(m):
    classname = m.__class__.__name__
    if classname.find('Linear') != -1:
        nn.init.kaiming_normal_(m.weight, a=0, mode='fan_out')
        nn.init.constant_(m.bias, 0.0)
    elif classname.find('Conv') != -1:
        nn.init.kaiming_normal_(m.weight, a=0, mode='fan_in')
        if m.bias is not None:
            nn.init.constant_(m.bias, 0.0)
    elif classname.find('BatchNorm') != -1:
        if m.affine:
            nn.init.constant_(m.weight, 1.0)
            nn.init.constant_(m.bias, 0.0)

# Bilinear Attention Pooling
class BAP(nn.Module):
    def __init__(self, pool='GAP'):
        super(BAP, self).__init__()
        assert pool in ['GAP', 'GMP']
        if pool == 'GAP':
            self.pool = None
        else:
            self.pool = nn.AdaptiveMaxPool2d(1)

    def forward(self, features, attentions):
        B, C, H, W = features.size()
        _, M, AH, AW = attentions.size()

        # match size
        if AH != H or AW != W:
            attentions = F.upsample_bilinear(attentions, size=(H, W))

        # feature_matrix: (B, M, C) -> (B, M * C)
        if self.pool is None:
            feature_matrix = (torch.einsum('imjk,injk->imn', (attentions, features)) / float(H * W)).view(B, -1)
        else:
            feature_matrix = []
            for i in range(M):
                AiF = self.pool(features * attentions[:, i:i + 1, ...]).view(B, -1)
                feature_matrix.append(AiF)
            feature_matrix = torch.cat(feature_matrix, dim=1)

        # sign-sqrt
        feature_matrix_raw = torch.sign(feature_matrix) * torch.sqrt(torch.abs(feature_matrix) + EPSILON)

        # l2 normalization along dimension M and C
        feature_matrix = F.normalize(feature_matrix_raw, dim=-1)

        if self.training:
            fake_att = torch.zeros_like(attentions).uniform_(0, 2)
        else:
            fake_att = torch.ones_like(attentions)
        counterfactual_feature = (torch.einsum('imjk,injk->imn', (fake_att, features)) / float(H * W)).view(B, -1)

        counterfactual_feature = torch.sign(counterfactual_feature) * torch.sqrt(torch.abs(counterfactual_feature) + EPSILON)

        counterfactual_feature = F.normalize(counterfactual_feature, dim=-1)
        return feature_matrix, counterfactual_feature

def batch_augment(images, attention_map, mode='crop', theta=0.5, padding_ratio=0.1):
    batches, _, imgH, imgW = images.size()

    if mode == 'crop':
        crop_images = []
        for batch_index in range(batches):
            atten_map = attention_map[batch_index:batch_index + 1]
            if isinstance(theta, tuple):
                theta_c = random.uniform(*theta) * atten_map.max()
            else:
                theta_c = theta * atten_map.max()

            crop_mask = F.upsample_bilinear(atten_map, size=(imgH, imgW)) >= theta_c
            nonzero_indices = torch.nonzero(crop_mask[0, 0, ...])
            height_min = max(int(nonzero_indices[:, 0].min().item() - padding_ratio * imgH), 0)
            height_max = min(int(nonzero_indices[:, 0].max().item() + padding_ratio * imgH), imgH)
            width_min = max(int(nonzero_indices[:, 1].min().item() - padding_ratio * imgW), 0)
            width_max = min(int(nonzero_indices[:, 1].max().item() + padding_ratio * imgW), imgW)

            crop_images.append(
                F.upsample_bilinear(images[batch_index:batch_index + 1, :, height_min:height_max, width_min:width_max],
                                    size=(imgH, imgW)))
        crop_images = torch.cat(crop_images, dim=0)
        return crop_images

    elif mode == 'drop':
        drop_masks = []
        for batch_index in range(batches):
            atten_map = attention_map[batch_index:batch_index + 1]
            if isinstance(theta, tuple):
                theta_d = random.uniform(*theta) * atten_map.max()
            else:
                theta_d = theta * atten_map.max()

            drop_masks.append(F.upsample_bilinear(atten_map, size=(imgH, imgW)) < theta_d)
        drop_masks = torch.cat(drop_masks, dim=0)
        drop_images = images * drop_masks.float()
        return drop_images

    else:
        raise ValueError('Expected mode in [\'crop\', \'drop\'], but received unsupported augmentation method %s' % mode)

class WSDAN_CAL(nn.Module):
    def __init__(self, num_classes, M=32, net='inception_mixed_6e', pretrained=False):
        super(WSDAN_CAL, self).__init__()
        self.num_classes = num_classes
        self.M = M
        self.net = net

        # Network Initialization
        if 'inception' in net:
            if net == 'inception_mixed_6e':
                self.features = inception_v3(pretrained=pretrained).get_features_mixed_6e()
                self.num_features = 768
            elif net == 'inception_mixed_7c':
                self.features = inception_v3(pretrained=pretrained).get_features_mixed_7c()
                self.num_features = 2048
            else:
                raise ValueError('Unsupported net: %s' % net)
        elif 'resnet' in net:
            self.features = getattr(resnet, net)(pretrained=pretrained).get_features()
            self.num_features = 512 * self.features[-1][-1].expansion
        elif 'coat' in net:
            self.features = getattr(coatnet, net)().get_features()
            if '0' in net or '1' in net:
                self.num_features = 768
            elif '2' in net:
                self.num_features = 1026
            elif '3' in net or '4' in net:
                self.num_features = 1536
            else:
                raise ValueError('Not given valid CoAtNet size.')
        else:
            raise ValueError('Unsupported net: %s' % net)

        # Attention Maps
        self.attentions = BasicConv2d(self.num_features, self.M, kernel_size=1)

        # Bilinear Attention Pooling
        self.bap = BAP(pool='GAP')

        # Classification Layer
        self.fc = nn.Linear(self.M * self.num_features, self.num_classes, bias=False)

        logging.info('WSDAN: using {} as feature extractor, num_classes: {}, num_attentions: {}'.format(net, self.num_classes, self.M))

    def visualize(self, x):
        batch_size = x.size(0)

        # Feature Maps, Attention Maps and Feature Matrix
        feature_maps = self.features(x)
        if self.net != 'inception_mixed_7c':
            attention_maps = self.attentions(feature_maps)
        else:
            attention_maps = feature_maps[:, :self.M, ...]
        # print(feature_maps.shape)
        # print(attention_maps.shape)

        feature_matrix = self.bap(feature_maps, attention_maps)[0]
        p = self.fc(feature_matrix * 100.)

        return p, attention_maps

    def forward(self, x):
        batch_size = x.size(0)

        # Feature Maps, Attention Maps and Feature Matrix
        feature_maps = self.features(x)
        if self.net != 'inception_mixed_7c':
            attention_maps = self.attentions(feature_maps)
        else:
            attention_maps = feature_maps[:, :self.M, ...]

        
        feature_matrix, feature_matrix_hat = self.bap(feature_maps, attention_maps)

        # Classification
        p = self.fc(feature_matrix * 100.)

        # Generate Attention Map
        if self.training:
            # Randomly choose one of attention maps Ak
            attention_map = []
            for i in range(batch_size):
                attention_weights = torch.sqrt(attention_maps[i].sum(dim=(1, 2)).detach() + EPSILON)
                attention_weights = F.normalize(attention_weights, p=1, dim=0)
                k_index = np.random.choice(self.M, 2, p=attention_weights.cpu().numpy())
                attention_map.append(attention_maps[i, k_index, ...])
            attention_map = torch.stack(attention_map)  # (B, 2, H, W) - one for cropping, the other for dropping
        else:
            attention_map = torch.mean(attention_maps, dim=1, keepdim=True)  # (B, 1, H, W)
        
        return p, p - self.fc(feature_matrix_hat * 100.), feature_matrix, attention_map

    def load_state_dict(self, state_dict, strict=True):
        model_dict = self.state_dict()
        pretrained_dict = {k: v for k, v in state_dict.items()
                           if k in model_dict and model_dict[k].size() == v.size()}

        if len(pretrained_dict) == len(state_dict):
            print('%s: All params loaded' % type(self).__name__)
        else:
            print('%s: Some params were not loaded:' % type(self).__name__)
            not_loaded_keys = [k for k in state_dict.keys() if k not in pretrained_dict.keys()]
            print(('%s, ' * (len(not_loaded_keys) - 1) + '%s') % tuple(not_loaded_keys))

        model_dict.update(pretrained_dict)
        super(WSDAN_CAL, self).load_state_dict(model_dict)

[(26, 'torch.nn.init.normal_', 'nn.init.normal_', (['m.weight'], {'std': '(0.001)'}), True, 'import torch.nn as nn\n'), (33, 'torch.nn.init.kaiming_normal_', 'nn.init.kaiming_normal_', (['m.weight'], {'a': '(0)', 'mode': '"""fan_out"""'}), True, 'import torch.nn as nn\n'), (34, 'torch.nn.init.constant_', 'nn.init.constant_', (['m.bias', '(0.0)'], {}), True, 'import torch.nn as nn\n'), (76, 'torch.nn.functional.normalize', 'F.normalize', (['feature_matrix_raw'], {'dim': '(-1)'}), True, 'import torch.nn.functional as F\n'), (86, 'torch.nn.functional.normalize', 'F.normalize', (['counterfactual_feature'], {'dim': '(-1)'}), True, 'import torch.nn.functional as F\n'), (111, 'torch.cat', 'torch.cat', (['crop_images'], {'dim': '(0)'}), False, 'import torch\n'), (165, 'models.inception.BasicConv2d', 'BasicConv2d', (['self.num_features', 'self.M'], {'kernel_size': '(1)'}), False, 'from models.inception import inception_v3, BasicConv2d\n'), (171, 'torch.nn.Linear', 'nn.Linear', (['(self.M * self.num_features)', 'self.num_classes'], {'bias': '(False)'}), True, 'import torch.nn as nn\n'), (28, 'torch.nn.init.constant_', 'nn.init.constant_', (['m.bias', '(0.0)'], {}), True, 'import torch.nn as nn\n'), (36, 'torch.nn.init.kaiming_normal_', 'nn.init.kaiming_normal_', (['m.weight'], {'a': '(0)', 'mode': '"""fan_in"""'}), True, 'import torch.nn as nn\n'), (52, 'torch.nn.AdaptiveMaxPool2d', 'nn.AdaptiveMaxPool2d', (['(1)'], {}), True, 'import torch.nn as nn\n'), (60, 'torch.nn.functional.upsample_bilinear', 'F.upsample_bilinear', (['attentions'], {'size': '(H, W)'}), True, 'import torch.nn.functional as F\n'), (70, 'torch.cat', 'torch.cat', (['feature_matrix'], {'dim': '(1)'}), False, 'import torch\n'), (73, 'torch.sign', 'torch.sign', (['feature_matrix'], {}), False, 'import torch\n'), (81, 'torch.ones_like', 'torch.ones_like', (['attentions'], {}), False, 'import torch\n'), (84, 'torch.sign', 'torch.sign', (['counterfactual_feature'], {}), False, 'import torch\n'), (102, 'torch.nonzero', 'torch.nonzero', (['crop_mask[0, 0, ...]'], {}), False, 'import torch\n'), (124, 'torch.cat', 'torch.cat', (['drop_masks'], {'dim': '(0)'}), False, 'import torch\n'), (217, 'torch.stack', 'torch.stack', (['attention_map'], {}), False, 'import torch\n'), (219, 'torch.mean', 'torch.mean', (['attention_maps'], {'dim': '(1)', 'keepdim': '(True)'}), False, 'import torch\n'), (38, 'torch.nn.init.constant_', 'nn.init.constant_', (['m.bias', '(0.0)'], {}), True, 'import torch.nn as nn\n'), (101, 'torch.nn.functional.upsample_bilinear', 'F.upsample_bilinear', (['atten_map'], {'size': '(imgH, imgW)'}), True, 'import torch.nn.functional as F\n'), (109, 'torch.nn.functional.upsample_bilinear', 'F.upsample_bilinear', (['images[batch_index:batch_index + 1, :, height_min:height_max, width_min:\n    width_max]'], {'size': '(imgH, imgW)'}), True, 'import torch.nn.functional as F\n'), (214, 'torch.nn.functional.normalize', 'F.normalize', (['attention_weights'], {'p': '(1)', 'dim': '(0)'}), True, 'import torch.nn.functional as F\n'), (41, 'torch.nn.init.constant_', 'nn.init.constant_', (['m.weight', '(1.0)'], {}), True, 'import torch.nn as nn\n'), (42, 'torch.nn.init.constant_', 'nn.init.constant_', (['m.bias', '(0.0)'], {}), True, 'import torch.nn as nn\n'), (73, 'torch.abs', 'torch.abs', (['feature_matrix'], {}), False, 'import torch\n'), (79, 'torch.zeros_like', 'torch.zeros_like', (['attentions'], {}), False, 'import torch\n'), (82, 'torch.einsum', 'torch.einsum', (['"""imjk,injk->imn"""', '(fake_att, features)'], {}), False, 'import torch\n'), (84, 'torch.abs', 'torch.abs', (['counterfactual_feature'], {}), False, 'import torch\n'), (97, 'random.uniform', 'random.uniform', (['*theta'], {}), False, 'import random\n'), (64, 'torch.einsum', 'torch.einsum', (['"""imjk,injk->imn"""', '(attentions, features)'], {}), False, 'import torch\n'), (119, 'random.uniform', 'random.uniform', (['*theta'], {}), False, 'import random\n'), (123, 'torch.nn.functional.upsample_bilinear', 'F.upsample_bilinear', (['atten_map'], {'size': '(imgH, imgW)'}), True, 'import torch.nn.functional as F\n'), (141, 'models.inception.inception_v3', 'inception_v3', ([], {'pretrained': 'pretrained'}), False, 'from models.inception import inception_v3, BasicConv2d\n'), (144, 'models.inception.inception_v3', 'inception_v3', ([], {'pretrained': 'pretrained'}), False, 'from models.inception import inception_v3, BasicConv2d\n')]