from torch.utils.data import DataLoader, Sampler, Dataset, RandomSampler, DistributedSampler
import json
# from lsp_model import GPT2Tokenizer
from transformers import GPT2Tokenizer
from tqdm import tqdm
import torch
from torch.nn.utils.rnn import pad_sequence
import random
import pickle
import os

def shuffle_persona(persona_list, persona_label):
    id_list = list(range(len(persona_list)))
    shuffled_id_list = list(range(len(persona_list)))
    random.shuffle(shuffled_id_list)

    n_list = [_ for _ in range(len(persona_list))]
    n_label = [-1] * len(persona_label)
    k = 0
    for i, j in zip(id_list, shuffled_id_list):
        n_list[j] = persona_list[i]
        if k < len(persona_label) and persona_label[k] == i:
            n_label[k] = j
            k += 1

    # for check
    # print ('==' * 20)
    # print (id_list)
    # print (shuffled_id_list)
    # print (persona_list, persona_label)
    # print (n_list, n_label)

    return n_list, n_label

def read_file(path, with_persona_label, shuffle=False, mode='label', all_seq_loss=False, single_turn=False, small_data=False, only_persona_response=False):
    END_OF_TEXT_TOKEN = '<|endoftext|>'
    tokenizer = GPT2Tokenizer.from_pretrained('gpt2') # the smallest version of GPT-2, with 124M parameters.
    # tokenizer = GPT2Tokenizer.from_pretrained("microsoft/DialoGPT-small")
    eos = tokenizer.encoder[END_OF_TEXT_TOKEN]

    examples = []

    pick_id = random.randint(1, 7)
    if not shuffle and os.path.exists(path.strip('output') + 'cached'):
        # return torch.load(path.strip('output') + 'cached.npy').tolist()
        return pickle.load(open(path.strip('output') + 'cached', "rb"))

    if shuffle and os.path.exists(path.strip('output') + f'shuffle{pick_id}'):
        # return torch.load(path.strip('output') + 'cached.npy').tolist()
        return pickle.load(open(path.strip('output') + f'shuffle{pick_id}', "rb"))

    with open(path, 'r', encoding='utf-8') as r:
        data = json.load(r)
    # [persona_list, history, response, persona_label] # response is not a list
    for i, (persona_list, history, response, persona_label) in tqdm(enumerate(data)):
    # for i, (persona_list, history, response, persona_label) in enumerate(data):
        if only_persona_response and persona_label[0] == -1:
            continue
        if shuffle:
            persona_list, persona_label = shuffle_persona(persona_list, persona_label)
        if single_turn:
            history = [history[-1]]
        response = [response]
        persona_label_entry = [persona_list[e] for e in persona_label if e > -1]
        persona_label_entry = [' '.join(persona_label_entry).strip()]
        persona_label = [str(e + 1) for e in persona_label] # idx + 1

        if with_persona_label: # add to response list
            if mode == 'entry' and persona_label_entry != ['']:
                # persona_label_entry[0] = persona_label_entry[0] + '\t'
                persona_label_entry[0] = persona_label_entry[0]
                response = [persona_label_entry[0], response[0]] # persona entries + '\t' + response
                assert len(response) == 2
            # print (persona_list, history, response, persona_label, persona_label_entry)
            else:
                # label_str = ' '.join(persona_label).strip() + '\t'
                # label_str = ' '.join(persona_label).strip() + END_OF_TEXT_TOKEN # 直接加<eos> tokenizer不识别。。。
                label_str = ' '.join(persona_label).strip()
                response = [label_str, response[0]]
                assert len(response) == 2

        # tokenize
        persona_list = [tokenizer.encode(s) for s in persona_list]
        history = [tokenizer.encode(s) for s in history]
        response = [tokenizer.encode(s) for s in response]
        persona_label_entry = [tokenizer.encode(s) for s in persona_label_entry]
        persona_label = [tokenizer.encode(str(s)) for s in ' '.join(persona_label).strip()]
        # if not with_persona_label:
        #     persona_label = None

        # making input_ids, position_ids, lm_ids...
        example = make_example_inputs(i, persona_list, history, response, persona_label, persona_label_entry, eos, all_seq_loss=all_seq_loss)
        examples.append(example)

        if small_data and i >= 200:
            break

    if not shuffle and not os.path.exists(path.strip('output') + 'cached'):
        # torch.save(path.strip('output') + 'cached.npy', np.array(examples))
        pickle.dump(examples, open(path.strip('output') + 'cached', "wb"))

    if shuffle and not os.path.exists(path.strip('output') + f'shuffle{pick_id}'):
        # torch.save(path.strip('output') + 'cached.npy', np.array(examples))
        pickle.dump(examples, open(path.strip('output') + f'shuffle{pick_id}', "wb"))
    return examples

def make_example_inputs(id, personas, context, response, persona_label, persona_label_entry, eos, all_seq_loss=False):
    # print (personas , context , persona_label , persona_label_entry, response)

    # sents = None
    # if persona_label:
    #     sents = personas + context + persona_label_entry + response # 0 + 1 + 2 + 2
    # else:
    #     sents = personas + context + response
    sents = personas + context + response

    # 1. input_ids: 每个uttr加了eos，去掉了最后一位
    # print (sents)
    input_ids = [i for s in sents for i in s+[eos]][:-1]
    token_type_ids = []  # this becomes round ids
    lm_labels = []

    # 2. lm_labels: input_ids[1:] + [eos]
    #    token_type_ids: 0 for persona, 1 for context, 2 for persona_label and response
    for i, s in enumerate(sents):
        if i == 0:
            token_type_ids += [0] * len(s)
            lm_labels += [-1] * len(s) if not all_seq_loss else s[1:] + [eos]
        elif i < len(personas): # persona: 0
            token_type_ids += [0] * (len(s) + 1)
            lm_labels += [-1] * (len(s) + 1) if not all_seq_loss else s + [eos]
        elif i < len(personas) + len(context): # context: 1
            token_type_ids += [1] * (len(s) + 1)
            lm_labels += [-1] * (len(s) + 1) if not all_seq_loss else s + [eos]
        else: # persona_label/entry + '\t' + response: 2
            token_type_ids += [2] * (len(s) + 1)
            lm_labels += (s + [eos])

    # handle trailing -1's
    i = len(lm_labels) - 1
    while i >= 0:
        if lm_labels[i] != -1:
            break
        i -= 1
    input_ids = input_ids[:i+1]
    lm_labels = lm_labels[:i+1]
    token_type_ids = token_type_ids[:i+1]

    # pad to multiples of 8
    while len(input_ids) % 8 != 0:
        input_ids.append(0)
        token_type_ids.append(0)
        lm_labels.append(-1)

    # 3. position_ids
    position_ids = list(range(len(input_ids)))
    assert (len(input_ids) == len(position_ids) == len(token_type_ids)
            == len(lm_labels))
    assert len(input_ids) % 8 == 0

    # example = [id, input_ids, position_ids, token_type_ids,
                            # lm_labels]
    # print (all_seq_loss, input_ids, lm_labels)

    example = {
        'id': id,
        'input_ids': input_ids,
        'position_ids': position_ids,
        'token_type_ids': token_type_ids,
        'lm_labels': lm_labels,
        'input_len': len(input_ids)
    }
    return example

class PersonaDataset(Dataset):
    """ pytorch dataset for GPT2 training """
    def __init__(self, path, max_len=None, with_persona_label=True, shuffle=False, all_seq_loss=False, single_turn=False, small_data=False, only_persona_response=False, **kwargs):
        self.example_ids = read_file(path, with_persona_label, shuffle, all_seq_loss=all_seq_loss, single_turn=single_turn, small_data=small_data, only_persona_response=only_persona_response)
        # print ('data_num = ', len(self.example_ids))
        self.max_len = max_len  # this max_len do truncate

    def __getitem__(self, i):
        return self.example_ids[i]

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

    @staticmethod
    def collate(features):
        # print (features)
        input_ids = pad_sequence([torch.tensor(f['input_ids'], dtype=torch.long)
                                  for f in features],
                                 batch_first=True, padding_value=0)
        position_ids = pad_sequence([torch.tensor(f['position_ids'],
                                                  dtype=torch.long)
                                     for f in features],
                                    batch_first=True, padding_value=0)
        token_type_ids = pad_sequence([torch.tensor(f['token_type_ids'],
                                                    dtype=torch.long)
                                       for f in features],
                                      batch_first=True, padding_value=0)
        labels = pad_sequence([torch.tensor(f['lm_labels'], dtype=torch.long)
                               for f in features],
                              batch_first=True, padding_value=-1)
        return (input_ids, position_ids, token_type_ids, labels)

# test
# tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
# path = '/misc/kfdata01/kf_grp/lchen/D3/data_manipulation/data_distillation/predictions/test/output'
# path = '/misc/kfdata01/kf_grp/lchen/D3/data_manipulation/data_distillation/predictions/valid/output'
# path = '/misc/kfdata01/kf_grp/lchen/D3/data_manipulation/data_distillation/predictions/train/output'
# dataset = PersonaDataset(path, max_len=180, with_persona_label=False, shuffle=False)
# dataset = PersonaDataset(path, max_len=180, with_persona_label=False, shuffle=True)
# # dataset = PersonaDataset(path, max_len=256, with_persona_label=False, shuffle=False, single_turn=True)
# dataset = PersonaDataset(path, max_len=256, with_persona_label=False, shuffle=False, single_turn=True, only_persona_response=True)
# sampler = RandomSampler(dataset) if True else DistributedSampler(dataset)
# dataloader = DataLoader(dataset, sampler=sampler, batch_size=4, collate_fn=PersonaDataset.collate)

# for i, batch in enumerate(dataloader):
#     seq_len = batch[0].shape[1]
#     input_ids, position_ids, token_ids, label_ids, *_ = batch

#     if i > 5:
#         break

#     # visualize data
#     print ('=='*10 + ' visualize data ' + '=='*10)
#     print ('input_ids.shape, position_ids.shape, label_ids.shape = ', input_ids.shape, position_ids.shape, label_ids.shape) # torch.Size([4, 512]) torch.Size([4, 512])
#     print ('input_ids[0] = ', input_ids[0])
#     print ('position_ids[0] = ', position_ids[0])
#     print ('token_ids[0] = ', token_ids[0])
#     print ('label_ids[0] = ', label_ids[0])
#     # 'GPT2Tokenizer' object has no attribute 'batch_decode'???
#     # print (tokenizer.batch_decode(inputs[0], skip_special_tokens=True))
#     # print (tokenizer.batch_decode(labels[0], skip_special_tokens=True))

#     print ('input_ids = \n', tokenizer.decode(input_ids[0].tolist()))
#     # mask = token_ids.eq(2) | token_ids.eq(3)
#     mask = token_ids.eq(2)
#     print (mask.shape)
#     # print (mask[0])
#     print ('label_ids = \n', tokenizer.decode(label_ids[0][mask[0]].tolist()))
#     print ()
#     # break


from lib import binary_tree
from lib import stack
import operator

def parse(exp):
    root = binary_tree.BinaryTree()
    current_node = root
    st = stack.Stack()
    st.push(current_node)
    lst = list(exp)
    op_lst = ['+', '-', '*', '/']

    for i in lst:
        if i == "(":
            root.insert_left('')
            current_node = root.get_left_child()
        elif i.isdigit():
            current_node.set_value(int(i))
            parent = st.pop()
            current_node = parent
        elif i in op_lst:
            parent.set_value(i)
            st.push(parent)
            parent.insert_right('')
            current_node = parent.get_right_child()
        elif i == ")":
            current_node = st.pop()
        else:
            raise ValueError

    return root

def evaluate(self, parse_tree):
    """Evaluates the parser tree.

       :returns The evaluated expression.
    """
    opers = {'+': operator.add, '-': operator.sub, '*': operator.mul,
             '/': operator.truediv}

    left_c = parse_tree.getLeftChild()
    right_c = parse_tree.getRightChild()

    if left_c and right_c:
        fn = opers[parse_tree.getRootVal()]
        return fn(self.evaluate(left_c), self.evaluate(right_c))
    else:
        return parse_tree.getRootVal()

if __name__ == "__main__":
    inp = "(3d+(4*5))"
    parse(inp)


import numpy as np
from six.moves import xrange
import tensorflow as tf

class Memory(object):
  """Memory module."""

  def __init__(self, key_dim, memory_size, vocab_size,
               choose_k=256, alpha=0.1, correct_in_top=1, age_noise=8.0,
               var_cache_device='', nn_device=''):
    self.key_dim = key_dim
    self.memory_size = memory_size
    self.vocab_size = vocab_size
    self.choose_k = min(choose_k, memory_size)
    self.alpha = alpha
    self.correct_in_top = correct_in_top
    self.age_noise = age_noise
    self.var_cache_device = var_cache_device  # Variables are cached here.
    self.nn_device = nn_device  # Device to perform nearest neighbour matmul.

    caching_device = var_cache_device if var_cache_device else None
    self.update_memory = tf.constant(True)  # Can be fed "false" if needed.
    self.mem_keys = tf.get_variable(
        'memkeys', [self.memory_size, self.key_dim], trainable=False,
        initializer=tf.random_uniform_initializer(-0.0, 0.0),
        caching_device=caching_device)
    self.mem_vals = tf.get_variable(
        'memvals', [self.memory_size], dtype=tf.int32, trainable=False,
        initializer=tf.constant_initializer(0, tf.int32),
        caching_device=caching_device)
    self.mem_age = tf.get_variable(
        'memage', [self.memory_size], dtype=tf.float32, trainable=False,
        initializer=tf.constant_initializer(0.0), caching_device=caching_device)
    self.recent_idx = tf.get_variable(
        'recent_idx', [self.vocab_size], dtype=tf.int32, trainable=False,
        initializer=tf.constant_initializer(0, tf.int32))

    # variable for projecting query vector into memory key
    self.query_proj = tf.get_variable(
        'memory_query_proj', [self.key_dim, self.key_dim], dtype=tf.float32,
        initializer=tf.truncated_normal_initializer(0, 0.01),
        caching_device=caching_device)

  def get(self):
    return self.mem_keys, self.mem_vals, self.mem_age, self.recent_idx

  def set(self, k, v, a, r=None):
    return tf.group(
        self.mem_keys.assign(k),
        self.mem_vals.assign(v),
        self.mem_age.assign(a),
        (self.recent_idx.assign(r) if r is not None else tf.group()))

  def clear(self):
    return tf.variables_initializer([self.mem_keys, self.mem_vals, self.mem_age,
                                     self.recent_idx])

  def get_hint_pool_idxs(self, normalized_query):
    """Get small set of idxs to compute nearest neighbor queries on.

    This is an expensive look-up on the whole memory that is used to
    avoid more expensive operations later on.

    Args:
      normalized_query: A Tensor of shape [None, key_dim].

    Returns:
      A Tensor of shape [None, choose_k] of indices in memory
      that are closest to the queries.

    """
    # look up in large memory, no gradients
    with tf.device(self.nn_device):
      similarities = tf.matmul(tf.stop_gradient(normalized_query),
                               self.mem_keys, transpose_b=True, name='nn_mmul')
    _, hint_pool_idxs = tf.nn.top_k(
        tf.stop_gradient(similarities), k=self.choose_k, name='nn_topk')
    return hint_pool_idxs

  def make_update_op(self, upd_idxs, upd_keys, upd_vals,
                     batch_size, use_recent_idx, intended_output):
    """Function that creates all the update ops."""
    mem_age_incr = self.mem_age.assign_add(tf.ones([self.memory_size],
                                                   dtype=tf.float32))
    with tf.control_dependencies([mem_age_incr]):
      mem_age_upd = tf.scatter_update(
          self.mem_age, upd_idxs, tf.zeros([batch_size], dtype=tf.float32))

    mem_key_upd = tf.scatter_update(
        self.mem_keys, upd_idxs, upd_keys)
    mem_val_upd = tf.scatter_update(
        self.mem_vals, upd_idxs, upd_vals)

    if use_recent_idx:
      recent_idx_upd = tf.scatter_update(
          self.recent_idx, intended_output, upd_idxs)
    else:
      recent_idx_upd = tf.group()

    return tf.group(mem_age_upd, mem_key_upd, mem_val_upd, recent_idx_upd)

  def query(self, query_vec, intended_output, use_recent_idx=True):
    """Queries memory for nearest neighbor.

    Args:
      query_vec: A batch of vectors to query (embedding of input to model).
      intended_output: The values that would be the correct output of the
        memory.
      use_recent_idx: Whether to always insert at least one instance of a
        correct memory fetch.

    Returns:
      A tuple (result, mask, teacher_loss).
      result: The result of the memory look up.
      mask: The affinity of the query to the result.
      teacher_loss: The loss for training the memory module.
    """

    batch_size = tf.shape(query_vec)[0]
    output_given = intended_output is not None

    # prepare query for memory lookup
    query_vec = tf.matmul(query_vec, self.query_proj)
    normalized_query = tf.nn.l2_normalize(query_vec, dim=1)

    hint_pool_idxs = self.get_hint_pool_idxs(normalized_query)

    if output_given and use_recent_idx:  # add at least one correct memory
      most_recent_hint_idx = tf.gather(self.recent_idx, intended_output)
      hint_pool_idxs = tf.concat(
          axis=1,
          values=[hint_pool_idxs, tf.expand_dims(most_recent_hint_idx, 1)])
    choose_k = tf.shape(hint_pool_idxs)[1]

    with tf.device(self.var_cache_device):
      # create small memory and look up with gradients
      my_mem_keys = tf.stop_gradient(tf.gather(self.mem_keys, hint_pool_idxs,
                                               name='my_mem_keys_gather'))
      similarities = tf.matmul(tf.expand_dims(normalized_query, 1),
                               my_mem_keys, adjoint_b=True, name='batch_mmul')
      hint_pool_sims = tf.squeeze(similarities, [1], name='hint_pool_sims')
      hint_pool_mem_vals = tf.gather(self.mem_vals, hint_pool_idxs,
                                     name='hint_pool_mem_vals')
    # Calculate softmax mask on the top-k if requested.
    # Softmax temperature. Say we have K elements at dist x and one at (x+a).
    # Softmax of the last is e^tm(x+a)/Ke^tm*x + e^tm(x+a) = e^tm*a/K+e^tm*a.
    # To make that 20% we'd need to have e^tm*a ~= 0.2K, so tm = log(0.2K)/a.
    softmax_temp = max(1.0, np.log(0.2 * self.choose_k) / self.alpha)
    mask = tf.nn.softmax(hint_pool_sims[:, :choose_k - 1] * softmax_temp)

    # prepare hints from the teacher on hint pool
    teacher_hints = tf.to_float(
        tf.abs(tf.expand_dims(intended_output, 1) - hint_pool_mem_vals))
    teacher_hints = 1.0 - tf.minimum(1.0, teacher_hints)

    teacher_vals, teacher_hint_idxs = tf.nn.top_k(
        hint_pool_sims * teacher_hints, k=1)
    neg_teacher_vals, _ = tf.nn.top_k(
        hint_pool_sims * (1 - teacher_hints), k=1)

    # bring back idxs to full memory
    teacher_idxs = tf.gather(
        tf.reshape(hint_pool_idxs, [-1]),
        teacher_hint_idxs[:, 0] + choose_k * tf.range(batch_size))

    # zero-out teacher_vals if there are no hints
    teacher_vals *= (
        1 - tf.to_float(tf.equal(0.0, tf.reduce_sum(teacher_hints, 1))))

    # prepare returned values
    nearest_neighbor = tf.to_int32(
        tf.argmax(hint_pool_sims[:, :choose_k - 1], 1))
    no_teacher_idxs = tf.gather(
        tf.reshape(hint_pool_idxs, [-1]),
        nearest_neighbor + choose_k * tf.range(batch_size))

    # we'll determine whether to do an update to memory based on whether
    # memory was queried correctly
    sliced_hints = tf.slice(teacher_hints, [0, 0], [-1, self.correct_in_top])
    incorrect_memory_lookup = tf.equal(0.0, tf.reduce_sum(sliced_hints, 1))

    # loss based on triplet loss
    teacher_loss = (tf.nn.relu(neg_teacher_vals - teacher_vals + self.alpha)
                    - self.alpha)

    with tf.device(self.var_cache_device):
      result = tf.gather(self.mem_vals, tf.reshape(no_teacher_idxs, [-1]))

    # prepare memory updates
    update_keys = normalized_query
    update_vals = intended_output

    fetched_idxs = teacher_idxs  # correctly fetched from memory
    with tf.device(self.var_cache_device):
      fetched_keys = tf.gather(self.mem_keys, fetched_idxs, name='fetched_keys')
      fetched_vals = tf.gather(self.mem_vals, fetched_idxs, name='fetched_vals')

    # do memory updates here
    fetched_keys_upd = update_keys + fetched_keys  # Momentum-like update
    fetched_keys_upd = tf.nn.l2_normalize(fetched_keys_upd, dim=1)
    # Randomize age a bit, e.g., to select different ones in parallel workers.
    mem_age_with_noise = self.mem_age + tf.random_uniform(
        [self.memory_size], - self.age_noise, self.age_noise)

    _, oldest_idxs = tf.nn.top_k(mem_age_with_noise, k=batch_size, sorted=False)

    with tf.control_dependencies([result]):
      upd_idxs = tf.where(incorrect_memory_lookup,
                          oldest_idxs,
                          fetched_idxs)
      # upd_idxs = tf.Print(upd_idxs, [upd_idxs], "UPD IDX", summarize=8)
      upd_keys = tf.where(incorrect_memory_lookup,
                          update_keys,
                          fetched_keys_upd)
      upd_vals = tf.where(incorrect_memory_lookup,
                          update_vals,
                          fetched_vals)

    def make_update_op():
      return self.make_update_op(upd_idxs, upd_keys, upd_vals,
                                 batch_size, use_recent_idx, intended_output)

    update_op = tf.cond(self.update_memory, make_update_op, tf.no_op)

    with tf.control_dependencies([update_op]):
      result = tf.identity(result)
      mask = tf.identity(mask)
      teacher_loss = tf.identity(teacher_loss)

    return result, mask, tf.reduce_mean(teacher_loss)

class LSHMemory(Memory):
  """Memory employing locality sensitive hashing.

  Note: Not fully tested.
  """

  def __init__(self, key_dim, memory_size, vocab_size,
               choose_k=256, alpha=0.1, correct_in_top=1, age_noise=8.0,
               var_cache_device='', nn_device='',
               num_hashes=None, num_libraries=None):
    super(LSHMemory, self).__init__(
        key_dim, memory_size, vocab_size,
        choose_k=choose_k, alpha=alpha, correct_in_top=1, age_noise=age_noise,
        var_cache_device=var_cache_device, nn_device=nn_device)

    self.num_libraries = num_libraries or int(self.choose_k ** 0.5)
    self.num_per_hash_slot = max(1, self.choose_k // self.num_libraries)
    self.num_hashes = (num_hashes or
                       int(np.log2(self.memory_size / self.num_per_hash_slot)))
    self.num_hashes = min(max(self.num_hashes, 1), 20)
    self.num_hash_slots = 2 ** self.num_hashes

    # hashing vectors
    self.hash_vecs = [
        tf.get_variable(
            'hash_vecs%d' % i, [self.num_hashes, self.key_dim],
            dtype=tf.float32, trainable=False,
            initializer=tf.truncated_normal_initializer(0, 1))
        for i in xrange(self.num_libraries)]

    # map representing which hash slots map to which mem keys
    self.hash_slots = [
        tf.get_variable(
            'hash_slots%d' % i, [self.num_hash_slots, self.num_per_hash_slot],
            dtype=tf.int32, trainable=False,
            initializer=tf.random_uniform_initializer(maxval=self.memory_size,
                                                      dtype=tf.int32))
        for i in xrange(self.num_libraries)]

  def get(self):  # not implemented
    return self.mem_keys, self.mem_vals, self.mem_age, self.recent_idx

  def set(self, k, v, a, r=None):  # not implemented
    return tf.group(
        self.mem_keys.assign(k),
        self.mem_vals.assign(v),
        self.mem_age.assign(a),
        (self.recent_idx.assign(r) if r is not None else tf.group()))

  def clear(self):
    return tf.variables_initializer([self.mem_keys, self.mem_vals, self.mem_age,
                                     self.recent_idx] + self.hash_slots)

  def get_hash_slots(self, query):
    """Gets hashed-to buckets for batch of queries.

    Args:
      query: 2-d Tensor of query vectors.

    Returns:
      A list of hashed-to buckets for each hash function.
    """

    binary_hash = [
        tf.less(tf.matmul(query, self.hash_vecs[i], transpose_b=True), 0)
        for i in xrange(self.num_libraries)]
    hash_slot_idxs = [
        tf.reduce_sum(
            tf.to_int32(binary_hash[i]) *
            tf.constant([[2 ** i for i in xrange(self.num_hashes)]],
                        dtype=tf.int32), 1)
        for i in xrange(self.num_libraries)]
    return hash_slot_idxs

  def get_hint_pool_idxs(self, normalized_query):
    """Get small set of idxs to compute nearest neighbor queries on.

    This is an expensive look-up on the whole memory that is used to
    avoid more expensive operations later on.

    Args:
      normalized_query: A Tensor of shape [None, key_dim].

    Returns:
      A Tensor of shape [None, choose_k] of indices in memory
      that are closest to the queries.

    """
    # get hash of query vecs
    hash_slot_idxs = self.get_hash_slots(normalized_query)

    # grab mem idxs in the hash slots
    hint_pool_idxs = [
        tf.maximum(tf.minimum(
            tf.gather(self.hash_slots[i], idxs),
            self.memory_size - 1), 0)
        for i, idxs in enumerate(hash_slot_idxs)]

    return tf.concat(axis=1, values=hint_pool_idxs)

  def make_update_op(self, upd_idxs, upd_keys, upd_vals,
                     batch_size, use_recent_idx, intended_output):
    """Function that creates all the update ops."""
    base_update_op = super(LSHMemory, self).make_update_op(
        upd_idxs, upd_keys, upd_vals,
        batch_size, use_recent_idx, intended_output)

    # compute hash slots to be updated
    hash_slot_idxs = self.get_hash_slots(upd_keys)

    # make updates
    update_ops = []
    with tf.control_dependencies([base_update_op]):
      for i, slot_idxs in enumerate(hash_slot_idxs):
        # for each slot, choose which entry to replace
        entry_idx = tf.random_uniform([batch_size],
                                      maxval=self.num_per_hash_slot,
                                      dtype=tf.int32)
        entry_mul = 1 - tf.one_hot(entry_idx, self.num_per_hash_slot,
                                   dtype=tf.int32)
        entry_add = (tf.expand_dims(upd_idxs, 1) *
                     tf.one_hot(entry_idx, self.num_per_hash_slot,
                                dtype=tf.int32))

        mul_op = tf.scatter_mul(self.hash_slots[i], slot_idxs, entry_mul)
        with tf.control_dependencies([mul_op]):
          add_op = tf.scatter_add(self.hash_slots[i], slot_idxs, entry_add)
          update_ops.append(add_op)

    return tf.group(*update_ops)


"""
Simplemonitor logger for seq

Inspiration from
https://raw.githubusercontent.com/eifinger/appdaemon-scripts/master/seqSink/seqSink.py
"""

import datetime
import json
from typing import Optional, cast

import requests

from ..Monitors.monitor import Monitor
from .logger import Logger, register

@register
class SeqLogger(Logger):
    """Logging to seq"""

    logger_type = "seq"
    only_failures = False
    buffered = False
    dateformat = None

    def __init__(self, config_options: Optional[dict] = None) -> None:
        if config_options is None:
            config_options = {}
        super().__init__(config_options)

        # i.e. http://192.168.0.5:5341
        self.endpoint = cast(
            str, self.get_config_option("endpoint", required=True, allow_empty=False)
        )
        self.timeout = cast(
            int, self.get_config_option("timeout", required_type="int", default=5)
        )
        # Potentially, would need to add a header for ApiKey

        # Send message to indicate we have started logging
        self.log_to_seq(
            self.endpoint,
            "SeqLogger",
            "simpleMonitor",
            "__init__",
            None,
            "logging enabled for simpleMonitor",
            False,
        )

    def save_result2(self, name: str, monitor: Monitor) -> None:
        try:
            is_fail = monitor.test_success() is False

            self.log_to_seq(
                self.endpoint,
                name,
                monitor.name,
                monitor.monitor_type,
                str(monitor.get_params()),
                monitor.describe(),
                is_fail,
            )
        except Exception:
            self.logger_logger.exception("Error sending to seq in %s", monitor.name)

    def describe(self) -> str:
        return "Sends simple log to seq using raw endpoint"

    def log_to_seq(
        self, endpoint, name, app_name, monitor_type, params, description, is_fail
    ):
        """Send an event to seq"""
        event_data = {
            "Timestamp": str(datetime.datetime.now()),
            "Level": "Error" if is_fail is True else "Information",
            "MessageTemplate": str(description),
            "Properties": {
                "Type": "simpleMonitor",
                "Name": name,
                "Monitor": str(app_name),
                "MonitorType": monitor_type,
                # "Params": params
            },
        }
        if params is not None:
            event_data["Properties"]["Params"] = params

        request_body = {"Events": [event_data]}

        try:
            _ = json.dumps(request_body)  # This just checks it is valid...
        except TypeError:
            self.logger_logger.error("Could not serialise %s", request_body)
            return

        try:
            response = requests.post(
                self.endpoint, json=request_body, timeout=self.timeout
            )
            if not response.status_code == 200 and not response.status_code == 201:
                self.logger_logger.error(
                    "POST to seq failed with status code: %s", response
                )
        except requests.RequestException:
            self.logger_logger.exception("Failed to log to seq")


#pip install -q -U tensorflow-text

#import tensorflow_text as text

import seaborn as sns
import tensorflow as tf
import numpy as np
import pandas as pd
import os
import nltk
from nltk.stem import WordNetLemmatizer
from nltk.tokenize import WhitespaceTokenizer
wn = WordNetLemmatizer()
stopwords = nltk.corpus.stopwords.words('english')
import re
import utils
from config import *

os.environ["OMP_NUM_THREADS"] = str(NUM_THREADS)
os.environ["TF_NUM_INTRAOP_THREADS"] = str(NUM_THREADS)
os.environ["TF_NUM_INTEROP_THREADS"] = str(NUM_THREADS)

tf.config.threading.set_inter_op_parallelism_threads(
    NUM_THREADS
)
tf.config.threading.set_intra_op_parallelism_threads(
    NUM_THREADS
)
tf.config.set_soft_device_placement(True)

df, unique_labels = utils.preprocess_data(PATH)
df = utils.clean_data(df)

def label_freq(df, unique_labels):
    freq = dict()
    for unique_label in unique_labels:
        freq[unique_label] = df[unique_label].mean() * df[unique_label].shape[0]
    print("Label frequency ", freq)
    sns.set(rc={'figure.figsize': (11.7, 8.27)})
    plot = sns.barplot(list(freq.keys()), list(freq.values()))
    plot.set_xticklabels(plot.get_xticklabels(),
                         rotation=45,
                         horizontalalignment='right')

def save_model(model, base_fp):
    # save the model: first the weights then the arch
    model.save_weights('{}-weights.{}'.format(base_fp, FILE_TYPE))
    with open('{}-architecture.json'.format(base_fp), 'w') as f:
        f.write(model.to_json())

with tf.device('/cpu:0'):
    from tensorflow.keras.preprocessing.text import Tokenizer
    from tensorflow.keras.preprocessing.sequence import pad_sequences
    tokenizer = Tokenizer(num_words=MAX_WORDS, lower=True)
    tokenizer.fit_on_texts(df['content'])
    sequences = tokenizer.texts_to_sequences(df['content'])
    x = pad_sequences(sequences, maxlen=MAX_LEN)

    train = df
    x_train = train["content"].str.lower()
    y_train = train[unique_labels].values

    tokenizer = tf.keras.preprocessing.text.Tokenizer(num_words=MAX_WORDS, lower=True)

    tokenizer.fit_on_texts(x_train)

    x_train = tokenizer.texts_to_sequences(x_train)
    x_train = tf.keras.preprocessing.sequence.pad_sequences(x_train, maxlen=MAX_LEN)

    embeddings_index = {}

    with open(GLOVE_EMBEDDING, encoding='utf8') as f:
        for line in f:
            values = line.rstrip().rsplit(' ')
            word = values[0]
            embed = np.asarray(values[1:], dtype='float32')
            embeddings_index[word] = embed

    word_index = tokenizer.word_index

    embedding_matrix = np.zeros((MAX_WORDS, EMBEDDING_SIZE), dtype='float32')

    for word, i in word_index.items():

        if i >= MAX_WORDS:
            continue

        embedding_vector = embeddings_index.get(word)

        if embedding_vector is not None:
            embedding_matrix[i] = embedding_vector

    input = tf.keras.layers.Input(shape=(MAX_LEN,))
    x = tf.keras.layers.Embedding(MAX_WORDS, EMBEDDING_SIZE, weights=[embedding_matrix], trainable=False)(input)

    x = tf.keras.layers.Bidirectional(tf.keras.layers.GRU(128, return_sequences=True, dropout=0.1,
                                                          recurrent_dropout=0.1))(x)

    x = tf.keras.layers.Conv1D(100, kernel_size=1)(x)

    avg_pool = tf.keras.layers.GlobalAveragePooling1D()(x)
    max_pool = tf.keras.layers.GlobalMaxPooling1D()(x)

    x = tf.keras.layers.concatenate([avg_pool, max_pool])
    # x = tf.keras.layers.Dense(10, activation='relu')(x)
    preds = tf.keras.layers.Dense(len(unique_labels), activation="sigmoid")(x)

    model = tf.keras.Model(input, preds)

    model.summary()

    model.compile(loss='binary_crossentropy', optimizer=tf.keras.optimizers.Adam(lr=1e-3), metrics=['accuracy'])

    batch_size = 128

    checkpoint_path = "./training_1/cp.ckpt"
    checkpoint_dir = os.path.dirname(checkpoint_path)

    cp_callback = tf.keras.callbacks.ModelCheckpoint(checkpoint_path,
                                                     save_weights_only=True,
                                                     verbose=1)

    callbacks = [
        tf.keras.callbacks.EarlyStopping(patience=5, monitor='val_loss'),
        tf.keras.callbacks.TensorBoard(log_dir='./logs'),
        cp_callback
    ]

    model.fit(x_train, y_train, validation_split=0.2, batch_size=batch_size,
              epochs=EPOCHS, callbacks=callbacks, verbose=1)

    save_model(model, BASE_FP)

    model.save(checkpoint_dir, save_format='tf')

    latest = tf.train.latest_checkpoint(checkpoint_dir)

    model.load_weights(latest)

    predictions = model.predict(np.expand_dims(x_train[43], 0))

    print(tokenizer.sequences_to_texts([x_train[43]]))
    print(y_train[43])
    print(predictions)


from django import forms
from catalog.models import Products, Version

class MixinStyle:

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        for field_name, field in self.fields.items():
            field.widget.attrs['class'] = 'form-control'

class ProductsForms(MixinStyle, forms.ModelForm):
    class Meta:
        model = Products
        exclude = ('product_date_create', 'product_date_edit')

    def clean_product_name(self):
        cleaned_name = self.cleaned_data['product_name']
        forbidden_words = ["казино", "криптовалюта", "крипта", "биржа", "дешево", "бесплатно", "обман", "полиция",
                           "радар"]
        for word in forbidden_words:
            if word in cleaned_name:
                raise forms.ValidationError(f'Запрещено истольвозвать слово {word}.')
            if word.title() in cleaned_name:
                raise forms.ValidationError(f'Запрещено истольвозвать слово {word}.')
        return cleaned_name

    def clean_product_text(self):
        cleaned_text = self.cleaned_data['product_text']
        forbidden_words = ["казино", "криптовалюта", "крипта", "биржа", "дешево", "бесплатно", "обман", "полиция",
                           "радар"]
        for word in forbidden_words:
            if word.title() in cleaned_text:
                raise forms.ValidationError(f'Запрещено истольвозвать слово {word}.')
            if word in cleaned_text:
                raise forms.ValidationError(f'Запрещено истольвозвать слово {word}.')
        return cleaned_text

class VersionForms(MixinStyle, forms.ModelForm):
    class Meta:
        model = Version
        fields = '__all__'


from django.http import HttpResponse
from django.shortcuts import render
import string as st
ab = st.punctuation
print(ab)
def index(request):

   return render(request,'index.html')
    # return HttpResponse('''''')

def analyze(request):
    djtext = request.POST.get('text','default')

    removepunc = request.POST.get('removepunc', 'off')

    capital = request.POST.get('capital','off')
    newLine = request.POST.get('newline','off')
    spaceRemover = request.POST.get('spaceRemover','off')
    counter = request.POST.get('characterCount','off')

    if removepunc =="on":
       analyzed = ""
       for char in djtext:
           if char not in ab:
               analyzed = analyzed + char
       params = {"purpose":"remove punctuation","analyzed_text":analyzed}
       djtext = analyzed

       # return render(request,'analyze.html',params)

    if capital == "on":
        analyzed = ""
        for char in djtext:
             analyzed = analyzed + char.upper()
        params = {"purpose": "capitalisation", "analyzed_text": analyzed}
        # return render(request, 'analyze.html', params)
        djtext = analyzed
    if newLine =='on':
        analyzed = ""
        for char in djtext:
            if char !="\n" and char!="\r":
             analyzed = analyzed +char
        params = {"purpose": "removeNewLine", "analyzed_text": analyzed}
        # return render(request, 'analyze.html', params)
        djtext= analyzed
    if spaceRemover =='on':
        analyzed = ""
        for char in djtext:
            if char !=" ":
             analyzed = analyzed +char
        params = {"purpose": "extraSpacesRemover", "analyzed_text": analyzed}
        # return render(request, 'analyze.html', params)
        djtext=analyzed
    if counter =='on':
        count = 0
        for char in djtext:
            count +=1
        final = f"total no of characters in = {djtext} is {count}"

        params = {"purpose": "characterCounter", "analyzed_text": final}
        # return render(request, 'analyze.html', params)
        djtext=final
    if counter!='on' and spaceRemover !='on' and newLine !='on' and capital != "on" and removepunc !="on" :
        return HttpResponse("Sorry .........please select option")

    return render(request, 'analyze.html', params)


#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Date    : 2019-03-12 14:42:47
# @Author  : Your Name (you@example.org)
# @Link    : http://example.org
# @Version : $Id$
# @Desc    : 用CNN做的模型  初始化权重 ReLU
#
import tensorflow as tf
from numpy import float32
import datetime
from tensorflow.examples.tutorials.mnist import input_data

# 初始化w、b
def  weight_variable(shape):
    initial = tf.truncated_normal(shape, stddev = 0.1)
    return tf.Variable(initial)
def bias_variable(shape):
    initial = tf.constant(0.1, shape= shape)
    return tf.Variable(initial)

# 卷积和池化
def conv2d(x,W):
    return tf.nn.conv2d(x,W,strides=[1,1,1,1],padding="SAME")
def max_pool_2x2(x):
    return tf.nn.max_pool(x,ksize=[1,2,2,1], strides=[1, 2, 2,1],padding="SAME")

start = datetime.datetime.now()
# 载入数据
mnist = input_data.read_data_sets('./data',one_hot=True)

# None表示可以是任何维度
x = tf.placeholder(tf.float32,[None,784])
# y:正确值(正确label)
y = tf.placeholder(tf.float32, [None, 10])

############ 第一层CNN ###############
# 28*28  --->   24*24  --->   12*12
W_conv1 = weight_variable([5,5,1,32])
b_conv1 = bias_variable([32])
# input
x_image = tf.reshape(x,[-1,28,28,1])
h_conv1 = tf.nn.relu(conv2d(x_image,W_conv1) + b_conv1)
h_pool1 = max_pool_2x2(h_conv1)

############ 第二层CNN ###############
# 12*12   --->   8*8   --->  7*7

W_conv2 = weight_variable([5,5,32,64])
b_conv2 = bias_variable([64])

h_conv2 = tf.nn.relu(conv2d(h_pool1,W_conv2) + b_conv2)
h_pool2 = max_pool_2x2(h_conv2)

############ 全连接层 ###############
W_fc1 = weight_variable([7*7*64,1024])
b_fc1 = bias_variable([1024])
h_pool2_flat = tf.reshape(h_pool2,[-1,7*7*64])
h_fc1 = tf.nn.relu(tf.matmul(h_pool2_flat,W_fc1)+b_fc1)

############ Dropout 减少过拟合 ###############
keep_prob = tf.placeholder("float")
h_fc1_drop = tf.nn.dropout(h_fc1,keep_prob)

############ 输出层 ###############
W_fc2 = weight_variable([1024,10])
b_fc2 = bias_variable([10])
y_conv = tf.nn.softmax(tf.matmul(h_fc1_drop,W_fc2) + b_fc2)

# 损失函数 loss function
cross_entropy = -tf.reduce_sum(y * tf.log(y_conv))

# ADAM优化器做梯度下降
train_step = tf.train.AdamOptimizer(1e-4).minimize(cross_entropy)

# 找出预测正确的标签 tf.equal()返回的是true or false 数组
correct_prediction = tf.equal(tf.argmax(y_conv,1), tf.argmax(y,1))
# 求正确率
accuracy = tf.reduce_mean(tf.cast(correct_prediction, "float"))

#存储训练的模型
saver = tf.train.Saver(max_to_keep=1)

# 一批次大小
batch_size = 50
# 所有批次数量
n_batch = mnist.train.num_examples // batch_size
with tf.Session() as sess:
    saver_max_acc = 0
    sess.run(tf.initialize_all_variables())
    for epoch in range(1):
        for i in range(n_batch):
            batch = mnist.train.next_batch(50)
            sess.run(train_step,feed_dict={x:batch[0], y: batch[1], keep_prob: 0.5})
            if i%100 == 0:
                train_accuracy = accuracy.eval(feed_dict={x:batch[0], y: batch[1], keep_prob: 1.0})
                print("step %d acc : %g"%(i,train_accuracy))

                acc=accuracy.eval(feed_dict={x: mnist.test.images[0:2000], y: mnist.test.labels[0:2000], keep_prob: 1.0})
                print("test %d accuracy %g"%(epoch,acc))
                # 添加判断语句，选择保存精度最高的模型
                if acc > saver_max_acc:
                    saver.save(sess,'./model/mnist02_CNN/mnist02_CNN.ckpt',global_step=epoch+1)
                    saver_max_acc = acc

end = datetime.datetime.now()
print((end - start).seconds)


#!/usr/bin/python

import sqlite3
import os

def make_rsem_db(db_fname, samples_dir, filepath):
    conn = sqlite3.connect(db_fname)
    cur = conn.cursor()

    cur.execute(
        '''
        CREATE TABLE stringtie_rsem (
        sample TEXT NOT NULL,
        transcript_id TEXT NOT NULL,
        gene_id TEXT NOT NULL,
        length INT NOT NULL,
        eff_length REAL NOT NULL,
        expected_count REAL NOT NULL,
        tpm REAL NOT NULL,
        fpkm REAL NOT NULL,
        isopct REAL NOT NULL
        )
        '''
    )

    for sample in os.listdir(samples_dir):
        fname = os.path.join(samples_dir, sample, filepath)
        data = []
        with open(fname, 'r') as f:
            num_cols = len(f.readline().strip().split())+1
            for line in f:
                row = line.strip().split()
                row[2] = int(row[2])
                row[3:] = map(float, row[3:])
                row = [sample] + row
                data.append(row)
        cur.executemany(
        '''
        INSERT INTO stringtie_rsem VALUES ({})
        '''.format(",".join(["?"]*num_cols)),
        data
        )

    conn.commit()

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("db_fname")
    parser.add_argument("samples_dir")
    parser.add_argument("filepath")
    args = parser.parse_args()

    make_rsem_db(args.db_fname, args.samples_dir, args.filepath)


import os

import docker

from babysage.utils import find_project_root, get_input_dir, get_output_dir
from babysage.config import get_config

def run_local(dockerfile):
    client = docker.from_env()

    project_root = find_project_root()

    if project_root:
        dockerfile = dockerfile or 'Dockerfile'
        tag = f"babysage-{get_config('experiment_name')}",
        client.images.build(
            path=project_root,
            quiet=False,
            rm=True,
            tag=tag,
            dockerfile=dockerfile)
        client.containers.run(
            tag,
            'train',
            volumes={
                os.path.join(project_root, 'input'): {
                    'bind': get_input_dir(), 'mode': 'rw'},
                os.path.join(project_root, 'output'): {
                    'bind': get_output_dir(), 'mode': 'rw'}})
    else:
        raise Exception('''
        Current directry is not a babysage project directry (could not find `babysage.yml` file''')


import rclpy
from rclpy.node import Node
from std_msgs.msg import String
import socket
import threading
import json

class State:
    def __init__(self) -> None:
        self.pose = None
        self.gripper = None

class UrxNodePub(Node):
    def __init__(self, state):
        super().__init__('socket_robot')
        self.state = state
        self.robot_control_pub = self.create_publisher(String, 'robot_control', 1)
        self.robot_status = self.create_subscription(String, 'urx_status', self.robot_data_callback, 1)

    def robot_data_callback(self, data):
        payload = json.loads(data.data)
        self.state.pose = payload["position"]
        self.state.gripper = payload["gripper"]

    def send_command(self, data: str):
        send_data = String()
        send_data.data = data
        self.robot_control_pub.publish(send_data)

class SocketServer:
    def __init__(self, robot_control_pub, state, host: str = "0.0.0.0", port: int = 6666) -> None:
        self.robot_control_pub = robot_control_pub
        self.host = host
        self.port = port
        self.state = state

        self.server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.server_socket.bind((self.host, self.port))
        self.server_socket.listen(5)

    def handle_client(self, client_socket):
        while True:
            try:
                data = client_socket.recv(1024)
                if data:
                    try:
                        data = json.loads(data)
                        if data["type"] == "movel":
                            self.robot_control_pub.send_command(json.dumps({
                                "type": "movel",
                                "data": data["pose"]
                            }))
                        elif data["type"] == "getl":
                            client_socket.sendall(json.dumps({"pose": self.state.pose}).encode("utf-8"))

                        elif data["type"] == "gripper_pose":
                            client_socket.sendall(json.dumps({"pose": self.state.gripper}).encode("utf-8"))

                        elif data["type"] == "gripper_action":
                            self.robot_control_pub.send_command(json.dumps({
                                "type": "gripper",
                                "data": data["pose"]
                            }))
                    except json.decoder.JSONDecodeError:
                        pass

            except ConnectionResetError: # If client disconnect
                client_socket.close()
                break

    def clients_joiner(self):
        while True:
            client_socket, addr = self.server_socket.accept()
            threading.Thread(target=lambda: self.handle_client(client_socket)).start()

def main(args=None):
    rclpy.init(args=args)

    state = State()
    robot_control_pub = UrxNodePub(state)
    socket_server = SocketServer(robot_control_pub, state)
    threading.Thread(target=socket_server.clients_joiner).start()

    rclpy.spin(robot_control_pub)
    robot_control_pub.destroy_node()
    rclpy.shutdown()

if __name__ == "__main__":
    main()


import settings

SECRET_KEY = getattr(settings, "SECRET_KEY", "localhost")

from pymongo import MongoClient

client = MongoClient(SECRET_KEY, 27017, authSource="admin")
db = client.dbproducts

from flask import Blueprint, jsonify, request

product_get = Blueprint('product_get', __name__,)

# 상품 GET
@product_get.route('/product', methods=['GET'])
def get_product():
    item_receive = request.args.get('item_give')

    try:
        sort_receive = int(request.args.get('sort_give'))
    except:
        sort_receive = 0

    if item_receive == 'all':
        # 모든 상품의 list 전달
        coll_arr = db.list_collection_names()
        all_product = []
        for coll in coll_arr:
            all_product += list(db[coll].find({}, {'_id': False}))

        if sort_receive == 1:
            all_product = sorted(
                all_product, key=lambda product: product["like"])
        elif sort_receive == -1:
            all_product = sorted(
                all_product, key=lambda product: product["like"], reverse=True)

        return jsonify({'result': "success", 'documents': all_product})
    else:
        # 하나의 상품 list 전달
        if sort_receive == 0:
            one_product = list(db[item_receive].find({}, {'_id': False}))
        else:
            one_product = list(db[item_receive].find(
                {}, {'_id': False}).sort("like", sort_receive))

        return jsonify({'result': "success", 'documents': one_product})


from decor_timetable import *
import json
n = input('Для просмотра расписания гр. 18704 введите 1: ')
if n == '1':
    lst = []
    with open('time.txt', encoding='utf-8') as f:
        for line in f.readlines():
            a = Day(json.loads(line))
            lst.append(str(a).split(', '))
    print('|{:^20}|{:^20}|{:^20}|{:^20}|{:^20}|'.format('1 пара', '3 пара',
                                                        '2 пара', '4 пара', '5 пара'))
    print('-' *106)
    for i in lst:
        a = Lesson(i)
        print(a)
    print('-' * 106)
else:
    print('До свидания!')


from lxml import etree
import sys
"""
Find missing audio files in each lesson and print outline showing their location.

Author: Carolyn Anderson          Last Modified: 5/12/2015
"""

#Create xsl transformation functions
audio_xsl_raw = etree.parse("audio_find.xsl")
audio_xsl = etree.XSLT(audio_xsl_raw)

def processLessonSet(level, lessonset):
  f = open("../audio.html", "w") #create new file
  f.write(str(audio_xsl(lessonset))) #apply xslt transformation and write to file
  f.close()

doc = etree.parse(sys.argv[1]) #parse the xml doc
lessonset = doc.xpath("/lessonset")[0]#get lessonset
processLessonSet(0, lessonset)#process lessonset


from __future__ import absolute_import
import torch
from torch import nn
from torch.autograd import Variable
import numpy as np
import warnings
from .LOSSES import LOSSES

class DRO_TOPK(nn.Module):
    # Truncated DRO for p_choice = 1

    ####TODO: finish two methods, the third methods need to be done.
    def __init__(self, alpha=40, margin=0.5, beta=0,  K = 5,
                 select_TOPK_all = 1, loss = "margin_loss", **kwargs):

        # self.choose = 1, DRO over batch
        # self.choose = 2, DRO over class
        # self.choose = 3, DRO over each anchor.
        super(DRO_TOPK, self).__init__()
        self.margin = margin
        self.alpha = alpha
        self.beta = beta
        self.select_TOPK_all = select_TOPK_all
        self.K = K
        self.loss = loss
        self.LOSSES = LOSSES(alpha = self.alpha, beta = self.beta, margin = self.margin)
        print("loss:", self.loss, 'alpha:', self.alpha, 'beta:', self.beta,
              "margin:", self.margin, 'K', self.K)

    def forward(self, inputs, targets):
        n = inputs.size(0)
        sim_mat = torch.matmul(inputs, inputs.t())
        targets = targets.cuda()

        eyes_ = Variable(torch.eye(n, n)).cuda()

        #print("eyes_.dtype, targets.dtype:", eyes_.dtype, targets.dtype)

        pos_mask = targets.expand(n,n).eq(targets.expand(n,n).t())
        neg_mask = eyes_.eq(eyes_) - pos_mask # negative pairs index mat
        pos_mask = pos_mask - eyes_.eq(1) # positive pairs index mat

        pos_sim = torch.masked_select(sim_mat, pos_mask)
        neg_sim = torch.masked_select(sim_mat, neg_mask)

        #method_to_call = ,
        #result = method_to_call()

        myloss = getattr(self.LOSSES, self.loss)
        #print("function call", myloss)
        pos_loss, neg_loss = myloss(pos_sim, neg_sim)

        all_loss = torch.cat((pos_loss, neg_loss), 0)
        num_of_zeros = torch.sum(torch.eq(all_loss, 0))

        if self.select_TOPK_all == 1:

            if self.K * 4 > all_loss.size()[0]:
                loss = torch.mean(all_loss)
                warnings.warn("K larger than the total number of pairs in the batch. The loss is calculated using all the pairs by default.")
            else:
                top_loss, _ = torch.topk(all_loss, self.K * 4, largest=True)
                loss = torch.mean(top_loss)

        else:
            if self.K * 2 <= pos_loss.size()[0] and self.K * 2 <= neg_loss.size()[0]:
                top_pos_loss, _ = torch.topk(pos_loss, self.K * 2, largest=True)
                top_neg_loss, _ = torch.topk(neg_loss, self.K * 2, largest=True)
                loss = (torch.sum(top_pos_loss)+torch.sum(top_neg_loss))/(self.K*4)
            else:
                top_neg_loss, _ = torch.topk(neg_loss, pos_loss.size()[0], largest=True)
                loss = (torch.sum(pos_loss) + torch.sum(top_neg_loss)) / (2*pos_loss.size()[0])
                warnings.warn("K either larger than the number of positive pairs or larger than negative pairs. The loss is calculated using all the positive pairs and the same number amount of negative pairs which have the largest losses by default.")

        mean_neg_sim = torch.mean(pos_sim).item()
        mean_pos_sim = torch.mean(neg_sim).item()

        return loss, num_of_zeros.item(), mean_neg_sim, mean_pos_sim

    # def calibration_check_p(self, p):
    #     # remove nan, all o p, and make sure sum(p) = 1
    #
    #     if (torch.sum(torch.isnan(p)) != 0): # remove nan
    #         p[torch.isnan(p)] == 1 / torch.sum(torch.isnan(p))
    #     if (torch.sum(p) != 0):  # make sure the sum of p equal to 1.
    #         p = p / (torch.sum(p))
    #     return p
    #
    #
    # def truncate_p(self, np_p, K_thresh):
    #     order_np_p = -np.sort(-np_p)  # ordering from large to small
    #     np_p[np_p < order_np_p[K_thresh]] = 0
    #     np_p = np_p / sum(np_p)
    #     return np_p
def main():
    data_size = 32
    input_dim = 3
    output_dim = 2
    num_class = 4
    # margin = 0.5
    x = Variable(torch.rand(data_size, input_dim), requires_grad=False)
    # print(x)
    w = Variable(torch.rand(input_dim, output_dim), requires_grad=True)
    inputs = x.mm(w)
    y_ = 8*list(range(num_class))
    targets = Variable(torch.IntTensor(y_))

    DRO_TOPK(K = 1000, select_TOPK_all=2)(inputs, targets)

if __name__ == '__main__':
    main()
    print(LOSSES.__dict__.items())
    print('Congratulations to you!')


#!/usr/bin/env python

import os
# import shutil
import numpy as np
from math import pi, cos, sin, sqrt
import rospy
from std_msgs.msg import Float32
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
from std_msgs.msg import Float32MultiArray
from geometry_msgs.msg import PoseStamped, Twist
from crazyflie_game.msg import Mocap
from player_recorder import DataRecorder

class GameRecorder(object):

	def __init__(self, Ds='', Is='',
				 max_size=1e4,
				 rate=10,
				 logger_dir='res1'):

		self._player_dict = dict()
		for i, D in enumerate(Ds):
			if D != '':
				self._player_dict['D'+str(i+1)] = D
		for i, I in enumerate(Is):
			if I != '':
				self._player_dict['I'+str(i+1)] = I

		script_dir = os.path.dirname(__file__)
		self._results_dir = os.path.join(script_dir, logger_dir + '/')
		# self._results_dir = os.path.join('RAgame/exp_results/', logger_dir + '/')
		self._a_dirc = os.path.join(self._results_dir, 'a.csv')
		if os.path.exists(self._a_dirc):
			os.remove(self._a_dirc)

		# if os.path.exists(self._results_dir):
		# 	shutil.rmtree(self._results_dir)
		if not os.path.isdir(self._results_dir):
			os.makedirs(self._results_dir)

		self._init_time = self._get_time()
		self._save_interval = 50
		self.rate = rospy.Rate(rate)

		self._locations = {'D1': DataRecorder(max_size=max_size),
						   'D2': DataRecorder(max_size=max_size),
						   'I1': DataRecorder(max_size=max_size)}
		self._headings = {'D1': DataRecorder(max_size=max_size),
						  'D2': DataRecorder(max_size=max_size),
						  'I1': DataRecorder(max_size=max_size)}
		self._location_sub_callback_dict = {'D1': self._getLocD1, 'D2': self._getLocD2, 'I1': self._getLocI}
		self._location_subs = dict()
		for p_id, cf_frame in self._player_dict.items():
			self._location_subs.update({p_id: rospy.Subscriber('/' + cf_frame + '/mocap', Mocap, self._location_sub_callback_dict[p_id])})
		self._a_sub = rospy.Subscriber('/a', Float32, self._update_a)
		# self._heading_sub_callback_dict = {'D1': self._getHeadingD1, 'D2': self._getHeadingD2, 'I1': self._getHeadingI}
		# self._heading_subs = dict()
		# for p_id, cf_frame in self._player_dict.items():
		#	 self._heading_subs.update(
		#		 {p_id: rospy.Subscriber('/' + cf_frame + '/heading_anl', Float32, self._heading_sub_callback_dict[p_id])})

		self._locs_plot = self._init_locs_plot()
		# last_res_id = 0
		# for _, dirc, file in os.walk(script_dir):
		#	 for d in dirc:
		#		 if 'res_' in d:
		#			 if int(d.split('_')[-1]) > last_res_id:
		#				 last_res_id = int(d.split('_')[-1])

	def _get_time(self):
		t = rospy.Time.now()
		return t.secs + t.nsecs * 1e-9

	def _update_a(self, a):
		t = self._get_time() - self._init_time
		with open(self._a_dirc, 'a') as f:
			f.write('%.3f, %.3f\n'%(t, a.data))

	def _rotate(self, data):
		return np.array([-data[1], data[0]])

	def _getLocD1(self, data):
		self._locations['D1'].record(self._get_time() - self._init_time, self._rotate(np.array([data.position[0], data.position[1]])))

	def _getLocD2(self, data):
		self._locations['D2'].record(self._get_time() - self._init_time, self._rotate(np.array([data.position[0], data.position[1]])))

	def _getLocI(self, data):
		self._locations['I1'].record(self._get_time() - self._init_time, self._rotate(np.array([data.position[0], data.position[1]])))

	def _getHeadingD1(self, data):
		self._headings['D1'].record(self._get_time() - self._init_time, data.data-pi/2)

	def _getHeadingD2(self, data):
		self._headings['D2'].record(self._get_time() - self._init_time, data.data-pi/2)

	def _getHeadingI(self, data):
		self._headings['I1'].record(self._get_time() - self._init_time, data.data-pi/2)

	def _init_locs_plot(self):
		fig, ax = plt.subplots(tight_layout=True)
		ax.set_xlabel('x(m)')
		ax.set_ylabel('y(m)')
		ax.grid()
		# gs = gridspec.GridSpec(2, 2)
		# axs = [fig.add_subplot(gs[2, 0]) for i in range(2)]
		# axs.append(fig.add_subplot(gs[:, 1]))
		# axs[0].set_xlabel('t(s)')
		# axs[0].set_ylabel('ID1(m)')
		# axs[1].set_xlabel('t(s)')
		# axs[1].set_ylabel('ID2(m)')
		# axs[2].set_xlabel('x(m)')
		# axs[2].set_ylabel('y(m)')

		plt.show(block=False)
		return {'fig': fig, 'axs': ax}

	def plot_locs(self):
		xD1 = np.asarray(self._locations['D1'].data)
		xD2 = np.asarray(self._locations['D2'].data)
		xI = np.asarray(self._locations['I1'].data)

		hD1 = np.asarray(self._headings['D1'].data)
		hD2 = np.asarray(self._headings['D2'].data)
		hI = np.asarray(self._headings['I1'].data)
		# print(hD1)

		def get_cap_ring(xd):
			rs = []
			for tht in np.linspace(0, 6.28, 50):
				x = xd[0] + .25*cos(tht)
				y = xd[1] + .25*sin(tht)
				rs.append((np.array([x, y])))
			return np.asarray(rs)

		def draw_anl_arrow(ax, x, y, heading):
			vx, vy = cos(heading), sin(heading)
			lenv = 10*sqrt(vx**2 + vy**2)
			vx, vy = vx/lenv, vy/lenv
			ax.arrow(x, y, vx, vy, fc='b', ec='b', head_width=.01, zorder=10)

		if len(xD1) > 100 and len(xD2) > 100 and len(xI) > 100:
			self._locs_plot['axs'].clear()

			self._locs_plot['axs'].plot([.5, .5], [-1.2, 1.2,], 'r')
			self._locs_plot['axs'].plot(xD1[100:-1, 0], xD1[100:-1, 1], 'b')
			self._locs_plot['axs'].plot(xD2[100:-1, 0], xD2[100:-1, 1], 'g')
			self._locs_plot['axs'].plot(xI[100:-1, 0], xI[100:-1, 1], 'r')
			ring1 = get_cap_ring(xD1[-1, :])
			ring2 = get_cap_ring(xD2[-1, :])
			self._locs_plot['axs'].plot(ring1[:, 0], ring1[:, 1], 'b')
			self._locs_plot['axs'].plot(ring2[:, 0], ring2[:, 1], 'g')
			if len(hD1) > 1 and len(hD2) > 1 and len(hI) > 1:
				draw_anl_arrow(self._locs_plot['axs'], xD1[-1, 0], xD1[-1, 1], hD1[-1])
				draw_anl_arrow(self._locs_plot['axs'], xD2[-1, 0], xD2[-1, 1], hD2[-1])
				draw_anl_arrow(self._locs_plot['axs'], xI[-1, 0], xI[-1, 1], hI[-1])
			self._locs_plot['fig'].canvas.draw()
			self._locs_plot['axs'].grid()
			self._locs_plot['axs'].axis('equal')
			if int(len(xI)) % self._save_interval == 0:
				self._locs_plot['fig'].savefig(self._results_dir + 'traj.png')

if __name__ == '__main__':

	rospy.init_node('game_recorder', anonymous=True)

	Ds = rospy.get_param("~Ds", '').split(',')
	Is = rospy.get_param("~Is", '').split(',')
	logger_dir = rospy.get_param("~logger_dir", '')
	# vd = rospy.get_param("~vd", 0.)
	# vi = rospy.get_param("~vi", 0.)
	# r = rospy.get_param("~a", 0.)
	# r_close = rospy.get_param("~r_close", 1.)
	# k_close = rospy.get_param("~k_close", .9)

	recorder = GameRecorder(Ds=Ds, Is=Is, logger_dir=logger_dir)

	while not rospy.is_shutdown():
		recorder.plot_locs()


#!/usr/bin/env python

from torch.nn import Linear
import torch
import numpy as np
from json_plus import Serializable

class LinearExtended(Linear, Serializable):
    """Implementation of a linear neural net with sensitivity analysis."""

    def __init__(self, in_features, out_features, bias=None):
        super(LinearExtended, self).__init__(
            in_features, out_features, bias=True
        )

    def overall_sensitivity(self):
        """Returns the sensitivity to adversarial examples of the layer."""
        if self.mod1:
            s = torch.max(torch.max(self.weight, -1)[0], -1)[0].item()
        else:
            s = torch.max(torch.sqrt(torch.sum(self.weight * self.weight, -1)))[0].item()
        s *= np.sqrt(2. / np.e)
        return s

    def interval_forward(self, x_min, x_max):
        w_pos = (self.weight > 0.).float() * self.weight
        w_neg = (self.weight < 0.).float() * self.weight
        y_min = x_min.matmul(w_pos.t()) + x_max.matmul(w_neg.t()) + self.bias
        y_max = x_max.matmul(w_pos.t()) + x_min.matmul(w_neg.t()) + self.bias
        return y_min, y_max

    def sensitivity(self, previous_layer):
        """Given the sensitivity of the previous layer (a vector of length equal
        to the number of inputs), it computes the sensitivity to adversarial examples
         of the current layer, as a vector of length equal to the output size of the
         layer.  If the input sensitivity of the previous layer is None, then unit
         sensitivity is assumed."""
        if previous_layer is None:
            previous_layer = self.weight.new(1, self.in_features)
            previous_layer.fill_(1.)
        else:
            previous_layer = previous_layer.view(1, self.in_features)
        w = previous_layer * self.weight
        s = torch.sum(torch.abs(w), -1)
        return s

    def dumps(self):
        d = dict(
            in_features=self.in_features,
            out_features=self.out_features,
            weight=self.weight.data.cpu().numpy(),
            bias=None if self.bias is None else self.bias.data.cpu().numpy(),
        )
        return Serializable.dumps(d)

    @staticmethod
    def loads(s, device):
        d = Serializable.loads(s)
        m = LinearExtended(
            d['in_features'],
            d['out_features'],
            bias=d['bias'] is not None
        )
        m.weight.data = torch.from_numpy(d['weight']).to(device)
        if d['bias'] is not None:
            m.bias.data = torch.from_numpy(d['bias']).to(device)
        return m