"""
Instantaneous plots.

"""
__author__ = "Jack Goffinet"
__date__ = "May 2019"

import h5py
import matplotlib
import matplotlib.pyplot as plt
plt.switch_backend('agg')
import numpy as np
import os
from sklearn.decomposition import PCA
from scipy.interpolate import interp1d
from scipy.io import loadmat
from scipy.optimize import minimize
from scipy.spatial.distance import euclidean

from fastdtw import fastdtw

def get_imaging_info(loader):
	"""Get filenames and times from the loader & match them w/ spike data."""
	spike_data = np.load('09182018_data.npy').item()
	binary_spikes, real_spikes = [], []
	# For each collected song...
	count = 0
	for i in range(len(loader.dataset)):
		sample = loader.dataset.__getitem__(i)
		filename, file_time = sample['filename'], sample['file_time']
		# See if there'scorresponding imaging.
		flag = True
		for j in range(len(spike_data['filenames'])):
			if spike_data['filenames'][j] in filename and \
						abs(spike_data['times'][j] - file_time) < 0.5:
				binary_spikes.append(spike_data['binary_spikes'][j])
				real_spikes.append(spike_data['real_spikes'][j])
				flag = False
				count += 1
				break
		if flag:
			binary_spikes.append(None)
			real_spikes.append(None)
	print("portion identified:", count / len(loader.dataset))
	return {'binary':binary_spikes, 'real':real_spikes}

def apply_pca(latent_paths):
	pca = PCA(n_components=latent_paths.shape[-1])
	original_shape = latent_paths.shape
	latent_paths = pca.fit_transform(latent_paths.reshape(-1,latent_paths.shape[-1]))
	return latent_paths.reshape(original_shape)

# def fit_taus(latent_paths, ts):
# 	"""DTW time fits."""
# 	# Find tau fits.
# 	target_path = latent_paths[0] # Pick some actual path to be the target path.
# 	taus = np.zeros((len(latent_paths), len(ts)))
# 	for i in range(latent_paths.shape[0]):
# 		_, path = fastdtw(target_path, latent_paths[i], radius=2, dist=euclidean)
# 		for p in path:
# 			taus[i,p[1]] = ts[p[0]]
# 	return taus, target_path

def resample_paths(latent_paths, ts, taus):
	result = np.zeros(latent_paths.shape)
	for i in range(len(latent_paths)):
		interp = interp1d(taus[i], latent_paths[i], axis=0, fill_value='extrapolate', assume_sorted=True)
		result[i] = interp(ts)
	return result

# def fit_taus(latent_paths, ts):
# 	"""Quadratic time fits."""
# 	# Find tau fits.
# 	target_path = np.median(latent_paths, axis=0)
# 	# target_path = latent_paths[0]
# 	target_interp = interp1d(ts, target_path, axis=0, fill_value='extrapolate', assume_sorted=True)
# 	coeffs = np.zeros((len(latent_paths),2))
# 	ts_2 = np.power(ts, 2)
# 	for i in range(latent_paths.shape[0]):
# 		def loss(coeff):
# 			q_ts = coeff[0] * ts_2 + coeff[1] * ts
# 			return np.sum(np.abs(target_interp(q_ts) - latent_paths[i]))
# 		res = minimize(loss, np.array([0.0,1.0]))
# 		coeffs[i] = res.x
# 	taus = np.zeros((len(latent_paths), len(ts)))
# 	for i in range(len(taus)):
# 		taus[i] = coeffs[i,0] * ts_2 + coeffs[i,1] * ts
# 	return latent_paths, taus, target_path

def fit_taus(latent_paths, ts):
	"""Linear time fits."""
	# Find tau fits.
	target_path = latent_paths[0]
	target_interp = interp1d(ts, target_path, axis=0, fill_value='extrapolate', assume_sorted=True)
	coeffs = np.zeros((len(latent_paths),2))
	# ts_2 = np.power(ts, 2)
	for i in range(latent_paths.shape[0]):
		def loss(coeff):
			q_ts = coeff[1] * ts + coeff[0]
			return np.sum(np.abs(target_interp(q_ts) - latent_paths[i]))
		res = minimize(loss, np.array([0.0,1.0]))
		coeffs[i] = res.x
	taus = np.zeros((len(latent_paths), len(ts)))
	for i in range(len(taus)):
		taus[i] = coeffs[i,1] * ts + coeffs[i,0]
	return taus, target_path

def subtract_target(latent_paths, taus, target_path, ts):
	target_interp = interp1d(ts, target_path, axis=0, fill_value='extrapolate', assume_sorted=True)
	for i in range(latent_paths.shape[0]):
		for j in range(latent_paths.shape[1]):
			latent_paths[i,j] -= target_interp(taus[i,j])
	return latent_paths

# def plot_paths(latent_paths, ts):
# 	taus, target_path = fit_taus(latent_paths, ts)
#
# 	mean_path = np.mean(latent_paths, axis=0)
# 	# for i in range(len(latent_paths)):
# 		# latent_paths[i] -= mean_path
#
# 	latent_paths = subtract_target(latent_paths, taus, mean_path, ts) # subtract target/mean?
# 	latent_paths = apply_pca(latent_paths)
#
# 	_, axarr = plt.subplots(5,1, sharex=True)
# 	axarr[0].imshow(np.load('blk215_spec.npy'), origin='lower', aspect='auto', \
# 			extent=[np.min(ts), np.max(ts), 300, 12000])
# 	for j in range(4):
# 		for i in range(len(latent_paths)):
# 				axarr[j+1].plot(taus[i], latent_paths[i,:,j], c='b', lw=0.5, alpha=0.05)
# 		axarr[j+1].plot(ts, target_path[:,j], c='r', lw=1)
# 		axarr[j+1].set_ylabel('PC'+str(j+1))
# 	axarr[-1].set_xlabel('tau = f(t)')
# 	axarr[0].set_title("blk215 instantaneous song variability")
# 	plt.savefig('temp.pdf')
# 	plt.close('all')

def plot_paths_imaging(latent_paths, ts, loader, unit_num=4, filename='temp.pdf'):
	taus, target_path = fit_taus(latent_paths, ts)
	latent_paths = resample_paths(latent_paths, ts, taus)

	latent_paths -= np.mean(latent_paths, axis=0)
	# latent_paths /= np.std(latent_paths, axis=0)
	latent_paths = apply_pca(latent_paths)

	cmap = matplotlib.cm.get_cmap('viridis')

	imaging_info = get_imaging_info(loader)
	imaging_n = len(imaging_info['binary'])

	_, axarr = plt.subplots(3,1, sharex=True)
	axarr[0].imshow(np.load('blk215_spec.npy'), origin='lower', aspect='auto', \
			extent=[np.min(ts), np.max(ts), 300, 12000])
	spikes = np.array([spike for spike in imaging_info['real'] if spike is not None])
	max_spike = np.sqrt(np.max(spikes[:,unit_num]))

	no_spike_indices, spike_indices = [], []
	for i in range(imaging_n):
		if imaging_info['real'][i] is None:
			continue
		if imaging_info['real'][i][unit_num] > 0:
			spike_indices.append(i)
		else:
			no_spike_indices.append(i)
	no_spike_indices, spike_indices = np.array(no_spike_indices), np.array(spike_indices)
	if min(len(no_spike_indices), len(spike_indices)) < 10:
		return

	for j in range(2):
		# for i in range(imaging_n):
		# 	spikes = imaging_info['real'][i]
		# 	spike = 0.0 if spikes is None else np.sqrt(spikes[unit_num])/max_spike
		# 	alpha = 0.0 if spikes is None else 0.5
		# 	rgba = cmap(spike)
		# 	# c = 'r' if spikes is not None and spikes[unit_num] == 1 else 'k'
		# 	# alpha = 0.4 if c == 'r' else 0.1
		# 	axarr[j+1].plot(ts, latent_paths[i,:,j], c=rgba, lw=0.5, alpha=alpha)
		# for i in range(imaging_n, len(latent_paths)):
		# 	axarr[j+1].plot(ts, latent_paths[i,:,j], c='k', lw=0.5, alpha=0.01)

		# Quantiles
		index_set, colors = [no_spike_indices, spike_indices], ['b', 'r']
		for indices, color in zip(index_set, colors):
			temp = np.quantile(latent_paths[indices,:,j], 0.5, axis=0)
			axarr[j+1].plot(ts, temp, c=color, ls='-', alpha=0.8, lw=0.5)

			std = np.std(latent_paths[indices,:,j], axis=0, ddof=1)
			temp_1 = temp - std
			temp_2 = temp + std
			# temp_1 = np.quantile(latent_paths[indices,:,j], 0.25, axis=0)
			# temp_2 = np.quantile(latent_paths[indices,:,j], 0.75, axis=0)
			axarr[j+1].fill_between(ts, temp_1, temp_2, color=color, alpha=0.1)

		axarr[j+1].set_ylabel('PC'+str(j+1))
	axarr[-1].set_xlabel('tau = f(t)')
	axarr[0].set_title("blk215 instantaneous song variability")
	plt.savefig(filename)
	plt.close('all')

if __name__ == '__main__':
	pass

###


import objsocket
import queue
import threading
import pickle
import time

class Controller:
    '''
    Это класс написанный для примера, в реальности на его месте
    скорее всего будет программа на VB.net, которая будет подбирать
    параметры. Она должна будет посылать по сокетам то же, что и этот
    класс.
    '''

    def _init__(self, ip, task_port, result_port):
        self.task_socket = objsocket.objsocket()
        self.result_socket = objsocket.objsocket()

        self.task_socket.connect(ip, task_port)
        self.result_socket.connect(ip, result_port)

    def main(self):
        tasks = ['spam', 'eggs', 'ham']
        for i in range(3):
            task_socket.send_obj(tasks.pop())
            answer = result_socket.recv_obj()
            print(answer)

class Master:
    '''
    Класс, который получает задачи от чего-нибудь вроде Controller,
    распределяет их по Node-ам, получает результаты и посылает их
    обратно в Controller.
    '''

    def __init__(self, task_port, result_port, manager_port, allowed_node_addrs):
        self.tasks = queue.Queue()
        self.results = queue.Queue()
        self.active_task_stream = False
        self.controller = _ControllerConnection(task_port, result_port, self)
        self.node_manager = _NodeManager(manager_port, allowed_node_addrs, self)

    def start(self):
        self.controller.start()
        self.node_manager.start()

    def grab_result(self, result):
        '''Эту функцию вызывает _NodeConnection чтобы передать результаты в Master'''
        self.results.put(result)
        self.tasks.task_done()

    def give_task(self):
        '''Эту функцию вызывает _NodeConnection чтобы взять задание из Master'''
        try:
            task = self.tasks.get_nowait()
        except queue.Empty:
            task = None

        return task

    def reuse_task(self, task):
        '''Эту функцию вызывает _NodeConnection чтобы вернуть задание в Master в случае обрыва соединения с Node'''
        self.tasks.put(task)

    def put_task(self, task):
        '''Эту функцию будет вызывать _ControllerConnection чтобы положить задание в очередь Master'''
        self.tasks.put(task)

    def get_result(self):
        '''Эту функцию будет вызывать _ControllerConnection чтобы взять результат из очереди Master'''
        try:
            return self.results.get(timeout=1)
        except queue.Empty:
            return None

class _NodeManager:
    def __init__(self, manager_port, allowed_node_ips, master):
        self.master = master
        self.allowed_node_ips = allowed_node_ips

        self.thread = threading.Thread(target=self.main)
        self.nodes = {} # {ip: node, }

        self.sock = objsocket.objsocket()
        self.sock.bind_and_listen(manager_port, len(allowed_node_ips))

    def start(self):
        self.thread.start()

    def main(self):
        while True:
            conn, addr = self.sock.accept()
            ip = addr[0]
            print('Node with addr {} is trying to connect.'. format(addr))
            if ip in self.allowed_node_ips:
                if ip in self.nodes:
                    if self.nodes[ip].stop():
                        print('Stopped old node with ip {}'.format(ip))
                    else:
                        print('Old node with ip {} was already dead'.format(ip))
                node_connection = _NodeConnection(conn, self.master, self)
                self.nodes[ip] = node_connection
                node_connection.start()
                print('Connected to node at {}'.format(ip))
            else:
                print('Refused')

class _NodeConnection:
    def __init__(self, sock, master, manager):
        self.master = master
        self.sock = sock
        self.thread = threading.Thread(target=self.main)
        self.is_alive = True
        self.is_dead = False

    def start(self):
        self.thread.start()

    def stop(self):
        if not self.is_dead:
            self.is_alive = False
            while not self.is_dead:
                time.sleep(1)
            return True
        else:
            return False

    def main(self):
        while self.is_alive:
            task = self.master.give_task()

            if task == None:
                time.sleep(1)

            try:
                self.sock.send_obj(task)
            except objsocket.ServerUnavailibleException:
                print('Disconnected from node.')
                if task != None:
                    self.master.reuse_task(task)
                    print('Task returned to queue')
                break

            if task != None:
                print('Sent task {}'.format(task))

            try:
                result = self.sock.recv_obj()
            except objsocket.EndOfInputException:
                print('Disconnected from node.')
                if task != None:
                    self.master.reuse_task(task)
                    print('Task returned to queue')
                break

            if task != None:
                print('Got result')
                self.master.grab_result(result)

        self.is_dead = True

class _ControllerConnection:
    def __init__(self, task_port, result_port, master):
        self.master = master
        self.thread = threading.Thread(target=self.main)

        self.task_socket = objsocket.objsocket()
        self.result_socket = objsocket.objsocket()

        self.task_socket.bind_and_listen(task_port)
        self.result_socket.bind_and_listen(result_port)

    def start(self):
        self.thread.start()

    def main(self):
        while True:
            task_connection, task_adress = self.task_socket.accept()
            self.master.active_task_stream = True
            print('Connected to task_socket {}'.format(task_adress))
            result_connection, result_adress = self.result_socket.accept()
            print('Connected to result_socket {}'.format(result_adress))

            task_thread = threading.Thread(target=self._task_stream, args=(task_connection, ))
            result_thread = threading.Thread(target=self._result_stream, args=(result_connection, ))

            task_thread.start()
            result_thread.start()

            task_thread.join()
            result_thread.join()

    def _task_stream(self, task_connection):
        while True:
            try:
                task = task_connection.recv_obj()
            except objsocket.EndOfInputException:
                task_connection.close()
                self.master.active_task_stream = False
                print('Disconnected from task_socket')
                break
            print('Got task from cont {}'.format(task))
            self.master.put_task(task)

    def _result_stream(self, result_connection):
        while True:
            result = self.master.get_result()
            if result == None:
                pass
                if not self.master.active_task_stream:
                    result_connection.close()
                    print('Disconnected from result_socket and closed connection to cont')
                    break
            else:
                result_connection.send_obj(result)
                print('Sent result to cont')

class Node:
    '''
    Класс, который получает задания от Master и возвращает ответы.
    '''

    def __init__(self, function, master_ip, master_port):
        self.function = function
        self.master_ip = master_ip
        self.master_port = master_port

    def start(self):
        while True:
            sock = objsocket.objsocket()
            successful = sock.connect(self.master_ip, self.master_port)
            if successful:
                print('Connected to master {}@{}'.format(self.master_ip, self.master_port))
            else:
                time.sleep(1)
                continue

            while True:
                try:
                    task = sock.recv_data()
                except objsocket.EndOfInputException:
                    print('Disconnected from master {}@{}'.format(self.master_ip, self.master_port))
                    sock.close()
                    break

                task = pickle.loads(task)
                if task != None:
                    print('Got task {}'.format(task))
                    sock.send_obj(self.function(task))
                    print('Sent result')
                else:
                    sock.send_obj(None)


# Copyright 2019 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
An example of fuzz testing and then enhance non-robustness model.
"""
import random
import numpy as np

import mindspore
from mindspore import Model
from mindspore import context
from mindspore import Tensor
from mindspore.train.serialization import load_checkpoint, load_param_into_net
from mindspore.nn import SoftmaxCrossEntropyWithLogits
from mindspore.nn.optim.momentum import Momentum

from mindarmour.adv_robustness.defenses import AdversarialDefense
from mindarmour.fuzz_testing import Fuzzer
from mindarmour.fuzz_testing import ModelCoverageMetrics
from mindarmour.utils.logger import LogUtil

from examples.common.dataset.data_processing import generate_mnist_dataset
from examples.common.networks.lenet5.lenet5_net import LeNet5

LOGGER = LogUtil.get_instance()
TAG = 'Fuzz_testing and enhance model'
LOGGER.set_level('INFO')

def example_lenet_mnist_fuzzing():
    """
    An example of fuzz testing and then enhance the non-robustness model.
    """
    # upload trained network
    ckpt_path = '../common/networks/lenet5/trained_ckpt_file/lenet_m1-10_1250.ckpt'
    net = LeNet5()
    load_dict = load_checkpoint(ckpt_path)
    load_param_into_net(net, load_dict)
    model = Model(net)
    mutate_config = [{'method': 'Blur',
                      'params': {'auto_param': [True]}},
                     {'method': 'Contrast',
                      'params': {'auto_param': [True]}},
                     {'method': 'Translate',
                      'params': {'auto_param': [True]}},
                     {'method': 'Brightness',
                      'params': {'auto_param': [True]}},
                     {'method': 'Noise',
                      'params': {'auto_param': [True]}},
                     {'method': 'Scale',
                      'params': {'auto_param': [True]}},
                     {'method': 'Shear',
                      'params': {'auto_param': [True]}},
                     {'method': 'FGSM',
                      'params': {'eps': [0.3, 0.2, 0.4], 'alpha': [0.1]}}
                     ]

    # get training data
    data_list = "../common/dataset/MNIST/train"
    batch_size = 32
    ds = generate_mnist_dataset(data_list, batch_size, sparse=False)
    train_images = []
    for data in ds.create_tuple_iterator(output_numpy=True):
        images = data[0].astype(np.float32)
        train_images.append(images)
    train_images = np.concatenate(train_images, axis=0)

    # initialize fuzz test with training dataset
    model_coverage_test = ModelCoverageMetrics(model, 10, 1000, train_images)

    # fuzz test with original test data
    # get test data
    data_list = "../common/dataset/MNIST/test"
    batch_size = 32
    init_samples = 5000
    max_iters = 50000
    mutate_num_per_seed = 10
    ds = generate_mnist_dataset(data_list, batch_size, num_samples=init_samples,
                                sparse=False)
    test_images = []
    test_labels = []
    for data in ds.create_tuple_iterator(output_numpy=True):
        images = data[0].astype(np.float32)
        labels = data[1]
        test_images.append(images)
        test_labels.append(labels)
    test_images = np.concatenate(test_images, axis=0)
    test_labels = np.concatenate(test_labels, axis=0)
    initial_seeds = []

    # make initial seeds
    for img, label in zip(test_images, test_labels):
        initial_seeds.append([img, label])

    model_coverage_test.calculate_coverage(
        np.array(test_images[:100]).astype(np.float32))
    LOGGER.info(TAG, 'KMNC of test dataset before fuzzing is : %s',
                model_coverage_test.get_kmnc())
    LOGGER.info(TAG, 'NBC of test dataset before fuzzing is : %s',
                model_coverage_test.get_nbc())
    LOGGER.info(TAG, 'SNAC of test dataset before fuzzing is : %s',
                model_coverage_test.get_snac())

    model_fuzz_test = Fuzzer(model, train_images, 10, 1000)
    gen_samples, gt, _, _, metrics = model_fuzz_test.fuzzing(mutate_config,
                                                             initial_seeds,
                                                             eval_metrics='auto',
                                                             max_iters=max_iters,
                                                             mutate_num_per_seed=mutate_num_per_seed)

    if metrics:
        for key in metrics:
            LOGGER.info(TAG, key + ': %s', metrics[key])

    def split_dataset(image, label, proportion):
        """
        Split the generated fuzz data into train and test set.
        """
        indices = np.arange(len(image))
        random.shuffle(indices)
        train_length = int(len(image) * proportion)
        train_image = [image[i] for i in indices[:train_length]]
        train_label = [label[i] for i in indices[:train_length]]
        test_image = [image[i] for i in indices[:train_length]]
        test_label = [label[i] for i in indices[:train_length]]
        return train_image, train_label, test_image, test_label

    train_image, train_label, test_image, test_label = split_dataset(
        gen_samples, gt, 0.7)

    # load model B and test it on the test set
    ckpt_path = '../common/networks/lenet5/trained_ckpt_file/lenet_m2-10_1250.ckpt'
    net = LeNet5()
    load_dict = load_checkpoint(ckpt_path)
    load_param_into_net(net, load_dict)
    model_b = Model(net)
    pred_b = model_b.predict(Tensor(test_image, dtype=mindspore.float32)).asnumpy()
    acc_b = np.sum(np.argmax(pred_b, axis=1) == np.argmax(test_label, axis=1)) / len(test_label)
    print('Accuracy of model B on test set is ', acc_b)

    # enhense model robustness
    lr = 0.001
    momentum = 0.9
    loss_fn = SoftmaxCrossEntropyWithLogits(Sparse=True)
    optimizer = Momentum(net.trainable_params(), lr, momentum)

    adv_defense = AdversarialDefense(net, loss_fn, optimizer)
    adv_defense.batch_defense(np.array(train_image).astype(np.float32),
                              np.argmax(train_label, axis=1).astype(np.int32))
    preds_en = net(Tensor(test_image, dtype=mindspore.float32)).asnumpy()
    acc_en = np.sum(np.argmax(preds_en, axis=1) == np.argmax(test_label, axis=1)) / len(test_label)
    print('Accuracy of enhensed model on test set is ', acc_en)

if __name__ == '__main__':
    # device_target can be "CPU", "GPU" or "Ascend"
    context.set_context(mode=context.GRAPH_MODE, device_target="Ascend")
    example_lenet_mnist_fuzzing()


import os
from copy import copy, deepcopy
from typing import TYPE_CHECKING, Dict, List, Optional, Union

from funcy import first

from dvc.exceptions import DvcException
from dvc.render import Renderer
from dvc.render.utils import get_files

if TYPE_CHECKING:
    from dvc.repo.plots.template import PlotTemplates
    from dvc.types import StrPath

REVISION_FIELD = "rev"
INDEX_FIELD = "step"

class PlotMetricTypeError(DvcException):
    def __init__(self, file):
        super().__init__(
            "'{}' - file type error\n"
            "Only JSON, YAML, CSV and TSV formats are supported.".format(file)
        )

class PlotDataStructureError(DvcException):
    def __init__(self):
        super().__init__(
            "Plot data extraction failed. Please see "
            "https://man.dvc.org/plots for supported data formats."
        )

def _filter_fields(
    datapoints: List[Dict], filename, revision, fields=None
) -> List[Dict]:
    if not fields:
        return datapoints
    assert isinstance(fields, set)

    new_data = []
    for data_point in datapoints:
        new_dp = copy(data_point)

        keys = set(data_point.keys())
        if keys & fields != fields:
            raise DvcException(
                "Could not find fields: '{}' for '{}' at '{}'.".format(
                    ", ".join(fields), filename, revision
                )
            )

        to_del = keys - fields
        for key in to_del:
            del new_dp[key]
        new_data.append(new_dp)
    return new_data

def _lists(dictionary):
    for _, value in dictionary.items():
        if isinstance(value, dict):
            yield from _lists(value)
        elif isinstance(value, list):
            yield value

def _find_data(data: Union[Dict, List], fields=None) -> List[Dict]:
    if not isinstance(data, dict):
        return data

    if not fields:
        # just look for first list of dicts
        fields = set()

    for lst in _lists(data):
        if (
            all(isinstance(dp, dict) for dp in lst)
            and set(first(lst).keys()) & fields == fields
        ):
            return lst
    raise PlotDataStructureError()

def _append_index(datapoints: List[Dict], append_index=False) -> List[Dict]:
    if not append_index or INDEX_FIELD in first(datapoints).keys():
        return datapoints

    for index, data_point in enumerate(datapoints):
        data_point[INDEX_FIELD] = index
    return datapoints

def _append_revision(datapoints: List[Dict], revision) -> List[Dict]:
    for data_point in datapoints:
        data_point[REVISION_FIELD] = revision
    return datapoints

class VegaRenderer(Renderer):
    DIV = """
    <div id = "{id}">
        <script type = "text/javascript">
            var spec = {partial};
            vegaEmbed('#{id}', spec);
        </script>
    </div>
    """

    def __init__(self, data: Dict, templates: "PlotTemplates"):
        super().__init__(data)
        self.templates = templates

    def _squash_props(self) -> Dict:
        resolved: Dict[str, str] = {}
        for rev_data in self.data.values():
            for file_data in rev_data.get("data", {}).values():
                props = file_data.get("props", {})
                resolved = {**resolved, **props}
        return resolved

    def _datapoints(self, props: Dict):
        fields = props.get("fields", set())
        if fields:
            fields = {*fields, props.get("x"), props.get("y")} - {None}

        datapoints = []
        for revision, rev_data in self.data.items():
            for filename, file_data in rev_data.get("data", {}).items():
                if "data" in file_data:
                    tmp = deepcopy(file_data.get("data"))
                    tmp = _find_data(tmp, fields=fields - {INDEX_FIELD})
                    tmp = _append_index(
                        tmp, append_index=props.get("append_index", False)
                    )
                    tmp = _filter_fields(
                        tmp,
                        filename=filename,
                        revision=revision,
                        fields=fields,
                    )
                    tmp = _append_revision(tmp, revision=revision)
                    datapoints.extend(tmp)
        return datapoints

    def get_vega(self) -> Optional[str]:
        props = self._squash_props()

        template = self.templates.load(props.get("template", None))

        if not props.get("x") and template.has_anchor("x"):
            props["append_index"] = True
            props["x"] = INDEX_FIELD

        datapoints = self._datapoints(props)

        if datapoints:
            if not props.get("y") and template.has_anchor("y"):
                fields = list(first(datapoints))
                skip = (REVISION_FIELD, props.get("x"))
                props["y"] = first(
                    f for f in reversed(fields) if f not in skip
                )
            return template.render(datapoints, props=props)
        return None

    def _convert(self, path: "StrPath"):
        return self.get_vega()

    @staticmethod
    def matches(data):
        files = get_files(data)
        extensions = set(map(lambda f: os.path.splitext(f)[1], files))
        return extensions.issubset({".yml", ".yaml", ".json", ".csv", ".tsv"})


from django.urls import path

from contracts import views

urlpatterns = [

    path("manage_job_contracts", views.manage_job_contracts, name="manage_job_contracts"),
    path("terminate_contract/<int:contract_id>/", views.terminate_contract, name="terminate_contract"),
    path("edit_contract_page/<int:contract_id>/", views.edit_contract_page, name="edit_contract_page"),
    path("terminated_contracts", views.terminated_contracts_page, name="terminated_contracts_page"),
    path("activate_contract/<int:contract_id>/", views.activate_contract, name="activate_contract"),
    path("user_contracts_page", views.user_contracts_page, name="user_contracts_page"),
    path("manage_offences_page", views.manage_offences_page, name="manage_offences_page"),
    path("edit_offence_page/<int:offence_id>/", views.edit_offence_page, name="edit_offence_page"),
    path("delete_offence/<int:offence_id>/", views.delete_offence, name="delete_offence"),
    path("manage_penalties_page", views.manage_penalties_page, name="manage_penalties_page"),
    path("edit_penalty_page/<int:penalty_id>/", views.edit_penalty_page, name="edit_penalty_page"),
    path("delete_penalty/<int:penalty_id>/", views.delete_penalty, name="delete_penalty"),
    path("manage_terminations_page", views.manage_terminations_page, name="manage_terminations_page"),
    path("edit_termination/<int:termination_id>/", views.edit_termination_page, name="edit_termination_page"),
    path("delete_termination/<int:termination_id>/", views.delete_termination, name="delete_termination"),
]


import scrapy
import dateparser
from datetime import datetime
import pytz
from cmate.items import CmateItem

class EspnRss(scrapy.spiders.XMLFeedSpider):

  name = "EspnRss"
  allowed_domains = ["espn.com"]

  start_urls = [
    'http://www.espn.com/espn/rss/news',
  ]
  itertag = 'item'

  # Parse a node of the RSS tree
  def parse_node(self, response, node):
    ## Open python shell here
    #vars = globals().copy()
    #vars.update(locals())
    #shell = code.InteractiveConsole(vars)
    #shell.interact()

    link = node.xpath('//item//link/text()').extract_first()
    return scrapy.http.Request(link, callback=self.parse_article)

  # Parse ESPN article
  def parse_article(self, response):
    item = CmateItem()

    item['source'] = 'ESPN'
    item['category'] = 'sports'
    item['title'] = response.css('header.article-header').css('h1::text').extract_first()

    bodyList = response.css('div.article-body').css('p::text').extract()
    item['body'] = '\n'.join(bodyList)

    lastUpdated = response.css('span.timestamp::attr(data-date)').extract_first()
    item['posted_time'] = dateparser.parse(lastUpdated)
    item['scrape_time'] = datetime.now(pytz.utc)

    item['url'] = response.url

    yield item


# -*- coding: utf-8 -*-
"""
Created on Sat May 18 20:23:11 2020

@author: nikic
"""
import numpy
def FancyPrint(A, B, selected):
    for row in range(len(B)):
        print("(", end='')
        for col in range(len(A[row])):
            print("\t{1:10.2f}{0}".format(" " if (selected is None or selected != (row, col)) else "*", A[row][col]), end='')
            print("\t) * (\tX{0}) = (\t{1:10.2f})".format(row + 1,
            B[row]))
data = numpy.genfromtxt('/nympy-gauss-slau.csv', delimiter=';')
matrix_list = []
matrix = []
for row in data:
    first_col = row[0]
    if numpy.isnan(first_col):
        matrix_list.append(matrix)
        matrix = []
        continue
    mask = ~numpy.isnan(row)
    matrix.append(row[mask])
matrix_list.append(matrix)
f = open('numpy-gauss-slv.csv', 'wb+')
f.truncate()
for matrix in matrix_list:
    M = numpy.array(matrix)
    myA = numpy.delete(M, M.shape[1] - 1, axis=1)
    myB = M[:, [-1]].flatten()
    print("Исходная система:")
    FancyPrint(myA, myB, None)
    slv = numpy.linalg.solve(myA, myB)
    print("Решение:")
    print(slv)
    numpy.savetxt(f, numpy.array([slv]), delimiter=',')
f.close()


#!/usr/bin/env python
# -*- coding:utf-8 -*-
############################################
# File Name    : __init__.py
# Created By   : Suluo - sampson.suluo@gmail.com
# Creation Date: 2019-07-29
# Last Modified: 2019-08-11 19:56:46
# Descption    :
# Version      : Python 3.7
############################################
import argparse
import torch
from torch import nn
from Net import Constants

def cal_correct(pred, gold):
    """cal correct"""
    pred = torch.argmax(pred, dim=-1)
    non_pak_mask = gold.ne(Constants.PAD)
    n_correct = pred.eq(gold)
    n_correct = n_correct.masked_select(non_pak_mask).sum().item()
    return n_correct

class Loss(nn.Module):
    """Loss 基类"""
    def __init__(self, opt):
        super().__init__()
        self.opt = opt

    def cal_performance(self, pred, gold):
        """cal loss and correct"""
        loss = self.forward(pred, gold)
        n_correct = cal_correct(pred, gold)
        return loss, n_correct

    def forward(self, pred, gold):
        """cal Loss"""
        raise NotImplementedError()

def main(args):
    """main"""
    return args

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('-e', '--num', type=int, default=100, help='input num')
    args = parser.parse_args()
    main(args)


# -*- encoding: utf-8 -*-
import sys
sys.path.append("/Users/wang 1/researchDemo/researchDemo/Research/")

import twitterAPIKey
import stop_words
import re
import numpy as np
from langdetect import detect
from janome.tokenizer import Tokenizer
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.decomposition import LatentDirichletAllocation
from sklearn.metrics.pairwise import cosine_similarity
from sklearn.feature_extraction.text import CountVectorizer
import nltk.stem

from sklearn.model_selection import GridSearchCV
# import mglearn
# from sklearn.pipeline import make_pipeline
# from sklearn.linear_model import LogisticRegression
stop = stop_words.stop_words

###################################################
# Link to Twitter
try:
    api = twitterAPIKey.linkToTwitter()
except RuntimeError:
    print('Fail to access to Twitter!\n')
else:
    print("Access to Twitter successfully!\n")

def number_of_following(user_name):
    twitter_name = user_name
    following_list = api.friends_ids(screen_name=twitter_name)
    return user_name + " has followed " + str(len(following_list)) + " accounts."

###################################################
# Get tweets from target user
def Tweets(user_name):
    user_tweets = api.user_timeline(screen_name=user_name, count=500)
    return user_tweets

###################################################
# Get tweets from users' following list
class TimeLineContent:

    def __init__(self, id):
        self.id = id
        self.name = api.get_user(id=self.id).name

    def get_timeline(self, count):
        print("\nGetting " + self.name + "'s tweets...")
        try:
            timelines = api.user_timeline(id=self.id, count=count)
        except:
            print("Failed to get " + self.name + "'s tweets!")
        else:
            print("Got " + self.name + "'s tweets successfully!")
            return timelines

    def output_timeline(self, count=500):
        timelines = self.get_timeline(count)
        timeline_new = ""
        try:
            for timeline in timelines:
                timeline_new += timeline.text.replace('\n', ' ')
        except:
            pass
        else:
            return timeline_new

###################################################
# use Jnome to process Japanese
def process_jp(timeline):
    t = Tokenizer()
    new_time_line = [token.surface for token in t.tokenize(timeline) if token.part_of_speech.split(',')[0] in ["名詞", "形容詞"]]
    return new_time_line

###################################################
# delete not English words
def englishlize(textList):
    temp = []
    for i in textList:
        temp2 = []
        print(i)
        for n in i.split(" "):
            if re.match("^[@#]*[A-Za-z0-9]+$", n):
                temp2.append(n)
            else:
                try:
                    detect(n)
                except:
                    pass
                else:
                    if detect(n) == 'jp':
                        temp.append(" ".join(process_jp(n)))

        temp.append(" ".join(temp2))
    return temp

###################################################
# Get favorite contents
def favorites(account_id):
    favorites_list = []

    for i in api.favorites(id=account_id):
        if i.text:
            if detect(i.text) == 'en':
                vect = CountVectorizer(max_features=140, stop_words=stop)
                i = englishlize(i.text.replace('\n', ' ').split(' '))
                vect.fit_transform(i)
                favorites_list.append(" ".join(vect.get_feature_names()))

            elif detect(i.text) == 'jp':
                temp = process_jp(i)
                favorites_list.append(temp)

            else:
                pass

    print(favorites_list)
    return favorites_list

###################################################
# LDA model
def lda_clustering(new_timeline, n=5):
    print('\nStart to do clustering . . .')
    #################
    # Bag of Words
    # Delete stop words
    # Delete non-english words
    # parameters = {'learning_method': ['batch', 'online'],
    #               'n_topics': [5, 10, 15, 20, 25],
    #               'perp_tol': [0.001, 0.01, 0.1],
    #               'doc_topic_prior': [0.001, 0.01, 0.05, 0.1, 0.2],
    #               'topic_word_prior': [0.001, 0.01, 0.05, 0.1, 0.2],
    #               'max_iter': [1000]}
    # print(englishlize(new_timeline))

    # for i in new_timeline:
    #     for_lda = []
    #     if i:
    #         if detect(i) == 'en':
    #             temp = englishlize(new_timeline)
    #             for_lda.append(temp)
    #             print(temp)
    #
    #         elif detect(i) == 'jp':
    #             temp = " ".join(process_jp(i))
    #             for_lda.append(temp)
    #         else:
    #             pass

    temp = englishlize(new_timeline)
    vect = CountVectorizer(max_features=40000, min_df=.1, stop_words=stop)
    X = vect.fit_transform(temp)
    #################
    # LDA
    lda = LatentDirichletAllocation(n_topics=n, learning_method='batch', max_iter=25, random_state=0)
    lda.fit_transform(X)
    # lda = LatentDirichletAllocation()
    # model = GridSearchCV(lda, parameters)
    # model.fit(X)

    # print(model.best_estimator_)
    # print(model.best_params_)

    feature_names = np.array(vect.get_feature_names())

    topic_list = []
    index = 0
    for topic_idx, topic in enumerate(lda.components_):
        # message = "Topic #%d: " % topic_idx
        message = ""
        message += " ".join([feature_names[i] for i in topic.argsort()[:-10 - 1:-1]])
        topic_list.append(message)
        print('\n', "#Topic", index, ":", message)
        index += 1
    print('\nFinishing clustering . . .')
    return topic_list

###################################################
# cosine similarity
def cos_similarity(key_word_list, profile):
    train_set = [key_word_list]
    train_set.append(" ".join([x for x in profile]))
    count_vectorizer = CountVectorizer()
    count_matrix_train = count_vectorizer.fit_transform(train_set)
    similarity = cosine_similarity(count_matrix_train[0:1], count_matrix_train)
    return similarity

###################################################
# main
def runLDA(user_name):
    new_timeline = []
    following_list = api.friends_ids(screen_name=user_name)
    print(len(following_list), "in following list.")
    following_count = 0

    for i in following_list:
        following_count += 1
        try:
            temp = TimeLineContent(i)
        except:
            print('\nFailed to get information from user id:', i)
            continue
        else:
            new_timeline.append(str(temp.output_timeline(200)))
        print("(", following_count, "/", len(following_list), ") Done")
        # if following_count == 3:
        #     break

    topic_list = lda_clustering(new_timeline)
    return topic_list

###################################################
# calculate cosine similarity
def runCos(user_name, topic_list):
    topic_index = 0
    favorites_list = favorites(user_name)
    result = []
    for i in topic_list:
        print('\nTopic #', topic_index, '***************')
        result.append("Topic " + str(topic_index) + "'s Cosine similarity is " + str(cos_similarity(i, favorites_list)[0][1]))
        print("Cosine similarity is", cos_similarity(i, favorites_list)[0][1])
        topic_index += 1
    return result

###################################################
# calculate tf
def runTF(user_name, topic_list):
    topic_index = 0
    favorites_list = favorites(user_name)
    result = []
    temp = " ".join(favorites_list).split(' ')
    for i in topic_list:
        temp_topic_list = i.split(' ')
        # print(temp)
        # print(temp_topic_list)
        sum_tf = 0
        for j in temp_topic_list:
            if temp.count(j) != 0:
                sum_tf += temp_topic_list.count(j) / len(temp)
        print('\nTopic #', topic_index, '***************')
        result.append("Topic " + str(topic_index) + "'s TF is " + str(sum_tf))
        print("TF is", str(sum_tf))
        topic_index += 1
    return result

###################################################
# read local test data
def local_test():
    f = open("/Users/wang 1/researchDemo/researchDemo/Research/timeline.txt", "r+")
    text = f.read()
    timeline = text[2:-2].split("', '")
    f.close()
    return timeline

# twitter_name = input("アカウントを入力してください：")
# twitter_name = "MyTopicTest"
# new_timeline = []
# following_list = api.friends_ids(screen_name=twitter_name)
# print(len(followingList), "in following list.")
# followingCount = 0
#
# for i in following_list:
#     following_count += 1
#     try:
#         temp = TimeLineContent(i)
#     except:
#         print('\nFailed to get information from user id:', i)
#         continue
#     else:
#         new_timeline.append(str(temp.output_timeline(200)))
#     print("(", following_count, "/", len(following_list), ") Done")
#     # if following_count == 3:
#     #     break

# f = open('timeline.txt', "r+")
# f.write(str(new_timeline))
# text = f.read()
# timeline = text[2:-2].split("', '")
#
# lda_clustering(timeline)
#
# f.close()
# topic_list = lda_clustering(new_timeline, 10)

# print('\nFavorites:')
# favorites_list = favorites(twitter_name)
# print()

# topic_index = 0
# for i in topic_list:
#     print('\nTopic #', topic_index, '***************')
#     print("Cosine similarity is", cos_similarity(i, favorites_list)[0][1])
#     topic_index += 1

# a = api.get_user(id='PokemonGoApp')
# print('\nPokemonGpApp: ' + a.description)


#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import shutil
import logging
import datetime
import time
import numpy as np

np.set_printoptions(precision=6, suppress=True)

from matplotlib import pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

from utils import *
from evaluation import *
import argparse

# setup logging
for handler in logging.root.handlers[:]:
    logging.root.removeHandler(handler)
if not os.path.isdir('logs'):
    os.makedirs('logs')
log_file = 'logs/log_{}.log'.format(datetime.datetime.now().strftime("%Y_%B_%d_%I-%M-%S%p"))
open(log_file, 'a').close()

# create logger
logger = logging.getLogger('main')
logger.setLevel(logging.INFO)
formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')

# create console handler and set level to debug
ch = logging.StreamHandler()
ch.setLevel(logging.INFO)
ch.setFormatter(formatter)
logger.addHandler(ch)

# add to log file
fh = logging.FileHandler(log_file)
fh.setLevel(logging.INFO)
fh.setFormatter(formatter)
logger.addHandler(fh)

def log(str): logger.info(str)

import torch
from torch.utils.data import DataLoader
from torch.autograd import Variable
import torch.nn as nn
import torch.nn.functional as F

log('is GPU available? {}'.format(torch.cuda.is_available()))
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

parser = argparse.ArgumentParser(description="TGGAN")

datasets = ['taxi', 'authorship', 'checkin', 'sim_100', 'sim_500', 'sim_2000']
parser.add_argument("-d", "--dataset", default="checkin", type=str,
                    help="one of: {}".format(", ".join(sorted(datasets))))
processes = ['rand_binomial', 'rand_poisson']
parser.add_argument("-sp", "--simProcess", default="rand_binomial", type=str,
                    help="one of: {}".format(", ".join(sorted(processes))))
parser.add_argument("-nn", "--numberNode", default=10, type=int,
                    help="if run simulation data, this is the number of nodes")
parser.add_argument("-p", "--probability", default=0.5, type=int,
                    help="if run simulation data, this is the number of time slices")
parser.add_argument("-nt", "--numberTime", default=10, type=int,
                    help="this is the number of time slices for both real data and simulation data")
parser.add_argument("-sc", "--scale", default=0.1, type=float,
                    help="scale of gaussian prior for kernel density estimation in DeepTemporalWalk")
parser.add_argument("-iw", "--init_walk_method", default='uniform', type=str,
                    help="TemporalWalk sampler")

# DeepTemporalWalk
parser.add_argument("-bs", "--batch_size", default=32, type=int,
                    help="random walks batch size in DeepTemporalWalk")
parser.add_argument("-lr", "--learningrate", default=1e-3, type=float,
                    help="if this run should run all evaluations")
parser.add_argument("-rl", "--rw_len", default=4, type=int,
                    help="random walks maximum length in DeepTemporalWalk")
parser.add_argument("--num-G-layer", default=15, type=int,
                    help="The number of layers in the Generator")
parser.add_argument("-ud", "--use_decoder", default='normal', type=str,
                    help="if decoder function")
parser.add_argument("-es", "--embedding_size", default=16, type=int,
                    help="embedding size of nodes, W_down")
parser.add_argument("-td", "--time_deconv", default=8, type=int,
                    help="deconv output channels number")
parser.add_argument("-ts", "--time_sample_num", default=4, type=int,
                    help="time sampling number")
parser.add_argument("-cm", "--constraint_method", default='min_max', type=str,
                    help="time constraint computing method")
parser.add_argument("-ne", "--n_eval_loop", default=40, type=int,
                    help="number of walk loops")
parser.add_argument("--W_down_generator_size", default=128, type=int,
                    help="The size of W_down_generator_size")
parser.add_argument("--W_down_discriminator_size", default=128, type=int,
                    help="The size of W_down_discriminator_size")
parser.add_argument("--noise-dim", default=64,
                    help="The dim of the random noise that is used as input.")
parser.add_argument("--hidden-unit", default=64,
                    help="The dim of the random noise that is used as input.")
parser.add_argument("--x-mode", default='uniform', type=str,
                    help="time constraint computing method")
parser.add_argument("--noise-type", choices=["Gaussian", "Uniform"], default="Gaussian",
                    help="The noise type to feed into the generator.")
parser.add_argument("--constraint-method", choices=["min_max", "relu", "clip"], default="min_max",
                    help="The constraint method for time budget.")
parser.add_argument("--spatial-constraint-method", choices=["nearest_point"], default="nearest_point",
                    help="The constraint method for geo location.")
parser.add_argument("--n-epochs", default=5, type=int,
                    help="max epoches")
parser.add_argument("--clip-value", type=float, default=0.05,
                    help="lower and upper clip value for disc. weights")
parser.add_argument("--n-critic", type=int, default=5,
                    help="number of training steps for discriminator per iter")
parser.add_argument("-mi", "--max_iters", default=100000, type=int,
                    help="max iterations")
parser.add_argument("-ev", "--eval_every", default=1000, type=int,
                    help="evaluation interval of epochs")

parser.add_argument("-te", "--is_test", default=True, type=bool,
                    help="if this is a testing period.")
parser.add_argument("--contact_time", default=0.01, type=float,
                    help="the contact time when evaluating the temporal network")
parser.add_argument("--early_stopping", default=1e-5, type=float,
                    help="stop training if evaluation metrics are good enough")

args = parser.parse_args()

class TemporalDataset(torch.utils.data.Dataset):
    def __init__(self, edges, n_nodes, t_end, scale, args):
        self.edges = edges
        self.n_nodes = n_nodes
        self.t_end = t_end
        self.scale = scale

        self.rw_len = args.rw_len
        self.init_walk_method = args.init_walk_method
        # Initialize a temporal walker
        self.walker = TemporalWalker(self.n_nodes, self.edges, self.t_end, self.scale, self.rw_len, 1,
                                init_walk_method=self.init_walk_method)

        self.data = self.cache(self.walker)

    def cache(self, walker):
        real_walk_data = []
        for i in range(20000):
            temp = walker.walk().__next__()[0]
            real_walk_data.append(temp)

        print("Done!")
        return real_walk_data

    def __getitem__(self, item):
        walks = self.data[item]
        return walks

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

class Generator(nn.Module):
    def __init__(self, N, n_samples, t_end, args):
        super(Generator, self).__init__()
        self.N = N
        self.t_end = t_end

        self.x_mode = args.x_mode
        self.rw_len = args.rw_len
        self.num_G_layer = args.num_G_layer
        self.use_decoder = args.use_decoder
        self.batch_size = args.batch_size
        self.noise_dim = args.noise_dim
        self.hidden_unit = args.hidden_unit
        self.W_down_generator_size = args.W_down_generator_size
        self.constraint_method = args.constraint_method
        self.spatial_constraint_method = args.spatial_constraint_method

        #self.polygon, self.zone = read_geojson(args.dataset)
        self.init_lin_1 = nn.Linear(self.noise_dim, self.W_down_generator_size)
        self.init_lin_2_h = nn.Linear(self.W_down_generator_size, self.W_down_generator_size)
        self.init_lin_2_c = nn.Linear(self.W_down_generator_size, self.W_down_generator_size)

        self.lin_input_x = nn.Sequential(
                            nn.Linear(2, 1, bias=False),
                            nn.Linear(1, self.W_down_generator_size),
                            nn.Tanh()
                            )

        self.lin_t0_loc = nn.Sequential(
                            nn.Linear(self.W_down_generator_size, self.hidden_unit),
                            nn.Linear(self.hidden_unit, 1)
                            )
        self.lin_t0_scale = nn.Sequential(
                            nn.Linear(self.W_down_generator_size, self.hidden_unit),
                            nn.Linear(self.hidden_unit, 1)
                            )
        self.lin_t0_last = nn.Linear(1, self.W_down_generator_size)

        self.lin_tau_loc = nn.Sequential(
                            nn.Linear(self.W_down_generator_size, self.hidden_unit),
                            nn.Linear(self.hidden_unit, 1)
                            )
        self.lin_tau_scale = nn.Sequential(
                            nn.Linear(self.W_down_generator_size, self.hidden_unit),
                            nn.Linear(self.hidden_unit, 1)
                            )
        self.lin_tau_last = nn.Linear(1, self.W_down_generator_size)

        self.lin_coord_x_loc = nn.Sequential(
                            nn.Linear(self.W_down_generator_size, self.hidden_unit),
                            nn.Linear(self.hidden_unit, 1)
                            )
        self.lin_coord_x_scale = nn.Sequential(
                            nn.Linear(self.W_down_generator_size, self.hidden_unit),
                            nn.Linear(self.hidden_unit, 1)
                            )

        self.lin_coord_y_loc = nn.Sequential(
                            nn.Linear(self.W_down_generator_size, self.hidden_unit),
                            nn.Linear(self.hidden_unit, 1)
                            )
        self.lin_coord_y_scale = nn.Sequential(
                            nn.Linear(self.W_down_generator_size, self.hidden_unit),
                            nn.Linear(self.hidden_unit, 1)
                            )

        self.lin_coord_last = nn.Linear(2, self.W_down_generator_size)

        self.lin_w_down = nn.Linear(self.W_down_generator_size, self.N)
        self.lin_w_up = nn.Linear(self.N, self.W_down_generator_size)

        self.lin_end = nn.Sequential(
                            nn.Linear(self.W_down_generator_size, self.hidden_unit),
                            nn.Tanh(),
                            nn.Linear(self.hidden_unit, 2)
                            )

        self.lstm = nn.LSTM(self.W_down_generator_size, self.W_down_generator_size, self.num_G_layer)

    def generate_time(self, output, name):
        if name == 't0':
            loc_t0 = self.lin_t0_loc(output)
            scale_t0 = self.lin_t0_scale(output)
            #time = [self.truncated_normal_(torch.ones(1).to(device), mean=loc_t0[i, 0], std=scale_t0[i, 0]) for i in range(self.batch_size)]
            time = torch.distributions.normal.Normal(loc_t0, torch.relu(scale_t0)+1e-7).rsample()
        elif name == 'tau':
            loc_t0 = self.lin_tau_loc(output)
            scale_t0 = self.lin_tau_scale(output)
            #time = [self.truncated_normal_(torch.ones(1).to(device), mean=loc_t0[i, 0], std=scale_t0[i, 0]) for i in range(self.batch_size)]
            time = torch.distributions.normal.Normal(loc_t0, torch.relu(scale_t0)+1e-7).rsample()
        #time = torch.stack(time)
        return time

    def generate_coordinate(self, output):
        x_loc = self.lin_coord_x_loc(output)
        x_scale = self.lin_coord_x_scale(output)
        y_loc = self.lin_coord_y_loc(output)
        y_scale = self.lin_coord_y_scale(output)
        x_coord = torch.distributions.normal.Normal(x_loc, torch.relu(x_scale)+1e-7).rsample()
        y_coord = torch.distributions.normal.Normal(y_loc, torch.relu(y_scale)+1e-7).rsample()
        return torch.cat((x_coord, y_coord), dim=1)

    def truncated_normal_(self, tensor, mean=0, std=0.03):
        size = tensor.shape
        tmp = tensor.new_empty(size+(4,)).normal_()
        valid = (tmp < 2) & (tmp > -2)
        ind = valid.max(-1, keepdim=True)[1]
        tensor.data.copy_(tmp.gather(-1, ind).squeeze(-1))
        tensor.data.mul_(std).add_(mean)
        return tensor

    def time_constraint(self, t, epsilon=1e-2, method='min_max'):
        if method == 'relu':
            t = torch.relu(t) - torch.relu(t - 1.)
        elif method == 'clip':
            t = torch.clamp(t, 0., 1.)
        elif method == 'min_max':
            min_ = torch.min(t)
            max_ = torch.max(t)

            if min_ < epsilon:
                t = t-min_
            if max_ > 1.0:
                t = t/max_
        return t

    def spatial_constraint(self, coordinates, polygon, zone, method='nearest_point'):
        with torch.no_grad():
            cloest_point = []
            for coordinate in coordinates:
                point = Point(utm.from_latlon(coordinate.cpu().numpy()[0], coordinate.cpu().numpy()[1]))
                if polygon.contains(point):
                    cloest_point.append(coordinate)
                else:
                    p1, p2 = nearest_points(polygon, point)
                    nearest_point = list(map(float, re.findall("\d+\.\d+", p1.wkt)))
                    nearest_point = utm.to_latlon(nearest_point[0], nearest_point[1], zone[0], zone[1])
                    cloest_point.append(torch.Tensor(nearest_point))
            return torch.stack(cloest_point)

    def forward(self, z, x_input=None, t0_input=None, edge_input=None, coord_input=None, tau_input=None):
        init_c, init_h = [], []
        for _ in range(self.num_G_layer):
            intermediate = torch.tanh(self.init_lin_1(z))
            init_c.append(torch.tanh(self.init_lin_2_c(intermediate)))
            init_h.append(torch.tanh(self.init_lin_2_h(intermediate)))
        # Initialize an input tensor
        hidden = (torch.stack(init_c, dim=0), torch.stack(init_h, dim=0))

        # generate start x, and its residual time
        if self.x_mode == 'uniform':
            if x_input is not None:
                x_output = x_input
            else:
                # generate start node binary if not need
                x_output = torch.randint(low=0, high=2, size=(self.batch_size,))
                x_output = F.one_hot(x_output, 2).to(device)
        elif self.x_mode == 'generate':
            print("This method is not implemented yet")

        inputs = self.lin_input_x(x_output.float())

        # LSTM steps
        node_outputs = []
        tau_outputs = []
        coord_outputs = []

        # generate the first three start elements: start x, residual time, and maximum possible length
        out, hidden = self.lstm(inputs.unsqueeze(0), hidden)
        if t0_input is not None: # for evaluation generation
            t0_res_output = t0_input
        else:
            t0_res_output = self.generate_time(out.squeeze(0), 't0')

            if self.constraint_method != "none":
                t0_res_output = self.time_constraint(
                    t0_res_output, method=self.constraint_method) * self.t_end

            t0_reference = torch.ones(self.batch_size).to(device)
            condition = torch.eq(torch.argmax(x_output, axis=-1), t0_reference)
            t0_res_output = torch.where(condition, t0_reference, t0_res_output.squeeze(-1)).unsqueeze(-1)

        res_time = t0_res_output
        # convert to input
        inputs = torch.tanh(self.lin_t0_last(t0_res_output))

        # LSTM time steps
        for i in range(self.rw_len):
            # generate temporal edge part
            for j in range(2):
                out, hidden = self.lstm(inputs.unsqueeze(0), hidden)
                if edge_input is not None and i <= self.rw_len-2:  # for evaluation generation
                    output = edge_input[:, i*2 + j]
                else:
                    # Decrease to dimension N
                    logit = self.lin_w_down(out.squeeze(0))
                    output = F.gumbel_softmax(logit, dim=1, tau=3, hard=True)

                node_outputs.append(output)
                # Back to dimension d
                inputs = self.lin_w_up(output)

                # Calculate the coordinates
                if coord_input is not None and i <= self.rw_len-2:  # for evaluation generation
                    coord_output = coord_input[:, i*2 + j]
                else:
                    out, hidden = self.lstm(inputs.unsqueeze(0), hidden)
                    coord_output = self.generate_coordinate(out.squeeze(0))
                    #print(coord_output[0, :])
                    #if self.spatial_constraint_method != "none":
                        #coord_output = self.spatial_constraint(coord_output, self.polygon, self.zone).to(device)

                coord_outputs.append(coord_output)
                # Back to dimension d
                inputs = self.lin_coord_last(coord_output)

            # LSTM for tau
            out, hidden = self.lstm(inputs.unsqueeze(0), hidden)
            if tau_input is not None and i <= self.rw_len-2:  # for evaluation generation
                tau = tau_input[:, i]
            else:
                tau = self.generate_time(out.squeeze(0), 'tau')

                if self.constraint_method != "none":
                    tau = self.time_constraint(tau, method=self.constraint_method) * res_time

            res_time = tau
            tau_outputs.append(tau)
            # convert to input
            inputs = torch.tanh(self.lin_tau_last(tau))

        # LSTM for end indicator
        out, hidden = self.lstm(inputs.unsqueeze(0), hidden)
        # generate end binary
        end_logit = self.lin_end(out.squeeze(0))
        end_output = F.gumbel_softmax(end_logit, dim=1, tau=3, hard=True)

        node_outputs = torch.stack(node_outputs, dim=1)
        tau_outputs = torch.stack(tau_outputs, dim=1)
        coord_outputs = torch.stack(coord_outputs, dim=1)
        return x_output, t0_res_output, node_outputs, tau_outputs, end_output, coord_outputs

class Discriminator(nn.Module):
    def __init__(self, n_samples, n_nodes, args):
        super(Discriminator, self).__init__()
        self.N = n_nodes
        self.W_down_discriminator_size = args.W_down_discriminator_size
        self.hidden_unit = args.hidden_unit
        self.rw_len = args.rw_len
        self.batch_size = n_samples

        self.W_down_x_discriminator_1 = nn.Linear(2, 1, bias=False)
        self.W_down_x_discriminator_2 = nn.Linear(1, self.W_down_discriminator_size)

        self.W_down_res_discriminator = nn.Linear(1, self.W_down_discriminator_size)

        self.W_down_discriminator = nn.Linear(self.N, self.W_down_discriminator_size, bias=False)

        self.W_down_tau_discriminator = nn.Linear(1, self.W_down_discriminator_size)

        self.W_down_end_discriminator_1 = nn.Linear(2, 1, bias=False)
        self.W_down_end_discriminator_2 = nn.Linear(1, self.W_down_discriminator_size)

        self.W_down_coord_discriminator = nn.Linear(2, self.W_down_discriminator_size, bias=False)

        self.lstm_all = nn.LSTM(self.W_down_discriminator_size, self.W_down_discriminator_size)
        self.lstm_node = nn.LSTM(self.W_down_discriminator_size, self.W_down_discriminator_size)
        self.lstm_time = nn.LSTM(self.W_down_discriminator_size, self.W_down_discriminator_size)
        self.lstm_loc = nn.LSTM(self.W_down_discriminator_size, self.W_down_discriminator_size)
        self.lstm_node_time = nn.LSTM(self.W_down_discriminator_size, self.W_down_discriminator_size)
        self.lstm_time_loc = nn.LSTM(self.W_down_discriminator_size, self.W_down_discriminator_size)
        self.lstm_node_loc = nn.LSTM(self.W_down_discriminator_size, self.W_down_discriminator_size)

        self.linear_out_all = nn.Sequential(
                            nn.Linear(self.W_down_discriminator_size, self.hidden_unit),
                            nn.Tanh(),
                            nn.Linear(self.hidden_unit, 1)
                            )
        self.linear_out_time = nn.Sequential(
                            nn.Linear(self.W_down_discriminator_size, self.hidden_unit),
                            nn.Tanh(),
                            nn.Linear(self.hidden_unit, 1)
                            )
        self.linear_out_node = nn.Sequential(
                            nn.Linear(self.W_down_discriminator_size, self.hidden_unit),
                            nn.Tanh(),
                            nn.Linear(self.hidden_unit, 1)
                            )
        self.linear_out_coord = nn.Sequential(
                            nn.Linear(self.W_down_discriminator_size, self.hidden_unit),
                            nn.Tanh(),
                            nn.Linear(self.hidden_unit, 1)
                            )
        self.linear_out_node_time = nn.Sequential(
                            nn.Linear(self.W_down_discriminator_size, self.hidden_unit),
                            nn.Tanh(),
                            nn.Linear(self.hidden_unit, 1)
                            )
        self.linear_out_node_coord = nn.Sequential(
                            nn.Linear(self.W_down_discriminator_size, self.hidden_unit),
                            nn.Tanh(),
                            nn.Linear(self.hidden_unit, 1)
                            )
        self.linear_out_time_coord = nn.Sequential(
                            nn.Linear(self.W_down_discriminator_size, self.hidden_unit),
                            nn.Tanh(),
                            nn.Linear(self.hidden_unit, 1)
                            )

    def forward(self, z):
        x, t0_res, node_inputs, tau_inputs, end, coords = z[0], z[1], z[2], z[3], z[4], z[5]
        # Reshape x
        x = torch.tanh(self.W_down_x_discriminator_2(self.W_down_x_discriminator_1(x.view(-1, 2).float())))
        # Reshape t0_res
        t0_res = torch.tanh(self.W_down_res_discriminator(t0_res.view(-1, 1).float()))
        # Reshape node_inputs
        node_inputs = self.W_down_discriminator(node_inputs.view(-1, self.N).float())
        node_inputs = node_inputs.view(-1, self.rw_len*2, self.W_down_discriminator_size)
        # Reshape tau_inputs
        tau_inputs = torch.tanh(self.W_down_tau_discriminator(tau_inputs.reshape(-1, 1).float()))
        tau_inputs = tau_inputs.view(-1, self.rw_len, self.W_down_discriminator_size)
        # Reshape end
        end = self.W_down_end_discriminator_2(self.W_down_end_discriminator_1(end.view(-1, 2).float()))
        # Reshape coordinates
        coords = self.W_down_coord_discriminator(coords.view(-1, 2).float())
        coords = coords.view(-1, self.rw_len*2, self.W_down_discriminator_size)

        inputs_all = [x] + [t0_res]
        inputs_time = [x] + [t0_res]
        inputs_node = []
        inputs_coord = []
        inputs_node_time = [x] + [t0_res]
        inputs_node_coord = []
        inputs_time_coord = [x] + [t0_res]

        for i in range(self.rw_len):
            inputs_all += [tau_inputs[:, i, :]]+[node_inputs[:, i*2, :]]+[coords[:, i*2, :]]+[node_inputs[:, i*2+1, :]]+[coords[:, i*2+1, :]]
            inputs_time += [tau_inputs[:, i, :]]
            inputs_node += [node_inputs[:, i*2, :]]+[node_inputs[:, i*2+1, :]]
            inputs_coord += [coords[:, i*2, :]]+[coords[:, i*2+1, :]]
            inputs_node_time += [tau_inputs[:, i, :]]+[node_inputs[:, i*2, :]]+[node_inputs[:, i*2+1, :]]
            inputs_node_coord += [node_inputs[:, i*2, :]]+[coords[:, i*2, :]]+[node_inputs[:, i*2+1, :]]+[coords[:, i*2+1, :]]
            inputs_time_coord += [tau_inputs[:, i, :]]+[coords[:, i*2, :]]+[coords[:, i*2+1, :]]

        inputs_all += [end]
        inputs_time += [end]
        inputs_node_time += [end]
        inputs_time_coord += [end]

        inputs_all = torch.stack(inputs_all)
        inputs_node = torch.stack(inputs_node)
        inputs_time = torch.stack(inputs_time)
        inputs_coord = torch.stack(inputs_coord)
        inputs_node_time = torch.stack(inputs_node_time)
        inputs_node_coord = torch.stack(inputs_node_coord)
        inputs_time_coord = torch.stack(inputs_time_coord)

        # LSTM outputs
        out_all, _ = self.lstm_all(inputs_all)
        out_node, _ = self.lstm_node(inputs_node)
        out_time, _ = self.lstm_time(inputs_time)
        out_coord, _ = self.lstm_loc(inputs_coord)
        out_node_time, _ = self.lstm_node_time(inputs_node_time)
        out_time_coord, _ = self.lstm_time_loc(inputs_time_coord)
        out_node_coord, _ = self.lstm_node_loc(inputs_node_coord)

        return (self.linear_out_all(out_all[-1]),self.linear_out_time(out_time[-1]),self.linear_out_node(out_node[-1]),
                self.linear_out_coord(out_coord[-1]),self.linear_out_node_time(out_node_time[-1]),
                self.linear_out_node_coord(out_node_coord[-1]),self.linear_out_time_coord(out_time_coord[-1]))

def generate_discrete(generator, n_samples, t_end, args):
    """
        Generate a random walk in index representation (instead of one hot). This is faster but prevents the gradients
        from flowing into the generator, so we only use it for evaluation purposes.
    """
    start_x_0 = Variable(F.one_hot(torch.zeros(n_samples,).to(torch.int64), 2).float()).to(device)
    start_x_1 = Variable(F.one_hot(torch.ones(n_samples,).to(torch.int64), 2).float()).to(device)
    start_t0 = Variable(torch.ones(n_samples, 1).float()).to(device)

    initial_noise = make_noise((n_samples, args.noise_dim), args.noise_type).to(device)

    fake_x, fake_t0, fake_e, fake_tau, fake_end, fake_coordinate = [], [], [], [], [], []
    with torch.no_grad():
        for i in range(args.n_eval_loop):
            if i == 0:
                fake_x_output, fake_t0_res_output, fake_node_output, fake_tau_output, fake_end_output, fake_coord_output = generator(
                        initial_noise, x_input=start_x_1, t0_input=start_t0)
            else:
                if args.rw_len == 1:
                    t0_input = fake_tau_output[:, -1, :]
                    fake_x_output, fake_t0_res_output, fake_node_output, fake_tau_output, fake_end_output, fake_coord_output = generator(
                        initial_noise, x_input=start_x_0, t0_input=t0_input)
                else:
                    t0_input = fake_tau_output[:, 0, :]
                    edge_input = fake_node_output[:, 2:, :]
                    coord_input = fake_coord_output[:, 2:, :]
                    tau_input = fake_tau_output[:, 1:, :]
                    fake_x_output, fake_t0_res_output, fake_node_output, fake_tau_output, fake_end_output, fake_coordinate_output = generator(
                            initial_noise, x_input=start_x_0, t0_input=t0_input, edge_input=edge_input, tau_input=tau_input)

            fake_x_outputs_discrete = torch.argmax(fake_x_output, dim=-1)
            fake_node_outputs_discrete = torch.argmax(fake_node_output, dim=-1)
            fake_end_discretes = torch.argmax(fake_end_output, dim=-1)

            fake_x.append(fake_x_outputs_discrete.detach().cpu().numpy())
            fake_t0.append(fake_t0_res_output.detach().cpu().numpy())
            fake_e.append(fake_node_outputs_discrete.detach().cpu().numpy())
            fake_tau.append(fake_tau_output.detach().cpu().numpy())
            fake_end.append(fake_end_discretes.detach().cpu().numpy())
            fake_coordinate.append(fake_coord_output.detach().cpu().numpy())

    return fake_x, fake_t0, fake_e, fake_tau, fake_end, fake_coordinate

def weights_calculation(fake_scores, real_scores, eps=1.75):
    with torch.no_grad():
        fake_scores = torch.Tensor([sum(i) for i in fake_scores])
        real_scores = torch.Tensor([sum(i) for i in real_scores])

        abs_value = torch.abs(fake_scores-real_scores)
        alpha = eps*torch.softmax(abs_value, dim=-1)
        return alpha

def train(dataloader, generator, discriminator, optimizer_G, optimizer_D, args, n_nodes, t_end,
          edge_contact_time, train_edges, test_edges, save_directory, output_directory, timing_directory,
          eval_every=4, max_iter=1000, eval_transitions=1e3, max_patience=5, is_test=True):
    n_eval_iters = int(eval_transitions / args.batch_size)
    # Getting a single data iterator
    iterloader = iter(dataloader)
    batches_done = 0

    # Read Polygons and the corresponding zones
    if args.dataset != "authorship":
        polygon, zone = read_geojson(args.dataset)
    else:
        polygons, zones = read_geojson(args.dataset)

    # Initialize the time
    starting_time = time.time()
    timestr = time.strftime("%Y%m%d-%H%M%S")

    log("Start Training...")
    for epoch in range(max_iter):
        for i, real_walk in enumerate(dataloader):
            # Getting real edge inputs
            real_edge_inputs_discrete = real_walk[:, 1:, 0:2].to(torch.long)
            real_node_inputs_discrete = real_edge_inputs_discrete.view(args.batch_size, -1)
            real_node_inputs = F.one_hot(real_node_inputs_discrete, n_nodes).to(device)
            real_tau_inputs = real_walk[:, 1:, 2:3].to(device)
            real_coords_inputs = real_walk[:, 1:, 3:]
            real_coords_inputs = real_coords_inputs.reshape(args.batch_size, 2*args.rw_len, -1).to(device)
            # Getting real indicator inputs
            real_x_input_discretes = real_walk[:, 0, 0].to(torch.long)
            real_x_inputs = F.one_hot(real_x_input_discretes, 2).to(device)
            real_end_discretes = real_walk[:, 0, 1].to(torch.long)
            real_ends = F.one_hot(real_end_discretes, 2).to(device)
            real_t0_res_inputs = real_walk[:, 0:1, 2].to(device)

            initial_noise = make_noise((args.batch_size, args.noise_dim), args.noise_type).to(device)
            # ---------------------
            #  Train Discriminator
            # ---------------------
            discriminator.train()
            optimizer_D.zero_grad()

            fake_output = generator(initial_noise)

            fake_x_inputs = fake_output[0].detach()
            fake_t0_res_inputs = fake_output[1].detach()
            fake_node_inputs = fake_output[2].detach()
            fake_tau_inputs = fake_output[3].detach()
            fake_ends = fake_output[4].detach()
            fake_coordinates = fake_output[5].detach()

            fake_scores = discriminator((fake_x_inputs, fake_t0_res_inputs, fake_node_inputs, fake_tau_inputs, fake_ends, fake_coordinates))
            real_scores = discriminator((real_x_inputs, real_t0_res_inputs, real_node_inputs, real_tau_inputs, real_ends, real_coords_inputs))

            alpha = weights_calculation(fake_scores, real_scores)
            fake_scores = [x*alpha[i] for i, x in enumerate(fake_scores)]
            real_scores = [x*alpha[i] for i, x in enumerate(real_scores)]

            loss_D = torch.mean(sum(fake_scores)) - torch.mean(sum(real_scores))

            loss_D.backward()
            optimizer_D.step()

            # Clip weights of discriminator
            for p in discriminator.parameters():
                p.data.clamp_(-args.clip_value, args.clip_value)

            # Train the generator every n_critic iterations
            if i % args.n_critic == 0:
                # -----------------
                #  Train Generator
                # -----------------
                optimizer_G.zero_grad()

                # Generate a batch of random temporal walks
                syn_x_inputs, syn_t0_res_inputs, syn_node_inputs, syn_tau_inputs, syn_ends, syn_coordinates = generator(initial_noise)
                # Adversarial loss
                syn_scores = discriminator((syn_x_inputs, syn_t0_res_inputs, syn_node_inputs, syn_tau_inputs, syn_ends, syn_coordinates))
                syn_scores = [x*alpha[i] for i, x in enumerate(syn_scores)]
                loss_G = -torch.mean(sum(syn_scores))

                # Loss back-propagation
                loss_G.backward()
                optimizer_G.step()

                log("[Epoch %d/%d] [Batch %d/%d] [D loss: %f] [G loss: %f]"
                    % (epoch+1, max_iter, batches_done % len(dataloader), len(dataloader), loss_D.item(), loss_G.item()))

            batches_done += 1

        # Evaluate the model's progress.
        if (epoch+1) % eval_every == 0:
            log('**** Starting Evaluating: {}-th Epoch ****'.format(epoch+1))

            # Creating instances for fake and real graphs
            fake_graphs, fake_x_t0, real_walks, real_x_t0 = [], [], [], []

            for q in range(n_eval_iters):
                fake_output = generate_discrete(generator, args.batch_size, t_end, args)
                fake_x = fake_output[0]
                fake_t0 = fake_output[1]
                fake_edges = fake_output[2]
                fake_t = fake_output[3]
                fake_end = fake_output[4]
                fake_coord = fake_output[5]

                smpls = None
                stop = [False] * args.batch_size
                for i in range(args.n_eval_loop):
                    x, t0, e, tau, le, coo = \
                    fake_x[i], fake_t0[i], fake_edges[i], fake_t[i], fake_end[i], fake_coord[i]

                    if q == 0 and i >= args.n_eval_loop-3:
                        log('eval_iters: {} eval_loop: {}'.format(q, i))
                        log('generated [x, t0, e, tau, end, coord]\n[{}, {}, {}, {}, {}]'.format(
                            x[0], t0[0, 0], e[0, :], tau[0, :, 0], le[0],
                            ))
                        log('generated [x, t0, e, tau, end, coord]\n[{}, {}, {}, {}, {}]'.format(
                            x[1], t0[1, 0], e[1, :], tau[1, :, 0], le[1],
                            ))
                    e = e.reshape(-1, args.rw_len, 2)
                    tau = tau.reshape(-1, args.rw_len, 1)
                    coo = coo.reshape(-1, args.rw_len, 4)
                    if i == 0:
                        smpls = np.concatenate([e, tau, coo],axis=-1)
                    else:
                        new_pred = np.concatenate([e[:, -1:], tau[:, -1:], coo[:, -1:]], axis=-1)
                        smpls = np.concatenate([smpls, new_pred], axis=1)

                    # judge if reach max length
                    for b in range(args.batch_size):
                        b_le = le[b]

                        if i == 0 and b_le == 1:  # end
                            stop[b] = True
                        if i > 0 and stop[b]:  # end
                            smpls[b, -1, :] = -1
                        if i > 0 and not stop[b] and b_le == 1:
                            stop[b] = True

                fake_x = np.array(fake_x).reshape(-1, 1)
                fake_t0 = np.array(fake_t0).reshape(-1, 1)
                fake_len = np.array(fake_end).reshape(-1, 1)
                fake_start = np.c_[fake_x, fake_t0, fake_len]
                fake_x_t0.append(fake_start)
                fake_graphs.append(smpls)

                if q%5 == 0:
                    print("Done, {} of {} eval iterations".format(q, n_eval_iters))

            fake_graphs = np.array(fake_graphs)
            fake_coords = fake_graphs[:, :, :, 3:]
            fake_coords = fake_coords[fake_coords[:, :, :, 0] != -1].reshape(-1, 2)
            constrained_coords = fake_coords.reshape(-1, 2)
            if args.dataset != "authorship":
                constrained_coords = spatial_constraint(constrained_coords, polygon, zone)
            else:
                constrained_coords = spatial_constraint(constrained_coords, polygons=polygons, zones=zones)

            fake_graphs = fake_graphs[:, :, :, :3]

            for _ in range(n_eval_iters):
                try:
                    true_walk = next(iterloader)
                except StopIteration:
                    iterloader = iter(dataloader)
                    true_walk = next(iterloader)

                # Getting real edge inputs
                real_edge = true_walk[:, 1:, 0:2].to(torch.long).numpy()
                real_x = true_walk[:, 0, 0].to(torch.long).numpy()
                real_length = true_walk[:, 0, 1].to(torch.long).numpy()
                real_tau = true_walk[:, 1:, 2:3].numpy()
                real_t0 = true_walk[:, 0:1, 2].numpy()

                walk = np.c_[real_edge.reshape(-1, 2), real_tau.reshape(-1, 1)]
                real_walks.append(walk)
                real_start = np.stack([real_x, real_t0[:, 0], real_length], axis=1)
                real_x_t0.append(real_start)

            if is_test:
                try:
                    fake_walks = fake_graphs.reshape(-1, 3)
                    fake_mask = fake_walks[:, 0] > -1
                    fake_walks = fake_walks[fake_mask]
                    fake_x_t0 = np.array(fake_x_t0).reshape(-1, 3)

                    real_walks = np.array(real_walks).reshape(-1, 3)
                    real_mask = real_walks[:, 0] > -1
                    real_walks = real_walks[real_mask]
                    real_x_t0 = np.array(real_x_t0).reshape(-1, 3)

                    truth_train_time = train_edges[:, 3:4]
                    truth_train_res_time = t_end - truth_train_time
                    truth_train_walks = np.concatenate([train_edges[:, 1:3], truth_train_res_time], axis=1)
                    truth_train_x_t0 = np.c_[np.zeros((len(train_edges), 1)), truth_train_res_time]
                    truth_train_x_t0 = np.r_[truth_train_x_t0, np.ones((len(train_edges), 2))]

                    truth_test_time = test_edges[:, 3:4]
                    truth_test_res_time = t_end - truth_test_time
                    truth_test_walks = np.c_[test_edges[:, 1:3], truth_test_res_time]
                    truth_test_x_t0 = np.c_[np.zeros((len(test_edges), 1)), truth_test_res_time]
                    truth_test_x_t0 = np.r_[truth_test_x_t0, np.ones((len(test_edges), 2))]

                    fake_e_list, fake_e_counts = np.unique(fake_walks[:, 0:2], return_counts=True, axis=0)
                    real_e_list, real_e_counts = np.unique(real_walks[:, 0:2], return_counts=True, axis=0)

                    truth_train_e_list, truth_train_e_counts = np.unique(truth_train_walks[:, 0:2],
                                                                         return_counts=True, axis=0)
                    truth_test_e_list, truth_test_e_counts = np.unique(truth_test_walks[:, 0:2],
                                                                       return_counts=True, axis=0)
                    truth_e_list, truth_e_counts = np.unique(
                        np.r_[truth_test_walks[:, 0:2], truth_test_walks[:, 0:2]], return_counts=True, axis=0)
                    n_e = len(truth_e_list)

                    real_x_list, real_x_counts = np.unique(real_x_t0[:, 0], return_counts=True)
                    fake_x_list, fake_x_counts = np.unique(fake_x_t0[:, 0], return_counts=True)
                    truth_x_list, truth_x_counts = real_x_list, real_x_counts

                    real_len_list, real_len_counts = np.unique(real_x_t0[:, 2], return_counts=True)
                    fake_len_list, fake_len_counts = np.unique(fake_x_t0[:, 2], return_counts=True)
                    truth_len_list, truth_len_counts = real_len_list, real_len_counts

                    fig = plt.figure(figsize=(2 * 9, 2 * 9))
                    fig.suptitle('Truth, Real, and Fake edges comparisons')
                    dx = 0.3
                    dy = dx
                    zpos = 0

                    fake_ax = fig.add_subplot(221, projection='3d')
                    fake_ax.bar3d(fake_e_list[:, 0], fake_e_list[:, 1], zpos, dx, dy, fake_e_counts)
                    fake_ax.set_xlim([0, n_nodes])
                    fake_ax.set_ylim([0, n_nodes])
                    fake_ax.set_xticks(range(n_nodes))
                    fake_ax.set_yticks(range(n_nodes))
                    fake_ax.set_xticklabels([str(n) if n % 5 == 0 else '' for n in range(n_nodes)])
                    fake_ax.set_yticklabels([str(n) if n % 5 == 0 else '' for n in range(n_nodes)])
                    fake_ax.set_title('fake edges number: {}'.format(len(fake_e_list)))

                    real_ax = fig.add_subplot(222, projection='3d')
                    real_ax.bar3d(real_e_list[:, 0], real_e_list[:, 1], zpos, dx, dy, real_e_counts)
                    real_ax.set_xlim([0, n_nodes])
                    real_ax.set_ylim([0, n_nodes])
                    real_ax.set_xticks(range(n_nodes))
                    real_ax.set_yticks(range(n_nodes))
                    real_ax.set_xticklabels([str(n) if n % 5 == 0 else '' for n in range(n_nodes)])
                    real_ax.set_yticklabels([str(n) if n % 5 == 0 else '' for n in range(n_nodes)])
                    real_ax.set_title('real edges number: {}'.format(len(real_e_list)))

                    truth_ax = fig.add_subplot(223, projection='3d')
                    truth_ax.bar3d(truth_train_e_list[:, 0], truth_train_e_list[:, 1], zpos, dx, dy,
                                   truth_train_e_counts)
                    truth_ax.set_xlim([0, n_nodes])
                    truth_ax.set_ylim([0, n_nodes])
                    truth_ax.set_xticks(range(n_nodes))
                    truth_ax.set_yticks(range(n_nodes))
                    truth_ax.set_xticklabels([str(n) if n % 5 == 0 else '' for n in range(n_nodes)])
                    truth_ax.set_yticklabels([str(n) if n % 5 == 0 else '' for n in range(n_nodes)])
                    truth_ax.set_title('truth train edges number: {}'.format(len(truth_train_e_list)))

                    truth_ax = fig.add_subplot(222, projection='3d')
                    truth_ax.bar3d(truth_test_e_list[:, 0], truth_test_e_list[:, 1], zpos, dx, dy, truth_test_e_counts)
                    truth_ax.set_xlim([0, n_nodes])
                    truth_ax.set_ylim([0, n_nodes])
                    truth_ax.set_xticks(range(n_nodes))
                    truth_ax.set_yticks(range(n_nodes))
                    truth_ax.set_xticklabels([str(n) if n % 5 == 0 else '' for n in range(n_nodes)])
                    truth_ax.set_yticklabels([str(n) if n % 5 == 0 else '' for n in range(n_nodes)])
                    truth_ax.set_title('truth test edges number: {}'.format(len(truth_test_e_list)))

                    plt.tight_layout()
                    #plt.show()
                    plt.savefig('{}/iter_{}_edges_counts_validation.png'.format(output_directory, epoch+1), dpi=90)
                    plt.close()

                except ValueError as e:
                    print(e)
                    log('reshape fake walks got error. Fake graphs shape: {} \n{}'.format(fake_graphs[0].shape, fake_walks[:3]))
                    continue

            fake_graphs = convert_graphs(fake_graphs)
            fake_graph_file = "{}/{}_assembled_graph_iter_{}.npz".format(output_directory, timestr, epoch+1)
            fake_graphs[:, 3] = t_end - fake_graphs[:, 3]
            np.savez_compressed(fake_graph_file, fake_graphs=fake_graphs, real_walks=real_walks)

            real_coords = test_edges[:, 4:].reshape(-1, 2)
            real_coords = np.sort(real_coords, axis=0)
            fake_coords = np.sort(constrained_coords, axis=0)

            Gs = Graphs(test_edges[:, :4], N=n_nodes, tmax=t_end, edge_contact_time=edge_contact_time)
            FGs = Graphs(fake_graphs, N=n_nodes, tmax=t_end, edge_contact_time=edge_contact_time)
            mmd_avg_degree = MMD_Average_Degree_Distribution(Gs, FGs)
            mmd_avg_group_size = MMD_Average_Group_Size(Gs, FGs)
            mmd_avg_group_number = MMD_Mean_Group_Number(Gs, FGs)
            mmd_avg_coordination_number = MMD_Mean_Coordination_Number(Gs, FGs)
            mmd_mean_group_number = MMD_Mean_Group_Number(Gs, FGs)
            mmd_group_duration = MMD_Mean_Group_Duration(Gs, FGs)
            mmd_avg_coordinates = MMD(real_coords[:2000], fake_coords[:2000])
            log('mmd_avg_degree: {}'.format(mmd_avg_degree))
            log('mmd_avg_group_size: {}'.format(mmd_avg_group_size))
            log('mmd_avg_group_number: {}'.format(mmd_avg_group_number))
            log('mmd_avg_coordination_number: {}'.format(mmd_avg_coordination_number))
            log('mmd_mean_group_number: {}'.format(mmd_mean_group_number))
            log('mmd_group_duration: {}'.format(mmd_group_duration))
            log('mmd_avg_coordinates: {}'.format(mmd_avg_coordinates))

            log('Real Mean_Average_Degree_Distribution: \n{}'.format(Gs.Mean_Average_Degree_Distribution()))
            log('Fake Mean_Average_Degree_Distribution: \n{}'.format(FGs.Mean_Average_Degree_Distribution()))

            if args.early_stopping is not None:
                if mmd_avg_degree < args.early_stopping:
                    log('**** end training because evaluation is reached ****')
                    break

            t = time.time() - starting_time
            log('**** end evaluation **** took {} seconds so far...'.format(int(t)))
        log('The weights on each discriminator {} at {}-th Epoch'.format(alpha, epoch+1))

def main(args):
    save_directory = "snapshots-{}".format(args.dataset)
    output_directory='outputs-{}'.format(args.dataset)
    timing_directory='timings-{}'.format(args.dataset)

    if not os.path.isdir(save_directory):
        os.makedirs(save_directory)

    if not os.path.isdir(output_directory):
        os.makedirs(output_directory)

    if not os.path.isdir(timing_directory):
        os.makedirs(timing_directory)

    log('use {} data'.format(args.dataset))
    log('-'*40)

    if args.dataset == "taxi":
        scale = 0.1
        train_ratio = 0.85
        t_end = 1.
        n_nodes = 66
        t0 = 0.1
        edge_contact_time = t0*0.01
    elif args.dataset == "checkin":
        scale = 0.1
        train_ratio = 0.85
        t_end = 1.
        n_nodes = 70
        t0 = 0.1
        edge_contact_time = t0*0.01
    elif args.dataset == "authorship":
        scale = 0.1
        train_ratio = 0.85
        t_end = 1.
        n_nodes = 628
        t0 = 0.1
        edge_contact_time = t0*0.1
    elif args.dataset == "sim_100":
        scale = 0.1
        train_ratio = 0.85
        t_end = 1.
        t0 = 0.1
        n_nodes = 99
        edge_contact_time = t0*0.5
    elif args.dataset == "sim_500":
        scale = 0.1
        train_ratio = 0.85
        t_end = 1.
        t0 = 0.1
        n_nodes = 495
        edge_contact_time = t0*0.5
    else:
        scale = 0.1
        train_ratio = 0.85
        t_end = 1.
        t0 = 0.1
        n_nodes = 2420
        edge_contact_time = t0*0.5

    edges = np.loadtxt('data/{}_dataset/{}_data.txt'.format(args.dataset, args.dataset))
    train_edges, test_edges = Split_Train_Test(edges, train_ratio)

    real_walks = TemporalDataset(train_edges, n_nodes, t_end, scale, args)
    dataloader = DataLoader(real_walks, batch_size=args.batch_size, shuffle=True, drop_last=True)

    generator = Generator(n_nodes, args.batch_size, t_end, args).to(device)
    discriminator = Discriminator(args.batch_size, n_nodes, args).to(device)

    optimizer_G = torch.optim.RMSprop(generator.parameters(), lr=args.learningrate)
    optimizer_D = torch.optim.RMSprop(discriminator.parameters(), lr=args.learningrate)

    train(dataloader, generator, discriminator, optimizer_G, optimizer_D, args, n_nodes, t_end,
          edge_contact_time, train_edges, test_edges, save_directory, output_directory, timing_directory)

if __name__ == "__main__":
    print("Start Training...")
    main(args)


#모듈
# import time
# print(time.localtime().tm_year,'년',time.localtime().tm_mon,'월',
#       time.localtime().tm_mday,'일',time.localtime().tm_hour,'시',
#       time.localtime().tm_min,'분',time.localtime().tm_sec,'초')

# from datetime import datetime
# now=datetime.now()
# print(now)
# print(now.strftime('%Y-%m-%d %H:%M:%S'))

#sleep 함수 : 프로그램 실행 속도를 제어
# import time
# print('시작')
# time.sleep(3)
# print('3초지남')

#타이머
# from time import sleep
# sec=int(input('몇초?'))
# print('타이머시작')
# # t=[1,2,3]
# for a in range(1,sec+1) :
#     sleep(1)
#     print(a,'초')
# print('타이머 종료')

#난수모듈
#주사위 게임
# import random
# awin=0 #a이 긴횟수
# bwin=0 #b가 이긴횟수
# while True :
#     a= random.randint(1,6)
#     b= random.randint(1,6)
#     print('1번 주사위 :',a,'2번 주사위 :',b)
#     if a>b :
#         print('1번 주사위',a,'가 이겼다.')
#
#     elif a<b :
#         print('2번 주사위',b, '가 이겼다.')
#
#     else :
#         print('무승부')

# import random
#sample()
# print(random.sample(range(1,46),6))

#choice()
# print(random.choice(['가위','바위','보']))

#shuffle() :섞는다
# data=['나비','나비','벌','벌','꽃','꽃']
# random.shuffle(data)
# print(data)

import turtle
turtle.shape('turtle')
turtle.color('red')
for x in range(100):
    turtle.forward(11)
    turtle.right(90)
    turtle.forward(11)
    turtle.left(90)
    turtle.forward(11)
    turtle.left(90)

turtle.done()


import sys

digits = ["ZERO", "ONE", "TWO", "THREE", "FOUR", "FIVE", "SIX", "SEVEN", "EIGHT", "NINE"]
keys   = {"Z" : 0, "W" : 2 , "U" : 4 , "X" : 6, "G" : 8 }

def update(s, digit):
	# and remove letters
	for dc in digits[digit]:
		for i in range(len(s)):
			if s[i] == dc:
				s[i] = ' '
				break

numcases = int(sys.stdin.readline())
for _c in range(numcases):
	count  = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
	s = [x for x in sys.stdin.readline()]
	s = s[:-1]
	# first pass: eliminate non-ambiguous ones
	for c in s:
		if c in keys:
			# count digit
			count[keys[c]] += 1
			update(s, keys[c])
	# now look for FIVE: F is key
	for c in s:
		if c == 'F':
			count[5] += 1
			update(s, 5)
	# NINE: I
	for c in s:
		if c == 'I':
			count[9] += 1
			update(s, 9)
	# SEVEN: S
	for c in s:
		if c == 'S':
			count[7] += 1
			update(s, 7)
	# ONE: O
	for c in s:
		if c == 'O':
			count[1] += 1
			update(s, 1)
	# THREE: T
	for c in s:
		if c == 'T':
			count[3] += 1
			update(s, 3)
	# now all should be replaced
	for c in s:
		assert(c == ' ')

	# output
	s = ""
	for i in range(10):
		for j in range(count[i]): s += str(i)
	print("Case #%d: %s" % (_c+1, s))


num_names = int(input('What is the number of names wanted on the list?'))
max_len = 0
names = []

for i in range(1, num_names+1):

    name = input(f"Please enter name number {i}. ").capitalize()
    while name.isalpha() == False:
        name = input(f"Silly you, names don't have integers. Please enter name number {i}. ").capitalize()
    names.append(name)

    if len(name) > max_len:

        max_len = len(name)

names.sort()

# Creates index for the characters in the name strings.
# max_len is the length of the largest name.

for i in range(max_len):

    line = ''

    # Creates index for which name string in names list we are referencing
    for x in range(num_names):

        # Checks if name is shorter than our current index to avoid indexing errors
        if len(names[x]) >= i+1:

            # Adds the character and 3 spaces to the line
            line += names[x][i] + ' '*3

        else:
            # Adds 4 spaces
            line += ' '*4

    print(line)


# -*- coding: utf-8 -*-
#
#% $Id$
#
#
# Copyright (C) 2002-2011
# The MeqTree Foundation &
# ASTRON (Netherlands Foundation for Research in Astronomy)
# P.O.Box 2, 7990 AA Dwingeloo, The Netherlands
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, see <http://www.gnu.org/licenses/>,
# or write to the Free Software Foundation, Inc.,
# 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
#

import sys

from Tigger.Models.Formats import load, save, listFormats
from Kittens.widgets import BusyIndicator
import Kittens.pixmaps
import Kittens.utils
import Kittens.config
import os.path

__version__ = "1.3.8"

release_string = __version__
svn_revision_string = __version__
svn_revision_html = __version__

matplotlib_nuked = False

startup_dprint = startup_dprintf = lambda *dum:None
_verbosity = Kittens.utils.verbosity(name="tigger")
dprint = _verbosity.dprint
dprintf = _verbosity.dprintf
pixmaps = Kittens.pixmaps.PixmapCache("Tigger")
ConfigFileName = ".tigger.conf"
ConfigFile = Kittens.config.DualConfigParser("tigger.conf",["/usr/lib/Tigger", os.path.dirname(__file__)])
Config = Kittens.config.SectionParser(ConfigFile,"Tigger")

def import_pyfits ():
  """Helper function to import pyfits and return it. Provides a workaround for
  pyfits-2.3, which is actually arrogant enough (fuck you with a bargepole, pyfits!)
  to replace the standard warnings.formatwarning function with its own BROKEN version,
  thus breaking all other code that uses the warnings module."""
  if 'pyfits' not in sys.modules:
    import pyfits
    import warnings
    if getattr(pyfits,'formatwarning',None) is warnings.formatwarning:
      def why_is_pyfits_overriding_warnings_formatwarning_with_a_broken_one_damn_you_pyfits (message,category,  filename,lineno,line=None):
        return str(message)+'\n'
      warnings.formatwarning = why_is_pyfits_overriding_warnings_formatwarning_with_a_broken_one_damn_you_pyfits
    if getattr(pyfits,'showwarning',None) is warnings.showwarning:
      def showwarning_damn_you_pyfits_damn_you_sincerely (message,category,filename,lineno,file=None,line=None):
        pyfits.showwarning(message,category,filename,lineno,file=file)
      warnings.showwarning = showwarning_damn_you_pyfits_damn_you_sincerely
  return pyfits

def nuke_matplotlib ():
  """Some people think nothing of importing matplotlib at every opportunity, with no regard
  to consequences. Tragically, some of these people also write Python code, and some of them
  are responsible for astLib. Seriously man, if I just want to pull in WCS support, why the fuck
  do I need the monstrous entirety of matplotlib to come along with it, especially since it
  kills things like Qt outright?
  This function prevents such perversitities from happening, by inserting dummy modules
  into the sys.modules dict. Call nuke_matplotlib() once, and all further attempts to
  import matplotlib by any other code will be cheerfully ignored.
  """
  if 'pylab' not in sys.modules:
    # replace the modules referenced by astLib by dummy_module objects, which return a dummy callable for every attribute
    class dummy_module (object):
      def __getattr__ (self,name):
        return 'nowhere' if name == '__file__' else (lambda *args,**kw:True)
    sys.modules['pylab'] = sys.modules['matplotlib'] = sys.modules['matplotlib.patches'] = dummy_module()
    matplotlib_nuked = True


# num = int(input('num: '))

# 三元运算符（三目运算符）：就是简写if...else...结构
# if num > 10:
#     print('num值大于10')
# else:
#     print('num值不大于10')

# 语法：结果1 if 条件 else 结果2
# print('num值大于10') if num > 10 else print('num值不大于10')

# 案例：获得两个数中的大值 | 小者
n1 = int(input('n1: '))
n2 = int(input('n2: '))
res = n1 if n1 > n2 else n2
print(res)
res = n2 if n1 > n2 else n1
print(res)