import logging
from typing import Dict, Tuple

import pandas as pd
import numpy as np
from numpy.random import default_rng

from . import utils
from .. import const
from .helper import DatasetHelper

logger = logging.getLogger(__name__)

class Dataset(DatasetHelper):
    def __init__(self,
                 size: tuple,
                 name: str = "",
                 d_type: str = const.DATASET,
                 excludes: list = None,
                 random_seed: int = None,
                 rebuild: bool = False,
                 save: bool = False):
        """Base class to interact with the original dataset to build custom datasets.

        Args:
            size (tuple): size of the output images in the dataset.
            name (str): Name of dataset used to build paths for
            storing generated files. Defaults to ''.
            d_type (str, optional): [description]. Defaults to 'dataset'.
            excludes (list, optional): [description]. Defaults to None.
            random_seed (int, optional): [description]. Defaults to None.
            rebuild (bool, optional): [description]. Defaults to False.
            save (bool, optional): Whether to save the generated files for
            faster builds in the future. Defaults to True.

        Raises:
            ValueError: Will fail to build on non-supported dataset type.
        """
        assert len(size) == 2
        self.size = size
        rows = self.size[0]
        cols = self.size[1]

        self.name = name
        self.d_type = d_type
        self.rng = default_rng(seed=random_seed)
        self.excludes = excludes

        # raw labels.xlsx
        self.ant_data = pd.read_excel(
            f'./dataset/labels.xlsx', header=0)

        self.img_meta_path = f'./dataset/{self.name}_{self.d_type}_{rows}_{cols}_img_meta.npy'
        self.images_path = f'./dataset/{self.name}_{self.d_type}_{rows}_{cols}_images.npy'
        self.labels_path = f'./dataset/{self.name}_{self.d_type}_{rows}_{cols}_labels.npy'
        self.ids_path = f'./dataset/{self.name}_{self.d_type}_{rows}_{cols}_ids.npy'

        # load img ids and labels
        self._build_labels(rebuild, save)
        logger.info('Loaded image metadata.')

        self._build_images(rebuild, save)
        logger.info(f'Loaded sub-image dataset of size {rows} by {cols}.')
        logger.info(
            f'Total of {len(self.labels)} sub-images built from {len(np.unique(self.ids))} images.')

        logger.info(f'Class data:')
        for k, v in self.class_data().items():
            logger.info(f'{k}: {v}')

    def class_data(self) -> Dict:
        """Return the class data for the dataset.
        Intended to be overridden by the dataset type.

        Returns:
            Dict: Class data of the dataset.
        """
        uniques = np.unique(self.labels, return_counts=True)
        res = {}
        for i in range(len(uniques[0])):
            res[uniques[0][i]] = uniques[1][i]

        return res

    def _build_labels(self, rebuild: bool, save: bool):
        """Helper for Dataset.build_labels, it will use the rebuild variable to
        force a rebuild to new files. It will also rebuild on failure-to-load
        the expected files.

        Args:
            rebuild (bool): Force rebuilding the dataset.
            save (bool): Save the generated files.

        Updates:
            img_meta: img metadata used for getting sample info (label, etc.).
        """
        if rebuild:
            self.build_labels(save)

        try:
            self.img_meta = np.load(self.img_meta_path)
        except Exception as e:
            logger.debug(e)
            # build on failure-to-load
            self.img_meta = self.build_labels(save)

    def build_labels(self, save: bool) -> np.ndarray:
        """Builds the img_meta.npy file based on data from the labels.xlsx file.

        Args:
            save (bool): Save the generated files.
            dtype (str): If const.DATASET_RS, then convert the labels to 0 or 1.

        Returns:
            np.ndarray: Image metadata scraped from original dataset.
        """
        logger.info(f'Generating labels.')

        # convert human expert labels using majority voting
        label_cols = self.ant_data[['Jp', 'Becca', 'Katy']]
        if self.d_type == const.DATASET_RS:
            label = utils.convert_labels_rs(label_cols.mode(axis=1)[0])
        elif self.d_type == const.DATASET_ALL:
            label = utils.convert_labels(label_cols.mode(axis=1)[0])

        label['class'] = label['class'] + 1

        # clip between the number of classes
        m = int(label['class'].max())
        idx = label['class'].isin(range(1, m+1))
        label = label.loc[idx]

        id_col = self.ant_data['Photo_number']
        id_col = id_col.to_frame('id')
        id_col = id_col.loc[idx.values]

        # store as arrays
        img_labels = label['class'].values
        img_ids = id_col['id'].values

        assert len(img_labels) == len(img_ids)

        # save image ids and labels
        res = np.stack([img_ids.astype(int), img_labels.astype(int)])

        if save:
            np.save(self.img_meta_path, res)

        return res

    def _build_images(self, rebuild: bool, save: bool):
        """Helper for Dataset.build_images, it will use the rebuild variable to
        force a rebuild to new files. It will also rebuild on failure-to-load
        the expected files.

        Args:
            rebuild (bool): Force rebuilding the dataset.
            save (bool): Save the generated files.

        Updates:
            subimages: Array of generated images.
            labels: Ground-truth label of each image.
            ids: Original sample id of each image.
        """
        if rebuild:
            self.images, self.labels, self.ids = self.build_images(save)
            return

        try:
            self.images = np.load(self.images_path)
            self.labels = np.load(self.labels_path)
            self.ids = np.load(self.ids_path)
        except Exception as e:
            logger.debug(e)

            # build on failure-to-load
            self.images, self.labels, self.ids = self.build_images(save)

    def build_images(self, save: bool):
        'Must be implemented by the dataset type.'
        raise NotImplementedError

    def get_label(self, _id: int) -> int:
        """Given an image ID, return the label.

        Args:
            _id (int): ID of the original sample

        Raises:
            ValueError: NA label for iamge

        Returns:
            int: Original sample label
        """
        idx = np.where(self.img_meta[0] == _id)
        label = self.img_meta[1][idx]
        try:
            return int(label)
        except Exception as e:
            logger.debug(e)
            try:
                return int(label[0])
            except Exception as e:
                logger.debug(e)
                raise ValueError(
                    f'Class of ID[{_id}] is NA or not considered in this version.')

    def is_included(self, _id: int) -> bool:
        """Given an image id, return if the image is included from the dataset.

        Args:
            _id (int): ID of the sample.

        Returns:
            bool: True if included, else false
        """
        if _id in self.ids:
            return True
        return False

    def stratified_split(self, n: int) -> Tuple[np.ndarray, np.ndarray]:
        """Stratified split with $n samples per class.

        Args:
            n (int): Number of samples per class

        Returns:
            train_x: Training images
            train_y: Training labels
            test_x: Test images
            test_y: Training labels
        """
        labels = self.labels-1
        uniques = np.unique(labels, return_counts=True)

        # stratified sample n for each class in self.labels
        train_idxs = np.zeros((len(uniques[0])*n), dtype=np.int)
        for _class in range(len(uniques[0])):
            idx = np.where(labels == uniques[0][_class])[0]
            samples = self.rng.choice(idx, size=n, replace=False)
            train_idxs[_class*n:(_class+1)*n] = samples

        test_idxs = np.array([idx for idx in range(
            len(labels)) if idx not in train_idxs])

        self.train_x = self.images[(train_idxs)]
        self.train_y = labels[(train_idxs)]
        self.test_x = self.images[(test_idxs)]
        self.test_y = labels[(test_idxs)]

        assert len(self.train_y) + len(self.test_y) == len(self.labels)

        return self.train_x, self.train_y

    def build_validation_set(self, split: float = 0.5) -> Tuple[np.ndarray, np.ndarray]:
        try:
            # see if test_y exists
            self.test_y
        except Exception as e:
            msg = "Test samples have yet to be made."
            logger.error(msg)
            raise e

        val_idxs = self.rng.choice(
            np.arange(len(self.test_y)),
            size=int(split*len(self.test_y)),
            replace=False
        )

        test_idxs = np.array([idx for idx in range(
            len(self.test_y)) if idx not in val_idxs])

        self.val_x = self.test_x[(val_idxs)]
        self.val_y = self.test_y[(val_idxs)]
        self.test_x = self.test_x[(test_idxs)]
        self.test_y = self.test_y[(test_idxs)]

        return self.val_x, self.val_y

    def train_set(self) -> Tuple[np.ndarray, np.ndarray]:
        return self.train_x, self.train_y

    def test_set(self) -> Tuple[np.ndarray, np.ndarray]:
        return self.test_x, self.test_y

    def val_set(self) -> Tuple[np.ndarray, np.ndarray]:
        return self.val_x, self.val_y


import cv2
import numpy as np
import tensorflow.keras
from PIL import Image
from numpy import asarray
import random as rd

class DataGenerator(tensorflow.keras.utils.Sequence):
    'Generates data for Keras'
    def __init__(self, list_IDs, labels,samples=1,frames=100, batch_size=32, dim=(32,32,32), n_channels=3,
                 n_classes=8, shuffle=False):
        'Initialization'
        self.dim = dim
        self.samples = samples
        self.frames = frames
        self.batch_size = batch_size
        self.labels = labels
        self.list_IDs = list_IDs
        self.n_channels = n_channels
        self.n_classes = n_classes
        self.shuffle = shuffle
        self.on_epoch_end()
        self.root = 'D:/data1/'
        #self.root = '/data1/shengjun/handWashDB/data1/'

    def __len__(self):
        'Denotes the number of batches per epoch'
        return int(np.floor(len(self.list_IDs))*14)
        #return int(np.floor(len(self.list_IDs) / self.batch_size))

    def __getitem__(self, index):
        'Generate one batch of data'
        # Generate indexes of the batch
        indexes = self.indexes[index*self.batch_size:(index+1)*self.batch_size]

        # Find list of IDs
        k = int(index/14)
        samples= self.list_IDs[k]
        IDs = index%14 *4

        # Accessing a text file - www.101computing.net/mp3-playlist/
        fileNameTrain = self.root + samples + '.dat'
        file = open(fileNameTrain, "r")

        # Repeat for each song in the text file
        fields=[]
        for line in file:
            # Let's split the line into an array called "fields" using the ";" as a separator:
            field = line.split(",")
            for i in range(0,28,4):
                fields.append('BKA')
                fields.append('0' if i==0 else field[i-1])
                fields.append('BKB')
                fields.append(field[i+1])
                fields.append(field[i])
                fields.append(field[i+1])
                fields.append(field[i+2])
                fields.append(field[i+3])

        # It is good practice to close the file at the end to free up resources
        file.close()

        #list_IDs_temp = [self.list_IDs[k] for k in indexes]
        #label =
        # Generate data
        X, y = self.__data_generation(samples, fields, IDs)

        return X, y

    def on_epoch_end(self):
        'Updates indexes after each epoch'
        self.indexes = np.arange(len(self.list_IDs))
        if self.shuffle == True:
            np.random.shuffle(self.indexes)

    def __data_generation(self, samples, fields, IDs):
        'Generates data containing batch_size samples' # X : (n_samples, *dim, n_channels)
        # Initialization

        # Generate data
        NameA = fields[IDs]
        va = np.long(fields[IDs + 1])
        NameB = fields[IDs + 2]
        vb = np.long(fields[IDs + 3])
        Name = NameA[0:2]
        x = vb - va
        dis = 100
        win = x if x < dis else dis

        diff = x - dis

        start = 0 if diff <= 0 else rd.randrange(0, diff, 30)
        print(' diff, start, name: %d,%d, %s \n' % (diff, start, Name))

        X = np.empty((self.samples, vb-va, *self.dim, self.n_channels))
        y = np.empty((self.samples), dtype=int)
        for i in range(0,win):
            # load the image
            fileName = self.root + samples + '/_Color_'+str(va+i+start)+'.png'
            image = cv2.imread(fileName)
            # convert image to numpy array
            # Store sample
            gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
            gray_image2 = cv2.resize(gray_image,(100,100))
            data = np.expand_dims(gray_image2, axis=-1)
            X[0,i,] = data

        # Store class
        y[0] = self.labels[Name]

        return X, tensorflow.keras.utils.to_categorical(y, num_classes=self.n_classes)


import numpy as np

import gbkfit.math
import gbkfit.driver.native.libgbkfit_host as native_module
from gbkfit.driver import backend

__all__ = ['Backend']

def _ptr(x):
    return x.__array_interface__['data'][0] if x is not None else 0

def _shape(x):
    return x.__array_interface__['shape'] if x is not None else (0,)

def _size(x):
    return gbkfit.math.prod(_shape(x))

def _get_class(classes, dtype):
    if dtype not in classes:
        requested = np.dtype(dtype).name
        supported = ', '.join([np.dtype(dt).name for dt in classes])
        raise RuntimeError(
            f"the requested dtype is not supported "
            f"(requested: {requested}; supported: {supported})")
    return classes[dtype]

class Backend(backend.Backend):

    def make_dmodel_dcube(self, dtype):
        return _DModelDCube(dtype)

    def make_dmodel_mmaps(self, dtype):
        return _DModelMMaps(dtype)

    def make_gmodel(self, dtype):
        return _GModel(dtype)

    def make_gmodel_mcdisk(self, dtype):
        return _GModelMCDisk(dtype)

    def make_gmodel_smdisk(self, dtype):
        return _GModelSMDisk(dtype)

    def make_objective(self, dtype):
        return _Objective(dtype)

class _DModelDCube(backend.DModelDCube):

    _CLASSES = {
        np.float32: native_module.DModelDCubef32}

    def __deepcopy__(self, memodict):
        return self.__class__(self._dtype)

    def __init__(self, dtype):
        self._backend = _get_class(self._CLASSES, dtype)()
        self._dtype = dtype

    def convolve(
            self,
            size_hi,
            scube_hi, scube_hi_fft,
            wcube_hi, wcube_hi_fft,
            psf3d_hi, psf3d_hi_fft):
        self._backend.convolve(
            size_hi,
            _ptr(scube_hi), _ptr(scube_hi_fft),
            _ptr(wcube_hi), _ptr(wcube_hi_fft),
            _ptr(psf3d_hi), _ptr(psf3d_hi_fft))

    def downscale(
            self,
            scale,
            edge_hi,
            size_hi, size_lo,
            scube_hi, scube_lo):
        self._backend.downscale(
            scale,
            edge_hi,
            size_hi, size_lo,
            _ptr(scube_hi), _ptr(scube_lo))

    def make_mask(self, mask_spat, mask_spec, mask_coef, size, cube, mask):
        self._backend.make_mask(
            mask_spat, mask_spec, mask_coef, size, _ptr(cube), _ptr(mask))

class _DModelMMaps(backend.DModelMMaps):

    _CLASSES = {
        np.float32: native_module.DModelMMapsf32}

    def __deepcopy__(self, memodict):
        return self.__class__(self._dtype)

    def __init__(self, dtype):
        self._backend = _get_class(self._CLASSES, dtype)()
        self._dtype = dtype

    def moments(self, size, step, zero, scube, mmaps, masks, orders):
        self._backend.moments(
            size, step, zero,
            _ptr(scube),
            _ptr(mmaps),
            _ptr(masks),
            _ptr(orders),
            _size(orders))

class _GModel(backend.GModel):

    _CLASSES = {
        np.float32: native_module.GModelf32}

    def __deepcopy__(self, memodict):
        return self.__class__(self._dtype)

    def __init__(self, dtype):
        self._backend = _get_class(self._CLASSES, dtype)()
        self._dtype = dtype

    def make_wcube(self, spat_size, spec_size, spat_data, spec_data):
        self._backend.make_wcube(
            spat_size[0], spat_size[1], spat_size[2],
            spec_size,
            _ptr(spat_data),
            _ptr(spec_data))

class _GModelMCDisk(backend.GModelMCDisk):

    _CLASSES = {
        np.float32: native_module.GModelMCDiskf32}

    def __deepcopy__(self, memodict):
        return self.__class__(self._dtype)

    def __init__(self, dtype):
        self._backend = _get_class(self._CLASSES, dtype)()
        self._dtype = dtype

    def evaluate(
            self,
            cflux, nclouds, ncloudspt, hasordint,
            loose, tilted, rnodes,
            vsys, xpos, ypos, posa, incl,
            rpt_uids, rpt_cvalues, rpt_ccounts, rpt_pvalues, rpt_pcounts,
            rht_uids, rht_cvalues, rht_ccounts, rht_pvalues, rht_pcounts,
            vpt_uids, vpt_cvalues, vpt_ccounts, vpt_pvalues, vpt_pcounts,
            vht_uids, vht_cvalues, vht_ccounts, vht_pvalues, vht_pcounts,
            dpt_uids, dpt_cvalues, dpt_ccounts, dpt_pvalues, dpt_pcounts,
            dht_uids, dht_cvalues, dht_ccounts, dht_pvalues, dht_pcounts,
            zpt_uids, zpt_cvalues, zpt_ccounts, zpt_pvalues, zpt_pcounts,
            spt_uids, spt_cvalues, spt_ccounts, spt_pvalues, spt_pcounts,
            wpt_uids, wpt_cvalues, wpt_ccounts, wpt_pvalues, wpt_pcounts,
            spat_size, spat_step, spat_zero, spat_rota,
            spec_size, spec_step, spec_zero,
            image, scube, rcube, wcube,
            rdata, vdata, ddata):
        self._backend.evaluate(
            cflux, nclouds, _ptr(ncloudspt), _size(ncloudspt), _ptr(hasordint),
            loose, tilted,
            _size(rnodes),
            _ptr(rnodes),
            _ptr(vsys), _ptr(xpos), _ptr(ypos), _ptr(posa), _ptr(incl),
            _size(rpt_uids),
            _ptr(rpt_uids),
            _ptr(rpt_cvalues), _ptr(rpt_ccounts),
            _ptr(rpt_pvalues), _ptr(rpt_pcounts),
            _ptr(rht_uids),
            _ptr(rht_cvalues), _ptr(rht_ccounts),
            _ptr(rht_pvalues), _ptr(rht_pcounts),
            _size(vpt_uids),
            _ptr(vpt_uids),
            _ptr(vpt_cvalues), _ptr(vpt_ccounts),
            _ptr(vpt_pvalues), _ptr(vpt_pcounts),
            _ptr(vht_uids),
            _ptr(vht_cvalues), _ptr(vht_ccounts),
            _ptr(vht_pvalues), _ptr(vht_pcounts),
            _size(dpt_uids),
            _ptr(dpt_uids),
            _ptr(dpt_cvalues), _ptr(dpt_ccounts),
            _ptr(dpt_pvalues), _ptr(dpt_pcounts),
            _ptr(dht_uids),
            _ptr(dht_cvalues), _ptr(dht_ccounts),
            _ptr(dht_pvalues), _ptr(dht_pcounts),
            _size(zpt_uids),
            _ptr(zpt_uids),
            _ptr(zpt_cvalues), _ptr(zpt_ccounts),
            _ptr(zpt_pvalues), _ptr(zpt_pcounts),
            _size(spt_uids),
            _ptr(spt_uids),
            _ptr(spt_cvalues), _ptr(spt_ccounts),
            _ptr(spt_pvalues), _ptr(spt_pcounts),
            _size(wpt_uids),
            _ptr(wpt_uids),
            _ptr(wpt_cvalues), _ptr(wpt_ccounts),
            _ptr(wpt_pvalues), _ptr(wpt_pcounts),
            spat_size[0], spat_size[1], spat_size[2],
            spat_step[0], spat_step[1], spat_step[2],
            spat_zero[0], spat_zero[1], spat_zero[2],
            spec_size,
            spec_step,
            spec_zero,
            _ptr(image), _ptr(scube), _ptr(rcube), _ptr(wcube),
            _ptr(rdata), _ptr(vdata), _ptr(ddata))

class _GModelSMDisk(backend.GModelSMDisk):

    _CLASSES = {
        np.float32: native_module.GModelSMDiskf32}

    def __deepcopy__(self, memodict):
        return self.__class__(self._dtype)

    def __init__(self, dtype):
        self._backend = _get_class(self._CLASSES, dtype)()
        self._dtype = dtype

    def evaluate(
            self,
            loose, tilted, rnodes,
            vsys, xpos, ypos, posa, incl,
            rpt_uids, rpt_cvalues, rpt_ccounts, rpt_pvalues, rpt_pcounts,
            rht_uids, rht_cvalues, rht_ccounts, rht_pvalues, rht_pcounts,
            vpt_uids, vpt_cvalues, vpt_ccounts, vpt_pvalues, vpt_pcounts,
            vht_uids, vht_cvalues, vht_ccounts, vht_pvalues, vht_pcounts,
            dpt_uids, dpt_cvalues, dpt_ccounts, dpt_pvalues, dpt_pcounts,
            dht_uids, dht_cvalues, dht_ccounts, dht_pvalues, dht_pcounts,
            zpt_uids, zpt_cvalues, zpt_ccounts, zpt_pvalues, zpt_pcounts,
            spt_uids, spt_cvalues, spt_ccounts, spt_pvalues, spt_pcounts,
            wpt_uids, wpt_cvalues, wpt_ccounts, wpt_pvalues, wpt_pcounts,
            spat_size, spat_step, spat_zero, spat_rota,
            spec_size, spec_step, spec_zero,
            image, scube, rcube, wcube,
            rdata, vdata, ddata):
        self._backend.evaluate(
            loose, tilted,
            _size(rnodes),
            _ptr(rnodes),
            _ptr(vsys), _ptr(xpos), _ptr(ypos), _ptr(posa), _ptr(incl),
            _size(rpt_uids),
            _ptr(rpt_uids),
            _ptr(rpt_cvalues), _ptr(rpt_ccounts),
            _ptr(rpt_pvalues), _ptr(rpt_pcounts),
            _ptr(rht_uids),
            _ptr(rht_cvalues), _ptr(rht_ccounts),
            _ptr(rht_pvalues), _ptr(rht_pcounts),
            _size(vpt_uids),
            _ptr(vpt_uids),
            _ptr(vpt_cvalues), _ptr(vpt_ccounts),
            _ptr(vpt_pvalues), _ptr(vpt_pcounts),
            _ptr(vht_uids),
            _ptr(vht_cvalues), _ptr(vht_ccounts),
            _ptr(vht_pvalues), _ptr(vht_pcounts),
            _size(dpt_uids),
            _ptr(dpt_uids),
            _ptr(dpt_cvalues), _ptr(dpt_ccounts),
            _ptr(dpt_pvalues), _ptr(dpt_pcounts),
            _ptr(dht_uids),
            _ptr(dht_cvalues), _ptr(dht_ccounts),
            _ptr(dht_pvalues), _ptr(dht_pcounts),
            _size(zpt_uids),
            _ptr(zpt_uids),
            _ptr(zpt_cvalues), _ptr(zpt_ccounts),
            _ptr(zpt_pvalues), _ptr(zpt_pcounts),
            _size(spt_uids),
            _ptr(spt_uids),
            _ptr(spt_cvalues), _ptr(spt_ccounts),
            _ptr(spt_pvalues), _ptr(spt_pcounts),
            _size(wpt_uids),
            _ptr(wpt_uids),
            _ptr(wpt_cvalues), _ptr(wpt_ccounts),
            _ptr(wpt_pvalues), _ptr(wpt_pcounts),
            spat_size[0], spat_size[1], spat_size[2],
            spat_step[0], spat_step[1], spat_step[2],
            spat_zero[0], spat_zero[1], spat_zero[2],
            spec_size,
            spec_step,
            spec_zero,
            _ptr(image), _ptr(scube), _ptr(rcube), _ptr(wcube),
            _ptr(rdata), _ptr(vdata), _ptr(ddata))

class _Objective(backend.Objective):

    _CLASSES = {
        np.float32: native_module.Objectivef32}

    def __deepcopy__(self, memodict):
        return self.__class__(self._dtype)

    def __init__(self, dtype):
        self._backend = _get_class(self._CLASSES, dtype)()
        self._dtype = dtype

    def count_pixels(self, data, model, size, count):
        self._backend.count_pixels(_ptr(data), _ptr(model), size, _ptr(count))


import torch
from transformers import AutoModel, AutoTokenizer, RobertaTokenizer, PhobertTokenizer

# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team and authors from University of Illinois at Chicago.
# Copyright (c) 2018, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import copy
import os
import logging
import argparse
import random
import json
import time
import numpy as np
import torch.nn as nn
import torch
from torch.utils.data import RandomSampler, SequentialSampler
from transformers import AdamW, \
    get_linear_schedule_with_warmup, RobertaTokenizer, RobertaModel, RobertaPreTrainedModel
from transformers.models.roberta.modeling_roberta import RobertaLayer

from common import get_labels, get_predictions, \
    get_predictions2, write_result, PreprocessConfig, combine_sentences2, predict, \
    read_preprocess_load, LoaderConfig, convert_to_features, to_data_loader, FILE_ENCODING
import modelconfig
from torchcrf import CRF

logger = logging.getLogger(__name__)

def warmup_linear(x, warmup=0.002):
    if x < warmup:
        return x / warmup
    return 1.0 - x

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

class HSUM(nn.Module):
    def __init__(self, count, config, num_labels):
        super(HSUM, self).__init__()
        self.count = count
        self.num_labels = num_labels
        self.pre_layers = torch.nn.ModuleList()
        self.crf_layers = torch.nn.ModuleList()
        self.classifier = torch.nn.Linear(config.hidden_size, num_labels)
        for i in range(count):
            self.pre_layers.append(RobertaLayer(config))
            self.crf_layers.append(CRF(num_labels))

    def forward(self, layers, attention_mask, labels):
        losses = []
        logitses = []
        output = torch.zeros_like(layers[0])
        total_loss = torch.Tensor(0)
        for i in range(self.count):
            output = output + layers[-i - 1]
            output = self.pre_layers[i](output, attention_mask)[0]
            logits = self.classifier(output)
            if labels is not None:
                loss = self.crf_layers[i](logits.view(100, -1, self.num_labels), labels.view(100, -1))
                losses.append(loss)
            logitses.append(logits)
        if labels is not None:
            total_loss = torch.sum(torch.stack(losses), dim=0)
        avg_logits = torch.sum(torch.stack(logitses), dim=0) / self.count
        return -total_loss, avg_logits

class GRoIE(nn.Module):
    def __init__(self, count, config, num_labels):
        super(GRoIE, self).__init__()
        self.count = count
        self.num_labels = num_labels
        self.pre_layers = torch.nn.ModuleList()
        self.crf_layers = torch.nn.ModuleList()
        self.classifier = torch.nn.Linear(config.hidden_size, num_labels)
        self.dropout = torch.nn.Dropout(config.hidden_dropout_prob)
        for i in range(count):
            self.pre_layers.append(RobertaLayer(config))
            self.crf_layers.append(CRF(num_labels))

    def forward(self, layers, attention_mask, labels):
        losses = []
        logitses = []
        for i in range(self.count):
            layer = self.pre_layers[i](layers[-i - 1], attention_mask)[0]
            layer = self.dropout(layer)
            logits = self.classifier(layer)
            if labels is not None:
                loss = self.crf_layers[i](logits.view(100, -1, self.num_labels), labels.view(100, -1))
                losses.append(loss)
            logitses.append(logits)
        if labels is not None:
            total_loss = torch.sum(torch.stack(losses), dim=0)
        else:
            total_loss = torch.Tensor(0)
        avg_logits = torch.sum(torch.stack(logitses), dim=0) / self.count
        return -total_loss, avg_logits

class PhoBertaForAE(RobertaPreTrainedModel):
    def __init__(self, config, num_labels=3, model_name='P-SUM'):
        super(PhoBertaForAE, self).__init__(config)
        self.num_labels = num_labels
        self.roberta = RobertaModel(config)
        self.groie = GRoIE(4, config, num_labels) if model_name == 'P-SUM' else HSUM(4, config, num_labels)
        self.init_weights()

    def forward(self, input_ids, token_type_ids=None, attention_mask=None, labels=None):
        out_puts = self.roberta(input_ids=input_ids,  # token_type_ids=token_type_ids,
                                attention_mask=attention_mask,
                                output_hidden_states=True,
                                output_attentions=True)
        loss, logits = self.groie(out_puts.hidden_states,
                                  self.get_extended_attention_mask(attention_mask, input_ids.size(), None), labels)
        if labels is not None:
            return loss
        else:
            return logits

def train(args):
    label_list = get_labels()
    tokenizer = PhobertTokenizer.from_pretrained(modelconfig.MODEL_ARCHIVE_MAP[args.bert_model], add_prefix_space=True)

    preprocess_config = PreprocessConfig(tokenizer, seq_len=args.max_seq_length, no_context=args.no_context,
                                         seq_start=args.seq_start)
    vectorize_config = LoaderConfig(batch_size=args.train_batch_size, sampler=RandomSampler)
    train_dataloader, train_data = read_preprocess_load(args.data_dir, "train",
                                                        preprocess_config, vectorize_config)

    num_train_steps = int(len(train_data.sentences) / args.train_batch_size) * args.num_train_epochs

    logger.info("***** Running training *****")
    logger.info("  Num examples = %d", len(train_data.sentences))
    logger.info("  Batch size = %d", args.train_batch_size)
    logger.info("  Num steps = %d", num_train_steps)

    # >>>>> validation
    if args.no_valid is False:
        vectorize_config = LoaderConfig(batch_size=args.train_batch_size, sampler=SequentialSampler)
        valid_dataloader, valid_data = read_preprocess_load(args.data_dir, "dev",
                                                            preprocess_config, vectorize_config)
        logger.info("***** Running validations *****")
        logger.info("  Num orig examples = %d", len(valid_data.sentences))
        logger.info("  Batch size = %d", args.train_batch_size)

        best_valid_loss = float('inf')
        valid_losses = []
    # <<<<< end of validation declarati on

    model = PhoBertaForAE.from_pretrained(modelconfig.MODEL_ARCHIVE_MAP[args.bert_model], num_labels=len(label_list),
                                         model_name='P-SUM')
    model.to(device)
    # Prepare optimizer
    param_optimizer = [(k, v) for k, v in model.named_parameters() if v.requires_grad == True]
    param_optimizer = [n for n in param_optimizer if 'pooler' not in n[0]]
    no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']
    optimizer_grouped_parameters = [
        {'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)], 'weight_decay': 0.01},
        {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}
    ]
    # t_total = num_train_steps
    ### In 🤗 Transformers, optimizer and schedules are split and instantiated like this:
    optimizer = AdamW(optimizer_grouped_parameters, lr=args.learning_rate, weight_decay=0.01,
                      correct_bias=False)  # To reproduce BertAdam specific behavior set correct_bias=False
    scheduler = get_linear_schedule_with_warmup(optimizer,
                                                num_warmup_steps=int(num_train_steps / 10),
                                                num_training_steps=num_train_steps)  # PyTorch scheduler
    global_step = 0
    model.train()
    for epoch in range(args.num_train_epochs):
        for step, batch in enumerate(train_dataloader):
            batch = tuple(t.to(device) for t in batch)
            input_ids, segment_ids, input_mask, label_ids = batch

            loss = model(input_ids, segment_ids, input_mask, label_ids)
            loss.backward()

            torch.nn.utils.clip_grad_norm_(model.parameters(),
                                           1.0)  # Gradient clipping is not in AdamW anymore (so you can use amp without issue)

            optimizer.step()
            scheduler.step()

            optimizer.zero_grad()
            global_step += 1
            # >>>> perform validation at the end of each epoch.
        print("training loss: ", loss.item(), epoch + 1)
        logger.info("training loss: %f, epoch: %d", loss.item(), epoch + 1)
        new_dirs = os.path.join(args.output_dir, str(epoch + 1))
        os.mkdir(new_dirs)
        if args.no_valid is False:
            model.eval()
            with torch.no_grad():
                losses = []
                valid_size = 0
                for step, batch in enumerate(valid_dataloader):
                    batch = tuple(t.to(device) for t in batch)  # multi-gpu does scattering it-self
                    input_ids, segment_ids, input_mask, label_ids = batch
                    loss = model(input_ids, segment_ids, input_mask, label_ids)
                    losses.append(loss.data.item())
                    valid_size += input_ids.size(0)
                valid_loss = sum(losses) / valid_size
                logger.info("validation loss: %f, epoch: %d", valid_loss, epoch + 1)
                valid_losses.append(valid_loss)
                test(args, dev_as_test=True, output_dir=new_dirs, model=model)
            if valid_loss < best_valid_loss:
                best_valid_loss = valid_loss
            model.train()

    if args.no_valid is False:
        with open(os.path.join(args.output_dir, "valid.json"), "w", encoding=FILE_ENCODING) as fw:
            json.dump({"valid_losses": valid_losses}, fw)

    torch.save(model, os.path.join(args.output_dir, "model.pt"))

def test(args, dev_as_test=None, output_dir=None, model=None,
         model_dir=None):  # Load a trained model that you have fine-tuned (we assume evaluate on cpu)
    if output_dir is None:
        output_dir = args.output_dir
    if model_dir is None:
        model_dir = args.output_dir

    tokenizer = PhobertTokenizer.from_pretrained(modelconfig.MODEL_ARCHIVE_MAP[args.bert_model], add_prefix_space=True)
    if dev_as_test:
        data_dir = os.path.join(args.data_dir, 'dev_as_test')
    else:
        data_dir = args.data_dir

    preprocess_config = PreprocessConfig(tokenizer, seq_len=args.max_seq_length, no_context=args.no_context)
    loader_config = LoaderConfig(batch_size=args.eval_batch_size)
    eval_dataloader, eval_data = read_preprocess_load(data_dir, "test",
                                                      preprocess_config, loader_config)

    logger.info("***** Running evaluation *****")
    logger.info("  Num examples = %d", len(eval_data.sentences))
    logger.info("  Batch size = %d", args.eval_batch_size)

    if model is None:
        _model = torch.load(os.path.join(model_dir, "model.pt"))
        _model.to(device)
        _model.eval()
    else:
        _model = model

    probs = np.array(predict(_model, eval_dataloader, device))
    preds = np.argmax(probs, axis=-1)

    pr_ensemble, pr_test_first = get_predictions(preds, eval_data.sentences, eval_data.sentence_numbers)
    prob_ensemble, prob_test_first = get_predictions2(probs, eval_data.sentences, eval_data.sentence_numbers)

    if args.no_context:
        ens = [pr_test_first]
        method_names = ['NC']
    else:
        ens = [pr_ensemble, prob_ensemble, pr_test_first, prob_test_first]
        method_names = ['CMV', 'CMVP', 'F', 'FP']

    for ensem, method_name in zip(ens, method_names):
        output_eval_json = os.path.join(output_dir, "predictions_{}.json".format(method_name))
        write_result(output_eval_json, eval_data.orig_sentences, eval_data.lengths,
                     eval_data.sentences, eval_data.labels, ensem)

    # test with sentence in context
    if args.no_context is False:
        seq_len = args.max_seq_length
        tag_map = {l: i for i, l in enumerate(get_labels())}
        starting_pos = np.arange(0, seq_len, 32)
        # starting_pos[0] = 1
        for start_p in starting_pos:
            tt_lines, tt_tags, line_nos, line_starts = combine_sentences2(eval_data.sentences, eval_data.labels,
                                                                          seq_len - 1, start_p, tokenizer)

            input_ids, segment_ids, masks, label_ids = convert_to_features(tt_lines, tt_tags, tag_map, tokenizer,
                                                                           seq_len)
            data_loader = to_data_loader(input_ids, segment_ids, masks, label_ids, batch_size=args.eval_batch_size)
            probs = np.array(predict(_model, data_loader, device))
            preds = np.argmax(probs, axis=-1).tolist()

            pred_tags = []
            for i, pred in enumerate(preds):
                idx = line_nos[i].index(i)
                pred_tags.append([t for t in
                                  pred[line_starts[i][idx] + 1:line_starts[i][idx] + 1 + len(eval_data.sentences[i])]])

            output_eval_json = os.path.join(output_dir, "predictions_start_position_{}.json".format(start_p))
            write_result(output_eval_json, eval_data.orig_sentences, eval_data.lengths,
                         eval_data.sentences, eval_data.labels, pred_tags)

    print("##############")

def main():
    parser = argparse.ArgumentParser()

    parser.add_argument("--bert_model", default='bert-base', type=str)

    parser.add_argument("--data_dir",
                        default=None,
                        type=str,
                        required=True,
                        help="The input data dir containing json files.")

    parser.add_argument("--output_dir",
                        default=None,
                        type=str,
                        required=True,
                        help="The output directory where the model predictions and checkpoints will be written.")

    ## Other parameters
    parser.add_argument("--max_seq_length",
                        default=100,
                        type=int,
                        help="The maximum total input sequence length after WordPiece tokenization. \n"
                             "Sequences longer than this will be truncated, and sequences shorter \n"
                             "than this will be padded.")
    parser.add_argument("--no_valid",
                        default=False,
                        action='store_true',
                        help="Whether to run validation.")
    parser.add_argument("--only_test",
                        default=False,
                        action='store_true',
                        help="Whether to run training.")
    parser.add_argument("--do_save",
                        default=False,
                        action='store_true',
                        help="Whether to save model after training.")
    parser.add_argument("--no_context",
                        default=False,
                        action='store_true',
                        help="Whether we're training a sentence-crossed model.")
    parser.add_argument("--combine_strategy",
                        default="sequential",
                        type=str,
                        help="Whether we're training a sentence-crossed using combination strategy model.")
    parser.add_argument("--seq_start",
                        default=0,
                        type=int,
                        help="Window Bert start position for processing context data")
    parser.add_argument("--train_batch_size",
                        default=16,
                        type=int,
                        help="Total batch size for training.")
    parser.add_argument("--eval_batch_size",
                        default=8,
                        type=int,
                        help="Total batch size for eval.")
    parser.add_argument("--learning_rate",
                        default=3e-5,
                        type=float,
                        help="The initial learning rate for Adam.")

    parser.add_argument("--num_train_epochs",
                        default=4,
                        type=int,
                        help="Total number of training epochs to perform.")
    parser.add_argument("--warmup_proportion",
                        default=0.1,
                        type=float,
                        help="Proportion of training to perform linear learning rate warmup for. "
                             "E.g., 0.1 = 10%% of training.")
    parser.add_argument('--seed',
                        type=int,
                        default=0,
                        help="random seed for initialization")
    parser.add_argument('--log_file',
                        default=None,
                        type=str,
                        required=True,
                        help="log file")

    args = parser.parse_args()

    # args.seed = 3
    #
    # print(args.seed)

    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    os.makedirs(args.output_dir, exist_ok=True)

    logging.basicConfig(
        filename=args.log_file,
        format='%(asctime)s - %(levelname)s - %(name)s -   %(message)s',
        datefmt='%m/%d/%Y %H:%M:%S',
        level=logging.INFO)

    if args.only_test is False:
        train(args)

    test(args)

    if args.do_save is False:
        os.remove(os.path.join(args.output_dir, "model.pt"))

if __name__ == "__main__":
    start_time = time.time()
    main()
    print("--- %s total seconds ---" % (time.time() - start_time))


# -*- coding: utf-8 -*-
"""
TODO :
    > fix the default output channel number

A completely dongle-free approach to control the thermal cameras
and mics based on the sounddevice library !!

The current script is tailored for use in the audio-thermal recording array
to be used in the field to record bats.

Based on the AVR Soundmexpro based scripts written with
Holger R. Goerlitz

Usage instructions :

Press F5 to run the script start the sync signal playback, you should now have
camera live feed

    To begin recording press any key

    To stop recording press any key again

    To end the session , either :
        a) Interrupt with a system exit : CTRL+c (creates a tiny file at the end because of keyboard triggering)
        b) Press the interrupt button with the mouse (doesn't create an extra file )
            This only works on the IPython shell

The program records N input channels simultaneously and outputs 3 signals when
triggered for recording:

    channel 1: sync. a square 25 Hz signal . The rising edge causes all Thermalcapture
                cameras to capture a frame
    channel 2: trigger. a 20KHz sine wave. when this signal is played the frames
                captured by the cameras are saved to disk.
    channel 3: cross-device sync signal. A copy of the sync signal is played back
                through split BNC cables fed into two Fireface UCs. This allows
                an estimation of the AD conversion delay between the devices

By default, even though data is being collected from all channels, only some channels
are saved into the wav file.

Created on Mon Nov 06 18:10:01 2017

 Version 0.0.0 ( semantic versioning number)

@author: Thejasvi Beleyur
"""
import os
import Queue
import datetime as dt
import time
import numpy as np
import sounddevice as sd
from scipy import signal
import soundfile
import matplotlib.pyplot as plt
plt.rcParams['agg.path.chunksize'] = 10000
from pynput.keyboard import  Listener

class fieldrecorder_trigger():

    def __init__(self,rec_durn, duty_cycle = None, device_name=None,input_output_chs=(2,2),
                 target_dir = '~\\Desktop\\',**kwargs):
        '''

        Inputs:
        rec_durn : float. duration of the whole session in seconds
        device_name : string. name of the device as shown by sd.query_devices()
                    Defaults to None - which will throw an error if there are not at
                    least 3 output channels

        input_output_chs: tuple with integers. Number of channels for recording and playback.

        target_dir : file path. place where the output WAV files will be saved

        **kwargs:
            exclude_channels: list with integers. These channels will not be saved
                              into the WAV file. Defaults to the digital channels
                              in the double Fireface UC setup

            rec_bout : integer. Number of seconds each recording bout should be
                       Defaults to 10 seconds.

            trigger_level : integer <=0. The dB rms for the level at which the recordings
                        get triggered.
            monitor_channels : array like with integers. Channel indices that will be
                               used for monitoring

            bandpass_freqs : tuple. Highpass and lowpass frequencies for the trigger calculation
                       Defaults to the whole frequency spectrum.

        '''
        self.rec_durn = rec_durn
        self.press_count = 0
        self.recording = False
        self.sync_freq = 25
        self.device_name = device_name
        self.input_output_chs = input_output_chs
        self.target_dir = target_dir
        self.FFC_interval = 2
        self.fs = 192000
        self.duty_cycle = duty_cycle

        if self.device_name  is None:
            self.tgt_ind = None
        else:
            self.get_device_indexnumber(self.device_name)

        try:
            expanded_path = os.path.expanduser(target_dir)
            os.chdir(expanded_path)
        except:
            raise ValueError('Unable to find the target directory: ' + target_dir)

        # unless stated specifically exclude the digital input channels
        # from both Firefaces;
        #Device 1: 9:12 , Device 2: 21:24

        self.all_recchannels = range(self.input_output_chs[0])

        if 'exclude_channels' not in kwargs.keys():
            self.exclude_channels = [8,9,10,11,20,21,22,23]
        else:
            self.exclude_channels = kwargs['exclude_channels']

        self.save_channels  = list(set(self.all_recchannels) - set(self.exclude_channels))

        # set the recording bout duration :
        if 'rec_bout' not in kwargs.keys():
            self.rec_bout = 10
        else :
            self.rec_bout = kwargs['rec_bout']

        if 'trigger_level' not in kwargs.keys():
            self.trigger_level = -50 # dB level ref max
        else:
            self.trigger_level = kwargs['trigger_level']

        if 'monitor_channels' not in kwargs.keys():
            self.monitor_channels = [0,1,2,3]
        else:
            self.monitor_channels = kwargs['monitor_channels']

        if 'bandpass_freqs' in kwargs.keys():
            self.highpass_freq, self.lowpass_freq = kwargs['bandpass_freqs']
            nyq_freq = self.fs/2.0
            self.b, self.a = signal.butter(4, [self.highpass_freq/nyq_freq,
                                               self.lowpass_freq/nyq_freq],
                                btype='bandpass')
            self.bandpass = True
        else:
            self.bandpass = False

        if duty_cycle is None:
            self.minimum_interval = 0
        else:
            self.minimum_interval = ((1-duty_cycle)/duty_cycle)*self.rec_bout

    def thermoacousticpy(self):
        '''
        Performs the synchronised recording of thermal cameras and audio.

        '''

        one_cycledurn = 1.0/self.sync_freq
        num_cycles = 1
        sig_durn = num_cycles*one_cycledurn
        t = np.linspace(0,sig_durn,int(self.fs*sig_durn))
        sine_fn = 2*np.pi*self.sync_freq*t + np.pi

        self.sync_signal = np.float32( signal.square(sine_fn,0.5) )
        self.sync_signal *= 0.25

        trigger_freq = 20*10**3

        # generate the FFC signal - which is a 20 KHz sine wave
        ffc_duration = 0.035
        t_ffc = np.linspace(0,ffc_duration, int(self.fs*ffc_duration))
        ffc_freq = 20000
        sine_ffc = np.sin(2*np.pi*ffc_freq*t_ffc)
        self.ffc_signal = np.float32(np.zeros(self.sync_signal.size))
        self.ffc_signal[:sine_ffc.size]  = sine_ffc

        # generate the signals across all channels that will be delivered

        # conv to 32 bit so sounddevice can take the signals as inputs
        self.trigger_signal = np.float32(np.sin(2*np.pi*t*trigger_freq))
        self.empty_signal = np.float32(np.zeros(self.sync_signal.size))

        self.only_sync = np.column_stack((self.sync_signal, self.empty_signal,
                                          self.empty_signal, self.empty_signal,
                                          self.empty_signal))

        self.trig_and_sync = np.column_stack((self.sync_signal, self.trigger_signal,
                                              self.sync_signal,self.empty_signal,
                                              self.empty_signal,))

        self.sync_and_FFC = np.column_stack((self.sync_signal, self.empty_signal,
                                          self.empty_signal, self.empty_signal,
                                          self.ffc_signal))

        self.S = sd.Stream(samplerate=self.fs,blocksize=self.sync_signal.size,
                           channels=self.input_output_chs,device=self.tgt_ind)

        start_time = np.copy(self.S.time)
        rec_time = np.copy(self.S.time)
        end_time =  start_time + self.rec_durn

        self.q = Queue.Queue()

        self.S.start()
        num_recordings = 0
        ffc_recnum = -999
        prev_rectime = 0.0

        try:

            while rec_time < end_time:

                self.mic_inputs = self.S.read(self.trig_and_sync.shape[0])
                self.ref_channels = self.mic_inputs[0][:,self.monitor_channels]
                self.ref_channels_bp = self.bandpass_sound(self.ref_channels)
                self.above_level = self.check_if_above_level(self.ref_channels_bp)

                # if duty cycle recording implemented:
                if self.above_level:
                    self.start_recording = self.minimum_interval_passed(self.S.time,
                                                                    prev_rectime,
                                                                    self.minimum_interval)
                else:
                    self.start_recording = False

                if self.start_recording:
                    print('starting_recording')
                    self.recbout_start_time = np.copy(self.S.time)
                    self.recbout_end_time = self.recbout_start_time + self.rec_bout
                    i = 0
                    print(self.S.time)
                    while  self.recbout_end_time >= self.S.time:
                        if i != 0:
                            self.q.put(self.S.read(self.trig_and_sync.shape[0]))
                        else:
                            self.q.put(self.mic_inputs)

                        self.S.write(self.trig_and_sync)
                        i += 1

                    print(self.S.time)
                    self.empty_qcontentsintolist()
                    self.save_qcontents_aswav()
                    self.start_recording = False

                    num_recordings += 1
                    prev_rectime = np.copy(self.S.time)
                else :

                    ffc_initiate = np.remainder(num_recordings,
                                                self.FFC_interval) == 0
                    if ffc_initiate:
                        # check if FFC has already taken place:
                        if ffc_recnum != num_recordings:
                            self.S.write(self.sync_and_FFC)
                            ffc_recnum = np.copy(num_recordings)

                    self.S.write(self.only_sync)

        except (KeyboardInterrupt, SystemExit):
            print('Stopping recording ..exiting ')

        self.S.stop()
        print('Queue size is',self.q.qsize())
        return(self.fs,self.rec)

    def minimum_interval_passed(self,timenow,last_recordingtime,
                                minimum_interval):
        ''' Calculates the time difference between the time at which the threshold
        is crossed and the last made recording to see if the elapsed time is
        beyond a minimum interval
        '''
        interval = timenow - last_recordingtime
        if interval < 0 :
            raise ValueError('Time elapsed cannot be <0!')

        if interval >= minimum_interval:
            return(True)
        else:
            return(False)

    def bandpass_sound(self, rec_buffer):
        """
        """
        if self.bandpass:
            rec_buffer_bp = np.apply_along_axis(lambda X : signal.lfilter(self.b, self.a, X),
                                                0, rec_buffer)
            return(rec_buffer_bp)
        else:
            return(rec_buffer)

    def check_if_above_level(self, mic_inputs):
        """Checks if the dB rms level of the input recording buffer is above
        threshold. If any of the microphones are above the given level then
        recording is initiated.

        Inputs:

            mic_inputs : Nsamples x Nchannels np.array. Data from soundcard

            level : integer <=0. dB rms ref max . If the input data buffer has an
                    dB rms >= this value then True will be returned.

        Returns:

            above_level : Boolean. True if the buffer dB rms is >= the trigger_level
        """

        dBrms_channel = np.apply_along_axis(self.calc_dBrms, 0, mic_inputs)
        above_level = np.any( dBrms_channel >= self.trigger_level)
        return(above_level)

    def calc_dBrms(self, one_channel_buffer):
        """
        """
        squared = np.square(one_channel_buffer)
        mean_squared = np.mean(squared)
        root_mean_squared = np.sqrt(mean_squared)
        try:
            dB_rms = 20.0*np.log10(root_mean_squared)
        except:
            dB_rms = -999.
        return(dB_rms)

    def empty_qcontentsintolist(self):
        try:
            self.q_contents = [ self.q.get()[0] for i in range(self.q.qsize()) ]

        except:
            raise IOError('Unable to empty queue object contents')

        pass

    def save_qcontents_aswav(self):

        print('Saving file now...')

        self.rec = np.concatenate(self.q_contents)

        self.rec2besaved = self.rec[:,self.save_channels]

        timenow = dt.datetime.now()
        self.timestamp = timenow.strftime('%Y-%m-%d_%H-%M-%S')
        self.idnumber =  int(time.mktime(timenow.timetuple())) #the unix time which provides a 10 digit unique identifier

        main_filename = 'MULTIWAV_' + self.timestamp+'_'+str(self.idnumber) +'.WAV'

        try:
            print('trying to save file... ')

            soundfile.write(main_filename,self.rec2besaved,self.fs)

            print('File saved')

            pass

        except:
            raise IOError('Could not save file !!')

        pass

    def get_device_indexnumber(self,device_name):
        '''
        Check for the device name in all of the recognised devices and
        return the index number within the list.

        '''
        self.device_list = sd.query_devices()

        self.tgt_dev_name = device_name
        self.tgt_dev_bool = [self.tgt_dev_name in each_device['name'] for each_device in self.device_list]

        if not True in self.tgt_dev_bool:

            print (sd.query_devices())

            raise ValueError('The input device \n' + self.tgt_dev_name+
            '\n could not be found, please look at the list above'+
                             ' for all the recognised devices'+
                             ' \n Please use sd.query_devices to check the  recognised'
                             +' devices on this computer')

        if sum(self.tgt_dev_bool) > 1 :
           raise ValueError('Multiple devices with the same string found'
           + ' please enter a more specific device name'
           ' \n Please use sd.query_devices to check the recognised'+
           ' devices on this computer')

        else:
            self.tgt_ind = int(np.argmax(3np.array(self.tgt_dev_bool)))

if __name__ == '__main__':

    dev_name = 'Fireface USB'
    in_out_channels = (24,5)
    tgt_directory = 'C:\\Users\\tbeleyur\\Documents\\fieldwork_2018\\actrackdata\\wav\\2018-07-25_002\\'
    #tgt_directory = 'C:\\Users\\tbeleyur\\Desktop\\test\\'

    a = fieldrecorder_trigger(3500, input_output_chs= in_out_channels,
                              device_name= dev_name, target_dir= tgt_directory,
                              trigger_level=-1.0, duty_cycle=0.18,
                              monitor_channels=[0,1,4], rec_bout = 15.0,
                              bandpass_freqs = [20000.0, 60000.0]
                              )
    fs,rec= a.thermoacousticpy()


import pybullet as p
import pybullet_data
from time import sleep
import random
import numpy as np

p.connect(p.GUI)
p.setGravity(0,0,-1)

p.setAdditionalSearchPath(pybullet_data.getDataPath())
p.loadURDF("plane.urdf")
hands = []
for i in range(5):
    objects = p.loadMJCF("MPL/MPL.xml")
    hand = objects[0]
    hands.append(hand)
    p.resetBasePositionAndOrientation(hand, posObj=[0,0,0.5*i], ornObj=[0,0,0,1])
    for i in range(p.getNumJoints(hand)):
        p.setJointMotorControl2(hand, i, controlMode=p.VELOCITY_CONTROL, targetPosition=0)

for i in range(100000):

    for j,hand in enumerate(hands):
        if (i-j*10)%50==0:
            num_joints = p.getNumJoints(hand)
            # print(num_joints, " Joints.")
            jointIndices = np.arange(num_joints)
            targetPositions = np.random.random(size=[num_joints])
            forces = np.random.random(size=[num_joints])
            # targetPositions[::5] = 0
            p.setJointMotorControlArray(hand, controlMode=p.POSITION_CONTROL, jointIndices=jointIndices, targetPositions=targetPositions, forces=forces)

    p.stepSimulation()
    sleep(.003)


import numpy as np
import pandas as pd

import pysam
import click

BINS = [0, 10, 25, 50, 75, 100, 150, 200, 500, np.inf]

def add_pa_tag(bam_fn, output_bam_fn, polya_lengths_fn, bins):
    polya_lengths = pd.read_table(
        polya_lengths_fn, sep='\t',
        usecols=['readname', 'contig', 'position', 'polya_length', 'qc_tag'],
        na_values=['-1'],
        dtype={'readname': str, 'contig': str, 'position': int, 'polya_length': float, 'qc_tag': 'category'},
        index_col='readname'
    )
    polya_lengths = polya_lengths[polya_lengths.qc_tag == 'PASS'].dropna()

    bins_str = np.array([f'{i:04.0f}-{j:04.0f}' for i, j in zip(bins[:-1], bins[1:])])
    polya_lengths['pa_binned'] = bins_str[np.digitize(polya_lengths.polya_length, bins=bins) - 1]
    polya_lengths_dict = polya_lengths.polya_length.to_dict()
    polya_lengths_bins = polya_lengths.pa_binned.to_dict()
    with pysam.AlignmentFile(bam_fn) as inbam, pysam.AlignmentFile(output_bam_fn, 'wb', template=inbam) as outbam:
        for aln in inbam.fetch():
            try:
                pa = polya_lengths_dict[aln.query_name]
                ba = polya_lengths_bins[aln.query_name]
            except KeyError:
                continue
            aln.set_tag('pA', pa, value_type='f')
            aln.set_tag('bA', ba, value_type='Z')
            outbam.write(aln)

@click.command()
@click.option('-b', '--bam-fn', required=True)
@click.option('-o', '--output-bam-fn', required=True)
@click.option('-p', '--nanopolish-polya-tsv-fn', required=True)
def label_bam(bam_fn, output_bam_fn, nanopolish_polya_tsv_fn):
    '''
    Take TSV output from Nanopolish polya and use it to annotate a BAM file
    using pA tag (float, original polyA length from nanopolish) and bA tag
    (str, binned polyA length). Reads which are unmapped or QC filtered by
    nanopolish are removed.
    '''
    add_pa_tag(bam_fn, output_bam_fn, nanopolish_polya_tsv_fn, bins=BINS)

if __name__ == '__main__':
    label_bam()


import os
import random
import time

import torch
import torch.utils.data

from args import get_args
from hmdb_loader import HMDB_DataLoader
import s3dg
from tqdm import tqdm
import numpy as np

from sklearn import preprocessing
from sklearn.svm import LinearSVC

def main():
    args = get_args()
    assert args.eval_video_root != ''
    model = s3dg.S3D(
        args.num_class, space_to_depth=False, word2vec_path=''
    )
    model = torch.nn.DataParallel(model)
    checkpoint_path = args.pretrain_cnn_path
    print("=> loading checkpoint '{}'".format(checkpoint_path))
    checkpoint = torch.load(checkpoint_path)
    model.load_state_dict(checkpoint["state_dict"])
    model.eval()
    model.cuda()

    # Data loading code
    dataset = HMDB_DataLoader(
        data=os.path.join(os.path.dirname(__file__), 'csv/hmdb51.csv'),
        num_clip=args.num_windows_test,
        video_root=args.eval_video_root,
        num_frames=args.num_frames,
        size=args.video_size,
        crop_only=False,
        center_crop=True,
        with_flip=True,
    )
    dataloader = torch.utils.data.DataLoader(
        dataset,
        batch_size=args.batch_size,
        shuffle=False,
        drop_last=False,
        num_workers=args.num_thread_reader,
    )
    # train for one epoch
    evaluate(dataloader, model, args)

def evaluate(train_loader, model, args):
    all_video_embd = []
    labels = []
    split1 = []
    split2 = []
    split3 = []
    with torch.no_grad():
        for i_batch, data in tqdm(enumerate(train_loader)):
            split1.append(data['split1'])
            split2.append(data['split2'])
            split3.append(data['split3'])
            labels.append(np.array(data['label']))
            video = data['video'].float().cuda()
            video = video / 255.0
            video = video.view(-1, video.shape[2], video.shape[3], video.shape[4], video.shape[5])
            video_embd = model(video, None, mode='video', mixed5c=True)
            video_embd = video_embd.view(len(data['label']), -1, video_embd.shape[1])
            video_embd  = video_embd.cpu().numpy()
            all_video_embd.append(video_embd)
        split1 = torch.cat(split1).cpu().numpy()
        split2 = torch.cat(split2).cpu().numpy()
        split3 = torch.cat(split3).cpu().numpy()
    all_video_embd = np.concatenate(all_video_embd, axis=0)
    labels = np.concatenate(labels)
    le = preprocessing.LabelEncoder()
    labels = le.fit_transform(labels)
    for reg in [100.0]:
        c = LinearSVC(C=reg)
        for split in range(3):
            if split == 0:
                s = split1
            elif split == 1:
                s = split2
            else:
                s = split3
            X_train, X_test = all_video_embd[np.where(s == 1)[0]].reshape((-1, 1024)), all_video_embd[np.where(s == 2)[0]].reshape((-1, 1024))
            label_train, label_test = labels[np.where(s == 1)[0]].repeat(args.num_windows_test), labels[np.where(s == 2)[0]]
            print('Fitting SVM for split {} and C: {}'.format(split + 1, reg))
            c.fit(X_train, label_train)
            X_pred = c.decision_function(X_test)
            X_pred = np.reshape(X_pred, (len(label_test), args.num_windows_test, -1))
            X_pred = X_pred.sum(axis=1)
            X_pred = np.argmax(X_pred, axis=1)
            acc = np.sum(X_pred == label_test) / float(len(X_pred))
            print("Top 1 accuracy split {} and C {} : {}".format(split + 1, reg, acc))

if __name__ == "__main__":
    main()


import sys
import argparse
import subprocess
import re
import numpy as np

sigma_dict = {'VSLAM': {'APE': 72.46, 'RPE': 6.72, 'ARE': 7.41, 'RRE': 0.26, 'Badness': 20.68, 'InitQuality': 2.79, 'robustness': 2.27, 'Relocalization Time': 0.65},
              'VISLAM': {'APE': 55.83, 'RPE': 2.92, 'ARE': 2.48, 'RRE': 0.17, 'Badness': 2.38, 'InitQuality': 1.85, 'robustness': 0.95, 'Relocalization Time': 1.42}}

def get_value(res_str, key):
    reg_float = '\D*([-+]?[0-9]*\.?[0-9]+)'
    try:
        return float(re.search(r'' + key + reg_float, res_str).group(1))
    except:
        print('Error in get_value', key, '\nOutput : ', res_str)
        return -1

def get_score(res_str, method, key, verbose=True):
    raw_value = 0
    if key == 'Badness':
        raw_value = 100 - get_value(res_str, 'CMPL')
    elif key == 'InitQuality':
        raw_value = get_value(res_str, 'E_init')
    else:
        raw_value = get_value(res_str, key)
    sigma = sigma_dict[method][key]
    score = (sigma * sigma) / (sigma * sigma + raw_value * raw_value)
    if raw_value < 0:
        score = 0
    if verbose:
        print('---->', key, ': %.3f' % (raw_value) if raw_value >= 0 else ': invalid', ', Score :', "%.4f" % score)
    return score

def select_round(total_round, gt_folder, traj_folder):
    ape_list = []
    for i in range(total_round):
        traj_file = traj_folder + '/' + str(i) + '-pose.txt'
        try:
            text_result = subprocess.check_output(
                ['bin/accuracy', gt_folder, traj_file, fix_scale]).decode(sys.stdout.encoding)
            ape_list.append(get_score(text_result, method, 'APE', False))
        except Exception:
            print('\nFail to run: ', 'bin/accuracy', gt_folder, traj_file, fix_scale)
            return -1
    return ape_list.index(np.percentile(ape_list, 50, interpolation='nearest'))

# e.g. python3 ismar-score.py --round 5 --is_vislam 1 --trajectory_base_dir /path/to/your/trajectory --gt_base_dir /path/to/your/dataset/folder
if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description='ISMAR2019 SLAM Challenge Timing')
    parser.add_argument('--round', type=str)
    parser.add_argument('--is_vislam', type=str)
    parser.add_argument('--trajectory_base_dir', type=str)
    parser.add_argument('--gt_base_dir', type=str)
    parser.add_argument('--time_ratio', type=str, default='1.0')
    args = parser.parse_args()
    trajectory_base_dir = args.trajectory_base_dir
    gt_base_dir = args.gt_base_dir
    total_round = int(args.round)
    time_ratio = args.time_ratio
    is_vislam = True if args.is_vislam == '1' else '0'
    fix_scale = '1' if is_vislam == True else '0'
    method = 'VISLAM' if is_vislam == True else 'VSLAM'

    # test-round1
    accuracy_cmpl_init_eval_list = ['C1_test', 'C3_test', 'C5_test', 'C7_test', 'C9_test', 'C11_test', 'D8_test', 'D10_test']
    robustness_eval_list = ['D0_test', 'D3_test', 'D4_test']
    reloc_time_eval_list = ['D6_test']

    # test-full
    accuracy_cmpl_init_eval_list = ['C0_test', 'C1_test', 'C2_test', 'C3_test', 'C4_test', 'C5_test', 'C6_test', 'C7_test', 'C8_test', 'C9_test', 'C10_test', 'C11_test', 'D8_test', 'D9_test', 'D10_test']
    robustness_eval_list = ['D0_test', 'D1_test', 'D2_test', 'D3_test', 'D4_test']
    reloc_time_eval_list = ['D5_test', 'D6_test', 'D7_test']

    full_list = list(set(accuracy_cmpl_init_eval_list +
                         robustness_eval_list + reloc_time_eval_list))
    full_list.sort()

    slam_run_full_count = 0
    slam_run_full_time = 0.0

    for seq_name in full_list:
        print('------------------------------')
        has_static_segment = '1' if (seq_name.find('A') != -1 or seq_name.find('C') != -1) else '0'
        gt_folder = gt_base_dir + '/' + seq_name
        round = select_round(total_round, gt_folder,
                             trajectory_base_dir + '/' + seq_name)
        print('Sequence Name', seq_name, ', Select Round', round if round >= 0 else 'invalid')
        # invalid round detected
        if round < 0:
            print('invalid round detected!')
            continue
        time_file = trajectory_base_dir + '/' + \
            seq_name + '/' + str(round) + '-time.txt'

        text_result = subprocess.check_output(
            ['bin/extract_time', gt_folder, time_file]).decode(sys.stdout.encoding)
        slam_run_count = int(get_value(text_result, 'slam_run_count'))
        slam_run_time = get_value(text_result, 'slam_run_time')
        local_frame_rate = slam_run_count / slam_run_time
        print('Frame Rate : %.4f' %(local_frame_rate))
        slam_run_full_count += slam_run_count
        slam_run_full_time += slam_run_time
    if slam_run_full_count > 0:
        frame_rate = slam_run_full_count / slam_run_full_time
        print('Average Frame Rate : %.4f\nNon-Realtime Penalty Factor %.4f' % (frame_rate, min(30.0, frame_rate)/30.0))


import argparse
import glob
import json
import math
import os
import shutil
import sys
import cv2
import numpy as np
import random

def iou(boxA, boxB):
    boxA = [int(x) for x in boxA]
    boxB = [int(x) for x in boxB]

    xA = max(boxA[0], boxB[0])
    yA = max(boxA[1], boxB[1])
    xB = min(boxA[2], boxB[2])
    yB = min(boxA[3], boxB[3])

    interArea = max(0, xB - xA + 1) * max(0, yB - yA + 1)

    boxAArea = (boxA[2] - boxA[0] + 1) * (boxA[3] - boxA[1] + 1)
    boxBArea = (boxB[2] - boxB[0] + 1) * (boxB[3] - boxB[1] + 1)

    iou = interArea / float(boxAArea + boxBArea - interArea)

    return iou

def plot_one_box(x, img, color=None, label=None, line_thickness=None, ground_truth=False):
    """
    description: Plots one bounding box on image img,
                 this function comes from YoLov5 project.
    param:
        x:      a box likes [x1,y1,x2,y2]
        img:    a opencv image object
        color:  color to draw rectangle, such as (0,255,0)
        label:  str
        line_thickness: int
    return:
        no return
    """
    tl = (
        line_thickness or round(0.002 * (img.shape[0] + img.shape[1]) / 2) + 1
    )  # line/font thickness
    color = color or [random.randint(0, 255) for _ in range(3)]
    c1, c2 = (int(x[0]), int(x[1])), (int(x[2]), int(x[3]))
    cv2.rectangle(img, c1, c2, color, thickness=tl, lineType=cv2.LINE_AA)
    if label:
        tf = max(tl - 1, 1)  # font thickness
        if ground_truth:
            cv2.putText(
                img,
                label,
                (c1[0], c2[1] - 2),
                0,
                tl / 3,
                [225, 0, 0],
                thickness=tf,
                lineType=cv2.LINE_AA,
            )
        else:
            cv2.putText(
                img,
                label,
                (c1[0], c1[1] - 2),
                0,
                tl / 3,
                [0, 0, 255],
                thickness=tf,
                lineType=cv2.LINE_AA,
            )

def save_error_image(
    image_dir, basename, save_dir, detection_results, ground_truths, is_num_error=False
):
    img_path = os.path.join(image_dir, basename)
    save_path = os.path.join(save_dir, basename)
    print(img_path)
    img = cv2.imread(img_path)
    tl = round(0.002 * (img.shape[0] + img.shape[1]) / 2) + 1
    if is_num_error:
        cv2.putText(
            img,
            "num error",
            (0, 30),
            0,
            tl / 3,
            [0, 0, 0],
            thickness=max(tl - 1, 1),
            lineType=cv2.LINE_AA,
        )
    else:
        cv2.putText(
            img,
            "cls error",
            (0, 30),
            0,
            tl / 3,
            [0, 0, 0],
            thickness=max(tl - 1, 1),
            lineType=cv2.LINE_AA,
        )
    for detection_result in detection_results:
        plot_one_box(
            detection_result[2:6],
            img,
            (0, 0, 255),
            "{}: {:.2f}".format(detection_result[0], detection_result[1]),
        )
    for ground_truth in ground_truths:
        plot_one_box(
            ground_truth[1:5], img, (255, 0, 0), ground_truth[0], ground_truth=True
        )
    cv2.imwrite(save_path, img)

def main(args):
    os.makedirs(args.save_dir, exist_ok=True)
    ground_truth_txts = glob.glob(os.path.join(args.ground_truth, "*.txt"))
    num_error = 0
    cls_error = 0
    for ground_truth_txt in ground_truth_txts:
        basename = os.path.basename(ground_truth_txt)
        image_id = os.path.splitext(basename)[0]
        detection_results = []
        with open(
            os.path.join(args.detection_result, image_id + ".txt"),
            "r",
            encoding="utf-8",
        ) as f:
            for line in f:
                items = line.strip().split(" ")
                class_name = items[0]
                score = float(items[1])
                bbox_dt = list(map(int, items[2:6]))
                detection_results.append([class_name, score, *bbox_dt])
        ground_truths = []
        with open(ground_truth_txt, "r", encoding="utf-8") as f:
            for line in f:
                items = line.strip().split(" ")
                class_name = items[0]
                bbox_gt = list(map(int, items[1:]))
                ground_truths.append([class_name, *bbox_gt])
        # different count
        if len(detection_results) != len(ground_truths):
            save_error_image(
                args.image_dir,
                image_id + ".jpg",
                args.save_dir,
                detection_results,
                ground_truths,
                is_num_error=True,
            )
            num_error += 1
            continue

        for ground_truth in ground_truths:
            max_iou = -1
            idx = -1
            for i, detection_result in enumerate(detection_results):
                cur_iou = iou(ground_truth[1:5], detection_result[2:6])
                if cur_iou > max_iou:
                    max_iou = cur_iou
                    idx = i
            if ground_truth[0] == detection_results[idx][0] and max_iou > args.overlap:
                max_iou = -1
                continue
            else:
                save_error_image(
                    args.image_dir,
                    image_id + ".jpg",
                    args.save_dir,
                    detection_results,
                    ground_truths,
                )
                cls_error += 1
                print(max_iou, ground_truth[0], detection_results[idx][0])
                break
    print("total", len(ground_truth_txts))
    print("num error", num_error)
    print("cls error", cls_error)

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("-i", "--image_dir", type=str, help="original image dir")
    parser.add_argument("-d", "--detection_result", type=str, help="detection result")
    parser.add_argument("-g", "--ground_truth", type=str, help="ground truth")
    parser.add_argument("-s", "--save_dir", type=str, help="path to save result")
    parser.add_argument("-o", "--overlap", type=float, default=0.5, help="min overlap")
    args = parser.parse_args()
    main(args)


# PYTHON 2.7
import Image, ImageOps
import ImageTk
import Tkinter
import tkMessageBox
import numpy as np
import glob
import os
import face_recognition
from scipy import ndimage
import cv2
import cv

########################################################################################################
# parameters
model_dir = "./trump_train" # trained model
current = 0
size_x = 256 # for the trained model input
size_y = 256
dilation_size = 1
delay = 100 # for video capture
########################################################################################################

# for video capturing
global last_frame                                      #creating global variable
last_frame = np.zeros((size_x, size_y, 3), dtype=np.uint8)
global cap
cap = cv2.VideoCapture(0)
print("*** LOADED VIDEO CAPTURE MODULE ***")
run_gan = False
f_img = np.zeros((size_x, size_y))
gen_img = np.zeros((size_x, size_y))

# load generator from disk
import tensorflow as tf
sess = tf.Session()
saver = tf.train.import_meta_graph(os.path.join(model_dir, "model-270000.meta"))
saver.restore(sess, tf.train.latest_checkpoint(model_dir))
b = sess.graph.get_tensor_by_name("batch:1")
op = sess.graph.get_tensor_by_name("convert_outputs/convert_image:0")
feed = {b: None}
print("*** LOADED MODEL FROM DISK ***")

# move between available images
def move(delta, face_ind=0, run_model=True):

    # collect frame
    orig_img = lmain.img#img
    orig_img = np.array(orig_img)
    #print("Shape: ", orig_img.shape)

    # ready for processing with model
    if run_model:
        #################################################################################
        def create_flat_list(list, ind):
            o = []
            if (len(list) == 0):
                print("No face detected.")
                return None
            fields = list[0].keys()
            for f in fields:
                o.append(list[ind][f])

            flat_list = np.asarray([item for sublist in o for item in sublist])

            return flat_list

        # extract facial landmarks
        face_landmarks_list = face_recognition.face_landmarks(orig_img)
        flfl = create_flat_list(face_landmarks_list, face_ind)
        if flfl is None: return np.zeros((size_x, size_y)), np.zeros((size_x, size_y))

        # create binary image
        bw_img = np.zeros(orig_img.shape[0:2])
        bw_img[tuple(np.vstack((flfl[:, 1], flfl[:, 0])))] = 1  # [255, 0, 0]
        bw_img = np.array(ndimage.grey_dilation(bw_img, size=(dilation_size, dilation_size)), dtype=np.uint8)
        bw_img = cv2.cvtColor(bw_img, cv.CV_GRAY2RGB) * 255

        f_img = bw_img#np.ones((256,256, 3), dtype=np.uint8) * 255
        print("Finished face recognition.")
        #################################################################################

        pre = (f_img.astype(np.float32) / 255) * 2 - 1 # ! important or the output will be corrupted
        feed[b] = np.reshape(pre, (1, size_x, size_y, 3))
        r = sess.run(op, feed)
        gen_img = np.reshape(r, (size_x, size_y, 3)).astype(np.uint8)
        print("Finished generation.")
        #################################################################################

    # prepare for visualization in GUI
    last_gen = gen_img
    last_f = f_img
    gen_img = Image.fromarray(gen_img)
    f_img = Image.fromarray(f_img)
    gen_img = ImageTk.PhotoImage(gen_img)
    f_img = ImageTk.PhotoImage(f_img)
    lmain['text'] = "Original"
    label_gen["text"] = "Generated"
    label_gen["image"] = gen_img
    label_gen.photo = gen_img
    label_f["text"] = "Facial Landmarks"
    label_f["image"] = f_img
    label_f.photo = f_img

    return last_gen, last_f

def show_vid(run_model=True):                                        #creating a function
    if not cap.isOpened():                             #checks for the opening of camera
        print("cant open the camera")
    flag, frame = cap.read()
    frame = cv2.flip(frame, 1)
    if flag is None:
        print("Major error!")
    elif flag:
        global last_frame
        last_frame = frame.copy()

    pic = cv2.cvtColor(last_frame, cv2.COLOR_BGR2RGB)
    img = Image.fromarray(pic)
    img = img.resize((size_x, size_y), Image.ANTIALIAS)
    imgtk = ImageTk.PhotoImage(image=img)
    lmain.img = img
    lmain.imgtk = imgtk
    if (run_model and run_gan): lmain.run_gan = False
    elif (run_model or run_gan): lmain.run_gan = True
    lmain.configure(image=imgtk)

    # run model inference
    last_gen, last_f = move(0, 0, run_model)

    lmain.after(delay, show_vid)

    return last_gen, last_f

if __name__ == '__main__':
    # root frame
    root = Tkinter.Tk()
    root.title("GAN Live Tool")

    # 1st frame
    im_frame = Tkinter.Frame(root, width=1300, height=300)
    im_frame.pack(side=Tkinter.TOP)

    title_label = Tkinter.Label(im_frame, font=('bold'), compound=Tkinter.TOP)
    title_label.pack(side=Tkinter.TOP)
    title_label["text"] = "Video fed from Camera"
    model_label = Tkinter.Label(im_frame, font=('bold'), compound=Tkinter.TOP)
    model_label.pack(side=Tkinter.TOP)
    model_label["text"] = "Model from: " + model_dir
    lmain = Tkinter.Label(im_frame, font=('bold'), compound=Tkinter.TOP)
    lmain.pack(side=Tkinter.LEFT)
    label_f = Tkinter.Label(im_frame, font=('bold'), compound=Tkinter.TOP)
    label_f.pack(side=Tkinter.LEFT)
    label_gen = Tkinter.Label(im_frame, font=('bold'), compound=Tkinter.TOP)
    label_gen.pack(side=Tkinter.LEFT)

    #run_frame = Tkinter.Frame(root)
    #run_frame.pack(side=Tkinter.BOTTOM)
    # for activating real time transformation
    #Tkinter.Button(run_frame, text="START / STOP", command=lambda: show_vid(True)).pack(side=Tkinter.BOTTOM)

    # just run webcam
    last_gen, last_f = show_vid(True)

    # update
    gen_img = last_gen
    f_img = last_f

    # run
    #move(0)

    root.mainloop()
    cap.release()


import pickle
import sys
import time
import numpy as np

import torch
from fvcore.nn import FlopCountAnalysis
import xnas.core.config as config
import xnas.logger.logging as logging
from xnas.core.config import cfg
from xnas.core.builder import setup_env, space_builder, optimizer_builder

import xnas.algorithms.RMINAS.sampler.sampling as sampling
from xnas.algorithms.RMINAS.utils.RMI_torch import RMI_loss
from xnas.algorithms.RMINAS.sampler.RF_sampling import RF_suggest
from xnas.algorithms.RMINAS.utils.random_data import get_random_data

from xnas.spaces.DARTS.utils import geno_from_alpha, reformat_DARTS
from xnas.spaces.NASBench201.utils import dict2config, get_cell_based_tiny_net, CellStructure

config.load_configs()
logger = logging.get_logger(__name__)

# RMINAS hyperparameters initialization
RF_space = None
api = None

def rminas_hp_builder():
    global RF_space, api
    if cfg.SPACE.NAME == 'infer_nb201':
        from nas_201_api import NASBench201API
        api = NASBench201API(cfg.BENCHMARK.NB201PATH)
        RF_space = 'nasbench201'
    elif cfg.SPACE.NAME == 'infer_darts':
        RF_space = 'darts'
    elif cfg.SPACE.NAME == 'nasbenchmacro':
        RF_space = 'nasbenchmacro'
        from xnas.evaluations.NASBenchMacro.evaluate import evaluate, data
        api = data
    elif cfg.SPACE.NAME == 'proxyless':
        RF_space = 'proxyless'
            # for example : arch = '00000000'
            # arch = ''
            # evaluate(arch)

def main():
    device = setup_env()

    rminas_hp_builder()

    assert cfg.SPACE.NAME in ['infer_nb201', 'infer_darts',"nasbenchmacro", "proxyless"]
    assert cfg.LOADER.DATASET in ['cifar10', 'cifar100', 'imagenet', 'imagenet16_120'], 'dataset error'

    if cfg.LOADER.DATASET == 'cifar10':
        from xnas.algorithms.RMINAS.teacher_model.resnet20_cifar10.resnet import resnet20
        checkpoint_res = torch.load('xnas/algorithms/RMINAS/teacher_model/resnet20_cifar10/resnet20.th')
        network = torch.nn.DataParallel(resnet20())
        network.load_state_dict(checkpoint_res['state_dict'])
        network = network.module

    elif cfg.LOADER.DATASET == 'cifar100':
        from xnas.algorithms.RMINAS.teacher_model.resnet101_cifar100.resnet import resnet101
        network = resnet101()
        network.load_state_dict(torch.load('xnas/algorithms/RMINAS/teacher_model/resnet101_cifar100/resnet101.pth'))

    elif cfg.LOADER.DATASET == 'imagenet':
        assert cfg.SPACE.NAME in ('infer_darts', 'proxyless')
        logger.warning('Our method does not directly search in ImageNet.')
        logger.warning('Only partial tests have been conducted, please use with caution.')
        import xnas.algorithms.RMINAS.teacher_model.fbresnet_imagenet.fbresnet as fbresnet
        network = fbresnet.fbresnet152()

    elif cfg.LOADER.DATASET == 'imagenet16_120':
        assert cfg.SPACE.NAME == 'infer_nb201'
        from nas_201_api import ResultsCount

        """Teacher Network: using best arch searched from cifar10 and weight from nb201."""
        filename = 'xnas/algorithms/RMINAS/teacher_model/nb201model_imagenet16120/009930-FULL.pth'
        xdata = torch.load(filename)
        odata  = xdata['full']['all_results'][('ImageNet16-120', 777)]
        result = ResultsCount.create_from_state_dict(odata)
        result.get_net_param()
        arch_config = result.get_config(CellStructure.str2structure) # create the network with params
        net_config = dict2config(arch_config, None)
        network = get_cell_based_tiny_net(net_config)
        network.load_state_dict(result.get_net_param())

    network.cuda()

    """selecting well-performed data."""
    more_data_X, more_data_y = get_random_data(cfg.LOADER.BATCH_SIZE, cfg.LOADER.DATASET)
    with torch.no_grad():
        ce_loss = torch.nn.CrossEntropyLoss(reduction='none').cuda()
        more_logits = network(more_data_X)
        _, indices = torch.topk(-ce_loss(more_logits, more_data_y).cpu().detach(), cfg.LOADER.BATCH_SIZE)
    data_y = more_data_y.detach()
    data_X = more_data_X.detach()
    with torch.no_grad():
        feature_res = network.feature_extractor(data_X)

    RFS = RF_suggest(space=RF_space, logger=logger, api=api, thres_rate=cfg.RMINAS.RF_THRESRATE, seed=cfg.RNG_SEED)

    # loss function
    loss_fun_cka = RMI_loss(data_X.size()[0])
    loss_fun_cka = loss_fun_cka.requires_grad_()
    loss_fun_cka.cuda()
    loss_fun_log = torch.nn.CrossEntropyLoss().cuda()

    def train_arch(modelinfo):
        flops = None
        if cfg.SPACE.NAME == 'infer_nb201':
            # get arch
            arch_config = {
                'name': 'infer.tiny',
                'C': 16, 'N': 5,
                'arch_str':api.arch(modelinfo),
                'num_classes': cfg.LOADER.NUM_CLASSES}
            net_config = dict2config(arch_config, None)
            model = get_cell_based_tiny_net(net_config).cuda()
        elif cfg.SPACE.NAME == 'infer_darts':
            cfg.TRAIN.GENOTYPE = str(modelinfo)
            model = space_builder().cuda()
        elif cfg.SPACE.NAME == 'nasbenchmacro':
            model = space_builder().cuda()
            optimizer = optimizer_builder("SGD", model.parameters())
        elif cfg.SPACE.NAME == 'proxyless':
            model = space_builder(stage_width_list=[16, 24, 40, 80, 96, 192, 320],depth_param=modelinfo[:6],ks=modelinfo[6:27][modelinfo[6:27]>0],expand_ratio=modelinfo[27:][modelinfo[27:]>0],dropout_rate=0).cuda()
            optimizer = optimizer_builder("SGD", model.parameters())
            with torch.no_grad():
                tensor = (torch.rand(1, 3, 224, 224).cuda(),)
                flops = FlopCountAnalysis(model, tensor).total()
            # lr_scheduler = lr_scheduler_builder(optimizer)

            # nbm_trainer = OneShotTrainer(
            #     supernet=model,
            #     criterion=criterion,
            #     optimizer=optimizer,
            #     lr_scheduler=lr_scheduler,
            #     train_loader=train_loader,
            #     test_loader=valid_loader,
            #     sample_type='iter'
            # )

        model.train()
        # weights optimizer
        optimizer = optimizer_builder("SGD", model.parameters())
        epoch_losses = []
        for cur_epoch in range(1, cfg.OPTIM.MAX_EPOCH+1):
            optimizer.zero_grad()

            features, logits = model.forward_with_features(data_X, modelinfo)
            loss_cka = loss_fun_cka(features, feature_res)
            loss_logits = loss_fun_log(logits, data_y)
            loss = cfg.RMINAS.LOSS_BETA * loss_cka + (1-cfg.RMINAS.LOSS_BETA)*loss_logits
            loss.backward()

            optimizer.step()
            epoch_losses.append(loss.detach().cpu().item())
            if cur_epoch == cfg.OPTIM.MAX_EPOCH:

                return loss.cpu().detach().numpy(), {'epoch_losses':epoch_losses, 'flops':flops}

    trained_arch_darts, trained_loss = [], []
    def train_procedure(sample):
        if cfg.SPACE.NAME == 'infer_nb201':
            mixed_loss, epoch_losses = train_arch(sample)[0]
            mixed_loss = np.inf if np.isnan(mixed_loss) else mixed_loss
            trained_loss.append(mixed_loss)
            arch_arr = sampling.nb201genostr2array(api.arch(sample))
            RFS.trained_arch.append({'arch':arch_arr, 'loss':mixed_loss})
            RFS.trained_arch_index.append(sample)
        elif cfg.SPACE.NAME == 'infer_darts':
            sample_geno = geno_from_alpha(sampling.darts_sug2alpha(sample))  # type=Genotype
            trained_arch_darts.append(str(sample_geno))
            mixed_loss, epoch_losses = train_arch(sample_geno)[0]
            mixed_loss = np.inf if np.isnan(mixed_loss) else mixed_loss
            trained_loss.append(mixed_loss)
            RFS.trained_arch.append({'arch':sample, 'loss':mixed_loss})
        elif cfg.SPACE.NAME == 'nasbenchmacro':
            sample_geno = ''.join(sample.astype('str'))  # type=Genotype
            trained_arch_darts.append((sample_geno))
            mixed_loss, info = train_arch(sample)
            mixed_loss = np.inf if np.isnan(mixed_loss) else mixed_loss
            trained_loss.append(mixed_loss)
            RFS.trained_arch.append({'arch':sample, 'loss':mixed_loss,'gt':api[sample_geno]['mean_acc'],'losses':info["epoch_losses"]})
        elif cfg.SPACE.NAME == 'proxyless':
            sample_geno = ''.join(sample.astype('str'))  # type=Genotype
            trained_arch_darts.append((sample_geno))
            mixed_loss, info = train_arch(sample)
            mixed_loss = np.inf if np.isnan(mixed_loss) else mixed_loss
            trained_loss.append(mixed_loss)
            RFS.trained_arch.append({'arch':sample, 'loss':mixed_loss,'gt':info["flops"],'losses':info["epoch_losses"]})

        logger.info("sample: {}, loss:{}".format(sample, mixed_loss))

    start_time = time.time()
    # ====== Warmup ======
    warmup_samples = RFS.warmup_samples(cfg.RMINAS.RF_WARMUP)
    logger.info("Warming up with {} archs".format(cfg.RMINAS.RF_WARMUP))
    for sample in warmup_samples:
        train_procedure(sample)
    RFS.Warmup()
    # ====== RF Sampling ======
    sampling_time = time.time()
    sampling_cnt= 0
    while sampling_cnt < cfg.RMINAS.RF_SUCC:
        print(sampling_cnt)
        sample = RFS.fitting_samples()
        train_procedure(sample)
        sampling_cnt += RFS.Fitting()
    if sampling_cnt >= cfg.RMINAS.RF_SUCC:
        logger.info('successfully sampling good archs for {} times'.format(sampling_cnt))
    else:
        logger.info('failed sampling good archs for only {} times'.format(sampling_cnt))
    logger.info('RF sampling time cost:{}'.format(str(time.time() - sampling_time)))

    # ====== Evaluation ======
    logger.info('Total time cost: {}'.format(str(time.time() - start_time)))
    logger.info('Actual training times: {}'.format(len(RFS.trained_arch_index)))
    if cfg.SPACE.NAME == 'infer_nb201':
        logger.info('Searched architecture:\n{}'.format(str(RFS.optimal_arch(method='sum', top=50))))
        logger.info('Searched architecture:\n{}'.format(str(RFS.optimal_arch(method='greedy', top=50))))
    elif cfg.SPACE.NAME == 'infer_darts':
        op_sample = RFS.optimal_arch(method='sum', top=50)
        op_alpha = torch.from_numpy(np.r_[op_sample, op_sample])
        op_geno = reformat_DARTS(geno_from_alpha(op_alpha))
        logger.info('Searched architecture@top50:\n{}'.format(str(op_geno)))
    elif cfg.SPACE.NAME == 'nasbenchmacro':
        op_sample = RFS.optimal_arch(method='sum', top=50)
        # op_alpha = torch.from_numpy(np.r_[op_sample, op_sample])
        # op_geno = reformat_DARTS(geno_from_alpha(op_alpha))
        logger.info('Searched architecture@top50:\n{}'.format(str(op_sample)))
        print(api[op_sample]['mean_acc'])
    elif cfg.SPACE.NAME == 'proxyless':
        op_sample = RFS.optimal_arch(method='sum', top=100)
        # op_alpha = torch.from_numpy(np.r_[op_sample, op_sample])
        # op_geno = reformat_DARTS(geno_from_alpha(op_alpha))
        logger.info('Searched architecture@top100:\n{}'.format(str(op_sample)))
        print(api[op_sample]['mean_acc'])

if __name__ == '__main__':
    main()


#!/usr/bin/env python

# numpy package
import numpy as np

# torch package
import torch
import torchvision
from torch.nn.functional import cross_entropy, softmax, log_softmax

# basic package
import os
import sys
sys.path.append('.')
import argparse
from tqdm import tqdm
from datetime import datetime

# custom package
from loader.argument_print import argument_print
from loader.loader import dataset_loader, network_loader, attack_loader

# cudnn enable
torch.backends.cudnn.benchmark = True
torch.backends.cudnn.enabled = True

# argument parser
parser = argparse.ArgumentParser(description='Joint Adversarial Defense')
parser.add_argument('--lr', default=0.01, type=float, help='learning rate')
parser.add_argument('--steps', default=10, type=int, help='adv. steps')
parser.add_argument('--eps', required=True, type=float, help='max norm')
parser.add_argument('--dataset', required=True, type=str, help='dataset name')
parser.add_argument('--network', required=True, type=str, help='network name')
parser.add_argument('--data_root', required=True, type=str, help='path to dataset')
parser.add_argument('--epoch', default=60, type=int, help='epoch number')
parser.add_argument('--attack', default='pgd', type=str, help='attack type')
parser.add_argument('--save_dir', default='./experiment', type=str, help='save directory')
args = parser.parse_args()

# loading dataset, network, and attack
trainloader, testloader = dataset_loader(args)
net = network_loader(args, mean=args.mean, std=args.std).cuda()
attack = attack_loader(args, net)

# Adam Optimizer with KL divergence, and Scheduling Learning rate
optimizer = torch.optim.SGD(net.parameters(), lr=args.lr, momentum=0.9, weight_decay=1e-2)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=20, gamma=0.2)

# Setting checkpoint date time
date_time = datetime.today().strftime("%m%d%H%M")

# checkpoint_name
checkpoint_name = 'TRADE_'+args.network+'_'+args.dataset+'_'+date_time+'.pth'

# argument print
argument_print(args, checkpoint_name)

def train():

    # Modeling Adversarial Loss
    for epoch in range(args.epoch):

        # train environment
        net.train()

        print('\n\n[TRADE/Epoch] : {}'.format(epoch+1))

        total_cross_loss = 0
        correct = 0
        total = 0

        for batch_idx, (inputs, targets) in enumerate(tqdm(trainloader)):

            # dataloader parsing and generate adversarial examples
            inputs, targets = inputs.cuda(), targets.cuda()

            # learning network parameters
            optimizer.zero_grad()
            adv_x = attack(inputs, targets) if args.eps != 0 else inputs
            cross_entropy_loss = cross_entropy(net(adv_x), targets)-(softmax(net(inputs).detach(),dim=1)*log_softmax(net(adv_x),dim=1)).sum(dim=1).mean()
            cross_entropy_loss.backward()
            optimizer.step()

            # validation
            pred = torch.max(net(adv_x).detach(), dim=1)[1]
            correct += torch.sum(pred.eq(targets)).item()
            total += targets.numel()

            # logging two types loss and total loss
            total_cross_loss += cross_entropy_loss.item()

            if batch_idx % 50 == 0 and batch_idx != 0:
                print('[TRADE/Train] Iter: {}, Acc: {:.3f}, CE: {:.3f}'.format(
                    batch_idx, # Iter
                    100.*correct / total, # Acc
                    total_cross_loss / (batch_idx+1) # CrossEntropy
                    )
                )

        # Scheduling learning rate by stepLR
        scheduler.step()

        # Adversarial validation
        adversarial_test()

        # Save checkpoint file
        torch.save({
            'epoch': epoch+1,
            'model_state_dict': net.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'total_cross_entropy_loss' : total_cross_loss / (batch_idx+1)
            }, os.path.join(args.save_dir,checkpoint_name))

        # argument print
        argument_print(args, checkpoint_name)

def adversarial_test():

    correct = 0
    total = 0

    print('\n\n[TRADE/Test] Under Testing ... Wait PLZ')
    for batch_idx, (inputs, targets) in enumerate(tqdm(testloader)):

        # dataloader parsing and generate adversarial examples
        inputs, targets = inputs.cuda(), targets.cuda()
        adv_x = attack(inputs, targets) if args.eps != 0 else inputs

        # Evaluation
        outputs = net(adv_x).detach()

        # Test
        pred = torch.max(outputs, dim=1)[1]
        correct += torch.sum(pred.eq(targets)).item()
        total += targets.numel()

    print('[TRADE/Test] Acc: {:.3f}'.format(100.*correct / total))

if __name__ == "__main__":
    train()