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


from PIL import Image, ImageDraw, ImageFont

# Create Image object with the input image

try:
    image = Image.open('background.png')
    print("All done loading!")
except Exception as e:
    print("Something went wrong", e)

# Initialize the drawing context
# The image object as background

draw = ImageDraw.Draw(image)

# Create font object with the font file and specify
# desired size

font = ImageFont.truetype('UDDigiKyokashoN-B.ttc', size=46)

# Tuple that represents the position of the message

(x, y) = (100, 75)

message = "宜しくお願い致します！"
black_color = "rgb(100, 0, 100)" # Black

# draw the image on the background

draw.text((x,y), message, fill=black_color, font=font, align="center")

(x2, y2) = (150, 150)
name = "ケーシーさん"
white_color = "rgb(255, 255, 225)" # White
draw.text((x2, y2), name, fill=white_color, font=font, align="center")

# Save the edited image!

image.save('greeting_card.png')


""" Test file for the operations of the pipeline.

For now, all operations will be tested in this file. It is forseen though, that
the file will grow to a point where it is no longer practical and the
testcases need to be split into different files. How this will be done is not
entirely clear right now.

@author Max Weise
created 23.07.2022
"""

from src.pipetex import operations
from src.pipetex import enums, exceptions
from tests import util_functions

import os
import os.path
import pytest
import shutil
from typing import Optional, Tuple

# === Fixtures ===
FILE_PREFIX = "[piped]"

@pytest.fixture
def simple_testfile():
    """Generates a simple tex file to run tests with."""
    file_name = "test_file"
    with open(f"{file_name}.tex", "w+", encoding="utf-8") as f:
        f.write(
            "\\documentclass[a4paper, draft, 12pt]{article}\n"
            "\\begin{document}\n"
            "This is a testfile\n"
            "\\end{document}\n"
        )

    yield file_name

    util_functions.remove_files(file_name)

@pytest.fixture
def simple_testfile_no_draft():
    """Generates a simple tex file to run tests with."""
    file_name = "test_file"
    with open(f"{file_name}.tex", "w+", encoding="utf-8") as f:
        f.write(
            "\\documentclass[a4paper, 12pt]{article}\n"
            "\\begin{document}\n"
            "This is a testfile\n"
            "\\end{document}\n"
        )

    yield file_name

    util_functions.remove_files(file_name)

@pytest.fixture
def dirty_working_dir():
    """Create some auxiliary files to 'pollute' the working dir."""
    file_name = f"{FILE_PREFIX}_test_file"
    ex = ['tex', 'aux', 'pdf', 'glo', 'bib']
    util_functions.write_empty_file(file_name, ex)

    yield file_name

    util_functions.remove_files(file_name)

    util_functions.remove_files("not_removed")

    if "DEPLOY" in os.listdir():
        shutil.rmtree("DEPLOY")

@pytest.fixture
def config_dict():
    return {"file_prefix": "[piped]"}

# === Type Def ===
Monad = Tuple[bool, Optional[exceptions.InternalException]]

def test_remove_draft_option(simple_testfile, config_dict):
    """Tests that the draft option from the classdefinition gets removed. """
    file_name = simple_testfile

    success, error = operations.remove_draft_option(file_name, config_dict)

    with open(f"{file_name}.tex", "r") as f:
        lines_in_testfile: list[str] = [line.rstrip() for line in f]

    assert success
    assert not error
    assert "draft" not in lines_in_testfile[0]
    assert len(lines_in_testfile) == 4

def test_remove_draft_option_fileNotFound(config_dict):
    """Tests that the draft option from the classdefinition gets removed. """
    file_name = "not_a_file"
    success, error = operations.remove_draft_option(file_name, config_dict)

    assert not success
    assert error
    assert type(error.severity_level) == int
    assert 20 < error.severity_level <= 30

def test_remove_draft_option_draftOptionNotFound(simple_testfile_no_draft,
                                                 config_dict):
    """Tests that the draft option from the classdefinition gets removed. """
    file_name = simple_testfile_no_draft

    with open(f"{file_name}.tex", "r") as f:
        lines_befor_runing_test: list[str] = [line.rstrip() for line in f]

    success, error = operations.remove_draft_option(
        simple_testfile_no_draft,
        config_dict
    )

    with open(f"{file_name}.tex", "r") as f:
        lines_after_runing_test: list[str] = [line.rstrip() for line in f]

    assert lines_after_runing_test[1:] == lines_befor_runing_test[1:]
    assert not success
    assert error
    assert type(error.severity_level) == int
    assert error.severity_level <= 10

def test_copy_file(simple_testfile, config_dict):
    """Tests that the working file is copied correctly."""
    file_name = simple_testfile
    new_file_name = f"{FILE_PREFIX}_{file_name}"

    success, error = operations.copy_latex_file(file_name, config_dict)
    expectedNewName: str = config_dict[enums.ConfigDictKeys.NEW_NAME.value]

    assert success
    assert not error
    assert f"{new_file_name}.tex" in os.listdir()
    assert f"{new_file_name}" == expectedNewName
    assert f"{file_name}.tex" in os.listdir()

def test_copy_file_fileNotFound(config_dict):
    """Tests that the working file is copied correctly."""

    success, error = operations.copy_latex_file("Not a testfile", config_dict)

    assert not success
    assert error
    assert type(error.severity_level) == int
    assert 20 < error.severity_level <= 30

def test_clean_working_dir(dirty_working_dir, config_dict):
    """Tests that the working dir gets cleaned properly."""

    # write a testfile which should not be removed by the method
    not_removed = "not_removed"
    with open(f"{not_removed}.tex", "w+"):
        pass

    test_file = dirty_working_dir

    success, error = operations.clean_working_dir(
        test_file,
        config_dict,
        new_file_name=test_file
    )
    current_dir = os.listdir()

    assert success
    assert not error

    for ex in ['tex', 'aux', 'pdf', 'glo', 'bib']:
        assert f"{test_file}.{ex}" not in current_dir
    assert f"{not_removed}.tex" in current_dir
    assert "DEPLOY" in current_dir
    assert f"{test_file}.pdf" in os.listdir("./DEPLOY")

def test_clean_working_dir_FolderAlreadyExists(dirty_working_dir, config_dict):
    """Tests that the function executes properly."""

    test_file = dirty_working_dir
    # Create the folder as part of the env setup
    os.mkdir("DEPLOY")

    success, error = operations.clean_working_dir(
        dirty_working_dir,
        config_dict,
        new_file_name=test_file
    )
    current_dir = os.listdir()

    assert not success
    assert error
    assert type(error.severity_level) == int
    assert error.severity_level <= 10

    for ex in ['tex', 'aux', 'pdf', 'glo', 'bib']:
        assert f"{test_file}.{ex}" not in current_dir
    assert "DEPLOY" in current_dir
    assert f"{test_file}.pdf" in os.listdir("./DEPLOY")

def test_clean_working_dir_FileInFolderAlreadyExists(dirty_working_dir,
                                                     config_dict):
    """Tests that the function executes properly."""

    # Create folder and file
    os.mkdir("DEPLOY")
    with open("DEPLOY/test_file.pdf", "w+", encoding="utf-8"):
        pass

    success, error = operations.clean_working_dir(
        dirty_working_dir,
        config_dict
    )

    assert not success
    assert error
    assert type(error.severity_level) == int
    assert error.severity_level <= 10

    files_in_deploy = os.listdir("./DEPLOY")
    print(files_in_deploy)
    assert len(files_in_deploy) == 2

def test_clean_working_dir_PdfFileNotFound(dirty_working_dir, config_dict):
    """Tests that the correct error type is thrown."""
    test_file = dirty_working_dir

    os.remove(f"{test_file}.pdf")
    success, error = operations.clean_working_dir(
        dirty_working_dir,
        config_dict
    )

    assert not success
    assert error
    assert type(error.severity_level) == int
    assert 20 < error.severity_level <= 30


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 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()


'''
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 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 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.")


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


#!/usr/bin/env python
# coding: utf-8

# In[1]:

import pandas as pd
import numpy as np
import math
import os

# In[2]:

##創資料夾
def mkdir(path):
    folder = os.path.exists(path)
    #判斷結果
    if not folder:
        #如果不存在，則建立新目錄
        os.makedirs(path)
        print('-----建立成功-----')
    else:
        #如果目錄已存在，則不建立，提示目錄已存在
        print(path+'目錄已存在')

## 資料匯出
def OutputCSV(df):
    dir_path = 'Rank'
    csv_name = 'Rank'
    Result = dir_path + '/' + csv_name + '.csv'
    df_SAMPLE = pd.DataFrame.from_dict(df)
    df_SAMPLE.to_csv(Result, index=False)
    print('成功產出'+ Result)

def insert_df(name, buy_percentage, sell_trade_num, stop_loss, take_profit, profit):
    dir_path = 'Rank'
    csv_name = 'Rank'
    mkdir(dir_path)
    path = dir_path + '/' + csv_name + '.csv'
    try:
        df = pd.read_csv(path)
        temp=pd.DataFrame({'name':name,
                  'buy_percentage':buy_percentage,
                  'sell_trade_num':sell_trade_num,
                  'stop_loss':stop_loss,
                  'take_profit':take_profit,
                  'profit':profit},
                index=[0])
        df=df.append(temp, ignore_index=True)
        OutputCSV(df)
    except OSError:
        df = pd.DataFrame(columns=['name','buy_percentage','sell_trade_num','stop_loss','take_profit','profit'])
        OutputCSV(df)
        df = pd.read_csv(path)
        temp=pd.DataFrame({'name':name,
                  'buy_percentage':buy_percentage,
                  'sell_trade_num':sell_trade_num,
                  'stop_loss':stop_loss,
                  'take_profit':take_profit,
                  'profit':profit},
                index=[0])
        df=df.append(temp, ignore_index=True)
        OutputCSV(df)

# In[3]:

class Trade():
    trade_count = 0

    # 建構式
    def __init__(self, price, num):
        Trade.trade_count += 1
        self.price = price #買進價格
        self.num = num #買幾張
        self.total_cost = price*num*1000 #尚未考慮手續費

    def get_sell_profit(self, sell_price):
        return sell_price*self.num*1000 - self.total_cost

    def get_current_profit_pc(self, current_price):
        profit = self.get_sell_profit(current_price)
        profit_pc = profit / self.total_cost
        return profit_pc*100

def buy(buy_price):
    global cash
    cash_ = principal * buy_percentage/100
    if cash < cash_:
        cash_ = cash
    can_buy_num = cash_ // (buy_price*1000)
    if can_buy_num > 0:
        buy_trade = Trade(buy_price,can_buy_num)
        portfolio.append(buy_trade)
        cash -= buy_trade.total_cost #尚未考慮手續費
        print('以 %.1f元，買進 %d 張 2330，現金餘額 %.2f 元' % (buy_price,can_buy_num, cash))
    else:
        print('有買進訊號，但無法購買')

def sell(trade, sell_price):
    global cash, profit_count, loss_count, total_profit_rate
    sell_gain = trade.num*1000*sell_price #尚未考慮手續費
    cash += sell_gain
    profit_loss = trade.get_sell_profit(sell_price)
    profit_rate = trade.get_current_profit_pc(sell_price)
    total_profit_rate += profit_rate

    if profit_loss >= 0:
        profit_count += 1
    else:
        loss_count += 1

    portfolio.remove(trade)
    print('以 %.1f元，賣出 %d 張 2330，利損：%.2f元，現金餘額 %.2f 元'          % (sell_price,trade.num,profit_loss,cash))

def meet_strategy_sell(sell_price):
    if sell_trade_num == 'all':
        for _ in range(len(portfolio)):
            sell(portfolio[0], sell_price)
    else:
        tmp = min(sell_trade_num, len(portfolio))
        for _ in range(tmp):
            sell(portfolio[0], sell_price)

def cross_over(col1, col2, idx):
    # col1向上突破col2
    if idx == 0: return False
    return bt_data[col1][idx] > bt_data[col2][idx] and bt_data[col1][idx-1] < bt_data[col2][idx-1]

def cross_under(col1, col2, idx):
    # col1向上突破col2
    if idx == 0: return False
    return bt_data[col1][idx] < bt_data[col2][idx] and bt_data[col1][idx-1] > bt_data[col2][idx-1]

# In[4]:

def backtest(take_profit, stop_loss, buy_pc, sell_trade_n):
    global buy_percentage, sell_trade_num, ROI
    buy_percentage = buy_pc
    sell_trade_num = sell_trade_n
#     stop_loss = 5  # (%)  不停損 math.inf
#     take_profit = 10 # (%) 不停利 math.inf
    stock_num_2330 = 0
    principal = 1000000 * 10

    # 計算交易策略績效
    Trade.trade_count = 0

    for i in range(len(bt_data)-1):
        close = bt_data['Close'][i]
        next_close = bt_data['Close'][i+1]
        next_open = bt_data['Open'][i+1]
        rsi6 = bt_data['RSI6'][i]
        uprate_1day = bt_data['predict_1day_uprate'][i]

        # new
        msg = 'day:{}\n'.format(i)

        ### 先檢查有無 達到停利停損標準的 Trade
        for trade in portfolio:
            if trade.get_current_profit_pc(next_open) <= -stop_loss:
                msg += '(賣出)[停損賣出]\n'
                sell(trade, next_open)
            elif trade.get_current_profit_pc(next_open) >= take_profit:
                msg += '(買入)[停利賣出]\n'
                sell(trade, next_open)

        ### 買進策略
        if uprate_1day == 2 and next_open <= close:
            msg += '(買入)[Model 預測說要漲]\n'
            buy(next_close)

        if cross_over('K_value', 'D_value', i):
            msg += '(買入)[KD 黃金交叉]\n'
            buy(next_open)

        if cross_over('MA_5days', 'MA_20days', i):
            msg += '(買入)[MA 黃金交叉]\n'
            buy(next_open)

        if rsi6 <= 25:
            msg += '(買入)[RSI6 <= 25]\n'
            buy(next_open)

        ### 賣出策略

            #模型預測跌，當沖
        if len(portfolio)>0 and uprate_1day == 0 and next_open >= close:
            msg += '(當沖)[Model 預測說要跌，隔日開盤買進，收盤賣出]\n'
            buy(next_open)
            sell(portfolio[-1], next_close)

        if len(portfolio)>0 and cross_under('K_value', 'D_value', i):
            msg += '(賣出)[KD 死亡交叉]\n'
            meet_strategy_sell(next_open)

        if len(portfolio)>0 and cross_under('MA_5days', 'MA_20days', i):
            msg += '(賣出)[MA 死亡交叉]\n'
            meet_strategy_sell(next_open)

        if len(portfolio)>0 and rsi6 >= 70:
            msg += '(賣出)[RSI6 >= 70]\n'
            meet_strategy_sell(next_open)

        # new
        fornt_end_json['line_bot_push_msg'].append(msg)

        value_of_current_stock = np.sum([x.num*1000*next_open for x in portfolio])
        total_asset = value_of_current_stock + cash
        fornt_end_json['total_asset'].append(total_asset)

        print(msg)
        print('目前總資產 %d' % total_asset)

#         time.sleep(0.1)

    print('*'*25, '最後一天，全部平倉', '*'*25)
    last_close = bt_data['Close'][len(bt_data)-1]
    for _ in range(len(portfolio)):
        sell(portfolio[0], last_close)

    ROI = (cash/principal-1)*100

# In[7]:

bt_data = pd.read_csv('backtest_2330_data.csv')
# bt_data = bt_data[:120]

profit_count, loss_count, total_profit_rate, ROI = 0,0,0,0
principal = 1000000 * 10
cash = principal
portfolio = []
sell_trade_num = None

# new
fornt_end_json = dict()
fornt_end_json['total_asset'] = []
fornt_end_json['line_bot_push_msg'] = []

import json
def start_backtest_and_save_record(name, take_profit, stop_loss, buy_pc, sell_trade_num):
    ### buy_pc 0~100
    global profit_count, loss_count, total_profit_rate, ROI, fornt_end_json, principal, cash
    #init
    profit_count, loss_count, total_profit_rate, ROI = 0,0,0,0
    fornt_end_json = dict()
    fornt_end_json['total_asset'] = []
    fornt_end_json['line_bot_push_msg'] = []
    cash = principal

    Trade.trade_count = 0

    if take_profit == 0:
        # 不停利
        take_profit = math.inf

    if stop_loss == 0:
        # 不停損
        stop_loss = math.inf

    if sell_trade_num == 0:
        # 符合交易策略，一次全賣
        sell_trade_num = 'all'

    backtest(take_profit, stop_loss, buy_pc, sell_trade_num)

    trade_count = profit_count + loss_count
    win_rate = round(profit_count / trade_count*100,2)

    fornt_end_json['backtest_metric'] = dict()
    fornt_end_json['backtest_metric']['total_profit'] = round(ROI,2)
    fornt_end_json['backtest_metric']['avg_profit'] = round(total_profit_rate / trade_count,2)
    fornt_end_json['backtest_metric']['win_rate'] = win_rate

    fornt_end_json['backtest_metric']['trade'] = trade_count
    fornt_end_json['backtest_metric']['profit_trade'] = profit_count
    fornt_end_json['backtest_metric']['loss_trade'] = loss_count

    print('總報酬率: %.2f%%' % fornt_end_json['backtest_metric']['total_profit'])
    print('平均報酬率: %.2f%%' % fornt_end_json['backtest_metric']['avg_profit'])
    print('勝率: %.2f%%' % (fornt_end_json['backtest_metric']['win_rate']))

    print('總交易次數: %d次' % fornt_end_json['backtest_metric']['trade'])
    print('獲利次數: %d次' % fornt_end_json['backtest_metric']['profit_trade'])
    print('虧損次數: %d次' % fornt_end_json['backtest_metric']['loss_trade'])

#     print('總報酬率:', ROI)

    tmp_msg = fornt_end_json['line_bot_push_msg']
    # fornt_end_json['line_bot_push_msg'] = [m.replace('\n', '<br/>') for m in tmp_msg]
    return json.dumps(fornt_end_json)


import wx, os
from pvg_common import previewPanel, options

import wx.lib.newevent
ThumbnailClickedEvt, EVT_THUMBNAIL_CLICKED = wx.lib.newevent.NewCommandEvent()

from wx.lib.filebrowsebutton import FileBrowseButton, DirBrowseButton

class thumbnailPanel(previewPanel):
    """
    A small preview Panel to be used as thumbnail
    """

    def __init__( self, parent, monitor_number, thumbnailSize=(320,240) ):
        previewPanel.__init__(self, parent, size=thumbnailSize, keymode=False)

        self.number = int(monitor_number)
        self.allowEditing = False

        self.displayNumber()

        self.Bind(wx.EVT_LEFT_UP, self.onLeftClick)

    def displayNumber(self):
        """
        """
        # font type: wx.DEFAULT, wx.DECORATIVE, wx.ROMAN, wx.SCRIPT, wx.SWISS, wx.MODERN
        # slant: wx.NORMAL, wx.SLANT or wx.ITALIC
        # weight: wx.NORMAL, wx.LIGHT or wx.BOLD
        #font1 = wx.Font(10, wx.SWISS, wx.ITALIC, wx.NORMAL)
        # use additional fonts this way ...
        pos = int(self.size[0]/2 - 20), int(self.size[1]/2 - 20),
        font1 = wx.Font(35, wx.SWISS, wx.NORMAL, wx.NORMAL)
        text1 = wx.StaticText( self, wx.ID_ANY, '%s' % (self.number+1), pos)
        text1.SetFont(font1)

    def onLeftClick(self, evt):
        """
        Send signal around that the thumbnail was clicked
        """
        event = ThumbnailClickedEvt(self.GetId())

        event.id = self.GetId()
        event.number = self.number
        event.thumbnail = self

        self.GetEventHandler().ProcessEvent(event)

class panelGridView(wx.ScrolledWindow):
    """
    """
    def __init__(self, parent, gridSize, thumbnailSize=(320,240) ):
        """
        """
        wx.ScrolledWindow.__init__(self, parent, wx.ID_ANY, size=(-1,600))
        self.SetScrollbars(1, 1, 1, 1)
        self.SetScrollRate(10, 10)
        self.parent = parent
        self.thumbnailSize = thumbnailSize

        grid_mainSizer = wx.GridSizer(6,3,2,2)

        #Populate the thumbnail grid
        self.previewPanels = []
        for i in range (int(gridSize)):
            self.previewPanels.append ( thumbnailPanel(self, monitor_number=i, thumbnailSize=self.thumbnailSize) )
            grid_mainSizer.Add(self.previewPanels[i])#, 0, wx.EXPAND|wx.FIXED_MINSIZE, 0)

        self.SetSizer(grid_mainSizer)

        self.Bind(EVT_THUMBNAIL_CLICKED, self.onThumbnailClicked)

    def onThumbnailClicked(self, event):
        """
        Relay event to sibling panel
        """
        wx.PostEvent(self.parent.lowerPanel, event)
        event.Skip()

class panelConfigure(wx.Panel):
    """
    """
    def __init__(self, parent):
        """
        """
        wx.Panel.__init__(self, parent, wx.ID_ANY, size=(-1,300), style=wx.SUNKEN_BORDER|wx.TAB_TRAVERSAL)
        self.parent = parent

        self.thumbnail = None
        self.mask_file = None
        self.source = None
        self.sourceType = None
        self.track = None
        self.trackType = None

        lowerSizer = wx.BoxSizer(wx.HORIZONTAL)

        #Static box1 (LEFT)
        sb_1 = wx.StaticBox(self, -1, "Select Monitor")#, size=(250,-1))
        sbSizer_1 = wx.StaticBoxSizer (sb_1, wx.VERTICAL)

        n_monitors = options.GetOption("Monitors")
        self.MonitorList = ['Monitor %s' % (int(m) + 1) for m in range( n_monitors )]
        self.thumbnailNumber = wx.ComboBox(self, -1, size=(-1,-1) , choices=self.MonitorList, style=wx.CB_DROPDOWN | wx.CB_READONLY | wx.CB_SORT)
        self.Bind ( wx.EVT_COMBOBOX, self.onChangingMonitor, self.thumbnailNumber)

        self.currentSource = wx.TextCtrl (self, -1, "No Source Selected", style=wx.TE_READONLY)

        btnSizer_1 = wx.BoxSizer(wx.HORIZONTAL)
        self.btnPlay = wx.Button( self, wx.ID_FORWARD, label="Play")
        self.btnStop = wx.Button( self, wx.ID_STOP, label="Stop")
        self.Bind(wx.EVT_BUTTON, self.onPlay, self.btnPlay)
        self.Bind(wx.EVT_BUTTON, self.onStop, self.btnStop)
        self.btnPlay.Enable(False); self.btnStop.Enable(False)
        self.applyButton = wx.Button( self, wx.ID_APPLY )
        self.applyButton.SetToolTip(wx.ToolTip("Apply and Save to file"))
        self.Bind(wx.EVT_BUTTON, self.onApplySource, self.applyButton)

        btnSizer_1.Add ( self.btnPlay , 0, wx.ALIGN_LEFT|wx.ALL, 5 )
        btnSizer_1.Add ( self.btnStop , 0, wx.ALIGN_CENTER|wx.LEFT|wx.TOP|wx.DOWN, 5 )
        btnSizer_1.Add ( self.applyButton, 0, wx.ALIGN_RIGHT|wx.LEFT|wx.RIGHT|wx.TOP, 5 )

        sbSizer_1.Add ( self.thumbnailNumber, 0, wx.ALIGN_CENTRE|wx.LEFT|wx.RIGHT|wx.TOP, 5 )
        sbSizer_1.Add ( self.currentSource, 0, wx.EXPAND|wx.ALIGN_CENTRE|wx.LEFT|wx.RIGHT|wx.TOP, 5 )
        sbSizer_1.Add ( btnSizer_1, 0, wx.EXPAND|wx.ALIGN_BOTTOM|wx.TOP, 5 )

        lowerSizer.Add (sbSizer_1, 0, wx.EXPAND|wx.ALL, 5)

        #Static box2 (CENTER)
        sb_2 = wx.StaticBox(self, -1, "Select Video input" )
        sbSizer_2 = wx.StaticBoxSizer (sb_2, wx.VERTICAL)
        grid2 = wx.FlexGridSizer( 0, 2, 0, 0 )

        n_cams = options.GetOption("Webcams")
        self.WebcamsList = [ 'Webcam %s' % (int(w) +1) for w in range( n_cams ) ]
        rb1 = wx.RadioButton(self, -1, 'Camera', style=wx.RB_GROUP)
        source1 = wx.ComboBox(self, -1, size=(285,-1) , choices = self.WebcamsList, style=wx.CB_DROPDOWN | wx.CB_READONLY | wx.CB_SORT)
        self.Bind(wx.EVT_COMBOBOX, self.sourceCallback, source1)

        rb2 = wx.RadioButton(self, -1, 'File' )
        source2 = FileBrowseButton(self, -1, labelText='', size=(300,-1), changeCallback = self.sourceCallback)

        rb3 = wx.RadioButton(self, -1, 'Folder' )
        source3 = DirBrowseButton (self, style=wx.DD_DIR_MUST_EXIST, labelText='', size=(300,-1), changeCallback = self.sourceCallback)

        self.controls = []
        self.controls.append((rb1, source1))
        self.controls.append((rb2, source2))
        self.controls.append((rb3, source3))

        for radio, source in self.controls:
            grid2.Add( radio , 0, wx.ALIGN_LEFT|wx.ALIGN_CENTER_VERTICAL|wx.ALL, 2 )
            grid2.Add( source , 0, wx.ALIGN_CENTRE|wx.ALIGN_CENTER_VERTICAL|wx.ALL, 2 )
            self.Bind(wx.EVT_RADIOBUTTON, self.onChangeSource, radio )
            source.Enable(False)

        self.controls[0][1].Enable(True)

        #grid2.Add(wx.StaticText(self, -1, ""))

        sbSizer_2.Add( grid2 )
        lowerSizer.Add(sbSizer_2, 0, wx.EXPAND|wx.ALL, 5)

        #Static box3 (RIGHT)
        sb_3 = wx.StaticBox(self, -1, "Set Tracking Parameters")
        sbSizer_3 = wx.StaticBoxSizer (sb_3, wx.VERTICAL)

        sbSizer_31 = wx.BoxSizer (wx.HORIZONTAL)

        self.activateTracking = wx.CheckBox(self, -1, "Activate Tracking")
        self.activateTracking.SetValue(False)
        self.activateTracking.Bind ( wx.EVT_CHECKBOX, self.onActivateTracking)

        self.isSDMonitor = wx.CheckBox(self, -1, "Sleep Deprivation Monitor")
        self.isSDMonitor.SetValue(False)
        self.isSDMonitor.Bind ( wx.EVT_CHECKBOX, self.onSDMonitor)
        self.isSDMonitor.Enable(False)

        sbSizer_31.Add (self.activateTracking, 0, wx.ALIGN_LEFT|wx.LEFT|wx.RIGHT|wx.TOP, 5 )
        sbSizer_31.Add (self.isSDMonitor, 0, wx.ALIGN_LEFT|wx.LEFT|wx.RIGHT|wx.TOP, 5 )

        self.pickMaskBrowser = FileBrowseButton(self, -1, labelText='Mask File')

        #sbSizer_3.Add ( self.activateTracking , 0, wx.ALIGN_LEFT|wx.LEFT|wx.RIGHT|wx.TOP, 5 )
        sbSizer_3.Add ( sbSizer_31 , 0, wx.ALIGN_LEFT|wx.LEFT|wx.RIGHT|wx.TOP, 5 )
        sbSizer_3.Add ( self.pickMaskBrowser , 0, wx.ALIGN_LEFT|wx.LEFT|wx.RIGHT|wx.TOP|wx.EXPAND, 5 )

        #trackingTypeSizer = wx.Sizer(wx.HORIZONTAL)
        self.trackDistanceRadio = wx.RadioButton(self, -1, "Activity as distance traveled", style=wx.RB_GROUP)
        self.trackVirtualBM = wx.RadioButton(self, -1, "Activity as midline crossings count")
        self.trackPosition = wx.RadioButton(self, -1, "Only position of flies")
        sbSizer_3.Add (wx.StaticText ( self, -1, "Calculate fly activity as..."), 0, wx.ALIGN_LEFT|wx.LEFT|wx.RIGHT|wx.TOP, 5 )
        sbSizer_3.Add (self.trackDistanceRadio, 0, wx.ALIGN_LEFT|wx.LEFT|wx.RIGHT|wx.TOP, 2 )
        sbSizer_3.Add (self.trackVirtualBM, 0, wx.ALIGN_LEFT|wx.LEFT|wx.RIGHT|wx.TOP, 2 )
        sbSizer_3.Add (self.trackPosition, 0, wx.ALIGN_LEFT|wx.LEFT|wx.RIGHT|wx.TOP, 2 )

        lowerSizer.Add(sbSizer_3, -1, wx.EXPAND|wx.ALL, 5)

        self.SetSizer(lowerSizer)
        self.Bind(EVT_THUMBNAIL_CLICKED, self.onThumbnailClicked)

    def __getSource(self):
        """
        check which source is ticked and what is the associated value
        """

        for (r, s), st in zip(self.controls,range(3)):
            if r.GetValue():
                source = s.GetValue()
                sourceType = st

        return source, sourceType

    def __getTrackingType(self):
        """
        return which tracking we are chosing
        ['DISTANCE','VBS','XY_COORDS']
        """
        if self.trackDistanceRadio.GetValue(): trackType = "DISTANCE"
        elif self.trackVirtualBM.GetValue(): trackType = "VBS"
        elif self.trackPosition.GetValue(): trackType = "XY_COORDS"

        return trackType

    def onPlay (self, event=None):
        """
        """
        if self.thumbnail:
            self.thumbnail.Play()
            self.btnStop.Enable(True)

    def onStop (self, event=None):
        """
        """
        if self.thumbnail and self.thumbnail.isPlaying:
            self.thumbnail.Stop()
            self.btnStop.Enable(False)

    def onThumbnailClicked(self, evt):
        """
        Picking thumbnail by clicking on it
        """
        self.monitor_number = evt.number + 1
        self.thumbnail = evt.thumbnail
        self.thumbnailNumber.SetValue(self.MonitorList[self.monitor_number -1 ])
        self.updateThumbnail()

    def onChangingMonitor(self, evt):
        """
        Picking thumbnail by using the dropbox
        """
        sel = evt.GetString()
        self.monitor_number = self.MonitorList.index(sel) + 1
        self.thumbnail = self.parent.scrollThumbnails.previewPanels[self.monitor_number]         #this is not very elegant
        self.updateThumbnail()

    def updateThumbnail(self):
        """
        Refreshing thumbnail data
        """
        if options.HasMonitor(self.monitor_number):
            sourceType, source, track, mask_file, trackType, isSDMonitor = options.GetMonitor(self.monitor_number)
        else:
            sourceType, source, track, mask_file, trackType, isSDMonitor = [0, '', False, '', 1, False]

        if sourceType == 0 and source != '':
            source = self.WebcamsList[source]

        self.source = self.thumbnail.source = source
        self.sourceType = self.thumbnail.sourceType = sourceType
        self.thumbnail.track = track
        if self.thumbnail.hasMonitor():
                self.thumbnail.mon.isSDMonitor = isSDMonitor

        #update first static box
        active = self.thumbnail.hasMonitor()
        self.applyButton.Enable ( active )
        self.btnPlay.Enable ( active )
        self.btnStop.Enable ( active and self.thumbnail.isPlaying )

        text = os.path.split(str(self.source))[1] or "No Source Selected"
        self.currentSource.SetValue( text )

        #update second static box
        for radio, src in self.controls:
            src.Enable(False); src.SetValue('')

        radio, src = self.controls[self.sourceType]
        radio.SetValue(True); src.Enable(True)
        src.SetValue(self.source)

        #update third static box
        self.activateTracking.SetValue(self.thumbnail.track)
        self.isSDMonitor.SetValue(isSDMonitor)
        self.pickMaskBrowser.SetValue(mask_file or '')
        [self.trackDistanceRadio, self.trackVirtualBM, self.trackPosition][trackType].SetValue(True)

    def sourceCallback (self, event):
        """
        """
        self.applyButton.Enable(True)

    def onChangeSource(self, event):
        """

        """

        radio_selected = event.GetEventObject()

        for radio, source in self.controls:
            if radio is radio_selected:
                source.Enable(True)
            else:
                source.Enable(False)

        self.applyButton.Enable(True)

    def onApplySource(self, event):
        """
        """

        source, sourceType = self.__getSource()
        track = self.activateTracking.GetValue()
        self.mask_file = self.pickMaskBrowser.GetValue()
        self.trackType = self.__getTrackingType()

        if self.thumbnail:

            if sourceType > 0: camera = source
            else: camera = self.WebcamsList.index(source)

            self.thumbnail.source = camera
            self.thumbnail.sourceType = sourceType

            #Change the source text
            self.currentSource.SetValue( os.path.split(source)[1] )

            #Set thumbnail's source
            self.thumbnail.setMonitor(camera)

            #Enable buttons
            self.btnPlay.Enable(True)
            self.activateTracking.Enable(True)
            self.pickMaskBrowser.Enable(True)

            self.saveMonitorConfiguration()

    def saveMonitorConfiguration(self):
        """
        """

        options.SetMonitor(self.monitor_number,
                           self.thumbnail.sourceType,
                           self.thumbnail.source+1,
                           self.thumbnail.track,
                           self.mask_file,
                           self.trackType,
                           self.thumbnail.mon.isSDMonitor
                           )
        options.Save()

    def onActivateTracking(self, event):
        """
        """
        if self.thumbnail:
            self.thumbnail.track = event.IsChecked()

    def onSDMonitor(self, event):
        """
        """
        if self.thumbnail:
            self.thumbnail.mon.isSDMonitor = event.IsChecked()

class panelOne(wx.Panel):
    """
    Panel number One
    All the thumbnails
    """
    def __init__(self, parent):

        wx.Panel.__init__(self, parent)

        monitor_number = options.GetOption("Monitors")
        tn_size = options.GetOption("ThumbnailSize")

        self.temp_source  = ''
        self.source = ''
        self.sourceType = -1

        self.scrollThumbnails = panelGridView(self, gridSize=monitor_number, thumbnailSize=tn_size)
        self.lowerPanel = panelConfigure(self)

        self.PanelOneSizer = wx.BoxSizer(wx.VERTICAL)
        self.PanelOneSizer.Add(self.scrollThumbnails, 1, wx.EXPAND, 0)
        self.PanelOneSizer.Add(self.lowerPanel, 0, wx.EXPAND, 0)
        self.SetSizer(self.PanelOneSizer)

    def StopPlaying(self):
        """
        """
        self.lowerPanel.onStop()


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}'))


"""
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 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 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)


# 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)


import copy
import json
import os
import warnings
from absl import app, flags

import torch
from tensorboardX import SummaryWriter
from torchvision.datasets import CIFAR10
from torchvision.datasets import ImageFolder
from torchvision.utils import make_grid, save_image
from torchvision import transforms
from tqdm import trange
import torch.distributed as dist

from diffusion import GaussianDiffusionTrainer, GaussianDiffusionSampler
from model import UNet

FLAGS = flags.FLAGS
flags.DEFINE_enum('dataset', 'cifar10', ['cifar10', 'imagenet64'], help='dataset')
# UNet
flags.DEFINE_integer('ch', 128, help='base channel of UNet')
flags.DEFINE_multi_integer('ch_mult', [1, 2, 2, 2], help='channel multiplier')
flags.DEFINE_multi_integer('attn', [1], help='add attention to these levels')
flags.DEFINE_integer('num_res_blocks', 2, help='# resblock in each level')
flags.DEFINE_float('dropout', 0.1, help='dropout rate of resblock')
# Gaussian Diffusion
flags.DEFINE_integer('T', 1024, help='total diffusion steps')
flags.DEFINE_integer('time_scale', 1, help='total diffusion steps')
flags.DEFINE_enum('mean_type', 'xstart', ['xstart', 'epsilon'], help='predict variable')
flags.DEFINE_enum('var_type', 'fixedlarge', ['fixedlarge', 'fixedsmall'], help='variance type')
flags.DEFINE_enum('loss_type', 'both', ['x', 'eps', 'both'], help='loss type')
# Training
flags.DEFINE_float('lr', 2e-4, help='target learning rate')
flags.DEFINE_float('wd', 0.001, help='target learning rate')
flags.DEFINE_float('grad_clip', 1., help="gradient norm clipping")
flags.DEFINE_integer('total_steps', 800000, help='total training steps')
flags.DEFINE_integer('img_size', 32, help='image size')
flags.DEFINE_integer('warmup', 5000, help='learning rate warmup')
flags.DEFINE_integer('batch_size', 128, help='batch size')
flags.DEFINE_integer('num_workers', 4, help='workers of Dataloader')
flags.DEFINE_float('ema_decay', 0.9999, help="ema decay rate")
flags.DEFINE_string('gpu_id', '4,5,6,7', help='multi gpu training')
flags.DEFINE_integer('local_rank', 0, help='local rank')
flags.DEFINE_bool('distributed', False, help='multi gpu training')
flags.DEFINE_integer('num_gpus', 4, help='multi gpu training')
flags.DEFINE_bool('conditional', False, help='use conditional or not')
flags.DEFINE_integer('class_num', 10, help='class num')
# Logging & Sampling
flags.DEFINE_string('logdir', './logs/CIFAR10/1024', help='log directory')
flags.DEFINE_integer('sample_size', 64, "sampling size of images")
flags.DEFINE_integer('sample_step', 1000, help='frequency of sampling')
flags.DEFINE_integer('save_step', 1000, help='frequency of saving checkpoints, 0 to disable during training')

def ema(source, target, decay):
    source_dict = source.state_dict()
    target_dict = target.state_dict()
    for key in source_dict.keys():
        target_dict[key].data.copy_(
            target_dict[key].data * decay +
            source_dict[key].data * (1 - decay))

def infiniteloop(dataloader):
    while True:
        for x, y in iter(dataloader):
            yield x, y

def get_rank():
    if not dist.is_available():
        return 0
    if not dist.is_initialized():
        return 0
    return dist.get_rank()

def warmup_lr(step):
    return min(step, FLAGS.warmup) / FLAGS.warmup

def train():
    if get_rank() == 0:
        if not os.path.exists(os.path.join(FLAGS.logdir, 'ddim_clip')):
            os.makedirs(os.path.join(FLAGS.logdir, 'ddim_clip'))
    # dataset
    if FLAGS.dataset == 'cifar10':
        dataset = CIFAR10(
            root='../../data', train=True, download=True,
            transform=transforms.Compose([
                transforms.RandomHorizontalFlip(),
                transforms.ToTensor(),
                transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
            ]))
    elif FLAGS.dataset == 'imagenet64':
        dataset = ImageFolder(
            '../../data/ImageNet64/train',
            transforms.Compose([
                transforms.RandomHorizontalFlip(),
                transforms.ToTensor(),
                transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
            ]))

    batch_size = int(FLAGS.batch_size / FLAGS.num_gpus)
    if FLAGS.distributed:
        train_sampler = torch.utils.data.distributed.DistributedSampler(dataset)
        train_loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size,
                                               num_workers=FLAGS.num_workers, pin_memory=True, sampler=train_sampler, drop_last=True)
    else:
        train_loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=True,
                                                num_workers=FLAGS.num_workers, drop_last=True)
    train_looper = infiniteloop(train_loader)

    # model setup
    net_model = UNet(
        T=FLAGS.T*FLAGS.time_scale, ch=FLAGS.ch, ch_mult=FLAGS.ch_mult, attn=FLAGS.attn,
        num_res_blocks=FLAGS.num_res_blocks, dropout=FLAGS.dropout,
        conditional=FLAGS.conditional, class_num=FLAGS.class_num)
    ema_model = copy.deepcopy(net_model)

    trainer = GaussianDiffusionTrainer(
        net_model, FLAGS.T, FLAGS.time_scale, FLAGS.loss_type, FLAGS.mean_type).cuda(FLAGS.local_rank)
    net_sampler = GaussianDiffusionSampler(
        net_model, FLAGS.T, FLAGS.time_scale, img_size=FLAGS.img_size,
        mean_type=FLAGS.mean_type, var_type=FLAGS.var_type).cuda(FLAGS.local_rank)
    ema_sampler = GaussianDiffusionSampler(
        ema_model, FLAGS.T, FLAGS.time_scale, img_size=FLAGS.img_size,
        mean_type=FLAGS.mean_type, var_type=FLAGS.var_type).cuda(FLAGS.local_rank)
    if FLAGS.distributed:
        trainer = torch.nn.parallel.DistributedDataParallel(trainer, device_ids=[FLAGS.local_rank], output_device=FLAGS.local_rank)
        net_sampler = torch.nn.parallel.DistributedDataParallel(net_sampler, device_ids=[FLAGS.local_rank], output_device=FLAGS.local_rank)
        ema_sampler = torch.nn.parallel.DistributedDataParallel(ema_sampler, device_ids=[FLAGS.local_rank], output_device=FLAGS.local_rank)

    optim = torch.optim.Adam(net_model.parameters(), lr=FLAGS.lr, weight_decay=FLAGS.wd)
    sched = torch.optim.lr_scheduler.LambdaLR(optim, lr_lambda=warmup_lr)

    # log setup
    x_T = torch.randn(int(FLAGS.sample_size / FLAGS.num_gpus), 3, FLAGS.img_size, FLAGS.img_size)
    x_T = x_T.cuda(FLAGS.local_rank)
    grid = (make_grid(next(iter(train_loader))[0][:int(FLAGS.sample_size / FLAGS.num_gpus)]) + 1) / 2
    if get_rank() == 0:
        writer = SummaryWriter(FLAGS.logdir)
        writer.add_image('real_sample', grid)
        writer.flush()
        # backup all arguments
        with open(os.path.join(FLAGS.logdir, "flagfile.txt"), 'w') as f:
            f.write(FLAGS.flags_into_string())
    # show model size
    model_size = 0
    for param in net_model.parameters():
        model_size += param.data.nelement()
    if get_rank() == 0:
        print('Model params: %.2f M' % (model_size / 1024 / 1024))

    # start training
    for step in range(1, FLAGS.total_steps + 1):
        if FLAGS.distributed:
            train_sampler.set_epoch(step)
        # train
        samples = next(train_looper)
        x_0, y = samples[0].cuda(FLAGS.local_rank), samples[1].cuda(FLAGS.local_rank)
        loss = trainer(x_0, y)
        optim.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(net_model.parameters(), FLAGS.grad_clip)
        optim.step()
        sched.step()
        ema(net_model, ema_model, FLAGS.ema_decay)

        # log
        if get_rank() == 0:
            writer.add_scalar('loss', loss, step)

        # sample
        if FLAGS.sample_step > 0 and (step % FLAGS.sample_step == 0 or step == 1):
            ema_model.eval()
            with torch.no_grad():
                y_target = torch.randint(FLAGS.class_num, size=(x_T.shape[0],), device=x_T.device)
                x_0 = ema_sampler.module.ddim(x_T, 1, True, y=y_target)
                grid = (make_grid(x_0) + 1) / 2
                path = os.path.join(FLAGS.logdir, 'ddim_clip', '%d.png' % step)
                if get_rank() == 0:
                    save_image(grid, path)
                    writer.add_image('ddim_clip', grid, step)
            ema_model.train()

        # save
        if FLAGS.save_step > 0 and step % FLAGS.save_step == 0:
            if get_rank() == 0:
                ckpt = {
                    'net_model': net_model.state_dict(),
                    'ema_model': ema_model.state_dict(),
                    'x_T': x_T,
                    'T': trainer.module.T,
                    'time_scale': trainer.module.time_scale,
                }
                torch.save(ckpt, os.path.join(FLAGS.logdir, 'ckpt.pt'))
    torch.distributed.barrier()
    if get_rank() == 0:
        writer.close()

def main(argv):
    os.environ["CUDA_VISIBLE_DEVICES"] = FLAGS.gpu_id
    FLAGS.num_gpus = int(os.environ["WORLD_SIZE"]) if "WORLD_SIZE" in os.environ else 1
    FLAGS.distributed = FLAGS.num_gpus > 1
    if FLAGS.distributed:
        torch.cuda.set_device(FLAGS.local_rank)
        torch.distributed.init_process_group(backend="nccl", init_method="env://")
    # suppress annoying inception_v3 initialization warning
    warnings.simplefilter(action='ignore', category=FutureWarning)
    train()

if __name__ == '__main__':
    app.run(main)