import os
import torch
import torch.nn as nn
import torch.nn.functional as F
from layers.modules.l2norm import L2Norm
from models.mobilenet_v2_xiaomi import MobileNetV2, InvertedResidual
from models.FairNAS_A import FairNasA
from models.FairNAS_B import FairNasB
import numpy as np
from collections import namedtuple
GraphPath = namedtuple("GraphPath", ['s0', 'name', 's1'])  #

def add_extras(cfg, i, batch_norm=False):
    # Extra layers added to VGG for feature scaling
    layers = []
    in_channels = i
    flag = False
    for k, v in enumerate(cfg):
        if in_channels != 'S':
            if v == 'S':
                layers += [nn.Conv2d(in_channels, cfg[k + 1],
                           kernel_size=(1, 3)[flag], stride=2, padding=1)]
            else:
                layers += [nn.Conv2d(in_channels, v, kernel_size=(1, 3)[flag])]
            flag = not flag
        in_channels = v
    return layers

def vgg(cfg, i, batch_norm=False):
    layers = []
    in_channels = i
    for v in cfg:
        if v == 'M':
            layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
        elif v == 'C':
            layers += [nn.MaxPool2d(kernel_size=2, stride=2, ceil_mode=True)]
        else:
            conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1)
            if batch_norm:
                layers += [conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)]
            else:
                layers += [conv2d, nn.ReLU(inplace=True)]
            in_channels = v
    pool5 = nn.MaxPool2d(kernel_size=2, stride=2, padding=0)
    conv6 = nn.Conv2d(512, 1024, kernel_size=3, padding=1)
    conv7 = nn.Conv2d(1024, 1024, kernel_size=1)
    layers += [pool5, conv6,
               nn.ReLU(inplace=True), conv7, nn.ReLU(inplace=True)]
    return layers

class S3FD(nn.Module):

    def __init__(self, phase, size, num_classes):
        super(S3FD, self).__init__()
        self.phase = phase
        self.num_classes = num_classes
        self.size = size

        vgg_cfg = [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'C', 512, 512, 512, 'M',
                512, 512, 512]
        self.vgg = nn.ModuleList(vgg(vgg_cfg, 3))

        extras_cfg = [256, 'S', 512, 128, 'S', 256]
        self.extras = nn.ModuleList(add_extras(extras_cfg, 1024))

        self.conv3_3_L2Norm = L2Norm(256, 10)
        self.conv4_3_L2Norm = L2Norm(512, 8)
        self.conv5_3_L2Norm = L2Norm(512, 5)

        self.loc, self.conf = self.multibox(self.num_classes)

        if self.phase == 'test':
            self.softmax = nn.Softmax(dim=-1)

        if self.phase == 'train':
            for m in self.modules():
                if isinstance(m, nn.Conv2d):
                    if m.bias is not None:
                        nn.init.xavier_normal_(m.weight.data)
                        m.bias.data.fill_(0.02)
                    else:
                        m.weight.data.normal_(0, 0.01)
                elif isinstance(m, nn.BatchNorm2d):
                    m.weight.data.fill_(1)
                    m.bias.data.zero_()

    def multibox(self, num_classes):
        loc_layers = []
        conf_layers = []
        # Max-out BG label
        loc_layers += [nn.Conv2d(256, 1 * 4, kernel_size=3, padding=1)]
        conf_layers += [nn.Conv2d(256, 1 * 4, kernel_size=3, padding=1)]
        #conf_layers += [nn.Conv2d(256, 1 * num_classes, kernel_size=3, padding=1)]

        loc_layers += [nn.Conv2d(512, 1 * 4, kernel_size=3, padding=1)]
        conf_layers += [nn.Conv2d(512, 1 * num_classes, kernel_size=3, padding=1)]
        loc_layers += [nn.Conv2d(512, 1 * 4, kernel_size=3, padding=1)]
        conf_layers += [nn.Conv2d(512, 1 * num_classes, kernel_size=3, padding=1)]
        loc_layers += [nn.Conv2d(1024, 1 * 4, kernel_size=3, padding=1)]
        conf_layers += [nn.Conv2d(1024, 1 * num_classes, kernel_size=3, padding=1)]
        loc_layers += [nn.Conv2d(512, 1 * 4, kernel_size=3, padding=1)]
        conf_layers += [nn.Conv2d(512, 1 * num_classes, kernel_size=3, padding=1)]
        loc_layers += [nn.Conv2d(256, 1 * 4, kernel_size=3, padding=1)]
        conf_layers += [nn.Conv2d(256, 1 * num_classes, kernel_size=3, padding=1)]
        return nn.Sequential(*loc_layers), nn.Sequential(*conf_layers)

    def load_weights(self, base_file):
        other, ext = os.path.splitext(base_file)
        if ext == '.pkl' or '.pth':
            print('Loading weights into state dict...')
            self.load_state_dict(torch.load(base_file,
                                 map_location=lambda storage, loc: storage))
            print('Finished!')
        else:
            print('Sorry only .pth and .pkl files supported.')

    def forward(self, x):
        sources = list()
        loc = list()
        conf = list()
        detection_dimension = list()

        # apply vgg up to conv4_3 relu and conv5_3 relu
        for k in range(30):
            x = self.vgg[k](x)
            if 15 == k:
                s = self.conv3_3_L2Norm(x)
                sources.append(s)
                detection_dimension.append(x.shape[2:])
            elif 22 == k:
                s = self.conv4_3_L2Norm(x)
                sources.append(s)
                detection_dimension.append(x.shape[2:])
            elif 29 == k:
                s = self.conv5_3_L2Norm(x)
                sources.append(s)
                detection_dimension.append(x.shape[2:])

        # apply vgg up to fc7
        for k in range(30, len(self.vgg)):
            x = self.vgg[k](x)
        sources.append(x)
        detection_dimension.append(x.shape[2:])

        # apply extra layers and cache source layer outputs
        for k, v in enumerate(self.extras):
            x = F.relu(v(x), inplace=True)
            if k % 2 == 1:
                sources.append(x)
                detection_dimension.append(x.shape[2:])

        detection_dimension = torch.Tensor(detection_dimension)
        detection_dimension = detection_dimension.cuda()

        for index, (x, l, c) in enumerate(zip(sources, self.loc, self.conf)):
            #print(x.size())
            if index != 0:
                loc.append(l(x).permute(0, 2, 3, 1).contiguous())
                conf.append(c(x).permute(0, 2, 3, 1).contiguous())
            else:
                loc.append(l(x).permute(0, 2, 3, 1).contiguous())
                conf_t = c(x)
                max_conf, _ = conf_t[:, 0:3, :, :].max(1, keepdim=True)
                lab_conf = conf_t[:, 3:, :, :]
                out_conf = torch.cat((max_conf, lab_conf), dim=1)
                conf.append(out_conf.permute(0, 2, 3, 1).contiguous())

        '''for index, (x, l, c) in enumerate(zip(sources, self.loc, self.conf)):
            loc.append(l(x).permute(0, 2, 3, 1).contiguous())
            conf.append(c(x).permute(0, 2, 3, 1).contiguous())'''

        loc = torch.cat([o.view(o.size(0), -1) for o in loc], 1)
        conf = torch.cat([o.view(o.size(0), -1) for o in conf], 1)

        if self.phase == "test":
            output = (loc.view(loc.size(0), -1, 4),
                        self.softmax(conf.view(-1, self.num_classes)),
                        detection_dimension)
        else:
            output = (loc.view(loc.size(0), -1, 4),
                        conf.view(conf.size(0), -1, self.num_classes),
                        detection_dimension)

        return output

class S3FD_MV2(nn.Module):

    def __init__(self, phase, size, num_classes):
        super(S3FD_MV2, self).__init__()
        self.phase = phase
        self.num_classes = num_classes
        self.size = size
        self.base_net = MobileNetV2(width_mult=1.0).features
        self.source_layer_indexes = [
            GraphPath(7, 'conv', 3),
            GraphPath(14, 'conv', 3),
            19,
        ]
        # print(self.base_net)

        self.extras = nn.ModuleList([
            InvertedResidual(1280, 512, stride=2, expand_ratio=0.2),
            InvertedResidual(512, 256, stride=2, expand_ratio=0.25),
            InvertedResidual(256, 256, stride=2, expand_ratio=0.5)
        ])

        self.conv3_3_L2Norm = L2Norm(192, 10)
        self.conv4_3_L2Norm = L2Norm(576, 8)
        self.conv5_3_L2Norm = L2Norm(1280, 5)

        self.loc, self.conf = self.multibox(self.num_classes)

        if self.phase == 'test':
            self.softmax = nn.Softmax(dim=-1)

        if self.phase == 'train':
            for m in self.modules():
                if isinstance(m, nn.Conv2d):
                    if m.bias is not None:
                        nn.init.xavier_normal_(m.weight.data)
                        m.bias.data.fill_(0.02)
                    else:
                        m.weight.data.normal_(0, 0.01)
                elif isinstance(m, nn.BatchNorm2d):
                    m.weight.data.fill_(1)
                    m.bias.data.zero_()

    def multibox(self, num_classes):
        loc_layers = []
        conf_layers = []
        # Max-out BG label
        loc_layers += [nn.Conv2d(192, 1 * 4, kernel_size=3, padding=1)]
        conf_layers += [nn.Conv2d(192, 1 * 4, kernel_size=3, padding=1)]
        # conf_layers += [nn.Conv2d(256, 1 * num_classes, kernel_size=3, padding=1)]
        # conv4_3
        loc_layers += [nn.Conv2d(576, 1 * 4, kernel_size=3, padding=1)]
        conf_layers += [nn.Conv2d(576, 1 * num_classes, kernel_size=3, padding=1)]
        # conv5_3
        loc_layers += [nn.Conv2d(1280, 1 * 4, kernel_size=3, padding=1)]
        conf_layers += [nn.Conv2d(1280, 1 * num_classes, kernel_size=3, padding=1)]
        # fc6
        loc_layers += [nn.Conv2d(512, 1 * 4, kernel_size=3, padding=1)]
        conf_layers += [nn.Conv2d(512, 1 * num_classes, kernel_size=3, padding=1)]
        # fc7
        loc_layers += [nn.Conv2d(256, 1 * 4, kernel_size=3, padding=1)]
        conf_layers += [nn.Conv2d(256, 1 * num_classes, kernel_size=3, padding=1)]
        # conv7_2
        loc_layers += [nn.Conv2d(256, 1 * 4, kernel_size=3, padding=1)]
        conf_layers += [nn.Conv2d(256, 1 * num_classes, kernel_size=3, padding=1)]
        return nn.Sequential(*loc_layers), nn.Sequential(*conf_layers)

    def load_weights(self, base_file):
        other, ext = os.path.splitext(base_file)
        if ext == '.pkl' or '.pth':
            print('Loading weights into state dict...')
            self.load_state_dict(torch.load(base_file,
                                            map_location=lambda storage, loc: storage))
            print('Finished!')
        else:
            print('Sorry only .pth and .pkl files supported.')

    def forward(self, x: torch.Tensor):
        confidences = []
        locations = []
        start_layer_index = 0
        header_index = 0
        detection_dimension = list()

        for end_layer_index in self.source_layer_indexes:
            if isinstance(end_layer_index, GraphPath):
                path = end_layer_index
                end_layer_index = end_layer_index.s0
                added_layer = None
            elif isinstance(end_layer_index, tuple):
                added_layer = end_layer_index[1]
                end_layer_index = end_layer_index[0]
                path = None
            else:
                added_layer = None
                path = None
            for layer in self.base_net[start_layer_index: end_layer_index]:
                x = layer(x)
            if added_layer:
                y = added_layer(x)
            else:
                y = x
            if path:
                sub = getattr(self.base_net[end_layer_index], path.name)
                for layer in sub[:path.s1]:
                    x = layer(x)
                y = x
                for layer in sub[path.s1:]:
                    x = layer(x)
                end_layer_index += 1
            start_layer_index = end_layer_index
            confidence, location, dims = self.compute_header(header_index, y)
            header_index += 1
            confidences.append(confidence)
            locations.append(location)
            detection_dimension.append(dims)

        for layer in self.base_net[end_layer_index:]:
            x = layer(x)

        for layer in self.extras:
            x = layer(x)
            confidence, location, dims = self.compute_header(header_index, x)
            header_index += 1
            confidences.append(confidence)
            locations.append(location)
            detection_dimension.append(dims)

        confidences = torch.cat(confidences, 1)
        locations = torch.cat(locations, 1)

        detection_dimension = torch.Tensor(detection_dimension)

        if self.phase == "test":
            output = (locations,
                      self.softmax(confidences),
                      detection_dimension)
        else:
            output = (locations,
                      confidences,
                      detection_dimension)

        return output

    def compute_header(self, i, x):
        # add extra normalization
        if i == 0:
            x =  self.conv3_3_L2Norm(x)
        elif i == 1:
            x = self.conv4_3_L2Norm(x)
        elif i == 2:
            x = self.conv5_3_L2Norm(x)

        # maxout
        if i == 0:
            conf_t = self.conf[i](x)
            max_conf, _ = conf_t[:, 0:3, :, :].max(1, keepdim=True)
            lab_conf = conf_t[:, 3:, :, :]
            confidence = torch.cat((max_conf, lab_conf), dim=1)
        else:
            confidence = self.conf[i](x)

        confidence = confidence.permute(0, 2, 3, 1).contiguous()
        confidence = confidence.view(confidence.size(0), -1, self.num_classes)

        location = self.loc[i](x)
        location = location.permute(0, 2, 3, 1).contiguous()
        location = location.view(location.size(0), -1, 4)

        return confidence, location, x.shape[2:]

class S3FD_FairNAS_A(nn.Module):

    def __init__(self, phase, size, num_classes):
        super(S3FD_FairNAS_A, self).__init__()
        self.phase = phase
        self.num_classes = num_classes
        self.size = size
        self.base_net = FairNasA().features
        self.source_layer_indexes = [
            GraphPath(8, 'conv', 3),
            GraphPath(16, 'conv', 3),
            22,
        ]
        # print(self.base_net)

        self.extras = nn.ModuleList([
            InvertedResidual(1280, 512, stride=2, expand_ratio=0.2),
            InvertedResidual(512, 256, stride=2, expand_ratio=0.25),
            InvertedResidual(256, 256, stride=2, expand_ratio=0.5)
        ])

        # self.expand = InvertedResidual(120, 192, stride=1, expand_ratio=1)

        self.conv3_3_L2Norm = L2Norm(120, 10)
        self.conv4_3_L2Norm = L2Norm(576, 8)
        self.conv5_3_L2Norm = L2Norm(1280, 5)

        self.loc, self.conf = self.multibox(self.num_classes)

        if self.phase == 'test':
            self.softmax = nn.Softmax(dim=-1)

        if self.phase == 'train':
            for m in self.modules():
                if isinstance(m, nn.Conv2d):
                    if m.bias is not None:
                        nn.init.xavier_normal_(m.weight.data)
                        m.bias.data.fill_(0.02)
                    else:
                        m.weight.data.normal_(0, 0.01)
                elif isinstance(m, nn.BatchNorm2d):
                    m.weight.data.fill_(1)
                    m.bias.data.zero_()

    def multibox(self, num_classes):
        loc_layers = []
        conf_layers = []
        # Max-out BG label
        loc_layers += [nn.Conv2d(120, 1 * 4, kernel_size=3, padding=1)]
        conf_layers += [nn.Conv2d(120, 1 * 4, kernel_size=3, padding=1)]
        # conf_layers += [nn.Conv2d(256, 1 * num_classes, kernel_size=3, padding=1)]
        # conv4_3
        loc_layers += [nn.Conv2d(576, 1 * 4, kernel_size=3, padding=1)]
        conf_layers += [nn.Conv2d(576, 1 * num_classes, kernel_size=3, padding=1)]
        # conv5_3
        loc_layers += [nn.Conv2d(1280, 1 * 4, kernel_size=3, padding=1)]
        conf_layers += [nn.Conv2d(1280, 1 * num_classes, kernel_size=3, padding=1)]
        # fc6
        loc_layers += [nn.Conv2d(512, 1 * 4, kernel_size=3, padding=1)]
        conf_layers += [nn.Conv2d(512, 1 * num_classes, kernel_size=3, padding=1)]
        # fc7
        loc_layers += [nn.Conv2d(256, 1 * 4, kernel_size=3, padding=1)]
        conf_layers += [nn.Conv2d(256, 1 * num_classes, kernel_size=3, padding=1)]
        # conv7_2
        loc_layers += [nn.Conv2d(256, 1 * 4, kernel_size=3, padding=1)]
        conf_layers += [nn.Conv2d(256, 1 * num_classes, kernel_size=3, padding=1)]
        return nn.Sequential(*loc_layers), nn.Sequential(*conf_layers)

    def load_weights(self, base_file):
        other, ext = os.path.splitext(base_file)
        if ext == '.pkl' or '.pth':
            print('Loading weights into state dict...')
            self.load_state_dict(torch.load(base_file,
                                            map_location=lambda storage, loc: storage))
            print('Finished!')
        else:
            print('Sorry only .pth and .pkl files supported.')

    def forward(self, x: torch.Tensor):
        confidences = []
        locations = []
        start_layer_index = 0
        header_index = 0
        detection_dimension = list()

        for end_layer_index in self.source_layer_indexes:
            if isinstance(end_layer_index, GraphPath):
                path = end_layer_index
                end_layer_index = end_layer_index.s0
                added_layer = None
            elif isinstance(end_layer_index, tuple):
                added_layer = end_layer_index[1]
                end_layer_index = end_layer_index[0]
                path = None
            else:
                added_layer = None
                path = None
            for layer in self.base_net[start_layer_index: end_layer_index]:
                x = layer(x)
            if added_layer:
                y = added_layer(x)
            else:
                y = x
            if path:
                sub = getattr(self.base_net[end_layer_index], path.name)
                for layer in sub[:path.s1]:
                    x = layer(x)
                y = x
                for layer in sub[path.s1:]:
                    x = layer(x)
                end_layer_index += 1
            start_layer_index = end_layer_index
            confidence, location, dims = self.compute_header(header_index, y)
            header_index += 1
            confidences.append(confidence)
            locations.append(location)
            detection_dimension.append(dims)

        for layer in self.base_net[end_layer_index:]:
            x = layer(x)

        for layer in self.extras:
            x = layer(x)
            confidence, location, dims = self.compute_header(header_index, x)
            header_index += 1
            confidences.append(confidence)
            locations.append(location)
            detection_dimension.append(dims)

        confidences = torch.cat(confidences, 1)
        locations = torch.cat(locations, 1)

        detection_dimension = torch.Tensor(detection_dimension)

        if self.phase == "test":
            output = (locations,
                      self.softmax(confidences),
                      detection_dimension)
        else:
            output = (locations,
                      confidences,
                      detection_dimension)

        return output

    def compute_header(self, i, x):
        # add extra normalization
        if i == 0:
            # x = self.expand(x)
            x =  self.conv3_3_L2Norm(x)
        elif i == 1:
            x = self.conv4_3_L2Norm(x)
        elif i == 2:
            x = self.conv5_3_L2Norm(x)

        # maxout
        if i == 0:
            conf_t = self.conf[i](x)
            max_conf, _ = conf_t[:, 0:3, :, :].max(1, keepdim=True)
            lab_conf = conf_t[:, 3:, :, :]
            confidence = torch.cat((max_conf, lab_conf), dim=1)
        else:
            confidence = self.conf[i](x)

        confidence = confidence.permute(0, 2, 3, 1).contiguous()
        confidence = confidence.view(confidence.size(0), -1, self.num_classes)

        location = self.loc[i](x)
        location = location.permute(0, 2, 3, 1).contiguous()
        location = location.view(location.size(0), -1, 4)

        return confidence, location, x.shape[2:]

class S3FD_FairNAS_B(nn.Module):

    def __init__(self, phase, size, num_classes):
        super(S3FD_FairNAS_B, self).__init__()
        self.phase = phase
        self.num_classes = num_classes
        self.size = size
        self.base_net = FairNasB().features
        self.source_layer_indexes = [
            GraphPath(8, 'conv', 3),
            GraphPath(16, 'conv', 3),
            22,
        ]
        # print(self.base_net)

        self.extras = nn.ModuleList([
            InvertedResidual(1280, 512, stride=2, expand_ratio=0.2),
            InvertedResidual(512, 256, stride=2, expand_ratio=0.25),
            InvertedResidual(256, 256, stride=2, expand_ratio=0.5)
        ])

        # self.expand = InvertedResidual(120, 192, stride=1, expand_ratio=1)

        self.conv3_3_L2Norm = L2Norm(120, 10)
        self.conv4_3_L2Norm = L2Norm(576, 8)
        self.conv5_3_L2Norm = L2Norm(1280, 5)

        self.loc, self.conf = self.multibox(self.num_classes)

        if self.phase == 'test':
            self.softmax = nn.Softmax(dim=-1)

        if self.phase == 'train':
            for m in self.modules():
                if isinstance(m, nn.Conv2d):
                    if m.bias is not None:
                        nn.init.xavier_normal_(m.weight.data)
                        m.bias.data.fill_(0.02)
                    else:
                        m.weight.data.normal_(0, 0.01)
                elif isinstance(m, nn.BatchNorm2d):
                    m.weight.data.fill_(1)
                    m.bias.data.zero_()

    def multibox(self, num_classes):
        loc_layers = []
        conf_layers = []
        # Max-out BG label
        loc_layers += [nn.Conv2d(120, 1 * 4, kernel_size=3, padding=1)]
        conf_layers += [nn.Conv2d(120, 1 * 4, kernel_size=3, padding=1)]
        # conf_layers += [nn.Conv2d(256, 1 * num_classes, kernel_size=3, padding=1)]
        # conv4_3
        loc_layers += [nn.Conv2d(576, 1 * 4, kernel_size=3, padding=1)]
        conf_layers += [nn.Conv2d(576, 1 * num_classes, kernel_size=3, padding=1)]
        # conv5_3
        loc_layers += [nn.Conv2d(1280, 1 * 4, kernel_size=3, padding=1)]
        conf_layers += [nn.Conv2d(1280, 1 * num_classes, kernel_size=3, padding=1)]
        # fc6
        loc_layers += [nn.Conv2d(512, 1 * 4, kernel_size=3, padding=1)]
        conf_layers += [nn.Conv2d(512, 1 * num_classes, kernel_size=3, padding=1)]
        # fc7
        loc_layers += [nn.Conv2d(256, 1 * 4, kernel_size=3, padding=1)]
        conf_layers += [nn.Conv2d(256, 1 * num_classes, kernel_size=3, padding=1)]
        # conv7_2
        loc_layers += [nn.Conv2d(256, 1 * 4, kernel_size=3, padding=1)]
        conf_layers += [nn.Conv2d(256, 1 * num_classes, kernel_size=3, padding=1)]
        return nn.Sequential(*loc_layers), nn.Sequential(*conf_layers)

    def load_weights(self, base_file):
        other, ext = os.path.splitext(base_file)
        if ext == '.pkl' or '.pth':
            print('Loading weights into state dict...')
            self.load_state_dict(torch.load(base_file,
                                            map_location=lambda storage, loc: storage))
            print('Finished!')
        else:
            print('Sorry only .pth and .pkl files supported.')

    def forward(self, x: torch.Tensor):
        confidences = []
        locations = []
        start_layer_index = 0
        header_index = 0
        detection_dimension = list()

        for end_layer_index in self.source_layer_indexes:
            if isinstance(end_layer_index, GraphPath):
                path = end_layer_index
                end_layer_index = end_layer_index.s0
                added_layer = None
            elif isinstance(end_layer_index, tuple):
                added_layer = end_layer_index[1]
                end_layer_index = end_layer_index[0]
                path = None
            else:
                added_layer = None
                path = None
            for layer in self.base_net[start_layer_index: end_layer_index]:
                x = layer(x)
            if added_layer:
                y = added_layer(x)
            else:
                y = x
            if path:
                sub = getattr(self.base_net[end_layer_index], path.name)
                for layer in sub[:path.s1]:
                    x = layer(x)
                y = x
                for layer in sub[path.s1:]:
                    x = layer(x)
                end_layer_index += 1
            start_layer_index = end_layer_index
            confidence, location, dims = self.compute_header(header_index, y)
            header_index += 1
            confidences.append(confidence)
            locations.append(location)
            detection_dimension.append(dims)

        for layer in self.base_net[end_layer_index:]:
            x = layer(x)

        for layer in self.extras:
            x = layer(x)
            confidence, location, dims = self.compute_header(header_index, x)
            header_index += 1
            confidences.append(confidence)
            locations.append(location)
            detection_dimension.append(dims)

        confidences = torch.cat(confidences, 1)
        locations = torch.cat(locations, 1)

        detection_dimension = torch.Tensor(detection_dimension)

        if self.phase == "test":
            output = (locations,
                      self.softmax(confidences),
                      detection_dimension)
        else:
            output = (locations,
                      confidences,
                      detection_dimension)

        return output

    def compute_header(self, i, x):
        # add extra normalization
        if i == 0:
            # x = self.expand(x)
            x =  self.conv3_3_L2Norm(x)
        elif i == 1:
            x = self.conv4_3_L2Norm(x)
        elif i == 2:
            x = self.conv5_3_L2Norm(x)

        # maxout
        if i == 0:
            conf_t = self.conf[i](x)
            max_conf, _ = conf_t[:, 0:3, :, :].max(1, keepdim=True)
            lab_conf = conf_t[:, 3:, :, :]
            confidence = torch.cat((max_conf, lab_conf), dim=1)
        else:
            confidence = self.conf[i](x)

        confidence = confidence.permute(0, 2, 3, 1).contiguous()
        confidence = confidence.view(confidence.size(0), -1, self.num_classes)

        location = self.loc[i](x)
        location = location.permute(0, 2, 3, 1).contiguous()
        location = location.view(location.size(0), -1, 4)

        return confidence, location, x.shape[2:]

if __name__ == '__main__':
    net = S3FD_MV2('train', 640, 2)
    image = np.zeros((1,3,640,640), dtype=np.float32)
    output = net.forward(torch.from_numpy(image))
    print(output)

    net = S3FD_FairNAS_A('train', 640, 2)
    image = np.zeros((1,3,640,640), dtype=np.float32)
    output = net.forward(torch.from_numpy(image))
    print(output)

    net = S3FD_FairNAS_B('train', 640, 2)
    image = np.zeros((1,3,640,640), dtype=np.float32)
    output = net.forward(torch.from_numpy(image))
    print(output)


'''
  Brought from https://github.com/zju3dv/EasyMocap/blob/master/easymocap/socket/o3d.py
'''
import os,sys
import open3d as o3d
from .vis_utils_o3d13 import Timer, get_rgb_01, CritRange, Config, Vector3dVector, Vector2iVector, get_uvs,\
                                    load_mesh, create_mesh, create_mesh_with_uvmap, convert_verts_to_cam_space
from .socket_utils import BaseSocket, log
import json
import numpy as np
from os.path import join
import copy
import importlib
from config import args
import constants
import torch
import random

def load_object(module_name, module_args):
    module_path = '.'.join(module_name.split('.')[:-1])
    module = importlib.import_module(module_path)
    name = module_name.split('.')[-1]
    obj = getattr(module, name)(**module_args)
    return obj

def merge_params(params):
    params_batch = {}
    for key in params[0]:
        if isinstance(params[0][key], np.ndarray):
            if len(params[0][key].shape)>0:
                params_batch[key] = torch.from_numpy(np.concatenate([param[key] for param in params], 0)).float()
    return params_batch

rotate = False
def o3d_callback_rotate(vis=None):
    global rotate
    rotate = not rotate
    return False

class VisOpen3DSocket(BaseSocket):
    def __init__(self, host, port, cfg) -> None:
        # output
        self.write = cfg.write
        self.out = cfg.out
        self.cfg = cfg
        if self.write:
            print('[Info] capture the screen to {}'.format(self.out))
            os.makedirs(self.out, exist_ok=True)
        # scene
        vis = o3d.visualization.VisualizerWithKeyCallback()
        vis.register_key_callback(ord('A'), o3d_callback_rotate)
        if cfg.rotate:
            o3d_callback_rotate()
        vis.create_window(window_name='Visualizer', width=cfg.width, height=cfg.height)
        self.vis = vis
        # load the scene
        for key, mesh_args in cfg.scene.items():
            mesh = load_object(key, mesh_args)
            self.vis.add_geometry(mesh)
        for key, val in cfg.extra.items():
            mesh = load_mesh(val["path"])
            trans = np.array(val['transform']).reshape(4, 4)
            mesh.transform(trans)
            self.vis.add_geometry(mesh)
        # create vis => update => super() init
        super().__init__(host, port, debug=cfg.debug)
        self.block = cfg.block
        if os.path.exists(cfg.body_model_template):
            body_template = o3d.io.read_triangle_mesh(cfg.body_model_template)
            self.body_template = body_template
        else:
            self.body_template = None
        self.body_model = load_object(cfg.body_model.module, {'model_path':args().smpl_model_path})
        zero_params = {
            'poses': torch.zeros(1, 72).float(),
            'betas': torch.zeros(1, 10).float(),
        }
        self.max_human = cfg.max_human
        self.track = cfg.track
        self.filter = cfg.filter
        self.camera_pose = cfg.camera.camera_pose
        self.init_camera(cfg.camera.camera_pose)
        self.zero_vertices = Vector3dVector(np.zeros((6890, 3)))
        self.filter_dict = {}

        self.vertices, self.meshes = [], []
        self.verts_available_ids, self.current_mesh_num = [], 0
        self.pid_verts_dict = {}
        for i in range(self.max_human):
            self.add_human(zero_params)
        self.verts_available_ids = sorted(self.verts_available_ids,reverse = True)
        self.set_meshes_zero()
        self.verts_change_cacher = []

        self.count = 0
        self.previous = {}
        self.critrange = CritRange(**cfg.range)
        self.new_frames  = cfg.new_frames

    def set_meshes_zero(self):
        for ind in range(self.current_mesh_num):
            self.meshes[ind].vertices = self.zero_vertices

    def add_human(self, zero_params):
        smpl_outs = self.body_model(**zero_params)
        vertices = smpl_outs['verts'].cpu().numpy()[0]
        self.vertices.append(vertices)
        self.verts_available_ids.append(self.current_mesh_num)
        self.current_mesh_num+=1
        faces = self.body_model.faces

        if args().mesh_cloth in constants.wardrobe or args().mesh_cloth=='random':
            uvs = get_uvs(args().smpl_uvmap)
            if args().mesh_cloth=='random':
                cloth_id = random.sample(list(constants.wardrobe.keys()), 1)[0]
                print('choose cloth: ', cloth_id, constants.wardrobe[cloth_id])
                texture_file = os.path.join(args().wardrobe, constants.wardrobe[cloth_id])
            else:
                texture_file = os.path.join(args().wardrobe, constants.wardrobe[args().mesh_cloth])
            mesh = create_mesh_with_uvmap(vertices, faces, texture_path=texture_file, uvs=Vector2iVector(uvs))
        elif args().mesh_cloth in constants.mesh_color_dict:
            mesh_color = np.array(constants.mesh_color_dict[args().mesh_cloth])/255.
            mesh = create_mesh(vertices=vertices, faces=faces, colors=mesh_color)
        else:
            mesh = create_mesh(vertices=vertices, faces=faces)

        self.meshes.append(mesh)
        self.vis.add_geometry(mesh)
        self.init_camera(self.camera_pose)

    @staticmethod
    def set_camera(cfg, camera_pose):
        theta, phi = np.deg2rad(-(cfg.camera.theta + 90)), np.deg2rad(cfg.camera.phi)
        theta = theta + np.pi
        st, ct = np.sin(theta), np.cos(theta)
        sp, cp = np.sin(phi), np.cos(phi)
        rot_x = np.array([
            [1., 0., 0.],
            [0., ct, -st],
            [0, st, ct]
        ])
        rot_z = np.array([
            [cp, -sp, 0],
            [sp, cp, 0.],
            [0., 0., 1.]
        ])
        camera_pose[:3, :3] = rot_x @ rot_z
        return camera_pose

    def init_camera(self, camera_pose):
        ctr = self.vis.get_view_control()
        init_param = ctr.convert_to_pinhole_camera_parameters()
        # init_param.intrinsic.set_intrinsics(init_param.intrinsic.width, init_param.intrinsic.height, fx, fy, cx, cy)
        init_param.extrinsic = np.array(camera_pose)
        ctr.convert_from_pinhole_camera_parameters(init_param)

    def get_camera(self):
        ctr = self.vis.get_view_control()
        init_param = ctr.convert_to_pinhole_camera_parameters()
        return np.array(init_param.extrinsic)

    def main(self, datas):
        if self.debug:log('[Info] Load data {}'.format(self.count))
        if isinstance(datas, str):
            datas = json.loads(datas)
        for data in datas:
            for key in data.keys():
                if key == 'id':
                    continue
                data[key] = np.array(data[key])

        with Timer('forward'):
            params = []
            for i, data in enumerate(datas):
                if i >= len(self.meshes):
                    print('[Error] the number of human exceeds!')
                    self.add_human(data)
                if 'vertices' in data.keys():
                    vertices = data['vertices']
                    self.vertices[i] = Vector3dVector(vertices)
                else:
                    params.append(data)
            if len(params) > 0:
                params = merge_params(params)
                vertices = self.body_model(**params)['verts'].cpu().numpy()
                for i in range(vertices.shape[0]):
                    verts_processed = convert_verts_to_cam_space(vertices[i])
                    self.vertices[i] = Vector3dVector(verts_processed)
            for i in range(len(datas), len(self.meshes)):
                self.vertices[i] = self.zero_vertices
        # Open3D will lock the thread here
        with Timer('set vertices'):
            for i in range(len(self.vertices)):
                self.meshes[i].vertices = self.vertices[i]

    def o3dcallback(self):
        if rotate:
            self.cfg.camera.phi += np.pi/10
            camera_pose = self.set_camera(self.cfg, self.get_camera())
            self.init_camera(camera_pose)

    def update(self):
        if self.disconnect and not self.block:
            self.previous.clear()
        if not self.queue.empty():
            if self.debug:log('Update' + str(self.queue.qsize()))
            datas = self.queue.get()
            if not self.block:
                while self.queue.qsize() > 0:
                    datas = self.queue.get()
            self.main(datas)
            with Timer('update geometry'):
                for mesh in self.meshes:
                    mesh.compute_triangle_normals()
                    self.vis.update_geometry(mesh)
                self.o3dcallback()
                self.vis.poll_events()
                self.vis.update_renderer()
            if self.write:
                outname = join(self.out, '{:06d}.jpg'.format(self.count))
                with Timer('capture'):
                    self.vis.capture_screen_image(outname)
            self.count += 1
        else:
            with Timer('update renderer', True):
                self.o3dcallback()
                self.vis.poll_events()
                self.vis.update_renderer()

'''
                if i >= len(self.meshes):
                    print('[Error] the number of human exceeds!')
                    self.add_human(data)

            #if 'keypoints3d' not in data.keys() and self.filter:
            #    smpl_outs = self.body_model(**data)
            #    data['keypoints3d'] = smpl_outs['j3d'].cpu().numpy()[0]
        #if self.filter:
        #   datas = self.filter_human(datas)
    def filter_human(self, datas):
        datas_new = []
        for data in datas:
            kpts3d = np.array(data['keypoints3d'])
            data['keypoints3d'] = kpts3d
            pid = data['id']
            if pid not in self.previous.keys():
                if not self.critrange(kpts3d):
                    continue
                self.previous[pid] = 0
            self.previous[pid] += 1
            if self.previous[pid] > self.new_frames:
                datas_new.append(data)
        return datas_new

'''


#   Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import numpy as np
import paddle.fluid as fluid
import parl
from parl import layers
from parl.utils import machine_info

class AtariAgent(parl.Agent):
    def __init__(self, algorithm, obs_shape, act_dim,
                 learn_data_provider=None):
        assert isinstance(obs_shape, (list, tuple))
        assert isinstance(act_dim, int)
        self.obs_shape = obs_shape
        self.act_dim = act_dim
        super(AtariAgent, self).__init__(algorithm)

        if learn_data_provider:
            self.learn_reader.decorate_tensor_provider(learn_data_provider)
            self.learn_reader.start()

    def build_program(self):
        self.sample_program = fluid.Program()
        self.predict_program = fluid.Program()
        self.learn_program = fluid.Program()

        with fluid.program_guard(self.sample_program):
            obs = layers.data(
                name='obs', shape=self.obs_shape, dtype='float32')
            self.sample_actions, self.behaviour_logits = self.alg.sample(obs)

        with fluid.program_guard(self.predict_program):
            obs = layers.data(
                name='obs', shape=self.obs_shape, dtype='float32')
            self.predict_actions = self.alg.predict(obs)

        with fluid.program_guard(self.learn_program):
            obs = layers.data(
                name='obs', shape=self.obs_shape, dtype='float32')
            actions = layers.data(name='actions', shape=[], dtype='int64')
            behaviour_logits = layers.data(
                name='behaviour_logits', shape=[self.act_dim], dtype='float32')
            rewards = layers.data(name='rewards', shape=[], dtype='float32')
            dones = layers.data(name='dones', shape=[], dtype='float32')
            lr = layers.data(
                name='lr', shape=[1], dtype='float32', append_batch_size=False)
            entropy_coeff = layers.data(
                name='entropy_coeff',
                shape=[1],
                dtype='float32',
                append_batch_size=False)

            self.learn_reader = fluid.layers.create_py_reader_by_data(
                capacity=32,
                feed_list=[
                    obs, actions, behaviour_logits, rewards, dones, lr,
                    entropy_coeff
                ])

            obs, actions, behaviour_logits, rewards, dones, lr, entropy_coeff = fluid.layers.read_file(
                self.learn_reader)

            vtrace_loss, kl = self.alg.learn(obs, actions, behaviour_logits,
                                             rewards, dones, lr, entropy_coeff)
            self.learn_outputs = [
                vtrace_loss.total_loss, vtrace_loss.pi_loss,
                vtrace_loss.vf_loss, vtrace_loss.entropy, kl
            ]
        self.learn_program = parl.compile(self.learn_program,
                                          vtrace_loss.total_loss)

    def sample(self, obs_np):
        """
        Args:
            obs_np: a numpy float32 array of shape ([B] + observation_space).
            Format of image input should be NCHW format.

        Returns:
            sample_ids: a numpy int64 array of shape [B]
        """
        obs_np = obs_np.astype('float32')

        sample_actions, behaviour_logits = self.fluid_executor.run(
            self.sample_program,
            feed={'obs': obs_np},
            fetch_list=[self.sample_actions, self.behaviour_logits])
        return sample_actions, behaviour_logits

    def predict(self, obs_np):
        """
        Args:
            obs_np: a numpy float32 array of shape ([B] + observation_space)
            Format of image input should be NCHW format.

        Returns:
            sample_ids: a numpy int64 array of shape [B]
        """
        obs_np = obs_np.astype('float32')

        predict_actions = self.fluid_executor.run(
            self.predict_program,
            feed={'obs': obs_np},
            fetch_list=[self.predict_actions])[0]
        return predict_actions

    def learn(self):
        total_loss, pi_loss, vf_loss, entropy, kl = self.fluid_executor.run(
            self.learn_program, fetch_list=self.learn_outputs)
        return total_loss, pi_loss, vf_loss, entropy, kl


#!/usr/bin/env python
"""Pre-schedule DDF sequences
"""
# pylint: disable=no-member

# imports
import sys
import logging
from argparse import ArgumentParser

import yaml
import numpy as np
import pandas as pd
import astropy.coordinates
import astropy.units as u

import lsst.sims.utils

# constants

# exception classes

# interface functions

def schedule_all(mag_limit, location, config):
    """Schedule one field on one band.

    Parameters
    ----------
    m5 : `pandas.DataFrame`
        Has a multilevel index with the following levels:
        field_name : `str`
            the field name
        band : `str`
            the band

        Includes the following columns:
        mjd : `float`
             MJD of candidate time
        m5 : `float`
             5-sigma limiting magnitude of the field if observed at that time
    `location` : `astropy.coordinates.EarthLocation`
        the location of the observatory
    config : `dict`
        Configuration parameters

    Return
    ------
    schedule : `pandas.DataFrame`
        includes three columns:
        mjd : `float`
            the best time at which to start the sequence of exposures
        why : `str`
            an indicator of why this sequence was scheduled
        night : `int`
            the MJD of the night (at midnight) on which the sequence is to be scheduled
        sequence : `str`
            which sequence this is

    """
    seq_schedules = []

    for seq_config in config["sequences"]:
        logger.info(f'Scheduling {seq_config["label"]}')
        seq_schedule = schedule_sequence(mag_limit, location, seq_config)
        seq_schedule["sequence"] = seq_config["label"]
        logger.info(f'Computing scheduled for {seq_config["label"]}')
        mag_limit["scheduled"] = _compute_scheduled(
            mag_limit, seq_schedule, seq_config["sequence_duration"]
        )
        seq_schedules.append(seq_schedule)

    logger.info("Compiling full schedule")
    full_schedule = (
        pd.concat(seq_schedules).sort_values("mjd").set_index("mjd", drop=False)
    )
    return full_schedule

def schedule_sequence(mag_limit, location, config):
    """Schedule one set of sequences.

    Parameters
    ----------
    m5 : `pandas.DataFrame`
        Has a multilevel index with the following levels:
        field_name : `str`
            the field name
        band : `str`
            the band

        Includes the following columns:
        mjd : `float`
             MJD of candidate time
        m5 : `float`
             5-sigma limiting magnitude of the field if observed at that time
    `location` : `astropy.coordinates.EarthLocation`
        the location of the observatory
    config : `dict`
        Configuration parameters, with the following contents:
        field_name : `str`
            the name of the field to schedule
        mag_lim_band : `str`
            the name of the filter to schedule
        sequence_duration : `astropy.units.Quantity`
            the duration of a block of one sequence of exposures
        caninocal_gap : `astropy.units.Quantity`
            the desired time between sequences of exposures
        min_gap: `astropy.units.Quantity`
            the minimum gap for which "bridge" exposures should be scheduled
        max_gap: `astropy.units.Quantity`
            the target maximum time between sequences of exposures
        season_gap :  `astropy.units.Quantity`
            the gap time greater than which no bridges should be attempted
        mag_limit : `dict` of `str`: `float`
            target magnitude limits in each band

    Return
    ------
    schedule : `pandas.DataFrame`
        includes three columns:
        mjd : `float`
            the best time at which to start the sequence of exposures
        why : `str`
            an indicator of why this sequence was scheduled
        night : `int`
            the MJD of the night (at midnight) on which the sequence is to be scheduled

    """
    # pylint: disable=too-many-locals
    these_m5 = (
        mag_limit.sort_index()
        .loc[(config["field_name"], config["mag_lim_band"])]
        .sort_index()
        .copy()
    )

    min_m5 = _compute_rolling_m5(these_m5, config["sequence_duration"]).set_index(
        "mjd", drop=False
    )

    min_m5["night_mjd"] = compute_night_mjd(min_m5["mjd"], location)

    bridge_nights = _find_bridge_nights(mag_limit, location, config)
    bridge_gap = config["bridge_gap"]
    maintain_cadence = config["maintain_cadence_in_gap"]

    scheduled_sequences = []
    for night_mjd in range(min_m5.night_mjd.min(), min_m5.night_mjd.max()):
        if night_mjd in bridge_nights["night_before_mjd"].values:
            why = "pregap"
            attempt_tonight = True
            force_tonight = True
        elif bridge_gap and (night_mjd in bridge_nights["bridge_night_mjd"].values):
            why = "bridge"
            attempt_tonight = True
            force_tonight = True
        elif night_mjd in bridge_nights["night_after_mjd"].values:
            why = "postgap"
            attempt_tonight = True
            force_tonight = True
        elif len(scheduled_sequences) == 0:
            # We are just starting
            why = "start"

            attempt_tonight = True
            force_tonight = False
        elif (night_mjd - scheduled_sequences[-1]["night_mjd"]) * u.day >= config[
            "canonical_gap"
        ]:
            why = "cadence"
            attempt_tonight = True
            force_tonight = maintain_cadence
        else:
            continue

        if not attempt_tonight:
            continue

        candidate_times = min_m5.query(f"night_mjd == {night_mjd}")
        if len(candidate_times) < 1:
            assert maintain_cadence or not force_tonight
            continue

        best_time = min_m5.loc[candidate_times["m5"].idxmax()]
        if isinstance(best_time, pd.DataFrame):
            best_time = best_time.sort_values("count", ascending=True).iloc[-1]
        if (not force_tonight) and (best_time.m5 < config["mag_limit"]):
            continue

        if best_time.m5 < config["gap_mag_limit"]:
            continue

        scheduled_sequences.append({"mjd": best_time.mjd, "why": why})
        scheduled_sequences[-1]["night_mjd"] = compute_night_mjd(
            best_time.mjd, location
        )

    schedule = pd.DataFrame(scheduled_sequences)

    return schedule

def compute_night_mjd(mjd, location):
    """Convert the floating point mjd to the integer local Julian date for the night.

    Parameters
    ----------
    mjd : `float`, `pandas.Series`, or `numpy.ndarray`

    Returns
    -------
    jd : `int`, `pandas.Series`, or `numpy.ndarray`
    """
    # add longitude to get into the local timezone,
    # round to find the nearest midnight
    night_mjd = np.round(mjd + (location.lon.deg / 360.0)).astype(int)
    return night_mjd

def read_config(fname):
    """Read m5 configuration file

    Parameters
    ----------
    fname: `str`
        The name of the file to read configuration from.

    Return
    ------
    config: `dict`
        Dictionary of configuration values
    """
    logger.debug("Reading configuration from %s", fname)

    with open(fname, "r") as config_file:
        config = yaml.load(config_file.read(), Loader=yaml.FullLoader)

    # Apply units
    for seq_config in config["sequences"]:
        seq_config["sequence_duration"] = u.Quantity(
            seq_config["sequence_duration"]
        ).to(u.second)
        seq_config["max_gap"] = u.Quantity(seq_config["max_gap"]).to(u.day)
        seq_config["min_gap"] = u.Quantity(seq_config["min_gap"]).to(u.day)
        seq_config["season_gap"] = u.Quantity(seq_config["season_gap"]).to(u.day)
        seq_config["canonical_gap"] = u.Quantity(seq_config["canonical_gap"]).to(u.day)

    site_name = "LSST" if config["site_name"] == "LSST" else config["site_name"]
    site = lsst.sims.utils.Site(site_name)
    config["location"] = astropy.coordinates.EarthLocation(
        lat=site.latitude, lon=site.longitude, height=site.height
    )

    return config

# classes

# internal functions & classes

def _infer_time_sampling(mag_limit):
    mjds = pd.Series(mag_limit["mjd"].unique()).sort_values()
    timestep_duration = ((mjds - mjds.shift(1)).median() * u.day).to(u.minute)
    return timestep_duration

def _compute_rolling_m5(mag_limit, roll_window):
    mag_limit = mag_limit.query("not scheduled").copy().sort_index()
    mag_limit["datetime"] = pd.to_datetime(
        mag_limit.mjd + 2400000.5, origin="julian", unit="D"
    )
    mag_limit["counter"] = 1
    mag_limit.set_index("datetime", inplace=True, drop=False)
    roll_seconds = roll_window.to("second").value
    mag_limit_roll = mag_limit.rolling(f"{int(roll_seconds)}s")
    min_mag_limit = mag_limit_roll[["mjd", "moon_angle", "night", "m5"]].min()
    min_mag_limit["start_datetime"] = pd.to_datetime(
        min_mag_limit.mjd + 2400000.5, origin="julian", unit="D"
    )
    min_mag_limit["count"] = mag_limit_roll["counter"].sum().astype(int)
    min_mag_limit = (
        min_mag_limit.reset_index()
        .rename(columns={"datetime": "end_datetime"})
        .set_index("start_datetime", drop=False)
    )
    min_mag_limit["m5"] = min_mag_limit["m5"].fillna(-np.inf)

    # Infer which windows do not have a full set of samples, and toss them
    sample_dt = _infer_time_sampling(mag_limit)
    expected_samples = int(np.floor((roll_window.to(sample_dt.unit) / sample_dt).value))

    min_mag_limit.query(
        f"(count == {expected_samples}) or (count == {expected_samples+1})",
        inplace=True,
    )
    min_mag_limit.sort_values("count", ascending=False).groupby(
        level="start_datetime"
    ).first()

    return min_mag_limit

def _find_gaps(mjds, min_gap, season_gap, location, night_epoch_mjd=0):
    gaps = pd.DataFrame({"start": np.unique(np.sort(mjds))})
    gaps["end"] = gaps.start.shift(-1)
    gaps.dropna(inplace=True)
    gaps["duration"] = gaps["end"] - gaps["start"]
    gaps["mjd"] = 0.5 * (gaps["end"] + gaps["start"])
    gaps["night_before"] = compute_night_mjd(gaps["start"], location) - night_epoch_mjd
    gaps["night_after"] = compute_night_mjd(gaps["end"], location) - night_epoch_mjd
    gaps["gap_nights"] = gaps["night_after"] - gaps["night_before"]
    gaps.query(
        f"({min_gap} <= gap_nights) and ({season_gap} > gap_nights)", inplace=True
    )
    gaps.set_index("mjd", inplace=True)
    gaps.sort_index(inplace=True)
    return gaps

def _find_bridge_nights(all_mag_limit, location, config):
    oversampled_mag_limit = (
        all_mag_limit.sort_index()
        .loc[(config["field_name"], config["mag_lim_band"])]
        .sort_index()
        .copy()
    )
    mag_limit = _compute_rolling_m5(oversampled_mag_limit, config["sequence_duration"])
    good_mag_limit = mag_limit.query(f'm5>{config["mag_limit"]}')
    night_epoch_mjd = (
        compute_night_mjd(mag_limit.iloc[0].mjd, location) - mag_limit.iloc[0].night
    )
    gaps = _find_gaps(
        good_mag_limit.mjd,
        config["min_gap"].to(u.day).value,
        config["season_gap"].to(u.day).value,
        location,
        night_epoch_mjd,
    )
    gaps["bridge_mjd"] = np.nan
    gaps["has_bridge"] = False
    max_gap = config["max_gap"].to(u.day).value
    for mjd, gap in gaps.iterrows():
        candidate_bridges = mag_limit.query(
            f"(night > {gap.night_before}) and (night < {gap.night_after})"
        ).query(f"(mjd < {gap.start+max_gap}) and (mjd > {gap.end-max_gap})")
        if len(candidate_bridges) == 0:
            continue
        best_bridge = candidate_bridges.loc[candidate_bridges["m5"].idxmax()]
        # Sometimes there can be two time windows with the same starting,
        # differing by a sample time.
        if isinstance(best_bridge, pd.DataFrame):
            best_bridge = best_bridge.sort_values("count").iloc[-1]

        gaps["has_bridge"] = True
        gaps.loc[mjd, "bridge_mjd"] = best_bridge["mjd"]

    gaps["bridge_night_mjd"] = compute_night_mjd(gaps["bridge_mjd"].fillna(0), location)
    gaps["night_before_mjd"] = (gaps["night_before"] + night_epoch_mjd).astype(int)
    gaps["night_after_mjd"] = (gaps["night_after"] + night_epoch_mjd).astype(int)
    return gaps

def _compute_scheduled(m5_limits, schedule, sequence_duration):
    scheduled = (
        m5_limits["scheduled"]
        .reset_index()
        .set_index("mjd", drop=False)
        .sort_index()
        .copy()
    )

    seq_days = sequence_duration.to(u.day).value
    for _, obs_seq in schedule.iterrows():
        start_mjd = obs_seq.mjd
        end_mjd = obs_seq.mjd + seq_days
        scheduled.loc[start_mjd:end_mjd, "scheduled"] = True

    scheduled.set_index(m5_limits.index.names, inplace=True)
    return scheduled["scheduled"]

def main():
    """Parse command line arguments and config file, and run"""
    parser = ArgumentParser()
    parser.add_argument("config", help="configuration file")
    parser.add_argument("m5", help="file from which to load limiting magnitudes")
    parser.add_argument("output", help="file in which to write results")

    args = parser.parse_args()
    config_fname = args.config
    m5_fname = args.m5
    output_fname = args.output

    config = read_config(config_fname)

    logger.info("Reading m5 from %s", m5_fname)
    m5_limits = (
        pd.read_hdf(m5_fname)
        .reset_index()
        .query("sun_alt < -18")
        .set_index(["field_name", "band", "mjd"], drop=False)
        .assign(scheduled=False)
    )

    schedule = schedule_all(m5_limits, config["location"], config)
    schedule.to_csv(output_fname, sep="\t", index=False, header=True)

    return 0

def _init_logger(log_level=logging.DEBUG):
    """Create the ddfpresched logger and set initial configuration"""
    ddfpresched_logger = logging.getLogger("ddfpresched")
    ddfpresched_logger.setLevel(log_level)
    handler = logging.StreamHandler()
    handler.setLevel(log_level)
    formatter = logging.Formatter("%(asctime)s\t%(name)s\t%(levelname)s\t%(message)s")
    handler.setFormatter(formatter)
    ddfpresched_logger.addHandler(handler)
    return ddfpresched_logger

if __name__ == "__main__":
    logger = _init_logger()
    status = main()  # pylint: disable=invalid-name
    sys.exit(status)


from rdkit import Chem
from rdkit import DataStructs
from random import shuffle
import numpy as np
import time
from rdkit.Chem import Descriptors
from tqdm import tqdm
from multiprocessing import Process
import os
import subprocess

def get_(similarity_lib, scale, to_file=False, n_dev=10000, show=True, ids=None):
    if type(similarity_lib) == str:
        with open(similarity_lib,'r') as f:
            libm = []
            libs = []
            for l in f:
                m, sm = l.strip().split(':')
                libm.append(m)
                libs.append(sm.split(','))
    else:
        libm = [i[0] for i in similarity_lib]
        libs = [i[1] for i in similarity_lib]
    libmq = []
    print('cal ref QEDs')
    with tqdm(total=len(libm)) as ref_pbar:
        for i in libm:
            libmq.append(Descriptors.qed(Chem.MolFromSmiles(i)))
            if show:
                ref_pbar.update()
    libsq = []
    libss = []
    print('cal candidate QEDs')
    with tqdm(total=len(libs)) as cdd_pbar:
        for lidx,i in enumerate(libs):
            temp_ = []
            k = 0
            tp = libm[lidx]
            while len(temp_)<scale and k<len(i):
                if i[k] != tp:
                    temp_.append(i[k])
                k += 1
            libss.append(temp_)
            libsq.append([Descriptors.qed(Chem.MolFromSmiles(j)) for j in temp_])
            if show:
                cdd_pbar.update()
    opt = []
    optv = []
    print('build pair')
    with tqdm(total=len(libm)) as bd_pbar:
        for midx in range(len(libm)):
            diff = [abs(libmq[midx]-libsq[midx][cidx]) for cidx in range(len(libsq[midx]))]
            sel = np.argmax(diff)
            optv.append(max(diff))
            if libmq[midx]<libsq[midx][sel]:
                opt.append([libm[midx], libss[midx][sel]])
            else:
                opt.append([libss[midx][sel], libm[midx]])
            if show:
                bd_pbar.update()
    print('remove repeats')
    opt = ['&'.join(i) for i in opt]
    opt = list(set(opt))
    opt = [i.split('&') for i in opt]

    if to_file:
        with open(to_file,'w') as f:
            for r in opt:
                f.write(','.join([str(i) for i in r])+'\n')
    simv = []
    print('cal pair similarity')
    with tqdm(total=len(libm)) as sv_pbar:
        for r in opt:
            simv.append(DataStructs.TanimotoSimilarity(Chem.RDKFingerprint(Chem.MolFromSmiles(r[0])),Chem.RDKFingerprint(Chem.MolFromSmiles(r[1]))))
            if show:
                sv_pbar.update()
    optv = np.mean(optv)
    simv = np.mean(simv)
    print('split data')
    idx= list(range(len(opt)))
    shuffle(idx)
    train_idx = idx[:-10000]
    dev_idx = idx[-10000:]
    train_opt = [opt[i] for i in train_idx]
    dev_opt = [opt[i] for i in dev_idx]

    if ids == None:
        return train_opt, dev_opt, optv, simv
    else:
        with open('train_subprocess_{0}.tsv'.format(ids), 'w') as f:
            for r in train_opt:
                f.write('{0}\t{1}\n'.format(r[0], r[1]))
        with open('dev_subprocess_{0}.tsv'.format(ids), 'w') as f:
            for r in dev_opt:
                f.write('{0}\t{1}\n'.format(r[0], r[1]))
        with open('rec_subprocess_{0}'.format(ids),'w') as f:
            f.write('{0}\n'.format(optv))
            f.write('{0}\n'.format(simv))

def get_s(similarity_lib, scale, to_file=False, n_dev=10000, show=True, ids=None):
    if type(similarity_lib) == str:
        with open(similarity_lib,'r') as f:
            libm = []
            libs = []
            for l in f:
                m, sm = l.strip().split(':')
                libm.append(m)
                libs.append(sm.split(','))
    else:
        libm = [i[0] for i in similarity_lib]
        libs = [i[1] for i in similarity_lib]
    libmq = []
    libmfp = []
    print('cal ref QEDs')
    with tqdm(total=len(libm)) as ref_pbar:
        for i in libm:
            rmol = Chem.MolFromSmiles(i)
            libmfp.append(Chem.RDKFingerprint(rmol))
            libmq.append(Descriptors.qed(rmol))
            if show:
                ref_pbar.update()
    opt = []
    optv = []
    simv = []
    print('build pair')
    with tqdm(total=len(libm)) as bd_pbar:
        for midx in range(len(libm)):
            rfp = libmfp[midx]
            rq = libmq[midx]
            max_d = 0
            csmi = 'C1CCCCC1'
            sim_v = 0
            for cdd in libs[midx]:
                cmol = Chem.MolFromSmiles(cdd)
                cfp = Chem.RDKFingerprint(cmol)
                sim = DataStructs.TanimotoSimilarity(rfp, cfp)
                if sim<scale[1] and sim>scale[0]:
                    cq = Descriptors.qed(cmol)
                    diff = cq - rq
                    if diff > max_d:
                        csmi = cdd
                        max_d = diff
                        sim_v = sim
            if max_d > 0:
                opt.append([libm[midx], csmi])
                optv.append(max_d)
                simv.append(sim_v)
            if show:
                bd_pbar.update()
    if to_file:
        with open(to_file,'w') as f:
            for r in opt:
                f.write(','.join([str(i) for i in r])+'\n')

    print('split data')
    idx= list(range(len(opt)))
    shuffle(idx)
    if len(opt)<n_dev:
        n = len(str(len(opt)))-1
        kn = '1'+'0'*n
        kn = int(int(kn)/10)
    else:
        kn = n_dev
    train_idx = idx[:-kn]
    dev_idx = idx[-kn:]
    train_opt = [opt[i] for i in train_idx]
    dev_opt = [opt[i] for i in dev_idx]

    if ids == None:
        return train_opt, dev_opt, optv, simv
    else:
        with open('train_subprocess_{0}.tsv'.format(ids), 'w') as f:
            for r in train_opt:
                f.write('{0}\t{1}\n'.format(r[0], r[1]))
        with open('dev_subprocess_{0}.tsv'.format(ids), 'w') as f:
            for r in dev_opt:
                f.write('{0}\t{1}\n'.format(r[0], r[1]))
        optv = np.array(optv)
        simv = np.array(simv)
        np.save('simv_subprocess_{0}.npy'.format(ids), simv)
        np.save('optv_subprocess_{0}.npy'.format(ids), optv)

def multi(similarity_lib, n_jobs, scale, n_dev=10000):
    lib = []
    with open(similarity_lib,'r') as f:
        for l in f:
            m, sm = l.strip().split(':')
            lib.append([m, sm.split(',')])
    n_jobs = max(n_jobs, 1)
    recn_per_job = round(len(lib)/n_jobs)
    rec_lists = [lib[i*recn_per_job:(i+1)*recn_per_job] for i in range(n_jobs-1)]
    rec_lists.append(lib[(n_jobs-1)*recn_per_job:])
    n_dev = int(n_dev/n_jobs)

    sub_process = []
    for sp in range(n_jobs):
        sub_process.append(Process(target=get_s, args=(rec_lists[sp], scale, False, n_dev, False, sp)))
    for sp in sub_process:
        sp.start()
    for sp in sub_process:
        sp.join()
    # merge files and remove temporary files
    train_opt = []
    dev_opt = []
    simv = []
    optv = []

    for spf in range(n_jobs):
        with open('train_subprocess_{0}.tsv'.format(spf)) as f:
            train_sp = f.readlines()
            train_sp = [i.strip().split('\t') for i in train_sp]
        train_opt += train_sp
        with open('dev_subprocess_{0}.tsv'.format(spf)) as f:
            dev_sp = f.readlines()
            dev_sp = [i.strip().split('\t') for i in dev_sp]
        dev_opt += dev_sp
        simv_sp = np.load('simv_subprocess_{0}.npy'.format(spf))
        simv += list(simv_sp)
        optv_sp = np.load('optv_subprocess_{0}.npy'.format(spf))
        optv += list(optv_sp)
        subprocess.call('rm train_subprocess_{0}.tsv'.format(spf), shell=True)
        subprocess.call('rm dev_subprocess_{0}.tsv'.format(spf), shell=True)
        subprocess.call('rm simv_subprocess_{0}.npy'.format(spf), shell=True)
        subprocess.call('rm optv_subprocess_{0}.npy'.format(spf), shell=True)

    return train_opt, dev_opt, optv, simv


# main code that contains the neural network setup
# policy + critic updates
# see ddpg.py for other details in the network
import importlib
# from tensorboardX import SummaryWriter
from torch.optim import Adam

import ddpg
importlib.reload(ddpg)
import numpy as np
import random

# import model
# importlib.reload(model)
import model2 as model
importlib.reload(model)
# import model2Noisy as model
# importlib.reload(model)

import torch
import buffer
importlib.reload(buffer)
import torch.nn.functional as F

import utilities
importlib.reload(utilities)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# device = 'cpu'
# device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

BUFFER_SIZE = int(1e6)  # replay buffer size
BATCH_SIZE = 256#128         # minibatch size
GAMMA = 0.99            # discount factor
TAU = 1e-2#1e-3              # for soft update of target parameters
LR_ACTOR =  5e-4         # learning rate of the actor
MIN_LR_ACTOR = 1e-4
LR_CRITIC = 1e-3 #3e-4        # learning rate of the critic
WEIGHT_DECAY = 0#0.0001   # L2 weight decay
# log_path = os.getcwd()+"/log"
# logger = SummaryWriter(log_dir=log_path)

class MADDPG:
    def __init__(self, state_size, action_size, random_seed, warm_up=BATCH_SIZE,  lr_actor=LR_ACTOR, lr_critic=LR_CRITIC, num_agents=2,):
        super(MADDPG, self).__init__()
        # critic input = obs_full + actions = 14+2+2+2=20
        self.shared_critic_local =  model.Critic(state_size, action_size, random_seed).to(device)
        self.shared_critic_target = model.Critic(state_size, action_size, random_seed).to(device)
        self.shared_critic_optimizer = Adam(self.shared_critic_local.parameters(), lr=lr_critic, weight_decay=0)
        self.maddpg_agent = [ddpg.DDPGAgent(state_size, action_size, 12, warm_up, lr_actor,lr_critic, self.shared_critic_local, self.shared_critic_target, self.shared_critic_optimizer),
                             ddpg.DDPGAgent(state_size, action_size, 0, warm_up, lr_actor,lr_critic,self.shared_critic_local, self.shared_critic_target, self.shared_critic_optimizer)]

        self.discount_factor = GAMMA
        self.tau = TAU
        self.iter = 0
        self.num_agents = num_agents
        self.memory=buffer.ReplayBuffer(action_size, BUFFER_SIZE, random_seed)
#         self.memory_2=buffer.ReplayBuffer(action_size, BUFFER_SIZE, random_seed)

    def reset(self):
        for agent in self.maddpg_agent:
            agent.reset()

    def reduce_LR(self, reduce_factor=0.8, min_lr_actor=MIN_LR_ACTOR):
        for agent in self.maddpg_agent:
            agent.reduce_LR(reduce_factor,MIN_LR_ACTOR)

    def reduce_tau(self, reduce_factor=0.9, min_tau=1e-3):
        if self.tau > min_tau:
            self.tau = min(self.tau*reduce_factor, min_tau)
            print('!!! update self.tau', self.tau)

    def add_to_memory(self, states, actions, rewards, next_states, dones, i_episode):
        """Save experience in replay memory, and use random sample from buffer to learn."""
        # Save experience / reward
        obs_full = np.reshape(states, newshape=(1, -1))
        next_obs_full = np.reshape(next_states, newshape=(1, -1))
        self.memory.push(states, obs_full, actions, rewards, next_states,next_obs_full, dones)

    def learn_and_update(self, episode_i, episode_t,  update_iteration, start_update_episode=0,ts_per_update=1 , updates_per_ts=1, episodes_per_update=1):
        if episode_t % ts_per_update != 0:
            return update_iteration

        if len(self.memory) > BATCH_SIZE and episode_i % episodes_per_update ==0 and episode_i > start_update_episode:
            samples = self.memory.sample(BATCH_SIZE)
#             if random.random() < 0.01:
#                 print('samples', samples)
            if updates_per_ts < 1:
                if random.random() < updates_per_ts:
                    update_iteration += 1
                    for agent_i in range(self.num_agents):
#                         samples = self.memory.sample(BATCH_SIZE)
                        self.update(samples, agent_i)
                        self.update_targets(True)

            else:
                for updat4e_i in range(updates_per_ts):
                    update_iteration += 1
                    for agent_i in range(self.num_agents):
#                         samples = self.memory.sample(BATCH_SIZE)
                        self.update(samples, agent_i)
                        self.update_targets(True)
#             self.update_targets(update_critic=True)

        return update_iteration

    def get_actors(self):
        """get actors of all the agents in the MADDPG object"""
        actors = [ddpg_agent.actor for ddpg_agent in self.maddpg_agent]
        return actors

    def get_target_actors(self):
        """get target_actors of all the agents in the MADDPG object"""
        target_actors = [ddpg_agent.target_actor for ddpg_agent in self.maddpg_agent]
        return target_actors

    def act(self, obs_all_agents, noise=0.0):
        """
        get actions from all agents in the MADDPG object
        obs_all_agents : 2 * 24 nparray => tensor
        """
        actions = [agent.act(torch.tensor(obs, dtype=torch.float), noise) for agent, obs in zip(self.maddpg_agent, obs_all_agents)]
        return actions

    def target_act(self, obs_all_agents, noise=0.0):
        """get target network actions from all the agents in the MADDPG object """
        target_actions = [ddpg_agent.target_act(obs, noise) for ddpg_agent, obs in zip(self.maddpg_agent, obs_all_agents)]
        return target_actions

    def update(self, samples, agent_number):
        """update the critics and actors of all the agents """

        # need to transpose each element of the samples
        # to flip obs[parallel_agent][agent_number] to
        obs, obs_full, action, reward, next_obs, next_obs_full, done = samples
#         obs_full = torch.stack(obs_full)
#         next_obs_full = torch.stack(next_obs_full)
        agent = self.maddpg_agent[agent_number]
        agent.actor_target.eval()
        agent.critic_target.eval()
        #critic loss = batch mean of (y- Q(s,a) from target network)^2
        #y = reward of this timestep + discount * Q(st+1,at+1) from target network
        target_actions = self.target_act(next_obs) # a list of size 2, each torch inside Size([128, 2])
        target_actions = torch.cat(target_actions, dim=1)
        target_critic_input = torch.cat((next_obs_full.view(BATCH_SIZE, -1),target_actions.view(BATCH_SIZE, -1)), dim=1).to(device)

        with torch.no_grad():
            q_next = agent.critic_target(target_critic_input)

        y = reward[agent_number].view(-1, 1) + self.discount_factor * q_next * (1 - done[agent_number].view(-1, 1))
        action=action.t()
        critic_input = torch.cat((obs_full.view(BATCH_SIZE, -1), action.view(BATCH_SIZE, -1)), dim=1).to(device)
        agent.actor_target.train()
        agent.critic_target.train()
        q = agent.critic_local(critic_input)

        huber_loss = torch.nn.SmoothL1Loss()
        critic_loss = huber_loss(q, y.detach())

#         critic_loss = huber_loss(q, y)
#         critic_loss = F.mse_loss(q, y)
        agent.critic_optimizer.zero_grad()
        critic_loss.backward()
        torch.nn.utils.clip_grad_norm_(agent.critic_local.parameters(), 1.0)
        agent.critic_optimizer.step()

        #update actor network using policy gradient
        # make input to agent
        # detach the other agents to save computation
        # saves some time for computing derivative
        q_input = [ self.maddpg_agent[i].actor_local(ob) if i == agent_number \
                   else self.maddpg_agent[i].actor_local(ob).detach()
                   for i, ob in enumerate(obs) ]

        q_input = torch.cat(q_input, dim=1)
        # combine all the actions and observations for input to critic
        # many of the obs are redundant, and obs[1] contains all useful information already
        q_input2 = torch.cat((obs_full, q_input), dim=1)

        # get the policy gradient
        actor_loss = -agent.critic_local(q_input2).mean()
        agent.actor_optimizer.zero_grad()
        actor_loss.backward()
#         torch.nn.utils.clip_grad_norm_(agent.actor_local.parameters(), 0.5)
        agent.actor_optimizer.step()

#         al = actor_loss.cpu().detach().item()
#         cl = critic_loss.cpu().detach().item()
#         if random.random() < 0.008:
#             print('agent{} :  actor_loss {}'.format(agent_number, al))
#         logger.add_scalars('agent%i/losses' % agent_number,
#                            {'critic loss': cl,
#                             'actor_loss': al},
#                            self.iter)

    def update_targets(self, update_critic):
        """soft update targets"""
        self.iter += 1
        for ddpg_agent in self.maddpg_agent:
            utilities.soft_update(ddpg_agent.actor_target, ddpg_agent.actor_local, self.tau)
            if update_critic:
                utilities.soft_update(ddpg_agent.critic_target, ddpg_agent.critic_local, self.tau)


# Borrows extensively from https://github.com/tensorflow/tensorflow/blob/master/tensorflow/models/image/imagenet/classify_image.py

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os
import sys
import tarfile
import shutil
import tempfile
import subprocess
import re

import numpy as np
from scipy.misc import imread
from six.moves import urllib
import tensorflow as tf
import h5py

FLAGS = tf.app.flags.FLAGS
# classify_image_graph_def.pb:
#   Binary representation of the GraphDef protocol buffer.
# imagenet_synset_to_human_label_map.txt:
#   Map from synset ID to a human readable string.
# imagenet_2012_challenge_label_map_proto.pbtxt:
#   Text representation of a protocol buffer mapping a label to synset ID.
tf.app.flags.DEFINE_string(
    'model_dir', '/tmp/imagenet',
    """Path to classify_image_graph_def.pb, """
    """imagenet_synset_to_human_label_map.txt, and """
    """imagenet_2012_challenge_label_map_proto.pbtxt.""")
tf.app.flags.DEFINE_string('image_file', '',
                           """Absolute path to image file.""")
tf.app.flags.DEFINE_integer('num_top_predictions', 5,
                            """Display this many predictions.""")

# pylint: disable=line-too-long
DATA_URL = 'http://download.tensorflow.org/models/image/imagenet/inception-2015-12-05.tgz'
# pylint: enable=line-too-long

def create_graph():
  """Creates a graph from saved GraphDef file and returns a saver."""
  # Creates graph from saved graph_def.pb.
  with tf.gfile.FastGFile(os.path.join(
      FLAGS.model_dir, 'classify_image_graph_def.pb'), 'rb') as f:
    graph_def = tf.GraphDef()
    graph_def.ParseFromString(f.read())
    _ = tf.import_graph_def(graph_def, name='')

def run_inference_on_dataset(dataset, tmp_dir='/tmp/'):
  """Runs inference on an image.
  Args:
    dataset (DatasetPrep): Dataset
    tmp_dir (str): Directory to store images temporarily
  Returns:
    Nothing
  """
  # Creates graph from saved GraphDef.
  create_graph()
  image_keys = dataset.list_keys()
  features = np.zeros((len(image_keys), 2048), dtype=np.float16)
  config = tf.ConfigProto()
  config.gpu_options.allow_growth = True
  with tf.Session(config=config) as sess:
    for ind, img_record in enumerate(dataset):
        if ind % 1000 == 0:
            print ('Completed %d of %d'%(ind, len(image_keys)))

        new_fname = os.path.join(tmp_dir, os.path.basename(img_record.image_name))
        dataset.extract_image_to_location(img_record.id, new_fname)

        try:
            if not tf.gfile.Exists(new_fname):
                tf.logging.fatal('File does not exist %s', new_fname)
            image_data = tf.gfile.FastGFile(new_fname, 'rb').read()

            pool_3_tensor = sess.graph.get_tensor_by_name('pool_3:0')
            predictions = sess.run(pool_3_tensor,
                                   {'DecodeJpeg/contents:0': image_data})
        except: # Not a jpeg, use file to find extension, try to read with scipy
            try:
                filetype = subprocess.Popen(["file", new_fname], stdout=subprocess.PIPE).stdout.read()
                extension = re.search(r':[ ]+([A-Z]+) ', filetype).group(1).lower()
                new_new_fname = new_fname + '.{}'.format(extension)
                print('Renaming to {}'.format(new_new_fname))
                shutil.move(new_fname, new_new_fname)
                image = imread(new_new_fname) #Image.open(new_fname)
                image_data = np.array(image)[:, :, 0:3]  # Select RGB channels only.
                pool_3_tensor = sess.graph.get_tensor_by_name('pool_3:0')
                predictions = sess.run(pool_3_tensor,
                                       {'DecodeJpeg:0': image_data})
            except:
                filetype = subprocess.Popen(["file", new_new_fname], stdout=subprocess.PIPE).stdout.read()
                print('Expected PNG/JPEG, received:  {}'.format(filetype))
                print('Image data size: {}'.format(np.array(image).size()))
                raise

        features[ind, :] = np.squeeze(predictions)
        if os.path.exists(new_fname):
            os.remove(new_fname)

    return features

def save_hdf5(local_working_dir, hdf5_fname, image_features, image_ids):
    '''
    Create hdf5 file from features and filename list
    '''
    bname = os.path.basename(hdf5_fname)
    temp_fname = os.path.join(local_working_dir, bname)

    fOut = h5py.File(temp_fname, 'w')
    fOut.create_dataset('ids', data=image_ids)
    fOut.create_dataset('feats', data=image_features, dtype=np.float32)
    fOut.close()

    shutil.move(temp_fname, hdf5_fname)

def maybe_download_and_extract():
  """Download and extract model tar file."""
  dest_directory = FLAGS.model_dir
  if not os.path.exists(dest_directory):
    os.makedirs(dest_directory)
  filename = DATA_URL.split('/')[-1]
  filepath = os.path.join(dest_directory, filename)
  if not os.path.exists(filepath):
    def _progress(count, block_size, total_size):
      sys.stdout.write('\r>> Downloading %s %.1f%%' % (
          filename, float(count * block_size) / float(total_size) * 100.0))
      sys.stdout.flush()
    filepath, _ = urllib.request.urlretrieve(DATA_URL, filepath, _progress)
    print()
    statinfo = os.stat(filepath)
    print('Succesfully downloaded', filename, statinfo.st_size, 'bytes.')
  tarfile.open(filepath, 'r:gz').extractall(dest_directory)

def process_dataset(dataset_prep, output_fname, working_dir=tempfile.gettempdir()):
  """

  Args:
      dataset_prep (attalos.dataset.DatasetPrep): Dataset to convert
      output_fname: Output filename to extract to
      working_dir: Working directory to use for intermediate files

  Returns:

  """
  # Download Inception weights if not already present and extract for use
  maybe_download_and_extract()

  # Extract image features using inception network
  # TODO: Maybe this should batch in some way for large jobs?
  features = run_inference_on_dataset(dataset_prep)

  # Save computed features to file
  image_ids = [str(record.id) for record in dataset_prep]
  save_hdf5(working_dir, output_fname, features, image_ids)

def main(_):
  import argparse

  parser = argparse.ArgumentParser(description='Extract image features using Inception model.')
  parser.add_argument('--dataset_dir',
                      dest='dataset_dir',
                      type=str,
                      help='Directory with input images')
  parser.add_argument('--dataset_type',
                      dest='dataset_type',
                      default='mscoco',
                      choices=['mscoco', 'visualgenome', 'iaprtc', 'generic', 'espgame', 'nuswide'])
  parser.add_argument('--split',
                      dest='split',
                      default='train',
                      choices=['train', 'test', 'val'])
  parser.add_argument('--output_fname',
                      dest='output_fname',
                      default='image_features.hdf5',
                      type=str,
                      help='Output hd5f filename')
  parser.add_argument('--working_dir',
                      dest='working_dir',
                      default=tempfile.gettempdir(),
                      type=str,
                      help='Working directory for hdf5 file creation')
  args = parser.parse_args()

  if args.dataset_type == 'mscoco':
    print('Processing MSCOCO Data')
    from attalos.dataset.mscoco_prep import MSCOCODatasetPrep
    dataset_prep = MSCOCODatasetPrep(args.dataset_dir, split=args.split)
  elif args.dataset_type == 'visualgenome':
    print('Processing Visual Genome Data')
    from attalos.dataset.vg_prep import VGDatasetPrep
    dataset_prep = VGDatasetPrep(args.dataset_dir, split=args.split)
  elif args.dataset_type == 'iaprtc':
    print('Processing IAPRTC-12 data')
    from attalos.dataset.iaprtc12_prep import IAPRTC12DatasetPrep
    dataset_prep = IAPRTC12DatasetPrep(args.dataset_dir, split=args.split)
  elif args.dataset_type == 'generic':
    print('Processing espgame data')
    from attalos.dataset.generic_prep import GenericDatasetPrep
    dataset_prep = GenericDatasetPrep(args.dataset_dir, split=args.split)
  elif args.dataset_type == 'espgame':
    print('Processing espgame data')
    from attalos.dataset.espgame_prep import ESPGameDatasetPrep
    dataset_prep = ESPGameDatasetPrep(args.dataset_dir, split=args.split)
  elif args.dataset_type == 'nuswide':
    print('Processing nuswide data')
    from attalos.dataset.nuswide_prep import NUSWideDatasetPrep
    dataset_prep = NUSWideDatasetPrep(args.dataset_dir, split=args.split)
  else:
      raise NotImplementedError('Dataset type {} not supported'.format(args.dataset_type))
  process_dataset(dataset_prep, args.output_fname, working_dir=args.working_dir)

if __name__ == '__main__':
  tf.app.run()


import os
import typing

import moviepy.editor as mpy
import numpy as np
import requests
import torch
import torch.nn.functional
import torch.nn.functional
from hsh.library.hash import Hasher
from tqdm import tqdm

from .u2net import u2net

DEVICE = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

def iter_frames(path):
    return mpy.VideoFileClip(path).resize(height=320).iter_frames(dtype="uint8")

class Net(torch.nn.Module):
    def __init__(self, model_name):
        super(Net, self).__init__()
        hasher = Hasher()

        model, hash_val, drive_target, env_var = {
            'u2netp':          (u2net.U2NETP,
                                'e4f636406ca4e2af789941e7f139ee2e',
                                '1rbSTGKAE-MTxBYHd-51l2hMOQPT_7EPy',
                                'U2NET_PATH'),
            'u2net':           (u2net.U2NET,
                                '09fb4e49b7f785c9f855baf94916840a',
                                '1-Yg0cxgrNhHP-016FPdp902BR-kSsA4P',
                                'U2NET_PATH'),
            'u2net_human_seg': (u2net.U2NET,
                                '347c3d51b01528e5c6c071e3cff1cb55',
                                '1ao1ovG1Qtx4b7EoskHXmi2E9rp5CHLcZ',
                                'U2NET_PATH')
            }[model_name]
        path = os.environ.get(env_var, os.path.expanduser(os.path.join("~", ".u2net", model_name + ".pth")))
        net = model(3, 1)
        if not os.path.exists(path) or hasher.md5(path) != hash_val:
            head, tail = os.path.split(path)
            os.makedirs(head, exist_ok=True)

            URL = "https://docs.google.com/uc?export=download"

            session = requests.Session()
            response = session.get(URL, params={"id": drive_target}, stream=True)

            token = None
            for key, value in response.cookies.items():
                if key.startswith("download_warning"):
                    token = value
                    break

            if token:
                params = {"id": drive_target, "confirm": token}
                response = session.get(URL, params=params, stream=True)

            total = int(response.headers.get("content-length", 0))

            with open(path, "wb") as file, tqdm(
                desc=f"Downloading {tail} to {head}",
                total=total,
                unit="iB",
                unit_scale=True,
                unit_divisor=1024,
                ) as bar:
                for data in response.iter_content(chunk_size=1024):
                    size = file.write(data)
                    bar.update(size)
        net.load_state_dict(torch.load(path, map_location=torch.device(DEVICE)))
        net.to(device=DEVICE, dtype=torch.float32, non_blocking=True)
        net.eval()
        self.net = net

    def forward(self, block_input: torch.Tensor):
        image_data = block_input.permute(0, 3, 1, 2)
        original_shape = image_data.shape[2:]
        image_data = torch.nn.functional.interpolate(image_data, (320, 320), mode='bilinear')
        image_data = (image_data / 255 - 0.485) / 0.229
        out = self.net(image_data)[:, 0:1]
        ma = torch.max(out)
        mi = torch.min(out)
        out = (out - mi) / (ma - mi) * 255
        out = torch.nn.functional.interpolate(out, original_shape, mode='bilinear')
        out = out[:, 0]
        out = out.to(dtype=torch.uint8, device=torch.device('cpu'), non_blocking=True).detach()
        return out

@torch.no_grad()
def remove_many(image_data: typing.List[np.array], net: Net):
    image_data = np.stack(image_data)
    image_data = torch.as_tensor(image_data, dtype=torch.float32, device=DEVICE)
    return net(image_data).numpy()


from .util import cos_sim, dot_score, normalize, save_dict_to_tsv, load_tsv_to_dict
from .faiss_index import FaissBinaryIndex, FaissPQIndex, FaissHNSWIndex, FaissPCAIndex, FaissIndex
import logging
import sys
import torch
import faiss
import numpy as np
import os
from typing import Dict, List

logger = logging.getLogger(__name__)

#Parent class for any faiss search
class DenseRetrievalFaissSearch:

    def __init__(self, model, batch_size: int = 128, corpus_chunk_size: int = 50000, **kwargs):
        self.model = model
        self.batch_size = batch_size
        self.corpus_chunk_size = corpus_chunk_size
        self.score_functions = ['cos_sim','dot']
        self.mapping_tsv_keys = ["beir-docid", "faiss-docid"]
        self.faiss_index = None
        self.dim_size = 0
        self.results = {}
        self.mapping = {}
        self.rev_mapping = {}

    def _create_mapping_ids(self, corpus_ids):
        if not all(isinstance(doc_id, int) for doc_id in corpus_ids):
            for idx in range(len(corpus_ids)):
                self.mapping[corpus_ids[idx]] = idx
                self.rev_mapping[idx] = corpus_ids[idx]

    def _load(self, input_dir: str, prefix: str, ext: str):

        # Load ID mappings from file
        input_mappings_path = os.path.join(input_dir, "{}.{}.tsv".format(prefix, ext))
        logger.info("Loading Faiss ID-mappings from path: {}".format(input_mappings_path))
        self.mapping = load_tsv_to_dict(input_mappings_path, header=True)
        self.rev_mapping = {v: k for k, v in self.mapping.items()}
        passage_ids = sorted(list(self.rev_mapping))

        # Load Faiss Index from disk
        input_faiss_path = os.path.join(input_dir, "{}.{}.faiss".format(prefix, ext))
        logger.info("Loading Faiss Index from path: {}".format(input_faiss_path))

        return input_faiss_path, passage_ids

    def save(self, output_dir: str, prefix: str, ext: str):

        # Save BEIR -> Faiss ids mappings
        save_mappings_path = os.path.join(output_dir, "{}.{}.tsv".format(prefix, ext))
        logger.info("Saving Faiss ID-mappings to path: {}".format(save_mappings_path))
        save_dict_to_tsv(self.mapping, save_mappings_path, keys=self.mapping_tsv_keys)

        # Save Faiss Index to disk
        save_faiss_path = os.path.join(output_dir, "{}.{}.faiss".format(prefix, ext))
        logger.info("Saving Faiss Index to path: {}".format(save_faiss_path))
        self.faiss_index.save(save_faiss_path)
        logger.info("Index size: {:.2f}MB".format(os.path.getsize(save_faiss_path)*0.000001))

    def _index(self, corpus: Dict[str, Dict[str, str]], score_function: str = None):

        logger.info("Encoding Corpus in batches... Warning: This might take a while!")
        corpus_ids = list(corpus.keys())
        self._create_mapping_ids(corpus_ids)
        corpus = [corpus[cid] for cid in corpus_ids]
        normalize_embeddings = True if score_function == "cos_sim" else False

        itr = range(0, len(corpus), self.corpus_chunk_size)

        for batch_num, corpus_start_idx in enumerate(itr):
            logger.info("Encoding Batch {}/{}...".format(batch_num+1, len(itr)))
            corpus_end_idx = min(corpus_start_idx + self.corpus_chunk_size, len(corpus))

            #Encode chunk of corpus
            sub_corpus_embeddings = self.model.encode_corpus(
                corpus[corpus_start_idx:corpus_end_idx],
                batch_size=self.batch_size,
                show_progress_bar=True,
                normalize_embeddings=normalize_embeddings)

            if not batch_num:
                corpus_embeddings = sub_corpus_embeddings
            else:
                corpus_embeddings = np.vstack([corpus_embeddings, sub_corpus_embeddings])

        #Index chunk of corpus into faiss index
        logger.info("Indexing Passages into Faiss...")

        faiss_ids = [self.mapping.get(corpus_id) for corpus_id in corpus_ids]
        self.dim_size = corpus_embeddings.shape[1]

        del sub_corpus_embeddings

        return faiss_ids, corpus_embeddings

    def search(self,
               corpus: Dict[str, Dict[str, str]],
               queries: Dict[str, str],
               top_k: int,
               score_function = str, **kwargs) -> Dict[str, Dict[str, float]]:

        assert score_function in self.score_functions

        if not self.faiss_index: self.index(corpus, score_function)

        logger.info("Encoding Queries...")
        query_ids = list(queries.keys())
        queries = [queries[qid] for qid in queries]
        query_embeddings = self.model.encode_queries(
            queries, show_progress_bar=True, batch_size=self.batch_size)

        faiss_scores, faiss_doc_ids = self.faiss_index.search(query_embeddings, top_k, **kwargs)

        for idx in range(len(query_ids)):
            scores = [float(score) for score in faiss_scores[idx]]
            if len(self.rev_mapping) != 0:
                doc_ids = [self.rev_mapping[doc_id] for doc_id in faiss_doc_ids[idx]]
            else:
                doc_ids = [str(doc_id) for doc_id in faiss_doc_ids[idx]]
            self.results[query_ids[idx]] = dict(zip(doc_ids, scores))

        return self.results

class BinaryFaissSearch(DenseRetrievalFaissSearch):

    def load(self, input_dir: str, prefix: str = "my-index", ext: str = "bin"):
        input_faiss_path, passage_ids = super()._load(input_dir, prefix, ext)
        base_index = faiss.read_index_binary(input_faiss_path)
        passage_embeddings = base_index.reconstruct(0)
        for idx in range(1, len(passage_ids)):
            passage_embeddings = np.vstack((passage_embeddings, base_index.reconstruct(idx)))
        self.faiss_index = FaissBinaryIndex(base_index, passage_ids, passage_embeddings)

    def index(self, corpus: Dict[str, Dict[str, str]], score_function: str = None):
        faiss_ids, corpus_embeddings = super()._index(corpus, score_function)
        logger.info("Using Binary Hashing in Flat Mode!")
        base_index = faiss.IndexBinaryFlat(self.dim_size * 8)
        self.faiss_index = FaissBinaryIndex.build(faiss_ids, corpus_embeddings, base_index)

    def save(self, output_dir: str, prefix: str = "my-index", ext: str = "bin"):
        super().save(output_dir, prefix, ext)

    def search(self,
            corpus: Dict[str, Dict[str, str]],
            queries: Dict[str, str],
            top_k: int,
            score_function = str, **kwargs) -> Dict[str, Dict[str, float]]:

        return super().search(corpus, queries, top_k, score_function, **kwargs)

class PQFaissSearch(DenseRetrievalFaissSearch):
    def __init__(self, model, batch_size: int = 128, corpus_chunk_size: int = 50000, num_of_centroids: int = 96,
                 code_size: int = 8, similarity_metric=faiss.METRIC_INNER_PRODUCT, **kwargs):
        super(PQFaissSearch, self).__init__(model, batch_size, corpus_chunk_size, **kwargs)
        self.num_of_centroids = num_of_centroids
        self.code_size = code_size
        self.similarity_metric = similarity_metric

    def load(self, input_dir: str, prefix: str = "my-index", ext: str = "pq"):
        input_faiss_path, passage_ids = super()._load(input_dir, prefix, ext)
        base_index = faiss.read_index(input_faiss_path)
        self.faiss_index = FaissPQIndex(base_index, passage_ids)

    def index(self, corpus: Dict[str, Dict[str, str]], score_function: str = None, **kwargs):
        faiss_ids, corpus_embeddings = super()._index(corpus, score_function, **kwargs)
        logger.info("Using Product Quantization (PQ) in Flat mode!")
        logger.info("Parameters Used: num_of_centroids: {} ".format(self.num_of_centroids))
        logger.info("Parameters Used: code_size: {}".format(self.code_size))

        base_index = faiss.IndexPQ(self.dim_size, self.num_of_centroids, self.code_size, self.similarity_metric)
        self.faiss_index = FaissPQIndex.build(faiss_ids, corpus_embeddings, base_index)

    def save(self, output_dir: str, prefix: str = "my-index", ext: str = "pq"):
        super().save(output_dir, prefix, ext)

    def search(self,
            corpus: Dict[str, Dict[str, str]],
            queries: Dict[str, str],
            top_k: int,
            score_function = str, **kwargs) -> Dict[str, Dict[str, float]]:

        return super().search(corpus, queries, top_k, score_function, **kwargs)

class HNSWFaissSearch(DenseRetrievalFaissSearch):
    def __init__(self, model, batch_size: int = 128, corpus_chunk_size: int = 50000, hnsw_store_n: int = 512,
                 hnsw_ef_search: int = 128, hnsw_ef_construction: int = 200, similarity_metric=faiss.METRIC_INNER_PRODUCT, **kwargs):
        super(HNSWFaissSearch, self).__init__(model, batch_size, corpus_chunk_size, **kwargs)
        self.hnsw_store_n = hnsw_store_n
        self.hnsw_ef_search = hnsw_ef_search
        self.hnsw_ef_construction = hnsw_ef_construction
        self.similarity_metric = similarity_metric

    def load(self, input_dir: str, prefix: str = "my-index", ext: str = "hnsw"):
        input_faiss_path, passage_ids = super()._load(input_dir, prefix, ext)
        base_index = faiss.read_index(input_faiss_path)
        self.faiss_index = FaissHNSWIndex(base_index, passage_ids)

    def index(self, corpus: Dict[str, Dict[str, str]], score_function: str = None, **kwargs):
        faiss_ids, corpus_embeddings = super()._index(corpus, score_function, **kwargs)

        logger.info("Using Approximate Nearest Neighbours (HNSW) in Flat Mode!")
        logger.info("Parameters Required: hnsw_store_n: {}".format(self.hnsw_store_n))
        logger.info("Parameters Required: hnsw_ef_search: {}".format(self.hnsw_ef_search))
        logger.info("Parameters Required: hnsw_ef_construction: {}".format(self.hnsw_ef_construction))

        base_index = faiss.IndexHNSWFlat(self.dim_size + 1, self.hnsw_store_n, self.similarity_metric)
        base_index.hnsw.efSearch = self.hnsw_ef_search
        base_index.hnsw.efConstruction = self.hnsw_ef_construction
        self.faiss_index = FaissHNSWIndex.build(faiss_ids, corpus_embeddings, base_index)

    def save(self, output_dir: str, prefix: str = "my-index", ext: str = "hnsw"):
        super().save(output_dir, prefix, ext)

    def search(self,
            corpus: Dict[str, Dict[str, str]],
            queries: Dict[str, str],
            top_k: int,
            score_function = str, **kwargs) -> Dict[str, Dict[str, float]]:

        return super().search(corpus, queries, top_k, score_function, **kwargs)

class FlatIPFaissSearch(DenseRetrievalFaissSearch):
    def load(self, input_dir: str, prefix: str = "my-index", ext: str = "flat"):
        input_faiss_path, passage_ids = super()._load(input_dir, prefix, ext)
        base_index = faiss.read_index(input_faiss_path)
        self.faiss_index = FaissIndex(base_index, passage_ids)

    def index(self, corpus: Dict[str, Dict[str, str]], score_function: str = None, **kwargs):
        faiss_ids, corpus_embeddings = super()._index(corpus, score_function, **kwargs)
        base_index = faiss.IndexFlatIP(self.dim_size)
        self.faiss_index = FaissIndex.build(faiss_ids, corpus_embeddings, base_index)

    def save(self, output_dir: str, prefix: str = "my-index", ext: str = "flat"):
        super().save(output_dir, prefix, ext)

    def search(self,
            corpus: Dict[str, Dict[str, str]],
            queries: Dict[str, str],
            top_k: int,
            score_function = str, **kwargs) -> Dict[str, Dict[str, float]]:

        return super().search(corpus, queries, top_k, score_function, **kwargs)

class PCAFaissSearch(DenseRetrievalFaissSearch):
    def __init__(self, model, base_index: faiss.Index, output_dimension: int, batch_size: int = 128,
                corpus_chunk_size: int = 50000, **kwargs):
        super(PCAFaissSearch, self).__init__(model, batch_size, corpus_chunk_size, **kwargs)
        self.base_index = base_index
        self.output_dim = output_dimension

    def load(self, input_dir: str, prefix: str = "my-index", ext: str = "pca"):
        input_faiss_path, passage_ids = super()._load(input_dir, prefix, ext)
        base_index = faiss.read_index(input_faiss_path)
        self.faiss_index = FaissPCAIndex(base_index, passage_ids)

    def index(self, corpus: Dict[str, Dict[str, str]], score_function: str = None, **kwargs):
        faiss_ids, corpus_embeddings = super()._index(corpus, score_function, **kwargs)
        logger.info("Creating PCA Matrix...")
        logger.info("Input Dimension: {}, Output Dimension: {}".format(self.dim_size, self.output_dim))
        pca_matrix = faiss.PCAMatrix(self.dim_size, self.output_dim, 0, True)
        final_index = faiss.IndexPreTransform(pca_matrix, self.base_index)
        self.faiss_index = FaissPCAIndex.build(faiss_ids, corpus_embeddings, final_index)

    def save(self, output_dir: str, prefix: str = "my-index", ext: str = "pca"):
        super().save(output_dir, prefix, ext)

    def search(self,
            corpus: Dict[str, Dict[str, str]],
            queries: Dict[str, str],
            top_k: int,
            score_function = str, **kwargs) -> Dict[str, Dict[str, float]]:

        return super().search(corpus, queries, top_k, score_function, **kwargs)