import cv2
import os
from moviepy.editor import VideoFileClip
import ffmpeg

def extractImageFromVideo(pathVideo):
    # Read the video from specified path
    cam = cv2.VideoCapture(pathVideo)
    fps = cam.get(cv2.CAP_PROP_FPS)
    print('fps = '+fps)
    try:
        # creating a folder named data
        if not os.path.exists('data'):
            os.makedirs('data')
    # if not created then raise error
    except OSError:
        print('Error: Creating directory of data')
    # frame
    currentframe = 0
    while (True):
        # reading from frame
        ret, frame = cam.read()
        if ret:
            # if video is still left continue creating images
            name = './data/frame' + str(currentframe) + '.jpg'
            print('Creating...' + name)
            # writing the extracted images
            cv2.imwrite(name, frame)
            # increasing counter so that it will
            # show how many frames are created
            currentframe += 1
        else:
            break
    # Release all space and windows once done
    cam.release()
    cv2.destroyAllWindows()

def extractAudioFromVideo(pathVideo, pathOutput):
    clip = VideoFileClip(pathVideo)
    clip.audio.write_audiofile(pathOutput)

def createVideoFromImages(pathImages):
    img_array = []
    for filename in os.listdir(pathImages):
        img = cv2.imread(pathImages+'\\'+filename)
        height, width, layers = img.shape
        size = (width, height)
        img_array.append(img)

    out = cv2.VideoWriter('project.avi', cv2.VideoWriter_fourcc(*'DIVX'), 15, size)

    for i in range(len(img_array)):
        out.write(img_array[i])
    out.release()

def test():
    input_video = ffmpeg.input(r"C:\Users\HuyBin\PycharmProjects\Steganography\Videos\project.mp4")
    input_audio = ffmpeg.input(r"C:\Users\HuyBin\PycharmProjects\Steganography\Videos\test.wav")
    ffmpeg.concat(input_video, input_audio, v=1, a=1).output('final.mp4').run()

if __name__ == '__main__':
    # extractImageFromVideo(r"C:\Users\HuyBin\Downloads\test.mp4")
    # extractAudioFromVideo(r"C:\Users\HuyBin\Downloads\test.mp4",'test.wav')
    # createVideoFromImages(r'C:\Users\HuyBin\PycharmProjects\Steganography\Videos\data')
    test()


import math
from typing import List

def problem_3a(data: List[str], horizontal_skips: List[int], vertical_skip: int = 1) -> List[int]:
    vals = []

    for skip in horizontal_skips:
        trees = 0
        pointer = 0
        horizontal_length = len(data[0])
        for row in data[::vertical_skip]:
            if row[pointer] == "#":
                trees += 1

            pointer += skip
            if pointer >= horizontal_length:
                pointer = pointer % horizontal_length

        vals.append(trees)
    return vals

def problem_3b(data: List[str]) -> List[int]:
    vals_1 = problem_3a(data, horizontal_skips=[1, 3, 5, 7], vertical_skip=1)
    vals_2 = problem_3a(data, horizontal_skips=[1], vertical_skip=2)
    return math.prod(vals_1) * math.prod(vals_2)


import matplotlib.pyplot as plt

data=[3,4,2,4,3,5,3,6,4,3]

if False:
    plt.plot(data)
    print(plt.show())

# Plotting a line
if False:
    # x axis values
	x = [1,2,3]
	# corresponding y axis values
	y = [2,4,1]

	# plotting the points
	plt.plot(x, y)

	# naming the x axis
	plt.xlabel('x - axis')
	# naming the y axis
	plt.ylabel('y - axis')

	# giving a title to my graph
	plt.title('My first graph!')

	# function to show the plot
	plt.show()

# Sample bar chart
if False:
    names = ['group_a', 'group_b', 'group_c']
    values = [1, 10, 100]

    plt.figure(figsize=(9, 3))
    plt.subplot(131)
    plt.bar(names, values)
    plt.subplot(132)
    plt.scatter(names, values)
    plt.subplot(133)
    plt.plot(names, values)
    plt.suptitle('Categorical Plotting')
    plt.show()

# Histogram
if False:
	# frequencies
	ages = [2,5,70,40,30,45,50,45,43,40,44,
	        60,7,13,57,18,90,77,32,21,20,40]

	# setting the ranges and no. of intervals
	range = (0, 100)
	bins = 10

	# plotting a histogram
	plt.hist(ages, bins, range, color = 'green',
	        histtype = 'bar', rwidth = 0.8)

	# x-axis label
	plt.xlabel('age')
	# frequency label
	plt.ylabel('No. of people')
	# plot title
	plt.title('My histogram')

	# function to show the plot
	plt.show()

# Scatter plot
if False:
	# x-axis values
	x = [1,2,3,4,5,6,7,8,9,10]
	# y-axis values
	y = [2,4,5,7,6,8,9,11,12,12]

	# plotting points as a scatter plot
	plt.scatter(x, y, label= "stars", color= "green",
	            marker= "*", s=30)

	# x-axis label
	plt.xlabel('x - axis')
	# frequency label
	plt.ylabel('y - axis')
	# plot title
	plt.title('My scatter plot!')
	# showing legend
	plt.legend()

	# function to show the plot
	plt.show()

# Pie-chart
if True:
	# defining labels
	activities = ['eat', 'sleep', 'work', 'play']

	# portion covered by each label
	slices = [3, 7, 8, 6]

	# color for each label
	colors = ['r', 'y', 'g', 'b']

	# plotting the pie chart
	plt.pie(slices, labels = activities, colors=colors,
	        startangle=90, shadow = True, explode = (0, 0, 0.1, 0),
	        radius = 1.2, autopct = '%1.1f%%')

	# plotting legend
	plt.legend()

	# showing the plot
	plt.show()


import sys

from setuptools import setup, find_packages
from setuptools.command.test import test as TestCommand

class Tox(TestCommand):
    user_options = [('tox-args=', 'a', "Arguments to pass to tox")]

    def initialize_options(self):
        TestCommand.initialize_options(self)
        self.tox_args = None

    def finalize_options(self):
        TestCommand.finalize_options(self)
        self.test_args = []
        self.test_suite = True

    def run_tests(self):
        # Import here since eggs aren't loaded outside of this scope
        import tox
        import shlex

        args = self.tox_args
        if args:
            args = shlex.split(self.tox_args)

        errno = tox.cmdline(args=args)
        sys.exit(errno)

setup(
    name = 'victor',
    version = '0.1.1',
    description = "A simple tool for debugging and profiling applications",
    url = 'https://github.com/jcomo/victor',
    author = 'Jonathan Como',
    author_email = 'jonathan.como@gmail.com',
    packages = find_packages(exclude=['docs', 'tests', 'scripts']),
    install_requires = [
        'six>=1.10',
    ],
    tests_require = ['tox'],
    cmdclass = {
        'test': Tox,
    },
    classifiers = [
        'Programming Language :: Python :: 2',
        'Programming Language :: Python :: 2.6',
        'Programming Language :: Python :: 2.7',
        'Programming Language :: Python :: 3.3',
        'Programming Language :: Python :: 3.4',
    ],
    keywords = 'debug profile python test'
)


import ud_dataloader
import mxnet as mx
from mxnet import nd, autograd, gluon
from config import train_data_fn as train_data
import config
import random

def getWordPos(data):
    words = {}
    pos_tag = set()
    for sen in data:
        for token in sen.tokens:
            w = token.form
            t = token.pos_tag
            if not w in words:
                words[w] = 1
            else:
                words[w] = words[w] + 1
            if not t in pos_tag:
                pos_tag.add(t)
    pos_tag = list(pos_tag)
    return words, pos_tag

class TaggerModel(gluon.Block):
    def __init__(self, vocab_size, num_embed, num_hidden, tag_count, **kwargs):
        super(TaggerModel, self).__init__(**kwargs)
        with self.name_scope():
            self.embed = gluon.nn.Embedding(vocab_size, num_embed, weight_initializer=mx.init.Uniform(0.1))
            self.lstm = gluon.rnn.LSTM(num_hidden, 1, bidirectional=True, input_size=num_embed)
            self.tag_cls = gluon.nn.Dense(tag_count, in_units=num_hidden*2)
        self.num_hidden = num_embed
        self.tag_count = tag_count

    def forward(self, inputs):
        embed = self.embed(inputs)
        s1, s2 = embed.shape
        embed = embed.reshape((s1, 1, s2))
        hidden = self.lstm(embed)
        batch_size, __, hn_size = hidden.shape
        hidden.reshape((batch_size, hn_size))
        cls = self.tag_cls(hidden)
        return cls

    def begin_state(self, *args, **kwargs):
        return self.rnn.begin_state(*args, **kwargs)

def mapTokenToId(sen: ud_dataloader.UDSentence, word_map:dict):
    ret = []
    for item in sen.tokens:
        ret.append(word_map[item.form])
    return ret

def mapTagToId(sen: ud_dataloader.UDSentence, tag_map:dict):
    ret = []
    for item in sen.tokens:
        ret.append(tag_map[item.pos_tag])
    return ret

data = ud_dataloader.parseDocument(train_data)
words, pos_tag = getWordPos(data)
word_list = sorted(list(words.keys()))
word_map = {}

for i, w in enumerate(word_list):
    word_map[w] = i
pos_tag_map = {}
for i, t in enumerate(pos_tag):
    pos_tag_map[t] = i

ctx = mx.gpu(0)
tagger = TaggerModel(len(word_list), 50, 50, len(pos_tag))
tagger.collect_params().initialize(mx.init.Xavier(), ctx=ctx)
trainer = gluon.Trainer(tagger.collect_params(), 'adam', {'learning_rate': 0.01})
loss = gluon.loss.SoftmaxCrossEntropyLoss()

for epoch in range(1, 10+1):
    random.shuffle(data)
    avg_loss = 0.0
    acc_accu = 0.0
    acc_total = 0
    for i, sen in enumerate(data):
        tokens = mapTokenToId(sen, word_map)
        tokens = mx.nd.array(tokens, ctx)
        tags = mapTagToId(sen, pos_tag_map)
        tags = mx.nd.array(tags, ctx)
        with autograd.record():
            outputs = tagger(tokens)
            pred = outputs.argmax(axis=1)
            acc_accu += (tags==pred).sum().asscalar()
            acc_total += outputs.shape[0]
            L = loss(outputs, tags)
            L = L.mean()
            L.backward()
        trainer.step(1)
        avg_loss += L.asscalar()
        if i % config.prompt_inteval == 0:
            avg_loss /= config.prompt_inteval
            acc = acc_accu / acc_total
            print("Epoch {} sen {} loss={} train acc={}".format(epoch, i, avg_loss, acc))
            avg_loss = 0
            acc_accu = 0
            acc_total = 0


from flask import Flask
from google.cloud import datastore, pubsub

publish_client = pubsub.PublisherClient()

topic = 'projects/gcpaceproject1/topics/visitorinfo'

app = Flask(__name__)

dsclient = datastore.Client()
def savevisitorinfo(visitor):
  entity = datastore.Entity(key=dsclient.key('visitorinfo'))
  entity.update({
    'visitorname': visitor
  })
  dsclient.put(entity)

@app.route('/')
def home():
  return '
I AM SERVICE1
'

@app.route('/visitor/<visitor>')
def visitorinfo(visitor):
    savevisitorinfo(visitor)
    bvisitor = visitor.encode("utf-8")
    publish_client.publish(topic, bvisitor)
    return '
Hello %s
' %visitor


import logging
import os
import time
import random
import json
from tqdm import tqdm
import sys
import torch
from itertools import chain
import torch.nn as nn
from torch.nn.utils import clip_grad_norm_
from torch.utils.data import DataLoader
from tensorboardX import SummaryWriter
from torch.optim.lr_scheduler import StepLR, MultiStepLR
import numpy as np
from configs.opts import parser
from model.main_model_2 import AV_VQVAE_Encoder, AT_VQVAE_Encoder, AV_VQVAE_Decoder, AT_VQVAE_Decoder, AVT_VQVAE_Encoder, AVT_VQVAE_Decoder
from model.CLUB import CLUBSample_group
from model.CPC import Cross_CPC, Cross_CPC_AVT
from utils import AverageMeter, Prepare_logger, get_and_save_args
from utils.container import metricsContainer
from utils.Recorder import Recorder
import torch.nn.functional as F
from bert_embedding import BertEmbedding
import pickle
# =================================  seed config ============================
SEED = 43
random.seed(SEED)
np.random.seed(seed=SEED)
torch.manual_seed(seed=SEED)
torch.cuda.manual_seed(seed=SEED)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False

# =============================================================================
def transpose(x):
    return x.transpose(-2, -1)

def normalize(*xs):
    return [None if x is None else F.normalize(x, dim=-1) for x in xs]

def AVPSLoss(av_simm, soft_label):
    """audio-visual pair similarity loss for fully supervised setting,
    please refer to Eq.(8, 9) in our paper.
    """
    # av_simm: [bs, 10]
    relu_av_simm = F.relu(av_simm)
    sum_av_simm = torch.sum(relu_av_simm, dim=-1, keepdim=True)
    avg_av_simm = relu_av_simm / (sum_av_simm + 1e-8)
    loss = nn.MSELoss()(avg_av_simm, soft_label)
    return loss

bert_embedding = BertEmbedding()
with open('../../cnt.pkl', 'rb') as fp:
    id2idx = pickle.load(fp)

def collate_func_AT(samples):
        bsz = len(samples)
        result = bert_embedding([sample['text_fea'] for sample in samples])
        query = []
        query_words = []
        for a, b in result:
            words = []
            words_emb = []
            for word, emb in zip(a, b):
                idx = bert_embedding.vocab.token_to_idx[word]
                if idx in id2idx and idx != 0:
                    words_emb.append(emb)
                    words.append(id2idx[idx])
            query.append(np.asarray(words_emb))
            query_words.append(words)

        query_len = []
        for i, sample in enumerate(query):
            # query_len.append(min(len(sample), 10))#max_num_words:10
            query_len.append(10)#max_num_words:10
        query1 = np.zeros([bsz, max(query_len), 768]).astype(np.float32)
        query_idx = np.zeros([bsz, max(query_len)]).astype(np.float32)
        for i, sample in enumerate(query):
            keep = min(sample.shape[0], query1.shape[1])
            query1[i, :keep] = sample[:keep]
            query_idx[i, :keep] = query_words[i][:keep]
        query_len = np.asarray(query_len)
        query, query_len = torch.from_numpy(query1).float(), torch.from_numpy(query_len).long()
        query_idx = torch.from_numpy(query_idx).long()

        return {
            'query': query,
            'query_idx': query_idx,
            'query_len': query_len,
            'audio_fea': torch.from_numpy(np.asarray([sample['audio_fea'] for sample in samples])).float()
        }

def collate_func_AVT(samples):
        bsz = len(samples)
        result = bert_embedding([sample['text_fea'] for sample in samples])
        query = []
        query_words = []
        for a, b in result:
            words = []
            words_emb = []
            for word, emb in zip(a, b):
                idx = bert_embedding.vocab.token_to_idx[word]
                if idx in id2idx and idx != 0:
                    words_emb.append(emb)
                    words.append(id2idx[idx])
            query.append(np.asarray(words_emb))
            query_words.append(words)

        query_len = []
        for i, sample in enumerate(query):
            # query_len.append(min(len(sample), 10))#max_num_words:10
            query_len.append(10)#max_num_words:10
        query1 = np.zeros([bsz, max(query_len), 768]).astype(np.float32)
        query_idx = np.zeros([bsz, max(query_len)]).astype(np.float32)
        for i, sample in enumerate(query):
            keep = min(sample.shape[0], query1.shape[1])
            """
            There may be cases where the sample length is 0,
            for example if your text happens to not be seen before in this BERT model.
            If that happens, you can
            1) clean the text before it enters BERT,
            2) add an if statement here,
            3) discard idx and directly import all embeddings after.
            """
            query1[i, :keep] = sample[:keep]
            query_idx[i, :keep] = query_words[i][:keep]
        query_len = np.asarray(query_len)
        query, query_len = torch.from_numpy(query1).float(), torch.from_numpy(query_len).long()
        query_idx = torch.from_numpy(query_idx).long()

        return {
            'query': query,
            'audio_fea': torch.from_numpy(np.asarray([sample['audio_fea'] for sample in samples])).float(),
            'video_fea': torch.from_numpy(np.asarray([sample['video_fea'] for sample in samples])).float()
        }

def main():
    # utils variable
    global args, logger, writer, dataset_configs
    # statistics variable
    global best_accuracy, best_accuracy_epoch
    best_accuracy, best_accuracy_epoch = 0, 0
    # configs
    dataset_configs = get_and_save_args(parser)
    parser.set_defaults(**dataset_configs)
    args = parser.parse_args()
    # select GPUs
    os.environ['CUDA_DEVICE_ORDER'] = "PCI_BUS_ID"
    # os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu

    '''Create snapshot_pred dir for copying code and saving models '''
    if not os.path.exists(args.snapshot_pref):
        os.makedirs(args.snapshot_pref)

    if os.path.isfile(args.resume):
        args.snapshot_pref = os.path.dirname(args.resume)

    logger = Prepare_logger(args, eval=args.evaluate)

    if not args.evaluate:
        logger.info(f'\nCreating folder: {args.snapshot_pref}')
        logger.info('\nRuntime args\n\n{}\n'.format(json.dumps(vars(args), indent=4)))
    else:
        logger.info(f'\nLog file will be save in a {args.snapshot_pref}/Eval.log.')

    '''dataset selection'''
    if args.dataset_name == 'ave':
        from dataset.AVE_dataset import AVEDataset as AVEDataset
    elif args.dataset_name =='vggsound':
        from dataset.VGGSOUND_dataset import VGGSoundDataset as AVEDataset
    elif args.dataset_name =='vggsound_AT':
        from dataset.VGGSOUND_dataset import VGGSoundDataset_AT as AVEDataset
    elif args.dataset_name =='vggsound_AVT':
        from dataset.VGGSOUND_dataset import VGGSoundDataset_AVT as AVEDataset
    elif args.dataset_name =='vggsound179k' or args.dataset_name =='vggsound81k':
        from dataset.VGGSOUND_dataset179k import VGGSoundDataset as AVEDataset
    else:
        raise NotImplementedError

    '''Dataloader selection'''
    if args.dataset_name == 'ave':
        data_root = 'data'
        train_dataloader = DataLoader(
            AVEDataset(data_root, split='train'),
            batch_size=args.batch_size,
            shuffle=True,
            num_workers=8,
            pin_memory=True
        )
        val_dataloader = DataLoader(
            AVEDataset(data_root, split='val'),
            batch_size=args.batch_size,
            shuffle=False,
            num_workers=8,
            pin_memory=True
        )
        test_dataloader = DataLoader(
            AVEDataset(data_root, split='test'),
            batch_size=args.batch_size,
            shuffle=False,
            num_workers=8,
            pin_memory=True
        )
    elif args.dataset_name == 'vggsound':
        meta_csv_path = 'vggsound-avel40k.csv'
        audio_fea_base_path = 'audio/zip'
        video_fea_base_path = 'video/zip'
        avc_label_base_path = 'label/zip'
        train_dataloader = DataLoader(
            AVEDataset(meta_csv_path, audio_fea_base_path, video_fea_base_path, avc_label_base_path, split='train'),
            batch_size=args.batch_size,
            shuffle=True,
            num_workers=8,
            pin_memory=True
        )
    elif args.dataset_name == 'vggsound_AT':
        meta_csv_path = 'vggsound-avel40k.csv'
        audio_fea_base_path = 'audio/zip'
        train_dataloader = DataLoader(
            AVEDataset(meta_csv_path, audio_fea_base_path, split='train'),
            batch_size=args.batch_size,
            shuffle=True,
            num_workers=8,
            pin_memory=True,
            collate_fn=collate_func_AT
        )
    elif args.dataset_name == 'vggsound_AVT':
        meta_csv_path = 'vggsound-avel40k.csv'
        audio_fea_base_path = 'vggsound40k/feature/audio/zip'
        video_fea_base_path = 'vggsound40k/feature/video/zip'
        train_dataloader = DataLoader(
            AVEDataset(meta_csv_path, audio_fea_base_path, video_fea_base_path, split='train'),
            batch_size=args.batch_size,
            shuffle=True,
            num_workers=8,
            pin_memory=False,
            collate_fn=collate_func_AVT
        )
    elif args.dataset_name == 'vggsound81k':
        meta_csv_path = 'video_name_vggsound81k_checked.csv'
        audio_fea_base_path = 'audio/zip'
        video_fea_base_path = 'video/zip'
        avc_label_base_path = '...'
        train_dataloader = DataLoader(
            AVEDataset(meta_csv_path, audio_fea_base_path, video_fea_base_path, avc_label_base_path, split='train'),
            batch_size=args.batch_size,
            shuffle=True,
            num_workers=8,
            pin_memory=True
        )
    elif args.dataset_name == 'vggsound179k':
        meta_csv_path = 'video_name_vggsound179k_checked.csv'
        audio_fea_base_path = 'audio/zip'
        video_fea_base_path = 'video/zip'
        avc_label_base_path = '...'
        train_dataloader = DataLoader(
            AVEDataset(meta_csv_path, audio_fea_base_path, video_fea_base_path, avc_label_base_path, split='train'),
            batch_size=args.batch_size,
            shuffle=True,
            num_workers=8,
            pin_memory=True
        )
    else:
        raise NotImplementedError

    '''model setting'''
    video_dim = 512
    text_dim = 768
    audio_dim = 128
    text_lstm_dim = 128
    video_output_dim = 2048
    text_output_dim = 256
    audio_output_dim = 256
    n_embeddings = 400
    embedding_dim = 256
    start_epoch = -1
    model_resume = False
    total_step = 0
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    Text_ar_lstm = nn.LSTM(text_dim, text_lstm_dim, num_layers=2, batch_first=True, bidirectional=True)

    if args.dataset_name == 'vggsound_AT':
        Encoder = AT_VQVAE_Encoder(text_lstm_dim*2, audio_dim, text_output_dim, audio_output_dim, n_embeddings, embedding_dim)
    if args.dataset_name == 'vggsound_AVT':
        Encoder = AVT_VQVAE_Encoder(audio_dim, video_dim, text_lstm_dim*2, audio_output_dim, video_output_dim, text_output_dim, n_embeddings, embedding_dim)
    else:
        Encoder = AV_VQVAE_Encoder(video_dim, audio_dim, video_output_dim, audio_output_dim, n_embeddings, embedding_dim)

    if args.dataset_name == 'vggsound_AVT':
        CPC = Cross_CPC_AVT(embedding_dim, hidden_dim=256, context_dim=256, num_layers=2)
    else:
        CPC = Cross_CPC(embedding_dim, hidden_dim=256, context_dim=256, num_layers=2)
    Video_mi_net = CLUBSample_group(x_dim=embedding_dim, y_dim=video_dim, hidden_size=256)
    Text_mi_net = CLUBSample_group(x_dim=embedding_dim, y_dim=text_output_dim, hidden_size=256)
    Audio_mi_net = CLUBSample_group(x_dim=embedding_dim, y_dim=audio_output_dim, hidden_size=256)

    if args.dataset_name == 'vggsound_AT':
        Decoder = AT_VQVAE_Decoder(text_lstm_dim*2, audio_dim, text_output_dim, audio_output_dim)
    if args.dataset_name == 'vggsound_AVT':
        Decoder = AVT_VQVAE_Decoder(audio_dim, video_dim, text_lstm_dim*2, audio_output_dim, video_output_dim, text_output_dim)
    else:
        Decoder = AV_VQVAE_Decoder(video_dim, audio_dim, video_output_dim, audio_output_dim)

    Text_ar_lstm.double()
    Encoder.double()
    CPC.double()
    Video_mi_net.double()
    Text_mi_net.double()
    Audio_mi_net.double()
    Decoder.double()

    '''optimizer setting'''
    Text_ar_lstm.to(device)
    Encoder.to(device)
    CPC.to(device)
    Video_mi_net.to(device)
    Text_mi_net.to(device)
    Audio_mi_net.to(device)
    Decoder.to(device)
    optimizer = torch.optim.Adam(chain(Text_ar_lstm.parameters(), \
                                       Encoder.parameters(), CPC.parameters(), Decoder.parameters()), lr=args.lr)
    optimizer_video_mi_net = torch.optim.Adam(Video_mi_net.parameters(), lr=args.mi_lr)
    optimizer_text_mi_net = torch.optim.Adam(Text_mi_net.parameters(), lr=args.mi_lr)
    optimizer_audio_mi_net = torch.optim.Adam(Audio_mi_net.parameters(), lr=args.mi_lr)
    scheduler = MultiStepLR(optimizer, milestones=[10, 20, 30], gamma=0.5)

    '''loss'''
    criterion = nn.BCEWithLogitsLoss().cuda()
    criterion_event = nn.CrossEntropyLoss().cuda()

    if model_resume is True:
        path_checkpoints = "..."
        checkpoints = torch.load(path_checkpoints)
        Encoder.load_state_dict(checkpoints['Encoder_parameters'])
        CPC.load_state_dict(checkpoints['CPC_parameters'])
        Video_mi_net.load_state_dict(checkpoints['Video_mi_net_parameters'])
        Audio_mi_net.load_state_dict(checkpoints['Audio_mi_net_parameters'])
        Decoder.load_state_dict(checkpoints['Decoder_parameters'])
        optimizer.load_state_dict(checkpoints['optimizer'])
        optimizer_audio_mi_net.load_state_dict(checkpoints['optimizer_audio_mi_net'])
        optimizer_video_mi_net.load_state_dict(checkpoints['optimizer_video_mi_net'])
        start_epoch = checkpoints['epoch']
        total_step = checkpoints['total_step']
        logger.info("Resume from number {}-th model.".format(start_epoch))

    '''Tensorboard and Code backup'''
    writer = SummaryWriter(args.snapshot_pref)
    recorder = Recorder(args.snapshot_pref, ignore_folder="Exps/")
    recorder.writeopt(args)

    '''Training and Evaluation'''

    for epoch in range(start_epoch+1, args.n_epoch):
        loss, total_step = train_epoch(CPC, Encoder,Text_ar_lstm, Audio_mi_net, Video_mi_net,Text_mi_net, Decoder, train_dataloader, criterion, criterion_event,
                                       optimizer, optimizer_audio_mi_net, optimizer_video_mi_net, optimizer_text_mi_net, epoch, total_step, args)

        save_path = os.path.join(args.model_save_path, 'your-model-{}.pt'.format(epoch))
        save_models(CPC, Encoder, Text_ar_lstm, Audio_mi_net, Video_mi_net, Text_mi_net, Decoder, optimizer, optimizer_audio_mi_net, optimizer_video_mi_net, optimizer_text_mi_net, epoch, total_step, save_path)
        logger.info(f"epoch: ******************************************* {epoch}")
        logger.info(f"loss: {loss}")
        scheduler.step()

def _export_log(epoch, total_step, batch_idx, lr, loss_meter):
    msg = 'Epoch {}, Batch {}, lr = {:.5f}, '.format(epoch, batch_idx, lr)
    for k, v in loss_meter.items():
        msg += '{} = {:.4f}, '.format(k, v)
    logger.info(msg)
    sys.stdout.flush()
    loss_meter.update({"batch": total_step})

def to_eval(all_models):
    for m in all_models:
        m.eval()

def to_train(all_models):
    for m in all_models:
        m.train()

# If resuming training is not required, downstream tasks only need to save the encoder & epoch & Text_ar_lstm, as these are the only components needed for inference.
def save_models(CPC, Encoder,Text_ar_lstm, Audio_mi_net, Video_mi_net, Text_mi_net, Decoder, optimizer, optimizer_audio_mi_net, optimizer_video_mi_net,optimizer_text_mi_net, epoch_num, total_step, path):
    state_dict = {
        'Encoder_parameters': Encoder.state_dict(),
        'CPC_parameters': CPC.state_dict(),
        'Text_ar_lstm_parameters': Text_ar_lstm.state_dict(),
        'Video_mi_net_parameters': Video_mi_net.state_dict(),
        'Text_mi_net_parameters': Text_mi_net.state_dict(),
        'Audio_mi_net_parameters': Audio_mi_net.state_dict(),
        'Decoder_parameters': Decoder.state_dict(),
        'optimizer': optimizer.state_dict(),
        'optimizer_video_mi_net': optimizer_video_mi_net.state_dict(),
        'optimizer_text_mi_net': optimizer_text_mi_net.state_dict(),
        'optimizer_audio_mi_net': optimizer_audio_mi_net.state_dict(),
        'epoch': epoch_num,
        'total_step': total_step
    }
    torch.save(state_dict, path)
    logging.info('save model to {}'.format(path))

def train_epoch_check(train_dataloader, epoch, total_step, args):
    train_dataloader = tqdm(train_dataloader)
    for n_iter, batch_data in enumerate(train_dataloader):

        '''Feed input to model'''
        visual_feature, audio_feature = batch_data
        visual_feature.cuda()
        audio_feature.cuda()

    return torch.zeros(1)

def train_epoch(CPC,Encoder,Text_ar_lstm, Audio_mi_net, Video_mi_net, Text_mi_net, Decoder,train_dataloader, criterion, criterion_event, optimizer, optimize_audio_mi_net, optimizer_video_mi_net, optimizer_text_mi_net, epoch, total_step, args):
    batch_time = AverageMeter()
    data_time = AverageMeter()
    losses = AverageMeter()
    train_acc = AverageMeter()
    end_time = time.time()
    models = [CPC,Encoder,Text_ar_lstm, Audio_mi_net, Video_mi_net, Text_mi_net, Decoder]
    to_train(models)
    # Note: here we set the model to a double type precision,
    # since the extracted features are in a double type.
    # This will also lead to the size of the model double increases.

    Encoder.cuda()
    Text_ar_lstm.cuda()
    Text_mi_net.cuda()
    Audio_mi_net.cuda()
    Video_mi_net.cuda()
    Decoder.cuda()
    CPC.cuda()
    optimizer.zero_grad()
    mi_iters = 5

    # train_dataloader = tqdm(train_dataloader)

    for n_iter, batch_data in enumerate(train_dataloader):

        data_time.update(time.time() - end_time)
        '''Feed input to model'''

        # Adjust the input as needed according to the requirements of the model being trained.
        # vggsound_AVT
        query, audio_feature, video_feature = batch_data['query'], batch_data['audio_fea'], batch_data['video_fea']

        # vggsound_AV
        # visual_feature, audio_feature, labels = batch_data    # vggsound40k
        # visual_feature, audio_feature = batch_data            # vggsound179k or vggsound81k
        query = query.double().cuda()
        audio_feature = audio_feature.to(torch.float64)
        batch_dim = query.size()[0]
        hidden_dim = 128
        num_layers = 2
        text_hidden = (torch.zeros(2*num_layers, batch_dim, hidden_dim).double().cuda(),
                  torch.zeros(2*num_layers, batch_dim, hidden_dim).double().cuda())
        text_feature, text_hidden = Text_ar_lstm(query, text_hidden)

        text_feature = text_feature.cuda().to(torch.float64)
        audio_feature = audio_feature.cuda().to(torch.float64)
        video_feature = video_feature.cuda().to(torch.float64)

        for i in range(mi_iters):
            optimizer_video_mi_net, lld_video_loss, optimizer_text_mi_net, lld_text_loss, optimize_audio_mi_net, lld_audio_loss = \
                mi_first_forward(audio_feature, video_feature, text_feature, Encoder, Audio_mi_net, Video_mi_net, Text_mi_net, optimize_audio_mi_net,optimizer_video_mi_net,optimizer_text_mi_net, epoch)

        audio_embedding_loss, video_embedding_loss,text_embedding_loss, mi_audio_loss, mi_video_loss, mi_text_loss, \
        accuracy1, accuracy2, accuracy3, accuracy4, accuracy5, accuracy6, accuracy7, accuracy8, accuracy9,\
        cpc_loss, audio_recon_loss, video_recon_loss, text_recon_loss, \
        audio_class, video_class, text_class, cmcm_loss, equal_num = mi_second_forward(CPC,audio_feature, video_feature, text_feature, Encoder,Audio_mi_net, Video_mi_net, Text_mi_net, Decoder,epoch)

        if n_iter % 20 == 0:
            logger.info("equal_num is {} in {}-th iteration.".format(equal_num, n_iter))

        loss_items = {
            "audio_recon_loss": audio_recon_loss.item(),
            "lld_audio_loss": lld_audio_loss.item(),
            "audio_embed_loss": audio_embedding_loss.item(),
            "audio_mine_loss": mi_audio_loss.item(),
            "text_recon_loss": text_recon_loss.item(),
            "lld_text_loss": lld_text_loss.item(),
            "text_embed_loss": text_embedding_loss.item(),
            "text_mine_loss": mi_text_loss.item(),
            "video_recon_loss": video_recon_loss.item(),
            "lld_video_loss": lld_video_loss.item(),
            "video_embed_loss": video_embedding_loss.item(),
            "video_mine_loss": mi_video_loss.item(),
            "acc_av": accuracy1.item(),
            "acc_at": accuracy2.item(),
            "acc_vt": accuracy3.item(),
            "acc_va": accuracy4.item(),
            "acc_ta": accuracy5.item(),
            "acc_tv": accuracy6.item(),
            "acc_aa": accuracy7.item(),
            "acc_vv": accuracy8.item(),
            "acc_tt": accuracy9.item(),
            "cpc_loss": cpc_loss.item(),
            "cmcm_loss": cmcm_loss.item()
        }

        metricsContainer.update("loss", loss_items)
        loss = audio_recon_loss + video_recon_loss + text_recon_loss + audio_embedding_loss +  video_embedding_loss\
                + text_embedding_loss+ mi_audio_loss + mi_video_loss + mi_text_loss + cpc_loss + cmcm_loss

        if n_iter % 20 == 0:
            _export_log(epoch=epoch, total_step=total_step+n_iter, batch_idx=n_iter, lr=0.0004, loss_meter=metricsContainer.calculate_average("loss"))

        loss.backward()

        '''Clip Gradient'''
        if args.clip_gradient is not None:
            for model in models:
                total_norm = clip_grad_norm_(model.parameters(), args.clip_gradient)

        '''Update parameters'''
        optimizer.step()
        optimizer.zero_grad()

        losses.update(loss.item(), text_feature.size(0) * 10)
        batch_time.update(time.time() - end_time)
        end_time = time.time()

        '''Add loss of a iteration in Tensorboard'''
        writer.add_scalar('Train_data/loss', losses.val, epoch * len(train_dataloader) + n_iter + 1)

        '''Add loss of an epoch in Tensorboard'''
        writer.add_scalar('Train_epoch_data/epoch_loss', losses.avg, epoch)

    return losses.avg, n_iter + total_step

def mi_first_forward(audio_feature, video_feature, text_feature, Encoder, Audio_mi_net, Video_mi_net,Text_mi_net, optimizer_audio_mi_net,optimizer_video_mi_net,optimizer_text_mi_net, epoch):

    optimizer_video_mi_net.zero_grad()
    optimizer_text_mi_net.zero_grad()
    optimizer_audio_mi_net.zero_grad()

    audio_semantic_result, video_semantic_result, text_semantic_result, \
    audio_encoder_result, video_encoder_result, video_club_feature, text_encoder_result, \
    audio_vq, video_vq, text_vq, audio_embedding_loss, video_embedding_loss, text_embedding_loss, cmcm_loss, equal_num\
    = Encoder(audio_feature, video_feature, text_feature, epoch)

    video_club_feature = video_club_feature.detach()
    text_encoder_result = text_encoder_result.detach()
    audio_encoder_result = audio_encoder_result.detach()
    video_vq = video_vq.detach()
    text_vq = text_vq.detach()
    audio_vq = audio_vq.detach()

    # video processing is different from audio and text modalities because video is more complex and feature extraction is more difficult.
    lld_video_loss = -Video_mi_net.loglikeli(video_vq, video_club_feature)
    lld_video_loss.backward()
    optimizer_video_mi_net.step()

    lld_text_loss = -Text_mi_net.loglikeli(text_vq, text_encoder_result)
    lld_text_loss.backward()
    optimizer_text_mi_net.step()

    lld_audio_loss = -Audio_mi_net.loglikeli(audio_vq, audio_encoder_result)
    lld_audio_loss.backward()
    optimizer_audio_mi_net.step()

    return optimizer_audio_mi_net, lld_audio_loss, optimizer_video_mi_net, lld_video_loss, optimizer_text_mi_net, lld_text_loss

def VQ_audio_forward(audio_feature, visual_feature, Encoder, optimizer,epoch):

    audio_vq_forward_loss = Encoder.Audio_vq_forward(audio_feature, visual_feature,epoch)
    audio_vq_forward_loss.backward()
    optimizer.step()
    optimizer.zero_grad()
    return audio_vq_forward_loss, optimizer

def VQ_video_forward(audio_feature, visual_feature, Encoder, optimizer,epoch):
    optimizer.zero_grad()
    video_vq_forard_loss = Encoder.Video_vq_forward(audio_feature, visual_feature,epoch)
    video_vq_forard_loss.backward()
    optimizer.step()
    optimizer.zero_grad()
    return video_vq_forard_loss, optimizer

def mi_second_forward(CPC, audio_feature, video_feature, text_feature, Encoder, Audio_mi_net, Video_mi_net, Text_mi_net, Decoder,epoch):
    audio_semantic_result, video_semantic_result, text_semantic_result, \
    audio_encoder_result, video_encoder_result, video_club_feature, text_encoder_result, \
    audio_vq, video_vq, text_vq, audio_embedding_loss, video_embedding_loss, text_embedding_loss, cmcm_loss, equal_num \
    = Encoder(audio_feature, video_feature, text_feature, epoch)

    mi_video_loss = Video_mi_net.mi_est(video_vq, video_club_feature)
    mi_text_loss = Text_mi_net.mi_est(text_vq, text_encoder_result)
    mi_audio_loss = Audio_mi_net.mi_est(audio_vq, audio_encoder_result)

    """Cross_CPC"""
    accuracy1, accuracy2, accuracy3, accuracy4, \
    accuracy5, accuracy6, accuracy7, accuracy8, accuracy9, \
    cpc_loss = CPC(audio_semantic_result, video_semantic_result, text_semantic_result)

    audio_recon_loss, video_recon_loss, text_recon_loss, audio_class, video_class, text_class \
        = Decoder(audio_feature, video_feature, text_feature, audio_encoder_result, video_encoder_result, text_encoder_result, audio_vq, video_vq, text_vq)

    return audio_embedding_loss, video_embedding_loss, text_embedding_loss, mi_audio_loss, mi_video_loss, mi_text_loss, \
           accuracy1, accuracy2, accuracy3, accuracy4, accuracy5, accuracy6, accuracy7, accuracy8, accuracy9, cpc_loss,  \
           audio_recon_loss, video_recon_loss, text_recon_loss, audio_class, video_class, text_class, cmcm_loss, equal_num, zero_num, same_layer_fusion_loss, adjacent_layer_separation_loss

def compute_accuracy_supervised(event_scores, labels):
    labels_foreground = labels[:, :, :-1]
    labels_BCE, labels_evn = labels_foreground.max(-1)
    labels_event, _ = labels_evn.max(-1)
    _, event_class = event_scores.max(-1)
    correct = event_class.eq(labels_event)
    correct_num = correct.sum().double()
    acc = correct_num * (100. / correct.numel())
    return acc

def save_checkpoint(state_dict, top1, task, epoch):
    model_name = f'{args.snapshot_pref}/model_epoch_{epoch}_top1_{top1:.3f}_task_{task}_best_model.pth.tar'
    torch.save(state_dict, model_name)

if __name__ == '__main__':
    main()


import time
import argparse
import numpy as np

# Include the parent direcotry of GEP in python path (not nice looking)
import os, sys
sys.path.append(os.path.join(os.getcwd(), ".."))

import torch

from Envs.lidar_V01 import Grid as grid
from Models.dqn_model import DQN

import collections

DEFAULT_ENV_NAME = "PongNoFrameskip-v4"

SIZE_X = 30
SIZE_Y = 30
INPUT_CHANNELS = 1
INPUT_SHAPE = (INPUT_CHANNELS, SIZE_X, SIZE_Y)
ACTION_SHAPE = 4

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("-m", "--model", required=True,
                        help="Model file to load")
    parser.add_argument("-e", "--env", default=DEFAULT_ENV_NAME,
                        help="Environment name to use, default=" +
                             DEFAULT_ENV_NAME)
    parser.add_argument("-r", "--record", help="Directory for video")
    parser.add_argument("--no-vis", default=True, dest='vis',
                        help="Disable visualization",
                        action='store_false')
    args = parser.parse_args()

    env = grid(size=[SIZE_X, SIZE_Y])

    net = DQN(INPUT_SHAPE ,ACTION_SHAPE)
    state = torch.load(args.model, map_location=lambda stg, _: stg)
    net.load_state_dict(state)

    state = env.reset()
    total_reward = 0.0
    c = collections.Counter()
    timeStep = 0

    while True:
        start_ts = time.time()
        state_v = torch.tensor(np.array([state], copy=False))
        q_vals = net(state_v).data.numpy()[0]
        action = np.argmax(q_vals)
        c[action] += 1
        state, reward, done = env.step(action)
        total_reward += reward
        timeStep += 1
        print(f"step:{timeStep} reward:{reward} q:{q_vals}")
        if done:
            break

    t = time.localtime()
    timestamp = time.strftime('%b%d_%H_%M', t)

    result_folder = os.path.join(os.getcwd(), f"experiments_{timestamp}")
    os.makedirs(f"experiments_{timestamp}")
    env.render(path = result_folder)

    print(f"Game finished after {time.time() - start_ts}[sec]")
    print(f"Total reward: {total_reward}")
    print("Action counts:", c)
    print(f"Inspect the frames at dir:{result_folder}")


from time import sleep

import pytest

from Common.Baseui import baseUI

class Test_mall():

    # def test_login(self,driver):
    #     base = baseUI(driver)
    #     base.driver.get('http://192.168.60.132/#/login')
    #     base.send_keys("输入用户名","//input[@name='username']",'')
    #     base.send_keys("输入密码", "//input[@name='password']", '')
    #     base.click('点击登录','//span[contains(text(),"登录")]')
    #     sleep(2)
    @pytest.mark.Ship
    def test_order(self,base):
        # base = baseUI(driver)
        base.click('点击订单列表',"(//span[contains(text(),'订单列表')])")
        base.click('点击订单状态','//label[contains(text(),"订单状态：")]/following-sibling::div//input')
        base.click('点击待发货',"//span[contains(text(),'待发货')]")
        base.click('点击搜索',"//span[contains(text(),'查询搜索')]")
        base.click('点击订单发货',"(//span[contains(text(),'订单发货')])[1]")
        base.click('选择物流公司','//input[@placeholder="请选择物流公司"]')
        base.click('选择物流公司','//span[contains(text(),"圆通快递")]')
        base.click('点击确定',"(//span[contains(text(),'确定')])[1]")
        base.click('点击确定',"(//span[contains(text(),'确定')])[2]")
        text=base.get_text('获取页面文本','''//div[@role="alert"]//p''')
        assert '成功'in text

    @pytest.mark.Return
    def test_return(self,base):
        base.click('点击退货原因设置','''//span[contains(text(),'退货原因设置')]''')
        # 点击添加
        base.click('点击添加', '(//button[@type="button"])[1]')
        # 输入原因类型
        base.send_keys('输入原因类型','''//*[contains(text(),'原因类型：')]/following-sibling::div//input''','不好看')
        # 点击确定
        base.click('点击确定','//div[@class="el-dialog__footer"]//span//button[2]/span')
        # 获取成功提示
        assert '成功'in base.driver.page_source

    @pytest.mark.reasons_for_return
    def test_teturn_do(self,base):
        # base=baseUI(driver)
        base.driver.get('http://192.168.60.132/#/oms/returnApply')
        # 点击处理状态
        base.click('点击处理状态','''//label[contains(text(),'处理状态：')]/following-sibling::div//input''')
        # 点击待处理
        base.click('点击待处理',"//span[contains(text(),'待处理')]")
        # 点击查询搜索
        base.click('点击查询搜索',"//span[contains(text(),'查询搜索')]")
        # 点击第一列信息
        base.click('点击第一列信息',"(//span[contains(text(),'查看详情')])[1]")
        # 下拉滚动
        base.scroll_screen('下拉滚动')
        # 获取金额
        money=base.get_text('获取金额','''//div[contains(text(),'订单金额')]/following-sibling::div''')
        money=money[1:]
        print(money)
        # 输入金额
        base.send_keys('输入金额','''//div[contains(text(),'确认退款金额')]/following-sibling::div//input''',str(money))
        base.click('点击选择','//div[contains(text(),"选择收货点")]/following-sibling::div//input')
        base.click('选择发货地址','''//span[contains(text(),'北京发货点')]''')
        base.send_keys('处理备注','//div[contains(text(),"处理备注")]/following-sibling::div//input','已处理')
        base.click('点击确认退货','''//span[contains(text(),'确认退货')]''')
        base.click('点击确定',"//span[contains(text(),'确定')]")
        text=base.get_text('获取页面文本', '''//div[@role="alert"]//p''')

        assert '成功' in text

        sleep(5)


# -*- coding: utf-8 -*-
from __future__ import print_function

import json

from clu.constants.consts import PROJECT_NAME
# from clu.config.keymap import FrozenNested
from clu.config.base import Nested
from clu.config.filebase import FileBase
from clu.exporting import Exporter

exporter = Exporter(path=__file__)
export = exporter.decorator()

options = { 'separators' : (',', ' : '),
             'sort_keys' : True,
                'indent' : 4 }

@export
class JsonFileBase(FileBase, Nested):

    """ The base class for “clu.config.jsonfile.JsonFile”. Override this
        class in your own project to use JSON file data in your Schema
        pipelines as a NamespacedMutableMapping – q.v. the docstring for
        “clu.config.jsonfile.JsonFile” sub.

        This class uses two mixins: both “clu.config.base.AppName” and
        “clu.config.filebase.FileName” are part of its inheritance chain.
        The “AppName” mixin acts on the “appname” class keyword, and the
        “FileName” mixin acts on the “filename” class keyword (furnishing
        many related class methods).
    """

    def loads(self, loaded):
        """ Load nested namespaced dictionary data from a JSON-encoded string """
        self.tree = json.loads(loaded)

    def dumps(self):
        """ Dump a JSON-encoded string from nested namespaced dictionary data """
        return json.dumps(self.tree, **options)

json_appname  = PROJECT_NAME
json_filename = f'{PROJECT_NAME}-config.json'

@export
class JsonFile(JsonFileBase, appname=json_appname,
                            filename=json_filename):

    """ A representation of a JSON file’s data as a NamespacedMutableMapping.

        This class is specifically germane to the CLU project – note
        that the “appname” and “filename” class keywords are used to
        assign values that are CLU-specific.

        CLU users who wish to use JSON files as NamespacedMutableMappings
        in their own projects should create a subclass of JsonFileBase of
        their own. Like this one, it needs to assign both the “appname”
        and the “filename” class keywords; it is unnecessary (but OK!) to
        define further methods, properties, class constants, and whatnot.
    """
    pass

# Assign the modules’ `__all__` and `__dir__` using the exporter:
__all__, __dir__ = exporter.all_and_dir()


from collections import deque
def solution(numbers, target):
    answer = 0
    dq = deque([(0,0)]) # 현재 숫자, 길이
    while dq:
        now = dq.popleft()
        if now[1]==len(numbers):# 길이가 numbers랑 같으면 체크하고 continue
            if now[0] == target:
                answer+=1
        else :
            plus = [now[0] + numbers[now[1]],now[1]+1]
            minus = [now[0] - numbers[now[1]],now[1]+1]
            dq.append(plus)
            dq.append(minus)
    return answer

프로그래머스/타켓 넘버.py


from turtle import forward
from typing import List

import torch
import torch.nn as nn

class MLPLayer(nn.Module):
    def __init__(
        self,
        in_dims: int = 768,
        hidden_dims: List[int] = [768],
        activation: str = "GELU"
    ) -> None:
        super().__init__()

        if activation == "gelu":
            activation_fn = nn.GELU()
        elif activation == "relu":
            activation_fn = nn.ReLU()
        elif activation == "mish":
            activation_fn = nn.Mish()
        elif activation == "leaky_relu":
            activation_fn == nn.LeakyReLU()

        layers = [
            nn.Linear(in_dims, hidden_dims[0]),
            # nn.LayerNorm(hidden_dims[0]),
            activation_fn
        ]

        for i in range(1, len(hidden_dims)):
            layers += [
                nn.Linar(hidden_dims[i - 1], hidden_dims[i]),
                # nn.LayerNorm(hidden_dims[i]),
                activation_fn
            ]

        self.net = nn.Sequential(*layers)

    def forward(self, x: torch.Tensor):
        return self.net(x)


import logging
import os
import time
import random
import json
from tqdm import tqdm
import sys
# import wandb
import torch
from itertools import chain
import torch.nn as nn
from torch.nn.utils import clip_grad_norm_
from torch.utils.data import DataLoader
from tensorboardX import SummaryWriter
from torch.optim.lr_scheduler import StepLR, MultiStepLR

import numpy as np
from configs.opts import parser
from model.main_model_2 import AV_VQVAE_Encoder
from model.main_model_2 import AV_VQVAE_Decoder
from model.main_model_2 import Semantic_Decoder, AVT_VQVAE_Encoder

from utils import AverageMeter, Prepare_logger, get_and_save_args
from utils.container import metricsContainer
from utils.Recorder import Recorder

import torch.nn.functional as F
from torch.nn.utils.rnn import pad_sequence
from dataset.UCF_VGGSOUND_datasets import VGGSoundDataset, UCFDataset

# =================================  seed config ============================
SEED = 43
random.seed(SEED)
np.random.seed(seed=SEED)
torch.manual_seed(seed=SEED)
torch.cuda.manual_seed(seed=SEED)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False

# =============================================================================

def main():
    # utils variable
    global args, logger, writer, dataset_configs
    # statistics variable

    global best_precision
    best_precision=0

    # configs
    dataset_configs = get_and_save_args(parser)
    parser.set_defaults(**dataset_configs)
    args = parser.parse_args()
    # select GPUs
    # os.environ['CUDA_DEVICE_ORDER'] = "PCI_BUS_ID"
    os.environ['CUDA_VISIBLE_DEVICES'] = "0"

    '''Create snapshot_pred dir for copying code and saving models '''
    if not os.path.exists(args.snapshot_pref):
        os.makedirs(args.snapshot_pref)

    if os.path.isfile(args.resume):
        args.snapshot_pref = os.path.dirname(args.resume)

    logger = Prepare_logger(args, eval=args.evaluate)

    if not args.evaluate:
        logger.info(f'\nCreating folder: {args.snapshot_pref}')
        logger.info('\nRuntime args\n\n{}\n'.format(json.dumps(vars(args), indent=4)))
    else:
        logger.info(f'\nLog file will be save in a {args.snapshot_pref}/Eval.log.')

    '''Dataloader selection'''
    if args.dataset_name == 'ucfv_vgga':
        train_csv_path = 'ucf2vggsound.csv'
        val_csv_path = 'vggsound2ucf.csv'
        audio_fea_base_path = 'vggsound40k/feature/audio/zip'
        video_fea_base_path = 'UCF101/feature/video/zip'

        train_dataloader = DataLoader(
            UCFDataset(train_csv_path, video_fea_base_path, split='train', modality='video'),
            batch_size=args.batch_size,
            shuffle=True,
            num_workers=8,
            pin_memory=False
        )
        val_dataloader = DataLoader(
            VGGSoundDataset(val_csv_path, audio_fea_base_path, split='val', modality='audio'),
            batch_size=args.batch_size,
            shuffle=False,
            num_workers=8,
            pin_memory=False
        )
    elif args.dataset_name == 'vgga_ucfv':
        train_csv_path = 'vggsound2ucf.csv'
        val_csv_path = 'ucf2vggsound.csv'
        audio_fea_base_path = 'vggsound40k/feature/audio/zip'
        video_fea_base_path = 'UCF101/feature/video/zip'
        train_dataloader = DataLoader(
            VGGSoundDataset(train_csv_path, audio_fea_base_path, split='train', modality='audio'),
            batch_size=args.batch_size,
            shuffle=True,
            num_workers=8,
            pin_memory=False
        )
        val_dataloader = DataLoader(
            UCFDataset(val_csv_path, video_fea_base_path, split='val', modality='video'),
            batch_size=args.batch_size,
            shuffle=False,
            num_workers=8,
            pin_memory=False
        )
    else:
        raise NotImplementedError

    '''model setting'''
    video_dim = 512
    audio_dim = 128
    video_output_dim = 2048
    audio_output_dim = 256
    text_lstm_dim = 128
    text_output_dim = 256
    n_embeddings = 400
    embedding_dim = 256
    start_epoch = -1
    model_resume = True
    total_step = 0
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # AV
    # Encoder = AV_VQVAE_Encoder(video_dim, audio_dim, video_output_dim, audio_output_dim, n_embeddings, embedding_dim)

    # AVT
    Encoder = AVT_VQVAE_Encoder(audio_dim, video_dim, text_lstm_dim*2, audio_output_dim, video_output_dim, text_output_dim, n_embeddings, embedding_dim)

    Decoder = Semantic_Decoder(input_dim=embedding_dim, class_num = 16)
    Encoder.double()
    Decoder.double()
    '''optimizer setting'''
    Encoder.to(device)
    Decoder.to(device)
    optimizer = torch.optim.Adam(chain(Encoder.parameters(), Decoder.parameters()), lr=args.lr)

    scheduler = MultiStepLR(optimizer, milestones=[10, 20, 30], gamma=0.5)

    '''loss'''
    criterion = nn.BCEWithLogitsLoss().cuda()
    criterion_event = nn.CrossEntropyLoss().cuda()

    if model_resume is True:
        path_checkpoints = "..."
        checkpoints = torch.load(path_checkpoints)
        Encoder.load_state_dict(checkpoints['Encoder_parameters'])
        start_epoch = checkpoints['epoch']
        logger.info("Resume from number {}-th model.".format(start_epoch))

    '''Tensorboard and Code backup'''
    writer = SummaryWriter(args.snapshot_pref)
    recorder = Recorder(args.snapshot_pref, ignore_folder="Exps/")
    recorder.writeopt(args)

    '''Training and Evaluation'''

    for epoch in range(start_epoch+1, args.n_epoch):

        loss, total_step = train_epoch(Encoder, Decoder, train_dataloader, criterion, criterion_event,
                                       optimizer, epoch, total_step, args)
        logger.info(f"epoch: *******************************************{epoch}")

        if ((epoch + 1) % args.eval_freq == 0) or (epoch == args.n_epoch - 1):
            loss = validate_epoch(Encoder, Decoder, val_dataloader, criterion, criterion_event, epoch, args)
            logger.info("-----------------------------")
            logger.info(f"evaluate loss:{loss}")
            logger.info("-----------------------------")
        scheduler.step()

def _export_log(epoch, total_step, batch_idx, lr, loss_meter):
    msg = 'Epoch {}, Batch {}, lr = {:.5f}, '.format(epoch, batch_idx, lr)
    for k, v in loss_meter.items():
        msg += '{} = {:.4f}, '.format(k, v)
    # msg += '{:.3f} seconds/batch'.format(time_meter)
    logger.info(msg)
    sys.stdout.flush()
    loss_meter.update({"batch": total_step})

def to_eval(all_models):
    for m in all_models:
        m.eval()

def to_train(all_models):
    for m in all_models:
        m.train()

def save_models(Encoder, optimizer, epoch_num, total_step, path):
    state_dict = {
        'Encoder_parameters': Encoder.state_dict(),
        'optimizer': optimizer.state_dict(),
        'epoch': epoch_num,
        'total_step': total_step,
    }
    torch.save(state_dict, path)
    logging.info('save model to {}'.format(path))

def train_epoch_check(train_dataloader, epoch, total_step, args):
    # train_dataloader = tqdm(train_dataloader)
    for n_iter, batch_data in enumerate(train_dataloader):

        '''Feed input to model'''
        feature, labels, mask = batch_data['feature'],batch_data['label'],batch_data['mask']
    return torch.zeros(1),torch.zeros(1)

def train_epoch(Encoder, Decoder, train_dataloader, criterion, criterion_event, optimizer, epoch, total_step, args):
    batch_time = AverageMeter()
    data_time = AverageMeter()
    losses = AverageMeter()
    train_precision = AverageMeter()
    end_time = time.time()
    models = [Encoder, Decoder]
    to_train(models)
    # Note: here we set the model to a double type precision,
    # since the extracted features are in a double type.
    # This will also lead to the size of the model double increases.

    Encoder.cuda()
    Decoder.cuda()
    optimizer.zero_grad()

    for n_iter, batch_data in enumerate(train_dataloader):

        data_time.update(time.time() - end_time)
        '''Feed input to model'''
        feat, labels = batch_data
        feat = feat.to(torch.float64).cuda()
        bs = feat.size(0)
        labels = labels.double().cuda()
        labels_foreground = labels[:, :, :-1]
        labels_BCE, labels_evn = labels_foreground.max(-1)
        labels_event, _ = labels_evn.max(-1)

        with torch.no_grad():
            if (args.dataset_name == 'ucfv_vgga'):
                vq = Encoder.Video_VQ_Encoder(feat)
            elif (args.dataset_name == 'vgga_ucfv'):
                vq = Encoder.Audio_VQ_Encoder(feat)
            else:
                raise NotImplementedError
        _class = Decoder(vq)
        event_loss = criterion_event(_class, labels_event.cuda())
        precision = compute_accuracy_supervised(_class, labels)
        loss_items = {
            "train_event_loss":event_loss.item(),
            "train_precision": precision.item(),
        }
        train_precision.update(precision.item(), bs * 10)
        metricsContainer.update("loss", loss_items)
        loss = event_loss

        if n_iter % 10 == 0:
            _export_log(epoch=epoch, total_step=total_step+n_iter, batch_idx=n_iter, lr=optimizer.state_dict()['param_groups'][0]['lr'], loss_meter=metricsContainer.calculate_average("loss"))
        loss.backward()

        '''Clip Gradient'''
        if args.clip_gradient is not None:
            for model in models:
                total_norm = clip_grad_norm_(model.parameters(), args.clip_gradient)

        '''Update parameters'''
        optimizer.step()
        optimizer.zero_grad()

        losses.update(loss.item(), feat.size(0) * 10)
        batch_time.update(time.time() - end_time)
        end_time = time.time()

        '''Add loss of a iteration in Tensorboard'''
        writer.add_scalar('Train_data/loss', losses.val, epoch * len(train_dataloader) + n_iter + 1)

        '''Add loss of an epoch in Tensorboard'''
        writer.add_scalar('Train_epoch_data/epoch_loss', losses.avg, epoch)

    logger.info(f'Train results (precision): {train_precision.avg:.4f}')
    return losses.avg, n_iter + total_step

@torch.no_grad()
def validate_epoch(Encoder,Decoder,val_dataloader, criterion, criterion_event, epoch, args, eval_only=False):
    Sigmoid_fun = nn.Sigmoid()
    batch_time = AverageMeter()
    data_time = AverageMeter()
    losses = AverageMeter()
    val_precision = AverageMeter()
    end_time = time.time()

    Encoder.eval()
    Decoder.eval()
    Encoder.cuda()
    Decoder.cuda()

    for n_iter, batch_data in enumerate(val_dataloader):
        data_time.update(time.time() - end_time)

        '''Feed input to model'''
        feat, labels = batch_data
        feat = feat.to(torch.float64).cuda()
        bs = feat.size(0)
        labels = labels.double().cuda()
        labels_foreground = labels[:, :, :-1]
        labels_BCE, labels_evn = labels_foreground.max(-1)
        labels_event, _ = labels_evn.max(-1)

        with torch.no_grad():
            if (args.dataset_name == 'ucfv_vgga'):
                vq = Encoder.Audio_VQ_Encoder(feat)
            elif (args.dataset_name == 'vgga_ucfv'):
                vq = Encoder.Video_VQ_Encoder(feat)
            else:
                raise NotImplementedError
        _class = Decoder(vq)
        event_loss = criterion_event(_class, labels_event.cuda())
        precision = compute_accuracy_supervised(_class, labels)
        loss_items = {
            "val_event_loss":event_loss.item(),
            "val_precision": precision.item(),
        }
        val_precision.update(precision.item(), bs * 10)
        metricsContainer.update("loss", loss_items)
        loss = event_loss
        losses.update(loss.item(), bs * 10)

    global best_precision
    if val_precision.avg > best_precision:
        best_precision = val_precision.avg

    logger.info(f'Eval results (precision and loss): {val_precision.avg:.4f} {losses.avg:.4f}')
    logger.info(f"Best results (precision): {best_precision:.4f}")
    return losses.avg

def compute_accuracy_supervised(event_scores, labels):
    labels_foreground = labels[:, :, :-1]
    labels_BCE, labels_evn = labels_foreground.max(-1)
    labels_event, _ = labels_evn.max(-1)
    _, event_class = event_scores.max(-1)
    correct = event_class.eq(labels_event)
    correct_num = correct.sum().double()
    acc = correct_num * (100. / correct.numel())
    return acc

if __name__ == '__main__':
    main()


from flask import Flask, request
from textblob import Word

app = Flask(__name__)

@app.route('/plural')
def define():
  word_str = request.args['word']
  word_obj = Word(word_str)
  return word_obj.pluralize() + "\n"

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


import torch
import torch.nn as nn
import numpy as np

class Self_Attn(nn.Module):
    """ Self attention Layer"""

    def __init__(self, in_dim, activation):
        """
        in_dim   -- input feature's channel dim
        activation    -- activation function type
        """
        super(Self_Attn, self).__init__()
        self.input_channel = in_dim
        self.activation = activation
        self.k = 8
        self.query = nn.Conv2d(self.input_channel, self.input_channel // self.k, kernel_size=1)
        self.key = nn.Conv2d(self.input_channel, self.input_channel // self.k, kernel_size=1)
        self.value = nn.Conv2d(self.input_channel, self.input_channel, kernel_size=1)
        self.h = nn.Conv2d(self.input_channel, self.input_channel, kernel_size=1)
        self.gamma = nn.Parameter(torch.zeros(1), requires_grad=True)
        self.softmax = nn.Softmax(dim=-1)

    def forward(self, x):
        batch_size, channel, width, height = x.size()

        q = self.query(x)
        q = q.view(batch_size, -1, width*height)
        print('q.shape',q.shape)
        k = self.key(x)
        k = k.view(batch_size, -1, width*height)
        print('k.shape',k.shape)
        s = torch.bmm(q.transpose(1,2), k)
        print('s.shape',s.shape)
        attention_map = self.softmax(s)
        print('attention_map.shape',attention_map.shape)
        h = self.h(x)
        h = h.view(batch_size, -1, width*height)
        print('h.shape',h.shape)
        attention_map_h = torch.bmm(h, attention_map.transpose(1,2))
        print('attention_map_h.shape',attention_map_h.shape)
        attention_map_h = attention_map_h.view(batch_size, -1, width, height)
        print('attention_map_h.shape',attention_map_h.shape)
        v = self.value(attention_map_h)
        print('v.shape',v.shape)
        out = x + self.gamma*v
        print('out.shape',out.shape)
        return out

        return None

image = np.random.randn(1, 64, 4, 4)
image = image.astype('float32')
image = torch.from_numpy(image)
model = Self_Attn(64, 'relu')
out = model(image)


import input_data
import matplotlib.pyplot as plt
import matplotlib.cm as cm
import matplotlib

# Helper function to display a zip code with predictions
def display_zip(images,labels):
	# Get each digit
	tmp1 = images[0]
	tmp1 = tmp1.reshape((28,28))

	tmp2 = images[1]
	tmp2 = tmp2.reshape((28,28))

	tmp3 = images[2]
	tmp3 = tmp3.reshape((28,28))

	tmp4 = images[3]
	tmp4 = tmp4.reshape((28,28))

	tmp5 = images[4]
	tmp5 = tmp5.reshape((28,28))

	# Add each digit to the figure
	fig = plt.figure()
	a = fig.add_subplot(1,5,1)
	plt.imshow(tmp1, cmap = cm.Greys)
	a.set_title(labels[0])

	a = fig.add_subplot(1,5,2)
	plt.imshow(tmp2, cmap = cm.Greys)
	a.set_title(labels[1])

	a = fig.add_subplot(1,5,3)
	plt.imshow(tmp3, cmap = cm.Greys)
	a.set_title(labels[2])

	a = fig.add_subplot(1,5,4)
	plt.imshow(tmp4, cmap = cm.Greys)
	a.set_title(labels[3])

	a = fig.add_subplot(1,5,5)
	plt.imshow(tmp5, cmap = cm.Greys)
	a.set_title(labels[4])

	# Show the figure
	plt.show()


import pygame
from .base import BaseState

class Gameplay(BaseState):
    def __init__(self):
        super(Gameplay, self).__init__()
        self.rect = pygame.Rect((0, 0), (80, 80))
        self.rect.center = self.screen_rect.center
        self.next_state = "GAME_OVER"

    def get_event(self, event):
        if event.type == pygame.QUIT:
            self.quit = True
        elif event.type == pygame.KEYUP:
            if event.key == pygame.K_UP:
                self.rect.move_ip(0, -10)
            if event.key == pygame.K_DOWN:
                self.rect.move_ip(0, 10)
            if event.key == pygame.K_LEFT:
                self.rect.move_ip(-10, 0)
            if event.key == pygame.K_RIGHT:
                self.rect.move_ip(10, 0)
            if event.key == pygame.K_SPACE:
                self.done = True

    def draw(self, surface):
        surface.fill(pygame.Color("black"))
        pygame.draw.rect(surface, pygame.Color("blue"), self.rect)