import warnings
import torch
import torch.nn as nn
import torch.nn.functional as F
from mmcv.cnn import ConvModule, xavier_init
from mmcv.runner import auto_fp16
import numpy as np
from timm.models.layers import DropPath, to_2tuple, trunc_normal_

from ..builder import NECKS

class AdaptiveAngleConv(nn.Module):
    def __init__(self, in_channel, out_channel, kernel_size=3, stride=1, padding=1, angle_list=[0]):
        super(AdaptiveAngleConv, self).__init__()
        self.kernel_size = kernel_size
        self.stride = stride
        self.padding = padding
        self.angle_list = angle_list
        self.branches = len(angle_list)
        self.baseline_conv = nn.Conv2d(in_channel, out_channel, (kernel_size, kernel_size), (stride, stride), (padding, padding))

    def forward(self, x):
        y = self.baseline_conv(x)
        y_45 = F.conv2d(x, self._rotate_kernel(self.baseline_conv.weight, self.angle_list[1]), bias=self.baseline_conv.bias, stride=self.stride, padding=self.padding)
        y_90 = F.conv2d(x, self._rotate_kernel(self.baseline_conv.weight, self.angle_list[2]), bias=self.baseline_conv.bias, stride=self.stride, padding=self.padding)
        y_135 = F.conv2d(x, self._rotate_kernel(self.baseline_conv.weight, self.angle_list[3]), bias=self.baseline_conv.bias, stride=self.stride, padding=self.padding)
        y_180 = F.conv2d(x, self._rotate_kernel(self.baseline_conv.weight, self.angle_list[4]), bias=self.baseline_conv.bias, stride=self.stride, padding=self.padding)
        y_225 = F.conv2d(x, self._rotate_kernel(self.baseline_conv.weight, self.angle_list[5]), bias=self.baseline_conv.bias, stride=self.stride, padding=self.padding)
        y_270 = F.conv2d(x, self._rotate_kernel(self.baseline_conv.weight, self.angle_list[6]), bias=self.baseline_conv.bias, stride=self.stride, padding=self.padding)
        y_315 = F.conv2d(x, self._rotate_kernel(self.baseline_conv.weight, self.angle_list[7]), bias=self.baseline_conv.bias, stride=self.stride, padding=self.padding)
        return [y, y_45, y_90, y_135, y_180, y_225, y_270, y_315]

    def _rotate_kernel(self, original_kernel, angle):  # only 3*3 kernel
        n = angle // 45
        orig_tran = original_kernel.permute(2, 3, 0, 1)
        new_kernel = torch.zeros_like(orig_tran)
        if n == 1 :  # 45 degree
            new_kernel[0][1:][:][:] = orig_tran[0][:2][:][:]
            new_kernel[1][2][:][:] = orig_tran[0][2][:][:]
            new_kernel[2][2][:][:] = orig_tran[1][2][:][:]
            new_kernel[2][:2][:][:] = orig_tran[2][1:][:][:]
            new_kernel[0][0][:][:] = orig_tran[1][0][:][:]
            new_kernel[1][0][:][:] = orig_tran[2][0][:][:]
            new_kernel[1][1][:][:] = orig_tran[1][1][:][:]
            new_kernel = new_kernel.permute(2, 3, 0, 1)

        if n == 2:  # 90 degree
            l = len(orig_tran)
            for i in range(l):
                for j in range(l):
                    new_kernel[j][l - 1 - i][:][:] = orig_tran[i][j][:][:]
            new_kernel = new_kernel.permute(2, 3, 0, 1)

        if n == 3:  #  135 degree
            new_kernel[1][2][:][:] = orig_tran[0][0][:][:]
            new_kernel[2][2][:][:] = orig_tran[0][1][:][:]
            new_kernel[2][0][:][:] = orig_tran[1][2][:][:]
            new_kernel[2][1][:][:] = orig_tran[0][2][:][:]
            new_kernel[0][0][:][:] = orig_tran[2][1][:][:]
            new_kernel[1][0][:][:] = orig_tran[2][2][:][:]
            new_kernel[0][1][:][:] = orig_tran[2][0][:][:]
            new_kernel[0][2][:][:] = orig_tran[1][0][:][:]
            new_kernel[1][1][:][:] = orig_tran[1][1][:][:]
            new_kernel = new_kernel.permute(2, 3, 0, 1)

        if n == 4:  # 180 degree
            l = len(orig_tran)
            for i in range(l):
                for j in range(l):
                    new_kernel[i][j][:][:] = orig_tran[l - 1 - i][l - 1 - j][:][:]
            new_kernel = new_kernel.permute(2, 3, 0, 1)

        if n == 5:
            new_kernel[2][1][:][:] = orig_tran[0][0][:][:]
            new_kernel[2][0][:][:] = orig_tran[0][1][:][:]
            new_kernel[1][0][:][:] = orig_tran[0][2][:][:]
            new_kernel[0][0][:][:] = orig_tran[1][2][:][:]
            new_kernel[0][1][:][:] = orig_tran[2][2][:][:]
            new_kernel[0][2][:][:] = orig_tran[2][1][:][:]
            new_kernel[1][2][:][:] = orig_tran[2][0][:][:]
            new_kernel[2][2][:][:] = orig_tran[1][0][:][:]
            new_kernel[1][1][:][:] = orig_tran[1][1][:][:]
            new_kernel = new_kernel.permute(2, 3, 0, 1)

        if n == 6:  # 270 degree
            l = len(orig_tran)
            for i in range(l):
                for j in range(l):
                    new_kernel[l - 1 - j][i][:][:] = orig_tran[i][j][:][:]
            new_kernel = new_kernel.permute(2, 3, 0, 1)

        if n == 7:  # 315 degree
            new_kernel[1][0][:][:] = orig_tran[0][0][:][:]
            new_kernel[0][0][:][:] = orig_tran[0][1][:][:]
            new_kernel[0][1][:][:] = orig_tran[0][2][:][:]
            new_kernel[0][2][:][:] = orig_tran[1][2][:][:]
            new_kernel[1][2][:][:] = orig_tran[2][2][:][:]
            new_kernel[2][2][:][:] = orig_tran[2][1][:][:]
            new_kernel[2][1][:][:] = orig_tran[2][0][:][:]
            new_kernel[2][0][:][:] = orig_tran[1][0][:][:]
            new_kernel[1][1][:][:] = orig_tran[1][1][:][:]
            new_kernel = new_kernel.permute(2, 3, 0, 1)

        return new_kernel

class SKConv(nn.Module):
    def __init__(self, features, M, r, L=32):
        """ Constructor
        Args:
            features: input channel dimensionality.
            WH: input spatial dimensionality, used for GAP kernel size.
            M: the number of branchs.
            G: num of convolution groups.
            r: the radio for compute d, the length of z.
            stride: stride, default 1.
            L: the minimum dim of the vector z in paper, default 32.
        """
        super(SKConv, self).__init__()
        d = max(int(features / r), L)
        self.features = features
        self.bn_relus = nn.ModuleList()
        for i in range(M-1):  # 45, 90, 135, 180
            self.bn_relus.append(nn.Sequential(
                nn.BatchNorm2d(features),
                nn.ReLU(inplace=False)
            ))
        # self.gap = nn.AvgPool2d(int(WH/stride))
        self.fc = nn.Linear(features, d)
        self.fcs = nn.ModuleList([])
        for i in range(M-1):
            self.fcs.append(
                nn.Linear(d, features)
            )
        self.softmax = nn.Softmax(dim=1)
        self.branch = M-1

    def forward(self, X):
        feas = []
        for i in range(self.branch):
            feas.append(self.bn_relus[i](X[i+1]).unsqueeze_(dim=1))
        feas = torch.cat(feas, dim=1)
        fea_U = torch.sum(feas, dim=1)
        fea_s = fea_U.mean(-1).mean(-1)
        fea_z = self.fc(fea_s)
        for i, fc in enumerate(self.fcs):
            vector = fc(fea_z).unsqueeze_(dim=1)
            if i == 0:
                attention_vectors = vector
            else:
                attention_vectors = torch.cat([attention_vectors, vector], dim=1)
        attention_vectors = self.softmax(attention_vectors)
        attention_vectors = attention_vectors.unsqueeze(-1).unsqueeze(-1)
        fea_v = (feas * attention_vectors).sum(dim=1)
        fea_v += X[0]  # residual structure
        return fea_v

@NECKS.register_module()
class FPN_AAR1(nn.Module):
    def __init__(self,
                 in_channels,
                 out_channels,
                 num_outs,
                 start_level=0,
                 end_level=-1,
                 add_extra_convs=False,
                 extra_convs_on_inputs=True,
                 relu_before_extra_convs=False,
                 no_norm_on_lateral=False,
                 conv_cfg=None,
                 norm_cfg=None,
                 act_cfg=None,
                 upsample_cfg=dict(mode='nearest')):
        super(FPN_AAR1, self).__init__()
        assert isinstance(in_channels, list)
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.num_ins = len(in_channels)
        self.num_outs = num_outs
        self.relu_before_extra_convs = relu_before_extra_convs
        self.no_norm_on_lateral = no_norm_on_lateral
        self.fp16_enabled = False
        self.upsample_cfg = upsample_cfg.copy()

        if end_level == -1:
            self.backbone_end_level = self.num_ins
            assert num_outs >= self.num_ins - start_level
        else:
            # if end_level < inputs, no extra level is allowed
            self.backbone_end_level = end_level
            assert end_level <= len(in_channels)
            assert num_outs == end_level - start_level
        self.start_level = start_level
        self.end_level = end_level
        self.add_extra_convs = add_extra_convs
        assert isinstance(add_extra_convs, (str, bool))
        if isinstance(add_extra_convs, str):
            # Extra_convs_source choices: 'on_input', 'on_lateral', 'on_output'
            assert add_extra_convs in ('on_input', 'on_lateral', 'on_output')
        elif add_extra_convs:  # True
            if extra_convs_on_inputs:
                # TODO: deprecate `extra_convs_on_inputs`
                warnings.simplefilter('once')
                warnings.warn(
                    '"extra_convs_on_inputs" will be deprecated in v2.9.0,'
                    'Please use "add_extra_convs"', DeprecationWarning)
                self.add_extra_convs = 'on_input'
            else:
                self.add_extra_convs = 'on_output'

        self.lateral_convs = nn.ModuleList()
        self.fpn_rotate_convs = nn.ModuleList()  # rotate_convolution , M angle branch
        self.fusions = nn.ModuleList()  # add fusion module SK-Net

        for i in range(self.start_level, self.backbone_end_level):
            l_conv = ConvModule(
                in_channels[i],
                out_channels,
                1,
                conv_cfg=conv_cfg,
                norm_cfg=norm_cfg if not self.no_norm_on_lateral else None,
                act_cfg=act_cfg,
                inplace=False)

            fpn_rotate_conv = AdaptiveAngleConv(
                out_channels,
                out_channels,
                3,
                1,
                1,
                angle_list=[0, 45, 90, 135, 180, 225, 270, 315]
            )
            fusion = SKConv(256, 8, 2, 32)

            self.lateral_convs.append(l_conv)
            self.fpn_rotate_convs.append(fpn_rotate_conv)
            self.fusions.append(fusion)

        # add extra conv layers (e.g., RetinaNet)
        extra_levels = num_outs - self.backbone_end_level + self.start_level
        if self.add_extra_convs and extra_levels >= 1:
            for i in range(extra_levels):
                if i == 0 and self.add_extra_convs == 'on_input':
                    in_channels = self.in_channels[self.backbone_end_level - 1]
                else:
                    in_channels = out_channels

                extra_fpn_conv = ConvModule(
                    in_channels,
                    out_channels,
                    3,
                    stride=2,
                    padding=1,
                    conv_cfg=conv_cfg,
                    norm_cfg=norm_cfg,
                    act_cfg=act_cfg,
                    inplace=False)
                self.fpn_rotate_convs.append(extra_fpn_conv)

    # default init_weights for conv(msra) and norm in ConvModule
    def init_weights(self):
        """Initialize the weights of FPN module."""
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                xavier_init(m, distribution='uniform')
            elif isinstance(m, nn.LayerNorm):
                nn.init.constant_(m.bias, 0)
                nn.init.constant_(m.weight, 1.0)
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1.0)
                m.bias.data.zero_()

    @auto_fp16()
    def forward(self, inputs):
        """Forward function."""
        assert len(inputs) == len(self.in_channels)

        # build laterals
        laterals = [
            lateral_conv(inputs[i + self.start_level])
            for i, lateral_conv in enumerate(self.lateral_convs)
        ]

        # build top-down path
        used_backbone_levels = len(laterals)
        for i in range(used_backbone_levels - 1, 0, -1):
            # In some cases, fixing `scale factor` (e.g. 2) is preferred, but
            #  it cannot co-exist with `size` in `F.interpolate`.
            if 'scale_factor' in self.upsample_cfg:
                laterals[i - 1] += F.interpolate(laterals[i],
                                                 **self.upsample_cfg)
            else:
                prev_shape = laterals[i - 1].shape[2:]
                laterals[i - 1] += F.interpolate(
                    laterals[i], size=prev_shape, **self.upsample_cfg)

        # build outputs
        # part 1: from original levels
        # convolution: local
        conv_outs = [
            self.fpn_rotate_convs[i](laterals[i]) for i in range(used_backbone_levels)
        ]
        # local-global fusion
        outs = [
            self.fusions[i](conv_outs[i]) for i in range(used_backbone_levels)
        ]

        # part 2: add extra levels
        if self.num_outs > len(outs):
            # use max pool to get more levels on top of outputs
            # (e.g., Faster R-CNN, Mask R-CNN)
            if not self.add_extra_convs:
                for i in range(self.num_outs - used_backbone_levels):
                    outs.append(F.max_pool2d(outs[-1], 1, stride=2))
            # add conv layers on top of original feature maps (RetinaNet)
            else:
                if self.add_extra_convs == 'on_input':
                    extra_source = inputs[self.backbone_end_level - 1]
                elif self.add_extra_convs == 'on_lateral':
                    extra_source = laterals[-1]
                elif self.add_extra_convs == 'on_output':
                    extra_source = outs[-1]
                else:
                    raise NotImplementedError
                outs.append(self.fpn_rotate_convs[used_backbone_levels](extra_source))
                for i in range(used_backbone_levels + 1, self.num_outs):
                    if self.relu_before_extra_convs:
                        outs.append(self.fpn_rotate_convs[i](F.relu(outs[-1])))
                    else:
                        outs.append(self.fpn_rotate_convs[i](outs[-1]))
        return tuple(outs)


"""
This module has an implementation of A2C, which is described in
- OpenAI Baselines: ACKTR & A2C
  - URL: https://openai.com/blog/baselines-acktr-a2c/

A2C is a synchronous version of A3C, which is described in
- http://proceedings.mlr.press/v48/mniha16.pdf
"""
from typing import Optional, Tuple

import numpy as np
import torch
from torch import nn

from ..config import Config
from ..lib.rollout import RolloutStorage
from ..net import ActorCriticNet, Policy, RnnState
from ..prelude import Action, Array, State
from .base import A2CLikeAgent, Netout

class A2CAgent(A2CLikeAgent[State]):
    SAVED_MEMBERS = "net", "lr_cooler"

    def __init__(self, config: Config) -> None:
        super().__init__(config)
        self.storage: RolloutStorage[State] = RolloutStorage(
            config.nsteps, config.nworkers, config.device
        )
        self.net: ActorCriticNet = config.net("actor-critic")  # type: ignore
        self.optimizer = config.optimizer(self.net.parameters())
        self.lr_cooler = config.lr_cooler(self.optimizer.param_groups[0]["lr"])
        self.eval_rnns = self.rnn_init()

    def set_mode(self, train: bool = True) -> None:
        self.net.train(mode=train)

    def rnn_init(self) -> RnnState:
        return self.net.recurrent_body.initial_state(
            self.config.nworkers, self.config.device
        )

    def _reset(self, initial_states: Array[State]) -> None:
        self.storage.set_initial_state(initial_states, self.rnn_init())

    def eval_reset(self) -> None:
        self.eval_rnns.fill_(0.0)

    def eval_action(self, state: Array, net_outputs: Optional[Netout] = None) -> Action:
        if state.ndim == len(self.net.state_dim):
            # treat as batch_size == 1
            state = np.stack([state])
        with torch.no_grad():
            policy, self.eval_rnns[0] = self.net.policy(
                state, self.eval_rnns[0].unsqueeze()
            )
        if net_outputs is not None:
            net_outputs["policy"] = policy
        return policy.eval_action(self.config.eval_deterministic)

    def eval_action_parallel(self, states: Array) -> Array[Action]:
        with torch.no_grad():
            policy, self.eval_rnns = self.net.policy(states, self.eval_rnns)
        return policy.eval_action(self.config.eval_deterministic)

    def _network_in(self, states: Array[State]) -> Tuple[Array, RnnState, torch.Tensor]:
        return (
            self.penv.extract(states),
            self.storage.rnn_states[-1],
            self.storage.masks[-1],
        )

    @torch.no_grad()
    def actions(self, states: Array[State]) -> Tuple[Array[Action], dict]:
        policy, value, rnns = self.net(*self._network_in(states))
        actions = policy.action().squeeze().cpu().numpy()
        return actions, dict(rnn_state=rnns, policy=policy, value=value)

    def _pre_backward(self, _policy: Policy, _value: torch.Tensor) -> None:
        pass

    def _step_optimizer(self) -> None:
        nn.utils.clip_grad_norm_(self.net.parameters(), self.config.grad_clip)
        self.optimizer.step()
        self.lr_cooler.lr_decay(self.optimizer)

    def train(self, last_states: Array[State]) -> None:
        with torch.no_grad():
            next_value = self.net.value(*self._network_in(last_states))
        if self.config.use_gae:
            self.storage.set_gae_returns(
                next_value, self.config.discount_factor, self.config.gae_lambda
            )
        else:
            self.storage.set_ac_returns(next_value, self.config.discount_factor)
        policy, value, _ = self.net(
            self.storage.batch_states(self.penv),
            self.storage.rnn_states[0],
            self.storage.batch_masks(),
        )

        advantage = self.storage.returns[:-1].flatten() - value
        actions = self.storage.batch_actions()
        policy_loss = -(policy.log_prob(actions) * advantage.detach()).mean()
        value_loss = advantage.pow(2).mean()
        entropy_loss = policy.entropy().mean()
        self._pre_backward(policy, value)
        self.optimizer.zero_grad()
        (
            policy_loss
            + self.config.value_loss_weight * 0.5 * value_loss
            - self.config.entropy_weight * entropy_loss
        ).backward()
        self._step_optimizer()
        self.network_log(
            policy_loss=policy_loss.item(),
            value_loss=value_loss.item(),
            entropy_loss=entropy_loss.item(),
        )
        self.storage.reset()


import numpy as np
import pandas as pd
from pydub import AudioSegment
from subprocess import call, check_call, CalledProcessError
import threading

import os
from datetime import datetime
import tensorflow as tf

from sklearn.preprocessing import LabelBinarizer

explosion_sounds = [
'Fireworks',
'Burst, pop',
'Eruption',
'Gunshot, gunfire',
'Explosion',
'Boom'
]

motor_sounds = [
'Chainsaw',
'Medium engine (mid frequency)',
'Light engine (high frequency)',
'Heavy engine (low frequency)',
'Engine starting',
'Engine',
'Motor vehicle (road)',
'Vehicle'
]

wood_sounds = [
'Wood',
'Chop',
'Splinter',
'Crack'
]

human_sounds = [
'Chatter',
'Conversation'
]

nature_sounds = [
"Silence",
"Stream",
"Wind noise (microphone)",
"Wind",
"Rustling leaves",
"Howl",
"Raindrop",
"Rain on surface",
"Rain",
"Thunderstorm",
"Thunder",
"Fire",
"Crackle"
]

ambient_sounds = nature_sounds[:-2]
impact_sounds = explosion_sounds + motor_sounds + \
                wood_sounds + human_sounds + \
                nature_sounds[-2:]

def get_csv_data():
    label_names = pd.read_csv("data/audioset/class_labels_indices.csv")

    balanced_train = pd.read_csv("data/audioset/balanced_train_segments.csv",quotechar='"', skipinitialspace=True, skiprows=2)
    balanced_train.columns = [balanced_train.columns[0][2:]] + balanced_train.columns[1:].values.tolist()

    unbalanced_train = pd.read_csv("data/audioset/unbalanced_train_segments.csv",quotechar='"', skipinitialspace=True, skiprows=2)
    unbalanced_train.columns = [unbalanced_train.columns[0][2:]] + unbalanced_train.columns[1:].values.tolist()

    train = pd.concat([unbalanced_train, balanced_train], axis=0, ignore_index=True)

    # Filter out sounds we don't care about
    forest_sounds = pd.Series(ambient_sounds + impact_sounds, name='display_name')
    forest_sounds = pd.DataFrame(forest_sounds).merge(label_names, how='inner', left_on='display_name', right_on='display_name')

    # Binarizer for labels of interest.  We add ;None' as a placeholder for sounds that
    # are NOT of interest.
    name_bin = LabelBinarizer()
    name_bin.fit(forest_sounds.mid.values.tolist() + ['None'])

    # Binarize labels for dataset
    # There are multiple labels per row, so we have to split them and do a binarization
    # on each column.  We then aggregate all the one-hot variables to get our final encoding.
    train_labels_split = train.positive_labels.str.split(',', expand=True)
    train_labels_split.fillna('None', inplace=True)
    train_label_binarized = name_bin.transform(train_labels_split[train_labels_split.columns[0]])
    for column in train_labels_split.columns:
        train_label_binarized |= name_bin.transform(train_labels_split[column])
    train_label_binarized = pd.DataFrame(train_label_binarized, columns = name_bin.classes_)
    del train_label_binarized['None']

    # Remove rows for uninteresting sounds
    train = train.loc[train_label_binarized.sum(axis=1) > 0]
    train.index = range(train.shape[0])
    train_label_binarized = train_label_binarized.loc[train_label_binarized.sum(axis=1) > 0]
    train_label_binarized.index = range(train_label_binarized.shape[0])

    # Translate mid to display name
    new_column_names = []
    for column in train_label_binarized.columns:
        new_column_names += [label_names.loc[label_names.mid == column].iloc[0].display_name]
    train_label_binarized.columns = new_column_names

    return train, train_label_binarized

def download_clip(YTID, start_seconds, end_seconds):
    url = "https://www.youtube.com/watch?v=" + YTID
    target_file = 'sounds/audioset/' + YTID + '-' + str(start_seconds) + '-' + str(end_seconds)+".wav"

    # No need to download audio file that already has been downloaded
    if os.path.isfile(target_file):
        return

    tmp_filename = 'sounds/tmp_clip_' + YTID + '-' + str(start_seconds) + '-' + str(end_seconds)
    tmp_filename_w_extension = tmp_filename +'.wav'

    try:
        check_call(['youtube-dl', url, '--audio-format', 'wav', '-x', '-o', tmp_filename +'.%(ext)s'])
    except CalledProcessError:
        # do nothing
        print "Exception CalledProcessError!"
        return

    try:
        aud_seg = AudioSegment.from_wav(tmp_filename_w_extension)[start_seconds*1000:end_seconds*1000]

        aud_seg.export(target_file, format="wav")
    except:
        print "Error while reading file!"

    os.remove(tmp_filename_w_extension)

def download_data():

    df, labels_binarized = get_csv_data()

    if not os.path.exists('sounds'):
        os.makedirs('sounds')

    if not os.path.exists('sounds/audioset'):
        os.makedirs('sounds/audioset')

    threads = []

    for index in range(df.shape[0]):
        row = df.iloc[index]

        print "Downloading", str(index), "of", df.shape[0], ":", row.YTID

        # download_clip(row.YTID, row.start_seconds, row.end_seconds)
        t = threading.Thread(target=download_clip, args=(row.YTID, row.start_seconds, row.end_seconds))
        threads.append(t)
        t.start()

        if len(threads) > 4:
            threads[0].join()
            threads = threads[1:]

    for t in threads:
        t.join()

# slightly modified from https://stackoverflow.com/questions/42703849/audioset-and-tensorflow-understanding
def read_audio_record(audio_record, output_to_file=None):
    vid_ids = []
    labels = []
    start_time_seconds = [] # in secondes
    end_time_seconds = []
    feat_audio = []
    count = 0

    with tf.device("/cpu:0"):

        config = tf.ConfigProto()
        config.gpu_options.per_process_gpu_memory_fraction = 0.0001

        sess = tf.InteractiveSession(config=config)

        for example in tf.python_io.tf_record_iterator(audio_record):

            tf_example = tf.train.Example.FromString(example)

            vid_ids.append(tf_example.features.feature['video_id'].bytes_list.value[0].decode(encoding='UTF-8'))
            labels.append(list(tf_example.features.feature['labels'].int64_list.value))
            start_time_seconds.append(tf_example.features.feature['start_time_seconds'].float_list.value[0])
            end_time_seconds.append(tf_example.features.feature['end_time_seconds'].float_list.value[0])

            tf_seq_example = tf.train.SequenceExample.FromString(example)
            n_frames = len(tf_seq_example.feature_lists.feature_list['audio_embedding'].feature)

            rgb_frame = []
            audio_frame = []
            # iterate through frames
            for i in range(n_frames):
                audio_frame.append(tf.cast(tf.decode_raw(
                        tf_seq_example.feature_lists.feature_list['audio_embedding'].feature[i].bytes_list.value[0],tf.uint8)
                               ,tf.float32).eval())

            feat_audio.append([])

            feat_audio[count].append(audio_frame)
            count+=1
        sess.close()

    df = pd.DataFrame(zip(vid_ids, labels, start_time_seconds, end_time_seconds, feat_audio),
                      columns=['video_id', 'labels', 'start_time_seconds', 'end_time_seconds', 'features'])
    df['features'] = df.features.apply(lambda x: np.array(x[0]))

    if output_to_file:
        df.to_pickle(audio_record.replace('.tfrecord', '.pkl'))

    return df

def get_recursive_sound_names(designated_sound_names):

    ontology_dict = pd.read_json("data/audioset/ontology.json")

    ontology_dict_from_name = ontology_dict.copy()
    ontology_dict_from_name.index = ontology_dict_from_name['name']
    ontology_dict_from_name = ontology_dict_from_name.T

    ontology_dict_from_id = ontology_dict.copy()
    ontology_dict_from_id.index = ontology_dict_from_id.id
    ontology_dict_from_id = ontology_dict_from_id.T

    del ontology_dict

    designated_sound_ids = map(lambda sound: ontology_dict_from_name[sound]['id'], designated_sound_names)

    def get_ids(id):
        id_list = [id]

        for child_id in ontology_dict_from_id[id].child_ids:
            id_list += get_ids(child_id)

        return id_list

    recursive_sound_ids = []
    for id in designated_sound_ids:
        recursive_sound_ids += get_ids(id)
    recursive_sound_ids = list(set(recursive_sound_ids))

    recursive_sound_names = map(lambda sound_id: ontology_dict_from_id[sound_id]['name'], recursive_sound_ids)

    return recursive_sound_names

def get_all_sound_names():
    return get_recursive_sound_names(ambient_sounds), get_recursive_sound_names(impact_sounds)

def get_data():

    label_names = pd.read_csv("data/audioset/class_labels_indices.csv")

    audio_files = map(lambda x: 'bal_train/'+x, os.listdir('data/audioset/audioset_v1_embeddings/bal_train')) + \
                  map(lambda x: 'unbal_train/'+x, os.listdir('data/audioset/audioset_v1_embeddings/unbal_train'))
    audio_files = filter(lambda x: x.endswith('.tfrecord'), audio_files)
    audio_files.sort()

    path_prefix = 'data/audioset/audioset_v1_embeddings/'

    for audio_record in audio_files:

        if os.path.isfile(path_prefix + audio_record.replace('.tfrecord', '.pkl')):
            continue

        print "Reading", audio_record, "..."

        try:
            pid = os.fork()
        except OSError:
            sys.stderr.write("Could not create a child process\n")
            continue

        if pid == 0:
            read_audio_record(path_prefix + audio_record, True)
            os._exit(0)
        else:
            os.waitpid(pid, 0)

    pickle_files = map(lambda x: 'bal_train/'+x, os.listdir('data/audioset/audioset_v1_embeddings/bal_train')) + \
                  map(lambda x: 'unbal_train/'+x, os.listdir('data/audioset/audioset_v1_embeddings/unbal_train'))
    pickle_files = filter(lambda x: x.endswith('.pkl'), pickle_files)
    pickle_files.sort()

    ambient_sounds, impact_sounds = get_all_sound_names()

    print "Reading pickles..."

    df = []
    for audio_record in pickle_files:

        print "Reading", audio_record, "pickle"

        sub_df = pd.read_pickle(path_prefix + audio_record)

        def check_sounds(x):
            for sound in (impact_sounds+ambient_sounds):
                if sound in x:
                    return True

            return False

        sub_df['labels'] = sub_df['labels'].map(lambda arr: [label_names.iloc[x].display_name for x in arr])
        sub_df = sub_df.loc[sub_df['labels'].map(check_sounds)]
        df += [sub_df]

    df = pd.concat(df, ignore_index=True)

    # Binarize the labels
    name_bin = LabelBinarizer().fit(ambient_sounds + impact_sounds)
    labels_split = df['labels'].apply(pd.Series).fillna('None')
    labels_binarized = name_bin.transform(labels_split[labels_split.columns[0]])
    for column in labels_split.columns:
        labels_binarized |= name_bin.transform(labels_split[column])
    labels_binarized = pd.DataFrame(labels_binarized, columns = name_bin.classes_)

    return df, labels_binarized


from __future__ import print_function
import numpy as np
import cv2
from rcnn.processing.bbox_transform import clip_boxes
from rcnn.processing.generate_anchor import generate_anchors_fpn, anchors_plane
from rcnn.processing.nms import gpu_nms_wrapper, cpu_nms_wrapper
from networks.retinaface_network import RetinaFaceNetwork

class RetinaFace:
    def __init__(self, model_weights, use_gpu_nms=True, nms=0.4, decay4=0.5):
        self.decay4 = decay4
        self.nms_threshold = nms
        self.fpn_keys = []
        self.anchor_cfg = None
        self.preprocess = False
        _ratio = (1.,)
        self._feat_stride_fpn = [32, 16, 8]
        self.anchor_cfg = {
          '32': {'SCALES': (32,16), 'BASE_SIZE': 16, 'RATIOS': _ratio, 'ALLOWED_BORDER': 9999},
          '16': {'SCALES': (8,4), 'BASE_SIZE': 16, 'RATIOS': _ratio, 'ALLOWED_BORDER': 9999},
          '8': {'SCALES': (2,1), 'BASE_SIZE': 16, 'RATIOS': _ratio, 'ALLOWED_BORDER': 9999},
        }
        for s in self._feat_stride_fpn:
            self.fpn_keys.append('stride%s'%s)
        self._anchors_fpn = dict(zip(self.fpn_keys, generate_anchors_fpn(dense_anchor=False, cfg=self.anchor_cfg)))
        for k in self._anchors_fpn:
            v = self._anchors_fpn[k].astype(np.float32)
            self._anchors_fpn[k] = v
        self._num_anchors = dict(zip(self.fpn_keys, [anchors.shape[0] for anchors in self._anchors_fpn.values()]))
        if use_gpu_nms:
            self.nms = gpu_nms_wrapper(self.nms_threshold, 0)
        else:
            self.nms = cpu_nms_wrapper(self.nms_threshold)
        pixel_means=[0.0, 0.0, 0.0]
        pixel_stds=[1.0, 1.0, 1.0]
        pixel_scale = 1.0
        self.pixel_means = np.array(pixel_means, dtype=np.float32)
        self.pixel_stds = np.array(pixel_stds, dtype=np.float32)
        self.pixel_scale = float(pixel_scale)
        self.bbox_stds = [1.0, 1.0, 1.0, 1.0]
        self.scales = [1024, 1980]
        self.model = RetinaFaceNetwork(model_weights).model

    def detect(self, img, threshold=0.5):
        """
        Detect all the faces and landmarks in an image
        :param img: input image
        :param threshold: detection threshold
        :return: tuple faces, landmarks
        """
        proposals_list = []
        scores_list = []
        landmarks_list = []

        im_shape = img.shape
        target_size = self.scales[0]
        max_size = self.scales[1]
        im_size_min = np.min(im_shape[0:2])
        im_size_max = np.max(im_shape[0:2])
        im_scale = float(target_size) / float(im_size_min)
        if np.round(im_scale * im_size_max) > max_size:
            im_scale = float(max_size) / float(im_size_max)
        if im_scale != 1.0:
            img = cv2.resize(img, None, None, fx=im_scale, fy=im_scale, interpolation=cv2.INTER_LINEAR)

        img = img.astype(np.float32)
        im_info = [img.shape[0], img.shape[1]]
        im_tensor = np.zeros((1, 3, img.shape[0], img.shape[1]))
        for i in range(3):
            im_tensor[0, i, :, :] = (img[:, :, 2 - i] / self.pixel_scale - self.pixel_means[2 - i]) / self.pixel_stds[
                2 - i]
        net_out = self.model.predict(im_tensor.transpose(0, 2, 3, 1))
        net_out_reshaped = []
        for elt in net_out:
            net_out_reshaped.append(elt.transpose(0, 3, 1, 2))
        net_out = net_out_reshaped
        sym_idx = 0

        for _idx,s in enumerate(self._feat_stride_fpn):
            _key = 'stride%s'%s
            stride = int(s)
            scores = net_out[sym_idx]
            scores = scores[:, self._num_anchors['stride%s'%s]:, :, :]

            bbox_deltas = net_out[sym_idx + 1]
            height, width = bbox_deltas.shape[2], bbox_deltas.shape[3]

            A = self._num_anchors['stride%s'%s]
            K = height * width
            anchors_fpn = self._anchors_fpn['stride%s'%s]
            anchors = anchors_plane(height, width, stride, anchors_fpn)
            anchors = anchors.reshape((K * A, 4))
            scores = scores.transpose((0, 2, 3, 1)).reshape((-1, 1))

            bbox_deltas = bbox_deltas.transpose((0, 2, 3, 1))
            bbox_pred_len = bbox_deltas.shape[3]//A
            bbox_deltas = bbox_deltas.reshape((-1, bbox_pred_len))
            bbox_deltas[:, 0::4] = bbox_deltas[:,0::4] * self.bbox_stds[0]
            bbox_deltas[:, 1::4] = bbox_deltas[:,1::4] * self.bbox_stds[1]
            bbox_deltas[:, 2::4] = bbox_deltas[:,2::4] * self.bbox_stds[2]
            bbox_deltas[:, 3::4] = bbox_deltas[:,3::4] * self.bbox_stds[3]
            proposals = self.bbox_pred(anchors, bbox_deltas)

            proposals = clip_boxes(proposals, im_info[:2])

            if stride==4 and self.decay4<1.0:
                scores *= self.decay4

            scores_ravel = scores.ravel()
            order = np.where(scores_ravel>=threshold)[0]
            proposals = proposals[order, :]
            scores = scores[order]

            proposals[:, 0:4] /= im_scale
            proposals_list.append(proposals)
            scores_list.append(scores)

            landmark_deltas = net_out[sym_idx + 2]
            landmark_pred_len = landmark_deltas.shape[1]//A
            landmark_deltas = landmark_deltas.transpose((0, 2, 3, 1)).reshape((-1, 5, landmark_pred_len//5))
            landmarks = self.landmark_pred(anchors, landmark_deltas)
            landmarks = landmarks[order, :]

            landmarks[:, :, 0:2] /= im_scale
            landmarks_list.append(landmarks)
            sym_idx += 3

        proposals = np.vstack(proposals_list)
        if proposals.shape[0]==0:
            landmarks = np.zeros( (0,5,2) )
            return np.zeros( (0,5) ), landmarks
        scores = np.vstack(scores_list)
        scores_ravel = scores.ravel()
        order = scores_ravel.argsort()[::-1]

        proposals = proposals[order, :]
        scores = scores[order]
        landmarks = np.vstack(landmarks_list)
        landmarks = landmarks[order].astype(np.float32, copy=False)

        pre_det = np.hstack((proposals[:,0:4], scores)).astype(np.float32, copy=False)
        keep = self.nms(pre_det)
        det = np.hstack( (pre_det, proposals[:,4:]) )
        det = det[keep, :]
        landmarks = landmarks[keep]

        return det, landmarks

    @staticmethod
    def bbox_pred(boxes, box_deltas):
        """
        Transform the set of class-agnostic boxes into class-specific boxes
        by applying the predicted offsets (box_deltas)
        :param boxes: !important [N 4]
        :param box_deltas: [N, 4 * num_classes]
        :return: [N 4 * num_classes]
        """
        if boxes.shape[0] == 0:
            return np.zeros((0, box_deltas.shape[1]))

        boxes = boxes.astype(np.float, copy=False)
        widths = boxes[:, 2] - boxes[:, 0] + 1.0
        heights = boxes[:, 3] - boxes[:, 1] + 1.0
        ctr_x = boxes[:, 0] + 0.5 * (widths - 1.0)
        ctr_y = boxes[:, 1] + 0.5 * (heights - 1.0)

        dx = box_deltas[:, 0:1]
        dy = box_deltas[:, 1:2]
        dw = box_deltas[:, 2:3]
        dh = box_deltas[:, 3:4]

        pred_ctr_x = dx * widths[:, np.newaxis] + ctr_x[:, np.newaxis]
        pred_ctr_y = dy * heights[:, np.newaxis] + ctr_y[:, np.newaxis]
        pred_w = np.exp(dw) * widths[:, np.newaxis]
        pred_h = np.exp(dh) * heights[:, np.newaxis]

        pred_boxes = np.zeros(box_deltas.shape)
        # x1
        pred_boxes[:, 0:1] = pred_ctr_x - 0.5 * (pred_w - 1.0)
        # y1
        pred_boxes[:, 1:2] = pred_ctr_y - 0.5 * (pred_h - 1.0)
        # x2
        pred_boxes[:, 2:3] = pred_ctr_x + 0.5 * (pred_w - 1.0)
        # y2
        pred_boxes[:, 3:4] = pred_ctr_y + 0.5 * (pred_h - 1.0)

        if box_deltas.shape[1]>4:
            pred_boxes[:,4:] = box_deltas[:,4:]

        return pred_boxes

    @staticmethod
    def landmark_pred(boxes, landmark_deltas):
        if boxes.shape[0] == 0:
          return np.zeros((0, landmark_deltas.shape[1]))
        boxes = boxes.astype(np.float, copy=False)
        widths = boxes[:, 2] - boxes[:, 0] + 1.0
        heights = boxes[:, 3] - boxes[:, 1] + 1.0
        ctr_x = boxes[:, 0] + 0.5 * (widths - 1.0)
        ctr_y = boxes[:, 1] + 0.5 * (heights - 1.0)
        pred = landmark_deltas.copy()
        for i in range(5):
            pred[:,i,0] = landmark_deltas[:,i,0]*widths + ctr_x
            pred[:,i,1] = landmark_deltas[:,i,1]*heights + ctr_y
        return pred


import logging
import pathlib
from typing import Dict, List
from typing import Tuple
from collections import defaultdict

import numpy as np
import aspect_based_sentiment_analysis as absa
from sklearn.metrics import confusion_matrix
from aspect_based_sentiment_analysis import PatternRecognizer
from aspect_based_sentiment_analysis import Pipeline
from joblib import Memory

from . import extension
from . import utils
from . import plots
from .recognition_key_token import mask_tokens

logger = logging.getLogger('analysis.recognition-key-token-pair')
HERE = pathlib.Path(__file__).parent
memory = Memory(HERE / 'outputs')

def mask_examples(nlp: Pipeline, domain: str, part_parts: Tuple[int, int]):
    dataset = absa.load_examples('semeval', domain, test=True)
    # Split a dataset because it's better to cache more freq.
    part, parts = part_parts
    chunks = utils.split(dataset, n=parts)
    dataset_chunk = chunks[part]

    for i, example in enumerate(dataset_chunk):
        yield i, -1, -1, example    # Predict without a mask.

        [tokenized_example] = nlp.tokenize([example])
        n = len(tokenized_example.text_tokens)

        x, y = np.triu_indices(n, k=1)  # Exclude the diagonal.
        for ij in zip(x, y):
            new_example = mask_tokens(nlp, tokenized_example, indices=ij)
            yield i, *ij, new_example

@memory.cache(ignore=['nlp'])
def _retrieve_labels(
        nlp: Pipeline,
        domain: str,
        part_parts: Tuple[int, int]
) -> np.ndarray:
    partial_results = []
    examples = mask_examples(nlp, domain, part_parts)
    batches = absa.utils.batches(examples, batch_size=32)
    for batch in batches:
        indices, *masked_tokens_ij, batch_examples = zip(*batch)
        predictions = nlp.transform(batch_examples)
        y_hat = [e.sentiment.value for e in predictions]
        partial_results.extend(zip(indices, *masked_tokens_ij, y_hat))
    return np.array(partial_results)

def retrieve_labels(nlp: Pipeline, domain: str, parts: int) -> np.ndarray:
    results = []
    d = defaultdict(list)
    for part in range(parts):
        partial_results = _retrieve_labels(nlp, domain, (part, parts))
        results.extend(partial_results)

        n = len(d)
        for i, *ij, y_hat in partial_results:
            i += n  # We have parts so the index needs to be shifted.
            d[i].append(y_hat)
    results = np.array(results)

    example_indices, mask_i, mask_j, y_new = results.T
    # The index -1 means predictions without masking.
    y_ref = y_new[mask_i == -1]

    y_new = []
    for ref, classes in zip(y_ref, d.values()):
        classes = np.array(classes)
        available = np.where(classes != ref)
        available_classes = classes[available].tolist()

        if not available_classes:
            y_new.append(ref)
            continue

        new = np.bincount(available_classes).argmax()
        y_new.append(new)
    y_new = np.array(y_new)
    mask = y_ref != y_new
    return y_ref, y_new, mask

@memory.cache(ignore=['nlp'])
# The pattern recognizer name is used to distinguish function calls (caching).
def _evaluate(nlp: Pipeline, domain: str, name: str) -> np.ndarray:
    partial_results = []
    dataset = absa.load_examples('semeval', domain, test=True)
    batches = absa.utils.batches(dataset, batch_size=32)
    for batch in batches:
        predictions = nlp.transform(batch)
        predictions = list(predictions)  # Keep in memory to append at the end.

        new_batch = [mask_tokens(nlp, prediction, k=2) for prediction in predictions]
        new_predictions = nlp.transform(new_batch)

        y_ref = [e.sentiment.value for e in predictions]
        y_new = [e.sentiment.value for e in new_predictions]
        partial_results.extend(zip(y_ref, y_new))
    # It's not a generator because we cache function results.
    return np.array(partial_results)

def evaluate(
        recognizer: PatternRecognizer,
        domain: str,
        name: str
) -> List[Tuple[float, np.ndarray, np.ndarray]]:
    results = _evaluate(recognizer, domain, name)
    y_ref, y_new = results.T
    matrix = confusion_matrix(y_ref, y_new)
    acc = 1 - utils.accuracy(matrix)  # The aim is to change a prediction.
    return acc, matrix, y_ref != y_new

def experiment(models: Dict[str, str], save_confusion_matrix: bool = True):
    utils.setup_logger(HERE / 'logs' / 'recognition-key-token-pair.log')
    logger.info('Begin Evaluation: the Key Token Pair Recognition')

    for domain in ['restaurant', 'laptop']:
        name = models[domain]
        nlp = absa.load(name)

        y_ref, y_new, mask = retrieve_labels(nlp, domain, parts=10)
        max_score = sum(mask) / len(mask)
        max_score_matrix = confusion_matrix(y_ref, y_new)

        random = extension.RandomPatternRecognizer()
        attention = extension.AttentionPatternRecognizer(max_patterns=5)
        gradient = extension.GradientPatternRecognizer(max_patterns=5)
        basic = absa.BasicPatternRecognizer(max_patterns=5)
        recognizers = [random, attention, gradient, basic]

        results = []
        for recognizer in recognizers:
            nlp.professor.pattern_recognizer = recognizer
            result = evaluate(nlp, domain, repr(recognizer))
            results.append(result)

        if save_confusion_matrix:
            plots.confusion_matrix.save_figure(results, domain)

        logger.info(
            f'{domain.upper()} DOMAIN\n'
            f'Examples that have at least one key token pair: '
            f'{max_score:.4f}\ny_ref means a prediction without a mask, '
            f'and y_new with a single mask.\n'
            f'Confusion Matrix (y_ref, y_new):\n{max_score_matrix}\n\n'
            f'Random Pattern Recognizer\n'
            f'Acc.: {results[0][0]/max_score:.4f} ({results[0][0]:.4f})\n'
            f'Confusion Matrix (y_ref, y_new):\n{results[0][1]}\n\n'
            f'Attention Pattern Recognizer\n'
            f'Acc.: {results[1][0]/max_score:.4f} ({results[1][0]:.4f})\n'
            f'Confusion Matrix (y_ref, y_new):\n{results[1][1]}\n\n'
            f'Gradient Pattern Recognizer\n'
            f'Acc.: {results[2][0]/max_score:.4f} ({results[2][0]:.4f})\n'
            f'Confusion Matrix (y_ref, y_new):\n{results[2][1]}\n\n'
            f'Basic Pattern Recognizer\n'
            f'Acc.: {results[3][0]/max_score:.4f} ({results[3][0]:.4f})\n'
            f'Confusion Matrix (y_ref, y_new):\n{results[3][1]}\n\n')


#! /usr/bin/python3
# -*- coding: utf-8 -*-

"""
    ********
    Crosslet
    ********

    :class:`~nenupy.crosslet.crosslet.Crosslet` is the main class
    for both :class:`~nenupy.crosslet.xstdata.XST_Data` and
    :class:`~nenupy.crosslet.tvdata.TV_Data`, which inherit from
    it.

    This enables *beamforming* with the
    :meth:`~nenupy.crosslet.crosslet.Crosslet.beamform` method
    (see also :ref:`tuto_beamforming` for a detailed tutorial)
    and *imaging* with the
    :meth:`~nenupy.crosslet.crosslet.Crosslet.image` method
    (see also :ref:`tuto_tv` for a detailed tutorial) from
    cross-correlation statistics data.
"""

__author__ = 'Alan Loh'
__copyright__ = 'Copyright 2020, nenupy'
__credits__ = ['Alan Loh']
__maintainer__ = 'Alan'
__email__ = 'alan.loh@obspm.fr'
__status__ = 'Production'
__all__ = [
    'Crosslet'
]

import numpy as np
import astropy.units as un
try:
    from tqdm import tqdm
except ModuleNotFoundError:
    def tqdm(func):
        """ Overdefine tqdm to return `range` or `enumerate`
        """
        return func

from nenupy.astro import wavelength, HpxSky, eq_zenith
from nenupy.instru import nenufar_loc, read_cal_table, ma_pos
from nenupy.crosslet import UVW
from nenupy.beamlet.sdata import SData

import logging
log = logging.getLogger(__name__)

# ============================================================= #
# ------------------------- Functions ------------------------- #
# ============================================================= #
try:
    from numba import jit,prange
except ModuleNotFoundError:
    from functools import wraps
    def jit(**kwargs):
        """ Override numba.jit function by defining a decorator
            which does nothing but returning the decorated
            function.
        """
        def inner_function(function):
            @wraps(function)
            def wrapper(*args, **kwargs):
                function(*args, **kwargs)
            return wrapper
        return inner_function
    log.warning(
        'numba module not found, operations will take longer.'
    )

@jit(nopython=True, parallel=True, fastmath=True)
def ft_mul(x, y):
    return x * y

@jit(nopython=True, parallel=True, fastmath=True)
def ft_phase(ul, vm):
    return np.exp(
        - 2.j * np.pi * (ul + vm)
    )

@jit(nopython=True, parallel=True, fastmath=True)
def ft_sum(vis, exptf):
    return np.mean(
        vis * exptf,
    )
# ============================================================= #

# ============================================================= #
# ------------------------- Crosslet -------------------------- #
# ============================================================= #
class Crosslet(object):
    """ :class:`~nenupy.crosslet.crosslet.Crosslet` class is not
        designed to be called directly but rather as a base class
        for both :class:`~nenupy.crosslet.xstdata.XST_Data` and
        :class:`~nenupy.crosslet.tvdata.TV_Data` (see their
        related documentation for further instructions on how
        to load these data sets).
    """

    def __init__(self):
        self.freqs = None
        self.mas = None
        self.dt = None
        self.times = None
        self.vis = None

    # --------------------------------------------------------- #
    # --------------------- Getter/Setter --------------------- #
    @property
    def mas(self):
        """ Mini-Arrays used to get the cross-correlation
            statistics data.

            :setter: Mini-Array list

            :getter: Mini-Array list

            :type: :class:`~numpy.ndarray`
        """
        return self._mas
    @mas.setter
    def mas(self, m):
        if m is not None:
            self._mas_idx = np.arange(m.size)
            self._ant1, self._ant2 = np.tril_indices(m.size, 0)
        self._mas = m
        return

    @property
    def xx(self):
        """ Extracts the XX polarization from the cross-correlation
            statistics data. Also refered to as ``'NW'``.
            Array is shaped like (time, frequencies, baselines).

            :getter: XX polarization

            :type: :class:`~numpy.ndarray`
        """
        xx = self.vis[:, :, self._get_cross_idx('X', 'X')]
        log.info(
            'XX loaded.'
        )
        return xx

    @property
    def yy(self):
        """ Extracts the YY polarization from the cross-correlation
            statistics data. Also refered to as ``'NE'``.
            Array is shaped like (time, frequencies, baselines).

            :getter: YY polarization

            :type: :class:`~numpy.ndarray`
        """
        yy = self.vis[:, :, self._get_cross_idx('Y', 'Y')]
        log.info(
            'YY loaded.'
        )
        return yy

    @property
    def yx(self):
        """ Extracts the YX polarization from the cross-correlation
            statistics data.
            Array is shaped like (time, frequencies, baselines).

            :getter: YX polarization

            :type: :class:`~numpy.ndarray`
        """
        yx = self.vis[:, :, self._get_cross_idx('Y', 'X')]
        log.info(
            'YX loaded.'
        )
        return yx

    @property
    def xy(self):
        """ Extracts the XY polarization from the cross-correlation
            statistics data. This polarization is not recorded by
            default for the NenuFAR auto-correlations. However
            this can be computed as it is the complex conjugate
            of :attr:`~nenupy.crosslet.crosslet.Crosslet.yx`
            auto-correlations.
            Array is shaped like (time, frequencies, baselines).

            :getter: XY polarization

            :type: :class:`~numpy.ndarray`
        """
        # Deal with lack of auto XY cross in XST-like data
        ma1, ma2 = np.tril_indices(self.mas.size, 0)
        auto = ma1 == ma2
        cross = ~auto
        _xy = np.zeros(
            (list(self.yx.shape[:-1]) + [ma1.size]),
            dtype=np.complex
        )
        _xy[:, :, auto] = self.yx[:, :, auto].conj()
        # Get XY correlations
        indices = self._get_cross_idx('X', 'Y')
        _xy[:, :, cross] = self.vis[:, :, indices]
        log.info(
            'XY loaded.'
        )
        return _xy

    @property
    def stokes_i(self):
        r""" Computes the stokes parameter I from the
            cross-correlation statistics data using
            :attr:`~nenupy.crosslet.crosslet.Crosslet.xx` and
            :attr:`~nenupy.crosslet.crosslet.Crosslet.yy`.

            .. math::
                \mathbf{I}(t, \nu) = \frac{1}{2} \left[
                    \mathbf{XX}(t, \nu) + \mathbf{YY}(t, \nu)
                \right]

            Array is shaped like (time, frequencies, baselines).

            :getter: Stokes I data

            :type: :class:`~numpy.ndarray`
        """
        return 0.5*(self.xx + self.yy)

    @property
    def mask_auto(self):
        """ Masks auto-correlations. In particular, this is used
            to compute the image (see
            :meth:`~nenupy.crosslet.crosslet.Crosslet.image`).

            :getter: Auto-correlation mask

            :type: :class:`~numpy.ndarray`
        """
        ma1, ma2 = np.tril_indices(self.mas.size, 0)
        auto = ma1 == ma2
        return ~auto

    # --------------------------------------------------------- #
    # ------------------------ Methods ------------------------ #
    # def plot(self):
    #     """
    #     """
    #     mat = np.zeros(
    #         (self.mas.size, self.mas.size),
    #         dtype=np.float
    #     )
    #     mat[np.tril_indices(self.mas.size, 0)] = np.abs(self.xy[0, 0, :])
    #     fig = plt.figure(figsize=(10, 10))
    #     plt.imshow(mat, origin='lower')
    #     return

    def beamform(self, az, el, pol='NW', ma=None, calibration='default'):
        r""" Converts cross correlation statistics data XST,
            :math:`\mathbf{X}(t, \nu)`, in beamformed data BST,
            :math:`B(t, \nu)`, where :math:`t` and :math:`\nu` are
            the time and the frequency respectively.
            :math:`\mathbf{X}(t, \nu)` is a subset of XST data at
            the required polarization ``pol``.

            This is done for a given phasing direction in local
            sky coordinates :math:`\varphi` (azimuth, ``az``) and
            :math:`\theta` (elevation, ``el``), with a selection
            of Mini-Arrays ``ma`` (numbered :math:`a`).

            .. math::
                B (t, \nu) = \operatorname{Re} \left\{
                \sum
                    \left[
                        \underset{\scriptscriptstyle a \times 1}{\mathbf{C}}
                        \cdot
                        \underset{\scriptscriptstyle 1 \times a}{\mathbf{C}^{H}}
                    \right](\nu)
                    \cdot
                    \underset{\scriptscriptstyle a \times a}{\mathbf{X}} (t, \nu)
                    \cdot
                    \left[
                        \underset{\scriptscriptstyle a \times 1}{\mathbf{P}}
                        \cdot
                        \underset{\scriptscriptstyle 1 \times a}{\mathbf{P}^{H}}
                    \right](\nu)
                \right\}

            .. math::
                \rm{with} \quad
                \cases{
                    \mathbf{C}(\nu ) = e^{2 \pi i \nu \mathbf{ t }} \quad \rm{the~calibration~ file}\\
                    \mathbf{P} (\nu) = e^{-2 \pi i \frac{\nu}{c} (\mathbf{b} \cdot \mathbf{u})} \quad \rm{phasing}\\
                    \underset{\scriptscriptstyle a \times 3}{\mathbf{b}} = \mathbf{a}_{1} - \mathbf{a}_{2} \quad \rm{baseline~positions}\\
                    \underset{\scriptscriptstyle 3 \times 1}{\mathbf{u}} = \left[
                        \cos(\theta)\cos(\varphi),
                        \cos(\theta)\sin(\varphi),
                        \sin(\theta) \right]
                }

            :param az:
                Azimuth coordinate used for beamforming (default
                unit is degrees in `float` input).
            :type az: `float` or :class:`~astropy.units.Quantity`
            :param el:
                Elevation coordinate used for beamforming (default
                unit is degrees in `float` input).
            :type el: `float` or :class:`~astropy.units.Quantity`
            :param pol:
                Polarization (either ``'NW'`` or ``'NE'``.
            :type pol: `str`
            :param ma:
                Subset of Mini-Arrays (minimum 2) used for
                beamforming.
            :type ma: `list` or :class:`~numpy.ndarray`
            :param calibration:
                Antenna delay calibration file (i.e, :math:`\mathbf{C}`).
                If ``'none'``, no calibration is applied.
                If ``'default'``, the standard calibration file is
                used, otherwise the calibration file name should
                be given (see also :func:`~nenupy.instru.instru.read_cal_table`).
            :type calibration: `str`

            :returns: Beamformed data.
            :rtype: :class:`~nenupy.beamlet.sdata.SData`

            :Example:
                >>> from nenupy.crosslet import XST_Data
                >>> xst = XST_Data('20191129_141900_XST.fits')
                >>> bf = xst.beamform(
                        az=180,
                        el=90,
                        pol='NW',
                        ma=[17, 44],
                        calibration='default'
                    )
        """
        log.info(
            'Beamforming towards az={}, el={}, pol={}'.format(
                az,
                el,
                pol
            )
        )
        # Mini-Array selection
        if ma is None:
            ma = self.mas.copy()
        mas = self._mas_idx[np.isin(self.mas, ma)]
        # Calibration table
        if calibration.lower() == 'none':
            # No calibration
            cal = np.ones(
                (self.sb_idx.size, mas.size)
            )
        else:
            pol_idx = {'NW': [0], 'NE': [1]}
            cal = read_cal_table(
                calfile=calibration
            )
            cal = cal[np.ix_(
                self.sb_idx,
                mas,
                pol_idx[pol]
            )].squeeze()
        # Matrix of BSTs
        c = np.zeros(
            (
                self.times.size,
                self.freqs.size,
                mas.size,
                mas.size
            ),
            dtype=np.complex
        )
        # Matrix of phasings
        p = np.ones(
            (
                self.freqs.size,
                mas.size,
                mas.size
            ),
            dtype=np.complex
        )
        # Pointing direction
        if isinstance(az, un.Quantity):
            az = az.to(un.deg).value
        if isinstance(el, un.Quantity):
            el = el.to(un.deg).value
        az = np.radians(az)
        el = np.radians(el)
        u = np.array([
            np.cos(el) * np.cos(az),
            np.cos(el) * np.sin(az),
            np.sin(el)
        ])
        # Polarization selection
        mask = np.isin(self._ant1, mas) & np.isin(self._ant2, mas)
        log.info(
            'Loading data...'
        )
        if pol.upper() == 'NW':
            cpol = self.xx[:, :, mask]
        else:
            cpol = self.yy[:, :, mask]
        log.info(
            'Data of shape {} loaded for beamforming'.format(
                cpol.shape
            )
        )
        # Put the Xcorr in a matrix
        trix, triy = np.tril_indices(mas.size, 0)
        c[:, :, trix, triy] = cpol
        c[:, :, triy, trix] = c[:, :, trix, triy].conj()
        # Calibrate the Xcorr with the caltable
        for fi in prange(c.shape[1]):
            cal_i = np.expand_dims(cal[fi], axis=1)
            cal_i_h = np.expand_dims(cal[fi].T.conj(), axis=0)
            mul = np.dot(cal_i, cal_i_h)
            c[:, fi, :, :] *= mul[np.newaxis, :, :]
        # Phase the Xcorr
        dphi = np.dot(
            ma_pos[self._ant1[mask]] - ma_pos[self._ant2[mask]],
            u
        )
        wavel = wavelength(self.freqs).value
        p[:, trix, triy] = np.exp(
            -2.j*np.pi/wavel[:, None] * dphi
        )
        p[:, triy, trix] = p[:, trix, triy].conj()
        data = np.sum((c * p).real, axis=(2, 3))
        log.info(
            'Beamforming complete.'
        )
        return SData(
            data=np.expand_dims(data, axis=2),
            time=self.times,
            freq=self.freqs,
            polar=np.array([pol.upper()])
        )

    def image(self, resolution=1, fov=50):
        r""" Converts NenuFAR-TV-like data sets containing
            visibilities (:math:`V(u,v,\nu , t)`) into images
            :math:`I(l, m, \nu)` phase-centered at the local
            zenith while time averaging the visibilities.
            The Field of View ``fov`` argument defines the
            diameter angular size (zenith-centered) above which
            the image is not computed.

            .. math::
                I(l, m, \nu) = \int
                    \langle V(u, v, \nu, t) \rangle_t e^{
                        2 \pi i \frac{\nu}{c} \left(
                            \langle u(t) \rangle_t l + \langle v(t) \rangle_t m
                        \right)
                    }
                    \, du \, dv

            :param resolution:
                Resoltion (in degrees if a `float` is given) of
                the HEALPix grid (passed to initialize the
                :class:`~nenupy.astro.hpxsky.HpxSky` object).
            :type resolution: `float` or :class:`~astropy.units.Quantity`
            :param fov:
                Field of view diameter of the image (in degrees
                if a `float` is given).
            :type fov: `float` or :class:`~astropy.units.Quantity`

            :returns: HEALPix sky object embedding the computed
                image.
            :rtype: :class:`~nenupy.astro.hpxsky.HpxSky`

            :Example:
                >>> from nenupy.crosslet import TV_Data
                >>> import astropy.units as u
                >>> tv = TV_Data('20191204_132113_nenufarTV.dat')
                >>> im = tv.image(
                        resolution=0.2*u.deg,
                        fov=60*u.deg
                    )

            .. seealso::
                :class:`~nenupy.astro.hpxsky.HpxSky`,
                :meth:`~nenupy.astro.hpxsky.HpxSky.lmn`,
                :meth:`~nenupy.crosslet.uvw.UVW.from_tvdata`

            .. warning::
                This method is intended to be used for NenuFAR-TV
                data and relatively small XST datasets. It is not
                suited to long observations for which a MS
                conversion is required before using imaging
                dedicated softwares.
        """
        if not isinstance(fov, un.Quantity):
            fov *= un.deg
        f_idx = 0 # Frequency index
        # Sky preparation
        sky = HpxSky(resolution=resolution)
        exposure = self.times[-1] - self.times[0]
        sky.time = self.times[0] + exposure/2.
        sky._is_visible = sky._ho_coords.alt >= 90*un.deg - fov/2.
        phase_center = eq_zenith(sky.time)
        l, m, n = sky.lmn(phase_center=phase_center)
        # UVW coordinates
        uvw = UVW.from_tvdata(self)
        u = np.mean( # Mean in time
            uvw.uvw[:, :, 0],
            axis=0
        )[self.mask_auto]/wavelength(self.freqs[f_idx]).value
        v = np.mean(
            uvw.uvw[:, :, 1],
            axis=0
        )[self.mask_auto]/wavelength(self.freqs[f_idx]).value
        w = np.mean( # Mean in time
            uvw.uvw[:, :, 2],
            axis=0
        )[self.mask_auto]/wavelength(self.freqs[f_idx]).value
        # Mulitply (u, v) by (l, m) and compute FT exp
        ul = ft_mul(
            x=np.tile(u, (l.size, 1)).T,
            y= np.tile(l, (u.size, 1))
        )
        vm = ft_mul(
            x=np.tile(v, (m.size, 1)).T,
            y=np.tile(m, (v.size, 1))
        )
        phase = ft_phase(ul, vm)
        # Phase visibilities
        vis = np.mean( # Mean in time
            self.stokes_i,
            axis=0
        )[f_idx, :][self.mask_auto]
        im = np.zeros(l.size)
        for i in tqdm(prange(l.size)):
            im[i] = np.real(
                ft_sum(vis, phase[:, i])
            )
        sky.skymap[sky._is_visible] = im
        return sky

    def nearfield(self, resolution):
        """
        """
        return

    # --------------------------------------------------------- #
    # ----------------------- Internal ------------------------ #
    def _get_cross_idx(self, c1='X', c2='X'):
        """
        """
        corr = np.array(['X', 'Y']*self.mas.size)
        i_ant1, i_ant2 = np.tril_indices(self.mas.size*2, 0)
        corr_mask = (corr[i_ant1] == c1) & (corr[i_ant2] == c2)
        indices = np.arange(i_ant1.size)[corr_mask]
        return indices
# ============================================================= #


import cv2
import numpy as np
import os

from random import shuffle
from tqdm import tqdm

import matplotlib.pyplot as plt

import tensorflow as tf
import tflearn
from tflearn.layers.conv import conv_2d, max_pool_2d
from tflearn.layers.core import input_data, dropout, fully_connected
from tflearn.layers.normalization import local_response_normalization
from tflearn.layers.estimator import regression
import time
import subprocess
prefix = 'home/pallab/gestures-cnn/'
hand_cascade = cv2.CascadeClassifier(prefix + 'classifiers/' + 'cascade_30_15.xml')

TEST_DIR = "/home/pallab/Desktop/test"
IMG_SIZE = 64
LR = 1e-3
MODEL_NAME = 'gesture-{}-{}-{}.model'.format(LR, IMG_SIZE, '5-layer-3-conv-3FC')

def skinMask(frame, x0, y0, width, height ):
    global guessGesture, visualize, mod, lastgesture, saveImg
    # HSV values
    skinkernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(5,5))
    low_range = np.array([0, 50, 80])
    upper_range = np.array([30, 200, 255])

    cv2.rectangle(frame, (x0,y0),(x0+width,y0+height),(0,255,0),1)
    roi = frame[y0:y0+height, x0:x0+width]

    hsv = cv2.cvtColor(roi, cv2.COLOR_BGR2HSV)

    #Apply skin color range
    mask = cv2.inRange(hsv, low_range, upper_range)

    mask = cv2.erode(mask, skinkernel, iterations = 2)
    mask = cv2.dilate(mask, skinkernel, iterations = 1)
    mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, skinkernel)

    #blur
    mask = cv2.GaussianBlur(mask, (15,15),1)
    #cv2.imshow("Blur", mask)

    #bitwise and mask original frame
    res = cv2.bitwise_and(roi, roi, mask = mask)
    # color to grayscale
    res = cv2.cvtColor(res, cv2.COLOR_BGR2GRAY)

    return res

def soft(num):
    pass
    # if num==1:
    #     count=2
    #     while count>0:
    #             #p = subprocess.Popen(["C:/Program Files (x86)/VideoLAN/VLC/vlc.exe","file:///C:/Users/SAURADIP/Desktop/projvid/peace.mp4"])
    #             count=count-1
    #             print(count)
    # elif num==2:
    #     print(2)
    # elif num==3:
    #     print(3)
    # elif num==4:
    #     print(4)

def process_test_data():
    testing_data = []
    for img in tqdm(os.listdir(TEST_DIR)):
        path = os.path.join(TEST_DIR,img)
        img_num = img[0]
        img = cv2.imread(path,cv2.IMREAD_GRAYSCALE)
        img = cv2.resize(img, (IMG_SIZE,IMG_SIZE))
        testing_data.append([np.array(img), img_num])

    shuffle(testing_data)
    np.save('test_data.npy', testing_data)
    return testing_data

tf.reset_default_graph()
convnet = input_data(shape=[None, IMG_SIZE, IMG_SIZE, 1], name='input')

convnet = conv_2d(convnet, 32, 5, activation='relu')
convnet = max_pool_2d(convnet, 2)
convnet = local_response_normalization(convnet)
print(convnet.shape)

convnet = conv_2d(convnet, 64, 5, activation='relu')
convnet = local_response_normalization(convnet)
convnet = max_pool_2d(convnet, 2)
print(convnet.shape)

convnet = conv_2d(convnet, 128, 5, activation='relu')
convnet = max_pool_2d(convnet, 2)
convnet = local_response_normalization(convnet)
print(convnet.shape)

convnet = fully_connected(convnet, 1024, activation='relu')
convnet = dropout(convnet, 0.5)

convnet = fully_connected(convnet, 512, activation='relu')
convnet = dropout(convnet, 0.5)

convnet = fully_connected(convnet, 5, activation='softmax')
convnet = regression(convnet, optimizer='adam', learning_rate=LR, loss='categorical_crossentropy', name='targets')

model = tflearn.DNN(convnet)
model.load("/home/pallab/gestures-cnn/tfmodels/"+MODEL_NAME)

cap = cv2.VideoCapture(0)
count = 1
font = cv2.FONT_HERSHEY_DUPLEX
# while True:
#     ret, frame = cap.read()
#     #cv2.rectangle(frame, (300,300), (200,200), (0,255,0),0)
#     gray=cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
#     hands = hand_cascade.detectMultiScale(gray, 1.7, 30, minSize=(100, 100),maxSize=(150, 150))# minSize=(50, 50),maxSize=(80, 80))
#     for (x,y,w,h) in hands:
#         #print(x)
#         #print(y)
#         #crop_img = frame[x:x+h, y:y+w]
#         cv2.rectangle(frame,(x,y),(x+w,y+h),(200,30,200),2)
#         # value = (33, 33)
#         # hsv = cv2.cvtColor(crop_img,cv2.COLOR_BGR2GRAY)
#         # blur = cv2.GaussianBlur(hsv, (5, 5), 0)
#         # thres = cv2.adaptiveThreshold(blur,255,cv2.ADAPTIVE_THRESH_GAUSSIAN_C,cv2.THRESH_BINARY_INV,11,2)
#         # ret, res = cv2.threshold(thres, 0, 255, cv2.THRESH_BINARY_INV+cv2.THRESH_OTSU)
#         # # mask = cv2.inRange(hsv, np.array([0, 5, 3]), np.array([60, 255, 255]))
#         # # gaussian = cv2.GaussianBlur(mask, (11,11), 0)
#         # erosion = cv2.erode(res, None, iterations = 1)
#         # dilated = cv2.dilate(erosion,None,iterations = 1)
#         # median = cv2.medianBlur(dilated, 7)
#         # median = cv2.resize(median, (IMG_SIZE, IMG_SIZE))
#         # data = median.reshape(IMG_SIZE,IMG_SIZE,1)
#         # model_out = (model.predict([data])[0])
#         # probabilities = model_out
#         # prediction = (model_out == model_out.max(axis=0, keepdims=True)).astype(int)
#         # label = ""
#         # if np.array_equal((prediction),np.array([1.,0.,0.,0.,0.])):
#         #     label = "INDX"
#         # elif np.array_equal((prediction),np.array([0.,1.,0.,0.,0.]) ):
#         #     label = "VSHP"
#         # elif np.array_equal((prediction) , np.array([0.,0.,1.,0.,0.])):
#         #     label = "FIST"
#         # elif np.array_equal((prediction) , np.array([0.,0.,0.,1.,0.])):
#         #     label = "THMB"
#         # elif np.array_equal((prediction) , np.array([0.,0.,0.,0.,1.])):
#         #     label = "NOGS"

#         # gests = ["INDX", "VSHP", "FIST", "THMB", "NOGS"]
#         # plt.barh(gests, probabilities)
#         # plt.savefig("foobar.jpg")
#         # plot = cv2.imread("foobar.jpg")
#         # cv2.putText(frame, label, (200,200), font, 0.5, (0, 255, 0), 1, cv2.LINE_AA)
#         cv2.imshow('cropped', frame)
#         # cv2.imshow('mask', median)
#         # cv2.imshow("Scores", plot)
#         # os.remove("foobar.jpg")
#         # # write_img = cv2.resize(median, (50,50))
#         # # cv2.imwrite('images_data/peace/'+str(count)+'.jpg',write_img)
#         # # print count
#         # # count += 1
#         k = cv2.waitKey(1) & 0xFF
#         if k == 27:
#             break
kernel = np.ones((15,15),np.uint8)
kernel2 = np.ones((1,1),np.uint8)
skinkernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(5,5))
low_range = np.array([0, 50, 80])
upper_range = np.array([30, 200, 255])
while count != 1001:
    plt.clf()
    ret, frame = cap.read()
    ret = cap.set(3,640)
    ret = cap.set(4,480)
    #cv2.rectangle(frame, (400,400), (200,200), (0,255,0),0)
    crop_img = frame[200:400, 200:400]
    value = (33, 33)
    roi=skinMask(frame,400,200,100,100)
    cv2.imshow('win',roi)
    # #hsv = cv2.cvtColor(roi, cv2.COLOR_BGR2HSV)
    # #cv2.imshow('hsv',hsv)
    # # #Apply skin color range
    # #mask = cv2.inRange(roi, low_range, upper_range)

    # mask = cv2.erode(roi, skinkernel, iterations = 1)
    # mask = cv2.dilate(mask, skinkernel, iterations = 1)

    # #blur
    # mask = cv2.GaussianBlur(mask, (15,15), 1)
    # #cv2.imshow("Blur", mask)

    # #bitwise and mask original frame
    # res = cv2.bitwise_and(frame, frame, mask = mask)
    # # color to grayscale
    # #res = cv2.cvtColor(res, cv2.COLOR_BGR2GRAY)
    # #hands = hand_cascade.detectMultiScale(hsv, 1.7, 30, minSize=(100, 100),maxSize=(150, 150))# minSize=(50, 50),maxSize=(80, 80))
    # #hsv = cv2.cvtColor(crop_img,cv2.COLOR_BGR2GRAY)
    # #for (x,y,w,h) in hands:
    #     #cv2.rectangle(frame,(x,y),(x+w,y+h),(200,30,200),2)
    # gray=roi
    blur = cv2.GaussianBlur(roi, (5, 5), 0)
    thres = cv2.adaptiveThreshold(blur,255,cv2.ADAPTIVE_THRESH_GAUSSIAN_C,cv2.THRESH_BINARY_INV,11,3)
    ret, res = cv2.threshold(thres, 0, 255, cv2.THRESH_BINARY_INV+cv2.THRESH_OTSU)
    # mask = cv2.inRange(hsv, np.array([0, 5, 3]), np.array([60, 255, 255]))
    # gaussian = cv2.GaussianBlur(mask, (11,11), 0)
    erosion = cv2.erode(res, None, iterations = 2)
    dilated = cv2.dilate(erosion,None,iterations = 1)
    median = cv2.medianBlur(dilated, 7)
    median = cv2.resize(median, (IMG_SIZE, IMG_SIZE))
    data = median.reshape(IMG_SIZE,IMG_SIZE,1)
    model_out = (model.predict([data])[0])
    probabilities = model_out
    prediction = (model_out == model_out.max(axis=0, keepdims=True)).astype(int)
    label = ""
    if np.array_equal((prediction),np.array([1.,0.,0.,0.,0.])):
        label = "INDEX"
        #time.sleep(1)
        soft(1) # invoke software
    elif np.array_equal((prediction),np.array([0.,1.,0.,0.,0.]) ):
        label = "V-SHAPE"
        #time.sleep(1)
        soft(2)
    elif np.array_equal((prediction) , np.array([0.,0.,1.,0.,0.])):
        label = "FIST"
        #time.sleep(1)
        soft(3)
    elif np.array_equal((prediction) , np.array([0.,0.,0.,1.,0.])):
        label = "THUMB"
        #time.sleep(1)
        soft(4)
    elif np.array_equal((prediction) , np.array([0.,0.,0.,0.,1.])):
        label = "NO-GESTURE"

    #print(probabilities)
    gests = ["INDEX ("+str(round(probabilities[0]*100,2))+" % )", "V-SHAPE ("+str(round(probabilities[1]*100,2))+" % )", "FIST ("+str(round(probabilities[2]*100,2))+" % )", "THUMB ("+str(round(probabilities[3]*100,2))+" % )", "NO-GES ("+str(round(probabilities[4]*100,2))+" % )"]
    lists=plt.barh(gests, (np.around(probabilities,decimals=1)))
    if  np.around(probabilities,decimals=1).any() > 0.7:
        lists[probabilities.argmax()].set_color('g')

   # plt.set_xticklabels(np.around(np.multiply(probabilities,100),decimals=2), minor=False)
    plt.xticks(np.arange(0, 1, step=0.2))
    # plt.savefig("foobar.jpg")
    plt.title('Gesture Recognition')
    plt.xlabel('Probabilities of Detection')
    plt.ylabel('Gestures')
    plt.figure(figsize=(14,7))
    # plot = cv2.imread("foobar.jpg")
    cv2.putText(frame, label, (400,200), font, 0.5, (0, 255, 0), 1, cv2.LINE_AA)
    cv2.imshow('cropped', frame)
    cv2.imshow('mask', median)
    # cv2.imshow("Scores", plot)
    # os.remove("foobar.jpg")
    # write_img = cv2.resize(median, (50,50))
    # cv2.imwrite('images_data/peace/'+str(count)+'.jpg',write_img)
    # print(count)
    count += 1
    k = cv2.waitKey(5) & 0xFF
    if k == 27:
        break

cv2.destroyAllWindows()
cap.release()


import torch
import torchaudio
from librosa.util import find_files
from pathlib import Path
import os
from Preprocessing.AudioPreprocessor import AudioPreprocessor
import soundfile as sf
from speechbrain.pretrained import EncoderClassifier
from numpy import trim_zeros

wav_dir="../data/aridialect/aridialect_wav16000"
wav_files="../data/aridialect/train-text-pac.txt"
device="cpu"
cut_silences=True
speaker_embedding_function_ecapa = EncoderClassifier.from_hparams(source="speechbrain/spkrec-ecapa-voxceleb",run_opts={"device": str(device)},savedir="Models/speechbrain_speaker_embedding_ecapa")
speaker_embedding_function_xvector = EncoderClassifier.from_hparams(source="speechbrain/spkrec-xvect-voxceleb",run_opts={"device": str(device)},savedir="Models/speechbrain_speaker_embedding_xvector")
wav2mel = torch.jit.load("Models/SpeakerEmbedding/wav2mel.pt")
dvector = torch.jit.load("Models/SpeakerEmbedding/dvector-step250000.pt").to(device).eval()
ap = AudioPreprocessor(input_sr=16000, output_sr=16000, melspec_buckets=80, hop_length=256, n_fft=1024, cut_silence=cut_silences)

wav_list=[]
with open(wav_files, encoding="utf8") as f:
    transcriptions = f.read()
    trans_lines = transcriptions.split("\n")
    for line in trans_lines:
        if line.strip() != "":
            wav_list.append(line.split("|")[0] + ".wav")

print(wav_list)

#spkrs_path = list(os.path.basename(x) for x in Path(wav_dir).iterdir())
#spkrs = list(set(str(x)[0:str(x).find("_")] for x in spkrs_path))
spkrs = list(set(str(x)[0:str(x).find("_")] for x in wav_list))
#spkrs = list(os.path.basename(x) for x in Path(wav_dir).iterdir())
print(spkrs)

for spk in spkrs:
#for spk in ['spo']:
    print(spk)
    audio_paths = [os.path.join(wav_dir,filename) for filename in wav_list if filename.startswith(spk)]
    print(audio_paths)
    combined_spemb = torch.zeros([960])
    print(len(audio_paths))
    for audio_path in audio_paths:
        #read audio to datapoint
        #continue
        print(audio_path)
        try:
            wave, sr = sf.read(audio_path)
            norm_wave = ap.audio_to_wave_tensor(normalize=True, audio=wave)
        except ValueError:
            print("ValueError "+audio_path)
            continue
        norm_wave = torch.tensor(trim_zeros(norm_wave.numpy()))
        norm_wave.cpu().detach().numpy()

        ecapa_spemb = speaker_embedding_function_ecapa.encode_batch(torch.Tensor(norm_wave).to(device)).flatten().detach().cpu()
        xvector_spemb = speaker_embedding_function_xvector.encode_batch(torch.Tensor(norm_wave).to(device)).flatten().detach().cpu()
        dvector_spemb = dvector.embed_utterance(wav2mel(torch.Tensor(norm_wave).unsqueeze(0), 16000).to(device)).flatten().detach().cpu()
        print("Embedding dimensions")
        print(ecapa_spemb.size())
        print(xvector_spemb.size())
        print(dvector_spemb.size())
        combi = torch.cat([ecapa_spemb, xvector_spemb, dvector_spemb], dim=0)
        #print(combi)
        #print(combi.size())
        combined_spemb = combined_spemb + combi
        #print(combined_spemb)

    combined_spemb_mean = torch.div(combined_spemb,len(audio_paths))
    print(combined_spemb_mean)
    torch.save(combined_spemb_mean, "Models/SpeakerEmbedding/"+spk+".pt")
    #ecapa embedding 192 entries
    torch.save(combined_spemb_mean[0:192-1], "Models/SpeakerEmbedding/"+spk+"_ecapa.pt")
    #xvector embedding 512 entries
    torch.save(combined_spemb_mean[192:192+512-1], "Models/SpeakerEmbedding/"+spk+"_xvector.pt")
    #dvector embedding 256 entries
    torch.save(combined_spemb_mean[192+512:192+512+256-1], "Models/SpeakerEmbedding/"+spk+"_dvector.pt")


import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.autograd import Variable
import argparse
import json
from tensorboardX import SummaryWriter
import os

from image_loader import get_loader
from AE import AutoEncoder, Decoder

if __name__=='__main__':
    parser = argparse.ArgumentParser(description='BEGAN')
    parser.add_argument('--datapath', dest='datapath', required=True,
                        type=str, help='root folder for data')
    parser.add_argument('--batch_size', dest='batch_size', default=16,
                        type=int)
    parser.add_argument('--num_workers', dest='num_workers', default=4,
                        type=int, help='number of workers for data loader')
    parser.add_argument('--config', dest='config', default='began_config.json',
                        type=str, help='config file for BEGAN')
    parser.add_argument('--g', dest='netG', default='',
                        type=str, help='load Generative network')
    parser.add_argument('--d', dest='netD', default='',
                        type=str, help='load Discriminative network')
    parser.add_argument('--out', dest='outf', default='output',
                        type=str, help='output folder')
    args = parser.parse_args()

    # dataloader setting
    loader = get_loader(args.datapath, args.batch_size, args.num_workers)

    # load config
    with open(args.config) as f:
        config = json.loads(f.read())

    # print configuration
    print(json.dumps(config, indent=4, sort_keys=True))
    for arg in vars(args):
        print("{}:{}".format(arg, getattr(args, arg)))

    # set Networks
    d = AutoEncoder(channel=config['model']['channels'],
                    n=config['model']['N'],
                    h=config['model']['h'])

    g = Decoder(channel=config['model']['channels'],
                n=config['model']['N'],
                h=config['model']['h'])
    # load Network
    if args.netD != '':
        d.load_state_dict(torch.load(args.netD))
    if args.netG != '':
        g.load_state_dict(torch.load(args.netG))

    d.cuda()
    g.cuda()
    # hparams
    lambda_     = config['train']['lambda']
    g_lr = d_lr = config['train']['lr']
    num_iter    = config['train']['iter']
    num_epoch   = config['train']['epoch']
    gamma       = config['train']['gamma']
    beta1       = config['train']['beta1']
    beta2       = config['train']['beta2']

    k = 0.0
    #TODO beta
    d_optimizer = optim.Adam(d.parameters(), lr=d_lr, betas=(beta1, beta2))
    g_optimizer = optim.Adam(g.parameters(), lr=g_lr, betas=(beta1, beta2))

    d_scheduler = optim.lr_scheduler.StepLR(d_optimizer, 100, 0.98)
    g_scheduler = optim.lr_scheduler.StepLR(g_optimizer, 100, 0.98)

    def loss_func(i, o):
        return torch.mean(torch.abs(i-o))
    # tensorboardX
    writer = SummaryWriter()
    # make Directory
    if not os.path.exists(args.outf):
        os.makedirs(args.outf)
    try:
        for epoch in range(num_epoch):
            for i, data in enumerate(loader):
                step = epoch * len(loader) + i
                d_scheduler.step()
                g_scheduler.step()

                ### Discriminator Training ###
                d.zero_grad()
                # 1. real image
                real_img = Variable(data[0].cuda())
                restored_real_img = d(real_img)
                # 2. fake image loss
                noise = (torch.rand(args.batch_size, config['model']['h']) - 0.5)*2
                z = Variable(noise).cuda()
                fake_img = g(z).detach()  # Generator에 대해서는 학습하면 안된다. 이미 optimizer에서 학습할 파라미터를 정해주지만, 계산을 빨리 하기위해 이런다.
                restored_fake_img = d(fake_img)

                fake_loss = loss_func(fake_img, restored_fake_img)
                real_loss = loss_func(real_img, restored_real_img)
                # 3. compute discriminator loss
                loss_D = real_loss - k * fake_loss
                loss_D.backward()

                d_optimizer.step()

                ### Generator Training ###
                g.zero_grad()

                noise = (torch.rand(args.batch_size, config['model']['h']) - 0.5)*2
                z = Variable(noise).cuda()
                fake_img = g(z)
                restored_fake_img = d(fake_img)  #여기는 detach하면 backpropagation이 안흐를 듯!
                loss_G = loss_func(fake_img, restored_fake_img)
                loss_G.backward()

                g_optimizer.step()

                ### compute others ###
                k = k + lambda_*(gamma*real_loss.data[0] - fake_loss.data[0])
                k = max(min(k, 1), 0)
                M_global = real_loss.data[0] + abs(gamma*real_loss.data[0] - fake_loss.data[0])

                if step % 1000 == 0:
                    print(M_global, flush=True)
                    writer.add_scalar('information/M_global', M_global, step)
                    writer.add_scalar('information/k', k, step)
                    writer.add_scalar('information/real_loss', real_loss.data[0], step)
                    writer.add_scalar('information/fake_loss', fake_loss.data[0], step)
                    writer.add_image('RealImg/image', (real_img+1)/2, step)
                    writer.add_image('RealImg/restored', (restored_real_img+1)/2, step)
                    writer.add_image('FakeImg/image', (fake_img+1)/2, step)
                    writer.add_image('FakeImg/restored', (restored_fake_img+1)/2, step)
                    # save model
                    torch.save(g.state_dict(), '%s/netG_epoch_%d.pth' % (args.outf, step))
                    torch.save(d.state_dict(), '%s/netD_epoch_%d.pth' % (args.outf, step))
                if num_iter != -1 and step == num_iter:
                    raise StopIteration
    except StopIteration:
        pass
    writer.close()


#!/usr/bin/env python
import numpy as np
from onshape_api.client import Client
from copy import copy
from robot import RobotURDF, RobotSDF
import sys
import os
import json
import csg

# Loading configuration
robot = 'robots/demo/'
if len(sys.argv) > 1:
    robot = sys.argv[1]

configFile = robot+'/config.json'
client = Client(logging=False, creds=configFile)
config = json.load(open(configFile))

def configGet(name, default=None):
    global config
    if name in config:
        return config[name]
    else:
        if default is None:
            print('! ERROR missing key "'+name+'" in config')
            exit()
        else:
            return default

documentId = configGet('documentId')
drawFrames = configGet('drawFrames')
drawCollisions = configGet('drawCollisions', False)
useScads = configGet('useScads', True)
assemblyName = configGet('assemblyName', False)
outputFormat = configGet('outputFormat', 'urdf')
jointMaxEffort = configGet('jointMaxEffort', 1)
jointMaxVelocity = configGet('jointMaxVelocity', 20)
noDynamics =configGet('noDynamics', False)
outputDirectory = robot
tmp = configGet('dynamics', {})
dynamicsOverride = {}
for key in tmp:
    dynamicsOverride[key.lower()] = tmp[key]

try:
    os.makedirs(outputDirectory)
except OSError:
    pass

print('* Retrieving workspace ID ...')
document = client.get_document(documentId).json()
workspaceId = document['defaultWorkspace']['id']
print('- Workspace id: '+workspaceId)

print('* Retrieving elements in the document, searching for the assembly...')
elements = client.list_elements(documentId).json()
assemblyId = None
for element in elements:
    if element['type'] == 'Assembly' and (assemblyName is False or element['name'].lower() == assemblyName):
        print("- Found assembly, id: "+element['id']+', name: "'+element['name']+'"')
        assemblyId = element['id']

if assemblyId == None:
    print("! Unable to find assembly in this document")
    exit(1)

print('* Retrieving assembly')
assembly = client.get_assembly(documentId, workspaceId, assemblyId)

# Collecting parts instance from assembly and subassemblies
instances = {}
firstInstance = None
def collectParts(instancesToWalk):
    global firstInstance
    for instance in instancesToWalk:
        if firstInstance is None:
            firstInstance = instance['id']
        instances[instance['id']] = instance

root = assembly['rootAssembly']
collectParts(root['instances'])
for asm in assembly['subAssemblies']:
    collectParts(asm['instances'])

# Collecting occurences
occurrences = {}
for occurrence in root['occurrences']:
    occurrence['assignation'] = None
    occurrence['instance'] = instances[occurrence['path'][-1]]
    occurrence['transform'] = np.matrix(np.reshape(occurrence['transform'], (4, 4)))
    occurrences[tuple(occurrence['path'])] = occurrence

# Gets an occurrence given its path
def getOccurrence(path):
    return occurrences[tuple(path)]

assignations = {}
frames = {}
def assignParts(root, parent):
    assignations[root] = parent
    for occurrence in occurrences.values():
        if occurrence['path'][0] == root:
            occurrence['assignation'] = parent

print('* Getting assembly features, scanning for DOFs...')
trunk = None
relations = {}
topLevels = set()
features = root['features']
for feature in features:
    data = feature['featureData']

    child = data['matedEntities'][0]['matedOccurrence'][0]
    parent = data['matedEntities'][1]['matedOccurrence'][0]

    if data['name'][0:3] == 'dof':
        parts = data['name'].split('_')
        del parts[0]
        data['inverted'] = False
        if parts[-1] == 'inv' or parts[-1] == 'inverted':
            data['inverted'] = True
            del parts[-1]
        name = '_'.join(parts)
        if name == '':
            print('! Error: a DOF dones\'t have any name ("'+data['name']+'" should be "dof_...")')
            exit()

        relations[child] = [parent, data, name]
        assignParts(child, child)
        assignParts(parent, parent)
        if child not in frames:
            frames[child] = []
        if parent not in frames:
            frames[parent] = []

print('- Found '+str(len(relations))+' DOFs')

# If we have no DOF
if len(relations) == 0:
    trunk = firstInstance
    assignParts(firstInstance, firstInstance)

# Spreading parts assignations
changed = True
while changed:
    changed = False
    for feature in features:
        data = feature['featureData']
        occurrenceA = data['matedEntities'][0]['matedOccurrence'][0]
        occurrenceB = data['matedEntities'][1]['matedOccurrence'][0]

        if (occurrenceA not in assignations) or (occurrenceB not in assignations):
            if data['name'][0:5] == 'frame':
                name = '_'.join(data['name'].split('_')[1:])
                if occurrenceA in assignations:
                    frames[occurrenceA].append([name, data['matedEntities'][1]['matedOccurrence']])
                    assignParts(occurrenceB, {True: assignations[occurrenceA], False: 'frame'}[drawFrames])
                    changed = True
                if occurrenceB in assignations:
                    frames[occurrenceB].append([name, data['matedEntities'][0]['matedOccurrence']])
                    assignParts(occurrenceA, {True: assignations[occurrenceB], False: 'frame'}[drawFrames])
                    changed = True
            else:
                if occurrenceA in assignations:
                    assignParts(occurrenceB, assignations[occurrenceA])
                    changed = True
                if occurrenceB in assignations:
                    assignParts(occurrenceA, assignations[occurrenceB])
                    changed = True

print('* Building robot tree')
def collect(id):
    part = {}
    part['id'] = id
    part['children'] = []
    for childId in relations:
        entry = relations[childId]
        if entry[0] == id:
            child = collect(childId)
            child['mate'] = entry[1]
            child['relation'] = entry[2]
            part['children'].append(child)
    return part

# Searching for the trunk, assuming it is an element that has
# not children
for childId in relations:
    entry = relations[childId]
    if entry[0] not in relations:
        trunk = entry[0]
        break
trunkOccurrence = getOccurrence([trunk])
print('Trunk is '+trunkOccurrence['instance']['name'])

for occurrence in occurrences.values():
    if occurrence['assignation'] is None:
        print('! WARNING, part ('+occurrence['instance']['name']+') has no assignation, connecting it with trunk')
        occurrence['assignation'] = trunk

tree = collect(trunk)

if outputFormat == 'urdf':
    robot = RobotURDF()
elif outputFormat == 'sdf':
    robot = RobotSDF()
else:
    print('! ERROR Unknown output format: '+outputFormat+' (supported are urdf and sdf')
    exit()
robot.drawCollisions = drawCollisions
robot.jointMaxEffort = jointMaxEffort
robot.jointMaxVelocity = jointMaxVelocity
robot.noDynamics = noDynamics

# Adds a part to the current robot link
def addPart(occurrence, matrix):
    global noDynamics, dynamicsOverride
    part = occurrence['instance']

    # Importing STL file for this part
    prefix = extractPartName(part['name'])
    stlFile = prefix+'.stl'
    stl = client.part_studio_stl_m(part['documentId'], part['documentMicroversion'], part['elementId'], part['partId'])
    f = open(outputDirectory+'/'+stlFile, 'wb')
    f.write(stl)
    f.close()

    # Import the SCAD files pure shapes
    shapes = None
    if useScads:
        scadFile = prefix+'.scad'
        if os.path.exists(outputDirectory+'/'+scadFile):
            shapes = csg.process(outputDirectory+'/'+scadFile)

    # Obtain metadatas about part to retrieve color
    metadata = client.part_get_metadata(part['documentId'], part['documentMicroversion'], part['elementId'], part['partId'])
    if 'appearance' in metadata:
        colors = metadata['appearance']['color']
        color = np.array([colors['red'], colors['green'], colors['blue']])/255.0
    else:
        color = [0.5, 0.5, 0.5]

    # Obtain mass properties about that part
    if noDynamics:
        mass = 0
        com = [0]*3
        inertia = [0]*12
    else:
        if prefix in dynamicsOverride:
            entry = dynamicsOverride[prefix]
            mass = entry['mass']
            com = entry['com']
            inertia = entry['inertia']
        else:
            massProperties = client.part_mass_properties(part['documentId'], part['documentMicroversion'], part['elementId'], part['partId'])
            massProperties = massProperties['bodies'][part['partId']]
            mass = massProperties['mass'][0]
            com = massProperties['centroid']
            inertia = massProperties['inertia']

        if abs(mass) < 1e-9:
            print('! WARNING Part '+part['name']+' has no mass, maybe you should assign a material to it ?')

    pose = occurrence['transform']
    if robot.relative:
        pose = np.linalg.inv(matrix)*pose
    robot.addPart(pose, stlFile, mass, com, inertia, color, shapes)

partNames = {}
def extractPartName(name):
    parts = name.split(' ')
    del parts[-1]
    return '_'.join(parts).lower()

def processPartName(name):
    global partNames
    name = extractPartName(name)

    if name in partNames:
        partNames[name] += 1
    else:
        partNames[name] = 1

    return name+'_'+str(partNames[name])

def buildRobot(tree, matrix, linkPart=None):
    occurrence = getOccurrence([tree['id']])
    instance = occurrence['instance']
    print('~ Adding top-level instance ['+instance['name']+']')

    link = processPartName(instance['name'])

    # Collecting all children in the tree assigned to this top-level part
    robot.startLink(link, matrix)
    for occurrence in occurrences.values():
        if occurrence['assignation'] == tree['id'] and occurrence['instance']['type'] == 'Part':
            name = '_'.join(occurrence['path'])
            print('- Adding part '+occurrence['instance']['name'])
            addPart(occurrence, matrix)
    robot.endLink()

    # Adding the frames (linkage is relative to parent)
    if tree['id'] in frames:
        for name, part in frames[tree['id']]:
            frame = getOccurrence(part)['transform']
            if robot.relative:
                frame = np.linalg.inv(matrix)*frame
            robot.addFrame(name, frame)

    # Calling the function with recursion for children
    k = 0
    for child in tree['children']:
        mate = child['mate']
        if mate['matedEntities'][0]['matedOccurrence'][0] == tree['id']:
            matedOccurrence = mate['matedEntities'][1]
        else:
            matedOccurrence = mate['matedEntities'][0]
        childLinkPart = getOccurrence(matedOccurrence['matedOccurrence'])

        worldAxisFrame = childLinkPart['transform']
        origin = matedOccurrence['matedCS']['origin']
        zAxis = np.array(matedOccurrence['matedCS']['zAxis'])

        if mate['inverted']:
            zAxis = -zAxis

        translation = np.matrix(np.identity(4))
        translation[0, 3] += origin[0]
        translation[1, 3] += origin[1]
        translation[2, 3] += origin[2]
        worldAxisFrame = worldAxisFrame * translation

        childMatrix = matrix
        axisFrame = worldAxisFrame
        if robot.relative:
            axisFrame = np.linalg.inv(matrix)*axisFrame
            childMatrix = worldAxisFrame

        subLink = buildRobot(child, childMatrix, '_'.join(childLinkPart['path']))
        robot.addJoint(link, subLink, axisFrame, child['relation'], zAxis)

    return link

# Start building the robot
buildRobot(tree, np.matrix(np.identity(4)))
robot.finalize()
# print(tree)

print("* Writing "+robot.ext+" file")
f = open(outputDirectory+'/robot.'+robot.ext, 'w')
f.write(robot.xml)
f.close()