from importlib import import_module
import os
from pathlib import Path
from unittest.mock import call, patch

from django.conf import settings
from django.core.files import File
from django.core.files.uploadedfile import SimpleUploadedFile
from django.test import Client, TestCase
from django.urls import reverse
import plotly.graph_objects as go
from rest_framework import status

from src.chat_network import ChatNetwork

class WebTests(TestCase):
    def setUp(self):
        # https://stackoverflow.com/questions/4453764/how-do-i-modify-the-session-in-the-django-test-framework
        # http://code.djangoproject.com/ticket/10899
        settings.SESSION_ENGINE = 'django.contrib.sessions.backends.file'
        engine = import_module(settings.SESSION_ENGINE)
        store = engine.SessionStore()
        store.save()
        self.session = store
        self.client.cookies[settings.SESSION_COOKIE_NAME] = store.session_key

    @classmethod
    def tearDownClass(cls):
        if Path(os.path.join("data", "MyChat.txt")).is_file():
            os.remove(os.path.join("data", "MyChat.txt"))

        if Path(os.path.join("data", "MyChat2.txt")).is_file():
            os.remove(os.path.join("data", "MyChat2.txt"))

    def add_data_to_session(self, data):
        session = self.session
        session.update(self.client.session)
        session.update(data)
        session.save()

    def test_upload_view_is_up_and_running(self):
        # When
        response = Client().get("/upload_chat/")
        # Then
        self.assertEqual(response.status_code, status.HTTP_200_OK)
        self.assertContains(response, "<title>Upload Chat Export File</title>")

    def test_save_chat_file_name_in_session_information(self):
        # Given
        filename = 'MyChat.txt'
        file = File(open('tests/helpers/ChatExample.txt', 'rb'))
        uploaded_file = SimpleUploadedFile(filename, file.read(), 'text/plain')
        c = Client()
        # When
        c.post("/upload_chat/", {"chat_file": uploaded_file})
        # Then
        self.assertEqual("MyChat.txt", c.session["chat_file_name"])

    def test_identify_chat_file_in_visualization_view(self):
        # Given
        filename = 'MyChat.txt'
        file = File(open('tests/helpers/ChatExample.txt', 'rb'))
        uploaded_file = SimpleUploadedFile(filename, file.read(), 'text/plain')
        c = Client()
        # When
        response = c.post("/upload_chat/", {"chat_file": uploaded_file})
        # Then
        self.assertRedirects(
            response,
            reverse("chat_statistics"),
            status_code=302,
            target_status_code=200,
        )
        # When
        self.assertContains(c.get(response.url), "MyChat.txt")

    def test_home_page_redirects_to_stats_page_if_session_file_is_not_empty(self):
        # Given
        file1 = File(open('tests/helpers/ChatExample.txt', 'rb'))
        uploaded_file = SimpleUploadedFile('MyChat.txt', file1.read(), 'text/plain')
        c = Client()
        c.post("/upload_chat/", {"chat_file": uploaded_file})
        # When
        response = c.get("/")
        # Then
        self.assertRedirects(
            response,
            reverse("chat_statistics"),
            status_code=302,
            target_status_code=200,
        )

    def test_home_page_redirects_to_upload_page_if_session_file_is_empty(self):
        # When
        response = Client().get("/")
        # Then
        self.assertRedirects(
            response,
            reverse("upload_chat"),
            status_code=302,
            target_status_code=200,
        )

    def test_chat_statistics_page_redirects_to_home_page_and_then_to_upload_page_if_traces_are_missing(self):
        # Given
        session = self.session
        session["chat_file_name"] = "example.txt"
        session.save()
        expected_redirections = [(reverse("home"), 302), (reverse('upload_chat'), 302)]
        # When
        response = self.client.get("/stats/", follow=True)
        # Then
        self.assertEqual(expected_redirections, response.redirect_chain)

    def test_chat_statistics_page_redirects_to_home_page_and_then_to_upload_page_if_chat_file_name_is_missing(self):
        # Given
        expected_redirections = [(reverse("home"), 302), (reverse('upload_chat'), 302)]
        # When
        response = Client().get("/stats/", follow=True)
        # Then
        self.assertEqual(expected_redirections, response.redirect_chain)

    def test_stats_page_redirects_to_home_and_then_upload_page_and_remove_file_if_error_in_decoding(self):
        # Given
        file1 = File(open('tests/helpers/ChatExampleWrong.PNG', 'rb'))
        uploaded_file = SimpleUploadedFile('MyChat.txt', file1.read(), 'text/plain')
        expected_redirections = [
            (reverse("home") + "?decodingerror=true", 302),
            (reverse('upload_chat') + "?decodingerror=true", 302)
        ]
        # When
        response = self.client.post("/upload_chat/", {"chat_file": uploaded_file}, follow=True)
        # Then
        self.assertEqual(expected_redirections, response.redirect_chain)
        self.assertIsNone(self.client.session.get("chat_file_name", None))
        self.assertContains(response, "Something went wrong. Try again")

    @patch.object(
        ChatNetwork,
        "draw",
        side_effect=[(go.Figure(), "my_node_traces", "my_edge_traces"), go.Figure()],
        autospec=True
    )
    def test_use_traces_if_they_are_saved_in_session(self, draw_mock):
        # Given
        file1 = File(open('tests/helpers/ChatExample.txt', 'rb'))
        uploaded_file = SimpleUploadedFile('MyChat.txt', file1.read(), 'text/plain')
        self.client.post("/upload_chat/", {"chat_file": uploaded_file}, follow=False)

        self.add_data_to_session({
            "node_traces": "my_node_traces",
            "edge_traces": "my_edge_traces",
        })
        # When
        self.client.get("/stats/")
        # Then
        draw_mock.assert_has_calls([
            call(draw_mock.call_args_list[0][0][0], return_traces=True),
            call(draw_mock.call_args_list[1][0][0], node_traces="my_node_traces", edge_traces="my_edge_traces"),
        ])

    @patch.object(
        ChatNetwork,
        "draw",
        side_effect=[(go.Figure(), "node_traces", "edge_traces"), go.Figure()],
        autospec=True
    )
    def test_calculate_traces_and_then_use_them_next_time_you_visit_the_page(self, draw_mock):
        # Given
        file1 = File(open('tests/helpers/ChatExample.txt', 'rb'))
        uploaded_file = SimpleUploadedFile('MyChat.txt', file1.read(), 'text/plain')
        c = Client()
        c.post("/upload_chat/", {"chat_file": uploaded_file}, follow=False)
        # When
        c.get("/stats/")
        # Then
        draw_mock.assert_has_calls([
            call(draw_mock.call_args_list[0][0][0], return_traces=True),
            call(draw_mock.call_args_list[1][0][0], node_traces="node_traces", edge_traces="edge_traces")
        ])
        self.assertEqual("node_traces", c.session["node_traces"])
        self.assertEqual("edge_traces", c.session["edge_traces"])

    @patch.object(
        ChatNetwork,
        "draw",
        return_value=(go.Figure(), "node_traces", "edge_traces"),
        autospec=True
    )
    def test_delete_traces_from_session_when_new_chat_is_uploaded(self, draw_mock):
        # Given
        self.add_data_to_session({
            "node_traces": "old_node_traces",
            "edge_traces": "old_edge_traces",
        })

        file1 = File(open('tests/helpers/ChatExample.txt', 'rb'))
        uploaded_file = SimpleUploadedFile('MyChat.txt', file1.read(), 'text/plain')
        file1.close()
        # When
        self.client.post("/upload_chat/", {"chat_file": uploaded_file}, follow=False)
        # Then
        draw_mock.assert_called_once_with(
            draw_mock.call_args[0][0], return_traces=True
        )
        self.assertEqual("node_traces", self.client.session["node_traces"])
        self.assertEqual("edge_traces", self.client.session["edge_traces"])

    @patch.object(
        ChatNetwork,
        "draw",
        return_value=(go.Figure(), "node_traces", "edge_traces"),
        autospec=True
    )
    def test_generate_traces_and_discard_file_after_uploading(self, draw_mock):
        # Given
        file1 = File(open('tests/helpers/ChatExample.txt', 'rb'))
        uploaded_file = SimpleUploadedFile('MyChat.txt', file1.read(), 'text/plain')
        file1.close()
        # When
        self.client.post("/upload_chat/", {"chat_file": uploaded_file}, follow=False)
        # Then
        self.assertFalse(Path("data/MyChat.txt").is_file())
        draw_mock.assert_called_once_with(
            draw_mock.call_args[0][0], return_traces=True,
        )
        self.assertEqual("node_traces", self.client.session["node_traces"])
        self.assertIsNotNone("edge_traces", self.client.session["edge_traces"])


import torch
from torch.optim import Adam
from dgl.dataloading import GraphDataLoader
from torch.utils.data.sampler import SubsetRandomSampler
import dgl
from model import CG, CosineDecayScheduler, LogReg
from torch_geometric import seed_everything
import numpy as np
import warnings
import yaml
import argparse
from sklearn.model_selection import StratifiedKFold, GridSearchCV
from sklearn.svm import SVC
from sklearn.metrics import f1_score
from dataset import load_graph_classification_dataset
import seaborn as sns
import matplotlib.pyplot as plt
import pandas as pd

warnings.filterwarnings('ignore')

def load_best_configs(args, path):
    with open(path, "r") as f:
        configs = yaml.load(f, yaml.FullLoader)

    configs = configs[args.dataname]

    for k, v in configs.items():
        if "lr" in k or "w" in k:
            v = float(v)
        setattr(args, k, v)
    return args

def evaluate_graph_embeddings_using_svm(embeddings, labels):
    result = []

    kf = StratifiedKFold(n_splits=10, shuffle=True)
    for train_index, test_index in kf.split(embeddings, labels):
        x_train = embeddings[train_index]
        x_test = embeddings[test_index]
        y_train = labels[train_index]
        y_test = labels[test_index]
        params = {"C": [1e-3, 1e-2, 1e-1, 1, 10, 100, 1000]}
        svc = SVC(random_state=42)
        clf = GridSearchCV(svc, params)
        clf.fit(x_train, y_train)
        preds = clf.predict(x_test)
        f1 = f1_score(y_test, preds, average="micro")
        result.append(f1)
    test_f1 = np.mean(result)
    test_std = np.std(result)

    return test_f1, test_std

def collate_fn(batch):
    # graphs = [x[0].add_self_loop() for x in batch]
    graphs = [x[0] for x in batch]
    labels = [x[1] for x in batch]
    batch_g = dgl.batch(graphs)
    labels = torch.cat(labels, dim=0)
    return batch_g, labels

seed_everything(35536)
parser = argparse.ArgumentParser(description="GraphOP")
parser.add_argument("--dataname", type=str, default="ENZYMES")
parser.add_argument("--cuda", type=int, default=0)
args = parser.parse_args()
args = load_best_configs(args, "config.yaml")
dataname = args.dataname

graphs, (n_feat, num_classes) = load_graph_classification_dataset(dataname)
train_idx = torch.arange(len(graphs))
batch_size = 256
train_sampler = SubsetRandomSampler(train_idx)
train_loader = GraphDataLoader(graphs, collate_fn=collate_fn,
                               batch_size=batch_size, shuffle=True)
eval_loader = GraphDataLoader(graphs, collate_fn=collate_fn, batch_size=batch_size,
                              shuffle=True)

device = torch.device(f'cuda:{args.cuda}' if torch.cuda.is_available() else 'cpu')

def train():
    model = CG(n_feat, args.hidden, args.rate, 32, args.alpha, args.layer).to(device)
    optimizer = Adam(model.trainable_parameters(), lr=args.lr, weight_decay=args.w)
    lr_scheduler = CosineDecayScheduler(args.lr, args.warmup, args.epochs)
    mm_scheduler = CosineDecayScheduler(1 - 0.99, 0, args.epochs)
    for epoch in range(1, args.epochs + 1):
        model.train()
        # update momentum
        mm = 1 - mm_scheduler.get(epoch - 1)
        # mm = 0.99
        # update learning rate
        lr = lr_scheduler.get(epoch - 1)
        for param_group in optimizer.param_groups:
            param_group['lr'] = lr
        for g, _ in train_loader:
            g = g.to(device)
            loss = model(g, g.ndata["attr"])
            print(f"Epoch: {epoch}, Loss: {loss.item()}")
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            model.update_target_network(mm)
    x_list = []
    y_list = []
    model.eval()
    for g, label in eval_loader:
        g = g.to(device)
        z1 = model.get_embed(g, g.ndata['attr'])
        y_list.append(label.numpy())
        x_list.append(z1.detach().cpu().numpy())
    x = np.concatenate(x_list, axis=0)
    y = np.concatenate(y_list, axis=0)
    test_f1, test_std = evaluate_graph_embeddings_using_svm(x, y)
    print(args)
    print(f"Acc: {test_f1}, Std: {test_std}")
    # train_embs = torch.from_numpy(x)[idx_train].to(device)
    # test_embs = torch.from_numpy(x)[idx_test].to(device)
    #
    # train_lbls = torch.from_numpy(y)[idx_train].to(device)
    # test_lbls = torch.from_numpy(y)[idx_test].to(device)
    # b_xent = torch.nn.CrossEntropyLoss()
    # accs = []
    # for i in range(20):
    #
    #     log = LogReg(train_embs.shape[1], 2)
    #     opt = torch.optim.Adam(log.parameters(), lr=0.01, weight_decay=5e-4)
    #     if torch.cuda.is_available():
    #         log.to(device)
    #
    #     for _ in range(100):
    #         log.train()
    #         opt.zero_grad()
    #         logits = log(train_embs)
    #         loss = b_xent(logits, train_lbls)
    #         loss.backward()
    #         opt.step()
    #
    #     log.eval()
    #     logits = log(test_embs)
    #     preds = torch.argmax(logits, dim=1)
    #     f1 = f1_score(test_lbls.cpu(), preds.cpu(), average='micro')  # f1 score
    #     accs.append(f1.item() * 100)
    #
    # accs = np.array(accs)
    #
    # best_acc = accs.mean().item()
    # best_std = accs.std().item()
    # if accs.mean().item() > best_acc:
    #     best_acc = accs.mean().item()
    #     best_std = accs.std().item()
    #     best_epoch = epoch
    #
    # print('avg_f1: {0:.2f}, f1_std: {1:.2f}\n'.format(best_acc, best_std))

train()

# PTC = np.array([[65.6, 67.1, 64.4, 63, 62.2, 61.6, 62.7, 61.9, 61.3]])
# NCI109 = np.array([[79.57, 79.7, 80.2, 79.1, 79.9, 79.7, 79.5, 79.4, 79.3]])
# ENZYMES = np.array([[58.5, 56.8, 55.9, 57.6, 58, 57, 58.3, 59.3, 58]])
# MUTAG = np.array([[86, 88.4, 90.5, 90.5, 88.9, 88.4, 87.8, 88.3, 88.9]])
# z = pd.DataFrame(np.concatenate([NCI109, MUTAG], axis=0),
#                  columns=[0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9],
#                  index=["NCI109", "MUTAG"])
# sns.set_palette("pastel")
# # sns.set(style='darkgrid')
# sns.lineplot(data=z.transpose(), markers=True, linewidth=6)
# plt.xlabel(r"$\beta$")
# plt.ylabel("Accuracy")
# # sns.despine()
# # plt.title("Graph Classification")
# plt.grid()
# plt.savefig("./A2.svg", dpi=800)
# plt.show()


import multiprocessing as mp
import time

output = []

def callback_func(result):
    global ouput
    print(f'callback_func: {result}')
    output.append(result)

def sum_all_values(x):
    total = 0
    for i in range(1, x + 1):
        total += i
    #print(f'{total=}')
    return total

if __name__ == "__main__":
    pool = mp.Pool(4)
    results = [pool.apply_async(sum_all_values, args=(x,), callback=callback_func) for x in range(10000, 10000 + 10)]
    pool.close()

    # do something else

    # collect all of the results into a list
    #output = [p.get() for p in results]
    pool.join()
    print(f'the output is {output}')


import torch
import torch.nn as nn
from torch import Tensor

class GaGNet(nn.Module):
    def __init__(self,
                 cin: int,
                 k1: tuple,
                 k2: tuple,
                 c: int,
                 kd1: int,
                 cd1: int,
                 d_feat: int,
                 p: int,
                 q: int,
                 dilas: list,
                 fft_num: int,
                 is_u2: bool,
                 is_causal: bool,
                 is_squeezed: bool,
                 acti_type: str,
                 intra_connect: str,
                 norm_type: str,
                 ):
        """
        :param cin: input channels, default 2 for RI input
        :param k1: kernel size of 2-D GLU, (2, 3) by default
        :param k2: kernel size of the UNet-block, (1, 3) by default
        :param c: channels of the 2-D Convs, 64 by default
        :param kd1: kernel size of the dilated convs in the squeezedTCM, 3 by default
        :param cd1: channels of the dilated convs in the squeezedTCM, 64 by default
        :param d_feat: channels in the regular 1-D convs, 256 by default
        :param p: number of SqueezedTCMs within a group, 2 by default
        :param q: number of GGMs, 3 by default
        :param dilas: dilation rates, [1, 2, 5, 9] by default
        :param fft_num: fft number, 320 by default
        :param is_u2: whether U^{2} is set, True by default
        :param is_causal: whether causal setting, True by default
        :param is_squeezed: whether to squeeze the complex residual modeling path, False by default
        :param acti_type: the activation type in the glance block, "sigmoid" by default
        :param intra_connect: skip-connection type within the UNet-block , "cat" by default
        :param norm_type: "IN" by default.

        Be careful!!! track_running_stats in IN is False by default, i.e., in both training and validation phases, the
        statistics from the batch data are calculated, which causes the non-causality. For BN, however, track_running_sta
        ts is True by default, i.e., in the training phase, the batch statistics are calculated and in the inference
        phase, the global statistics will be fixed and lead to causal inference. If you choose IN, donot forget to switch
        the param track_running_stats into True!
        """
        super(GaGNet, self).__init__()
        self.cin = cin
        self.k1 = k1
        self.k2 = k2
        self.c = c
        self.kd1 = kd1
        self.cd1 = cd1
        self.d_feat = d_feat
        self.p = p
        self.q = q
        self.dilas = dilas
        self.fft_num = fft_num
        self.is_u2 = is_u2
        self.is_causal = is_causal
        self.is_squeezed = is_squeezed
        self.acti_type = acti_type
        self.intra_connect = intra_connect
        self.norm_type = norm_type

        if is_u2:
            self.en = U2Net_Encoder(cin, k1, k2, c, intra_connect, norm_type)
        else:
            self.en = UNet_Encoder(cin, k1, c, norm_type)
        self.gags = nn.ModuleList([GlanceGazeModule(kd1, cd1, d_feat, p, dilas, fft_num, is_causal, is_squeezed,
                                                    acti_type, norm_type) for _ in range(q)])

    def forward(self, inpt: Tensor) -> list:
        if inpt.ndim == 3:
            inpt = inpt.unsqueeze(dim=1)
        b_size, _, seq_len, _ = inpt.shape
        feat_x = self.en(inpt)
        x = feat_x.transpose(-2, -1).contiguous()
        x = x.view(b_size, -1, seq_len)
        pre_x = inpt.transpose(-2, -1).contiguous()
        out_list = []
        for i in range(len(self.gags)):
            tmp = self.gags[i](x, pre_x)
            pre_x = tmp
            out_list.append(tmp)
        return out_list

class GlanceGazeModule(nn.Module):
    def __init__(self,
                 kd1: int,
                 cd1: int,
                 d_feat: int,
                 p: int,
                 dilas: list,
                 fft_num: int,
                 is_causal: bool,
                 is_squeezed: bool,
                 acti_type: str,
                 norm_type: str,
                 ):
        super(GlanceGazeModule, self).__init__()
        self.kd1 = kd1
        self.cd1 = cd1
        self.d_feat = d_feat
        self.p = p
        self.dilas = dilas
        self.fft_num = fft_num
        self.is_causal = is_causal
        self.is_squeezed = is_squeezed
        self.acti_type = acti_type
        self.norm_type = norm_type

        self.glance_block = GlanceBlock(kd1, cd1, d_feat, p, dilas, fft_num, is_causal, acti_type, norm_type)
        self.gaze_block = GazeBlock(kd1, cd1, d_feat, p, dilas, fft_num, is_causal, is_squeezed, norm_type)

    def forward(self, feat_x: Tensor, pre_x: Tensor) -> Tensor:
        """
        :param feat_x: (B, C1, T)
        :param pre_x: (B, 2, F, T)
        :return: (B, 2, F, T)
        """
        gain_filter = self.glance_block(feat_x, pre_x)
        com_resi = self.gaze_block(feat_x, pre_x)
        # crm
        pre_mag, pre_phase = torch.norm(pre_x, dim=1), torch.atan2(pre_x[:,-1,...], pre_x[:,0,...])
        filtered_x = pre_mag * gain_filter
        coarse_x = torch.stack((filtered_x*torch.cos(pre_phase), filtered_x*torch.sin(pre_phase)), dim=1)  # coarse filtering
        x = coarse_x + com_resi
        return x

class GlanceBlock(nn.Module):
    def __init__(self,
                 kd1: int,
                 cd1: int,
                 d_feat: int,
                 p: int,
                 dilas: list,
                 fft_num: int,
                 is_causal: bool,
                 acti_type: str,
                 norm_type: str,
                 ):
        super(GlanceBlock, self).__init__()
        self.kd1 = kd1
        self.cd1 = cd1
        self.d_feat = d_feat
        self.p = p
        self.dilas = dilas
        self.fft_num = fft_num
        self.is_causal = is_causal
        self.acti_type = acti_type
        self.norm_type = norm_type

        ci = (fft_num // 2 + 1) * 2 + d_feat
        self.in_conv_main = nn.Conv1d(ci, d_feat, 1)
        self.in_conv_gate = nn.Sequential(
            nn.Conv1d(ci, d_feat, 1),
            nn.Sigmoid()
        )
        tcn_g_list = []
        for _ in range(p):
            tcn_g_list.append(SqueezedTCNGroup(kd1, cd1, d_feat, dilas, is_causal, norm_type))
        self.tcn_g = nn.Sequential(*tcn_g_list)
        if acti_type == "sigmoid":
            acti = nn.Sigmoid()
        elif acti_type == "tanh":
            acti = nn.Tanh()
        elif acti_type == "relu":
            acti = nn.ReLU()
        else:
            raise RuntimeError("a activation function must be assigned!")
        self.linear_g = nn.Sequential(
            nn.Conv1d(d_feat, fft_num//2+1, 1),
            acti
        )

    def forward(self, feat_x: Tensor, pre_x: Tensor) -> Tensor:
        """
        :param feat_x: (B, C1, T)
        :param pre_x: (B, 2, F, T)
        :return: filter gain, (B, 1, F, T)
        """
        b_size, _, freq_num, seq_len = pre_x.shape
        pre_x = pre_x.view(b_size, -1, seq_len)
        inpt = torch.cat((feat_x, pre_x), dim=1)
        x = self.in_conv_main(inpt) * self.in_conv_gate(inpt)
        x = self.tcn_g(x)
        gain = self.linear_g(x)
        return gain

class GazeBlock(nn.Module):
    def __init__(self,
                 kd1: int,
                 cd1: int,
                 d_feat: int,
                 p: int,
                 dilas: list,
                 fft_num: int,
                 is_causal: bool,
                 is_squeezed: bool,
                 norm_type: str,
                 ):
        super(GazeBlock, self).__init__()
        self.kd1 = kd1
        self.cd1 = cd1
        self.d_feat = d_feat
        self.p = p
        self.dilas = dilas
        self.fft_num = fft_num
        self.is_causal = is_causal
        self.is_squeezed = is_squeezed
        self.norm_type = norm_type

        # Components
        ci = (fft_num//2+1)*2 + d_feat
        self.in_conv_main = nn.Conv1d(ci, d_feat, 1)
        self.in_conv_gate = nn.Sequential(
            nn.Conv1d(ci, d_feat, 1),
            nn.Sigmoid()
        )

        if not is_squeezed:
            tcn_r_list, tcn_i_list = [], []
            for _ in range(p):
                tcn_r_list.append(SqueezedTCNGroup(kd1, cd1, d_feat, dilas, is_causal, norm_type))
                tcn_i_list.append(SqueezedTCNGroup(kd1, cd1, d_feat, dilas, is_causal, norm_type))
            self.tcm_r, self.tcm_i = nn.Sequential(*tcn_r_list), nn.Sequential(*tcn_i_list)
        else:
            tcn_ri_list = []
            for _ in range(p):
                tcn_ri_list.append(SqueezedTCNGroup(kd1, cd1, d_feat, dilas, is_causal, norm_type))
            self.tcm_ri = nn.Sequential(*tcn_ri_list)

        self.linear_r, self.linear_i = nn.Conv1d(d_feat, fft_num//2+1, 1), nn.Conv1d(d_feat, fft_num//2+1, 1)

    def forward(self, feat_x: Tensor, pre_x: Tensor) -> Tensor:
        """
        :param feat_x: (B, C1, T)
        :param pre_x:  (B, 2, F, T)
        :return: complex residual, (B, 2, F, T)
        """
        b_size, _, freq_num, seq_len = pre_x.shape
        pre_x = pre_x.view(b_size, -1, seq_len)
        inpt = torch.cat((feat_x, pre_x), dim=1)
        x = self.in_conv_main(inpt) * self.in_conv_gate(inpt)
        if not self.is_squeezed:
            x_r, x_i = self.tcm_r(x), self.tcm_i(x)
        else:
            x = self.tcm_ri(x)
            x_r, x_i = x, x
        x_r, x_i = self.linear_r(x_r), self.linear_i(x_i)
        return torch.stack((x_r, x_i), dim=1)

class SqueezedTCNGroup(nn.Module):
    def __init__(self,
                 kd1: int,
                 cd1: int,
                 d_feat: int,
                 dilas: list,
                 is_causal: bool,
                 norm_type: str,
                 ):
        super(SqueezedTCNGroup, self).__init__()
        self.kd1 = kd1
        self.cd1 = cd1
        self.d_feat = d_feat
        self.dilas = dilas
        self.is_causal = is_causal
        self.nomr_type = norm_type

        tcn_list = []
        for i in range(len(dilas)):
            tcn_list.append(SqueezedTCM(kd1, cd1, d_feat, dilas[i], is_causal, norm_type))
        self.tcns = nn.Sequential(*tcn_list)

    def forward(self, x: Tensor) -> Tensor:
        return self.tcns(x)

class SqueezedTCM(nn.Module):
    def __init__(self,
                 kd1: int,
                 cd1: int,
                 d_feat: int,
                 dilation: int,
                 is_causal: bool,
                 norm_type: str,
                 ):
        super(SqueezedTCM, self).__init__()
        self.kd1 = kd1
        self.cd1 = cd1
        self.d_feat = d_feat
        self.dilation = dilation
        self.is_causal = is_causal
        self.norm_type = norm_type

        self.in_conv = nn.Conv1d(d_feat, cd1, 1, bias=False)
        if is_causal:
            padding = ((kd1-1)*dilation, 0)
        else:
            padding = ((kd1-1)*dilation//2, (kd1-1)*dilation//2)
        self.d_conv = nn.Sequential(
            nn.PReLU(cd1),
            NormSwitch(norm_type, "1D", cd1),
            nn.ConstantPad1d(padding, value=0.),
            nn.Conv1d(cd1, cd1, kernel_size=kd1, dilation=dilation, bias=False)
            )
        self.out_conv = nn.Sequential(
            nn.PReLU(cd1),
            NormSwitch(norm_type, "1D", cd1),
            nn.Conv1d(cd1, d_feat, 1, bias=False)
            )
    def forward(self, x: Tensor) -> Tensor:
        resi = x
        x = self.in_conv(x)
        x = self.d_conv(x)
        x = self.out_conv(x)
        x = x + resi
        return x

class U2Net_Encoder(nn.Module):
    def __init__(self,
                 cin: int,
                 k1: tuple,
                 k2: tuple,
                 c: int,
                 intra_connect: str,
                 norm_type: str,
                 ):
        super(U2Net_Encoder, self).__init__()
        self.cin = cin
        self.k1 = k1
        self.k2 = k2
        self.c = c
        self.intra_connect = intra_connect
        self.norm_type = norm_type
        k_beg = (2, 5)
        c_end = 64
        meta_unet = []
        meta_unet.append(
            En_unet_module(cin, c, k_beg, k2, intra_connect, norm_type, scale=4))
        meta_unet.append(
            En_unet_module(c, c, k1, k2, intra_connect, norm_type, scale=3))
        meta_unet.append(
            En_unet_module(c, c, k1, k2, intra_connect, norm_type, scale=2))
        meta_unet.append(
            En_unet_module(c, c, k1, k2, intra_connect, norm_type, scale=1))
        self.meta_unet_list = nn.ModuleList(meta_unet)
        self.last_conv = nn.Sequential(
            GateConv2d(c, c_end, k1, (1,2)),
            NormSwitch(norm_type, "2D", c_end),
            nn.PReLU(c_end)
        )
    def forward(self, x: Tensor) -> Tensor:
        for i in range(len(self.meta_unet_list)):
            x = self.meta_unet_list[i](x)
        x = self.last_conv(x)
        return x

class UNet_Encoder(nn.Module):
    def __init__(self,
                 cin: int,
                 k1: tuple,
                 c: int,
                 norm_type: str,
                 ):
        super(UNet_Encoder, self).__init__()
        self.cin = cin
        self.k1 = k1
        self.c = c
        self.norm_type = norm_type
        k_beg = (2, 5)
        c_end = 64
        unet = []
        unet.append(nn.Sequential(
            GateConv2d(cin, c, k_beg, (1,2)),
            NormSwitch(norm_type, "2D", c),
            nn.PReLU(c)))
        unet.append(nn.Sequential(
            GateConv2d(c, c, k1, (1,2)),
            NormSwitch(norm_type, "2D", c),
            nn.PReLU(c)))
        unet.append(nn.Sequential(
            GateConv2d(c, c, k1, (1,2)),
            NormSwitch(norm_type, "2D", c),
            nn.PReLU(c)))
        unet.append(nn.Sequential(
            GateConv2d(c, c, k1, (1,2)),
            NormSwitch(norm_type, "2D", c),
            nn.PReLU(c)))
        unet.append(nn.Sequential(
            GateConv2d(c, c_end, k1, (1,2)),
            NormSwitch(norm_type, "2D", c_end),
            nn.PReLU(c_end)))
        self.unet_list = nn.ModuleList(unet)

    def forward(self, x: Tensor) -> Tensor:
        for i in range(len(self.unet_list)):
            x = self.unet_list[i](x)
        return x

class En_unet_module(nn.Module):
    def __init__(self,
                 cin: int,
                 cout: int,
                 k1: tuple,
                 k2: tuple,
                 intra_connect: str,
                 norm_type,
                 scale: int,
                 ):
        super(En_unet_module, self).__init__()
        self.k1 = k1
        self.k2 = k2
        self.cin = cin
        self.cout = cout
        self.intra_connect = intra_connect
        self.norm_type = norm_type
        self.scale = scale

        in_conv_list = []
        in_conv_list.append(GateConv2d(cin, cout, k1, (1, 2)))
        in_conv_list.append(NormSwitch(norm_type, "2D", cout))
        in_conv_list.append(nn.PReLU(cout))
        self.in_conv = nn.Sequential(*in_conv_list)

        enco_list, deco_list = [], []
        for _ in range(scale):
            enco_list.append(Conv2dunit(k2, cout, norm_type))
        for i in range(scale):
            if i == 0:
                deco_list.append(Deconv2dunit(k2, cout, "add", norm_type))
            else:
                deco_list.append(Deconv2dunit(k2, cout, intra_connect, norm_type))
        self.enco = nn.ModuleList(enco_list)
        self.deco = nn.ModuleList(deco_list)
        self.skip_connect = Skip_connect(intra_connect)

    def forward(self, x: Tensor):
        x_resi = self.in_conv(x)
        x = x_resi
        x_list = []
        for i in range(len(self.enco)):
            x = self.enco[i](x)
            x_list.append(x)

        for i in range(len(self.deco)):
            if i == 0:
                x = self.deco[i](x)
            else:
                x_con = self.skip_connect(x, x_list[-(i+1)])
                x = self.deco[i](x_con)
        x_resi = x_resi + x
        del x_list
        return x_resi

class Conv2dunit(nn.Module):
    def __init__(self,
                 k: tuple,
                 c: int,
                 norm_type: str,
                 ):
        super(Conv2dunit, self).__init__()
        self.k = k
        self.c = c
        self.norm_type = norm_type
        self.conv = nn.Sequential(
            nn.Conv2d(c, c, k, (1, 2)),
            NormSwitch(norm_type, "2D", c),
            nn.PReLU(c))

    def forward(self, x):
        return self.conv(x)

class Deconv2dunit(nn.Module):
    def __init__(self,
                 k: tuple,
                 c: int,
                 intra_connect: str,
                 norm_type: str,
                 ):
        super(Deconv2dunit, self).__init__()
        self.k = k
        self.c = c
        self.intra_connect = intra_connect
        self.norm_type = norm_type
        deconv_list = []
        if self.intra_connect == "add":
            deconv_list.append(nn.ConvTranspose2d(c, c, k, (1, 2)))
        elif self.intra_connect == "cat":
            deconv_list.append(nn.ConvTranspose2d(2*c, c, k, (1, 2)))
        deconv_list.append(NormSwitch(norm_type, "2D", c))
        deconv_list.append(nn.PReLU(c))
        self.deconv = nn.Sequential(*deconv_list)

    def forward(self, x):
        return self.deconv(x)

class GateConv2d(nn.Module):
    def __init__(self,
                 in_channels: int,
                 out_channels: int,
                 kernel_size: tuple,
                 stride: tuple):
        super(GateConv2d, self).__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = kernel_size
        self.stride = stride
        k_t = kernel_size[0]
        if k_t > 1:
            padding = (0, 0, k_t-1, 0)
            self.conv = nn.Sequential(
                nn.ConstantPad2d(padding, value=0.),
                nn.Conv2d(in_channels=in_channels, out_channels=out_channels*2, kernel_size=kernel_size, stride=stride))
        else:
            self.conv = nn.Conv2d(in_channels=in_channels, out_channels=out_channels*2, kernel_size=kernel_size,
                                  stride=stride)

    def forward(self, inputs: Tensor) -> Tensor:
        if inputs.dim() == 3:
            inputs = inputs.unsqueeze(dim=1)
        x = self.conv(inputs)
        outputs, gate = x.chunk(2, dim=1)
        return outputs * gate.sigmoid()

class Skip_connect(nn.Module):
    def __init__(self,
                 connect: str):
        super(Skip_connect, self).__init__()
        self.connect = connect

    def forward(self, x_main, x_aux):
        if self.connect == "add":
            x = x_main + x_aux
        elif self.connect == "cat":
            x = torch.cat((x_main, x_aux), dim=1)
        return x

class Chomp_T(nn.Module):
    def __init__(self,
                 t: int):
        super(Chomp_T, self).__init__()
        self.t = t

    def forward(self, x):
        return x[:, :, 0:-self.t, :]

class NormSwitch(nn.Module):
    """
    Currently, only BN and IN are considered
    """
    def __init__(self,
                 norm_type: str,
                 dim_size: str,
                 c: int,
                 ):
        super(NormSwitch, self).__init__()
        self.norm_type = norm_type
        self.dim_size = dim_size
        self.c = c

        assert norm_type in ["BN", "IN"] and dim_size in ["1D", "2D"]
        if norm_type == "BN":
            if dim_size == "1D":
                self.norm = nn.BatchNorm1d(c)
            else:
                self.norm = nn.BatchNorm2d(c)
        else:
            if dim_size == "1D":
                self.norm = nn.InstanceNorm1d(c, affine=True)
            else:
                self.norm = nn.InstanceNorm2d(c, affine=True)

    def forward(self, x):
        return self.norm(x)

def numParams(net):
    import numpy as np
    num = 0
    for param in net.parameters():
        if param.requires_grad:
            num += int(np.prod(param.size()))
    return num

def stagewise_com_mag_mse_loss(esti_list, label, frame_list):
    alpha_list = [0.1 for _ in range(len(esti_list))]
    alpha_list[-1] = 1
    mask_for_loss = []
    utt_num = label.size()[0]
    with torch.no_grad():
        for i in range(utt_num):
            tmp_mask = torch.ones((frame_list[i], label.size()[-2]), dtype=label.dtype)
            mask_for_loss.append(tmp_mask)
        mask_for_loss = nn.utils.rnn.pad_sequence(mask_for_loss, batch_first=True).to(label.device)
        mask_for_loss = mask_for_loss.transpose(-2, -1).contiguous()
        com_mask_for_loss = torch.stack((mask_for_loss, mask_for_loss), dim=1)
    loss1, loss2 = 0., 0.
    mag_label = torch.norm(label, dim=1)
    for i in range(len(esti_list)):
        mag_esti = torch.norm(esti_list[i], dim=1)
        loss1 = loss1 + alpha_list[i] * (((esti_list[i] - label) ** 2.0) * com_mask_for_loss).sum() / com_mask_for_loss.sum()
        loss2 = loss2 + alpha_list[i] * (((mag_esti - mag_label) ** 2.0) * mask_for_loss).sum() / mask_for_loss.sum()
    return 0.5 * (loss1 + loss2)

def main(args, net):
    batch_size = args.batch_size
    sr = args.sr
    wav_len = int(args.wav_len * sr)
    win_size = int(args.win_size * sr)
    win_shift = int(args.win_shift * sr)
    fft_num = args.fft_num
    noisy_list, target_list, frame_list = [], [], []
    for i in range(batch_size):
        noisy_list.append(torch.rand(wav_len))
        target_list.append(torch.rand(wav_len))
        frame_list.append((wav_len - win_size + win_size) // win_shift + 1)
    noisy_wav, target_wav = nn.utils.rnn.pad_sequence(noisy_list, batch_first=True), \
                            nn.utils.rnn.pad_sequence(target_list, batch_first=True)

    noisy_stft = torch.stft(noisy_wav, fft_num, win_shift, win_size, torch.hann_window(win_size).to(noisy_wav.device))
    target_stft = torch.stft(target_wav, fft_num, win_shift, win_size, torch.hann_window(win_size).to(target_wav.device))
    _, freq_num, seq_len, _ = noisy_stft.shape
    noisy_stft, target_stft = noisy_stft.permute(0,3,2,1).cuda(), target_stft.permute(0,3,1,2).cuda()
    # conduct sqrt power-compression
    noisy_mag, noisy_phase = torch.norm(noisy_stft, dim=1) ** 0.5, torch.atan2(noisy_stft[:, -1, ...], noisy_stft[:, 0, ...])
    target_mag, target_phase = torch.norm(target_stft, dim=1) ** 0.5, torch.atan2(target_stft[:, -1, ...], target_stft[:, 0, ...])
    noisy_stft = torch.stack((noisy_mag * torch.cos(noisy_phase), noisy_mag * torch.sin(noisy_phase)), dim=1).cuda()
    target_stft = torch.stack((target_mag * torch.cos(target_phase), target_mag * torch.sin(target_phase)), dim=1).cuda()

    esti_list = net(noisy_stft)
    print('input size:{} -> output size:{}, label size:{}'.format(noisy_stft.shape, esti_list[-1].shape, target_stft.shape))
    loss = stagewise_com_mag_mse_loss(esti_list, target_stft, frame_list)
    print("Calculated loss value:{}".format(loss.item()))

if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser("This script provides the network code of gagnet and a simple testing, you can"
                                     "train the network according to your own pipeline")
    parser.add_argument("--batch_size", type=int, default=4)
    parser.add_argument("--sr", type=int, default=16000)
    parser.add_argument("--wav_len", type=float, default=6.0)
    parser.add_argument("--win_size", type=float, default=0.020)
    parser.add_argument("--win_shift", type=float, default=0.010)
    parser.add_argument("--fft_num", type=int, default=320)
    parser.add_argument("--k1", type=tuple, default=(2, 3))
    parser.add_argument("--k2", type=tuple, default=(1, 3))
    parser.add_argument("--c", type=int, default=64)
    parser.add_argument("--kd1", type=int, default=3)
    parser.add_argument("--cd1", type=int, default=64)
    parser.add_argument("--d_feat", type=int, default=256)
    parser.add_argument("--p", type=int, default=2)
    parser.add_argument("--q", type=int, default=3)
    parser.add_argument("--dilas", type=list, default=[1,2,5,9])
    parser.add_argument("--is_u2", type=bool, default=True, choices=[True, False])
    parser.add_argument("--is_causal", type=bool, default=True, choices=[True, False])
    parser.add_argument("--is_squeezed", type=bool, default=False, choices=[True, False])
    parser.add_argument("--acti_type", type=str, default="sigmoid", choices=["sigmoid", "tanh", "relu"])
    parser.add_argument("--intra_connect", type=str, default="cat", choices=["cat", "add"])
    parser.add_argument("--norm_type", type=str, default="IN", choices=["BN", "IN"])

    args = parser.parse_args()
    net = GaGNet(
        cin=2,
        k1=args.k1,
        k2=args.k2,
        c=args.c,
        kd1=args.kd1,
        cd1=args.cd1,
        d_feat=args.d_feat,
        p=args.p,
        q=args.q,
        dilas=args.dilas,
        fft_num=args.fft_num,
        is_u2=args.is_u2,
        is_causal=args.is_causal,
        is_squeezed=args.is_squeezed,
        acti_type=args.acti_type,
        intra_connect=args.intra_connect,
        norm_type=args.norm_type,
        ).cuda()
    net.eval()
    print('The number of parameters of the model is:%.5d' % numParams(net))
    main(args, net)


import tkinter as tkr
from tkinter import ttk
from tkinter import messagebox
import tkinter.filedialog as tkfd
import smtplib
import imghdr
from email.message import EmailMessage
import imap_tools
from imap_tools import MailBox

app = tkr.Tk()
app.geometry('1000x800')
app.title('EmailApp')

##def show():
##    print('')
with MailBox('imap.gmail.com').login('yudhisthir1998@gmail.com','pass','INBOX') as mailbox:
    bodies = [msg.text or msg.html for msg in mailbox.fetch()]

###Mail_list###
selected = ""
lb = tkr.Variable(app)
mail_box = tkr.Listbox(app,width=30,height=600,activestyle='dotbox',fg='white',bg='black')
mail_list = ['yudhisthir1998@gmail.com' ,'yudhisthirsingh579@gmail.com']

for i in mail_list:
    mail_box.insert(tkr.END,i)
mail_box.select_set(0)
curs = mail_box.curselection
mail_box.place(x=0,y=0)

msg_box = tkr.Listbox(app , width=30, height=600,activestyle='dotbox',fg='white',bg='black')

for i in range(10):
    msg_box.insert(tkr.END,bodies[i])
msg_box.select_set(0)
curs = msg_box.curselection
msg_box.place(x=600,y=10)

###label###

tkr.Label(app,text='To:',font=('Arial',18),fg='black').place(x=340,y=50)
tkr.Label(app,text='Sub:',font=('Arial',18),fg='black').place(x=340,y=100)
tkr.Label(app,text='Msg:',font=('Arial',18),fg='black').place(x=340,y=150)

##Entry###

toEnt = tkr.Variable(app)
toEnt.set(mail_box.get(curs()))
tkr.Entry(app,textvariable=toEnt,font=(20),width=25).place(x=300,y=55)

SubEnt = tkr.Variable(app)
tkr.Entry(app,textvariable=SubEnt,font=(20),width=25).place(x=300,y=105)

msgEnt = tkr.Variable(app)
tkr.Entry(app,textvariable=msgEnt,font=(20),width=25).place(x=300,y=155)

#######
def addAttachment():
    with open('demo.txt','rb') as f1:
        file_data = f1.read()
        file_type = imghdr.what(f1.name)
        file_name = f1.name

###attachment###

def attach():
    global img
    file = tkfd.askopenfiles(mode='r',filetypes=[('Images','*.jpeg')])
    print(file)
    Img_name =file[0].name
    img1 = Image.open(Img_name)
    img1 = img1.resize((200,200))
    img = ImageTk.PhotoImage(img1)

    if img != null:
        addAttachment()

##EmailMessage####
def sendMail():
    toEnt.set(mail_box.get(curs()))

    msg = EmailMessage()
    msg['To'] = toEnt.get()
    msg['Subject'] = SubEnt.get()
    msg['from'] = 'yudhister@gmail.com'
    msg.set_content(msgEnt.get())
    msg.add_attachment(file_data,maintype='application',subtype='octet_stream',filename=file_name)
    server = smtplib.SMTP_SSL('smtp.gmail.com',465)
    server.login('yudhister@gmail.com','pass')
    server.send_message(msg)
    server.quit()

tkr.Button(app,text='attach',command=attach,bg='lightblue').place(x=450,y=300)
tkr.Button(app,text='Send',command=sendMail,bg='green').place(x=600,y=300)

### inbox ###

##app.mailloop()


import logging
import signal

# import zmq

from create_base_weights import ensure_base_weights
from learn_servers import learner_block_server
import parse_params as pp
from remote_logging import configure_logging, cleanup_logging

HELP_STRING = "Multi round learning"

logger = logging.getLogger("Multi round learning")

def build_args():
    parser = pp.create_parser(HELP_STRING)
    pp.add_remote_log(parser)
    pp.add_epochs(parser)
    pp.add_jid(parser)
    pp.add_inproc(parser)
    pp.add_verbose(parser)

    pp.add_data_provider(parser)
    pp.add_model_provider(parser)
    pp.add_report_storage_provider(parser)
    pp.add_storage_provider(parser)
    pp.add_storage_dirs(parser)
    pp.add_fedavg_provider(parser)
    pp.add_pubsub_bind(parser)
    pp.add_pubsub_address(parser)
    pp.add_partitions(parser)
    pp.add_iterations(parser)
    pp.add_learn_partitions(parser)

    args = pp.get_args(parser)
    pp.bind_args_to_context(args)
    return args

def main():
    signal.signal(signal.SIGINT, signal.SIG_DFL)

    print("Starting multi partition learning")

    args = build_args()
    configure_logging(args.remote_log)

    ensure_base_weights()

    try:
        server = learner_block_server(args.inproc)
        if not server.wait_ready():
            exit(1)
        iterations = args.iterations
        if args.parts:
            server.learn_multi_and_close(args.parts, iterations)
        else:
            server.learn_all_and_close(iterations)
    finally:
        cleanup_logging()
        # zmq.Context().instance().destroy()

if __name__ == "__main__":
    main()


import socket
import threading

def receive_messages(client_socket):
    while True:
        try:
            message = client_socket.recv(1024).decode("utf-8")
            if not message:
                break
            print("Повідомлення від сервера: {}".format(message))

        except ConnectionResetError:
            break

    client_socket.close()
    print("З'єднання з сервером закрите.")

def send_message(client_socket):
    while True:
        message = input("Введіть повідомлення: ")
        client_socket.send(message.encode("utf-8"))

        if message.lower() == "закрити":
            break

    client_socket.close()
    print("З'єднання з сервером закрите.")

def start_client():
    host = "127.0.0.1"
    port = 12345

    client_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    client_socket.connect((host, port))
    print("Підключено до сервера.")

    receive_thread = threading.Thread(target=receive_messages, args=(client_socket,))
    send_thread = threading.Thread(target=send_message, args=(client_socket,))

    receive_thread.start()
    send_thread.start()

start_client()


import re

from c7n_gcp.provider import resources
from c7n_gcp.query import QueryResourceManager, TypeInfo, ChildResourceManager, ChildTypeInfo
from c7n.utils import local_session

@resources.register('app-engine')
class AppEngineApp(QueryResourceManager):
    """GCP resource: https://cloud.google.com/appengine/docs/admin-api/reference/rest/v1/apps
    """
    class resource_type(TypeInfo):
        service = 'appengine'
        version = 'v1'
        component = 'apps'
        enum_spec = ('get', '[@]', None)
        scope = None
        id = 'id'
        name = 'name'
        default_report_fields = [
            'id', 'locationId', 'servingStatus', 'authDomain', 'defaultHostName']
        asset_type = "appengine.googleapis.com/Application"
        permissions = ('appengine.applications.get',)

        @staticmethod
        def get(client, resource_info):
            return client.execute_query(
                'get', {'appsId': re.match('apps/(.*)',
                    resource_info['resourceName']).group(1)})

    def get_resource_query(self):
        return {'appsId': local_session(self.session_factory).get_default_project()}

@resources.register('app-engine-certificate')
class AppEngineCertificate(ChildResourceManager):
    """GCP resource:
    https://cloud.google.com/appengine/docs/admin-api/reference/rest/v1/apps.authorizedCertificates
    """
    def _get_parent_resource_info(self, child_instance):
        return {'resourceName': re.match(
            '(apps/.*?)/authorizedCertificates/.*', child_instance['name']).group(1)}

    class resource_type(ChildTypeInfo):
        service = 'appengine'
        version = 'v1'
        component = 'apps.authorizedCertificates'
        enum_spec = ('list', 'certificates[]', None)
        scope = None
        name = 'displayName'
        id = 'id'
        parent_spec = {
            'resource': 'app-engine',
            'child_enum_params': {
                ('id', 'appsId')
            }
        }
        default_report_fields = ['displayName', 'expireTime']
        permissions = ('appengine.applications.get',)

        @staticmethod
        def get(client, resource_info):
            apps_id, cert_id = re.match('apps/(.*?)/authorizedCertificates/(.*)',
                                        resource_info['resourceName']).groups()
            return client.execute_query('get', {'appsId': apps_id,
                                                'authorizedCertificatesId': cert_id})

@resources.register('app-engine-domain')
class AppEngineDomain(ChildResourceManager):
    """GCP resource:
    https://cloud.google.com/appengine/docs/admin-api/reference/rest/v1/apps.authorizedDomains/list#AuthorizedDomain
    """
    class resource_type(ChildTypeInfo):
        service = 'appengine'
        version = 'v1'
        component = 'apps.authorizedDomains'
        enum_spec = ('list', 'domains[]', None)
        scope = None
        id = 'id'
        name = "name"
        default_report_fields = [id, name]
        parent_spec = {
            'resource': 'app-engine',
            'child_enum_params': {
                ('id', 'appsId')
            }
        }
        permissions = ('appengine.applications.get',)

@resources.register('app-engine-domain-mapping')
class AppEngineDomainMapping(ChildResourceManager):
    """GCP resource:
    https://cloud.google.com/appengine/docs/admin-api/reference/rest/v1/apps.domainMappings
    """
    def _get_parent_resource_info(self, child_instance):
        return {'resourceName': re.match(
            '(apps/.*?)/domainMappings/.*', child_instance['name']).group(1)}

    class resource_type(ChildTypeInfo):
        service = 'appengine'
        version = 'v1'
        component = 'apps.domainMappings'
        enum_spec = ('list', 'domainMappings[]', None)
        scope = None
        name = "name"
        id = 'id'
        default_report_fields = [id, name]
        parent_spec = {
            'resource': 'app-engine',
            'child_enum_params': {
                ('id', 'appsId')
            }
        }
        permissions = ('appengine.applications.get',)

        @staticmethod
        def get(client, resource_info):
            apps_id, mapping_id = re.match('apps/(.*?)/domainMappings/(.*)',
                                           resource_info['resourceName']).groups()
            return client.execute_query('get', {'appsId': apps_id,
                                                'domainMappingsId': mapping_id})

@resources.register('app-engine-firewall-ingress-rule')
class AppEngineFirewallIngressRule(ChildResourceManager):
    """GCP resource:
    https://cloud.google.com/appengine/docs/admin-api/reference/rest/v1/apps.firewall.ingressRules
    """
    def _get_parent_resource_info(self, child_instance):
        return {'resourceName': 'apps/%s' %
                                local_session(self.session_factory).get_default_project()}

    class resource_type(ChildTypeInfo):
        service = 'appengine'
        version = 'v1'
        component = 'apps.firewall.ingressRules'
        enum_spec = ('list', 'ingressRules[]', None)
        scope = None
        name = id = 'priority'
        parent_spec = {
            'resource': 'app-engine',
            'child_enum_params': {
                ('id', 'appsId')
            }
        }
        default_report_fields = ['priority', 'action', 'sourceRange', 'description']
        permissions = ('appengine.applications.get',)

        @staticmethod
        def get(client, resource_info):
            apps_id, ingress_rules_id = re.match('apps/(.*?)/firewall/ingressRules/(.*)',
                                                 resource_info['resourceName']).groups()
            return client.execute_query(
                'get', {'appsId': apps_id,
                        'ingressRulesId': ingress_rules_id})


import os

import pytest

from archai.discrete_search.algos.regularized_evolution import (
    RegularizedEvolutionSearch,
)

@pytest.fixture(scope="session")
def output_dir(tmp_path_factory):
    return tmp_path_factory.mktemp("out_re")

def test_regularized_evolution(output_dir, search_space, search_objectives, surrogate_model):
    algo = RegularizedEvolutionSearch(
        search_space=search_space,
        search_objectives=search_objectives,
        output_dir=output_dir,
        num_iters=5,
        init_num_models=4,
        pareto_sample_size=4,
        history_size=10,
        seed=1,
    )

    search_results = algo.search()
    assert len(os.listdir(output_dir)) > 0

    df = search_results.get_search_state_df()
    assert all(0 <= x <= 0.4 for x in df["Random1"].tolist())

    all_models = [m for iter_r in search_results.results for m in iter_r["models"]]

    # Checks if all registered models satisfy constraints
    _, valid_models = search_objectives.validate_constraints(all_models)
    assert len(valid_models) == len(all_models)


'''Train RAVDESS with PyTorch.'''
# 10 crop for data enhancement
from __future__ import print_function

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import torch.backends.cudnn as cudnn
import numpy as np
import os
import time
import argparse
import utils
from data.IEMOCAP import IEMOCAP
from data.COMMANDS import COMMANDS
from torch.autograd import Variable
from EdgeRNN import EdgeRNN
from EdgeRNNG import EdgeRNNG
from thop import profile
#from audiomentations import Compose, AddGaussianNoise, TimeStretch, PitchShift, Shift     ##Data augmentation

parser = argparse.ArgumentParser(description='PyTorch RAVDESS RAVDESS Training')
parser.add_argument('--model', type=str, default='EdgeRNN', help='CNN architecture')
parser.add_argument('--dataset', type=str, default='IEMOCAP', help='CNN dataset')
parser.add_argument('--train_bs', default=64, type=int, help='learning rate')
parser.add_argument('--test_bs', default=64, type=int, help='learning rate')
parser.add_argument('--lr', default=0.01, type=float, help='learning rate')
opt = parser.parse_args()

use_cuda = torch.cuda.is_available()
best_PrivateTest_acc = 0  # best PrivateTest accuracy
best_PrivateTest_acc_epoch = 0
start_epoch = 0  # start from epoch 0 or last checkpoint epoch

learning_rate_decay_start = 80  # 50
learning_rate_decay_every = 5 # 5
learning_rate_decay_rate = 0.9 # 0.9

total_epoch = 10000

total_prediction_fps = 0
total_prediction_n = 0

# Data
print('==> Preparing data..')
transform_train = Compose([
    AddGaussianNoise(min_amplitude=0.001, max_amplitude=0.015, p=0.5),
    TimeStretch(min_rate=0.8, max_rate=1.25, p=0.5),
    PitchShift(min_semitones=-4, max_semitones=4, p=0.5),
    Shift(min_fraction=-0.5, max_fraction=0.5, p=0.5),
])

if opt.dataset  == 'COMMANDS':
    print ("This is COMMANDS")
    NUM_CLASSES = 11
    trainset = COMMANDS(split = 'Training', transform=transform_train) #If you do not use data augmentation, transform = None.
    trainloader = torch.utils.data.DataLoader(trainset, batch_size=opt.train_bs, shuffle=True, num_workers=4)
    PrivateTestset = COMMANDS(split = 'PrivateTest')
    PrivateTestloader = torch.utils.data.DataLoader(PrivateTestset, batch_size=opt.test_bs, shuffle=False, num_workers=1)

else:
    print ("This is IEMOCAP")
    NUM_CLASSES = 4
    trainset = IEMOCAP(split = 'Training', transform=transform_train)  #If you do not use data augmentation, transform = None.
    trainloader = torch.utils.data.DataLoader(trainset, batch_size=opt.train_bs, shuffle=True, num_workers=4)
    PrivateTestset = IEMOCAP(split = 'PrivateTest')
    PrivateTestloader = torch.utils.data.DataLoader(PrivateTestset, batch_size=opt.test_bs, shuffle=False, num_workers=1)

if opt.model == 'EdgeRNN':
    net = EdgeRNN(num_class=NUM_CLASSES)
elif opt.model == 'EdgeRNNG':
    net = EdgeRNNG(num_class=NUM_CLASSES)

if opt.dataset  == 'COMMANDS':
    flops, params = profile(net, input_size=(1, 152, 32))
else:
    flops, params = profile(net, input_size=(1, 152, 181))

print("The FLOS of this model is  %0.3f M" % float(flops/1024/1024))
print("The params of this model is  %0.3f M" % float(params/1024/1024))

if opt.dataset == 'IEMOCAP' and opt.model == 'EdgeRNNG' and os.path.exists('EdgeRNNG_IEMOCAP_model.pt'):
    # Load checkpoint.
    print('==> Resuming from checkpoint EdgeRNNG_IEMOCAP_model.pt ..')
    Private_checkpoint = torch.load('EdgeRNNG_IEMOCAP_model.pt')
    best_PrivateTest_acc = Private_checkpoint['best_PrivateTest_acc']
    best_PrivateTest_acc_epoch = Private_checkpoint['best_PrivateTest_acc_epoch']
    print ('best_PrivateTest_acc is '+ str(best_PrivateTest_acc))
    net.load_state_dict(Private_checkpoint['net'], strict=False)
    start_epoch = Private_checkpoint['best_PrivateTest_acc_epoch'] + 1

if opt.dataset == 'IEMOCAP' and opt.model == 'EdgeRNN' and os.path.exists('EdgeRNN_IEMOCAP_model.pt'):
    # Load checkpoint.
    print('==> Resuming from checkpoint EdgeRNN_IEMOCAP_model.pt ..')
    Private_checkpoint = torch.load('EdgeRNN_IEMOCAP_model.pt')
    best_PrivateTest_acc = Private_checkpoint['best_PrivateTest_acc']
    best_PrivateTest_acc_epoch = Private_checkpoint['best_PrivateTest_acc_epoch']
    print ('best_PrivateTest_acc is '+ str(best_PrivateTest_acc))
    net.load_state_dict(Private_checkpoint['net'], strict=False)
    start_epoch = Private_checkpoint['best_PrivateTest_acc_epoch'] + 1

if opt.dataset == 'COMMANDS' and opt.model == 'EdgeRNN' and os.path.exists('EdgeRNN_COMMANDS_model.pt'):
    # Load checkpoint.
    print('==> Resuming from checkpoint EdgeRNN_COMMANDS_model.pt ..')
    Private_checkpoint = torch.load('EdgeRNN_COMMANDS_model.pt')
    best_PrivateTest_acc = Private_checkpoint['best_PrivateTest_acc']
    best_PrivateTest_acc_epoch = Private_checkpoint['best_PrivateTest_acc_epoch']
    print ('best_PrivateTest_acc is '+ str(best_PrivateTest_acc))
    net.load_state_dict(Private_checkpoint['net'], strict=False)
    start_epoch = Private_checkpoint['best_PrivateTest_acc_epoch'] + 1

if opt.dataset == 'COMMANDS' and opt.model == 'EdgeRNNG' and os.path.exists('EdgeRNNG_COMMANDS_model.pt'):
    # Load checkpoint.
    print('==> Resuming from checkpoint EdgeRNNG_COMMANDS_model.pt ..')
    Private_checkpoint = torch.load('EdgeRNNG_COMMANDS_model.pt')
    best_PrivateTest_acc = Private_checkpoint['best_PrivateTest_acc']
    best_PrivateTest_acc_epoch = Private_checkpoint['best_PrivateTest_acc_epoch']
    print ('best_PrivateTest_acc is '+ str(best_PrivateTest_acc))
    net.load_state_dict(Private_checkpoint['net'], strict=False)
    start_epoch = Private_checkpoint['best_PrivateTest_acc_epoch'] + 1

else:
    print('==> Building model..')

if use_cuda:
    net.cuda()

criterion = nn.CrossEntropyLoss()
criterion.cuda()
# optimizer = optim.Adam(net.parameters(), weight_decay=5e-4)
optimizer = optim.SGD(net.parameters(), lr=opt.lr, momentum=0.9, weight_decay=5e-4)

# Training
def train(epoch):
    print('\nEpoch: %d' % epoch)
    global Train_acc
    net.train()
    train_loss = 0
    correct = 0
    total = 0
    conf_mat = np.zeros((NUM_CLASSES, NUM_CLASSES))

    if epoch > learning_rate_decay_start and learning_rate_decay_start >= 0:
        frac = (epoch - learning_rate_decay_start) // learning_rate_decay_every
        decay_factor = learning_rate_decay_rate ** frac
        current_lr = opt.lr * decay_factor
        utils.set_lr(optimizer, current_lr)  # set the decayed rate
    else:
        current_lr = opt.lr
    print('learning_rate: %s' % str(current_lr))

    for batch_idx, (inputs, targets) in enumerate(trainloader):
        if use_cuda:
            inputs, targets = inputs.cuda(), targets.cuda()
        optimizer.zero_grad()
        inputs, targets = Variable(inputs), Variable(targets)
        outputs = net(inputs)
        loss = criterion(outputs, targets)
        loss.backward()
        utils.clip_gradient(optimizer, 0.1)
        optimizer.step()
        train_loss += loss.item()

        conf_mat += utils.confusion_matrix(outputs, targets, NUM_CLASSES)

        acc = sum([conf_mat[i, i] for i in range(conf_mat.shape[0])])/conf_mat.sum()
        uacc_per_class = [conf_mat[i, i]/conf_mat[i].sum() for i in range(conf_mat.shape[0])]
        unweighted_acc = sum(uacc_per_class)/len(uacc_per_class)

        prec_per_class = [conf_mat[i, i] / conf_mat[:, i].sum() for i in range(conf_mat.shape[0])]
        average_precision = sum(prec_per_class)/len(prec_per_class)

        utils.progress_bar(batch_idx, len(trainloader), 'Loss: %.3f | Acc: %.3f%% | unweighted_Acc: %.3f%%'
            % (train_loss/(batch_idx+1), 100.* acc, 100.* unweighted_acc))

    Train_acc = 100.* acc

def PrivateTest(epoch):
    global PrivateTest_acc
    global best_PrivateTest_acc
    global best_PrivateTest_acc_epoch
    global total_prediction_fps
    global total_prediction_n
    net.eval()
    PrivateTest_loss = 0
    correct = 0
    total = 0
    t_prediction = 0
    conf_mat = np.zeros((NUM_CLASSES, NUM_CLASSES))
    for batch_idx, (inputs, targets) in enumerate(PrivateTestloader):
        t = time.time()
        if use_cuda:
            inputs, targets = inputs.cuda(), targets.cuda()
        inputs, targets = Variable(inputs), Variable(targets)
        outputs = net(inputs)
        t_prediction += (time.time() - t)

        loss = criterion(outputs, targets)
        PrivateTest_loss += loss.item()

        conf_mat += utils.confusion_matrix(outputs, targets, NUM_CLASSES)

        acc = sum([conf_mat[i, i] for i in range(conf_mat.shape[0])])/conf_mat.sum()
        uacc_per_class = [conf_mat[i, i]/conf_mat[i].sum() for i in range(conf_mat.shape[0])]
        unweighted_acc = sum(uacc_per_class)/len(uacc_per_class)

        prec_per_class = [conf_mat[i, i] / conf_mat[:, i].sum() for i in range(conf_mat.shape[0])]
        average_precision = sum(prec_per_class)/len(prec_per_class)

        utils.progress_bar(batch_idx, len(PrivateTestloader), 'Loss: %.3f | Acc: %.3f%% | unweighted_Acc: %.3f%%'
            % (PrivateTest_loss / (batch_idx + 1), 100. * acc, 100.* unweighted_acc))
    total_prediction_fps = total_prediction_fps + (1 / (t_prediction / len(PrivateTestloader)))
    total_prediction_n = total_prediction_n + 1
    print('Prediction time: %.2f' % t_prediction + ', Average : %.5f/image' % (t_prediction / len(PrivateTestloader))
         + ', Speed : %.2fFPS' % (1 / (t_prediction / len(PrivateTestloader))))

    # Save checkpoint.
    PrivateTest_acc = 100.* acc
    if PrivateTest_acc > best_PrivateTest_acc:
        print('Saving..')
        print("best_PrivateTest_acc: %0.3f" % PrivateTest_acc)
        state = {
            'net': net.state_dict() if use_cuda else net,
            'best_PrivateTest_acc': PrivateTest_acc,
            'best_PrivateTest_acc_epoch': epoch,
        }
        if opt.dataset == 'COMMANDS' and opt.model == 'EdgeRNN':
            torch.save(state, 'EdgeRNN_COMMANDS_model.pt')
        if opt.dataset == 'COMMANDS' and opt.model == 'EdgeRNNG':
            torch.save(state, 'EdgeRNNG_COMMANDS_model.pt')
        if opt.dataset == 'IEMOCAP' and opt.model == 'EdgeRNN':
            torch.save(state, 'EdgeRNN_IEMOCAP_model.pt')
        if opt.dataset == 'IEMOCAP' and opt.model == 'EdgeRNNG':
            torch.save(state, 'EdgeRNNG_IEMOCAP_model.pt')
        best_PrivateTest_acc = PrivateTest_acc
        best_PrivateTest_acc_epoch = epoch

for epoch in range(start_epoch, total_epoch):
    train(epoch)
    PrivateTest(epoch)

print("best_PrivateTest_acc: %0.3f" % best_PrivateTest_acc)
print("best_PrivateTest_acc_epoch: %d" % best_PrivateTest_acc_epoch)

print("total_prediction_fps: %0.2f" % total_prediction_fps)
print("total_prediction_n: %d" % total_prediction_n)
print('Average speed: %.2f FPS' % (total_prediction_fps / total_prediction_n))


# for image processing and reading/writing
from PIL import Image
# For matrix operations
import numpy as np

np.seterr(divide='ignore', invalid='ignore')
# for cmd line arguments
import sys

# for plotting the image
import matplotlib as mpl
import matplotlib.pyplot as plt

# Define Colour Conversion, and Drawing Functions
def RGB_HSL(R, G, B):
    # Convert RGB matrixes to HSL and return HSL

    # calculation variables
    r_prime = R / 255
    g_prime = G / 255
    b_prime = B / 255
    C_max_1 = np.maximum(r_prime, np.maximum(g_prime, b_prime))
    C_min_1 = np.minimum(r_prime, np.minimum(g_prime, b_prime))
    Delta = C_max_1 - C_min_1

    # intialize Hue
    Hue_r = np.select([Delta != 0], [60 * (((g_prime - b_prime) / Delta) % 6)])
    Hue_g = np.select([Delta != 0], [60 * (((b_prime - r_prime) / Delta) + 2)])
    Hue_b = np.select([Delta != 0], [60 * (((r_prime - g_prime) / Delta) + 4)])
    condlist = [C_max_1 == r_prime, C_max_1 == g_prime, C_max_1 == b_prime]
    choicelist = [Hue_r, Hue_g, Hue_b]
    Hue = np.select(condlist, choicelist)

    # lightness
    Lightness = (C_max_1 + C_min_1) / 2

    # Saturation (For HSL)
    Saturation_L = np.select([Delta != 0], [Delta / (1 - abs(2 * Lightness - 1))])

    return Hue, Saturation_L, Lightness

def RGB_HSV(R, G, B):
    # Convert RGB matrixes to HSV and return HSV

    # calculation variables
    r_prime = R / 255
    g_prime = G / 255
    b_prime = B / 255
    C_max_1 = np.maximum(r_prime, np.maximum(g_prime, b_prime))
    C_min_1 = np.minimum(r_prime, np.minimum(g_prime, b_prime))
    Delta = C_max_1 - C_min_1

    # intialize Hue
    Hue_r = np.select([Delta != 0], [60 * (((g_prime - b_prime) / Delta) % 6)])
    Hue_g = np.select([Delta != 0], [60 * (((b_prime - r_prime) / Delta) + 2)])
    Hue_b = np.select([Delta != 0], [60 * (((r_prime - g_prime) / Delta) + 4)])
    condlist = [C_max_1 == r_prime, C_max_1 == g_prime, C_max_1 == b_prime]
    choicelist = [Hue_r, Hue_g, Hue_b]
    Hue = np.select(condlist, choicelist)

    # Saturation (For HSV)
    Saturation_V = np.select([C_max_1 != 0], [(Delta) / (C_max_1)])

    # Value
    Value = C_max_1

    return Hue, Saturation_V, Value

def HSV_to_RGB(H, S, V):
    # Convert HSV matrixes to RGB and return RGB

    C = V * S
    X = C * (1 - abs((H / 60) % 2 - 1))
    m = V - C

    Cond = np.array([(H >= 0) & (H < 60), (H >= 60) & (H < 120), (H >= 120) & (H < 180), (H >= 180) & (H < 240),
                     (H >= 240) & (H < 300), (H >= 300) & (H < 360)])
    R = np.floor((np.select(Cond, [C, X, 0, 0, X, C]) + m) * 255)
    G = np.floor((np.select(Cond, [X, C, C, X, 0, 0]) + m) * 255)
    B = np.floor((np.select(Cond, [0, 0, X, C, C, X]) + m) * 255)

    return R, G, B

def HSL_to_RGB(H, S, L):
    # Convert HSL matrixes to RGB and return RGB

    C = (1 - abs(2 * L - 1)) * S
    X = C * (1 - abs((H / 60) % 2 - 1))
    m = L - C / 2

    Cond = np.array([(H >= 0) & (H < 60), (H >= 60) & (H < 120), (H >= 120) & (H < 180), (H >= 180) & (H < 240),
                     (H >= 240) & (H < 300), (H >= 300) & (H < 360)])
    R = np.floor((np.select(Cond, [C, X, 0, 0, X, C]) + m) * 255)
    G = np.floor((np.select(Cond, [X, C, C, X, 0, 0]) + m) * 255)
    B = np.floor((np.select(Cond, [0, 0, X, C, C, X]) + m) * 255)

    return R, G, B

def drawing_image(image):
    # updates matplotlib figure
    # these commands are reqired to keep the figure from pausing the script
    ax.imshow(image, aspect='auto')
    plt.ion()
    plt.draw()
    plt.pause(0.001)

# Image grayscale, blurring, rotating, etc functions
def greyscale(im_matrix):
    temp_image = np.empty_like(im_matrix)
    r_grey = im_matrix[:, :, 0]
    g_grey = im_matrix[:, :, 1]
    b_grey = im_matrix[:, :, 2]

    temp = (0.3 * r_grey) + (0.59 * g_grey) + (0.11 * b_grey)

    temp_image[:, :, 0] = temp
    temp_image[:, :, 1] = temp
    temp_image[:, :, 2] = temp
    return temp_image

def image_blur(matrix):
    # return blurred matrix based on user input
    temp_image = np.empty_like(matrix)

    # make count matrix for averaging
    blur_dict_count = {}
    for a in range(9):
        blur_dict_count[a] = np.ones_like(matrix[:, :, 0])
    index = 0
    for i in range(3):
        for j in range(3):
            blur_dict_count[index] = np.pad(blur_dict_count[index], ((i, 2 - i), (j, 2 - j)))
            # print(blur_dict_count[index][0:10,0:10])
            index = index + 1

    count = np.zeros_like(blur_dict_count[0])
    for b in range(9):
        count = count + blur_dict_count[b]

    count_1 = count
    # final count matrix for proper division
    count = count[1:-1, 1:-1]

    # determine average for each channel
    for color in range(3):
        # print(color)
        color_matrix = matrix[:, :, color]
        blur_dict_colour = {}
        for a in range(9):
            blur_dict_colour[a] = color_matrix

        index = 0
        for i in range(3):
            for j in range(3):
                blur_dict_colour[index] = np.pad(blur_dict_colour[index], ((i, 2 - i), (j, 2 - j)))
                index = index + 1

        sum_color = np.zeros_like(blur_dict_colour[0])
        for c in range(9):
            # print(blur_dict_colour[c][0:10,0:10], "\n\n")
            sum_color = sum_color + (blur_dict_colour[c] / count_1)
            # print(sum_color[0:10,0:10], "\n\n")

        sum_color_final = sum_color[1:-1, 1:-1]
        average_color = sum_color_final  # / count

        # print(sum_color_final[0:10,0:10])
        # print(count[0:10,0:10])
        temp_image[:, :, color] = average_color.astype(int)

    return temp_image

    # 4 is original matrix, padded with zeroes on all sides
    # 0 1 2
    # 3 4 5
    # 6 7 8

def rotate(matrix):
    # return rotated matrix based on user input
    direction = input("Direction Left or right?: ").upper()
    if direction[0] == "L":
        direct = (0, 1)
    elif direction[0] == "R":
        direct = (1, 0)
    else:
        exit(1)

    try:
        rotations = input("Number of rotations?: ")
        rotations = abs(int(rotations))

    except:
        print("Use an integer")
        rotations = 0

    shape = matrix.shape
    if len(shape) == 3:
        temp_rotate = np.rot90(matrix, rotations, direct)
    if len(shape) == 2:
        temp_rotate = np.rot90(matrix, rotations, direct)
    return temp_rotate

# Main script starts here
arguments = sys.argv
if len(arguments) != 2:
    print("CMD line argument needed,include file name")
    exit(1)
else:
    image_name = arguments[1]
im = Image.open(image_name)

# initialize changes to dictionary storing changes
changes_RGB = {"R": 1, "G": 1, "B": 1}
changes_HSV = {"H": 0, "S": 1, "V": 1}
changes_HSL = {"H": 0, "S": 1, "L": 1}

# Get R,G,B Matrixes
data = np.array(im)
r_matrix = data[:, :, 0]
g_matrix = data[:, :, 1]
b_matrix = data[:, :, 2]

# initialize "preview" Window
fig = plt.figure(frameon=False)
ax = plt.Axes(fig, [0., 0., 1., 1.])
ax.set_axis_off()
fig.add_axes(ax)
drawing_image(data)

# intialize HSV, and HSL image matrices
H_HSV, S_HSV, V_HSV = RGB_HSV(r_matrix, g_matrix, b_matrix)
H_HSL, S_HSL, L_HSL = RGB_HSL(r_matrix, g_matrix, b_matrix)

# Default change to RGB, channel R
current_change = "RGB"
changing = "R"

# Initialize Temporary Matrices for each of the channels (Default to original image)
img_temp = data
img_temp_HSV_tmp = data
img_temp_HSL_tmp = data

# starting command line interface loops
while True:
    print("Currently Changing:", current_change)
    if input("Done making changes?: ") == "yes":
        break

    # Modify image using RGB channels
    while current_change == "RGB":
        # Change Channel
        if changing == "R":
            print("\n(R:", changes_RGB["R"], ") G:", changes_RGB["G"], " B:", changes_RGB["B"], sep='')
            try:
                changes_RGB["R"] = abs(float(input("Change R to:")))
            except:
                print("Use a valid number")
        if changing == "G":
            print("\nR:", changes_RGB["R"], " (G:", changes_RGB["G"], ") B:", changes_RGB["B"], sep='')

            try:
                changes_RGB["G"] = abs(float(input("Change G to:")))
            except:
                print("Use a valid number")

        if changing == "B":
            print("\nR:", changes_RGB["R"], " G:", changes_RGB["G"], " (B:", changes_RGB["B"], ")", sep='')

            try:
                changes_RGB["B"] = abs(float(input("Change B to:")))
            except:
                print("Use a valid number")

        # Update each of the temporary channels
        R_temp = np.minimum(255, r_matrix.astype(int) * changes_RGB["R"])
        G_temp = np.minimum(255, g_matrix.astype(int) * changes_RGB["G"])
        B_temp = np.minimum(255, b_matrix.astype(int) * changes_RGB["B"])
        img_temp[:, :, 0] = R_temp
        img_temp[:, :, 1] = G_temp
        img_temp[:, :, 2] = B_temp

        # draw updated image
        drawing_image(img_temp)

        # give user change to change "channel" or "channel type"
        tmp = input("change colour stream (back to change colour style): ")
        if tmp.upper() == "R" or tmp.upper() == "G" or tmp.upper() == "B":
            changing = tmp.upper()
        elif tmp.upper() == "BACK":
            tmp_2 = input("new style 'HSV' or 'HSL': ")
            if tmp_2.upper() == "HSV" or tmp_2.upper() == "HSL":
                current_change = tmp_2.upper()
                changing = "H"
            break
    print("Currently Changing:", current_change)
    if input("Done making changes?: ") == "yes":
        break

    # Modify image using HSV channels
    while current_change == "HSV":
        change = False
        if changing == "H":
            print("\n(H:", changes_HSV["H"], ") S:", changes_HSV["S"], " V:", changes_HSV["V"], sep='')
            try:
                tmp = changes_HSV["H"]
                changes_HSV["H"] = abs(float(input("Change H to: ")))
                if tmp != changes_HSV["H"]:
                    change = True
            except:
                print("Use a valid number")
        if changing == "S":
            print("\nH:", changes_HSV["H"], " (S:", changes_HSV["S"], ") V:", changes_HSV["V"], sep='')
            try:
                tmp = changes_HSV["S"]
                changes_HSV["S"] = abs(float(input("Change S to: ")))
                if tmp != changes_HSV["S"]:
                    change = True
            except:
                print("Use a valid number")

        if changing == "V":
            print("\nH:", changes_HSV["H"], " S:", changes_HSV["S"], " (V:", changes_HSV["V"], ")", sep='')
            try:
                tmp = changes_HSV["V"]
                changes_HSV["V"] = abs(float(input("Change V to: ")))
                if tmp != changes_HSV["V"]:
                    change = True
            except:
                print("Use a valid number")

        # Update each of the temporary channels,
        # if it has changed from the last instance
        if change == True:
            H_HSV_tmp = (H_HSV + changes_HSV["H"]) % 360
            S_HSV_tmp = np.minimum(S_HSV * changes_HSV["S"], 1)
            V_HSV_tmp = np.minimum(V_HSV * changes_HSV["V"], 1)

            R_HSV_tmp, G_HSV_tmp, B_HSV_tmp = HSV_to_RGB(H_HSV_tmp, S_HSV_tmp, V_HSV_tmp)

            R_HSV_tmp = R_HSV_tmp.astype(int)
            G_HSV_tmp = G_HSV_tmp.astype(int)
            B_HSV_tmp = B_HSV_tmp.astype(int)

            img_temp_HSV_tmp[:, :, 0] = R_HSV_tmp
            img_temp_HSV_tmp[:, :, 1] = G_HSV_tmp
            img_temp_HSV_tmp[:, :, 2] = B_HSV_tmp

            # draw updated image
            drawing_image(img_temp_HSV_tmp)

        # give user change to change "channel" or "channel type"
        tmp = input("change HSV stream (back to change colour style): ")
        if tmp.upper() == "H" or tmp.upper() == "S" or tmp.upper() == "V":
            changing = tmp.upper()
        elif tmp.upper() == "BACK":
            tmp_2 = input("new style 'RGB' or 'HSL': ")
            if tmp_2.upper() == "RGB":
                current_change = tmp_2.upper()
                changing = "R"
            elif tmp_2.upper() == "HSL":
                current_change = tmp_2.upper()
                changing = "H"
            break

    print("Currently Changing:", current_change)
    if input("Done making changes?: ") == "yes":
        break

    # Modify image using HSL channels
    while current_change == "HSL":
        change = False
        if changing == "H":
            print("\n(H:", changes_HSL["H"], ") S:", changes_HSL["S"], " L:", changes_HSL["L"], sep='')
            try:
                tmp = changes_HSL["H"]
                changes_HSL["H"] = abs(float(input("Change H to: ")))
                if tmp != changes_HSL["H"]:
                    change = True
            except:
                print("Use a valid number")
        if changing == "S":
            print("\nH:", changes_HSL["H"], " (S:", changes_HSL["S"], ") L:", changes_HSL["L"], sep='')
            try:
                tmp = changes_HSL["S"]
                changes_HSL["S"] = abs(float(input("Change S to: ")))
                if tmp != changes_HSL["S"]:
                    change = True
            except:
                print("Use a valid number")

        if changing == "L":
            print("\nH:", changes_HSL["H"], " S:", changes_HSL["S"], " (L:", changes_HSL["L"], ")", sep='')
            try:
                tmp = changes_HSL["L"]
                changes_HSL["L"] = abs(float(input("Change L to: ")))
                if tmp != changes_HSL["L"]:
                    change = True
            except:
                print("Use a valid number")

        # Update each of the temporary channels,
        # if it has changed from the last instance
        if change == True:
            H_HSL_tmp = (H_HSL + changes_HSL["H"]) % 360
            S_HSL_tmp = np.minimum(S_HSL * changes_HSL["S"], 1)
            L_HSL_tmp = np.minimum(L_HSL * changes_HSL["L"], 1)

            R_HSL_tmp, G_HSL_tmp, B_HSL_tmp = HSL_to_RGB(H_HSL_tmp, S_HSL_tmp, L_HSL_tmp)

            R_HSL_tmp = R_HSL_tmp.astype(int)
            G_HSL_tmp = G_HSL_tmp.astype(int)
            B_HSL_tmp = B_HSL_tmp.astype(int)

            img_temp_HSL_tmp[:, :, 0] = R_HSL_tmp
            img_temp_HSL_tmp[:, :, 1] = G_HSL_tmp
            img_temp_HSL_tmp[:, :, 2] = B_HSL_tmp

            # draw updated image
            drawing_image(img_temp_HSL_tmp)

        # give user change to change "channel" or "channel type"
        tmp = input("change HSL stream (back to change colour style): ")
        if tmp.upper() == "H" or tmp.upper() == "S" or tmp.upper() == "L":
            changing = tmp.upper()
        elif tmp.upper() == "BACK":
            tmp_2 = input("new style 'RGB' or 'HSV': ")
            if tmp_2.upper() == "RGB":
                current_change = tmp_2.upper()
                changing = "R"
            elif tmp_2.upper() == "HSV":
                current_change = tmp_2.upper()
                changing = "H"
            break

# update Actual Image Matrix, after these changes are made

if current_change == "RGB":
    img_final = img_temp

elif current_change == "HSV":
    img_final = img_temp_HSV_tmp

elif current_change == "HSL":
    img_final = img_temp_HSL_tmp

else:
    print("Error Updating true image matrix")
    exit(1)

drawing_image(img_final)

# final updates, ie: rotate, grayscale
while True:
    print("Options are: BW, Blur, Rotate, Done")
    action = input("What would you like to do?: ").upper()

    if action == "BW":
        img_final = greyscale(img_final)
        drawing_image(img_final)

    elif action == "BLUR":
        img_final = image_blur(img_final)
        drawing_image(img_final)

    elif action == "ROTATE":
        img_final = rotate(img_final)
        drawing_image(img_final)

    elif action == "DONE":
        break
    else:
        print("Invalid Input, Please try again")

# file saving stuff
save = input("Save file? (Y or N): ").upper()
# name sanitization
if save[0] == "Y" or save[0] == "1":
    f_name = input("File name :")
    print(f_name[-5:])
    if f_name[-5:-1] == ".jpeg":
        f_name = f_name

    else:
        f_name = f_name.split('.')[0] + ".jpeg"

    # actually save name now
    test = Image.fromarray(img_final)
    test.save(f_name, "JPEG")
    print("File saved as", f_name)

print("Done")


import pandas as pd
from tqdm import tqdm
import copy

def create_lookup_table():
    # get all allowed words
    with open('allowed_words.txt', 'r') as file:
        possible_words = file.read().split('\n')
    lookup_table = pd.DataFrame(index=possible_words, columns=possible_words)
    char_lookup_table = {}
    progress_bar = tqdm(total=len(possible_words), desc='Lookup table')
    for i in range(len(possible_words)):  # loop through all words
        word_new_row, list_word_new_row = get_charlist_from_word(possible_words, i)
        char_lookup_table[word_new_row] = list_word_new_row
        for j in range(i):  # for all already processed words
            word_other_row = possible_words[j]  # get word
            list_word_other_row = copy.deepcopy(char_lookup_table[word_other_row])
            value = calculate_value(copy.deepcopy(list_word_new_row), list_word_other_row)
            lookup_table.loc[word_new_row][word_other_row] = value
            lookup_table.loc[word_other_row][word_new_row] = value
        lookup_table.loc[word_new_row][word_new_row] = '500'
        progress_bar.update(1)
    progress_bar.close()
    split_lookup_table(lookup_table)

# input: word like 'house'
# output: ('house', [('h', 0), ('o', 1), ('u', 2), ('s', 3), ('e', 4)])
def get_charlist_from_word(possible_words, i):
    word_new_row = possible_words[i]  # get word
    list_word_new_row = []  # get char from word
    for c in range(len(word_new_row)):
        list_word_new_row.append((word_new_row[c], c))  # c indicates position in word
    return word_new_row, list_word_new_row

# input: 2 words
# output: amount of greens, yellows, and reds
def calculate_value(copy_word_new_row, list_word_other_row):
    green = 0
    yellow = 0
    while len(copy_word_new_row) > 0:  # while still char needs to be checked
        c1, pos1 = copy_word_new_row.pop(0)  # get first char
        # define a few variables
        lowest_pos = 100
        c3 = ''
        for c2, pos2 in list_word_other_row:  # for every char still in list_word_other_row
            if c1 == c2 and pos1 == pos2:  # check for greens
                green += 1
                list_word_other_row.remove((c2, pos2))
                lowest_pos = 100  # reset value in case we did find a yellow
                break
            elif c1 == c2 and (c2, pos2) not in copy_word_new_row:  # check for yellows and making sure there is not a green later on
                lowest_pos = pos2
                c3 = c2
        if lowest_pos != 100:  # if yellows but no green found
            yellow += 1
            list_word_other_row.remove((c3, lowest_pos))
    # adding data to both tables
    return str(green) + str(yellow) + str(5 - green - yellow)

# Because github has a file limit of 100MB, I split up the files to be accessable if someone runs my files from github
def split_lookup_table(lookup_table):
    progress_bar = tqdm(total=13, desc='Creating parts')
    for i in range(13):
        df_part = lookup_table.iloc[1000 * i:1000 * (i + 1)]
        df_part.to_csv('../lookup_table_part/part' + str(i) + '.txt')
        progress_bar.update(1)
    progress_bar.close()

def load_lookup_table(folder=False):
    file = ''
    if folder:
        file += '../'
    file += 'lookup_table_part/part0.txt'
    lookup_table = pd.read_csv(file, index_col=[0], dtype=str)
    for i in range(1, 13):
        file = ''
        if folder:
            file += '../'
        file += 'lookup_table_part/part' + str(i) + '.txt'
        part = pd.read_csv(file, index_col=[0], dtype=str)
        lookup_table = pd.concat([lookup_table, part])
    return lookup_table

if __name__ == '__main__':
    create_lookup_table()
    split_lookup_table()


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

import logging
import yaml
import jinja2

import admiral.exception as exc

# Try LibYAML first and if unavailable, fall back to pure Python implementation
try:
    from yaml import CLoader as Loader, CDumper as Dumper
except ImportError:
    from yaml import Loader, Dumper

LOG = logging.getLogger(__name__)

def _load_config(config_file):
    """Load configuration

    Args:
        config_file: path to the file that should be read

    Returns:
        Configuration hash - config file after parsing.
    """
    try:
        with open(config_file) as fh:
            config = yaml.safe_load(fh)
    except yaml.YAMLError as e:
        msg = "Failed to parse config file {0}: {1}"
        raise exc.UserInputException(msg.format(config_file, str(e)))

    return config

def _validate_config(config):
    # TODO - write config validation using i.e. cerberus library:
    # http://docs.python-cerberus.org/en/latest/
    pass

def _apply_envvars(config, env):
    """Fill in all template placeholders with environment data

    Some of the variables depend on envrionment. In order to provide
    different setup for different environments, we template the configuration
    and with environment data.

    The substituion is simple - change the configuration hash back to plaintext,
    use jinja to template with the data from environments seciton, and then
    convet it back to hash.

    Args:
        config: configuration hash that should be templated
        env: environent data to use

    Returns:
        Configuration hash destined for particular environment
    """
    LOG.debug("Pre-templated config: {0}".format(config))
    tmp = yaml.dump(config, Dumper=Dumper, default_flow_style=False)
    try:
        template = jinja2.Template(tmp)
        rendered_template = template.render(config["environments"][env])
        res = yaml.load(rendered_template, Loader=Loader)
    except (yaml.YAMLError, jinja2.TemplateError) as e:
        msg = "Syntax errors found while templating variables: {0}"
        raise exc.UserInputException(msg.format(str(e)))
    LOG.info("Templated pods configuration: {0}".format(res))
    return res

def _normalize_container_names(config, env):
    """Expand container names with pod and environement data

    Container names defined in configuration file do not take into account
    different envrionments where the configuration can be deployed. The purpose
    of this function is to rename all the container names that can be found in
    configuration file so that pod and environemnt is reflected. I.E.

        redis becomes dev_giant-weather_redis (env: dev, pod: giant-weather)

    Places where this renaming is necessary include:
        - `machine_of` field
        - `binds_to` list
        - `links` list
        - container names in each pod hash

    Args:
        config: hash containing the configuration for processing
        env: target environment for the script

    Returns:
        Configuration with all the names fixed
    """
    for pod in config['pods']:
        new_h = {}
        h = config['pods'][pod]['containers']
        for old_name in h:
            new_name_f = env + "_" + pod + "_{0}"
            # Fix links:
            if 'links' in h[old_name]:
                new_links = []
                old_links = h[old_name]['links']
                for link in old_links:
                    c, alias = link.split(':')
                    new_links.append(new_name_f.format(c) + ':' + alias)
                h[old_name]['links'] = new_links

            # Fix machine_of:
            if 'machine_of' in h[old_name]:
                h[old_name]['machine_of'] = new_name_f.format(h[old_name]['machine_of'])

            # Fix binds_to:
            if 'binds_to' in h[old_name]:
                h[old_name]['binds_to'] = [new_name_f.format(unit) for unit in h[old_name]['binds_to']]

            # Replace the hash:
            new_h[new_name_f.format(old_name)] = h[old_name]
        config['pods'][pod]['containers'] = new_h

def load_config(config_file, environment):
    """Load and process the configuration

    This function takes care of all processing neccesary for the config file,
    such as templating, validation, loading, etc...

    Args:
        config_file: location of the config file
        environment: target environment that will be used while reconfiguring
                     the cluster
    Returns:
        Parsed config file
    """
    config = _load_config(config_file)
    _validate_config(config)
    config = _apply_envvars(config, environment)
    _normalize_container_names(config, environment)
    return config


import random
import time
import csv
import cv2
from keras.models import load_model
import numpy as np

# WORK IN PROGRESS
# functional programming paradigm for RPS game using only the camera

print("about to load model")
model = load_model('keras_model.h5') #load model ready to make prediction
print("load model complete")
cap = cv2.VideoCapture(0) # Open a camera for video capturing. To open default camera using default backend just pass 0.
data = np.ndarray(shape=(1, 224, 224, 3), dtype=np.float32) # construct a numpy array
label_file = 'labels.txt'

with open(label_file) as f:
        reader = csv.reader(f, delimiter=' ')
        label_dict = dict(reader)
        print(label_dict)

def get_normalized_image():
    ret, frame = cap.read()
    resized_frame = cv2.resize(frame, (224, 224), interpolation = cv2.INTER_AREA)
    image_np = np.array(resized_frame)
    normalized_image = (image_np.astype(np.float32) / 127.0) - 1 # Normalize the image
    data[0] = normalized_image

    return frame, data

def get_prediction():
    frame, data = get_normalized_image()
    prediction = model.predict(data)
    cv2.imshow('frame', frame)
    # Press q to close the window
    print(prediction)
    prediction_index = np.argmax(prediction)
    print(prediction_index)

    return prediction_index

def get_user_choice():
    prediction_index = get_prediction()
    moves = []
    for k, v in label_dict.items():
        moves.append(v)
    user_choice = moves[prediction_index]
    print(user_choice)

    return user_choice

def get_winner(prediction):
    pass
    # take predicted player move
    # take computer move
    # compute & return winner

#def get_countdown():
    # pass

while True:
    start = time.time()
    get_user_choice()
    end = time.time()
    time_elapsed = end - start
    if time_elapsed >= 3:
        print(f"you chose {user_choice}!")

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

# After the loop release the cap object
cap.release()
# Destroy all the windows
cv2.destroyAllWindows()


import torch
from torch import nn

from data import consts
from models import video_cnn, VGG, transformer, video_decoder, flow_net
from utils import sequence_decoder

class VSADisNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.video_cnn = video_cnn.ThreeDimCNN(out_feature=True)
        self.seq2seq = transformer.Transformer(consts.EMBED_SIZE, num_layers=3, dim_feedforward=1024)
        vocab_size = len(consts.VOCAB)
        self.embed = transformer.TokenEmbedding(vocab_size, consts.EMBED_SIZE)

        self.spatial_net = VGG.vgg13_bn()
        self.static_id_classifier = nn.Linear(consts.EMBED_SIZE, len(consts.VSA_TRAIN_USERS))

        self.flow_net = flow_net.FlowNet()
        self.motion_id_classifier = nn.Linear(consts.EMBED_SIZE, len(consts.VSA_TRAIN_USERS))

        self.video_decoder = video_decoder.VideoDecoder2(norm_mode="adin")

    def forward(self, inputs):
        """
        :param inputs: [original, same_id, same_content]
               -> element: [img_sequence, token, user_id, flow_sequence]
        :return:
        """
        bs = inputs[0][0].shape[0]
        features = self.encoder(inputs)
        concat_tokens = torch.cat((inputs[0][1], inputs[1][1], inputs[2][1]), dim=0)
        content_probs = self.content_recognizer(features[0]["pooled_feature"], concat_tokens)
        static_id_probs = self.static_id_classifier(features[1]["pooled_feature"])
        motion_id_probs = self.motion_id_classifier(features[2])

        features_for_decode = []
        probs = {"content_prob_list": content_probs.split(bs, dim=0),
                 "static_id_prob_list": static_id_probs.split(bs * consts.SAMPLE_IMG_NUM, dim=0),
                 "motion_id_prob_list": motion_id_probs.split(bs, dim=0),
                 "content_feature": features[0]["pooled_feature"].split(bs, dim=0),
                 "static_feature": features[1]["pooled_feature"].split(bs * consts.SAMPLE_IMG_NUM, dim=0),
                 "motion_feature": features[2].split(bs, dim=0),
                 }
        for i in range(3):
            features_for_decode.append({"content": features[0]["mid_feature"][bs * i: bs * (i + 1), ...],
                                        "static": features[1]["mid_feature"][
                                                  bs * consts.SAMPLE_IMG_NUM * i: bs * consts.SAMPLE_IMG_NUM * (i + 1),
                                                  ...]
                                       .reshape(bs, consts.SAMPLE_IMG_NUM, consts.EMBED_SIZE, 3, 6).mean(1),
                                        "flow": features[2][bs * i: bs * (i + 1), ...]})

        videos = self.decoder(features_for_decode)
        return probs, videos

    def encoder(self, inputs):
        """
        Inputs: [original, same_id, same_content]
        Returns: [content_features, static_features, flow_features]
        """
        content_features = self.video_cnn(torch.cat((inputs[0][0], inputs[1][0], inputs[2][0]), dim=0))
        pooled_feature_list = []
        mid_feature_list = []
        for i in range(3):
            img_sequence = inputs[i][0]
            b, c, l, h, w = img_sequence.shape
            img_sequence = img_sequence.permute(0, 2, 1, 3, 4).reshape(b * consts.SAMPLE_IMG_NUM, c, h, w)
            concat_static_features = self.spatial_net(img_sequence)
            pooled_feature_list.append(concat_static_features["pooled_feature"])
            mid_feature_list.append(concat_static_features["mid_feature"])
        static_id_features = {"pooled_feature": torch.cat(pooled_feature_list, dim=0),
                              "mid_feature": torch.cat(mid_feature_list, dim=0)}

        flow_features = self.flow_net(torch.cat((inputs[0][3], inputs[1][3], inputs[2][3])))
        return content_features, static_id_features, flow_features

    def decoder(self, features):
        content_mid_features = []
        static_mid_features = []
        flow_features = []
        for feature in features:
            content_mid_features.append(feature["content"])
            static_mid_features.append(feature["static"])
            flow_features.append(feature["flow"])
        id_a_content_a = [self.decode_one_video(content_mid_features[0], static_mid_features[0], flow_features[0]),
                          self.decode_one_video(content_mid_features[2], static_mid_features[0], flow_features[0]),
                          self.decode_one_video(content_mid_features[0], static_mid_features[0], flow_features[1]),
                          self.decode_one_video(content_mid_features[0], static_mid_features[1], flow_features[0])
                          ]

        id_a_content_b = [self.decode_one_video(content_mid_features[1], static_mid_features[1], flow_features[1]),
                          self.decode_one_video(content_mid_features[1], static_mid_features[1], flow_features[0]),
                          self.decode_one_video(content_mid_features[1], static_mid_features[0], flow_features[1])
                          ]

        id_b_content_a = [self.decode_one_video(content_mid_features[2], static_mid_features[2], flow_features[2]),
                          self.decode_one_video(content_mid_features[0], static_mid_features[2], flow_features[2])
                          ]

        return {"id_a_content_a": id_a_content_a, "id_a_content_b": id_a_content_b, "id_b_content_a": id_b_content_a}

    def decode_one_video(self, content, static, motion):
        static = static.unsqueeze(2)
        content = content + static  # not inplace
        video = self.video_decoder(content, latent=motion)
        return video

    def content_recognizer(self, content_feature, tokens):
        tokens = tokens[:, :-1]
        token_embed = self.embed(tokens)
        attn_mask = sequence_decoder.generate_square_subsequent_mask(tokens.shape[1])
        padding_mask = (tokens == consts.PAD)
        prob = self.seq2seq(content_feature, token_embed, src_mask=None, tgt_mask=attn_mask,
                            src_padding_mask=None, tgt_padding_mask=padding_mask,
                            memory_key_padding_mask=None, batch_first=True)
        prob = torch.matmul(prob, torch.transpose(self.embed.weight, 1, 0))
        return prob

    def predict(self, image_sequence):
        lip_feature = self.video_cnn(image_sequence)["pooled_feature"]
        lip_feature = lip_feature.permute(1, 0, 2)  # (frames, batch_size, embed_size)
        memory = self.seq2seq.encode(lip_feature.cuda(), src_mask=None)
        predict = sequence_decoder.beam_decode(self.seq2seq, self.embed, memory, consts.TOKEN_MAX_LEN)
        return predict