#!/usr/bin/env python
"""Example for the tool linear generator"""

import time
import sys
from math3d import Transform

from urlibs.system import URTrajectoryGenerator

__author__ = "Johannes Schrimpf"
__copyright__ = "Copyright 2012, NTNU"
__credits__ = ["Johannes Schrimpf"]
__license__ = "GPL"
__maintainer__ = "Johannes Schrimpf"
__email__ = "johannes.schrimpf(_at_)itk.ntnu.no"
__status__ = "Development"

def example_tlg(robot_ip):
    # Initialize the trajectory generator
    robot = URTrajectoryGenerator(robot_ip, log_q_des=True)
    robot.start()
    time.sleep(1)
    if not robot._is_running:
        print("Robot not initialized, exiting...")
        exit()
    # Test tool linear generator
    robot.start_tlg()
    # Save current pose
    start_pose = robot.get_tool_pose()
    # Create a target pose based on the starting pose
    # Care: target_pose = start_pose would not create
    # a new object, just a pointer to start_pose.
    target_pose = Transform(start_pose)
    target_pose.pos.x += 0.05
    # Wait until tlg is initialized
    time.sleep(0.01)
    # Set target pose
    robot.trajectory_generator.set_target_pose(target_pose)
    # Wait before returning
    time.sleep(1)

if __name__ == "__main__":
    if len(sys.argv) != 2:
        print("Please give robot ip as argument")
        exit()
    example_tlg(sys.argv[1])


from django import forms

class IFriendlyForm(forms.Form):
	CHOICES = (
		('U6', 'U6'),
		('U7', 'U7'),
    	('U8', 'U8'),
    	('U9', 'U9'),
    	('U10', 'U10'),
    	('U11', 'U11'),
    	('U12', 'U12'),
    	('U13', 'U13'),
        ('U14', 'U14(IP)'),
    	('U15', 'U15'),
    	('U17', 'U17'),
    )
	leeftijd = forms.ChoiceField(choices=CHOICES, required=True, label='Leeftijd')
	datum = forms.CharField(label='datum', widget=forms.TextInput())


from threading import Thread
from mutable import MutableStateHolder
from time import sleep

__author__ = 'ygintsyak'

class RunMe(Thread):

    def __init__(self, name):
        super(RunMe, self).__init__(name="Thread-" + str(name))

    def run(self):

        MutableStateHolder().reset()
        MutableStateHolder().state_2 = self.name
        while True:
            sleep(5)
            MutableStateHolder().state_1 += 5
            MutableStateHolder().print_me()

'''
  Runs forever
'''
RunMe("\033[95mCall-1\033[0m").start()
sleep(10)
RunMe("\033[94mCall-2\033[0m").start()
sleep(10)
RunMe("'\033[92mCall-3\033[0m").start()


import MAUS
import sys
import os
import numpy
import scipy.stats
import argparse

PARTICLE_ID_MAP = { -13 : 2, 13 : -2 }

if __name__ == "__main__" :
  parser = argparse.ArgumentParser("KDE Sampler")

  parser.add_argument( "data_file", help="An icool03 type beam file of particles to perform the resampling")

  parser.add_argument( "-N", "--number_particles", default=100, help="Number of particles to generate" )

  parser.add_argument( "-O", "--output_filename", default="resampled_beam", help="Output File Name" )

  parser.add_argument( "-D", "--output_directory", default="./", help="Output Directory" )

  parser.add_argument( "-P", "--particle_id", default=-13, type=int, help="Specify the PID of the beam to produce (Geant4 notation)" )

  namespace = parser.parse_args()

  vectors = None
  weights = None

  vectors = numpy.loadtxt(namespace.data_file, delimiter=' ', usecols=(6, 7, 8, 9, 10, 11), skiprows=2)
  weights = numpy.loadtxt(namespace.data_file, delimiter=' ', usecols=(5), skiprows=2)

#  kernel = scipy.stats.gaussian_kde(numpy.transpose(vectors), weights=weights)
  kernel = scipy.stats.gaussian_kde(numpy.transpose(vectors))

  outfilename = os.path.join(namespace.output_directory, namespace.output_filename)+".txt"
  with open( outfilename, "w" ) as outfile :

    outfile.write( "Custom Input file\n" )
    outfile.write( "0. 0. 0. 0. 0. 0. 0. 0.\n" )

    for count in range( namespace.number_particles ) :
      particle = kernel.resample(1)
#      print particle

      # 0  : Event Number
      value_1 = 1 # Particle Number
      value_2 = PARTICLE_ID_MAP[namespace.particle_id] # Particle ID
      value_3 = 0 # Status
      value_4 = 0.0 # Time
      value_5 = 0.0 # Local Weight
      value_6 = float(particle[0]) # x [m]
      value_7 = float(particle[1]) # y [m]
      value_8 = float(particle[2]) # z [m]
      value_9 = float(particle[3]) # px [GeV/c]
      value_10 = float(particle[4]) # py [GeV/c]
      value_11 = float(particle[5]) # pz [GeV/c]
      value_12 = 0 # sx
      value_13 = 0 # sy
      value_14 = 0 # sz

      outfile.write( str( count ) + ' ' +\
                     str( value_1 ) + ' ' + \
                     str( value_2 ) + ' ' + \
                     str( value_3 ) + ' ' + \
                     str( value_4 ) + ' ' +\
                     str( value_5 ) + ' ' +\
                     str( value_6 ) + ' ' +\
                     str( value_7 ) + ' ' +\
                     str( value_8 ) + ' ' +\
                     str( value_9 ) + ' ' +\
                     str( value_10 ) + ' ' +\
                     str( value_11 ) + ' ' +\
                     str( value_12 ) + ' ' +\
                     str( value_13 ) + ' ' +\
                     str( value_14 ) + '\n' )


"""Tests for letsencrypt_apache.parser."""
import os
import shutil
import sys
import unittest

import augeas
import mock
import zope.component

from letsencrypt import errors
from letsencrypt.display import util as display_util

from letsencrypt_apache.tests import util

class ApacheParserTest(util.ApacheTest):
    """Apache Parser Test."""

    def setUp(self):
        super(ApacheParserTest, self).setUp()

        zope.component.provideUtility(display_util.FileDisplay(sys.stdout))

        from letsencrypt_apache.parser import ApacheParser
        self.aug = augeas.Augeas(flags=augeas.Augeas.NONE)
        self.parser = ApacheParser(self.aug, self.config_path, self.ssl_options)

    def tearDown(self):
        shutil.rmtree(self.temp_dir)
        shutil.rmtree(self.config_dir)
        shutil.rmtree(self.work_dir)

    def test_root_normalized(self):
        from letsencrypt_apache.parser import ApacheParser
        path = os.path.join(self.temp_dir, "debian_apache_2_4/////"
                            "two_vhost_80/../two_vhost_80/apache2")
        parser = ApacheParser(self.aug, path, None)
        self.assertEqual(parser.root, self.config_path)

    def test_root_absolute(self):
        from letsencrypt_apache.parser import ApacheParser
        parser = ApacheParser(self.aug, os.path.relpath(self.config_path), None)
        self.assertEqual(parser.root, self.config_path)

    def test_root_no_trailing_slash(self):
        from letsencrypt_apache.parser import ApacheParser
        parser = ApacheParser(self.aug, self.config_path + os.path.sep, None)
        self.assertEqual(parser.root, self.config_path)

    def test_parse_file(self):
        """Test parse_file.

        letsencrypt.conf is chosen as the test file as it will not be
        included during the normal course of execution.

        """
        file_path = os.path.join(
            self.config_path, "sites-available", "letsencrypt.conf")

        self.parser._parse_file(file_path)  # pylint: disable=protected-access

        # search for the httpd incl
        matches = self.parser.aug.match(
            "/augeas/load/Httpd/incl [. ='%s']" % file_path)

        self.assertTrue(matches)

    def test_find_dir(self):
        from letsencrypt_apache.parser import case_i
        test = self.parser.find_dir(case_i("Listen"), "443")
        # This will only look in enabled hosts
        test2 = self.parser.find_dir(case_i("documentroot"))
        self.assertEqual(len(test), 2)
        self.assertEqual(len(test2), 3)

    def test_add_dir(self):
        aug_default = "/files" + self.parser.loc["default"]
        self.parser.add_dir(aug_default, "AddDirective", "test")

        self.assertTrue(
            self.parser.find_dir("AddDirective", "test", aug_default))

        self.parser.add_dir(aug_default, "AddList", ["1", "2", "3", "4"])
        matches = self.parser.find_dir("AddList", None, aug_default)
        for i, match in enumerate(matches):
            self.assertEqual(self.parser.aug.get(match), str(i + 1))

    def test_add_dir_to_ifmodssl(self):
        """test add_dir_to_ifmodssl.

        Path must be valid before attempting to add to augeas

        """
        from letsencrypt_apache.parser import get_aug_path
        self.parser.add_dir_to_ifmodssl(
            get_aug_path(self.parser.loc["default"]),
            "FakeDirective", "123")

        matches = self.parser.find_dir("FakeDirective", "123")

        self.assertEqual(len(matches), 1)
        self.assertTrue("IfModule" in matches[0])

    def test_get_aug_path(self):
        from letsencrypt_apache.parser import get_aug_path
        self.assertEqual("/files/etc/apache", get_aug_path("/etc/apache"))

    def test_set_locations(self):
        with mock.patch("letsencrypt_apache.parser.os.path") as mock_path:

            mock_path.isfile.return_value = False

            # pylint: disable=protected-access
            self.assertRaises(errors.PluginError,
                              self.parser._set_locations, self.ssl_options)

            mock_path.isfile.side_effect = [True, False, False]

            # pylint: disable=protected-access
            results = self.parser._set_locations(self.ssl_options)

            self.assertEqual(results["default"], results["listen"])
            self.assertEqual(results["default"], results["name"])

if __name__ == "__main__":
    unittest.main()  # pragma: no cover


from hangman.Dictionary import WordServer

class GuessWord(object):
    word = ""
    good_guesses = []
    bad_guesses = []
    display_string = ""
    won = False
    wordServer = WordServer()

    def __init__(self, max_guesses, difficultyLevel):
        self.word = self.getRandomWord(difficultyLevel)
        self.max_guesses = max_guesses
        self.display_string = "_"*len(self.word)
        self.getDefinition()

    def guess(self, charGuess):
        if charGuess not in self.good_guesses and charGuess not in self.bad_guesses:
            found = False
            for c in range(0, len(self.word)):
                if self.word[c] == charGuess:
                    self.good_guesses.append(charGuess)
                    self.buildDisplayString()
                    found = True
            if found == False:
                self.bad_guesses.append(charGuess)

    def guesses_left(self):
        guessCount = len(self.good_guesses) + len(self.bad_guesses)
        return self.max_guesses - guessCount

    def buildDisplayString(self):
        self.display_string = ""
        for cw in self.word:
            for cg in self.good_guesses:
                if cg == cw:
                    self.display_string += cg
                    break
            else:
                self.display_string += "_"

        self.isWinner(self.display_string)

    def isWinner(self, guessWord):
        if self.word == guessWord:
            self.won = True
            self.display_string = guessWord

    def getRandomWord(self, difficultyLevel):
        return self.wordServer.getRandomWord(difficultyLevel)

    def getFormattedWord(self):
        formattedWord = []
        for c in self.display_string:
            formattedWord.append(c)
            formattedWord.append(" ")
        else:
            del formattedWord[-1]

        return "".join(formattedWord)

    def getDefinition(self):
        defintions = self.wordServer.getDefinitions(self.word)


import keras
from keras.models import model_from_json, Sequential
from keras.layers import Conv2D, Dense, Flatten, Activation, BatchNormalization, GlobalAveragePooling2D
from keras.callbacks import EarlyStopping, ReduceLROnPlateau, ModelCheckpoint
from keras.utils import get_custom_objects, CustomObjectScope
from keras.preprocessing.image import ImageDataGenerator
from keras.datasets import mnist
from keras import backend as K

import os
import sys
import argparse
import json
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split

sys.path.append('..')
import functions
from callbacks import EpochLog
from my_json_encoder import MyJsonEncoder

parser = argparse.ArgumentParser()
parser.add_argument('activation', help='Activation function name to be approximated')
parser.add_argument('rg', help='Approximate range of x')
parser.add_argument('deg', help='Approximate degree of polynomial')
parser.add_argument('--bn', action='store_true')
parser.add_argument('--gap', action='store_true')
parser.add_argument('--full', action='store_true')
parser.add_argument('--da', action='store_true')
args = parser.parse_args()
activation_str = args.activation
rg_str = args.rg
deg_str = args.deg
bn_enabled = args.bn
gap_enabled = args.gap
full_enabled = args.full
da_enabled = args.da

ACTIVATION_FUNCTION_MAP = {
    'swish_rg3_deg4': functions.swish_rg3_deg4,
    'swish_rg4_deg4': functions.swish_rg4_deg4,
    'swish_rg5_deg4': functions.swish_rg5_deg4,
    'swish_rg6_deg4': functions.swish_rg6_deg4,
    'swish_rg7_deg4': functions.swish_rg7_deg4,
    'swish_rg8_deg4': functions.swish_rg8_deg4,
    'mish_rg4_deg4': functions.mish_rg4_deg4,
    'mish_rg5_deg4': functions.mish_rg5_deg4,
    'mish_rg6_deg4': functions.mish_rg6_deg4
}

APPROX_ACTIVATION_NAME = f'{activation_str}_rg{rg_str}_deg{deg_str}'
if APPROX_ACTIVATION_NAME not in ACTIVATION_FUNCTION_MAP:
    print(f'\'{APPROX_ACTIVATION_NAME}\' is not registered.')
    sys.exit(1)
APPROX_ACTIVATION = ACTIVATION_FUNCTION_MAP[APPROX_ACTIVATION_NAME]

# Register custom activation function
get_custom_objects().update({APPROX_ACTIVATION_NAME: APPROX_ACTIVATION})

# Declare constants
CUR_DIR = os.path.dirname(os.path.abspath(__file__))
BASE_FILE_NAME = f'{os.path.splitext(os.path.basename(__file__))[0]}-{APPROX_ACTIVATION_NAME}'
if bn_enabled:
    BASE_FILE_NAME+='-BN'
if gap_enabled:
    BASE_FILE_NAME+='-GAP'
if full_enabled:
    BASE_FILE_NAME+='-full'
if da_enabled:
    BASE_FILE_NAME+='-DA'
LOG_FILE_PATH = f'{CUR_DIR}/epoch_logs/{BASE_FILE_NAME}_log.txt'
# MODEL_FILE_PATH = f'{CUR_DIR}/saved_models/{BASE_FILE_NAME}_model.h5'
MODEL_WEIGHTS_PATH = f'{CUR_DIR}/saved_models/{BASE_FILE_NAME}_weights.h5'
MODEL_STRUCTURE_PATH = f'{CUR_DIR}/saved_models/{BASE_FILE_NAME}_structure.json'
HISTORY_FILE_PATH = f'{CUR_DIR}/histories/{BASE_FILE_NAME}_history.json'
HISTORY_GRAPH_PATH = f'{CUR_DIR}/graphs/{BASE_FILE_NAME}_history.png'

BATCH_SIZE = 128
NUM_CLASSES = 10
EPOCHS = 150
IMG_HIGHT, IMG_WIDTH = 28, 28

def preprocess_input_image(train_input, test_input):
    if K.image_data_format() == 'channels_first':
        train_input = train_input.reshape(
            train_input.shape[0], 1, IMG_HIGHT, IMG_WIDTH)
        test_input = test_input.reshape(
            test_input.shape[0], 1, IMG_HIGHT, IMG_WIDTH)
    else:
        train_input = train_input.reshape(
            train_input.shape[0], IMG_HIGHT, IMG_WIDTH, 1)
        test_input = test_input.reshape(
            test_input.shape[0], IMG_HIGHT, IMG_WIDTH, 1)

    # Normalization
    train_input = train_input.astype('float32')
    test_input = test_input.astype('float32')
    train_input /= 255.0
    test_input /= 255.0

    return train_input, test_input

def mnist_model(input_shape):
    model = Sequential()

    model.add(Conv2D(5, (5, 5), strides=(2, 2), input_shape=input_shape))
    if bn_enabled:
        model.add(BatchNormalization())
    model.add(Activation(APPROX_ACTIVATION_NAME))

    model.add(Conv2D(50, (5, 5), strides=(2, 2)))
    if bn_enabled:
        model.add(BatchNormalization())
    model.add(Activation(APPROX_ACTIVATION_NAME))

    if gap_enabled:
        model.add(GlobalAveragePooling2D())
    else:
        model.add(Flatten())
    model.add(Dense(NUM_CLASSES, activation='softmax'))

    return model

def main():
    if os.path.exists(LOG_FILE_PATH):
        os.remove(LOG_FILE_PATH)

    # Load mnist data
    (x_train, y_train), (x_test, y_test) = mnist.load_data()

    x_train, x_test = preprocess_input_image(x_train, x_test)
    input_shape = x_train[0].shape

    print('input shape:', input_shape)
    print(x_train.shape[0], 'train samples')
    print(x_test.shape[0], 'test samples')

    # Convert class vectors to binary class metrices (One Hot Encoding)
    y_train = keras.utils.to_categorical(y_train, NUM_CLASSES)
    y_test = keras.utils.to_categorical(y_test, NUM_CLASSES)

    # Split data for training and validation
    if da_enabled:
        X_train, X_val, Y_train, Y_val = train_test_split(x_train, y_train, test_size=0.1, random_state=42)

    # Define CNN model
    model = mnist_model(input_shape)
    # Compile the model
    model.compile(loss='categorical_crossentropy',
                  optimizer=keras.optimizers.Adadelta(),
                  metrics=['accuracy'])

    print(model.summary())
    with open(LOG_FILE_PATH, 'a') as f:
        model.summary(print_fn=lambda x: f.write(x + "\r\n"))
        f.write("\n\n")

    model_structure_json = model.to_json()
    with open(MODEL_STRUCTURE_PATH, 'w') as f:
        f.write(model_structure_json)

    # Define callbacks
    # early_stopping = EarlyStopping(monitor='val_loss', patience=10, verbose=1)
    learning_rate_reduction = ReduceLROnPlateau(monitor='val_loss',
                                                patience=3,
                                                verbose=1,
                                                factor=0.5,
                                                min_lr=0.0001)
    model_checkpoint = ModelCheckpoint(MODEL_WEIGHTS_PATH,
                                       monitor='val_loss',
                                       save_best_only=True,
                                       save_weights_only=True)
    epoch_log = EpochLog(LOG_FILE_PATH)

    callbacks = [
        # early_stopping,
        learning_rate_reduction,
        model_checkpoint,
        epoch_log
    ]

    if da_enabled:
        # Data augumentation
        datagen = ImageDataGenerator(
            rotation_range=15,
            zoom_range = 0.1,
            width_shift_range=0.1,
            height_shift_range=0.1,
            horizontal_flip=False,
            vertical_flip=False
        )

        steps_per_epoch = X_train.shape[0] // BATCH_SIZE

        # Train model
        history = model.fit_generator(datagen.flow(X_train, Y_train, batch_size=BATCH_SIZE),
                                      epochs=EPOCHS,
                                      steps_per_epoch=steps_per_epoch,
                                      validation_data=(X_val, Y_val),
                                      verbose=1,
                                      callbacks=callbacks)
    else:
        # Train model
        history = model.fit(x_train, y_train,
                            batch_size=BATCH_SIZE,
                            epochs=EPOCHS,
                            verbose=1,
                            validation_split=0.1,
                            callbacks=callbacks)

    # Save history
    with open(HISTORY_FILE_PATH, 'w') as f:
        json.dump(history.history, f, cls=MyJsonEncoder)

    # Visualize training history
    functions.plot_training_history(history.history, HISTORY_GRAPH_PATH)

    # Load best model
    with CustomObjectScope({APPROX_ACTIVATION_NAME: APPROX_ACTIVATION}):
        with open(MODEL_STRUCTURE_PATH, 'r') as f:
            model_structure_json = f.read()
        best_model = model_from_json(model_structure_json)
        best_model.load_weights(MODEL_WEIGHTS_PATH)
        best_model.compile(loss='categorical_crossentropy',
                           optimizer=keras.optimizers.Adadelta(),
                           metrics=['accuracy'])

    # Evaluate model
    test_loss, test_acc = best_model.evaluate(x_test, y_test, verbose=1)

    print('Test loss:', test_loss)
    print('Test accuracy:', test_acc)
    with open(LOG_FILE_PATH, 'a') as f:
        f.write("\n\nTest loss: {}, Test accuracy: {}\n".format(
            test_loss, test_acc))

    # Predict for test samples
    preds = best_model.predict(x_test, verbose=1)
    print("Prediction result:")
    print(preds)

if __name__ == '__main__':
    main()


# Studentnumber : 1716390
# Class : V2C

import random
import sys

class quickSort:

	"""
	Constructor of the quickSort class.

	Parameters
	----------
	unsortedList : list:
		This list contains the values you want to quicksort.
	"""
	def __init__(self, unsortedList):
		assert len(unsortedList) > 0
		self.counter = 0
		self.quickSortList = unsortedList

	"""
	Function to swap two values in a array using indices.

	Parameters
	----------
	initialList : list
		This list you want to swap the indices of.

	firstIndex : integer
		The index of the first element you want to swap.

	secondIndex : integer
		The index of the second element you want to swap.
	"""
	def swap(self, initialList, firstIndex, secondIndex):
		temp = initialList[secondIndex]
		initialList[secondIndex] = initialList[firstIndex]
		initialList[firstIndex] = temp

	"""
	Start the quicksort algorithm.
	"""
	def startSort(self):
		self.counter = 0
		self.quickSort(self.quickSortList)

	"""
	Method used to quicksort a list using recusion.

	Parameters
	----------
	quickSortList : List
		The list you want to quickSort.

	lowIndex: Integer
		The lowest index of the list.
		In the recusion process this will be the lowest index of the subList.

	highIndex: Integer
		The highest index of the list.
		In the recusion process this will be the highest index of the subList.
	"""
	def quickSort(self, quickSortList, lowIndex = 0, highIndex = -1):
		if highIndex == -1:
			highIndex = len(quickSortList) - 1

		if(lowIndex >= highIndex):
			return quickSortList

		# middle index
		# pivotIndex = int((lowIndex + highIndex) / 2)

		# random pivot index
		# pivotIndex = random.randrange(lowIndex, highIndex)

		# worst case pivotIndex
		pivotIndex = quickSortList.index(min(quickSortList[lowIndex:highIndex]))

		newPivotIndex = self.partition(quickSortList, lowIndex, highIndex, quickSortList[pivotIndex])
		self.quickSort(quickSortList, lowIndex, newPivotIndex -1)
		self.quickSort(quickSortList, newPivotIndex, highIndex)

	"""
	Method used to quicksort a list using recusion.
	This method swaps two indices to the right subList.
	The items on the left of the pivot index are lower than the pivot.
	The items on the right of the pivot index are higher than the pivot.

	Parameters
	----------
	quickSortList : List
		The list you want to quickSort.

	lowIndex: Integer
		The lowest index of the list.
		In the recusion process this will be the lowest index of the subList.

	highIndex: Integer
		The highest index of the list.
		In the recusion process this will be the highest index of the subList.

	pivotValue:
		The value of the selected pivot.
	"""
	def partition(self, quickSortList, lowIndex, highIndex, pivotValue):
		while(lowIndex <= highIndex):
			while(quickSortList[lowIndex] < pivotValue):
				lowIndex += 1
				self.counter += 1
			while(quickSortList[highIndex] > pivotValue):
				highIndex -= 1
				self.counter += 1

			if(lowIndex <= highIndex):
				self.swap(quickSortList, lowIndex, highIndex)
				lowIndex += 1
				highIndex -= 1

		# new index of the pivot
		return lowIndex

amountOfListElements = 10000
global minValue
global maxValue
global loopCounter
global recursionCounter
loopCounter = 0
minValue = None
maxValue = None
looping = True

sys.setrecursionlimit(10000)
for i in range(10):
	try:
		recursionCounter = 0
		loopCounter += 1
		random.seed()
		unsortedList = [random.randrange(1, amountOfListElements * 2, 1) for x in range(amountOfListElements)]
		qs = quickSort(unsortedList)

		qs.startSort()

		for i in range(len(qs.quickSortList) - 2):
			if not qs.quickSortList[i] <= qs.quickSortList[i]:
				print("not sorted")

		print(qs.counter)
		if maxValue is None or qs.counter > maxValue:
			maxValue = qs.counter
		elif minValue is None or qs.counter < minValue:
			minValue = qs.counter

		#print(recursionCounter)
	except KeyboardInterrupt:
		break

print("max: ", maxValue)
print("min: ", minValue)
print("amount of tests: ", loopCounter)

# middle of list pivot
#	max :  151746
#	min :  90690
#	loop:  6738


from collections import defaultdict
from enum import Enum
from os.path import join, realpath
import sys; sys.path.insert(0, realpath(join("/home/ubuntu/hummingbot/hummingbot/connector/derivative/dydx_perpetual/dydx_perpetual_derivative.py", "../../../")))

from hummingbot.connector.derivative.position import Position

from hummingbot.core.network_iterator import NetworkStatus
from hummingbot.core.utils.estimate_fee import estimate_fee

from async_timeout import timeout

from hummingbot.core.clock import Clock
from hummingbot.core.data_type.cancellation_result import CancellationResult
from hummingbot.connector.derivative.dydx_perpetual.dydx_perpetual_in_flight_order import DydxPerpetualsInFlightOrder

from hummingbot.core.utils.tracking_nonce import get_tracking_nonce

import asyncio
import hashlib
import hmac
import time
import logging
from decimal import Decimal
from typing import Optional, List, Dict, Any, AsyncIterable
from urllib.parse import urlencode

import aiohttp

from hummingbot.core.data_type.order_book import OrderBook
from hummingbot.core.event.events import (
    OrderType,
    TradeType,
    MarketOrderFailureEvent,
    MarketEvent,
    OrderCancelledEvent,
    BuyOrderCompletedEvent,
    BuyOrderCreatedEvent,
    SellOrderCreatedEvent,
    OrderFilledEvent,
    SellOrderCompletedEvent, PositionSide, PositionMode, PositionAction)
from hummingbot.core.utils.async_utils import safe_ensure_future, safe_gather
from hummingbot.core.utils.asyncio_throttle import Throttler
from hummingbot.logger import HummingbotLogger
from hummingbot.connector.derivative.dydx_perpetual.dydx_perpetual_order_book_tracker import DydxPerpetualOrderBookTracker
from hummingbot.connector.derivative.dydx_perpetual.dydx_perpetual_user_stream_tracker import DydxPerpetualUserStreamTracker
# from hummingbot.connector.derivative.dydx_perpetual.dydx_perpetual_utils import convert_from_exchange_trading_pair, convert_to_exchange_trading_pair
from hummingbot.connector.derivative.dydx_perpetual.constants import (
    PERPETUAL_BASE_URL,
    TESTNET_BASE_URL,
    DIFF_STREAM_URL,
    TESTNET_STREAM_URL
)
from hummingbot.connector.derivative_base import DerivativeBase, s_decimal_NaN
from hummingbot.connector.trading_rule import TradingRule

from dydx3 import Client
import configparser
from dydx3.constants import NETWORK_ID_MAINNET, NETWORK_ID_ROPSTEN
import json
from datetime import datetime
from hummingbot.core.event.events import TradeFee, TradeFeeType
from hummingbot.core.clock_mode import ClockMode

class MethodType(Enum):
    GET = "GET"
    POST = "POST"
    DELETE = "DELETE"
    PUT = "PUT"

bpm_logger = None

BROKER_ID = "x-3QreWesy"

def get_client_order_id(order_side: str, trading_pair: object):
    nonce = get_tracking_nonce()
    symbols: str = trading_pair.split("-")
    base: str = symbols[0].upper()
    quote: str = symbols[1].upper()
    return f"{BROKER_ID}-{order_side.upper()[0]}{base[0]}{base[-1]}{quote[0]}{quote[-1]}{nonce}"

class DydxPerpetualDerivative(DerivativeBase):
    MARKET_RECEIVED_ASSET_EVENT_TAG = MarketEvent.ReceivedAsset
    MARKET_BUY_ORDER_COMPLETED_EVENT_TAG = MarketEvent.BuyOrderCompleted
    MARKET_SELL_ORDER_COMPLETED_EVENT_TAG = MarketEvent.SellOrderCompleted
    MARKET_ORDER_CANCELLED_EVENT_TAG = MarketEvent.OrderCancelled
    MARKET_TRANSACTION_FAILURE_EVENT_TAG = MarketEvent.TransactionFailure
    MARKET_ORDER_FAILURE_EVENT_TAG = MarketEvent.OrderFailure
    MARKET_ORDER_FILLED_EVENT_TAG = MarketEvent.OrderFilled
    MARKET_BUY_ORDER_CREATED_EVENT_TAG = MarketEvent.BuyOrderCreated
    MARKET_SELL_ORDER_CREATED_EVENT_TAG = MarketEvent.SellOrderCreated

    API_CALL_TIMEOUT = 10.0
    SHORT_POLL_INTERVAL = 5.0
    UPDATE_ORDER_STATUS_MIN_INTERVAL = 10.0
    LONG_POLL_INTERVAL = 120.0
    ORDER_NOT_EXIST_CONFIRMATION_COUNT = 3

    @classmethod
    def logger(cls) -> HummingbotLogger:
        global bpm_logger
        if bpm_logger is None:
            bpm_logger = logging.getLogger(__name__)
        return bpm_logger

    def __init__(self,
                 dydx_perpetual_api_key: str = None,
                 dydx_perpetual_api_secret: str = None,
                 dydx_perpetual_api_passphrase: str = None,
                 dydx_perpetual_eth_wallet: str = None,
                 trading_pairs: Optional[List[str]] = None,
                 trading_required: bool = True,
                 **domain):
#         self._testnet = True if len(domain) > 0 else False
        self._testnet = True
        super().__init__()
        self._api_key = dydx_perpetual_api_key
        self._api_secret = dydx_perpetual_api_secret

        # dydx_perpetual authentication credentials are obtained from config file rather than input variables due to unconventional credentials different from centralized exchanges.
        self._eth_address = dydx_perpetual_eth_wallet
        self._api_key_credentials = {'secret': dydx_perpetual_api_secret, 'key': dydx_perpetual_api_key, 'passphrase': dydx_perpetual_api_passphrase}

        self._trading_required = trading_required
        self._account_balances = {}
        self._account_available_balances = {}

        self._base_url = PERPETUAL_BASE_URL if self._testnet is False else TESTNET_BASE_URL
        self._stream_url = DIFF_STREAM_URL if self._testnet is False else TESTNET_STREAM_URL
        self._network_id = NETWORK_ID_MAINNET if self._testnet is False else NETWORK_ID_ROPSTEN

        self._user_stream_tracker = DydxPerpetualUserStreamTracker(
            base_url=self._base_url,
            stream_url=self._stream_url,
            network_id=self._network_id,
            eth_address=self._eth_address,
            api_key_credentials=self._api_key_credentials
        )
        print(domain)
        domain = domain['domain']
        self._order_book_tracker = DydxPerpetualOrderBookTracker(trading_pairs=trading_pairs, domain=domain)
        self._ev_loop = asyncio.get_event_loop()
        self._poll_notifier = asyncio.Event()
        self._in_flight_orders = {}
        self._order_not_found_records = {}
        self._last_timestamp = 0
        self._trading_rules = {}
        # self._trade_fees = {} # To be implemented or removed
        # self._last_update_trade_fees_timestamp = 0 # To be implemented or removed
        self._status_polling_task = None
        self._user_stream_event_listener_task = None
        self._trading_rules_polling_task = None
        self._last_poll_timestamp = 0
        self._throttler = Throttler((10.0, 1.0))
        self._funding_rate = 0
        self._account_positions = {}
        self._position_mode = None
        self._leverage = 1

    @property
    def name(self) -> str:
        return "dydx_perpetual_testnet" if self._testnet else "dydx_perpetual"

    @property
    def order_books(self) -> Dict[str, OrderBook]:
        return self._order_book_tracker.order_books

    @property
    def ready(self):
        return all(self.status_dict.values())

    @property
    def in_flight_orders(self) -> Dict[str, DydxPerpetualsInFlightOrder]:
        return self._in_flight_orders

    @property
    def status_dict(self):
        return {
            "order_books_initialized": self._order_book_tracker.ready,
            "account_balance": len(self._account_balances) > 0 if self._trading_required else True,
            "trading_rule_initialized": len(self._trading_rules) > 0,

            # TODO: Uncomment when figured out trade fees
            # "trade_fees_initialized": len(self._trade_fees) > 0
        }

    @property
    def limit_orders(self):
        return [in_flight_order.to_limit_order() for in_flight_order in self._in_flight_orders.values()]

    def start(self, clock: Clock, timestamp: float):
        super().start(clock, timestamp)

    def stop(self, clock: Clock):
        super().stop(clock)

    async def start_network(self):
        self._order_book_tracker.start()
        self._trading_rules_polling_task = safe_ensure_future(self._trading_rules_polling_loop())
        if self._trading_required:
            self._status_polling_task = safe_ensure_future(self._status_polling_loop())
            self._user_stream_tracker_task = safe_ensure_future(self._user_stream_tracker.start())
            self._user_stream_event_listener_task = safe_ensure_future(self._user_stream_event_listener())

    def _stop_network(self):

        self._order_book_tracker.stop()
        if self._status_polling_task is not None:
            self._status_polling_task.cancel()
        if self._user_stream_tracker_task is not None:
            self._user_stream_tracker_task.cancel()
        if self._user_stream_event_listener_task is not None:
            self._user_stream_event_listener_task.cancel()
        if self._trading_rules_polling_task is not None:
            self._trading_rules_polling_task.cancel()
        self._status_polling_task = self._user_stream_tracker_task = \
            self._user_stream_event_listener_task = None

    async def stop_network(self):
        self._stop_network()

    async def check_network(self) -> NetworkStatus:
        try:
            client = Client(host=self._base_url,)
            client.public.get_time()
        except asyncio.CancelledError:
            raise
        except Exception:
            return NetworkStatus.NOT_CONNECTED
        return NetworkStatus.CONNECTED

    def supported_order_types(self) -> List[OrderType]:
        """
        :return a list of OrderType supported by this connector.
        Note that Market order type is no longer required and will not be used.
        """
#         return [OrderType.LIMIT, OrderType.MARKET]
        return [OrderType.LIMIT]

    # ORDER PLACE AND CANCEL EXECUTIONS ---
    async def create_order(self,
                           trade_type: TradeType,
                           order_id: str,
                           trading_pair: str,
                           amount: Decimal,
                           order_type: OrderType,
                           position_action: PositionAction,
                           price: Optional[Decimal] = Decimal("NaN")):

        trading_rule: TradingRule = self._trading_rules[trading_pair]
        if position_action not in [PositionAction.OPEN, PositionAction.CLOSE]:
            raise ValueError("Specify either OPEN_POSITION or CLOSE_POSITION position_action.")

        amount = self.quantize_order_amount(trading_pair, amount)
        price = self.quantize_order_price(trading_pair, price)

        if amount < trading_rule.min_order_size:
            raise ValueError(f"Buy order amount {amount} is lower than the minimum order size "
                             f"{trading_rule.min_order_size}")

        order_result = None
        api_params = {"symbol": trading_pair,
                      "side": "BUY" if trade_type is TradeType.BUY else "SELL",
                      "type": "LIMIT" if order_type is OrderType.LIMIT else "MARKET",
                      "quantity": f"{amount}",
                      "newClientOrderId": order_id
                      }
        if order_type == OrderType.LIMIT:
            api_params["price"] = f"{price}"
            api_params["timeInForce"] = "GTC"

        if self._position_mode == PositionMode.HEDGE:
            if position_action == PositionAction.OPEN:
                api_params["positionSide"] = "LONG" if trade_type is TradeType.BUY else "SHORT"
            else:
                api_params["positionSide"] = "SHORT" if trade_type is TradeType.BUY else "LONG"

        self.start_tracking_order(order_id, "", trading_pair, trade_type, price, amount, order_type, self._leverage, position_action.name)

        # To be updated!!! (Priority)
        try:
            client = Client(
                network_id=self._network_id,
                host=self._base_url,
                default_ethereum_address = self._eth_address,
                api_key_credentials=self._api_key_credentials,
            )
            account_response = client.private.get_account()
            position_id = account_response['account']['positionId']

            # Post an bid at a price that is unlikely to match.
            order_params = {
                'position_id': position_id,
                'market': api_params['symbol'],
                'side': api_params['side'],
                'order_type': api_params['type'],
                'post_only': False,
                'size': api_params['quantity'],
                'price': api_params['price'],
                'limit_fee': '0.0015',
                'expiration_epoch_seconds': time.time() + 99999999,
            }
            result_raw = client.private.create_order(**order_params)['order']
            exchange_order_id = str(result_raw["id"])

            # Check if order is filled immediately
            time.sleep(0.05)
            result_fill = client.private.get_fills(**{'order_id':exchange_order_id})['fills']
            avgPrice = "0.00000"
            executedQty = "0.00000"
            if result_fill and len(result_fill) > 0:
                avgPrice = result_fill[0]['price']
                executedQty = result_fill[0]['size']

            # Reformat order_result
            order_result = {
                'orderId': exchange_order_id,
                'symbol': result_raw['market'],
                'status': 'NEW',
                'clientOrderId': order_id,
                'price': result_raw['price'],
                'avgPrice': avgPrice,
                'origQty': result_raw['size'],
                'executedQty': executedQty,
                'cumQty': '0',
                'cumQuote': '0',
                'timeInForce': 'GTC',
                'type': result_raw['type'],
                'reduceOnly': False,
                'closePosition': False,
                'side': result_raw['side'],
                'positionSide': 'BOTH',
                'stopPrice': '0',
                'workingType': 'CONTRACT_PRICE',
                'priceProtect': False,
                'origType': 'LIMIT',
                'updateTime': time.time()
            }

            tracked_order = self._in_flight_orders.get(order_id)
            if tracked_order is not None:
                self.logger().info(f"Created {order_type.name.lower()} {trade_type.name.lower()} order {order_id} for "
                                   f"{amount} {trading_pair}.")
                tracked_order.exchange_order_id = exchange_order_id

            event_tag = self.MARKET_BUY_ORDER_CREATED_EVENT_TAG if trade_type is TradeType.BUY \
                else self.MARKET_SELL_ORDER_CREATED_EVENT_TAG
            event_class = BuyOrderCreatedEvent if trade_type is TradeType.BUY else SellOrderCreatedEvent
            self.trigger_event(event_tag,
                               event_class(self.current_timestamp,
                                           order_type,
                                           trading_pair,
                                           amount,
                                           price,
                                           order_id,
                                           leverage=self._leverage,
                                           position=position_action.name))
            return order_result
        except asyncio.CancelledError:
            raise
        except Exception as e:
            self.stop_tracking_order(order_id)
            self.logger().network(
                f"Error submitting order to Binance Perpetuals for {amount} {trading_pair} "
                f"{'' if order_type is OrderType.MARKET else price}.",
                exc_info=True,
                app_warning_msg=str(e)
            )
            self.trigger_event(self.MARKET_ORDER_FAILURE_EVENT_TAG,
                               MarketOrderFailureEvent(self.current_timestamp, order_id, order_type))

    async def execute_buy(self,
                          order_id: str,
                          trading_pair: str,
                          amount: Decimal,
                          order_type: OrderType,
                          position_action: PositionAction,
                          price: Optional[Decimal] = s_decimal_NaN):
        return await self.create_order(TradeType.BUY, order_id, trading_pair, amount, order_type, position_action, price)

    def buy(self, trading_pair: str, amount: object, order_type: object = OrderType.MARKET,
            price: object = s_decimal_NaN, **kwargs) -> str:

        t_pair: str = trading_pair
        order_id: str = get_client_order_id("sell", t_pair)
        safe_ensure_future(self.execute_buy(order_id, trading_pair, amount, order_type, kwargs["position_action"], price))
        return order_id

    async def execute_sell(self,
                           order_id: str,
                           trading_pair: str,
                           amount: Decimal,
                           order_type: OrderType,
                           position_action: PositionAction,
                           price: Optional[Decimal] = s_decimal_NaN):
        return await self.create_order(TradeType.SELL, order_id, trading_pair, amount, order_type, position_action, price)

    def sell(self, trading_pair: str, amount: object, order_type: object = OrderType.MARKET,
             price: object = s_decimal_NaN, **kwargs) -> str:

        t_pair: str = trading_pair
        order_id: str = get_client_order_id("sell", t_pair)
        safe_ensure_future(self.execute_sell(order_id, trading_pair, amount, order_type, kwargs["position_action"], price))
        return order_id

    async def cancel_all(self, timeout_seconds: float):
        incomplete_orders = [order for order in self._in_flight_orders.values() if not order.is_done]
        tasks = [self.execute_cancel(order.trading_pair, order.client_order_id) for order in incomplete_orders]
        order_id_set = set([order.client_order_id for order in incomplete_orders])
        successful_cancellations = []

        try:
            async with timeout(timeout_seconds):
                cancellation_results = await safe_gather(*tasks, return_exceptions=True)
                for cancel_result in cancellation_results:
                    # TODO: QUESTION --- SHOULD I CHECK FOR THE BinanceAPIException CONSIDERING WE ARE MOVING AWAY FROM BINANCE-CLIENT?
                    if isinstance(cancel_result, dict) and "clientOrderId" in cancel_result:
                        client_order_id = cancel_result.get("clientOrderId")
                        order_id_set.remove(client_order_id)
                        successful_cancellations.append(CancellationResult(client_order_id, True))
        except Exception:
            self.logger().network(
                "Unexpected error cancelling orders.",
                exc_info=True,
                app_warning_msg="Failed to cancel order with Binance Perpetual. Check API key and network connection."
            )
        failed_cancellations = [CancellationResult(order_id, False) for order_id in order_id_set]
        return successful_cancellations + failed_cancellations

    async def cancel_all_account_orders(self, trading_pair: str):
        try:
            client = Client(
                network_id=self._network_id,
                host=self._base_url,
                default_ethereum_address = self._eth_address,
                api_key_credentials=self._api_key_credentials,
            )
            response = client.private.cancel_all_orders(market=trading_pair)
            if response['cancelOrders'] and len(response['cancelOrders']) >= 0:
                for order_id in list(self._in_flight_orders.keys()):
                    self.stop_tracking_order(order_id)
            else:
                raise IOError(f"Error cancelling all account orders. Server Response: {response}")
        except Exception as e:
            self.logger().error("Could not cancel all account orders.")
            raise e

    def cancel(self, trading_pair: str, client_order_id: str):
        safe_ensure_future(self.execute_cancel(trading_pair, client_order_id))
        return client_order_id

    async def execute_cancel(self, trading_pair: str, client_order_id: str):
        try:
            # Find exchange_order_id using client_order_id
            open_orders = [order for order in self._in_flight_orders.values() if not order.is_done]
            order_id_set = {order.client_order_id:order.exchange_order_id for order in open_orders}
            print(order_id_set)
            exchange_order_id = order_id_set[client_order_id]

            client = Client(
                network_id=self._network_id,
                host=self._base_url,
                default_ethereum_address = self._eth_address,
                api_key_credentials=self._api_key_credentials,
            )
            response = client.private.cancel_order(order_id)

            if response['cancelOrder']['status']=='OPEN':
                self.logger().info(f"Successfully canceled order {client_order_id}")
                self.stop_tracking_order(client_order_id)
                self.trigger_event(self.MARKET_ORDER_CANCELLED_EVENT_TAG,
                                   OrderCancelledEvent(self.current_timestamp, client_order_id))
                return response

        except Exception as e:
            try:
                error_response = json.loads(e.response.text)
                error_msg = error_response['errors'][0]['msg']
                if 'it is already canceled' in error_msg or 'No order exists with id' in error_msg:
                    self.logger().debug(f"The order {client_order_id} does not exist on Binance Perpetuals. "
                                        f"No cancellation needed.")
                    self.stop_tracking_order(client_order_id)
                    self.trigger_event(self.MARKET_ORDER_CANCELLED_EVENT_TAG,
                                       OrderCancelledEvent(self.current_timestamp, client_order_id))
                    return {"origClientOrderId": client_order_id}
            except:
                pass
            self.logger().error(f"Could not cancel order {client_order_id} (on Binance Perp. {trading_pair})")
            raise e
        return response

    # TODO: Implement
    async def close_position(self, trading_pair: str):
        pass

    def quantize_order_amount(self, trading_pair: str, amount: object, price: object = Decimal(0)):
        trading_rule: TradingRule = self._trading_rules[trading_pair]
        # current_price: object = self.get_price(trading_pair, False)
        notional_size: object
        quantized_amount = DerivativeBase.quantize_order_amount(self, trading_pair, amount)
        if quantized_amount < trading_rule.min_order_size:
            return Decimal(0)
        """
        if price == Decimal(0):
            notional_size = current_price * quantized_amount
        else:
            notional_size = price * quantized_amount
        """

        # TODO: NOTIONAL MIN SIZE DOES NOT EXIST
        # if notional_size < trading_rule.min_notional_size * Decimal("1.01"):
        #     return Decimal(0)

        return quantized_amount

    def get_order_price_quantum(self, trading_pair: str, price: object):
        trading_rule: TradingRule = self._trading_rules[trading_pair]
        return trading_rule.min_price_increment

    def get_order_size_quantum(self, trading_pair: str, order_size: object):
        trading_rule: TradingRule = self._trading_rules[trading_pair]
        return Decimal(trading_rule.min_base_amount_increment)

    # ORDER TRACKING ---
    def start_tracking_order(self, order_id: str, exchange_order_id: str, trading_pair: str, trading_type: object,
                             price: object, amount: object, order_type: object, leverage: int, position: str):
        self._in_flight_orders[order_id] = DydxPerpetualsInFlightOrder(
            client_order_id=order_id,
            exchange_order_id=exchange_order_id,
            trading_pair=trading_pair,
            order_type=order_type,
            trade_type=trading_type,
            price=price,
            amount=amount,
            leverage=leverage,
            position=position

        )

    def stop_tracking_order(self, order_id: str):
        if order_id in self._in_flight_orders:
            del self._in_flight_orders[order_id]
        if order_id in self._order_not_found_records:
            del self._order_not_found_records[order_id]

    # Delete this??
    async def _iter_user_event_queue(self) -> AsyncIterable[Dict[str, any]]:
        while True:
            try:
#                 yield await self._user_stream_tracker.start()
#                 yield await self._user_stream_tracker.user_stream
                yield await self._user_stream_tracker.user_stream.get()
            except asyncio.CancelledError:
                raise
            except Exception:
                self.logger().network(
                    "Unknown error. Retrying after 1 seconds.",
                    exc_info=True,
                    app_warning_msg="Could not fetch user events from Binance. Check API key and network connection."
                )
                await asyncio.sleep(1.0)
    # This function:
    # Get order trade update. if order has been filled or If order has been cancelled
    # ACCOUNT_UPDATE, # Update balance # update position
    # MARGIN_CALL
    async def _user_stream_event_listener(self):
        async for event_message in self._iter_user_event_queue():
            try:
                if event_message['type'] != 'subscribed':
                    continue
                # If user stream event is an order update
                if 'orders' in event_message['contents']:
                    orders = event_message['contents']['orders']

                    # For each order
                    for order in orders:
                        exchange_order_id = order['id']

                        # Get open orders
                        open_orders = [order for order in self._in_flight_orders.values() if not order.is_done]
                        order_id_set = {order.exchange_order_id:order.client_order_id for order in open_orders}

                        # If the order has been cancelled
                        if exchange_order_id not in order_id_set:
                            continue

                        # Find client_order_id using exchange_order_id
                        client_order_id = order_id_set[exchange_order_id]

                        # Update _in_flight_orders
                        tracked_order = self._in_flight_orders.get(client_order_id)
                        trade_fee = self.get_fee(
                            base_currency=order['market'].split('-')[0],
                            quote_currency=order['market'].split('-')[1],
                            order_type=order['type'],
                            order_side=order['side'],
                            amount=Decimal(order['size']),
                            price=Decimal(order['price'])
                        )
                        order['trade_fee'] = trade_fee
                        tracked_order.update_with_execution_report(order)

                        # Send trigger_event for order fill event
                        trade_type = TradeType.BUY if order['side'] == "BUY" else TradeType.SELL
                        if order['status'] in ["ENTIRELY_FILLED", "FILLED", "CLOSED"]:
                            order_filled_event = OrderFilledEvent(
                                timestamp=int(time.time()),
                                order_id=client_order_id,
                                trading_pair=order['market'],
                                trade_type=trade_type,
                                order_type = OrderType.LIMIT if order['type'] == "LIMIT" else OrderType.MARKET,
                                price=Decimal(order['price']),
                                amount=Decimal(order['size']),
                                leverage=self._leverage,
                                trade_fee=trade_fee,
                                exchange_trade_id=order_message.get("t"),
                                position=tracked_order.position
                            )
                            self.trigger_event(self.MARKET_ORDER_FILLED_EVENT_TAG, order_filled_event)

                        # Send trigger_event if order is done
                        if tracked_order.is_done:
                            if not tracked_order.is_failure:
                                event_tag = None
                                event_class = None
                                if trade_type is TradeType.BUY:
                                    event_tag = self.MARKET_BUY_ORDER_COMPLETED_EVENT_TAG
                                    event_class = BuyOrderCompletedEvent
                                else:
                                    event_tag = self.MARKET_SELL_ORDER_COMPLETED_EVENT_TAG
                                    event_class = SellOrderCompletedEvent
                                self.logger().info(f"The {tracked_order.order_type.name.lower()} {trade_type} order {client_order_id} has completed "
                                                   f"according to websocket delta.")
                                self.trigger_event(event_tag,
                                                   event_class(self.current_timestamp,
                                                               client_order_id,
                                                               tracked_order.base_asset,
                                                               tracked_order.quote_asset,
                                                               (tracked_order.fee_asset or tracked_order.quote_asset),
                                                               tracked_order.executed_amount_base,
                                                               tracked_order.executed_amount_quote,
                                                               tracked_order.fee_paid,
                                                               tracked_order.order_type))
                            else:
                                if tracked_order.is_cancelled:
                                    if tracked_order.client_order_id in self._in_flight_orders:
                                        self.logger().info(f"Successfully cancelled order {tracked_order.client_order_id} according to websocket delta.")
                                        self.trigger_event(self.MARKET_ORDER_CANCELLED_EVENT_TAG,
                                                           OrderCancelledEvent(self.current_timestamp,
                                                                               tracked_order.client_order_id))
                                    else:
                                        self.logger().info(f"The {tracked_order.order_type.name.lower()} order {tracked_order.client_order_id} has failed "
                                                           f"according to websocket delta.")
                                        self.trigger_event(self.MARKET_ORDER_FAILURE_EVENT_TAG,
                                                           MarketOrderFailureEvent(self.current_timestamp,
                                                                                tracked_order.client_order_id,
                                                                                tracked_order.order_type))
                            self.stop_tracking_order(tracked_order.client_order_id)

                # if user stream event is an account update
                if 'account' in event_message['contents']:
                    update_data = event_message['contents']['account']

                    # Update balances - dydx_perpetual only has USDC collateral currently
                    asset_name = "USDC"
                    self._account_balances[asset_name] = Decimal(update_data["freeCollateral"]) # Trading wallet balance
                    self._account_available_balances[asset_name] = Decimal(update_data["freeCollateral"]) # Cross wallet balance. Not available in dydx_perpetual.

                    # update positions
                    open_positions = update_data['openPositions']
                    for market in open_positions:
                        open_position = open_positions[market]
                        position = self._account_positions.get(market, None)
                        if position is not None:
                            position.update_position(position_side=open_position['side'],
                                                     unrealized_pnl = Decimal(open_position['unrealizedPnl']),
                                                     entry_price = Decimal(open_position['entryPrice']),
                                                     amount = Decimal(open_position['size']))
                        else:
                            await self._update_positions()

            except asyncio.CancelledError:
                raise
            except Exception as e:
                self.logger().error(f"Unexpected error in user stream listener loop: {e}", exc_info=True)
                await asyncio.sleep(5.0)

    def tick(self, timestamp: float):
        """
        Is called automatically by the clock for each clock's tick (1 second by default).
        It checks if status polling task is due for execution.
        """
        now = time.time()
        poll_interval = (self.SHORT_POLL_INTERVAL
                         if now - self._user_stream_tracker.last_recv_time > 60.0
                         else self.LONG_POLL_INTERVAL)
        last_tick = int(self._last_timestamp / poll_interval)
        current_tick = int(timestamp / poll_interval)
        if current_tick > last_tick:
            if not self._poll_notifier.is_set():
                self._poll_notifier.set()
        self._last_timestamp = timestamp

    # MARKET AND ACCOUNT INFO ---
    def get_fee(self, base_currency: str, quote_currency: str, order_type: object, order_side: object,
                amount: object, price: object):
        is_maker = order_type is OrderType.LIMIT
        fee = estimate_fee("dydx_perpetual", is_maker)
        try:
            client = Client(
                network_id=self._network_id,
                host=self._base_url,
                default_ethereum_address = self._eth_address,
                api_key_credentials=self._api_key_credentials,
            )
            user_info = client.private.get_user()
            fee = Decimal(user_info['user']['makerFeeRate']) if is_maker else Decimal(user_info['user']['takerFeeRate'])
            fee = TradeFee(percent=fee, flat_fees=[])
        except Exception as e:
            print(e)
        return fee

    def get_order_book(self, trading_pair: str) -> OrderBook:
        order_books: dict = self._order_book_tracker.order_books
        if trading_pair not in order_books:
            raise ValueError(f"No order book exists for '{trading_pair}'.")
        return order_books[trading_pair]

    async def _update_trading_rules(self):
        last_tick = self._last_timestamp / 60.0
        current_tick = self.current_timestamp / 60.0
        if current_tick > last_tick or len(self._trading_rules) < 1:
            client = Client(host=self._base_url,)
            exchange_info = client.public.get_markets()['markets']
            trading_rules_list = self._format_trading_rules(exchange_info)
            self._trading_rules.clear()
            for trading_rule in trading_rules_list:
                self._trading_rules[trading_rule.trading_pair] = trading_rule

    def _format_trading_rules(self, exchange_info: Dict[str, Any]) -> List[TradingRule]:
        return_val = {key:{} for key in exchange_info}
        return_val = []
        for key in exchange_info:
            try:
                row = exchange_info[key]
                trading_pair = key
                min_order_size = row['minOrderSize']
                tick_size = row['tickSize']
                step_size = row['stepSize']

                return_val.append(TradingRule(
                    trading_pair,
                    min_order_size=Decimal(min_order_size),
                    min_price_increment=Decimal(tick_size),
                    min_base_amount_increment=Decimal(step_size),
                ))
            except Exception as e:
                self.logger().error(f"Error parsing the trading pair rule {rule}. Error: {e}. Skipping...",
                                    exc_info=True)
        return return_val

    async def _trading_rules_polling_loop(self):
        while True:
            try:
                await safe_gather(
                    self._update_trading_rules()

                    # TODO: Uncomment when implemented
                    # self._update_trade_fees()
                )
                await asyncio.sleep(60)
            except asyncio.CancelledError:
                raise
            except Exception:
                self.logger().network("Unexpected error while fetching trading rules.", exc_info=True,
                                      app_warning_msg="Could not fetch new trading rules from Binance Perpetuals. "
                                                      "Check network connection.")
                await asyncio.sleep(0.5)

    async def _status_polling_loop(self):
        while True:
            try:
                self._poll_notifier = asyncio.Event()
                await self._poll_notifier.wait()
                await safe_gather(
                    self._update_balances(),
                    self._update_positions()
                )
                await self._update_order_fills_from_trades(),
                await self._update_order_status()
                self._last_poll_timestamp = self.current_timestamp
            except asyncio.CancelledError:
                raise
            except Exception:
                self.logger().network("Unexpected error while fetching account updates.", exc_info=True,
                                      app_warning_msg="Could not fetch account updates from Binance Perpetuals. "
                                                      "Check API key and network connection.")
                await asyncio.sleep(0.5)

    # To be updated!!!!
    async def _update_balances(self):
        local_asset_names = {'USDC'}
        client = Client(
            network_id=self._network_id,
            host=self._base_url,
            default_ethereum_address = self._eth_address,
            api_key_credentials=self._api_key_credentials,
        )
        account = client.private.get_account()
        self._account_available_balances['USDC'] = account['account']['freeCollateral']
        self._account_balances['USDC'] = account['account']['equity']

    # TODO: Note --- Data Structure Assumes One-way Position Mode [not hedge position mode] (see Binance Futures Docs)
    # Note --- Hedge Mode allows for Both Long and Short Positions on a trading pair
    async def _update_positions(self):
        client = Client(
            network_id=self._network_id,
            host=self._base_url,
            default_ethereum_address = self._eth_address,
            api_key_credentials=self._api_key_credentials,
        )
        positions = client.private.get_positions(status='OPEN')['positions']
        for position in positions:
            trading_pair = position['market']
            position_side = PositionSide[position['side']]
            unrealized_pnl = Decimal(position['unrealizedPnl'])
            entry_price = Decimal(position['entryPrice'])
            amount = Decimal(position['size'])
            leverage = 25 # Leverage refers to the account's maximum leverage. Dydx does not have this feature yet, so set this to the maximum leverage possible.
            if amount != 0:
                self._account_positions[trading_pair + position_side.name] = Position(
                    trading_pair=trading_pair,
                    position_side=position_side,
                    unrealized_pnl=unrealized_pnl,
                    entry_price=entry_price,
                    amount=amount,
                    leverage=leverage
                )
            else:
                if (trading_pair + position_side.name) in self._account_positions:
                    del self._account_positions[trading_pair + position_side.name]

    async def _update_order_fills_from_trades(self):
        last_tick = int(self._last_poll_timestamp / self.UPDATE_ORDER_STATUS_MIN_INTERVAL)
        current_tick = int(self.current_timestamp / self.UPDATE_ORDER_STATUS_MIN_INTERVAL)

        if current_tick > last_tick and len(self._in_flight_orders) > 0:

            # Create exchange_order_id_to_order_map
            exchange_order_id_to_order_map = defaultdict(lambda: {})
            for order in self._in_flight_orders.values():
                exchange_order_id_to_order_map[order.exchange_order_id] = order

            # Get filled orders
            client = Client(
                network_id=self._network_id,
                host=self._base_url,
                default_ethereum_address = self._eth_address,
                api_key_credentials=self._api_key_credentials,
            )
            fills = client.private.get_fills(limit=100)['fills']
            self.logger().debug(f"Polling for order fills of trading_pairs.")

            # Update order status
            for fill in fills:
                fill_exchange_order_id = fill['orderId']
                if fill_exchange_order_id in exchange_order_id_to_order_map:
                    order = exchange_order_id_to_order_map[fill_exchange_order_id]
                    trade = {
                        'id':filled_order.exchange_order_id, # exchange_order_id
                        'order_id':filled_order.client_order_id, # order_id
                        'qty':fill['size'],
                        'quoteQty':Decimal(fill['size']) * Decimal(fill['price']),
                        'commission':fill['fee'],
                        'commissionAsset': 'USDC', #optional
                    }
                    print('fill', fill)
                    print('trade', trade)
                    applied_trade = order.update_with_trade_updates(trade)
                    if applied_trade:
                        self.trigger_event(
                            self.MARKET_ORDER_FILLED_EVENT_TAG,
                            OrderFilledEvent(
                                self.current_timestamp,
                                order.client_order_id,
                                order.trading_pair,
                                order.trade_type,
                                order_type,
                                Decimal(trade.get("price")),
                                Decimal(trade.get("qty")),
                                self.get_fee(
                                    order.base_asset,
                                    order.quote_asset,
                                    order_type,
                                    order.trade_type,
                                    Decimal(trade["price"]),
                                    Decimal(trade["qty"])),
                                exchange_trade_id=trade["id"],
                                leverage=self._leverage,
                                position=order.position
                            )
                        )

    # To be updated!!!
    async def _update_order_status(self):
        last_tick = int(self._last_poll_timestamp / self.UPDATE_ORDER_STATUS_MIN_INTERVAL)
        current_tick = int(self.current_timestamp / self.UPDATE_ORDER_STATUS_MIN_INTERVAL)
        if current_tick > last_tick and len(self._in_flight_orders) > 0:
            tracked_orders = list(self._in_flight_orders.values())
            print('tracked_orders', tracked_orders)

            client = Client(
                network_id=self._network_id,
                host=self._base_url,
                default_ethereum_address = self._eth_address,
                api_key_credentials=self._api_key_credentials,
            )
            results = []
            for order in tracked_orders:
                try:
                    order_update = client.private.get_order_by_id(order_id=order.exchange_order_id)['order']
                except:
                    order_update = {'error':'order not found'}
                results.append(order_update)
            self.logger().debug(f"Polling for order status updates of {len(results)} orders.")
            for order_update, tracked_order in zip(results, tracked_orders):
                print('order_update',order_update)
                print('tracked_order',tracked_order)
                client_order_id = tracked_order.client_order_id
                if client_order_id not in self._in_flight_orders:
                    continue
                if isinstance(order_update, Exception):
                    # NO_SUCH_ORDER code
                    if order_update["error"] == 'order not found':
                        self._order_not_found_records[client_order_id] = \
                            self._order_not_found_records.get(client_order_id, 0) + 1
                        if self._order_not_found_records[client_order_id] < self.ORDER_NOT_EXIST_CONFIRMATION_COUNT:
                            continue
                        self.trigger_event(
                            self.MARKET_ORDER_FAILURE_EVENT_TAG,
                            MarketOrderFailureEvent(self.current_timestamp, client_order_id, tracked_order.order_type)
                        )
                        self.stop_tracking_order(client_order_id)
                    else:
                        self.logger().network(f"Error fetching status update for the order {client_order_id}: "
                                              f"{order_update}.")
                    continue

                tracked_order.last_state = order_update.get("status")
                order_type = OrderType.LIMIT if order_update.get("type") == "LIMIT" else OrderType.MARKET
                executed_amount_base = Decimal(order_update.get("size", "0")) - Decimal(order_update.get("remainingSize", "0"))
                executed_amount_quote = executed_amount_base * Decimal(order_update.get("price", "0"))

                if tracked_order.is_done:
                    if not tracked_order.is_failure:
                        event_tag = None
                        event_class = None
                        if tracked_order.trade_type is TradeType.BUY:
                            event_tag = self.MARKET_BUY_ORDER_COMPLETED_EVENT_TAG
                            event_class = BuyOrderCompletedEvent
                        else:

                            event_tag = self.MARKET_SELL_ORDER_COMPLETED_EVENT_TAG
                            event_class = SellOrderCompletedEvent
                        self.logger().info(f"The {order_type.name.lower()} {tracked_order.trade_type.name.lower()} order {client_order_id} has "
                                           f"completed according to order status API.")
                        self.trigger_event(event_tag,
                                           event_class(self.current_timestamp,
                                                       client_order_id,
                                                       tracked_order.base_asset,
                                                       tracked_order.quote_asset,
                                                       (tracked_order.fee_asset or tracked_order.base_asset),
                                                       executed_amount_base,
                                                       executed_amount_quote,
                                                       tracked_order.fee_paid,
                                                       order_type))
                    else:
                        if tracked_order.is_cancelled:
                            self.logger().info(f"Successfully cancelled order {client_order_id} according to order status API.")
                            self.trigger_event(self.MARKET_ORDER_CANCELLED_EVENT_TAG,
                                               OrderCancelledEvent(self.current_timestamp,
                                                                   client_order_id))
                        else:
                            self.logger().info(f"The {order_type.name.lower()} order {client_order_id} has failed according to "
                                               f"order status API.")
                            self.trigger_event(self.MARKET_ORDER_FAILURE_EVENT_TAG,
                                               MarketOrderFailureEvent(self.current_timestamp,
                                                                       client_order_id,
                                                                       order_type))
                    self.stop_tracking_order(client_order_id)

    # Dydx currently does not allow users to set maximum margin.
    async def _set_margin(self, trading_pair: str, leverage: int = 1):
        params = {
            "symbol": trading_pair,
            "leverage": leverage
        }
        self._leverage = leverage
        self.logger().info(f"Leverage Successfully set to {leverage}.")
#         if set_leverage["leverage"] == leverage:
#             self._leverage = leverage
#             self.logger().info(f"Leverage Successfully set to {leverage}.")
#         else:
#             self.logger().error("Unable to set leverage.")
        return leverage

    def set_margin(self, trading_pair: str, leverage: int = 1):
        safe_ensure_future(self._set_margin(trading_pair, leverage))

    """
    async def get_position_pnl(self, trading_pair: str):
        await self._update_positions()
        return self._account_positions.get(trading_pair)
    """

    async def _get_funding_rate(self, trading_pair):
        # TODO: Note --- the "premiumIndex" endpoint can get markPrice, indexPrice, and nextFundingTime as well
        client = Client(host=self._base_url,)
        market_stats = client.public.get_markets(market=trading_pair)
        self._funding_rate = Decimal(market_stats['markets'][trading_pair]['nextFundingRate'])

    # William's comment: currently self._funding_rate is a Decimal variable for storing funding rate for one trading_pair. Should upgrade it into a json list of funding rates for every trading_pair.
    def get_funding_rate(self, trading_pair):
        safe_ensure_future(self._get_funding_rate(trading_pair))
        return self._funding_rate

    async def _set_position_mode(self, position_mode: PositionMode):
        initial_mode = await self._get_position_mode()
        if initial_mode != position_mode:
            self.logger().error(f"Unable to set postion mode to {position_mode.name}.")
        else:
            self.logger().info(f"Using {position_mode.name} position mode.")

    async def _get_position_mode(self):
        # Set _position_mode to ONEWAY only
        if self._position_mode is None:
            self._position_mode = PositionMode.ONEWAY
        return self._position_mode

    def set_position_mode(self, position_mode: PositionMode):
        safe_ensure_future(self._set_position_mode(position_mode))


#! /usr/bin/env python
import sys
from roc_model import ROCModel
from heat_problem import HeatProblem
from analysis_utils import *

size = int(sys.argv[1])
N = size
mesh_size = size
source = (0, 0, 1, size)
sink = [(1, 0, size-1, 1), (size-1, 0, 1, size-1), (1, size-1, size-1, 1)]
conductance = 10**-3  # this'll be used as resistance directly
hp = HeatProblem(N, source, sink, conductance, src_val=100.)

m = ROCModel(mesh_size)
m.load_problem(hp)
m.run_spice_solver()

surface3d = True

print('Node Potentials')
print('---------------')
print_node_potentials(m)
print()
print('Node Currents')
print('-------------')
print_node_currents(m)
print()
# print('Node Current Splits')
# print('-------------')
# splits = run_current_split_analysis(m)
# print_current_splits(splits)
# generate_sparams_from_splits(splits)

if not surface3d:
    plot_heatmap(m, current_flow_plot=None)
else:
    plot_surface(m)


# -*- coding: utf-8 -*-
# BSD 3-Clause License
#
# Copyright (c) 2017
# All rights reserved.
# Copyright 2022 Huawei Technologies Co., Ltd
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# * Redistributions of source code must retain the above copyright notice, this
#   list of conditions and the following disclaimer.
#
# * Redistributions in binary form must reproduce the above copyright notice,
#   this list of conditions and the following disclaimer in the documentation
#   and/or other materials provided with the distribution.
#
# * Neither the name of the copyright holder nor the names of its
#   contributors may be used to endorse or promote products derived from
#   this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
# ==========================================================================

#! -*- coding:utf-8 -*-
# 原项目：https://kexue.fm/archives/8847

from bert4torch.tokenizers import Tokenizer
from bert4torch.models import build_transformer_model, BaseModel
from bert4torch.snippets import sequence_padding, Callback, ListDataset, get_pool_emb, seed_everything
import torch.nn as nn
import torch
import torch.optim as optim
from torch.utils.data import DataLoader
from sklearn.metrics.pairwise import paired_cosine_distances
from scipy.stats import spearmanr
from tqdm import tqdm
import sys

# =============================基本参数=============================
# pooling, task_name = sys.argv[1:]  # 传入参数
pooling, task_name = 'cls', 'ATEC'  # debug使用
print('pooling: ', pooling, ' task_name: ', task_name)
assert task_name in ['ATEC', 'BQ', 'LCQMC', 'PAWSX', 'STS-B']
assert pooling in {'first-last-avg', 'last-avg', 'cls', 'pooler'}

maxlen = 64 if task_name != 'PAWSX' else 128
batch_size = 32
config_path = 'F:/Projects/pretrain_ckpt/bert/[google_tf_base]--chinese_L-12_H-768_A-12/bert_config.json'
checkpoint_path = 'F:/Projects/pretrain_ckpt/bert/[google_tf_base]--chinese_L-12_H-768_A-12/pytorch_model.bin'
dict_path = 'F:/Projects/pretrain_ckpt/bert/[google_tf_base]--chinese_L-12_H-768_A-12/vocab.txt'
device = 'cuda' if torch.cuda.is_available() else 'cpu'
seed_everything(42)

# 建立分词器
tokenizer = Tokenizer(dict_path, do_lower_case=True)

class MyDataset(ListDataset):
    @staticmethod
    def load_data(filename):
        """加载数据
        单条格式：(文本1, 文本2, 标签id)
        """
        D = []
        with open(filename, encoding='utf-8') as f:
            for l in f:
                l = l.strip().split('\t')
                if len(l) == 3:
                    D.append((l[0], l[1], int(l[2])))
        return D

def collate_fn(batch):
    batch_token_ids, batch_labels = [], []
    for text1, text2, label in batch:
        for text in [text1, text2]:
            token_ids, _ = tokenizer.encode(text, maxlen=maxlen)
            batch_token_ids.append(token_ids)
            batch_labels.append([label])

    batch_token_ids = torch.tensor(sequence_padding(batch_token_ids), dtype=torch.long, device=device)
    batch_labels = torch.tensor(batch_labels, dtype=torch.float, device=device)
    return batch_token_ids, batch_labels.flatten()

# 加载数据集
train_dataloader = DataLoader(MyDataset(f'F:/Projects/data/corpus/sentence_embedding/{task_name}/{task_name}.train.data'), batch_size=batch_size, shuffle=True, collate_fn=collate_fn)
valid_dataloader = DataLoader(MyDataset(f'F:/Projects/data/corpus/sentence_embedding/{task_name}/{task_name}.valid.data'), batch_size=batch_size, collate_fn=collate_fn)
test_dataloader = DataLoader(MyDataset(f'F:/Projects/data/corpus/sentence_embedding/{task_name}/{task_name}.test.data'), batch_size=batch_size, collate_fn=collate_fn)

# 定义bert上的模型结构
class Model(BaseModel):
    def __init__(self, pool_method='cls'):
        super().__init__()
        self.pool_method = pool_method
        with_pool = 'linear' if pool_method == 'pooler' else True
        output_all_encoded_layers = True if pool_method == 'first-last-avg' else False
        self.bert = build_transformer_model(config_path, checkpoint_path, segment_vocab_size=0,
                                            with_pool=with_pool, output_all_encoded_layers=output_all_encoded_layers)

    def forward(self, token_ids):
        hidden_state, pooler = self.bert([token_ids])
        sem_emb = get_pool_emb(hidden_state, pooler, token_ids.gt(0).long(), self.pool_method)
        return sem_emb

model = Model().to(device)

class MyLoss(nn.Module):
    def forward(self, y_pred, y_true):
        # 1. 取出真实的标签
        y_true = y_true[::2]    # tensor([1, 0, 1]) 真实的标签

        # 2. 对输出的句子向量进行l2归一化   后面只需要对应为相乘  就可以得到cos值了
        norms = (y_pred ** 2).sum(axis=1, keepdims=True) ** 0.5
        # y_pred = y_pred / torch.clip(norms, 1e-8, torch.inf)
        y_pred = y_pred / norms

        # 3. 奇偶向量相乘
        y_pred = torch.sum(y_pred[::2] * y_pred[1::2], dim=1) * 20

        # 4. 取出负例-正例的差值
        y_pred = y_pred[:, None] - y_pred[None, :]  # 这里是算出所有位置 两两之间余弦的差值
        # 矩阵中的第i行j列  表示的是第i个余弦值-第j个余弦值
        y_true = y_true[:, None] < y_true[None, :]   # 取出负例-正例的差值
        y_true = y_true.float()
        y_pred = y_pred - (1 - y_true) * 1e12
        y_pred = y_pred.view(-1)
        y_pred = torch.cat((torch.tensor([0.0], device=device), y_pred), dim=0)  # 这里加0是因为e^0 = 1相当于在log中加了1
        return torch.logsumexp(y_pred, dim=0)

# 定义使用的loss和optimizer，这里支持自定义
model.compile(
    loss=MyLoss(),
    optimizer=optim.Adam(model.parameters(), lr=2e-5),
)

class Evaluator(Callback):
    """评估与保存
    """
    def __init__(self):
        self.best_val_consine = 0.

    def on_epoch_end(self, global_step, epoch, logs=None):
        val_consine = self.evaluate(valid_dataloader)
        test_consine = self.evaluate(test_dataloader)

        if val_consine > self.best_val_consine:
            self.best_val_consine = val_consine
            # model.save_weights('best_model.pt')
        print(f'valid_consine: {val_consine:.5f}, test_consine: {test_consine:.5f}, best_val_consine: {self.best_val_consine:.5f}\n')

    # 定义评价函数
    def evaluate(self, data):
        embeddings1, embeddings2, labels = [], [], []
        for batch_token_ids, batch_labels in tqdm(data, desc='Evaluate'):
            embeddings = model.predict(batch_token_ids)
            embeddings1.append(embeddings[::2])
            embeddings2.append(embeddings[1::2])
            labels.append(batch_labels[::2])
        embeddings1 = torch.cat(embeddings1).cpu().numpy()
        embeddings2 = torch.cat(embeddings2).cpu().numpy()
        labels = torch.cat(labels).cpu().numpy()
        cosine_scores = 1 - (paired_cosine_distances(embeddings1, embeddings2))
        eval_pearson_cosine, _ = spearmanr(labels, cosine_scores)
        return eval_pearson_cosine

if __name__ == '__main__':
    evaluator = Evaluator()
    model.fit(train_dataloader, epochs=5, steps_per_epoch=None, callbacks=[evaluator])
else:
    model.load_weights('best_model.pt')


"""CPU functionality."""

import sys

# #List operations
LDI  = 0b10000010
PRN  = 0b01000111
MUL  = 0b10100010
PUSH = 0b01000101
POP  = 0b01000110
HLT  = 0b00000001
CALL = 0b01010000
RET  = 0b00010001
ADD  = 0b10100000
JMP  = 0b01010100
JEQ  = 0b01010101
JNE  = 0b01010110
CMP  = 0b10100111

class CPU:
    """Main CPU class."""

    def __init__(self):
        """Construct a new CPU."""
        self.registers = [0] * 8  # r0 - r7
        self.running = False
        self.ram = [0] * 256
        self.pc = 0
        self.sp = 7
        self.flag = 0b00000000

    def ram_read(self, MAR):
        """Read the RAM. MAR = memory address register"""
        try:
            return self.ram[MAR]
        except IndexError:
            print("index out of range for RAM read")

    def ram_write(self, MDR, MAR):
        """write to the RAM. MDR = Memory Data Register"""
        try:
            self.ram[MAR] = MDR
        except IndexError:
            print("index out of range for RAM write")

    def increment_pc(self, op_code):
        add_to_pc = (op_code >> 6) + 1
        self.pc += add_to_pc

    def load(self, filename):
        """Load a program into memory."""

        address = 0

        try:
            with open(filename) as f:
                for line in f:
                    comment_split = line.split('#')
                    number = comment_split[0].strip()

                    if number == "":
                        continue  # ignore blank lines

                    val = int(number, 2)

                    self.ram_write(val, address)

                    address += 1

        except FileNotFoundError:
            print(f"{sys.argv[0]}: {filename} not found")
            sys.exit(2)

    def alu(self, op, reg_a, reg_b):
        """ALU operations."""

        if op == "ADD":
            self.registers[reg_a] += self.registers[reg_b]
        # elif op == "SUB": etc
        elif op == "MUL":
            self.registers[reg_a] = self.registers[reg_a] * \
                self.registers[reg_b]
        elif op == "CMP":
            if self.registers[reg_a] == self.registers[reg_b]:
                self.flag = HLT
            elif self.registers[reg_a] > self.registers[reg_b]:
                self.flag = LDI
            else:
                self.flag = 0b00000100
        else:
            raise Exception("Unsupported ALU operation")

    def trace(self):
        """
        Handy function to print out the CPU state. You might want to call this
        from run() if you need help debugging.
        """

        print(f"TRACE: %02X | %02X %02X %02X |" % (
            self.pc,
            # self.fl,
            # self.ie,
            self.ram_read(self.pc),
            self.ram_read(self.pc + 1),
            self.ram_read(self.pc + 2)
        ), end='')

        for i in range(8):
            print(" %02X" % self.reg[i], end='')

        print()

########### ADDED TO CLEAN UP RUN. RAN OUT OF TIME #########
    # def ldi_op(self, op_code):
    #     address = self.ram_read(self.pc + 1)
    #     data = self.ram_read(self.pc + 2)
    #     self.registers[address] = data
    #     self.increment_pc(op_code)

    # def htl_op(self, op_code):
    #     self.running = False
    #     sys.exit(1)

    # def prn_op(self, op_code):
    #     address_a = self.ram_read(self.pc + 1)
    #     print(self.registers[address_a])
    #     self.increment_pc(op_code)

    # def comp_op(self, op_code):
    #     address_a = self.ram_read(self.pc + 1)
    #     address_b = self.ram_read(self.pc + 2)
    #     self.alu('CMP', address_a, address_b)
    #     self.increment_pc(op_code)

    # def add_op(self, op_code):
    #     address_a = self.ram_read(self.pc + 1)
    #     address_b = self.ram_read(self.pc + 2)
    #     self.alu('ADD', address_a, address_b)
    #     self.increment_pc(op_code)

    # def mul_op(self, op_code):
    #     address_a = self.ram_read(self.pc + 1)
    #     address_b = self.ram_read(self.pc + 2)
    #     self.alu('MUL', address_a, address_b)
    #     self.increment_pc(op_code)

    # def jmp_op(self, op_code):
    #     register_address = self.ram_read(self.pc + 1)
    #     self.pc = self.registers[register_address]

    # def jeq_op(self, op_code):
    #     register_address = self.ram_read(self.pc + 1)
    #     # check if equal is true
    #     if self.flag == HLT:
    #         self.pc = self.registers[register_address]
    #     # else increment by op code
    #     else:
    #         self.increment_pc(op_code)

    # def jne_op(self, op_code):
    #     register_address = self.ram_read(self.pc + 1)
    #     if self.flag != HLT:
    #         self.pc = self.registers[register_address]
    #     else:
    #         self.increment_pc(op_code)

    # def push_op(self, op_code):
    #     register_address = self.ram_read(self.pc + 1)
    #     val = self.registers[register_address]
    #     self.registers[self.sp] -= 1  # decrement the stack pointer
    #     self.ram[self.registers[self.sp]] = val
    #     self.increment_pc(op_code)

    # def pop_op(self, op_code):
    #     register_address = self.ram_read(self.pc + 1)
    #     val = self.ram[self.registers[self.sp]]
    #     self.registers[register_address] = val
    #     self.registers[self.sp] += 1
    #     self.increment_pc(op_code)

    # def call_op(self, op_code):
    #     self.registers[self.sp] -= 1
    #     self.ram[self.registers[self.sp]] = self.pc + 2
    #     address_of_subroutine = self.ram[self.pc + 1]
    #     self.pc = self.registers[address_of_subroutine]

    # def ret_op(self, op_code):
    #     self.pc = self.ram[self.registers[self.sp]]
    #     self.registers[self.sp] += 1

    # def operation_switch(self, arg):
    #     operations = {
    #         LDI: self.ldi_op(arg),
    #         PRN: self.prn_op(arg),
    #         MUL: self.mul_op(arg),
    #         PUSH: self.push_op(arg),
    #         POP: self.pop_op(arg),
    #         HLT: self.htl_op(arg),
    #         CALL: self.call_op(arg),
    #         RET: self.ret_op(arg),
    #         ADD: self.add_op(arg),
    #         JMP: self.jmp_op(arg),
    #         JEQ: self.jeq_op(arg),
    #         JNE: self.jne_op(arg),
    #         CMP: self.comp_op(arg),
    #     }

    #     return operations.get(arg, print("error fetching arg"))

    def run(self):
        """Run the CPU."""
        self.running = True

        while self.running:
            op_code = self.ram_read(self.pc)

            # self.operation_switch(op_code)

            # HLT
            if op_code == HLT:
                self.running = False
                sys.exit(1)
            # LDI
            elif op_code == LDI:
                address = self.ram_read(self.pc + 1)
                data = self.ram_read(self.pc + 2)
                self.registers[address] = data
                self.increment_pc(op_code)
            # PRN
            elif op_code == PRN:
                address_a = self.ram_read(self.pc + 1)
                print(self.registers[address_a])
                self.increment_pc(op_code)
                pass
            # CMP
            elif op_code == CMP:
                address_a = self.ram_read(self.pc + 1)
                address_b = self.ram_read(self.pc + 2)
                self.alu('CMP', address_a, address_b)
                self.increment_pc(op_code)
            # ADD
            elif op_code == ADD:
                address_a = self.ram_read(self.pc + 1)
                address_b = self.ram_read(self.pc + 2)
                self.alu('ADD', address_a, address_b)
                self.increment_pc(op_code)
            # MUL
            elif op_code == MUL:
                address_a = self.ram_read(self.pc + 1)
                address_b = self.ram_read(self.pc + 2)
                self.alu('MUL', address_a, address_b)
                self.increment_pc(op_code)
            # JMP
            elif op_code == JMP:
                register_address = self.ram_read(self.pc + 1)
                self.pc = self.registers[register_address]
            # JEQ
            elif op_code == JEQ:
                register_address = self.ram_read(self.pc + 1)
                if self.flag == HLT:
                    self.pc = self.registers[register_address]
                else:
                    self.increment_pc(op_code)
            # JNE
            elif op_code == JNE:
                register_address = self.ram_read(self.pc + 1)
                if self.flag != HLT:
                    self.pc = self.registers[register_address]
                else:
                    self.increment_pc(op_code)
            # PUSH
            elif op_code == PUSH:
                register_address = self.ram_read(self.pc + 1)
                val = self.registers[register_address]
                self.registers[self.sp] -= 1
                self.ram[self.registers[self.sp]] = val
                self.increment_pc(op_code)
            # POP
            elif op_code == POP:
                register_address = self.ram_read(self.pc + 1)
                val = self.ram[self.registers[self.sp]]
                self.registers[register_address] = val
                self.registers[self.sp] += 1
                self.increment_pc(op_code)
            # CALL
            elif op_code == CALL:
                self.registers[self.sp] -= 1
                self.ram[self.registers[self.sp]] = self.pc + 2
                address_of_subroutine = self.ram[self.pc + 1]
                self.pc = self.registers[address_of_subroutine]
            # RET
            elif op_code == RET:
                self.pc = self.ram[self.registers[self.sp]]
                self.registers[self.sp] += 1
            else:
                print('here is the else')