import re
import torch
import torch.nn as nn
from torch.nn import init
import functools
from torch.optim import lr_scheduler
import torch.nn.functional as F
from torch.nn.utils.spectral_norm import spectral_norm as SpectralNorm
######################################################################################
# base function for network structure
######################################################################################

# def init_weights(net, init_type='normal', gain=0.02):
#     """Get different initial method for the network weights"""
#     def init_func(m):
#         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, gain)
#             elif init_type == 'xavier':
#                 init.xavier_normal_(m.weight.data, gain=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=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:
#             init.normal_(m.weight.data, 1.0, 0.02)
#             init.constant_(m.bias.data, 0.0)

#     print('initialize network with %s' % init_type)
#     net.apply(init_func)

# Creates SPADE normalization layer based on the given configuration
# SPADE consists of two steps. First, it normalizes the activations using
# your favorite normalization method, such as Batch Norm or Instance Norm.
# Second, it applies scale and bias to the normalized output, conditioned on
# the segmentation map.
# The format of |config_text| is spade(norm)(ks), where
# (norm) specifies the type of parameter-free normalization.
#       (e.g. syncbatch, batch, instance)
# (ks) specifies the size of kernel in the SPADE module (e.g. 3x3)
# Example |config_text| will be spadesyncbatch3x3, or spadeinstance5x5.
# Also, the other arguments are
# |norm_nc|: the #channels of the normalized activations, hence the output dim of SPADE
# |label_nc|: the #channels of the input semantic feature map, hence the input dim of SPADE

def calc_mean_std(feat, eps=1e-5):
    # eps is a small value added to the variance to avoid divide-by-zero.
    size = feat.size()
    assert (len(size) == 4)
    N, C = size[:2]
    feat_var = feat.view(N, C, -1).var(dim=2) + eps
    feat_std = feat_var.sqrt().view(N, C, 1, 1)
    feat_mean = feat.view(N, C, -1).mean(dim=2).view(N, C, 1, 1)
    return feat_mean, feat_std

def adaptive_instance_normalization(content_feat, style_feat):
    assert (content_feat.size()[:2] == style_feat.size()[:2])
    size = content_feat.size()
    style_mean, style_std = calc_mean_std(style_feat)
    content_mean, content_std = calc_mean_std(content_feat)

    normalized_feat = (content_feat - content_mean.expand(
        size)) / content_std.expand(size)
    return normalized_feat * style_std.expand(size) + style_mean.expand(size)

def _calc_feat_flatten_mean_std(feat):
    # takes 3D feat (C, H, W), return mean and std of array within channels
    assert (feat.size()[0] == 3)
    assert (isinstance(feat, torch.FloatTensor))
    feat_flatten = feat.view(3, -1)
    mean = feat_flatten.mean(dim=-1, keepdim=True)
    std = feat_flatten.std(dim=-1, keepdim=True)
    return feat_flatten, mean, std

def get_norm_layer(norm_type='batch'):
    """Get the normalization layer for the networks"""
    if norm_type == 'batch':
        norm_layer = functools.partial(nn.BatchNorm2d, momentum=0.1, affine=True)
    elif norm_type == 'instance':
        norm_layer = functools.partial(nn.InstanceNorm2d, affine=True)
    elif norm_type == 'adain':
        norm_layer = functools.partial(ADAIN)
    elif norm_type == 'spade':
        norm_layer = functools.partial(SPADE, config_text='spadeinstance3x3')
    elif norm_type == 'none':
        norm_layer = None
    else:
        raise NotImplementedError('normalization layer [%s] is not found' % norm_type)

    if norm_type != 'none':
        norm_layer.__name__ = norm_type

    return norm_layer

def get_nonlinearity_layer(activation_type='PReLU'):
    """Get the activation layer for the networks"""
    if activation_type == 'ReLU':
        nonlinearity_layer = nn.ReLU()
    elif activation_type == 'SELU':
        nonlinearity_layer = nn.SELU()
    elif activation_type == 'LeakyReLU':
        nonlinearity_layer = nn.LeakyReLU(0.1)
    elif activation_type == 'PReLU':
        nonlinearity_layer = nn.PReLU()
    else:
        raise NotImplementedError('activation layer [%s] is not found' % activation_type)
    return nonlinearity_layer

def get_scheduler(optimizer, opt):
    """Get the training learning rate for different epoch"""
    if opt.lr_policy == 'lambda':
        def lambda_rule(epoch):
            lr_l = 1.0 - max(0, epoch+1+1+opt.iter_count-opt.niter) / float(opt.niter_decay+1)
            return lr_l
        scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda_rule)
    elif opt.lr_policy == 'step':
        scheduler = lr_scheduler.StepLR(optimizer, step_size=opt.lr_decay_iters, gamma=0.1)
    elif opt.lr_policy == 'exponent':
        scheduler = lr_scheduler.ExponentialLR(optimizer, gamma=0.95)
    else:
        raise NotImplementedError('learning rate policy [%s] is not implemented', opt.lr_policy)
    return scheduler

def print_network(net):
    """print the network"""
    num_params = 0
    for param in net.parameters():
        num_params += param.numel()
    print(net)
    print('total number of parameters: %.3f M' % (num_params/1e6))

def init_net(net, init_type='normal', activation='relu', gpu_ids=[]):
    """print the network structure and initial the network"""
    print_network(net)

    if len(gpu_ids) > 0:
        assert(torch.cuda.is_available())
        net.cuda()
        net = torch.nn.DataParallel(net, gpu_ids)
    init_weights(net, init_type)
    return net

def _freeze(*args):
    """freeze the network for forward process"""
    for module in args:
        if module:
            for p in module.parameters():
                p.requires_grad = False

def _unfreeze(*args):
    """ unfreeze the network for parameter update"""
    for module in args:
        if module:
            for p in module.parameters():
                p.requires_grad = True

def _freeze_flow(*args):
    """ unfreeze the network for parameter update"""
    for module in args:
        if module:
            for p in module.parameters():
                p.requires_grad = True

def spectral_norm(module, use_spect=True):
    """use spectral normal layer to stable the training process"""
    if use_spect:
        return SpectralNorm(module)
    else:
        return module

def coord_conv(input_nc, output_nc, use_spect=False, use_coord=False, with_r=False, **kwargs):
    """use coord convolution layer to add position information"""
    if use_coord:
        return CoordConv(input_nc, output_nc, with_r, use_spect, **kwargs)
    else:
        return spectral_norm(nn.Conv2d(input_nc, output_nc, **kwargs), use_spect)

######################################################################################
# Network basic function
######################################################################################
class AddCoords(nn.Module):
    """
    Add Coords to a tensor
    """
    def __init__(self, with_r=False):
        super(AddCoords, self).__init__()
        self.with_r = with_r

    def forward(self, x):
        """
        :param x: shape (batch, channel, x_dim, y_dim)
        :return: shape (batch, channel+2, x_dim, y_dim)
        """
        B, _, x_dim, y_dim = x.size()

        # coord calculate
        xx_channel = torch.arange(x_dim).repeat(B, 1, y_dim, 1).type_as(x)
        yy_cahnnel = torch.arange(y_dim).repeat(B, 1, x_dim, 1).permute(0, 1, 3, 2).type_as(x)
        # normalization
        xx_channel = xx_channel.float() / (x_dim-1)
        yy_cahnnel = yy_cahnnel.float() / (y_dim-1)
        xx_channel = xx_channel * 2 - 1
        yy_cahnnel = yy_cahnnel * 2 - 1

        ret = torch.cat([x, xx_channel, yy_cahnnel], dim=1)

        if self.with_r:
            rr = torch.sqrt(xx_channel ** 2 + yy_cahnnel ** 2)
            ret = torch.cat([ret, rr], dim=1)

        return ret

class CoordConv(nn.Module):
    """
    CoordConv operation
    """
    def __init__(self, input_nc, output_nc, with_r=False, use_spect=False, **kwargs):
        super(CoordConv, self).__init__()
        self.addcoords = AddCoords(with_r=with_r)
        input_nc = input_nc + 2
        if with_r:
            input_nc = input_nc + 1
        self.conv = spectral_norm(nn.Conv2d(input_nc, output_nc, **kwargs), use_spect)

    def forward(self, x):
        ret = self.addcoords(x)
        ret = self.conv(ret)

        return ret

class EncoderBlock(nn.Module):
    def __init__(self, input_nc, output_nc, norm_layer=nn.BatchNorm2d, nonlinearity=nn.LeakyReLU(),
                 use_spect=False, use_coord=False):
        super(EncoderBlock, self).__init__()

        kwargs_down = {'kernel_size': 4, 'stride': 2, 'padding': 1}
        kwargs_fine = {'kernel_size': 3, 'stride': 1, 'padding': 1}

        conv1 = coord_conv(input_nc,  output_nc, use_spect, use_coord, **kwargs_down)
        conv2 = coord_conv(output_nc, output_nc, use_spect, use_coord, **kwargs_fine)

        if type(norm_layer) == type(None):
            self.model = nn.Sequential(conv1, nonlinearity, nonlinearity, conv2)
        else:
            self.model = nn.Sequential(norm_layer(input_nc),  nonlinearity, conv1,
                                       norm_layer(output_nc), nonlinearity, conv2)

    def forward(self, x):
        out = self.model(x)
        return out

class NormalBlock(nn.Module):
    def __init__(self, input_nc, output_nc, norm_layer=nn.BatchNorm2d, nonlinearity=nn.LeakyReLU(),
                 use_spect=False, use_coord=False):
        super(NormalBlock, self).__init__()

        kwargs_fine = {'kernel_size': 3, 'stride': 1, 'padding': 1}

        conv = coord_conv(input_nc, output_nc, use_spect, use_coord, **kwargs_fine)

        if type(norm_layer) == type(None):
            self.model = nn.Sequential(nonlinearity, conv)
        else:
            self.model = nn.Sequential(norm_layer(input_nc),  nonlinearity, conv)

    def forward(self, x):
        out = self.model(x)
        return out

class ResBlock(nn.Module):
    """
    Define an Residual block for different types
    """
    def __init__(self, input_nc, output_nc=None, hidden_nc=None, norm_layer=nn.BatchNorm2d, nonlinearity= nn.LeakyReLU(),
                learnable_shortcut=False, use_spect=False, use_coord=False):
        super(ResBlock, self).__init__()

        hidden_nc = input_nc if hidden_nc is None else hidden_nc
        output_nc = input_nc if output_nc is None else output_nc
        self.learnable_shortcut = True if input_nc != output_nc else learnable_shortcut

        kwargs = {'kernel_size': 3, 'stride': 1, 'padding': 1}
        kwargs_short = {'kernel_size': 1, 'stride': 1, 'padding': 0}

        conv1 = coord_conv(input_nc, hidden_nc, use_spect, use_coord, **kwargs)
        conv2 = coord_conv(hidden_nc, output_nc, use_spect, use_coord, **kwargs)

        if type(norm_layer) == type(None):
            self.model = nn.Sequential(nonlinearity, conv1, nonlinearity, conv2,)
        else:
            self.model = nn.Sequential(norm_layer(input_nc), nonlinearity, conv1,
                                       norm_layer(hidden_nc), nonlinearity, conv2,)

        if self.learnable_shortcut:
            bypass = coord_conv(input_nc, output_nc, use_spect, use_coord, **kwargs_short)
            self.shortcut = nn.Sequential(bypass,)

    def forward(self, x):
        if self.learnable_shortcut:
            out = self.model(x) + self.shortcut(x)
        else:
            out = self.model(x) + x
        return out

class ResBlocks(nn.Module):
    """docstring for ResBlocks"""
    def __init__(self, num_blocks, input_nc, output_nc=None, hidden_nc=None, norm_layer=nn.BatchNorm2d,
                 nonlinearity= nn.LeakyReLU(), learnable_shortcut=False, use_spect=False, use_coord=False):
        super(ResBlocks, self).__init__()
        hidden_nc = input_nc if hidden_nc is None else hidden_nc
        output_nc = input_nc if output_nc is None else output_nc

        self.model=[]
        if num_blocks == 1:
            self.model += [ResBlock(input_nc, output_nc, hidden_nc,
                           norm_layer, nonlinearity, learnable_shortcut, use_spect, use_coord)]

        else:
            self.model += [ResBlock(input_nc, hidden_nc, hidden_nc,
                           norm_layer, nonlinearity, learnable_shortcut, use_spect, use_coord)]
            for i in range(num_blocks-2):
                self.model += [ResBlock(hidden_nc, hidden_nc, hidden_nc,
                               norm_layer, nonlinearity, learnable_shortcut, use_spect, use_coord)]
            self.model += [ResBlock(hidden_nc, output_nc, hidden_nc,
                            norm_layer, nonlinearity, learnable_shortcut, use_spect, use_coord)]

        self.model = nn.Sequential(*self.model)

    def forward(self, inputs):
        return self.model(inputs)

class ResBlockDecoder(nn.Module):
    """
    Define a decoder block
    """
    def __init__(self, input_nc, output_nc, hidden_nc=None, norm_layer=nn.BatchNorm2d, nonlinearity= nn.LeakyReLU(),
                 use_spect=False, use_coord=False):
        super(ResBlockDecoder, self).__init__()

        hidden_nc = input_nc if hidden_nc is None else hidden_nc

        conv1 = spectral_norm(nn.Conv2d(input_nc, hidden_nc, kernel_size=3, stride=1, padding=1), use_spect)
        conv2 = spectral_norm(nn.ConvTranspose2d(hidden_nc, output_nc, kernel_size=3, stride=2, padding=1, output_padding=1), use_spect)
        bypass = spectral_norm(nn.ConvTranspose2d(input_nc, output_nc, kernel_size=3, stride=2, padding=1, output_padding=1), use_spect)

        if type(norm_layer) == type(None):
            self.model = nn.Sequential(nonlinearity, conv1, nonlinearity, conv2,)
        else:
            self.model = nn.Sequential(norm_layer(input_nc), nonlinearity, conv1, norm_layer(hidden_nc), nonlinearity, conv2,)

        self.shortcut = nn.Sequential(bypass)

    def forward(self, x):
        out = self.model(x) + self.shortcut(x)
        return out

class ResBlockEncoder(nn.Module):
    """
    Define a decoder block
    """
    def __init__(self, input_nc, output_nc, hidden_nc=None, norm_layer=nn.BatchNorm2d, nonlinearity= nn.LeakyReLU(),
                 use_spect=False, use_coord=False):
        super(ResBlockEncoder, self).__init__()

        hidden_nc = input_nc if hidden_nc is None else hidden_nc

        conv1 = spectral_norm(nn.Conv2d(input_nc, hidden_nc, kernel_size=3, stride=1, padding=1), use_spect)
        conv2 = spectral_norm(nn.Conv2d(hidden_nc, output_nc, kernel_size=4, stride=2, padding=1), use_spect)
        bypass = spectral_norm(nn.Conv2d(input_nc, output_nc, kernel_size=1, stride=1, padding=0), use_spect)

        if type(norm_layer) == type(None):
            self.model = nn.Sequential(nonlinearity, conv1, nonlinearity, conv2,)
        else:
            self.model = nn.Sequential(norm_layer(input_nc), nonlinearity, conv1,
                                       norm_layer(hidden_nc), nonlinearity, conv2,)
        self.shortcut = nn.Sequential(nn.AvgPool2d(kernel_size=2, stride=2),bypass)

    def forward(self, x):
        out = self.model(x) + self.shortcut(x)
        return out

class Output(nn.Module):
    """
    Define the output layer
    """
    def __init__(self, input_nc, output_nc, kernel_size = 3, norm_layer=nn.BatchNorm2d, nonlinearity= nn.LeakyReLU(),
                 use_spect=False, use_coord=False):
        super(Output, self).__init__()

        kwargs = {'kernel_size': kernel_size, 'padding':0, 'bias': True}

        self.conv1 = coord_conv(input_nc, output_nc, use_spect, use_coord, **kwargs)

        if type(norm_layer) == type(None):
            self.model = nn.Sequential(nonlinearity, nn.ReflectionPad2d(int(kernel_size/2)), self.conv1, nn.Tanh())
        else:
            self.model = nn.Sequential(norm_layer(input_nc), nonlinearity, nn.ReflectionPad2d(int(kernel_size / 2)), self.conv1, nn.Tanh())

    def forward(self, x):
        out = self.model(x)

        return out

class Jump(nn.Module):
    """
    Define the output layer
    """
    def __init__(self, input_nc, output_nc, kernel_size = 3, norm_layer=nn.BatchNorm2d, nonlinearity= nn.LeakyReLU(),
                 use_spect=False, use_coord=False):
        super(Jump, self).__init__()

        kwargs = {'kernel_size': kernel_size, 'padding':0, 'bias': True}

        self.conv1 = coord_conv(input_nc, output_nc, use_spect, use_coord, **kwargs)

        if type(norm_layer) == type(None):
            self.model = nn.Sequential(nonlinearity, nn.ReflectionPad2d(int(kernel_size/2)), self.conv1)
        else:
            self.model = nn.Sequential(norm_layer(input_nc), nonlinearity, nn.ReflectionPad2d(int(kernel_size / 2)), self.conv1)

    def forward(self, x):
        out = self.model(x)
        return out

class LinearBlock(nn.Module):
    def __init__(self, input_nc, output_nc, norm_layer=nn.BatchNorm2d,  nonlinearity= nn.LeakyReLU(),
                 use_spect=False):
        super(LinearBlock, self).__init__()
        use_bias = True

        self.fc = spectral_norm(nn.Linear(input_nc, output_nc, bias=use_bias), use_spect)

        if type(norm_layer) == type(None):
            self.model = nn.Sequential(nonlinearity, self.fc)
        else:
            self.model = nn.Sequential(norm_layer(input_nc), nonlinearity, self.fc)

    def forward(self, x):
        out = self.model(x)
        return out

def extract_image_patches(images, ksizes, strides, rates, padding='same'):
    """
    Extract patches from images and put them in the C output dimension.
    :param padding:
    :param images: [batch, channels, in_rows, in_cols]. A 4-D Tensor with shape
    :param ksizes: [ksize_rows, ksize_cols]. The size of the sliding window for
     each dimension of images
    :param strides: [stride_rows, stride_cols]
    :param rates: [dilation_rows, dilation_cols]
    :return: A Tensor
    """
    assert len(images.size()) == 4
    assert padding in ['same', 'valid']

    if padding == 'same':
        images = same_padding(images, ksizes, strides, rates)
    elif padding == 'valid':
        pass
    else:
        raise NotImplementedError('Unsupported padding type: {}.\
                Only "same" or "valid" are supported.'.format(padding))

    unfold = torch.nn.Unfold(kernel_size=ksizes,
                             dilation=rates,
                             padding=0,
                             stride=strides)
    patches = unfold(images)
    return patches  # [N, C*k*k, L], L is the total number of such blocks

def reduce_mean(x, axis=None, keepdim=False):
    if not axis:
        axis = range(len(x.shape))
    for i in sorted(axis, reverse=True):
        x = torch.mean(x, dim=i, keepdim=keepdim)
    return x

def reduce_sum(x, axis=None, keepdim=False):
    if not axis:
        axis = range(len(x.shape))
    for i in sorted(axis, reverse=True):
        x = torch.sum(x, dim=i, keepdim=keepdim)
    return x

def same_padding(images, ksizes, strides, rates):
    assert len(images.size()) == 4
    batch_size, channel, rows, cols = images.size()
    out_rows = (rows + strides[0] - 1) // strides[0]
    out_cols = (cols + strides[1] - 1) // strides[1]
    effective_k_row = (ksizes[0] - 1) * rates[0] + 1
    effective_k_col = (ksizes[1] - 1) * rates[1] + 1
    padding_rows = max(0, (out_rows-1)*strides[0]+effective_k_row-rows)
    padding_cols = max(0, (out_cols-1)*strides[1]+effective_k_col-cols)
    # Pad the input
    padding_top = int(padding_rows / 2.)
    padding_left = int(padding_cols / 2.)
    padding_bottom = padding_rows - padding_top
    padding_right = padding_cols - padding_left
    paddings = (padding_left, padding_right, padding_top, padding_bottom)
    images = torch.nn.ZeroPad2d(paddings)(images)
    return images

class TCAM(nn.Module):
    def __init__(self, ksize=3, stride=1, rate=2, in_channel=512):
        super(TCAM, self).__init__()
        self.ksize = ksize
        self.stride = stride
        self.rate = rate
        self.get_mask = nn.Sequential(nn.Conv2d(in_channel, 1, kernel_size=3, stride=1, padding=1, bias=True),
                                          nn.Sigmoid())

    def forward(self, f_ct, f_cs, f_xs, test=False):
        raw_int_f_ct = list(f_ct.size())  # b*c*h*w
        raw_int_f_cs = list(f_cs.size())  # b*c*h*w

        kernel = 2 * self.rate

        phi_s = extract_image_patches(f_xs, ksizes=[kernel, kernel],
                                      strides=[self.rate * self.stride,
                                               self.rate * self.stride],
                                      rates=[1, 1],
                                      padding='same')  # [N, C*k*k, L]
        phi_s = phi_s.view(raw_int_f_cs[0], raw_int_f_cs[1], kernel, kernel, -1)
        phi_s = phi_s.permute(0, 4, 1, 2, 3)  # raw_shape: [N, L, C, k, k]
        phi_s_groups = torch.split(phi_s, 1, dim=0)

        f_ct = F.interpolate(f_ct, scale_factor=1. / self.rate, mode='nearest')
        f_cs = F.interpolate(f_cs, scale_factor=1. / self.rate, mode='nearest')
        int_f_ct = list(f_ct.size())
        int_f_cs = list(f_cs.size())
        f_ct_groups = torch.split(f_ct, 1, dim=0)

        omiga_s = extract_image_patches(f_cs, ksizes=[self.ksize, self.ksize],
                                  strides=[self.stride, self.stride],
                                  rates=[1, 1],
                                  padding='same')

        omiga_s = omiga_s.view(int_f_cs[0], int_f_cs[1], self.ksize, self.ksize, -1)
        omiga_s = omiga_s.permute(0, 4, 1, 2, 3)
        omiga_s_groups = torch.split(omiga_s, 1, dim=0)

        f_xt = []
        offsets = []

        for f_ct_i, omiga_s_i, phi_s_i in zip(f_ct_groups, omiga_s_groups, phi_s_groups):
            escape_NaN = torch.FloatTensor([1e-4]).cuda()
            omiga_s_i = omiga_s_i[0]
            max_omiga_s_i = torch.max(torch.sqrt(reduce_sum(torch.pow(omiga_s_i, 2),
                                                     axis=[1, 2, 3],
                                                     keepdim=True)),
                               escape_NaN)
            omiga_s_i_normed = omiga_s_i / max_omiga_s_i

            f_ct_i = same_padding(f_ct_i, [self.ksize, self.ksize], [1, 1], [1, 1])

            f_xt_i = F.conv2d(f_ct_i, omiga_s_i_normed, stride=1)

            f_xt_i = f_xt_i.view(1, int_f_cs[2] * int_f_cs[3], int_f_ct[2], int_f_ct[3])

            f_xt_i = F.softmax(f_xt_i * 10, dim=1)

            if not test:
                offset = torch.argmax(f_xt_i, dim=1, keepdim=True)

                offset = torch.cat([offset % int_f_ct[3], offset // int_f_ct[3]], dim=1)

            # deconv for patch pasting
            omiga_s_i_center = phi_s_i[0]
            if self.rate == 1:
                f_xt_i = F.pad(f_xt_i, [0, 1, 0, 1])
            f_xt_i = F.conv_transpose2d(f_xt_i, omiga_s_i_center, stride=self.rate, padding=1) / 4.
            f_xt.append(f_xt_i)
            if not test:
                offsets.append(offset)

        f_xt = torch.cat(f_xt, dim=0)

        f_xt.contiguous().view(raw_int_f_ct)
        mask = self.get_mask(f_xt)

        flow = None
        grids = None
        if not test:
            offsets = torch.cat(offsets, dim=0)
            offsets = offsets.view(int_f_ct[0], 2, *int_f_ct[2:])

            grids = offsets.type(torch.FloatTensor)
            grids = grids.cuda()

            h_add = torch.arange(int_f_ct[2]).view([1, 1, int_f_ct[2], 1]).expand(int_f_ct[0], -1, -1, int_f_ct[3])
            w_add = torch.arange(int_f_ct[3]).view([1, 1, 1, int_f_ct[3]]).expand(int_f_ct[0], -1, int_f_ct[2], -1)
            ref_coordinate = torch.cat([h_add, w_add], dim=1)
            ref_coordinate = ref_coordinate.cuda()

            offsets = offsets - ref_coordinate
            flow = offsets.type(torch.FloatTensor).cuda()

        return f_xt, mask, flow, grids

class ResBlockEncoderOptimized(nn.Module):
    """
    Define an Encoder block for the first layer of the discriminator
    """
    def __init__(self, input_nc, output_nc, hidden_nc=None, norm_layer=nn.BatchNorm2d, nonlinearity= nn.LeakyReLU(),
                 use_spect=False, use_coord=False):
        super(ResBlockEncoderOptimized, self).__init__()

        hidden_nc = input_nc if hidden_nc is None else hidden_nc

        conv1 = spectral_norm(nn.Conv2d(input_nc, hidden_nc, kernel_size=3, stride=1, padding=1), use_spect)
        conv2 = spectral_norm(nn.Conv2d(hidden_nc, output_nc, kernel_size=4, stride=2, padding=1), use_spect)
        bypass = spectral_norm(nn.Conv2d(input_nc, output_nc, kernel_size=1, stride=1, padding=0), use_spect)

        if type(norm_layer) == type(None):
            self.model = nn.Sequential(conv1, nonlinearity, conv2,)
        else:
            self.model = nn.Sequential(conv1, norm_layer(hidden_nc), nonlinearity, conv2,)
        self.shortcut = nn.Sequential(nn.AvgPool2d(kernel_size=2, stride=2), bypass)

    def forward(self, x):
        out = self.model(x) + self.shortcut(x)
        return out

class EncoderBlockOptimized(nn.Module):
    """
        Define an Encoder block for the first layer of the generator
    """
    def __init__(self, input_nc, output_nc, norm_layer=nn.BatchNorm2d, nonlinearity=nn.LeakyReLU(),
                 use_spect=False, use_coord=False):
        super(EncoderBlockOptimized, self).__init__()

        kwargs_down = {'kernel_size': 4, 'stride': 2, 'padding': 1}
        kwargs_fine = {'kernel_size': 3, 'stride': 1, 'padding': 1}

        conv1 = coord_conv(input_nc,  output_nc, use_spect, use_coord, **kwargs_down)
        conv2 = coord_conv(output_nc, output_nc, use_spect, use_coord, **kwargs_fine)

        if type(norm_layer) == type(None):
            self.model = nn.Sequential(conv1, nonlinearity, conv2)
        else:
            self.model = nn.Sequential(conv1, norm_layer(output_nc), nonlinearity, conv2)

    def forward(self, x):
        out = self.model(x)
        return out


#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os
import sys
import argparse
import numpy as np
from easydict import EasyDict as edict
import sklearn
from sklearn.preprocessing import normalize
import time

import signal
import traceback

reload(sys)
sys.setdefaultencoding("utf-8")

if sys.version_info.major == 2:
    import Queue as queue
    from Queue import Queue
else:
    import queue as Queue
import threading

global signal_stop
global read_ths
global infer_ths
global write_ths
global read_done
global write_done
global infer_done

signal_stop = False

def sigint_handler(signum, frame):
    global signal_stop
    global read_ths
    global infer_ths
    global write_ths

    signal_stop = True
    print("main-thread exit")

    # wait for finish
    for t in read_ths:
        t.join()
    for t in infer_ths:
        t.join()
    write_ths.join()

    sys.exit()

signal.signal(signal.SIGINT, sigint_handler)

class writeThread(threading.Thread):
    def __init__(self, out_q, write_func, tasks):
        threading.Thread.__init__(self)
        self.out_q = out_q
        self.write_func = write_func
        self.tasks = tasks

    def run(self):
        global signal_stop
        global write_done
        global infer_done

        done = 0
        while not signal_stop:
            try:
                # embedding, outpaths = self.out_q.get(block=True)
                embedding, outpaths = self.out_q.get(timeout=1)
                if done % 10000 == 0:
                    print(done, self.tasks)
                for idx, outp in enumerate(outpaths):
                    self.write_func(outp, embedding[idx].flatten())
                    done += 1
            except queue.Empty as e:
                if infer_done:
                    break
                pass
        print("write Thread done")
        write_done = True

class readThread(threading.Thread):
    def __init__(self, in_q, out_q, read_func, is_flip, shape):
        threading.Thread.__init__(self)
        self.in_q = in_q
        self.out_q = out_q
        self.read_func = read_func
        self.is_flip = is_flip
        self.shape = shape

    def run(self):
        global signal_stop
        unfinished = True
        while not signal_stop and unfinished:
            try:
                image_path, outpath = self.in_q.get(timeout=1)
                out_img = []
                img = self.read_func(image_path, self.shape[0] == 1)
                if len(img.shape) == 3:
                    img = img[:, :, ::-1]  # to rgb
                    img = np.transpose(img, (2, 0, 1))
                else:
                    img = img[np.newaxis, :]
                attempts = [0, 1] if self.is_flip else [0]
                for flipid in attempts:
                    _img = np.copy(img)
                    if flipid == 1:
                        self.do_flip(_img)
                    out_img.append(_img)

                self.out_q.put((out_img, outpath))

            except queue.Empty as e:
                unfinished = False
                break
        print("read Thread done")

    def do_flip(self, data):
        for idx in range(data.shape[0]):
            data[idx, :, :] = np.fliplr(data[idx, :, :])

class inferThread(threading.Thread):
    def __init__(self, net, fwd_func, in_q, out_q, image_shape, is_flip):
        threading.Thread.__init__(self)
        self.net = net
        self.fwd_func = fwd_func
        self.in_q = in_q
        self.out_q = out_q
        self.batchsz = net.batch_size
        self.image_shape = image_shape
        self.is_flip = is_flip
        self.end = False

    def run(self):
        global signal_stop
        global read_done

        bz = self.batchsz

        if self.is_flip:
            bz = self.batchsz // 2 * 2

        input_blob = np.zeros(
            shape=(bz, self.image_shape[0], self.image_shape[1], self.image_shape[2])
        )

        while not signal_stop and not self.end:
            idx = 0
            outpaths = []
            while idx < bz:
                try:
                    images, outpath = self.in_q.get(timeout=1)
                    input_blob[idx] = images[0]
                    idx += 1
                    if self.is_flip:
                        input_blob[idx] = images[1]
                        idx += 1
                    outpaths.append(outpath)
                except queue.Empty as e:
                    if read_done:
                        self.end = True
                    break

            if idx == 0:
                continue

            embedding = self.fwd_func(self.net, input_blob)

            if self.is_flip:
                embedding1 = embedding[0::2]
                embedding2 = embedding[1::2]
                embedding = embedding1 + embedding2

            # assert len(outpaths) == embedding.shape[0], "%d vs %d"%(len(outpaths), embedding.shape[0])
            if self.is_flip:
                assert len(outpaths) == idx / 2
            else:
                assert len(outpaths) == idx

            embedding = sklearn.preprocessing.normalize(embedding[0 : len(outpaths)])

            self.out_q.put((embedding, outpaths))

        print("infer thread done")

"""
class binThread (threading.Thread):
    def __init__(self, net, in_q, image_shape, is_flip, read_func, write_func):
        threading.Thread.__init__(self)
        self.net = net
        self.in_q = in_q;
        self.batchsz = net.batch_size;
        self.image_shape = image_shape
        self.is_flip = is_flip
        self.read_func = read_func
        self.write_func = write_func
        self.end = False

    def run(self):
        global signal_stop

        bz = self.batchsz

        if (self.is_flip):
            bz = self.batchsz//2 * 2

        input_blob = np.zeros(shape=(bz, self.image_shape[0],
                    self.image_shape[1], self.image_shape[2]))

        while not signal_stop and not self.end:
            # load bin
            inpath, outpath = self.in_q.get(timeout=1)
            data_set = self.read_func(inpath, self.image_shape)

            data_list = data_set[0]
            issame_list = data_set[1]
            embeddings_list = []

            for i in range( len(data_list) ):
                data = data_list[i]
                embeddings = None
                ba = 0
                while ba<data.shape[0]:
                    bb = min(ba+bz, data.shape[0])
                    count = bb-ba
                    _data = nd.slice_axis(data, axis=0, begin=bb-bz, end=bb)
                    db = mx.io.DataBatch(data=(_data,))
                    self.net.model.forward(db, is_train=False)
                    net_out = self.net.model.get_outputs()
                    _embeddings = net_out[0].asnumpy()
                    if embeddings is None:
                        embeddings = np.zeros( (data.shape[0], _embeddings.shape[1]) )
                    embeddings[ba:bb,:] = _embeddings[(bz-count):,:]
                    ba = bb
                embeddings_list.append(embeddings)

            embeddings = embeddings_list[0].copy()
            if self.is_flip:
                embeddings = embeddings_list[0] + embeddings_list[1]
            embeddings = sklearn.preprocessing.normalize(embeddings)

            # save output numpy
            self.write_func(outpath, embeddings)
            self.write_func(outpath+".issame", issame_list)

        print("infer thread done")
"""

class CitrusBaseInfer(object):
    def __init__(self, args, dtype="fp32"):
        self._args = args
        self._dtype = dtype
        self._net = None
        self._ctx = []
        self._image_shape = []
        self._batch_size = 8
        self._net_scale_type = "large"
        self._flops = None
        self._epoch = None
        self._gpu_mem_use = 0
        self._nets = []
        self.fwd_func = None

        self._load_model()

    def _load_model(self):
        raise NotImplementedError()

    #

    def infer_embedding(
        self, imgs=None, read_func=None, write_func=None, is_flip=False
    ):
        global read_ths
        global infer_ths
        global write_ths
        global read_done
        global write_done
        global infer_done

        read_done = False
        write_done = False
        infer_done = False
        assert 0 != len(imgs) and imgs is not None, "images to infer must not be empty!"
        assert write_func is not None
        assert read_func is not None

        try:
            in_q = Queue()
            for items in imgs:
                in_q.put(items)
            tasks = in_q.qsize()

            print("begin thread")

            out_q = Queue()
            img_q = Queue()

            read_ths = []
            infer_ths = []

            read_thread_num = len(self._nets) * 2
            if read_thread_num > 30:
                read_thread_num = 30
            for n in range(read_thread_num):
                read_ths.append(
                    readThread(in_q, img_q, read_func, is_flip, self._image_shape)
                )

            for n in self._nets:
                infer_ths.append(
                    inferThread(
                        n, self.fwd_func, img_q, out_q, self._image_shape, is_flip
                    )
                )

            for t in read_ths:
                t.start()

            for t in infer_ths:
                t.start()

            write_ths = writeThread(out_q, write_func, tasks)
            write_ths.start()

            while not read_done:
                done = True
                for t in read_ths:
                    if t.isAlive():
                        done = False
                if done:
                    read_done = True
                    print("read threads done!\n")
                time.sleep(1)

            while not infer_done:
                done = True
                for t in infer_ths:
                    if t.isAlive():
                        done = False
                if done:
                    infer_done = True
                    print("infer threads done!\n")
                time.sleep(1)

            while not write_done:
                # if use join(), ctrl-c can not work
                time.sleep(1)

            if in_q.qsize() == 0 and img_q.qsize() == 0 and out_q.qsize() == 0:
                return True
            else:
                print(
                    "not done! in_q.qsize()",
                    in_q.qsize(),
                    "img_q.qsize()",
                    img_q.qsize(),
                    "out_q.qsize()",
                    out_q.qsize(),
                )
                return False
        except:
            print("error", traceback.print_exc())
            signal_stop = True
            return False

    #

    """
    def infer_bin(self, imgs=None, read_func = None, write_func = None, is_flip=False):
        global infer_ths
        global infer_done

        infer_done = False
        assert 0 != len(imgs) and imgs is not None, "images to infer must not be empty!"
        assert write_func is not None
        assert read_func is not None

        try:
          in_q = Queue()
          for items in imgs:
            in_q.put(items)

          infer_ths = []

          for n in self._nets:
            infer_ths.append(binThread(n, in_q, self._image_shape, is_flip, read_func, write_func))

          for t in infer_ths:
              t.start()

          while not infer_done:
            done = True
            for t in infer_ths:
              if t.isAlive():
                  done = False
            if done:
                infer_done = True
                print("infer threads done!\n")
            time.sleep(1)

          return True
        except:
          print("error", traceback.print_exc())
          signal_stop = True
          return False
    """
    #


from datetime import datetime

from sqlalchemy import Column, Integer, String, DateTime, Date, Boolean, ForeignKey
from sqlalchemy.orm import relationship

from server import db, ma

class Experience(db.Model):
    __tablename__ = "tb_experience"

    id = Column(Integer, primary_key=True)

    institution_name = Column(String)

    id_type_of_experience = Column(Integer, ForeignKey("tb_domain.id"))
    type_of_experience = relationship("Domain", foreign_keys=[id_type_of_experience])

    current_experience = Column(Boolean)

    start_year = Column(Date)

    finish_year = Column(Date)

    details = Column(String)

    created_at = Column(DateTime, default=datetime.now)
    created_by = Column(String)
    updated_at = Column(DateTime, default=datetime.now, onupdate=datetime.now)
    updated_by = Column(String)

class ExperienceSchema(ma.ModelSchema):
    class Meta:
        fields = (
            "id",
            "name",
            "email",
            "description",
            "birth_date",
            "created_at",
            "created_by"
        )


#coding: utf-8

from grading_tools import TestConfiguration, ProgramTest, CheckOutputMixin, CheckStderrMixin, CheckMultiCorePerformance
import sys
import numpy as np
from sklearn.metrics.pairwise import euclidean_distances
import os
import io
import os.path as osp
import sys
from contextlib import redirect_stdout
import numba as nb
from numba.typed import List as nbList
import pprint

from colorama import init, Fore

def green(s):
    return Fore.GREEN + str(s) + Fore.RESET

def red(s):
    return Fore.RED + str(s) + Fore.RESET

def test_result(b):
    if b:
        return green(b)
    else:
        return red(b)

class BaseTSP:
    def parse_input(self, inp):
        lines = inp.split('\n')
        N = int(lines[0])
        points = np.zeros((N, 2))
        for i, l in enumerate(lines[1:-1]):
            p = l.split()
            points[i, 0] = float(p[0])
            points[i, 1] = float(p[1])

        distances = euclidean_distances(points)
        return N, points, distances

    def parse_output(self, stdout):
        try:
            lines = stdout.split('\n')
            distance, opt = lines[0].split()
            distance = float(distance)
            order = [int(i) for i in lines[1].split()]
        except:
            for i, l in enumerate(lines):
                print('linha', i, ':', l)
            raise Exception("Formato de saída inválido")

        return distance, opt, order

    def tamanho_tour(self, order, distances):
        tour = 0
        last = order[0]
        N = len(order)
        for i in range(1, N):
            tour += distances[last, order[i]]
            last = order[i]
        tour += distances[order[N-1], order[0]]

        return tour

    def test_caminho_tem_distancia_mostrada(self, test, stdout, stderr):
        N, points, distances = self.parse_input(test.input)
        distance, opt, order = self.parse_output(stdout)
        tour = self.tamanho_tour(order, distances)
        res = np.isclose(tour, distance)
        if not res:
            print('Distância calculada:', tour, 'Mostrada:', distance)
        return res

    def test_passa_por_todas_cidades(self, test, stdout, stderr):
        lines = stdout.split('\n')
        cities = [int(i) for i in lines[1].split()]
        cities.sort()
        return cities == list(range(len(cities)))

@nb.jit
def tamanho_tour_nb(order, distances):
        tour = 0
        last = order[0]
        N = len(order)
        for i in range(1, N):
            tour += distances[last, order[i]]
            last = order[i]
        tour += distances[order[N-1], order[0]]

        return tour

def tem_troca_py(order, distance, distances):
    N = len(order)
    for i in range(N):
        for j in range(i+1, N):
            new_order = order[:]
            new_order[i], new_order[j] = new_order[j], new_order[i]
            new_distance = tamanho_tour_nb(new_order, distances)
            is_close = abs(distance - new_distance)
            if is_close > 0.01 and distance > new_distance:
                return True, i, j, new_distance

    return False, 0, 0, 0.0

tem_troca_nb = nb.jit(tem_troca_py)

class TesteBuscaLocal(ProgramTest, BaseTSP):
    def tem_troca(self, order, distance, distances):
        distance = self.tamanho_tour(order, distances)
        ret, i, j, new_distance = tem_troca_nb(nbList(order), distance, distances)
        if ret:
            print('Troca encontrada: ', i, j)
            print('Melhoria:', distance, '->', new_distance)
        return ret

    def test_nao_tem_troca(self, test, stdout, stderr):
        N, points, distances = self.parse_input(test.input)
        distance, opt, order = self.parse_output(stdout)

        return not self.tem_troca(order, distance, distances)

    def parse_stderr(self, stderr):
        lines = stderr.split('\n')
        solutions = []
        for l in lines:
            l = l.strip()
            try:
                distance = float(l.split(' ')[1])
            except:
                print(l.split())
                raise Exception('Formato de saída de erros inválido')
            try:
                order = [int(i) for i in l.split(' ')[2:]]
            except:
                print(l.split(' ')[2:])
                raise Exception('Formato de saída de erros inválido')

            solutions.append((distance, order))

        return solutions

    def test_fez_10N_solucoes(self, test, stdout, stderr):
        N, points, distances = self.parse_input(test.input)
        distance, opt, order = self.parse_output(stdout)
        solutions = self.parse_stderr(stderr)
        return len(solutions) == 10 * N

    def test_toda_solucao_tem_caminho_correto(self, test, stdout, stderr):
        N, points, distances = self.parse_input(test.input)
        distance, opt, order = self.parse_output(stdout)
        solutions = self.parse_stderr(stderr)
        for i, sol in enumerate(solutions):
            dist_calculada = self.tamanho_tour(sol[1], distances)
            if not np.isclose(dist_calculada, sol[0]):
                print(f'Caminho incorreto na linha {i}. Mostrada {sol[0]}, calculada {dist_calculada}')
                print(sol[1], self.tamanho_tour(sol[1], distances), tamanho_tour_nb(sol[1], distances))
                return False
        return True

    def test_toda_solucao_eh_otimo_local(self, test, stdout, stderr):
        N, points, distances = self.parse_input(test.input)
        distance, opt, order = self.parse_output(stdout)
        solutions = self.parse_stderr(stderr)
        for i, sol in enumerate(solutions):
            if self.tem_troca(sol[1], sol[0], distances):
                print('Troca acima encontrada na linha', i)
                return False
        return True

class TesteHeuristico(ProgramTest, CheckOutputMixin, BaseTSP):
    pass

class TesteBuscaExaustiva(ProgramTest, CheckStderrMixin, BaseTSP):
    def test_tour_otimo(self, test, stdout, stderr):
        N, points, distances = self.parse_input(test.input)
        distance, opt, order = self.parse_output(stdout)
        correct_distance, correct_opt, correct_order = self.parse_output(test.output)
        return np.isclose(correct_distance, distance)

    def test_opt_1(self, test, stdout, stderr):
        distance, opt, order = self.parse_output(stdout)
        return int(opt) == 1

class TesteBuscaExaustivaPerf(ProgramTest, BaseTSP):
    def test_tour_otimo(self, test, stdout, stderr):
        N, points, distances = self.parse_input(test.input)
        distance, opt, order = self.parse_output(stdout)
        correct_distance, correct_opt, correct_order = self.parse_output(test.output)
        return np.isclose(correct_distance, distance)

    def test_opt_1(self, test, stdout, stderr):
        distance, opt, order = self.parse_output(stdout)
        return int(opt) == 1

class TestePerformance(ProgramTest, BaseTSP):
    pass

class TesteMultiCorePequeno(TesteBuscaLocal, CheckMultiCorePerformance):
    pass

def compila_programa(ext, nome, flags, nvcc=False):
    compilador = 'g++'
    if nvcc:
        compilador = 'nvcc'
    arquivos = ' '.join([arq for arq in os.listdir('.') if arq.split('.')[-1] == ext])
    ret = os.system(f'{compilador} {arquivos} -O3 {flags} -o {nome} 2>1 > /dev/null')
    if ret != 0 :
        raise IOError(f'Erro de compilação em {os.getcwd()}!')

def testa_heuristico():
    os.chdir('heuristico')
    compila_programa('cpp', 'heuristico', '')
    tests = TestConfiguration.from_pattern('.', 'in*.txt', 'out*txt')
    tester = TesteHeuristico('./heuristico', tests)
    res = tester.main()
    os.chdir('..')
    return res

def testa_busca_local_sequencial():
    os.chdir('busca-local')
    compila_programa('cpp', 'busca-local', '')
    tests = TestConfiguration.from_pattern('.', 'in*.txt', 'out*txt',
                                           environ={
                                               'DEBUG': '1'
                                           }, time_limit=1800)
    tester = TesteBuscaLocal('./busca-local', tests)
    res = tester.main()
    os.chdir('..')
    return res

def testa_busca_exaustiva():
    os.chdir('busca-exaustiva')
    compila_programa('cpp', 'busca-global', '')
    tests = TestConfiguration.from_pattern('.', 'in*.txt', 'out*txt',
                                           'err*txt', environ={
                                               'DEBUG': '1'
                                           })
    tester = TesteBuscaExaustiva('./busca-global', tests)
    res = tester.main()
    os.chdir('..')
    return res

def testa_busca_local_omp():
    os.chdir('busca-local')
    compila_programa('cpp', 'busca-local-paralela', '-fopenmp')
    tests = TestConfiguration.from_pattern('.', 'in*.txt', 'out*txt',
                                           environ={
                                               'DEBUG': '1'
                                           }, time_limit=1800)
    files_multicore = [f'./in-{i}.txt' for i in range(5, 10)]
    omp_tests = {k:v for k, v in tests.items() if k in files_multicore}
    del omp_tests['./in-7.txt']
    del omp_tests['./in-8.txt']

    teste_simples = TesteMultiCorePequeno('./busca-local-paralela', omp_tests)
    res = teste_simples.main()
    os.chdir('..')
    return res

def testa_busca_local_gpu():
    os.chdir('busca-local')
    compila_programa('cu', 'busca-local-gpu', '', True)
    tests = TestConfiguration.from_pattern('.', 'in*.txt', 'out*txt',
                                        environ={
                                            'DEBUG': '1'
                                        })
    tester = TesteBuscaLocal('./busca-local-gpu', tests)

    perf_tests = {}
    perf_tests['in6.txt'] = TestConfiguration.from_file('in-6.txt',
                                                '',
                                                check_stderr=False,
                                                environ={'DEBUG': '0'},
                                                time_limit=30)
    perf_tests['in11.txt'] = TestConfiguration.from_file('perf-in-11.txt',
                                                '',
                                                check_stderr=False,
                                                environ={'DEBUG': '0'},
                                                time_limit=30)

    perf_tester = TestePerformance('./busca-local-gpu', perf_tests)
    res = tester.main() and perf_tester.main()
    os.chdir('..')
    return res

def testa_busca_local_gpu2():
    os.chdir('busca-local')
    compila_programa('cu', 'busca-local-gpu2', '', True)
    tests = TestConfiguration.from_pattern('.', 'in*.txt', 'out*txt',
                                        environ={
                                            'DEBUG': '1'
                                        })
    tester = TesteBuscaLocal('./busca-local-gpu2', tests)

    perf_tests = {}
    perf_tests['in6.txt'] = TestConfiguration.from_file('in-6.txt',
                                                '',
                                                check_stderr=False,
                                                environ={'DEBUG': '0'},
                                                time_limit=10)
    perf_tests['in11.txt'] = TestConfiguration.from_file('perf-in-11.txt',
                                                    '',
                                                    check_stderr=False,
                                                    environ={'DEBUG': '0'},
                                                    time_limit=9)

    perf_tests['in12.txt'] = TestConfiguration.from_file('perf-in-12.txt',
                                                    '',
                                                    check_stderr=False,
                                                    environ={'DEBUG': '0'},
                                                    time_limit=65)

    perf_tester = TestePerformance('./busca-local-gpu2', perf_tests)
    res = tester.main() and perf_tester.main()
    os.chdir('..')
    return res

def testa_busca_local_perf():
    os.chdir('busca-local')
    compila_programa('cpp', 'busca-local-perf', '')
    performance_tests = {}
    performance_tests['in10.txt'] = TestConfiguration.from_file('in-10.txt',
                                                    '',
                                                    check_stderr=False,
                                                    environ={'DEBUG': '0'},
                                                    time_limit=0.5)

    performance_tests['in11.txt'] = TestConfiguration.from_file('perf-in-11.txt',
                                                    '',
                                                    check_stderr=False,
                                                    environ={'DEBUG': '0'},
                                                    time_limit=9)

    performance_tests['in12.txt'] = TestConfiguration.from_file('perf-in-12.txt',
                                                    '',
                                                    check_stderr=False,
                                                    environ={'DEBUG': '0'},
                                                    time_limit=180)

    perf_tester = TestePerformance('./busca-local-perf', performance_tests)
    res = perf_tester.main()
    os.chdir('..')
    return res

def testa_busca_exaustiva_perf1():
    os.chdir('busca-exaustiva')
    compila_programa('cpp', 'busca-global-perf', '-fopenmp')
    performance_tests = {}
    performance_tests['in-5.txt'] = TestConfiguration.from_file('in-5.txt',
                                                    'out-5.txt',
                                                    check_stderr=False,
                                                    environ={'DEBUG': '0'},
                                                    time_limit=1)
    performance_tests['in-6.txt'] = TestConfiguration.from_file('in-6.txt',
                                                    'out-6.txt',
                                                    check_stderr=False,
                                                    environ={'DEBUG': '0'},
                                                    time_limit=1)
    performance_tests['in-7.txt'] = TestConfiguration.from_file('in-7.txt',
                                                    'out-7.txt',
                                                    check_stderr=False,
                                                    environ={'DEBUG': '0'},
                                                    time_limit=5)

    perf_tester = TesteBuscaExaustivaPerf('./busca-global-perf', performance_tests)
    res = perf_tester.main()
    os.chdir('..')
    return res

def testa_busca_exaustiva_perf2():
    os.chdir('busca-exaustiva')
    compila_programa('cpp', 'busca-global-perf2', '-fopenmp')
    performance_tests = {}
    performance_tests['in-7.txt'] = TestConfiguration.from_file('in-7.txt',
                                                    'out-7.txt',
                                                    check_stderr=False,
                                                    environ={'DEBUG': '0'},
                                                    time_limit=1)
    performance_tests['perf-in-8.txt'] = TestConfiguration.from_file('perf-in-8.txt',
                                                    'perf-out-8.txt',
                                                    check_stderr=False,
                                                    environ={'DEBUG': '0'},
                                                    time_limit=2)
    performance_tests['perf-in-9.txt'] = TestConfiguration.from_file('perf-in-9.txt',
                                                    'perf-out-9.txt',
                                                    check_stderr=False,
                                                    environ={'DEBUG': '0'},
                                                    time_limit=30)

    perf_tester = TesteBuscaExaustivaPerf('./busca-global-perf2', performance_tests)
    res = perf_tester.main()
    os.chdir('..')
    return res

def testa_busca_exaustiva_perf3():
    os.chdir('busca-exaustiva')
    compila_programa('cpp', 'busca-global-perf3', '-fopenmp')
    performance_tests = {}
    performance_tests['in-7.txt'] = TestConfiguration.from_file('in-7.txt',
                                                    'out-7.txt',
                                                    check_stderr=False,
                                                    environ={'DEBUG': '0'},
                                                    time_limit=1)
    performance_tests['perf-in-8.txt'] = TestConfiguration.from_file('perf-in-8.txt',
                                                    'perf-out-8.txt',
                                                    check_stderr=False,
                                                    environ={'DEBUG': '0'},
                                                    time_limit=1)
    performance_tests['perf-in-9.txt'] = TestConfiguration.from_file('perf-in-9.txt',
                                                    'perf-out-9.txt',
                                                    check_stderr=False,
                                                    environ={'DEBUG': '0'},
                                                    time_limit=10)
    performance_tests['perf-in-10.txt'] = TestConfiguration.from_file('perf-in-10.txt',
                                                    'perf-out-10.txt',
                                                    check_stderr=False,
                                                    environ={'DEBUG': '0'},
                                                    time_limit=25)

    perf_tester = TesteBuscaExaustivaPerf('./busca-global-perf3', performance_tests)
    res = perf_tester.main()
    os.chdir('..')
    return res

def testa_busca_exaustiva_perf4():
    os.chdir('busca-exaustiva')
    compila_programa('cpp', 'busca-global-perf4', '-fopenmp')
    performance_tests = {}
    performance_tests['in-7.txt'] = TestConfiguration.from_file('in-7.txt',
                                                    'out-7.txt',
                                                    check_stderr=False,
                                                    environ={'DEBUG': '0'},
                                                    time_limit=1)
    performance_tests['perf-in-8.txt'] = TestConfiguration.from_file('perf-in-8.txt',
                                                    'perf-out-8.txt',
                                                    check_stderr=False,
                                                    environ={'DEBUG': '0'},
                                                    time_limit=1)
    performance_tests['perf-in-9.txt'] = TestConfiguration.from_file('perf-in-9.txt',
                                                    'perf-out-9.txt',
                                                    check_stderr=False,
                                                    environ={'DEBUG': '0'},
                                                    time_limit=2)
    performance_tests['perf-in-10.txt'] = TestConfiguration.from_file('perf-in-10.txt',
                                                    'perf-out-10.txt',
                                                    check_stderr=False,
                                                    environ={'DEBUG': '0'},
                                                    time_limit=3)
    performance_tests['perf-in-11.txt'] = TestConfiguration.from_file('perf-in-11.txt',
                                                    'perf-out-11.txt',
                                                    check_stderr=False,
                                                    environ={'DEBUG': '0'},
                                                    time_limit=3)
    perf_tester = TesteBuscaExaustivaPerf('./busca-global-perf4', performance_tests)
    res = perf_tester.main()
    os.chdir('..')
    return res

def testa_busca_exaustiva_perf5():
    os.chdir('busca-exaustiva')
    compila_programa('cpp', 'busca-global-perf5', '-fopenmp')
    performance_tests = {}
    performance_tests['in-7.txt'] = TestConfiguration.from_file('in-7.txt',
                                                    'out-7.txt',
                                                    check_stderr=False,
                                                    environ={'DEBUG': '0'},
                                                    time_limit=1)
    performance_tests['perf-in-8.txt'] = TestConfiguration.from_file('perf-in-8.txt',
                                                    'perf-out-8.txt',
                                                    check_stderr=False,
                                                    environ={'DEBUG': '0'},
                                                    time_limit=1)
    performance_tests['perf-in-9.txt'] = TestConfiguration.from_file('perf-in-9.txt',
                                                    'perf-out-9.txt',
                                                    check_stderr=False,
                                                    environ={'DEBUG': '0'},
                                                    time_limit=2)
    performance_tests['perf-in-10.txt'] = TestConfiguration.from_file('perf-in-10.txt',
                                                    'perf-out-10.txt',
                                                    check_stderr=False,
                                                    environ={'DEBUG': '0'},
                                                    time_limit=3)
    performance_tests['perf-in-11.txt'] = TestConfiguration.from_file('perf-in-11.txt',
                                                    'perf-out-11.txt',
                                                    check_stderr=False,
                                                    environ={'DEBUG': '0'},
                                                    time_limit=3)
    performance_tests['perf-in-12.txt'] = TestConfiguration.from_file('perf-in-12.txt',
                                                    'perf-out-12.txt',
                                                    check_stderr=False,
                                                    environ={'DEBUG': '0'},
                                                    time_limit=45)
    performance_tests['perf-in-13.txt'] = TestConfiguration.from_file('perf-in-13.txt',
                                                    'perf-out-13.txt',
                                                    check_stderr=False,
                                                    environ={'DEBUG': '0'},
                                                    time_limit=70)

    perf_tester = TesteBuscaExaustivaPerf('./busca-global-perf5', performance_tests)
    res = perf_tester.main()
    os.chdir('..')
    return res

def testa_busca_local_omp_perf():
    os.chdir('busca-local')
    compila_programa('cpp', 'busca-local-perf-omp', '-fopenmp')
    performance_tests = {}
    performance_tests['in10.txt'] = TestConfiguration.from_file('in-10.txt',
                                                    '',
                                                    check_stderr=False,
                                                    environ={'DEBUG': '0'},
                                                    time_limit=0.5)

    performance_tests['in11.txt'] = TestConfiguration.from_file('perf-in-11.txt',
                                                    '',
                                                    check_stderr=False,
                                                    environ={'DEBUG': '0'},
                                                    time_limit=3)

    performance_tests['in12.txt'] = TestConfiguration.from_file('perf-in-12.txt',
                                                    '',
                                                    check_stderr=False,
                                                    environ={'DEBUG': '0'},
                                                    time_limit=60)

    perf_tester = TestePerformance('./busca-local-perf-omp', performance_tests)
    res = perf_tester.main()
    os.chdir('..')
    return res

if __name__ == "__main__":

    ignorar = io.StringIO()

    testesD = {
        'heuristico': ('Heuristico (sequencial)', testa_heuristico),
        'local': ('Busca local (sequencial)', testa_busca_local_sequencial),
        'global': ('Busca exaustiva (sequencial)', testa_busca_exaustiva),
        'local-paralela': ('Busca local (paralela)', testa_busca_local_omp),
        'local-gpu': ('Busca local (GPU)', testa_busca_local_gpu)
    }

    if len(sys.argv) > 1:
        tst = sys.argv[1]
        if tst in testesD:
            tst = testesD[tst]
            print(tst[0], ':', tst[1]())
        else:
            print('Testes disponíveis:')
            pprint.pprint(testesD)
        sys.exit(0)

    print(f'''
===========================================
Rubrica D
===========================================''')

    with open('feedback-heuristico.txt', 'w') as f:
        with redirect_stdout(f):
            res_heuristico = testa_heuristico()

    print(f'Heuristico (sequencial): {test_result(res_heuristico)}')

    with open('feedback-local.txt', 'w') as f:
        with redirect_stdout(f):
            res_busca_local_sequencial = testa_busca_local_sequencial()

    print(f'Busca local (sequencial): {test_result(res_busca_local_sequencial)}')

    with open('feedback-global.txt', 'w') as f:
        with redirect_stdout(f):
            res_global = testa_busca_exaustiva()
    print(f'Busca exaustiva (sequencial): {test_result(res_global)}')

    with open('feedback-local-paralela.txt', 'w') as f:
        with redirect_stdout(f):
            res_busca_local_paralela = testa_busca_local_omp()

    print(f'Busca local (paralela): {test_result(res_busca_local_paralela)}')

    with open('feedback-local-gpu.txt', 'w') as f:
        with redirect_stdout(f):
            res_busca_local_gpu = testa_busca_local_gpu()

    print(f'Busca local (GPU): {test_result(res_busca_local_gpu)}')

    print('''===========================================
Rubrica A+
===========================================''')

    with open('feedback-local-perf.txt', 'w') as f:
        with redirect_stdout(f):
            res_busca_local_perf = testa_busca_local_perf()

    print(f'Busca local (desempenho sequencial): {test_result(res_busca_local_perf)}')

    with open('feedback-local-perf-omp.txt', 'w') as f:
        with redirect_stdout(f):
            res_busca_local_omp_perf = testa_busca_local_omp_perf()

    print(f'Busca local (desempenho paralelo 1): {test_result(res_busca_local_omp_perf)}')

    with open('feedback-local-perf-gpu2.txt', 'w') as f:
        with redirect_stdout(f):
            res_busca_local_perf_gpu = testa_busca_local_gpu2()

    print(f'Busca local (desempenho GPU 1): {test_result(res_busca_local_perf_gpu)}')

    with open('feedback-global-perf.txt', 'w') as f:
        with redirect_stdout(f):
            res_busca_global_p1 = testa_busca_exaustiva_perf1()

    print(f'Busca global (desempenho nível 1): {test_result(res_busca_global_p1)}')

    with open('feedback-global-perf2.txt', 'w') as f:
        with redirect_stdout(f):
            res_busca_global_p2 = testa_busca_exaustiva_perf2()

    print(f'Busca global (desempenho nível 2): {test_result(res_busca_global_p2)}')

    with open('feedback-global-perf3.txt', 'w') as f:
        with redirect_stdout(f):
            res_busca_global_p3 = testa_busca_exaustiva_perf3()

    print(f'Busca global (desempenho nível 3): {test_result(res_busca_global_p3)}')

    with open('feedback-global-perf4.txt', 'w') as f:
        with redirect_stdout(f):
            res_busca_global_p4 = testa_busca_exaustiva_perf4()

    print(f'Busca global (desempenho nível 4): {test_result(res_busca_global_p4)}')

    with open('feedback-global-perf5.txt', 'w') as f:
        with redirect_stdout(f):
            res_busca_global_p5 = testa_busca_exaustiva_perf5()

    print(f'Busca global (desempenho nível 5): {test_result(res_busca_global_p5)}')