import pandas as pd
from keras.models import Model
from keras.layers import Dense, Input, Dropout, MaxPooling1D, Conv1D
from keras.layers import LSTM, Lambda
from keras.layers import TimeDistributed, Bidirectional
from keras.layers.normalization import BatchNormalization
import numpy as np
import tensorflow as tf
import re
import keras.callbacks
import sys
import os

def binarize(x, sz=71):
    return tf.to_float(tf.one_hot(x, sz, on_value=1, off_value=0, axis=-1))

def binarize_outshape(in_shape):
    return in_shape[0], in_shape[1], 71

def striphtml(html):
    p = re.compile(r'<.*?>')
    return p.sub('', html)

def clean(s):
    return re.sub(r'[^\x00-\x7f]', r'', s)

# record history of training
class LossHistory(keras.callbacks.Callback):
    def on_train_begin(self, logs={}):
        self.losses = []
        self.accuracies = []

    def on_batch_end(self, batch, logs={}):
        self.losses.append(logs.get('loss'))
        self.accuracies.append(logs.get('acc'))

total = len(sys.argv)
cmdargs = str(sys.argv)

print ("Script name: %s" % str(sys.argv[0]))
checkpoint = None
if len(sys.argv) == 2:
    if os.path.exists(str(sys.argv[1])):
        print ("Checkpoint : %s" % str(sys.argv[1]))
        checkpoint = str(sys.argv[1])

data = pd.read_csv("./input/train.csv")
headers = ['id', 'comment_text', 'toxic', 'severe_toxic', 'obscene', 'threat', 'insult', 'identity_hate']
meaningful_headers =  ['comment_text', 'toxic', 'severe_toxic', 'obscene', 'threat', 'insult', 'identity_hate']
data = pd.DataFrame(data, columns = headers)
txt = ''
docs = []
sentences = []
sentiments = []
print("data size,",data.shape)
count = 0
# print (data[meaningful_headers])
length = data.shape[0]
for index, row in data.iterrows():
  entry = row
  cont = entry[1]
  sentiment = entry[2:]
  # print (entry)
  # print (cont, "\n", sentiment)
  sentences = re.split(r'(?<!\w\.\w.)(?<![A-Z][a-z]\.)(?<=\.|\?)\s', clean(striphtml(cont)))
  # sentences = [sent.lower() for sent in sentences]
  # print(sentences)
  # break
  count+=1
  docs.append(sentences)
  sentiments.append(sentiment)
print ("count:", count)
num_sent = []
for doc in docs:
    num_sent.append(len(doc))
    for s in doc:
        txt += s

chars = set(txt)

print('total chars:', len(chars))
char_indices = dict((c, i) for i, c in enumerate(chars))
indices_char = dict((i, c) for i, c in enumerate(chars))
print("docs size", len(docs))
# print ("!!!!docs::", docs)
print('Sample doc{}'.format(docs[1200]))

maxlen = 200
max_sentences = 15

X = np.ones((len(docs), max_sentences, maxlen), dtype=np.int64) * -1
y = np.array(sentiments)

for i, doc in enumerate(docs):
    for j, sentence in enumerate(doc):
        if j < max_sentences:
            for t, char in enumerate(sentence[-maxlen:]):
                X[i, j, (maxlen - 1 - t)] = char_indices[char]

print('Sample chars in X:{}'.format(X[1200, 2]))
print('y:{}'.format(y[1200,:]))

ids = np.arange(len(X))
np.random.shuffle(ids)

# shuffle
X = X[ids]
y = y[ids]
rows = X.shape[0]
X_train = X[:rows]
X_test = X[-100:]

y_train = y[:rows]
y_test = y[-100:]

filter_length = [2]
nb_filter = [100]
pool_length = 2
# document input
document = Input(shape=(max_sentences, maxlen), dtype='int64')
# sentence input
in_sentence = Input(shape=(maxlen,), dtype='int64')
# char indices to one hot matrix, 1D sequence to 2D
embedded = Lambda(binarize, output_shape=binarize_outshape)(in_sentence)
# embedded: encodes sentence
for i in range(len(nb_filter)):
    embedded = Conv1D(filters=nb_filter[i],
                      kernel_size=filter_length[i],
                      padding='valid',
                      activation='relu',
                      kernel_initializer='glorot_normal',
                      strides=1)(embedded)

    embedded = Dropout(0.1)(embedded)
    embedded = MaxPooling1D(pool_size=pool_length)(embedded)

bi_lstm_sent = \
    Bidirectional(LSTM(128, return_sequences=False, dropout=0.15, recurrent_dropout=0.15, implementation=0))(embedded)

# sent_encode = merge([forward_sent, backward_sent], mode='concat', concat_axis=-1)
sent_encode = Dropout(0.3)(bi_lstm_sent)

# sentence encoder
encoder = Model(inputs=in_sentence, outputs=sent_encode)
encoder.summary()

encoded = TimeDistributed(encoder)(document)
# encoded: sentences to bi-lstm for document encoding
b_lstm_doc = \
    Bidirectional(LSTM(128, return_sequences=False, dropout=0.15, recurrent_dropout=0.15, implementation=0))(encoded)

output = Dropout(0.3)(b_lstm_doc)
output = Dense(128, activation='relu')(output)
output = Dropout(0.3)(output)
output = Dense(6, activation='sigmoid')(output)

model = Model(inputs=document, outputs=output)

model.summary()

if checkpoint:
    model.load_weights(checkpoint)

file_name = os.path.basename(sys.argv[0]).split('.')[0]
#check_cb = keras.callbacks.ModelCheckpoint('checkpoints/' + file_name + '.{epoch:02d}-{val_loss:.2f}.hdf5',
check_cb = keras.callbacks.ModelCheckpoint('checkpoints/' + file_name + '.hdf5',
                                           monitor='val_loss',
                                           verbose=0, save_best_only=True, mode='min')
earlystop_cb = keras.callbacks.EarlyStopping(monitor='val_loss', patience=7, verbose=1, mode='auto')
history = LossHistory()
model.compile(loss='binary_crossentropy', optimizer='adam', metrics=['accuracy'])

model.fit(X_train, y_train, validation_data=(X_test, y_test), batch_size=10,
          epochs=2, shuffle=True, callbacks=[earlystop_cb, check_cb, history])

# just showing access to the history object
print (history.losses)
print (history.accuracies)


import datetime
import time
import pandas as pd

import numpy as np

greetings = {
  "early_morning": ["wow, so early in the morning, nice {}", "you have a great start {}"],
  "morning": ["good morning {}", "happy morning {}", "have a nice morning {}", "such a lovely morning {}"],
  "noon": ["good afternoon {}", "happy afternoon {}", "have a nice afternoon {}", "such a lovely afternoon {}"],
  "evening": ["good evening {}", "happy evening {}", "have a nice evening {}", "such a lovely evening {}"],
  "night":  ["good night {}", "have a nice sleep {}", "you should rest now {}", "go and sleep {}"],
  "late_night": ["it's so dark now, you should sleep {}", "you did it, you have shown your hardwork, go and rest {}"]
}

def checkGreetingType(author):
  # current_time = datetime.datetime.now()
  try:
    current_hour = datetime.time.hour.getter
  except AttributeError:
    print(">>> AttributeError found. The command -ds greet has some invalid method which is not working now.")
    return """
        `-ds greet` command is not working because of some backend bugs. We are trying to fix it for you
    """

  print(current_hour)

  # if current_hour <= 14:
    # current_hour += 10          # write from here
  # current_hour = pd.datetime.now()
  # current_minutes = current_time.minute

  greeting_type = ""

  if (current_hour > 4) and (current_hour <= 8): greeting_type = "early_morning" # early morning
  elif (current_hour > 8) and (current_hour <= 12): greeting_type = "morning"    # morning
  elif (current_hour > 12) and (current_hour <= 16): greeting_type = "noon"          # noon
  elif (current_hour > 16) and (current_hour <= 20): greeting_type = "evening"            # evening
  elif (current_hour > 20) and (current_hour <= 24): greeting_type = "night"     # night
  elif (current_hour <= 4): greeting_type = "late_night"      # late night

  message = ""
  message_list = []

  if greeting_type in greetings.keys():
    message_list = greetings.get(greeting_type)

  if greeting_type != "":
    messageIndex = np.random.randint(0, len(message_list))
    message = message_list[messageIndex]

  return """
    {}
  """.format(message).format(author)
# checkGreetingType()     # testing this feature


import torch
import torchvision
from faster_rcnn_utils.engine import evaluate
from faster_rcnn_utils.AIZOODataset import AIZOODataset
from faster_rcnn_utils.transforms import get_transform
from faster_rcnn_utils import utils
import os
import numpy as np
# os.environ["CUDA_VISIBLE_DEVICES"] = '0'

def eva_rcnn(model,data_loader_val, device):
    test_path = './val'

    # train on the GPU or on the CPU, if a GPU is not available
    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
    print(device)

    model.to(device)

    # evaluate on the test dataset
    coco_evaluator=evaluate(model, data_loader_val, device=device)
    mAP_one=None
    for coco_Evaluator in coco_evaluator.coco_eval.values():
        mAP_one=np.mean([coco_Evaluator.eval['precision'][:,0:101:10,0,0,2].mean(),coco_Evaluator.eval['precision'][:,0:101:10,1,0,2].mean()])
    return mAP_one


# -*- coding: utf-8 -*-

from __future__ import absolute_import
from __future__ import print_function
from __future__ import division

import numpy as np
import tensorflow as tf
import os

from . import image_preprocess
from net.box_utils import train_bbox_anchor_tf_op
import glob

class Data_Manager(object):

    def __init__(self, data_record_path, batch_size , is_training, image_size, anchors, class_numer, scale_factors, anchor_pos_iou, data_format = "NHWC"):

        self.data_record_path = data_record_path
        self.batch_size = batch_size
        self.is_training = is_training
        self.image_size = image_size
        self.anchor_number = len(np.concatenate(anchors, axis=0))
        self.anchors_tensor = tf.constant(np.concatenate(anchors, axis=0),dtype=tf.float32)
        self.class_number = class_numer
        self.anchor_scale = scale_factors
        self.anchor_pos_iou = anchor_pos_iou
        self.total_sample_number = self.countTFRecordSampleNumber(self.data_record_path)
        # self.total_sample_number = 0

        self.img_name_batch, self.img_batch, self.gtboxes_and_label_batch_float, self.num_obs_batch, self.img_height, self.img_width = \
            self.init_data_manager(self.batch_size, self.image_size, self.is_training, self.anchors_tensor, data_format)

    def read_single_example_and_decode(self, filename_queue):

        reader = tf.TFRecordReader()

        _, serialized_example = reader.read(filename_queue)

        features = tf.parse_single_example(
            serialized=serialized_example,
            features={
                'img_name': tf.FixedLenFeature([], tf.string),
                'img_height': tf.FixedLenFeature([], tf.int64),
                'img_width': tf.FixedLenFeature([], tf.int64),
                'img': tf.FixedLenFeature([], tf.string),
                'gtboxes_and_label': tf.FixedLenFeature([], tf.string),
                'num_objects': tf.FixedLenFeature([], tf.int64)
            }
        )

        img_name = features['img_name']
        img_height = tf.cast(features['img_height'], tf.int32)
        img_width = tf.cast(features['img_width'], tf.int32)
        img = tf.decode_raw(features['img'], tf.uint8)

        img = tf.reshape(img, shape=[img_height, img_width, 3])

        gtboxes_and_label = tf.decode_raw(features['gtboxes_and_label'], tf.int32)
        gtboxes_and_label = tf.reshape(gtboxes_and_label, [-1, 5])
        # gtboxes_and_label.set_shape([None, 5])

        num_objects = tf.cast(features['num_objects'], tf.int32)

        return img_name, img, gtboxes_and_label, num_objects , img_height, img_width

    #在训练阶段,去预处理训练用到的图片
    def process_image_for_training(self, filename_queue, image_size, data_format):

        #读一张图片的信息
        img_name, img, gtboxes_and_label, num_objects, img_height, img_width = self.read_single_example_and_decode(filename_queue)

        #box坐标信息 转换成 相对于图像 宽度和高度的信息
        #gt_box_and_label_tensor ymin xmin ymax xmax 相对原始宽度和高度
        gt_box_and_label_tensor = image_preprocess.box_info_normilization(gtboxes_and_label, img_height, img_width)

        #随机的采样
        #经过随机采样后的形成的box
        image_tensor, labels, bboxes, distort_bbox =\
            image_preprocess.sample_distorted_bounding_box_crop(img, tf.reshape(gt_box_and_label_tensor[:,4], shape=[-1,1]), gt_box_and_label_tensor[:,:4] ,min_object_covered=0.8, area_range=(0.75, 1.0), aspect_ratio_range=(0.75, 1.33))

        #随机左右翻转操作
        #此时box 坐标会变化
        image_tensor, bboxes = image_preprocess.random_flip_left_right(image_tensor, bboxes)

        #在最后一个步骤进行resize操作
        image_tensor = image_preprocess.resize_image_with_fixed_size(image_tensor, image_size, image_size,
                                                                          method=tf.image.ResizeMethod.BILINEAR,
                                                                          align_corners=False)

        # #随机进行对图片的修正 ,像素变更为0 到 1之间的类型
        image_tensor = image_preprocess.convert_image_format(image_tensor)

        image_tensor = image_preprocess.random_adjust_image_pixes(image_tensor)

        image_tensor = image_preprocess.convert_image_format(image_tensor,type="0_255")

        #去中心化操作
        image_tensor = image_preprocess.image_whitened(image_tensor)

        #是否转换
        if data_format == "NCHW":

            image_tensor = tf.transpose(image_tensor, (2,0,1))

        #合并label和gt box信息
        gt_box_and_label_tensor = tf.concat([bboxes, labels],axis=1)

        return img_name, image_tensor, gt_box_and_label_tensor, num_objects , img_height, img_width

    def process_image_for_testing(self,filename_queue, image_size, data_format):

        #读一张图片的信息
        img_name, img, gtboxes_and_label, num_objects, img_height, img_width = self.read_single_example_and_decode(filename_queue)

        image_tensor = img

        #box坐标信息 转换成 相对于图像 宽度和高度的信息
        #gt_box_and_label_tensor ymin xmin ymax xmax 相对原始宽度和高度
        gt_box_and_label_tensor = image_preprocess.box_info_normilization(gtboxes_and_label, img_height, img_width)

        #在最后一个步骤进行resize操作
        image_tensor = image_preprocess.resize_image_with_fixed_size(image_tensor, image_size, image_size,
                                                                          method=tf.image.ResizeMethod.BILINEAR,
                                                                          align_corners=False)

        #去中心化操作
        image_tensor = image_preprocess.image_whitened(image_tensor)

        if data_format == "NCHW":

            image_tensor = tf.transpose(image_tensor, (2,0,1))

        return img_name, image_tensor, gt_box_and_label_tensor, num_objects, img_height, img_width

    def read_and_prepocess_single_img(self, filename_queue, image_size, is_training=True, data_format ="NHWC"):

        if is_training == True:

            img_name, distort_image_tensor, gt_box_and_label_tensor, num_objects, img_height, img_width = \
                self.process_image_for_training(filename_queue, image_size, data_format)

        else:

            img_name, distort_image_tensor, gt_box_and_label_tensor, num_objects, img_height, img_width = \
                self.process_image_for_testing(filename_queue, image_size, data_format)

        return img_name, distort_image_tensor, gt_box_and_label_tensor, num_objects, img_height, img_width

    def encode_ground_truth_box_label_for_train(self, anchors_tensor, anchor_number, gtboxes_and_label_tensor, scale_factors , class_number , pos_iou_threshold = 0.5):

        return train_bbox_anchor_tf_op.generate_train_encoding_labels_tf_operation(anchors_tensor, anchor_number, gtboxes_and_label_tensor[:,:4], gtboxes_and_label_tensor[:,4], scale_factors, class_number, pos_iou_threshold=pos_iou_threshold)

    def init_data_manager(self, batch_size, image_size, is_training, anchors, data_format):

        #这种写法只能是训练过程中不能进行验证或者test操作.
        pattern = os.path.join(self.data_record_path, '*.tfrecord')

        print('tfrecord path is -->', os.path.abspath(pattern))

        filename_tensorlist = tf.train.match_filenames_once(pattern)

        filename_queue = tf.train.string_input_producer(filename_tensorlist)

        img_name, img, gtboxes_and_label_float, num_obs, img_height, img_width = self.read_and_prepocess_single_img(filename_queue, image_size, is_training=is_training, data_format= data_format)

        #将label信息进行编码操作
        encode_gt_box_label_tensor = self.encode_ground_truth_box_label_for_train(self.anchors_tensor, self.anchor_number, gtboxes_and_label_float, self.anchor_scale, self.class_number, self.anchor_pos_iou)

        #gtboxes_and_label_float 临时添加
        img_name_batch, img_batch, encode_gt_box_label_tensor, num_obs_batch, img_height, img_width = \
            tf.train.batch(
                           [img_name, img, encode_gt_box_label_tensor, num_obs, img_height, img_width],
                           batch_size=batch_size,
                           capacity=50,
                           num_threads=4,
                           dynamic_pad=True)

        # img_tensor_with_bound_box_tensor = tf.image.draw_bounding_boxes(img_batch, gtboxes_and_label_batch_float[:,:4])

        return img_name_batch, img_batch, encode_gt_box_label_tensor, num_obs_batch, img_height, img_width

    def next_batch(self):

        return self.img_name_batch, self.img_batch, self.gtboxes_and_label_batch_float, self.num_obs_batch, self.img_height, self.img_width

    def countTFRecordSampleNumber(self, tfrecordPath):

        tf_record_path_pattern = tfrecordPath + "/*.tfrecord"

        recordList = glob.glob(tf_record_path_pattern)

        if recordList is None:

            raise Exception("Count The Number Of Sample Failure")

        number = 0

        for tfrecord in recordList:

            for record in tf.python_io.tf_record_iterator(tfrecord):
                number = number + 1

        return number


#!/usr/bin/env python3
import os
import json
import argparse
import requests
import sqlite3
import math
from astropy.time import Time
from zchecker import ZChecker, Config

parser = argparse.ArgumentParser(epilog='Requires environment variables ZBROWSER_BASE_URL, '
                                 'SLACK_USER_POST_URL, and SLACK_ZTFSSWG_POST_URL')
parser.add_argument('--date', help='generate report for this date')
parser.add_argument('-n', action='store_true', dest='debug', help='no-op mode for testing')
parser.add_argument('--user', action='store_true', help='post message to user for testing')
args = parser.parse_args()

base_url = os.getenv('ZBROWSER_BASE_URL')
user_post_url = os.getenv('SLACK_USER_POST_URL')
ztfsswg_post_url = os.getenv('SLACK_ZTFSSWG_POST_URL')

config = Config.from_file()
with ZChecker(config=config, save_log=False) as z:
    z.db.create_function('sqrt', 1, math.sqrt)
    z.db.create_function('pow', 2, math.pow)

    targets = {}
    for i, d in z.db.execute('SELECT objid,desg FROM obj').fetchall():
        targets[i] = d

    rows = z.db.execute('''
SELECT objid,desg FROM obj
WHERE desg GLOB '[CP]*'
   OR desg GLOB '*[0-9]P'
   OR desg in ('2060', '4015', '7968', '60558', '118401',
               '323137', '300163', '457175')
''').fetchall()
    comets = {}
    for i, d in rows:
        comets[i] = d

    is_comet = 'objid IN ({})'.format(','.join(
        [str(i) for i in comets.keys()]))

    if args.date is None:
        last_nightid, last_night, exposures = z.db.execute('''
        SELECT nightid,date,exposures FROM ztf_nights ORDER BY date DESC LIMIT 1
        ''').fetchone()
    else:
        last_nightid, last_night, exposures = z.db.execute('''
        SELECT nightid,date,exposures FROM ztf_nights WHERE date=:date
        ''', {'date': args.date}).fetchone()

    in_last_night = 'nightid={}'.format(last_nightid)

    observed = []
    for vmin, vmax in zip([22, 20, 18, -99], [23, 22, 20, 18]):
        observed.append(z.db.execute('''
        SELECT COUNT(DISTINCT objid) FROM ztf_found
        WHERE ''' + is_comet + '''
          AND vmag > ?
          AND vmag <= ?
          AND (ra3sig*ra3sig + dec3sig*dec3sig) < 22500
          AND ''' + in_last_night, [vmin, vmax]).fetchone()[0])

    total = z.db.execute('''
        SELECT COUNT(DISTINCT objid) FROM ztf_found
        WHERE ''' + is_comet + '''
          AND (ra3sig*ra3sig + dec3sig*dec3sig) < 22500
          AND ''' + in_last_night).fetchone()[0]

    rows = z.db.execute('''
    SELECT objid,AVG(vmag) AS meanv FROM ztf_found
    WHERE ''' + is_comet.format(col='desg') + '''
    AND vmag IS NOT NULL
    AND (ra3sig*ra3sig + dec3sig*dec3sig) < 22500
    AND ''' + in_last_night + '''
    GROUP BY objid
    ORDER BY meanv LIMIT 5
    ''').fetchall()

    if len(rows) == 0:
        brightest_comets = '(not applicable)'
    else:
        brightest_comets = '\n'.join([
            ('<' + base_url + '?obs-by-target={target}|{target}> ({v:.1f})')
            .format(target=comets[row[0]], v=row[1])
        for row in rows])

    rows = z.db.execute('''
    SELECT objid,count() AS c FROM ztf_found
    WHERE ''' + is_comet.format(col='desg') + '''
    AND ''' + in_last_night + '''
    GROUP BY objid
    ORDER BY c DESC
    ''').fetchall()

    if len(rows) == 0:
        best_observed_comets = '(not applicable)'
    else:
        n = max(sum([row[1] > 9 for row in rows]), 5)
        best_observed_comets = '\n'.join([
            ('<' + base_url + '?obs-by-target={target}|{target}> ({c})')
            .format(target=comets[row[0]], c=row[1])
            for row in rows[:n]])

    rows = z.db.execute('''
    SELECT objid,MAX(ostat) AS max_ostat FROM ztf_phot
    INNER JOIN ztf_found USING (foundid)
    WHERE ''' + in_last_night + '''
      AND ostat NOT NULL
      AND ostat >= 3
    GROUP BY objid
    ''').fetchall()

    if len(rows) == 0:
        outbursts = 'No suspected outbursts.'
    else:
        outbursts = 'Possible outbursts:\n  ' + '\n  '.join([
            ('<' + base_url + '?obs-by-target={target}|{target}> ({ostat:.1f})')
            .format(target=targets[row[0]], ostat=row[1])
            for row in rows])

    # check for possible ephemeris updates
    rows = z.db.execute('''
    SELECT objid,POW(POW(ra3sig - last_ra3sig, 2) + POW(dec3sig - last_dec3sig, 2), 0.5) / (obsjd - last_obsjd) AS d
    FROM (
        SELECT nightid,objid,obsjd,
               tmtp,LAG(obsjd) OVER (PARTITION BY objid ORDER BY obsjd) AS last_obsjd,
               ra3sig,LAG(ra3sig) OVER (PARTITION BY objid ORDER BY obsjd) AS last_ra3sig,
               dec3sig,LAG(dec3sig) OVER (PARTITION BY objid ORDER BY obsjd) AS last_dec3sig
        FROM ztf_found ORDER BY obsjd DESC
    ) WHERE ''' + in_last_night + ''' AND last_ra3sig > ra3sig AND last_dec3sig > dec3sig AND d > 0.3
    GROUP BY objid ORDER BY d DESC;
    ''').fetchall()
    if len(rows) == 0:
        ephemeris_updates = '';
    else:
        ephemeris_updates = '\nPossible ephemeris updates:\n  ' + '\n  '.join([
            ('<' + base_url + '?obs-by-target={target}|{target}> ({d:.1f})')
            .format(target=targets[row[0]], d=row[1])
            for row in rows])

summary = f'On the night of <{base_url}?obs-by-date={last_night}|{last_night}>, out of {exposures} exposure{"" if exposures == 1 else "s"}, {total} comet{" was" if total == 1 else "s were"} found.\n{outbursts}{ephemeris_updates}\n\n'

if exposures == 0:
    text = 'No exposures were taken on {}.'.format(last_night)
    msg = {
        'attachments': [
            {
                'fallback': text,
                'title': f'<{base_url}|ZChecker> Daily Summary',
                'text': text
            }
        ]
    }
else:
    msg = {
        'attachments': [
            {
                'fallback': 'ZChecker daily summary for ' + last_night,
                'title': f'<{base_url}|ZChecker> Daily Summary',
                'text': summary,
                'fields': [
                    {
                        'title': 'V (mag)',
                        'value': '22<V≤23\n20<V≤22\n18<V≤20\n≤18',
                        'short': True
                    },
                    {
                        'title': 'Comets covered',
                        'value': '\n'.join([str(n) for n in observed]),
                        'short': True
                    },
                    {
                        'title': 'Brightest comets:',
                        'value': brightest_comets,
                        'short': False
                    },
                    {
                        'title': 'Best-observed comets:',
                        'value': best_observed_comets,
                        'short': False
                    }
                ]
            }
        ]
    }

if args.debug:
    print(json.dumps(msg, indent='  '))
elif args.user:
    r = requests.post(user_post_url, data=json.dumps(msg))
else:
    r = requests.post(ztfsswg_post_url, data=json.dumps(msg))


import copy
import random
import torch
import numpy as np
import torch.nn.functional as F
from train_cifar10 import build_model, load_model, prepare_dataset, prepare_dataset_without_normalization
from test import load_trigger, make_SP_test_dataset, init_trigger, apply_trigger_on_inputs
from test import get_box_trigger_func, get_test_preprocess_func, get_pattern_trigger_func
from tqdm import tqdm
from IBAU import hypergrad as hg

device = 'cuda' if torch.cuda.is_available() else 'cpu'

def get_results(model, criterion, data_loader, device):
    model.eval()
    val_loss = 0
    correct = 0
    total = 0
    with torch.no_grad():
        for batch_idx, (inputs, targets) in enumerate(data_loader):
            inputs, targets = inputs.to(device), targets.to(device)
            outputs = model(inputs)
            loss = criterion(outputs, targets.long())

            val_loss += loss.item()
            _, predicted = outputs.max(1)
            total += targets.size(0)
            correct += predicted.eq(targets).sum().item()
        return correct / total

def main():
    model_path = './checkpoint/ckpt_trojan_4vs1.pth'
    # model_path = './checkpoint/box_4x4_resnet18.pth'
    # model_path = './checkpoint/trojan_0.8.pth'
    trigger_path = './checkpoint/trigger_first_try_0.8.pth'

    trainset, testset = prepare_dataset()

    att_val_set = copy.deepcopy(testset)
    src_lb = 3
    tgt_lb = 5
    mask_tanh_tensor, pattern_tanh_tensor, src_lb, tgt_lb = load_trigger(trigger_path)
    trigger_func = get_pattern_trigger_func(mask_tanh_tensor, pattern_tanh_tensor)

    # trigger_func = get_box_trigger_func()
    att_val_set = make_SP_test_dataset(att_val_set, src_lb=src_lb, tgt_lb=tgt_lb, trigger_func=trigger_func)

    # data loader for verifying the clean test accuracy
    clnloader = torch.utils.data.DataLoader(
        testset, batch_size=200, shuffle=False, num_workers=2)

    # data loader for verifying the attack success rate
    poiloader_cln = torch.utils.data.DataLoader(
        testset, batch_size=200, shuffle=False, num_workers=2)

    poiloader = torch.utils.data.DataLoader(
        att_val_set, batch_size=200, shuffle=False, num_workers=2)

    _, testset_wo = prepare_dataset_without_normalization()
    preprocess_func = get_test_preprocess_func()

    # data loader for the unlearning step
    unlloader = torch.utils.data.DataLoader(
        testset_wo, batch_size=100, shuffle=False, num_workers=2)

    classes = ('plane', 'car', 'bird', 'cat', 'deer',
               'dog', 'frog', 'horse', 'ship', 'truck')

    def get_next_data_func(dataloader):
        _iter = iter(dataloader)

        def _getnext():
            try:
                data = next(_iter)
            except:
                _iter = iter(dataloader)
                data = next(_iter)
            return data

        return _getnext

    data_func = get_next_data_func(unlloader)
    mask_tanh_tensor, pattern_tanh_tensor, l1_target = init_trigger(size=32)

    # define the inner loss L2
    def loss_inner(inner_params, model_params):
        images, labels = data_func()
        tgt_labs = (labels == src_lb)
        labels[tgt_labs] = tgt_lb
        images, labels = images.cuda(), labels.long().cuda()
        _mask_tanh_tensor, _pattern_tanh_tensor = inner_params
        inputs, l1_loss = apply_trigger_on_inputs(images, labels=tgt_labs, \
                                                  mask_tanh_tensor=_mask_tanh_tensor,
                                                  pattern_tanh_tensor=_pattern_tanh_tensor, \
                                                  l1_target=l1_target)
        per_img = preprocess_func(inputs)
        per_logits = model.forward(per_img)
        loss = F.cross_entropy(per_logits, labels, reduction='none')
        loss_regu = torch.mean(loss) + 0.001 * l1_loss
        # print(l1_loss)
        return loss_regu

    # define the outer loss L1
    def loss_outer(inner_params, model_params):
        images, labels = data_func()
        tgt_labs = (labels == src_lb)
        images, labels = images.cuda(), labels.long().cuda()
        _mask_tanh_tensor, _pattern_tanh_tensor = inner_params
        poi_labels = torch.zeros_like(labels, device='cuda')
        inputs, _ = apply_trigger_on_inputs(images, labels=tgt_labs, \
                                            mask_tanh_tensor=_mask_tanh_tensor,
                                            pattern_tanh_tensor=_pattern_tanh_tensor, \
                                            l1_target=l1_target)
        unlearn_imgs = preprocess_func(inputs)
        logits = model.forward(unlearn_imgs)
        loss = F.cross_entropy(logits, labels)
        return loss

    inner_opt = hg.GradientDescent(loss_inner, 0.1)

    net = build_model()
    model, best_acc, start_epoch, _ = load_model(net, model_path)
    # model.eval()
    # outer_optimizer = torch.optim.Adam(params=model.parameters(), lr=0.001)
    outer_optimizer = torch.optim.SGD(params=model.parameters(), lr=0.002)
    criterion = torch.nn.CrossEntropyLoss()

    ACC = get_results(model, criterion, clnloader, device)
    ASR = get_results(model, criterion, poiloader, device)

    print('Original ASR:', ASR)
    print('Original ACC:', ACC)

    ASR_list = [ASR]
    ACC_list = [ACC]

    T = 50
    K = 50
    l1_limit = 0.1 * (32 * 32)

    for round in range(10):
        mask_tanh_tensor, pattern_tanh_tensor, l1_target = init_trigger(size=32)
        # batch_pert = torch.zeros([1, 3, 32, 32], requires_grad=True, device='cuda')
        # batch_opt = torch.optim.SGD(params=[mask_tanh_tensor, pattern_tanh_tensor], lr=10)
        inner_optimizer = torch.optim.Adam(params=[mask_tanh_tensor, pattern_tanh_tensor],
                                           lr=0.1)  # you can fine-tune this

        inner_params = [mask_tanh_tensor, pattern_tanh_tensor]
        outer_params = list(model.parameters())

        model.eval()
        for t in range(T):
            inner_optimizer.zero_grad()
            loss_regu = loss_inner(inner_params, outer_params)
            loss_regu.backward(retain_graph=True)
            inner_optimizer.step()

            # l1-ball
            mask_tanh_tensor = mask_tanh_tensor * min(1, l1_limit / torch.norm(mask_tanh_tensor, p=1).data)

        # l2-ball
        # pert = batch_pert * min(1, 10 / torch.norm(batch_pert))

        # unlearn step
        # for batchnum in range(len(images_list)):
        #     print('=================', batchnum)
        outer_optimizer.zero_grad()
        hg.fixed_point(inner_params, outer_params, K, inner_opt, loss_outer, stochastic=True)
        outer_optimizer.step()

        model.eval()
        asr = get_results(model, criterion, poiloader, device)
        acc = get_results(model, criterion, clnloader, device)
        ASR_list.append(asr)
        ACC_list.append(acc)
        print('Round:', round)
        print('ASR:', asr)
        print('ACC:', acc)

if __name__ == '__main__':
    main()


from __future__ import absolute_import, division, print_function

import argparse
from datetime import datetime
import numpy as np
import tensorflow as tf
import os

from keras import backend as K
from keras.callbacks import EarlyStopping, TensorBoard
import numpy as np
from sklearn.metrics import roc_auc_score, confusion_matrix

from prediction.GRUDModel.models import create_grud_model, load_grud_model
from prediction.GRUDModel.callbacks import ModelCheckpointwithBestWeights
from datautils.dataset import Dataset
from datautils.helper import save_output

#%%
# set GPU usage for tensorflow backend
if K.backend() == 'tensorflow':
    import tensorflow as tf
    config = tf.ConfigProto()
    config.gpu_options.per_process_gpu_memory_fraction = .1
    config.gpu_options.allow_growth = True
    K.set_session(tf.Session(config=config))

#%%
# parse arguments
## general
arg_parser = argparse.ArgumentParser()
arg_parser.add_argument('--working_path', default=os.path.join("prediction","GRUDModel"))

## data
arg_parser.add_argument('dataset_name', default='physionet',
                        help='The data files should be saved in [working_path]/data/[dataset_name] directory.')
arg_parser.add_argument('label_name', default='sepsis')
# arg_parser.add_argument('--max_timesteps', type=int, default=200,
#                         help='Time series of at most # time steps are used. Default: 200.')
# arg_parser.add_argument('--max_timestamp', type=int, default=48*60*60,
#                         help='Time series of at most # seconds are used. Default: 48 (hours).')

## model
arg_parser.add_argument('--recurrent_dim', type=lambda x: x and [int(xx) for xx in x.split(',')] or [], default='64')
arg_parser.add_argument('--hidden_dim', type=lambda x: x and [int(xx) for xx in x.split(',')] or [], default='64')
arg_parser.add_argument('--model', default='GRUD', choices=['GRUD', 'GRUforward', 'GRU0', 'GRUsimple'])
arg_parser.add_argument('--use_bidirectional_rnn', default=False)

## training
arg_parser.add_argument('--pretrained_model_file', default=None,
                        help='If pre-trained model is provided, training will be skipped.') # e.g., [model_name]_[i_fold].h5
arg_parser.add_argument('--epochs', type=int, default=100)
arg_parser.add_argument('--early_stopping_patience', type=int, default=10)
arg_parser.add_argument('--batch_size', type=int, default=100)
arg_parser.add_argument('--impute_forward', type=int, default=0)
arg_parser.add_argument('--seed', type=int, default=42)
arg_parser.add_argument('--data', default="sepsis_data")
arg_parser.add_argument('--imputation_mode', default="decay")

## set the actual arguments if running in notebook
if not (__name__ == '__main__' and '__file__' in globals()):
    ARGS = arg_parser.parse_args([
        'physionet',
        'sepsis',
        '--model', 'GRUD',
        '--hidden_dim', '',
        '--epochs', '100'
    ])
else:
    ARGS = arg_parser.parse_args()

print('Arguments:', ARGS)

#%%
dataset = Dataset("data/"+ARGS.data,
                    batchSize=100,
                    train_ratio=0.8,
                    normalize=True,
                    padding=True,
                    imputeForward=(True if ARGS.impute_forward else False),
                    seed=ARGS.seed)

dataset.test_data.batchSize = 1

#%%
def roc_auc_score_mod(y_true, y_pred):
    try:
        auc_val = roc_auc_score(np.reshape(y_true, (-1,1)), np.reshape(y_pred, (-1,1)))
    except ValueError:
        auc_val = 0.0
    return np.float32(auc_val)

def auroc(y_true, y_pred):
    return tf.py_func(roc_auc_score_mod, (y_true, y_pred), tf.float32)

def remove_padding(y_true, y_pred, data):
    targets = []
    predictions = []
    predictions_ind = []
    files = []

    for i in range(y_pred.shape[0]):
        targets.extend(list(y_true[i, 0:data.x_lengths[i]]))
        predictions.extend(list(y_pred[i, 0:data.x_lengths[i]]))
        predictions_ind.append(list(y_pred[i, 0:data.x_lengths[i]]))
        files.append(data.files[i])

    targets = np.reshape(np.array(targets), (-1,1))
    predictions = np.reshape(np.array(predictions), (-1,1))

    return targets, predictions, predictions_ind, files

def get_auc(y_true, y_pred, data):
    targets = []
    predictions = []
    files = []

    for i in range(len(data.x_lengths)):
        targets.extend(list(y_true[i, 0:data.x_lengths[i]]))
        predictions.extend(list(y_pred[i, 0:data.x_lengths[i]]))
        files.append(data.files[i])

    targets = np.reshape(np.array(targets), (-1,1))
    predictions = np.reshape(np.array(predictions), (-1,1))

    return roc_auc_score(targets, predictions)

# k-fold cross-validation
pred_y_list_all = []
auc_score_list_all = []

timestamp = datetime.now().strftime('%Y%m%d_%H%M%S_%f')
print('Timestamp: {}'.format(timestamp))

#for i_fold in range(1):
i_fold = 0
print('{}-th fold...'.format(i_fold))

# Load or train the model.
if ARGS.pretrained_model_file is not None:
    model = load_grud_model(os.path.join(ARGS.working_path,
                                            ARGS.pretrained_model_file.format(i_fold=i_fold)))
else:
    model = create_grud_model(input_dim=dataset.train_data.input_dim,
                                output_dim=dataset.train_data.output_dim,
                                output_activation="sigmoid",
                                recurrent_dim=ARGS.recurrent_dim,
                                hidden_dim=ARGS.hidden_dim,
                                predefined_model=ARGS.model,
                                use_bidirectional_rnn=ARGS.use_bidirectional_rnn,
                                input_decay=('exp_relu' if ARGS.imputation_mode == "decay" else None),
                                x_imputation=('raw' if ARGS.imputation_mode == "mean" else 'forward')
                                )
    if i_fold == 0:
        model.summary()
    model.compile(optimizer='adam', loss="binary_crossentropy", metrics=['accuracy', auroc], sample_weight_mode="temporal")
    model.fit_generator(
        generator=dataset.train_data.getBatchGenerator(shuffle=True),
        steps_per_epoch=dataset.train_data.getSteps(),
        epochs=ARGS.epochs,
        verbose=1,
        validation_data=dataset.val_data.getBatchGenerator(shuffle=True),
        validation_steps=dataset.val_data.getSteps(),
        callbacks=[
            EarlyStopping(monitor="val_loss", patience=ARGS.early_stopping_patience, restore_best_weights=True),
            ModelCheckpointwithBestWeights(
                file_dir=os.path.join(ARGS.working_path, 'model', ARGS.imputation_mode + '_seed-' + str(ARGS.seed) + '_' + timestamp + '_' + str(i_fold))
            ),
            TensorBoard(
                log_dir=os.path.join(ARGS.working_path, 'tb_logs', timestamp + '_' + str(i_fold))
            )
        ]
        )
    model.save(os.path.join(ARGS.working_path, 'model',
                            ARGS.imputation_mode + '_seed-' + str(ARGS.seed) + '_' + timestamp + '_' + str(i_fold), 'model.h5'))

# Evaluate the model
true_y_list = [
    dataset.train_data.getTargets(), dataset.val_data.getTargets(), dataset.test_data.getTargets()
]
pred_y_list = [
    model.predict_generator(dataset.train_data.getInputGenerator(),
                            steps=dataset.train_data.getSteps()),
    model.predict_generator(dataset.val_data.getInputGenerator(),
                            steps=dataset.val_data.getSteps()),
    model.predict_generator(dataset.test_data.getInputGenerator(),
                            steps=dataset.test_data.getSteps()),
]
dataset_list = [
    dataset.train_data,
    dataset.val_data,
    dataset.test_data
]
processed_list = [remove_padding(ty, py, d) for ty, py, d in zip(true_y_list, pred_y_list, dataset_list)]
#auc_score_list = [get_auc(ty, py, d) for ty, py, d in zip(true_y_list, pred_y_list, dataset_list)] # [3, n_task]
auc_score_list = [roc_auc_score(t, p) for t,p,_,_ in processed_list] # [3, n_task]
print('AUC scores (train,val,test): {}'.format(auc_score_list))
pred_y_list_all.append(pred_y_list)
auc_score_list_all.append(auc_score_list)

auc_score_list_all = np.stack(auc_score_list_all, axis=0)
# print('Mean AUC score: {}; Std AUC score: {}'.format(
#     np.mean(auc_score_list_all, axis=0),
#     np.std(auc_score_list_all, axis=0)))

result_path = os.path.join(ARGS.working_path, 'results', timestamp)
if not os.path.exists(result_path):
    os.makedirs(result_path)
np.savez_compressed(os.path.join(result_path, 'predictions.npz'),
                    pred_y_list_all=pred_y_list_all)
np.savez_compressed(os.path.join(result_path, 'auroc_score.npz'),
                    auc_score_list_all=auc_score_list_all)

for (ty, py, _, _) in processed_list:
    print('-'*20)
    c_m = confusion_matrix(ty, np.array(py > 0.5).astype(int))
    print(c_m)
    PPV = c_m[1,1] / (c_m[1,1] + c_m[0,1])
    print(f"PPV/Precision = {PPV}")
    TPR = c_m[1,1] / c_m[1].sum()
    print(f"TPR/Sensitivity/Recall = {TPR}")
    TNR = c_m[0,0] / c_m[0].sum()
    print(f"TNR/Specificity = {TNR}")

print("Saving output...")

test_ty, test_py, test_py_ind, test_files = processed_list[2]
save_output(test_py_ind, test_files, "results", "GRU-D", ARGS.imputation_mode, seed=ARGS.seed, threshold=0.5)

print('Finished!', '='*20)


#!/usr/bin/env python

import rospy
import rospkg
import numpy as np
from os import path
from time import time
from datetime import datetime
from ackermann_msgs.msg import AckermannDrive, AckermannDriveStamped
from geometry_msgs.msg import (
    Point,
    Pose,
    PoseWithCovariance,
    PoseWithCovarianceStamped,
    PoseStamped,
    Quaternion,
)
from tf.transformations import quaternion_from_euler
from util import quaternion_to_angle
from geometry import oval_path, simple_track
from geometry import calculate_theta, calculate_diff_theta

class Controller(object):
    def __init__(
            self,
            track,
            init_pose,
            noise_factor=0.068,
            recording_on=True):
        rospack = rospkg.RosPack()

        # Constants
        self.PACKAGE_NAME = 'can_racecar'
        self.PACKAGE_DIR = rospack.get_path(self.PACKAGE_NAME)
        self.UPDATE_RATE = 20.0
        self.MAX_STEERING_ANGLE = 0.34
        self.MAX_VELOCITY = 2.0
        self.NOISE_FACTOR = noise_factor
        self.DATA_FILE = path.join(self.PACKAGE_DIR, 'data', datetime.now().strftime('%Y%m%d%H%M%S'))

        # Controller State
        self.car_pose = None
        self.track = track
        self.last_diff = None

        # Data Collection
        self.data = []
        self._recording_on = recording_on

        # ----- Read Parameters -----
        control_topic = rospy.get_param(
            '~control_topic', '/mux/ackermann_cmd_mux/input/navigation')
        init_pose_topic = rospy.get_param(
            '~init_pose_topic', '/initialpose')
        car_pose_topic = rospy.get_param(
            '~car_pose_topic', '/car_pose')

        # ----- Publishers -----
        self.pub_controls = rospy.Publisher(
            control_topic, AckermannDriveStamped, queue_size=1)
        self.pub_init_pose = rospy.Publisher(
            init_pose_topic, PoseWithCovarianceStamped, queue_size=1)

        # ----- Subscribers -----
        self.sub_car_pose = rospy.Subscriber(
            car_pose_topic, PoseStamped, self.car_pose_cb, queue_size=1)

        # initialization
        rospy.sleep(rospy.Duration.from_sec(3))
        self.send_init_pose(init_pose)

        # ----- Timer -----
        self.update_timer = rospy.Timer(
            rospy.Duration.from_sec(1.0 / self.UPDATE_RATE), self.timer_cb)

    def car_pose_cb(self, msg):
        rotation = quaternion_to_angle(msg.pose.orientation)
        self.car_pose = np.array([
            msg.pose.position.x,
            msg.pose.position.y,
            rotation
        ], dtype=np.float)

    def timer_cb(self, event):
        self.apply_pid_controller()
        if self._recording_on:
            if len(self.data) % 100 == 0:
                # data = np.asarray(self.data)
                np.save(self.DATA_FILE, self.data)

    def apply_pid_controller(self, model='PD'):
        velocity = 0.4 * self.MAX_VELOCITY
        steering_angle = 0.0
        if model == 'P':
            diff = np.array(self.calculcate_diff_drive(*self.car_pose))
            K = np.array([1, 0.1])
            err_str = 'Error(dis, psi): {:.2f}, {:.2f}'.format(*diff)
        elif model == 'PD':
            diff = np.array(self.calculate_pd_diff_drive(*self.car_pose))
            K = np.array([1, 0.1, 0.1, 0.1])
            err_str = 'Error(dis, psi, dot_dis, dot_psi): {:.2f}, {:.2f}, {:.2f}, {:.2f}'.format(*diff)
        elif model == 'PID':
            raise NotImplementedError
        else:
            raise RuntimeError('Choose a model from \'P\', \'PD\' and \'PID\'')
        steering_angle = np.clip(K.dot(diff), -self.MAX_STEERING_ANGLE, self.MAX_STEERING_ANGLE)
        print('{} Steering: {:.2f}'.format(err_str, steering_angle))
        env_noise = np.sin(time()) * self.NOISE_FACTOR
        actual_steering = np.clip(
            steering_angle + env_noise,
            -self.MAX_STEERING_ANGLE,
            self.MAX_STEERING_ANGLE)
        if self._recording_on:
            self.data.append(diff.tolist() + [steering_angle])
        self.send_command((velocity, actual_steering))

    def send_init_pose(self, pose_data):
        assert len(pose_data) == 3
        x, y, theta = float(pose_data[0]), float(pose_data[1]), float(pose_data[2])
        q = Quaternion(*quaternion_from_euler(0, 0, theta))
        point = Point(x=x, y=y)
        pose = PoseWithCovariance(pose=Pose(position=point, orientation=q))
        self.pub_init_pose.publish(PoseWithCovarianceStamped(pose=pose))

    def send_command(self, cmd):
        assert len(cmd) == 2
        v, delta = float(cmd[0]), float(cmd[1])
        drive = AckermannDrive(steering_angle=delta, speed=v)
        self.pub_controls.publish(AckermannDriveStamped(drive=drive))

    def estimate_future_pose(self, epslion=0.1):
        car_x, car_y, car_ori = self.car_pose
        est_x = car_x + epslion * np.cos(car_ori)
        est_y = car_y + epslion * np.sin(car_ori)
        return est_x, est_y, car_ori

    def calculcate_diff_drive(self, x, y, psi):
        """ Calculate (d - d_des) and (psi - psi_des) for a proportional controller. """
        diff_dis, track_point = self.track.distance_to(x, y)
        if not self.track.contains(x, y):
            diff_dis = -diff_dis
        # calculate diff orientation
        psi_des = self.track.estimate_derivative(calculate_theta(*track_point))
        diff_psi = calculate_diff_theta(psi, psi_des)
        return diff_dis, diff_psi

    def calculate_pd_diff_drive(self, x, y, psi):
        diff_dis, diff_psi = self.calculcate_diff_drive(x, y, psi)
        if self.last_diff is None:
            dot_diff_dis, dot_diff_psi = 0, 0
        else:
            dot_diff_dis = (diff_dis - self.last_diff[0]) * self.UPDATE_RATE
            dot_diff_psi = calculate_diff_theta(diff_psi, self.last_diff[1]) * self.UPDATE_RATE
        self.last_diff = (diff_dis, diff_psi)
        return diff_dis, diff_psi, dot_diff_dis, dot_diff_psi

if __name__ == '__main__':
    rospy.init_node('pub_control')

    pid_controller = Controller(
        track=simple_track, init_pose=(0, 2.5, 0))

    # ----- Main control loop -----
    rospy.spin()


#=============================================================================#
#                                                                             #
# Copyright (c) 2017                                                          #
#                                                                             #
# Permission is hereby granted, free of charge, to any person obtaining a     #
# copy of this software and associated documentation files (the "Software"),  #
# to deal in the Software without restriction, including without limitation   #
# the rights to use, copy, modify, merge, publish, distribute, sublicense,    #
# and/or sell copies of the Software, and to permit persons to whom the       #
# Software is furnished to do so, subject to the following conditions:        #
#                                                                             #
# The above copyright notice and this permission notice shall be included in  #
# all copies or substantial portions of the Software.                         #
#                                                                             #
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR  #
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,    #
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE #
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER      #
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING     #
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER         #
# DEALINGS IN THE SOFTWARE.                                                   #
#                                                                             #
#=============================================================================#

from math import pi, sin, cos
from numpy import array, zeros, ones, identity

import os

os.environ['PATH'] += os.getcwd()+'\\assimp;'

import pyassimp
import pyassimp.postprocess

from src.matrix_transforms import *

from src.models import AssimpModel, UiModel
from src.game import GameInstance as game

import sys

class Entity(object):
	"""
	This class is for representing a 3d game object.
	"""

	def __init__(self, model_class):
		self._matrix = identity(4)
		self._modelClass = model_class

	@property
	def pos(self):
		return self._matrix[3][0:3]

	@property
	def matrix(self):
		""" Returns the current model matrix. """
		return self._matrix

	#@property
	#def normalMatrix(self):
	#	return self.model.normalMatrix

	@property
	def modelClass(self):
		return self._modelClass

	def translate(self, t):
		"""
		Move the entity relative to its current position /and rotation/.
		"""
		m_translate_in_place(self._matrix, t)

	def rotate(self, r):
		""" Rotate the entity relative to its current rotation. """
		m_rotate_in_place(self._matrix, r)

class Character(Entity):
	"""
	This class extends the #Entity class to add mobility to the object.
	"""

	def __init__(self, model_class):
		super(Character, self).__init__(model_class)

		self.movement = zeros(3)
		self.movementSpeed = 0.
		self.heading = 0.

	def update(self, time_passed):
		""" Update the position of the object. """
		movement = self.movement.dot(self.movementSpeed).dot(time_passed)
		self.translate(movement)

class PlayerCharacter(Character):
	"""
	This class extends the #Character class to add control methods.
	"""

	def __init__(self, model_class):
		super(PlayerCharacter, self).__init__(model_class)

		self.camera = None

	def setCamera(self, camera):
		self.camera = camera

	def forwardsBackwards(self, mag):
		self.movement[2] -= mag
		if abs(self.movement[2]) > abs(mag):
			self.movement[2] = -mag

	def leftRight(self, mag):
		self.movement[0] += mag
		if abs(self.movement[0]) > abs(mag):
			self.movement[0] = mag

	def update(self, time_passed):
		if (self.movement[0] != 0 or self.movement[2] != 0):
			if (self.camera):
				self.rotate((0., self.camera.az-self.heading, 0.))

			self.heading = self.camera.az

		super(PlayerCharacter, self).update(time_passed)

class UiModelVerticalAlignment:
	Top    = 0
	Centre = 1
	Bottom = 2

class UiModelHorizontalAlignment:
	Left   = 0
	Centre = 1
	Right  = 2

from re import compile, search

ui_meas_re = compile("^([-+]?(?:\d+(?:\.\d*)?|\.\d+)(?:[eE][-+]?\d+)?)(?:\s*)([%(?:px)]*)")

class UiMeas:
	def __init__(self, val, type="sc"):
		self.val  = val
		self.type = type

class UiEntity(Entity):

	def __init__(self, horiz_offset = UiMeas(0, "%"),
		                 vert_offset  = UiMeas(0, "%"),
		                 width        = UiMeas(50, "%"),
		                 height       = UiMeas(50, "%"),
		                 horiz_align  = UiModelHorizontalAlignment.Left,
		                 vert_align   = UiModelVerticalAlignment.Top):
		super(UiEntity, self).__init__(UiModel)

		translation = [0., 0., 0.]
		scale = [1., 1., 1.]
		pre_scale_translation  = [0., 0., 0.]
		post_scale_translation = [0., 0., 0.]

		if horiz_align == UiModelHorizontalAlignment.Left:
			pre_scale_translation[0] = -1.
			post_scale_translation[0] = .5
			horiz_offset_scl = 1.
		elif horiz_align == UiModelHorizontalAlignment.Right:
			pre_scale_translation[0] = 1.
			post_scale_translation[0] = -.5
			horiz_offset_scl = -1.

		if vert_align == UiModelVerticalAlignment.Top:
			pre_scale_translation[1] = 1.
			post_scale_translation[1] = -.5
			vert_offset_scl = -1.
		elif vert_align == UiModelVerticalAlignment.Bottom:
			pre_scale_translation[1] = -1.
			post_scale_translation[1] = .5
			vert_offset_scl = 1.

		if horiz_offset.type == "sc":
			translation[0] = horiz_offset_scl*horiz_offset.val
		elif horiz_offset.type == "%":
			translation[0] = horiz_offset_scl*horiz_offset.val/50.

		if vert_offset.type == "sc":
			translation[1] = vert_offset_scl*vert_offset.val
		elif vert_offset.type == "%":
			translation[1] = vert_offset_scl*vert_offset.val/50.

		if width.type == "sc":
			scale[0] = width.val
		elif width.type == "%":
			scale[0] = width.val/50.

		if height.type == "sc":
			scale[1] = height.val
		elif height.type == "%":
			scale[1] = height.val/50.

		m_translate_in_place(self._matrix, translation)
		m_translate_in_place(self._matrix, pre_scale_translation)
		m_scale_in_place(self._matrix, scale)
		m_translate_in_place(self._matrix, post_scale_translation)

game.renderer.aLight.setColour([1., 1., 1.])
game.renderer.aLight.setAmplitude(0.1)

dLight = game.renderer.getDirectionalLight()
dLight.setDirection([-0.707107, 0.0, -0.707107])
#dLight.setDirection([0.,0.,-1.])
dLight.setColour([0.75, 0.75, 0.75])

pLight = game.renderer.getPointLight()
pLight.setPosition([-3, 0, -3])
pLight.setAmplitude(5)
pLight.setColour([1,0.25,0])

pyassimp_processing_flags = (pyassimp.postprocess.aiProcess_Triangulate |
                             pyassimp.postprocess.aiProcess_OptimizeMeshes)

class Cube(AssimpModel):
	def __init__(self):
		asset = pyassimp.load("assets/obj/cube.obj", processing=pyassimp_processing_flags)
		super(Cube, self).__init__(asset)

class Sphere(AssimpModel):
	def __init__(self):
		asset = pyassimp.load("assets/obj/sphere.dae", processing=pyassimp_processing_flags)
		super(Sphere, self).__init__(asset)

class Spider(AssimpModel):
	def __init__(self):
		asset = pyassimp.load("assets/obj/spider.obj", processing=pyassimp_processing_flags)
		super(Spider, self).__init__(asset)

		self.rotate((0., pi/2, 0.))
		self.translate((-0.05, 0., 0.))
		self.scale((.01, .01, .01 ))

class Duck(AssimpModel):
	def __init__(self):
		asset = pyassimp.load("assets/obj/duck.dae", processing=pyassimp_processing_flags)
		super(Duck, self).__init__(asset)

		self.scale((.01, .01, .01))

game.renderer.addModel(Cube)
game.renderer.addModel(Sphere)
game.renderer.addModel(Spider)
game.renderer.addModel(Duck)
game.renderer.addModel(UiModel)

cube_1 = Entity(Cube)
cube_1.translate(( 0., 4., 0.))

cube_2 = Entity(Cube)
cube_2.translate(( -1.,  0., -3. ))

sphere_1 = Entity(Sphere)
sphere_1.translate((3., 0., -4.))

#sys.stderr.write(str(sphere.meshes[0].__dict__)+'\n')
#sys.stderr.write(str(sphere_q.meshes[0].__dict__)+'\n')

spider_1 = Entity(Spider)
spider_1.translate(( -3.,  0.,  0. ))
spider_1.rotate(( 0., -pi/2, 0. ))
spider_1.rotate(( pi/6, 0.,  0. ))

#sys.stderr.write(str(duck.meshes[0].__dict__)+'\n')

duck_1 = Entity(Duck)
duck_1.translate((5., 0., -1.))

spider_char = PlayerCharacter(Spider)
spider_char.movementSpeed = 3.
spider_char.translate((2., 0., 0.))

#cube_1.setColour(ones(108, dtype='float32'))
#cube_2.setColour(ones(108, dtype='float32'))

game.addEntity(cube_1)
game.addEntity(cube_2)
game.addEntity(spider_1)
game.addEntity(sphere_1)
game.addEntity(duck_1)

game.setPlayerCharacter(spider_char)

#render_screen = UiEntity(horiz_offset = UiMeas(5, "%"),
#                         vert_offset  = UiMeas(5, "%"),
#                         width        = UiMeas(90, "%"),
#                         height       = UiMeas(90, "%"))
#game.renderer.addUiEntity(render_screen)

game.run()

del game


'''
File: icd9Formatter.py
Author: Adam Pah
Description:
Handling icd9 formatting and translation problems
Needs outside database to handle translations
'''

def icd9_normalizer(icd9, check_decimal=False):
    '''
    input:
        icd9 - the icd9 code as a string
    output:
        formatted icd9 code, standardized
    '''
    icd9 = str(icd9)
    #First switch, is the icd9 in decimal format?
    if check_decimal:
        print icd9
        #Havent seen an icd9 code in deciimal format with a leading 0 yet
        if icd9[0]=='0':
            dicd9 = icd9[:3] + '.' + icd9[3:]
            sicd9 = dicd9.lstrip('0')
            return str(float(sicd9))
        #Handle the V codes now
        elif icd9[0]=='V':
            dicd9 = icd9[:3] + '.' + icd9[3:]
            return dicd9
        else:
            dicd9 = icd9[:3] + '.' + icd9[3:]
            return str(float(dicd9))
    else:
        try:
            return str(float(icd9))
        except ValueError:
            return icd9

def icd9_string(sicd9):
    '''
    input: icd9 string normalized
    output: icd9 description form phewas database
    '''
    import gale.databases.mongoConnect as mcxn

    tdb = mcxn.MongoConnection(dbs['Phewas'])
    doc = tdb.collection.find_one({'icd9': sicd9}, {'icd9_string': 1})
    tdb.tearDown()
    if doc:
        return doc['icd9_string']
    else:
        return 'Description not found'

def phewas_translator():
    '''
    input:
    output:
        dictionary of icd9 to phewas groups
    dependencies: gale
    '''
    import gale.databases.mongoConnect as mcxn

    tdb = mcxn.MongoConnection(dbs['Phewas'])
    docs = tdb.collection.find({}, {'icd9': 1, 'phewas_code': 1})
    #Condense to dictionary
    data = {}
    for i in docs:
        data[i['icd9']] = i['phewas_code']
    return data

def return_phewas_group(icd9, tdb=None):
    '''
    input:
        icd9, the icd9 code to return
        mongoConnect database connection (optional, increases performance)
    output: phewas_group

    if item not found in database then 'None' value is returned.
    '''
    if not tdb:
        import gale.databases.mongoConnect as mcxn
        tdb = mcxn.MongoConnection(dbs['Phewas'])
    #Get it
    doc = tdb.collection.find_one({'icd9': icd9}, {'phewas_code': 1})
    if doc!=None:
        return doc['phewas_code']
    else:
        return doc

def return_phewas_group_string(icd9, tdb=None):
    '''
    input:
        icd9, the icd9 code to return
        mongoConnect database connection (optional, increases performance)
    output: phewas_group

    if item not found in database then 'None' value is returned.
    '''
    if not tdb:
        import gale.databases.mongoConnect as mcxn
        tdb = mcxn.MongoConnection(dbs['Phewas'])
    #Get it
    doc = tdb.collection.find_one({'icd9': icd9}, {'phewas_string': 1})
    if doc!=None:
        return doc['phewas_string']
    else:
        return doc

def icd9_validator(icd9):
    '''
    input:
        icd9 code
    output:
        True/False if it is valid
    '''
    import numpy as np
    try:
        ficd9 = float(icd9)
        if np.isnan(ficd9)==False:
            return True
        else:
            return False
    except ValueError:
        return False


# Licensed under a 3-clause BSD style license - see LICENSE.rst
import pytest
import numpy as np
from numpy.testing import assert_allclose
import astropy.units as u
from gammapy.datasets import Datasets, SpectrumDatasetOnOff
from gammapy.estimators.flux import FluxEstimator
from gammapy.modeling.models import PowerLawSpectralModel, SkyModel
from gammapy.utils.testing import requires_data, requires_dependency

@pytest.fixture(scope="session")
def fermi_datasets():
    filename = "$GAMMAPY_DATA/fermi-3fhl-crab/Fermi-LAT-3FHL_datasets.yaml"
    filename_models = "$GAMMAPY_DATA/fermi-3fhl-crab/Fermi-LAT-3FHL_models.yaml"

    return Datasets.read(filename=filename, filename_models=filename_models)

@pytest.fixture(scope="session")
def hess_datasets():
    datasets = Datasets()
    pwl = PowerLawSpectralModel(amplitude="3.5e-11 cm-2s-1TeV-1", index=2.7)
    model = SkyModel(spectral_model=pwl, name="Crab")

    for obsid in [23523, 23526]:
        dataset = SpectrumDatasetOnOff.read(
            f"$GAMMAPY_DATA/joint-crab/spectra/hess/pha_obs{obsid}.fits"
        )
        dataset.models = model
        datasets.append(dataset)

    return datasets

@requires_data()
@requires_dependency("iminuit")
def test_flux_estimator_fermi_no_reoptimization(fermi_datasets):
    estimator = FluxEstimator(
        0,
        energy_min="1 GeV",
        energy_max="100 GeV",
        norm_n_values=5,
        norm_min=0.5,
        norm_max=2,
        reoptimize=False,
        selection_optional="all"
    )

    result = estimator.run(fermi_datasets)

    assert_allclose(result["norm"], 0.98949, atol=1e-3)
    assert_allclose(result["ts"], 25083.75408, atol=1e-3)
    assert_allclose(result["norm_err"], 0.01998, atol=1e-3)
    assert_allclose(result["norm_errn"], 0.0199, atol=1e-3)
    assert_allclose(result["norm_errp"], 0.0199, atol=1e-3)
    assert len(result["norm_scan"]) == 5
    assert_allclose(result["norm_scan"][0], 0.5)
    assert_allclose(result["norm_scan"][-1], 2)
    assert_allclose(result["e_min"], 10 * u.GeV, atol=1e-3)
    assert_allclose(result["e_max"], 83.255 * u.GeV, atol=1e-3)

@requires_data()
@requires_dependency("iminuit")
def test_flux_estimator_fermi_with_reoptimization(fermi_datasets):
    estimator = FluxEstimator(
        0,
        energy_min="1 GeV",
        energy_max="100 GeV",
        reoptimize=True,
        selection_optional=None,
    )
    result = estimator.run(fermi_datasets)

    assert_allclose(result["norm"], 0.989989, atol=1e-3)
    assert_allclose(result["ts"], 18729.907481, atol=1e-3)
    assert_allclose(result["norm_err"], 0.01998, atol=1e-3)

@requires_data()
@requires_dependency("iminuit")
def test_flux_estimator_1d(hess_datasets):
    estimator = FluxEstimator(
        source="Crab",
        energy_min=1 * u.TeV,
        energy_max=10 * u.TeV,
        selection_optional=["errn-errp", "ul"],
    )
    result = estimator.run(hess_datasets)

    assert_allclose(result["norm"], 1.218139, atol=1e-3)
    assert_allclose(result["ts"], 527.492959, atol=1e-3)
    assert_allclose(result["norm_err"], 0.095496, atol=1e-3)
    assert_allclose(result["norm_errn"], 0.093204, atol=1e-3)
    assert_allclose(result["norm_errp"], 0.097818, atol=1e-3)
    assert_allclose(result["norm_ul"], 1.525773, atol=1e-3)
    assert_allclose(result["e_min"], 1 * u.TeV, atol=1e-3)
    assert_allclose(result["e_max"], 10 * u.TeV, atol=1e-3)

@requires_data()
@requires_dependency("iminuit")
def test_flux_estimator_incorrect_energy_range(fermi_datasets):
    with pytest.raises(ValueError):
        FluxEstimator(source="Crab", energy_min=10 * u.TeV, energy_max=1 * u.TeV)

    fe = FluxEstimator(
        source="Crab Nebula", energy_min=0.18 * u.TeV, energy_max=0.2 * u.TeV
    )

    result = fe.run(fermi_datasets)

    assert np.isnan(result["norm"])

@requires_data()
@requires_dependency("iminuit")
def test_inhomogeneous_datasets(fermi_datasets, hess_datasets):

    for dataset in hess_datasets:
        dataset.models = fermi_datasets.models

    datasets = Datasets()

    datasets.extend(fermi_datasets)
    datasets.extend(hess_datasets)

    estimator = FluxEstimator(
        source="Crab Nebula",
        energy_min=1 * u.TeV,
        energy_max=10 * u.TeV,
        selection_optional=None,
    )
    result = estimator.run(datasets)

    assert_allclose(result["norm"], 1.190622, atol=1e-3)
    assert_allclose(result["ts"], 612.50171, atol=1e-3)
    assert_allclose(result["norm_err"], 0.090744, atol=1e-3)
    assert_allclose(result["e_min"], 0.693145 * u.TeV, atol=1e-3)
    assert_allclose(result["e_max"], 2 * u.TeV, atol=1e-3)


import mxnet as mx
import os
import time
from concurrent import futures
import numpy as np

path='http://data.mxnet.io/models/imagenet/'
[mx.test_utils.download(path+'resnet/50-layers/resnet-50-0000.params'),
 mx.test_utils.download(path+'resnet/50-layers/resnet-50-symbol.json'),
 mx.test_utils.download(path+'synset.txt')]

ctx = mx.gpu(0)
ngpu = 1
group2ctx = {'embed': mx.gpu(0),\
             'decode': mx.gpu(ngpu - 1)}

with open('synset.txt', 'r') as f:
    labels = [l.rstrip() for l in f]

sym, args, aux = mx.model.load_checkpoint('resnet-50',0)

#fname = mx.test_utils.download('https://github.com/dmlc/web-data/blob/master/mxnet/doc/tutorials/python/predict_image/cat.jpg?raw=true')
fname = mx.test_utils.download('https://raw.githubusercontent.com/awslabs/mxnet-model-server/master/docs/images/kitten_small.jpg?raw=true')
img = mx.image.imread(fname)
# convert into format (batch, RGB, width, height)
img = mx.image.imresize(img, 224, 224) # resize
img = img.transpose((2, 0, 1)) # Channel first
img = img.expand_dims(axis=0) # batchify
img = img.astype(dtype='float32')
args['data'] = img

softmax = mx.nd.random_normal(shape=(1,))
args['softmax_label'] = softmax

exe = sym.bind(ctx=ctx, args=args, aux_states=aux, grad_req='null',group2ctx=group2ctx)

exe.forward()
prob = exe.outputs[0].asnumpy()
# print the top-5
prob = np.squeeze(prob)
a = np.argsort(prob)[::-1]
for i in a[0:5]:
    print('probability=%f, class=%s' %(prob[i], labels[i]))


import keras
import tensorflow as tf
import keras.backend.tensorflow_backend as K
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
# config.gpu_options.per_process_gpu_memory_fraction = 0.9
sess = tf.Session(config=config)
K.set_session(sess)

import os
import sys

sys.path.insert(0, '../')
from models.crnn import crnn
from data_utils.transform import reshape_to_target, pre_processing
from .ctc_decode import ctc_decode as cd

import yaml
import cv2
import numpy as np
from easydict import EasyDict as ET

def main(args):
    f = open(args.config)
    cfgs = yaml.load(f)
    f.close()
    cfgs = ET(cfgs)
    image_size = cfgs.IMAGE_SIZE
    charset = cfgs.CHARSET
    weight = cfgs.WEIGHT
    img_path = args.img

    h, w, c = image_size.split(',')
    image_size = (int(h), int(w), int(c))

    with open(charset) as f:
        charset = f.readline().strip('\n')
        f.close()
    nb_classes = len(charset) + 1

    model, *_ = crnn(nb_classes, image_size)
    model.load_weights(weight, by_name=True)

    ori_img = cv2.imread(img_path)
    img = reshape_to_target(ori_img, image_size)
    img = pre_processing(img)
    img = np.expand_dims(img, axis=0)

    prob = model.predict(img)
    result_str = cd(prob, charset)
    print(result_str)
    if args.viz:
        cv2.imshow('display', ori_img)
        cv2.waitKey(0)


#!/usr/bin/python3
import rospy
from sensor_msgs.msg import JointState
from random import random, choice
import time
import numpy as np
from roboy_control_msgs.msg import Emotion

rospy.init_node("idle_moves")
topic_root = "/roboy/pinky/"
publisher = rospy.Publisher(topic_root + 'joint_targets', JointState, queue_size=1)
emotion_pub = rospy.Publisher(topic_root + "cognition/face/emotion", Emotion, queue_size=1)
axis1_range_upper = 0.3
shoulder_range_upper = -0.4
elbow_range_upper =  0.5
head_range_upper = 0.5
sleep_range_upper = 3.0

emotions = ["shy", "kiss", "hypno", "hearts", "lucky", "lookleft", "lookright", "surprise", "smileblink", "rolling"]

msg = JointState()
msg.name = ["head_axis1","shoulder_right_axis0", "shoulder_right_axis1", "shoulder_left_axis0", "shoulder_left_axis1", "elbow_right_axis0", "elbow_left_axis0"]

#msg.name = ["shoulder_right_axis0", "shoulder_right_axis1", "shoulder_left_axis0", "shoulder_left_axis1", "elbow_right_axis0", "elbow_left_axis0"]
old_targets = None

while not rospy.is_shutdown():
    msg.position = []

    targets = [random()*head_range_upper-0.25,random()*shoulder_range_upper, random()*shoulder_range_upper, random()*axis1_range_upper, random()*shoulder_range_upper, random()*elbow_range_upper,random()*elbow_range_upper]
    #targets = [random()*shoulder_range_upper, random()*shoulder_range_upper, random()*axis1_range_upper, random()*shoulder_range_upper, random()*elbow_range_upper,random()*elbow_range_upper]
    if old_targets is None:
        old_targets = [0]*len(targets)
    interp_targets = []
    for i in range(len(targets)):
        interp_targets.append(np.linspace(old_targets[i],targets[i],100))
    for i in range(len(interp_targets[0])):
        msg.position = []
        for j in range(len(targets)):#joints
            msg.position.append(interp_targets[j][i])
        msg.velocity = [0]*len(msg.name)
        msg.effort = [0]*len(msg.name)
        publisher.publish(msg)
        rospy.loginfo(msg.position)
        rospy.sleep(0.01)

    #msg.velocity = [0]*len(msg.name)
    #msg.effort = [0]*len(msg.name)
    #publisher.publish(msg)
    if random() > 0.5:
        rospy.loginfo("emotion")
        emotion_msg = Emotion(emotion=choice(emotions))
        emotion_pub.publish(emotion_msg)
    rospy.sleep(random()*sleep_range_upper)
    old_targets = targets


import os, sys
import numpy as np
import os.path as osp
import time
import shutil as sh
from collections import defaultdict
import ipdb

g_file_buff = []

def load_dataset(folder):
    dataset = defaultdict(dict)
    for dirname, _, files in os.walk(folder):
        for fi in files:
            if fi.endswith('.npy'):
                frame_no = fi.split('_')[0]
                filename = osp.join(dirname, fi)
                vec = np.load(filename)
                if frame_no not in dataset.keys():
                    dataset[frame_no]['file'] = []
                    dataset[frame_no]['featu'] = []
                # vec = vec / np.linalg.norm(vec)
                dataset[frame_no]['file'].append(filename)
                dataset[frame_no]['featu'].append(vec)
    print('frame count: ', len(dataset))

    for k in dataset.keys():
        dataset[k]['featu'] = np.array(dataset[k]['featu'])
    return dataset

def parse_fname(fname):
    fname = osp.splitext(osp.basename(fname))[0]
    frame_no, x, y = fname.split('_')
    frame_no = int(frame_no)
    x = int(x[1:])
    y = int(y[1:])
    return frame_no, x, y

def match(data_gallery, data_probe, threshold=0.5):
    featu_g = data_gallery['featu']
    featu_p = data_probe['featu']
    file_g = data_gallery['file']
    file_p = data_probe['file']

    mmap = featu_p.dot(featu_g.T)
    col_max = np.argsort(mmap,0)[-1,:]
    col_pairs = [(col_max[i], i) for i in range(len(col_max))]
    row_max = np.argsort(mmap,1)[:,-1]
    row_pairs = [(i, row_max[i]) for i in range(len(row_max))]

    cross_matched_pairs = []
    for d in col_pairs:
        if d in row_pairs and mmap[d] > threshold:
            pf = file_p[d[0]]
            gf = file_g[d[1]]
            p_frame_no, p_x, p_y = parse_fname(pf)
            g_frame_no, g_x, g_y = parse_fname(gf)

            if abs(p_frame_no - g_frame_no) < 750: continue
            if abs(p_y - g_y) < 200: continue
            if pf in g_file_buff or gf in g_file_buff: continue

            pair = {}
            pair['id'] = d
            pair['file'] = (pf, gf)
            pair['frame_no'] = (p_frame_no, g_frame_no)
            pair['coord'] = ((p_x, p_y), (g_x, g_y))
            pair['similary'] = mmap[d]

            g_file_buff.append(pf)
            g_file_buff.append(gf)
            cross_matched_pairs.append(pair)
    return cross_matched_pairs

def save_matched_pairs(cross_matched_pairs, save_dir):
    for i, pair in enumerate(cross_matched_pairs):
        folder = osp.join(save_dir, '%06d_%.3f'%(i, pair['similary']))
        if not osp.exists(folder):
            os.makedirs(folder)
        pf, gf = pair['file']
        pf = pf.replace('.npy', '.jpg')
        gf = gf.replace('.npy', '.jpg')
        dst_pf = osp.join(folder, osp.basename(pf))
        dst_gf = osp.join(folder, osp.basename(gf))
        sh.copy(pf, dst_pf)
        sh.copy(gf, dst_gf)

if __name__=='__main__':
    data_dir = '/home/ct/Data/queue/crops'
    save_dir = '/home/ct/Data/queue/matched'

    dataset = load_dataset(data_dir)
    fn_list = list(dataset.keys())
    fn_list.sort()
    cross_matched_pairs = []
    for i in range(len(fn_list))[1:]:
        data_p = dataset[fn_list[i]]
        for j in range(i):
            data_g = dataset[fn_list[j]]
            # ipdb.set_trace()
            cross_matched_pairs += match(data_g, data_p, threshold=0.75)
    # ipdb.set_trace()
    save_matched_pairs(cross_matched_pairs, save_dir)


#!/usr/bin/env python3

import numpy as np
import torch

from utils import Visualizer
from models.memory import TartanAirMemory, NordlandMemory, RobotCarMemory
from utils import Projector, PairwiseCosine, gen_probe

from .lifelong import get_ll_loss

class MemReplayLoss():
    def __init__(self, writer=None, viz_start=float('inf'), viz_freq=200, counter=None, args=None):
        super().__init__()
        self.args = args
        self.writer, self.counter, self.viz_start, self.viz_freq = writer, counter, viz_start, viz_freq
        self.viz = Visualizer('tensorboard', writer=self.writer)

        if args.dataset == 'tartanair':
            self.memory = TartanAirMemory(capacity=args.mem_size, n_probe=1200, out_device=args.device)
        elif args.dataset == 'nordland':
            self.memory = NordlandMemory(capacity=args.mem_size, out_device=args.device)
        elif args.dataset == 'robotcar':
            self.memory = RobotCarMemory(capacity=args.mem_size, dist_tol=20, head_tol=15, out_device=args.device)

        self.n_triplet, self.n_recent, self.n_pair = 4, 0, 1
        self.min_sample_size = 32

        self.gd_match = GlobalDescMatchLoss(n_triplet=self.n_triplet, n_pair=self.n_pair, writer=writer, viz=self.viz, viz_start=viz_start, viz_freq=viz_freq, counter=self.counter)
        self.ll_loss = get_ll_loss(args, writer=writer, viz=self.viz, viz_start=viz_start, viz_freq=viz_freq, counter=self.counter)

    def __call__(self, net, img, aux, env):
        device = img.device
        self.store_memory(img, aux, env)

        if len(self.memory) < self.min_sample_size:
            return torch.zeros(1).to(device)

        _, (ank_batch, pos_batch, neg_batch), (pos_rel, neg_rel) = \
            self.memory.sample_frames(self.n_triplet, self.n_recent, self.n_pair)

        # no suitable triplet
        if ank_batch is None:
            return torch.zeros(1).to(device)

        img = recombine('img', ank_batch, pos_batch, neg_batch)

        loss = 0
        gd = net(img=img)
        gd_match = self.gd_match(gd)
        loss += gd_match

        # forgetting prevention
        if self.ll_loss is not None:
            for ll_loss in self.ll_loss:
                loss_name = ll_loss.name.lower()
                if loss_name in ['mas', 'rmas', 'ewc', 'cewc', 'si']:
                    loss += ll_loss(model=net, gd=gd)
                elif loss_name in ['kd', 'rkd', 'ifgir']:
                    loss += ll_loss(model=net, gd=gd, img=img)
                else:
                    raise ValueError(f'Unrecognized lifelong loss: {ll_loss}')

        # logging and visualization
        if self.writer is not None:
            n_iter = self.counter.steps if self.counter is not None else 0
            self.writer.add_scalars('Loss', {'global': gd_match}, n_iter)
            self.writer.add_histogram('Misc/RelN', neg_rel, n_iter)
            self.writer.add_histogram('Misc/RelP', pos_rel, n_iter)
            self.writer.add_scalars('Misc/MemoryUsage', {'len': len(self.memory)}, n_iter)

            # show triplets
            if self.viz is not None and n_iter >= self.viz_start and n_iter % self.viz_freq == 0:
                H, W = img.shape[2:]
                if isinstance(self.memory, TartanAirMemory):
                    N = ank_batch['pos'].shape[1]

                    # project points from pos to ank
                    mem_pts_scr = Projector.world2pix(pos_batch['pos'].reshape(-1, N, 3), (H, W), ank_batch['pose'], ank_batch['K'], ank_batch['depth_map'])[0]
                    B_total = self.n_triplet * (self.n_pair * 2 + 1)
                    mem_pts_scr_ = mem_pts_scr.reshape(self.n_triplet, self.n_pair * N, 2)
                    proj_pts_ = torch.cat([
                        torch.zeros_like(mem_pts_scr_).fill_(np.nan),
                        mem_pts_scr_,
                        torch.zeros_like(mem_pts_scr_).fill_(np.nan)], 1).reshape(B_total, self.n_pair * N, 2)

                    proj_pts_color = torch.arange(self.n_pair)[None, :, None].expand(B_total, self.n_pair, N) + 1
                    proj_pts_color = proj_pts_color.reshape(B_total, self.n_pair * N).detach().cpu().numpy()
                else:
                    proj_pts_ = proj_pts_color = None

                ank_img, pos_img, neg_img = img.split([self.n_triplet, self.n_triplet * self.n_pair, self.n_triplet * self.n_pair])
                self.viz.show(
                    torch.cat([pos_img.reshape(self.n_triplet, self.n_pair, 3, H, W), ank_img[:, None], neg_img.reshape(self.n_triplet, self.n_pair, 3, H, W)], dim=1).reshape(-1, 3, H, W),
                    proj_pts_, 'tab10', values=proj_pts_color, vmin=0, vmax=10, name='Misc/GlobalDesc/Triplet', step=n_iter,
                    nrow=(self.n_pair * 2 + 1))

        return loss

    def store_memory(self, imgs, aux, env):
        self.memory.swap(env[0])
        if isinstance(self.memory, TartanAirMemory):
            depth_map, pose, K = aux
            points_w = Projector.pix2world(gen_probe(depth_map), depth_map, pose, K)
            self.memory.store_fifo(pos=points_w, img=imgs, depth_map=depth_map, pose=pose, K=K)
        elif isinstance(self.memory, NordlandMemory):
            offset = aux
            self.memory.store_fifo(img=imgs, offset=offset)
        elif isinstance(self.memory, RobotCarMemory):
            location, heading = aux
            self.memory.store_fifo(img=imgs, location=location, heading=heading)

def recombine(key, *batches):
    tensors = [batch[key] for batch in batches]
    reversed_shapes = [list(reversed(tensor.shape[1:])) for tensor in tensors]
    common_shapes = []
    for shape in zip(*reversed_shapes):
        if all(s == shape[0] for s in shape):
            common_shapes.insert(0, shape[0])
        else:
            break
    return torch.cat([tensor.reshape(-1, *common_shapes) for tensor in tensors])

class GlobalDescMatchLoss():

    def __init__(self, n_triplet=8, n_pair=1, writer=None,
                 viz=None, viz_start=float('inf'), viz_freq=200, counter=None, debug=False):
        super().__init__()
        self.counter, self.writer = counter, writer
        self.viz, self.viz_start, self.viz_freq = viz, viz_start, viz_freq
        self.cosine = PairwiseCosine()
        self.n_triplet, self.n_pair = n_triplet, n_pair

    def __call__(self, gd):
        gd_a, gd_p, gd_n = gd.split([self.n_triplet, self.n_triplet * self.n_pair, self.n_triplet * self.n_pair])
        gd_a = gd_a[:, None]
        gd_p = gd_p.reshape(self.n_triplet, self.n_pair, *gd_p.shape[1:])
        gd_n = gd_n.reshape(self.n_triplet, self.n_pair, *gd_n.shape[1:])

        sim_ap = self.cosine(gd_a, gd_p)
        sim_an = self.cosine(gd_a, gd_n)
        triplet_loss = (sim_an - sim_ap + 1).clamp(min=0)

        # logging
        n_iter = self.counter.steps
        if self.writer is not None:
            self.writer.add_histogram('Misc/RelevanceLoss', triplet_loss, n_iter)
            self.writer.add_histogram('Misc/SimAP', sim_ap, n_iter)
            self.writer.add_histogram('Misc/SimAN', sim_an, n_iter)
            self.writer.add_scalars('Misc/GD', {'2-Norm': torch.norm(gd, dim=-1).mean()}, n_iter)

        return triplet_loss.mean()


import copy
import numpy as np
from collections import defaultdict

extra_distance = defaultdict(int, {
    36: 12  # Command Center
})

class LocalAgentExtractor:
    def __init__(self, id, summary_writer=None, resolution=24, map_boundary=48):
        self.inited = False
        self.done = False
        self.id = id
        self.unit_tag = None
        self.x = 0
        self.y = 0
        self.raw_unit = None
        self.resolution = resolution
        self.map_boundary = map_boundary
        self.last_total_damage_dealt = 0
        self.last_total_damage_taken = 0
        self.enemies = []
        self.action_step = 0
        self.summary_writer = summary_writer
        pass

    def reset(self):
        self.done = False
        self.inited = False

    def get_init_info(self, ob):
        feature_units = [
            unit for unit in ob.feature_units if unit.alliance == 1 and unit.unit_type == 35]
        if len(feature_units) <= self.id:
            # raise NotImplementedError("len(feature_units) < extractor id")
            return None
        unit = feature_units[self.id]
        if unit.tag == 0:
            raise NotImplementedError("unit tag == 0")
        self.unit_tag = unit.tag
        self.inited = True
        return 1

    def get_unit(self, ob):
        for unit in ob.feature_units:
            if unit.tag == self.unit_tag:
                return unit
        return None

    def get_raw_unit(self, ob):
        for unit in ob.raw_units:
            if unit.tag == self.unit_tag:
                return unit
        return None

    def get_enemy(self, ob):
        self.enemies = [
            unit for unit in ob.raw_units if unit.alliance == 4]

    def get_action(self, action_id, action_arg):
        if action_id != 2:
            return action_id, action_arg, 0
        xy = action_arg[2]
        for e in self.enemies:
            x_diff = xy[0] - e.x
            y_diff = xy[1] - e.y
            distance = y_diff**2 + x_diff**2
            if distance < 16+extra_distance[e.unit_type]:
                print(f"hit. Distance = {distance}")
                self.action_step += 1
                if self.summary_writer:
                    self.summary_writer.add_scalar(
                        "hit_distance", distance, self.action_step)
                return 3, [action_arg[0], action_arg[1], [e.tag]], 5
            elif distance < 36+extra_distance[e.unit_type]:
                print(f"not hit. Distance = {distance}")
        return action_id, action_arg, 0

    def get_reward(self, obs):
        reward = 0
        cur_damage_dealt = obs.score_by_vital.total_damage_dealt.life + \
            obs.score_by_vital.total_damage_dealt.shields
        reward += (cur_damage_dealt - self.last_total_damage_dealt)*10
        if reward < 0:
            self.reward = 0
            return 0
        self.last_total_damage_dealt = cur_damage_dealt
        cur_damage_taken = obs.score_by_vital.total_damage_taken.life
        self.last_total_damage_taken = cur_damage_taken
        # reward = reward // 10
        self.reward = reward
        return reward

    def extract(self, obs):
        observation = obs.observation
        # Init if needed
        if not self.inited:
            if self.get_init_info(observation) is None:
                return None

        if self.done:
            return None
        unit = self.get_unit(observation)
        if unit is None:
            self.done = True
            return None
        self.raw_unit = self.get_raw_unit(observation)
        self.get_enemy(observation)
        x = unit.x
        y = unit.y
        self.x = x
        self.y = y
        x_left = x - self.resolution // 2
        x_right = x + self.resolution // 2
        y_up = y - self.resolution // 2
        y_down = y + self.resolution // 2
        if x_left <= 0 and y_up <= 0:
            feature_screen = observation.feature_screen[:, :y_down, :x_right]
            feature_screen = np.pad(feature_screen, pad_width=(
                (0, 0), (-y_up, 0), (-x_left, 0)), mode='constant', constant_values=0)
        elif x_left <= 0 and y_down >= self.map_boundary:
            feature_screen = observation.feature_screen[:,
                                                        y_up:, :x_right]
            feature_screen = np.pad(feature_screen, pad_width=(
                (0, 0), (0, y_down-self.map_boundary), (-x_left, 0)), mode='constant', constant_values=0)
        elif x_right >= self.map_boundary and y_up <= 0:
            feature_screen = observation.feature_screen[:,
                                                        :y_down, x_left:]
            feature_screen = np.pad(feature_screen, pad_width=(
                (0, 0), (-y_up, 0), (0, x_right-self.map_boundary)), mode='constant', constant_values=0)
        elif x_right >= self.map_boundary and y_down >= self.map_boundary:
            feature_screen = observation.feature_screen[:,
                                                        y_up:, x_left:]
            feature_screen = np.pad(feature_screen, pad_width=(
                (0, 0), (0, y_down-self.map_boundary), (0, x_right-self.map_boundary)), mode='constant', constant_values=0)
        elif x_left <= 0:
            feature_screen = observation.feature_screen[:,
                                                        y_up:y_down, :x_right]
            feature_screen = np.pad(feature_screen, pad_width=(
                (0, 0), (0, 0), (-x_left, 0)), mode='constant', constant_values=0)
        elif y_down >= self.map_boundary:
            feature_screen = observation.feature_screen[:,
                                                        y_up:, x_left:x_right]
            feature_screen = np.pad(feature_screen, pad_width=(
                (0, 0), (0, y_down-self.map_boundary), (0, 0)), mode='constant', constant_values=0)
        elif y_up <= 0:
            feature_screen = observation.feature_screen[:,
                                                        :y_down, x_left:x_right]
            feature_screen = np.pad(feature_screen, pad_width=(
                (0, 0), (-y_up, 0), (0, 0)), mode='constant', constant_values=0)
        elif x_right >= self.map_boundary:
            feature_screen = observation.feature_screen[:,
                                                        y_up:y_down, x_left:]
            feature_screen = np.pad(feature_screen, pad_width=(
                (0, 0), (0, 0), (0, x_right-self.map_boundary)), mode='constant', constant_values=0)
        else:
            feature_screen = observation.feature_screen[:,
                                                        y_up:y_down, x_left:x_right]

        # TODO: calc reward
        reward = self.get_reward(observation)
        if feature_screen.shape != (27, self.resolution, self.resolution):
            return {
            "feature_screen": np.zeros((27, 24, 24)),
            "reward": 0,
            "done": False,
            "unit": np.zeros((46,)),
            "x": 0,
            "y": 0
        }
        return {
            "feature_screen": feature_screen,
            "reward": reward,
            "done": self.done,
            "unit": self.raw_unit,
            "x": self.x,
            "y": self.y
        }

# from nsnn.obsextractor import BaseObsExtractor
# import copy
# import numpy as np

# class LocalAgentExtractor(BaseObsExtractor):
#     def __init__(self, id, resolution = 30):
#         self.inited = False
#         self.done = False
#         self.id = id
#         self.unit_tag = None
#         self.x = 0
#         self.y = 0
#         self.raw_unit = None
#         self.resolution = resolution
#         self.last_total_damage_dealt = 0
#         self.last_total_damage_taken = 0
#         self.enemy_num = 0
#         pass

#     def reset(self):
#         self.done = False
#         self.inited = False

#     def get_init_info(self, ob):
#         feature_units = [
#             unit for unit in ob.feature_units if unit.alliance == 1 and unit.unit_type == 35]
#         self.enemy_num = len([unit for unit in ob.feature_units if unit.alliance == 4])
#         if len(feature_units) <= self.id:
#             # raise NotImplementedError("len(feature_units) < extractor id")
#             return None
#         unit = feature_units[self.id]
#         if unit.tag == 0:
#             raise NotImplementedError("unit tag == 0")
#         self.unit_tag = unit.tag
#         self.inited = True
#         return 1

#     def get_unit(self, ob):
#         for unit in ob.feature_units:
#             if unit.tag == self.unit_tag:
#                 return unit
#         return None

#     def get_raw_unit(self, ob):
#         for unit in ob.raw_units:
#             if unit.tag == self.unit_tag:
#                 return unit
#         return None

#     def get_reward(self, obs):
#         reward = 0
#         cur_damage_dealt = obs.score_by_vital.total_damage_dealt.life
#         self.last_total_damage_dealt = cur_damage_dealt
#         cur_damage_taken = obs.score_by_vital.total_damage_taken.life
#         reward += (cur_damage_dealt - cur_damage_taken)
#         self.last_total_damage_taken = cur_damage_taken
#         enemy_num = len([unit for unit in obs.feature_units if unit.alliance == 4])
#         if enemy_num - self.enemy_num > 0:
#             reward += 100*(enemy_num - self.enemy_num)
#             self.enemy_num = enemy_num
#         if reward == 0:
#             reward -= 10
#         self.reward = reward
#         return reward

#     def extract(self, obs):
#         observation = obs.observation
#         # Init if needed
#         if not self.inited:
#             if self.get_init_info(observation) is None:
#                 return None

#         if self.done:
#             return None
#         unit = self.get_unit(observation)
#         if unit is None:
#             self.done = True
#             return None
#         self.raw_unit = self.get_raw_unit(observation)
#         x = unit.x
#         y = unit.y
#         self.x = x
#         self.y = y
#         x_left = x - 15
#         x_right = x + 15
#         y_up = y - 15
#         y_down = y + 15
#         if x_left <= 0 and y_up <= 0:
#             feature_screen = observation.feature_screen[:, :y_down, :x_right]
#             feature_screen = np.pad(feature_screen, pad_width=(
#                 (0, 0), (-y_up, 0), (-x_left, 0)), mode='constant', constant_values=0)
#         elif x_left <= 0 and y_down >= self.map_boundary:
#             feature_screen = observation.feature_screen[:,
#                                                         y_up:, :x_right]
#             feature_screen = np.pad(feature_screen, pad_width=(
#                 (0, 0), (0, y_down-self.map_boundary), (-x_left, 0)), mode='constant', constant_values=0)
#         elif x_right >= self.map_boundary and y_up <= 0:
#             feature_screen = observation.feature_screen[:,
#                                                         :y_down, x_left:]
#             feature_screen = np.pad(feature_screen, pad_width=(
#                 (0, 0), (-y_up, 0), (0, x_right-self.map_boundary)), mode='constant', constant_values=0)
#         elif x_right >= self.map_boundary and y_down >= self.map_boundary:
#             feature_screen = observation.feature_screen[:,
#                                                         y_up:, x_left:]
#             feature_screen = np.pad(feature_screen, pad_width=(
#                 (0, 0), (0, y_down-self.map_boundary), (0, x_right-self.map_boundary)), mode='constant', constant_values=0)
#         elif x_left <= 0:
#             feature_screen = observation.feature_screen[:,
#                                                         y_up:y_down, :x_right]
#             feature_screen = np.pad(feature_screen, pad_width=(
#                 (0, 0), (0, 0), (-x_left, 0)), mode='constant', constant_values=0)
#         elif y_down >= self.map_boundary:
#             feature_screen = observation.feature_screen[:,
#                                                         y_up:, x_left:x_right]
#             feature_screen = np.pad(feature_screen, pad_width=(
#                 (0, 0), (0, y_down-self.map_boundary), (0, 0)), mode='constant', constant_values=0)
#         elif y_up <= 0:
#             feature_screen = observation.feature_screen[:,
#                                                         :y_down, x_left:x_right]
#             feature_screen = np.pad(feature_screen, pad_width=(
#                 (0, 0), (-y_up, 0), (0, 0)), mode='constant', constant_values=0)
#         elif x_right >= self.map_boundary:
#             feature_screen = observation.feature_screen[:,
#                                                         y_up:y_down, x_left:]
#             feature_screen = np.pad(feature_screen, pad_width=(
#                 (0, 0), (0, 0), (0, x_right-self.map_boundary)), mode='constant', constant_values=0)
#         else:
#             feature_screen = observation.feature_screen[:,
#                                                         y_up:y_down, x_left:x_right]

#         obs.observation.feature_screen = feature_screen
#         # TODO: calc reward
#         reward = self.get_reward(observation)
#         return {
#             "feature_screen": feature_screen,
#             "reward": reward,
#             "done": self.done,
#             "unit": self.raw_unit
#         }