import numpy as np
import json
import random
import pdb
from math import ceil
from tqdm import tqdm

class Ontology:

    def __init__(self, dictionary):
        self.dictionary = dictionary

    def add_raw(self, fname, mutual_slot):
        '''Read from json'''
        with open(fname, 'r', encoding='utf8') as f:
            self.ontology_dict = json.load(f)
        self.ontology_list = []
        for s in self.ontology_dict:
            self.ontology_list.append([s, self.ontology_dict[s]])
        self.mutual_slot = mutual_slot
        self._count()
        self._get_descartes()

    def add_examples(self, fname):
        with open(fname, 'r', encoding='utf8') as f:
            self.examples = json.load(f)

    def _get_descartes(self):
        self.descartes = []
        for slot, values in self.ontology_list:
            if slot == self.mutual_slot:
                continue
            for value in values:
                for status in self.ontology_dict[self.mutual_slot]:
                    item = '{}:{}-{}:{}'.format(slot, value, self.mutual_slot, status)
                    self.descartes.append(item)

    def _count(self):

        n_status = len(self.ontology_dict[self.mutual_slot])
        self.num = {
            'global': sum([len(item[1]) for item in self.ontology_list \
                if item[0] != self.mutual_slot]) * n_status
        }
        for item in self.ontology_list:
            slot, values = item
            if slot == self.mutual_slot:
                continue
            self.num[slot] = len(values) * n_status

    def add(self, fname):

        with open(fname, 'r', encoding='utf8') as f:
            ontology = json.load(f)
        self.ontology_list = ontology['ontology_list']
        self.ontology_dict = ontology['ontology_dict']
        self.mutual_slot = ontology['mutual_slot']
        self.examples = ontology['examples']
        self._count()
        self._get_descartes()

    def save(self, fname):

        with open(fname, 'w', encoding='utf8') as f:
            ontology = {
                'ontology_list': self.ontology_list,
                'ontology_dict': self.ontology_dict,
                'mutual_slot': self.mutual_slot,
                'examples': self.examples
            }
            json.dump(ontology, f, indent=4, ensure_ascii=False)

    def _get_offsets(self):
        if getattr(self, 'offsets', None) is None:
            offsets = []
            for i, item in enumerate(self.ontology_list):
                slot, values = item
                if slot == self.mutual_slot:
                    continue
                offsets.append(len(values))
                if i > 0:
                    offsets[i] += offsets[i - 1]
            self.offsets = offsets
            self.offsets = [0] + self.offsets[:-1]

    def _pad(self, labels_seqs):
        # 如果每种编号有n个，则第n个为<start>，n+1为<end>，n+2为<pad>
        labels_seqs_len = dict()
        for slot in labels_seqs.keys():
            pad_id = self.num[slot] + 2
            n = len(labels_seqs[slot])
            labels_seqs_len[slot] = n
            labels_seqs[slot] = labels_seqs[slot] + \
                [pad_id] * (self.num[slot] - n)

        return labels_seqs, labels_seqs_len

    def _window_labels2indexes(self, texts):

        n_status = len(self.ontology_dict[self.mutual_slot])
        labels_indexes_dict = dict()
        for slot in self.ontology_dict:
            if slot == self.mutual_slot:
                continue
            labels_indexes_dict[slot] = [0] * n_status * len(self.ontology_dict[slot])

        for text in texts:
            # 疾病:心律不齐-状态:阳性
            sv1, sv2 = text.split('-')
            s1, v1 = sv1.split(':')
            s2, v2 = sv2.split(':')
            v1_idx = self.ontology_dict[s1].index(v1)
            v2_idx = self.ontology_dict[s2].index(v2)
            slot_idx = v1_idx * n_status + v2_idx
            labels_indexes_dict[s1][slot_idx] = 1

        labels_indexes = []
        for slot, values in self.ontology_list:
            if slot == self.mutual_slot:
                continue
            labels_indexes += labels_indexes_dict[slot]

        return labels_indexes

    def labels2indexes(self, texts, style):

        self._get_offsets()
        return self._window_labels2indexes(texts)

    def vec2label(self, label_v):
        label_w = []
        for i, item in enumerate(label_v):
            if item == 1:
                label_w.append(self.descartes[i])
        return label_w

    def onto2ids(self):
        # max_len = max([max([len(slot) + len(value) + 1 for value in self.ontology_dict]) \
        #     for slot in self.ontology_dict])
        max_len = 35
        slots_values = dict()
        slots_values_lens = dict()
        for slot in self.ontology_dict.keys():
            for value in self.ontology_dict[slot]:
                seq = slot + ':' + value
                if len(self.examples[slot][value]) != 0:
                    seq += '('
                    for example in self.examples[slot][value]:
                        seq += example + ','
                    seq += ')'
                length = len(seq)
                seq = self.dictionary.words2ids(seq, max_len)
                try:
                    slots_values[slot].append(seq)
                    slots_values_lens[slot].append(length)
                except:
                    slots_values[slot] = [seq]
                    slots_values_lens[slot] = [length]
            slots_values[slot] = np.array(slots_values[slot], np.int32)
            slots_values_lens[slot] = np.array(slots_values_lens[slot], np.int32)

        return slots_values, slots_values_lens

class Corp:

    def __init__(self, fname):
        self.fname = fname

    def __iter__(self):
        with open(self.fname, 'r', encoding='utf8') as f:
            for line in f:
                if len(line) > 10:
                    yield list(line)

class Dictionary:

    def __init__(self):
        pass

    def load(self, fname):
        '''Load the dictionary'''
        import pdb
        with open(fname, 'r', encoding='utf8') as f:
            first_line = next(f)
            self.vocab_size = int(first_line.split(' ')[0])
            self.emb_size = int(first_line.split(' ')[1])
            self.emb = []
            self.vocab_list = []
            for line in tqdm(f):
                line_list = line.split(' ')
                if line[0] == ' ':
                    word = ' '
                    emb = np.array(list(map(lambda x: float(x), line_list[2:-1])))
                else:
                    word = line_list[0]
                    emb = np.array(list(map(lambda x: float(x), line_list[1:-1])))
                self.vocab_list.append(word)
                self.emb.append(emb)
        self.emb = np.array(self.emb, np.float32)
        self._create_index()

    def _create_index(self):

        self._id2word = self.vocab_list
        self._word2id = dict([(word,i) for i,word in enumerate(self._id2word)])

    def train(self, fname, emb_size=300):
        '''Train the word embedding'''
        self.emb_size = emb_size
        counter = Counter()
        for sentence in Corp(fname):
            counter.update(sentence)

        # 去除出现数量小于10的字
        sorted_vocab = sorted(counter.items(), key=itemgetter(1), reverse=True)
        sorted_vocab = filter(lambda x: x[1] >= 10, sorted_vocab)
        # 创建 vocab_list
        self.vocab_list = [x[0] for x in sorted_vocab]
        self.vocab_list.append('<pad>')
        self.vocab_list.append('<unk>')
        self.vocab_size = len(self.vocab_list)
        self._create_index()

        # 训练词向量
        model = Word2Vec(Corp(fname), size=emb_size, sg=1, min_count=1)
        self.emb = []
        for word in self.vocab_list[:-2]:
            self.emb.append(model[word])
        self.emb = np.array(self.emb, np.float32)
        self.emb = np.concatenate([self.emb, np.zeros([1, emb_size])], 0)
        self.emb = np.concatenate([self.emb, np.random.rand(1, emb_size) / 2], 0)
        self.emb = self.emb.astype(np.float32)

    def save(self, fname):
        with open(fname, 'w', encoding='utf8') as f:
            f.write('{} {}\n'.format(self.vocab_size, self.emb_size))
            for i, word in enumerate(self.vocab_list):
                if word == '\n':
                    word = '\\n'
                emb = list(map(lambda x: str(x), self.emb[i].tolist()))
                emb = ' '.join(emb)
                content = word + ' ' + emb + ' \n'
                f.write(content)

    def _build_norm(self):
        if getattr(self, 'emb_norm', None) is None:
            self.emb_norm = np.linalg.norm(self.emb, axis=1) # [vocab_size]

    def id2word(self, idx):
        return self._id2word[idx]

    def word2id(self, word, return_unk=False):
        try:
            return self._word2id[word]
        except KeyError:
            if return_unk:
                return self._word2id['<unk>']
            else:
                raise

    def id2emb(self, idx):
        return self.emb[idx]

    def word2emb(self, word, return_unk=False):
        idx = self.word2id(word, return_unk)
        return self.emb[idx]

    def words2ids(self, words, max_len=50):
        n = len(words)
        ids = [self.word2id(word, True) for word in words]
        if n < max_len:
            ids += [self.word2id('<pad>')] * (max_len - n)
        else:
            ids = ids[: max_len]
        return ids

    def most_similar(self, word, topn=10):
        # 获得词向量的L2范数
        self._build_norm()

        # 计算consine
        idx = self.word2id(word, False)
        emb = self.emb[idx]
        emb_norm = self.emb_norm[idx]
        dot = np.dot(self.emb, emb)
        cosine = dot / (emb_norm * self.emb_norm)
        cosine = -cosine

        # 前topn+1个数的序号数组，未排序
        most_extreme = np.argpartition(cosine, topn+1)[: topn+1]

        # 前topn+1个数排序后的序号数组
        sorted_idx = most_extreme.take(np.argsort(cosine.take(most_extreme)))

        similar_list = [(self.id2word(idx), -cosine[idx]) for idx in sorted_idx[1:]]
        return similar_list

class Data:
    def __init__(self, max_len, dictionary, ontology):
        self.max_len = max_len
        self.datasets = dict()
        self.dictionary = dictionary
        self.ontology = ontology

    def _add_window(self, name, fname):
        with open(fname, 'r', encoding='utf8') as f:
            dialogs = json.load(f)
        self.datasets[name] = {
            'origin': dialogs,
            'windows_utts': [],
            'windows_utts_len': [],
            'labels': []
        }
        for dialog in dialogs:
            for window in dialog:
                window_utts = window['utterances']
                self.window_size = len(window_utts)
                window_utts_len = [len(utt) for utt in window_utts]
                self.datasets[name]['windows_utts_len'].append(window_utts_len)
                window_utts = [self.dictionary.words2ids(utt, self.max_len) \
                    for utt in window_utts]
                self.datasets[name]['windows_utts'].append(window_utts)
                label = window['label']
                label = filter(lambda x: '药物' not in x, label)
                label_indexes = self.ontology.labels2indexes(label, 'window')
                self.datasets[name]['labels'].append(label_indexes)
        self.datasets[name]['windows_utts'] = np.array(
            self.datasets[name]['windows_utts'], np.int32)
        self.datasets[name]['windows_utts_len'] = np.array(
            self.datasets[name]['windows_utts_len'], np.int32)
        self.datasets[name]['labels'] = np.array(
            self.datasets[name]['labels'], np.int32)
        self.datasets[name]['style'] = 'window'
        self.datasets[name]['num'] = self.datasets[name]['windows_utts'].shape[0]

    def add_raw(self, name, fname, style):
        if style == 'window':
            self._add_window(name, fname)
        elif style == 'window-chunk':
            self._add_window_chunk(name, fname)

    def _add_window_chunk(self, name, fname):
        with open(fname, 'r', encoding='utf8') as f:
            dialogs = json.load(f)
        self.datasets[name] = {
            'origin': dialogs,
            'chunk_seqs': [],
            'chunk_seqs_len': [],
            'labels': []
        }
        for dialog in dialogs:
            for window in dialog:
                utts = window['utterances']
                chunk_seq = []
                for utt in utts:
                    if len(utt) != 0:
                        chunk_seq += list(utt) + ['<seg>']
                    else:
                        chunk_seq += list(utt)
                chunk_seq = chunk_seq[:-1]
                chunk_seq_len = len(chunk_seq)
                chunk_seq_ids = self.dictionary.words2ids(chunk_seq, self.max_len)

                self.datasets[name]['chunk_seqs'].append(chunk_seq_ids)
                self.datasets[name]['chunk_seqs_len'].append(chunk_seq_len)

                label = window['label']
                label = filter(lambda x: '药物' not in x, label)
                label_indexes = self.ontology.labels2indexes(label, 'window')
                self.datasets[name]['labels'].append(label_indexes)

        self.datasets[name]['chunk_seqs'] = np.array(
            self.datasets[name]['chunk_seqs'], np.int32)
        self.datasets[name]['chunk_seqs_len'] = np.array(
            self.datasets[name]['chunk_seqs_len'], np.int32)
        self.datasets[name]['labels'] = np.array(
            self.datasets[name]['labels'], np.int32)
        self.datasets[name]['style'] = 'window-chunk'
        self.datasets[name]['num'] = self.datasets[name]['chunk_seqs'].shape[0]

    def _batch_window(self, name, batch_size, shuffle=True):
        n = self.datasets[name]['num']
        batch_num = int(ceil(n / batch_size))
        indexes = list(range(n))
        if shuffle:
            random.shuffle(indexes)
        for i in range(batch_num):
            indexes_batch = indexes[i * batch_size: (i + 1) * batch_size]
            windows_utts_batch = self.datasets[name]['windows_utts'].\
                take(indexes_batch, 0)
            windows_utts_lens_batch = self.datasets[name]['windows_utts_len'].\
                take(indexes_batch, 0)
            labels_batch = self.datasets[name]['labels'].\
                take(indexes_batch, 0)
            yield windows_utts_batch, windows_utts_lens_batch, labels_batch

    def _batch_window_chunk(self, name, batch_size, shuffle=True):
        n = self.datasets[name]['num']
        batch_num = int(ceil(n / batch_size))
        indexes = list(range(n))
        if shuffle:
            random.shuffle(indexes)
        for i in range(batch_num):
            indexes_batch = indexes[i * batch_size: (i + 1) * batch_size]
            chunk_seqs_batch = self.datasets[name]['chunk_seqs'].\
                take(indexes_batch, 0)
            chunk_seqs_len_batch = self.datasets[name]['chunk_seqs_len'].\
                take(indexes_batch, 0)
            labels_batch = self.datasets[name]['labels'].\
                take(indexes_batch, 0)
            yield chunk_seqs_batch, chunk_seqs_len_batch, labels_batch

    def batch(self, name, batch_size, shuffle=True):
        if self.datasets[name]['style'] == 'window':
            yield from self._batch_window(name, batch_size, shuffle)
        elif self.datasets[name]['style'] == 'window-chunk':
            yield from self._batch_window_chunk(name, batch_size, shuffle)


import numpy as np
import pandas as pd
from joblib import Parallel, delayed
import gensim.utils as gu
from sklearn.preprocessing import OneHotEncoder, LabelEncoder
import os
import argparse

def doWork(i, lyrics, genre):
    if not i%1000:
        print("song number", i)
        #print("LYRICS: ", lyrics)
        #print("GENRE", genre)
    if not pd.isnull(lyrics):
        tokens = gu.simple_preprocess(lyrics)
        if tokens is not None and len(tokens) > 0:
            return tokens, genre

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('-o', '--outdir', action='store', help='Output directory', required=True)
    parser.add_argument('-i', '--infile', action='store', help='Input CSV file with lyrics and labels', required=True)
    parser.add_argument('--songs-per-genre', action='store', type=int, help='Number of songs per genre', required=True)
    args = parser.parse_args()

    print("reading lyrics csv")
    chosen_genres = ["Rock","R&B","Pop","Metal","Jazz","Indie","Hip-Hop","Folk","Electronic","Country"]
    input_df = pd.read_csv(args.infile)
    input_df = input_df[(input_df.genre.isin(chosen_genres))].dropna(0, subset=["lyrics"])

    seed=0
    by_genre = []
    for g in chosen_genres:
        by_genre.append(input_df[(input_df.genre == g)].sample(n=args.songs_per_genre,replace=True,random_state=seed))
    input_df = pd.concat(by_genre)

# PREPARE Word list Corpus and Genres
    corpus_with_genres = Parallel(n_jobs=4)(delayed(doWork)(i, series['lyrics'], series['genre']) for (i, series) in input_df.iterrows())
    #print("separating corpus")
    #corpus = [t[0] for t in corpusAndGenres]
    #print("separating genres")
    #genres = df["genre"][[t[1] for t in corpusAndGenres]]
    #print("label encoding genres")
    #leGenres = LabelEncoder().fit_transform(genres)
    corpus = []
    genres = []
    for c in corpus_with_genres:
        if c is not None and len(c[0]) > 0:
            corpus.append(c[0])
            genres.append(c[1])
    print("saving files")
    np.save(os.path.join(args.outdir, "corpus.npy"), corpus)
    np.save(os.path.join(args.outdir, "genres.npy"), genres)
    #np.save(os.path.join(outdir,"../data/genres.npy"))
    #np.save(os.path.join(outdir,"../data/leGenres.npy"))

if __name__ == "__main__":
    main()


import sys
import webrtcvad
import numpy as np
from mic_array import MicArray
from pixel_ring import pixel_ring

RATE = 16000
CHANNELS = 4
VAD_FRAMES = 10     # ms
DOA_FRAMES = 200    # ms

def main():
    vad = webrtcvad.Vad(3)

    speech_count = 0
    chunks = []
    doa_chunks = int(DOA_FRAMES / VAD_FRAMES)

    try:
        with MicArray(RATE, CHANNELS, RATE * VAD_FRAMES / 1000)  as mic:
            for chunk in mic.read_chunks():
                # Use single channel audio to detect voice activity
                if vad.is_speech(chunk[0::CHANNELS].tobytes(), RATE):
                    speech_count += 1
                    sys.stdout.write('1')
                else:
                    sys.stdout.write('0')

                sys.stdout.flush()

                chunks.append(chunk)
                if len(chunks) == doa_chunks:
                    if speech_count > (doa_chunks / 2):
                        frames = np.concatenate(chunks)
                        direction = mic.get_direction(frames)
                        pixel_ring.set_direction(direction)
                        print('\n{}'.format(int(direction)))

                    speech_count = 0
                    chunks = []

    except KeyboardInterrupt:
        pass

    pixel_ring.off()

if __name__ == '__main__':
    main()


#!/usr/bin/python3

import numpy as np
import torch
from torch import Tensor
from torch.utils.data import Dataset, DataLoader
import pickle
import os
import librosa
from torch.utils.data.dataloader import default_collate
import torchaudio

class ASVspoof2019LA(Dataset):
    def __init__(self, path_to_audio='/data/neil/DS_10283_3336/', path_to_features='/data2/neil/ASVspoof2019LA/',
                 part='train', feature='LFCC', feat_len=750, genuine_only=False):
        super(ASVspoof2019LA, self).__init__()
        self.path_to_audio = path_to_audio
        self.path_to_features = path_to_features
        self.part = part
        self.ptf = os.path.join(path_to_features, self.part)
        self.feat_len = feat_len
        self.feature = feature
        self.genuine_only = genuine_only
        self.tag = {"-": 0, "A01": 1, "A02": 2, "A03": 3, "A04": 4, "A05": 5, "A06": 6, "A07": 7, "A08": 8, "A09": 9,
                      "A10": 10, "A11": 11, "A12": 12, "A13": 13, "A14": 14, "A15": 15, "A16": 16, "A17": 17, "A18": 18,
                      "A19": 19}
        self.label = {"spoof": 1, "bonafide": 0}
        self.all_files = librosa.util.find_files(os.path.join(self.ptf, self.feature), ext="pt")
        if self.genuine_only:
            if self.part in ["train", "dev"]:
                num_bonafide = {"train": 2580, "dev": 2548}
                self.all_files = self.all_files[:num_bonafide[self.part]]
            else:
                res = []
                for item in self.all_files:
                    if "bonafide" in item:
                        res.append(item)
                self.all_files = res
                assert len(self.all_files) == 7355

    def __len__(self):
        return len(self.all_files)

    def __getitem__(self, idx):
        filepath = self.all_files[idx]
        basename = os.path.basename(filepath)
        all_info = basename.split(".")[0].split("_")
        filename = "_".join(all_info[1:4])
        if self.feature != "Raw":
            featureTensor = torch.load(filepath)
            this_feat_len = featureTensor.shape[1]
            if this_feat_len > self.feat_len:
                startp = np.random.randint(this_feat_len - self.feat_len)
                featureTensor = featureTensor[:, startp:startp + self.feat_len, :]
            elif this_feat_len < self.feat_len:
                featureTensor = repeat_padding_Tensor(featureTensor, self.feat_len)
        else:
            # file_path = os.path.join(self.path_to_audio, "LA/ASVspoof2019_LA_" + self.part, "flac", filename+".flac")
            # featureTensor, sr = torchaudio.load(file_path)
            featureTensor = torch.load(filepath)
            this_feat_len = featureTensor.shape[1]
            if this_feat_len > self.feat_len:
                startp = np.random.randint(this_feat_len - self.feat_len)
                featureTensor = featureTensor[:, startp:startp + self.feat_len]
            elif this_feat_len < self.feat_len:
                featureTensor = repeat_padding_RawTensor(featureTensor, self.feat_len)

        tag = self.tag[all_info[4]]
        label = self.label[all_info[5]]
        return featureTensor, filename, tag, label, 2019

    def collate_fn(self, samples):
        return default_collate(samples)

class VCC2020(Dataset):
    def __init__(self, path_to_features="/data2/neil/VCC2020/", feature='LFCC',
                 feat_len=750, genuine_only=False):
        super(VCC2020, self).__init__()
        self.ptf = path_to_features
        self.feat_len = feat_len
        self.feature = feature
        self.tag = {"-": 0, "SOU": 20, "T01": 21, "T02": 22, "T03": 23, "T04": 24, "T05": 25, "T06": 26, "T07": 27, "T08": 28, "T09": 29,
                    "T10": 30, "T11": 31, "T12": 32, "T13": 33, "T14": 34, "T15": 35, "T16": 36, "T17": 37, "T18": 38, "T19": 39,
                    "T20": 40, "T21": 41, "T22": 42, "T23": 43, "T24": 44, "T25": 45, "T26": 46, "T27": 47, "T28": 48, "T29": 49,
                    "T30": 50, "T31": 51, "T32": 52, "T33": 53, "TAR": 54}
        self.label = {"spoof": 1, "bonafide": 0}
        self.genuine_only = genuine_only
        self.all_files = librosa.util.find_files(os.path.join(self.ptf, self.feature), ext="pt")

    def __len__(self):
        if self.genuine_only:
            return 220
        return len(self.all_files)

    def __getitem__(self, idx):
        filepath = self.all_files[idx]
        basename = os.path.basename(filepath)
        all_info = basename.split(".")[0].split("_")
        featureTensor = torch.load(filepath)
        this_feat_len = featureTensor.shape[1]
        if this_feat_len > self.feat_len:
            startp = np.random.randint(this_feat_len - self.feat_len)
            featureTensor = featureTensor[:, startp:startp + self.feat_len, :]
        elif this_feat_len < self.feat_len:
            featureTensor = repeat_padding_Tensor(featureTensor, self.feat_len)
        tag = self.tag[all_info[-2]]
        label = self.label[all_info[-1]]
        return featureTensor, basename, tag, label, 2020

    def collate_fn(self, samples):
        return default_collate(samples)

class ASVspoof2015(Dataset):
    def __init__(self, path_to_features, part='train', feature='LFCC', feat_len=750):
        super(ASVspoof2015, self).__init__()
        self.path_to_features = path_to_features
        self.part = part
        self.ptf = os.path.join(path_to_features, self.part)
        self.feat_len = feat_len
        self.feature = feature
        self.tag = {"human": 0, "S1": 1, "S2": 2, "S3": 3, "S4": 4, "S5": 5,
                    "S6": 6, "S7": 7, "S8": 8, "S9": 9, "S10": 10}
        self.label = {"spoof": 1, "human": 0}
        self.all_files = librosa.util.find_files(os.path.join(self.ptf, self.feature), ext="pt")

    def __len__(self):
        return len(self.all_files)

    def __getitem__(self, idx):
        filepath = self.all_files[idx]
        basename = os.path.basename(filepath)
        all_info = basename.split(".")[0].split("_")
        featureTensor = torch.load(filepath)
        this_feat_len = featureTensor.shape[1]
        if this_feat_len > self.feat_len:
            startp = np.random.randint(this_feat_len - self.feat_len)
            featureTensor = featureTensor[:, startp:startp + self.feat_len, :]
        elif this_feat_len < self.feat_len:
            featureTensor = repeat_padding_Tensor(featureTensor, self.feat_len)
        filename =  all_info[1]
        tag = self.tag[all_info[-2]]
        label = self.label[all_info[-1]]
        return featureTensor, filename, tag, label, 2015

    def collate_fn(self, samples):
        return default_collate(samples)

class ASVspoof2019LASim(Dataset):
    def __init__(self, path_to_features="/data2/neil/ASVspoof2019LA/",
                 path_to_devicedaudio="/data/shared/ASVspoof2019LA-Sim",
                 path_to_deviced="/dataNVME/neil/ASVspoof2019LADevice",
                 part="train", feature='LFCC', feat_len=750):
        super(ASVspoof2019LASim, self).__init__()
        self.path_to_features = path_to_features
        suffix = {"train" : "", "dev":"Dev", "eval": "Eval"}
        self.path_to_deviced = path_to_deviced + suffix[part]
        self.path_to_devicedaudio = os.path.join(path_to_devicedaudio, part)
        self.path_to_features = path_to_features
        self.ptf = os.path.join(path_to_features, part)
        self.feat_len = feat_len
        self.feature = feature
        self.tag = {"-": 0, "A01": 1, "A02": 2, "A03": 3, "A04": 4, "A05": 5, "A06": 6, "A07": 7, "A08": 8, "A09": 9,
                    "A10": 10, "A11": 11, "A12": 12, "A13": 13, "A14": 14, "A15": 15, "A16": 16, "A17": 17, "A18": 18,
                    "A19": 19}
        self.label = {"spoof": 1, "bonafide": 0}
        self.devices = ['AKSPKRS80sUk002-16000', 'AKSPKRSVinUk002-16000', 'Doremi-16000', 'RCAPB90-16000',
                        'ResloRBRedLabel-16000', 'AKSPKRSSpeaker002-16000', 'BehritoneirRecording-16000',
                        'OktavaML19-16000', 'ResloRB250-16000', 'SonyC37Fet-16000', ""]
        if part == "eval":
            self.devices = ['AKSPKRS80sUk002-16000', 'AKSPKRSVinUk002-16000', 'Doremi-16000', 'RCAPB90-16000',
                            'ResloRBRedLabel-16000', 'AKSPKRSSpeaker002-16000', 'BehritoneirRecording-16000',
                            'OktavaML19-16000', 'ResloRB250-16000', 'SonyC37Fet-16000',
                            'iPadirRecording-16000', 'iPhoneirRecording-16000', ""]
        self.original_all_files = librosa.util.find_files(os.path.join(self.ptf, self.feature), ext="pt")
        self.deviced_all_files = [librosa.util.find_files(os.path.join(self.path_to_deviced, devicex), ext="pt") for devicex in self.devices]

    def __len__(self):
        return len(self.original_all_files) * len(self.devices)

    def __getitem__(self, idx):
        device_idx = idx % len(self.devices)
        filename_idx = idx // len(self.devices)
        if device_idx == 0:
            filepath = self.original_all_files[filename_idx]
        else:
            filepath = self.deviced_all_files[device_idx-1][filename_idx]

        basename = os.path.basename(filepath)
        all_info = basename.split(".")[0].split("_")
        filename = "_".join(all_info[1:4])

        if self.feature != "Raw":
            featureTensor = torch.load(filepath)
            this_feat_len = featureTensor.shape[1]
            if this_feat_len > self.feat_len:
                startp = np.random.randint(this_feat_len - self.feat_len)
                featureTensor = featureTensor[:, startp:startp + self.feat_len, :]
            elif this_feat_len < self.feat_len:
                featureTensor = repeat_padding_Tensor(featureTensor, self.feat_len)
        else:
            if device_idx == 0:
                featureTensor = torch.load(filepath)
            else:
                device_name = self.devices[device_idx-1]
                file_path = os.path.join(self.path_to_devicedaudio, device_name, filename+".wav")
                featureTensor, sr = torchaudio.load(file_path)
            this_feat_len = featureTensor.shape[1]
            if this_feat_len > self.feat_len:
                startp = np.random.randint(this_feat_len - self.feat_len)
                featureTensor = featureTensor[:, startp:startp + self.feat_len]
            elif this_feat_len < self.feat_len:
                featureTensor = repeat_padding_RawTensor(featureTensor, self.feat_len)

        tag = self.tag[all_info[4]]
        label = self.label[all_info[5]]
        return featureTensor, filename, tag, label, device_idx

    def collate_fn(self, samples):
        return default_collate(samples)

class ASVspoof2021LAeval(Dataset):
    def __init__(self, path_to_features="/dataNVME/neil/ASVspoof2021LAFeatures",
                 feature='LFCC', feat_len=500):
        super(ASVspoof2021LAeval, self).__init__()
        self.path_to_features = path_to_features
        self.ptf = path_to_features
        self.feat_len = feat_len
        self.feature = feature
        self.all_files = librosa.util.find_files(os.path.join(self.ptf, self.feature), ext="pt")

    def __len__(self):
        return len(self.all_files)

    def __getitem__(self, idx):
        filepath = self.all_files[idx]
        basename = os.path.basename(filepath)
        all_info = basename.split(".")[0].split("_")
        assert len(all_info) == 4
        featureTensor = torch.load(filepath)
        this_feat_len = featureTensor.shape[1]

        if this_feat_len > self.feat_len:
            startp = np.random.randint(this_feat_len - self.feat_len)
            featureTensor = featureTensor[:, startp:startp + self.feat_len, :]
        elif this_feat_len < self.feat_len:
            featureTensor = repeat_padding_Tensor(featureTensor, self.feat_len)

        filename =  "_".join(all_info[1:])
        return featureTensor, filename

    def collate_fn(self, samples):
        return default_collate(samples)

class ASVspoof2019Transm_aug(Dataset):
    def __init__(self, path_to_ori="/data2/neil/ASVspoof2019LA/",
                 path_to_augFeatures="/dataNVME/neil/ASVspoof2019LA_augFeatures",
                 part="train", feature='LFCC', feat_len=500):
        super(ASVspoof2019Transm_aug, self).__init__()
        if feature == "Melspec":
            path_to_augFeatures = "/data3/neil/ASVspoof2019LA_augFeatures"
        self.path_to_features = path_to_augFeatures
        self.part = part
        self.ori = os.path.join(path_to_ori, part)
        self.ptf = os.path.join(path_to_augFeatures, part)
        self.feat_len = feat_len
        self.feature = feature
        self.ori_files = librosa.util.find_files(os.path.join(self.ori, self.feature), ext="pt")
        self.all_files = librosa.util.find_files(os.path.join(self.ptf, self.feature), ext="pt")
        self.tag = {"-": 0, "A01": 1, "A02": 2, "A03": 3, "A04": 4, "A05": 5, "A06": 6}
        self.label = {"spoof": 1, "bonafide": 0}
        self.channel = ['no_channel', 'amr[br=10k2,nodtx]', 'amr[br=5k9]', 'amr[br=6k7,nodtx]',
                        'amr[br=7k95,nodtx]', 'amrwb[br=12k65]', 'amrwb[br=15k85]', 'g711[law=a]',
                        'g711[law=u]', 'g722[br=64k]', 'g726[law=a,br=16k]', 'g726[law=a,br=24k]',
                        'g726[law=u,40k]', 'g726[law=u,br=24k]', 'g726[law=u,br=32k]', 'g728',
                        'silk[br=10k,loss=10]', 'silk[br=15k,loss=5]', 'silk[br=15k]',
                        'silk[br=20k,loss=5]', 'silk[br=5k,loss=10]', 'silk[br=5k]', 'amr[br=12k2]',
                        'amr[br=5k9,nodtx]', 'amrwb[br=6k6,nodtx]', 'g722[br=56k]', 'g726[law=a,br=32k]',
                        'g726[law=a,br=40k]', 'silk[br=15k,loss=10]', 'silk[br=20k]',
                        'silkwb[br=10k,loss=5]', 'amr[br=10k2]', 'amr[br=4k75]', 'amr[br=7k95]',
                        'amrwb[br=15k85,nodtx]', 'amrwb[br=23k05]', 'g726[law=u,br=16k]', 'g729a',
                        'gsmfr', 'silkwb[br=10k,loss=10]', 'silkwb[br=20k]', 'silkwb[br=30k,loss=10]',
                        'amr[br=7k4,nodtx]', 'amrwb[br=6k6]', 'silk[br=10k]', 'silk[br=5k,loss=5]',
                        'silkwb[br=30k,loss=5]', 'amr[br=4k75,nodtx]', 'amr[br=7k4]', 'g722[br=48k]',
                        'silk[br=20k,loss=10]', 'silkwb[br=30k]', 'amr[br=5k15]',
                        'silkwb[br=20k,loss=5]', 'amrwb[br=23k05,nodtx]', 'amrwb[br=12k65,nodtx]',
                        'silkwb[br=20k,loss=10]', 'amr[br=6k7]', 'silkwb[br=10k]', 'silk[br=10k,loss=5]']
        self.channel_dict = dict(zip(iter(self.channel), range(len(self.channel))))

    def __len__(self):
        return len(self.ori_files) + len(self.all_files)

    def __getitem__(self, idx):
        if idx < len(self.ori_files):
            filepath = self.ori_files[idx]
            basename = os.path.basename(filepath)
            all_info = basename.split(".")[0].split("_")
            assert len(all_info) == 6
            channel = "no_channel"
        else:
            filepath = self.all_files[idx - len(self.ori_files)]
            basename = os.path.basename(filepath)
            all_info = basename.split(".")[0].split("_")
            assert len(all_info) == 7
            channel = all_info[6]
        featureTensor = torch.load(filepath)
        this_feat_len = featureTensor.shape[1]
        if this_feat_len > self.feat_len:
            startp = np.random.randint(this_feat_len - self.feat_len)
            featureTensor = featureTensor[:, startp:startp + self.feat_len, :]
        elif this_feat_len < self.feat_len:
            featureTensor = repeat_padding_Tensor(featureTensor, self.feat_len)
        filename = "_".join(all_info[1:4])
        tag = self.tag[all_info[4]]
        label = self.label[all_info[5]]
        return featureTensor, filename, tag, label, self.channel_dict[channel]

    def collate_fn(self, samples):
        return default_collate(samples)

class ASVspoof2019TransmDevice_aug(Dataset):
    def __init__(self, path_to_ori="/data2/neil/ASVspoof2019LA/",
                 path_to_augFeatures="/data3/neil/ASVspoof2019LAPA_augFeatures",
                 part="train", feature='LFCC', feat_len=500):
        super(ASVspoof2019TransmDevice_aug, self).__init__()
        self.path_to_features = path_to_augFeatures
        self.part = part
        self.ori = os.path.join(path_to_ori, part)
        self.ptf = os.path.join(path_to_augFeatures, part)
        self.feat_len = feat_len
        self.feature = feature
        self.ori_files = librosa.util.find_files(os.path.join(self.ori, self.feature), ext="pt")
        self.all_files = librosa.util.find_files(os.path.join(self.ptf, self.feature), ext="pt")
        self.tag = {"-": 0, "A01": 1, "A02": 2, "A03": 3, "A04": 4, "A05": 5, "A06": 6}
        self.label = {"spoof": 1, "bonafide": 0}
        self.channel = ['no_channel', 'amr[br=10k2,nodtx]', 'amr[br=5k9]', 'amr[br=6k7,nodtx]',
                        'amr[br=7k95,nodtx]', 'amrwb[br=12k65]', 'amrwb[br=15k85]', 'g711[law=a]',
                        'g711[law=u]', 'g722[br=64k]', 'g726[law=a,br=16k]', 'g726[law=a,br=24k]',
                        'g726[law=u,40k]', 'g726[law=u,br=24k]', 'g726[law=u,br=32k]', 'g728',
                        'silk[br=10k,loss=10]', 'silk[br=15k,loss=5]', 'silk[br=15k]',
                        'silk[br=20k,loss=5]', 'silk[br=5k,loss=10]', 'silk[br=5k]', 'amr[br=12k2]',
                        'amr[br=5k9,nodtx]', 'amrwb[br=6k6,nodtx]', 'g722[br=56k]', 'g726[law=a,br=32k]',
                        'g726[law=a,br=40k]', 'silk[br=15k,loss=10]', 'silk[br=20k]',
                        'silkwb[br=10k,loss=5]', 'amr[br=10k2]', 'amr[br=4k75]', 'amr[br=7k95]',
                        'amrwb[br=15k85,nodtx]', 'amrwb[br=23k05]', 'g726[law=u,br=16k]', 'g729a',
                        'gsmfr', 'silkwb[br=10k,loss=10]', 'silkwb[br=20k]', 'silkwb[br=30k,loss=10]',
                        'amr[br=7k4,nodtx]', 'amrwb[br=6k6]', 'silk[br=10k]', 'silk[br=5k,loss=5]',
                        'silkwb[br=30k,loss=5]', 'amr[br=4k75,nodtx]', 'amr[br=7k4]', 'g722[br=48k]',
                        'silk[br=20k,loss=10]', 'silkwb[br=30k]', 'amr[br=5k15]',
                        'silkwb[br=20k,loss=5]', 'amrwb[br=23k05,nodtx]', 'amrwb[br=12k65,nodtx]',
                        'silkwb[br=20k,loss=10]', 'amr[br=6k7]', 'silkwb[br=10k]', 'silk[br=10k,loss=5]']
        self.channel_dict = dict(zip(iter(self.channel), range(len(self.channel))))
        self.devices = ['OktavaML19-16000.ir', 'iPhoneirRecording-16000.ir', 'iPadirRecording-16000.ir',
                        'ResloRB250-16000.ir', 'telephonehornT65C-16000.ir', 'ResloSR1-16000.ir', 'RCAPB90-16000.ir',
                        'ResloRBRedLabel-16000.ir', 'telephone90sC-16000.ir', 'SonyC37Fet-16000.ir', 'Doremi-16000.ir',
                        'BehritoneirRecording-16000.ir', ""]
        self.device_dict = dict(zip(iter(self.devices), range(len(self.devices))))

    def __len__(self):
        return len(self.ori_files) + len(self.all_files)

    def __getitem__(self, idx):
        if idx < len(self.ori_files):
            filepath = self.ori_files[idx]
            basename = os.path.basename(filepath)
            all_info = basename.split(".")[0].split("_")
            assert len(all_info) == 6
            channel = "no_channel"
            device = ""
        else:
            filepath = self.all_files[idx - len(self.ori_files)]
            basename = os.path.basename(filepath)
            all_info = basename[:-3].split("_")
            assert len(all_info) == 8
            channel = all_info[6]
            device = all_info[7]
        featureTensor = torch.load(filepath)

        this_feat_len = featureTensor.shape[1]
        if this_feat_len > self.feat_len:
            startp = np.random.randint(this_feat_len - self.feat_len)
            featureTensor = featureTensor[:, startp:startp + self.feat_len, :]
        elif this_feat_len < self.feat_len:
            featureTensor = repeat_padding_Tensor(featureTensor, self.feat_len)

        filename = "_".join(all_info[1:4])
        tag = self.tag[all_info[4]]
        label = self.label[all_info[5]]
        return featureTensor, filename, tag, label, \
               np.array([self.channel_dict[channel], self.device_dict[device]])

    def collate_fn(self, samples):
        return default_collate(samples)

def repeat_padding_Tensor(spec, ref_len):
    mul = int(np.ceil(ref_len / spec.shape[1]))
    spec = spec.repeat(1, mul, 1)[:, :ref_len, :]
    return spec

def repeat_padding_RawTensor(raw, ref_len):
    mul = int(np.ceil(ref_len / raw.shape[1]))
    raw = raw.repeat(1, mul)[:, :ref_len]
    return raw


import os
import numpy as np
import raisimpy as raisim
import time

raisim.World.setLicenseFile(os.path.dirname(os.path.abspath(__file__)) + "/../../rsc/activation.raisim")
anymal_urdf_file = os.path.dirname(os.path.abspath(__file__)) + "/../../rsc/anymal/urdf/anymal.urdf"
laikago_urdf_file = os.path.dirname(os.path.abspath(__file__)) + "/../../rsc/laikago/laikago.urdf"
atlas_urdf_file = os.path.dirname(os.path.abspath(__file__)) + "/../../rsc/atlas/robot.urdf"
monkey_file = os.path.dirname(os.path.abspath(__file__)) + "/../../rsc/monkey/monkey.obj"
dummy_inertia = np.zeros([3, 3])
np.fill_diagonal(dummy_inertia, 0.1)

world = raisim.World()
world.setTimeStep(0.001)
server = raisim.RaisimServer(world)
ground = world.addGround()

anymal = world.addArticulatedSystem(anymal_urdf_file)
anymal.setName("anymal")
anymal_nominal_joint_config = np.array([0, -1.5, 0.54, 1.0, 0.0, 0.0, 0.0, 0.03, 0.4, -0.8,
                                        -0.03, 0.4, -0.8, 0.03, -0.4, 0.8, -0.03, -0.4, 0.8])
anymal.setGeneralizedCoordinate(anymal_nominal_joint_config)
anymal.setPdGains(200*np.ones([18]), np.ones([18]))
anymal.setPdTarget(anymal_nominal_joint_config, np.zeros([18]))

laikago = world.addArticulatedSystem(laikago_urdf_file)
laikago.setName("laikago")
laikago_nominal_joint_config = np.array([0, 1.5, 0.48, 1, 0.0, 0.0, 0.0, 0.0, 0.5, -1, 0, 0.5, -1,
                                         0.00, 0.5, -1, 0, 0.5, -0.7])
laikago.setGeneralizedCoordinate(laikago_nominal_joint_config)
laikago.setPdGains(200*np.ones([18]), np.ones([18]))
laikago.setPdTarget(laikago_nominal_joint_config, np.zeros([18]))

atlas = world.addArticulatedSystem(atlas_urdf_file)
atlas.setName("atlas")
atlas_nominal_joint_config = np.zeros(atlas.getGeneralizedCoordinateDim())
atlas_nominal_joint_config[2] = 1.5
atlas_nominal_joint_config[3] = 1
atlas.setGeneralizedCoordinate(atlas_nominal_joint_config)

server.launchServer(8080)

for i in range(5):
    for j in range(5):
        object_type = (i + j*6) % 5

        if object_type == 0:
            obj = world.addMesh(monkey_file, 5.0, dummy_inertia, np.array([0, 0, 0]), 0.3)
        elif object_type == 1:
            obj = world.addCylinder(0.2, 0.3, 2.0)
        elif object_type == 2:
            obj = world.addCapsule(0.2, 0.3, 2.0)
        elif object_type == 3:
            obj = world.addBox(0.4, 0.4, 0.4, 2.0)
        else:
            obj = world.addSphere(0.3, 2.0)

        obj.setPosition(i-2.5, j-2.5, 5)

time.sleep(2)
world.integrate1()

### get dynamic properties
# mass matrix
mass_matrix = anymal.getMassMatrix()
# non-linear term (gravity+coriolis)
non_linearities = anymal.getNonlinearities([0,0,-9.81])
# Jacobians
jaco_foot_lh_linear = anymal.getDenseFrameJacobian("LF_ADAPTER_TO_FOOT")
jaco_foot_lh_angular = anymal.getDenseFrameRotationalJacobian("LF_ADAPTER_TO_FOOT")

for i in range(500000):
    time.sleep(0.001)
    world.integrate()

server.killServer()


import os
import logging
import datetime
import json
import click
import numpy as np
import joblib

from collections import defaultdict
from rllab.misc import tensor_utils
from rllab.algos.trpo import TRPO
from rllab.baselines.linear_feature_baseline import LinearFeatureBaseline
from rllab.envs.website_env import WebsiteEnv
from rllab.envs.normalized_env import normalize
from rllab.policies.gaussian_mlp_policy import GaussianMLPPolicy

# Flask
from flask import Flask
from flask_cors import CORS
from flask import request
from flask import jsonify

logger = logging.getLogger(__name__)

CONFIG = {
    'num_actions': 1,
    'num_observations': 1,
    'snapshot_dir_itr': 0,
    'snapshot_dir': '',
    'agent_dict': {},
    'batch_size': 20,
    'batch': [],
    'max_bad_updates_before_purge': 60,
    'num_bad_updates': 0,
    'noise': 0.5,
    'picked_correct': [],
    'max_bad_buttons_to_update': 5,
    'file_to_load_policy': ''
}
home = os.path.expanduser('~')
SNAPSHOT_DIR = os.path.join(home, 'rllab/experiments_{}_{}')

def listify_dict(d):
    """
    Recursively convert arrays to lists in a python object
    """

    if type(d) is np.ndarray:
        return d.tolist()

    if isinstance(d, list):
        for idx, k in enumerate(d):
            d[idx] = listify_dict(k)
        return d

    if isinstance(d, dict):
        for k in d.keys():
            d[k] = listify_dict(d[k])

    return d

def arrayify_dict(d):
    """
    Assume dict is made of values that are lists, that need to be
    converted to arrays
    """
    for k in d.keys():
        d[k] = np.array(d[k])

def validate_action_dim(action):
    """
    Validate that the observations are properly formatted
    """
    if not action:
        return False

    if len(action) == CONFIG['num_actions']:
        return True
    return False

def validate_obs_dim(obs):
    """
    Validate that the observations are properly formatted
    """
    if not obs:
        return False

    if len(obs) == CONFIG['num_observations']:
        return True
    return False

def make_path(r, obs, action, agent_info):
    path = dict()
    path["rewards"] = tensor_utils.stack_tensor_list(r)
    path["observations"] = tensor_utils.stack_tensor_list(obs)
    path["actions"] = tensor_utils.stack_tensor_list(action)
    path["env_infos"] = {}
    path["agent_infos"] = tensor_utils.stack_tensor_dict_list(agent_info)
    return path

def update_policy(obs, action, agent_info, reward):
    """
    Update the policy
    """
    if len(np.array(obs).shape) == 1:
        obs = np.array([obs])
    else:
        obs = np.array(obs)

    if len(np.array(action).shape) == 1:
        action = np.array([action])
    else:
        action = np.array(action)

    if isinstance(agent_info, list):
        [arrayify_dict(a) for a in agent_info]
    else:
        arrayify_dict(agent_info)
        agent_info = [agent_info]
    r = np.array(reward)

    path = make_path(r, obs, action, agent_info)
    update_in_batch(path)

def update_in_batch(path):
    """
    update the policy in batch (asynchronously if possible)
    """
    CONFIG['batch'].append(path)

    if len(CONFIG['batch']) > CONFIG['batch_size']:
        algo = CONFIG['agent_dict']['agent']
        good_update = algo.train_from_single_sample(
            CONFIG['batch'], log_dir=CONFIG['snapshot_dir'])

        if not good_update:
            CONFIG['num_bad_updates'] += 1

            if CONFIG['num_bad_updates'] > CONFIG['max_bad_updates_before_purge']:
                print('Did max bad updates, purging batch data')
                CONFIG['batch'] = []
                return

            print('Update step was bad, skipping, but saving data')
            return

        mean_reward = np.mean(CONFIG['picked_correct'])
        print('Mean reward at itr {} is {}'.format(
            CONFIG['snapshot_dir_itr'], mean_reward))
        f_str = '{}\t{}\n'.format(CONFIG['snapshot_dir_itr'], mean_reward)
        with open(CONFIG['reward_file'], 'a') as f:
            f.write(f_str)
        CONFIG['snapshot_dir_itr'] += 1
        if len(CONFIG['picked_correct']) >= 5 * CONFIG['batch_size']:
            CONFIG['picked_correct'] = CONFIG['picked_correct'][-4 * CONFIG['batch_size']:]

        print('Saving data to ', CONFIG['data_file'])
        f_str = ''
        for path in CONFIG['batch']:
            path = listify_dict(path)
            f_str += json.dumps(path) + '\n'

        with open(CONFIG['data_file'], 'a') as f:
            f.write(f_str)

        CONFIG['batch'] = []

def init_agent(num_obs=1, num_actions=1, network_name='', hidden_sizes=(32, 32)):

    msg = 'initing network with num_actions={}, num_obs={}, network_name={}, hidden_sizes={}'
    print(msg.format(num_actions, num_obs, network_name, hidden_sizes))

    # set the snapshot dir stuff
    datetime_str = datetime.datetime.now().strftime("%Y-%m-%d-%H:%M:%S")
    CONFIG['snapshot_dir'] = SNAPSHOT_DIR.format(
        network_name, datetime_str)

    os.makedirs(CONFIG['snapshot_dir'], exist_ok=True)
    print('Making new network on ', datetime_str,
          'with {} actions and {} obs'.format(num_actions, num_obs))

    # create data file
    file = 'data{}_act{}_obs{}.json'.format(
        CONFIG['snapshot_dir_itr'], num_actions, num_obs)
    CONFIG['data_file'] = os.path.join(CONFIG['snapshot_dir'], file)
    file = open(CONFIG['data_file'], 'w')
    file.close()
    print('Made data file at ', CONFIG['data_file'])
    CONFIG['snapshot_dir_itr'] += 1

    # make reward monitoring file
    CONFIG['reward_file'] = os.path.join(CONFIG['snapshot_dir'], 'reward_monitoring.txt')

    CONFIG['num_actions'] = num_actions
    CONFIG['num_observations'] = num_obs

    # Create the env and agent
    env = WebsiteEnv(num_actions=num_actions, action_bounds=[1.] * num_actions,
                     num_observations=num_obs,
                     observation_bound=1.)
    env = normalize(env)

    policy = GaussianMLPPolicy(
        env_spec=env.spec,
        hidden_sizes=hidden_sizes,
        adaptive_std=True
    )

    baseline = LinearFeatureBaseline(env_spec=env.spec)

    algo = TRPO(
        env=env,
        policy=policy,
        baseline=baseline,
        discount=0.99,
        step_size=0.01,
        optimizer_args={"accept_violation": False, "cg_iters": 20},
        # center_adv=False,  # we will scale rewards appropriately
    )
    algo.init_opt()

    CONFIG['agent_dict']['env'] = env
    CONFIG['agent_dict']['agent'] = algo
    CONFIG['agent_dict']['policy'] = algo.policy

    # train from a single samples to load the graph
    obs = env.observation_space.sample()
    action, agent_info = algo.policy.get_action(obs)
    path = make_path(np.array([1.]), [obs], [action], [agent_info])
    algo.train_from_single_sample([path])

def init_agent_from_file():
    filename = CONFIG['file_to_load_policy']
    print('loading model from file', filename)
    params = joblib.load(filename)
    CONFIG['agent_dict']['env'] = params['env']
    CONFIG['agent_dict']['agent'] = params['algo']
    CONFIG['agent_dict']['policy'] = params['policy']
    CONFIG['snapshot_dir_itr'] = params['itr'] + 1
    CONFIG['snapshot_dir'] = os.path.dirname(filename)
    CONFIG['reward_file'] = os.path.join(CONFIG['snapshot_dir'], 'reward_monitoring.txt')
    CONFIG['num_actions'] = params['env'].action_dim
    CONFIG['num_observations'] = params['env'].observation_space.shape[0]

    # create data file
    file = 'data{}_act{}_obs{}.json'.format(
        CONFIG['snapshot_dir_itr'], CONFIG['num_actions'],
        CONFIG['num_observations'])
    CONFIG['data_file'] = os.path.join(CONFIG['snapshot_dir'], file)
    file = open(CONFIG['data_file'], 'w')
    file.close()

def mse(x, y):
    return np.mean((x - y)**2)

def add_noise_to_action_array(x, noise):
    # orig_x = np.copy(x)
    x = np.copy(x)
    noise = abs(noise)
    noise = min(1., noise)
    noise = max(1e-8, noise)

    # get the new mean button to sample from
    # for each dimension, go left or right, whichever is valid
    for idx, xx in enumerate(x):
        if xx + noise > 1:
            x[idx] = xx - noise
        elif xx - noise < -1:
            x[idx] = xx + noise
        else:
            if np.random.rand(1) > .5:
                x[idx] = xx - noise
            else:
                x[idx] = xx + noise

    # give the mean point, generate a random point close to it, with the std being `noise` / 2.
    x_prime = np.random.multivariate_normal(x, np.diag([noise / 2.] * len(x)))
    x_prime = np.clip(x_prime, -1, 1)

    # any -1, or 1 values should be moved back into the space by some random value
    edge_interval = noise * 0.1
    x_prime[x_prime == 1] = np.random.uniform(1 - edge_interval, 1., len(x_prime[x_prime == 1]))
    x_prime[x_prime == -1] = np.random.uniform(-1, -1 + edge_interval, len(x_prime[x_prime == -1]))

    # d = np.mean((x_prime - orig_x)**2)

    return x_prime

"""
FLASK APP
"""

def configure_app(flask_app):
    CORS(flask_app)

    if CONFIG['file_to_load_policy'] == '':
        init_agent()
    else:
        init_agent_from_file()

# Needs to be external for gunicorn
app = Flask(__name__)
configure_app(app)

@app.route('/')
def health():
    return jsonify({'status': 'ok'}), 200

@app.route('/init_networks', methods=['POST'])
def init_networks():
    req = request.get_json()
    num_actions = req.get('num_actions')
    num_obs = req.get('num_observations')
    network_name = req.get('network_name', 'default')
    hidden_sizes = req.get('hidden_sizes', (32, 32))

    if not isinstance(network_name, str):
        msg = 'network_name is not a string: {}'.format(network_name)
        return jsonify({'status': msg}), 400

    if not isinstance(hidden_sizes, list) and\
            not isinstance(hidden_sizes, tuple):
        msg = 'hidden_sizes is not a tuple: {}'.format(hidden_sizes)
        return jsonify({'status': msg}), 400

    init_agent(num_obs, num_actions, network_name, hidden_sizes)
    return jsonify({'status': 'ok'}), 200

@app.route('/get_multi_action', methods=['POST'])
def get_multiple_actions():
    """
    Get multiple actions by feeding each action as the next observation
    """
    req = request.get_json()
    obs = req.get('observation')
    noise = req.get('noise', CONFIG['noise'])
    n_bads = req.get('n_bads', 1)
    n_actions = req.get('n_actions', 1)

    if not validate_obs_dim(obs):
        return jsonify({'status': 'invalid input obs'}), 400

    # get n_actions by feeding in each subsequent action as the next observation
    observations = []
    agent_infos = []
    next_obs = obs.copy()
    good_dict = defaultdict(list)
    for n in range(n_actions):
        observations.append(list(next_obs))
        action, agent_info = CONFIG['agent_dict']['policy'].get_action(next_obs)
        agent_info = listify_dict(agent_info)
        agent_infos.append(agent_info)
        action = np.clip(action, -1, 1)
        next_obs = np.copy(action)
        good_dict['actions'].append(list(action))
    good_dict['agent_info'] = agent_infos
    good_dict['action'] = list(np.concatenate(good_dict['actions']))
    good_dict['observations'] = observations

    # format bad response
    action = good_dict['action']
    bad_button_list = []
    for i in range(n_bads):
        bad_action = add_noise_to_action_array(action, noise)
        bad_dict = dict()
        bad_dict['action'] = list(bad_action)
        bad_dict['agent_info'] = agent_infos
        bad_dict['observations'] = observations
        bad_dict['actions'] = bad_action.reshape(n_actions, len(obs)).tolist()
        bad_button_list.append(bad_dict)

    resp = {
        'good': good_dict,
        'bad': bad_button_list
    }
    return jsonify(resp), 200

@app.route('/get_action', methods=['POST'])
def get_action():
    """
    Get a single action
    """
    req = request.get_json()
    obs = req.get('observation')
    noise = req.get('noise', CONFIG['noise'])
    n_bads = req.get('n_bads', 1)

    if not validate_obs_dim(obs):
        return jsonify({'status': 'invalid input obs'}), 400

    action, agent_info = CONFIG['agent_dict']['policy'].get_action(obs)
    agent_info = listify_dict(agent_info)
    action = np.clip(action, -1, 1)

    # format response
    good_dict = dict()
    good_dict['action'] = list(action)
    good_dict['actions'] = list(action)
    good_dict['agent_info'] = agent_info
    good_dict['observations'] = [obs]

    # format bad response
    bad_button_list = []
    for i in range(n_bads):
        bad_action = add_noise_to_action_array(action, noise)
        bad_dict = dict()
        bad_dict['action'] = list(bad_action)
        bad_dict['actions'] = list(bad_action)
        bad_dict['agent_info'] = agent_info
        bad_dict['observations'] = [obs]
        bad_button_list.append(bad_dict)

    resp = {
        'good': good_dict,
        'bad': bad_button_list
    }
    return jsonify(resp), 200

@app.route('/get_random_obs', methods=['GET'])
def get_random_obs():
    obs = CONFIG['agent_dict']['env'].observation_space.sample()
    return jsonify({'observation': list(obs)}), 200

@app.route('/update_policy_from_game', methods=['POST'])
def update_policy_from_game():
    """
    Update the policy from a sequence of actions taken
    """
    req = request.get_json()
    button_list = req.get('button_list')

    # get the button that was picked
    good_button = list(filter(lambda x: x['isGood'], button_list))[0]
    other_buttons = filter(lambda x: not x['isGood'], button_list)

    picked_correct = 0
    if good_button['picked']:
        picked_correct = 1
    # record whether the user picked the button we predicted given the context
    CONFIG['picked_correct'].append(picked_correct)

    picked_button = None
    unpicked_buttons = []
    for b in other_buttons:
        if b['picked'] is True:
            picked_button = b
        else:
            unpicked_buttons.append(b)

    if good_button['picked']:
        picked_button = good_button

    if picked_button is None:
        return 'no button was picked', 400

    num_steps = len(picked_button['observations'])
    if good_button['picked']:
        # reward the good button since it was picked
        mses = map(
            lambda x: mse(np.array(x['action']), good_button['action']),
            unpicked_buttons)

        r = np.mean(list(mses))
        update_policy(
            good_button['observations'], good_button['actions'],
            good_button['agent_info'], [r] * num_steps)
    elif not good_button['picked']:
        # reward the picked button
        mses = map(
            lambda x: mse(np.array(x['action']), picked_button['action']),
            unpicked_buttons + [good_button])
        r = np.mean(list(mses))
        print('reward for picked button is ', r)
        update_policy(
            picked_button['observations'], picked_button['actions'],
            picked_button['agent_info'], [r] * num_steps)

        # penalize the good button
        mses = mse(np.array(picked_button['action']), good_button['action'])
        mses *= -1
        update_policy(
            good_button['observations'], good_button['actions'],
            good_button['agent_info'], [mses] * num_steps)

    # penalize button that were unpicked and that are maximally far from the button
    # that was picked
    mses = map(
        lambda x: mse(np.array(x['action']), picked_button['action']),
        unpicked_buttons)
    unpicked_buttons = sorted(
        zip(unpicked_buttons, mses),
        key=lambda x: x[1],
        reverse=True)
    mses = list(map(lambda x: x[1], unpicked_buttons))
    unpicked_buttons = list(map(lambda x: x[0], unpicked_buttons))

    num_unpicked = len(unpicked_buttons)
    num_to_update = min(num_unpicked, CONFIG['max_bad_buttons_to_update'])
    for idx, u in enumerate(unpicked_buttons[:num_to_update]):
        update_policy(
            u['observations'], u['actions'],
            u['agent_info'], [-mses[idx]] * num_steps)

    return jsonify({'status': 'ok'}), 200

@click.command()
@click.option('--server-port', default=8080, help='server port')
def start_app(server_port):
    logger.info('Starting on ', server_port)
    app.run(port=server_port)

if __name__ == '__main__':
    start_app()


import numpy as np
from grabscreen import grab_screen
import cv2
import time
from directkeys import PressKey,ReleaseKey, W, A, S, D
# from models import inception_v3 as googlenet
from getkeys import key_check
from collections import deque, Counter
import random
from statistics import mode,mean
import numpy as np
from motion import motion_detection
from xception import get_model
from keras.models import load_model

GAME_WIDTH = 1920
GAME_HEIGHT = 1080

how_far_remove = 800
rs = (20,15)
log_len = 25

motion_req = 800
motion_log = deque(maxlen=log_len)

WIDTH = 480
HEIGHT = 270
LR = 1e-3
EPOCHS = 10

choices = deque([], maxlen=5)
hl_hist = 250
choice_hist = deque([], maxlen=hl_hist)

w = [1,0,0,0,0,0,0,0,0]
s = [0,1,0,0,0,0,0,0,0]
a = [0,0,1,0,0,0,0,0,0]
d = [0,0,0,1,0,0,0,0,0]
wa = [0,0,0,0,1,0,0,0,0]
wd = [0,0,0,0,0,1,0,0,0]
sa = [0,0,0,0,0,0,1,0,0]
sd = [0,0,0,0,0,0,0,1,0]
nk = [0,0,0,0,0,0,0,0,1]

t_time = 0.25

def straight():
    PressKey(W)
    ReleaseKey(A)
    ReleaseKey(D)
    ReleaseKey(S)

def left():
    if random.randrange(0,3) == 1:
        PressKey(W)
    else:
        ReleaseKey(W)
    PressKey(A)
    ReleaseKey(S)
    ReleaseKey(D)
    #ReleaseKey(S)

def right():
    if random.randrange(0,3) == 1:
        PressKey(W)
    else:
        ReleaseKey(W)
    PressKey(D)
    ReleaseKey(A)
    ReleaseKey(S)

def reverse():
    PressKey(S)
    ReleaseKey(A)
    ReleaseKey(W)
    ReleaseKey(D)

def forward_left():
    PressKey(W)
    PressKey(A)
    ReleaseKey(D)
    ReleaseKey(S)

def forward_right():
    PressKey(W)
    PressKey(D)
    ReleaseKey(A)
    ReleaseKey(S)

def reverse_left():
    PressKey(S)
    PressKey(A)
    ReleaseKey(W)
    ReleaseKey(D)

def reverse_right():
    PressKey(S)
    PressKey(D)
    ReleaseKey(W)
    ReleaseKey(A)

def no_keys():

    if random.randrange(0,3) == 1:
        PressKey(W)
    else:
        ReleaseKey(W)
    ReleaseKey(A)
    ReleaseKey(S)
    ReleaseKey(D)

# model = get_model()
MODEL_NAME = 'gta5-{}'.format('xception')
model = load_model(MODEL_NAME)

print('We have loaded a previous model!!!!')

def main():
    last_time = time.time()
    for i in list(range(4))[::-1]:
        print(i+1)
        time.sleep(1)

    paused = False
    mode_choice = 0

    screen = grab_screen(region=(0,40,GAME_WIDTH,GAME_HEIGHT+40))
    screen = cv2.cvtColor(screen, cv2.COLOR_BGR2RGB)
    prev = cv2.resize(screen, (WIDTH,HEIGHT))

    t_minus = prev
    t_now = prev
    t_plus = prev

    while(True):

        if not paused:
            screen = grab_screen(region=(0,0,GAME_WIDTH,GAME_HEIGHT))
            screen = cv2.cvtColor(screen, cv2.COLOR_BGR2RGB)

            last_time = time.time()
            screen = cv2.resize(screen, (WIDTH,HEIGHT))

            delta_count_last = motion_detection(t_minus, t_now, t_plus, screen)

            t_minus = t_now
            t_now = t_plus
            t_plus = screen
            t_plus = cv2.blur(t_plus,(4,4))

            # prediction = model.predict([screen.reshape(HEIGHT,WIDTH,3)])[0]
            print(screen.reshape(1,HEIGHT,WIDTH,3).shape)
            prediction = model.predict(screen.reshape(1,HEIGHT,WIDTH,3))[0]
            prediction = np.array(prediction) * np.array([4.5, 0.1, 0.1, 0.1,  1.8,   1.8, 0.5, 0.5, 0.2])

            mode_choice = np.argmax(prediction)

            if mode_choice == 0:
                straight()
                choice_picked = 'straight'

            elif mode_choice == 1:
                reverse()
                choice_picked = 'reverse'

            elif mode_choice == 2:
                left()
                choice_picked = 'left'
            elif mode_choice == 3:
                right()
                choice_picked = 'right'
            elif mode_choice == 4:
                forward_left()
                choice_picked = 'forward+left'
            elif mode_choice == 5:
                forward_right()
                choice_picked = 'forward+right'
            elif mode_choice == 6:
                reverse_left()
                choice_picked = 'reverse+left'
            elif mode_choice == 7:
                reverse_right()
                choice_picked = 'reverse+right'
            elif mode_choice == 8:
                no_keys()
                choice_picked = 'nokeys'

            motion_log.append(delta_count_last)
            motion_avg = round(mean(motion_log),3)
            print('loop took {} seconds. Motion: {}. Choice: {}'.format( round(time.time()-last_time, 3) , motion_avg, choice_picked))

            if motion_avg < motion_req and len(motion_log) >= log_len:
                print('WERE PROBABLY STUCK FFS, initiating some evasive maneuvers.')

                # 0 = reverse straight, turn left out
                # 1 = reverse straight, turn right out
                # 2 = reverse left, turn right out
                # 3 = reverse right, turn left out

                quick_choice = random.randrange(0,4)

                if quick_choice == 0:
                    reverse()
                    time.sleep(random.uniform(1,2))
                    forward_left()
                    time.sleep(random.uniform(1,2))

                elif quick_choice == 1:
                    reverse()
                    time.sleep(random.uniform(1,2))
                    forward_right()
                    time.sleep(random.uniform(1,2))

                elif quick_choice == 2:
                    reverse_left()
                    time.sleep(random.uniform(1,2))
                    forward_right()
                    time.sleep(random.uniform(1,2))

                elif quick_choice == 3:
                    reverse_right()
                    time.sleep(random.uniform(1,2))
                    forward_left()
                    time.sleep(random.uniform(1,2))

                for i in range(log_len-2):
                    del motion_log[0]

        keys = key_check()

        # p pauses game and can get annoying.
        if 'T' in keys:
            if paused:
                paused = False
                time.sleep(1)
            else:
                paused = True
                ReleaseKey(A)
                ReleaseKey(W)
                ReleaseKey(D)
                time.sleep(1)

main()


import os
import cv2
import sys
import numpy as np

import torch
import torch.nn as nn
import torch.backends.cudnn as cudnn
import torchvision
import argparse
import config as cf
import operator
import csv

from torchvision import datasets, models, transforms
from networks import *
from torch.autograd import Variable
from PIL import Image

parser = argparse.ArgumentParser(description='Pytorch Cell Classification weight upload')
parser.add_argument('--net_type', default='resnet', type=str, help='model')
parser.add_argument('--depth', default=50, type=int, help='depth of model')
args = parser.parse_args()

# Phase 1 : Model Upload
print('\n[Test Phase] : Model Weight Upload')
use_gpu = torch.cuda.is_available()

# upload labels
data_dir = cf.aug_base
trainset_dir = cf.data_base.split("/")[-1]+os.sep
dsets = datasets.ImageFolder(data_dir, None)

H = datasets.ImageFolder(os.path.join(cf.aug_base, 'train'))
dset_classes = H.classes

def softmax(x):
    return np.exp(x) / np.sum(np.exp(x), axis=0)

def getNetwork(args):
    if (args.net_type == 'alexnet'):
        file_name = 'alexnet'
    elif (args.net_type == 'vggnet'):
        file_name = 'vgg-%s' %(args.depth)
    elif (args.net_type == 'resnet'):
        file_name = 'resnet-%s' %(args.depth)
    else:
        print('[Error]: Network should be either [alexnet / vgget / resnet]')
        sys.exit(1)

    return file_name

# uploading the model
print("| Loading checkpoint model for crop inference...")
assert os.path.isdir('../3_classifier/checkpoint'),'[Error]: No checkpoint directory found!'
assert os.path.isdir('../3_classifier/checkpoint/'+trainset_dir),'[Error]: There is no model weight to upload!'
file_name = getNetwork(args)
checkpoint = torch.load('../3_classifier/checkpoint/'+trainset_dir+file_name+'.t7')
model = checkpoint['model']

if use_gpu:
    model.cuda()
    cudnn.benchmark = True

model.eval()

sample_input = Variable(torch.randn(1,3,224,224), volatile=False)
if use_gpu:
    sampe_input = sample_input.cuda()

test_transform = transforms.Compose([
    transforms.Scale(224),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(cf.mean, cf.std)
])

def check_and_mkdir(in_dir):
    if not os.path.exists(in_dir):
        os.makedirs(in_dir)

for file_number in range(args.start, args.finish+1):
    print("| Predicting Box Inference for TEST%d..." %file_number)
    original_img = cv2.imread('/home/bumsoo/Data/test/MICCAI_img/TEST%d.png' %file_number)
    mask_img = cv2.imread('./results/masks/TEST%d.png' %file_number)
    ret, threshed_img = cv2.threshold(cv2.cvtColor(mask_img, cv2.COLOR_BGR2GRAY), 100, 255, cv2.THRESH_BINARY)
    kernel = np.ones((3,3), np.uint8)
    closing = cv2.morphologyEx(threshed_img, cv2.MORPH_CLOSE, kernel, iterations=4)

    _, contours, _ = cv2.findContours(closing, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    count = 0

    with open('results/inferenced/TEST%d/TEST%d.csv' %(file_number, file_number), 'w') as csvfile:
        fieldnames = ['prediction', 'x', 'y', 'w', 'h']
        writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
        for cnt in contours:
            area = cv2.contourArea(cnt)
            #print(area)

            if (area > 30**2):
                x, y, w, h = cv2.boundingRect(cnt)
                crop = original_img[y:y+h, x:x+w]
                crop = cv2.cvtColor(crop, cv2.COLOR_BGR2RGB) # Swap the image into RGB input

                if test_transform is not None:
                    img = test_transform(Image.fromarray(crop, mode='RGB'))

                inputs = img
                inputs = Variable(inputs, volatile=True)

                if use_gpu:
                    inputs = inputs.cuda()
                inputs = inputs.view(1, inputs.size(0), inputs.size(1), inputs.size(2))

                outputs = model(inputs)
                softmax_res = softmax(outputs.data.cpu().numpy()[0])
                index, score = max(enumerate(softmax_res), key=operator.itemgetter(1))

                count += 1
                if ('RBC' in dset_classes[index]):
                    print("\tRBC_%d : %f" %(count, score))
                else:
                    if ("Neutrophil" in dset_classes[index] and score < 0.9):
                        if h1_transform is not None:
                            img = h1_transform(Image.fromarray(crop, mode='RGB'))
                        inputs = img
                        inputs = Variable(inputs, volatile=True)

                        if use_gpu:
                            inputs = inputs.cuda()
                        inputs = inputs.view(1, inputs.size(0), inputs.size(1), inputs.size(2))

                        H1_outputs = model_NH(inputs)
                        hr_softmax = softmax(H1_outputs.data.cpu().numpy()[0])
                        idx, sc = max(enumerate(hr_softmax), key=operator.itemgetter(1))

                        answ = H1_classes[idx]
                    elif "Lymphocyte" in dset_classes[index] and score < 0.9:
                        if h2_transform is not None:
                            img = h2_transform(Image.fromarray(crop, mode='RGB'))
                        inputs = img
                        inputs = Variable(inputs, volatile=True)

                        if use_gpu:
                            inputs = inputs.cuda()
                        inputs = inputs.view(1, inputs.size(0), inputs.size(1), inputs.size(2))

                        H2_outputs = model_NH(inputs)
                        hr_softmax = softmax(H2_outputs.data.cpu().numpy()[0])
                        idx, sc = max(enumerate(hr_softmax), key=operator.itemgetter(1))

                        answ = H2_classes[idx]
                    else:
                        answ = dset_classes[index]

                    crop = cv2.cvtColor(crop, cv2.COLOR_RGB2BGR)
                    # cv2.imwrite("./results/cropped/TEST%d/%s_%d.png" %(file_number, H1_classes[index], count), crop)
                    writer.writerow({
                        'prediction': answ,
                        'x': x,
                        'y': y,
                        'w': w,
                        'h': h
                    })
                    print("\t%s_%d : %f" %(answ, count, score))
                    #print('\t%s_%d : %f' %(dset_classes[index], count, score))


import os
import os.path as osp
import argparse
import numpy as np
import torch
import torch.backends.cudnn as cudnn

# user-defined libs
from dataset.dataset import VideoDataset
from lib.networks import AttentionGenerator, FCEncoder
# The config folder should be in consistent with each other
from configs.config import Config
from util.utils import LossAccumulator
from train import train
from test import test
from util.utils import construct_res_dict
# from adabelief_pytorch import AdaBelief
import time
import json

def parse_args():
    # Set up a parser for command line arguments
    parser = argparse.ArgumentParser("Weakly few-shot learning training script")
    parser.add_argument('-v', '--verbose', action='store_true', default=True, help="increase output verbosity")
    # =============================== Data =============================================================================
    parser.add_argument('--dataset', type=str, default='Thumos14reduced', choices=['Thumos14reduced', 'ActivityNet1.2'], help="select training dataset")
    parser.add_argument('--num_workers', type=int, default=4, help="number of CPU workers used in data loading")
    parser.add_argument('--shuffle', dest='shuffle', action='store_true', help="Shuffle data")
    parser.set_defaults(shuffle=True)
    # =============================== Training =========================================================================
    # c-way k-shot n-query [[5, 1, 19], [5, 5, 15], [20, 1, 10], [20, 5, 5]]
    parser.add_argument("--pick_class_num", type=int, default=5)
    parser.add_argument("--sample_num_per_class", type=int, default=1)
    parser.add_argument("--batch_num_per_class", type=int, default=8)
    parser.add_argument("--num_dim", type=int, default=128)
    parser.add_argument("--num_episodes", type=int, default=5001)

    parser.add_argument('--optim', type=str, default='adam', choices=['sgd', 'adam', 'adabelief'])
    parser.add_argument('--weight_decay', type=float, default=5e-4, help="Weight decay")
    parser.add_argument('--momentum', type=float, default=0.9, help="momentum value for optimisation")
    parser.add_argument('--ckp_freq', default=2500, type=int, help='set number iterations per checkpoint model saving')
    parser.add_argument('--paral', dest='paral', action='store_true',
                        help="use data parallelism (multiple GPUs) during training")
    parser.set_defaults(paral=True)
    parser.add_argument('--no-paral', dest='paral', action='store_false',
                        help="do NOT use data parallelism (multiple GPUs) during training -- use first GPU available")
    parser.add_argument('--cuda', dest='cuda', action='store_true', help="use CUDA during training")
    parser.add_argument('--no-cuda', dest='cuda', action='store_false', help="do NOT use CUDA during training")
    parser.set_defaults(cuda=True)
    parser.add_argument('--distance', type=str, choices=['cosine', 'l2', 'eu'], default='eu')
    parser.add_argument('--tsm', type=str, choices=['cosine', 'eu', 'ip'], default='cosine', help="The way to make  comparison in calculating segment-to-segment similarity")
    parser.add_argument('--lambda_cent', default=0, type=float, help='lambda for center loss')
    parser.add_argument('--lambda_ent', default=0, type=float, help='lambda for entropy loss')
    parser.add_argument('--lambda_contr', default=0, type=float, help='lambda for contrastive loss')
    parser.add_argument('--norm', dest='norm', action='store_true')
    parser.add_argument('--no-norm', dest='norm', action='store_false')
    parser.set_defaults(norm=False)
    parser.add_argument('--encoder', dest='encoder', action='store_true', help='Using video encoder')
    parser.add_argument('--no-encoder', dest='encoder', action='store_false', help='Using video encoder')
    parser.set_defaults(encoder=True)
    parser.add_argument('--bn', default=True, action='store_true', help='Decide if batch normalization is needed')
    parser.add_argument('--split', default='cvpr18', choices=[None, 'cvpr18', 'reverse'], help='This is used to split training and testing class')
    parser.add_argument('--num_in', choices=[1, 2, 4, 6], type=int, help='Number of attention generator input')
    # ================================ Testing =========================================================================
    parser.add_argument('--mode', type=str, dest='mode', default='training', choices=['training', 'testing'])
    parser.add_argument('--test_ep', default=1000, type=int)
    parser.add_argument('--test_batch_size', default=1, type=int)
    parser.add_argument('--load', type=str, help="load pre-trained model")
    parser.add_argument('--models_root', default="models/", help="root directory for saving trained models")
    parser.add_argument('--postfix', default='test', help="postfix to discriminate different iterations")
    parser.add_argument('--eval_one', default=True, action='store_true', help='If evaluate on one video or many')
    args = parser.parse_args()

    return args

def save_result(checkpoint_model_filename, epis, loss_accumulator, fcencoder, args, cfg, attgen=[], acc=[], map=[]):

    checkpoint_results_dict = construct_res_dict(iteration=epis,
                                                 loss_cls=loss_accumulator.get_loss_cls())

    attgen_state_dict = attgen.module.state_dict() if args.paral else attgen.state_dict() if args.num_in > 1 else []
    fcencoder_state_dict = fcencoder.module.state_dict() if args.paral else fcencoder.state_dict() if args.encoder else []

    torch.save({'state_dict': fcencoder_state_dict,
                'state_dict_filter': attgen_state_dict,
                'args': args,
                'cfg': cfg,
                'acc': np.array(acc),
                'map': np.array(map),
                'results': checkpoint_results_dict}, checkpoint_model_filename)

def main():
    """training/testing script.
    Options:
        --dataset       : set training dataset ('ActivityNet1.2', 'thumos14')
        --models_root   : set root directory of models directories
        --batch_size    : set training batch size
        --num_workers   : set number of CPU workers during data loading
        --cuda          : use CUDA during training
        --no-cuda       : do NOT use CUDA during training
    """

    ################################################################################
    ##                                [ Load Data ]                               ##
    ################################################################################
    args = parse_args()
    activityNet = False if args.dataset == 'Thumos14reduced' else True

    cfg = Config(args.dataset)
    lr_schedule = cfg.lr_schedule

    if args.verbose:
        print("#.Load dataset: {}...".format(args.dataset))

    # Whole dataset, needs to sample in each episode
    videodb = VideoDataset(args.dataset, cfg, split=args.split)
    num_class = len(videodb.classlist)
    train_class = videodb.train_cls
    num_class_train = len(train_class)
    test_class = videodb.test_cls
    num_class_test = len(test_class)

    # Create output models directory under args.models_root
    # Models dir format: <models_root>/
    model_basename = "{}_{}way_{}shot_split_{}_enc_{}_tsm_{}_lcent_{}_indim_{}_{}"\
        .format(args.dataset,
                args.pick_class_num,
                args.sample_num_per_class,
                args.split,
                args.encoder,
                args.tsm,
                args.lambda_cent,
                args.num_in,
                args.postfix)

    models_dir = osp.join(args.models_root, model_basename)
    if not osp.exists(models_dir):
        os.makedirs(models_dir, exist_ok=True)

    ################################################################################
    ##                             [ Build Network ]                              ##
    ################################################################################
    # Set default tensor type
    if torch.cuda.is_available():
        if args.cuda:
            torch.set_default_tensor_type('torch.cuda.FloatTensor')
        if not args.cuda:
            print("*** WARNING ***: It looks like you have a CUDA device, but aren't using CUDA.\n"
                  "                 Run with --cuda for optimal training speed.")
            torch.set_default_tensor_type('torch.FloatTensor')
    else:
        torch.set_default_tensor_type('torch.FloatTensor')

    # Build network
    if args.encoder:
        fcencoder = FCEncoder(n_feature=2048, out_dim=args.num_dim)
        # Print parameter numbers
        if args.verbose:
            model_params = filter(lambda p: p.requires_grad, fcencoder.parameters())
            num_params = sum([np.prod(p.size()) for p in model_params])
            print("  \\__.Number of parameters encoder: {}m".format(num_params / pow(10, 6)))
    else:
        fcencoder = []
        args.num_dim = 1024
    if args.num_in>1:
        attgen = AttentionGenerator(args.num_in, args.bn)
        if args.verbose:
            model_params_filter = filter(lambda p: p.requires_grad, attgen.parameters())
            num_params_filter = sum([np.prod(p.size()) for p in model_params_filter])
            print("  \\__.Number of parameters filter : {}".format(num_params_filter))
    else:
        attgen = []

    # Initialize network with pre-trained model
    if args.load:
        if args.verbose:
            print("#. Load network with pre-trained model: {}".format(args.load))
        model = torch.load(args.load, map_location=lambda storage, loc: storage)

        model_args = model['args']
        model_cfg = model['cfg']
        if args.encoder:
            model_state_dict = model['state_dict']
            fcencoder.load_state_dict(model_state_dict)
            fcencoder.eval()
        # import pdb
        # pdb.set_trace()
        if args.num_in>1:
            model_state_dict_filter = model['state_dict_filter']
            attgen.load_state_dict(model_state_dict_filter)
            attgen.eval()

        if args.mode=='testing':
            args = model_args
            args.mode = 'testing'
        else:
            args = model_args
        cfg = model_cfg

    # Print options
    if args.verbose:
        print("#.xxx")
        print("  \\__.Dataset")
        print("       \\__Name                            : {}".format(args.dataset))
        print("       \\__CPU workers                     : {}".format(args.num_workers))
        print("       \\__Class Num of all, train, test   : {}/{}/{}".format(num_class,
                                                                             num_class_train,
                                                                             num_class_test))
        print("      \\__training class is                : {}".format(train_class))
        print("      \\__testing class is                 : {}".format(test_class))
        print("  \\__.Training parameters for each training episode")
        print("       \\__Pick class Num                  : {}".format(args.pick_class_num))
        print("       \\__Sample Num                      : {}".format(args.sample_num_per_class))
        print("       \\__Batch  Num                      : {}".format(args.batch_num_per_class))
        print("       \\__Number of episodes              : {}".format(args.num_episodes))
        print("       \\__LR schedule                     : {}".format(lr_schedule))
        print("       \\__Weight decay                    : {}".format(args.weight_decay))
        print("       \\__Use CUDA                        : {}".format(args.cuda))
        print("  \\__.Output")
        print("       \\__Models directory                : {}{}".format(args.models_root, model_basename))
        print("  \\__.Options")
        print("       \\__Use Encoder                     : {}".format(args.encoder))
        print("       \\__Use entropy loss                : {}".format(args.lambda_ent))
        print("       \\__Use Attention generator         : {}".format(args.num_in))
        print("       \\__TSM                             : {}".format(args.tsm))
        print("       \\__Distance metric                 : {}".format(args.distance))

    # Parallelize data over multiple GPUs in the batch dimension (if available) and enable benchmark mode in cudnn.
    if args.encoder:
        if args.cuda and torch.cuda.is_available():
            cudnn.benchmark = True
            if args.paral:
                fcencoder = torch.nn.DataParallel(fcencoder).cuda()
            else:
                fcencoder = fcencoder.cuda()

    if args.num_in > 1:
        if args.cuda and torch.cuda.is_available():
            cudnn.benchmark = True
            attgen = torch.nn.DataParallel(attgen).cuda()
        else:
            attgen = attgen.cuda()

    if args.mode == 'training':
        if args.encoder:
            optimizer = torch.optim.Adam(fcencoder.parameters(), lr=lr_schedule[1], weight_decay=args.weight_decay)
        else:
            optimizer = []
        if args.num_in > 1:
            optimizer_filter = torch.optim.Adam(attgen.parameters(), lr=lr_schedule[1], weight_decay=args.weight_decay)
        else:
            optimizer_filter = []

        loss_accumulator = LossAccumulator()
        best_acc = 0

        if args.verbose:
            print("#.Training on {}...".format(args.dataset))

        for epis in range(args.num_episodes):
            if args.encoder:
                fcencoder.train()
            if args.num_in > 1:
                attgen.train()
            train(fcencoder,
                  lr_schedule,
                  optimizer,
                  videodb,
                  epis, args,
                  loss_accumulator,
                  attgen,
                  optimizer_filter,
                  train_class)
            if epis % args.ckp_freq == 0 and epis > 0:
                checkpoint_model_filename = osp.join(models_dir, "cp.pth")
                save_result(checkpoint_model_filename, epis, loss_accumulator, fcencoder, args, cfg, attgen)
                if args.encoder:
                    fcencoder.eval()
                if args.num_in > 1:
                    attgen.eval()
                ap_100, acc_100 = test(args,
                                       videodb,
                                       fcencoder,
                                       attgen,
                                       models_dir,
                                       test_class,
                                       activityNet)
                if best_acc < np.mean(ap_100):
                    best_acc = np.mean(ap_100)
                    checkpoint_model_filename = osp.join(models_dir, "best.pth")
                    save_result(checkpoint_model_filename,
                                epis,
                                loss_accumulator,
                                fcencoder,
                                args,
                                cfg,
                                attgen,
                                acc_100,
                                ap_100)

    else:
        print('I am doing testing')
        ap_100, acc_100 = test(args,
                               videodb,
                               fcencoder,
                               attgen,
                               models_dir,
                               test_class,
                               activityNet)

if __name__ == '__main__':
    main()