"""
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_', True, 'import torch.nn as nn\n'), (33, 'torch.nn.init.kaiming_normal_', 'nn.init.kaiming_normal_', True, 'import torch.nn as nn\n'), (34, 'torch.nn.init.constant_', 'nn.init.constant_', True, 'import torch.nn as nn\n'), (76, 'torch.nn.functional.normalize', 'F.normalize', True, 'import torch.nn.functional as F\n'), (86, 'torch.nn.functional.normalize', 'F.normalize', True, 'import torch.nn.functional as F\n'), (111, 'torch.cat', 'torch.cat', False, 'import torch\n'), (165, 'models.inception.BasicConv2d', 'BasicConv2d', False, 'from models.inception import inception_v3, BasicConv2d\n'), (171, 'torch.nn.Linear', 'nn.Linear', True, 'import torch.nn as nn\n'), (28, 'torch.nn.init.constant_', 'nn.init.constant_', True, 'import torch.nn as nn\n'), (36, 'torch.nn.init.kaiming_normal_', 'nn.init.kaiming_normal_', True, 'import torch.nn as nn\n'), (52, 'torch.nn.AdaptiveMaxPool2d', 'nn.AdaptiveMaxPool2d', True, 'import torch.nn as nn\n'), (60, 'torch.nn.functional.upsample_bilinear', 'F.upsample_bilinear', True, 'import torch.nn.functional as F\n'), (70, 'torch.cat', 'torch.cat', False, 'import torch\n'), (73, 'torch.sign', 'torch.sign', False, 'import torch\n'), (81, 'torch.ones_like', 'torch.ones_like', False, 'import torch\n'), (84, 'torch.sign', 'torch.sign', False, 'import torch\n'), (102, 'torch.nonzero', 'torch.nonzero', False, 'import torch\n'), (124, 'torch.cat', 'torch.cat', False, 'import torch\n'), (217, 'torch.stack', 'torch.stack', False, 'import torch\n'), (219, 'torch.mean', 'torch.mean', False, 'import torch\n'), (38, 'torch.nn.init.constant_', 'nn.init.constant_', True, 'import torch.nn as nn\n'), (101, 'torch.nn.functional.upsample_bilinear', 'F.upsample_bilinear', True, 'import torch.nn.functional as F\n'), (109, 'torch.nn.functional.upsample_bilinear', 'F.upsample_bilinear', True, 'import torch.nn.functional as F\n'), (214, 'torch.nn.functional.normalize', 'F.normalize', True, 'import torch.nn.functional as F\n'), (41, 'torch.nn.init.constant_', 'nn.init.constant_', True, 'import torch.nn as nn\n'), (42, 'torch.nn.init.constant_', 'nn.init.constant_', True, 'import torch.nn as nn\n'), (73, 'torch.abs', 'torch.abs', False, 'import torch\n'), (79, 'torch.zeros_like', 'torch.zeros_like', False, 'import torch\n'), (82, 'torch.einsum', 'torch.einsum', False, 'import torch\n'), (84, 'torch.abs', 'torch.abs', False, 'import torch\n'), (97, 'random.uniform', 'random.uniform', False, 'import random\n'), (64, 'torch.einsum', 'torch.einsum', False, 'import torch\n'), (119, 'random.uniform', 'random.uniform', False, 'import random\n'), (123, 'torch.nn.functional.upsample_bilinear', 'F.upsample_bilinear', True, 'import torch.nn.functional as F\n'), (141, 'models.inception.inception_v3', 'inception_v3', False, 'from models.inception import inception_v3, BasicConv2d\n'), (144, 'models.inception.inception_v3', 'inception_v3', False, 'from models.inception import inception_v3, BasicConv2d\n')]