import torch
import torch.nn as nn
from torch.nn import init
import torch.nn.functional as F
from torch.optim import lr_scheduler
import torchvision
from torchvision.models import resnet34

import functools
from einops import rearrange

import numpy as np

import models
from models.help_funcs import Transformer, TransformerDecoder, TwoLayerConv2d
from torch.nn.modules.padding import ReplicationPad2d
from models.ChangeFormer import ChangeFormerV1, ChangeFormerV6
###############################################################################
# Helper Functions
###############################################################################

def get_scheduler(optimizer, args):
    """Return a learning rate scheduler

    Parameters:
        optimizer          -- the optimizer of the network
        args (option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions．
                              opt.lr_policy is the name of learning rate policy: linear | step | plateau | cosine

    For 'linear', we keep the same learning rate for the first <opt.niter> epochs
    and linearly decay the rate to zero over the next <opt.niter_decay> epochs.
    For other schedulers (step, plateau, and cosine), we use the default PyTorch schedulers.
    See https://pytorch.org/docs/stable/optim.html for more details.
    """
    if args.lr_policy == 'linear':
        def lambda_rule(epoch):
            lr_l = 1.0 - epoch / float(args.max_epochs + 1)
            return lr_l
        scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda_rule)
    elif args.lr_policy == 'step':
        step_size = args.max_epochs//3
        # args.lr_decay_iters
        scheduler = lr_scheduler.StepLR(optimizer, step_size=step_size, gamma=0.1)
    elif args.lr_policy == 'multistep':
        scheduler = lr_scheduler.MultiStepLR(optimizer,
        milestones=[2, 4, 7, 11, 15, 25, 35, 47, 60, 70, 90, 110, 130, 150, 170, 180, 190], gamma=0.5)
    else:
        return NotImplementedError('learning rate policy [%s] is not implemented', args.lr_policy)
    return scheduler

class Identity(nn.Module):
    def forward(self, x):
        return x

def get_norm_layer(norm_type='instance'):
    """Return a normalization layer

    Parameters:
        norm_type (str) -- the name of the normalization layer: batch | instance | none

    For BatchNorm, we use learnable affine parameters and track running statistics (mean/stddev).
    For InstanceNorm, we do not use learnable affine parameters. We do not track running statistics.
    """
    if norm_type == 'batch':
        norm_layer = functools.partial(nn.BatchNorm2d, affine=True, track_running_stats=True)
    elif norm_type == 'instance':
        norm_layer = functools.partial(nn.InstanceNorm2d, affine=False, track_running_stats=False)
    elif norm_type == 'none':
        norm_layer = lambda x: Identity()
    else:
        raise NotImplementedError('normalization layer [%s] is not found' % norm_type)
    return norm_layer

def init_weights(net, init_type='normal', init_gain=0.02):
    """Initialize network weights.

    Parameters:
        net (network)   -- network to be initialized
        init_type (str) -- the name of an initialization method: normal | xavier | kaiming | orthogonal
        init_gain (float)    -- scaling factor for normal, xavier and orthogonal.

    We use 'normal' in the original pix2pix and CycleGAN paper. But xavier and kaiming might
    work better for some applications. Feel free to try yourself.
    """
    def init_func(m):  # define the initialization function
        classname = m.__class__.__name__
        if hasattr(m, 'weight') and (classname.find('Conv') != -1 or classname.find('Linear') != -1):
            if init_type == 'normal':
                init.normal_(m.weight.data, 0.0, init_gain)
            elif init_type == 'xavier':
                init.xavier_normal_(m.weight.data, gain=init_gain)
            elif init_type == 'kaiming':
                init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
            elif init_type == 'orthogonal':
                init.orthogonal_(m.weight.data, gain=init_gain)
            else:
                raise NotImplementedError('initialization method [%s] is not implemented' % init_type)
            if hasattr(m, 'bias') and m.bias is not None:
                init.constant_(m.bias.data, 0.0)
        elif classname.find('BatchNorm2d') != -1:  # BatchNorm Layer's weight is not a matrix; only normal distribution applies.
            init.normal_(m.weight.data, 1.0, init_gain)
            init.constant_(m.bias.data, 0.0)

    print('initialize network with %s' % init_type)
    net.apply(init_func)  # apply the initialization function <init_func>

def init_net(net, init_type='normal', init_gain=0.02, gpu_ids=[]):
    """Initialize a network: 1. register CPU/GPU device (with multi-GPU support); 2. initialize the network weights
    Parameters:
        net (network)      -- the network to be initialized
        init_type (str)    -- the name of an initialization method: normal | xavier | kaiming | orthogonal
        gain (float)       -- scaling factor for normal, xavier and orthogonal.
        gpu_ids (int list) -- which GPUs the network runs on: e.g., 0,1,2

    Return an initialized network.
    """
    if len(gpu_ids) > 0:
        assert(torch.cuda.is_available())
        net.to(gpu_ids[0])
        if len(gpu_ids) > 1:
            net = torch.nn.DataParallel(net, gpu_ids)  # multi-GPUs
    init_weights(net, init_type, init_gain=init_gain)
    return net

def define_G(args, init_type='normal', init_gain=0.02, gpu_ids=[]):
    if args.net_G == 'base_resnet18':
        net = ResNet(input_nc=3, output_nc=2, output_sigmoid=False)

    elif args.net_G == 'base_transformer_pos_s4':
        net = BASE_Transformer(input_nc=3, output_nc=2, token_len=4, resnet_stages_num=4,
                             with_pos='learned')

    elif args.net_G == 'base_transformer_pos_s4_dd8':
        net = BASE_Transformer(input_nc=3, output_nc=2, token_len=4, resnet_stages_num=4,
                             with_pos='learned', enc_depth=1, dec_depth=8)
    elif args.net_G == 'base_transformer_pos_s4_dd8_o5':
        net = BASE_Transformer(input_nc=3, output_nc=5, token_len=4, resnet_stages_num=4,
                             with_pos='learned', enc_depth=1, dec_depth=8)

    elif args.net_G == 'base_transformer_pos_s4_dd8_dedim8':
        net = BASE_Transformer(input_nc=3, output_nc=2, token_len=4, resnet_stages_num=4,
                             with_pos='learned', enc_depth=1, dec_depth=8, decoder_dim_head=8)
    elif args.net_G == 'base_transformer_pos_s4_dd8_t8_e2d4':
       net = BASE_Transformer(input_nc=3, output_nc=2, token_len=8, resnet_stages_num=4,
                              with_pos='learned', enc_depth=2, dec_depth=4, decoder_dim_head=8)
    elif args.net_G == 'unet_coupled_trans_256':
       net = UNet_Change_Transformer()
    elif args.net_G == 'unet_coupled_two_trans_256':
       net = UNet_Change_Two_Transformer()
    elif args.net_G == 'siamUnet_conc':
       net = SiamUnet_conc(input_nbr=3, label_nbr=2)
    elif args.net_G == 'siamUnet':
       net = Res34_Unet_Double()
    elif args.net_G == 'changeFormer':
       net = ChangeFormerV1()
    elif args.net_G == 'changeFormerV6':
       net = ChangeFormerV6()
    elif args.net_G == 'newUNetTrans':
       net =  BASE_Transformer_UNet(input_nc=3, output_nc=2, token_len=4, resnet_stages_num=4,
                              with_pos='learned', with_decoder_pos='learned', enc_depth=1, dec_depth=8)
    else:
        raise NotImplementedError('Generator model name [%s] is not recognized' % args.net_G)
    return init_net(net, init_type, init_gain, gpu_ids)

###############################################################################
# main Functions
###############################################################################

class ResNet(torch.nn.Module):
    def __init__(self, input_nc, output_nc,
                 resnet_stages_num=5, backbone='resnet18',
                 output_sigmoid=False, if_upsample_2x=True):
        """
        In the constructor we instantiate two nn.Linear modules and assign them as
        member variables.
        """
        super(ResNet, self).__init__()
        expand = 1
        if backbone == 'resnet18':
            self.resnet = models.resnet18(pretrained=True,
                                          replace_stride_with_dilation=[False,True,True])
        elif backbone == 'resnet34':
            self.resnet = models.resnet34(pretrained=True,
                                          replace_stride_with_dilation=[False,True,True])
        elif backbone == 'resnet50':
            self.resnet = models.resnet50(pretrained=True,
                                          replace_stride_with_dilation=[False,True,True])
            expand = 4
        else:
            raise NotImplementedError
        self.relu = nn.ReLU()
        self.upsamplex2 = nn.Upsample(scale_factor=2)
        self.upsamplex4 = nn.Upsample(scale_factor=4, mode='bilinear')

        self.classifier = TwoLayerConv2d(in_channels=32, out_channels=output_nc)

        self.resnet_stages_num = resnet_stages_num

        self.if_upsample_2x = if_upsample_2x
        if self.resnet_stages_num == 5:
            layers = 512 * expand
        elif self.resnet_stages_num == 4:
            layers = 256 * expand
        elif self.resnet_stages_num == 3:
            layers = 128 * expand
        else:
            raise NotImplementedError
        self.conv_pred = nn.Conv2d(layers, 32, kernel_size=3, padding=1)

        self.output_sigmoid = output_sigmoid
        self.sigmoid = nn.Sigmoid()

    def forward(self, x1, x2):
        x1 = self.forward_single(x1)
        x2 = self.forward_single(x2)
        x = torch.abs(x1 - x2)
        if not self.if_upsample_2x:
            x = self.upsamplex2(x)
        x = self.upsamplex4(x)
        x = self.classifier(x)

        if self.output_sigmoid:
            x = self.sigmoid(x)
        return x

    def forward_single(self, x):
        # resnet layers
        x = self.resnet.conv1(x)
        x = self.resnet.bn1(x)
        x = self.resnet.relu(x)
        x = self.resnet.maxpool(x)

        x_4 = self.resnet.layer1(x) # 1/4, in=64, out=64
        x_8 = self.resnet.layer2(x_4) # 1/8, in=64, out=128

        if self.resnet_stages_num > 3:
            x_8 = self.resnet.layer3(x_8) # 1/8, in=128, out=256

        if self.resnet_stages_num == 5:
            x_8 = self.resnet.layer4(x_8) # 1/32, in=256, out=512
        elif self.resnet_stages_num > 5:
            raise NotImplementedError

        if self.if_upsample_2x:
            x = self.upsamplex2(x_8)
        else:
            x = x_8
        # output layers
        x = self.conv_pred(x)
        return x

class BASE_Transformer(ResNet):
    """
    Resnet of 8 downsampling + BIT + bitemporal feature Differencing + a small CNN
    """
    def __init__(self, input_nc, output_nc, with_pos, resnet_stages_num=5,
                 token_len=4, token_trans=True,
                 enc_depth=1, dec_depth=1,
                 dim_head=64, decoder_dim_head=64,
                 tokenizer=True, if_upsample_2x=True,
                 pool_mode='max', pool_size=2,
                 backbone='resnet18',
                 decoder_softmax=True, with_decoder_pos=None,
                 with_decoder=True):
        super(BASE_Transformer, self).__init__(input_nc, output_nc,backbone=backbone,
                                             resnet_stages_num=resnet_stages_num,
                                               if_upsample_2x=if_upsample_2x,
                                               )
        self.token_len = token_len
        self.conv_a = nn.Conv2d(32, self.token_len, kernel_size=1,
                                padding=0, bias=False)
        self.tokenizer = tokenizer
        if not self.tokenizer:
            #  if not use tokenzier，then downsample the feature map into a certain size
            self.pooling_size = pool_size
            self.pool_mode = pool_mode
            self.token_len = self.pooling_size * self.pooling_size

        self.token_trans = token_trans
        self.with_decoder = with_decoder
        dim = 32
        mlp_dim = 2*dim

        self.with_pos = with_pos
        if with_pos == 'learned':
            self.pos_embedding = nn.Parameter(torch.randn(1, self.token_len*2, 32))
        decoder_pos_size = 256//4
        self.with_decoder_pos = with_decoder_pos
        if self.with_decoder_pos == 'learned':
            self.pos_embedding_decoder =nn.Parameter(torch.randn(1, 32,
                                                                 decoder_pos_size,
                                                                 decoder_pos_size))
        self.enc_depth = enc_depth
        self.dec_depth = dec_depth
        self.dim_head = dim_head
        self.decoder_dim_head = decoder_dim_head
        self.transformer = Transformer(dim=dim, depth=self.enc_depth, heads=8,
                                       dim_head=self.dim_head,
                                       mlp_dim=mlp_dim, dropout=0)
        self.transformer_decoder = TransformerDecoder(dim=dim, depth=self.dec_depth,
                            heads=8, dim_head=self.decoder_dim_head, mlp_dim=mlp_dim, dropout=0,
                                                      softmax=decoder_softmax)

    def _forward_semantic_tokens(self, x):
        b, c, h, w = x.shape
        spatial_attention = self.conv_a(x)
        spatial_attention = spatial_attention.view([b, self.token_len, -1]).contiguous()
        spatial_attention = torch.softmax(spatial_attention, dim=-1)
        x = x.view([b, c, -1]).contiguous()
        tokens = torch.einsum('bln,bcn->blc', spatial_attention, x)
        return tokens

    def _forward_reshape_tokens(self, x):
        # b,c,h,w = x.shape
        if self.pool_mode == 'max':
            x = F.adaptive_max_pool2d(x, [self.pooling_size, self.pooling_size])
        elif self.pool_mode == 'ave':
            x = F.adaptive_avg_pool2d(x, [self.pooling_size, self.pooling_size])
        else:
            x = x
        tokens = rearrange(x, 'b c h w -> b (h w) c')
        return tokens

    def _forward_transformer(self, x):
        if self.with_pos:
            x += self.pos_embedding
        x = self.transformer(x)
        return x

    def _forward_transformer_decoder(self, x, m):
        b, c, h, w = x.shape
        if self.with_decoder_pos == 'fix':
            x = x + self.pos_embedding_decoder
        elif self.with_decoder_pos == 'learned':
            x = x + self.pos_embedding_decoder
        x = rearrange(x, 'b c h w -> b (h w) c')
        x = self.transformer_decoder(x, m)
        x = rearrange(x, 'b (h w) c -> b c h w', h=h)
        return x

    def _forward_simple_decoder(self, x, m):
        b, c, h, w = x.shape
        b, l, c = m.shape
        m = m.expand([h,w,b,l,c])
        m = rearrange(m, 'h w b l c -> l b c h w')
        m = m.sum(0)
        x = x + m
        return x

    def forward(self, x1, x2):
        # forward backbone resnet
        x1 = self.forward_single(x1)
        x2 = self.forward_single(x2)

        #  forward tokenzier
        if self.tokenizer:
            token1 = self._forward_semantic_tokens(x1)
            token2 = self._forward_semantic_tokens(x2)
        else:
            token1 = self._forward_reshape_tokens(x1)
            token2 = self._forward_reshape_tokens(x2)

        # forward transformer encoder
        if self.token_trans:
            self.tokens_ = torch.cat([token1, token2], dim=1)
            self.tokens = self._forward_transformer(self.tokens_)
            token1, token2 = self.tokens.chunk(2, dim=1)
        # forward transformer decoder
        if self.with_decoder:
            x1 = self._forward_transformer_decoder(x1, token1)
            x2 = self._forward_transformer_decoder(x2, token2)
        else:
            x1 = self._forward_simple_decoder(x1, token1)
            x2 = self._forward_simple_decoder(x2, token2)
        # feature differencing
        x = torch.abs(x1 - x2)
        if not self.if_upsample_2x:
            x = self.upsamplex2(x)
        x = self.upsamplex4(x)
        # forward small cnn
        x = self.classifier(x)
        if self.output_sigmoid:
            x = self.sigmoid(x)
        return x

class ConvReluBN(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size=3):
        super(ConvReluBN, self).__init__()
        self.layer = nn.Sequential(
            nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True)
        )
    def forward(self, x):
        return self.layer(x)

class ConvRelu(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size=3):
        super(ConvRelu, self).__init__()
        self.layer = nn.Sequential(
            nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, padding=1),
            nn.ReLU(inplace=True)
        )
    def forward(self, x):
        return self.layer(x)

class ChannelAttention(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size=3):
        super(ChannelAttention, self).__init__()

        assert kernel_size in (3, 7), 'kernel size must be 3 or 7'
        padding = 3 if kernel_size == 7 else 1
        self.conv1 = nn.Conv2d(in_channels*2, out_channels, kernel_size, padding=padding, bias=False)
        self.tanh = nn.Tanh()

    def forward(self, x_1, x_2):
        x = torch.cat([x_1, x_2], dim=1)
        x = self.conv1(x)
        return self.tanh(x)

class PreNorm(nn.Module):
    def __init__(self, dim, fn):
        super().__init__()
        self.norm = nn.LayerNorm(dim)
        self.fn = fn
    def forward(self, x, **kwargs):
        return self.fn(self.norm(x), **kwargs)

class FeedForward(nn.Module):
    def __init__(self, dim, hidden_dim, dropout = 0.):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(dim, hidden_dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, dim),
            nn.Dropout(dropout)
        )
    def forward(self, x):
        return self.net(x)

class Attention(nn.Module):
    def __init__(self, dim, heads = 8, dim_head = 64, dropout = 0.):
        super().__init__()
        inner_dim = dim_head *  heads
        self.heads = heads
        self.scale = dim ** -0.5

        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias = False)
        self.to_out = nn.Sequential(
            nn.Linear(inner_dim, dim),
            nn.Dropout(dropout)
        )

    def forward(self, x, mask = None):
        b, n, _, h = *x.shape, self.heads
        qkv = self.to_qkv(x).chunk(3, dim = -1)
        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = h), qkv)
        dots = torch.einsum('bhid,bhjd->bhij', q, k) * self.scale
        mask_value = -torch.finfo(dots.dtype).max

        if mask is not None:
            mask = F.pad(mask.flatten(1), (1, 0), value = True)
            assert mask.shape[-1] == dots.shape[-1], 'mask has incorrect dimensions'
            mask = mask[:, None, :] * mask[:, :, None]
            dots.masked_fill_(~mask, mask_value)
            del mask

        attn = dots.softmax(dim=-1)
        out = torch.einsum('bhij,bhjd->bhid', attn, v)
        out = rearrange(out, 'b h n d -> b n (h d)')
        out = self.to_out(out)
        return out

class Residual(nn.Module):
    def __init__(self, fn):
        super().__init__()
        self.fn = fn
    def forward(self, x, **kwargs):
        return self.fn(x, **kwargs) + x

class Transformer(nn.Module):
    def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout):
        super().__init__()
        self.layers = nn.ModuleList([])
        for _ in range(depth):
            self.layers.append(nn.ModuleList([
                Residual(PreNorm(dim, Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout))),
                Residual(PreNorm(dim, FeedForward(dim, mlp_dim, dropout = dropout)))
            ]))
    def forward(self, x, mask = None):
        for attn, ff in self.layers:
            x = attn(x, mask = mask)
            x = ff(x)
        return x

# Attention for bottleneck layers
class ChannelAttention_OnBottle(nn.Module):
    def __init__(self, in_planes, ratio=16, att_type='max'):
        super(ChannelAttention_OnBottle, self).__init__()
        self.att_type = att_type
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.max_pool = nn.AdaptiveMaxPool2d(1)
        # self.min_pool = nn.AdaptiveMinPool2d(1)

        self.fc1 = nn.Conv2d(in_planes, in_planes // ratio, 1, bias=False)
        self.relu = nn.ReLU()
        self.fc2 = nn.Conv2d(in_planes // ratio, in_planes, 1, bias=False)
        self.sigmoid = nn.Sigmoid()
        self.fc3 = nn.Linear(64, 512)
        self.fc4 = nn.Linear(96, 512)

    def forward(self, x):
        if self.att_type == 'max':
            out = self.fc2(self.relu(self.fc1(self.max_pool(x))))
        elif self.att_type == 'max_avg':
            max_out = self.relu(self.fc1(self.max_pool(x)))
            avg_out = self.relu(self.fc1(self.avg_pool(x)))
            out = torch.cat([max_out, avg_out], 1).squeeze()
            out = self.fc3(out).unsqueeze(0).unsqueeze(2).unsqueeze(3)
        elif self.att_type == 'avg_max_min':
            avg_out = self.relu(self.fc1(self.avg_pool(x)))
            # min_out = self.relu(self.fc1(self.min_pool(x)))
            max_out = self.relu(self.fc1(self.max_pool(x)))
            out = torch.cat([avg_out, min_out, max_out], 1).squeeze()
            out = self.fc4(out).unsqueeze(0).unsqueeze(2).unsqueeze(3)
        return self.relu(out)

class UNet_Change_Transformer(nn.Module):
    def __init__(self, pretrained=False, **kwargs):
        super(UNet_Change_Transformer, self).__init__()

        encoder_filters = [64, 64, 128, 256, 512]
        decoder_filters = np.asarray([48, 64, 96, 128, 320])

        self.encoder_filters = [64, 64, 128, 256, 512]
        self.decoder_filters = np.asarray([48, 64, 96, 160, 320])

        self.conv6 = ConvRelu(encoder_filters[-1], decoder_filters[-1])
        self.conv6_2 = ConvRelu(decoder_filters[-1] + encoder_filters[-2]*2 , decoder_filters[-1])
        self.conv7 = ConvRelu(decoder_filters[-1], decoder_filters[-2])
        self.conv7_2 = ConvRelu(decoder_filters[-2] + encoder_filters[-3]*2 , decoder_filters[-2])
        self.conv8 = ConvRelu(decoder_filters[-2], decoder_filters[-3])
        self.conv8_2 = ConvRelu(decoder_filters[-3] + encoder_filters[-4]*2 , decoder_filters[-3])
        self.conv9 = ConvRelu(decoder_filters[-3], decoder_filters[-4])
        self.conv9_2 = ConvRelu(decoder_filters[-4] + encoder_filters[-5]*2 , decoder_filters[-4])
        self.conv10 = ConvRelu(decoder_filters[-4], decoder_filters[-5])

        self.res = nn.Conv2d(decoder_filters[-5], 2, 1, stride=1, padding=0)

        self._initialize_weights()

        encoder = torchvision.models.resnet34(pretrained=pretrained)
        self.conv0 = ConvRelu(6,3)
        self.conv1 = nn.Sequential(
                        encoder.conv1,
                        encoder.bn1,
                        encoder.relu)
        self.conv2 = nn.Sequential(
                        encoder.maxpool,
                        encoder.layer1)
        self.conv3 = encoder.layer2
        self.conv4 = encoder.layer3
        self.conv5 = encoder.layer4

        self.ca_skip_5 = ChannelAttention(encoder_filters[-1], encoder_filters[-1])
        self.ca_skip_4 = ChannelAttention(encoder_filters[-2], encoder_filters[-2])
        self.ca_skip_3 = ChannelAttention(encoder_filters[-3], encoder_filters[-3])
        self.ca_skip_2 = ChannelAttention(encoder_filters[-4], encoder_filters[-4])
        self.ca_skip_1 = ChannelAttention(encoder_filters[-5], encoder_filters[-5])

        self.ca_bottle_max = ChannelAttention_OnBottle(512, att_type='max')
        self.ca_bottle_avg_min = ChannelAttention_OnBottle(512, att_type='max_avg')
        self.sigmoid = nn.Sigmoid()
        self.linearb = nn.Linear(1024, 512)

        dim = 64
        mlp_dim = 2*dim
        enc_depth = 2
        dim_head = 64
        decoder_dim_head = 64
        decoder_softmax = True
        self.transformer = Transformer(dim=dim, depth=3, heads=4,
                                       dim_head=dim_head,
                                       mlp_dim=mlp_dim, dropout=0.05)

        self.transformer_decoder = TransformerDecoder(dim=dim, depth=2,
                            heads=8, dim_head=decoder_dim_head, mlp_dim=mlp_dim, dropout=0,
                                                      softmax=decoder_softmax)
        self.classifier = TwoLayerConv2d(in_channels=512, out_channels=2)

    def forward_1(self, x1, x2):

        # Encoder 1
        x_1 = x1
        enc1_1 = self.conv1(x_1)
        enc2_1 = self.conv2(enc1_1)
        enc3_1 = self.conv3(enc2_1)
        enc4_1 = self.conv4(enc3_1)
        enc5_1 = self.conv5(enc4_1)

        # Encoder 2
        x_2 = x2
        enc1_2 = self.conv1(x_2)
        enc2_2 = self.conv2(enc1_2)
        enc3_2 = self.conv3(enc2_2)
        enc4_2 = self.conv4(enc3_2)
        enc5_2 = self.conv5(enc4_2)

        # Bottleneck
        enc5_1 = (self.ca_bottle_max(enc5_1)*enc5_1)
        enc5_2 = (self.ca_bottle_max(enc5_2)*enc5_2)

        enc5 = self.ca_skip_5(enc5_1,enc5_2)

        B_, C_, H_, W_ = enc5.shape
        enc5_i = enc5.view([B_, C_, H_*W_])
        enc5_i = self.transformer(enc5_i)
        enc5_i = enc5_i.view([B_, C_, H_, W_])
        enc5 = self.ca_skip_5(enc5_i,enc5)

        # Decoder
        enc4 = self.ca_skip_4(enc4_1, enc4_2)
        # enc4 = attention_block(enc4_1, enc4_2, self.encoder_filters[-2])
        dec6 = self.conv6(F.interpolate(enc5, scale_factor=2))
        dec6 = self.conv6_2(torch.cat([dec6, enc4
                ], 1))

        enc3 = self.ca_skip_3(enc3_1, enc3_2)
        # enc3 = attention_block(enc3_1, enc3_2, self.encoder_filters[-3])
        dec7 = self.conv7(F.interpolate(dec6, scale_factor=2))
        dec7 = self.conv7_2(torch.cat([dec7, enc3
                ], 1))

        enc2 = self.ca_skip_2(enc2_1, enc2_2)
        # enc2 = attention_block(enc2_1, enc2_2, self.encoder_filters[-4])
        dec8 = self.conv8(F.interpolate(dec7, scale_factor=2))
        dec8 = self.conv8_2(torch.cat([dec8, enc2
                ], 1))

        enc1 = self.ca_skip_2(enc1_1, enc1_2)
        # enc1 = attention_block(enc1_1, enc1_2, self.encoder_filters[-5])
        dec9 = self.conv9(F.interpolate(dec8, scale_factor=2))
        dec9 = self.conv9_2(torch.cat([dec9,
                enc1
                ], 1))

        dec10 = self.conv10(F.interpolate(dec9, scale_factor=2))
        out = self.res(dec10)
        return out

    def forward(self, x1, x2):
        # Encoder 1
        x_1 = x1
        enc1_1 = self.conv1(x_1)
        enc2_1 = self.conv2(enc1_1)
        enc3_1 = self.conv3(enc2_1)
        enc4_1 = self.conv4(enc3_1)
        enc5_1 = self.conv5(enc4_1)

        # Encoder 2
        x_2 = x2
        enc1_2 = self.conv1(x_2)
        enc2_2 = self.conv2(enc1_2)
        enc3_2 = self.conv3(enc2_2)
        enc4_2 = self.conv4(enc3_2)
        enc5_2 = self.conv5(enc4_2)

        # Bottleneck
        enc5_1 = (self.ca_bottle_max(enc5_1)*enc5_1)
        enc5_2 = (self.ca_bottle_max(enc5_2)*enc5_2)
        enc5_c = self.ca_skip_5(enc5_1,enc5_2)

        B_, C_, H_, W_ = enc5_c.shape
        enc5_i = enc5_c.view([B_, C_, H_*W_])
        enc5_i = self.transformer(enc5_i)
        enc5 = enc5_i.view([B_, C_, H_, W_])
        enc5 = self.ca_skip_5(enc5,enc5_c)

        # Decoder
        # enc4 = self.ca_skip_4(enc4_1, enc4_2)
        # enc4 = attention_block(enc4_1, enc4_2, self.encoder_filters[-2])
        dec6 = self.conv6(F.interpolate(enc5, scale_factor=2))
        dec6 = self.conv6_2(torch.cat([dec6, enc4_1, enc4_2
                ], 1))

        # enc3 = self.ca_skip_3(enc3_1, enc3_2)
        # enc3 = attention_block(enc3_1, enc3_2, self.encoder_filters[-3])
        dec7 = self.conv7(F.interpolate(dec6, scale_factor=2))
        dec7 = self.conv7_2(torch.cat([dec7, enc3_1, enc3_2
                ], 1))

        # enc2 = self.ca_skip_2(enc2_1, enc2_2)
        # enc2 = attention_block(enc2_1, enc2_2, self.encoder_filters[-4])
        dec8 = self.conv8(F.interpolate(dec7, scale_factor=2))
        dec8 = self.conv8_2(torch.cat([dec8, enc2_1, enc2_2
                ], 1))

        # enc1 = self.ca_skip_2(enc1_1, enc1_2)
        # enc1 = attention_block(enc1_1, enc1_2, self.encoder_filters[-5])
        dec9 = self.conv9(F.interpolate(dec8, scale_factor=2))
        dec9 = self.conv9_2(torch.cat([dec9, enc1_1, enc1_2], 1))

        dec10 = self.conv10(F.interpolate(dec9, scale_factor=2))

        out = self.res(dec10)

        # enc5_c = enc5_c.view([B_, C_, H_*W_])
        # interim_out = self.transformer_decoder(enc5_i, enc5_c)
        # interim_out = interim_out.view([B_, C_, H_, W_])
        # interim_out = self.classifier(enc5)

        return out #, interim_out

    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d) or isinstance(m, nn.Linear):
                m.weight.data = nn.init.kaiming_normal_(m.weight.data)
                if m.bias is not None:
                    m.bias.data.zero_()
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()

class UNet_Change_Two_Transformer(nn.Module):
    def __init__(self, pretrained=False, **kwargs):
        super(UNet_Change_Two_Transformer, self).__init__()

        encoder_filters = [64, 64, 128, 256, 512]
        decoder_filters = np.asarray([48, 64, 96, 128, 320])

        self.encoder_filters = [64, 64, 128, 256, 512]
        self.decoder_filters = np.asarray([48, 64, 96, 160, 320])

        self.conv6 = ConvRelu(encoder_filters[-1], decoder_filters[-1])
        self.conv6_2 = ConvRelu(decoder_filters[-1] + encoder_filters[-2]*2, decoder_filters[-1])
        self.conv7 = ConvRelu(decoder_filters[-1], decoder_filters[-2])
        self.conv7_2 = ConvRelu(decoder_filters[-2] + encoder_filters[-3]*2, decoder_filters[-2])
        self.conv8 = ConvRelu(decoder_filters[-2], decoder_filters[-3])
        self.conv8_2 = ConvRelu(decoder_filters[-3] + encoder_filters[-4] , decoder_filters[-3])
        self.conv9 = ConvRelu(decoder_filters[-3], decoder_filters[-4])
        self.conv9_2 = ConvRelu(decoder_filters[-4] + encoder_filters[-5]*2 , decoder_filters[-4])
        self.conv10 = ConvRelu(decoder_filters[-4], decoder_filters[-5])

        self.res = nn.Conv2d(decoder_filters[-5], 2, 1, stride=1, padding=0)

        self._initialize_weights()

        encoder = torchvision.models.resnet34(pretrained=pretrained)
        self.conv1 = nn.Sequential(
                        encoder.conv1,
                        encoder.bn1,
                        encoder.relu)
        self.conv2 = nn.Sequential(
                        encoder.maxpool,
                        encoder.layer1)
        self.conv3 = encoder.layer2
        self.conv4 = encoder.layer3
        self.conv5 = encoder.layer4

        self.ca_skip_5 = ChannelAttention(encoder_filters[-1], encoder_filters[-1])
        # self.ca_skip_4 = ChannelAttention(encoder_filters[-2], encoder_filters[-2])
        # self.ca_skip_3 = ChannelAttention(encoder_filters[-3], encoder_filters[-3])
        # self.ca_skip_2 = ChannelAttention(encoder_filters[-4], encoder_filters[-4])
        self.ca_skip_1 = ChannelAttention(encoder_filters[-5], encoder_filters[-5])

        # self.ca_bottle_max = ChannelAttention_OnBottle(512, att_type='max')
        # self.ca_bottle_avg_min = ChannelAttention_OnBottle(512, att_type='max_avg')
        # self.sigmoid = nn.Sigmoid()
        # self.linearb = nn.Linear(1024, 512)

        dim = 64
        dim2 = 4096
        dim3 = 1024
        mlp_dim = 2*dim
        enc_depth = 2
        dim_head = 64
        decoder_dim_head = 64
        decoder_softmax = True
        self.transformer = Transformer(dim=dim, depth=3, heads=4,
                                       dim_head=dim_head,
                                       mlp_dim=mlp_dim, dropout=0.01)

        self.transformer2 = Transformer(dim=dim2, depth=2, heads=2,
                                       dim_head=dim_head,
                                       mlp_dim=dim2, dropout=0.001)

        # self.transformer3 = Transformer(dim=dim3, depth=2, heads=2,
        #                                dim_head=dim_head,
        #                                mlp_dim=dim3, dropout=0.001)

        self.classifier = TwoLayerConv2d(in_channels=512, out_channels=2)
        # self.convT = nn.Conv2d(128, 64, kernel_size=3, stride=1, padding=1)

    def forward(self, x1, x2):
        # Encoder 1
        x_1 = x1
        enc1_1 = self.conv1(x_1)
        enc2_1 = self.conv2(enc1_1)
        enc3_1 = self.conv3(enc2_1)
        enc4_1 = self.conv4(enc3_1)
        enc5_1 = self.conv5(enc4_1)

        # Encoder 2
        x_2 = x2
        enc1_2 = self.conv1(x_2)
        enc2_2 = self.conv2(enc1_2)
        enc3_2 = self.conv3(enc2_2)
        enc4_2 = self.conv4(enc3_2)
        enc5_2 = self.conv5(enc4_2)

        # Bottleneck
        # enc5_1 = (self.ca_bottle_max(enc5_1)*enc5_1)
        # enc5_2 = (self.ca_bottle_max(enc5_2)*enc5_2)
        # enc5_c = self.ca_skip_5(enc5_1,enc5_2)

        ## run 1: updating channel attention
        enc5 = self.ca_skip_5(enc5_1,enc5_2)
        B_, C_, H_, W_ = enc5.shape
        enc5_i = enc5.view([B_, C_, H_*W_]).contiguous()

        enc5_diff = (enc5_1 - enc5_2)
        spatial_attention = enc5_diff.view([B_, C_, H_*W_]).contiguous()
        spatial_attention = torch.softmax(spatial_attention, dim=-1)

        enc5_i = torch.einsum('bln,bln->bln', spatial_attention, enc5_i)
        enc5_t = self.transformer(enc5_i)
        enc5 = enc5_t.view([B_, C_, H_, W_]).contiguous()
        # enc5 = self.ca_skip_5(enc5_t,enc5)

        # Decoder
        # enc4 = self.ca_skip_4(enc4_1, enc4_2)

        dec6 = self.conv6(F.interpolate(enc5, scale_factor=2))
        dec6 = self.conv6_2(torch.cat([dec6, enc4_1, enc4_2], 1))

        # run4: depth=2, heads=2
        # enc3 = self.ca_skip_3(enc3_1, enc3_2)
        # B_, C_, H_, W_ = enc3.shape
        # enc3_i = enc3.view([B_, C_, H_*W_]).contiguous()

        # enc3_diff = (enc3_1 - enc3_2)
        # spatial_attention = enc3_diff.view([B_, C_, H_*W_]).contiguous()
        # spatial_attention = torch.softmax(spatial_attention, dim=-1)

        # enc3_i = torch.einsum('bln,bln->bln', spatial_attention, enc3_i)
        # enc3_t = self.transformer3(enc3_i)
        # enc3_t = enc3_t.view([B_, C_, H_, W_]).contiguous()
        # enc3 = self.ca_skip_3(enc3, enc3_t)

        dec7 = self.conv7(F.interpolate(dec6, scale_factor=2))
        dec7 = self.conv7_2(torch.cat([dec7, enc3_1, enc3_2], 1))

        ## run0: depth=2, heads=1
        enc2 = self.ca_skip_2(enc2_1, enc2_2)
        B_, C_, H_, W_ = enc2.shape
        enc2 = enc2.view([B_, C_, H_*W_]).contiguous()

        enc2_diff = (enc2_1 - enc2_2)
        spatial_attention = enc2_diff.view([B_, C_, H_*W_]).contiguous()
        spatial_attention = torch.softmax(spatial_attention, dim=-1)

        enc2 = torch.einsum('bln,bln->bln', spatial_attention, enc2)
        enc2 = self.transformer2(enc2)
        enc2 = enc2.view([B_, C_, H_,W_]).contiguous()

        dec8 = self.conv8(F.interpolate(dec7, scale_factor=2))
        dec8 = self.conv8_2(torch.cat([dec8, enc2], 1))

        # enc1 = self.ca_skip_2(enc1_1, enc1_2)
        dec9 = self.conv9(F.interpolate(dec8, scale_factor=2))
        dec9 = self.conv9_2(torch.cat([dec9, enc1_1, enc1_2], 1))

        dec10 = self.conv10(F.interpolate(dec9, scale_factor=2))
        out = self.res(dec10)

        return out #, interim_out

    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d) or isinstance(m, nn.Linear):
                m.weight.data = nn.init.kaiming_normal_(m.weight.data)
                if m.bias is not None:
                    m.bias.data.zero_()
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()

class SiamUnet_conc(nn.Module):
    """SiamUnet_conc segmentation network."""

    def __init__(self, input_nbr, label_nbr):
        super(SiamUnet_conc, self).__init__()

        self.input_nbr = input_nbr

        self.conv11 = nn.Conv2d(input_nbr, 16, kernel_size=3, padding=1)
        self.bn11 = nn.BatchNorm2d(16)
        self.do11 = nn.Dropout2d(p=0.2)
        self.conv12 = nn.Conv2d(16, 16, kernel_size=3, padding=1)
        self.bn12 = nn.BatchNorm2d(16)
        self.do12 = nn.Dropout2d(p=0.2)

        self.conv21 = nn.Conv2d(16, 32, kernel_size=3, padding=1)
        self.bn21 = nn.BatchNorm2d(32)
        self.do21 = nn.Dropout2d(p=0.2)
        self.conv22 = nn.Conv2d(32, 32, kernel_size=3, padding=1)
        self.bn22 = nn.BatchNorm2d(32)
        self.do22 = nn.Dropout2d(p=0.2)

        self.conv31 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        self.bn31 = nn.BatchNorm2d(64)
        self.do31 = nn.Dropout2d(p=0.2)
        self.conv32 = nn.Conv2d(64, 64, kernel_size=3, padding=1)
        self.bn32 = nn.BatchNorm2d(64)
        self.do32 = nn.Dropout2d(p=0.2)
        self.conv33 = nn.Conv2d(64, 64, kernel_size=3, padding=1)
        self.bn33 = nn.BatchNorm2d(64)
        self.do33 = nn.Dropout2d(p=0.2)

        self.conv41 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
        self.bn41 = nn.BatchNorm2d(128)
        self.do41 = nn.Dropout2d(p=0.2)
        self.conv42 = nn.Conv2d(128, 128, kernel_size=3, padding=1)
        self.bn42 = nn.BatchNorm2d(128)
        self.do42 = nn.Dropout2d(p=0.2)
        self.conv43 = nn.Conv2d(128, 128, kernel_size=3, padding=1)
        self.bn43 = nn.BatchNorm2d(128)
        self.do43 = nn.Dropout2d(p=0.2)

        self.upconv4 = nn.ConvTranspose2d(128, 128, kernel_size=3, padding=1, stride=2, output_padding=1)

        self.conv43d = nn.ConvTranspose2d(384, 128, kernel_size=3, padding=1)
        self.bn43d = nn.BatchNorm2d(128)
        self.do43d = nn.Dropout2d(p=0.2)
        self.conv42d = nn.ConvTranspose2d(128, 128, kernel_size=3, padding=1)
        self.bn42d = nn.BatchNorm2d(128)
        self.do42d = nn.Dropout2d(p=0.2)
        self.conv41d = nn.ConvTranspose2d(128, 64, kernel_size=3, padding=1)
        self.bn41d = nn.BatchNorm2d(64)
        self.do41d = nn.Dropout2d(p=0.2)

        self.upconv3 = nn.ConvTranspose2d(64, 64, kernel_size=3, padding=1, stride=2, output_padding=1)

        self.conv33d = nn.ConvTranspose2d(192, 64, kernel_size=3, padding=1)
        self.bn33d = nn.BatchNorm2d(64)
        self.do33d = nn.Dropout2d(p=0.2)
        self.conv32d = nn.ConvTranspose2d(64, 64, kernel_size=3, padding=1)
        self.bn32d = nn.BatchNorm2d(64)
        self.do32d = nn.Dropout2d(p=0.2)
        self.conv31d = nn.ConvTranspose2d(64, 32, kernel_size=3, padding=1)
        self.bn31d = nn.BatchNorm2d(32)
        self.do31d = nn.Dropout2d(p=0.2)

        self.upconv2 = nn.ConvTranspose2d(32, 32, kernel_size=3, padding=1, stride=2, output_padding=1)

        self.conv22d = nn.ConvTranspose2d(96, 32, kernel_size=3, padding=1)
        self.bn22d = nn.BatchNorm2d(32)
        self.do22d = nn.Dropout2d(p=0.2)
        self.conv21d = nn.ConvTranspose2d(32, 16, kernel_size=3, padding=1)
        self.bn21d = nn.BatchNorm2d(16)
        self.do21d = nn.Dropout2d(p=0.2)

        self.upconv1 = nn.ConvTranspose2d(16, 16, kernel_size=3, padding=1, stride=2, output_padding=1)

        self.conv12d = nn.ConvTranspose2d(48, 16, kernel_size=3, padding=1)
        self.bn12d = nn.BatchNorm2d(16)
        self.do12d = nn.Dropout2d(p=0.2)
        self.conv11d = nn.ConvTranspose2d(16, label_nbr, kernel_size=3, padding=1)

        self.sm = nn.LogSoftmax(dim=1)

    def forward(self, x1, x2):

        """Forward method."""
        # Stage 1
        x11 = self.do11(F.relu(self.bn11(self.conv11(x1))))
        x12_1 = self.do12(F.relu(self.bn12(self.conv12(x11))))
        x1p = F.max_pool2d(x12_1, kernel_size=2, stride=2)

        # Stage 2
        x21 = self.do21(F.relu(self.bn21(self.conv21(x1p))))
        x22_1 = self.do22(F.relu(self.bn22(self.conv22(x21))))
        x2p = F.max_pool2d(x22_1, kernel_size=2, stride=2)

        # Stage 3
        x31 = self.do31(F.relu(self.bn31(self.conv31(x2p))))
        x32 = self.do32(F.relu(self.bn32(self.conv32(x31))))
        x33_1 = self.do33(F.relu(self.bn33(self.conv33(x32))))
        x3p = F.max_pool2d(x33_1, kernel_size=2, stride=2)

        # Stage 4
        x41 = self.do41(F.relu(self.bn41(self.conv41(x3p))))
        x42 = self.do42(F.relu(self.bn42(self.conv42(x41))))
        x43_1 = self.do43(F.relu(self.bn43(self.conv43(x42))))
        x4p = F.max_pool2d(x43_1, kernel_size=2, stride=2)

        ####################################################
        # Stage 1
        x11 = self.do11(F.relu(self.bn11(self.conv11(x2))))
        x12_2 = self.do12(F.relu(self.bn12(self.conv12(x11))))
        x1p = F.max_pool2d(x12_2, kernel_size=2, stride=2)

        # Stage 2
        x21 = self.do21(F.relu(self.bn21(self.conv21(x1p))))
        x22_2 = self.do22(F.relu(self.bn22(self.conv22(x21))))
        x2p = F.max_pool2d(x22_2, kernel_size=2, stride=2)

        # Stage 3
        x31 = self.do31(F.relu(self.bn31(self.conv31(x2p))))
        x32 = self.do32(F.relu(self.bn32(self.conv32(x31))))
        x33_2 = self.do33(F.relu(self.bn33(self.conv33(x32))))
        x3p = F.max_pool2d(x33_2, kernel_size=2, stride=2)

        # Stage 4
        x41 = self.do41(F.relu(self.bn41(self.conv41(x3p))))
        x42 = self.do42(F.relu(self.bn42(self.conv42(x41))))
        x43_2 = self.do43(F.relu(self.bn43(self.conv43(x42))))
        x4p = F.max_pool2d(x43_2, kernel_size=2, stride=2)

        ####################################################
        # Stage 4d
        x4d = self.upconv4(x4p)
        pad4 = ReplicationPad2d((0, x43_1.size(3) - x4d.size(3), 0, x43_1.size(2) - x4d.size(2)))
        x4d = torch.cat((pad4(x4d), x43_1, x43_2), 1)
        x43d = self.do43d(F.relu(self.bn43d(self.conv43d(x4d))))
        x42d = self.do42d(F.relu(self.bn42d(self.conv42d(x43d))))
        x41d = self.do41d(F.relu(self.bn41d(self.conv41d(x42d))))

        # Stage 3d
        x3d = self.upconv3(x41d)
        pad3 = ReplicationPad2d((0, x33_1.size(3) - x3d.size(3), 0, x33_1.size(2) - x3d.size(2)))
        x3d = torch.cat((pad3(x3d), x33_1, x33_2), 1)
        x33d = self.do33d(F.relu(self.bn33d(self.conv33d(x3d))))
        x32d = self.do32d(F.relu(self.bn32d(self.conv32d(x33d))))
        x31d = self.do31d(F.relu(self.bn31d(self.conv31d(x32d))))

        # Stage 2d
        x2d = self.upconv2(x31d)
        pad2 = ReplicationPad2d((0, x22_1.size(3) - x2d.size(3), 0, x22_1.size(2) - x2d.size(2)))
        x2d = torch.cat((pad2(x2d), x22_1, x22_2), 1)
        x22d = self.do22d(F.relu(self.bn22d(self.conv22d(x2d))))
        x21d = self.do21d(F.relu(self.bn21d(self.conv21d(x22d))))

        # Stage 1d
        x1d = self.upconv1(x21d)
        pad1 = ReplicationPad2d((0, x12_1.size(3) - x1d.size(3), 0, x12_1.size(2) - x1d.size(2)))
        x1d = torch.cat((pad1(x1d), x12_1, x12_2), 1)
        x12d = self.do12d(F.relu(self.bn12d(self.conv12d(x1d))))
        x11d = self.conv11d(x12d)

        return self.sm(x11d)

class ResNet_UNet(torch.nn.Module):
    def __init__(self, input_nc, output_nc,
                 resnet_stages_num=5, backbone='resnet18',
                 output_sigmoid=False, if_upsample_2x=True):
        """
        In the constructor we instantiate two nn.Linear modules and assign them as
        member variables.
        """
        super(ResNet_UNet, self).__init__()
        expand = 1
        if backbone == 'resnet18':
            self.resnet = models.resnet18(pretrained=True, replace_stride_with_dilation=[False,True,True])
        elif backbone == 'resnet34':
            self.resnet = models.resnet34(pretrained=True, replace_stride_with_dilation=[False,True,True])
        else:
            raise NotImplementedError
        self.relu = nn.ReLU()
        self.upsamplex2 = nn.Upsample(scale_factor=2)
        self.upsamplex4 = nn.Upsample(scale_factor=4, mode='bilinear')

        self.resnet_stages_num = resnet_stages_num

        self.if_upsample_2x = if_upsample_2x
        if self.resnet_stages_num == 5:
            layers = 512 * expand
        elif self.resnet_stages_num == 4:
            layers = 256 * expand
        elif self.resnet_stages_num == 3:
            layers = 128 * expand
        else:
            raise NotImplementedError
        self.conv_pred = nn.Conv2d(384, 32, kernel_size=3, padding=1)

    def forward_single(self, x):
        # resnet layers
        x = self.resnet.conv1(x)
        x = self.resnet.bn1(x)
        x_2 = self.resnet.relu(x)
        x_2_pool = self.resnet.maxpool(x)

        x_4 = self.resnet.layer1(x_2_pool) # 1/4, in=64, out=64

        x_8 = self.resnet.layer2(x_4) # 1/8, in=64, out=128
        x_8_pool = self.resnet.maxpool(x_8)

        x_10 = self.resnet.layer3(x_8_pool) # 1/8, in=128, out=256

        if self.resnet_stages_num > 4:
            raise NotImplementedError

        # print(x_2.shape, x_4.shape, x_8.shape, x_10.shape)
        x = self.upsamplex2(x_10)

        return x_2, x_4, x_8, x_10

class BASE_Transformer_UNet(ResNet_UNet):
    """
    Resnet of 8 downsampling + BIT + bitemporal feature Differencing + a small CNN
    """
    def __init__(self, input_nc, output_nc, with_pos, resnet_stages_num=5,
                 token_len=4, token_trans=True,
                 enc_depth=1, dec_depth=1,
                 dim_head=64, decoder_dim_head=64,
                 tokenizer=True, if_upsample_2x=True,
                 pool_mode='max', pool_size=2,
                 backbone='resnet18',
                 decoder_softmax=True, with_decoder_pos=None,
                 with_decoder=True):
        super(BASE_Transformer_UNet, self).__init__(input_nc, output_nc,backbone=backbone,
                                             resnet_stages_num=resnet_stages_num,
                                               if_upsample_2x=if_upsample_2x,
                                               )

        print("using UNet Transformer !!!!")

        self.token_len = token_len
        self.tokenizer = tokenizer
        self.token_trans = token_trans
        self.with_decoder = with_decoder
        self.with_pos = with_pos

        if not self.tokenizer:
            #  if not use tokenzier，then downsample the feature map into a certain size
            self.pooling_size = pool_size
            self.pool_mode = pool_mode
            self.token_len = self.pooling_size * self.pooling_size

        # conv squeeze layers before transformer
        dim_5, dim_4, dim_3, dim_2 = 32, 32, 32, 32
        self.conv_squeeze_5 = nn.Sequential(nn.Conv2d(256, dim_5, kernel_size=1, padding=0, bias=False),
                                            nn.ReLU())
        self.conv_squeeze_4 = nn.Sequential(nn.Conv2d(128, dim_4, kernel_size=1, padding=0, bias=False),
                                            nn.ReLU())
        self.conv_squeeze_3 = nn.Sequential(nn.Conv2d(64, dim_3, kernel_size=1, padding=0, bias=False),
                                            nn.ReLU())
        self.conv_squeeze_2 = nn.Sequential(nn.Conv2d(64, dim_2, kernel_size=1, padding=0, bias=False),
                                            nn.ReLU())
        self.conv_squeeze_layers = [self.conv_squeeze_2, self.conv_squeeze_3, self.conv_squeeze_4, self.conv_squeeze_5]

        self.conv_token_5 = nn.Conv2d(dim_5, self.token_len, kernel_size=1, padding=0, bias=False)
        self.conv_token_4 = nn.Conv2d(dim_4, self.token_len, kernel_size=1, padding=0, bias=False)
        self.conv_token_3 = nn.Conv2d(dim_3, self.token_len, kernel_size=1, padding=0, bias=False)
        self.conv_token_2 = nn.Conv2d(dim_2, self.token_len, kernel_size=1, padding=0, bias=False)
        self.conv_tokens_layers = [self.conv_token_2, self.conv_token_3, self.conv_token_4, self.conv_token_5]

        self.conv_decode_5 = nn.Conv2d(64, 32, kernel_size=3, padding=1, bias=False)
        self.conv_decode_4 = nn.Conv2d(64, 32, kernel_size=3, padding=1, bias=False)
        self.conv_decode_3 = nn.Conv2d(64, 32, kernel_size=3, padding=1, bias=False)
        self.conv_decode_2 = nn.Conv2d(64, 32, kernel_size=3, padding=1, bias=False)
        self.conv_decode_layers = [self.conv_decode_2, self.conv_decode_3, self.conv_decode_4, self.conv_decode_5]

        if with_pos is 'learned':
            self.pos_embedding_5 = nn.Parameter(torch.randn(1, self.token_len*2, dim_5))
            self.pos_embedding_4 = nn.Parameter(torch.randn(1, self.token_len*2, dim_4))
            self.pos_embedding_3 = nn.Parameter(torch.randn(1, self.token_len*2, dim_3))
            self.pos_embedding_2 = nn.Parameter(torch.randn(1, self.token_len*2, dim_2))
            self.pos_embedding_layers = [self.pos_embedding_2, self.pos_embedding_3, self.pos_embedding_4, self.pos_embedding_5]

        decoder_pos_size = 256//4
        self.with_decoder_pos = with_decoder_pos
        if self.with_decoder_pos == 'learned':
            self.pos_embedding_decoder_5 =nn.Parameter(torch.randn(1, dim_5, 16, 16))
            self.pos_embedding_decoder_4 =nn.Parameter(torch.randn(1, dim_4, 32, 32))
            self.pos_embedding_decoder_3 =nn.Parameter(torch.randn(1, dim_3, 64, 64))
            self.pos_embedding_decoder_2 =nn.Parameter(torch.randn(1, dim_2, decoder_pos_size, decoder_pos_size))
            self.pos_embedding_decoder_layers = [self.pos_embedding_decoder_2, self.pos_embedding_decoder_3, self.pos_embedding_decoder_4, self.pos_embedding_decoder_5]

        self.enc_depth = enc_depth
        self.dec_depth = dec_depth
        self.dim_head = dim_head
        self.decoder_dim_head = decoder_dim_head
        self.transformer_5 = Transformer(dim=dim_5, depth=self.enc_depth, heads=4,
                                        dim_head=self.dim_head, mlp_dim=dim_5, dropout=0)
        self.transformer_decoder_5 = TransformerDecoder(dim=dim_5, depth=4, heads=4,
                                                        dim_head=self.decoder_dim_head, mlp_dim=dim_5, dropout=0, softmax=decoder_softmax)
        self.transformer_4 = Transformer(dim=dim_4, depth=self.enc_depth, heads=4,
                                        dim_head=self.dim_head, mlp_dim=dim_4, dropout=0)
        self.transformer_decoder_4 = TransformerDecoder(dim=dim_4, depth=4, heads=4, dim_head=self.decoder_dim_head,
                                                         mlp_dim=dim_4, dropout=0, softmax=decoder_softmax)
        self.transformer_3 = Transformer(dim=dim_3, depth=self.enc_depth, heads=8,
                                         dim_head=self.dim_head, mlp_dim=dim_3, dropout=0)
        self.transformer_decoder_3 = TransformerDecoder(dim=dim_3, depth=8, heads=8, dim_head=self.decoder_dim_head,
                                                        mlp_dim=dim_3, dropout=0, softmax=decoder_softmax)
        self.transformer_2 = Transformer(dim=dim_2, depth=self.enc_depth, heads=1,
                                         dim_head=32, mlp_dim=dim_2, dropout=0)
        self.transformer_decoder_2 = TransformerDecoder(dim=dim_2, depth=1, heads=1, dim_head=32,
                                                        mlp_dim=dim_2, dropout=0, softmax=decoder_softmax)
        self.transformer_layers = [self.transformer_2, self.transformer_3, self.transformer_4, self.transformer_5]
        self.transformer_decoder_layers = [self.transformer_decoder_2, self.transformer_decoder_3, self.transformer_decoder_4, self.transformer_decoder_5]

        self.conv_layer2_0 = TwoLayerConv2d(in_channels=128, out_channels=32, kernel_size=3)

        # # EXP NEW : Worked Better than BiT and changeformer
        self.conv_layer2 = nn.Sequential(nn.Conv2d(in_channels=32, out_channels=32, kernel_size=3, padding=1),
                                        nn.ReLU())
        self.conv_layer3 = nn.Sequential(nn.Conv2d(in_channels=32, out_channels=32, kernel_size=3, padding=1),
                                        nn.ReLU())
        self.conv_layer4 = nn.Sequential(nn.Conv2d(in_channels=32, out_channels=32, kernel_size=3, padding=1),
                                        nn.ReLU())
        self.classifier = nn.Conv2d(in_channels=32, out_channels=output_nc, kernel_size=3, padding=1)

        # FINALIZED CLASSIFIER: Works Best
        # self.classifier = nn.Conv2d(in_channels=32, out_channels=output_nc, kernel_size=3, padding=1)

        # # EXP NEW
        # self.conv_layer2 = nn.Sequential(nn.Conv2d(in_channels=40, out_channels=16, kernel_size=3, padding=1),
        #                                 nn.ReLU(),
        #                                 nn.Conv2d(in_channels=16, out_channels=8, kernel_size=3, padding=1),
        #                                 nn.ReLU())
        # self.conv_layer3 = nn.Sequential(nn.Conv2d(in_channels=48, out_channels=16, kernel_size=3, padding=1),
        #                                 nn.ReLU(),
        #                                 nn.Conv2d(in_channels=16, out_channels=8, kernel_size=3, padding=1),
        #                                 nn.ReLU())
        # self.conv_layer4 = nn.Sequential(nn.Conv2d(in_channels=64, out_channels=32, kernel_size=3, padding=1),
        #                                 nn.ReLU(),
        #                                 nn.Conv2d(in_channels=32, out_channels=16, kernel_size=3, padding=1),
        #                                 nn.ReLU())
        # self.classifier = nn.Sequential(nn.Conv2d(in_channels=8, out_channels=5, kernel_size=3, padding=1),
        #                                 nn.ReLU(),
        #                                 nn.Conv2d(in_channels=5, out_channels=output_nc, kernel_size=3, padding=1))

    def _forward_semantic_tokens(self, x, layer=None):
        b, c, h, w = x.shape
        spatial_attention = self.conv_tokens_layers[layer](x)
        spatial_attention = spatial_attention.view([b, self.token_len, -1]).contiguous()
        spatial_attention = torch.softmax(spatial_attention, dim=-1)
        x = x.view([b, c, -1]).contiguous()
        tokens = torch.einsum('bln,bcn->blc', spatial_attention, x)
        return tokens

    def _forward_transformer(self, x, layer):
        if self.with_pos:
            x += self.pos_embedding_layers[layer]
        x = self.transformer_layers[layer](x)
        return x

    def _forward_transformer_decoder(self, x, m, layer):
        b, c, h, w = x.shape
        if self.with_decoder_pos == 'learned':
            x = x + self.pos_embedding_decoder_layers[layer]
        x = rearrange(x, 'b c h w -> b (h w) c')
        x = self.transformer_decoder_layers[layer](x, m)
        x = rearrange(x, 'b (h w) c -> b c h w', h=h)
        return x

    def _forward_trans_module(self, x1, x2, layer):
        x1 = self.conv_squeeze_layers[layer](x1)
        x2 = self.conv_squeeze_layers[layer](x2)
        token1 = self._forward_semantic_tokens(x1, layer)
        token2 = self._forward_semantic_tokens(x2, layer)
        self.tokens_ = torch.cat([token1, token2], dim=1)
        self.tokens = self._forward_transformer(self.tokens_, layer)
        token1, token2 = self.tokens.chunk(2, dim=1)
        x1 = self._forward_transformer_decoder(x1, token1, layer)
        x2 = self._forward_transformer_decoder(x2, token2, layer)
        # return torch.abs(x1 - x2)

        # V1, V2
        # x1 = self._forward_transformer_decoder(x1, token2, layer)
        # x2 = self._forward_transformer_decoder(x2, token1, layer)
        # return torch.add(x1, x2)

        # V3
        diff_token = torch.abs(token2 - token1)
        diff_x = self.conv_decode_layers[layer](torch.cat([x1,x2], axis=1))
        x = self._forward_transformer_decoder(diff_x, diff_token, layer)
        return x

    def forward(self, x1, x2):
        # forward backbone resnet
        a_128, a_64, a_32, a_16 = self.forward_single(x1)
        b_128, b_64, b_32, b_16 = self.forward_single(x2)

        #  level 5 in=256x16x16 out=32x16x16
        x1, x2 = a_16, b_16
        out_5 = self._forward_trans_module(x1, x2, layer=3)
        out_5 = self.upsamplex2(out_5)

        # level 4: in=128x32x32 out=32x32x32
        x1, x2 = a_32, b_32
        out_4 = self._forward_trans_module(x1, x2, layer=2)
        out_4 = out_4 + out_5
        out_4 = self.upsamplex2(out_4)
        out_4 = self.conv_layer4(out_4)

        # level 3: in=64x64x64 out=32x64x64
        x1, x2 = a_64, b_64
        out_3 = self._forward_trans_module(x1, x2, layer=1)
        out_3 = out_3 + out_4
        # out_3 = self.conv_layer3(torch.cat([out_3, out_4], axis=1))
        out_3 = self.upsamplex2(out_3)
        out_3 = self.conv_layer3(out_3)

        # level 2: in=64x128x128
        out_2 = self.conv_layer2_0(torch.cat([a_128, b_128], 1))
        out_2 = out_2 + out_3
        # out_2 = self.conv_layer2(torch.cat([out_2, out_3], axis=1))
        out_2 = self.upsamplex2(out_2)
        out_2 = self.conv_layer2(out_2)

        # print(out_2.shape, out_3.shape, out_4.shape, out_5.shape)
        # forward small cnn
        x = self.classifier(out_2)

        return x

class Res34_Unet_Double(nn.Module):
    def __init__(self, pretrained=False, **kwargs):
        super(Res34_Unet_Double, self).__init__()

        encoder_filters = [64, 64, 128, 256, 512]
        decoder_filters = np.asarray([48, 64, 96, 160, 320])

        self.conv6 = ConvRelu(encoder_filters[-1], decoder_filters[-1])
        self.conv6_2 = ConvRelu(decoder_filters[-1] + encoder_filters[-2], decoder_filters[-1])
        self.conv7 = ConvRelu(decoder_filters[-1], decoder_filters[-2])
        self.conv7_2 = ConvRelu(decoder_filters[-2] + encoder_filters[-3], decoder_filters[-2])
        self.conv8 = ConvRelu(decoder_filters[-2], decoder_filters[-3])
        self.conv8_2 = ConvRelu(decoder_filters[-3] + encoder_filters[-4], decoder_filters[-3])
        self.conv9 = ConvRelu(decoder_filters[-3], decoder_filters[-4])
        self.conv9_2 = ConvRelu(decoder_filters[-4] + encoder_filters[-5], decoder_filters[-4])
        self.conv10 = ConvRelu(decoder_filters[-4], decoder_filters[-5])

        self.res = nn.Conv2d(decoder_filters[-5] * 2, 5, 1, stride=1, padding=0)
        #self.res = nn.Conv2d(decoder_filters[-5], 5, 1, stride=1, padding=0)
        self._initialize_weights()

        encoder = resnet34(pretrained=pretrained)
        self.conv0 = ConvRelu(6,3)
        self.conv1 = nn.Sequential(
                        encoder.conv1,
                        encoder.bn1,
                        encoder.relu)
        self.conv2 = nn.Sequential(
                        encoder.maxpool,
                        encoder.layer1)
        self.conv3 = encoder.layer2
        self.conv4 = encoder.layer3
        self.conv5 = encoder.layer4

    def forward1(self, x):
        batch_size, C, H, W = x.shape
        #x = self.conv0(x)
        enc1 = self.conv1(x)
        enc2 = self.conv2(enc1)
        enc3 = self.conv3(enc2)
        enc4 = self.conv4(enc3)
        enc5 = self.conv5(enc4)

        dec6 = self.conv6(F.interpolate(enc5, scale_factor=2))
        dec6 = self.conv6_2(torch.cat([dec6, enc4
                ], 1))

        dec7 = self.conv7(F.interpolate(dec6, scale_factor=2))
        dec7 = self.conv7_2(torch.cat([dec7, enc3
                ], 1))

        dec8 = self.conv8(F.interpolate(dec7, scale_factor=2))
        dec8 = self.conv8_2(torch.cat([dec8, enc2
                ], 1))

        dec9 = self.conv9(F.interpolate(dec8, scale_factor=2))
        dec9 = self.conv9_2(torch.cat([dec9,
                enc1
                ], 1))

        dec10 = self.conv10(F.interpolate(dec9, scale_factor=2))

        return dec10

    def forward(self, x1, x2):
        dec10_0 = self.forward1(x1)
        dec10_1 = self.forward1(x2)
        x = torch.cat([dec10_0, dec10_1], 1)
        #dec10 = self.sa(dec10) * dec10
        #x = self.forward1(x)
        return self.res(x)

    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d) or isinstance(m, nn.Linear):
                m.weight.data = nn.init.kaiming_normal_(m.weight.data)
                if m.bias is not None:
                    m.bias.data.zero_()
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()

### Three layers UNet bottleneck -> siamese ###
''' # Encoder 1
        x_1 = x1
        enc1_1 = self.conv1(x_1)
        enc2_1 = self.conv2(enc1_1)
        enc3_1 = self.conv3(enc2_1)

        # Encoder 2
        x_2 = x2
        enc1_2 = self.conv1(x_2)
        enc2_2 = self.conv2(enc1_2)
        enc3_2 = self.conv3(enc2_2)

        enc3 = self.ca_skip_3(enc3_1, enc3_2)
        B_, C_, H_, W_ = enc3.shape
        enc3_i = enc3.view([B_, C_, H_*W_])
        enc3_i = self.transformer(enc3_i)
        enc3_i = enc3_i.view([B_, C_, H_, W_])
        enc3 = self.ca_skip_3(enc3_i,enc3)

        enc2 = self.ca_skip_2(enc2_1, enc2_2)
        dec8 = self.conv8(F.interpolate(enc3, scale_factor=2))
        dec8 = self.conv8_2(torch.cat([dec8, enc2
                ], 1))

        enc1 = self.ca_skip_2(enc1_1, enc1_2)
        dec9 = self.conv9(F.interpolate(dec8, scale_factor=2))
        dec9 = self.conv9_2(torch.cat([dec9,
                enc1
                ], 1))

        dec10 = self.conv10(F.interpolate(dec9, scale_factor=2))
        out = self.res(dec10)
        return out
'''


from django.test import TestCase
from django.test import Client
from django.urls import reverse
from inventory.tests.factories import (
    StyleVersionFactory,
    UserFactory
)
from inventory.tests.functions import login_as

class TestActivateTheme(TestCase):
    view_name = "activate_theme"

    def setUp(self):
        self.client = Client()
        self.user = UserFactory()
        self.version = StyleVersionFactory()
        self.active_version = StyleVersionFactory(currently_live=True,
                                                  currently_test=True)
        self.url = reverse(
            self.view_name,
            urlconf="inventory.urls",
            args=[self.version.pk, "live"])
        self.return_url = reverse("themes_list", urlconf="inventory.urls")
        self.title = "List of Themes and Versions"

    def test_no_login(self):
        response = self.client.get(self.url)
        self.assertRedirects(response,
                             "/login/?next=%s" % self.url,
                             fetch_redirect_response=False)

    def test_activate_live(self):
        login_as(self.user, self)
        response = self.client.get(self.url, follow=True)
        self.assertRedirects(response, "%s?changed_id=%d" % (
            self.return_url,
            self.version.pk))
        self.assertContains(response, self.title)
        self.assertContains(response, "Activated %s on %s" % (
            self.version,
            "live"))
        self.assertContains(
            response,
            '' % reverse(
                self.view_name,
                urlconf="inventory.urls",
                args=[self.active_version.pk, "live"]))
        self.assertNotContains(
            response,
            '' % reverse(
                self.view_name,
                urlconf="inventory.urls",
                args=[self.active_version.pk, "test"]))

    def test_activate_test(self):
        login_as(self.user, self)
        self.url = reverse(
            self.view_name,
            urlconf="inventory.urls",
            args=[self.version.pk, "test"])
        response = self.client.get(self.url, follow=True)
        self.assertRedirects(response, "%s?changed_id=%d" % (
            self.return_url,
            self.version.pk))
        self.assertContains(response, self.title)
        self.assertContains(response, "Activated %s on %s" % (
            self.version,
            "test"))
        self.assertContains(
            response,
            '' % reverse(
                self.view_name,
                urlconf="inventory.urls",
                args=[self.active_version.pk, "test"]))
        self.assertNotContains(
            response,
            '' % reverse(
                self.view_name,
                urlconf="inventory.urls",
                args=[self.active_version.pk, "live"]))

    def test_activate_bad_id(self):
        login_as(self.user, self)
        response = self.client.get(reverse(
            self.view_name,
            urlconf="inventory.urls",
            args=[self.active_version.pk+1, "live"]))
        self.assertEqual(404, response.status_code)

    def test_activate_bad_system(self):
        login_as(self.user, self)
        response = self.client.get(reverse(
            self.view_name,
            urlconf="inventory.urls",
            args=[self.version.pk, "bad"]))
        self.assertEqual(404, response.status_code)


from typing import Dict, Iterable, List, Union
import torch
from pathlib import Path
from torch.utils.data import DataLoader

def create_dataloader(dataset, batch_size, shuffle=True):
    """
    Create a DataLoader instance from a dataset.

    Parameters:
    dataset: the dataset to load
    batch_size (int): the number of samples per batch
    shuffle (bool, optional): whether to shuffle the data. Default is True.

    Returns:
    DataLoader: a DataLoader instance
    """
    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=shuffle)
    return dataloader

class StepLR:
    """
    Step Learning Rate scheduler.
    Adjusts the learning rate according to the number of epochs passed.
    """

    def __init__(self, optimizer, learning_rate: float, total_epochs: int):
        # Initializes the StepLR class.

        self.optimizer = optimizer
        self.total_epochs = total_epochs
        self.base = learning_rate

    def __call__(self, epoch):
        # Adjusts the learning rate according to the current epoch.

        if epoch < self.total_epochs * 3/10:
            lr = self.base
        elif epoch < self.total_epochs * 6/10:
            lr = self.base * 0.2
        elif epoch < self.total_epochs * 8/10:
            lr = self.base * 0.2 ** 2
        else:
            lr = self.base * 0.2 ** 3
        # Apply the learning rate to all parameter groups in the optimizer
        for param_group in self.optimizer.param_groups:
            param_group["lr"] = lr

    def lr(self) -> float:
        # Returns the current learning rate.

        return self.optimizer.param_groups[0]["lr"]

def parse_output(input_string: str) -> List[str]:
    """
    Parses a string into a list of strings.

    Parameters:
    input_string (str): The input string to parse.

    Returns:
    List[str]: The parsed output list.
    """

    input_list = input_string.split("\n")  # split the input by newline characters
    output_list = []
    for i, item in enumerate(input_list):
        item = item.lstrip(
            "0123456789. "
        )  # remove numbers and any leading whitespace
        if item:  # skip empty items
            output_list.append(item)
    return output_list

def get_device() -> torch.device:
    """
    Get the current device.

    If a CUDA GPU is available, it is used. Else, it checks for MPS compatibility.
    If not, it defaults to the CPU.

    Returns:
    torch.device: The device to use for computations.
    """

    if torch.cuda.is_available():
        device = torch.device("cuda:0")
    elif torch.has_mps:
        device = torch.device("mps")
    else:
        device = torch.device("cpu")

    print(f"Using device: {device}")

    return device

def read_json(path: Path) -> List[Dict[str, Union[str, int]]]:
    """
    Reads a JSON file and returns it as a Python list of dictionaries.

    Parameters:
    path (Path): The path to the JSON file to read.

    Returns:
    List[Dict[str, Union[str, int]]]: The content of the JSON file as a Python list of dictionaries.
    """
    with open(path, "r") as f:
        data: List[Dict[str, Union[str, int]]] = json.load(f)
    return data

def save_json(path: Path, container: Iterable) -> None:
    """
    Saves a Python iterable (such as a list or a dictionary) into a JSON file.

    Parameters:
    path (Path): The path to the JSON file to write into.
    container (Iterable): The Python iterable to write into the JSON file.
    """
    print(f"Saving json to {path}")
    with open(path, "w") as outfile:
        json.dump(container, outfile, ensure_ascii=False, indent=4)


import torch
from scipy.signal import butter, lfilter, filtfilt

from scipy import signal
import numpy as np
import pywt

from .models import *

CUTOFF_FREQ = 100
SAMPLING_RATE = 5500
INPUT_SIZE = 27500

def STFT(x, f_s=SAMPLING_RATE, nperseg=4000):
    f, t, Zxx_2_0 = signal.stft(np.abs(x), f_s, nperseg=nperseg)
    return np.abs(Zxx_2_0)

def downsample(x, interval=250):
    num_iter = int(x.shape[0]/interval)
    downsampled = np.zeros(num_iter)
    for i in range(num_iter):
        downsampled[i] = x[i*interval]
    return downsampled

def DWT(signal, f_s=SAMPLING_RATE, interval_ms=200):
    """
      https://sigmobile.org/mobicom/2015/papers/p65-wangA.pdf,
      https://pywavelets.readthedocs.io/en/latest/ref/dwt-discrete-wavelet-transform.html
    """
    samples_200_ms = int(f_s*interval_ms/1000)
    num_iter = int(signal.shape[0]/samples_200_ms)
    energies_over_time = np.zeros((12, num_iter))
    for i in range(num_iter):
        signal_snippet = signal[i*samples_200_ms:i *
                                samples_200_ms + samples_200_ms]
        detail_components = pywt.wavedec(signal_snippet, 'haar', level=12)[1:]
        energies = np.array([np.sum(np.power(components, 2))
                            for components in detail_components])
        energies_over_time[:, i] = energies
    return energies_over_time

def butterworth(signal, cutoff_freq=CUTOFF_FREQ, sampling_rate=SAMPLING_RATE, order=4):
    """Butterworth low-pass filter"""
    nyquist = 1/2 * sampling_rate
    normal_cutoff_freq = cutoff_freq/nyquist
    b, a = butter(order, normal_cutoff_freq, btype="low", analog=False)
    signal_filtered = lfilter(b, a, signal)
    return signal_filtered

def normalize_min_max(signal: np.ndarray) -> np.ndarray:
    xmin, xmax = np.amin(signal), np.amax(signal)
    return (signal - xmin ) / (xmax-xmin)

class ModelHandler:
    def __init__(self, model_path):
        # Load model dict
        self.model_dict = torch.load(model_path)

        # Initialize the model
        self.base_model = self.model_dict['base_model']
        self.model = model_types[self.base_model]()

        # Loads in the weights of the model
        self.model.load_state_dict(self.model_dict['model_state_dict'])

        # Mean and standardeviation.
        #self.mean = self.model_dict['data_mean']
        #self.std = self.model_dict['data_std']
        if self.base_model == "2DCNN":
            self.dwt_mean = self.model_dict['dwt_mean']
            self.dwt_std = self.model_dict['dwt_std']

        # Datasets and training parameters
        self.datasets = self.model_dict['datasets']
        self.epochs = self.model_dict['epoch']

        # Metrics
        self.train_loss = self.model_dict['train_loss']
        self.train_accuracy = self.model_dict['train_accuracy']
        self.test_accuracy = self.model_dict['test_accuracy']
        self.f1_score = self.model_dict['f1_score']
        self.num_train_samples = self.model_dict['num_train_samples']
        self.num_test_samples = self.model_dict['num_test_samples']
        self.batch_size = self.model_dict['batch_size']

    def get_relevant_metrics(self):
        irrelevant_metrics = {"optimizer_state_dict", "model_state_dict"}
        return {key:val for key, val in self.model_dict.items() if key not in irrelevant_metrics}

    def predict(self, data):
        '''
        Predicts the outcome of the input data. It will give the predicted class and also the certainty.
        '''
        # Absolute value of data
        data = np.abs(data)

        # Normalize and pre-process
        data = normalize_min_max(data)
        data = butterworth(data)
        if self.base_model == "1DCNN":
            data = downsample(data)
        elif self.base_model == "2DCNN":
            data = DWT(data)
            data = (data - self.dwt_mean)/self.dwt_std

        sample = torch.tensor(data, dtype=torch.float)
        # second unsqueeze to add batch dimension (needed for model).
        sample = sample.unsqueeze(dim=0).unsqueeze(dim=0)
        # Assertion to check if dimension is supported.
        assert sample.dim() <= 4, \
            'The input data has too many dimensions, max is 2 but the given input has {}.'.format(
                sample.dim() - 2)
        sample = sample.permute((1, 0, 2)) if sample.dim(
        ) == 3 else sample.permute((1, 0, 2, 3))

        self.model.eval()
        prob = self.model(sample)

        return torch.argmax(prob).detach().numpy(), np.ravel(prob.detach().numpy())


import pymongo

myclient = pymongo.MongoClient("mongodb://localhost:27017/")
test_database = myclient["testdatabase_dec2021"]
new_collection = test_database["store_collection"]

oneStore = {
    "_id" : "0239218378213",
    "store_number": "0111",
    "store_name": "Duluth Parkway",
    "address": "0111 Duluth Pwky"
}
mylist = [
    { "store_number": "0141", "store_name": "Perimeter Circle", "address": "141 Perimeter Circle" },
    { "store_number": "0131", "store_name": "Spring Hill", "address": "131 Spring Hill Pkwy" }
]
x =  new_collection.insert_one(oneStore)
# x = new_collection.insert_many(mylist)

print("Done!", x.inserted_id)
myclient.close()


import json
import datetime
from concurrent.futures import ThreadPoolExecutor, ProcessPoolExecutor
from pygtail import Pygtail

tp_executor = ThreadPoolExecutor(3)

class EventTimedStorage:
    def __init__(self, duration):
          self.duration = datetime.timedelta(minutes=duration)
          self.events = []
          self.shops = {'shops.computeruniverse' : True}
          self.not_shops = ['worker', 'discord_webhook.webhook']

    def add_event(self, timestamp, type, event, shop=None):
        self.events.append((timestamp, type, shop, event))
        self.events = list(filter(lambda ts_ev: ts_ev[0] + self.duration >= timestamp, self.events))
        if shop != None and shop not in self.shops and shop not in self.not_shops:
            self.shops[shop] = True

    def _get_events(self, timestamp):
        self.events = list(filter(lambda ts_ev: ts_ev[0] + self.duration >= timestamp, self.events))
        return self.events

    def _get_last_shops(self, timestamp):
        self.events = list(filter(lambda ts_ev: ts_ev[0] + self.duration >= timestamp, self.events))
        return [e[2] for e in self.events]

class NotifierEventStorage(EventTimedStorage):
    def __init__(self, duration):
        EventTimedStorage.__init__(self, duration)
        self.is_broke_param = {
            False : 'Notifier is work',
            True : 'Notifier is broken',
        }

    def is_notifier_broken(self, timestmp):
        return sum(e[1] == 'ERROR' for e in self._get_events(timestmp)) > 1

class MonitorEventStorage(EventTimedStorage):
    def __init__(self, duration):
        EventTimedStorage.__init__(self, duration)
        self.is_broke_param = {
            True : 'Monitor is work',
            False : 'Monitor is broken',
        }
        self.message_param = {
            True : 'Monitor restarted',
            False : 'Cannot give items last 15 minute',
        }

    # def is_monitor_broken(self, timestmp):
    #     return sum(e[1] == 'ERROR' for e in self._get_events(timestmp)) > 1

    def is_monitor_broken(self, timestmp):
        result = {}
        shop = None
        last_shops = self._get_last_shops(timestmp)
        if last_shops == []:
            return {}
        for s in self.shops:
            if s not in last_shops:
                result[s] = False
                shop = s
            else:
                result[s] = True
        return result

class LogCheck():
    def __init__(self):
        self.path = None
        self.webhook = None
        self.timestamp_last_restart = None
        self.sleep_time = datetime.timedelta(minutes=15).total_seconds()

    def notify(self, log):
        log['webhooks'] = self.webhook
        self.send_notification(**log)
        print(log)

    def restart_docker(self, shop):
        try:
            if (self.timestamp_last_restart == None
                or (datetime.datetime.now() - self.timestamp_last_restart).total_seconds() > self.sleep_time
                or shop == 'shops.computeruniverse'):
                # open("/to_check/logs/restart_state.txt", "w")
                open("Production/MessiahNotify/logs/restart_state.txt", "w")
                self.timestamp_last_restart = datetime.datetime.now()
        except Exception as e:
            raise e

    def send_notification_impl(self, info, asctime, name, webhooks, levelname,
                                message, override=None, **kwargs):
        if not override:
            override = {}

        webhook = DiscordWebhook(url=webhooks)
        embed = DiscordEmbed(title='event: ' + info + ', time: ' + asctime + ' ' + ', file name: ' + name,
                            description=message,
                            color=override.get('color', 0),
                            )
        webhook.add_embed(embed)
        webhook.execute()

    def send_notification(self, **kwargs):
        self.send_notification_impl(**kwargs)

    def return_log(self):
        tail = Pygtail(self.path, paranoid=True, offset_file='log.offset')
        try:
            line = tail.next()
            log = json.loads(line)
            return log
        except:
            tail._update_offset_file()

    def check_on_error(self):
        pass

class LogCheckNotify(LogCheck):
    def __init__(self):
        LogCheck.__init__(self)
        self.path = 'daemon.log'
        self.webhook = 'https://discord.com/api/webhooks/871656292568162334/iOwyRZMwH-rDE_dkVPpsmaUXeHPxmmGT9RaxnqiSMVrzN2tM1vDDdT2WCBe7ZdrY5Bdt'
        self.is_broke = 0

    def check_on_error(self):
        notifier_checker = NotifierEventStorage(15)
        while not is_stopped():
            log = self.return_log()
            if log != None:
                try:
                    timestamp = datetime.datetime.strptime(log['asctime'], '%Y-%m-%d %H:%M:%S,%f')
                    notifier_checker.add_event(timestamp, log['levelname'], log['message'])
                    current_time = datetime.datetime.now()
                    is_broke = notifier_checker.is_notifier_broken(current_time)
                    if is_broke != self.is_broke:
                        self.is_broke = is_broke
                        log['info'] = notifier_checker.is_broke_param[self.is_broke]
                        self.notify(log)
                except:
                    raise

class LogCheckMonitors(LogCheck):
    def __init__(self):
        LogCheck.__init__(self)
        self.path = 'daemon.log'
        self.webhook = 'https://discord.com/api/webhooks/871656292568162334/iOwyRZMwH-rDE_dkVPpsmaUXeHPxmmGT9RaxnqiSMVrzN2tM1vDDdT2WCBe7ZdrY5Bdt'
        self.shops_state = {}

    def check_on_error(self):
        checker = MonitorEventStorage(15)
        while not is_stopped():
            log = self.return_log()
            if log != None:
                try:
                    timestamp = datetime.datetime.strptime(log['asctime'], '%Y-%m-%d %H:%M:%S,%f')
                    checker.add_event(timestamp, log['levelname'], log['message'], log['name'])
                    self.shops_state = checker.shops
                    current_time = datetime.datetime.now()
                    is_state_change= checker.is_monitor_broken(current_time)
                    if is_state_change != {} and is_state_change != self.shops_state:
                        broken_shop = [s for s in is_state_change if is_state_change[s] != self.shops_state[s]]
                        checker.shops = is_state_change
                        for shop in broken_shop:
                            log['name'] = shop
                            log['levelname'] = 'INFO'
                            log['info'] = checker.is_broke_param[is_state_change[shop]]
                            log['message'] = checker.message_param[is_state_change[shop]]
                            # self.notify(log)
                            if is_state_change[shop] == False:
                                self.restart_docker(shop)
                except:
                    raise

checker_stopped = [False]

checkers = [
    # LogCheckNotify,
    LogCheckMonitors,
]

pp_executor = ProcessPoolExecutor(3)

def is_stopped():
    return checker_stopped[0]

def start_checker():
    # apis = [api_cls() for api_cls in checkers]

    # futures = [(pp_executor.submit(api.check_on_error))
    #            for api in apis]
    # return futures
    LogCheckMonitors().check_on_error()

def stop_checker():
    print('Stopping checker')
    global checker_stopped
    checker_stopped[0] = True


from multiprocessing import Manager
from multiprocessing import Process
from imutils.video import WebcamVideoStream
import imutils
from servo.cobaObject import ObjCenter
from servo.pid import PID
import argparse
import signal
import time
import sys
import cv2 as cv
import servo.servoControl as sc

panServoRange = (-90, 90)
tiltServoRange = (-70, 70)
panServo = 27
tiltServo = 17

def signal_handler(sig, frame):
    print("[INFO] ctrl + c pressed, exiting...")

    # sc.disableServo(panServo)
    # sc.disableServo(tiltServo)

    sys.exit()

def obj_center(objX, objY, centerX, centerY):
    signal.signal(signal.SIGINT, signal_handler)

    vs = WebcamVideoStream(args.camera).start()
    # time.sleep(2.0)

    obj = ObjCenter()

    while True:
        frame = vs.read()
        frame = imutils.resize(frame,width=300)
        (H, W, _) = frame.shape
        centerX.value = W // 2
        centerY.value = H // 2

        objectLoc = obj.update(frame, (centerX.value, centerY.value))
        (objX.value, objY.value) = objectLoc
        cv.imshow("frame", frame)
        print(objectLoc)
        cv.waitKey(1)
    vs.stop()

def pid_process(output, p, i, d, objCoord, centerCoord):
    signal.signal(signal.SIGINT, signal_handler)

    p = PID(p.value, i.value, d.value)
    p.initialize()

    while True:

        error =  centerCoord.value - objCoord.value

        output.value = p.update(error)

def in_range(val,  start, end):
    return (val >= start and val <= end)

def set_servos(pan, tlt):
    signal.signal(signal.SIGINT, signal_handler)

    while True:
        panAngle = -1 * pan.value
        tiltAngle = -1 * tlt.value

        if in_range(panAngle, panServoRange[0], panServoRange[1]):
            sc.setServoAngle(panServo,panAngle)

        if in_range(tiltAngle, tiltServoRange[0], tiltServoRange[1]):
            sc.setServoAngle(tiltServo, tiltAngle)

if __name__ == "__main__":
    ap = argparse.ArgumentParser()

    ap.add_argument("-c", "--camera", type=int, default=0, help="change camera")

    args = ap.parse_args()

    with Manager() as manager:
        centerX = manager.Value("i", 0)
        centerY = manager.Value("i", 0)

        objX = manager.Value("i", 0)
        objY = manager.Value("i", 0)

        pan = manager.Value("i", 0)
        tlt = manager.Value("i", 0)

        panP = manager.Value("f", 0.09)
        panI = manager.Value("f", 0.08)
        panD = manager.Value("f", 0.002)

        tiltP = manager.Value("f", 0.11)
        tiltI = manager.Value("f", 0.10)
        tiltD = manager.Value("f", 0.002)

        processObjectCenter = Process(target=obj_center, args=(objX, objY, centerX, centerY) )
        processPanning = Process(target=pid_process, args=(pan, panP, panI, panD, objX, centerX))
        processTilting = Process(target=pid_process, args=(tlt, tiltP, tiltI, tiltD, objY, centerY))
        processSetServos = Process(target=set_servos, args=(pan, tlt))

        processObjectCenter.start()
        # processPanning.start()
        processTilting.start()
        processSetServos.start()

        processObjectCenter.join()
        # processPanning.join()
        processTilting.join()
        processSetServos.join()

        # sc.disableServo(panServo)
        # sc.disableServo(tiltServo)


import math
from typing import Optional

import torch

class AsyncGradientClipping:
    def __init__(
        self,
        device: torch.device,
        comm_group: Optional[torch.distributed.ProcessGroup] = None,
        norm_type: float = 2.0,
    ) -> None:
        self.comm_group = (
            comm_group if comm_group is not None else torch.distributed.group.WORLD
        )
        self.norm_type = norm_type
        self._norm_acc = torch.tensor(0.0, device=device)

    def get_clip_scale(self, max_norm: float, eps: float = 1e-6) -> float:
        grad_norm_acc = self._norm_acc.item()
        grad_norm = math.pow(grad_norm_acc + eps, 1.0 / self.norm_type)
        clip_scale = min(max_norm / grad_norm, 1.0)
        self._norm_acc.zero_()
        return clip_scale

def update_norm_from_buckets(
    state: AsyncGradientClipping,
    bucket: torch.distributed.GradBucket,
) -> torch.futures.Future[torch.Tensor]:
    grad = bucket.buffer()
    world_size = state.comm_group.size()
    grad.div_(world_size)

    def _acc_grad_norm(fut: torch.futures.Future[torch.Tensor]) -> torch.Tensor:
        synced_grad = fut.value()[0]  # fut.value() type: List[torch.Tensor]
        grad_to_power_p = synced_grad.detach().pow(state.norm_type)
        state._norm_acc += grad_to_power_p.sum()
        return synced_grad

    return (
        torch.distributed.all_reduce(grad, group=state.comm_group, async_op=True)
        .get_future()
        .then(_acc_grad_norm)
    )


from flask import request, render_template
from flask_mail import Message
from esocrow.models import FeedBackForm

from keys import EMAIL

from app import app, mail

@app.route('/contact', methods=['GET', 'POST'])
def registration():
    """
    Checks to see if data is posted or not, if so submits the form and emails
    it to site admin else displays contact form
    """
    form = FeedBackForm(request.form)

    if request.method == 'POST' and form.validate_on_submit():
        name = form.name.data
        email_address = form.email.data
        feedback = form.feedback.data
        message = form.message.data
        msg = Message(f"ESO-Crow: {feedback}", sender=(name, email_address), reply_to=email_address,
                      recipients=[EMAIL])
        msg.body = message
        mail.send(msg)

        return render_template('contact/feedback_sent.html')
    return render_template('contact/contact.html', form=form)

@ app.route('/feedback_sent')
def email_sent():
    """ Returns the template for feedback_sent.html """
    return render_template('contact/feedback_sent.html')


"""issue_idx2

Revision ID: 32bda2f94d52
Revises: 0c5432eab8c5
Create Date: 2023-04-28 11:20:55.298185

"""
from alembic import op

# Polar Custom Imports

# revision identifiers, used by Alembic.
revision = "32bda2f94d52"
down_revision = "0c5432eab8c5"
branch_labels: tuple[str] | None = None
depends_on: tuple[str] | None = None

def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_index(
        "idx_issues_id_closed_at", "issues", ["id", "issue_closed_at"], unique=False
    )
    # ### end Alembic commands ###

def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_index("idx_issues_id_closed_at", table_name="issues")
    # ### end Alembic commands ###


# -*- coding: utf-8 -*-
# Question 4
# Author: Weining Qu (wq2155)

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from evaluation import *

def NBClassifier(data, target):
    Y = data[target].values
    X_all = data[[col for col in data.columns if col != target]]

    # Categorize data by target class
    X_byclass = {}
    Y_byclass = {}
    phi = {}
    for i in np.unique(Y):
        X_byclass[i] = X_all[Y == i]
        Y_byclass[i] = Y[Y == i]
        phi[i] = (len(Y_byclass[i])+1)/ (len(Y)+2)

    # Calculate count probabilities of features
    prob_X = {}
    for i in np.unique(Y):
        prob_X[i] = {}
        for j in range(X_all.shape[1]):
            X_j = X_byclass[i].iloc[:,j]
            prob_X[i][j] = X_j.value_counts()/len(X_j)

        # Model "age" feature as continuous variable following Gaussian distribution
        prob_X[i][1] = [np.mean(X_byclass[i].iloc[:, 1]), np.var(X_byclass[i].iloc[:, 1])]
    return phi, prob_X

def NBpredict(testX, phi, prob_X):
    def GaussianProb(x, mean, var):
        # (np.exp(-(x-mean)**2/(2*var)))*(1/np.sqrt(2*np.pi*var))
        gaussian_prob = np.exp(-(x-mean)**2/(2*var)) * (1/np.sqrt(2*np.pi*var))
        return gaussian_prob

    predY = []
    predY_prob = []
    for i in range(testX.shape[0]):
        X_i = testX.iloc[i,:].values
        Cond_pos = 1
        for s in range(len(X_i)):
            if (s != 1) and (X_i[s] in prob_X[1][s].index.tolist()):
                Cond_pos = Cond_pos * prob_X[1][s][X_i[s]]
            else:  # "age"
                Cond_pos = Cond_pos * GaussianProb(X_i[s],prob_X[1][s][0], prob_X[1][s][1])
        prob_pos = phi[1] * Cond_pos

        Cond_neg = 1
        for s in range(len(X_i)):
            if (s != 1) and (X_i[s] in prob_X[0][s].index.tolist()):
                Cond_neg = Cond_neg * prob_X[0][s][X_i[s]]
            else:  # "age"
                Cond_neg = Cond_neg * GaussianProb(X_i[s], prob_X[1][s][0], prob_X[1][s][1])
        prob_neg = phi[0] * Cond_neg

        if prob_pos >= prob_neg:
            predY.append(1)
            predY_prob.append(1-prob_pos)
        else:
            predY.append(0)
            predY_prob.append(prob_neg)
    predY = np.array(predY)
    predY_prob = np.array(predY_prob)

    return predY, predY_prob

################ Evaluation ###########################################################
def Accuracy(pred, true):
    return sum(pred == true) / len(pred)

def Test_trainsize(traindata, testdata, target):
    accu_pair = []
    for s in range(1,11):
        N = int(traindata.shape[0] * s/10)
        tr_part = traindata.iloc[:N,:]
        # print(tr_part)
        phi, prob_X = NBClassifier(tr_part, target)
        testX = testdata[[col for col in testdata.columns if col != target]]
        testY = testdata[target]
        predY, predprobY = NBpredict(testX, phi, prob_X)
        accuracy_test = Accuracy(predY, testY)
        accu_pair.append((N, accuracy_test))
    x, y = zip(*accu_pair)
    y_max = np.argmax(y)
    plt.figure()
    plt.plot(x,y)
    plt.scatter(x[y_max],y[y_max], c='red')
    plt.xlabel("Training sample size")
    plt.ylabel("Accuracy")
    plt.title("Naive Bayes")
    plt.savefig("./Fig/NB_trainsizetest.png")
    plt.show()

    return accu_pair

if __name__ == "__main__":
    traindata = pd.read_csv("propublicaTrain.csv", sep=',')
    testdata = pd.read_csv("propublicaTest.csv", sep=',')
    target = 'two_year_recid'

    phi, prob_X = NBClassifier(traindata, target)
    testX = testdata[[col for col in testdata.columns if col != target]]
    testY = testdata[target]

    predY, predprobY = NBpredict(testX, phi, prob_X)
    accuracy_test = Accuracy(predY, testY)
    print("Testdata accuracy: {:.4f}".format(accuracy_test))

    plot_roc(testY, predprobY, "NB")  # ROC Curve and AUC

    accu_diffsize = Test_trainsize(traindata, testdata, target)  # Test training sample size effect

    ############## Fairness analysis ################################################################
    df = pd.concat([testX['race'], pd.Series(predY), testY], axis=1)
    df.columns = ['race', 'ypred', 'ytrue']
    # Demographic parity
    df_race0 = df[df['race'] == 0]
    df_race1 = df[df['race'] == 1]
    p_race0_0 = (df_race0.shape[0] - df_race0['ypred'].sum()) / df_race0.shape[0]
    p_race0_1 = df_race0['ypred'].sum() / df_race0.shape[0]
    p_race1_0 = (df_race1.shape[0] - df_race1['ypred'].sum()) / df_race1.shape[0]
    p_race1_1 = df_race1['ypred'].sum() / df_race1.shape[0]

    # Equalized Odds (EO)
    df_race0_y0 = df_race0[df_race0['ytrue'] == 0]
    eo_race0_00 = (df_race0_y0.shape[0] - df_race0_y0['ypred'].sum()) / df_race0_y0.shape[0]
    df_race0_y1 = df_race0[df_race0['ytrue'] == 1]
    eo_race0_11 = df_race0_y1['ypred'].sum() / df_race0_y1.shape[0]
    df_race1_y0 = df_race1[df_race1['ytrue'] == 0]
    eo_race1_00 = (df_race1_y0.shape[0] - df_race1_y0['ypred'].sum()) / df_race1_y0.shape[0]
    df_race1_y1 = df_race1[df_race1['ytrue'] == 1]
    eo_race1_11 = df_race1_y1['ypred'].sum() / df_race1_y1.shape[0]

    # Predictive Parity (PP)
    df_race0_hy0 = df_race0[df_race0['ypred'] == 0]
    pp_race0_00 = (df_race0_hy0.shape[0] - df_race0_hy0['ytrue'].sum()) / df_race0_hy0.shape[0]
    df_race0_hy1 = df_race0[df_race0['ypred'] == 1]
    pp_race0_11 = df_race0_hy1['ytrue'].sum() / df_race0_hy1.shape[0]
    df_race1_hy0 = df_race1[df_race1['ypred'] == 0]
    pp_race1_00 = (df_race1_hy0.shape[0] - df_race1_hy0['ytrue'].sum()) / df_race1_hy0.shape[0]
    df_race1_hy1 = df_race1[df_race1['ypred'] == 1]
    pp_race1_11 = df_race1_y1['ytrue'].sum() / df_race1_y1.shape[0]


"""
Base class
"""
import logging
from abc import abstractmethod, ABCMeta
from functools import lru_cache

from mininet.topo import Topo

from distriopt import VirtualNetwork
from distriopt.constants import *

_log = logging.getLogger(__name__)

class PackingSolver(object, metaclass=ABCMeta):
    def __init__(self, virtual, physical):
        """"""
        self.virtual = (
            VirtualNetwork.from_mininet(virtual)
            if isinstance(virtual, Topo)
            else virtual
        )
        self.physical = physical
        self.solution = None
        self.status = NotSolved
        self.lb = 0

    def _get_ub(self, vm_type):
        """Return an upper bound on the maximum number of EC2 instances of type vm_type needed to pack all the nodes"""

        cpu_cores_instance, memory_instance = (
            self.physical.cores(vm_type),
            self.physical.memory(vm_type),
        )
        n_instances_needed = remaining_cpu_cores = remaining_memory = 0
        for u in self.virtual.nodes():
            req_cores, req_memory = (
                self.virtual.req_cores(u),
                self.virtual.req_memory(u),
            )
            if req_cores <= remaining_cpu_cores and req_memory <= remaining_memory:
                remaining_cpu_cores -= req_cores
                remaining_memory -= req_memory
            elif req_cores <= cpu_cores_instance and req_memory <= memory_instance:
                remaining_cpu_cores = cpu_cores_instance - req_cores
                remaining_memory = memory_instance - req_memory
                n_instances_needed += 1
        return n_instances_needed

    def _get_feasible_instances(self, u):
        """Return a set with the instances types feasible for node u
        """
        return set(
            vm_type
            for vm_type in self.physical.vm_options
            if self.virtual.req_cores(u) <= self.physical.cores(vm_type)
            and self.virtual.req_memory(u) <= self.physical.memory(vm_type)
        )

    @lru_cache(maxsize=256)
    def _get_cheapest_feasible(self, cores, memory):
        """Given a demand in terms of number of cores and memory return the cheapest EC2 instance with enough resources.
        """
        # if (cores > self.physical.cores(self.vm_max_cores) or memory > self.physical.memory(self.vm_max_cores)) \
        #        and (
        #        cores > self.physical.cores(self.vm_max_memory) or memory > self.physical.memory(self.vm_max_memory)):
        #    return None

        return min(
            (
                (vm, self.physical.hourly_cost(vm))
                for vm in self.physical.vm_options
                if cores <= self.physical.cores(vm)
                and memory <= self.physical.memory(vm)
            ),
            key=lambda x: x[1],
        )[0]

    @abstractmethod
    def solve(self, **kwargs):
        """This method must be implemented"""


import MySQLdb

from datetime import datetime

class WikiaDB:
    def __init__(self):
        self.db = MySQLdb.connect(user="psawaya",db="wikia")

        c = self.db.cursor()

        #Create files table, if it doesn't already exist
        c.execute("""CREATE TABLE if not exists `wikia`.`wikia_files` (
          `id` int(11) NOT NULL AUTO_INCREMENT,
          `wiki_name` tinytext,
          `url` tinytext,
          `downloaded` tinyint(1) DEFAULT NULL,
          `last_modified` timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
          `filesize` bigint(20) DEFAULT NULL,
          PRIMARY KEY (`id`)
        ) ENGINE=MyISAM AUTO_INCREMENT=103 DEFAULT CHARSET=utf8;""")

        #Create dirs table, if it doesn't already exist
        c.execute("""CREATE TABLE if not exists `wikia`.`wikia_dirs` (
          `id` INT NOT NULL AUTO_INCREMENT,
          `url` TINYTEXT NOT NULL,
          `lastvisited` TIMESTAMP,
          PRIMARY KEY (`id`)
        ) CHARACTER SET utf8;""")

    def recordDirAsScraped(self,url):
        c = self.db.cursor()

        c.execute("INSERT INTO wikia_dirs (url,lastvisited) VALUES ('%s','%s')" % (url,datetime.now()))

    def dirAlreadyVisited(self,dir_url):
        c = self.db.cursor()

        c.execute("SELECT id FROM wikia_dirs WHERE url = '%s'" % dir_url)

        return c.fetchone() is not None

    def recordArchiveFile(self,wiki_name,url,downloaded,last_modified,filesize_str):
        c = self.db.cursor()

        # print "last_modified: %i" % last_modified

        c.execute("""INSERT INTO wikia_files (wiki_name,url,downloaded,last_modified,filesize)
            VALUES ('%s','%s',%i,'%s',%i)"""  % (wiki_name,url,int(downloaded), datetime.fromtimestamp(last_modified),self.parseFilesizeString(filesize_str)))

    def fileNotArchivedOrOld(self,url,last_modified):
        c = self.db.cursor()

        c.execute("""SELECT id FROM wikia_files WHERE
            url = '%s' AND (downloaded = 1 AND last_modified >= '%s')""" % (url, datetime.fromtimestamp(last_modified)))

        return c.fetchone() is None

    @staticmethod
    def parseFilesizeString(filesize_str):
        try:
            mult_factor = {'K' : 1000, 'M' : 1000000, 'G' : 1000000000}[filesize_str[-1].upper()]
        except KeyError:
            mult_factor = 1

        return float(filesize_str[:-1]) * mult_factor


from pip import List

class Solution:
    def topKFrequent(self, nums: List[int], k: int) -> List[int]:
        self.dict = {}
        s = set(nums)
        res = []

        for num in s:
            self.dict[num] = nums.count(num)

        res = sorted(self.dict.keys(), key=lambda item: self.dict.get(item), reverse=True)

        return res[:k]

백준/해시 테이블/상위 K 빈도 요소.py

상위 K 빈도 요소.py