import numpy as np
import pandas as pd
import re
import itertools
from larch import P, X, DataFrames, Model
from larch.model.model_group import ModelGroup
from larch.util import Dict
import larch
import os
import yaml
import importlib
import logging
from larch.log import logger_name
from pathlib import Path

from ...abm.models.util import cdap
from .general import apply_coefficients, explicit_value_parameters

_logger = logging.getLogger(logger_name)

def generate_alternatives(n_persons):
    """
    Generate a dictionary of CDAP alternatives.

    The keys are the names of the patterns, and
    the values are the alternative code numbers.

    Parameters
    ----------
    n_persons : int

    Returns
    -------
    Dict
    """
    basic_patterns = ["M", "N", "H"]
    alt_names = list(
        "".join(i) for i in itertools.product(basic_patterns, repeat=n_persons)
    )
    alt_codes = np.arange(1, len(alt_names) + 1)
    return Dict(zip(alt_names, alt_codes))

def apply_replacements(expression, prefix, tokens):
    """
    Convert general person terms to specific person terms for the CDAP model.

    Parameters
    ----------
    expression : str
        An expression from the "Expression" column
        of cdap_INDIV_AND_HHSIZE1_SPEC.csv, or similar.
    prefix : str
        A prefix to attach to each token in `expression`.
    tokens : list-like of str
        A list of tokens to edit within an expression.

    Returns
    -------
    expression : str
        The modified expression
    """
    for i in tokens:
        expression = re.sub(fr"\b{i}\b", f"{prefix}_{i}", expression)
    return expression

def cdap_base_utility_by_person(model, n_persons, spec, alts=None, value_tokens=()):
    """
    Build the base utility by person for each pattern.

    Parameters
    ----------
    model : larch.Model
    n_persons : int
    spec : pandas.DataFrame
        The base utility by person spec provided by
        the ActivitySim framework.
    alts : dict, optional
        The keys are the names of the patterns, and
        the values are the alternative code numbers,
        as created by `generate_alternatives`.  If not
        given, the alts are automatically regenerated
        using that function.
    value_tokens : list-like of str, optional
        A list of tokens to edit within an the expressions,
        generally the column names of the provided values
        from the estimation data bundle.  Only used when
        `n_persons` is more than 1.
    """
    if n_persons == 1:
        for i in spec.index:
            if not pd.isna(spec.loc[i, "M"]):
                model.utility_co[1] += X(spec.Expression[i]) * P(spec.loc[i, "M"])
            if not pd.isna(spec.loc[i, "N"]):
                model.utility_co[2] += X(spec.Expression[i]) * P(spec.loc[i, "N"])
            if not pd.isna(spec.loc[i, "H"]):
                model.utility_co[3] += X(spec.Expression[i]) * P(spec.loc[i, "H"])
    else:
        if alts is None:
            alts = generate_alternatives(n_persons)
        person_numbers = range(1, n_persons + 1)
        for pnum in person_numbers:
            for i in spec.index:
                for aname, anum in alts.items():
                    z = pnum - 1
                    if not pd.isna(spec.loc[i, aname[z]]):
                        x = apply_replacements(
                            spec.Expression[i], f"p{pnum}", value_tokens
                        )
                        model.utility_co[anum] += X(x) * P(spec.loc[i, aname[z]])

def interact_pattern(n_persons, select_persons, tag):
    """
    Compile a regex pattern to match CDAP alternatives.

    Parameters
    ----------
    n_persons : int
    select_persons : list-like of int
        The persons to be selected.
    tag : str
        The activity letter, currently one of {M,N,H}.

    Returns
    -------
    re.compile
    """

    pattern = ""
    p = 1
    while len(pattern) < n_persons:
        pattern += tag if p in select_persons else "."
        p += 1
    return re.compile(pattern)

def cdap_interaction_utility(model, n_persons, alts, interaction_coef, coefficients):

    person_numbers = list(range(1, n_persons + 1))

    matcher = re.compile("coef_[HMN]_.*")
    interact_coef_map = {}
    for c in coefficients.index:
        if matcher.search(c):
            c_split = c.split("_")
            for j in c_split[2:]:
                interact_coef_map[(c_split[1], j)] = c
                if all((i == "x" for i in j)):  # wildcards also map to empty
                    interact_coef_map[(c_split[1], "")] = c

    for (cardinality, activity), coefs in interaction_coef.groupby(
        ["cardinality", "activity"]
    ):
        _logger.info(
            f"{n_persons} person households, interaction cardinality {cardinality}, activity {activity}"
        )
        if cardinality > n_persons:
            continue
        elif cardinality == n_persons:
            this_aname = activity * n_persons
            this_altnum = alts[this_aname]
            for rowindex, row in coefs.iterrows():
                expression = "&".join(
                    f"(p{p}_ptype == {t})"
                    for (p, t) in zip(person_numbers, row.interaction_ptypes)
                    if t != "*"
                )
                if expression:
                    if (activity, row.interaction_ptypes) in interact_coef_map:
                        linear_component = X(expression) * P(
                            interact_coef_map[(activity, row.interaction_ptypes)]
                        )
                    else:
                        linear_component = X(expression) * P(row.coefficient)
                else:
                    if (activity, row.interaction_ptypes) in interact_coef_map:
                        linear_component = P(
                            interact_coef_map[(activity, row.interaction_ptypes)]
                        )
                    else:
                        linear_component = P(row.coefficient)
                _logger.debug(
                    f"utility_co[{this_altnum} {this_aname}] += {linear_component}"
                )
                model.utility_co[this_altnum] += linear_component
        elif cardinality < n_persons:
            for combo in itertools.combinations(person_numbers, cardinality):
                pattern = interact_pattern(n_persons, combo, activity)
                for aname, anum in alts.items():
                    if pattern.match(aname):
                        for rowindex, row in coefs.iterrows():
                            expression = "&".join(
                                f"(p{p}_ptype == {t})"
                                for (p, t) in zip(combo, row.interaction_ptypes)
                                if t != "*"
                            )
                            # interaction terms without ptypes (i.e. with wildcards)
                            # only apply when the household size matches the cardinality
                            if expression != "":
                                if (
                                    activity,
                                    row.interaction_ptypes,
                                ) in interact_coef_map:
                                    linear_component = X(expression) * P(
                                        interact_coef_map[
                                            (activity, row.interaction_ptypes)
                                        ]
                                    )
                                else:
                                    linear_component = X(expression) * P(
                                        row.coefficient
                                    )
                                _logger.debug(
                                    f"utility_co[{anum} {aname}] += {linear_component}"
                                )
                                model.utility_co[anum] += linear_component

def cdap_split_data(households, values):
    if "cdap_rank" not in values:
        raise ValueError("assign cdap_rank to values first")
    # only process the first 5 household members
    values = values[values.cdap_rank <= 5]
    cdap_data = {}
    for hhsize, hhs_part in households.groupby(households.hhsize.clip(1, 5)):
        if hhsize == 1:
            v = pd.merge(values, hhs_part.household_id, on="household_id").set_index(
                "household_id"
            )
        else:
            v = (
                pd.merge(values, hhs_part.household_id, on="household_id")
                .set_index(["household_id", "cdap_rank"])
                .unstack()
            )
            v.columns = [f"p{i[1]}_{i[0]}" for i in v.columns]
            for agglom in ["override_choice", "model_choice"]:
                v[agglom] = v[[f"p{p}_{agglom}" for p in range(1, hhsize + 1)]].sum(1)
        cdap_data[hhsize] = v
    return cdap_data

def cdap_dataframes(households, values):
    data = cdap_split_data(households, values)
    dfs = {}
    for hhsize in data.keys():
        alts = generate_alternatives(hhsize)
        dfs[hhsize] = DataFrames(
            co=data[hhsize],
            alt_names=alts.keys(),
            alt_codes=alts.values(),
            av=1,
            ch=data[hhsize].override_choice.map(alts),
        )
    return dfs

# def _cdap_model(households, values, spec1, interaction_coef, coefficients):
#     cdap_data = cdap_dataframes(households, values)
#     m = {}
#     _logger.info(f"building for model 1")
#     m[1] = Model(dataservice=cdap_data[1])
#     cdap_base_utility_by_person(m[1], n_persons=1, spec=spec1)
#     m[1].choice_any = True
#     m[1].availability_any = True
#
#     # Add cardinality into interaction_coef if not present
#     if 'cardinality' not in interaction_coef:
#         interaction_coef['cardinality'] = interaction_coef['interaction_ptypes'].str.len()
#     for s in [2, 3, 4, 5]:
#         _logger.info(f"building for model {s}")
#         m[s] = Model(dataservice=cdap_data[s])
#         alts = generate_alternatives(s)
#         cdap_base_utility_by_person(m[s], s, spec1, alts, values.columns)
#         cdap_interaction_utility(m[s], s, alts, interaction_coef, coefficients)
#         m[s].choice_any = True
#         m[s].availability_any = True
#
#     result = ModelGroup(m.values())
#     explicit_value_parameters(result)
#     apply_coefficients(coefficients, result)
#     return result

def cdap_data(
    name="cdap",
    edb_directory="output/estimation_data_bundle/{name}/",
    coefficients_file="{name}_coefficients.csv",
    interaction_coeffs_file="{name}_interaction_coefficients.csv",
    households_file="../../final_households.csv",
    persons_file="../../final_persons.csv",
    spec1_file="{name}_INDIV_AND_HHSIZE1_SPEC.csv",
    settings_file="{name}_model_settings.yaml",
    chooser_data_file="{name}_values_combined.csv",
):
    edb_directory = edb_directory.format(name=name)
    if not os.path.exists(edb_directory):
        raise FileNotFoundError(edb_directory)

    def read_csv(filename, **kwargs):
        filename = filename.format(name=name)
        return pd.read_csv(os.path.join(edb_directory, filename), **kwargs)

    def read_yaml(filename, **kwargs):
        filename = filename.format(name=name)
        with open(os.path.join(edb_directory, filename), "rt") as f:
            return yaml.load(f, Loader=yaml.SafeLoader, **kwargs)

    settings = read_yaml(settings_file)

    try:
        hhs = read_csv(households_file)
    except FileNotFoundError:
        hhs = pd.read_csv(households_file)

    try:
        persons = read_csv(persons_file)
    except FileNotFoundError:
        persons = pd.read_csv(persons_file)

    person_type_map = settings.get('PERSON_TYPE_MAP')
    if person_type_map is None:
        raise KeyError("PERSON_TYPE_MAP missing from cdap_settings.yaml")

    person_rank = cdap.assign_cdap_rank(persons, person_type_map)

    coefficients = read_csv(
        coefficients_file, index_col="coefficient_name", comment="#",
    )

    interaction_coef = read_csv(
        interaction_coeffs_file,
        dtype={"interaction_ptypes": str},
        keep_default_na=False,
        comment="#",
    )

    spec1 = read_csv(spec1_file, comment="#")
    values = read_csv(chooser_data_file, comment="#")
    values["cdap_rank"] = person_rank

    return Dict(
        edb_directory=Path(edb_directory),
        person_data=values,
        spec1=spec1,
        interaction_coef=interaction_coef,
        coefficients=coefficients,
        households=hhs,
        settings=settings,
    )

def cdap_model(
    edb_directory="output/estimation_data_bundle/{name}/",
    coefficients_file="{name}_coefficients.csv",
    interaction_coeffs_file="{name}_interaction_coefficients.csv",
    households_file="../../final_households.csv",
    persons_file="../../final_persons.csv",
    spec1_file="{name}_INDIV_AND_HHSIZE1_SPEC.csv",
    settings_file="{name}_model_settings.yaml",
    chooser_data_file="{name}_values_combined.csv",
    return_data=False,
):
    d = cdap_data(
        name="cdap",
        edb_directory=edb_directory,
        coefficients_file=coefficients_file,
        interaction_coeffs_file=interaction_coeffs_file,
        households_file=households_file,
        persons_file=persons_file,
        spec1_file=spec1_file,
        settings_file=settings_file,
        chooser_data_file=chooser_data_file,
    )

    households = d.households
    values = d.person_data
    spec1 = d.spec1
    interaction_coef = d.interaction_coef
    coefficients = d.coefficients

    cdap_dfs = cdap_dataframes(households, values)
    m = {}
    _logger.info(f"building for model 1")
    m[1] = Model(dataservice=cdap_dfs[1])
    cdap_base_utility_by_person(m[1], n_persons=1, spec=spec1)
    m[1].choice_any = True
    m[1].availability_any = True

    # Add cardinality into interaction_coef if not present
    if "cardinality" not in interaction_coef:
        interaction_coef["cardinality"] = interaction_coef[
            "interaction_ptypes"
        ].str.len()
    for s in [2, 3, 4, 5]:
        _logger.info(f"building for model {s}")
        m[s] = Model(dataservice=cdap_dfs[s])
        alts = generate_alternatives(s)
        cdap_base_utility_by_person(m[s], s, spec1, alts, values.columns)
        cdap_interaction_utility(m[s], s, alts, interaction_coef, coefficients)
        m[s].choice_any = True
        m[s].availability_any = True

    model = ModelGroup(m.values())
    explicit_value_parameters(model)
    apply_coefficients(coefficients, model)
    if return_data:
        return model, d
    return model


"""The CarliniWagnerL2 attack
"""
# pylint: disable=missing-docstring
import logging

import numpy as np
import tensorflow as tf

from cleverhans.attacks.attack_ae import Attack
from cleverhans.compat import reduce_sum, reduce_max
from cleverhans.model import CallableModelWrapper, Model, wrapper_warning_logits
from cleverhans import utils
import tensorflow.contrib.slim as slim
from tensorflow import keras
from keras.layers import Input, Dense, Conv2D, MaxPooling2D, UpSampling2D, BatchNormalization, Activation
from keras.callbacks import EarlyStopping, ModelCheckpoint
from keras.optimizers import Adam
from keras.models import Model
from keras.models import Sequential
from keras.layers import Dense, Dropout, Activation, Flatten
from keras.layers import Conv2D, MaxPooling2D
from keras.models import load_model
from tensorflow.python.framework import meta_graph
from keras import backend as K

np_dtype = np.dtype('float32')
tf_dtype = tf.as_dtype('float32')

_logger = utils.create_logger("cleverhans.attacks.carlini_wagner_l2")
_logger.setLevel(logging.INFO)

class CarliniWagnerAE_Keras(Attack):

  def __init__(self, model, cl_model, sess, dtypestr='float32', **kwargs):
    """
    Note: the model parameter should be an instance of the
    cleverhans.model.Model abstraction provided by CleverHans.
    """
    if not isinstance(cl_model, Model):
      wrapper_warning_logits()
      #cl_model = CallableModelWrapper(cl_model, 'logits')
    if not isinstance(model, Model):
      wrapper_warning_logits()
      #model = CallableModelWrapper(model, 'logits')

    super(CarliniWagnerAE_Keras, self).__init__(model, sess, dtypestr, **kwargs)
    self.cl_model = cl_model
    self.feedable_kwargs = ('y', 'y_target')

    self.structural_kwargs = [
        'batch_size', 'confidence', 'targeted', 'learning_rate',
        'binary_search_steps', 'max_iterations', 'abort_early',
        'initial_const', 'clip_min', 'clip_max'
    ]

  def generate(self, x,x_t, **kwargs):

    assert self.sess is not None, \
        'Cannot use `generate` when no `sess` was provided'
    self.parse_params(**kwargs)

    #labels, nb_classes = self.get_or_guess_labels(x, kwargs)
    nb_classes = 10
    attack = CWL2(self.sess, self.model, self.cl_model, self.batch_size, self.confidence,
                  'x_target' in kwargs, self.learning_rate,
                  self.binary_search_steps, self.max_iterations,
                  self.abort_early, self.initial_const, self.clip_min,
                  self.clip_max, nb_classes,
                  x.get_shape().as_list()[1:])

    def cw_wrap(x_val, x_targ_val):
      return np.array(attack.attack(x_val, x_targ_val), dtype=self.np_dtype)

    wrap = tf.py_func(cw_wrap, [x, x_t], self.tf_dtype)
    wrap.set_shape(x.get_shape())

    return wrap

  def parse_params(self,
                   y=None,
                   y_target = None,
                   x_target=None,
                   batch_size=1,
                   confidence=0,
                   learning_rate=5e-2,
                   binary_search_steps=10,
                   max_iterations=1000,
                   abort_early=False,
                   initial_const=0.5,
                   clip_min=0,
                   clip_max=1):

    # ignore the y and y_target argument
    self.batch_size = batch_size
    self.confidence = confidence
    self.learning_rate = learning_rate
    self.binary_search_steps = binary_search_steps
    self.max_iterations = max_iterations
    self.abort_early = abort_early
    self.initial_const = initial_const
    self.clip_min = clip_min
    self.clip_max = clip_max

def ZERO():
  return np.asarray(0., dtype=np_dtype)

def convert_to_pb(weight_file,input_fld='',output_fld='', model_type = None):

    import os
    import os.path as osp
    from tensorflow.python.framework import graph_util
    from tensorflow.python.framework import graph_io
    from keras.models import load_model

    # weight_file is a .h5 keras model file
    output_node_names_of_input_network = ["pred0"]
    output_node_names_of_final_network = 'output_node'

    # change filename to a .pb tensorflow file
    output_graph_name = weight_file[:-2]+'pb'
    weight_file_path = osp.join(input_fld, weight_file)
    '''
    if(model_type=='AE'):
      input_img = Input(shape=(32, 32, 3))
      x = Conv2D(64, (3, 3), padding='same')(input_img)
      x = BatchNormalization()(x)
      x = Activation('relu')(x)
      x = MaxPooling2D((2, 2), padding='same')(x)
      x = Conv2D(32, (3, 3), padding='same')(x)
      x = BatchNormalization()(x)
      x = Activation('relu')(x)
      x = MaxPooling2D((2, 2), padding='same')(x)
      x = Conv2D(16, (3, 3), padding='same')(x)
      x = BatchNormalization()(x)
      x = Activation('relu')(x)
      encoded = MaxPooling2D((2, 2), padding='same')(x)

      x = Conv2D(16, (3, 3), padding='same')(encoded)
      x = BatchNormalization()(x)
      x = Activation('relu')(x)
      x = UpSampling2D((2, 2))(x)
      x = Conv2D(32, (3, 3), padding='same')(x)
      x = BatchNormalization()(x)
      x = Activation('relu')(x)
      x = UpSampling2D((2, 2))(x)
      x = Conv2D(64, (3, 3), padding='same')(x)
      x = BatchNormalization()(x)
      x = Activation('relu')(x)
      x = UpSampling2D((2, 2))(x)
      x = Conv2D(3, (3, 3), padding='same')(x)
      x = BatchNormalization()(x)
      decoded = Activation('sigmoid')(x)

      net_model = Model(input_img, decoded)
      net_model.load_weights(weight_file_path)

    if(model_type =='Classifier'):
      num_classes = 10
      net_model = Sequential()
      net_model.add(Conv2D(32, (3, 3), padding='same',
                       input_shape=(32,32,3)))
      net_model.add(Activation('relu'))
      net_model.add(Conv2D(32, (3, 3)))
      net_model.add(Activation('relu'))
      net_model.add(MaxPooling2D(pool_size=(2, 2)))
      net_model.add(Dropout(0.25))

      net_model.add(Conv2D(64, (3, 3), padding='same'))
      net_model.add(Activation('relu'))
      net_model.add(Conv2D(64, (3, 3)))
      net_model.add(Activation('relu'))
      net_model.add(MaxPooling2D(pool_size=(2, 2)))
      net_model.add(Dropout(0.25))

      net_model.add(Flatten())
      net_model.add(Dense(512))
      net_model.add(Activation('relu'))
      net_model.add(Dropout(0.5))
      net_model.add(Dense(num_classes))
      net_model.add(Activation('softmax'))

      opt = keras.optimizers.rmsprop(lr=0.0001, decay=1e-6)
      net_model.compile(loss='categorical_crossentropy',
                  optimizer=opt,
                  metrics=['accuracy'])
      net_model.load_weights(weight_file_path)
    '''
    net_model = model.load(weight_file_path)
    #print("model.outputs: ", net_model.outputs)
    #print("model.inputs: ", net_model.inputs)
    num_output = len(output_node_names_of_input_network)
    pred = [None]*num_output
    pred_node_names = [None]*num_output

    for i in range(num_output):
        pred_node_names[i] = output_node_names_of_final_network+str(i)
        pred[i] = tf.identity(net_model.output[i], name=pred_node_names[i])

    sess = K.get_session()

    constant_graph = graph_util.convert_variables_to_constants(sess, sess.graph.as_graph_def(), pred_node_names)
    graph_io.write_graph(constant_graph, output_fld, output_graph_name, as_text=False)
    print('saved the constant graph (ready for inference) at: ', osp.join(output_fld, output_graph_name))

    return output_fld+'/'+output_graph_name

def load_graph(frozen_graph_filename):
    # We load the protobuf file from the disk and parse it to retrieve the
    # unserialized graph_def
    with tf.gfile.GFile(frozen_graph_filename, "rb") as f:
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())

    # Then, we can use again a convenient built-in function to import a graph_def into the
    # current default Graph
    with tf.Graph().as_default() as graph:
        tf.import_graph_def(
            graph_def,
            input_map=None,
            return_elements=None,
            name="prefix",
            op_dict=None,
            producer_op_list=None
        )

    input_name = graph.get_operations()[0].name+':0'
    output_name = graph.get_operations()[-1].name+':0'

    return graph, input_name, output_name
def freeze_session(session, keep_var_names=None, output_names=None, clear_devices=True):

    from tensorflow.python.framework.graph_util import convert_variables_to_constants
    graph = session.graph
    with graph.as_default():
        freeze_var_names = list(set(v.op.name for v in tf.global_variables()).difference(keep_var_names or []))
        output_names = output_names or []
        output_names += [v.op.name for v in tf.global_variables()]
        # Graph -> GraphDef ProtoBuf
        input_graph_def = graph.as_graph_def()
        if clear_devices:
            for node in input_graph_def.node:
                node.device = ""
        frozen_graph = convert_variables_to_constants(session, input_graph_def,
                                                      output_names, freeze_var_names)
        return frozen_graph

class CWL2(object):
  def __init__(self, sess, model,cl_model, batch_size, confidence, targeted,
               learning_rate, binary_search_steps, max_iterations,
               abort_early, initial_const, clip_min, clip_max, num_labels,
               shape):

    self.sess = sess
    self.TARGETED = targeted
    self.LEARNING_RATE = learning_rate
    self.MAX_ITERATIONS = max_iterations
    self.BINARY_SEARCH_STEPS = binary_search_steps
    self.ABORT_EARLY = abort_early
    self.CONFIDENCE = confidence
    self.initial_const = initial_const
    self.batch_size = batch_size
    self.clip_min = clip_min
    self.clip_max = clip_max
    self.model = model
    self.cl_model = cl_model

    #convert model to tensorflow model

    self.repeat = binary_search_steps >= 10

    self.shape = shape = tuple([batch_size] + list(shape))
    #print("shape: ", shape)

    # the variable we're going to optimize over
    modifier = tf.Variable(np.zeros(shape, dtype=np_dtype))

    # these are variables to be more efficient in sending data to tf
    self.timg = tf.Variable(np.zeros(shape), dtype=tf_dtype, name='timg')
    self.targimg = tf.Variable(np.zeros(shape), dtype=tf_dtype, name='targimg')
    #self.tlab = tf.Variable(
     #   np.zeros((batch_size, num_labels)), dtype=tf_dtype, name='tlab')
    self.const = tf.Variable(
        np.zeros(batch_size), dtype=tf_dtype, name='const')

    # and here's what we use to assign them
    self.assign_timg = tf.placeholder(tf_dtype, shape, name='assign_timg')
    self.assign_targimg = tf.placeholder(tf_dtype, shape, name='assign_targimg')
    #self.assign_tlab = tf.placeholder(
     #   tf_dtype, (batch_size, num_labels), name='assign_tlab')
    self.assign_const = tf.placeholder(
        tf_dtype, [batch_size], name='assign_const')

    # the resulting instance, tanh'd to keep bounded from clip_min
    # to clip_max
    self.newimg = (tf.tanh(modifier + self.timg) + 1) / 2
    self.newimg = self.newimg * (clip_max - clip_min) + clip_min

    #targimg_lat = latent_layer_model.predict(self.targimg)
    '''
    tf_model_path_ae = convert_to_pb('cifar10_AE.h5','../cleverhans_tutorials/models','../cleverhans_tutorials/models', 'AE')
    tf_model_path_cl = convert_to_pb('cifar10_CNN.h5','../cleverhans_tutorials/models','../cleverhans_tutorials/models', 'Classifier')
    tf_model,tf_input,tf_output = load_graph(tf_model_path_ae)
    tf_cl_model,tf_cl_input,tf_cl_output = load_graph(tf_model_path_cl)

    #self.x_hat = model.predict(self.newimg, steps = 1)
    with tf.Graph().as_default() as graph1:

      x_hat_output  = tf_model.get_tensor_by_name(tf_output)
      x_hat_input = tf_model.get_tensor_by_name(tf_input)
      #self.x_hat_lat = latent_layer_model.predict(self.newimg)
      #self.x_hat = graph1.run(self.x_hat, feed_dict = {x_1 : self.newimg})
      #self.y_hat_logit = cl_model.predict(self.x_hat, steps = 1)
    with tf.Graph().as_default() as graph2:
      y_hat_logit = tf_cl_model.get_tensor_by_name(tf_cl_output)
      y_hat_logit_input = tf_cl_model.get_tensor_by_name(tf_cl_input)
      #self.y_hat_logit = self.sess.run(self.y_hat_logit, feed_dict = {x_2 : self.x_hat})
      #self.y_hat_logit = cl_model.predict(self.x_hat, steps = 1)

      y_hat_output = tf.argmax(y_hat_logit, axis = 1)

    x_1 = tf.placeholder(tf.float32, (None, 32,32, 3))
    graph = tf.get_default_graph()
    meta_graph1 = tf.train.export_meta_graph(graph=graph1)
    meta_graph.import_scoped_meta_graph(meta_graph1, input_map={'x_hat_input': x_1}, import_scope='graph1',
    out1 = graph.get_tensor_by_name('graph1/tf_output:0'))

    meta_graph2 = tf.train.export_meta_graph(graph=graph2)
    meta_graph.import_scoped_meta_graph(meta_graph2, input_map={'y_hat_logit_input': out1}, import_scope='graph2')
    #self.y_targ_logit = cl_model.predict(self.targimg, steps = 1)
    self.y_targ_logit = tf_cl_model.get_tensor_by_name(tf_cls_output)
    self.y_targ_logit = sess.run(self.y_targ_logit, feed_dict = {tf_cl_model.get_tensor_by_name(tf_cl_input): self.targimg})
    self.y_targ = tf.argmax(self.y_targ_logit, axis = 1)
    '''
    # distance to the input data

    #print("model.outputs: ", model.outputs)
    #print("model.inputs: ", model.inputs)
    frozen_graph = freeze_session(K.get_session(),output_names=[out.op.name for out in model.outputs])
    tf.train.write_graph(frozen_graph, "../cleverhans_tutorials/models", "tf_model_AE.pb", as_text=False)

    from tensorflow.python.platform import gfile

    f = gfile.FastGFile("../cleverhans_tutorials/models/tf_model_AE.pb", 'rb')
    graph_def = tf.GraphDef()
    # Parses a serialized binary message into the current message.
    graph_def.ParseFromString(f.read())
    f.close()

    sess.graph.as_default()
    tf.import_graph_def(graph_def)
    reconstruction_tensor = sess.graph.get_tensor_by_name('import/activation_7/Sigmoid:0')
    #self.x_hat = reconstruction_tensor(self.newimg)

    #self.y_hat_logit = cl_model.predict(self.x_hat, steps=1)
    #self.y_hat = tf.argmax(self.y_hat_logit, axis = 1)
    #self.x_hat = sess.run(reconstruction_tensor, {'import/input_1:0': self.newimg})
    self.other = (tf.tanh(self.timg) + 1) / 2
    self.other =  self.other * (clip_max - clip_min) + clip_min
    self.l2dist = reduce_sum(
        tf.square(self.newimg - self.other), list(range(1, len(shape))))

    print("shape of l2_dist: ", np.shape(self.l2dist))

    epsilon = 10e-8

    loss1 = reduce_sum(tf.square(self.x_hat-self.targimg))

    # sum up the losses
    self.loss2 = reduce_sum(self.l2dist)
    self.loss1 = reduce_sum(self.const * loss1)
    self.loss = self.loss1 + self.loss2

    # Setup the adam optimizer and keep track of variables we're creating
    start_vars = set(x.name for x in tf.global_variables())
    optimizer = tf.train.AdamOptimizer(self.LEARNING_RATE)
    self.train = optimizer.minimize(self.loss, var_list=[modifier])
    end_vars = tf.global_variables()
    new_vars = [x for x in end_vars if x.name not in start_vars]

    # these are the variables to initialize when we run
    self.setup = []
    self.setup.append(self.timg.assign(self.assign_timg))
    self.setup.append(self.targimg.assign(self.assign_targimg))
    #self.setup.append(self.tlab.assign(self.assign_tlab))
    self.setup.append(self.const.assign(self.assign_const))

    self.init = tf.variables_initializer(var_list=[modifier] + new_vars)

  def attack(self, imgs, targ_imgs):

    print("batch_size in attack: ", self.batch_size)
    r = []
    for i in range(0, len(imgs), self.batch_size):
      _logger.debug(
          ("Running CWL2 attack on instance %s of %s", i, len(imgs)))
      r.extend(
          self.attack_batch(imgs[i:i + self.batch_size],
                            targ_imgs[i:i + self.batch_size]))
    return np.array(r)

  def attack_batch(self, imgs, targ_imgs):
    """
    Run the attack on a batch of instance and labels.
    """

    def compare(y1,y2):
      if(y1==y2):
        return True
      else:
        return False

    def compare_dist(recon, orig, targ):

      a = np.sum((recon - orig)**2)
      b = np.sum((recon-targ)**2)
      #if  (tf.math.greater(a,b)) :
      if(a+80>b):
        return True
      else:
        return False

    batch_size = self.batch_size

    #print("batch_size: ", batch_size)

    oimgs = np.clip(imgs, self.clip_min, self.clip_max)
    #oimgs_lat = self.model.get_layer(oimgs, 'LATENT')

    # re-scale instances to be within range [0, 1]
    imgs = (imgs - self.clip_min) / (self.clip_max - self.clip_min)
    imgs = np.clip(imgs, 0, 1)
    # now convert to [-1, 1]
    imgs = (imgs * 2) - 1
    # convert to tanh-space
    imgs = np.arctanh(imgs * .999999)

    # set the lower and upper bounds accordingly
    lower_bound = np.zeros(batch_size)
    CONST = np.ones(batch_size) * self.initial_const
    upper_bound = np.ones(batch_size) * 1e8

    # placeholders for the best l2, score, and instance attack found so far
    o_bestl2 = [1e10] * batch_size
    #o_bestrec = np.copy(oimgs)
    o_bestrec = np.copy(oimgs)
    o_bestattack = np.copy(oimgs)

    for outer_step in range(self.BINARY_SEARCH_STEPS):
      # completely reset adam's internal state.
      self.sess.run(self.init)
      batch = imgs[:batch_size]
      batchtarg = targ_imgs[:batch_size]

      bestl2 = [1e10] * batch_size
      #bestrec = np.copy(oimgs)
      bestrec = np.copy(oimgs)

      _logger.debug("  Binary search step %s of %s",
                    outer_step, self.BINARY_SEARCH_STEPS)

      # The last iteration (if we run many steps) repeat the search once.
      if self.repeat and outer_step == self.BINARY_SEARCH_STEPS - 1:
        CONST = upper_bound

      # set the variables so that we don't have to send them over again
      self.sess.run(
          self.setup, {
              self.assign_timg: batch,
              self.assign_targimg: batchtarg,
              self.assign_const: CONST
          })

      prev = 1e8
      for iteration in range(self.MAX_ITERATIONS):
        # perform the attack
        _, l, l2s, nrec, nimg, yhat, ytarg = self.sess.run([
            self.train, self.loss, self.l2dist, self.x_hat,
            self.newimg, self.y_hat, self.y_targ
        ])
        #print("shape of yhat: ", np.shape(yhat))
        if iteration % ((self.MAX_ITERATIONS // 10) or 1) == 0:
          _logger.debug(("    Iteration {} of {}: loss={:.3g} " +
                         "l2={:.3g} ").format(
                             iteration, self.MAX_ITERATIONS, l,
                             np.mean(l2s)))

        # check if we should abort search if we're getting nowhere.
        if self.ABORT_EARLY and \
           iteration % ((self.MAX_ITERATIONS // 10) or 1) == 0:
          if l > prev * .9999:
            msg = "    Failed to make progress; stop early"
            _logger.debug(msg)
            break
          prev = l

        # adjust the best result found so far
        for e, (l2, nr, ii, yh, yt) in enumerate(zip(l2s, nrec, nimg, yhat, ytarg)):
          #lab = np.argmax(batchlab[e])
          if l2 < bestl2[e] and compare(yh,yt):
          #if l2<bestl2[e] and compare_dist(nr, imgs[e], targ_imgs[e]):
            bestl2[e] = l2
            bestrec[e] = nr
          if l2 < o_bestl2[e] and compare(yh,yt):
          #if l2 < o_bestl2[e] and compare_dist(nr, imgs[e], targ_imgs[e]):
            o_bestl2[e] = l2
            o_bestrec[e] = nr
            o_bestattack[e] = ii

      # adjust the constant as needed
      for e in range(batch_size):
        if compare(yhat[e], ytarg[e]):
        #if compare_dist(nrec[e], imgs[e], targ_imgs[e]):
          upper_bound[e] = min(upper_bound[e], CONST[e])
          if upper_bound[e] < 1e9:
            CONST[e] = (lower_bound[e] + upper_bound[e]) / 2
        else:
          # failure, either multiply by 10 if no solution found yet
          #          or do binary search with the known upper bound
          lower_bound[e] = max(lower_bound[e], CONST[e])
          if upper_bound[e] < 1e9:
            CONST[e] = (lower_bound[e] + upper_bound[e]) / 2
          else:
            CONST[e] *= 10
      _logger.debug("  Successfully generated adversarial examples " +
                    "on {} of {} instances.".format(
                        sum(upper_bound < 1e9), batch_size))
      o_bestl2 = np.array(o_bestl2)
      mean = np.mean(np.sqrt(o_bestl2[o_bestl2 < 1e9]))
      _logger.debug("   Mean successful distortion: {:.4g}".format(mean))

    # return the best solution found
    o_bestl2 = np.array(o_bestl2)
    #print("o_bestl2: ", o_bestl2)
    #print("shape of o_bestattack: ", np.shape(o_bestattack))
    return o_bestattack


import numpy as np
from secrethitlergame.players import SecretHitlerPlayer

class SecretHitlerState():

    def __init__(self, possible_lib_cards=6, n_cards=15,
                 n_players=7, n_fasc=2, names=None):
        self.turn = 1

        self.possible_liberals = possible_lib_cards
        self.cards_in_deck = n_cards
        self.possible_fasc = self.cards_in_deck - self.possible_liberals
        self.played_cards = []
        self.chanc_claims = []
        self.pres_claims = []
        self.vote_results = []

        self.governments = []

        self.n_liberals_left = self.possible_liberals
        self.n_fasc_left = self.possible_fasc

        self.cards_drawn_on_gov = 3

        self.n_players = n_players

        self.n_alive_players = self.n_players

        self.standstill_count = 0

        self.pres_seat = 0

        if names is None:
            self.players = {str(x): SecretHitlerPlayer(str(x), x)
                            for x in range(self.n_players)}
        else:
            self.players = {name: SecretHitlerPlayer(name, seat)
                            for seat, name in names.items()}
            assert len(self.players) == n_players

    @property
    def chanc(self):
        last_vote = self.vote_results[-1]
        if last_vote['result'] == 'Passed':
            chanc = last_vote['chanc']
        else:
            chanc = None
        return chanc

    @property
    def pres(self):
        for player in self.players.values():
            if player.seat == self.pres_seat:
                return player

    def discovered_card(self, card):
        if card.upper() == 'L':
            self.discovered_liberal()
        else:
            self.discovered_fasc()

    def discovered_liberal(self):
        self.n_liberals_left -= 1

    def discovered_fasc(self):
        self.n_fasc_left -= 1

    def determine_vote(self, chanc):
        votes = [player_obj.vote_info[-1]['vote']
                 for player_obj in self.players.values()]
        pct_in_favor = np.mean([vote.upper() == 'JA' for vote in votes])

        result_dict = {
            'turn': self.turn,
            'in_favor': [player_obj for player_obj in self.players.values()
                         if player_obj.vote_info[-1]['vote'].upper() == 'JA'],
            'against': [player_obj for player_obj in self.players.values()
                        if player_obj.vote_info[-1]['vote'].upper() != 'JA'],
            'pres': self.pres,
            'chanc': chanc
        }

        if pct_in_favor > .5:
            result_dict.update({'result': 'Passed'})
        else:
            result_dict.update({'result': 'Failed'})

        self.vote_results.append(result_dict)

    def update_claims(self):
        self.pres_claims.append({self.pres: self.pres.move_info[-1]})
        self.chanc_claims.append({self.chanc: self.chanc.move_info[-1]})

    def move_pres(self):
        self.turn += 1
        self.pres_seat += 1

    def after_failed_government(self):
        self.standstill_count += 1
        self.move_pres()

    def enact(self, card):
        self.played_cards.append(card)

    def after_successful_gov(self, played_card):
        self.enact(played_card)
        self.begin_enactment_process()
        self.governments.append([{'turn': self.turn,
                                  'pres': self.pres,
                                  'chanc': self.chanc,
                                  'result': self.played_cards[-1]}])
        self.update_claims()
        self.standstill_count = 0
        self.move_pres()
        self.cards_in_deck -= self.cards_drawn_on_gov

    def determine_course_of_action(self, played_card=None):
        last_vote = self.vote_results[-1]

        if last_vote['result'] == 'Passed':
            self.after_successful_gov(played_card)
        else:
            self.push_to_next_president()

    def begin_enactment_process(self):
        last_card = self.played_cards[-1]
        self.discovered_card(last_card)

    def push_to_next_president(self):
        self.pres_seat += 1
        if self.standstill_count < 3:
            self.standstill_count += 1
        else:
            self.topdeck()

    def topdeck(self):
        p_lib = self.n_liberals_left/self.cards_in_deck
        p_fasc = 1 - p_lib
        top_decked_card = np.random.choice(['L', 'F'], p=[p_lib, p_fasc])

        if top_decked_card == 'L':
            self.discovered_liberal()
        else:
            self.discovered_fasc()
        self.cards_in_deck -= 1
        self.standstill_count = 0


import cv2
import numpy as np
from collections import Counter
from PIL import ImageGrab
import os
import time

def screenGrab():
    sc = ImageGrab.grab()
    name = "screenshot"+str(int(time.time()))+".png"
    sc.save(os.getcwd() + "\\screenshots\\" +name , "PNG")
    print(name)
    return name

itemlist = ['bf_sword.png', 'bloodthirster.png', 'botrk.png', 'cursed_blade.png', 'darkin.png', 'divine.png', 'dragons_claw.png', 'firecannon.png', 'force_of_nature.png', 'frozen_heart.png', 'frozen_mallet.png', 'ghostblade.png', 'guardian_angel.png', 'gunblade.png', 'hush.png', 'infinity_edge.png', 'ionic_spark.png', 'knights_vow.png', 'locket.png', 'ludens_echo.png', 'morellonomicon.png', 'phantom_dancer.png', 'rabadons.png', 'rageblade.png', 'redemption.png', 'red_buff.png', 'runaans_hurricane.png', 'seraphs.png', 'shojin.png', 'statikk_shiv.png', 'sword_breaker.png', 'thornmail.png', 'titanic_hydra.png', 'warmogs.png', 'yuumi.png', 'zekes_herald.png', 'zephyr.png']
rawitemlist = ['chain_vest.png', 'giants_belt.png', 'golden_spatula.png', 'needlessly.png', 'negatron_cloak.png', 'recurve_bow.png', 'tear.png']
someitemslist = itemlist + rawitemlist

while True:
    itemsowned = []

    mainimage = "screenshots\\"+screenGrab()
    img = cv2.imread(mainimage)

    w_screenshot, h_screenshot = img.shape[:-1]

    w = round(w_screenshot/(1080/22))
    h = round(h_screenshot/(1920/22))

    for item in someitemslist:
        queryimage = item
        query1 = cv2.imread("items/"+queryimage)

        #w1/w = 1080/21  so we resize query for the screenshot
        query = cv2.resize(query1, (w, h))
        #cv2.imshow('resized query', query)

        result = cv2.matchTemplate(img, query, cv2.TM_CCOEFF_NORMED)
        threshold = 0.72
        loc = np.where(result >= threshold)

        for pt in zip(*loc[::-1]):
            cv2.rectangle(img, pt, (pt[0] + w, pt[1] + h),(0,255,255),2)
            #print("found", x)
            itemsowned.append(item)

        #cv2.imshow('detected', img)
        #cv2.waitKey(0)
        #cv2.destroyWindow("detected")

        img = cv2.imread(mainimage)

    print(Counter(itemsowned))
    #time.sleep(3)


#!/usr/bin/env python

import os
import astropy.units as u
import argparse
import pandas as pd
import numpy as np

from gammapy.spectrum import cosmic_ray_flux, CrabSpectrum

from pyirf.io.io import load_config
from pyirf.perf import (CutsOptimisation,
                        CutsDiagnostic,
                        CutsApplicator,
                        IrfMaker,
                        SensitivityMaker,
                        )

def main():
    # Read arguments
    parser = argparse.ArgumentParser(description='Make performance files')
    parser.add_argument('--config_file', type=str, required=True, help='')
    parser.add_argument(
        '--obs_time',
        type=str,
        required=True,
        help='Observation time, should be given as a string, value and astropy unit separated by a dot. E.g:`50.h`'
    )
    mode_group = parser.add_mutually_exclusive_group()
    mode_group.add_argument('--wave', dest="mode", action='store_const',
                            const="wave", default="tail",
                            help="if set, use wavelet cleaning")
    mode_group.add_argument('--tail', dest="mode", action='store_const',
                            const="tail",
                            help="if set, use tail cleaning, otherwise wavelets")
    args = parser.parse_args()

    # Read configuration file
    cfg = load_config(args.config_file)

    # Add obs. time in configuration file
    str_obs_time = args.obs_time.split('.')
    cfg['analysis']['obs_time'] = {'value': float(str_obs_time[0]), 'unit': str(str_obs_time[-1])}
    print(cfg['analysis']['obs_time'])

    # Create output directory if necessary
    outdir = os.path.join(cfg['general']['outdir'], 'irf_{}_ThSq_{}_Time{:.2f}{}'.format(
        args.mode,
        cfg['analysis']['thsq_opt']['type'],
        cfg['analysis']['obs_time']['value'],
        cfg['analysis']['obs_time']['unit'])
    )
    if not os.path.exists(outdir):
        os.makedirs(outdir)

    indir = cfg['general']['indir']
    template_input_file = cfg['general']['template_input_file']

    # Load data
    particles = ['gamma', 'electron', 'proton']
    evt_dict = dict()  # Contain DL2 file for each type of particle
    for particle in particles:
        # template looks like dl2_{}_{}_merged.h5
        infile = os.path.join(indir, template_input_file.format(args.mode, particle))
        evt_dict[particle] = pd.read_hdf(infile, key='reco_events')

    # Apply offset cut to proton and electron
    for particle in ['electron', 'proton']:
        # print('Initial stat: {} {}'.format(len(evt_dict[particle]), particle))
        evt_dict[particle] = evt_dict[particle].query('offset <= {}'.format(
            cfg['particle_information'][particle]['offset_cut'])
        )

    # Add required data in configuration file for future computation
    for particle in particles:
        n_files = cfg['particle_information'][particle]['n_files']
        print(f"{n_files} files for {particle}")
        cfg['particle_information'][particle]['n_files'] = \
            len(np.unique(evt_dict[particle]['obs_id']))
        cfg['particle_information'][particle]['n_simulated'] = \
            cfg['particle_information'][particle]['n_files'] * cfg['particle_information'][particle]['n_events_per_file']

    # Define model for the particles
    model_dict = {'gamma': CrabSpectrum('hegra').model,
                  'proton': cosmic_ray_flux,
                  'electron': cosmic_ray_flux}

    # Reco energy binning
    cfg_binning = cfg['analysis']['ereco_binning']
    ereco = np.logspace(np.log10(cfg_binning['emin']),
                        np.log10(cfg_binning['emax']),
                        cfg_binning['nbin'] + 1) * u.TeV

    # Handle theta square cut optimisation
    # (compute 68 % containment radius PSF if necessary)
    thsq_opt_type = cfg['analysis']['thsq_opt']['type']
    if thsq_opt_type in 'fixed':
        thsq_values = np.array([cfg['analysis']['thsq_opt']['value']]) * u.deg
        print('Using fixed theta cut: {}'.format(thsq_values))
    elif thsq_opt_type in 'opti':
        thsq_values = np.arange(0.05, 0.40, 0.01) * u.deg
        print('Optimising theta cut for: {}'.format(thsq_values))
    elif thsq_opt_type in 'r68':
        print('Using R68% theta cut')
        print('Computing...')
        cfg_binning = cfg['analysis']['ereco_binning']
        ereco = np.logspace(np.log10(cfg_binning['emin']),
                            np.log10(cfg_binning['emax']),
                            cfg_binning['nbin'] + 1) * u.TeV
        radius = 68

        thsq_values = list()
        for ibin in range(len(ereco) - 1):
            emin = ereco[ibin]
            emax = ereco[ibin + 1]

            energy_query = 'reco_energy > {} and reco_energy <= {}'.format(
                emin.value, emax.value
            )
            data = evt_dict['gamma'].query(energy_query).copy()

            min_stat = 0
            if len(data) <= min_stat:
                print('  ==> Not enough statistics:')
                print('To be handled...')
                thsq_values.append(0.3)
                continue
                # import sys
                # sys.exit()

            psf = np.percentile(data['offset'], radius)
            psf_err = psf / np.sqrt(len(data))

            thsq_values.append(psf)
        thsq_values = np.array(thsq_values) * u.deg
        # Set 0.05 as a lower value
        idx = np.where(thsq_values.value < 0.05)
        thsq_values[idx] = 0.05 * u.deg
        print('Using theta cut: {}'.format(thsq_values))

    # Cuts optimisation
    print('### Finding best cuts...')
    cut_optimiser = CutsOptimisation(
        config=cfg,
        evt_dict=evt_dict,
        verbose_level=0
    )

    # Weight events
    print('- Weighting events...')
    cut_optimiser.weight_events(
        model_dict=model_dict,
        colname_mc_energy=cfg['column_definition']['mc_energy']
    )

    # Find best cutoff to reach best sensitivity
    print('- Estimating cutoffs...')
    cut_optimiser.find_best_cutoff(energy_values=ereco, angular_values=thsq_values)

    # Save results and auxiliary data for diagnostic
    print('- Saving results to disk...')
    cut_optimiser.write_results(
        outdir, '{}.fits'.format(cfg['general']['output_table_name']),
       format='fits'
    )

    # Cuts diagnostic
    print('### Building cut diagnostics...')
    cut_diagnostic = CutsDiagnostic(config=cfg, indir=outdir)
    cut_diagnostic.plot_optimisation_summary()
    cut_diagnostic.plot_diagnostics()

    # Apply cuts and save data
    print('### Applying cuts to data...')
    cut_applicator = CutsApplicator(config=cfg, evt_dict=evt_dict, outdir=outdir)
    cut_applicator.apply_cuts()

    # Irf Maker
    print('### Building IRF...')
    irf_maker = IrfMaker(config=cfg, evt_dict=evt_dict, outdir=outdir)
    irf_maker.build_irf()

    # Sensitivity maker
    print('### Estimating sensitivity...')
    sensitivity_maker = SensitivityMaker(config=cfg, outdir=outdir)
    sensitivity_maker.load_irf()
    sensitivity_maker.estimate_sensitivity()

if __name__ == '__main__':
    main()


import itertools
import torch
import torch.nn as nn
import torch.nn.functional as func
import os
import glob
import numpy as np
import utils.torchvision_wrappers as models_wrappers
import lib.hed_pytorch.hed as hed
import torch.distributed as dist

class MoCo(nn.Module):
    """
    Build a MoCo model with: a query encoder, a key encoder, and a queue
    https://arxiv.org/abs/1911.05722
    """
    def __init__(self, base_encoder, dim=128, K=65536, m=0.999, T=0.07, mlp=False, debug=False, args=None):
        """
        dim: feature dimension (default: 128)
        K: queue size; number of negative keys (default: 65536)
        m: moco momentum of updating key encoder (default: 0.999)
        T: softmax temperature (default: 0.07)
        """
        super(MoCo, self).__init__()

        self.debug = debug
        self.args = args

        self.K = K
        self.m = m
        self.T = T

        # create the encoders
        # num_classes is the output fc dimension
        self.encoder_q = base_encoder(num_classes=dim, pretrained=args.imagenet_pretrained)
        self.encoder_k = base_encoder(num_classes=dim, pretrained=args.imagenet_pretrained)

        if mlp:  # hack: brute-force replacement

            dim_mlp = self.encoder_q.fc.weight.shape[1]
            self.encoder_q.fc = nn.Sequential(nn.Linear(dim_mlp, dim_mlp), nn.ReLU(), self.encoder_q.fc)
            self.encoder_k.fc = nn.Sequential(nn.Linear(dim_mlp, dim_mlp), nn.ReLU(), self.encoder_k.fc)

        for param_q, param_k in zip(self.encoder_q.parameters(), self.encoder_k.parameters()):
            param_k.data.copy_(param_q.data)  # initialize
            param_k.requires_grad = False  # not update by gradient

        # create the queue
        self.qMult = 1
        if self.args.multi_q:
            self.qMult = len(self.args.data.split(','))
            self.register_buffer("queue", torch.randn(self.qMult, dim, K))
            self.queue = nn.functional.normalize(self.queue, dim=1)
            self.register_buffer("queue_idx", - torch.ones(self.qMult, K))
        else:
            self.register_buffer("queue", torch.randn(dim, K))
            self.queue = nn.functional.normalize(self.queue, dim=0)
            self.register_buffer("queue_idx", - torch.ones(K))

        self.Q_lim = [None] * self.qMult
        self.register_buffer("queue_ptr", torch.zeros(1, dtype=torch.long))

    @torch.no_grad()
    def _momentum_update_key_encoder(self):
        """
        Momentum update of the key encoder
        """
        for param_q, param_k in zip(self.encoder_q.parameters(), self.encoder_k.parameters()):
            param_k.data = param_k.data * self.m + param_q.data * (1. - self.m)

    @torch.no_grad()
    def _dequeue_and_enqueue(self, keys, q_selector=None, sample_idx=None):
        # gather keys before updating queue
        if not self.debug:
            keys = concat_all_gather(keys)
            if q_selector is not None:
                q_selector = concat_all_gather(q_selector)
            if sample_idx is not None:
                sample_idx = concat_all_gather(sample_idx)
        # also for simplicity, here we assume each domain contributes the same number of samples to the batch
        batch_size = keys.shape[0] // self.qMult

        ptr = int(self.queue_ptr)
        assert self.K % batch_size == 0, f'self.K={self.K}, batch_size={batch_size}'  # for simplicity

        # replace the keys at ptr (dequeue and enqueue)
        if q_selector is not None:
            for iQ in range(self.qMult):
                if self.Q_lim[iQ] is not None:
                    local_ptr = ptr % self.Q_lim[iQ]
                else:
                    local_ptr = ptr
                active_keys = keys[q_selector == iQ]

                self.queue[iQ, :, local_ptr:local_ptr + batch_size] = active_keys.T
                self.queue_idx[iQ, local_ptr:local_ptr + batch_size] = sample_idx[q_selector == iQ]
        else:
            if self.Q_lim[0] is not None:
                local_ptr = ptr % self.Q_lim[0]
            else:
                local_ptr = ptr

            self.queue[:, local_ptr:local_ptr + batch_size] = keys.T
            self.queue_idx[local_ptr:local_ptr + batch_size] = sample_idx

        ptr = (ptr + batch_size) % self.K  # move pointer
        self.queue_ptr[0] = ptr

    @torch.no_grad()
    def _batch_shuffle_ddp(self, x, ix=None):
        """
        Batch shuffle, for making use of BatchNorm.
        *** Only support DistributedDataParallel (DDP) model. ***
        """
        # gather from all gpus
        batch_size_this = x.shape[0]
        x_gather = concat_all_gather(x)
        if ix is not None:
            ix_gather = concat_all_gather(ix)
        batch_size_all = x_gather.shape[0]

        num_gpus = batch_size_all // batch_size_this

        # random shuffle index
        idx_shuffle = torch.randperm(batch_size_all).cuda()

        # broadcast to all gpus
        dist.broadcast(idx_shuffle, src=0)

        # index for restoring
        idx_unshuffle = torch.argsort(idx_shuffle)

        # shuffled index for this gpu
        gpu_idx = dist.get_rank()
        idx_this = idx_shuffle.view(num_gpus, -1)[gpu_idx]

        ix_ret = None
        if ix is not None:
            ix_ret = ix_gather[idx_this]

        return x_gather[idx_this], ix_ret, idx_unshuffle

    @torch.no_grad()
    def _batch_unshuffle_ddp(self, x, idx_unshuffle):
        """
        Undo batch shuffle.
        *** Only support DistributedDataParallel (DDP) model. ***
        """
        # gather from all gpus
        batch_size_this = x.shape[0]
        x_gather = concat_all_gather(x)
        batch_size_all = x_gather.shape[0]

        num_gpus = batch_size_all // batch_size_this

        # restored index for this gpu
        gpu_idx = dist.get_rank()
        idx_this = idx_unshuffle.view(num_gpus, -1)[gpu_idx]

        return x_gather[idx_this]

    def forward(self, im_q, im_k, isedge_q=None, isedge_k=None, q_selector=None, sample_idx=None, sample_pth=None):
        """
        Input:
            im_q: a batch of query images
            im_k: a batch of key images
            isedge_q: binary index, 0 = image, 1 = edge
            isedge_k: binary index, 0 = image, 1 = edge
        Output:
            logits, targets
        """

        extra_outputs = {}

        # compute query features
        q = self.encoder_q(im_q)  # queries: NxC

        if isinstance(q, tuple) or isinstance(q, list):  # handle the case features are returned too
            extra_outputs['q_fm'] = nn.functional.normalize(q[2], dim=1)
            extra_outputs['q'] = nn.functional.normalize(torch.mean(torch.mean(q[2], dim=3), dim=2), dim=1)
            q = q[0]
        q = nn.functional.normalize(q, dim=1)
        # compute key features
        with torch.no_grad():  # no gradient to keys
            self._momentum_update_key_encoder()  # update the key encoder

            # shuffle for making use of BN
            if not self.debug:
                im_k, isedge_k, idx_unshuffle = self._batch_shuffle_ddp(im_k, isedge_k)

            k = self.encoder_k(im_k)  # keys: NxC

            if isinstance(k, tuple) or isinstance(k, list):  # handle the case features are returned too
                k = k[0]
            k = nn.functional.normalize(k, dim=1)

            # undo shuffle
            if not self.debug:
                k = self._batch_unshuffle_ddp(k, idx_unshuffle) # no need to unshuffle the "isedge_k" as they are no longer used beyond this point

        # compute logits
        # Einstein sum is more intuitive
        # positive logits: Nx1
        l_pos = torch.einsum('nc,nc->n', [q, k]).unsqueeze(-1)
        # negative logits: NxK
        BIG_NUMBER = 10000.0 * self.T  # taking temp into account for a good measure
        l_neg = (torch.zeros((q.shape[0], self.K)).cuda() - BIG_NUMBER)
        if self.args.multi_q:
            for iQ in range(self.qMult):
                if self.Q_lim[iQ] is not None:
                    qlim = self.Q_lim[iQ]
                else:
                    qlim = self.queue[iQ].shape[1]
                ixx = (q_selector == iQ)
                _l_neg = torch.einsum('nc,ck->nk', [q[ixx], self.queue[iQ][:,:qlim].clone().detach()])
                if sample_idx is not None:
                    for ii, indx in enumerate(sample_idx[ixx]):
                        _l_neg[ii, self.queue_idx[iQ][:qlim] == indx] = - BIG_NUMBER
                l_neg[ixx, :qlim] = _l_neg
        else:
            if self.Q_lim[0] is not None:
                qlim = self.Q_lim[0]
            else:
                qlim = self.queue.shape[1]
            _l_neg = torch.einsum('nc,ck->nk', [q, self.queue[:, :qlim].clone().detach()])
            if sample_idx is not None:
                for ii in range(q.shape[0]):
                    _l_neg[ii, self.queue_idx[:qlim] == sample_idx[ii]] = - BIG_NUMBER
            l_neg[:, :qlim] = _l_neg

        # logits: Nx(1+K)
        logits = torch.cat([l_pos, l_neg], dim=1)

        # apply temperature
        logits /= self.T

        # labels: positive key indicators
        labels = torch.zeros(logits.shape[0], dtype=torch.long).cuda()

        # dequeue and enqueue
        self._dequeue_and_enqueue(k, q_selector if self.args.multi_q else None, sample_idx)

        return logits, labels, extra_outputs

    def params_lock(self):
        for p in self.parameters():
            p.requires_grad_b_ = p.requires_grad
            p.requires_grad = False

    def params_unlock(self):
        for p in self.parameters():
            p.requires_grad = p.requires_grad_b_

# edge detector net
class HedNet(nn.Module):
    def __init__(self, args=None, train_hed=None, pretrained_hed=True):
        super(HedNet, self).__init__()
        self.args = args
        if train_hed is None:
            train_hed = args.train_hed
        self.edgeMean = [0.5, 0.5, 0.5]
        self.edgeSTD = [0.229, 0.224, 0.225]
        self.edgeNet = hed.Network(pretrained_hed)
        if not train_hed:
            self.edgeNet.eval()
            for param_edge in self.edgeNet.parameters():
                param_edge.requires_grad = False  # not update by gradient

    def denormalize(self, im, isedge):
        return im[isedge > 0, 0] * self.edgeSTD[0] + self.edgeMean[0]

    def normalize(self, edge):
        return torch.cat([
            ((edge - self.edgeMean[0]) / self.edgeSTD[0]),
            ((edge - self.edgeMean[1]) / self.edgeSTD[1]),
            ((edge - self.edgeMean[2]) / self.edgeSTD[2])
        ], 1)

    def normalized_edges(self, im, isedge):
        im_out = im.clone()
        isedge = isedge.squeeze()
        edge = self.edgeNet(im_out[isedge > 0, :]).clamp(0.0, 1.0)
        if self.args.stretch_edge:
            edge -= edge.min(dim=1)[0].min(dim=1)[0][:, None, None]
            edge /= (edge.max(dim=1)[0].max(dim=1)[0][:, None, None] + 1e-14)
        # normalizing edge images
        im_out[isedge > 0, :] = self.normalize(edge)
        return im_out

    def forward(self, im_q, im_k, isedge_q, isedge_k):
        im_q_out = self.normalized_edges(im_q, isedge_q)
        im_k_out = self.normalized_edges(im_k, isedge_k)
        return im_q_out, im_k_out

class DomainDiscriminator(nn.Module):
    def __init__(self, args=None):
        super(DomainDiscriminator, self).__init__()
        self.args = args
        ddisc_layers = [int(x) for x in args.ddisc_layers.split(',')]
        layers = []
        prev_n_f = 2048  # Todo: Why hard coded 2048??
        for i_l, n_f in enumerate(ddisc_layers):
            if i_l != 0:
                layers.append(torch.nn.Dropout(p=0.25))
            layers.append(torch.nn.Linear(prev_n_f, n_f))
            layers.append(torch.nn.LeakyReLU())
            prev_n_f = n_f
        layers.append(torch.nn.Linear(prev_n_f, len(args.data.split(','))))
        self.net = torch.nn.Sequential(*layers)

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

    def params_lock(self):
        for p in self.parameters():
            p.requires_grad_b_ = p.requires_grad
            p.requires_grad = False

    def params_unlock(self):
        for p in self.parameters():
            p.requires_grad = p.requires_grad_b_

# utils
@torch.no_grad()
def concat_all_gather(tensor):
    """
    Performs all_gather operation on the provided tensors.
    *** Warning ***: torch.distributed.all_gather has no gradient.
    """
    tensors_gather = [torch.ones_like(tensor) for _ in range(dist.get_world_size())]
    dist.all_gather(tensors_gather, tensor, async_op=False)
    output = torch.cat(tensors_gather, dim=0)
    return output

@torch.no_grad()
def concat_all_gather_interlace(tensor):
    """
    Performs all_gather operation on the provided tensors.
    *** Warning ***: torch.distributed.all_gather has no gradient.
    """
    tensors_gather = [torch.ones_like(tensor)
        for _ in range(dist.get_world_size())]
    dist.all_gather(tensors_gather, tensor, async_op=False)
    output = torch.cat(tensors_gather, dim=0)
    inds = [torch.arange(i, len(g)*len(tensors_gather), len(tensors_gather)) for i, g in enumerate(tensors_gather)]
    inds = torch.cat(inds, dim=0)
    sort_inds = torch.argsort(inds)
    return output[sort_inds]

def concat_all_gather_object_interlace(obj):
    """
    Performs all_gather operation on the provided tensors.
    *** Warning ***: torch.distributed.all_gather has no gradient.
    """
    obj_gather = [None for _ in range(dist.get_world_size())]
    dist.all_gather_object(obj_gather, obj)
    if len(obj_gather[0]) > 1:
        return list(filter(None, itertools.chain(*itertools.zip_longest(*obj_gather))))
    else:
        return obj_gather

def load_model(args, return_epoch=False):
    # create model
    print("=> creating model '{}'".format(args.arch))

    model = models_wrappers.__dict__[args.arch](num_classes=args.moco_dim, pretrained=args.imagenet_pretrained)

    # freeze all layers but the last fc
    for name, param in model.named_parameters():
        if name not in ['fc.weight', 'fc.bias']:
            param.requires_grad = False

    # load from pre-trained, before DistributedDataParallel constructor
    checkpoint = {}
    if args.resume:
        if not os.path.isfile(args.resume):
            if args.resume != 'not a path':
                print("=> no checkpoint found at '{}'".format(args.resume))

            if os.path.isfile(os.path.join(args.work_folder, 'checkpoint_last.pth.tar')):
                args.resume = os.path.join(args.work_folder, 'checkpoint_last.pth.tar')
            else:
                gpp = os.path.join(args.work_folder, 'checkpoint_*.pth.tar')
                cpts = glob.glob(gpp)
                cpts_ix = [int(x.split('/')[-1].split('_')[1].split('.')[0]) for x in cpts]
                if len(cpts_ix) > 0:
                    mx_ix = np.argmax(cpts_ix)
                    args.resume = cpts[mx_ix]
                else:
                    print('=> no checkpoints found')

        path2load = args.resume

        if os.path.isfile(path2load):
            print("=> loading checkpoint '{}'".format(path2load))
            checkpoint = torch.load(path2load, map_location="cpu")

            # rename moco pre-trained keys
            state_dict = checkpoint['state_dict']
            for k in list(state_dict.keys()):
                # retain only encoder_q up to before the embedding layer
                if k.startswith('module.encoder_q'): # and not k.startswith('module.encoder_q.fc'):
                    # remove prefix
                    state_dict[k[len("module.encoder_q."):]] = state_dict[k]
                # delete renamed or unused k
                del state_dict[k]

            if args.mlp:
                dim_mlp = model.fc.weight.shape[1]
                model.fc = nn.Sequential(nn.Linear(dim_mlp, dim_mlp), nn.ReLU(), model.fc)

            args.start_epoch = 0
            msg = model.load_state_dict(state_dict, strict=False)
            # assert set(msg.missing_keys) == {"fc.weight", "fc.bias"}
            print(f'=> missing keys: {msg.missing_keys}')

            print("=> loaded pre-trained model '{}'".format(path2load))
        else:
            print("=> no checkpoint found at '{}'".format(path2load))

    # For multiprocessing distributed, DistributedDataParallel constructor
    # should always set the single device scope, otherwise,
    # DistributedDataParallel will use all available devices.
    if args.gpu is not None:
        torch.cuda.set_device(args.gpu)
        model.cuda(args.gpu)
        model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu])
    else:
        model.cuda()
        # DistributedDataParallel will divide and allocate batch_size to all
        # available GPUs if device_ids are not set
        model = torch.nn.parallel.DistributedDataParallel(model)

    model.eval()

    if return_epoch:
        return model, checkpoint.get('epoch', 0)
    else:
        return model, 0


'''Emulates the IsMore ReHand application (accepts velocity commands and
sends feedback data over UDP) using the simulated ReHand.
'''

import time
import numpy as np
from math import sin, cos
import rehand
import socket
import select

from ismore import settings
from utils.constants import *

MAX_MSG_LEN = 200  # characters

feedback_freq = 200  # Hz
feedback_period = 1./feedback_freq  # secs

sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
sock.bind(settings.REHAND_UDP_SERVER_ADDR)

# create and start ReHand object
rh = rehand.ReHand(tstep=0.005)
rh.daemon = True
rh.start()

starting_pos = settings.starting_pos[['rh_pthumb', 'rh_pindex', 'rh_pfing3', 'rh_pprono']]
rh._set_pos(np.mat(starting_pos))
print 'setting ReHand starting position in sim app'

t_last_feedback = time.time()

# only start counting after first command is received
n_feedback_packets_sent = 0
received_first_cmd = False

# Andreas: Print information only every at 0.5 Hz second
time_to_print = False
time_old = time.time()
print_buffer = []

while True:
    if time.time() - time_old >= .5:
        time_to_print = True

    if time_to_print:
        for item in print_buffer:
            print(item)

        print_buffer = []
        time_old = time.time()

    # check if there is data available to receive before calling recv
    r, _, _ = select.select([sock], [], [], 0)
    if r:  # if the list r is not empty
        command = sock.recv(MAX_MSG_LEN) # the command string
        print_buffer.append('received command: {}'.format(command.rstrip('\r')))
        if not received_first_cmd:
            received_first_cmd = True
            t_received_first_cmd = time.time()

        items = command.rstrip('\r').split(' ')
        cmd_id = items[0]
        dev_id = items[1]
        data_fields = items[2:]

        #print 'cmd_id', cmd_id
        #print 'dev_id', dev_id
        #print 'equal to ReHand?', dev_id == 'ReHand'
        #print 'len(dev_id)', len(dev_id)
        #print command
        #print items

        assert dev_id == 'ReHand'

        if cmd_id == 'SetSpeed':
            des_thumb_vel = float(data_fields[0]) * deg_to_rad  # convert from deg/s to rad/s
            des_index_vel = float(data_fields[1]) * deg_to_rad  # convert from deg/s to rad/s
            des_fing3_vel = float(data_fields[2]) * deg_to_rad  # convert from deg/s to rad/s
            des_prono_vel = float(data_fields[3]) * deg_to_rad  # convert from deg/s to rad/s
            des_vel = np.mat([des_thumb_vel, des_index_vel, des_fing3_vel, des_prono_vel]).T

            rh.set_vel(des_vel)

    t = time.time()
    if t - t_last_feedback > feedback_period:
        # print 'sending feedback'
        state = rh.get_state()
        pos = np.array(state['pos']).reshape((4,))
        vel = np.array(state['vel']).reshape((4,))
        ts = int(time.time() * 1e6)  # in microseconds

        # convert from rad to deg (and rad/s to deg/s)
        pos *= rad_to_deg
        vel *= rad_to_deg

        freq = -1

        data_fields = (freq,
                       vel[0], pos[0], 0,
                       vel[1], pos[1], 0,
                       vel[2], pos[2], 0,
                       vel[3], pos[3], 0,
                       ts)

        feedback = 'ReHand Status %f %f %f %f %f %f %f %f %f %f %f %f %f %d\r' % data_fields
        print_buffer.append('sending feedback: {}'.format(feedback.rstrip('\r\n')))

        sock.sendto(feedback, settings.REHAND_UDP_CLIENT_ADDR)

        if received_first_cmd:
            n_feedback_packets_sent += 1
            print_buffer.append('# feedback packets sent: {}'.format(n_feedback_packets_sent))
            packets_sent = (n_feedback_packets_sent / (time.time() - t_received_first_cmd))
            print_buffer.append('packets/sec: {}'.format(packets_sent))

        t_last_feedback = t

    if time_to_print:
        time_to_print = False

# stop ReHand process
rh.stop()


from enum import Enum
from typing import List
import pandas as pd
import numpy as np
import spacy
from sklearn.metrics import confusion_matrix
from heapq import nlargest
from collections import OrderedDict
from .texdf import tex_df
from .types import SelectionType, VTAColumnType, ActionType

class Aggregate:
    selection_type: SelectionType
    texdf: tex_df.TexDF
    col_name: str
    col_type: VTAColumnType

    def __init__(self, selection_type, texdf, column, column_type):
        self.selection_type = selection_type
        self.texdf = texdf
        self.col_name = column
        self.col_type = column_type

    def __str__(self):
        col_series = self.texdf.get_dataview_column(self.col_name)
        return col_series.to_string()

    def confusion_matrix(self, md_tag, y_test, predictions, action=ActionType.Add):
        cm = confusion_matrix(y_test, predictions)
        metadata = []
        metadata.append({"Value": cm[0][0], "first": "Not Spam", "second": "Not Spam"})
        metadata.append({"Value": cm[0][1], "first": "Not Spam", "second": "Spam"})
        metadata.append({"Value": cm[1][0], "first": "Spam", "second": "Not Spam"})
        metadata.append({"Value": cm[1][1], "first": "Spam", "second": "Spam"})
        print("confusion matrix metadata: ")
        print(metadata)
        self.texdf.add_metadata(
            self.col_name, "confusion_matrix", VTAColumnType.MAP, metadata
        )
        return "confusion_matrix"

    def count(self, md_tag, action=ActionType.Add):
        rows = self.texdf.get_dataview_column(self.col_name)
        counts = rows.value_counts().to_dict()
        print("counts are: ")
        print(counts)
        self.texdf.add_metadata(
            self.col_name, md_tag, VTAColumnType.MAP, counts
        )
        return md_tag

    def word_scores(self, md_tag, action=ActionType.Add):
        print("in word scores!")
        metadata_dict = self.texdf.get_column_metadata(self.col_name)
        names = metadata_dict.get_metadata_by_tag(md_tag).value
        vectors = self.texdf.get_dataview_column(self.col_name)
        tfidf_vectors = np.array([v.todense() for v in vectors])
        tfidf_vectors[tfidf_vectors == 0] = np.nan
        means = np.nanmean(tfidf_vectors, axis=0)
        means = dict(zip(names, means.tolist()[0]))

        tfidf_vectors = np.array([v.todense() for v in vectors])
        # print(means)
        ordered = np.argsort(tfidf_vectors * -1)
        # ordered = ordered
        # log.info(means)

        # log.info(ordered)

        top_k = 50
        result = {}
        for i in range(len(tfidf_vectors)):
            # Pick top_k from each argsorted matrix for each doc
            for t in range(top_k):
                # Pick the top k word, find its average tfidf from the
                # precomputed dictionary using nanmean and save it to later use
                result[names[ordered[i, 0, t]]] = means[names[ordered[i, 0, t]]]
        top_k_words = nlargest(top_k, result, key=result.get)
        result_ordered = OrderedDict({k: result[k] for k in top_k_words})
        # result_ordered = {k: result[k] for k in top_k_words}
        # new_label_idx = {k: idx for idx, k in enumerate(top_k_words)}

        print("top words dict: ")
        print(result)
        self.texdf.add_metadata(
            self.col_name, "top_scores", VTAColumnType.MAP, result_ordered
        )

        # add coordination index
        # tfidf_vectors = np.array([v.todense() for v in df[column]])
        # TODO: make this more efficient
        reverse_idx_src = OrderedDict({l: [] for l in result_ordered.keys()})
        reverse_idx_target = {i: [] for i in range(len(tfidf_vectors))}
        for i in range(len(tfidf_vectors)):
            for j in range(len(names)):
                tfidf_value = tfidf_vectors[i, 0, j]
                if names[j] in top_k_words:
                    if tfidf_value > 0.0:
                        reverse_idx_src[names[j]].append(i)

        label_idx = {k: idx for idx, k in enumerate(reverse_idx_src.keys())}
        for k, v in reverse_idx_src.items():
            for idx in v:
                if reverse_idx_target.get(idx) is None:
                    reverse_idx_target[idx] = [label_idx[k]]
                else:
                    reverse_idx_target[idx].append(label_idx[k])

        print("top_scores_src: ")
        print(reverse_idx_src)

        print("top_scores_target: ")
        print(reverse_idx_target)

        self.texdf.add_coord_idx("top_scores_src", reverse_idx_src)
        self.texdf.add_coord_idx("top_scores_target", reverse_idx_target)
        return "top_scores"


import torch.nn as nn
import numpy as np
import pdb
from model.utils.cython_bbox import bbox_overlaps
import torch
from torch.autograd import Variable

class box_Augment(nn.Module):
    def __init__(self, ratio, offset):
        super(box_Augment, self).__init__()
        self.ratio = np.array([ratio]).T
        self.offset = np.array([offset]).T

    def forward(self, rois, cls_prob, im_labels, im_data):

        rois_augment = np.empty((0,4), dtype=np.float32)
        rois_init = rois.clone()
        rois = rois[...,1:]
        proposals = self._get_highest_score_proposals(rois, cls_prob, im_labels)
        gt_boxes = proposals['gt_boxes']

        # add original postive samples first
        rois_augment = np.vstack((rois_augment, self._sample_rois(rois, proposals)))

        # base on offset and ratio to augment box
        for i in range(gt_boxes.shape[0]):
            gt_scale = np.array([[gt_boxes[i,3]-gt_boxes[i,1], gt_boxes[i,2]-gt_boxes[i,0]]])
            gt_ctr_ori = np.array([(gt_boxes[i,3]+gt_boxes[i,1])/2, (gt_boxes[i,2]+gt_boxes[i,0])/2])
            aug_scale = self.ratio.dot(gt_scale)
            offset_element = np.vstack((np.hstack((self.offset, np.zeros((self.offset.shape[0], 1)))),
                                                   np.vstack((np.hstack((np.zeros((self.offset.shape[0], 1)), self.offset)), np.hstack((self.offset, self.offset))))))
            gt_ctr = gt_ctr_ori + offset_element
            gt_ctr = np.vstack((gt_ctr_ori, gt_ctr))

            for j in range(gt_ctr.shape[0]):
                for k in range(aug_scale.shape[0]):
                    box = np.array([(gt_ctr[j, 1]-aug_scale[k, 1]/2, gt_ctr[j, 0]-aug_scale[k, 0]/2, gt_ctr[j, 1]+aug_scale[k, 1]/2, gt_ctr[j, 0]+aug_scale[k, 0]/2)])
                    if box[0, 0]>0 and box[0, 1]>0 and box[0, 2]<im_data.shape[3] and box[0, 3]<im_data.shape[2] and box[0,0]<box[0,2] and box[0,1]<box[0,3]:
                        rois_augment = np.vstack((rois_augment, box))

        rois_augment = np.array([np.hstack((np.zeros((rois_augment.shape[0], 1)), rois_augment))], dtype=np.float32)
        rois_augment = Variable(torch.from_numpy(rois_augment).cuda())
        # ret_prob = rois_augment.new().new_zeros(1,rois_augment.size(1),21)

        # return rois_augment, ret_prob
        return rois_augment

    def _get_highest_score_proposals(self, boxes, cls_prob, im_labels):
        """Get proposals with highest score."""

        num_images, num_classes = im_labels.shape
        assert num_images == 1, 'batch size shoud be equal to 1'
        im_labels_tmp = im_labels[0, :]
        gt_boxes = np.zeros((0, 4), dtype=np.float32)
        gt_classes = np.zeros((0, 1), dtype=np.int32)
        gt_scores = np.zeros((0, 1), dtype=np.float32)

        if 21 == cls_prob.shape[2] : # added 1016
            cls_prob = cls_prob[:,:,1:]

        for i in range(num_classes):
            if im_labels_tmp[i] == 1:
                cls_prob_tmp = cls_prob[:,:, i].data
                max_index = np.argmax(cls_prob_tmp)
                gt_boxes = np.vstack((gt_boxes, boxes[:,max_index, :].reshape(1, -1)))
                gt_classes = np.vstack((gt_classes, (i + 1) * np.ones((1, 1), dtype=np.int32))) # for pushing ground
                gt_scores = np.vstack((gt_scores,
                    cls_prob_tmp[:, max_index] ))  # * np.ones((1, 1), dtype=np.float32)))
                cls_prob[:, max_index, :] = 0 #in-place operation <- OICR code but I do not agree

        proposals = {'gt_boxes' : gt_boxes,
                    'gt_classes': gt_classes,
                    'gt_scores': gt_scores}

        return proposals

    def _sample_rois(self, all_rois, proposals):

        gt_boxes = proposals['gt_boxes']
        overlaps = bbox_overlaps(
            np.ascontiguousarray(all_rois[0], dtype=np.float),
            np.ascontiguousarray(gt_boxes, dtype=np.float))
        try :
            max_overlaps = overlaps.max(axis=1)
        except :
            pdb.set_trace()

        fg_inds = np.where(max_overlaps >= 0.5)[0]
        # gt_index = np.where(max_overlaps == 1.0)[0]
        # fg_inds = np.array(list(set(fg_inds)-set(gt_index)))
        pos_samples = np.empty((0,4))
        if fg_inds.shape[0] != 0:
            pos_samples = np.vstack((pos_samples, all_rois[0][fg_inds, :]))

        return pos_samples


"""
PyAudio Example: Make a wire between input and output (i.e., record a
few samples and play them back immediately).
"""

import pyaudio
import time
import numpy as np
from matplotlib import pyplot as plt

CHUNK = 1024
WIDTH = 2
CHANNELS = 2
RATE = 44100
RECORD_SECONDS = 5

p = pyaudio.PyAudio()

stream = p.open(format=p.get_format_from_width(WIDTH),
                channels=CHANNELS,
                rate=RATE,
                input=True,
                output=True,
                frames_per_buffer=CHUNK)

print("* recording")

frames = ""

time.sleep(RECORD_SECONDS)

print("* recording finished")
print("* playing")

for i in range(0, int(RATE / CHUNK * RECORD_SECONDS)):
    data = stream.read(CHUNK)
    frames = frames + data
    stream.write(data, CHUNK)

print("* done")

numpydata = np.fromstring(frames, dtype=np.int16)

print(len(numpydata))
print(numpydata[100000])

# plot data
plt.plot(numpydata)
plt.show()

stream.stop_stream()
stream.close()

p.terminate()