# PyHSS
# This serves as a basic 3GPP Home Subscriber Server implimenting a EIR & IMS HSS functionality
import logging
import yaml
import os
import sys
import socket
import socketserver
import binascii
import time
import _thread
import threading
import sctp
import traceback
import pprint
import diameter as DiameterLib
import systemd.daemon
from threading import Thread, Lock
from logtool import *
import contextlib
import queue

class ThreadJoiner:
    def __init__(self, threads, thread_event):
        self.threads = threads
        self.thread_event = thread_event

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        if exc_type is not None:
            self.thread_event.set()
        for thread in self.threads:
            while thread.is_alive():
                try:
                    thread.join(timeout=1)
                except Exception as e:
                    print(
                        f"ThreadJoiner Exception: failed to join thread {thread}: {e}"
                    )
                    break

class PyHSS:
    def __init__(self):
        # Load config from yaml file
        try:
            with open("config.yaml", "r") as config_stream:
                self.yaml_config = yaml.safe_load(config_stream)
        except:
            print(f"config.yaml not found, exiting PyHSS.")
            quit()

        # Setup logging
        self.logtool = LogTool(HSS_Init=True)
        self.logtool.setup_logger(
            "HSS_Logger",
            self.yaml_config["logging"]["logfiles"]["hss_logging_file"],
            level=self.yaml_config["logging"]["level"],
        )
        self.logger = logging.getLogger("HSS_Logger")
        if self.yaml_config["logging"]["log_to_terminal"]:
            logging.getLogger().addHandler(logging.StreamHandler())

        # Setup Diameter
        self.diameter_instance = DiameterLib.Diameter(
            str(self.yaml_config["hss"].get("OriginHost", "")),
            str(self.yaml_config["hss"].get("OriginRealm", "")),
            str(self.yaml_config["hss"].get("ProductName", "")),
            str(self.yaml_config["hss"].get("MNC", "")),
            str(self.yaml_config["hss"].get("MCC", "")),
        )

        self.max_diameter_retries = int(
            self.yaml_config["hss"].get("diameter_max_retries", 1)
        )

        try:
            assert(self.yaml_config['prometheus']['enabled'] == True)
            assert(self.yaml_config['prometheus']['async_subscriber_count'] == True)

            self.logger.info("Enabling Prometheus Async Sub thread")
            #Add Prometheus Async Calls
            prom_async_thread = threading.Thread(
                target=self.prom_async_function,
                name=f"prom_async_function",
                args=(),
            )
            prom_async_thread.start()
        except:
            self.logger.info("Prometheus Async Sub Count thread disabled")

    def terminate_connection(self, clientsocket, client_address, thread_event):
        thread_event.set()
        clientsocket.close()
        self.logtool.Manage_Diameter_Peer(client_address, client_address, "remove")

    def handle_new_connection(self, clientsocket, client_address):
        # Create our threading event, accessible by sibling threads in this connection.
        socket_close_event = threading.Event()
        try:
            send_queue = queue.Queue()
            self.logger.debug(f"New connection from {client_address}")
            if (
                "client_socket_timeout" not in self.yaml_config["hss"]
                or self.yaml_config["hss"]["client_socket_timeout"] == 0
            ):
                self.yaml_config["hss"]["client_socket_timeout"] = 120
            clientsocket.settimeout(
                self.yaml_config["hss"].get("client_socket_timeout", 120)
            )

            send_data_thread = threading.Thread(
                target=self.send_data,
                name=f"send_data_thread",
                args=(clientsocket, send_queue, socket_close_event),
            )
            self.logger.debug("handle_new_connection: Starting send_data thread")
            send_data_thread.start()

            self.logtool.Manage_Diameter_Peer(client_address, client_address, "add")
            manage_client_thread = threading.Thread(
                target=self.manage_client,
                name=f"manage_client_thread: client_address: {client_address}",
                args=(
                    clientsocket,
                    client_address,
                    self.diameter_instance,
                    socket_close_event,
                    send_queue,
                ),
            )
            self.logger.debug("handle_new_connection: Starting manage_client thread")
            manage_client_thread.start()

            threads_to_join = [manage_client_thread]
            threads_to_join.append(send_data_thread)

            # If Redis is enabled, start manage_client_async and manage_client_dwr threads.
            if self.yaml_config["redis"]["enabled"]:
                if (
                    "async_check_interval" not in self.yaml_config["hss"]
                    or self.yaml_config["hss"]["async_check_interval"] == 0
                ):
                    self.yaml_config["hss"]["async_check_interval"] = 10
                manage_client_async_thread = threading.Thread(
                    target=self.manage_client_async,
                    name=f"manage_client_async_thread: client_address: {client_address}",
                    args=(
                        clientsocket,
                        client_address,
                        self.diameter_instance,
                        socket_close_event,
                        send_queue,
                    ),
                )
                self.logger.debug(
                    "handle_new_connection: Starting manage_client_async thread"
                )
                manage_client_async_thread.start()

                manage_client_dwr_thread = threading.Thread(
                    target=self.manage_client_dwr,
                    name=f"manage_client_dwr_thread: client_address: {client_address}",
                    args=(
                        clientsocket,
                        client_address,
                        self.diameter_instance,
                        socket_close_event,
                        send_queue,
                    ),
                )
                self.logger.debug(
                    "handle_new_connection: Starting manage_client_dwr thread"
                )
                manage_client_dwr_thread.start()

                threads_to_join.append(manage_client_async_thread)
                threads_to_join.append(manage_client_dwr_thread)

            self.logger.debug(
                f"handle_new_connection: Total PyHSS Active Threads: {threading.active_count()}"
            )
            for thread in threading.enumerate():
                if "dummy" not in thread.name.lower():
                    self.logger.debug(f"Active Thread name: {thread.name}")

            with ThreadJoiner(threads_to_join, socket_close_event):
                socket_close_event.wait()
                self.terminate_connection(
                    clientsocket, client_address, socket_close_event
                )
                self.logger.debug(f"Closing thread for client; {client_address}")
                return

        except Exception as e:
            self.logger.error(f"Exception for client {client_address}: {e}")
            self.logger.error(f"Closing connection for {client_address}")
            self.terminate_connection(clientsocket, client_address, socket_close_event)
            return

    @prom_diam_response_time_diam.time()
    def process_Diameter_request(
        self, clientsocket, client_address, diameter, data, thread_event, send_queue
    ):
        packet_length = diameter.decode_diameter_packet_length(
            data
        )  # Calculate length of packet from start of packet
        if packet_length <= 32:
            self.logger.error("Received an invalid packet with length <= 32")
            self.terminate_connection(clientsocket, client_address, thread_event)
            return

        data_sum = data + clientsocket.recv(
            packet_length - 32
        )  # Recieve remainder of packet from buffer
        packet_vars, avps = diameter.decode_diameter_packet(
            data_sum
        )  # Decode packet into array of AVPs and Dict of Packet Variables (packet_vars)
        try:
            packet_vars["Source_IP"] = client_address[0]
        except:
            self.logger.debug("Failed to add Source_IP to packet_vars")

        start_time = time.time()
        origin_host = diameter.get_avp_data(avps, 264)[0]  # Get OriginHost from AVP
        origin_host = binascii.unhexlify(origin_host).decode("utf-8")  # Format it

        # label_values = str(packet_vars['ApplicationId']), str(packet_vars['command_code']), origin_host, 'request'
        prom_diam_request_count.labels(
            str(packet_vars["ApplicationId"]),
            str(packet_vars["command_code"]),
            origin_host,
            "request",
        ).inc()

        self.logger.info(
            "\n\nNew request with Command Code: "
            + str(packet_vars["command_code"])
            + ", ApplicationID: "
            + str(packet_vars["ApplicationId"])
            + ", flags "
            + str(packet_vars["flags"])
            + ", e2e ID: "
            + str(packet_vars["end-to-end-identifier"])
        )

        # Gobble up any Response traffic that is sent to us:
        if packet_vars["flags_bin"][0:1] == "0":
            self.logger.info("Got a Response, not a request - dropping it.")
            self.logger.info(packet_vars)
            return

        # Send Capabilities Exchange Answer (CEA) response to Capabilites Exchange Request (CER)
        elif (
            packet_vars["command_code"] == 257
            and packet_vars["ApplicationId"] == 0
            and packet_vars["flags"] == "80"
        ):
            self.logger.info(
                f"Received Request with command code 257 (CER) from {origin_host}"
                + "\n\tSending response (CEA)"
            )
            try:
                response = diameter.Answer_257(
                    packet_vars, avps, str(self.yaml_config["hss"]["bind_ip"][0])
                )  # Generate Diameter packet
                # prom_diam_response_count_successful.inc()
            except:
                response = diameter.Respond_ResultCode(
                    packet_vars, avps, 5012
                )  # Generate Diameter response with "DIAMETER_UNABLE_TO_COMPLY" (5012)
                # prom_diam_response_count_fail.inc()
            self.logger.info("Generated CEA")
            self.logtool.Manage_Diameter_Peer(origin_host, client_address, "update")
            prom_diam_connected_peers.labels(origin_host).set(1)

        # Send Credit Control Answer (CCA) response to Credit Control Request (CCR)
        elif (
            packet_vars["command_code"] == 272
            and packet_vars["ApplicationId"] == 16777238
        ):
            self.logger.info(
                f"Received 3GPP Credit-Control-Request from {origin_host}"
                + "\n\tGenerating (CCA)"
            )
            try:
                response = diameter.Answer_16777238_272(
                    packet_vars, avps
                )  # Generate Diameter packet
            except Exception as E:
                response = diameter.Respond_ResultCode(
                    packet_vars, avps, 5012
                )  # Generate Diameter response with "DIAMETER_UNABLE_TO_COMPLY" (5012)
                self.logger.error(f"Failed to generate response {str(E)}")
            self.logger.info("Generated CCA")

        # Send Device Watchdog Answer (DWA) response to Device Watchdog Requests (DWR)
        elif (
            packet_vars["command_code"] == 280
            and packet_vars["ApplicationId"] == 0
            and packet_vars["flags"] == "80"
        ):
            self.logger.info(
                f"Received Request with command code 280 (DWR) from {origin_host}"
                + "\n\tSending response (DWA)"
            )
            self.logger.debug(f"Total PyHSS Active Threads: {threading.active_count()}")
            try:
                response = diameter.Answer_280(
                    packet_vars, avps
                )  # Generate Diameter packet
            except:
                response = diameter.Respond_ResultCode(
                    packet_vars, avps, 5012
                )  # Generate Diameter response with "DIAMETER_UNABLE_TO_COMPLY" (5012)
            self.logger.info("Generated DWA")
            self.logtool.Manage_Diameter_Peer(origin_host, client_address, "update")

        # Send Disconnect Peer Answer (DPA) response to Disconnect Peer Request (DPR)
        elif (
            packet_vars["command_code"] == 282
            and packet_vars["ApplicationId"] == 0
            and packet_vars["flags"] == "80"
        ):
            self.logger.info(
                f"Received Request with command code 282 (DPR) from {origin_host}"
                + "\n\tForwarding request..."
            )
            response = diameter.Answer_282(
                packet_vars, avps
            )  # Generate Diameter packet
            self.logger.info("Generated DPA")
            self.logtool.Manage_Diameter_Peer(origin_host, client_address, "remove")
            prom_diam_connected_peers.labels(origin_host).set(0)

        # S6a Authentication Information Answer (AIA) response to Authentication Information Request (AIR)
        elif (
            packet_vars["command_code"] == 318
            and packet_vars["ApplicationId"] == 16777251
            and packet_vars["flags"] == "c0"
        ):
            self.logger.info(
                f"Received Request with command code 318 (3GPP Authentication-Information-Request) from {origin_host}"
                + "\n\tGenerating (AIA)"
            )
            try:
                response = diameter.Answer_16777251_318(
                    packet_vars, avps
                )  # Generate Diameter packet
                self.logger.info("Generated AIR")
            except Exception as e:
                self.logger.info("Failed to generate Diameter Response for AIR")
                self.logger.info(e)
                traceback.print_exc()
                response = diameter.Respond_ResultCode(
                    packet_vars, avps, 4100
                )  # DIAMETER_USER_DATA_NOT_AVAILABLE
                self.logger.info("Generated DIAMETER_USER_DATA_NOT_AVAILABLE AIR")

        # S6a Update Location Answer (ULA) response to Update Location Request (ULR)
        elif (
            packet_vars["command_code"] == 316
            and packet_vars["ApplicationId"] == 16777251
        ):
            self.logger.info(
                f"Received Request with command code 316 (3GPP Update Location-Request) from {origin_host}"
                + "\n\tGenerating (ULA)"
            )
            try:
                response = diameter.Answer_16777251_316(
                    packet_vars, avps
                )  # Generate Diameter packet
                self.logger.info("Generated ULA")
            except Exception as e:
                self.logger.info("Failed to generate Diameter Response for ULR")
                self.logger.info(e)
                traceback.print_exc()
                response = diameter.Respond_ResultCode(
                    packet_vars, avps, 4100
                )  # DIAMETER_USER_DATA_NOT_AVAILABLE
                self.logger.info("Generated error DIAMETER_USER_DATA_NOT_AVAILABLE ULA")

            # Send ULA data & clear tx buffer
            clientsocket.sendall(bytes.fromhex(response))
            response = ""
            if "Insert_Subscriber_Data_Force" in yaml_config["hss"]:
                if yaml_config["hss"]["Insert_Subscriber_Data_Force"] == True:
                    self.logger.debug("ISD triggered after ULA")
                    # Generate Insert Subscriber Data Request
                    response = diameter.Request_16777251_319(
                        packet_vars, avps
                    )  # Generate Diameter packet
                    self.logger.info("Generated IDR")
                    # Send ISD data
                    send_queue.put(bytes.fromhex(response))
                    # clientsocket.sendall(bytes.fromhex(response))
                    self.logger.info("Sent IDR")
            return
        # S6a inbound Insert-Data-Answer in response to our IDR
        elif (
            packet_vars["command_code"] == 319
            and packet_vars["ApplicationId"] == 16777251
        ):
            self.logger.info(
                f"Received response with command code 319 (3GPP Insert-Subscriber-Answer) from {origin_host}"
            )
            return
        # S6a Purge UE Answer (PUA) response to Purge UE Request (PUR)
        elif (
            packet_vars["command_code"] == 321
            and packet_vars["ApplicationId"] == 16777251
        ):
            self.logger.info(
                f"Received Request with command code 321 (3GPP Purge UE Request) from {origin_host}"
                + "\n\tGenerating (PUA)"
            )
            try:
                response = diameter.Answer_16777251_321(
                    packet_vars, avps
                )  # Generate Diameter packet
            except:
                response = diameter.Respond_ResultCode(
                    packet_vars, avps, 5012
                )  # Generate Diameter response with "DIAMETER_UNABLE_TO_COMPLY" (5012)
                self.logger.error("Failed to generate PUA")
            self.logger.info("Generated PUA")
        # S6a Notify Answer (NOA) response to Notify Request (NOR)
        elif (
            packet_vars["command_code"] == 323
            and packet_vars["ApplicationId"] == 16777251
        ):
            self.logger.info(
                f"Received Request with command code 323 (3GPP Notify Request) from {origin_host}"
                + "\n\tGenerating (NOA)"
            )
            try:
                response = diameter.Answer_16777251_323(
                    packet_vars, avps
                )  # Generate Diameter packet
            except:
                response = diameter.Respond_ResultCode(
                    packet_vars, avps, 5012
                )  # Generate Diameter response with "DIAMETER_UNABLE_TO_COMPLY" (5012)
                self.logger.error("Failed to generate NOA")
            self.logger.info("Generated NOA")
        # S6a Cancel Location Answer eater
        elif (
            packet_vars["command_code"] == 317
            and packet_vars["ApplicationId"] == 16777251
        ):
            self.logger.info("Received Response with command code 317 (3GPP Cancel Location Request) from " + str(origin_host))

        # Cx Authentication Answer
        elif (
            packet_vars["command_code"] == 300
            and packet_vars["ApplicationId"] == 16777216
        ):
            self.logger.info(
                f"Received Request with command code 300 (3GPP Cx User Authentication Request) from {origin_host}"
                + "\n\tGenerating (MAA)"
            )
            try:
                response = diameter.Answer_16777216_300(
                    packet_vars, avps
                )  # Generate Diameter packet
            except Exception as e:
                self.logger.info(
                    "Failed to generate Diameter Response for Cx Auth Answer"
                )
                self.logger.info(e)
                self.logger.info(traceback.print_exc())
                self.logger.info(
                    type(e).__name__,          # TypeError
                    __file__,                  # /tmp/example.py
                    e.__traceback__.tb_lineno  # 2
                )

                response = diameter.Respond_ResultCode(
                    packet_vars, avps, 4100
                )  # DIAMETER_USER_DATA_NOT_AVAILABLE
            self.logger.info("Generated Cx Auth Answer")

        # Cx Server Assignment Answer
        elif (
            packet_vars["command_code"] == 301
            and packet_vars["ApplicationId"] == 16777216
        ):
            self.logger.info(
                f"Received Request with command code 301 (3GPP Cx Server Assignemnt Request) from {origin_host}"
                + "\n\tGenerating (MAA)"
            )
            try:
                response = diameter.Answer_16777216_301(
                    packet_vars, avps
                )  # Generate Diameter packet
            except Exception as e:
                self.logger.info(
                    "Failed to generate Diameter Response for Cx Server Assignment Answer"
                )
                self.logger.info(e)
                traceback.print_exc()
                response = diameter.Respond_ResultCode(
                    packet_vars, avps, 4100
                )  # DIAMETER_USER_DATA_NOT_AVAILABLE
            self.logger.info("Generated Cx Server Assignment Answer")

        # Cx Location Information Answer
        elif (
            packet_vars["command_code"] == 302
            and packet_vars["ApplicationId"] == 16777216
        ):
            self.logger.info(
                f"Received Request with command code 302 (3GPP Cx Location Information Request) from {origin_host}"
                + "\n\tGenerating (MAA)"
            )
            try:
                response = diameter.Answer_16777216_302(
                    packet_vars, avps
                )  # Generate Diameter packet
            except Exception as e:
                self.logger.info(
                    "Failed to generate Diameter Response for Cx Location Information Answer"
                )
                self.logger.info(e)
                traceback.print_exc()
                response = diameter.Respond_ResultCode(
                    packet_vars, avps, 4100
                )  # DIAMETER_USER_DATA_NOT_AVAILABLE
            self.logger.info("Generated Cx Location Information Answer")

        # Cx Multimedia Authentication Answer
        elif (
            packet_vars["command_code"] == 303
            and packet_vars["ApplicationId"] == 16777216
        ):
            self.logger.info(
                f"Received Request with command code 303 (3GPP Cx Multimedia Authentication Request) from {origin_host}"
                + "\n\tGenerating (MAA)"
            )
            try:
                response = diameter.Answer_16777216_303(
                    packet_vars, avps
                )  # Generate Diameter packet
            except Exception as e:
                self.logger.info(
                    "Failed to generate Diameter Response for Cx Multimedia Authentication Answer"
                )
                self.logger.info(e)
                traceback.print_exc()
                response = diameter.Respond_ResultCode(
                    packet_vars, avps, 4100
                )  # DIAMETER_USER_DATA_NOT_AVAILABLE
            self.logger.info("Generated Cx Multimedia Authentication Answer")

        # Sh User-Data-Answer
        elif (
            packet_vars["command_code"] == 306
            and packet_vars["ApplicationId"] == 16777217
        ):
            self.logger.info(
                f"Received Request with command code 306 (3GPP Sh User-Data Request) from {origin_host}"
            )
            try:
                response = diameter.Answer_16777217_306(
                    packet_vars, avps
                )  # Generate Diameter packet
            except Exception as e:
                self.logger.info(
                    "Failed to generate Diameter Response for Sh User-Data Answer"
                )
                self.logger.info(e)
                traceback.print_exc()
                response = diameter.Respond_ResultCode(
                    packet_vars, avps, 5001
                )  # DIAMETER_ERROR_USER_UNKNOWN
                send_queue.put(bytes.fromhex(response))
                # clientsocket.sendall(bytes.fromhex(response))
                self.logger.info("Sent negative response")
                return
            self.logger.info("Generated Sh User-Data Answer")

        # Sh Profile-Update-Answer
        elif (
            packet_vars["command_code"] == 307
            and packet_vars["ApplicationId"] == 16777217
        ):
            self.logger.info(
                f"Received Request with command code 307 (3GPP Sh Profile-Update Request) from {origin_host}"
            )
            try:
                response = diameter.Answer_16777217_307(
                    packet_vars, avps
                )  # Generate Diameter packet
            except Exception as e:
                self.logger.info(
                    "Failed to generate Diameter Response for Sh User-Data Answer"
                )
                self.logger.info(e)
                traceback.print_exc()
                response = diameter.Respond_ResultCode(
                    packet_vars, avps, 5001
                )  # DIAMETER_ERROR_USER_UNKNOWN
                send_queue.put(bytes.fromhex(response))
                # clientsocket.sendall(bytes.fromhex(response))
                self.logger.info("Sent negative response")
                return
            self.logger.info("Generated Sh Profile-Update Answer")

        # S13 ME-Identity-Check Answer
        elif (
            packet_vars["command_code"] == 324
            and packet_vars["ApplicationId"] == 16777252
        ):
            self.logger.info(
                f"Received Request with command code 324 (3GPP S13 ME-Identity-Check Request) from {origin_host}"
                + "\n\tGenerating (MICA)"
            )
            try:
                response = diameter.Answer_16777252_324(
                    packet_vars, avps
                )  # Generate Diameter packet
            except Exception as e:
                self.logger.info(
                    "Failed to generate Diameter Response for S13 ME-Identity Check Answer"
                )
                self.logger.info(e)
                traceback.print_exc()
                response = diameter.Respond_ResultCode(
                    packet_vars, avps, 4100
                )  # DIAMETER_USER_DATA_NOT_AVAILABLE
            self.logger.info("Generated S13 ME-Identity Check Answer")

        # SLh LCS-Routing-Info-Answer
        elif (
            packet_vars["command_code"] == 8388622
            and packet_vars["ApplicationId"] == 16777291
        ):
            self.logger.info(
                f"Received Request with command code 324 (3GPP SLh LCS-Routing-Info-Answer Request) from {origin_host}"
                + "\n\tGenerating (MICA)"
            )
            try:
                response = diameter.Answer_16777291_8388622(
                    packet_vars, avps
                )  # Generate Diameter packet
            except Exception as e:
                self.logger.info(
                    "Failed to generate Diameter Response for SLh LCS-Routing-Info-Answer"
                )
                self.logger.info(e)
                traceback.print_exc()
                response = diameter.Respond_ResultCode(
                    packet_vars, avps, 4100
                )  # DIAMETER_USER_DATA_NOT_AVAILABLE
            self.logger.info("Generated SLh LCS-Routing-Info-Answer")

        # Handle Responses generated by the Async functions
        elif packet_vars["flags"] == "00":
            self.logger.info(
                "Got response back with command code "
                + str(packet_vars["command_code"])
            )
            self.logger.info("response packet_vars: " + str(packet_vars))
            self.logger.info("response avps: " + str(avps))
            response = ""
        else:
            self.logger.error(
                "\n\nRecieved unrecognised request with Command Code: "
                + str(packet_vars["command_code"])
                + ", ApplicationID: "
                + str(packet_vars["ApplicationId"])
                + " and flags "
                + str(packet_vars["flags"])
            )
            for keys in packet_vars:
                self.logger.error(keys)
                self.logger.error("\t" + str(packet_vars[keys]))
            self.logger.error(avps)
            self.logger.error("Sending negative response")
            response = diameter.Respond_ResultCode(
                packet_vars, avps, 3001
            )  # Generate Diameter response with "Command Unsupported" (3001)
            send_queue.put(bytes.fromhex(response))
            # clientsocket.sendall(bytes.fromhex(response))  # Send it

        prom_diam_response_time_method.labels(
            str(packet_vars["ApplicationId"]),
            str(packet_vars["command_code"]),
            origin_host,
            "request",
        ).observe(time.time() - start_time)

        # Diameter Transmission
        retries = 0
        while retries < self.max_diameter_retries:
            try:
                send_queue.put(bytes.fromhex(response))
                break
            except socket.error as e:
                self.logger.error(f"Socket error for client {client_address}: {e}")
                retries += 1
                if retries > self.max_diameter_retries:
                    self.logger.error(
                        f"Max retries reached for client {client_address}. Closing connection."
                    )
                    self.terminate_connection(
                        clientsocket, client_address, thread_event
                    )
                    break
                time.sleep(1)  # Wait for 1 second before retrying
            except Exception as e:
                self.logger.info("Failed to send Diameter Response")
                self.logger.debug(f"Diameter Response Body: {str(response)}")
                self.logger.info(e)
                traceback.print_exc()
                self.terminate_connection(clientsocket, client_address, thread_event)
                self.logger.info("Thread terminated to " + str(client_address))
                break

    def manage_client(
        self, clientsocket, client_address, diameter, thread_event, send_queue
    ):
        while True:
            try:
                data = clientsocket.recv(32)
                if not data:
                    self.logger.info(
                        f"manage_client: Connection closed by {str(client_address)}"
                    )
                    self.terminate_connection(
                        clientsocket, client_address, thread_event
                    )
                    return
                self.process_Diameter_request(
                    clientsocket,
                    client_address,
                    diameter,
                    data,
                    thread_event,
                    send_queue,
                )

            except socket.timeout:
                self.logger.warning(
                    f"manage_client: Socket timeout for client: {client_address}"
                )
                self.terminate_connection(clientsocket, client_address, thread_event)
                return

            except socket.error as e:
                self.logger.error(
                    f"manage_client: Socket error for client {client_address}: {e}"
                )
                self.terminate_connection(clientsocket, client_address, thread_event)
                return

            except KeyboardInterrupt:
                # Clean up the connection on keyboard interrupt
                response = (
                    diameter.Request_282()
                )  # Generate Disconnect Peer Request Diameter packet
                send_queue.put(bytes.fromhex(response))
                # clientsocket.sendall(bytes.fromhex(response))  # Send it
                self.terminate_connection(clientsocket, client_address, thread_event)
                self.logger.info(
                    "manage_client: Connection closed nicely due to keyboard interrupt"
                )
                sys.exit()

            except Exception as manage_client_exception:
                self.logger.error(
                    f"manage_client: Exception in manage_client: {manage_client_exception}"
                )
                self.terminate_connection(clientsocket, client_address, thread_event)
                return

    def manage_client_async(
        self, clientsocket, client_address, diameter, thread_event, send_queue
    ):
        # # Sleep for 10 seconds to wait for the connection to come up
        time.sleep(10)
        self.logger.debug("manage_client_async: Getting ActivePeerDict")
        self.logger.debug(
            f"manage_client_async: Total PyHSS Active Threads: {threading.active_count()}"
        )
        ActivePeerDict = self.logtool.GetDiameterPeers()
        self.logger.debug(
            f"manage_client_async: Got Active Peer dict in Async Thread: {str(ActivePeerDict)}"
        )
        if client_address[0] in ActivePeerDict:
            self.logger.debug(
                "manage_client_async: This is host: "
                + str(ActivePeerDict[str(client_address[0])]["DiameterHostname"])
            )
            DiameterHostname = str(
                ActivePeerDict[str(client_address[0])]["DiameterHostname"]
            )
        else:
            self.logger.debug("manage_client_async: No matching Diameter Host found.")
            return

        while True:
            try:
                if thread_event.is_set():
                    self.logger.debug(
                        f"manage_client_async: Closing manage_client_async thread for client: {client_address}"
                    )
                    self.terminate_connection(
                        clientsocket, client_address, thread_event
                    )
                    return
                time.sleep(self.yaml_config["hss"]["async_check_interval"])
                self.logger.debug(
                    f"manage_client_async: Sleep interval expired for Diameter Peer {str(DiameterHostname)}"
                )
                if int(self.yaml_config["hss"]["async_check_interval"]) == 0:
                    self.logger.error(
                        f"manage_client_async: No async_check_interval Timer set - Not checking Async Queue for host connection {str(DiameterHostname)}"
                    )
                    return
                try:
                    self.logger.debug(
                        "manage_client_async: Reading from request queue '"
                        + str(DiameterHostname)
                        + "_request_queue'"
                    )
                    data_to_send = self.logtool.RedisHMGET(
                        str(DiameterHostname) + "_request_queue"
                    )
                    for key in data_to_send:
                        data = data_to_send[key].decode("utf-8")
                        send_queue.put(bytes.fromhex(data))
                        self.logtool.RedisHDEL(
                            str(DiameterHostname) + "_request_queue", key
                        )
                except Exception as redis_exception:
                    self.logger.error(
                        f"manage_client_async: Redis exception in manage_client_async: {redis_exception}"
                    )
                    self.terminate_connection(
                        clientsocket, client_address, thread_event
                    )
                    return

            except socket.timeout:
                self.logger.warning(
                    f"manage_client_async: Socket timeout for client: {client_address}"
                )
                self.terminate_connection(clientsocket, client_address, thread_event)
                return

            except socket.error as e:
                self.logger.error(
                    f"manage_client_async: Socket error for client {client_address}: {e}"
                )
                self.terminate_connection(clientsocket, client_address, thread_event)
                return
            except Exception:
                self.logger.error(
                    f"manage_client_async: Terminating for host connection {str(DiameterHostname)}"
                )
                self.terminate_connection(clientsocket, client_address, thread_event)
                return

    def manage_client_dwr(
        self, clientsocket, client_address, diameter, thread_event, send_queue
    ):
        while True:
            try:
                if thread_event.is_set():
                    self.logger.debug(
                        f"Closing manage_client_dwr thread for client: {client_address}"
                    )
                    self.terminate_connection(
                        clientsocket, client_address, thread_event
                    )
                    return
                if (
                    int(self.yaml_config["hss"]["device_watchdog_request_interval"])
                    != 0
                ):
                    time.sleep(
                        self.yaml_config["hss"]["device_watchdog_request_interval"]
                    )
                else:
                    self.logger.info("DWR Timer to set to 0 - Not sending DWRs")
                    return

            except:
                self.logger.error(
                    "No DWR Timer set - Not sending Device Watchdog Requests"
                )
                return
            try:
                self.logger.debug("Sending Keepalive to " + str(client_address) + "...")
                request = diameter.Request_280()
                send_queue.put(bytes.fromhex(request))
                # clientsocket.sendall(bytes.fromhex(request))  # Send it
                self.logger.debug("Sent Keepalive to " + str(client_address) + "...")
            except socket.error as e:
                self.logger.error(
                    f"manage_client_dwr: Socket error for client {client_address}: {e}"
                )
                self.terminate_connection(clientsocket, client_address, thread_event)
                return
            except Exception as e:
                self.logger.error(
                    f"manage_client_dwr: General exception for client {client_address}: {e}"
                )
                self.terminate_connection(clientsocket, client_address, thread_event)

    def get_socket_family(self):
        if ":" in self.yaml_config["hss"]["bind_ip"][0]:
            self.logger.info("IPv6 Address Specified")
            return socket.AF_INET6
        else:
            self.logger.info("IPv4 Address Specified")
            return socket.AF_INET

    def send_data(self, clientsocket, send_queue, thread_event):
        while not thread_event.is_set():
            try:
                data = send_queue.get(timeout=1)
                # Check if data is bytes, otherwise convert it using bytes.fromhex()
                if not isinstance(data, bytes):
                    data = bytes.fromhex(data)

                clientsocket.sendall(data)
            except (
                queue.Empty
            ):  # Catch the Empty exception when the queue is empty and the timeout has expired
                continue
            except Exception as e:
                self.logger.error(f"send_data_thread: Exception: {e}")
                return

    def start_server(self):
        if self.yaml_config["hss"]["transport"] == "SCTP":
            self.logger.debug("Using SCTP for Transport")
            # Create a SCTP socket
            sock = sctp.sctpsocket_tcp(self.get_socket_family())
            sock.initparams.num_ostreams = 64
            # Loop through the possible Binding IPs from the config and bind to each for Multihoming
            server_addresses = []

            # Prepend each entry into list, so the primary IP is bound first
            for host in self.yaml_config["hss"]["bind_ip"]:
                self.logger.info("Seting up SCTP binding on IP address " + str(host))
                this_IP_binding = [
                    (str(host), int(self.yaml_config["hss"]["bind_port"]))
                ]
                server_addresses = this_IP_binding + server_addresses

            print("server_addresses are: " + str(server_addresses))
            sock.bindx(server_addresses)
            self.logger.info("PyHSS listening on SCTP port " + str(server_addresses))
            systemd.daemon.notify("READY=1")
            # Listen for up to 20 incoming SCTP connections
            sock.listen(20)
        elif self.yaml_config["hss"]["transport"] == "TCP":
            self.logger.debug("Using TCP socket")
            # Create a TCP/IP socket
            sock = socket.socket(self.get_socket_family(), socket.SOCK_STREAM)
            sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            # Bind the socket to the port
            server_address = (
                str(self.yaml_config["hss"]["bind_ip"][0]),
                int(self.yaml_config["hss"]["bind_port"]),
            )
            sock.bind(server_address)
            self.logger.debug(
                "PyHSS listening on TCP port "
                + str(self.yaml_config["hss"]["bind_ip"][0])
            )
            systemd.daemon.notify("READY=1")
            # Listen for up to 20 incoming TCP connections
            sock.listen(20)
        else:
            self.logger.error("No valid transports found (No SCTP or TCP) - Exiting")
            quit()

        while True:
            # Wait for a connection
            self.logger.info("Waiting for a connection...")
            connection, client_address = sock.accept()
            _thread.start_new_thread(
                self.handle_new_connection,
                (
                    connection,
                    client_address,
                ),
            )

    def prom_async_function(self):
        while True:
            self.logger.debug("Running prom_async_function")
            self.diameter_instance.Generate_Prom_Stats()
            time.sleep(120)

if __name__ == "__main__":
    pyHss = PyHSS()
    pyHss.start_server()


# -*- coding: utf-8 -*-
from openerp.report import report_sxw
from openerp.report.report_sxw import rml_parse

from datetime import datetime
DATETIME_FORMAT = "%Y-%m-%d %H:%M:%S"
DATE_FORMAT = "%Y-%m-%d"

class Parser(rml_parse):

    def __init__(self, cr, uid, name, context):
        super(self.__class__, self).__init__(cr, uid, name, context)
        self.start_date = False
        self.end_date = False
        self.asset_types = []

        self.company_name = False
        self.company_address = False
        self.vat = False
        self.total_residual = False
        self.depreciation_value = False
        self.acc_depreciation = False
        self.remain_value = False

        self.data = {}

        self.localcontext.update({
            'get_vietname_date':self.get_vietname_date,
            'get_header':self.get_header,
            'get_start_date':self.get_start_date,
            'get_end_date':self.get_end_date,
            'get_company_name':self.get_company_name,
            'get_company_address':self.get_company_address,
            'get_company_vat':self.get_company_vat,

            'get_date_type':self.get_date_type,

            'get_wh':self.get_wh,
            'get_dept': self.get_dept,
            'get_line':self.get_line,

            'get_depreciation_value':self.get_depreciation_value,
            'get_acc_depreciation':self.get_acc_depreciation,
            'get_remaining_value':self.get_remaining_value,
        })

    def get_company(self):
        if self.company_id:
            company_obj = self.pool.get('res.company').browse(self.cr, self.uid, self.company_id)
            self.company_name = company_obj.name or ''
            self.company_address = company_obj.street or ''
            self.vat = company_obj.vat or ''
        return True

    def get_company_name(self):
        self.get_header()
        return self.company_name

    def get_company_address(self):
        return self.company_address

    def get_date_type(self):
        if self.date_type == 'depreciation_date':
            return u'Ngày khấu hao'
        if self.date_type == 'using_date':
            return u'Ngày đưa vào sử dụng'
        if self.date_type == 'purchase_date':
            return u'Ngày mua'
        return ''

    def get_company_vat(self):
        return self.vat

    def get_header(self):
        wizard_data = self.localcontext['data']['form']

        self.company_id = wizard_data['company_id'][0]
        self.start_date = wizard_data['date_start']
        self.end_date   = wizard_data['date_end']
        self.date_type = wizard_data['date_type']
        self.show_all = wizard_data['show_all']
        self.asset_types = (','.join(map(str, wizard_data['asset_category_ids'])))
        self.get_company()

        #THANH: Load all data
        res, depreciation_value, acc_depreciation, remain_value = self.pool.get('sql.depreciation.allocation').get_line(self.cr, self.date_type, self.show_all, self.start_date, self.end_date, self.asset_types, self.company_id)
        self.depreciation_value = depreciation_value
        self.acc_depreciation = acc_depreciation
        self.remain_value = remain_value
        self.data = res

    def get_start_date(self):
        return self.get_vietname_date(self.start_date)

    def get_end_date(self):
        return self.get_vietname_date(self.end_date)

    def get_convert_double(self,value):
        return float(value)

    def get_vietname_date(self, date):
        if not date:
            return ''
        date = datetime.strptime(date, DATE_FORMAT)
        return date.strftime('%d/%m/%Y')

    def get_wh(self):
        wh = []
        for key, val in self.data.items():
            wh.append(key)
        return wh

    def get_dept(self, wh):
        dept = []
        for key, val in self.data.items():
            if key == wh:
                for dept_key, dept_vals in val.items():
                    dept.append(dept_key)
        return dept

    def get_line(self, wh, dept):
        return self.data[wh][dept]

    def get_depreciation_value(self):
        return self.depreciation_value

    def get_acc_depreciation(self):
        return self.acc_depreciation

    def get_remaining_value(self):
        return self.remain_value

# vim:expandtab:smartindent:tabstop=4:softtabstop=4:shiftwidth=4:


import os
import requests
import pandas as pd
import openpyxl
import re
import math

from PyQt6.QtCore import *
from PyQt6.QtGui import *
from PyQt6.QtWidgets import *

import sys
from bs4 import BeautifulSoup

country_websites = pd.read_excel('website_data.xlsx')
country_websites_dict = country_websites.set_index('Country').to_dict()['Url']

class CountrySelector(QWidget):

    def __init__(self, title, country_website_base):
        QWidget.__init__(self)
        layout = QHBoxLayout()
        self.setLayout(layout)
        self.downloadURL = '/ww/en/'

        self.label = QLabel()
        self.label.setText(title)
        self.label.setFixedWidth(130)
        self.label.setFont(QFont("Arial"))
        self.label.setAlignment(Qt.AlignmentFlag.AlignLeft | Qt.AlignmentFlag.AlignVCenter)
        layout.addWidget(self.label)

        self.countryList = QComboBox()
        self.countryList.addItems(country_websites_dict.keys())

        self.countryList.currentIndexChanged.connect(self.setCountry)
        layout.addWidget(self.countryList)
        layout.addStretch()

    def setCountry(self, i):
        self.downloadURL = list(country_websites_dict.values())[i]

    def setLabelWidth(self, width):
        self.label.setFixedWidth(width)
        # --------------------------------------------------------------------

    def setlineEditWidth(self, width):
        self.lineEdit.setFixedWidth(width)

    def setlineEditWidth(self, width):
        self.lineEdit.setFixedWidth(width)


from post.models import PostModel
from post.api.serializers import PostDetailSerializer
from rest_framework import serializers

class MyPostsSerializer(PostDetailSerializer):
    image = serializers.SerializerMethodField(method_name='get_post_images', source='post_images')

    class Meta:
        model = PostModel
        fields = ('id', 'is_favourite', 'image', 'price', 'price_type', 'brand', 'auto_model', 'year', 'engine_volume', 'mileage', 'city', 'creation_date')

    def get_post_images(self, obj):
        return self.context['request'].build_absolute_uri(obj.post_images.first().image.url)


#!/usr/bin/env python
#-*- coding: utf-8 -*-
#If this code works, it was written by Seunghyun Lee(www.bluebaynetworks.co.kr).
#If not, I don't know who wrote it

"""Functions for downloading and reading Korean Font data."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os
import cv2
import numpy as np
import tensorflow as tf

class DataSet(object):

  def __init__(self, images, labels, font_cnt, dtype=tf.float32):
    dtype = tf.as_dtype(dtype).base_dtype
    if dtype not in (tf.uint8, tf.float32):
      raise TypeError('Invalid image dtype %r, expected uint8 or float32' %
                      dtype)
    assert images.shape[0] == labels.shape[0], (
        'images.shape: %s labels.shape: %s' % (images.shape,
                                                labels.shape))
    self._num_examples = images.shape[0]

    # Convert shape from [num examples, rows, columns, depth]
    # to [num examples, rows*columns] (assuming depth == 1)
    assert images.shape[3] == 1
    images = images.reshape(images.shape[0], images.shape[1] * images.shape[2])
    if dtype == tf.float32:
      # Convert from [0, 255] -> [0.0, 1.0].
      images = images.astype(np.float32)
      images = np.multiply(images, 1.0 / 255.0)

    self._images = images
    self._labels = labels
    self._epochs_completed = 0
    self._index_in_epoch = 0
    self._font_cnt = font_cnt

  @property
  def images(self):
    return self._images

  @property
  def labels(self):
    return self._labels

  @property
  def num_examples(self):
    return self._num_examples

  @property
  def epochs_completed(self):
    return self._epochs_completed

  def find_label_index(self, val):
    for i in range(0, self._font_cnt):
      if(self._labels[i] == val):
        return i
    return None

  def next_image(self):
    ret_label = np.zeros((1, self._font_cnt), dtype=np.uint8)
    start = self._index_in_epoch
    self._index_in_epoch += 1

    if self._index_in_epoch == self._num_examples:
      self._index_in_epoch = 0
      #print('********** all char trained --> reset **********')

    ret_label[0][start] = 1

    return self._images[start:start + 1], ret_label

  def next_batch(self, batch_size):
    start = self._index_in_epoch
    self._index_in_epoch += batch_size

    if self._index_in_epoch >= self._num_examples:
      end = self._num_examples
      self._index_in_epoch = 0
      #print('********** all char trained --> reset **********')
    else:
      end = start + batch_size

    ret_label = np.zeros((end - start, self._font_cnt), dtype=np.uint8)
    index = 0
    #print('********** start:%d  end:%d **********'%(start, end))
    for i in range(start, end):
      ret_label[index][i] = 1
      index += 1

    return self._images[start:end], ret_label

#파일 이름에 포함된 index 숫자와 파일 내용을 리턴
def extract_image(filename):
  #Open font jpg files
  img = cv2.imread(filename)
  img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
  base=os.path.basename(filename)
  os.path.splitext(base)
  index = int(os.path.splitext(base)[0])
  return index, img

"""
모든 폰트 이미지를 검색해 파일 내용을 읽는다.
항상 같은 순서로 읽어야 하기 때문에 파일을 소팅해서 읽는 것이 안전하다.
특히 플랫폼이 바뀔 경우에 필요하다.
"""
def extract_images(dir_name, font_cnt):
  data = None
  labels = np.zeros(font_cnt, dtype=np.int32)
  count = 0
  for subdir, dirs, files in sorted(os.walk(dir_name)):
    for file in sorted(files):
      filename = os.path.join(subdir, file)
      index, ret = extract_image(filename)
      label = np.zeros(font_cnt, dtype=np.int32)
      labels[count] = index
      count += 1
      if(data is None):
        data = ret
      else:
        data = np.append(data, ret)
  data = data.reshape(count, 28, 28, 1)
  print('Data shape:', data.shape, ' Label shape:', labels.shape )
  return data, labels

"""
label은 MNIST와 달리 바로 사용하지 못한다.
실시간으로 batch 사이즈에 맞게 만들어 줘야 한다.
"""
def read_data_sets(train_dir, font_cnt):
  VALIDATION_SIZE = 500

  class DataSets(object):
    pass

  data_sets = DataSets()
  basicdir = train_dir + 'basic'
  shiftdir = train_dir + 'shift'
  rotatedir = train_dir + 'rotate'

  print('--------------- Loading basic font image --------------------')
  images, labels = extract_images(basicdir, font_cnt)
  data_sets.basic = DataSet(images, labels, font_cnt)
  print('--------------- Loading Shift1 font image --------------------')
  sdir = shiftdir + "1"
  images, labels = extract_images(sdir, font_cnt)
  data_sets.shift1 = DataSet(images, labels, font_cnt)
  print('--------------- Loading Shift2 font image --------------------')
  sdir = shiftdir + "2"
  images, labels = extract_images(sdir, font_cnt)
  data_sets.shift2 = DataSet(images, labels, font_cnt)
  print('--------------- Loading Shift3 font image --------------------')
  sdir = shiftdir + "3"
  images, labels = extract_images(sdir, font_cnt)
  data_sets.shift3 = DataSet(images, labels, font_cnt)
  print('--------------- Loading Shift4 font image --------------------')
  sdir = shiftdir + "4"
  images, labels = extract_images(sdir, font_cnt)
  data_sets.shift4 = DataSet(images, labels, font_cnt)

  print('--------------- Loading Rotate1 font image --------------------')
  sdir = rotatedir + "1"
  images, labels = extract_images(sdir, font_cnt)
  data_sets.rotate1 = DataSet(images, labels, font_cnt)
  print('--------------- Loading Rotate2 font image --------------------')
  sdir = rotatedir + "2"
  images, labels = extract_images(sdir, font_cnt)
  data_sets.rotate2 = DataSet(images, labels, font_cnt)
  print('--------------- Loading Rotate3 font image --------------------')
  sdir = rotatedir + "3"
  images, labels = extract_images(sdir, font_cnt)
  data_sets.rotate3 = DataSet(images, labels, font_cnt)
  print('--------------- Loading Rotate4 font image --------------------')
  sdir = rotatedir + "4"
  images, labels = extract_images(sdir, font_cnt)
  data_sets.rotate4 = DataSet(images, labels, font_cnt)

  return data_sets

08 machine learning/1.0 Version code/input_korean.py


import pickle
from omegaconf import OmegaConf
from skilltranslation.envs.blockstacking import BlockStackFloatPandaEnv
import gym
import numpy as np
from mani_skill2.utils.wrappers import ManiSkillActionWrapper, NormalizeActionWrapper
import os.path as osp
from skilltranslation.envs.blockstacking.env import BlockStackMagicPandaEnv
from skilltranslation.utils.sampling import resample_teacher_trajectory

def animate(imgs, filename="animation.mp4", _return=True, fps=10):
    if isinstance(imgs, dict):
        imgs = imgs["image"]
    print(f"animating {filename}")
    from moviepy.editor import ImageSequenceClip

    imgs = ImageSequenceClip(imgs, fps=fps)
    imgs.write_videofile(filename, fps=fps)
    if _return:
        from IPython.display import Video

        return Video(filename, embed=True)

def collect_magic_teacher_traj(env: BlockStackMagicPandaEnv, render=False, mode='human', change_in_height=0, change_in_goal_height=0):
    assert env.connected == False

    assert env.goal_coords is not None
    init_xyz = env.get_articulations()[0].get_qpos()[:3]
    magic_traj=dict(
        observations=[],
        attns=[],
    )
    vids=[]
    if render and mode == 'human':
        env.render()
        env.get_viewer().paused=True
    i_step=0

    def goto(xyz, gripper_action, attn=0):
        assert gripper_action is None
        nonlocal i_step, block
        target=np.array(xyz)
        observations=[]  # during this goto()
        attns = []
        while True:
            robot_qpos=env.get_articulations()[0].get_qpos()
            delta_xyz=target - robot_qpos[:3]
            robot_to_target_dist=np.linalg.norm(delta_xyz)
            vel=np.linalg.norm(env.unwrapped._agent._robot.get_qvel()[:3])
            if robot_to_target_dist < 0.01 and vel < 0.05:
                break
            action=np.zeros(4)
            action[:3]=delta_xyz.copy()
            # action[-1]=gripper_action
            # print(action)
            if 'NormalizeActionWrapper' in str(env):
                action[:3]=action[:3] / env.env.action_space.high[0]
                # gripper always +1
                action[-1]=+1
            # collect observations
            obs_dict=env.get_obs()
            dense_obs=np.zeros(12)
            dense_obs[:5]=obs_dict['agent']['qpos']
            # dense_obs[5]=obs_dict['agent']['qvel']
            dense_obs[5:]=obs_dict['extra']['obs_blocks']['absolute'][attn * 7: attn * 7 + 7]
            observations.append(dense_obs)
            attns.append(attn)
            # attentions.append(attn)
            # print("{}: dist {:.4f}, act {}, cur xyz {}, target xyz {}".format(
            #     i_step, robot_to_target_dist, action, robot_qpos[:3], xyz
            # ))
            assert action[-1] == +1
            env.step(action)

            if render:
                if mode == 'rgb_array':
                    vids.append(env.render(mode))
                else:
                    env.render(mode)
            i_step+=1
        magic_traj['observations']+=observations
        magic_traj["attns"] += attns

    for i in range(len(env.blocks)):
        block=env.blocks[i]
        block_pos=env.blocks[i].pose.p
        # above_pos=[block_pos[0], block_pos[1], 0.25]
        # grasp_pos=[block_pos[0], block_pos[1], 0.04 + 0.112]
        above_pos=[block_pos[0], block_pos[1], 0.25+change_in_height]
        grasp_pos=[block_pos[0], block_pos[1], 0.113]

        goto(xyz=above_pos, gripper_action=None, attn=i)
        goto(xyz=grasp_pos, gripper_action=None, attn=i)
        env.magic_grasp(block)
        goto(xyz=above_pos, gripper_action=None, attn=i)

        goal_pos=env.goal_coords[i]
        # goal_above_pos=[goal_pos[0], goal_pos[1], 0.25]
        # release_pos=[goal_pos[0], goal_pos[1], goal_pos[2] + 0.132]
        goal_above_pos=[goal_pos[0], goal_pos[1], goal_pos[2]+0.25+ change_in_goal_height]
        release_pos=[goal_pos[0], goal_pos[1], goal_pos[2] + 0.096]

        goto(xyz=goal_above_pos, gripper_action=None, attn=i)
        goto(xyz=release_pos, gripper_action=None, attn=i)
        env.magic_release()
        goto(xyz=goal_above_pos, gripper_action=None, attn=i)

        # back to init
        print("go back to init")
        goto(xyz = init_xyz, gripper_action=None, attn=i)
        print("finished init")
        # last observation and attention:
    obs_dict=env.get_obs()
    dense_obs=np.zeros(12)
    dense_obs[:5]=obs_dict['agent']['qpos']
    dense_obs[5:]=obs_dict['extra']['obs_blocks']['absolute'][i * 7:i * 7 + 7]
    magic_traj['observations'].append(dense_obs)

    magic_traj['attns'] = np.array(magic_traj['attns'], dtype=int)
    magic_traj['observations'] = np.array(magic_traj['observations'], dtype=np.float32)
    # assert len(float_traj['attentions']) == len(float_traj['observations'])
    if render and mode == 'rgb_array':
        return magic_traj, vids

    if not render:
        return magic_traj
    return magic_traj

def collect_teacher_traj(env: BlockStackFloatPandaEnv, render=False, mode='human', fixed_gripper=True):
    assert env.goal_coords is not None
    float_traj = dict(
        observations=[],
        attns=[],
    )
    vids = []
    init_xyz = env.get_articulations()[0].get_qpos()[:3]
    if render and mode=='human':
        env.render()
        env.get_viewer().paused = True

    i_step = 0

    def goto(xyz, gripper_action, attn=0):
        nonlocal i_step, block
        target = np.array(xyz)
        observations = []
        attns = []
        while True:
            ee_pos = env.unwrapped._agent.hand_link.get_pose().p
            delta_xyz = target - ee_pos
            robot_to_target_dist = np.linalg.norm(delta_xyz)
            vel = np.linalg.norm(env.unwrapped._agent._robot.get_qvel()[:3])

            grasp_flag = env.unwrapped._agent.check_grasp(block)
            gripper_flag = grasp_flag != (gripper_action > 0)

            if robot_to_target_dist < 0.01 and vel<0.05 and gripper_flag:
                break
            action = np.zeros(4)
            action[:3] = delta_xyz.copy()
            action[-1] = gripper_action
            # print(action)
            if 'NormalizeActionWrapper' in str(env):
                action[:3] = action[:3] / env.env.action_space.high[0]
                action[-1] = 1 if gripper_action > 0 else -1
            obs_dict=env.get_obs()
            dense_obs=np.zeros(12)
            dense_obs[:5]=obs_dict['agent']['qpos']
            if fixed_gripper:
                dense_obs[3:5] = np.array([0.04, 0.04])
            # dense_obs[5:10]=obs_dict['agent']['qvel']
            dense_obs[5:]=obs_dict['extra']['obs_blocks']['absolute'][attn * 7: attn * 7 + 7]
            observations.append(dense_obs)
            attns.append(attn)
            # print("{}: dist {:.4f}, act {}, cur xyz {}, target xyz {}".format(
            #     i_step, robot_to_target_dist, action, ee_pos, xyz
            # ))
            # print(env.unwrapped._agent.hand_link.get_pose())
            env.step(action)

            if render:
                if mode == 'rgb_array':
                    vids.append(env.render(mode))
                else:
                    env.render(mode)
            i_step += 1
        float_traj["observations"] += observations
        float_traj["attns"] += attns

    for i in range(len(env.blocks)):
        block = env.blocks[i]
        block_pos = env.blocks[i].pose.p
        goal_pos = env.goal_coords[i]
        transport_height = goal_pos[2] + 0.12

        above_pos = [block_pos[0], block_pos[1], 0.25]
        grasp_pos = [block_pos[0], block_pos[1], 0.113]
        start_transport_pos = [block_pos[0], block_pos[1], transport_height]
        end_transport_pos = [goal_pos[0] - 0.04, goal_pos[1], transport_height]

        goto(xyz=above_pos, gripper_action=0.04, attn=i)
        goto(xyz=grasp_pos, gripper_action=0.04, attn=i)
        goto(xyz=grasp_pos, gripper_action=0, attn=i)
        goto(xyz=start_transport_pos, gripper_action=0, attn=i)
        goto(xyz=end_transport_pos, gripper_action=0, attn=i)

        goal_above_pos = [goal_pos[0], goal_pos[1], transport_height]
        release_pos = [goal_pos[0], goal_pos[1], goal_pos[2] + 0.096]

        goto(xyz=goal_above_pos, gripper_action=0, attn=i)
        goto(xyz=release_pos, gripper_action=0, attn=i)
        goto(xyz=release_pos, gripper_action=0.04, attn=i)
        goto(xyz=goal_above_pos, gripper_action=0.04, attn=i)
        goto(xyz=end_transport_pos, gripper_action=0.04, attn=i)
        goto(xyz = init_xyz, gripper_action=0.04, attn=i)

    # if render and mode == 'rgb_array':
    #     return vids
    float_traj['attns'] = np.array(float_traj['attns'], dtype=int)
    float_traj['observations'] = np.array(float_traj['observations'], dtype=np.float32)
    return float_traj

if __name__ == '__main__':
    cfg = OmegaConf.from_cli()
    from tqdm import tqdm
    np.set_printoptions(suppress=True, precision=3)
    save_path = cfg.save_path
    # size = cfg.size
    goal = cfg.goal
    N = 2
    if "n" in cfg:
        N = cfg.n
    env_name = 'BlockStackMagic-v0'
    env = gym.make(
        env_name,
        reward_mode='sparse',
        obs_mode='state_dict',
        num_blocks=1 if "train" in goal else -1,
        goal=goal,
    )
    env = ManiSkillActionWrapper(env)
    env = NormalizeActionWrapper(env)
    dataset = dict(teacher=dict())
    for i in tqdm(range(N)):
        env.reset(seed=i)
        # change_height = np.random.uniform(0,0.36)
        # change_goal_height = np.random.uniform(0,0.15)
        traj = collect_magic_teacher_traj(env, render=True, mode='human', change_in_goal_height=0.1, change_in_height=0.18)
        max_block_id = traj["attns"].max()
        all_obs = []
        attns = []
        for b_id in sorted(np.unique(traj["attns"])):
            obs = traj["observations"][np.where(traj["attns"] == b_id)]
            obs = resample_teacher_trajectory(obs, max_dist=4e-2)
            dataset["teacher"][f"{i}-{b_id}"] = dict(observations=obs)
            attns.append(np.ones(len(obs), dtype=int)*b_id)
            all_obs.append(obs)
        all_obs = np.concatenate(all_obs)
        attns = np.concatenate(attns)
        dataset["teacher"][f"{i}"] = dict(observations=all_obs, attns=attns)
        # do even resampling of teacher trajectory.
        # traj = resample_teacher_trajectory(traj["observations"])
        # dataset["traj"]
        # import pdb;pdb.set_trace()
        # vids = collect_teacher_traj(env, render=True, mode='rgb_array')
        # animate(vids, 'videos/teacher_{:d}.mp4'.format(i), fps=24)
    with open(osp.join(save_path), "wb") as f:
        pickle.dump(dataset, f)


# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right
class Solution:
    def inorderTraversal(self, root: Optional[TreeNode]) -> List[int]:
        t = []

        def inorder(root):
            '''
            Traverse the left node to end of tree, add the value, return up a node, try the right traversal. repeat.
            '''
            if root is None:
                return
            # traverse the left node to the end of the tree
            inorder(root.left)
            # append the value of the end node
            t.append(root.val)
            # traverse the right of the tree
            inorder(root.right)

        inorder(root)

        return t


from shutil import move, copy
from os import path, mkdir, getlogin
from time import sleep

print("-="*50)

#Obter o nome do projeto e fazer a verificação da existência do arquivo.
while True:
    nome_Projeto = str(input("Digite o nome do projeto: "))
    diretorio = (f"c:/Users/{getlogin()}/Documents/{nome_Projeto}")
    verificacao = path.exists(diretorio)
    if verificacao == True:
        print("Esse projeto já existe!")
    else:
        break

#Opção de adicionar o arquivo Readme
while True:
    readme_File = str(input("Deseja adicionar o arquivo Readme (yes/no): "))
    if readme_File != "yes" and readme_File != "no":
        print("Tente novamente!")
    else:
        break

#Diretorios
assets = str(f"{diretorio}/assets")
css_Dir = str(f"{diretorio}/assets/css")
javaS_Dir = str(f"{diretorio}/assets/js")

#Criação das pastas
mkdir(diretorio)
mkdir(assets)
mkdir(css_Dir)
mkdir(javaS_Dir)

#Criação dos arquivos
if readme_File == "yes":
    with open("README.md", 'w') as readme:
        readme.write(f"# {nome_Projeto}")
    move("README.md", diretorio)

with open ("style.css", 'w') as css:
    pass

with open ("script.js", 'w') as js:
    pass

#Mover arquivos para o destino final.
move("style.css", css_Dir)
move("script.js", javaS_Dir)
copy("Executavel\index.html", diretorio)

print("-="*50)
print("Projeto criado com sucesso!")
print("-="*50)
sleep(5)


from flask import Flask, render_template, request, redirect, url_for
from flask_wtf import FlaskForm
from wtforms import StringField, BooleanField
from wtforms.validators import DataRequired, Email
from contentful import Client
from trycourier import Courier
from flaskext.markdown import Markdown
import pendulum
import os
import uuid

app = Flask(__name__)
app.secret_key = os.environ.get("secret_key")
Markdown(app)

CTFLclient = Client(
    os.environ.get("CTFL_SpaceID"),
    os.environ.get("CTFL_CDA"),
    environment="master",
)

courierClient = Courier()

@app.template_filter("strftime")
def _jinja2_filter_datetime(date, fmt="dddd MMMM Do YYYY"):
    date = pendulum.parse(f"{date.year}-{date.month}-{date.day}", strict=False)
    return date.format(fmt)

class subscribeForm(FlaskForm):
    email = StringField(
        "Your email",
        render_kw={"class": "u-full-width"},
    )
    phoneNumber = StringField(
        "Your phone number",
        render_kw={"class": "u-full-width"},
    )
    digest = BooleanField("Digest?",
                            render_kw={"class": "u-full-width"})

    def validate(self, extra_validators=None):
        if super().validate(extra_validators):

            # your logic here e.g.
            if not (self.email.data or self.phoneNumber.data):
                self.email.errors.append('At least one field must have a value')
                return False
            else:
                return True

        return False

@app.route("/", methods=["GET", "POST"])
def home():
    form = subscribeForm()
    if request.method == "POST" and form.validate():
        recipient_id = uuid.uuid4()
        resp = courierClient.profiles.merge(
            recipient_id,
            {
                "phone_number": form.phoneNumber.data,
                "email": form.email.data
            }
        )
        list_id = "thing.each"
        if form.digest.data:
            list_id = "thing.digest"
        courierClient.lists.subscribe(list_id, recipient_id)
        return(list_id)
    posts = CTFLclient.entries({"content_type": "blogPost", "order": "-fields.publishDate"})
    return render_template("index.html", form=form, posts=posts)

@app.route("/post/<string:slug>")
def post(slug):
    post = CTFLclient.entries({"content_type": "blogPost", "fields.slug": slug})
    if post:
        return render_template("post.html", post=post[0])
    return redirect(url_for("home"))


import tkinter as tk
from tkinter import ttk
import matplotlib
matplotlib.use("TkAgg") #backend
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg, NavigationToolbar2Tk
from matplotlib.figure import Figure

LARGE_FONT = ("Verdana", 12)

class SpectralGui(tk.Tk):

    def __init__(self,*args,**kwargs):

        tk.Tk.__init__(self, *args, **kwargs)
        tk.Tk.wm_title(self,"Spectral GUI")
        tk.Tk.iconbitmap(self,'c:/users/juanr/downloads/logo_lec.ico')

        container = tk.Frame(self)
        container.pack(side="top", fill="both",expand=True)
        container.grid_rowconfigure(0,weight=1)
        container.grid_columnconfigure(0, weight=1)

        self.frames = {} #DICCIONARIO

        for F in (StartPage,PageOne,PageTwo,PageThree):

            frame = F(container,self)

            self.frames[F] = frame

            frame.grid(row=0, column=0, sticky="nsew")

        self.show_frame(StartPage)

    def show_frame(self,cont):
        frame = self.frames[cont]
        frame.tkraise()

def gf(strintoprint):
    print(strintoprint)

class StartPage(tk.Frame):

    def __init__(self, parent, controller):
        tk.Frame.__init__(self, parent)
        label = tk.Label(self, text = "Start Page", font = LARGE_FONT)
        label.pack(pady=10, padx=10)

        button = ttk.Button(self, text = "Visit Page 1",
                            command = lambda: controller.show_frame(PageOne))
        button.pack()

        button1 = ttk.Button(self, text="Visit Page 2",
                            command=lambda: controller.show_frame(PageTwo))
        button1.pack()

        button2 = ttk.Button(self, text="Visit Graph Page",
                             command=lambda: controller.show_frame(PageThree))
        button2.pack()

class PageOne(tk.Frame):

    def __init__(self,parent,controller):
        tk.Frame.__init__(self, parent)
        label = tk.Label(self, text="Page One", font=LARGE_FONT)
        label.pack(pady=10, padx=10)

        button = ttk.Button(self, text="Back to Home",
                            command=lambda: controller.show_frame(StartPage))
        button.pack()

        button2 = ttk.Button(self, text="Page Two",
                            command=lambda: controller.show_frame(PageTwo))
        button2.pack()

class PageTwo(tk.Frame):

    def __init__(self,parent,controller):
        tk.Frame.__init__(self, parent)
        label = tk.Label(self, text="Page Two", font=LARGE_FONT)
        label.pack(pady=10, padx=10)

        button = ttk.Button(self, text="Back to Home",
                            command=lambda: controller.show_frame(StartPage))
        button.pack()

        button2 = ttk.Button(self, text="Page One",
                           command=lambda: controller.show_frame(PageOne))
        button2.pack()

class PageThree(tk.Frame):

    def __init__(self,parent,controller):
        tk.Frame.__init__(self, parent)
        label = tk.Label(self, text="Graph Page", font=LARGE_FONT)
        label.pack(pady=10, padx=10)

        button = ttk.Button(self, text="Back to Home",
                            command=lambda: controller.show_frame(StartPage))
        button.pack()

        f = Figure(figsize=(5,5),dpi=100)
        a = f.add_subplot(111)
        a.plot([1,2,3,4,5,6,7,8],[1,2,3,4,5,6,7,8])

        #graph
        canvas = FigureCanvasTkAgg(f,self)
        canvas.draw()
        canvas.get_tk_widget().pack(side=tk.TOP, fill=tk.BOTH, expand = True)

        #navigation toolbar
        toolbar = NavigationToolbar2Tk(canvas,self)
        toolbar.update()
        #canvas.get_tk_widget().pack(side=tk.TOP, fill=tk.BOTH, expand=True) #para que este arriba la toolbar
        canvas._tkcanvas.pack(side=tk.BOTTOM, fill=tk.BOTH, expand = True)

app = SpectralGui()
app.mainloop()


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

"""import os

from reportlab.lib import colors
from reportlab.lib.pagesizes import letter, inch
from reportlab.platypus import SimpleDocTemplate, Table, TableStyle

doc_save_path = os.getcwd()

if not os.path.exists(doc_save_path+"/prueba"):
    os.mkdir(doc_save_path+"/prueba")

doc = SimpleDocTemplate(doc_save_path+"/prueba/simple_table_grid3.pdf", pagesize=letter)
# container for the 'Flowable' objects
elements = []

data= [['00', '01', '02', '03', '04'],
       ['10', '11', '12', '13', '14'],
       ['20', '21', '22', '23', '24'],
       ['30', '31', '32', '33', '34']]

t=Table(data,5*[0.4*inch], 4*[0.4*inch])

t.setStyle(TableStyle([('ALIGN',(1,1),(-2,-2),'RIGHT'),
                       ('TEXTCOLOR',(1,1),(-2,-2),colors.red),
                       ('VALIGN',(0,0),(0,-1),'TOP'),
                       ('TEXTCOLOR',(0,0),(0,-1),colors.blue),
                       ('ALIGN',(0,-1),(-1,-1),'CENTER'),

                       ('VALIGN',(0,-1),(-1,-1),'MIDDLE'),
                       ('TEXTCOLOR',(0,-1),(-1,-1),colors.green),
                       ('INNERGRID', (0,0), (-1,-1), 0.25, colors.black),
                       ('BOX', (0,0), (-1,-1), 0.25, colors.black),
                       ]))

elements.append(t)
# write the document to disk
doc.build(elements)"""

from reportlab.lib import colors
from reportlab.lib.pagesizes import letter, inch, landscape
from reportlab.pdfgen import textobject
from reportlab.platypus import Image, Paragraph, SimpleDocTemplate, Table
from reportlab.lib.styles import getSampleStyleSheet
from reportlab.pdfgen import canvas
import os

"""P0 = Paragraph('''
               A paragraph
               <super>1</super>''',
               styleSheet["BodyText"])

P = Paragraph('''
    <para align=center spaceb=3>The ReportLab Left
    Logo
    Image</para>''',
    styleSheet["BodyText"])"""

"""data= [['A', 'B', 'C', P0, 'D'],
       ['00', '01', '02', [P], '04'],
       ['10', '11', '12', [P], '14'],
       ['20', '21', '22', '23', '24'],
       ['30', '31', '32', '33', '34']]

t=Table(data,style=[('GRID',(1,1),(-2,-2),1,colors.green),
                    ('BOX',(0,0),(1,-1),2,colors.red),
                    ('LINEABOVE',(1,2),(-2,2),1,colors.blue),
                    ('LINEBEFORE',(2,1),(2,-2),1,colors.pink),
                    ('BACKGROUND', (0, 0), (0, 1), colors.pink),
                    ('BACKGROUND', (1, 1), (1, 2), colors.lavender),
                    ('BACKGROUND', (2, 2), (2, 3), colors.orange),
                    ('BOX',(0,0),(-1,-1),2,colors.black),
                    ('GRID',(0,0),(-1,-1),0.5,colors.black),
                    ('VALIGN',(3,0),(3,0),'BOTTOM'),
                    ('BACKGROUND',(3,0),(3,0),colors.limegreen),
                    ('BACKGROUND',(3,1),(3,1),colors.khaki),
                    ('ALIGN',(3,1),(3,1),'CENTER'),
                    ('BACKGROUND',(3,2),(3,2),colors.beige),
                    ('ALIGN',(3,2),(3,2),'LEFT'),
])
t._argW[3]=1.5*inch

elements.append(t)
# write the document to disk
doc.build(elements)
"""
class saveReport():

    ecuacion = None
    intervalos = None
    tolerancia = None
    raiz = None
    data = None
    metodo = None
    iteraciones = None
    path_biseccion = os.getcwd()+"/Soluciones/Biseccion"
    path_punto_fijo = os.getcwd()+"/Soluciones/Punto_Fijo"

    def __init__(self, eq, inter, tol, raiz_, data_, met, itera):
        self.ecuacion = eq
        self.intervalos = inter
        self.tolerancia = tol
        self.raiz = raiz_
        self.data = data_
        self.metodo = met
        self.iteraciones = itera

        self.crear_carpeta()
        self.crear_carpeta()

    def crear_pdf_biseccion(self):

        header = None

        titulo = "Int_"+str(self.intervalos)+"_Tol_"+self.tolerancia

        if self.metodo == "Biseccion":
            header = "f(x) = "+ self.ecuacion
            titulo = "/Soluciones/Biseccion/"+titulo
        else:
            header = "g(x) = "+ self.ecuacion
            titulo = "/Soluciones/Punto_Fijo/"+titulo

        doc = SimpleDocTemplate(os.getcwd()+titulo+".pdf", pagesize=landscape(letter))
        elements = []

        styleSheet = getSampleStyleSheet()

        header_ = Paragraph("Método de "+self.metodo+"", styleSheet['Heading1'])
        texto1 = Paragraph("La Ecuación "+header+"", styleSheet['Heading2'])
        texto2 = Paragraph("Tiene una Raiz x = "+self.raiz+"", styleSheet['Heading2'])
        texto3 = Paragraph("En el intervalo "+str(self.intervalos)+"", styleSheet['Heading2'])
        texto4 = Paragraph("Con tolerancia "+str(self.tolerancia)+"", styleSheet['Heading2'])
        texto5 = Paragraph("Tardó "+str(self.iteraciones)+" iteraciones en encontrar dicha raíz.", styleSheet['Heading2'])

        elements.append(header_)
        elements.append(texto1)
        elements.append(texto2)
        elements.append(texto3)
        elements.append(texto4)
        elements.append(texto5)
        elements.append(Paragraph("",styleSheet['Heading1']))

        self.data.insert(0, (Paragraph("n", styleSheet["Heading2"]), Paragraph("a",styleSheet["Heading2"]),
                            Paragraph("b",styleSheet["Heading2"]), Paragraph("p",styleSheet["Heading2"]),
                            Paragraph("f(p)",styleSheet["Heading2"])))

        for i in range(1,len(self.data),1):
            self.data[i].insert(0,i)

        t=Table(self.data,style=[
                    ('GRID',(0,0),(-1,-1),1,colors.black),
                    ('BACKGROUND', (0, 0), (0, -1), colors.blue),
                    ('TEXTCOLOR',(0,0),(-1,0),colors.white),
                    ('TEXTCOLOR',(0,0),(0,-1),colors.white),
                    ('SIZE', (0,0), (-1,-1), 15),
                    ('VALIGN',(0,0),(-1,-1),'TOP'),
                    ('BACKGROUND', (0, 0), (-1, 0), colors.red),
                    ('SIZE', (1,1),(-1,-1), 12)
        ])

        t._argW[0]=0.6*inch
        t._argW[1]=1.9*inch
        t._argW[2]=1.9*inch
        t._argW[3]=1.9*inch
        t._argW[4]=2.5*inch

        for i in range(len(self.data)):
            t._argH[i]=0.4*inch

        elements.append(t)
        # write the document to disk
        doc.build(elements)

    def crear_carpeta(self):

        if not os.path.exists(os.getcwd()+"/Soluciones"):
            os.mkdir(os.getcwd()+"/Soluciones")
            os.mkdir(self.path_biseccion)
            os.mkdir(self.path_punto_fijo)

    def crear_pdf_punto(self):

        header = None

        titulo = "Xi_"+str(self.intervalos)+"_Tol_"+self.tolerancia

        if self.metodo == "Biseccion":
            header = "f(x) = "+ self.ecuacion
            titulo = "/Soluciones/Biseccion/"+titulo
        else:
            header = "g(x) = "+ self.ecuacion
            titulo = "/Soluciones/Punto_Fijo/"+titulo

        doc = SimpleDocTemplate(os.getcwd()+titulo+".pdf", pagesize=landscape(letter))
        elements = []

        styleSheet = getSampleStyleSheet()

        header_ = Paragraph("Método de "+self.metodo+"", styleSheet['Heading1'])
        texto1 = Paragraph("La Ecuación "+header+"", styleSheet['Heading2'])
        texto2 = Paragraph("Tiene una Raiz x = "+self.raiz+"", styleSheet['Heading2'])
        texto3 = Paragraph("Valor Inicial   X1 ="+str(self.intervalos)+"", styleSheet['Heading2'])
        texto4 = Paragraph("Con tolerancia "+str(self.tolerancia)+"", styleSheet['Heading2'])
        texto5 = Paragraph("Tardó "+str(self.iteraciones)+" iteraciones en encontrar dicha raíz.", styleSheet['Heading2'])

        elements.append(header_)
        elements.append(texto1)
        elements.append(texto2)
        elements.append(texto3)
        elements.append(texto4)
        elements.append(texto5)
        elements.append(Paragraph("",styleSheet['Heading1']))

        self.data.insert(0, (Paragraph("n", styleSheet["Heading2"]), Paragraph("Xn",styleSheet["Heading2"]),
                            Paragraph("X",styleSheet["Heading2"])))

        for i in range(1,len(self.data),1):
            self.data[i].insert(0,i)

        t=Table(self.data,style=[
                    ('GRID',(0,0),(-1,-1),1,colors.black),
                    ('BACKGROUND', (0, 0), (0, -1), colors.blue),
                    ('TEXTCOLOR',(0,0),(-1,0),colors.white),
                    ('TEXTCOLOR',(0,0),(0,-1),colors.white),
                    ('SIZE', (0,0), (-1,-1), 15),
                    ('VALIGN',(0,0),(-1,-1),'TOP'),
                    ('BACKGROUND', (0, 0), (-1, 0), colors.red),
                    ('SIZE', (1,1),(-1,-1), 12)
        ])

        t._argW[0]=0.6*inch
        t._argW[1]=2.5*inch
        t._argW[2]=2.5*inch

        for i in range(len(self.data)):
            t._argH[i]=0.4*inch

        elements.append(t)
        # write the document to disk
        doc.build(elements)


#!/usr/bin/env python
"""

Neural SPARQL Machines - Generator utils.

'SPARQL as a Foreign Language' by Tommaso Soru and Edgard Marx et al., SEMANTiCS 2017
https://arxiv.org/abs/1708.07624

Version 2.0.0

"""
import collections
import http.client
import json
import logging
import re
import sys
import urllib.request, urllib.parse, urllib.error
import urllib.request, urllib.error, urllib.parse
from functools import reduce

ENDPOINT = "http://dbpedia.org/sparql"
GRAPH = "http://dbpedia.org"

def log_statistics( used_resources, special_classes, not_instanced_templates ):
    total_number_of_resources = len(used_resources)
    total_number_of_filled_placeholder_positions = sum(used_resources.values())
    examples_per_instance = collections.Counter()
    for resource in used_resources:
        count = used_resources[resource]
        examples_per_instance.update([count])

    logging.info('{:6d} used resources in {} placeholder positions'.format(total_number_of_resources, total_number_of_filled_placeholder_positions))
    for usage in examples_per_instance:
        logging.info('{:6d} resources occur \t{:6d} times \t({:6.2f} %) '.format(examples_per_instance[usage], usage, examples_per_instance[usage]*100/total_number_of_resources))
    for cl in special_classes:
        logging.info('{} contains: {}'.format(cl, ', '.join(special_classes[cl])))
    logging.info('{:6d} not instanciated templates:'.format(sum(not_instanced_templates.values())))
    for template in not_instanced_templates:
        logging.info('{}'.format(template))

def save_cache ( file, cache ):
    ordered = collections.OrderedDict(cache.most_common())
    with open(file, 'w') as outfile:
        json.dump(ordered, outfile)

def query_dbpedia( query ):
    param = dict()
    param["default-graph-uri"] = GRAPH
    param["query"] = query
    param["format"] = "JSON"
    param["CXML_redir_for_subjs"] = "121"
    param["CXML_redir_for_hrefs"] = ""
    param["timeout"] = "600"
    param["debug"] = "on"
    try:
        resp = urllib.request.urlopen(ENDPOINT + "?" + urllib.parse.urlencode(param))
        j = resp.read()
        resp.close()
    except (urllib.error.HTTPError, http.client.BadStatusLine):
        logging.debug("*** Query error. Empty result set. ***")
        j = '{ "results": { "bindings": [] } }'
    sys.stdout.flush()
    return json.loads(j)

def strip_brackets(s):
    s = re.sub(r'\([^)]*\)', '', s)
    if "," in s:
        s = s[:s.index(",")]
    return s.strip().lower()

REPLACEMENTS = [
    ['dbo:', 'http://dbpedia.org/ontology/', 'dbo_'],
    ['dbp:', 'http://dbpedia.org/property/', 'dbp_'],
    ['dbc:', 'http://dbpedia.org/resource/Category:', 'dbc_'],
    ['dbr:', 'res:', 'http://dbpedia.org/resource/', 'dbr_'],
    ['dct:', 'dct_'],
    ['geo:', 'geo_'],
    ['georss:', 'georss_'],
    ['rdf:', 'rdf_'],
    ['rdfs:', 'rdfs_'],
    ['foaf:', 'foaf_'],
    ['owl:', 'owl_'],
    ['yago:', 'yago_'],
    ['skos:', 'skos_'],
    [' ( ', '  par_open  '],
    [' ) ', '  par_close  '],
    ['(', ' attr_open '],
    [') ', ')', ' attr_close '],
    ['{', ' brack_open '],
    ['}', ' brack_close '],
    [' . ', ' sep_dot '],
    ['. ', ' sep_dot '],
    ['?', 'var_'],
    ['*', 'wildcard'],
    [' <= ', ' math_leq '],
    [' >= ', ' math_geq '],
    [' < ', ' math_lt '],
    [' > ', ' math_gt ']
]

STANDARDS = {
        'dbo_almaMater': ['dbp_almaMater'],
        'dbo_award': ['dbp_awards'],
        'dbo_birthPlace': ['dbp_birthPlace', 'dbp_placeOfBirth'],
        'dbo_deathPlace': ['dbp_deathPlace', 'dbp_placeOfDeath'],
        'dbo_child': ['dbp_children'],
        'dbo_college': ['dbp_college'],
        'dbo_hometown': ['dbp_hometown'],
        'dbo_nationality': ['dbo_stateOfOrigin'],
        'dbo_relative': ['dbp_relatives'],
        'dbo_restingPlace': ['dbp_restingPlaces', 'dbp_placeOfBurial', 'dbo_placeOfBurial', 'dbp_restingplace'],
        'dbo_spouse': ['dbp_spouse'],
        'dbo_partner': ['dbp_partner']
}

def encode( sparql ):
    encoded_sparql = do_replacements(sparql)
    shorter_encoded_sparql = shorten_query(encoded_sparql)
    normalized = normalize_predicates(shorter_encoded_sparql)
    return normalized

def decode ( encoded_sparql ):
    short_sparql = reverse_replacements(encoded_sparql)
    sparql = reverse_shorten_query(short_sparql)
    return sparql

def normalize_predicates( sparql ):
    for standard in STANDARDS:
        for alternative in STANDARDS[standard]:
            sparql = sparql.replace(alternative, standard)

    return sparql

def do_replacements( sparql ):
    for r in REPLACEMENTS:
        encoding = r[-1]
        for original in r[:-1]:
            sparql = sparql.replace(original, encoding)
    return sparql

def reverse_replacements( sparql ):
    for r in REPLACEMENTS:
        original = r[0]
        encoding = r[-1]
        sparql = sparql.replace(encoding, original)
        stripped_encoding = str.strip(encoding)
        sparql = sparql.replace(stripped_encoding, original)
    return sparql

def shorten_query( sparql ):
    sparql = re.sub(r'order by desc\s+....?_open\s+([\S]+)\s+....?_close', '_obd_ \\1', sparql, flags=re.IGNORECASE)
    sparql = re.sub(r'order by asc\s+....?_open\s+([\S]+)\s+....?_close', '_oba_ \\1', sparql, flags=re.IGNORECASE)
    sparql = re.sub(r'order by\s+([\S]+)', '_oba_ \\1', sparql, flags=re.IGNORECASE)
    return sparql

def reverse_shorten_query( sparql ):
    sparql = re.sub(r'_oba_ ([\S]+)', 'order by asc (\\1)', sparql, flags=re.IGNORECASE)
    sparql = re.sub(r'_obd_ ([\S]+)', 'order by desc (\\1)', sparql, flags=re.IGNORECASE)
    return sparql

def read_template_file(file):
    annotations = list()
    line_number = 1
    with open(file) as f:
        for line in f:
            values = line[:-1].split(';')
            target_classes = [values[0] or None, values[1] or None, values[2] or None]
            question = values[3]
            query = values[4]
            generator_query = values[5]
            id = values[6] if (len(values) >= 7 and values[6]) else line_number
            line_number += 1
            annotation = Annotation(question, query, generator_query, id, target_classes)
            annotations.append(annotation)
    return annotations

class Annotation:
    def __init__(self, question, query, generator_query, id=None, target_classes=None):
        self.question = question
        self.query = query
        self.generator_query = generator_query
        self.id = id
        self.target_classes = target_classes if target_classes != None else []
        self.variables = extract_variables(generator_query)

def extract_variables(query):
    variables = []
    query_form_pattern = r'^.*?where'
    query_form_match = re.search(query_form_pattern, query, re.IGNORECASE)
    if query_form_match:
        letter_pattern = r'\?(\w)'
        variables = re.findall(letter_pattern, query_form_match.group(0))
    return variables

def extract_encoded_entities( encoded_sparql ):
    sparql = decode(encoded_sparql)
    entities = extract_entities(sparql)
    encoded_entities = list(map(encode, entities))
    return encoded_entities

def extract_entities( sparql ):
    triples = extractTriples(sparql)
    entities = set()
    for triple in triples:
        possible_entities = [triple['subject'], triple['object']]
        sorted_out = [e for e in possible_entities if not e.startswith('?') and ':' in e]
        entities = entities.union([re.sub(r'^optional{', '', e, flags=re.IGNORECASE) for e in sorted_out])
    return entities

def extract_predicates( sparql ):
    triples = extractTriples(sparql)
    predicates = set()
    for triple in triples:
        pred = triple['predicate']
        predicates.add(pred)
    return predicates

def extractTriples (sparqlQuery):
    triples = []
    whereStatementPattern = r'where\s*?{(.*?)}'
    whereStatementMatch = re.search(whereStatementPattern, sparqlQuery, re.IGNORECASE)
    if whereStatementMatch:
        whereStatement = whereStatementMatch.group(1)
        triples = splitIntoTriples(whereStatement)
    return triples

def splitIntoTriples (whereStatement):
    tripleAndSeparators = re.split(r'(\.[\s\?\<$])', whereStatement)
    trimmed = [str.strip() for str in tripleAndSeparators]

    def repair (list, element):
        if element not in ['.', '.?', '.<']:
            previousElement = list[-1]
            del list[-1]
            if previousElement in ['.', '.?', '.<']:
                cutoff = previousElement[1] if previousElement in ['.?', '.<'] else ''
                list.append(cutoff + element)
            else:
                list.append(previousElement + ' ' + element)
        else:
            list.append(element)

        return list

    tripleStatements = reduce(repair, trimmed, [''])
    triplesWithNones = list(map(splitIntoTripleParts, tripleStatements))
    triples = [triple for triple in triplesWithNones if triple != None]
    return triples

def splitIntoTripleParts (triple):
    statementPattern = r'(\S+)\s+(\S+)\s+(\S+)'
    statementPatternMatch = re.search(statementPattern, triple)

    if statementPatternMatch:
        return {
            'subject': statementPatternMatch.group(1),
            'predicate': statementPatternMatch.group(2),
            'object': statementPatternMatch.group(3)
        }
    else:
        return None

def fix_URI(query):
	query = re.sub(r"dbr:([^\s]+)" , r"<http://dbpedia.org/resource/\1>" , query)
	if query[-2:]=="}>":
		query = query[:-2]+">}"
	return query


class Node:
    def __init__(self, key, val, freq=1):
        self.k = key
        self.v = val
        self.f = freq
        self.next = None
        self.prev = None

class LFUCache(object):

    def __init__(self, capacity):
        """
        :type capacity: int

        Needed
            Dictionary to access nodes in O(1)
            Dictionary to store frequency

        freq =
                                 t
            1 : (h) <-> (1) <-> (4) <->   (t)
                                      (5)->
            2 : (h) <-> (3) <-> (t)
            3 : LRU

        nodes =
            1 : (1)
            3 : (3)

        """
        self.freq = {}
        self.nodes = {}
        self.capacity = capacity

    def remove(self, node):
        node.prev.next, node.next.prev = node.next, node.prev
        lru = self.freq[node.f]
        self.freq[node.f][2] -= 1
        self.lru_is_empty(node.f)
        del self.nodes[node.k]

    def add(self, node, new_freq):
        self.lru_exists(new_freq)
        self.nodes[node.k] = node
        lru = self.freq[new_freq]
        curr_tail = lru[1]
        temp = curr_tail.prev
        node.next = curr_tail
        curr_tail.prev.next = node
        curr_tail.prev = node
        node.prev = temp
        lru[2] += 1

    def get(self, key):
        """
        :type key: int
        :rtype: int
        0. If key doesn't exist, return -1
        1. Get a reference to the node in the nodes dict()
        2. Remove this node from the current freq LRU
            -> If curr_freq LRU is len(2), delete the LRU from the freq dict()
        3. Add this node to the curr_freq + 1 LRU
            -> If curr_freq + 1 does not exist in freq
            -> Initialize a LRU head/tail structure
        """
        if key not in self.nodes:
            return -1

        curr_node = self.nodes[key]
        curr_freq = curr_node.f
        new_freq = curr_freq + 1

        self.remove(curr_node)
        self.add(curr_node, new_freq)
        curr_node.f = new_freq

        return curr_node.v

    def put(self, key, value):
        """
        :type key: int
        :type value: int
        :rtype: None

        1. Check if this key already exists
            -> If it does, remove the node from the current LRU
            -> Add the node to the curr_freq + 1

        2. Else, node doesn't exist
            -> Check if we are at capacity
                -> If we are, evict the head.next node from the min LRU
                -> If min LRU is len(2), delete the LRU from freq dict()
            -> Add the current node to the 1 freq dict()
        """

        if key in self.nodes:
            curr_node = self.nodes[key]
            curr_freq = curr_node.f
            new_freq = curr_freq+1

            self.remove(curr_node)
            self.add(curr_node, new_freq)
            curr_node.f = new_freq
            curr_node.v = value

        else:
            if self.is_capacity() and self.evict_lfu_lru():
                return
            new_node = Node(key, value)
            self.add(new_node, 1)

    def evict_lfu_lru(self):
        if not self.freq.keys():
            return True
        min_freq = min(self.freq.keys())
        min_head, min_tail, f = self.freq[min_freq]
        self.remove(min_head.next)
        return False

    def is_capacity(self):
        return len(self.nodes) == self.capacity

    def lru_is_empty(self, curr_freq):
        if self.freq[curr_freq][2] == 2:
            del self.freq[curr_freq]

    def lru_exists(self, new_freq):
        if new_freq not in self.freq:
            self.create_lru(new_freq)

    def create_lru(self, new_freq):
        head = Node('head', 'head', new_freq)
        tail = Node('tail', 'tail', new_freq)
        head.next, tail.prev = tail, head
        self.freq[new_freq] = [head, tail, 2]

# Your LFUCache object will be instantiated and called as such:
# obj = LFUCache(capacity)
# param_1 = obj.get(key)
# obj.put(key,value)