import torch

import torch.nn as nn

import numpy as np

import torchvision

import torch.nn.functional as F

import math

import copy

import collections

from pytorchcv.model_provider import get_model as ptcv_get_model

from pytorchcv.models.common import conv3x3_block

import pretrainedmodels

class Flatten(nn.Module):

    def forward(self, input):

        return input.view(input.size(0), -1)

def l2_norm(input,axis=1):

    norm = torch.norm(input,2,axis,True)

    output = torch.div(input, norm)

    return output 

class Window(nn.Module):

    def forward(self, x):

        return torch.clamp(x,0,1)

class ArcMarginProduct(nn.Module):

    r"""Implement of large margin arc distance: :

        Args:

            in_features: size of each input sample

            out_features: size of each output sample

            s: norm of input feature

            m: margin

            cos(theta + m)

        """

    def __init__(self, in_features, out_features,weights=None):

        super(ArcMarginProduct, self).__init__()

        if weights is None:

            self.weight = nn.Parameter(torch.FloatTensor(out_features, in_features))

            self.reset_parameters()

        else:

            self.weight = nn.Parameter(weights)

    def reset_parameters(self):

        stdv = 1. / math.sqrt(self.weight.size(1))

        self.weight.data.uniform_(-stdv, stdv)

#        self.k.data=torch.ones(1,dtype=torch.float)

    def forward(self, features):

        cosine = F.linear(l2_norm(features), l2_norm(self.weight))

        return cosine

class ArcClassifier(nn.Module):

    def __init__(self,in_features, out_features,weights=None):

        super(ArcClassifier, self).__init__()

        self.classifier = ArcMarginProduct(in_features, out_features,weights=weights)

        self.dropout1=nn.Dropout(p=0.5, inplace=True)

    def forward(self, x,eq):

        out = self.dropout1(x-eq)

        out = self.classifier(out)

        return out

    def no_grad(self):

        for param in self.parameters():

            param.requires_grad=False

    def do_grad(self):

        for param in self.parameters():

            param.requires_grad=True

class MyDenseNet(nn.Module):

    def __init__(self,model,

                 num_classes,

                 num_channels=1,

                 strategy='copy',

                 add_noise=0.,

                 dropout=0.5,

                 arcface=False,

                 return_features=False,

                 norm=False,

                 intermediate=0,

                 extra_pool=1,

                 pool_type='max',

                 wso=None,

                 dont_do_grad=['wso'],

                 do_bn=False):

        super(MyDenseNet, self).__init__()

        self.features= torch.nn.Sequential()

        self.num_channels=num_channels

        self.dont_do_grad=dont_do_grad

        self.pool_type=pool_type

        self.norm=norm

        self.return_features=return_features

        self.num_classes=num_classes

        self.extra_pool=extra_pool

        if wso is not None:

            conv_ = nn.Conv2d(1,self.num_channels, kernel_size=(1, 1))

            if hasattr(wso, '__iter__'):

                conv_.weight.data.copy_(torch.tensor([[[[1./wso[0][1]]]],[[[1./wso[1][1]]]],[[[1./wso[2][1]]]]]))

                conv_.bias.data.copy_(torch.tensor([0.5 - wso[0][0]/wso[0][1],

                                                    0.5 - wso[1][0]/wso[1][1],

                                                    0.5 -wso[2][0]/wso[2][1]]))

            self.features.add_module('wso_conv',conv_)

            self.features.add_module('wso_window',nn.Sigmoid())

            if do_bn:

                self.features.add_module('wso_norm',nn.BatchNorm2d(self.num_channels))

            else:

                self.features.add_module('wso_norm',nn.InstanceNorm2d(self.num_channels))

        if (strategy == 'copy') or (num_channels!=3):

            base = list(list(model.children())[0].named_children())[1:]

            conv0 = model.state_dict()['features.conv0.weight']

            new_conv=nn.Conv2d(self.num_channels, conv0.shape[0], kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)

            a=(np.arange(3*(self.num_channels//3+1),dtype=np.int)%3)

            np.random.shuffle(a)

            for i in range(self.num_channels):

                new_conv.state_dict()['weight'][:,i,:,:]=conv0.clone()[:,a[i],:,:]*(1.0+torch.randn_like(conv0[:,a[i],:,:])*add_noise)

            self.features.add_module('conv0',new_conv)

        else:

            base = list(list(model.children())[0].named_children())

        for (n,l) in base:

            self.features.add_module(n,l)

        if intermediate==0:

            self.num_features=list(model.children())[-1].in_features

            self.intermediate=None

        else:

            self.num_features=intermediate

            self.intermediate=nn.Linear(list(model.children())[-1].in_features, self.num_features)

        self.dropout1=nn.Dropout(p=dropout, inplace=True)

        if arcface:

            self.classifier=ArcMarginProduct(self.num_features, num_classes)

        else:

            self.classifier = nn.Linear(self.num_features//self.extra_pool, self.num_classes)

    def forward(self, x):

        x = self.features(x)

        x = F.relu(x, inplace=True)

        if self.pool_type=='avg':

            x = F.avg_pool3d(x.unsqueeze(1), kernel_size=(self.extra_pool,)+x.size()[2:]).view(x.size(0), -1)

        else:

            x = F.max_pool3d(x.unsqueeze(1), kernel_size=(self.extra_pool,)+x.size()[2:]).view(x.size(0), -1)

#        x = F.max_pool1d(x.view(x.unsqueeze(1),self.extra_pool).squeeze()

        x = self.dropout1(x)

        if self.intermediate is not None:

            x = self.intermediate(x)

            x = F.relu(x)

        features = x

        if self.norm:

            features = l2_norm(features,axis=1)

        out = self.classifier(features)

        return out if not self.return_features else (out,features)

    def parameter_scheduler(self,epoch):

        do_first=['classifier','wso']

        if epoch>0:

            for n,p in self.named_parameters():

                p.requires_grad=True

        else:

            for n,p in self.named_parameters():

                p.requires_grad= any(nd in n for nd in do_first)

    def no_grad(self):

        for param in self.parameters():

            param.requires_grad=False

    def do_grad(self):

        for n,p in self.named_parameters():

            p.requires_grad=  not any(nd in n for nd in self.dont_do_grad)

    def get_optimizer_parameters(self,klr):

        zero_layer=['conv0','norm0','ws_norm']

        param_optimizer = list(self.named_parameters())

        num_blocks=4

        no_decay=['bias']

        optimizer_grouped_parameters=[

            {'params': 

             [p for n, p in param_optimizer if (not any(nd in n for nd in no_decay) and any(nd in n for nd in zero_layer))], 

                                                                                                     'lr':klr*2e-5,'weight_decay': 0.01},

            {'params': 

             [p for n, p in param_optimizer if any(nd in n for nd in no_decay)  and any(nd in n for nd in zero_layer)], 

                                                                                                     'lr':klr*2e-5, 'weight_decay': 0.0}]

        optimizer_grouped_parameters.extend([

            {'params': [p for n, p in param_optimizer if (not any(nd in n for nd in no_decay) and ('wso' in n))], 

                                                                                                     'lr':klr*1e-5,'weight_decay': 0.01},

            {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)  and ('wso' in n)],

                                                                                                     'lr':klr*1e-5, 'weight_decay': 0.0}])

        optimizer_grouped_parameters.extend([

            {'params': [p for n, p in param_optimizer if (not any(nd in n for nd in no_decay) and ('classifier' in n))], 

                                                                                                       'lr':klr*1e-3,'weight_decay': 0.01},

            {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)  and ('classifier' in n)], 

                                                                                                     'lr':klr*1e-3, 'weight_decay': 0.0}])

        for i in range(num_blocks):

            optimizer_grouped_parameters.extend([

                {'params': [p for n, p in param_optimizer if (not any(nd in n for nd in no_decay) and ('denseblock{}'.format(i+1) in n))], 

                                                                                             'lr':klr*(2.0**i)*2e-5,'weight_decay': 0.01},

                {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)  and ('denseblock{}'.format(i+1) in n)], 

                                                                                             'lr':klr*(2.0**i)*2e-5, 'weight_decay': 0.0}])

        optimizer_grouped_parameters.extend([

            {'params': [p for n, p in param_optimizer if (not any(nd in n for nd in no_decay) and ('norm5' in n))],

                                                                                             'lr':klr*1e-4,'weight_decay': 0.01},

            {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)  and ('norm5' in n)], 

                                                                                              'lr':klr*1e-4, 'weight_decay': 0.0}])

        return(optimizer_grouped_parameters)

class MySENet(nn.Module):

    def __init__(self,model,

                 num_classes,

                 num_channels=3,

                 dropout=0.2,

                 return_features=False,

                 wso=None,

                 full_copy=False,

                 dont_do_grad=['wso'],

                 extra_pool=1,

                 do_bn=False):

        super(MySENet, self).__init__()

        self.num_classes=num_classes

        self.return_features=return_features

        self.num_channels = num_channels

        self.features= torch.nn.Sequential()

        self.extra_pool=extra_pool

        self.dont_do_grad=dont_do_grad

        if full_copy:

            for (n,l) in list(list(model.children())[0].named_children()):

                self.features.add_module(n,l)

            if wso is not None:

                self.dont_do_grad=model.dont_do_grad

        else:

            if wso is not None:

                conv_ = nn.Conv2d(1,self.num_channels, kernel_size=(1, 1))

                if hasattr(wso, '__iter__'):

                    conv_.weight.data.copy_(torch.tensor([[[[1./wso[0][1]]]],[[[1./wso[1][1]]]],[[[1./wso[2][1]]]]]))

                    conv_.bias.data.copy_(torch.tensor([0.5 - wso[0][0]/wso[0][1],

                                                        0.5 - wso[1][0]/wso[1][1],

                                                        0.5 -wso[2][0]/wso[2][1]]))

                self.features.add_module('wso_conv',conv_)

                self.features.add_module('wso_relu',nn.Sigmoid())

                if do_bn:

                    self.features.add_module('wso_norm',nn.BatchNorm2d(self.num_channels))

                else:

                    self.features.add_module('wso_norm',nn.InstanceNorm2d(self.num_channels))

#            layer0= torch.nn.Sequential()

#            layer0.add_module('conv1',model.conv1)

#            layer0.add_module('bn1',model.bn1)                        

            se_layers={'layer0':model.layer0,

                       'layer1':model.layer1,

                       'layer2':model.layer2,

                       'layer3':model.layer3,

                       'layer4':model.layer4}

            for key in se_layers:

                self.features.add_module(key,se_layers[key])

        self.dropout = dropout if dropout is None else nn.Dropout(p=dropout, inplace=True)

        self.classifier=nn.Linear(model.last_linear.in_features//self.extra_pool, self.num_classes)

    def forward(self, x):

        x = self.features(x)

        x = F.max_pool3d(x.unsqueeze(1), kernel_size=(self.extra_pool,)+x.size()[2:]).view(x.size(0), -1)

        if self.dropout is not None:

            x = self.dropout(x) 

        features = x

        out = self.classifier(features)

        return out if not self.return_features else (out,features) 

    def parameter_scheduler(self,epoch):

        do_first=['classifier']

        if epoch>0:

            for n,p in self.named_parameters():

                p.requires_grad=True

        else:

            for n,p in self.named_parameters():

                p.requires_grad= any(nd in n for nd in do_first)

    def no_grad(self):

        for param in self.parameters():

            param.requires_grad=False

    def do_grad(self):

        for n,p in self.named_parameters():

            p.requires_grad=  not any(nd in n for nd in self.dont_do_grad)

    def get_optimizer_parameters(self,klr):

        param_optimizer = list(self.named_parameters())

        num_blocks=5

        no_decay=['bias']

        optimizer_grouped_parameters=[

            {'params': [p for n, p in param_optimizer if (not any(nd in n for nd in no_decay) and ('classifier' in n))],

                                                                                                     'lr':klr*2e-4,'weight_decay': 0.01},

            {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)  and ('classifier' in n)], 

                                                                                                     'lr':klr*2e-4, 'weight_decay': 0.0}]

        optimizer_grouped_parameters.extend([

            {'params': [p for n, p in param_optimizer if (not any(nd in n for nd in no_decay) and ('wso' in n))], 

                                                                                                     'lr':klr*5e-6,'weight_decay': 0.01},

            {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)  and ('wso' in n)],

                                                                                                     'lr':klr*5e-6, 'weight_decay': 0.0}])

        for i in range(num_blocks):

            optimizer_grouped_parameters.extend([

            {'params': [p for n, p in param_optimizer if (not any(nd in n for nd in no_decay) and ('layer{}'.format(i) in n))],

                                                                                             'lr':klr*(2.0**i)*1e-5,'weight_decay': 0.01},

            {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)  and ('layer{}'.format(i) in n)],

                                                                                             'lr':klr*(2.0**i)*1e-5, 'weight_decay': 0.0}])

        return(optimizer_grouped_parameters)

class MyEfficientNet(nn.Module):

    def __init__(self,model,num_classes,num_channels=3,dropout=0.5,return_features=False,wso=True,

                 full_copy=False,

                 dont_do_grad=['wso'],

                 extra_pool=1,

                 num_features=None):

        super(MyEfficientNet, self).__init__()

        self.num_classes=num_classes

        self.return_features=return_features

        self.num_channels = num_channels

        self.features= torch.nn.Sequential()

        self.extra_pool=extra_pool

        self.dont_do_grad=dont_do_grad

        if full_copy:

            for (n,l) in list(list(model.children())[0].named_children()):

                self.features.add_module(n,l)

            if wso is not None:

                self.dont_do_grad=model.dont_do_grad

        else:

            if wso is not None:

                conv_ = nn.Conv2d(1,self.num_channels, kernel_size=(1, 1))

                if hasattr(wso, '__iter__'):

                    conv_.weight.data.copy_(torch.tensor([[[[1./wso[0][1]]]],[[[1./wso[1][1]]]],[[[1./wso[2][1]]]]]))

                    conv_.bias.data.copy_(torch.tensor([0.5 - wso[0][0]/wso[0][1],

                                                        0.5 - wso[1][0]/wso[1][1],

                                                        0.5 -wso[2][0]/wso[2][1]]))

                self.features.add_module('wso_conv',conv_)

                self.features.add_module('wso_relu',nn.Sigmoid())

                self.features.add_module('wso_norm',nn.InstanceNorm2d(self.num_channels))

            for (n,l) in list(list(model.children())[0].named_children()):

                self.features.add_module(n,l)

        self.dropout = dropout if dropout is None else nn.Dropout(p=dropout, inplace=True)

        if num_features is None:

            self.classifier=nn.Linear(model.output.fc.in_features//self.extra_pool, self.num_classes)

        else:

            self.classifier=nn.Linear(num_features, self.num_classes)

    def forward(self, x):

        x = self.features(x)

        x = F.avg_pool2d(x, kernel_size=x.size(-1)).view(x.size(0), -1)

        if self.extra_pool>1:

            x = x.view(x.shape[0],x.shape[1]//self.extra_pool,self.extra_pool).mean(-1)

        if self.dropout is not None:

            x = self.dropout(x)

        features = x

        out = self.classifier(features)

        return out if not self.return_features else (out,features) 

    def parameter_scheduler(self,epoch):

        do_first=['classifier']

        if epoch>0:

            for n,p in self.named_parameters():

                p.requires_grad=True

        else:

            for n,p in self.named_parameters():

                p.requires_grad= any(nd in n for nd in do_first)

    def no_grad(self):

        for param in self.parameters():

            param.requires_grad=False

    def do_grad(self):

        for n,p in self.named_parameters():

            p.requires_grad=  not any(nd in n for nd in self.dont_do_grad)

class NeighborsNet(nn.Module):

    def __init__(self,num_classes,num_features=1024,num_neighbors=1,classifier_layer=None,intermidiate=None,dropout=0.2):

        super(NeighborsNet, self).__init__()

        self.num_classes=num_classes

        if classifier_layer is not None:

            self.num_features = classifier_layer.in_features

        else:

            self.num_features=num_features

        self.num_neighbors=num_neighbors

        layers=collections.OrderedDict()

        if dropout>0:

            layers['dropout']=nn.Dropout(p=dropout)

        if intermidiate is not None:

            layers['intermidiate']=nn.Linear(self.num_features*(2*self.num_neighbors+1), intermidiate)

            layers['relu']=nn.ReLU()

            layers['classifier']=nn.Linear(intermidiate, self.num_classes)

        else:

            layers['classifier']=nn.Linear(self.num_features*(2*self.num_neighbors+1), self.num_classes)

        if (classifier_layer is not None) and (intermidiate is None):

            _=layers['classifier'].bias.data.copy_((1.0+0.2*self.num_neighbors)*classifier_layer.bias.data)

            d = torch.cat([0.1*classifier_layer.weight.data for i in range(self.num_neighbors)]+\

                             [classifier_layer.weight.data]+\

                             [0.1*classifier_layer.weight.data for i in range(self.num_neighbors)],dim=1)

            _=layers['classifier'].weight.data.copy_(d)

        self.network= torch.nn.Sequential(layers)

    def forward(self, x):

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

        return self.network(x) 

    def parameter_scheduler(self,epoch):

        do_first=['classifier']

        if epoch>0:

            for n,p in self.named_parameters():

                p.requires_grad=True

        else:

            for n,p in self.named_parameters():

                p.requires_grad= any(nd in n for nd in do_first)

    def no_grad(self):

        for param in self.parameters():

            param.requires_grad=False

    def do_grad(self):

        for param in self.parameters():

            param.requires_grad=True

class ResModelPool(nn.Module):

    def __init__(self,in_size):

        super(ResModelPool, self).__init__()

        self.dont_do_grad=[]

        self.conv2d1=torch.nn.Conv2d(1, 64, (9,in_size),stride=(1,in_size), padding=(4,0))

        self.bn0=torch.nn.BatchNorm1d(64)

#        self.relu0=torch.nn.ReLU()

        self.conv1d1=torch.nn.Conv1d(64, 64, 7, padding=3)

        self.bn1=torch.nn.BatchNorm1d(64)

        self.relu1=torch.nn.ReLU()

        self.conv1d2=torch.nn.Conv1d(128, 64, 5, padding=2)

        self.bn2=torch.nn.BatchNorm1d(64)

        self.relu2=torch.nn.ReLU()

        self.conv1d3=torch.nn.Conv1d(192, 6, 3, padding=1)

    def forward(self, x):

        x=x.unsqueeze(1)

        x = self.conv2d1(x)

        x=F.max_pool2d(x,kernel_size=(1,x.shape[-1])).squeeze(-1)        

        x0 = self.bn0(x)

#        x0 = self.relu0(x)

        x = self.conv1d1(x0)

        x = self.bn1(x)

        x1 = self.relu1(x)

        x = torch.cat([x0,x1],1)

        x = self.conv1d2(x)

        x = self.bn2(x)

        x2 = self.relu2(x)

        x = torch.cat([x0,x1,x2],1)

        out = self.conv1d3(x).transpose(-1,-2)

        return out 

    def no_grad(self):

        for param in self.parameters():

            param.requires_grad=False

    def do_grad(self):

        for n,p in self.named_parameters():

            p.requires_grad=  not any(nd in n for nd in self.dont_do_grad)

def mean_model(models):

    model = copy.deepcopy(models[0])

    params=[]

    for model_ in models:

        params.append(dict(model_.named_parameters()))

    param_dict=dict(model.named_parameters())

    for name in param_dict.keys():

        _=param_dict[name].data.copy_(torch.cat([param[name].data[...,None] for param in params],-1).mean(-1))

    return model