# -*- coding: utf-8 -*-

# Copyright (C) 2020. Huawei Technologies Co., Ltd. All rights reserved.
# This program is free software; you can redistribute it and/or modify
# it under the terms of the MIT License.
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# MIT License for more details.

"""Import all torch operators."""
import math
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from vega.search_space.networks.net_utils import NetTypes
from torch.autograd import Function
from vega.search_space.networks.network_factory import NetworkFactory

def import_all_operators():
    """Import all torch operators."""
    for _name in dir(nn):
        if not _name.startswith("_"):
            _cls = getattr(nn, _name)
            NetworkFactory.register_custom_cls(NetTypes.Operator, _cls)

import_all_operators()

@NetworkFactory.register(NetTypes.Operator)
class Input(nn.Module):
    """Create Input for forward x."""

    def __init__(self, key=None, data=None, cuda=False):
        self.key = key
        self.cuda = cuda
        self.data = data
        super(Input, self).__init__()

    def forward(self, x):
        """Forward x."""
        if self.data:
            return self.data
        if self.key is not None:
            x = x[self.key]
        if self.cuda:
            if torch.is_tensor(x):
                return x.cuda()
            else:
                if isinstance(x, dict):
                    x = {key: value.cuda() if torch.is_tensor(x) else value for key, value in x.items()}
                else:
                    x = [item.cuda() if torch.is_tensor(x) else item for item in x]
        return x

@NetworkFactory.register(NetTypes.Operator)
class Lambda(nn.Module):
    """Create Lambda for forward x."""

    def __init__(self, func, data=None):
        self.func = func
        self.data = data
        super(Lambda, self).__init__()

    def forward(self, x):
        """Forward x."""
        if self.data:
            return self.func(self.data)
        out = self.func(x)
        return out

@NetworkFactory.register(NetTypes.Operator)
class Reshape(nn.Module):
    """Create Lambda for forward x."""

    def __init__(self, *args):
        self.args = args
        super(Reshape, self).__init__()

    def forward(self, x):
        """Forward x."""
        return x.reshape(*self.args)

@NetworkFactory.register(NetTypes.Operator)
class Rermute(nn.Module):
    """Create Lambda for forward x."""

    def __init__(self, *args):
        self.args = args
        super(Rermute, self).__init__()

    def forward(self, x):
        """Forward x."""
        return x.rermute(*self.args)

@NetworkFactory.register(NetTypes.Operator)
def conv3x3(inchannel, outchannel, groups=1, stride=1):
    """Create conv3x3 layer.

    :param inchannel: input channel.
    :type inchannel: int
    :param outchannel: output channel.
    :type outchannel: int
    :param stride: the number to jump, default 1
    :type stride: int
    """
    return nn.Conv2d(inchannel, outchannel, kernel_size=3, stride=stride,
                     padding=1, groups=groups, bias=False)

@NetworkFactory.register(NetTypes.Operator)
def conv1X1(inchannel, outchannel, stride=1):
    """Create conv1X1 layer.

    :param inchannel: input channel.
    :type inchannel: int
    :param outchannel: output channel.
    :type outchannel: int
    :param stride: the number to jump, default 1
    :type stride: int
    """
    return nn.Conv2d(inchannel, outchannel, kernel_size=1, stride=stride, bias=False)

@NetworkFactory.register(NetTypes.Operator)
def conv5x5(inchannel, outchannel, stride=1, bias=False, dilation=1):
    """Create Convolution 5x5.

    :param in_planes: input channels
    :param out_planes: output channels
    :param stride: stride of the convolution
    :param bias: whether bias is contained
    :param dilation: dilation of the convolution
    :return: a convolution module
    """
    return nn.Conv2d(inchannel, outchannel, kernel_size=5, stride=stride,
                     padding=2, dilation=dilation, bias=bias)

@NetworkFactory.register(NetTypes.Operator)
def conv7x7(inchannel, outchannel, stride=1, bias=False, dilation=1):
    """Create Convolution 7x7.

    :param in_planes: input channels
    :param out_planes: output channels
    :param stride: stride of the convolution
    :param bias: whether bias is contained
    :param dilation: dilation of the convolution
    :return: a convolution module
    """
    return nn.Conv2d(inchannel, outchannel, kernel_size=7, stride=stride,
                     padding=3, dilation=dilation, bias=bias)

@NetworkFactory.register(NetTypes.Operator)
def conv_bn_relu6(inchannel, outchannel, kernel=3, stride=1):
    """Create conv1X1 layer.

    :param inchannel: input channel.
    :type inchannel: int
    :param outchannel: output channel.
    :type outchannel: int
    :param stride: the number to jump, default 1
    :type stride: int
    """
    return nn.Sequential(
        nn.Conv2d(inchannel, outchannel, kernel, stride, kernel // 2, bias=False),
        nn.BatchNorm2d(outchannel),
        nn.ReLU6(inplace=True)
    )

@NetworkFactory.register(NetTypes.Operator)
def conv_bn_relu(inchannel, outchannel, kernel_size, stride, padding, affine=True):
    """Create group of Convolution + BN + Relu.

    :param C_in: input channel
    :param C_out: output channel
    :param kernel_size: kernel size of convolution layer
    :param stride: stride of convolution layer
    :param padding: padding of convolution layer
    :param affine: whether use affine in batchnorm
    :return: group of Convolution + BN + Relu
    """
    return nn.Sequential(
        nn.Conv2d(inchannel, outchannel, kernel_size, stride=stride, padding=padding,
                  bias=False),
        nn.BatchNorm2d(outchannel, affine=affine),
        nn.ReLU(inplace=False),
    )

@NetworkFactory.register(NetTypes.Operator)
class View(nn.Module):
    """Call torch.view."""

    def __init__(self):
        super(View, self).__init__()

    def forward(self, x):
        """Forward x."""
        return x.view(x.size(0), -1)

@NetworkFactory.register(NetTypes.Operator)
class ConvWS2d(nn.Conv2d):
    """Conv2d with weight standarlization.

    :param in_channels: input channels
    :param out_channels: output channels
    :param kernel_size: kernel size
    :param stride: stride
    :param padding: num of padding
    :param dilation: num of dilation
    :param groups: num of groups
    :param bias: bias
    :param eps: eps
    """

    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0,
                 dilation=1, groups=1, bias=True, eps=1e-5):
        """Init conv2d with weight standarlization."""
        super(ConvWS2d, self).__init__(
            in_channels,
            out_channels,
            kernel_size,
            stride=stride,
            padding=padding,
            dilation=dilation,
            groups=groups,
            bias=bias)
        self.eps = eps

    def conv_ws_2d(self, input, weight, bias=None, stride=1, padding=0,
                   dilation=1, groups=1, eps=1e-5):
        """Conv2d with weight standarlization.

        :param input: input feature map
        :type input: torch.Tensor
        :param weight: weight of conv layer
        :type weight: torch.Tensor
        :param bias: bias
        :type bias: torch.Tensor
        :param stride: conv stride
        :type stride: int
        :param padding: num of padding
        :type padding: int
        :param dilation: num of dilation
        :type dilation: int
        :param groups: num of group
        :type groups: int
        :param eps: weight eps
        :type eps: float
        """
        c_in = weight.size(0)
        weight_flat = weight.view(c_in, -1)
        mean = weight_flat.mean(dim=1, keepdim=True).view(c_in, 1, 1, 1)
        std = weight_flat.std(dim=1, keepdim=True).view(c_in, 1, 1, 1)
        weight = (weight - mean) / (std + eps)
        return F.conv2d(input, weight, bias, stride, padding, dilation, groups)

    def forward(self, x):
        """Forward function of conv2d with weight standarlization."""
        return self.conv_ws_2d(x, self.weight, self.bias, self.stride, self.padding,
                               self.dilation, self.groups, self.eps)

@NetworkFactory.register(NetTypes.Operator)
class ChannelShuffle(nn.Module):
    """Shuffle the channel of features.

    :param groups: group number of channels
    :type groups: int
    """

    def __init__(self, groups):
        super(ChannelShuffle, self).__init__()
        self.groups = groups

    def forward(self, x):
        """Forward x."""
        batchsize, num_channels, height, width = x.data.size()
        channels_per_group = num_channels // self.groups
        x = x.view(batchsize, self.groups, channels_per_group, height, width)
        x = torch.transpose(x, 1, 2).contiguous()
        x = x.view(batchsize, -1, height, width)
        return x

@NetworkFactory.register(NetTypes.Operator)
class Shrink_Conv(nn.Module):
    """Call torch.cat.

    :param InChannel: channel number of input
    :type InChannel: int
    :param OutChannel: channel number of output
    :type OutChannel: int
    :param growRate: growth rate of block
    :type growRate: int
    :param nConvLayers: the number of convlution layer
    :type nConvLayers: int
    :param kSize: kernel size of convolution operation
    :type kSize: int
    """

    def __init__(self, InChannel, OutChannel, growRate, nConvLayers, kSize):
        super(Shrink_Conv, self).__init__()
        self.InChan = InChannel
        self.OutChan = OutChannel
        self.G = growRate
        self.C = nConvLayers
        if self.InChan != self.G:
            self.InConv = nn.Conv2d(self.InChan, self.G, 1, padding=0, stride=1)
        if self.OutChan != self.G and self.OutChan != self.InChan:
            self.OutConv = nn.Conv2d(self.InChan, self.OutChan, 1, padding=0,
                                     stride=1)
        self.Convs = nn.ModuleList()
        self.ShrinkConv = nn.ModuleList()
        for i in range(self.C):
            self.Convs.append(nn.Sequential(*[
                nn.Conv2d(self.G, self.G, kSize, padding=(kSize - 1) // 2,
                          stride=1), nn.ReLU()]))
            if i == (self.C - 1):
                self.ShrinkConv.append(
                    nn.Conv2d((2 + i) * self.G, self.OutChan, 1, padding=0,
                              stride=1))
            else:
                self.ShrinkConv.append(
                    nn.Conv2d((2 + i) * self.G, self.G, 1, padding=0, stride=1))

    def forward(self, x):
        """Forward x."""
        if self.InChan != self.G:
            x_InC = self.InConv(x)
            x_inter = self.Convs[0](x_InC)
            x_conc = torch.cat((x_InC, x_inter), 1)
            x_in = self.ShrinkConv[0](x_conc)
        else:
            x_inter = self.Convs[0](x)
            x_conc = torch.cat((x, x_inter), 1)
            x_in = self.ShrinkConv[0](x_conc)
        for i in range(1, self.C):
            x_inter = self.Convs[i](x_in)
            x_conc = torch.cat((x_conc, x_inter), 1)
            x_in = self.ShrinkConv[i](x_conc)
        return x_in

@NetworkFactory.register(NetTypes.Operator)
class Cont_Conv(nn.Module):
    """Call torch.cat.

    :param InChannel: channel number of input
    :type InChannel: int
    :param OutChannel: channel number of output
    :type OutChannel: int
    :param growRate: growth rate of block
    :type growRate: int
    :param nConvLayers: the number of convlution layer
    :type nConvLayers: int
    :param kSize: kernel size of convolution operation
    :type kSize: int
    """

    def __init__(self, InChannel, OutChannel, growRate, nConvLayers, kSize=3):
        super(Cont_Conv, self).__init__()
        self.InChan = InChannel
        self.OutChan = OutChannel
        self.G = growRate
        self.C = nConvLayers
        self.shup = nn.PixelShuffle(2)
        self.Convs = nn.ModuleList()
        self.ShrinkConv = nn.ModuleList()
        for i in range(self.C):
            self.Convs.append(nn.Sequential(*[
                nn.Conv2d(self.G, self.G, kSize, padding=(kSize - 1) // 2,
                          stride=1), nn.ReLU()]))
            if i < (self.C - 1):
                self.ShrinkConv.append(
                    nn.Conv2d((2 + i) * self.G, self.G, 1, padding=0, stride=1))
            else:
                self.ShrinkConv.append(
                    nn.Conv2d(int((2 + i) * self.G / 4), self.OutChan, 1,
                              padding=0, stride=1))

    def forward(self, x):
        """Forward x."""
        x_conc = x
        for i in range(0, self.C):
            x_inter = self.Convs[i](x)
            x_inter = self.Convs[i](x_inter)
            x_inter = self.Convs[i](x_inter)
            x_conc = torch.cat((x_conc, x_inter), 1)
            if i == (self.C - 1):
                x_conc = self.shup(x_conc)
                x_in = self.ShrinkConv[i](x_conc)
            else:
                x_in = self.ShrinkConv[i](x_conc)
        return x_in

@NetworkFactory.register(NetTypes.Operator)
class Esrn_Cat(nn.Module):
    """Call torch.cat."""

    def __init__(self):
        super(Esrn_Cat, self).__init__()

    def forward(self, x):
        """Forward x."""
        return torch.cat(list(x), 1)

@NetworkFactory.register(NetTypes.Operator)
class MicroDecoder_Upsample(nn.Module):
    """Call torch.Upsample."""

    def __init__(self, collect_inds, agg_concat):
        self.collect_inds = collect_inds
        self.agg_concat = agg_concat
        super(MicroDecoder_Upsample, self).__init__()

    def forward(self, x):
        """Forward x."""
        out = x[self.collect_inds[0]]
        for i in range(1, len(self.collect_inds)):
            collect = x[self.collect_inds[i]]
            if out.size()[2] > collect.size()[2]:
                # upsample collect
                collect = nn.Upsample(size=out.size()[2:], mode='bilinear', align_corners=True)(collect)
            elif collect.size()[2] > out.size()[2]:
                out = nn.Upsample(size=collect.size()[2:], mode='bilinear', align_corners=True)(out)
            if self.agg_concat:
                out = torch.cat([out, collect], 1)
            else:
                out += collect
        out = F.relu(out)
        return out

@NetworkFactory.register(NetTypes.Operator)
class ContextualCell_v1(nn.Module):
    """New contextual cell design."""

    def __init__(self, op_names, config, inp, repeats=1, concat=False):
        """Construct ContextualCell_v1 class.

        :param op_names: list of operation indices
        :param config: list of config numbers
        :param inp: input channel
        :param repeats: number of repeated times
        :param concat: concat the result if set to True, otherwise add the result
        """
        super(ContextualCell_v1, self).__init__()
        self._ops = nn.ModuleList()
        self._pos = []
        self._collect_inds = [0]
        self._pools = ['x']
        for ind, op in enumerate(config):
            # first op is always applied on x
            if ind == 0:
                pos = 0
                op_id = op
                self._collect_inds.remove(pos)
                op_name = op_names[op_id]
                self._ops.append(OPS[op_name](inp, 1, True, repeats))  # turn-off scaling in batch norm
                self._pos.append(pos)
                self._collect_inds.append(ind + 1)
                self._pools.append('{}({})'.format(op_name, self._pools[pos]))
            else:
                pos1, pos2, op_id1, op_id2 = op
                # drop op_id from loose ends
                for ind2, (pos, op_id) in enumerate(zip([pos1, pos2], [op_id1, op_id2])):
                    if pos in self._collect_inds:
                        self._collect_inds.remove(pos)
                    op_name = op_names[op_id]
                    self._ops.append(OPS[op_name](inp, 1, True, repeats))  # turn-off scaling in batch norm
                    self._pos.append(pos)
                    # self._collect_inds.append(ind * 3 + ind2 - 1) # Do not collect intermediate
                    self._pools.append('{}({})'.format(op_name, self._pools[pos]))
                # summation
                op_name = 'sum'
                self._ops.append(AggregateCell(size_1=None, size_2=None, agg_size=inp, pre_transform=False,
                                               concat=concat))  # turn-off convbnrelu
                self._pos.append([ind * 3 - 1, ind * 3])
                self._collect_inds.append(ind * 3 + 1)
                self._pools.append('{}({},{})'.format(op_name, self._pools[ind * 3 - 1], self._pools[ind * 3]))

    def forward(self, x):
        """Do an inference on ContextualCell_v1.

        :param x: input tensor
        :return: output tensor
        """
        feats = [x]
        for pos, op in zip(self._pos, self._ops):
            if isinstance(pos, list):
                assert len(pos) == 2, "Two ops must be provided"
                feats.append(op(feats[pos[0]], feats[pos[1]]))
            else:
                feats.append(op(feats[pos]))
        out = 0
        for i in self._collect_inds:
            out += feats[i]
        return out

@NetworkFactory.register(NetTypes.Operator)
class AggregateCell(nn.Module):
    """Aggregate two cells and sum or concat them up."""

    def __init__(self, size_1, size_2, agg_size, pre_transform=True, concat=False):
        """Construct AggregateCell.

        :param size_1: channel of first input
        :param size_2: channel of second input
        :param agg_size: channel of aggregated tensor
        :param pre_transform: whether to do a transform on two inputs
        :param concat: concat the result if set to True, otherwise add the result
        """
        super(AggregateCell, self).__init__()
        self.pre_transform = pre_transform
        self.concat = concat
        if self.pre_transform:
            self.branch_1 = conv_bn_relu(size_1, agg_size, 1, 1, 0)
            self.branch_2 = conv_bn_relu(size_2, agg_size, 1, 1, 0)
        if self.concat:
            self.conv1x1 = conv_bn_relu(agg_size * 2, agg_size, 1, 1, 0)

    def forward(self, x1, x2):
        """Do an inference on AggregateCell.

        :param x1: first input
        :param x2: second input
        :return: output
        """
        if self.pre_transform:
            x1 = self.branch_1(x1)
            x2 = self.branch_2(x2)
        if x1.size()[2:] > x2.size()[2:]:
            x2 = nn.Upsample(size=x1.size()[2:], mode='bilinear', align_corners=True)(x2)
        elif x1.size()[2:] < x2.size()[2:]:
            x1 = nn.Upsample(size=x2.size()[2:], mode='bilinear', align_corners=True)(x1)
        if self.concat:
            return self.conv1x1(torch.cat([x1, x2], 1))
        else:
            return x1 + x2

@NetworkFactory.register(NetTypes.Operator)
class GAPConv1x1(nn.Module):
    """Global Average Pooling + conv1x1."""

    def __init__(self, C_in, C_out):
        """Construct GAPConv1x1 class.

        :param C_in: input channel
        :param C_out: output channel
        """
        super(GAPConv1x1, self).__init__()
        if '0.2' in torch.__version__:
            # !!!!!!!!!!used for input size 448 with overall stride 32!!!!!!!!!!
            self.globalpool = nn.AvgPool2d(14)
        self.conv1x1 = conv_bn_relu(C_in, C_out, 1, stride=1, padding=0)

    def forward(self, x):
        """Do an inference on GAPConv1x1.

        :param x: input tensor
        :return: output tensor
        """
        size = x.size()[2:]
        if '0.2' in torch.__version__:
            out = self.globalpool(x)
        else:
            out = x.mean(2, keepdim=True).mean(3, keepdim=True)
        out = self.conv1x1(out)
        if '0.2' in torch.__version__:
            out = nn.Upsample(size=size, mode='bilinear')(out)
        else:
            out = nn.functional.interpolate(out, size=size, mode='bilinear', align_corners=False)
        return out

@NetworkFactory.register(NetTypes.Operator)
class DilConv(nn.Module):
    """Separable dilated convolution block."""

    def __init__(self, C_in, C_out, kernel_size, stride, padding,
                 dilation, affine=True):
        """Construct DilConv class.

        :param C_in: input channel
        :param C_out: output channel
        :param kernel_size: kernel size of the first convolution layer
        :param stride: stride of the first convolution layer
        :param padding: padding of the first convolution layer
        :param dilation: dilation of the first convolution layer
        :param affine: whether use affine in BN
        """
        super(DilConv, self).__init__()
        self.op = nn.Sequential(
            nn.ReLU(inplace=False),
            nn.Conv2d(C_in, C_in, kernel_size=kernel_size, stride=stride,
                      padding=padding, dilation=dilation, groups=C_in,
                      bias=False),
            nn.Conv2d(C_in, C_out, kernel_size=1, padding=0, bias=False),
            nn.BatchNorm2d(C_out, affine=affine),
        )

    def forward(self, x):
        """Do an inference on DilConv.

        :param x: input tensor
        :return: output tensor
        """
        return self.op(x)

@NetworkFactory.register(NetTypes.Operator)
class SepConv(nn.Module):
    """Separable convolution block with repeats."""

    def __init__(self, C_in, C_out, kernel_size, stride, padding, dilation=1, affine=True, repeats=1):
        """Construct SepConv class.

        :param C_in: number of input channel
        :param C_out: number of output channel
        :param kernel_size: kernel size of the first conv
        :param stride: stride of the first conv
        :param padding: padding of the first conv
        :param dilation: dilation of the first conv
        :param affine: whether to use affine in BN
        :param repeats: number of repeat times
        """
        super(SepConv, self).__init__()

        def basic_op():
            return nn.Sequential(
                nn.Conv2d(C_in, C_in, kernel_size=kernel_size, stride=stride,
                          padding=padding, dilation=dilation, groups=C_in, bias=False),
                nn.Conv2d(C_in, C_in, kernel_size=1, padding=0, bias=False),
                nn.BatchNorm2d(C_in, affine=affine),
                nn.ReLU(inplace=False))

        self.op = nn.Sequential()
        for idx in range(repeats):
            self.op.add_module('sep_{}'.format(idx),
                               basic_op())

    def forward(self, x):
        """Do an inference on SepConv.

        :param x: input tensor
        :return: output tensor
        """
        return self.op(x)

@NetworkFactory.register(NetTypes.Operator)
class Zero(nn.Module):
    """Zero block."""

    def __init__(self, stride):
        """Construct Zero class.

        :param stride: stride of the output
        """
        super(Zero, self).__init__()
        self.stride = stride

    def forward(self, x):
        """Do an inference on Zero.

        :param x: input tensor
        :return: output tensor
        """
        if self.stride == 1:
            return x.mul(0.)
        return x[:, :, ::self.stride, ::self.stride].mul(0.)

@NetworkFactory.register(NetTypes.Operator)
class FactorizedReduce(nn.Module):
    """Factorized reduce block."""

    def __init__(self, C_in, C_out, affine=True):
        """Construct FactorizedReduce class.

        :param C_in: input channel
        :param C_out: output channel
        :param affine: whether to use affine in BN
        """
        super(FactorizedReduce, self).__init__()
        assert C_out % 2 == 0
        self.relu = nn.ReLU(inplace=False)
        self.conv_1 = nn.Conv2d(C_in, C_out // 2, 1, stride=2,
                                padding=0, bias=False)
        self.conv_2 = nn.Conv2d(C_in, C_out // 2, 1, stride=2,
                                padding=0, bias=False)
        self.bn = nn.BatchNorm2d(C_out, affine=affine)

    def forward(self, x):
        """Do an inference on FactorizedReduce.

        :param x: input tensor
        :return: output tensor
        """
        x = self.relu(x)
        out = torch.cat([self.conv_1(x), self.conv_2(x[:, :, 1:, 1:])], dim=1)
        out = self.bn(out)
        return out

@NetworkFactory.register(NetTypes.Operator)
class Pad(nn.Module):
    """Create Input for forward x."""

    def __init__(self, planes):
        super(Pad, self).__init__()
        self.planes = planes

    def forward(self, x):
        """Forward x."""
        return F.pad(x[:, :, ::2, ::2], (0, 0, 0, 0, self.planes // 4, self.planes // 4), "constant", 0)

OPS = {
    'none': lambda C, stride, affine, repeats=1: Zero(stride),
    'avg_pool_3x3': lambda C, stride, affine, repeats=1: nn.AvgPool2d(
        3, stride=stride, padding=1, count_include_pad=False),
    'max_pool_3x3': lambda C, stride, affine, repeats=1: nn.MaxPool2d(
        3, stride=stride, padding=1),
    'global_average_pool': lambda C, stride, affine, repeats=1: GAPConv1x1(C, C),
    'skip_connect': lambda C, stride, affine, repeats=1: Input() if stride == 1 else FactorizedReduce(
        C, C, affine=affine),
    'sep_conv_3x3': lambda C, stride, affine, repeats=1: SepConv(C, C, 3, stride, 1, affine=affine, repeats=repeats),
    'sep_conv_5x5': lambda C, stride, affine, repeats=1: SepConv(C, C, 5, stride, 2, affine=affine, repeats=repeats),
    'sep_conv_7x7': lambda C, stride, affine, repeats=1: SepConv(C, C, 7, stride, 3, affine=affine, repeats=repeats),
    'dil_conv_3x3': lambda C, stride, affine, repeats=1: DilConv(C, C, 3, stride, 2, 2, affine=affine),
    'dil_conv_5x5': lambda C, stride, affine, repeats=1: DilConv(C, C, 5, stride, 4, 2, affine=affine),
    'conv_7x1_1x7': lambda C, stride, affine, repeats=1: nn.Sequential(
        nn.ReLU(inplace=False),
        nn.Conv2d(C, C, (1, 7), stride=(1, stride), padding=(0, 3), bias=False),
        nn.Conv2d(C, C, (7, 1), stride=(stride, 1), padding=(3, 0), bias=False),
        nn.BatchNorm2d(C, affine=affine)),
    'conv1x1': lambda C, stride, affine, repeats=1: nn.Sequential(
        conv1X1(C, C, stride=stride),
        nn.BatchNorm2d(C, affine=affine),
        nn.ReLU(inplace=False)),
    'conv3x3': lambda C, stride, affine, repeats=1: nn.Sequential(
        conv3x3(C, C, stride=stride),
        nn.BatchNorm2d(C, affine=affine),
        nn.ReLU(inplace=False)),
    'conv5x5': lambda C, stride, affine, repeats=1: nn.Sequential(
        conv5x5(C, C, stride=stride),
        nn.BatchNorm2d(C, affine=affine),
        nn.ReLU(inplace=False)),
    'conv7x7': lambda C, stride, affine, repeats=1: nn.Sequential(
        conv7x7(C, C, stride=stride),
        nn.BatchNorm2d(C, affine=affine),
        nn.ReLU(inplace=False)),
    'conv3x3_dil2': lambda C, stride, affine, repeats=1: nn.Sequential(
        conv3x3(C, C, stride=stride, dilation=2),
        nn.BatchNorm2d(C, affine=affine),
        nn.ReLU(inplace=False)),
    'conv3x3_dil3': lambda C, stride, affine, repeats=1: nn.Sequential(
        conv3x3(C, C, stride=stride, dilation=3),
        nn.BatchNorm2d(C, affine=affine),
        nn.ReLU(inplace=False)),
    'conv3x3_dil12': lambda C, stride, affine, repeats=1: nn.Sequential(
        conv3x3(C, C, stride=stride, dilation=12),
        nn.BatchNorm2d(C, affine=affine),
        nn.ReLU(inplace=False)),
    'sep_conv_3x3_dil3': lambda C, stride, affine, repeats=1: SepConv(
        C, C, 3, stride, 3, affine=affine, dilation=3, repeats=repeats),
    'sep_conv_5x5_dil6': lambda C, stride, affine, repeats=1: SepConv(
        C, C, 5, stride, 12, affine=affine, dilation=6, repeats=repeats)
}


import sys, os, time
import wave
import numpy as np
from keras.models import Sequential
from keras.layers import Dense, LSTM
from keras.layers.convolutional import Conv1D, MaxPooling1D
from keras.layers.embeddings import Embedding
from keras.preprocessing import sequence
from keras.utils import np_utils
from bottle import Bottle, route, run, request, response, BaseRequest
import json

app = Bottle()
BaseRequest.MEMFILE_MAX = 1024000000

def pre():
	return "hello client!"

def predict(data):

	name = data['name']
	syllable = data['syllable']
	prediction_model_file = "C:\\jython\\Pronunciation\\Pronunciation-Matching\\main\\python\\H5\\" + name + ".h5"
	prediction_model_file = "C:\\jython\\Pronunciation\\Pronunciation-Matching\\main\\python\\H5\\" + "Bag" + ".h5"

	SAMPLE_RATE = 44100
	CHUNK_SIZE = 1024

	no_of_files = 15
	training_percentage = 0.7
	sequence_length = 50
	embedding_vecor_length = 32

	model = Sequential()
	model.add(Embedding(50, embedding_vecor_length, input_length=sequence_length))
	model.add(Conv1D(filters=32, kernel_size=3, padding='same', activation='relu'))
	model.add(MaxPooling1D(pool_size=2))
	model.add(LSTM(100))
	model.add(Dense(2, activation='softmax'))
	model.compile(loss='sparse_categorical_crossentropy', optimizer='adam', metrics=['accuracy'])
	model.load_weights(prediction_model_file)

	s_array = np.asarray(syllable)
	s_array_new = np.reshape(s_array, (-1, 1))
	s_use = np.transpose(s_array_new)

	predictions = model.predict(s_use)

	true = predictions[0][0]*100
	false = predictions[0][1]*100
	return true

@app.hook('after_request')
def enable_cors():
	response.headers['Access-Control-Allow-Origin'] = '*'
	response.headers['Access-Control-Allow-Methods'] = 'PUT, GET, POST, DELETE, OPTIONS'
	response.headers['Access-Control-Allow-Headers'] = 'Origin, Accept, Content-Type, X-Requested-With, X-CSRF-Token'

@app.route('/')
def hello():
	return 'Hello World!'

@app.route('/api/test', method=['OPTIONS', 'POST'])
def apitest():
	if request.method == 'POST':
		data = request.json
		print(type(data))
		print(data)
		res = {"true": predict(data)}
		return res
	else:
		return 'ok'

app.run(port=5000, debug=True)


def selection_0():

    # Library import
    import numpy
    import matplotlib
    import matplotlib.pyplot   as plt
    import matplotlib.gridspec as gridspec

    # Library version
    matplotlib_version = matplotlib.__version__
    numpy_version      = numpy.__version__

    # Histo binning
    xBinning = numpy.linspace(0.0,2000.0,201,endpoint=True)

    # Creating data sequence: middle of each bin
    xData = numpy.array([5.0,15.0,25.0,35.0,45.0,55.0,65.0,75.0,85.0,95.0,105.0,115.0,125.0,135.0,145.0,155.0,165.0,175.0,185.0,195.0,205.0,215.0,225.0,235.0,245.0,255.0,265.0,275.0,285.0,295.0,305.0,315.0,325.0,335.0,345.0,355.0,365.0,375.0,385.0,395.0,405.0,415.0,425.0,435.0,445.0,455.0,465.0,475.0,485.0,495.0,505.0,515.0,525.0,535.0,545.0,555.0,565.0,575.0,585.0,595.0,605.0,615.0,625.0,635.0,645.0,655.0,665.0,675.0,685.0,695.0,705.0,715.0,725.0,735.0,745.0,755.0,765.0,775.0,785.0,795.0,805.0,815.0,825.0,835.0,845.0,855.0,865.0,875.0,885.0,895.0,905.0,915.0,925.0,935.0,945.0,955.0,965.0,975.0,985.0,995.0,1005.0,1015.0,1025.0,1035.0,1045.0,1055.0,1065.0,1075.0,1085.0,1095.0,1105.0,1115.0,1125.0,1135.0,1145.0,1155.0,1165.0,1175.0,1185.0,1195.0,1205.0,1215.0,1225.0,1235.0,1245.0,1255.0,1265.0,1275.0,1285.0,1295.0,1305.0,1315.0,1325.0,1335.0,1345.0,1355.0,1365.0,1375.0,1385.0,1395.0,1405.0,1415.0,1425.0,1435.0,1445.0,1455.0,1465.0,1475.0,1485.0,1495.0,1505.0,1515.0,1525.0,1535.0,1545.0,1555.0,1565.0,1575.0,1585.0,1595.0,1605.0,1615.0,1625.0,1635.0,1645.0,1655.0,1665.0,1675.0,1685.0,1695.0,1705.0,1715.0,1725.0,1735.0,1745.0,1755.0,1765.0,1775.0,1785.0,1795.0,1805.0,1815.0,1825.0,1835.0,1845.0,1855.0,1865.0,1875.0,1885.0,1895.0,1905.0,1915.0,1925.0,1935.0,1945.0,1955.0,1965.0,1975.0,1985.0,1995.0])

    # Creating weights for histo: y1_PT_0
    y1_PT_0_weights = numpy.array([0.0,0.0,5285.39509429,14229.9095616,13823.3376312,11383.9260492,17482.4570042,12603.6338402,12603.6338402,13010.2017706,10977.3581189,8131.37460661,10164.2182583,11383.9260492,10164.2182583,12603.6338402,8537.94653694,9757.65032793,6505.09888529,7318.23874595,10570.7901886,10977.3581189,8944.51446727,8537.94653694,7724.80667628,6505.09888529,8944.51446727,7318.23874595,7318.23874595,6911.67081562,5285.39509429,4472.25523363,6911.67081562,4878.82716396,4472.25523363,2845.98151231,4472.25523363,2439.41278198,4065.6873033,3252.55024264,2439.41278198,4878.82716396,3252.55024264,4065.6873033,4065.6873033,4878.82716396,3659.11937297,3252.55024264,2439.41278198,2439.41278198,813.137460661,1219.70659099,1219.70659099,1626.27532132,2439.41278198,2439.41278198,2845.98151231,1626.27532132,1626.27532132,813.137460661,813.137460661,3252.55024264,1626.27532132,2439.41278198,813.137460661,1219.70659099,1219.70659099,1219.70659099,813.137460661,813.137460661,1626.27532132,2032.84405165,813.137460661,1219.70659099,0.0,406.56873033,813.137460661,1219.70659099,406.56873033,406.56873033,0.0,813.137460661,813.137460661,0.0,0.0,406.56873033,0.0,406.56873033,406.56873033,813.137460661,0.0,0.0,406.56873033,0.0,0.0,0.0,813.137460661,0.0,406.56873033,0.0,406.56873033,406.56873033,0.0,0.0,0.0,406.56873033,0.0,0.0,0.0,0.0,406.56873033,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,406.56873033,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,406.56873033,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,406.56873033,0.0,0.0,0.0,0.0])

    # Creating weights for histo: y1_PT_1
    y1_PT_1_weights = numpy.array([0.0,0.0,425.879703443,709.799413146,851.759295775,141.959882629,851.759295775,709.799413146,425.879703443,567.839530516,141.959882629,709.799413146,141.959882629,567.839530516,141.959882629,0.0,141.959882629,141.959882629,0.0,0.0,283.919820814,283.919820814,283.919820814,141.959882629,283.919820814,0.0,141.959882629,283.919820814,0.0,141.959882629,0.0,141.959882629,0.0,141.959882629,0.0,283.919820814,0.0,0.0,0.0,141.959882629,0.0,0.0,141.959882629,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,141.959882629,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,141.959882629,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,141.959882629,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,141.959882629,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,141.959882629,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0])

    # Creating a new Canvas
    fig   = plt.figure(figsize=(12,6),dpi=80)
    frame = gridspec.GridSpec(1,1,right=0.7)
    pad   = fig.add_subplot(frame[0])

    # Creating a new Stack
    pad.hist(x=xData, bins=xBinning, weights=y1_PT_0_weights+y1_PT_1_weights,\
             label="$signal4TeV$", histtype="stepfilled", rwidth=1.0,\
             color="#ce5e60", edgecolor="#ce5e60", linewidth=1, linestyle="solid",\
             bottom=None, cumulative=False, normed=False, align="mid", orientation="vertical")

    pad.hist(x=xData, bins=xBinning, weights=y1_PT_0_weights,\
             label="$signal1TeV$", histtype="stepfilled", rwidth=1.0,\
             color="#5954d8", edgecolor="#5954d8", linewidth=1, linestyle="solid",\
             bottom=None, cumulative=False, normed=False, align="mid", orientation="vertical")

    # Axis
    plt.rc('text',usetex=False)
    plt.xlabel(r"p_{T} [ j_{1} ]   ( GeV ) ",\
               fontsize=16,color="black")
    plt.ylabel(r"$\mathrm{Events}$ $(\mathcal{L}_{\mathrm{int}} = 40.0\ \mathrm{fb}^{-1})$ ",\
               fontsize=16,color="black")

    # Boundary of y-axis
    ymax=(y1_PT_0_weights+y1_PT_1_weights).max()*1.1
    ymin=0 # linear scale
    #ymin=min([x for x in (y1_PT_0_weights+y1_PT_1_weights) if x])/100. # log scale
    plt.gca().set_ylim(ymin,ymax)

    # Log/Linear scale for X-axis
    plt.gca().set_xscale("linear")
    #plt.gca().set_xscale("log",nonposx="clip")

    # Log/Linear scale for Y-axis
    plt.gca().set_yscale("linear")
    #plt.gca().set_yscale("log",nonposy="clip")

    # Legend
    plt.legend(bbox_to_anchor=(1.05,1), loc=2, borderaxespad=0.)

    # Saving the image
    plt.savefig('../../HTML/MadAnalysis5job_0/selection_0.png')
    plt.savefig('../../PDF/MadAnalysis5job_0/selection_0.png')
    plt.savefig('../../DVI/MadAnalysis5job_0/selection_0.eps')

# Running!
if __name__ == '__main__':
    selection_0()


# -*- coding: utf-8 -*-
"""
Created on Wed Nov  6 12:24:34 2019

@author: Z0014354
"""

import numpy as np
import random

import gym
from collections import namedtuple, deque
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.distributions import Normal, MultivariateNormal

import torch.optim as optim
import time
from torch.utils.tensorboard import SummaryWriter
import argparse

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

def hidden_init(layer):
    fan_in = layer.weight.data.size()[0]
    lim = 1. / np.sqrt(fan_in)
    return (-lim, lim)

class Actor(nn.Module):
    """Actor (Policy) Model."""

    def __init__(self, state_size, action_size, seed, hidden_size=32, init_w=3e-3, log_std_min=-20, log_std_max=2):
        """Initialize parameters and build model.
        Params
        ======
            state_size (int): Dimension of each state
            action_size (int): Dimension of each action
            seed (int): Random seed
            fc1_units (int): Number of nodes in first hidden layer
            fc2_units (int): Number of nodes in second hidden layer
        """
        super(Actor, self).__init__()
        self.seed = torch.manual_seed(seed)
        self.log_std_min = log_std_min
        self.log_std_max = log_std_max

        self.fc1 = nn.Linear(state_size, hidden_size)
        self.fc2 = nn.Linear(hidden_size, hidden_size)

        self.mu = nn.Linear(hidden_size, action_size)
        self.log_std_linear = nn.Linear(hidden_size, action_size)

    def reset_parameters(self):
        self.fc1.weight.data.uniform_(*hidden_init(self.fc1))
        self.fc2.weight.data.uniform_(*hidden_init(self.fc2))
        self.mu.weight.data.uniform_(-init_w, init_w)
        self.log_std_linear.weight.data.uniform_(-init_w, init_w)

    def forward(self, state):

        x = F.relu(self.fc1(state), inplace=True)
        x = F.relu(self.fc2(x), inplace=True)
        mu = self.mu(x)

        log_std = self.log_std_linear(x)
        log_std = torch.clamp(log_std, self.log_std_min, self.log_std_max)
        return mu, log_std

    def evaluate(self, state, epsilon=1e-6):
        mu, log_std = self.forward(state)
        std = log_std.exp()
        dist = Normal(0, 1)
        e = dist.sample().to(device)
        action = torch.tanh(mu + e * std)
        log_prob = Normal(mu, std).log_prob(mu + e * std) - torch.log(1 - action.pow(2) + epsilon)

        return action, log_prob

    def get_action(self, state):
        """
        returns the action based on a squashed gaussian policy. That means the samples are obtained according to:
        a(s,e)= tanh(mu(s)+sigma(s)+e)
        """
        #state = torch.FloatTensor(state).to(device) #.unsqzeeze(0)
        mu, log_std = self.forward(state)
        std = log_std.exp()
        dist = Normal(0, 1)
        e      = dist.sample().to(device)
        action = torch.tanh(mu + e * std).cpu()
        #action = torch.clamp(action*action_high, action_low, action_high)
        return action[0]

class Critic(nn.Module):
    """Critic (Value) Model."""

    def __init__(self, state_size, action_size, seed, hidden_size=32):
        """Initialize parameters and build model.
        Params
        ======
            state_size (int): Dimension of each state
            action_size (int): Dimension of each action
            seed (int): Random seed
            hidden_size (int): Number of nodes in the network layers

        """
        super(Critic, self).__init__()
        self.seed = torch.manual_seed(seed)
        self.fc1 = nn.Linear(state_size+action_size, hidden_size)
        self.fc2 = nn.Linear(hidden_size, hidden_size)
        self.fc3 = nn.Linear(hidden_size, 1)
        self.reset_parameters()

    def reset_parameters(self):
        self.fc1.weight.data.uniform_(*hidden_init(self.fc1))
        self.fc2.weight.data.uniform_(*hidden_init(self.fc2))
        self.fc3.weight.data.uniform_(-3e-3, 3e-3)

    def forward(self, state, action):
        """Build a critic (value) network that maps (state, action) pairs -> Q-values."""
        x = torch.cat((state, action), dim=1)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        return self.fc3(x)

class Agent():
    """Interacts with and learns from the environment."""

    def __init__(self, state_size, action_size, random_seed, hidden_size, action_prior="uniform"):
        """Initialize an Agent object.

        Params
        ======
            state_size (int): dimension of each state
            action_size (int): dimension of each action
            random_seed (int): random seed
        """
        self.state_size = state_size
        self.action_size = action_size
        self.seed = random.seed(random_seed)

        self.target_entropy = -action_size  # -dim(A)
        self.alpha = 1
        self.log_alpha = torch.tensor([0.0], requires_grad=True)
        self.alpha_optimizer = optim.Adam(params=[self.log_alpha], lr=LR_ACTOR)
        self._action_prior = action_prior

        print("Using: ", device)

        # Actor Network
        self.actor_local = Actor(state_size, action_size, random_seed, hidden_size).to(device)
        self.actor_optimizer = optim.Adam(self.actor_local.parameters(), lr=LR_ACTOR)

        # Critic Network (w/ Target Network)
        self.critic1 = Critic(state_size, action_size, random_seed, hidden_size).to(device)
        self.critic2 = Critic(state_size, action_size, random_seed, hidden_size).to(device)

        self.critic1_target = Critic(state_size, action_size, random_seed,hidden_size).to(device)
        self.critic1_target.load_state_dict(self.critic1.state_dict())

        self.critic2_target = Critic(state_size, action_size, random_seed,hidden_size).to(device)
        self.critic2_target.load_state_dict(self.critic2.state_dict())

        self.critic1_optimizer = optim.Adam(self.critic1.parameters(), lr=LR_CRITIC, weight_decay=0)
        self.critic2_optimizer = optim.Adam(self.critic2.parameters(), lr=LR_CRITIC, weight_decay=0)

        # Replay memory
        self.memory = ReplayBuffer(action_size, BUFFER_SIZE, BATCH_SIZE, random_seed)

    def step(self, state, action, reward, next_state, done, step):
        """Save experience in replay memory, and use random sample from buffer to learn."""
        # Save experience / reward
        self.memory.add(state, action, reward, next_state, done)

        # Learn, if enough samples are available in memory
        if len(self.memory) > BATCH_SIZE:
            experiences = self.memory.sample()
            self.learn(step, experiences, GAMMA)

    def act(self, state):
        """Returns actions for given state as per current policy."""
        state = torch.from_numpy(state).float().to(device)
        action = self.actor_local.get_action(state).detach()
        return action

    def learn(self, step, experiences, gamma, d=1):
        """Updates actor, critics and entropy_alpha parameters using given batch of experience tuples.
        Q_targets = r + γ * (min_critic_target(next_state, actor_target(next_state)) - α *log_pi(next_action|next_state))
        Critic_loss = MSE(Q, Q_target)
        Actor_loss = α * log_pi(a|s) - Q(s,a)
        where:
            actor_target(state) -> action
            critic_target(state, action) -> Q-value
        Params
        ======
            experiences (Tuple[torch.Tensor]): tuple of (s, a, r, s', done) tuples
            gamma (float): discount factor
        """
        states, actions, rewards, next_states, dones = experiences

        # ---------------------------- update critic ---------------------------- #
        # Get predicted next-state actions and Q values from target models
        next_action, log_pis_next = self.actor_local.evaluate(next_states)

        Q_target1_next = self.critic1_target(next_states.to(device), next_action.squeeze(0).to(device))
        Q_target2_next = self.critic2_target(next_states.to(device), next_action.squeeze(0).to(device))

        # take the mean of both critics for updating
        Q_target_next = torch.min(Q_target1_next, Q_target2_next)

        if FIXED_ALPHA == None:
            # Compute Q targets for current states (y_i)
            Q_targets = rewards.cpu() + (gamma * (1 - dones.cpu()) * (Q_target_next.cpu() - self.alpha * log_pis_next.squeeze(0).cpu()))
        else:
            Q_targets = rewards.cpu() + (gamma * (1 - dones.cpu()) * (Q_target_next.cpu() - FIXED_ALPHA * log_pis_next.squeeze(0).cpu()))
        # Compute critic loss
        Q_1 = self.critic1(states, actions).cpu()
        Q_2 = self.critic2(states, actions).cpu()
        critic1_loss = 0.5*F.mse_loss(Q_1, Q_targets.detach())
        critic2_loss = 0.5*F.mse_loss(Q_2, Q_targets.detach())
        # Update critics
        # critic 1
        self.critic1_optimizer.zero_grad()
        critic1_loss.backward()
        self.critic1_optimizer.step()
        # critic 2
        self.critic2_optimizer.zero_grad()
        critic2_loss.backward()
        self.critic2_optimizer.step()
        if step % d == 0:
        # ---------------------------- update actor ---------------------------- #
            if FIXED_ALPHA == None:
                alpha = torch.exp(self.log_alpha)
                # Compute alpha loss
                actions_pred, log_pis = self.actor_local.evaluate(states)
                alpha_loss = - (self.log_alpha.cpu() * (log_pis.cpu() + self.target_entropy).detach().cpu()).mean()
                self.alpha_optimizer.zero_grad()
                alpha_loss.backward()
                self.alpha_optimizer.step()

                self.alpha = alpha
                # Compute actor loss
                if self._action_prior == "normal":
                    policy_prior = MultivariateNormal(loc=torch.zeros(self.action_size), scale_tril=torch.ones(self.action_size).unsqueeze(0))
                    policy_prior_log_probs = policy_prior.log_prob(actions_pred)
                elif self._action_prior == "uniform":
                    policy_prior_log_probs = 0.0

                actor_loss = (alpha * log_pis.squeeze(0).cpu() - self.critic1(states, actions_pred.squeeze(0)).cpu() - policy_prior_log_probs ).mean()
            else:

                actions_pred, log_pis = self.actor_local.evaluate(states)
                if self._action_prior == "normal":
                    policy_prior = MultivariateNormal(loc=torch.zeros(self.action_size), scale_tril=torch.ones(self.action_size).unsqueeze(0))
                    policy_prior_log_probs = policy_prior.log_prob(actions_pred)
                elif self._action_prior == "uniform":
                    policy_prior_log_probs = 0.0

                actor_loss = (FIXED_ALPHA * log_pis.squeeze(0).cpu() - self.critic1(states, actions_pred.squeeze(0)).cpu()- policy_prior_log_probs ).mean()
            # Minimize the loss
            self.actor_optimizer.zero_grad()
            actor_loss.backward()
            self.actor_optimizer.step()

            # ----------------------- update target networks ----------------------- #
            self.soft_update(self.critic1, self.critic1_target, TAU)
            self.soft_update(self.critic2, self.critic2_target, TAU)

    def soft_update(self, local_model, target_model, tau):
        """Soft update model parameters.
        θ_target = τ*θ_local + (1 - τ)*θ_target
        Params
        ======
            local_model: PyTorch model (weights will be copied from)
            target_model: PyTorch model (weights will be copied to)
            tau (float): interpolation parameter
        """
        for target_param, local_param in zip(target_model.parameters(), local_model.parameters()):
            target_param.data.copy_(tau*local_param.data + (1.0-tau)*target_param.data)

class ReplayBuffer:
    """Fixed-size buffer to store experience tuples."""

    def __init__(self, action_size, buffer_size, batch_size, seed):
        """Initialize a ReplayBuffer object.
        Params
        ======
            buffer_size (int): maximum size of buffer
            batch_size (int): size of each training batch
        """
        self.action_size = action_size
        self.memory = deque(maxlen=buffer_size)  # internal memory (deque)
        self.batch_size = batch_size
        self.experience = namedtuple("Experience", field_names=["state", "action", "reward", "next_state", "done"])
        self.seed = random.seed(seed)

    def add(self, state, action, reward, next_state, done):
        """Add a new experience to memory."""
        e = self.experience(state, action, reward, next_state, done)
        self.memory.append(e)

    def sample(self):
        """Randomly sample a batch of experiences from memory."""
        experiences = random.sample(self.memory, k=self.batch_size)

        states = torch.from_numpy(np.vstack([e.state for e in experiences if e is not None])).float().to(device)
        actions = torch.from_numpy(np.vstack([e.action for e in experiences if e is not None])).float().to(device)
        rewards = torch.from_numpy(np.vstack([e.reward for e in experiences if e is not None])).float().to(device)
        next_states = torch.from_numpy(np.vstack([e.next_state for e in experiences if e is not None])).float().to(device)
        dones = torch.from_numpy(np.vstack([e.done for e in experiences if e is not None]).astype(np.uint8)).float().to(device)

        return (states, actions, rewards, next_states, dones)

    def __len__(self):
        """Return the current size of internal memory."""
        return len(self.memory)

def SAC(n_episodes=200, max_t=500, print_every=10):
    scores_deque = deque(maxlen=100)
    average_100_scores = []

    for i_episode in range(1, n_episodes+1):

        state = env.reset()
        state = state.reshape((1,state_size))
        score = 0
        for t in range(max_t):

            action = agent.act(state)
            action_v = action.numpy()
            action_v = np.clip(action_v*action_high, action_low, action_high)
            next_state, reward, done, info = env.step(action_v)
            next_state = next_state.reshape((1,state_size))
            agent.step(state, action, reward, next_state, done, t)
            state = next_state
            score += reward

            if done:
                break

        scores_deque.append(score)
        writer.add_scalar("Reward", score, i_episode)
        writer.add_scalar("average_X", np.mean(scores_deque), i_episode)
        average_100_scores.append(np.mean(scores_deque))

        print('\rEpisode {} Reward: {:.2f}  Average100 Score: {:.2f}'.format(i_episode, score, np.mean(scores_deque)), end="")
        if i_episode % print_every == 0:
            print('\rEpisode {}  Reward: {:.2f}  Average100 Score: {:.2f}'.format(i_episode, score, np.mean(scores_deque)))

    torch.save(agent.actor_local.state_dict(), args.info + ".pt")

def play():
    agent.actor_local.eval()
    for i_episode in range(1):

        state = env.reset()
        state = state.reshape((1,state_size))

        while True:
            action = agent.act(state)
            action_v = action[0].numpy()
            action_v = np.clip(action_v*action_high, action_low, action_high)
            next_state, reward, done, info = env.step(action_v)
            next_state = next_state.reshape((1,state_size))
            state = next_state

            if done:
                break

parser = argparse.ArgumentParser(description="")
parser.add_argument("-env", type=str,default="Pendulum-v0", help="Environment name")
parser.add_argument("-info", type=str, help="Information or name of the run")
parser.add_argument("-ep", type=int, default=100, help="The amount of training episodes, default is 100")
parser.add_argument("-seed", type=int, default=0, help="Seed for the env and torch network weights, default is 0")
parser.add_argument("-lr", type=float, default=5e-4, help="Learning rate of adapting the network weights, default is 5e-4")
parser.add_argument("-a", "--alpha", type=float, help="entropy alpha value, if not choosen the value is leaned by the agent")
parser.add_argument("-layer_size", type=int, default=256, help="Number of nodes per neural network layer, default is 256")
parser.add_argument("-repm", "--replay_memory", type=int, default=int(1e6), help="Size of the Replay memory, default is 1e6")
parser.add_argument("--print_every", type=int, default=100, help="Prints every x episodes the average reward over x episodes")
parser.add_argument("-bs", "--batch_size", type=int, default=256, help="Batch size, default is 256")
parser.add_argument("-t", "--tau", type=float, default=1e-2, help="Softupdate factor tau, default is 1e-2")
parser.add_argument("-g", "--gamma", type=float, default=0.99, help="discount factor gamma, default is 0.99")
parser.add_argument("--saved_model", type=str, default=None, help="Load a saved model to perform a test run!")
args = parser.parse_args()

if __name__ == "__main__":
    env_name = args.env
    seed = args.seed
    n_episodes = args.ep
    GAMMA = args.gamma
    TAU = args.tau
    HIDDEN_SIZE = args.layer_size
    BUFFER_SIZE = int(args.replay_memory)
    BATCH_SIZE = args.batch_size        # minibatch size
    LR_ACTOR = args.lr         # learning rate of the actor
    LR_CRITIC = args.lr        # learning rate of the critic
    FIXED_ALPHA = args.alpha
    saved_model = args.saved_model

    t0 = time.time()
    writer = SummaryWriter("runs/"+args.info)
    env = gym.make(env_name)
    action_high = env.action_space.high[0]
    action_low = env.action_space.low[0]
    torch.manual_seed(seed)
    env.seed(seed)
    np.random.seed(seed)
    state_size = env.observation_space.shape[0]
    action_size = env.action_space.shape[0]
    agent = Agent(state_size=state_size, action_size=action_size, random_seed=seed,hidden_size=HIDDEN_SIZE, action_prior="uniform") #"normal"

    if saved_model != None:
        agent.actor_local.load_state_dict(torch.load(saved_model))
        play()
    else:
        SAC(n_episodes=args.ep, max_t=500, print_every=args.print_every)
    t1 = time.time()
    env.close()
    print("training took {} min!".format((t1-t0)/60))


import face_model
import argparse
import cv2
import sys
import numpy as np

import os
import csv
import glob

# general
parser = argparse.ArgumentParser(description='face model test')
parser.add_argument('--image-size', default='112,112', help='')
parser.add_argument('--model', default='/models/model-r100-ii/model,0', help='path to load model.')
parser.add_argument('--ga-model', default='', help='path to load model.')
parser.add_argument('--gpu', default=0, type=int, help='gpu id')
parser.add_argument('--det', default=0, type=int, help='mtcnn option, 1 means using R+O, 0 means detect from begining')
parser.add_argument('--flip', default=0, type=int, help='whether do lr flip aug')
parser.add_argument('--threshold', default=1.24, type=float, help='ver dist threshold')
args = parser.parse_args()

model = face_model.FaceModel(args)

### CHANGE THIS LINE TO YOUR <INPUT> DIRECTORY OF IMAGES
input_dir = './images/'

### CHANGE THIS LINE TO YOUR <OUTPUT> DIRECTORY OF CSV'S
output_dir = './csv/'

fnames1 = glob.glob(input_dir+'*.jpeg')
fnames2 = glob.glob(input_dir+'.png')
fnames = fnames1+fnames2
fnames_count = len(fnames)

for i, fname in enumerate(fnames):

	img = cv2.imread(fname)

	inp = model.get_input(img)

	if inp.any():
		feat = model.get_feature(inp)
	else:
		print('No face found.')
		break

	if feat.any():

		csv_fname = fname.split('/')[-1].split('.')[0]+'.csv'

		### COMMENT NEXT LINE OUT IF YOU DO NOT CARE ABOUT SEEING PROGRESS
		print('Writing to \'{}\' - Image {}/{}'.format(output_dir+csv_fname, i+1, fnames_count))

		with open(output_dir+'{}.csv'.format(csv_fname), 'w', newline='') as csvfile:
			writer = csv.writer(csvfile, delimiter=',')
			writer.writerow(feat)

print('Done processing images.')

# gender, age = model.get_ga(img)
# print(gender)
# print(age)
# sys.exit(0)
# img = cv2.imread('/raid5data/dplearn/megaface/facescrubr/112x112/Tom_Hanks/Tom_Hanks_54733.png')
# f2 = model.get_feature(img)
# dist = np.sum(np.square(f1-f2))
# print(dist)
# sim = np.dot(f1, f2.T)
# print(sim)
# diff = np.subtract(source_feature, target_feature)
# dist = np.sum(np.square(diff),1)


import tensorflow
from model import *
import numpy as np

np.set_printoptions(edgeitems = 100000, suppress = True)
num_mcep = 24
num_speakers = 10

model = CCGAN(num_features = num_mcep, num_speakers = num_speakers)
model.load(filepath = os.path.join('./model', 'CCGAN.ckpt'))
model.print_weight()


# Start command
# uvicorn restapi:app
# gunicorn restapi:app -k uvicorn.workers.UvicornWorker

from threading import Thread
from concurrent.futures import ThreadPoolExecutor
from fastapi import FastAPI, HTTPException
from starlette.middleware.cors import CORSMiddleware
from starlette.responses import FileResponse, RedirectResponse
from starlette.requests import Request
from PIL import Image
import os
import sys
import json
import importlib
import numpy as np
from datetime import datetime
import skimage.io as io
from skimage.metrics import *
from src.image_loader import image_loader
sys.path.append('./src/')
import models
import secrets

### Parameters #################################################
pantie_dir = './dream/'
converted_dir = './converted/'
log_dir = './logs/'
################################################################

def make_display_names():
    with open('./webapp.json', mode='r') as f:
        options = json.load(f)['all']
    display_names = [importlib.import_module('models.' + model).patcher(options=options).name for model in models.models_namelist]
    return display_names

def update_log(model, image, flags):
    log.update({str(len(log)): {'model': model, 'image': image, 'flag': flags}})
    try:
        with open(log_dir + logfile, 'w') as f:
            json.dump(log, f, indent=2, separators=(',', ':'))
    except:
        pass

app = FastAPI()
app.add_middleware(
    CORSMiddleware,
    allow_credentials=True,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"]
)

@app.get("/")
def home():
    return f'Here is Quiche Pantie Patch Server! You can access the panties: https://pantie-patch.herokuapp.com/api/dream/****.png. When you convert the panties: https://pantie-patch.herokuapp.com/api/convert/specify_avatar_name/****.png'

os.makedirs(converted_dir, exist_ok=True)
os.makedirs(log_dir, exist_ok=True)
panties = sorted(os.listdir('./dream'))
database = {"models": models.models_namelist, "images": panties}
display_names = make_display_names()
logfile = datetime.today().strftime('%Y_%m_%d_%H_%M') + '.json'
log = {}
try:
    with open(log_dir + logfile, 'w') as f:
        json.dump(log, f)
except:
    pass

class score_processor:
    def __init__(self, workers=2):
        self.panties = sorted(os.listdir(pantie_dir))
        self.done = False
        self.startpoint = 0
        try:
            self.score_matrix = np.load(converted_dir + 'score_matrix.npy')
            if self.score_matrix.shape[0] == len(panties):
                self.done = True
            else:
                pad = len(panties) - self.score_matrix.shape[0]
                self.startpoint = self.score_matrix.shape[0]
                self.score_matrix = np.pad(self.score_matrix, [(0, pad), (0, pad)], mode='constant')
        except:
            self.score_matrix = np.zeros((len(self.panties), len(self.panties)))
        self.workers = workers

    # it is for triangular score matrix. when you set full, it calculate a line of the full matrix
    def score_row(self, args):
        num = args[0]
        full = args[1]
        edge = 100
        remains = self.panties.copy()
        if full:
            remains.pop(num)
        else:
            [remains.pop(0) for i in range(num + 1)]
        template_loader = image_loader(fdir=pantie_dir, queuesize=32)
        template_loader.flist = remains
        template_loader.start()
        scores = []
        ref = io.imread(pantie_dir + panties[num])[50:-edge * 2, edge:-edge, :3]
        for check_pantie in remains:
            tmp = np.array(template_loader.read())[50:-edge * 2, edge:-edge, :3]
            scores.append(mean_squared_error(ref, tmp))
        if full:
            scores = np.insert(scores, num, 0)
        else:
            scores = np.pad(scores, (num + 1, 0), mode='constant')
        return np.array(scores)

    def argument_generator(self, num_pantie, flag=True):
        for i in range(self.startpoint, num_pantie):
            yield (i, flag)

    def start(self):
        t = Thread(target=self.process, args=())
        t.daemon = True
        if self.done is not True:
            t.start()
        return t

    def process(self):
        with ThreadPoolExecutor(max_workers=self.workers) as executor:
            nums = len(panties)
            if self.startpoint != 0:
                scores = executor.map(self.score_row, self.argument_generator(nums, True))
            else:
                scores = executor.map(self.score_row, self.argument_generator(nums, False))

            for row, score in enumerate(scores):
                if self.startpoint != 0:
                    self.score_matrix[:, row + self.startpoint] = score
                else:
                    self.score_matrix[row + self.startpoint, :] = score
            self.score_matrix = np.triu(self.score_matrix)
            self.score_matrix += self.score_matrix.T
        np.save(converted_dir + 'score_matrix.npy', self.score_matrix)
        self.done = True

sp = score_processor()
sp_thred = sp.start()

@app.get("/api/")
async def request_apps():
    if os.path.exists('./zips/sharelinks.json'):
        return {"apps": ["dream", "convert", "suggest", "gacha", "zip"]}
    else:
        return {"apps": ["dream", "convert", "suggest", "gacha"]}

@app.get("/api/dream/")
async def request_pantie_list():
    return {"images": panties}

@app.get("/api/dream/{image}")
async def send_pantie(image: str):
    if image not in database['images']:
        raise HTTPException(status_code=404, detail=f"{image} doesn't exist")
    return FileResponse(pantie_dir + image, media_type="image/png")

@app.get("/api/convert/")
async def request_model_list():
    return {"display_names": display_names, "models": models.models_namelist}

@app.get("/api/convert/{model}/")
async def request_model_option_list(model: str):
    with open('./webapp.json', mode='r') as f:
        options = json.load(f)
    if model not in database['models']:
        raise HTTPException(status_code=404, detail=f"{model} doesn't exist")
    display_name = display_names[database['models'].index(model)]
    try:
        options = options[model]
    except:
        options = []
    return {"display_name": display_name, "images": panties, "options": options}

@app.get("/api/convert/{model}/{image}")
async def send_converted(model: str, image: str, request: Request):
    if image not in database['images']:
        raise HTTPException(status_code=404, detail=f"{image} doesn't exist")
    if model not in database['models']:
        raise HTTPException(status_code=404, detail=f"{model} doesn't exist")

    with open('./webapp.json', mode='r') as f:
        options = json.load(f)['all']

    def option_parser(options):
        flags = []
        for key in options.keys():
            if request.query_params.get(key) is not None:
                print(key + ':' + request.query_params.get(key))
                tmp = options[key]
                if request.query_params.get(key) in ['true', 'True', 'yes', 'Yes']:
                    options[key] = True
                else:
                    options[key] = False
                if tmp != options[key]:
                    flags.append(key)
        return options, flags

    options, flags = option_parser(options)
    fdir = f'{converted_dir}{model}/'
    if len(flags) == 0:
        fdir += 'default/'
    elif len(flags) == 1:
        fdir += f'{flags[0]}/'
    else:
        fdir += 'others/'
    if not os.path.exists(fdir + image) or 'others' in fdir:
        os.makedirs(fdir, exist_ok=True)
        module = importlib.import_module('models.' + model)
        options['model'] = model
        options['input'] = './body/body_' + model + '.png'
        options['output'] = fdir + image
        options['pantie'] = int(image.split('.')[0]) - 1
        patcher = module.patcher(options=options)
        patched = patcher.patch(Image.open(pantie_dir + panties[options['pantie']]), transparent=True)
        patcher.save(patched, options['output'])
    update_log(model, image, flags)
    return FileResponse(fdir + image, media_type="image/png")

@app.get("/api/suggest/")
async def request_pantie_list():
    return {"images": panties}

@app.get("/api/suggest/{image}")
async def request_suggest_list(image: str):
    if image not in database['images']:
        raise HTTPException(status_code=404, detail=f"{image} doesn't exist")
    pantie = int(image[:-4]) - 1
    if sp.done:
        scores = sp.score_matrix[pantie, :]
    else:
        scores = sp.score_row((pantie, True))
    rank = np.argsort(scores)
    suggests = [database['images'][index] for index in rank[1:]]
    scores = [scores[index] for index in rank[1:]]
    return {"suggests": suggests, "scores": scores}

@app.get("/api/gacha/")
async def request_gacha_list():
    return {"1ren", "10ren"}

@app.get("/api/gacha/1ren")
async def send_1ren_result():
    return {panties[secrets.randbelow(len(panties))]}

@app.get("/api/gacha/10ren")
async def send_10ren_result():
    return [panties[secrets.randbelow(len(panties))] for i in range(10)]

@app.get("/api/zip/")
async def request_zip_list():
    try:
        with open('./zips/sharelinks.json', mode='r') as f:
            zips = json.load(f)
    except:
        raise HTTPException(status_code=404, detail=f"Zip doesn't exist")
    return {"display_names": display_names, "zips": [k for k in sorted(zips.keys())]}

@app.get("/api/zip/{zipname}")
async def send_zip(zipname: str):
    try:
        with open('./zips/sharelinks.json', mode='r') as f:
            zips = json.load(f)
    except:
        raise HTTPException(status_code=404, detail=f"Zip doesn't exist")
    if zipname not in zips.keys():
        raise HTTPException(status_code=404, detail=f"{zipname} doesn't exist")
    update_log(zipname.split('.')[0], zipname, ['zip'])
    return RedirectResponse(url=zips[zipname])


#done changes
import os
import sys
import json
import numpy as np
import torch
from torch import nn
from torch import optim
from torch.optim import lr_scheduler
from threading import Thread
from threading import Lock
import time
import cv2
from torch.autograd import Variable
import torch.nn.functional as F

from opts import parse_opts
from model import generate_model
from mean import get_mean, get_std
from spatial_transforms import (
    Compose, Normalize, Scale, CenterCrop, CornerCrop, MultiScaleCornerCrop,
    MultiScaleRandomCrop, RandomHorizontalFlip, ToTensor)
from temporal_transforms import LoopPadding, TemporalRandomCrop
from target_transforms import ClassLabel, VideoID
from target_transforms import Compose as TargetCompose
from dataset import get_training_set, get_validation_set, get_test_set
from utils import Logger
from train import train_epoch
from validation import val_epoch
import test

def readwebcam():
    cap = cv2.VideoCapture(0)
    global frames

    while (True):
        # Capture frame-by-frame
        ret, frame = cap.read()

        # Our operations on the frame come here
        #gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

        print(frame.shape)
        # Display the resulting frame
        cv2.imshow('frame', frame)
        frame = np.transpose(frame, (2, 0, 1))

        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

        lock.acquire()
        frames = np.delete(frames, (0), axis=0)
        frames = np.append(frames, frame.reshape(1, 3, 480, -1), axis=0)
        lock.release()
        #time.sleep(0.1)
    # When everything done, release the capture
    cap.release()

if __name__ == '__main__':
    opt = parse_opts()
    if opt.root_path != '':
        opt.video_path = os.path.join(opt.root_path, opt.video_path)
        opt.annotation_path = os.path.join(opt.root_path, opt.annotation_path)
        opt.result_path = os.path.join(opt.root_path, opt.result_path)
        if opt.resume_path:
            opt.resume_path = os.path.join(opt.root_path, opt.resume_path)
        if opt.pretrain_path:
            opt.pretrain_path = os.path.join(opt.root_path, opt.pretrain_path)
    opt.scales = [opt.initial_scale]
    for i in range(1, opt.n_scales):
        opt.scales.append(opt.scales[-1] * opt.scale_step)
    opt.arch = '{}-{}'.format(opt.model, opt.model_depth)
    opt.mean = get_mean(opt.norm_value, dataset=opt.mean_dataset)
    opt.std = get_std(opt.norm_value)
    print(opt)
    with open(os.path.join(opt.result_path, 'opts.json'), 'w') as opt_file:
        json.dump(vars(opt), opt_file)

    torch.manual_seed(opt.manual_seed)

    model, parameters = generate_model(opt)
    print(model)
    criterion = nn.CrossEntropyLoss()
    if not opt.no_cuda:
        criterion = criterion.cuda()

    if opt.no_mean_norm and not opt.std_norm:
        norm_method = Normalize([0, 0, 0], [1, 1, 1])
    elif not opt.std_norm:
        norm_method = Normalize(opt.mean, [1, 1, 1])
    else:
        norm_method = Normalize(opt.mean, opt.std)

    """if opt.test:
        spatial_transform = Compose([
            Scale(int(opt.sample_size / opt.scale_in_test)),
            CornerCrop(opt.sample_size, opt.crop_position_in_test),
            ToTensor(opt.norm_value), norm_method
        ])
        temporal_transform = LoopPadding(opt.sample_duration)
        target_transform = VideoID()

        test_data = get_test_set(opt, spatial_transform, temporal_transform,
                                 target_transform)
        test_loader = torch.utils.data.DataLoader(
            test_data,
            batch_size=opt.batch_size,
            shuffle=False,
            num_workers=opt.n_threads,
            pin_memory=True)
        test.test(test_loader, model, opt, test_data.class_names)
    """

    lock = Lock()
    t = Thread(target=readwebcam)
    t.start()
    frames = np.random.rand(16, 3, 480, 640)

    model.eval()
    classes = ['abort', 'circle', 'hello', 'no', 'stop', 'turn', 'turn_left', 'turn_right', 'warn']

    while (True):
        lock.acquire()
        print(frames[1])
        lock.release()
        # inputs type and shape
        #<class 'torch.Tensor'>
        #torch.Size([10, 3, 16, 112, 112])
        #TODO add transformations
        inputs = torch.unsqueeze(torch.from_numpy(frames), 0).permute(0, 2, 1, 3, 4)
        print(inputs.size())
        inputs = Variable(inputs, volatile=True)
        outputs = model(inputs)
        outputs = F.softmax(outputs)
        _, ind = torch.max(outputs)
        print(classes[ind])
        #time.sleep(1)

    t.join()


# flake8: noqa
from contextlib import ExitStack
from pfb.workers.main import cli
import click
from omegaconf import OmegaConf
import pyscilog
pyscilog.init('pfb')
log = pyscilog.get_logger('PSF')

@cli.command()
@click.option('-ms', '--ms', required=True,
              help='Path to measurement set.')
@click.option('-dc', '--data-column',
              help="Data column to image."
              "Must be the same across MSs")
@click.option('-wc', '--weight-column',
              help="Column containing natural weights."
              "Must be the same across MSs")
@click.option('-iwc', '--imaging-weight-column',
              help="Column containing/to write imaging weights to."
              "Must be the same across MSs")
@click.option('-eps', '--epsilon', type=float, default=1e-5,
              help='Gridder accuracy')
@click.option('--wstack/--no-wstack', default=True)
@click.option('--double-accum/--no-double-accum', default=True)
@click.option('--mock/--no-mock', default=False)
@click.option('-fc', '--flag-column', default='FLAG',
              help="Column containing data flags."
              "Must be the same across MSs")
@click.option('-muc', '--mueller-column',
              help="Column containing Mueller terms."
              "Must be the same across MSs")
@click.option('-rchunk', '--row-chunks',
              help="Number of rows in a chunk.")
@click.option('-o', '--output-filename', type=str, required=True,
              help="Basename of output.")
@click.option('-nb', '--nband', type=int, required=True,
              help="Number of imaging bands")
@click.option('-fov', '--field-of-view', type=float,
              help="Field of view in degrees")
@click.option('-srf', '--super-resolution-factor', type=float,
              help="Will over-sample Nyquist by this factor at max frequency")
@click.option('-psfo', '--psf-oversize', type=float, default=2.0,
              help='Size of the PSF relative to the dirty image.')
@click.option('-cs', '--cell-size', type=float,
              help='Cell size in arcseconds')
@click.option('-nx', '--nx', type=int,
              help="Number of x pixels")
@click.option('-ny', '--ny', type=int,
              help="Number of x pixels")
@click.option('-otype', '--output-type', default='f4',
              help="Data type of output")
@click.option('-ha', '--host-address',
              help='Address where the distributed client lives. '
              'Will use a local cluster if no address is provided')
@click.option('-nw', '--nworkers', type=int,
              help='Number of workers for the client.')
@click.option('-ntpw', '--nthreads-per-worker', type=int,
              help='Number of dask threads per worker.')
@click.option('-nvt', '--nvthreads', type=int,
              help="Total number of threads to use for vertical scaling (eg. gridder, fft's etc.)")
@click.option('-mem', '--mem-limit', type=float,
              help="Memory limit in GB. Default uses all available memory")
@click.option('-nthreads', '--nthreads', type=int,
              help="Total available threads. Default uses all available threads")
@click.option('-roc', '--row-out-chunk', type=int, default=10000,
              help="Size of row chunks for output weights and uvw")
def psf(**kw):
    '''
    Create a psf image from a list of measurement setsand write out the
    Mueller weights.

    The psf image cube is not normalised by wsum as this destroyes
    information. The MFS image is written out in units of Jy/beam
    and should have a peak of one otherwise something has gone wrong.

    The --field-of-view and --super-resolution-factor options
    (equivalently --cell-size, --nx and --ny) pertain to the size of
    the image (eg. dirty and model). The size of the PSF output image
    is controlled by the --psf-oversize option.

    The Stokes I weights required to apply the Hessian are also written out
    to a zarr data set called output-filename.zarr. This data set does not
    adhere to the MSv2 specs and is only meant to be used to apply the
    Hessian. In particular, the weights written out are a combination of
    imaging weights and the "Mueller" weights.

    If a host address is provided the computation can be distributed
    over imaging band and row. When using a distributed scheduler both
    mem-limit and nthreads is per node and have to be specified.

    When using a local cluster, mem-limit and nthreads refer to the global
    memory and threads available, respectively. By default the gridder will
    use all available resources.

    Disclaimer - Memory budgeting is still very crude!

    On a local cluster, the default is to use:

        nworkers = nband
        nthreads-per-worker = 1

    They have to be specified in ~.config/dask/jobqueue.yaml in the
    distributed case.

    if LocalCluster:
        ngridder-threads = nthreads//(nworkers*nthreads_per_worker)
    else:
        ngridder-threads = nthreads//nthreads-per-worker

    '''
    args = OmegaConf.create(kw)
    pyscilog.log_to_file(args.output_filename + '.log')
    from glob import glob
    ms = glob(args.ms)
    try:
        assert len(ms) > 0
        args.ms = ms
    except:
        raise ValueError(f"No MS at {args.ms}")

    if args.nworkers is None:
        args.nworkers = args.nband

    OmegaConf.set_struct(args, True)

    with ExitStack() as stack:
        from pfb import set_client
        args = set_client(args, stack, log)

        # TODO - prettier config printing
        print('Input Options:', file=log)
        for key in args.keys():
            print('     %25s = %s' % (key, args[key]), file=log)

        return _psf(**args)

def _psf(**kw):
    args = OmegaConf.create(kw)
    from omegaconf import ListConfig
    if not isinstance(args.ms, list) and not isinstance(args.ms, ListConfig) :
        args.ms = [args.ms]
    OmegaConf.set_struct(args, True)

    import numpy as np
    from pfb.utils.misc import chan_to_band_mapping
    import dask
    # from dask.distributed import performance_report
    from dask.graph_manipulation import clone
    from daskms import xds_from_storage_ms as xds_from_ms
    from daskms import xds_from_storage_table as xds_from_table
    from daskms import Dataset
    from daskms.experimental.zarr import xds_to_zarr
    import dask.array as da
    from africanus.constants import c as lightspeed
    from africanus.gridding.wgridder.dask import dirty as vis2im
    from ducc0.fft import good_size
    from pfb.utils.misc import stitch_images, plan_row_chunk
    from pfb.utils.fits import set_wcs, save_fits

    # chan <-> band mapping
    ms = args.ms
    nband = args.nband
    freqs, freq_bin_idx, freq_bin_counts, freq_out, band_mapping, chan_chunks = chan_to_band_mapping(
        ms, nband=nband)

    # gridder memory budget
    max_chan_chunk = 0
    max_freq = 0
    for ims in args.ms:
        for spw in freqs[ims]:
            counts = freq_bin_counts[ims][spw].compute()
            freq = freqs[ims][spw].compute()
            max_chan_chunk = np.maximum(max_chan_chunk, counts.max())
            max_freq = np.maximum(max_freq, freq.max())

    # assumes measurement sets have the same columns,
    # number of correlations etc.
    xds = xds_from_ms(args.ms[0])
    ncorr = xds[0].dims['corr']
    nrow = xds[0].dims['row']
    # we still have to cater for complex valued data because we cast
    # the weights to complex but we not longer need to factor the
    # weight column into our memory budget
    data_bytes = getattr(xds[0], args.data_column).data.itemsize
    bytes_per_row = max_chan_chunk * ncorr * data_bytes
    memory_per_row = bytes_per_row

    # flags (uint8 or bool)
    memory_per_row += bytes_per_row / 8

    # UVW
    memory_per_row += xds[0].UVW.data.itemsize * 3

    # ANTENNA1/2
    memory_per_row += xds[0].ANTENNA1.data.itemsize * 2

    # TIME
    memory_per_row += xds[0].TIME.data.itemsize

    # data column is not actually read into memory just used to infer
    # dtype and chunking
    columns = (args.data_column,
               args.weight_column,
               args.flag_column,
               'UVW', 'ANTENNA1',
               'ANTENNA2',
               'TIME')

    # flag row
    if 'FLAG_ROW' in xds[0]:
        columns += ('FLAG_ROW',)
        memory_per_row += xds[0].FLAG_ROW.data.itemsize

    # imaging weights
    if args.imaging_weight_column is not None:
        columns += (args.imaging_weight_column,)
        memory_per_row += bytes_per_row / 2

    # Mueller term (complex valued)
    if args.mueller_column is not None:
        columns += (args.mueller_column,)
        memory_per_row += bytes_per_row

    # get max uv coords over all fields
    uvw = []
    u_max = 0.0
    v_max = 0.0
    for ims in args.ms:
        xds = xds_from_ms(ims, columns=('UVW'), chunks={'row': -1})

        for ds in xds:
            uvw = ds.UVW.data
            u_max = da.maximum(u_max, abs(uvw[:, 0]).max())
            v_max = da.maximum(v_max, abs(uvw[:, 1]).max())
            uv_max = da.maximum(u_max, v_max)

    uv_max = uv_max.compute()
    del uvw

    # image size
    cell_N = 1.0 / (2 * uv_max * max_freq / lightspeed)

    if args.cell_size is not None:
        cell_size = args.cell_size
        cell_rad = cell_size * np.pi / 60 / 60 / 180
        if cell_N / cell_rad < 1:
            raise ValueError("Requested cell size too small. "
                             "Super resolution factor = ", cell_N / cell_rad)
        print("Super resolution factor = %f" % (cell_N / cell_rad), file=log)
    else:
        cell_rad = cell_N / args.super_resolution_factor
        cell_size = cell_rad * 60 * 60 * 180 / np.pi
        print("Cell size set to %5.5e arcseconds" % cell_size, file=log)

    if args.nx is None:
        fov = args.field_of_view * 3600
        npix = int(args.psf_oversize * fov / cell_size)
        if npix % 2:
            npix += 1
        nx = npix
        ny = npix
    else:
        nx = args.nx
        ny = args.ny if args.ny is not None else nx

    print("PSF size set to (%i, %i, %i)" % (nband, nx, ny), file=log)

    # get approx image size
    # this is not a conservative estimate when multiple SPW's map to a single
    # imaging band
    pixel_bytes = np.dtype(args.output_type).itemsize
    band_size = nx * ny * pixel_bytes

    if args.host_address is None:
        # full image on single node
        row_chunk = plan_row_chunk(args.mem_limit/args.nworkers, band_size, nrow,
                                   memory_per_row, args.nthreads_per_worker)

    else:
        # single band per node
        row_chunk = plan_row_chunk(args.mem_limit, band_size, nrow,
                                   memory_per_row, args.nthreads_per_worker)

    if args.row_chunks is not None:
        row_chunk = int(args.row_chunks)
        if row_chunk == -1:
            row_chunk = nrow

    print("nrows = %i, row chunks set to %i for a total of %i chunks per node" %
          (nrow, row_chunk, int(np.ceil(nrow / row_chunk))), file=log)

    chunks = {}
    for ims in args.ms:
        chunks[ims] = []  # xds_from_ms expects a list per ds
        for spw in freqs[ims]:
            chunks[ims].append({'row': row_chunk,
                                'chan': chan_chunks[ims][spw]['chan']})

    psfs = []
    radec = None  # assumes we are only imaging field 0 of first MS
    out_datasets = []
    for ims in args.ms:
        xds = xds_from_ms(ims, chunks=chunks[ims], columns=columns)

        # subtables
        ddids = xds_from_table(ims + "::DATA_DESCRIPTION")
        fields = xds_from_table(ims + "::FIELD")
        spws = xds_from_table(ims + "::SPECTRAL_WINDOW")
        pols = xds_from_table(ims + "::POLARIZATION")

        # subtable data
        ddids = dask.compute(ddids)[0]
        fields = dask.compute(fields)[0]
        spws = dask.compute(spws)[0]
        pols = dask.compute(pols)[0]

        for ds in xds:
            field = fields[ds.FIELD_ID]

            # check fields match
            if radec is None:
                radec = field.PHASE_DIR.data.squeeze()

            if not np.array_equal(radec, field.PHASE_DIR.data.squeeze()):
                continue

            # this is not correct, need to use spw
            spw = ds.DATA_DESC_ID

            uvw = clone(ds.UVW.data)

            data_type = getattr(ds, args.data_column).data.dtype
            data_shape = getattr(ds, args.data_column).data.shape
            data_chunks = getattr(ds, args.data_column).data.chunks

            weights = getattr(ds, args.weight_column).data
            if len(weights.shape) < 3:
                weights = da.broadcast_to(
                    weights[:, None, :], data_shape, chunks=data_chunks)

            if args.imaging_weight_column is not None:
                imaging_weights = getattr(
                    ds, args.imaging_weight_column).data
                if len(imaging_weights.shape) < 3:
                    imaging_weights = da.broadcast_to(
                        imaging_weights[:, None, :], data_shape, chunks=data_chunks)

                weightsxx = imaging_weights[:, :, 0] * weights[:, :, 0]
                weightsyy = imaging_weights[:, :, -1] * weights[:, :, -1]
            else:
                weightsxx = weights[:, :, 0]
                weightsyy = weights[:, :, -1]

            # apply mueller term
            if args.mueller_column is not None:
                mueller = getattr(ds, args.mueller_column).data
                weightsxx *= da.absolute(mueller[:, :, 0])**2
                weightsyy *= da.absolute(mueller[:, :, -1])**2

            # weighted sum corr to Stokes I
            weights = weightsxx + weightsyy

            # MS may contain auto-correlations
            if 'FLAG_ROW' in xds[0]:
                frow = ds.FLAG_ROW.data | (ds.ANTENNA1.data == ds.ANTENNA2.data)
            else:
                frow = (ds.ANTENNA1.data == ds.ANTENNA2.data)

            # only keep data where both corrs are unflagged
            flag = getattr(ds, args.flag_column).data
            flagxx = flag[:, :, 0]
            flagyy = flag[:, :, -1]
            # ducc0 uses uint8 mask not flag
            mask = ~ da.logical_or((flagxx | flagyy), frow[:, None])

            psf = vis2im(uvw,
                         freqs[ims][spw],
                         weights.astype(data_type),
                         freq_bin_idx[ims][spw],
                         freq_bin_counts[ims][spw],
                         nx,
                         ny,
                         cell_rad,
                         flag=mask.astype(np.uint8),
                         nthreads=args.nvthreads,
                         epsilon=args.epsilon,
                         do_wstacking=args.wstack,
                         double_accum=args.double_accum)

            psfs.append(psf)

            data_vars = {
                'FIELD_ID':(('row',), da.full_like(ds.TIME.data,
                            ds.FIELD_ID, chunks=args.row_out_chunk)),
                'DATA_DESC_ID':(('row',), da.full_like(ds.TIME.data,
                            ds.DATA_DESC_ID, chunks=args.row_out_chunk)),
                'WEIGHT':(('row', 'chan'), weights.rechunk({0:args.row_out_chunk})),  # why no 'f4'?
                'UVW':(('row', 'uvw'), uvw.rechunk({0:args.row_out_chunk}))
            }

            coords = {
                'chan': (('chan',), freqs[ims][spw])
            }

            out_ds = Dataset(data_vars, coords)

            out_datasets.append(out_ds)

    writes = xds_to_zarr(out_datasets, args.output_filename + '.zarr', columns='ALL')

    # dask.visualize(writes, filename=args.output_filename + '_psf_writes_graph.pdf', optimize_graph=False)
    # dask.visualize(psfs, filename=args.output_filename + '_psf_graph.pdf', optimize_graph=False)

    if not args.mock:
        # psfs = dask.compute(psfs, writes, optimize_graph=False)[0]
        # with performance_report(filename=args.output_filename + '_psf_per.html'):
        psfs = dask.compute(psfs, writes, optimize_graph=False)[0]

        psf = stitch_images(psfs, nband, band_mapping)

        hdr = set_wcs(cell_size / 3600, cell_size / 3600, nx, ny, radec, freq_out)
        save_fits(args.output_filename + '_psf.fits', psf, hdr,
                dtype=args.output_type)

        psf_mfs = np.sum(psf, axis=0)
        wsum = psf_mfs.max()
        psf_mfs /= wsum

        hdr_mfs = set_wcs(cell_size / 3600, cell_size / 3600, nx, ny, radec, np.mean(freq_out))
        save_fits(args.output_filename + '_psf_mfs.fits', psf_mfs, hdr_mfs,
                  dtype=args.output_type)

    print("All done here.", file=log)


import cv2
import numpy as np
from numpy import exp, abs, angle

from annotations import draw_convex_hull, annotate_landmarks
from constants import OVERLAY_POINTS, FEATHER_AMOUNT, SCALE_FACTOR, FOREHEAD_POINTS, LEFT_BROW_POINTS, \
    RIGHT_BROW_POINTS, \
    SCREEN
from utils import resize_ff, running_mean

def prepare_image(frame):
    # Prepare capture for face detection
    im_resized = cv2.resize(frame, (0, 0), fx=SCALE_FACTOR, fy=SCALE_FACTOR, interpolation=cv2.INTER_AREA)
    cv2.flip(im_resized, 1, dst=im_resized)

    im_gray = cv2.cvtColor(im_resized, cv2.COLOR_BGR2GRAY)
    im_gray = cv2.equalizeHist(im_gray)
    return im_resized[0:], im_gray[0:]

def face_mask(im, landmarks, roi):
    # Creates a masked image where only the face is shown
    # Uses ROI as bounds of frame
    zero_mat = np.zeros(im.shape[:2], dtype=np.float64)
    landmarks_org = landmarks.copy()
    for group in OVERLAY_POINTS:
        draw_convex_hull(zero_mat, landmarks[group], color=1)

    # for x in LEFT_BROW_POINTS:
    #     for _ in group:
    #         landmarks[x, 0] -= 1.3  # extend ROI to include foreheads
    #         draw_convex_hull(zero_mat, landmarks[group], color=1)
    #         landmarks = landmarks_org
    #
    # for x in RIGHT_BROW_POINTS:
    #     for _ in group:
    #         landmarks[x, 0] += 1.3  # extend ROI to include foreheads
    #         draw_convex_hull(zero_mat, landmarks[group], color=1)
    #         landmarks = landmarks_org

    for group in FOREHEAD_POINTS:
        for _ in group:
            landmarks[_, 1] -= 30  # extend ROI to include foreheads
            draw_convex_hull(zero_mat, landmarks[group], color=1)
            landmarks = landmarks_org

    zero_mat = np.array([zero_mat, zero_mat, zero_mat]).transpose(
        (1, 2, 0))  # creates an image array out of zero_mat [rows, cols, ch]

    zero_mat = (zero_mat > 0) * 1.0  # make hull points = 1

    # zero_mat = cv2.GaussianBlur(zero_mat, (FEATHER_AMOUNT, FEATHER_AMOUNT), 0)
    # zero_mat = cv2.GaussianBlur(zero_mat, (FEATHER_AMOUNT, FEATHER_AMOUNT), 0)
    zero_mat = cv2.blur(zero_mat, (FEATHER_AMOUNT, FEATHER_AMOUNT))
    zero_mat = cv2.blur(zero_mat, (FEATHER_AMOUNT, FEATHER_AMOUNT))  # this is faster than gaussian

    # apply the image mask to original frame
    im = im * zero_mat
    im = np.uint8(im)

    # resize and crop to ROI
    x = roi[0][0]
    y = roi[0][1]
    w = roi[0][2]
    h = roi[0][3]
    # im = im[y:y + h, x:x + w]  # crop image to extent of ROI

    rows, cols, ch = im.shape
    pts1 = np.float32([[x, y], [x + w, y + h], [x + w, y]])  # ROI rectangle
    pts2 = np.float32([[0, 0], [im.shape[1], im.shape[0]], [im.shape[1], 0]])  # frame bounds
    m = cv2.getAffineTransform(pts1, pts2)
    im = cv2.warpAffine(im, m, (cols, rows))

    # To shrink an image, it will generally look best with CV_INTER_AREA
    # interpolation, whereas to enlarge an image, it will generally look best
    im = cv2.resize(im, (720, 720), interpolation=cv2.INTER_LINEAR)

    # increase brightness - unused for now since it creates white regions
    # im = cv2.cvtColor(im, cv2.COLOR_BGR2HSV)  # convert it to hsv
    # h, s, v = cv2.split(im)
    # v += 255
    # im = cv2.merge((h, s, v))
    # im = cv2.cvtColor(im, cv2.COLOR_HSV2BGR)

    return im

def mask2(im, landmarks):
    # Use landmarks as bounds for frame
    # landmarks[x,y] = xth element of landmarks, yth dimension (where 0=x,1=y)
    # Top is calculated by getting the smallest value of the landmark values since origin is top left.
    # 30% is added to this value to include forehead since dlib stops at eyebrows.

    left_x = landmarks[0, 0]
    right_x = landmarks[15, 0]
    top_y = landmarks[landmarks[:, 1].argmin(), 1]

    extend_val = 2 * int(landmarks[30, 1] - landmarks[27, 1])  # 2 * nose length

    # top_extended_y = int(top_y - 0.25 * (landmarks[landmarks[:,1].argmax(), 1] - landmarks[landmarks[:,1].argmin(), 1]))
    top_extended_y = int(top_y) - extend_val
    if top_extended_y <= 0:
        top_extended_y = 0
    bottom_y = landmarks[landmarks[:, 1].argmax(), 1]

    # Crop image and resize it while maintaining aspect ratio

    im = im[top_extended_y:bottom_y, left_x:right_x]  # crop image

    # # resize
    # rows, cols, ch = im.shape
    # aspect = SCREEN[0] / rows # make height match fullscreen height
    # dim = (int(aspect * cols), SCREEN[1])
    # im = cv2.resize(im, dim, interpolation=cv2.INTER_LINEAR)

    # Vignette - this needs work.
    # rows, cols, ch = im.shape
    # a = cv2.getGaussianKernel(cols, 300)
    # b = cv2.getGaussianKernel(rows, 300)
    # c = b * a.T
    # d = c / c.max()
    # im = im * d

    return im

def mask3(im, landmarks, roi):
    # Creates a masked image where only the face is shown
    zero_mat = np.zeros(im.shape[:2], dtype=np.float64)
    landmarks_org = landmarks.copy()
    for group in OVERLAY_POINTS:
        draw_convex_hull(zero_mat, landmarks[group], color=1)

    for group in FOREHEAD_POINTS:
        for _ in group:
            landmarks[_, 1] -= 30  # extend ROI to include foreheads
            draw_convex_hull(zero_mat, landmarks[group], color=1)
            landmarks = landmarks_org

    zero_mat = np.array([zero_mat, zero_mat, zero_mat]).transpose(
        (1, 2, 0))  # creates an image array out of zero_mat [rows, cols, ch]

    zero_mat = (zero_mat > 0) * 1.0  # make hull points = 1

    # zero_mat = cv2.GaussianBlur(zero_mat, (FEATHER_AMOUNT, FEATHER_AMOUNT), 0)
    # zero_mat = cv2.GaussianBlur(zero_mat, (FEATHER_AMOUNT, FEATHER_AMOUNT), 0)
    zero_mat = cv2.blur(zero_mat, (FEATHER_AMOUNT, FEATHER_AMOUNT))
    zero_mat = cv2.blur(zero_mat, (FEATHER_AMOUNT, FEATHER_AMOUNT))  # this is faster than gaussian

    # apply the image mask to original frame
    im = im * zero_mat

    im = np.uint8(im)

    # resize and crop, need to seperate colour components to do this
    im = cv2.cvtColor(im, cv2.COLOR_BGR2HSV)  # convert it to hsv
    h, s, v = cv2.split(im)

    ytop = np.nonzero(v)[0].min()
    ybottom = np.nonzero(v)[0].max()
    xleft = np.nonzero(v)[1].min()
    xright = np.nonzero(v)[1].max()

    im = cv2.merge((h, s, v))
    im = cv2.cvtColor(im, cv2.COLOR_HSV2BGR)

    im = im[ytop:ybottom, xleft:xright]  # crop image

    return im

def cart2pol(x, y):
    rho = np.sqrt(x ** 2 + y ** 2)
    phi = np.arctan2(y, x)
    return rho, phi

def pol2cart(rho, phi):
    x = rho * np.cos(phi)
    y = rho * np.sin(phi)
    return x, y

def barrel_distort(im, amplitude):
    # Attempting to apply barrel distortion. Adapted from https://au.mathworks.com/help/images/examples/creating-a-gallery-of-transformed-images.html (Image 10)
    # And this stackoverflow.com/questions/6199636/formulas-for-barrel-pincushion-distortion
    # Also this https://gist.github.com/tanmaykm/4395121

    # get dimensions of image and find midpoint
    rows, cols, ch = im.shape

    midpoint_x = cols / 2
    midpoint_y = rows / 2

    # setup vars
    correction_factor = midpoint_y / (midpoint_y - amplitude * midpoint_y ** 3)
    x_range = range(1, cols + 1)
    y_range = range(1, rows + 1)
    xi, yi = np.meshgrid(x_range, y_range)  # creates a 2D grid the same size as the input image

    xt = np.ravel(xi) - midpoint_x  # stacks the x values into a single column and subtracts the midpoint
    yt = np.ravel(yi) - midpoint_y  # same but for y values

    r, phi = cart2pol(xt, yt)  # convert to polar coordinates so we can find the radius from center
    s = r - amplitude * r ** 3  # apply transformation function
    ut, vt = pol2cart(s, phi)  # convert back to cartesian

    u = np.reshape(ut * correction_factor,
                   (rows, cols)) + midpoint_x  # change it back to a grid and add the midpoint again
    v = np.reshape(vt * correction_factor, (rows, cols)) + midpoint_y
    u_32 = u.astype('float32')
    v_32 = v.astype('float32')

    im_distorted = cv2.remap(im, u_32, v_32, interpolation=cv2.INTER_CUBIC)  # borderMode=cv2.BORDER_CONSTANT

    return im_distorted