import platform

EDGE_FOLDER = "results/edge/"
CPU_FOLDER = "results/cpu/"

EDGETPU_SHARED_LIB = {
    'Linux':   'libedgetpu.so.1',
    'Darwin':   'libedgetpu.1.dylib',
    'Windows':   'edgetpu.dll'
}[platform.system()]

def make_interpreter(model_file):
    """Creates the interpreter object needed to deploy a model onto the tpu.

    Parameters
    ----------
    model_file : String
    Path to the tflite model that will be deployed to the edge tpu.

    Returns
    -------
    tflite.Interpreter Object
    """
    import tflite_runtime.interpreter as tflite

    model_file, *device = model_file.split('@')

    device = {'device': device[0]} if device else {}
    shared_library = EDGETPU_SHARED_LIB
    experimental_delegates = [tflite.load_delegate(
        shared_library, device)]  # Returns loaded Delegate object

    return tflite.Interpreter(model_path=model_file,
                              model_content=None,
                              experimental_delegates=experimental_delegates)

def edge_group_tflite_deployment(models_folder, count=5):
    """Deploys a group/series of tflite models onto the tpu.

    Parameters
    ----------
    models_folder : String
    Indicates the path to the folder where the compiled tflite models are
    located.

    count : Integer
    Number of times each model will be deployed.
    """
    from utils import deduce_operations_from_folder

    for model_info in deduce_operations_from_folder(models_folder, beginning="quant_", ending="_edgetpu.tflite"):
        edge_tflite_deployment(
                count, model_info[0], model_info[1])

def edge_tflite_deployment(count, model_file, model_name=None):
    """The actual deployment on the edge tpu of a single model.

    Parameters
    ----------
    model_file : String
    Indicates the path to the model which will be deployed located.

    model_name : String
    Indicates the name of the operation regarding this model, needed to know
    which folder the results must be saved onto.

    count : Integer
    Number of times this model will be deployed.
    """
    import time
    import logging
    import numpy as np
    from utils import create_csv_file, deduce_filename, deduce_operation_from_file

    edge_results = []
    if model_name == None:
        model_name = model_file.split("/")[model_file.count("/")]
        model_name = deduce_operation_from_file(
                model_name, beginning="quant_", ending="_edgetpu")

    # Creates Interpreter Object.
    interpreter = make_interpreter(model_file)
    interpreter.allocate_tensors()

    # This is where model specific inputs/labels are fed to the model in
    # order to be run correctly.

    input_details = interpreter.get_input_details()
    output_details = interpreter.get_output_details()

    input_shape = input_details[0]['shape']
    input_dtype = input_details[0]['dtype']

    input_data = np.array(np.random.random_sample(
        input_shape), dtype=input_dtype)
    interpreter.set_tensor(input_details[0]['index'], input_data)

    log = logging.getLogger(__name__)
    logging.basicConfig(format='%(levelname)s:%(message)s', level=logging.INFO)

    if count > 1:
        log.info("EDGE DEPLOYMENT...")

    for i in range(count):
        log.info(f"{model_name}: {i+1}/{count}")

        # INFERENCE TIME

        start = time.perf_counter()
        interpreter.invoke()  # Runs the interpreter/inference.

        # END
        inference_time = time.perf_counter() - start

        output_data = interpreter.get_tensor(output_details[0]['index'])

        edge_results.append([i, inference_time])
        # log.info('Inference Time %.1f us' % (inference_time * 10**6))

    create_csv_file(EDGE_FOLDER, model_name, edge_results)

def cpu_group_tflite_deployment(models_folder, count=5):
    """Deploys a group/series of tflite models onto the cpu.

    Parameters
    ----------
    models_folder : String
    Indicates the path to the folder where the compiled tflite models are
    located.

    count : Integer
    Number of times each model will be deployed.

    log_performance : Bool
    Tells if the deployed edge models are to be timed or not.
    """
    from utils import deduce_operations_from_folder

    for model_info in deduce_operations_from_folder(models_folder, beginning=None, ending=".tflite"):
        cpu_tflite_deployment(count, model_info[0], model_info[1])

def cpu_tflite_deployment(count, model_file, model_name=None):
    """The actual deployment on the cpu of a single model.

    Parameters
    ----------
    model_file : String
    Indicates the path to the model which will be deployed located.

    model_name : String
    Indicates the name of the operation regarding this model, needed to know
    which folder the results must be saved onto.

    count : Integer
    Number of times each model will be deployed.

    """
    import time
    import logging
    import tensorflow as tf
    import numpy as np
    from utils import create_csv_file, deduce_filename, deduce_operation_from_file

    cpu_results = []
    if model_name == None:
        model_name = model_file.split("/")[model_file.count("/")]
        model_name = deduce_operation_from_file(model_name, ending=".tflite")

    # Creates Interpreter Object.
    interpreter = tf.lite.Interpreter(model_path=model_file)
    interpreter.allocate_tensors()

    input_details = interpreter.get_input_details()
    output_details = interpreter.get_output_details()

    input_shape = input_details[0]['shape']
    input_dtype = input_details[0]['dtype']

    input_data = np.array(
            np.random.random_sample(input_shape), dtype=input_dtype)

    interpreter.set_tensor(input_details[0]['index'], input_data)
    log = logging.getLogger(__name__)
    logging.basicConfig(format='%(levelname)s:%(message)s', level=logging.INFO)

    if count > 1:
        log.info("CPU DEPLOYMENT...")

    for i in range(count):
        log.info(f"{model_name}: {i+1}/{count}")

        # INFERENCE TIME
        start = time.perf_counter()
        interpreter.invoke()  # Runs the interpreter/inference.

        # END
        inference_time = time.perf_counter() - start

        output_data = interpreter.get_tensor(output_details[0]['index'])

        cpu_results.append([i, inference_time])
        # log.info('Inference Time %.1f us\n' % (inference_time * 10**6))

    create_csv_file(CPU_FOLDER, model_name, cpu_results)

def tflite_deployment(quant_folder, tflite_folder, count=1000):
    """Manager function responsible for preping and executing the deployment
    of the compiled tflite models.

    Starts the docker, sets the number of times (count) that they will be
    deployed, copies the necessary folders on the docker, deploys the models,
    copies the results back.

    Parameters
    ----------
    count : Integer
    Indicates the number of times each model will be deplyed.
    """
    import os
    import utils
    from docker import docker_copy, docker_exec
    from docker import TO_DOCKER, FROM_DOCKER, HOME

    path_to_tensorDSE = utils.retrieve_folder_path(os.getcwd(), "TensorDSE")
    path_to_docker_results = HOME + \
        "TensorDSE/benchmarking/usb_analysis/results/"

    path_to_results = "results/"

    docker_copy(path_to_tensorDSE, TO_DOCKER)

    docker_exec("edge_deploy", quant_folder, count)
    docker_copy(path_to_docker_results + "edge/",
                FROM_DOCKER, path_to_results)

    docker_exec("cpu_deploy", tflite_folder, count)
    docker_copy(path_to_docker_results + "cpu/",
                       FROM_DOCKER, path_to_results)

if __name__ == '__main__':
    """Entry point to execute this script.

    Flags
    ---------
    -m or --mode
    Mode in which the script should run. Group, Single or Debug.
        Group is supposed to be used to deploy a group of models in sequence(lazy).
        Single deploys a single model.
        Debug simply deploys model on the edge tpu specified by user input within
        the debug mode.

    -d or --delegate
        Specifies which hardware should be used to delegate a model.
            cpu => cpu
            edge => edgetpu

    -t or --target
        Should be used in conjunction with the 'Single' mode, where -t must be followed
        by the path to the model (target) that will be deployed.

    -f or --folder
        Should be used in conjunction with the 'Group' mode, where -f must be followed
        by the path to a folder containing all models to be deployed.

    -c or --count
        Should be followed by the number of times one wishes to deploy the group of models
        or the single target (Depends on mode).
    """
    import argparse

    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter)

    parser.add_argument('-m', '--mode', default="Deploy",
                        required=False, help='Group, Single, Debug.')

    parser.add_argument('-d', '--delegate', default="",
                        required=False, help='cpu or edge.')

    parser.add_argument('-t', '--target',
                        help='File path to the .tflite file.')

    parser.add_argument('-f', '--folder', default="models/compiled/",
                        help='Path to folder where the group of models is located.')

    parser.add_argument('-c', '--count', type=int, default=1,
                        help='Number of times to run inference.')

    args = parser.parse_args()

    if (args.mode == "Group"):
        if (args.delegate == "cpu"):
            cpu_group_tflite_deployment(args.folder, count=args.count)
        else:
            edge_group_tflite_deployment(args.folder, count=args.count)

    elif (args.mode == "Single"):
        if (args.delegate == "cpu"):
            cpu_tflite_deployment(args.count, args.target)
        else:
            edge_tflite_deployment(args.count, args.target)

    elif args.mode == "Deploy-Group":
        import os
        for model in os.listdir(args.folder):
            if args.delegate == "cpu": target = f"{args.folder}{model}/{model}.tflite"
            else: target = f"{args.folder}{model}"
            os.system(
            f"python deploy.py -m Deploy-Single -d {args.delegate} -t {target} -c {str(args.count)}")

    elif args.mode == "Deploy-Single":
        import os
        from utils import retrieve_folder_path
        from utils import deduce_operation_from_file, extend_directory
        from docker import docker_exec, docker_copy, HOME, TO_DOCKER, FROM_DOCKER

        path_to_tensorDSE = retrieve_folder_path(os.getcwd(), "TensorDSE")
        path_to_docker_results = HOME + \
        "TensorDSE/benchmarking/usb_analysis/results/"

        docker_copy(path_to_tensorDSE, TO_DOCKER)
        op = args.target.split("/")[args.target.count("/")]
        if args.delegate == "cpu":
            op = deduce_operation_from_file(op, ending=".tflite")
            extend_directory("results/cpu/", f"{op}")
            path_to_results = f"results/cpu/{op}/"
            docker_exec("cpu_single_deploy", args.target, args.count)
            docker_copy(f"{path_to_docker_results}cpu/{op}/Results.csv",
                            FROM_DOCKER, path_to_results)
            docker_exec("remove", "TensorDSE")
        else:
            op = deduce_operation_from_file(op, beginning="quant_", ending="_edgetpu")
            extend_directory("results/edge/", f"{op}")
            path_to_results = f"results/edge/{op}/"
            docker_exec("edge_single_deploy", args.target, args.count)
            docker_copy(f"{path_to_docker_results}edge/{op}/Results.csv",
                            FROM_DOCKER, path_to_results)
            docker_exec("remove", "TensorDSE")

    elif args.mode == "Debug":
        import os
        from utils import deduce_operations_from_folder, retrieve_folder_path
        from docker import docker_start, docker_exec, docker_copy, TO_DOCKER

        docker_start()
        models_info = deduce_operations_from_folder(
                args.folder, beginning="quant_",ending="_edgetpu.tflite")

        path_to_tensorDSE = retrieve_folder_path(os.getcwd(), "TensorDSE")
        docker_copy(path_to_tensorDSE, TO_DOCKER)

        for m_i in models_info:
            inp = input(f"Operation {m_i[1]}, Continue to Next? ")
            if inp == "": continue
            elif inp == "c": break
            else:
                inp = "NONE"
                while(inp != "" and inp != 'c'):
                    if args.delegate == "cpu":
                        docker_exec("cpu_single_deploy", m_i[0], args.count)
                    else:
                        docker_exec("edge_single_deploy", m_i[0], args.count)
                    inp = input("Continue: ")

        docker_exec("remove", "TensorDSE")

    else:
        print("INVALID delegate input.")


"""
Rough script. Playing around with parallelizing on local machine while datagenjob.sh
gets munched by the cluster.
"""

from joblib import Parallel, delayed
import multiprocessing
from qpt import Sqpt_protocol
import param_optimizer as po
from misc_utilities import *
from qpt import *
import pickle

def ptest_i(maxis, counter, noise, measurements, axis_1_index):
    """
    use this to parallelize the for loop over basis vectors during datagen
    """

    try:
        cs = po.QPTparaopt(qubit_number=1, qparalist=maxis).cgen()
        c1n, c2n, c3n = [], [], []
        for index, channel in enumerate([damp, dephase, dep(0.2)]):
            if noise == 0:
                booll = False
            else:
                booll = True

            for i in measurements:
                protocol = Sqpt_protocol(channel, qubits=1,
                                    measurements=i, noise_level=noise,
                                    noisy_axis=(booll, True, True))
                oic = protocol.oics(qplist=maxis, c=cs)

                c, cd = protocol.sqpt(oic)
                if index == 0:
                    c1n.append(dn(c, cd, channel, 1))
                if index == 1:
                    c2n.append(dn(c, cd, channel, 1))
                if index == 2:
                    c3n.append(dn(c, cd, channel, 1))

        with open("ax1_{}_ampdamp{}n_{}maxis4.pickle".format(axis_1_index, int(noise*100), counter), "wb") as handle:
            pickle.dump(c1n, handle)
        with open("ax1_{}_dephase{}n_{}maxis4.pickle".format(axis_1_index, int(noise*100), counter), "wb") as handle:
            pickle.dump(c2n, handle)

        with open("ax1_{}_depol{}n_{}maxis4.pickle".format(axis_1_index, int(noise*100), counter), "wb") as handle:
            pickle.dump(c3n, handle)
    except NotImplementedError as e:
        print(e)


import paho.mqtt.client as mqtt
import RPi.GPIO as GPIO
import ssl, time, sys, json

MQTT_HOST = "afio7p37diyt8-ats.iot.eu-west-1.amazonaws.com"
CA_ROOT_CERT_FILE = "/home/pi/AmazonRootCA1.pem"
THING_NAME = "ESP"
THING_CERT_FILE = "/home/pi/f88bbe537b-certificate.pem.crt"
THING_PRIVATE_KEY_FILE = "/home/pi/f88bbe537b-private.pem.key"

LED_PIN = 21
MQTT_PORT = 8883
MQTT_KEEPALIVE_INTERVAL = 45
SHADOW_UPDATE_TOPIC = "$aws/things/" + THING_NAME + "/shadow/update"
SHADOW_UPDATE_ACCEPTED_TOPIC = "$aws/things/" + THING_NAME + "/shadow/update/accepted"
SHADOW_UPDATE_REJECTED_TOPIC = "$aws/things/" + THING_NAME + "/shadow/update/rejected"
SHADOW_UPDATE_DELTA_TOPIC = "$aws/things/" + THING_NAME + "/shadow/update/delta"
SHADOW_GET_TOPIC = "$aws/things/" + THING_NAME + "/shadow/get"
SHADOW_GET_ACCEPTED_TOPIC = "$aws/things/" + THING_NAME + "/shadow/get/accepted"
SHADOW_GET_REJECTED_TOPIC = "$aws/things/" + THING_NAME + "/shadow/get/rejected"
SHADOW_STATE_DOC_LED_ON = """{"state" : {"reported" : {"LED" : "ON"}}}"""
SHADOW_STATE_DOC_LED_OFF = """{"state" : {"reported" : {"LED" : "OFF"}}}"""

GPIO.setmode(GPIO.BCM)
GPIO.setwarnings(False)
GPIO.setup(LED_PIN, GPIO.OUT)

mqttc = mqtt.Client("client2")

def LED_Status_Change(Shadow_State_Doc, Type):
	# Parse LED Status from Shadow
	DESIRED_LED_STATUS = ""
	print ("\nParsing Shadow Json...")
	SHADOW_State_Doc = json.loads(Shadow_State_Doc)
	if Type == "DELTA":
		DESIRED_LED_STATUS = SHADOW_State_Doc['state']['LED']
	elif Type == "GET_REQ":
		DESIRED_LED_STATUS = SHADOW_State_Doc['state']['desired']['LED']
	print ("Desired LED Status: " + DESIRED_LED_STATUS)

	# Control LED
	if DESIRED_LED_STATUS == "ON":
		# Turn LED ON
		print ("\nTurning ON LED...")
		GPIO.output(LED_PIN, GPIO.HIGH)
		# Report LED ON Status back to Shadow
		print ("LED Turned ON. Reporting ON Status to Shadow...")
		mqttc.publish(SHADOW_UPDATE_TOPIC,SHADOW_STATE_DOC_LED_ON,qos=1)
	elif DESIRED_LED_STATUS == "OFF":
		# Turn LED OFF
		print ("\nTurning OFF LED...")
		GPIO.output(LED_PIN, GPIO.LOW)
		# Report LED OFF Status back to Shadow
		print ("LED Turned OFF. Reporting OFF Status to Shadow...")
		mqttc.publish(SHADOW_UPDATE_TOPIC,SHADOW_STATE_DOC_LED_OFF,qos=1)
	else:
		print ("---ERROR--- Invalid LED STATUS.")

def on_connect(mqttc, obj,flags, rc):
	print ("Connected to AWS IoT...")
	# Subscribe to Delta Topic
	mqttc.subscribe(SHADOW_UPDATE_DELTA_TOPIC, 1)
	# Subscribe to Update Topic
	#mqttc.subscribe(SHADOW_UPDATE_TOPIC, 1)
	# Subscribe to Update Accepted and Rejected Topics
	mqttc.subscribe(SHADOW_UPDATE_ACCEPTED_TOPIC, 1)
	mqttc.subscribe(SHADOW_UPDATE_REJECTED_TOPIC, 1)
	# Subscribe to Get Accepted and Rejected Topics
	mqttc.subscribe(SHADOW_GET_ACCEPTED_TOPIC, 1)
	mqttc.subscribe(SHADOW_GET_REJECTED_TOPIC, 1)

def on_message(mqttc, obj, msg):
	if str(msg.topic) == SHADOW_UPDATE_DELTA_TOPIC:
		print ("\nNew Delta Message Received...")
		SHADOW_STATE_DELTA = str(msg.payload.decode())
		print (SHADOW_STATE_DELTA)
		LED_Status_Change(SHADOW_STATE_DELTA, "DELTA")
	elif str(msg.topic) == SHADOW_GET_ACCEPTED_TOPIC:
		print ("\nReceived State Doc with Get Request...")
		SHADOW_STATE_DOC = str(msg.payload.decode())
		print (SHADOW_STATE_DOC)
		LED_Status_Change(SHADOW_STATE_DOC, "GET_REQ")
	elif str(msg.topic) == SHADOW_GET_REJECTED_TOPIC:
		SHADOW_GET_ERROR = str(msg.payload.decode())
		print ("\n---ERROR--- Unable to fetch Shadow Doc...\nError Response: " + SHADOW_GET_ERROR)
	elif str(msg.topic) == SHADOW_UPDATE_ACCEPTED_TOPIC:
		print ("\nLED Status Change Updated SUCCESSFULLY in Shadow...")
		print ("Response JSON: " + str(msg.payload.decode()))
	elif str(msg.topic) == SHADOW_UPDATE_REJECTED_TOPIC:
		SHADOW_UPDATE_ERROR = str(msg.payload.decode())
		print ("\n---ERROR--- Failed to Update the Shadow...\nError Response: " + SHADOW_UPDATE_ERROR)
	else:
		print ("AWS Response Topic: " + str(msg.topic))
		print ("QoS: " + str(msg.qos))
		print ("Payload: " + str(msg.payload.decode()))

def on_subscribe(mqttc, obj, mid, granted_qos):
	#As we are subscribing to 3 Topics, wait till all 3 topics get subscribed
	#for each subscription mid will get incremented by 1 (starting with 1)
	if mid == 3:
		# Fetch current Shadow status. Useful for reconnection scenario.
		mqttc.publish(SHADOW_GET_TOPIC,"",qos=1)

def on_disconnect(client, userdata, rc):
    if rc != 0:
        print ("Diconnected from AWS IoT. Trying to auto-reconnect...")

mqttc.on_message = on_message
mqttc.on_connect = on_connect
mqttc.on_subscribe = on_subscribe
mqttc.on_disconnect = on_disconnect

mqttc.tls_set(CA_ROOT_CERT_FILE, certfile=THING_CERT_FILE, keyfile=THING_PRIVATE_KEY_FILE, cert_reqs=ssl.CERT_REQUIRED, tls_version=ssl.PROTOCOL_TLSv1_2, ciphers=None)

# Connect with MQTT Broker
mqttc.connect(MQTT_HOST, MQTT_PORT, MQTT_KEEPALIVE_INTERVAL)

# Continue monitoring the incoming messages for subscribed topic
mqttc.loop_forever()


import sys
sys.setrecursionlimit(10000)

n = int(input())
total_c = 0
adj = [[] for _ in range(n)]
costs = {}
for _ in range(n-1):
    u, v, c = list(map(int, input().split()))
    u -= 1
    v -= 1
    total_c += c
    adj[u].append(v)
    adj[v].append(u)
    costs[f"{u};{v}"] = c
    costs[f"{v};{u}"] = c

visited = [False for _ in range(n)]
ans = [float("inf"), [None, None]]
def dfs(parent, node):
    global ans
    if visited[node]: return 0
    visited[node] = True

    sm = 0
    for s in adj[node]:
        if not visited[s]: sm += costs[f"{node};{s}"]
        sm += dfs(node, s)

    if parent != None:
        my_sum = sm + costs[f"{node};{parent}"]
        if ans[0] > min(ans[0], max(sm, total_c - my_sum)):
            ans = [min(ans[0], max(sm, total_c - my_sum)), [node + 1, parent + 1]]
    return sm

dfs(None, 0)
print(ans[0])


from __future__ import print_function
import os
import subprocess as sp
import sys
import xml.etree.ElementTree as ET
from time import time
import isoparser
from pylxd.client import Client
from pylxd.exceptions import LXDAPIException

# variables
containers_dir = "/var/lib/lxd/containers/"
datastores = "/var/lib/one/datastores/"
xml_pre = '/VM/TEMPLATE/'
separator = 40 * "-"
# MISC

def log_function(message, severity="i"):
    'print message with a definded severity in stderr'
    sep = ': '
    case = {
        "e": "ERROR",
        "i": "INFO",
        "d": "DEBUG"
    }

    print(case[severity] + sep + os.path.basename(sys.argv[0]) + sep + str(message), file=sys.stderr)

def log_function_verbose(message):
    for i in separator, message, separator:
        log_function(i, "d")

def clock(t0):
    'Calculates and logs in the logfile the time passed since $t0'
    duration = str(time() - t0)
    log_function('Script executed in almost ' + duration + ' seconds')

def vnc_start(VM_ID, VNC_PORT, VNC_PASSWD):
    'Starts VNC server in the one-$VM_ID container shell'
    try:
        sp.Popen('bash %s/vnc.bash %s %s %s' %
                 (sys.path[0], VM_ID, VNC_PORT, VNC_PASSWD), shell=True)
    except Exception as e:
        log_function(e, 'e')

def dir_empty(directory):
    directory = str(directory)
    if not os.listdir(directory):
        status = "empty"
    else:
        status = "non_empty"
    return status

def container_wipe(container, dicc):
    'Deletes $container after unmounting and unmapping its related storage'
    DISK_TYPE = xml_query_list(xml_pre + 'DISK/TYPE', dicc)
    num_hdds = len(DISK_TYPE)
    if num_hdds > 1:
        DISK_TARGET = xml_query_list(xml_pre + 'DISK/TARGET', dicc)
        for x in xrange(1, num_hdds):
            try:
                storage_unmap(container.devices, DISK_TARGET[x], DISK_TYPE[x])
            except Exception as e:
                expected = "Disk %s live removed previously" % (DISK_TARGET[x])
                for info in e, expected:
                    log_function(info)

    storage_rootfs_umount(DISK_TYPE[0], container.config)
    status = dir_empty(containers_dir + str(container.name))
    if status == "non_empty":
        log_function("Cannot delete non_empty container rootfs", 'e')
        sys.exit(1)
    else:
        container.delete()

# XML RELATED

def xml_start(xml):
    'Stores $xml file in $dicc dictionary'

    def xml_parse(root, dicc, way=""):
        new_way = way + "/" + root.tag
        value = []
        try:
            value = dicc[new_way]
        except:
            value = []
        value.append(root.text)
        dicc.__setitem__(new_way, value)
        for new_root in root._children:
            xml_parse(new_root, dicc, new_way)

    tree = ET.parse(xml, parser=None)
    root = tree.getroot()
    dicc = {}
    xml_parse(root, dicc, way="")
    return dicc

def xml_query_list(value, dicc):  # iplement item query or list query
    'Tries to return a list with $value instances from $dicc'
    try:
        value = dicc[value]
        # if len(value) == 1:
        #     value = value[0]
    except:
        value = [None]
    return value

def xml_query_dict(value, id_list, dicc):
    'Tries to return a dict with $value instances from $dicc'
    try:
        value_list = xml_query_list(value, dicc)
        pos = 0
        value_dict = {}
        for iface in id_list:
            value_dict[iface] = value_list[pos]
            pos += 1
    except Exception:
        value_dict = {None}
    return value_dict

def xml_query_item(value, dicc):
    'Returns the first element of the list $value obtained by calling xml_query'
    value = xml_query_list(value, dicc)
    return value[0]

# STORAGE SIMPLE
# TODO read disk specific arguments in a list
def storage_sysmap(DISK_ID, DISK_TYPE, DISK_IMAGE, VM_ID, DS_ID, DISK_CLONE):
    'Maps a $DISK_IMAGE device into the host corresponding block device, CEPH, LVM or FILESYSTEM'

    def storage_pre(command):
        blockdev = sp.check_output(command, shell=True)
        return blockdev.strip('\n')

    disk = None
    if DISK_TYPE == "FILE":
        disk = datastores + DS_ID + "/" + \
            VM_ID + "/" + 'disk.' + DISK_ID  # TODO fix hardcoded
        disk = storage_pre("losetup -f --show " + disk)
    elif DISK_TYPE == "BLOCK":
        pass
    elif DISK_TYPE == "RBD":
        if DISK_CLONE == 'YES':
            DISK_IMAGE = DISK_IMAGE + '-' + VM_ID + '-' + DISK_ID
        disk = storage_pre('rbd --id libvirt map ' + DISK_IMAGE)
    return disk

def storage_sysunmap(DISK_TYPE, source):
    'Unmaps a $source device from the host corresponding block device, CEPH, LVM or FILESYSTEM'
    if DISK_TYPE == "FILE":
        sp.call("losetup -d " + source, shell=True)
    elif DISK_TYPE == "BLOCK":
        pass
    elif DISK_TYPE == "RBD":
        sp.call('rbd unmap ' + source, shell=True)

def storage_lazer(device):
    'Returns the system device based on container mapped device'
    return sp.check_output('udevadm info -rq name /sys/dev/block/' + device['major'] + ':' + device['minor'], shell=True).strip('\n')

# STORAGE COMPOSED

def storage_rootfs_mount(VM_ID, DISK_TYPE, DS_ID, DISK_SOURCE, DISK_CLONE):
    'Mounts rootfs for container one-$VM_ID'
    source = storage_sysmap('0', DISK_TYPE, DISK_SOURCE,
                            VM_ID, DS_ID, DISK_CLONE)
    target = containers_dir + "one-" + VM_ID
    if dir_empty(target) == "non_empty":
        log_function(
            "Cannot mount container image over populated container directory", 'e')
        sys.exit(1)
    sp.call("mount " + source + " " + target, shell=True)
    return {'user.rootfs': source}

def storage_rootfs_umount(DISK_TYPE, container_config):
    'Unmounts the rootfs in $container_config and unmaps it from the host'
    source = container_config['user.rootfs']
    sp.call("umount " + source, shell=True)
    storage_sysunmap(DISK_TYPE, source)

def storage_context(container, contextiso):
    for i in contextiso.record().children:
        if i.is_directory is True:
            # FIXME: implement directory copy
            continue
        container.files.put('/mnt/' + i.name, i.content)

def storage_unmap(devices, disk_name, disk_type):
    disk = unmap(devices, disk_name)
    disk = storage_lazer(disk)
    storage_sysunmap(disk_type, disk)

# LXD CONFIG MAPPING

def map_disk(DISK_TARGET, host_device):
    'Creates a dictionary for LXD disk configuration, $DISK_TARGET will be path inside the container'
    dev_stat = os.stat(host_device)
    major = str(os.major(dev_stat.st_rdev))
    minor = str(os.minor(dev_stat.st_rdev))
    return {DISK_TARGET: {'path': '/dev/' + DISK_TARGET, 'type': 'unix-block', 'minor': minor, 'major': major}}

def map_cpu(CPU):
    'Creates a dicitionary for LXD containing $CPU percentage per core assigned to the container'
    CPU = str(int(float(CPU) * 100)) + '%'
    return {'limits.cpu.allowance': CPU}

def map_vcpu(VCPU):
    'Creates a dicitionary for LXD containing $VCPU cores assigned to the container'
    if VCPU:
        return {'limits.cpu': VCPU}
    else:
        return {'user.warning': 'using_all_CPUs'}

def map_ram(MEMORY):
    'Creates a dicitionary for LXD containing $MEMORY RAM allocated'
    MEMORY = MEMORY + 'MB'
    return {'limits.memory': MEMORY}

def map_xml(xml):
    'Creates a dicitionary for LXD containing the path to $xml deployment file'
    return {'user.xml': xml}

def map_nic(nic_name, NIC_BRIDGE, NIC_MAC, NIC_TARGET):
    'Creates a dicitionary for LXD containing $nic_name network interface configuration'
    return {nic_name: {'name': nic_name, 'type': 'nic', 'hwaddr': NIC_MAC, 'nictype': 'bridged', 'parent': NIC_BRIDGE, 'host_name': NIC_TARGET}}

def unmap(devices, device_name):
    'Delete and returns $device_name from $container_devices'
    source = devices[device_name]
    del devices[device_name]
    return source


import sys
import logging
from cPickle import loads

from allura import model as M

log = logging.getLogger('allura.migrate-artifact-refs')

# Threads have artifact references that must be migrated to the new system

def main():
    test = sys.argv[-1] == 'test'
    log.info('Fixing artifact references in threads')
    db = M.project_doc_session.db
    for thread in db.thread.find():
        ref = thread.pop('artifact_reference', None)
        if ref is None:
            continue
        Artifact = loads(ref['artifact_type'])
        artifact = Artifact.query.get(_id=ref['artifact_id'])
        M.ArtifactReference.from_artifact(artifact)
        thread['ref_id'] = artifact.index_id()
        if not test:
            db.thread.save(thread)
            log.info('saving thread %s', thread['_id'])
        else:
            log.info('would save thread %s', thread['_id'])
        M.artifact_orm_session.clear()

if __name__ == '__main__':
    main()


#!/usr/bin/env python
import matplotlib
import matplotlib.pyplot as plt
import matplotlib.cm as cm
import tasks
import sys
import run_model

def plot(seq_length):
    i, o = tasks.copy(8, seq_length)
    weights, outputs = do_task(i)
    plt.figure(1, figsize=(20, 7))
    plt.subplot(311)
    plt.imshow(i.T, interpolation='nearest')
    plt.subplot(312)
    plt.imshow(o.T, interpolation='nearest')
    plt.subplot(313)
    plt.imshow(outputs.T, interpolation='nearest')
    plt.show()

def plot_weights(seq_length):
    i, o = tasks.copy(8, seq_length)
    weights, outputs = do_task(i)
    plt.figure(1, figsize=(20, 20))
    plt.imshow(weights.T[100:123], interpolation='nearest', cmap=cm.gray)
    plt.show()

if __name__ == "__main__":
    P, do_task = run_model.make_model()
    P.load('l2_low_learning_rate.mdl')
    plot(130)
    # plot_weights(8)


from .parse import Config, Parser
from .auth import client

def create_map_task(source, destination):
    return {
        'sourceFields': [source],
        'destinationField': destination,
        'taskType': 'Map',
        'taskProperties': {}
    }

def create_projection_filter_task(*sources):
    return{
        "connectorOperator": {
            "S3": "PROJECTION",
            "Salesforce": "PROJECTION"
        },
        "sourceFields": list(sources),
        "taskProperties": {},
        "taskType": "Filter"
    }

def create_destination_flow_config_list(connector_type, profile_name, sf_object):
    return [{
        'connectorType': connector_type,
        'connectorProfileName': profile_name,
        'destinationConnectorProperties': {
            'Salesforce': {
                'object': sf_object,
                # 'errorHandlingConfig': {
                #     'failOnFirstDestinationError': True,
                #     'bucketPrefix': 'errors-contact',
                #     'bucketName': 't-est-bucket'
                # },
                'writeOperationType': 'INSERT'  # | 'UPSERT' | 'UPDATE'
            }
        }
    }]

def create_source_flow_config(bucket_name, bucket_prefix):
    return {
        'connectorType': 'S3',
        'sourceConnectorProperties': {
            'S3': {
                'bucketName': bucket_name,
                'bucketPrefix': bucket_prefix
            }
        }
    }

class SetFlow:
    config: Config
    connector_type = 'Salesforce'
    destinationFlowConfigList = None
    sourceFlowConfig = None
    triggerConfig = {
        'triggerType': 'OnDemand',
    }
    tasks = []

    def __init__(self, flow_name: str, description: str, config: Config) -> None:
        self.config = config
        for o in config.mappings_obj:
            print("config", o.destination, o.source)
        self.flow_name = flow_name
        self.description = description

    def set_sf(self, profile_name, sf_object):
        self.destinationFlowConfigList = create_destination_flow_config_list(
            'Salesforce', profile_name, sf_object)

    def set_s3(self, bucket_name, bucket_prefix):
        self.sourceFlowConfig = create_source_flow_config(
            bucket_name, bucket_prefix)

    def __set_tasks(self):
        sources = []
        for obj in self.config.mappings_obj:
            sources.append(obj.source)
            map_task = create_map_task(obj.source, obj.destination)
            self.tasks.append(map_task)
        self.tasks.insert(0, create_projection_filter_task(*sources))

    def create_flow(self):
        self.__set_tasks()
        print(self.tasks)
        print(self.destinationFlowConfigList)
        response = client.create_flow(
            flowName=self.flow_name,
            description=self.description,
            triggerConfig=self.triggerConfig,
            sourceFlowConfig=self.sourceFlowConfig,
            destinationFlowConfigList=self.destinationFlowConfigList,
            tasks=self.tasks
        )
        return response

    def update_flow(self):
        response = client.update_flow(
            flowName=self.flow_name,
            description=self.description,
            triggerConfig=self.triggerConfig,
            sourceFlowConfig=self.sourceFlowConfig,
            destinationFlowConfigList=self.destinationFlowConfigList,
            tasks=self.tasks
        )
        return response


import pandas as pd
import numpy as np
import json
import os
from Solver import solver
sol = solver()

class dataLoader:

    def __init__(self):
        return

    def getMarkerPose(self, frame, marker_num, csv_data):
        '''
            Read the marker pose and list as numpy array [tx,ty,tz].
        '''
        marker_pose = csv_data.iloc[frame, 10*(marker_num+1):10*(marker_num+1)+3] * 1000
        marker_pose = marker_pose.to_numpy()
        return marker_pose

    def makeFrame(self, frame_idx, csv_data, marker_num=5):
        '''
        makes 5 markers coordinates an array of frame
        '''
        markers = []
        for i in range(marker_num):
            markers.append(self.getMarkerPose(frame_idx, i, csv_data))
        markers = np.vstack(markers)
        markers = markers.astype('float64')
        return markers

    def makePanelFrame(self, frame_idx, csv_data):
        '''
        makes 4 markers coordinates an array of frame
        '''
        markers = []
        marker_num = [5, 6, 7, 8]
        for i in marker_num:
            markers.append(self.getMarkerPose(frame_idx, i, csv_data))
        markers = np.vstack(markers)
        markers = markers.astype('float64')
        return markers

    def getCentroid(self, initFilePath, markerNum, panel=False):
        '''
        get the centroid of 5 markers
        centralize those markers
        '''
        csv_data = pd.read_csv(initFilePath)
        frames = []
        normal_frames = []
        norms = []
        for i in range(400):
            if panel == False:
                frames.append(self.makeFrame(10+i, csv_data, markerNum))
            else:
                frames.append(self.makePanelFrame(10 + i, csv_data))
        frames = np.array(frames)
        frame0 = frames.mean(0)
        frame = frame0.mean(0)
        print("center", frame)
        for i in range(len(frame0)):
            normal_frames.append(frame0[i] - frame)
        normal_frames = np.vstack(normal_frames)
        # for i in range(len(normal_frames)):
        #     for j in range(len(normal_frames)):
        #         norms.append(np.linalg.norm(normal_frames[i] - normal_frames[j]))
        # print(norms)
        return normal_frames

    def frames(self, num_frames, csv_data, panel=False):
        if panel == False:
            fs = [self.makeFrame(0, csv_data)]
            # print('len csv', len(csv_data))
            start = int(np.random.random() * 50)
            frame_num = np.linspace(start, len(csv_data), num_frames, dtype=int)
            print('frame:', len(frame_num))
            # for i in range(num_frames):
            for i in frame_num:
                # frame_num = int(np.random.random() * len(csv_data))
                # fs.append(self.makeFrame(frame_num, csv_data))
                fs.append(self.makeFrame(i-1, csv_data))
        else:
            fs = [self.makePanelFrame(0, csv_data)]
            start = int(np.random.random() * 50)
            frame_num = np.linspace(start, len(csv_data), num_frames, dtype=int)
            print('frame:', len(frame_num))
            for i in frame_num:
                fs.append(self.makePanelFrame(i - 1, csv_data))
        return fs

    def getToolPose(self, frame, csv_path):
        '''
        get the pose data from Surgical_tool topic as:
        [tx, ty, tz, qx, qy, qz, qw]
        @param frame: index of the frame get the data from
        '''
        csv_data = pd.read_csv(csv_path, header=None)
        tool_placement = csv_data.iloc[frame, 1:4] * 1000
        tool_orientation = csv_data.iloc[frame, 4:]
        tool_placement = tool_placement.to_numpy()
        tool_orientation = tool_orientation.to_numpy()
        tool_pose = np.hstack((tool_placement, tool_orientation))
        tool_pose = tool_pose.astype('float64')
        return tool_pose

    def loadJson(self, path):
        '''
        load json file from path.
        '''
        with open(path, 'r') as load_f:
            load_dict = json.load(load_f)
        # print('dict', load_dict)
        ls = [load_dict['pose']['position']['x']*1000, load_dict['pose']['position']['y']*1000,
              load_dict['pose']['position']['z']*1000, load_dict['pose']['orientation']['x'],
              load_dict['pose']['orientation']['y'], load_dict['pose']['orientation']['z'],
              load_dict['pose']['orientation']['w']]
        ls = np.array(ls)
        return ls

    def loadHandeyeJson(self, path):
        f = open(path + 'calibration_optimized.json')
        hand_eye_info = json.load(f)
        quaternion_X = []
        for j in hand_eye_info["translation"]:
            quaternion_X.append(float(hand_eye_info["translation"][j])*1000)
        for i in hand_eye_info["rotation"]:
            quaternion_X.append(float(hand_eye_info["rotation"][i]))

        _, X = sol.seven2trans(np.array(quaternion_X))
        np.save(path+'hand_eye_X.npy', X)
        return X

    def loadMaualPointCloud(self, dirPath, sub_dir):
        fileNum = 0
        p_list = []
        sub_path = os.path.join(dirPath, sub_dir)
        for name in os.listdir(sub_path):
            name = os.path.join(os.path.abspath(sub_path), name)
            data = self.loadJson(name)
            data[0] = data[0]/1000
            data[1] = data[1]/1000
            data[2] = data[2]/1000
            # print(data)
            p_list.append(data)
            fileNum = fileNum + 1
        p_list = np.vstack(p_list)
        print('point num:', fileNum)
        return p_list

    def getRealPose(self, idx, dirPath):
        df = pd.read_csv(dirPath)
        pose_x = df['pose.position.x'][idx] * 1000
        pose_y = df['pose.position.y'][idx] * 1000
        pose_z = df['pose.position.z'][idx] * 1000
        orin_x = df['pose.orientation.x'][idx]
        orin_y = df['pose.orientation.y'][idx]
        orin_z = df['pose.orientation.z'][idx]
        orin_w = df['pose.orientation.w'][idx]
        pose = np.array([pose_x, pose_y, pose_z, orin_x, orin_y, orin_z, orin_w])
        # print(real_pose, '\n', pose)
        return pose

    def getTimestamp_list(self, csv_path):
        csv_data = pd.read_csv(csv_path, header=None)
        timestamp_list = []
        for i in range(len(csv_data)):
            timestamp = round(csv_data.iloc[i, 0], 5)
            timestamp_list.append(timestamp)
        # timestamp_list = timestamp_list.astype('float64')
        return timestamp_list

if __name__ == '__main__':
    ld = dataLoader()
    # tool_pose = ld.getToolPose(0, '/home/shc/Desktop/data/915/handeye_imu_2/cam_pose.csv')
    # tl = ld.getTimestamp_list('/home/shc/Desktop/data/915/handeye_imu_2/cam_pose.csv')
    # print(tl)


import cv2
from sight import *

image = Image(cv2.imread("robots3.jpg", 1)).resize((800, 600))

while True:
    original = image.clone()
    output = image.clone()

    mask = output.get_mask(
        [0, 20, 100],
        [59, 43, 255]
    )

    blobs = output.mask(mask).smooth().blur(0.3).draw_targets(mask).get_blobs(mask)

    blobs = sorted(
                blobs,
                key=lambda t: t[2],
                reverse=True
                )[:3]

    for (x, y, radius) in blobs:
        if radius > 20:
            # output.draw_text((x-radius/output.get_width(), y), 'Red Bot')
            original.draw_text(
                (
                    x-radius/output.get_width(),
                    y-radius/output.get_height()
                ), 'Red Bot')

    original.show("overlay")

    if escape_key_pressed():
        break

exit()


"""
qfly | Qualisys Drone SDK Example Script: Drone Pong

UNSTABLE / EXPERIMENTAL / UNDER DEVELOPMENT
"""

import math
import pynput
from time import sleep, time

from qfly import ParallelContexts, Pose, QualisysCrazyflie, QualisysDeck, World

# SETTINGS
cf_body_name = 'E7E7E7E706'
cf_uri = 'radio://0/80/2M/E7E7E7E706'
cf_marker_ids = [41, 62, 63, 64]

deck_body_name = 'E7E7E7E7E7'
deck_uri = 'radio://0/80/2M/E7E7E7E7E7'
deck_marker_ids = [1, 2, 3, 4]

# Watch key presses with a global variable
last_key_pressed = None

# Set up keyboard callback
def on_press(key):
    """React to keyboard."""
    global last_key_pressed
    last_key_pressed = key
    if key == pynput.keyboard.Key.esc:
        fly = False

# Listen to the keyboard
listener = pynput.keyboard.Listener(on_press=on_press)
listener.start()

# Set up world - the World object comes with sane defaults
world = World(expanse=1.5, speed_limit=1, padding=0.3)

vx = 1
vy = 1
r_paddle = 0.5
bounce_angle_max = 1.5  # in radians

# Stack up context managers
qcf = QualisysCrazyflie(cf_body_name,
                        cf_uri,
                        world,
                        marker_ids=cf_marker_ids)

deck = QualisysDeck(deck_body_name,
                    deck_uri,
                    deck_marker_ids)

with ParallelContexts(*[qcf, deck]):

    print("Beginning maneuvers...")

    # Let there be time
    t = time()
    dt = 0

    # MAIN LOOP WITH SAFETY CHECK
    while(qcf.is_safe()):

        # Mind the clock
        dt = time() - t

        # Terminate upon Esc command
        if last_key_pressed == pynput.keyboard.Key.esc:
            break

        if dt < 5:
            qcf.rise_in_place()
            continue

        # Bouncing algorithm for paddle hits
        if deck.pose is not None and qcf.pose is not None:
            if deck.pose.distance_to(qcf.pose) < r_paddle:
                relative_intersect_y = deck.pose.y + r_paddle - qcf.pose.y
                normalized_relative_intersect_y = relative_intersect_y / r_paddle
                bounce_angle = normalized_relative_intersect_y * bounce_angle_max
                # vx = abs(world.speed_limit * math.cos(bounce_angle))
                vx = 1.0
                vy = world.speed_limit * math.cos(bounce_angle)

        # Bouncing Algorithm when the Ball hit the edge of the canvas
        if qcf.pose.x > world.expanse - world.padding and vx > 0:
            vx = -vx
        if qcf.pose.x < -world.expanse + world.padding and vx < 0:
            vx = -vx
        if qcf.pose.y > world.expanse - world.padding and vy > 0:
            vy = -vy
        if qcf.pose.y < -world.expanse + world.padding and vy < 0:
            vy = -vy

        # Set target
        x = qcf.pose.x + vx
        y = qcf.pose.y + vy
        print(x)
        print(vx)
        print(y)
        print(vy)
        target = Pose(x, y, 1)
        # Engage
        qcf.safe_position_setpoint(target)

    # Land calmly
    qcf.land_in_place()


import tensorflow as tf
import matplotlib.pyplot as plt
import os
from agents import Predictor, A3CNetwork
from agents.evaluation import Evaluator
from simulation import StockSimulation
from data.access import DataPipeLine
from data.preparation import split_np_arrays
from utils import Hyperparameter

plt.style.use('ggplot')
# region Hyperparameter

hps = Hyperparameter()
hps.load(os.path.join('files', 'logging', 'A3C', '7eval17', 'Hyperparameter.yaml'))

predictor_dir = os.path.join('files',  'models', 'PredictorV2', '02RegWG' + str(hps.warengruppe[0]))
available_weights = os.listdir(predictor_dir)
available_weights.sort()
predictor_path = os.path.join(predictor_dir, available_weights[-1])
# endregion

pipeline = DataPipeLine(ZielWarengruppen=hps.warengruppe, DetailWarengruppe=hps.detail_warengruppe)
simulation_data = pipeline.get_regression_data()

train_data, test_data = split_np_arrays(*simulation_data)

predictor = Predictor()
predictor.build_model(
    dynamic_state_shape=simulation_data[1].shape[2],
    static_state_shape=simulation_data[2].shape[1]
)
predictor.load_from_weights(predictor_path)

print('Predicting ', end='')
train_pred = predictor.predict(
    {
        'dynamic_input': train_data[1],
        'static_input': train_data[2]
    }
)
print('and done ;)')
print('Predicting ', end='')
test_pred = predictor.predict(
    {
        'dynamic_input': test_data[1],
        'static_input': test_data[2]
    }
)
print('and done ;)')
# endregion

with tf.device("/cpu:0"):
    global_episodes = tf.Variable(0, dtype=tf.int32, name='global_episodes', trainable=False)
    master_network = A3CNetwork('global', None, hps)
    saver = tf.train.Saver(max_to_keep=1, var_list=tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope='global'))

with tf.Session() as sess:
    print('Loading Model...')

    ckpt = tf.train.get_checkpoint_state(hps.model_dir)
    saver.restore(sess, ckpt.model_checkpoint_path)
    simulation = StockSimulation(train_data, train_pred, hps)
    validator = StockSimulation(test_data, test_pred, hps)

    evaluation = Evaluator(master_network, simulation, validator, hps, session=sess)
    evaluation.show()


#%%
import sys
import os
import csv
import argparse
import multiprocessing
import numpy as np
import random
import time
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import transforms
from tqdm import tqdm
from checkpoint import (
    default_checkpoint,
    load_checkpoint,
    save_checkpoint,
    init_tensorboard,
    write_tensorboard,
)
from model import Encoder, Decoder
from dataset import CrohmeDataset, START, PAD, collate_batch,load_vocab
from data_tools.inkmltopng import inkml2png
from data_tools.extract_groundtruth import create_tsv

#%%
#inkml2png([0,"bonus_inkml",128,0])
create_tsv('bonus_inkml','bonus_inkml.tsv')

#%%
from train import parse_args
gt_test = "bonus_inkml.tsv"
tokensfile = "./data/tokens.tsv"
root = "./data_test_png"
use_cuda = torch.cuda.is_available()

input_size = (128, 128)
low_res_shape = (684, input_size[0] // 16, input_size[1] // 16)
high_res_shape = (792, input_size[0] // 8, input_size[1] // 8)

batch_size = 1
num_epochs = 1
print_epochs = 1
learning_rate = 1e-3
lr_epochs = 20
lr_factor = 0.1
weight_decay = 1e-4
max_grad_norm = 5.0
dropout_rate = 0.2
teacher_forcing_ratio = 0.5
seed = 1234

torch.manual_seed(seed)
is_cuda = use_cuda
hardware = "cuda" if is_cuda else "cpu"
device = torch.device(hardware)
transformers = transforms.Compose(
    [
        # Resize so all images have the same size
        transforms.Resize(input_size),
        transforms.ToTensor(),
    ]
)
test_dataset = CrohmeDataset(
    gt_test, tokensfile, root=root, crop=False, transform=transformers
)

test_data_loader = DataLoader(
    test_dataset,
    batch_size=batch_size,
    shuffle=True,
    collate_fn=collate_batch,
)

#%%

enc = Encoder(
    img_channels=3,
    dropout_rate=dropout_rate,
    checkpoint=load_checkpoint("checkpoints/0009.pth")["model"].get("encoder"),
).to(device)
dec = Decoder(
    len(test_dataset.id_to_token),
    low_res_shape,
    high_res_shape,
    checkpoint=load_checkpoint("checkpoints/0009.pth")["model"].get("decoder"),
    device=device,
).to(device)
enc.eval()
dec.eval()

#%%
criterion = nn.CrossEntropyLoss().to(device)
losses = []
grad_norms = []
correct_symbols = 0
total_symbols = 0
token_to_id, id_to_token = load_vocab(tokensfile)
for d in test_data_loader:
    input = d["image"].to(device)
    # The last batch may not be a full batch
    curr_batch_size = len(input)
    expected = d["truth"]["encoded"].to(device)
    batch_max_len = expected.size(1)
    # Replace -1 with the PAD token
    #expected[expected == -1] = test_data_loader.dataset.token_to_id[PAD]
    enc_low_res, enc_high_res = enc(input)
    # Decoder needs to be reset, because the coverage attention (alpha)
    # only applies to the current image.
    dec.reset(curr_batch_size)
    hidden = dec.init_hidden(curr_batch_size).to(device)
    # Starts with a START token
    sequence = torch.full(
        (curr_batch_size, 1),
        test_data_loader.dataset.token_to_id[START],
        dtype=torch.long,
        device=device,
    )
    # The teacher forcing is done per batch, not symbol
    use_teacher_forcing = False
    decoded_values = []
    for i in range(batch_max_len - 1):
        previous = expected[:, i] if use_teacher_forcing else sequence[:, -1]
        previous = previous.view(-1, 1)
        out, hidden = dec(previous, hidden, enc_low_res, enc_high_res)
        hidden = hidden.detach()
        _, top1_id = torch.topk(out, 1)
        sequence = torch.cat((sequence, top1_id), dim=1)
        decoded_values.append(out)

    decoded_values = torch.stack(decoded_values, dim=2).to(device)

    # decoded_values does not contain the start symbol
    #loss = criterion(decoded_values, expected[:, 1:])

    #losses.append(loss.item())
    #correct_symbols += torch.sum(sequence == expected, dim=(0, 1)).item()
    #total_symbols += expected.numel()

    result = {

    #"loss": np.mean(losses),
    #"correct_symbols": correct_symbols,
    #"total_symbols": total_symbols,
    "shape_out": decoded_values.size(),
    #"shape_in": expected,
    "sequence":sequence,
    }

    filepath = os.path.join("./lg2", d["path"][0].split('/')[-1][0:-4]) + '.lg'

    lg = csv.writer(open(filepath, 'w', newline=''), delimiter = ',')

    start = 0
    end = 0
    labels =  result["sequence"][0].tolist()[1:]
    lg.writerow(['#', str(labels)])
    for id in labels:
        if (id == 117) | (id == 118):
            break

        end += 1
        if (end == len(labels)):
            ids = range(start, end)
            s_ids = []
            for _ in ids:
                s_ids.append(_)

            lg.writerow(['O','X', id_to_token[id] , 1.0] + s_ids)
            start = end
            break

        if (labels[end] != labels[start]):
            ids = range(start, end)
            s_ids = []
            for _ in ids:
                s_ids.append(_)

            lg.writerow(['O','X', id_to_token[id] , 1.0] + s_ids)
            start = end

# %%
c = range(1,2)
for i in c:
    print(i)


##############################################################################
# Copyright 2016 IBM Corp.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
##############################################################################

import unittest
from mock import MagicMock
from pyxcli.mirroring.mirrored_entities import \
    MirroredCachedEntities
from pyxcli.mirroring.errors import NoMirrorDefinedError
from pyxcli.mirroring.recovery_manager import \
    SlaveIsNotConsistentRecoveryException, \
    NoLastReplicatedSnapshotRecoveryException

class CIMCLIResponse(object):

    def __init__(self, generic_cim_objects):
        self.generic_cim_objects = generic_cim_objects

    def __iter__(self):
        return iter(self.generic_cim_objects)

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

    @property
    def as_single_element(self):
        return (self.generic_cim_objects)[0]

    @property
    def as_list(self):
        return self.generic_cim_objects

    def as_dict(self, key_name):
        key_name = str(key_name)
        ret_dict = {}
        if(len(self.generic_cim_objects) > 0):
            single = self.generic_cim_objects[0]
            if (not hasattr(single, key_name)):
                raise Exception('%s is_not_a_valid_key' % key_name)
        for cim_obj in self.generic_cim_objects:
            ret_dict[getattr(cim_obj, key_name)] = cim_obj
        return ret_dict

class TestBaseRecoveryManager(unittest.TestCase):

    __test__ = False

    def setUpRecoveryManager(self):
        self.recovery_manager = None

    def set_mirror_list(self, cvolumes, ccgs):
        self.xcli_client_mock.cmd.mirror_list.return_value = cvolumes

    def set_main_mirror(self, vol1, cg1):
        self.xcli_mirror = vol1
        self.master = 'vol1'
        self.slave = 'vol2'

    def setUp(self):
        self.xcli_client_mock = MagicMock()
        self.mirrored_entities = MirroredCachedEntities(
            self.xcli_client_mock)
        self.setUpMirroredEntities()
        self.setUpRecoveryManager()
        self.recovery_manager.set_action_entities(self.mirrored_entities)

    def create_mock_host(self, name):
        host = MagicMock()
        host.name = name
        host.fc_ports = '11'
        host.cluster = ''
        return host

    def create_mock_resource(self, name, is_master, sync_state):
        vol = MagicMock()
        vol.name = name
        vol.local_peer_name = name
        vol.current_role = 'Master'
        if (not is_master):
            vol.current_role = 'Slave'
        vol.sync_state = sync_state
        vol.sync_type = 'sync_best_effort'
        vol.pool_name = "pool1"
        vol.pool = "pool1"
        vol.__str__.return_value = name + "," + \
            str(vol.current_role) + "," + vol.sync_type + "," + vol.sync_state
        return vol

    def create_mock_vol(self, name, is_master, sync_state):
        vol = self.create_mock_resource(name, is_master, sync_state)
        return vol

    def create_mock_cg(self, name, is_master, sync_state):
        cg = self.create_mock_resource(name, is_master, sync_state)
        return cg

    def create_mock_pool(self, name, snapshot_size=100, used_by_snapshots=0):
        pool = MagicMock()
        pool.name = name
        pool.snapshot_size = snapshot_size
        pool.used_by_snapshots = used_by_snapshots
        pool.__str__.return_value = name + "," + \
            str(pool.snapshot_size) + "," + str(pool.used_by_snapshots)
        pool.as_single.return_value = "blabla"
        return pool

    def setUpMirroredEntities(self):
        vol1 = self.create_mock_vol("vol1", True, "synched")
        self.xcli_mirror = vol1
        vol2 = self.create_mock_vol("vol2", False, "Initializing")
        volumes = [vol1, vol2]
        cvolumes = CIMCLIResponse(volumes)

        cg1 = self.create_mock_cg("cg1", True, "synched")
        cg2 = self.create_mock_cg("cg2", False, "Initializing")
        cgs = [cg1, cg2]
        ccgs = CIMCLIResponse(cgs)

        self.set_main_mirror(vol1, cg1)

        self.set_mirror_list(cvolumes, ccgs)
        self.xcli_client_mock.cmd.vol_list.return_value = cvolumes
        self.xcli_client_mock.cmd.cg_list.return_value = ccgs

        pool1 = self.create_mock_pool("pool1")
        pool2 = self.create_mock_pool("pool2", 0, 0)
        pools = [pool1, pool2]
        cpools = CIMCLIResponse(pools)
        self.xcli_client_mock.cmd.pool_list.return_value = cpools

        host1 = self.create_mock_host('host1')
        chosts = CIMCLIResponse([host1])
        self.xcli_client_mock.cmd.host_list.return_value = chosts

        return_value = CIMCLIResponse([])
        self.xcli_client_mock.cmd.cluster_list.return_value = return_value

    def test_verify_readiness_for_failover(self):
        self.assertRaises(
            NoMirrorDefinedError,
            self.recovery_manager.verify_readiness_for_failover, 'dummy')
        self.assertRaises(SlaveIsNotConsistentRecoveryException,
                          self.recovery_manager.verify_readiness_for_failover,
                          self.slave)
        self.recovery_manager.verify_readiness_for_failover(self.master)

    def test_promote_bad_id(self):
        self.assertRaises(NoMirrorDefinedError,
                          self.recovery_manager.promote, 'dummy')

    def test_failover_resource(self):
        self.recovery_manager.promote(self.master)
        self.recovery_manager.promote(self.slave)

    def test_test_promote_start(self):
        SNAP_TIME_FORMAT = "%Y-%m-%dT%H-%M-%S.0"
        self.recovery_manager.test_promote_start(
            self.master, 'test_snapshot_prefix', SNAP_TIME_FORMAT)

    def test_test_promote_stop(self):
        self.recovery_manager.test_promote_stop(
            self.master, 'test_snapshot_prefix')

    def test_prepare_reverse_replication(self):
        self.recovery_manager.prepare_reverse_replication(self.master)

    def test_reactivate_mirror(self):
        self.recovery_manager.reactivate_mirror(self.master)

    def test_start_async_job(self):
        self.recovery_manager.start_async_job(self.xcli_mirror)

    def test_snap_target_before_possible_override(self):
        self.recovery_manager.snap_target_before_possible_override(
            self.master, 'snapshot_name')

    def test_duplicate_target_snapshot_before_possible_override(self):
        # only to fit in pep8
        rec_man = self.recovery_manager
        self.assertRaises(
            NoLastReplicatedSnapshotRecoveryException,
            rec_man.duplicate_target_snapshot_before_possible_override,
            self.master, 'snapshot_name')

    def test_snapshot_name_format(self):
        self.assertEqual(
            self.recovery_manager._get_snapshot_name(
                'pre', 'resource', 'time'),
            'pre_resource_time')


import pygame
from math import log
from common import get_filename, get_rgb_from_complex, iterations, newton_prime, tolerance

black = 0, 0, 0
red = 255, 0, 0

size = 1024, 1024
width, height = size

r_min = 0.75
r_max = 0.9
i_min = 0.75
i_max = 0.9

def smoothing(zlist: list[complex]) -> float:
	i = len(zlist)
	try:
		z0, z1, root = zlist[-3:]
		ld0 = log(abs(z0 - root))
		ld1 = log(abs(z1 - root))
		correction = (log(tolerance) - ld0)/(ld1 - ld0)
		value = 1 - (i + 0.5*correction)/iterations
		return min(1, max(0, value))
	except ValueError:
		return 1 - i/iterations

def map_to_range(start: float, end: float, fraction: float) -> float:
	"""map 0 to start, 1 to end, and intermediate values to linearly between them"""
	return (end-start) * fraction + start

def get_z_from_coords(coords: tuple[int, int]) -> complex:
	x, y = coords
	real = map_to_range(r_min, r_max, x/width)
	imag = map_to_range(i_min, i_max, (height-y)/height)
	return real+1j*imag

def plotting():
	for x in range(width):
		for y in range(height):
			coords = x, y
			point = get_z_from_coords(coords)
			zlist = newton_prime(point)
			color = get_rgb_from_complex(zlist[-1], smoothing(zlist))
			screen.set_at(coords, [int(255*c) for c in color])
		refresh()

def movie(f, from_val: complex, to_val: complex, frames: int):
	"""Creates a series of images with filename movie/n.png. f is a function which takes a complex number and returns a function to throw into plotting."""
	values = (from_val + i/frames * (to_val - from_val) for i in range(frames))
	for i, value in enumerate(values):
		plotting(f(value))
		pygame.image.save(screen, 'movie/{0}.png'.format(str(i).zfill(5)))

# PYGAME STUFF
pygame.init()
screen = pygame.display.set_mode(size)
refresh = pygame.display.flip
title = 'Newton Prime Fractal'
pygame.display.set_caption(title)

# MAIN

# movie(lambda c: lambda z: (z - (-1)**c)*(z - 1j**c)*(z - (-1j)**c), 0, 4, 300)
plotting()
pygame.image.save(screen, get_filename(f'{title} {r_min + 1j*i_min} to {r_max + 1j*i_max}'))


'''
Created on Nov 13, 2012

@author: raphaelholca
'''

import sys
import cProfile

#Force matplotlib to not use any Xwindows backend; comment out to Plot_Print on screen
import matplotlib
matplotlib.use('Agg')

import matplotlib.pyplot as plt

from CompSim import CompSim
#@profile
def main(argv=None):

        global sim
        sim = CompSim()
        sim.setParam()
        sim.createSynapses()
        sim.createPopulations()
        sim.createSynTracker()
        sim.createConnectionPattern()
        sim.createInput()
        sim.runSimulation()
        sim.pickleValue()
#         cProfile.run('sim.runSimulation()')
#         sim.showResults()
        plt.show()

if __name__ == "__main__":
    sys.exit(main())


import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import os
import copy
import numpy as np
import pickle
import math
from datetime import datetime as date
from glob import glob
import torch
import robosuite.utils.transform_utils as T
from dh_parameters import robot_attributes
from IPython import embed

# epsilon for testing whether a number is close to zero
_EPS = np.finfo(float).eps * 4.0

# some keys are robot-specifig!
BC_skip_state_keys = ['robot0_joint_pos_cos', 'robot0_joint_pos_sin','robot0_joint_vel', 'robot0_proprio-state', 'ball_quat_quat', 'ball_quat', 'robot0_proprio-state', 'object-state']
skip_state_keys = []
#skip_state_keys = [ 'robot0_proprio-state', 'ball_pos', 'robot0_gripper_qvel', 'robot0_gripper_qpos', 'robot0_eef_quat', 'ball_quat', 'robot0_proprio-state', 'object-state']

#c1 = tf.nn.l2_normalize(x[:, :3], axis=-1)
#c2 = tf.nn.l2_normalize(x[:, 3:] - self.dot(c1,x[:, 3:])*c1, axis=-1)
#x = tf.concat([c1,c2], axis=-1)
#self.add_metric(mean_angle_btw_vectors(inputs, self.get_rotated(x)),
#                    name='mean_angular_distance', aggregation='mean')
# ideas for rotation losses from
# https://towardsdatascience.com/better-rotation-representations-for-accurate-pose-estimation-e890a7e1317f
#def euler_loss(y_true, y_pred):
#    dist1 = tf.abs(y_true - y_pred)
#    dist2 = tf.abs(2*np.pi + y_true - y_pred)
#    dist3 = tf.abs(-2*np.pi + y_true - y_pred)
#    loss = tf.where(dist1<dist2, dist1, dist2)
#    loss = tf.where(loss<dist3, loss, dist3)
#    return tf.reduce_mean(loss)
#
#def quaternion_loss(y_true, y_pred):
#    dist1 = tf.reduce_mean(tf.abs(y_true-y_pred), axis=-1)
#    dist2 = tf.reduce_mean(tf.abs(y_true+y_pred), axis=-1)
#    loss = tf.where(dist1<dist2, dist1, dist2)
#    return tf.reduce_mean(loss)
#
#def mean_angle_btw_vectors(v1, v2, eps = 1e-8):
#    dot_product = tf.reduce_sum(v1*v2, axis=-1)
#    cos_a = dot_product / (tf.norm(v1, axis=-1) * tf.norm(v2, axis=-1))
#    cos_a = tf.clip_by_value(cos_a, -1 + eps, 1 - eps)
#    angle_dist = tf.math.acos(cos_a) / np.pi * 180.0
#    return tf.reduce_mean(angle_dist)

def normalize_joints(angles):
    """
    This removes the wrapping from joint angles and ensures joint vals are bt -pi < vals < pi
    angles: np.array of joint angles in radians
    """
    while angles.max() > np.pi:
        angles[angles>np.pi] -= 2*np.pi
    while angles.min() < -np.pi:
        angles[angles<-np.pi] += 2*np.pi
    return angles

def quaternion_from_matrix(matrix):
    """Return quaternion from rotation matrix.
    from: https://github.com/BerkeleyAutomation/autolab_core/blob/master/autolab_core/transformations.py
    >>> R = rotation_matrix(0.123, (1, 2, 3))
    >>> q = quaternion_from_matrix(R)
    >>> numpy.allclose(q, [0.0164262, 0.0328524, 0.0492786, 0.9981095])
    True
    """
    q = np.empty((4,), dtype=np.float64)
    M = np.array(matrix, dtype=np.float64, copy=False)[:4, :4]
    t = np.trace(M)
    if t > M[3, 3]:
        q[3] = t
        q[2] = M[1, 0] - M[0, 1]
        q[1] = M[0, 2] - M[2, 0]
        q[0] = M[2, 1] - M[1, 2]
    else:
        i, j, k = 0, 1, 2
        if M[1, 1] > M[0, 0]:
            i, j, k = 1, 2, 0
        if M[2, 2] > M[i, i]:
            i, j, k = 2, 0, 1
        t = M[i, i] - (M[j, j] + M[k, k]) + M[3, 3]
        q[i] = t
        q[j] = M[i, j] + M[j, i]
        q[k] = M[k, i] + M[i, k]
        q[3] = M[k, j] - M[j, k]
    q *= 0.5 / math.sqrt(t * M[3, 3])
    return q

def np_dh_transform(theta, d, a, alpha):
    bs = theta.shape[0]
    T = np.zeros((bs,4,4))
    T[:,0,0] = T[:,0,0] +  np.cos(theta)
    T[:,0,1] = T[:,0,1] + -np.sin(theta)*np.cos(alpha)
    T[:,0,2] = T[:,0,2] +  np.sin(theta)*np.sin(alpha)
    T[:,0,3] = T[:,0,3] +  a*np.cos(theta)
    T[:,1,0] = T[:,1,0] +    np.sin(theta)
    T[:,1,1] = T[:,1,1] +    np.cos(theta)*np.cos(alpha)
    T[:,1,2] = T[:,1,2] +   -np.cos(theta)*np.sin(alpha)
    T[:,1,3] = T[:,1,3] +  a*np.sin(theta)
    T[:,2,1] = T[:,2,1] +  np.sin(alpha)
    T[:,2,2] = T[:,2,2] +   np.cos(alpha)
    T[:,2,3] = T[:,2,3] +  d
    T[:,3,3] = T[:,3,3] +  1.0
    return T

def torch_dh_transform(theta, d, a, alpha, device):
    bs = theta.shape[0]
    T = torch.zeros((bs,4,4), device=device)
    T[:,0,0] = T[:,0,0] +  torch.cos(theta)
    T[:,0,1] = T[:,0,1] + -torch.sin(theta)*torch.cos(alpha)
    T[:,0,2] = T[:,0,2] +  torch.sin(theta)*torch.sin(alpha)
    T[:,0,3] = T[:,0,3] +  a*torch.cos(theta)
    T[:,1,0] = T[:,1,0] +  torch.sin(theta)
    T[:,1,1] = T[:,1,1] +   torch.cos(theta)*torch.cos(alpha)
    T[:,1,2] = T[:,1,2] +   -torch.cos(theta)*torch.sin(alpha)
    T[:,1,3] = T[:,1,3] +  a*torch.sin(theta)
    T[:,2,1] = T[:,2,1] +  torch.sin(alpha)
    T[:,2,2] = T[:,2,2] +   torch.cos(alpha)
    T[:,2,3] = T[:,2,3] +  d
    T[:,3,3] = T[:,3,3] +  1.0
    return T

def quaternion_matrix(quaternion):
    """Return homogeneous rotation matrix from quaternion.
    >>> R = quaternion_matrix([0.06146124, 0, 0, 0.99810947])
    >>> numpy.allclose(R, rotation_matrix(0.123, (1, 0, 0)))
    True
    """
    q = np.array(quaternion[:4], dtype=np.float64, copy=True)
    nq = np.dot(q, q)
    if nq < _EPS:
        return np.identity(4)
    q *= math.sqrt(2.0 / nq)
    q = np.outer(q, q)
    return np.array(
        (
            (
                1.0 - q[1, 1] - q[2, 2],
                q[0, 1] - q[2, 3],
                q[0, 2] + q[1, 3],
                0.0,
            ),
            (
                q[0, 1] + q[2, 3],
                1.0 - q[0, 0] - q[2, 2],
                q[1, 2] - q[0, 3],
                0.0,
            ),
            (
                q[0, 2] - q[1, 3],
                q[1, 2] + q[0, 3],
                1.0 - q[0, 0] - q[1, 1],
                0.0,
            ),
            (0.0, 0.0, 0.0, 1.0),
        ),
        dtype=np.float64,
    )

def so3_relative_angle(R1, R2, cos_angle: bool = False):
    """
    from: https://pytorch3d.readthedocs.io/en/latest/_modules/pytorch3d/transforms/so3.html#so3_relative_angle

    Calculates the relative angle (in radians) between pairs of
    rotation matrices `R1` and `R2` with `angle = acos(0.5 * (Trace(R1 R2^T)-1))`

    .. note::
        This corresponds to a geodesic distance on the 3D manifold of rotation
        matrices.

    Args:
        R1: Batch of rotation matrices of shape `(minibatch, 3, 3)`.
        R2: Batch of rotation matrices of shape `(minibatch, 3, 3)`.
        cos_angle: If==True return cosine of the relative angle rather than
                   the angle itself. This can avoid the unstable
                   calculation of `acos`.

    Returns:
        Corresponding rotation angles of shape `(minibatch,)`.
        If `cos_angle==True`, returns the cosine of the angles.

    Raises:
        ValueError if `R1` or `R2` is of incorrect shape.
        ValueError if `R1` or `R2` has an unexpected trace.
    """
    R12 = torch.bmm(R1, R2.permute(0, 2, 1))
    return so3_rotation_angle(R12, cos_angle=cos_angle)

def so3_rotation_angle(R, eps: float = 1e-4, cos_angle: bool = False):
    """
    Calculates angles (in radians) of a batch of rotation matrices `R` with
    `angle = acos(0.5 * (Trace(R)-1))`. The trace of the
    input matrices is checked to be in the valid range `[-1-eps,3+eps]`.
    The `eps` argument is a small constant that allows for small errors
    caused by limited machine precision.

    Args:
        R: Batch of rotation matrices of shape `(minibatch, 3, 3)`.
        eps: Tolerance for the valid trace check.
        cos_angle: If==True return cosine of the rotation angles rather than
                   the angle itself. This can avoid the unstable
                   calculation of `acos`.

    Returns:
        Corresponding rotation angles of shape `(minibatch,)`.
        If `cos_angle==True`, returns the cosine of the angles.

    Raises:
        ValueError if `R` is of incorrect shape.
        ValueError if `R` has an unexpected trace.
    """

    N, dim1, dim2 = R.shape
    if dim1 != 3 or dim2 != 3:
        raise ValueError("Input has to be a batch of 3x3 Tensors.")

    rot_trace = R[:, 0, 0] + R[:, 1, 1] + R[:, 2, 2]

    if ((rot_trace < -1.0 - eps) + (rot_trace > 3.0 + eps)).any():
        raise ValueError("A matrix has trace outside valid range [-1-eps,3+eps].")

    # clamp to valid range
    rot_trace = torch.clamp(rot_trace, -1.0, 3.0)

    # phi ... rotation angle
    phi = 0.5 * (rot_trace - 1.0)

    if cos_angle:
        return phi
    else:
        # pyre-fixme[16]: `float` has no attribute `acos`.
        return phi.acos()

def mean_angle_btw_vectors(v1, v2, eps = 1e-4):
    # https://towardsdatascience.com/better-rotation-representation    s-for-accurate-pose-estimation-e890a7e1317f
    dot_product = torch.sum(v1*v2, axis=-1)
    cos_a = dot_product / (torch.norm(v1, dim=-1) * torch.norm(v2,     dim=-1))
    cos_a = torch.clamp(cos_a, -1 + eps, 1 - eps)
    angle_dist = torch.acos(cos_a)
    return torch.mean(angle_dist)

class robotDH():
    def __init__(self, robot_name, device='cpu'):
        self.device = device
        self.robot_name = robot_name
        self.npdh = robot_attributes[self.robot_name]
        self.tdh = {}
        for key, item in self.npdh.items():
            self.tdh[key] = torch.FloatTensor(item).to(self.device)

    def np_angle2ee(self, base_matrix, angles):
        """
            convert np joint angle to end effector for for ts,angles (in radians)
        """
        # ts, bs, feat
        ts, fs = angles.shape
        ee_pred = np.zeros((ts,7))
        # TODO join the time/batch so i don't have to loop this
        #_T = self.np_dh_transform(0, angles[:,0])
        #_T = np.array([[1,0,0,0],[0,-1,0,0],[0,0,-1,0],[0,0,0,1]], dtype=np.float)
        _T = base_matrix
        for _a in range(fs):
            _T1 = self.np_dh_transform(_a, angles[:,_a])
            _T = np.matmul(_T, _T1)
            #print(_a, _T[0, :3, 3])
            #print(T.mat2euler(_T[0]))
        return _T

    def torch_angle2ee(self, base_matrix, angles):
        """
            convert joint angle to end effector for reacher for ts,bs,f
        """
        # ts, bs, feat
        ts, fs = angles.shape
        ee_pred = torch.zeros((ts,4,4)).to(self.device)
        _T = base_matrix
        for _a in range(fs):
            _T1 = self.torch_dh_transform(_a, angles[:,_a])
            _T = torch.matmul(_T, _T1)
        return _T

    def np_dh_transform(self, dh_index, angles):
        theta = self.npdh['DH_theta_sign'][dh_index]*angles+self.npdh['DH_theta_offset'][dh_index]
        d = self.npdh['DH_d'][dh_index]
        a = self.npdh['DH_a'][dh_index]
        alpha = self.npdh['DH_alpha'][dh_index]
        return np_dh_transform(theta, d, a, alpha)

    def torch_dh_transform(self, dh_index, angles):
        theta = self.tdh['DH_theta_sign'][dh_index]*angles+self.tdh['DH_theta_offset'][dh_index]
        d = self.tdh['DH_d'][dh_index]
        a = self.tdh['DH_a'][dh_index]
        alpha = self.tdh['DH_alpha'][dh_index]
        return torch_dh_transform(theta, d, a, alpha, self.device)

# I"M NOT CONVINCED THESE WORK -
#def sincos2angle(sin_theta, cos_theta, use_numpy=False):
#    """ robosuite outputs the joint angle in sin(theta) cos(angle)
#    This function converts it to angles in radians """
#    if not use_numpy:
#        return torch.arctan(sin_theta/cos_theta)
#    else:
#        return np.arctan(sin_theta/cos_theta)
#
#def angle2sincos(theta, use_numpy=False):
#    """ convert an angle to sin(angle) cos(theta)  in radians """
#    if not use_numpy:
#        return torch.sin(theta), torch.cos(theta)
#    else:
#        return np.sin(theta), np.cos(theta)
#

def get_data_norm_params(data, device='cpu'):
    input_size = data.shape[1]
    train_mean = torch.FloatTensor([data[:,x].mean() for x in range(input_size)]).to(device)
    train_std = torch.FloatTensor([data[:,x].std() for x in range(input_size)]).to(device)
    return train_mean, train_std

def load_robosuite_data(data_file, random_state):
    # each episode is 500 steps long
    # data is
    sq = np.load(data_file, allow_pickle=True)

    input_size = len(sq[0][0]['robot0_joint_pos_sin']) + len(sq[0][0]['robot0_joint_pos_cos'])
    sq_fmt = []
    for episode in sq:
        ep_fmt = [np.hstack((x['robot0_joint_pos_sin'], x['robot0_joint_pos_cos'], x['robot0_eef_pos'], x['robot0_joint_pos'])) for x in episode]
        #ep_fmt = [np.hstack((x['robot0_joint_pos_sin'], x['robot0_joint_pos_cos'], x['robot0_eef_pos'])) for x in episode]
        sq_fmt.append(ep_fmt)
    # t, batch, features
    sq_fmt = np.array(sq_fmt).swapaxes(0,1)
    n_traces = sq_fmt.shape[1]

    # EE is at end of values
    indexes = np.arange(n_traces, dtype=np.int)
    random_state.shuffle(indexes)
    ttb = max([1,int(n_traces*.15)])
    t_inds = indexes[ttb:]
    v_inds = indexes[:ttb]
    data = {'valid':{}, 'train':{}}
    data['train']['input'] = sq_fmt[:-1,   t_inds, :input_size]
    data['train']['target'] = sq_fmt[1:,   t_inds, :input_size]
    data['train']['ee_target'] = sq_fmt[1:,t_inds, input_size:input_size+3]
    data['train']['jt_target'] = sq_fmt[1:,t_inds, input_size+3:]

    data['valid']['input'] = sq_fmt[:-1,    v_inds, :input_size]
    data['valid']['target'] = sq_fmt[1:,    v_inds, :input_size]
    data['valid']['ee_target']= sq_fmt[1:, v_inds, input_size:input_size+3]
    data['valid']['jt_target'] = sq_fmt[1:, v_inds, input_size+3:]
    return data

def plot_losses(loss_path):
    losses = np.load(loss_path)
    plt.figure()
    for phase, ll in [ ('valid',losses['valid']),('train', losses['train'])]:
        plt.plot(ll[1:,0], ll[1:,1], label=phase, marker='o')
    plt.title('losses')
    plt.legend()
    fname = loss_path.replace('.npz', '.png')
    print("saving loss image: {}".format(fname))
    plt.savefig(fname)
    plt.close()

def find_latest_checkpoint(basedir):
    assert os.path.isdir(basedir)
    search = os.path.join(basedir, '*.pt')
    print('searching {} for models'.format(search))
    found_models = sorted(glob(search))
    print('found {} models'.format(len(found_models)))
    # this is the latest model
    load_path = found_models[-1]
    print('using most recent - {}'.format(load_path))
    return load_path

def create_results_dir(exp_name, results_dir='results'):
    today = date.today()
    today_str = today.strftime("%y-%m-%d")
    exp_cnt = 0
    savebase = os.path.join(results_dir, '%s_%s_%02d'%(today_str, exp_name, exp_cnt))
    while len(glob(os.path.join(savebase, '*.pt'))):
         exp_cnt += 1
         savebase = os.path.join(results_dir, '%s_%s_%02d'%(today_str, exp_name, exp_cnt))
    if not os.path.exists(savebase):
        os.makedirs(savebase)
    if not os.path.exists(os.path.join(savebase, '__python')):
        os.makedirs(os.path.join(savebase, '__python'))
        os.system('cp *.py %s/__python'%savebase)
    return savebase

def seed_everything(seed=1234):
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    np.random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)

if __name__ == '__main__':
    #fb = 'results/21-01-21_v1_lstm_ee_03'
    #pt_latest = find_latest_checkpoint(fb)
    #plot_losses(pt_latest.replace('.pt', '_losses.npz'))
    plot_losses('results/21-01-21_v1_lstm_ee_03/model_0000023136_losses.npz')


# class tutorial

class Employee:

    num_emps = 0
    raise_amount = 1.04

    def __init__(self, first, last, pay):
        self.first = first
        self.last = last
        self.pay = pay
        self.email = first + '.' + last + '@gmail.com'

        Employee.num_emps += 1

    def fullname(self):
        return '{} {}'.format(self.first, self.last)

    def apply_raise(self):
        self.pay = int(self.pay * self.raise_amount)

    @classmethod
    def set_raise_amt(cls, amount):
        cls.raise_amt = amount

    @classmethod
    def from_string(cls, emp_str):
        first, last, pay = emp_str.split('-')
        return cls(first, last, pay)

    @staticmethod
    def is_workday(day):
        if day.weekday() == 5 or day.weekday() == 6:
            return False
        return True

class Developer(Employee):
    raise_amt = 1.10

    def __init__(self, first, last, pay, prog_lang):
        super().__init__(first, last, pay)
        self.prog_lang = prog_lang

emp1 = Employee('Flavio', 'Ryu', 1500)
emp2 = Employee('Giorgio', 'Moroder', 2000)

emp_str1 = 'Alberto-Digiulio-1700'
emp_str2 = 'Gianna-Falconi-1600'
emp_str3 = 'Maria-Lorenzini-1900'

new_emp1 = Employee.from_string(emp_str1)

print(new_emp1.email)
print(new_emp1.pay)

import datetime
my_date = datetime.date(2020, 6, 15)

print(Employee.is_workday(my_date))

dev1 = Developer('Gianni', 'Morandi', 780, 'Python')
dev2 = Developer('Gesù', 'Cristo', 0, 'Java')

print(dev1.prog_lang)
print(dev2.prog_lang)


from itertools import product
import numpy as np

H, W = map(int, input().split())
grid = []
for _ in range(H):
    array = list(input())
    grid.append(array)
b = np.array(grid)
b = np.where(b == ".", 1, 0)

L = np.zeros((H, W), dtype=np.int)
R = np.zeros((H, W), dtype=np.int)
D = np.zeros((H, W), dtype=np.int)
U = np.zeros((H, W), dtype=np.int)

for i in range(1, H):
    L[i] += (L[i-1]+1)*b[i-1]

for i in range(H-2, -1, -1):
    R[i] += (R[i+1]+1)*b[i+1]

for i in range(1, W):
    U[:, i] += (U[:, i-1]+1)*b[:, i-1]

for i in range(W-2, -1, -1):
    D[:, i] += (D[:, i+1]+1)*b[:, i+1]

answer = ((L+R+U+D)*b).max()

print(answer+1)


## CodeEval
## String List

import sys

def next_n(n, b):
    if n[-1] < b-1:
        return n[:-1] + [n[-1]+1]
    else:
        return next_n(n[:-1], b) + [0]

f = open(sys.argv[1], 'r')

for line in f:

    line = line.strip()
    if line:
        line = line.split(',')
        size = int(line[0])
        alpha = ''.join(sorted(set(line[1])))

        b = len(alpha)
        n = [0] * size
        sys.stdout.write(''.join([alpha[i] for i in n]))
        for i in xrange(1, b**size):
            n = next_n(n, b)
            sys.stdout.write(',' + ''.join([alpha[i] for i in n]))
        sys.stdout.write('\n')

f.close()


"""OAuth2 scope generation for WebAPI resources."""

import logging
import threading
from collections import defaultdict
from importlib import import_module

from django.conf import settings
from django.core.exceptions import ImproperlyConfigured
from django.utils.translation import gettext as _

from djblets.extensions.manager import get_extension_managers
from djblets.extensions.signals import (extension_enabled,
                                        extension_disabled)
from djblets.webapi.resources.mixins.oauth2_tokens import (
    ResourceOAuth2TokenMixin)

logger = logging.getLogger(__name__)

_enable_lock = threading.Lock()
_scopes = None

def get_scope_dictionary():
    """Return the scope dictionary.

    This method requires :setting:`WEB_API_ROOT_RESOURCE` setting to point to
    an instance of the root resource for your WebAPI. Optionally, the
    :setting:`WEB_API_SCOPE_DICT_CLASS` (defaulting to
    :py:class:`djblets.webapi.oauth2_scopes.WebAPIScopeDictionary`) can be set
    to specialize scope generation.

    If the :py:mod:`djblets.extensions` app is being used, then the
    :setting:`WEB_API_SCOPE_DICT_CLASS` setting should be set to
    :py:class:`djblets.webapi.oauth2_scopes.ExtensionEnabledWebAPIScopeDictionary`.

    Returns:
        WebAPIScopeDictionary:
        The scope dictionary.
    """
    global _scopes

    if _scopes is None:
        scopes_cls_name = getattr(
            settings, 'WEB_API_SCOPE_DICT_CLASS',
            'djblets.webapi.oauth2_scopes.WebAPIScopeDictionary')

        scopes_cls_mod, scopes_cls_attr = scopes_cls_name.rsplit('.', 1)

        try:
            scopes_cls = getattr(import_module(scopes_cls_mod),
                                 scopes_cls_attr)
        except (AttributeError, ImportError) as e:
            raise ImproperlyConfigured(
                'settings.WEB_API_SCOPE_DICT_CLASS %r could not be imported: '
                '%s'
                % (scopes_cls_name, e)
            )

        try:
            root_resource_path = getattr(settings, 'WEB_API_ROOT_RESOURCE')
        except AttributeError:
            raise ImproperlyConfigured(
                'settings.WEB_API_ROOT_RESOURCE is required.'
            )

        root_resource_mod, root_resource_attr = root_resource_path.rsplit('.',
                                                                          1)

        try:
            root_resource = getattr(import_module(root_resource_mod),
                                    root_resource_attr)
        except (AttributeError, ImportError) as e:
            raise ImproperlyConfigured(
                'settings.WEB_API_ROOT_RESOURCE %r could not be imported: %s'
                % (root_resource_path, e)
            )

        _scopes = scopes_cls(root_resource)

    return _scopes

class WebAPIScopeDictionary(object):
    """A Web API scope dictionary.

    This class knows how to build a list of available scopes from the WebAPI
    resource tree at runtime.

    By default, it will only have to walk the API tree once, after which the
    value can be cached.
    """

    def __init__(self, root_resource):
        """Initialize the scope dictionary.

        Args:
            root_resource (djblets.webapi.resources.base.WebAPIResource):
                The root resource to walk in order to build scopes.
        """
        self.resource_trees = {root_resource}
        self._update_lock = threading.Lock()
        self._scope_dict = {}

    @property
    def scope_dict(self):
        """The dictionary of scopes defined by this dictionary.

        The value is cached so that it will only be recomputed when the
        dictionary is updated.
        """
        if not self._scope_dict:
            with self._update_lock:
                if not self._scope_dict:
                    self._walk_resources(self.resource_trees)

                    assert self._scope_dict

        return self._scope_dict

    def keys(self):
        """Iterate through all keys in the dictionary.

        This is used by oauth2_provider when on Python 3.x to get the list
        of scope keys.

        Yields:
            unicode:
            The key for each scope.
        """
        return self.scope_dict.keys()

    def clear(self):
        """Clear all scopes from the dictionary.

        The next attempt at fetching scopes will repopulate the dictionary
        from scratch.
        """
        self._scope_dict.clear()

    def _walk_resources(self, resources):
        """Traverse the given resource trees and add the appropriate scopes.

        Args:
            resources (list of djblets.webapi.resources.base.WebAPIResource):
                The resources to generate scopes for. The children of these
                resources will also be traversed.

        Returns:
            list of unicode:
            The list of scopes added by walking the given resources.
        """
        for resource in resources:
            self._walk_resources(resource.list_child_resources)
            self._walk_resources(resource.item_child_resources)

            scope_to_methods = defaultdict(list)

            if not isinstance(resource, ResourceOAuth2TokenMixin):
                logging.warning(
                    'Resource %r does not inherit from '
                    'djblets.webapi.resources.mixins.oauth2_tokens.'
                    'ResourceOAuth2TokenMixin: it will not be accessible with '
                    'OAuth2 tokens. It is recommended that all your resources '
                    'inherit from a base class that includes this mixin.',
                    type(resource),
                )
                continue

            if not resource.oauth2_token_access_allowed:
                continue

            for method in resource.allowed_methods:
                try:
                    suffix = resource.HTTP_SCOPE_METHOD_MAP[method]
                except KeyError:
                    logger.error('Unknown HTTP method %s not present in '
                                 'HTTP_SCOPE_METHOD_MAP: a scope will not be '
                                 'generated for this method.',
                                 method)
                    continue

                scope_to_methods[suffix].append(method)

            for suffix, methods in scope_to_methods.items():
                scope_name = '%s:%s' % (resource.scope_name, suffix)

                self._scope_dict[scope_name] = (
                    _('Ability to perform HTTP %(methods)s on the %(name)s '
                      'resource')
                    % {
                        'methods': ', '.join(sorted(methods)),
                        'name': resource.name,
                    }
                )

    def __getitem__(self, key):
        """Return the description of a given scope.

        Args:
            key (unicode):
                The scope's key.

        Returns:
            unicode:
            The scope's description.

        Raises:
            KeyError:
                The scope's key was not in the dictionary.
        """
        return self.scope_dict[key]

    def __contains__(self, key):
        """Return whether the dictionary has a particular scope.

        Args:
            key (unicode):
                The scope's key.

        Returns:
            bool:
            ``True`` if the scope is in the dictionary.
        """
        return key in self.scope_dict

    def __repr__(self):
        """Return a string representation of this object.

        Returns:
            unicode:
            A string representation of this object.
        """
        return '%s(%r)' % (type(self).__name__,
                           list(self.scope_dict.keys()))

class ExtensionEnabledWebAPIScopeDictionary(WebAPIScopeDictionary):
    """A Web API scopes dictionary that supports extensions.

    This scope dictionary is only required if your app is using the
    :py:mod:`djblets.extensions` app.
    """

    def __init__(self, *args, **kwargs):
        """Initialize the scope dictionary.

        This adds signal handlers to ensure the dictionary stays up to date
        when extensions are initialized and uninitialized.

        Args:
            *args (tuple):
                Positional arguments to pass to the parent class.

            **kwargs (dict):
                Keyword arguments to pass to the the parent class.
        """
        super(ExtensionEnabledWebAPIScopeDictionary, self).__init__(*args,
                                                                    **kwargs)

        extension_enabled.connect(self._on_extension_enabled)
        extension_disabled.connect(self._on_extension_disabled)

        for manager in get_extension_managers():
            for extension in manager.get_enabled_extensions():
                self._on_extension_enabled(extension=extension)

    def _on_extension_enabled(self, extension, **kwargs):
        """Traverse an extensions resource trees when it is enabled.

        Args:
            extension (djblets.extensions.extension.Extension):
                The extension class that was enabled.

            **kwargs (dict):
                Ignored keyword arguments from the signal.
        """
        if extension.resources:
            self.resource_trees.update(extension.resources)
            self.clear()

    def _on_extension_disabled(self, extension, **kwargs):
        """Remove an extensions scopes when it is uninitialized.

        Args:
            ext_class (type):
                The extension class that was uninitialized.

            **kwargs (dict):
                Ignored keyword arguments from the signal.
        """
        if extension.resources:
            self.resource_trees.difference_update(extension.resources)
            self.clear()

def enable_web_api_scopes(*args, **kwargs):
    """Enable WebAPI scopes.

    The ``oauth2_provider.settings.oauth2_settings`` object will be patched so
    that scopes are correctly cached. The cache will be lazily updated after
    the scopes are updated.

    Args:
        *args (tuple):
            Ignored positional arguments.

        **kwargs (dict):
            Ignored keyword arguments.
    """
    with _enable_lock:
        if not isinstance(settings.OAUTH2_PROVIDER['SCOPES'],
                          WebAPIScopeDictionary):
            settings.OAUTH2_PROVIDER['SCOPES'] = get_scope_dictionary()


from sudoku import Sudoku
from userinterface import UserInterface

class Controller:
    """
    Controller for the Sudoku game, acting as an intermediary between the UserInterface and the Sudoku game logic.

    Attributes:
        game (Sudoku): An instance of the Sudoku game.
        ui (UserInterface): The graphical user interface for the game.

    Methods:
        start(): Initializes and starts the game.
        handle_cell_input(row, col, value): Handles input from the user into a Sudoku cell.
        solve_puzzle(): Solves the current puzzle.
    """

    def __init__(self, game: Sudoku, ui: UserInterface) -> None:
        """
        Initializes the Controller with the game and UI instances.

        Args:
            game (Sudoku): An instance of the Sudoku game.
            ui (UserInterface): The graphical user interface for the game.
        """
        self.game = game
        self.ui = ui

        # Set up the UI to use this controller for handling events
        self.ui.set_controller(self)

    def start(self) -> None:
        """
        Starts the Sudoku game.
        """
        self.game.start_game(difficulty='medium')
        self.ui.start()

    def handle_cell_input(self, row: int, col: int, value: int) -> bool:
        """
        Handles input from the user into a cell of the Sudoku grid.

        Args:
            row (int): The row index of the cell.
            col (int): The column index of the cell.
            value (int): The value entered into the cell.
        """
        if self.game.grid.is_valid_move(row, col, value):
            self.game.grid.cells[row][col].value = value
            if self.game.check_solution():
                self.ui.show_message("Congratulations! You solved the puzzle.")
            return True
        else:
            self.ui.show_message("Invalid move.")
            return False

    def solve_puzzle(self) -> None:
        """
        Solves the current puzzle and updates the UI with the solution.
        """
        if self.game.grid.solve():
            self.ui.update_grid()
            self.ui.show_message("Puzzle solved.")
        else:
            self.ui.show_message("No solution exists for this puzzle.")

    def update_cell(self, row:int, col:int, value:int) -> None:
        """
        Updates the value of a cell and checks if the move is valid.

        Args:
            row (int): The row index of the cell.
            col (int): The column index of the cell.
            value (int): The value to be placed in the cell.
        """
        if value in range(1, 10) and self.game.grid.is_valid_move(row, col, value):
            self.game.grid.cells[row][col].value = value
            if self.game.grid.is_solved_correctly():
                self.ui.show_congratulations()
        else:
            self.ui.show_error("Invalid move")

    def show_congratulations(self)-> None:
        self.ui.show_message("Congratulations!", "You solved the Sudoku puzzle!")

    def show_error(self, message: str)-> None:
        self.ui.show_message("Error", message)