import struct

# Some references:
# https://tomcat.apache.org/connectors-doc/ajp/ajpv13a.html

#global helpers
def pack_string(s):
	if s is None:
		return struct.pack(">h", -1)

	l = len(s)
	return struct.pack(">H%dsb" % l, l, s.encode('utf8'), 0)

def unpack(stream, fmt):
	size = struct.calcsize(fmt)
	buf = stream.read(size)
	return struct.unpack(fmt, buf)

def unpack_string(stream):
	size, = unpack(stream, ">h")
	if size == -1: # null string
		return None
	res, = unpack(stream, "%ds" % size)
	stream.read(1) # \0
	return res

class NotFoundException(Exception):
	pass

class AjpBodyRequest(object):
	# server == web server, container == servlet
	SERVER_TO_CONTAINER, CONTAINER_TO_SERVER = range(2)
	MAX_REQUEST_LENGTH = 8186

	def __init__(self, data_stream, data_len, data_direction=None):
		self.data_stream = data_stream
		self.data_len = data_len
		self.data_direction = data_direction

	def serialize(self):
		data = self.data_stream.read(AjpBodyRequest.MAX_REQUEST_LENGTH)
		if len(data) == 0:
			return struct.pack(">bbH", 0x12, 0x34, 0x00)
		else:
			res = struct.pack(">H", len(data))
			res += data

		if self.data_direction == AjpBodyRequest.SERVER_TO_CONTAINER:
			header = struct.pack(">bbH", 0x12, 0x34, len(res))
		else:
			header = struct.pack(">bbH", 0x41, 0x42, len(res))

		return header + res

	def send_and_receive(self, socket, stream):
		while True:
			data = self.serialize()
			socket.send(data)
			r = AjpResponse.receive(stream)
			while r.prefix_code != AjpResponse.GET_BODY_CHUNK and r.prefix_code != AjpResponse.SEND_HEADERS:
				r = AjpResponse.receive(stream)

			if r.prefix_code == AjpResponse.SEND_HEADERS or len(data) == 4:
				break

class AjpForwardRequest(object):
	"""
	AJP13_FORWARD_REQUEST :=
		prefix_code	  (byte) 0x02 = JK_AJP13_FORWARD_REQUEST
		method		   (byte)
		protocol		 (string)
		req_uri		  (string)
		remote_addr	  (string)
		remote_host	  (string)
		server_name	  (string)
		server_port	  (integer)
		is_ssl		   (boolean)
		num_headers	  (integer)
		request_headers *(req_header_name req_header_value)
		attributes	  *(attribut_name attribute_value)
		request_terminator (byte) OxFF

	"""

	_, OPTIONS, GET, HEAD, POST, PUT, DELETE, TRACE, PROPFIND, PROPPATCH, MKCOL, COPY, MOVE, LOCK, UNLOCK, ACL, REPORT, VERSION_CONTROL, CHECKIN, CHECKOUT, UNCHECKOUT, SEARCH, MKWORKSPACE, UPDATE, LABEL, MERGE, BASELINE_CONTROL, MKACTIVITY = range(28)

	REQUEST_METHODS = {'GET': GET, 'POST': POST, 'HEAD': HEAD, 'OPTIONS': OPTIONS, 'PUT': PUT, 'DELETE': DELETE, 'TRACE': TRACE}

	# server == web server, container == servlet
	SERVER_TO_CONTAINER, CONTAINER_TO_SERVER = range(2)

	COMMON_HEADERS = ["SC_REQ_ACCEPT",
		"SC_REQ_ACCEPT_CHARSET", "SC_REQ_ACCEPT_ENCODING", "SC_REQ_ACCEPT_LANGUAGE", "SC_REQ_AUTHORIZATION",
		"SC_REQ_CONNECTION", "SC_REQ_CONTENT_TYPE", "SC_REQ_CONTENT_LENGTH", "SC_REQ_COOKIE", "SC_REQ_COOKIE2",
		"SC_REQ_HOST", "SC_REQ_PRAGMA", "SC_REQ_REFERER", "SC_REQ_USER_AGENT"
	]

	ATTRIBUTES = ["context", "servlet_path", "remote_user", "auth_type", "query_string", "route", "ssl_cert", "ssl_cipher", "ssl_session", "req_attribute", "ssl_key_size", "secret", "stored_method"]

	def __init__(self, data_direction=None):
		self.prefix_code = 0x02
		self.method = None
		self.protocol = None
		self.req_uri = None
		self.remote_addr = None
		self.remote_host = None
		self.server_name = None
		self.server_port = None
		self.is_ssl = None
		self.num_headers = None
		self.request_headers = None
		self.attributes = None

		self.data_direction = data_direction

	def pack_headers(self):
		"""
			req_header_name :=
				sc_req_header_name | (string)  [see below for how this is parsed]
			sc_req_header_name := 0xA0xx (integer)
			req_header_value := (string)

			accept  0xA001  SC_REQ_ACCEPT
			accept-charset  0xA002  SC_REQ_ACCEPT_CHARSET
			accept-encoding 0xA003  SC_REQ_ACCEPT_ENCODING
			accept-language 0xA004  SC_REQ_ACCEPT_LANGUAGE
			authorization   0xA005  SC_REQ_AUTHORIZATION
			connection  0xA006  SC_REQ_CONNECTION
			content-type	0xA007  SC_REQ_CONTENT_TYPE
			content-length  0xA008  SC_REQ_CONTENT_LENGTH
			cookie  0xA009  SC_REQ_COOKIE
			cookie2 0xA00A  SC_REQ_COOKIE2
			host	0xA00B  SC_REQ_HOST
			pragma  0xA00C  SC_REQ_PRAGMA
			referer 0xA00D  SC_REQ_REFERER
			user-agent  0xA00E  SC_REQ_USER_AGENT

			store headers as dict
		"""
		self.num_headers = len(self.request_headers)

		res = ""
		res = struct.pack(">h", self.num_headers)
		for h_name in self.request_headers:
			if h_name.startswith("SC_REQ"):
				code = AjpForwardRequest.COMMON_HEADERS.index(h_name) + 1
				res += struct.pack("BB", 0xA0, code)
			else:
				res += pack_string(h_name)

			res += pack_string(self.request_headers[h_name])

		return res

	def pack_attributes(self):
		"""
			Information Code Value  Note
			?context	0x01	Not currently implemented
			?servlet_path   0x02	Not currently implemented
			?remote_user	0x03
			?auth_type  0x04
			?query_string   0x05
			?route  0x06
			?ssl_cert   0x07
			?ssl_cipher 0x08
			?ssl_session	0x09
			?req_attribute  0x0A	Name (the name of the attribut follows)
			?ssl_key_size   0x0B
			?secret 0x0C
			?stored_method  0x0D
			are_done	0xFF	request_terminator
		"""

		res = b""

		for attr in self.attributes:
			a_name = attr['name']
			code = AjpForwardRequest.ATTRIBUTES.index(a_name) + 1
			res += struct.pack("b", code)
			if a_name == "req_attribute":
				aa_name, a_value = attr['value']
				res += pack_string(aa_name)
				res += pack_string(a_value)
			else:
				res += pack_string(attr['value'])

		res += struct.pack("B", 0xFF)
		return res

	def serialize(self):
		res = ""

		res = struct.pack("bb", self.prefix_code, self.method)
		res += pack_string(self.protocol)
		res += pack_string(self.req_uri)
		res += pack_string(self.remote_addr)
		res += pack_string(self.remote_host)
		res += pack_string(self.server_name)
		res += struct.pack(">h", self.server_port)
		res += struct.pack("?", self.is_ssl)

		res += self.pack_headers()

		res += self.pack_attributes()

		if self.data_direction == AjpForwardRequest.SERVER_TO_CONTAINER:
			header = struct.pack(">bbh", 0x12, 0x34, len(res))
		else:
			header = struct.pack(">bbh", 0x41, 0x42, len(res))
		return header + res

	def parse(self, raw_packet):
		stream = StringIO(raw_packet)
		self.magic1, self.magic2, data_len = unpack(stream, "bbH")
		self.prefix_code, self.method = unpack(stream, "bb")
		self.protocol = unpack_string(stream)
		self.req_uri = unpack_string(stream)
		self.remote_addr = unpack_string(stream)
		self.remote_host = unpack_string(stream)
		self.server_name = unpack_string(stream)
		self.server_port = unpack(stream, ">h")
		self.is_ssl = unpack(stream, "?")
		self.num_headers, = unpack(stream, ">H")
		self.request_headers = {}
		for i in range(self.num_headers):
			code, = unpack(stream, ">H")
			if code > 0xA000:
				h_name = AjpForwardRequest.COMMON_HEADERS[code - 0xA001]
			else:
				h_name = unpack(stream, "%ds" % code)
				stream.read(1) # \0

			h_value = unpack_string(stream)

			self.request_headers[h_name] = h_value

	def send_and_receive(self, socket, stream, save_cookies=False):
		res = []
		i = socket.sendall(self.serialize())
		if self.method == AjpForwardRequest.POST:
			return res

		r = AjpResponse.receive(stream)
		assert r.prefix_code == AjpResponse.SEND_HEADERS
		res.append(r)

		if save_cookies and 'Set-Cookie' in r.response_headers:
			self.headers['SC_REQ_COOKIE'] = r.response_headers['Set-Cookie']

		# read body chunks and end response packets
		while True:
			r = AjpResponse.receive(stream)
			res.append(r)
			if r.prefix_code == AjpResponse.END_RESPONSE:
				break
			elif r.prefix_code == AjpResponse.SEND_BODY_CHUNK:
				continue
			else:
				raise NotImplementedError
				break

		return res

class AjpResponse(object):
	"""
		AJP13_SEND_BODY_CHUNK :=
	  	  prefix_code   3
	  	  chunk_length  (integer)
	  	  chunk		*(byte)

		AJP13_SEND_HEADERS :=
	  	  prefix_code	   4
	  	  http_status_code  (integer)
	  	  http_status_msg   (string)
	  	  num_headers	   (integer)
	  	  response_headers *(res_header_name header_value)

		res_header_name :=
			sc_res_header_name | (string)   [see below for how this is parsed]

		sc_res_header_name := 0xA0 (byte)

		header_value := (string)

		AJP13_END_RESPONSE :=
	  	  prefix_code	   5
	  	  reuse			 (boolean)

		AJP13_GET_BODY_CHUNK :=
	  	  prefix_code	   6
	  	  requested_length  (integer)
	"""

	# prefix codes
	_,_,_,SEND_BODY_CHUNK, SEND_HEADERS, END_RESPONSE, GET_BODY_CHUNK = range(7)

	# send headers codes
	COMMON_SEND_HEADERS = [
			"Content-Type", "Content-Language", "Content-Length", "Date", "Last-Modified",
			"Location", "Set-Cookie", "Set-Cookie2", "Servlet-Engine", "Status", "WWW-Authenticate"
			]

	def parse(self, stream):
		# read headers
		self.magic, self.data_length, self.prefix_code = unpack(stream, ">HHb")

		if self.prefix_code == AjpResponse.SEND_HEADERS:
			self.parse_send_headers(stream)
		elif self.prefix_code == AjpResponse.SEND_BODY_CHUNK:
			self.parse_send_body_chunk(stream)
		elif self.prefix_code == AjpResponse.END_RESPONSE:
			self.parse_end_response(stream)
		elif self.prefix_code == AjpResponse.GET_BODY_CHUNK:
			self.parse_get_body_chunk(stream)
		else:
			raise NotImplementedError

	def parse_send_headers(self, stream):
		self.http_status_code, = unpack(stream, ">H")
		self.http_status_msg = unpack_string(stream)
		self.num_headers, = unpack(stream, ">H")
		self.response_headers = {}
		for i in range(self.num_headers):
			code, = unpack(stream, ">H")
			if code <= 0xA000: # custom header
				h_name, = unpack(stream, "%ds" % code)
				stream.read(1) # \0
				h_value = unpack_string(stream)
			else:
				h_name = AjpResponse.COMMON_SEND_HEADERS[code-0xA001]
				h_value = unpack_string(stream)
			self.response_headers[h_name] = h_value

	def parse_send_body_chunk(self, stream):
		self.data_length, = unpack(stream, ">H")
		self.data = stream.read(self.data_length+1)

	def parse_end_response(self, stream):
		self.reuse, = unpack(stream, "b")

	def parse_get_body_chunk(self, stream):
		rlen, = unpack(stream, ">H")
		return rlen

	@staticmethod
	def receive(stream):
		r = AjpResponse()
		r.parse(stream)
		return r


import pytest
import tensorflow as tf
from packaging import version
from tensorflow.python.distribute import strategy_combinations
from tensorflow.python.eager import context

from larq import context as lq_context

if version.parse(tf.__version__) >= version.parse("1.15"):
    strategy_combinations.set_virtual_cpus_to_at_least(3)
    distributed_devices = ["/cpu:1", "/cpu:2"]
else:
    distributed_devices = ["/cpu:0"]

@pytest.fixture
def eager_mode():
    """pytest fixture for running test in eager mode"""
    with context.eager_mode():
        yield

@pytest.fixture
def graph_mode():
    """pytest fixture for running test in graph mode"""
    with context.graph_mode():
        with tf.compat.v1.Session().as_default():
            yield
            tf.keras.backend.clear_session()

@pytest.fixture(params=["eager", "graph"])
def eager_and_graph_mode(request):
    """pytest fixture for running test in eager and graph mode"""
    if request.param == "graph":
        with context.graph_mode():
            with tf.compat.v1.Session().as_default():
                yield request.param
                tf.keras.backend.clear_session()
    else:
        with context.eager_mode():
            yield request.param

@pytest.fixture(params=["graph", "tf_eager", "tf_keras_eager"])
def keras_should_run_eagerly(request):
    """Fixture to run in graph and two eager modes.

    The modes are:
    - Graph mode
    - TensorFlow eager and Keras eager
    - TensorFlow eager and Keras not eager

    The `tf.context` sets graph/eager mode for TensorFlow. The yield is True if Keras
    should run eagerly.
    """

    if request.param == "graph":
        if version.parse(tf.__version__) >= version.parse("2"):
            pytest.skip("Skipping graph mode for TensorFlow 2+.")

        with context.graph_mode():
            yield
    else:
        with context.eager_mode():
            yield request.param == "tf_keras_eager"

@pytest.fixture(params=[False, True])
def distribute_scope(request):
    if request.param is True:
        with tf.distribute.MirroredStrategy(distributed_devices).scope():
            yield request.param
    else:
        yield request.param

@pytest.fixture(params=[True, False])
def quantized(request):
    """pytest fixture for running test quantized and non-quantized"""
    with lq_context.quantized_scope(request.param):
        yield request.param

@pytest.fixture(params=["channels_last", "channels_first"])
def data_format(request):
    return request.param


import botocore.exceptions

from ideasdk.protocols import SocaContextProtocol
from ideadatamodel.cluster_resources import (
    SocaVPC,
    SocaCloudFormationStack,
    SocaOpenSearchDomain,
    SocaDirectory,
    SocaSubnet,
    SocaFileSystem,
    SocaSecurityGroup,
    SocaIAMRole,
    SocaS3Bucket,
    SocaSSHKeyPair,
    SocaACMCertificate,
    SocaAWSProfile,
    SocaEC2PrefixList
)

from ideasdk.utils import Utils
from ideadatamodel import errorcodes, exceptions, constants
from ideasdk.shell import ShellInvoker
from ideasdk.aws import AWSUtil, AwsClientProvider

from cacheout import Cache
from typing import Optional, List, Dict
import time

class AwsResourcesInMemoryDB:

    def __init__(self, context: SocaContextProtocol):
        self._context = context
        self._db = Cache(maxsize=10000, ttl=0)

    def clear(self):
        for key in self._db.keys():
            if key == 'aws.cli.profiles':
                continue
            self._db.delete(key)

    def get_vpcs(self) -> Optional[List[SocaVPC]]:
        return self._db.get('aws.ec2.vpcs')

    def set_vpcs(self, vpcs: List[SocaVPC]):
        return self._db.set('aws.ec2.vpcs', vpcs)

    def get_subnets(self) -> Optional[List[SocaSubnet]]:
        return self._db.get('aws.ec2.subnets')

    def set_subnets(self, subnets: List[SocaSubnet]):
        return self._db.set('aws.ec2.subnets', subnets)

    def get_opensearch_domains(self) -> Optional[List[SocaOpenSearchDomain]]:
        return self._db.get('aws.opensearch.domains')

    def set_opensearch_domains(self, domains: List[SocaOpenSearchDomain]):
        return self._db.set('aws.opensearch.domains', domains)

    def get_directories(self) -> Optional[List[SocaDirectory]]:
        return self._db.get('aws.directories')

    def set_directories(self, directories: List[SocaDirectory]):
        return self._db.set('aws.directories', directories)

    def get_file_systems(self) -> Optional[List[SocaFileSystem]]:
        return self._db.get('aws.file-system')

    def set_file_systems(self, file_systems: List[SocaFileSystem]):
        return self._db.set('aws.file-system', file_systems)

    def get_security_groups(self) -> Optional[List[SocaSecurityGroup]]:
        return self._db.get('aws.ec2.security-groups')

    def set_security_groups(self, security_groups: List[SocaSecurityGroup]):
        return self._db.set('aws.ec2.security-groups', security_groups)

    def get_iam_roles(self) -> Optional[List[SocaIAMRole]]:
        return self._db.get('aws.iam.roles')

    def set_iam_roles(self, iam_roles: List[SocaIAMRole]):
        return self._db.set('aws.iam.roles', iam_roles)

    def get_s3_buckets(self) -> Optional[List[SocaS3Bucket]]:
        return self._db.get('aws.s3.buckets')

    def set_s3_buckets(self, s3_buckets: List[SocaS3Bucket]):
        return self._db.set('aws.s3.buckets', s3_buckets)

    def get_ssh_key_pairs(self) -> Optional[List[SocaSSHKeyPair]]:
        return self._db.get('aws.ec2.ssh_key_pairs')

    def set_ssh_key_pairs(self, ssh_key_pairs: List[SocaSSHKeyPair]):
        return self._db.set('aws.ec2.ssh_key_pairs', ssh_key_pairs)

    def get_cloudformation_stacks(self) -> Optional[List[SocaCloudFormationStack]]:
        return self._db.get('aws.cloudformation.stacks')

    def set_cloudformation_stacks(self, cloudformation_stacks: List[SocaCloudFormationStack]):
        return self._db.set('aws.cloudformation.stacks', cloudformation_stacks)

    def get_acm_certificates(self) -> Optional[List[SocaACMCertificate]]:
        return self._db.get('aws.acm.certificates')

    def set_acm_certificates(self, acm_certificates: List[SocaACMCertificate]):
        return self._db.set('aws.acm.certificates', acm_certificates)

    def get_aws_profiles(self) -> Optional[List[SocaAWSProfile]]:
        return self._db.get('aws.cli.profiles')

    def set_aws_profiles(self, aws_profiles: List[SocaAWSProfile]):
        return self._db.set('aws.cli.profiles', aws_profiles)

    def get_ec2_prefix_lists(self) -> Optional[List[SocaEC2PrefixList]]:
        return self._db.get('aws.ec2.prefix_lists')

    def set_ec2_prefix_lists(self, ec2_prefix_lists: List[SocaEC2PrefixList]):
        return self._db.set('aws.ec2.prefix_lists', ec2_prefix_lists)

class AwsResources:
    def __init__(self, context: SocaContextProtocol, aws: AwsClientProvider = None, aws_util: AWSUtil = None):
        self._context = context
        self._shell = ShellInvoker()
        self._db = AwsResourcesInMemoryDB(self._context)
        self._aws = aws
        self._aws_util = aws_util

    def reset(self):
        self._db.clear()

    @property
    def aws(self) -> AwsClientProvider:
        return self._aws

    @property
    def aws_util(self) -> AWSUtil:
        return self._aws_util

    @staticmethod
    def _get_tag_value(key: str, tags: List[Dict]) -> Optional[str]:
        if tags is None:
            return None
        for tag in tags:
            tag_key = Utils.get_value_as_string('Key', tag)
            if tag_key == key:
                return Utils.get_value_as_string('Value', tag)
        return None

    def is_any_soca_cluster_installed(self) -> bool:
        stacks = self.get_cloudformation_stacks()
        for stack in stacks:
            if not stack.is_cluster():
                continue
            if stack.status != 'CREATE_COMPLETE':
                continue
            cluster_name = stack.cluster_name
            if Utils.is_empty(cluster_name):
                continue
            return True
        return False

    def list_available_cluster_names(self) -> List[str]:
        stacks = self.get_cloudformation_stacks()
        result = set()
        for stack in stacks:
            if not stack.is_cluster():
                continue
            cluster_name = stack.cluster_name
            if Utils.is_empty(cluster_name):
                continue
            result.add(cluster_name)
        return list(result)

    def get_vpcs(self, refresh: bool = False) -> List[SocaVPC]:
        try:
            if not refresh:
                vpcs = self._db.get_vpcs()
                if vpcs is not None:
                    return vpcs

            def result_cb(result) -> List[SocaVPC]:
                results = []
                listing = Utils.get_value_as_list('Vpcs', result, [])
                for entry in listing:
                    tags = Utils.get_value_as_list('Tags', entry, [])
                    vpc_id = Utils.get_value_as_string('VpcId', entry)
                    cidr_block = Utils.get_value_as_string('CidrBlock', entry)
                    title = f'{vpc_id} ({cidr_block})'
                    if tags is not None:
                        name = self._get_tag_value('Name', tags)
                        if Utils.is_not_empty(name):
                            title = f'{title} [{name}]'
                    results.append(SocaVPC(
                        type='aws.ec2.vpc',
                        title=title,
                        ref=entry
                    ))
                return results

            vpcs = self.aws_util.invoke_aws_listing(
                fn=self.aws.ec2().describe_vpcs,
                result_cb=result_cb
            )
            self._db.set_vpcs(vpcs)
            return vpcs
        except Exception as e:
            self.aws_util.handle_aws_exception(e)

    def get_opensearch_clusters(self, vpc_id: str, refresh: bool = False) -> List[SocaOpenSearchDomain]:
        try:
            if not refresh:
                domains = self._db.get_opensearch_domains()
                if domains is not None:
                    return domains

            domains = []
            list_domain_names_result = self.aws.es().list_domain_names()
            domain_names = Utils.get_value_as_list('DomainNames', list_domain_names_result, [])
            for entry in domain_names:
                domain_name = Utils.get_value_as_string('DomainName', entry)
                describe_domain_result = self.aws.es().describe_elasticsearch_domain(DomainName=domain_name)
                domain_status = describe_domain_result['DomainStatus']
                vpc_options = Utils.get_value_as_dict('VPCOptions', domain_status)
                if vpc_options is None:
                    continue
                domain_vpc_id = vpc_options['VPCId']
                if domain_vpc_id != vpc_id:
                    continue
                elasticsearch_version = domain_status['ElasticsearchVersion']
                title = f'{domain_name} ({elasticsearch_version})'
                domains.append(SocaOpenSearchDomain(
                    type='aws.opensearch.domain',
                    title=title,
                    ref=describe_domain_result
                ))

            self._db.set_opensearch_domains(domains)
            return domains
        except Exception as e:
            self.aws_util.handle_aws_exception(e)

    def get_directories(self, vpc_id: str, refresh: bool = False) -> List[SocaDirectory]:
        try:
            if not refresh:
                directories = self._db.get_directories()
                if directories is not None:
                    return directories

            def result_cb(result) -> List[SocaDirectory]:
                results = []
                listing = Utils.get_value_as_list('DirectoryDescriptions', result, [])
                for entry in listing:
                    directory_type = Utils.get_value_as_string('Type', entry)
                    if directory_type != 'MicrosoftAD':
                        continue
                    vpc_settings = Utils.get_value_as_dict('VpcSettings', entry, {})
                    ds_vpc_id = Utils.get_value_as_string('VpcId', vpc_settings)
                    if ds_vpc_id != vpc_id:
                        continue
                    directory_id = Utils.get_value_as_string('DirectoryId', entry)
                    ds_subnets = Utils.get_value_as_list('SubnetIds', vpc_settings, [])
                    ds_name = Utils.get_value_as_string('Name', entry)
                    title = f'{ds_name} (DirectoryId: {directory_id}, SubnetIds: {", ".join(ds_subnets)})'
                    results.append(SocaDirectory(
                        type='aws.directory',
                        title=title,
                        ref=entry
                    ))
                return results

            directories = self.aws_util.invoke_aws_listing(
                fn=self.aws.ds().describe_directories,
                result_cb=result_cb
            )
            self._db.set_directories(directories)
            return directories
        except Exception as e:
            self.aws_util.handle_aws_exception(e)

    def get_subnets(self, vpc_id: str, refresh: bool = False) -> List[SocaSubnet]:
        try:
            if not refresh:
                subnets = self._db.get_subnets()
                if subnets is not None:
                    return subnets

            def result_cb(result) -> List[SocaSubnet]:
                results = []
                listing = Utils.get_value_as_list('Subnets', result, [])
                for entry in listing:
                    tags = Utils.get_value_as_list('Tags', entry)
                    title = self._get_tag_value('Name', tags)
                    subnet_id = entry['SubnetId']
                    availability_zone = entry['AvailabilityZone']
                    availability_zone_id = entry['AvailabilityZoneId']
                    # CidrBlock is only present on IPv4-enable subnets (not on IPv6-only subnets)
                    cidr_block = entry['CidrBlock']
                    suffix = f'(SubnetId: {subnet_id}, CIDR Block: {cidr_block}, AZ/AZID: {availability_zone}/{availability_zone_id})'
                    if Utils.is_not_empty(title):
                        title = f'{title} {suffix}'
                    else:
                        title = suffix
                    results.append(SocaSubnet(
                        type='aws.ec2.subnet',
                        title=title,
                        ref=entry
                    ))
                return results

            subnets = self.aws_util.invoke_aws_listing(
                fn=self.aws.ec2().describe_subnets,
                fn_kwargs={
                    'Filters': [
                        {
                            'Name': 'vpc-id',
                            'Values': [vpc_id]
                        },
                        {
                            'Name': 'ipv6-native',
                            'Values': ['false']
                        }
                    ]
                },
                result_cb=result_cb
            )
            subnets.sort(key=lambda s: s.sort_order)
            self._db.set_subnets(subnets)
            return subnets
        except Exception as e:
            self.aws_util.handle_aws_exception(e)

    def get_file_systems(self, vpc_id: str, refresh: bool = False) -> List[SocaFileSystem]:
        try:
            if not refresh:
                file_systems = self._db.get_file_systems()
                if file_systems is not None:
                    return file_systems

            def result_cb(result, fs_type: str) -> List[SocaFileSystem]:
                results = []
                # Amazon File Cache uses slightly different structures
                if fs_type == constants.STORAGE_PROVIDER_FSX_CACHE:
                    listing_key = 'FileCaches'
                    fsid_key = 'FileCacheId'
                else:
                    listing_key = 'FileSystems'
                    fsid_key = 'FileSystemId'

                allowed_lifecycle_values = ['available', 'updating']

                if fs_type == constants.STORAGE_PROVIDER_EFS:
                    lifecycle_key = 'LifeCycleState'
                else:
                    # FSx / File Cache
                    lifecycle_key = 'Lifecycle'

                listing = Utils.get_value_as_list(listing_key, result, [])
                for entry in listing:
                    # We only work on file systems that are in a healthy state
                    # normalized to .lower() for case insensitive comparison
                    if Utils.get_value_as_string(lifecycle_key, entry, default='error').lower() not in allowed_lifecycle_values:
                        continue
                    file_system_id = Utils.get_value_as_string(fsid_key, entry)
                    storage_virtual_machines = None
                    volumes = None

                    if fs_type == 'efs':
                        file_system_provider = constants.STORAGE_PROVIDER_EFS
                    elif fs_type == 'fsx_cache':
                        file_system_provider = constants.STORAGE_PROVIDER_FSX_CACHE
                        # Perform some checks - only supporting Lustre-based file caches for now
                        file_cache_type = Utils.get_value_as_string('FileCacheType', entry, default='unknown')
                        if file_cache_type != 'LUSTRE':
                            continue

                        lustre_configuration = Utils.get_value_as_dict('LustreConfiguration', entry, {})
                        if Utils.is_empty(lustre_configuration):
                            continue

                        file_cache_deployment_type = Utils.get_value_as_string('DeploymentType', lustre_configuration, default='unknown')
                        if file_cache_deployment_type != 'CACHE_1':
                            continue

                    else:
                        file_system_type = Utils.get_value_as_string('FileSystemType', entry)
                        if file_system_type == 'ONTAP':
                            file_system_provider = constants.STORAGE_PROVIDER_FSX_NETAPP_ONTAP
                        elif file_system_type == 'OPENZFS':
                            file_system_provider = constants.STORAGE_PROVIDER_FSX_OPENZFS
                        elif file_system_type == 'LUSTRE':
                            file_system_provider = constants.STORAGE_PROVIDER_FSX_LUSTRE
                        elif file_system_type == 'WINDOWS':
                            file_system_provider = constants.STORAGE_PROVIDER_FSX_WINDOWS_FILE_SERVER
                        else:
                            continue

                    # File Cache doesn't return Tags in the describe_file_caches
                    if fs_type == 'fsx_cache':
                        try:
                            resource_arn = Utils.get_value_as_string('ResourceARN', entry)
                            tags = self.aws.fsx().list_tags_for_resource(ResourceARN=resource_arn)['Tags']
                        except Exception as e:
                            raise e
                    else:
                        tags = Utils.get_value_as_list('Tags', entry)

                    title = self._get_tag_value('Name', tags)
                    if Utils.is_not_empty(title):
                        title = f'{title} ({fsid_key}: {file_system_id}, Provider: {file_system_provider})'
                    else:
                        title = f'{fsid_key}: {file_system_id}, Provider: {file_system_provider}'

                    if fs_type == 'efs':
                        mount_targets_result = self.aws.efs().describe_mount_targets(FileSystemId=file_system_id)
                        # Rapid invocation of describe_mount_targets can cause RateExceeded on the account
                        time.sleep(.100)
                        mount_targets = Utils.get_value_as_list('MountTargets', mount_targets_result, [])

                        vpc_found = False
                        for mount_target in mount_targets:
                            if mount_target['VpcId'] == vpc_id:
                                vpc_found = True
                                break

                        if not vpc_found:
                            continue

                    elif fs_type == 'fsx_cache':
                        if 'VpcId' not in entry:
                            continue
                        if entry['VpcId'] != vpc_id:
                            continue

                    elif fs_type == 'fsx':

                        if 'VpcId' not in entry:
                            continue
                        if entry['VpcId'] != vpc_id:
                            continue

                        if file_system_provider in (
                            constants.STORAGE_PROVIDER_FSX_NETAPP_ONTAP,
                            constants.STORAGE_PROVIDER_FSX_OPENZFS
                        ):
                            describe_volumes_result = self.aws.fsx().describe_volumes(
                                Filters=[
                                    {
                                        'Name': 'file-system-id',
                                        'Values': [file_system_id]
                                    }
                                ]
                            )
                            volumes = Utils.get_value_as_list('Volumes', describe_volumes_result, [])

                        if file_system_provider == constants.STORAGE_PROVIDER_FSX_NETAPP_ONTAP:
                            describe_svm_result = self.aws.fsx().describe_storage_virtual_machines(
                                Filters=[
                                    {
                                        'Name': 'file-system-id',
                                        'Values': [file_system_id]
                                    }
                                ]
                            )
                            storage_virtual_machines = Utils.get_value_as_list('StorageVirtualMachines', describe_svm_result, [])

                    else:
                        continue

                    results.append(SocaFileSystem(
                        type=f'aws.file-system.{fs_type}',
                        title=title,
                        provider=file_system_provider,
                        ref={
                            'file_system': entry,
                            'storage_virtual_machines': storage_virtual_machines,
                            'volumes': volumes
                        }
                    ))
                return results

            file_systems = []

            efs = self.aws_util.invoke_aws_listing(
                fn=self.aws.efs().describe_file_systems,
                result_cb=result_cb,
                marker_based_paging=True,
                fs_type='efs'
            )
            file_systems += efs

            # Amazon File Cache in a discrete try/except block as it is
            # still undergoing deployment to all regions, and we don't want to
            # maintain a static listing of supported regions.
            try:
                caches = self.aws_util.invoke_aws_listing(
                    fn=self.aws.fsx().describe_file_caches,
                    result_cb=result_cb,
                    fs_type='fsx_cache'
                )
                file_systems += caches
            except botocore.exceptions.ClientError as error:
                if error.response['Error']['Code'] == 'BadRequest':
                    pass
                else:
                    raise error

            fsx = self.aws_util.invoke_aws_listing(
                fn=self.aws.fsx().describe_file_systems,
                result_cb=result_cb,
                fs_type='fsx'
            )
            file_systems += fsx

            self._db.set_file_systems(file_systems)

            return file_systems
        except Exception as e:
            self.aws_util.handle_aws_exception(e)

    def get_security_groups(self, refresh: bool = False) -> List[SocaSecurityGroup]:
        try:

            if not refresh:
                security_groups = self._db.get_security_groups()
                if security_groups is not None:
                    return security_groups

            def result_cb(result) -> List[SocaSecurityGroup]:
                results = []
                listing = Utils.get_value_as_list('SecurityGroups', result, [])
                for entry in listing:
                    tags = Utils.get_value_as_list('Tags', entry, [])
                    title = self._get_tag_value('Name', tags)
                    group_id = entry['GroupId']
                    if Utils.is_not_empty(title):
                        title = f'{title} ({group_id})'
                    else:
                        title = group_id

                    results.append(SocaSecurityGroup(
                        type='aws.ec2.security-group',
                        title=title,
                        ref=entry
                    ))
                return results

            security_groups = self.aws_util.invoke_aws_listing(
                fn=self.aws.ec2().describe_security_groups,
                result_cb=result_cb
            )
            self._db.set_security_groups(security_groups)
            return security_groups
        except Exception as e:
            self.aws_util.handle_aws_exception(e)

    def get_iam_roles(self, refresh: bool = False) -> List[SocaIAMRole]:
        try:
            if not refresh:
                iam_roles = self._db.get_iam_roles()
                if iam_roles is not None:
                    return iam_roles

            def result_cb(result) -> List[SocaIAMRole]:
                results = []
                listing = Utils.get_value_as_list('Roles', result, [])
                for entry in listing:
                    role_name = Utils.get_value_as_string('RoleName', entry)
                    results.append(SocaIAMRole(
                        type='aws.iam-role',
                        title=role_name,
                        ref=entry
                    ))
                return results

            iam_roles = self.aws_util.invoke_aws_listing(
                fn=self.aws.ec2().list_roles,
                result_cb=result_cb,
                marker_based_paging=True
            )
            self._db.set_iam_roles(iam_roles)
            return iam_roles
        except Exception as e:
            self.aws_util.handle_aws_exception(e)

    def get_s3_buckets(self, refresh: bool = False) -> List[SocaS3Bucket]:
        try:

            if not refresh:
                s3_buckets = self._db.get_s3_buckets()
                if s3_buckets is not None:
                    return s3_buckets

            list_buckets_result = self.aws.s3().list_buckets()
            buckets = Utils.get_value_as_list('Buckets', list_buckets_result, [])

            s3_buckets = []
            for bucket in buckets:
                name = Utils.get_value_as_string('Name', bucket)
                s3_buckets.append(SocaS3Bucket(
                    type='s3.bucket',
                    title=f'{name}',
                    ref=bucket
                ))

            self._db.set_s3_buckets(s3_buckets)
            return s3_buckets
        except Exception as e:
            self.aws_util.handle_aws_exception(e)

    def get_ec2_key_pairs(self, refresh: bool = False) -> List[SocaSSHKeyPair]:
        try:

            if not refresh:
                ssh_key_pairs = self._db.get_ssh_key_pairs()
                if ssh_key_pairs is not None:
                    return ssh_key_pairs

            key_pairs_result = self.aws.ec2().describe_key_pairs()
            key_pairs = Utils.get_value_as_list('KeyPairs', key_pairs_result, [])

            ssh_key_pairs = []
            for key_pair in key_pairs:
                key_name = Utils.get_value_as_string('KeyName', key_pair)
                ssh_key_pairs.append(SocaSSHKeyPair(
                    type='aws.ec2.key_pair',
                    title=key_name,
                    ref=key_pair
                ))

            self._db.set_ssh_key_pairs(ssh_key_pairs)
            return ssh_key_pairs

        except Exception as e:
            self.aws_util.handle_aws_exception(e)

    def get_cloudformation_stacks(self, refresh: bool = False) -> List[SocaCloudFormationStack]:
        try:

            if not refresh:
                cloudformation_stacks = self._db.get_cloudformation_stacks()
                if cloudformation_stacks is not None:
                    return cloudformation_stacks

            def result_cb(result) -> List[SocaCloudFormationStack]:
                results = []
                listing = Utils.get_value_as_list('Stacks', result, [])
                for entry in listing:
                    stack_name = Utils.get_value_as_string('StackName', entry)
                    results.append(SocaCloudFormationStack(
                        type='aws.cloudformation-stack',
                        title=stack_name,
                        ref=entry
                    ))
                return results

            cloudformation_stacks = self.aws_util.invoke_aws_listing(
                fn=self.aws.cloudformation().describe_stacks,
                result_cb=result_cb,
                marker_based_paging=False
            )
            self._db.set_cloudformation_stacks(cloudformation_stacks)
            return cloudformation_stacks
        except Exception as e:
            self.aws_util.handle_aws_exception(e)

    def get_acm_certificates(self, refresh: bool = False) -> List[SocaACMCertificate]:
        try:

            if not refresh:
                acm_certificates = self._db.get_acm_certificates()
                if acm_certificates is not None:
                    return acm_certificates

            def result_cb(result) -> List[SocaACMCertificate]:
                results = []
                listing = Utils.get_value_as_list('CertificateSummaryList', result, [])
                for entry in listing:
                    domain_name = Utils.get_value_as_string('DomainName', entry)
                    results.append(SocaACMCertificate(
                        type='aws.acm.certificate',
                        title=domain_name,
                        ref=entry
                    ))
                return results

            acm_certificates = self.aws_util.invoke_aws_listing(
                fn=self.aws.acm().list_certificates,
                result_cb=result_cb,
                marker_based_paging=False
            )
            self._db.set_acm_certificates(acm_certificates)
            return acm_certificates
        except Exception as e:
            self.aws_util.handle_aws_exception(e)

    def get_aws_profiles(self, refresh: bool = False) -> List[SocaAWSProfile]:

        if not refresh:
            aws_profiles = self._db.get_aws_profiles()
            if aws_profiles is not None:
                return aws_profiles

        result = self._shell.invoke(cmd='aws configure list-profiles', shell=True, skip_error_logging=True)
        if result.returncode != 0:
            error = result.stderr
            if Utils.is_empty(error):
                error = result.stdout
            raise exceptions.soca_exception(error_code=errorcodes.GENERAL_ERROR,
                                            message=error)

        profiles = str(result.stdout).splitlines(keepends=False)
        if len(profiles) == 0:
            raise exceptions.soca_exception(error_code=errorcodes.AWS_CLI_NOT_INSTALLED,
                                            message='AWS CLI is not installed or AWS config not found.')

        aws_profiles = []
        for profile in profiles:
            aws_profiles.append(SocaAWSProfile(
                title=profile,
                ref=profile
            ))

        self._db.set_aws_profiles(aws_profiles)
        return aws_profiles

    def get_ec2_prefix_lists(self, refresh: bool = False) -> List[SocaEC2PrefixList]:
        try:

            if not refresh:
                prefix_lists = self._db.get_ec2_prefix_lists()
                if prefix_lists is not None:
                    return prefix_lists

            def result_cb(result) -> List[SocaEC2PrefixList]:
                results = []
                listing = Utils.get_value_as_list('PrefixLists', result, [])
                for entry in listing:
                    prefix_list_name = Utils.get_value_as_string('PrefixListName', entry)
                    if Utils.is_not_empty(prefix_list_name) and prefix_list_name.startswith('com.amazonaws'):
                        continue
                    prefix_list_id = Utils.get_value_as_string('PrefixListId', entry)
                    results.append(SocaEC2PrefixList(
                        type='aws.ec2.prefix_list',
                        title=f'{prefix_list_name} [{prefix_list_id}]',
                        ref=entry
                    ))
                return results

            prefix_lists = self.aws_util.invoke_aws_listing(
                fn=self.aws.ec2().describe_managed_prefix_lists,
                result_cb=result_cb,
                marker_based_paging=False
            )
            self._db.set_ec2_prefix_lists(prefix_lists)
            return prefix_lists
        except Exception as e:
            self.aws_util.handle_aws_exception(e)


from django.conf.urls import include, url

from views import HomePageView, FixedRoadsView, NotFixedRoadsView, desa_datasets, brokenroadunfixed_datasets, brokenroadfixed_datasets, DataInput

urlpatterns=[
    url(r'^$', HomePageView.as_view(), name = 'home'),
    url(r'^fixedroads$', FixedRoadsView.as_view(), name = 'fixed'),
    url(r'^roadsinfo$', NotFixedRoadsView.as_view(), name = 'notfixed'),
    url(r'^add/$', DataInput.as_view(), name = 'add'),
    url(r'^desa_data/$', desa_datasets, name = 'desa'),
    url(r'^jalanrusakfixed_data/$', brokenroadfixed_datasets, name = 'brokenroadfixed'),
    url(r'^jalanrusakunfixed_data/$', brokenroadunfixed_datasets, name = 'brokenroadunfixed'),

]


import argparse
import multiprocessing as mp

import yaml

import mlvision.distributed
import mlvision.train

parser = argparse.ArgumentParser()
parser.add_argument("--config", type=str, help="path of config file to load")
parser.add_argument(
    "--devices",
    type=str,
    nargs="+",
    default=["cuda:0"],
    help="which devices to use on the local machine",
)

def main(rank, config_path, world_size, devices):
    import os

    # Keep only one gpu or cpu
    if str(devices[rank]) == "cpu":
        os.environ["CUDA_VISIBLE_DEVICES"] = ""
    else:
        os.environ["CUDA_VISIBLE_DEVICES"] = str(devices[rank].split(":")[-1])

    import logging

    logging.basicConfig()
    logger = logging.getLogger()
    if rank == 0:
        logger.setLevel(logging.INFO)
    else:
        logger.setLevel(logging.ERROR)

    params = None
    with open(config_path, "r") as f:
        params = yaml.load(f, Loader=yaml.FullLoader)

    world_size, rank = mlvision.distributed.initialize(
        rank_and_world_size=(rank, world_size)
    )
    logger.info(f"Running on rank: {rank}/{world_size}")

    mlvision.train.main(params)

if __name__ == "__main__":
    args = parser.parse_args()

    num_gpus = len(args.devices)
    mp.set_start_method("spawn")  # Fork does not allow gpu

    for rank in range(num_gpus):
        mp.Process(
            target=main, args=(rank, args.config, num_gpus, args.devices)
        ).start()


import tkinter as tk
from tkinter import messagebox, Scale, HORIZONTAL, ttk
from tkinter import *
import random
from algorithms import run_algorithm

#Global variables
max_elements = 500
window_height = 500 + 50 # +50 is the bar below the graph
window_width = 1250
sides_offset = 50
array = []

sort_algorithms = ['Bubble Sort', 'HeapSort', 'Insertion Sort', 'Merge Sort', 'QuickSort', 'Radix Sort', 'Selection Sort']

# Display the UI functions
def generate_UI():
    #User Interface Area
    #Row[0]
    Label(UI_frame,text='Sorting Algorithm Viewer', font=('helvetica',20)).grid(row=0, column=0, padx=5, pady=5)

    #Row[1]
    sizeEntry = Scale(UI_frame, from_=2, to=max_elements, length=200, resolution=1, orient=HORIZONTAL, label='Number of elements: ')
    sizeEntry.grid(row=1, column=0, padx=50, pady=5)

    minEntry = Scale(UI_frame, from_=0, to=100, length=200, resolution=1, orient=HORIZONTAL, label="Min Value")
    minEntry.grid(row=1, column=1, padx=50, pady=5)

    maxEntry = Scale(UI_frame, from_=0, to=500, length=200, resolution=1, orient=HORIZONTAL, label="Max Value")
    maxEntry.grid(row=1, column=2, padx=50, pady=5)

    Button(UI_frame, text='Generate Array',
            command=lambda: generate_n_display(sizeEntry.get(),minEntry.get(),maxEntry.get()),
            bg='blue',fg='white', height=2,width=15,
            font=('helvetica',9,'bold')).grid(row=1, column=3, padx=50, pady=5)

    #Row[2]
    Label(UI_frame, text="Algorithm: ", font=('helvetica',10)).grid(row=2, column=0, padx=50, pady=5, sticky=E)
    algMenu = ttk.Combobox(UI_frame, values=sort_algorithms)
    algMenu.grid(row=2, column=1, padx=100, pady=5)
    algMenu.current(0)

    speedScale = Scale(UI_frame, from_=0.0, to=2.0, length=200, digits=2, resolution=0.1, orient=HORIZONTAL, label='Time between comparisons [s]:')
    speedScale.grid(row=2, column=2, padx=5, pady=5)

    Button(UI_frame, text='Start Algorithm',
            command=lambda: run_algorithm(array, speedScale.get(), display_array, algMenu.get()),
            bg='green',fg='white', height=2,width=15,
            font=('helvetica',9,'bold')).grid(row=2, column=3, padx=50, pady=5)

# Starts the program, creating the array and displaying it on the second window
def generate_n_display(array_n_elements, min_value, max_value):
    global array

    array = []

    try:
        if array_n_elements < 2 or min_value > max_value:
            raise Exception()

        array = create_array(array_n_elements,min_value,max_value)

        display_array(array, ['red' for x in range(len(array))])

    except Exception:
        messagebox.showerror("Error","Invalid input.\nMinimum value must be lower or equal to the maximum value.")

# Creates a random array
def create_array(array_n_elements,min_value,max_value):
    global array

    min_index = 0
    max_index = 0
    while min_index == max_index:
        # Defines the indexes in the array for the minimum and maximum values
        # This makes sure the given values will exist in the array
        min_index = random.randint(0,array_n_elements - 1)
        max_index = random.randint(0,array_n_elements - 1)

    # Creates the array with random values between min and max
    # Randomly ordered
    for i in range(0,array_n_elements):
        if i == min_index:
            array.append(min_value)
        elif i == max_index:
            array.append(max_value)
        else:
            array.append(random.randint(min_value,max_value))

    return array

# Displays the array in the window
def display_array(array, colors_array):
    global window_height
    global window_width
    global sides_offset

    canvas.delete("all")

    columns_size = (window_width - sides_offset*2) / len(array)

    for i in range(0, len(array)):
        canvas.create_rectangle(sides_offset + (i*columns_size),
                            -0.9*array[i] + 500,
                            sides_offset + (i*columns_size) + columns_size,
                            window_height-50,
                            fill=colors_array[i])

    root.update_idletasks()

root = tk.Tk()
root.title('Sorting Algorithm Viewer')
root.resizable(width=False,height=False)

#frame / base lauout
UI_frame = tk.Frame(root, width= 1250, height=100)
UI_frame.grid(row=0, column=0, padx=10, pady=5)

#Just a frame separator
tk.Frame(root, width= 1250, height=2, bg='black').grid(row=1, column=0, padx=10, pady=5)

canvas = tk.Canvas(root,width=1250,height=550)
canvas.grid(row=2, column=0, padx=10, pady=5)

generate_UI()

root.mainloop()


from torch.utils.data import DataLoader
import torch
import random
import numpy as np

class CategoriesSampler:

    def __init__(self, index, lenth, way, shot):
        self.lenth = lenth
        self.way = way
        self.shot = shot

        self.index = index

    def __len__(self):
        return self.lenth

    def __iter__(self):
        for i_batch in range(self.lenth):
            batch = []
            classes = list(self.index.keys())
            for c in classes:
                lenth_per = torch.from_numpy(self.index[c])
                way_per = len(lenth_per)
                shot = torch.randperm(way_per)[:self.shot]
                batch.append((c * way_per + shot).int())
            batch = torch.stack(batch).reshape(-1)
            yield batch

def prepare_loader_normal(dataset, args):
    loader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True, num_workers=4,
                              pin_memory=True)
    return loader

def prepare_loader_sample1(dataset, args):
    loader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True, num_workers=4,
                        pin_memory=True)
    train_sampler = CategoriesSampler(dataset.sub_indexes,
                                      len(loader),
                                      args.way + 3,
                                      args.shot)
    train_fsl_loader = DataLoader(dataset=dataset,
                                  batch_sampler=train_sampler,
                                  num_workers=4,
                                  pin_memory=True)
    return loader, train_fsl_loader

def prepare_loader_sample2(dataset, args):
    loader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True, num_workers=4,
                        pin_memory=True, drop_last=True)
    train_sampler = CategoriesSampler(dataset.sub_indexes,
                                      len(loader),
                                      max(args.way, args.base_class),
                                      args.shot)
    train_fsl_loader = DataLoader(dataset=dataset,
                                  batch_sampler=train_sampler,
                                  num_workers=4,
                                  pin_memory=True)
    return loader, train_fsl_loader


import os
import ffmpeg
import numpy as np

from functools import lru_cache
from typing import Union

import torch
import torch.nn.functional as F

from .utils import exact_div

# hard-coded audio hyperparameters
SAMPLE_RATE = 16000
N_FFT = 400
N_MELS = 80
HOP_LENGTH = 160
CHUNK_LENGTH = 30
N_SAMPLES = CHUNK_LENGTH * SAMPLE_RATE  # 480000: number of samples in a chunk
N_FRAMES = exact_div(N_SAMPLES, HOP_LENGTH)  # 3000: number of frames in a mel spectrogram input

def pad_or_trim(array:torch.Tensor, length:int = N_SAMPLES, *, axis:int=-1):
    if array.shape[axis] > length:
        array = array.index_select(dim=axis, index=torch.arange(length, device=array.device))
    elif array.shape[axis] < length:
        pad_widths = [(0,0)] * array.ndim
        pad_widths[axis] = (0, length - array.shape[axis])
        array = F.pad(array, [pad for sizes in pad_widths[::-1] for pad in sizes])

    return array

def load_audio(file:str, sr:int = SAMPLE_RATE):
    try:
        out, _ = (
            ffmpeg.input(file, threads=0)
            .output("-", format="s16le", acodec="pcm_s16le", ac=1, ar=sr)
            .run(cmd=["ffmpeg", "-nostdin"], capture_stdout=True, capture_stderr=True)
        )
    except ffmpeg.Error as e:
        raise RuntimeError(f"Failed to load audio: {e.stderr.decode()}") from e

    return np.frombuffer(out, np.int16).flatten().astype(np.float32) / 32768.0

@lru_cache(maxsize=None)
def mel_filters(device):
    with np.load(os.path.join(os.path.dirname(__file__), "assets", "mel_filters.npz")) as f:
        return torch.from_numpy(f[f"mel_{N_MELS}"]).to(device)

def log_mel_spec(audio:Union[str, torch.Tensor, np.ndarray]):
    if not torch.is_tensor(audio):
        if isinstance(audio, str): audio = load_audio(audio)
        audio = torch.from_numpy(audio)

    stft = torch.stft(
        audio,
        N_FFT,
        HOP_LENGTH,
        window=torch.hann_window(N_FFT).to(audio.device),
        return_complex=True
    )

    log_spec = torch.clamp(mel_filters(audio.device) @ (stft[:, :-1].abs() ** 2), min=1e-10).log10()
    log_spec = torch.maximum(log_spec, log_spec.max() - 8.0)
    return (log_spec + 4.0) / 4.0


#%%

"""
Algorithm to learn the signed distance based on a USDA mesh of the boundary
    1) for each normal and the corresponding point
    2) sample N points between the point and Lxpoint in two directions: the normal and the non-normal
    3) train the NN on the outside points with their loss (positive)
    4) do the same for the inside points with their own loss (negative)
    5) evaluate testerror on all boundary points
    6) Visualise the 0-levelset with Pyvista (results are not great)
"""

import jax
import jax.numpy as jnp
import numpy as np
import flax.linen as nn
from flax.training import train_state, checkpoints
import optax
import os
import matplotlib.pyplot as plt
import pyvista as pv
import skimage

from sphpc import USDAGeom, make_dir, plot

#%%

EXPERIMENET_ID = "LearningSDF"
DATAFOLDER = "./data/" + EXPERIMENET_ID +"/"
make_dir(DATAFOLDER)
KEY = jax.random.PRNGKey(42)     ## Use same random points for all iterations

SURFACE_ID = 0   ## id of the boundary to learn
SAMPLING_FACTOR = 2

SAMPLING_DIST = 15
SAMPLING_SIZE = 100

PRINT_EVERY = 20

INIT_LR = 1e-3
MAX_EPOCHS = 1000
TOLERANCE = 1e-5

#%%

geometry = USDAGeom("meshes/glass.usda")
geometry.visualize(how="matplotlib");

#%%

ground_truth = geometry.boundaries[SURFACE_ID]
# ground_truth = geometry.sources[SURFACE_ID]

points = jnp.array([ground_truth.points[i] for i in ground_truth.faces])
normals = ground_truth.normals

## Data for testing
x_bd = ground_truth.points

## Data for training
def make_data(point, normal, sampling_dist, sampling_size, key):
    """ Make data for training in and opposite the direction of the normals """
    point = point[jnp.newaxis, :]
    normal = normal[jnp.newaxis, :]

    t = jax.random.uniform(key=key, shape=(sampling_size, 1), minval=0, maxval=sampling_dist)

    x_in = point - t*normal
    dist_in = -jnp.linalg.norm(x_in-point, axis=-1)[:, jnp.newaxis]

    x_out = point + t*normal
    dist_out = jnp.linalg.norm(x_out-point, axis=-1)[:, jnp.newaxis]

    return x_in, dist_in, x_out, dist_out

make_data_vec = jax.vmap(make_data, in_axes=(0, 0, None, None, 0), out_axes=(0))

res = make_data_vec(points, normals, SAMPLING_DIST, SAMPLING_SIZE, jax.random.split(KEY, num=points.shape[0]))

data_fn_vec = jax.vmap(make_data, in_axes=(0, 0, None, None, 0), out_axes=(0))

def make_data_vec(points, normals, sampling_dist, sampling_size, key):
    keys = jax.random.split(key, num=points.shape[0])

    data = data_fn_vec(points, normals, sampling_dist, sampling_size, keys)

    return [jnp.reshape(arr, (-1, arr.shape[-1])) for arr in data]

#%%

class MLP(nn.Module):
    input_size: int=3
    nb_layers:int=3
    nb_neurons_per_layer:int=20

    @nn.compact
    def __call__(self, x):                              ## TODO no normalisation
        y = nn.Dense(self.input_size)(x)
        y = nn.Dense(self.nb_neurons_per_layer)(y)
        for _ in range(0, self.nb_layers-1):
            y = nn.tanh(y)
            y = nn.Dense(self.nb_neurons_per_layer)(y)
        return nn.Dense(1)(y)

def init_flax_params(net:nn.Module):
    init_data = jnp.ones((10,3))
    params = net.init(KEY, init_data)
    print(net.tabulate(KEY, init_data, depth=8, console_kwargs={"force_jupyter":False}))
    return params

sdf = MLP()

print("Surrogate sdf archtecture: ")
params = init_flax_params(sdf)

#%%

total_steps = MAX_EPOCHS     ## number of optimisation steps

## Optimizer
scheduler = optax.piecewise_constant_schedule(init_value=INIT_LR,
                                            boundaries_and_scales={int(total_steps*0.4):0.1, int(total_steps*0.8):0.1})

optimizer = optax.adam(learning_rate=scheduler)
# optimizer = optax.sgd(learning_rate=scheduler)

## Flax training state
state = train_state.TrainState.create(apply_fn=sdf.apply,
                                        params=params,
                                        tx=optimizer)

#%%

def loss_fn_train(params, x_in, dist_in, x_out, dist_out):

    ## Error in
    dist_in_pred = sdf.apply(params, x_in)
    agr_in = jnp.abs(dist_in_pred-dist_in)
    # pen_in = optax.l2_loss((dist_in_pred>0).astype(float))

    ## Error out
    dist_out_pred = sdf.apply(params, x_out)
    agr_out = jnp.abs(dist_out_pred-dist_out)
    # pen_out = optax.l2_loss((dist_out_pred<0).astype(float))

    return jnp.mean(agr_in) + jnp.mean(agr_out)

@jax.jit
def loss_fn_test(params, x_bd):
    dist_bd_pred = sdf.apply(params, x_bd)
    agr_bd = jnp.abs(dist_bd_pred)

    return jnp.mean(agr_bd)

#%%

@jax.jit
def train_step(state, x_in, dist_in, x_out, dist_out):
    loss, grads = jax.value_and_grad(loss_fn_train)(state.params, x_in, dist_in, x_out, dist_out)
    state = state.apply_gradients(grads=grads)

    return state, loss

# train_step_vec = jax.jit(jax.vmap(train_step, in_axes=(None, 0, 0, 0, 0), out_axes=(0)))

train_losses = []
test_losses = []

#%%

if len(os.listdir(DATAFOLDER)) > 2:
    print("Found existing network, loading & training")
    state = checkpoints.restore_checkpoint(ckpt_dir=DATAFOLDER, prefix="sdf_checkpoint_", target=state)

else:
    print("Training from scratch")

epoch = 1
L = SAMPLING_DIST
subkey = KEY
test_losses.append(TOLERANCE+1)

while test_losses[-1] > TOLERANCE and epoch < MAX_EPOCHS+1:

    x_in, dist_in, x_out, dist_out = make_data_vec(points, normals, sampling_dist=L, sampling_size=SAMPLING_SIZE, key=KEY)
    state, loss = train_step(state, x_in, dist_in, x_out, dist_out)

    train_losses.append(loss)
    test_losses.append(loss_fn_test(state.params, x_bd))

    if epoch<=3 or epoch%PRINT_EVERY==0:
        print("Epoch: %-5d      TrainLoss: %.6f     TestLoss: %.6f" % (epoch, train_losses[-1], test_losses[-1]))

    L /= SAMPLING_FACTOR
    epoch += 1

checkpoints.save_checkpoint(DATAFOLDER, prefix="sdf_checkpoint_", target=state, step=state.step, overwrite=True)
print("Training done, saved network")

#%%

fig, ax = plt.subplots(1, 1, figsize=(6*1,4))

plot(train_losses[:], label='Train', x_label='epochs', y_scale="log", ax=ax)
plot(test_losses[1:], label='Test', x_label='epochs', y_scale="log", title="MAE", ax=ax);

# %%

## Marching cubes for visualisation

## Test function: a sphere with radius 1
def shpere(x):
    return -1. + np.linalg.norm(x, axis=-1)

# create a uniform grid to sample the function with
n = 100
x_min, y_min, z_min = -1.5, -1.5, -2
grid = pv.UniformGrid(
    dimensions=(n, n, n),
    spacing=(abs(x_min)*2 / n, abs(y_min)*2 / n, abs(z_min)*2 / n),
    origin=(x_min, y_min, z_min),
)

## Scalar values for the isosurface
values = np.asarray(sdf.apply(state.params, grid.points)).flatten()
# values = shpere(grid.points)
mesh = grid.contour(1, values, method='marching_cubes', rng=[-TOLERANCE, TOLERANCE])

## Plot the mesh with distance-based colouring
dist = np.linalg.norm(mesh.points, axis=1)
mesh.plot(scalars=dist, smooth_shading=True, specular=1, cmap="plasma", show_scalar_bar=False)


"""Module for compiling two-level rules out of FSTs for the components.
"""
__author__ = "© Kimmo Koskenniemi, 2018 - 2020"
__version__ = "2020-02-01"

import re

import hfst_dev as hfst

import twol.cfg as cfg

import twol.fs as fs

import twol.twbt as twbt

import twol.twexamp as twexamp

def init():
    """Initializes the module by computing several common FSTs

    Assumes that twexamp.read_fst() has read in cfg.examples_fst and
    initialized sone symbol sets.
    """
    global pistar_fst, pistar_fsa, diamond_sym, diamond_fst
    global trim_pre_fst, trim_post_fst

    assert cfg.examples_fst, "cfg.examples_fst not loaded (by twexamp module)"

    cfg.definitions["PAIRS"] = cfg.all_pairs_fst.copy()
    cfg.definitions["PI"] = cfg.all_pairs_fst.copy()

    diamond_sym = 'DIAMOND'
    diamond_fst = hfst.regex(diamond_sym)
    pi_fst = cfg.all_pairs_fst.copy()
    pistar_fst = cfg.all_pairs_fst.copy()
    pistar_fst.repeat_star()
    pistar_fst.remove_epsilons()
    pistar_fst.minimize()
    pistar_fsa = hfst.fst_to_fsa(pistar_fst, separator='^')
    pi_in_fst = pi_fst.copy()
    pi_in_fst.input_project()
    pi_out_fst = pi_fst.copy()
    pi_out_fst.output_project()
    pi_in_star_fst = pistar_fst.copy()
    pi_in_star_fst.input_project()
    pi_out_star_fst = pistar_fst.copy()
    pi_out_star_fst.output_project()
    if cfg.verbosity >= 20:
        twbt.ppfst(pistar_fst, title="pistar_fst")

    fst1 = fs.star(fs.crossprod(fs.expr("ZERO"),
                                 pi_in_fst))
    fst2 = fs.star(fs.concat(fst1,
                             fs.expr("ZERO:BEGIN")))
    fst3 = fs.concat(fst2, pi_in_star_fst)
    fst4 = fs.star(fs.concat(fs.expr("ZERO:END"),
                             fs.star(fs.crossprod(fs.expr("ZERO"),
                                                  pi_in_fst))))
    trim_pre_fst = fs.concat(fst3, fst4)
    trim_pre_fst.set_name("trim_pre_fst")
    #trim_pre_fst =  XRC.compile(
    #    "[[ZERO .x. [PI].u]* ZERO:BEGIN]* " \
    #    "[[PI].u]* " \
    #    "[ZERO:END [ZERO .x. [PI].u]*]*"
    #)

    fst1 = fs.star(fs.crossprod(pi_out_fst, fs.expr("ZERO")))
    fst2 = fs.star(fs.concat(fst1, fs.expr("BEGIN:ZERO")))
    fst3 = fs.concat(fst2, pi_out_star_fst)
    fst4 = fs.star(fs.concat(fs.expr("END:ZERO"),
                             fs.star(fs.crossprod(pi_out_fst,
                                                  fs.expr("ZERO")))))
    trim_post_fst = fs.concat(fst3, fst4)
    trim_post_fst.set_name("trim_post_fst")
    #trim_post_fst = XRC.compile(
    #    "[[[PI].l .x. ZERO]* BEGIN:ZERO]* " \
    #    "[[PI].l]* " \
    #    "[END:ZERO [[PI].l .x. ZERO]*]*"
    #)
    if cfg.verbosity >= 20:
        twbt.ppfst(trim_pre_fst)
        twbt.ppfst(trim_post_fst)
    return

def quote(str):
    """Protect '{' and '}' with a % in xerox regular expressions.

    >>> quote("a {ij}:j [a:c | b]")
    "a %{ij%}:j [a:c | b]"
    """
    return(re.sub(r"([{'}])", r"%\1", str))

def e(str):
    """Convert a two-level component expression into a FST.

    str -- a string containing a (two-level) regular expression
    Returns an FST which performs the mapping represented by str
    corresponding to the expression.
    """
    global XRC
    # print("Regex string:", str) ##
    if str == "":
        return(XRC.compile("[]"))
    F = XRC.compile(str)
    F.minimize()
    F.set_name(str)
    if cfg.verbosity >= 5:
        twbt.ppfst(F) ##
    return(F)

def generalized_restriction(precondition_fst,
                            postcondition_fst):
    """Combines the precondition FST and the postcondition FST into a rule FST.

    Conditions are formed out of the center/contexts of a rule.
    Each condition contains exactly two diamond symbols.

    Returns the generalized restriction as an FST.
    """
    global pistar_fst, diamond_sym
    temp_fst = precondition_fst.copy()
    temp_fst.subtract(postcondition_fst)
    temp_fst.minimize()
    temp_fst.set_name("PRECOND-POSTCOND")
    # twbt.ppfst(temp_fst, True) ##
    temp_fst.substitute(diamond_sym, hfst.EPSILON)
    # print(temp_fst.get_properties().items()) ##
    temp_fst.minimize()
    temp_fst.set_name("Diamonds removed")
    # twbt.ppfst(temp_fst, True) ##
    result_fst = pistar_fst.copy()
    result_fst.minus(temp_fst)
    result_fst.minimize()
    # twbt.ppfst(result_fst, True) ##
    result_fst.set_name("Doubly negated")
    # twbt.ppfst(result_fst, True) ##
    return(result_fst)

def x_to_condition(x_fst):
    """Computes and returns a condition FST out of the center FST of a rule.

    x_fst -- the center or X-part of a rule (located in front of the
    operator)

    returns [PI* ¤ X ¤ PI*] where ¤ is the diamond (not in PI) as an FST
    """
    global pistar_fst, diamond_fst
    result_fst = pistar_fst.copy()
    result_fst.concatenate(diamond_fst)
    result_fst.concatenate(x_fst)
    result_fst.concatenate(diamond_fst)
    result_fst.concatenate(pistar_fst)
    result_fst.minimize()
    result_fst.set_name("PISTAR ¤ X ¤ PISTAR")
    # twbt.ppfst(result_fst, True) ##
    return(result_fst)

def begin_end(expr_fst):
    """Removes everything before BEGIN and after END from expr_fst

    expr_fst -- an FST representing a set of pair strings str

    Returns an FST which represents pair strings which are maximal
    substrings of the above which do not contain any BEGIN or END
    symbols.
    """
    global trim_pre_fst, trim_post_fst
    result_fst = trim_pre_fst.copy()
    result_fst.compose(expr_fst)
    # twbt.ppfst(result_fst, True) ##
    result_fst.compose(trim_post_fst)
    # twbt.ppfst(result_fst, True) ##
    result_fst.substitute("ZERO", hfst.EPSILON)
    result_fst.minimize()
    # twbt.ppfst(result_fst, True) ##
    return(result_fst)

def context_to_condition(left_context_fst,
                         right_context_fst):
    """Convert one context into a condition (for the generalized restriction)

    left_context_fst -- the left context as an FST

    right_context_fst -- the right context as an FST

    Returns [PI* LC ¤ PI* ¤ RC PI*] as an FST
    """
    global pistar_fst, diamond_fst
    leftc_fst = pistar_fst.copy()
    leftc_fst.concatenate(left_context_fst)
    leftc_fst.minimize()
    result_fst = begin_end(leftc_fst)
    result_fst.concatenate(diamond_fst)
    result_fst.concatenate(pistar_fst)
    result_fst.concatenate(diamond_fst)
    rightc_fst = right_context_fst.copy()
    rightc_fst.concatenate(pistar_fst)
    rightc_fst.minimize()
    rightc_fst = begin_end(rightc_fst)
    result_fst.concatenate(rightc_fst)
    result_fst.minimize()
    return(result_fst)

def contexts_to_condition(*contexts):
    """A list of contexsts is converted into a condition.

    Each context in the list is converted separately and
    the result is the union of these and is returned as an FST.
    """
    global pistar_fst
    result_fst = hfst.HfstTransducer()
    for leftc, rightc in contexts:
        #twbt.ppfst(leftc, title="leftc") ##
        #twbt.ppfst(rightc, title="rightc") ##
        context_fst = context_to_condition(leftc, rightc)
        result_fst.disjunct(context_fst)
        result_fst.minimize()
        leftc_name = leftc.get_name()
        rightc_name = rightc.get_name()
        result_fst.set_name(leftc_name + "_" + rightc_name)
        #twbt.ppfst(result_fst, title="result") ##
    return(result_fst)

def mix_output(x_fst):
    """Computes an FSA that is used when creating negative examples

    First, it computes an expression Y which represent all possible
    (correct and incorrect) realizations of the input side of X.  Then,
    Y is transformed into an encoded FSA which can be a component of the
    transformation of correct examples into incorrect ones.

    x_fst -- the center FST (X part) of a rule Returns [X.u .o. PI*]
    encoded as an FSA (i.e. maps pairs to themselves)
    """
    global pistar_fst
    result_fst = x_fst.copy()
    result_fst.input_project()
    result_fst.compose(pistar_fst)
    result_fst.minimize()
    result_encod_fsa = hfst.fst_to_fsa(result_fst, separator="^")
    # twbt.ppfst(result_fsa, True) ##
    return result_encod_fsa

def mix_input(x_fst):
    """Computes an FSA that is used when creating negative examples

    First, it computes an expression Y which represent all possible
    (correct and incorrect) inputs for the output side of X.  Then,
    Y is transformed into an encoded FSA which can be a component of the
    transformation of correct examples into incorrect ones.

    x_fst -- the center FST (X part) of a rule Returns [PI* .o. X.l]
    encoded as an FSA (i.e. maps pairs to themselves)
    """
    global pistar_fst
    result_fst = pistar_fst.copy()
    temp_fst = x_fst.copy()
    temp_fst.output_project()
    result_fst.compose(temp_fst)
    result_fst.minimize()
    result_encod_fsa = hfst.fst_to_fsa(result_fst, separator="^")
    # twbt.ppfst(result_fsa, True) ##
    return result_encod_fsa

def selector_from_x(x_fst):
    """Compute and return [PI* X PI*]"""
    selector_fst = pistar_fst.copy() # starting to build it
    selector_fst.concatenate(x_fst)
    selector_fst.concatenate(pistar_fst) # now complete
    selector_fst.set_name("Selector " + x_fst.get_name())
    return selector_fst

def correct_to_incorrect(x_fst, side):
    """used for creating negative examples for <= rules

    In order to make negative examples for <= rules we need to transform
    the examples so that some correct input:output pairs are
    changed so that the output part becomes different.  The computed
    encoded FST maps correct inputs to any possible outputs (correct or
    incorrect).

    x_fst -- the FST for the X part of the rule

    side -- either "input" or "output"

    returns: an fst (encoded as a fsa) which maps correct examples into
    incorrect exs
    """
    global pistar_fst, pistar_fsa
    if side == "input":
        mixed_fsa = mix_input(x_fst)
    else:
        mixed_fsa = mix_output(x_fst)
    temp_encod_fsa = hfst.fst_to_fsa(x_fst, separator="^")
    temp_encod_fsa.cross_product(mixed_fsa) # now maps corr X to all variations
    # twbt.ppfst(temp_encod_fsa, True) ##
    corr_to_incorr_encod_fst = pistar_fsa.copy()
    corr_to_incorr_encod_fst.concatenate(temp_encod_fsa)
    corr_to_incorr_encod_fst.concatenate(pistar_fsa)
    corr_to_incorr_encod_fst.minimize() # now complete
    corr_to_incorr_encod_fst.set_name("Correct to incorrect")
    return corr_to_incorr_encod_fst

def incorrect_to_correct(x_fst):
    """Compute a transformation for right-arrow (=>) rules

    In order to make negative examples for the => rules we need to
    modify the examples so that some correct occurrences of X are
    modified so that the output part of X becomes something else,
    i.e. incorrect because it is in an unexpected context.

    x_fst -- FST for the center part (X) of a rule

    Returns: scrambler_fst -- an encoded FST which maps encoded
    instances of X into all possible correct and incorrect pairs (where
    the input symbol is the same but the output symbol perhaps
    different).
    """
    global pistar_fst, pistar_fsa
    x_encod_fsa = hfst.fst_to_fsa(x_fst, separator="^")
    mix_fst = mix_output(x_fst) # still an encoded fsa
    mix_fst.cross_product(x_encod_fsa) # now fst
    scrambler_fst = pistar_fsa.copy()
    scrambler_fst.concatenate(mix_fst)
    scrambler_fst.concatenate(pistar_fsa)
    scrambler_fst.minimize() # now complete
    scrambler_fst.set_name("Scrambler " + x_fst.get_name())
    return scrambler_fst

def rightarrow(name, x_fst, *contexts):
    """Compiles rules like X => [LC1,RC1),...(LCk,RCk)]

    name -- name to be given to the rule FST

    x_fst -- the center (X) of the rule

    \*contexts -- list of contexts, i.e. pairs of left and right context

    Returns a triple:
    rule_fst -- the compiled rule

    selector_fst -- FST which selects examples which are relevant for
    this rule

    scrambler_fst -- an encoded FST which produces negative examples
    """
    precondition_fst = x_to_condition(x_fst)
    postcondition_fst = contexts_to_condition(*contexts)
    rule_fst = generalized_restriction(precondition_fst, postcondition_fst)
    rule_fst.set_name(name)
    # twbt.ppfst(rule_fst, True) ##
    selector_fst = selector_from_x(x_fst)
    scrambler_fst = incorrect_to_correct(x_fst)
    # twbt.ppfst(scrambler_fst, True) ##
    return rule_fst, selector_fst, scrambler_fst

def output_coercion(name, x_fst, *contexts):
    """Compiles rules like X <= LC1 _ RC1, ..., LCk _ RCk

    name -- name to be given to the rule FST

    x_fst -- the center (X) of the rule

    \*contexts -- list of contexts, i.e. pairs of left and right context

    Returns a triple:
    rule_fst -- the compiled rule

    selector_fst -- FST which selects examples which are relevant for
    this rule

    scrambler_fst -- an encoded FST which produces negative examples
    """
    global pistar_fst
    postcondition_fst = x_to_condition(x_fst)
    x_all_fst = x_fst.copy()
    x_all_fst.input_project()
    x_all_fst.compose(pistar_fst)
    precondition_fst = x_to_condition(x_all_fst)
    context_condition_fst = contexts_to_condition(*contexts)
    precondition_fst.intersect(context_condition_fst)
    rule_fst = generalized_restriction(precondition_fst, postcondition_fst)
    rule_fst.set_name(name)
    if cfg.verbosity >= 20:
        twbt.ppfst(rule_fst, True)
    x_any_fst = x_fst.copy()
    x_any_fst.input_project()
    x_any_fst.compose(pistar_fst)
    ###x_any_fst.minus(x_fst)
    selector_fst = selector_from_x(x_any_fst)
    scrambler_fst = correct_to_incorrect(x_fst, "output")
    return rule_fst, selector_fst, scrambler_fst

def input_coercion(name, x_fst, *contexts):
    """Compiles rules like X <-- LC1 _ RC1, ..., LCk _ RCk

    name -- name to be given to the rule FST

    x_fst -- the center (X) of the rule

    \*contexts -- list of contexts, i.e. pairs of left and right
    context

    Returns a triple:
    rule_fst -- the compiled rule

    selector_fst -- FST which selects examples which are relevant for
    this rule

    scrambler_fst -- an encoded FST which produces negative examples
    """
    global pistar_fst
    postcondition_fst = x_to_condition(x_fst)
    x_all_fst = pistar_fst.copy()
    temp_fst = x_fst.copy()
    temp_fst.output_project()
    x_all_fst.compose(temp_fst) # PI* .o. X.l
    precondition_fst = x_to_condition(x_all_fst)
    context_condition_fst = contexts_to_condition(*contexts)
    precondition_fst.intersect(context_condition_fst)
    rule_fst = generalized_restriction(precondition_fst, postcondition_fst)
    rule_fst.set_name(name)
    if cfg.verbosity >= 20:
        twbt.ppfst(rule_fst, True)
    selector_fst = selector_from_x(x_all_fst)
    scrambler_fst = correct_to_incorrect(x_fst, "input")
    return rule_fst, selector_fst, scrambler_fst

def doublearrow(name, x_fst, *contexts):
    """Compiles rules like X <=> [LC1,RC1),...(LCk,RCk)]

    name -- name to be given to the rule FST

    x_fst -- the center (X) of the rule

    \*contexts -- list of contexts, i.e. pairs of left and right
    context

    Returns a triple:

    rule_fst -- the compiled rule

    selector_fst -- FST which selects examples which are relevant for
    this rule

    scrambler_fst -- an encoded FST which produces negative examples

    """
    rule_fst, selector_fst, scrambler_fst = rightarrow(name, x_fst, *contexts)
    rule2_fst, selector2_fst, scrambler2_fst = output_coercion(name, x_fst, *contexts)
    rule_fst.intersect(rule2_fst)
    rule_fst.minimize()
    rule_fst.set_name(name)
    scrambler_fst.disjunct(scrambler2_fst)
    scrambler_fst.minimize()
    selector_fst.disjunct(selector2_fst)
    selector_fst.minimize()
    # twbt.ppfst(rule_fst, True) ##
    return rule_fst, selector_fst, scrambler_fst

def center_exclusion(name, x_fst, *contexts):
    """Compiles rules like X /<= [LC1,RC1),...(LCk,RCk)]

    name -- name to be given to the rule FST

    x_fst -- the center (X) of the rule

    \*contents -- list of contexts, i.e. pairs of left and right context

    Returns a triple:

    rule_fst -- the compiled rule

    selector_fst -- FST which selects examples which are relevant for this rule

    scrambler_fst -- empty_fst (negative examples not relevant for these rules)
    """
    context_condition_fst = contexts_to_condition(*contexts)
    x_condition_fst = x_to_condition(x_fst)
    context_condition_fst.intersect(x_condition_fst)
    null_fst = hfst.empty_fst()
    rule_fst = generalized_restriction(context_condition_fst, null_fst)
    rule_fst.set_name(name)
    # twbt.ppfst(rule_fst, True) ##
    selector_fst = selector_from_x(x_fst)
    scrambler_fst = hfst.empty_fst()
    return rule_fst, selector_fst, scrambler_fst

if __name__ == "__main__":
    """This module forms the rule FSTs out of the collection of component
FSTs which have been compiled out of the center parts, left and right
contexts of the rules.  The formulas used in the module are those from
[ylijyrä2006]_.  """

    twex.read_examples()
    init(1)
    #define("V", "PI .o.[a|e|i|o|ä|ö]")
    #define("C", "[PI .o. [h|l|n|s|t|v]] | %{ij%}:j")
    R1 = doublearrow("{ao}:o <=> _ {ij}:",
                     e("%{ao%}:o"),
                     (e("[]"), e("[%{ij%} .o. PI]")))
    twbt.ppfst(R1, True)
    rule2_fst = doublearrow("{ij}:j <=> V :Ø* _ :Ø* V",
                     "%{ij%}:j",
                     ("V [PI .o. Ø]*", "[PI .o. Ø]* V"))
    twbt.ppfst(rule2_fst, True)
    R3 = doublearrow("{tl}:l <=> _ CLOSED",
                     "%{tl%}:l",
                     ("[]", "V %{ij%}:i* C [C | [PI .o. Ø]* END]"))
    twbt.ppfst(R3, True)


from configure import ProductionConfig

__conf = ProductionConfig()

log_config = {
    'loggers': {
        'app': {
            'level': 'INFO',
            'propagate': False,
            'handlers': [
                'handler_error',
                'handler_debug',
                'handler_info'
            ]
        }
    },
    'version': 1,
    'disable_existing_loggers': True,
    'handlers': {
        'handler_error': {
            'formatter': 'simple',
            'backupCount': 7,
            'level': 'ERROR',
            'interval': 1,
            'when': 'midnight',
            'class': 'logging.handlers.TimedRotatingFileHandler',
            'filename': __conf.BASE_DIR + '/logs/error.log'},
        'handler_debug': {
            'formatter': 'simple',
            'backupCount': 7,
            'level': 'DEBUG',
            'interval': 1,
            'when': 'midnight',
            'class': 'logging.handlers.TimedRotatingFileHandler',
            'filename': __conf.BASE_DIR + '/logs/debug.log'
        },
        'handler_info': {
            'formatter': 'simple',
            'backupCount': 7,
            'level': 'INFO',
            'interval': 1,
            'when': 'midnight',
            'class': 'logging.handlers.TimedRotatingFileHandler',
            'filename': __conf.BASE_DIR + '/logs/info.log'
        }
    },
    'root': {
        'level': 'NOTSET',
        'handlers': ['handler_info']
    },
    'formatters': {
        'simple': {
            'format': '%(asctime)s [%(funcName)s:%(lineno)s] [%(levelname)s] %(message)s'
        },
        'full': {
            'format': '%(asctime)s [%(levelname)s] [%(funcName)s:%(lineno)s] %(process)d %(thread)d %(message)s'
            }
        }
    }


from PyQt5 import QtCore, QtGui, QtWidgets
from PyQt5.QtWidgets import QGraphicsView, QGraphicsScene, QApplication, QGraphicsEllipseItem, \
    QGraphicsSceneHoverEvent, QGraphicsSceneMouseEvent, QMessageBox
from PyQt5.QtCore import Qt, QPointF

class MovableImage(QtWidgets.QGraphicsPixmapItem):
    """Provide a movable image"""

    def __init__(self, image_address):
        super().__init__()  # Initial position must be (0,0) to avoid bias in coordinate system...
        # self.setPixmap(QtGui.QPixmap(image_address))
        self.setPixmap(QtGui.QPixmap())
        self.setAcceptHoverEvents(True)  # hover events are used to change mouse cursor

    def setNewPixmap(self, pixmap):
        self.setPixmap(pixmap)

    def hoverEnterEvent(self, event: 'QGraphicsSceneHoverEvent'):
        """When cursor enters the object, set cursor to open hand"""
        QApplication.instance().setOverrideCursor(Qt.OpenHandCursor)

    def hoverLeaveEvent(self, event: 'QGraphicsSceneHoverEvent'):
        """When cursor leaves the object, restore mouse cursor"""
        QApplication.instance().restoreOverrideCursor()

    def mouseMoveEvent(self, event: 'QGraphicsSceneMouseEvent'):
        """mouseMoveEvent is called whenever a mouse button is pressed and the cursor is moved. """
        new_cursor_position = event.scenePos()  # mouse cursor in scene coordinates
        old_cursor_position = event.lastScenePos()
        old_top_left_corner = self.scenePos()
        new_top_left_corner_x = new_cursor_position.x() - old_cursor_position.x() + old_top_left_corner.x()
        new_top_left_corner_y = new_cursor_position.y() - old_cursor_position.y() + old_top_left_corner.y()
        self.setPos(QPointF(new_top_left_corner_x, new_top_left_corner_y))  # update disk top left corner

    def mousePressEvent(self, event: 'QGraphicsSceneMouseEvent'): pass

    def mouseDoubleClickEvent(self, event: 'QGraphicsSceneMouseEvent'): pass

    def mouseReleaseEvent(self, event: 'QGraphicsSceneMouseEvent'): pass

class PhotoViewer(QtWidgets.QGraphicsView):

    def __init__(self, parent):
        super(PhotoViewer, self).__init__(parent)
        self._zoom = 0
        self._empty = True
        self._scene = QtWidgets.QGraphicsScene(self)
        self._photo = QtWidgets.QGraphicsPixmapItem()
        self._photo.setTransformationMode(Qt.SmoothTransformation)
        self._photo.setTransformOriginPoint()
        self.small_part = MovableImage('images/small_part.png')
        self.small_part.setTransformationMode(Qt.SmoothTransformation)
        self.big_part = MovableImage('images/big_part.png')

        self._scene.addItem(self._photo)
        self._scene.addItem(self.big_part)
        self._scene.addItem(self.small_part)

        self.setScene(self._scene)
        self.setTransformationAnchor(QtWidgets.QGraphicsView.AnchorUnderMouse)
        self.setResizeAnchor(QtWidgets.QGraphicsView.AnchorUnderMouse)
        self.setVerticalScrollBarPolicy(QtCore.Qt.ScrollBarAlwaysOff)
        self.setHorizontalScrollBarPolicy(QtCore.Qt.ScrollBarAlwaysOff)
        self.setBackgroundBrush(QtGui.QBrush(QtGui.QColor(30, 30, 30)))
        self.setFrameShape(QtWidgets.QFrame.NoFrame)

    def hasPhoto(self):
        return not self._empty

    def fitInView(self, scale=True):
        rect = QtCore.QRectF(self._photo.pixmap().rect())
        if not rect.isNull():
            self.setSceneRect(rect)
            if self.hasPhoto():
                unity = self.transform().mapRect(QtCore.QRectF(0, 0, 1, 1))
                self.scale(1 / unity.width(), 1 / unity.height())
                viewrect = self.viewport().rect()
                scenerect = self.transform().mapRect(rect)
                factor = min(viewrect.width() / scenerect.width(),
                             viewrect.height() / scenerect.height())
                self.scale(factor, factor)
            self._zoom = 0

    def setPhoto(self, pixmap=None):
        # pixmap2 = QtGui.QPixmap('t.png')
        self._zoom = 0
        if pixmap and not pixmap.isNull():
            self._empty = False
            self.setDragMode(QtWidgets.QGraphicsView.ScrollHandDrag)
            self._photo.setPixmap(pixmap)
            # self._photo2.setPixmap(pixmap2)
        else:
            self._empty = True
            self.setDragMode(QtWidgets.QGraphicsView.NoDrag)
            self._photo.setPixmap(QtGui.QPixmap())
            # self._photo2.setPixmap(QtGui.QPixmap())
        self.fitInView()

    def wheelEvent(self, event):
        if self.hasPhoto():
            if event.angleDelta().y() > 0:
                factor = 1.25
                self._zoom += 1
            else:
                factor = 0.8
                self._zoom -= 1
            if self._zoom > 0:
                self.scale(factor, factor)
            elif self._zoom == 0:
                self.fitInView()
            else:
                self._zoom = 0

class Window(QtWidgets.QWidget):
    def __init__(self):
        super(Window, self).__init__()
        self.viewer = PhotoViewer(self)
        # 'Load image' button
        self.btnLoad = QtWidgets.QToolButton(self)
        self.btnLoad.setText('Load image')
        self.btnLoad.clicked.connect(self.loadImage)
        # Button to change from drag/pan to getting pixel info
        self.btnOperate = QtWidgets.QToolButton(self)
        self.btnOperate.setText('Operate')
        self.btnOperate.clicked.connect(self.operate)

        # Arrange layout
        VBlayout = QtWidgets.QVBoxLayout(self)
        VBlayout.addWidget(self.viewer)
        HBlayout = QtWidgets.QHBoxLayout()
        HBlayout.setAlignment(QtCore.Qt.AlignLeft)
        HBlayout.addWidget(self.btnLoad)
        HBlayout.addWidget(self.btnOperate)
        VBlayout.addLayout(HBlayout)

    def loadImage(self):
        self.viewer.setPhoto(QtGui.QPixmap('images/hip.jpg'))

    def operate(self):
        if self.viewer.hasPhoto():
            self.viewer.small_part.setNewPixmap(QtGui.QPixmap('images/small_part.png'))
            self.viewer.big_part.setNewPixmap(QtGui.QPixmap('images/big_part.png'))
        else:
            msg = QMessageBox()
            msg.setIcon(QMessageBox.Information)

            msg.setText("Error")
            msg.setInformativeText("You should load an image before operation")
            msg.setWindowTitle("No Image")
            msg.setStandardButtons(QMessageBox.Ok)

            msg.exec_()

if __name__ == '__main__':
    import sys

    app = QtWidgets.QApplication(sys.argv)
    window = Window()
    window.setGeometry(500, 300, 800, 600)
    window.show()
    sys.exit(app.exec_())


"""
Used to randomly divide a dataset (in csv format) into:
    - training set (80%)
    - validation set (10%)
    - test set (10%)

Directions for use:
    1: set LOAD_FILENAME, FILE_PATH to the proper file path
    2: set TRAIN_FILENAME, VALIDATE_FILENAME, and TEST_FILENAME to desired name
    2: run "python3 train-validate-test.py"

"""
from random import randint

LOAD_FILENAME = "CleanedFeaturesData.csv" # load data from this file
FILE_PATH = "TrainValidateTest/" # location for the files you create

TRAIN_FILENAME = "xTrainFeatures.csv"
VALIDATE_FILENAME = "xValidateFeatures.csv"
TEST_FILENAME = "xTestFeatures.csv"

def main():
    with open(LOAD_FILENAME) as file:
        line = file.readline()
        headers = line.split(',')
        headers[-1] = headers[-1].rstrip('\n')
        train = [headers]
        validate = [headers]
        test = [headers]

        for line in file:
            split_line = line.split(',')
            split_line[-1] = split_line[-1].rstrip('\n')
            rand_num = randint(1,10)
            if rand_num == 1:
                validate.append(split_line)
            elif rand_num == 2:
                test.append(split_line)
            else:
                train.append(split_line)

        create_file(FILE_PATH + TRAIN_FILENAME, train)
        create_file(FILE_PATH + VALIDATE_FILENAME, validate)
        create_file(FILE_PATH + TEST_FILENAME, test)

# borrowed from split-cat-num.py
def create_filestring(data):
    # creates a string to write to a file based on the passed list
    string = ''
    for person in data:
        for attribute in person:
            string += str(attribute)
            string += ","
        string = string[:-1]
        string += "\n"
    return string

# borrowed from split-cat-num.py
def create_file(filename, data):
    # writes a csv file in `filename` based containing `data`
    string = create_filestring(data)
    with open(filename, 'w') as file:
        file.write(string)

if __name__ == "__main__":
    main()


#!/bin/env python3

import sys

def search(graph, start, end, marks, p):
    if start == end:
        print(p + ['end'])
        return 1

    result = 0
    adj_smalls = set(marks.keys())
    for adj in (graph[start] - adj_smalls):
        if adj.islower():
            result += search(graph, adj, end, (marks | {adj: 1} ), p+[adj])
        else:
            result += search(graph, adj, end, marks, p+[adj])

    if max(marks.values()) <= 1:
        for adj in (graph[start] & adj_smalls):
            if adj not in ['start','end']:
                result += search(graph, adj, end, (marks | {adj:2} ), p+[adj])

    return result

def main():
    lines = sys.stdin.readlines()
    graph = {}

    for line in lines:
        tokens = line.split('-')
        a = tokens[0].strip()
        b = tokens[1].strip()
        if a not in graph:
            graph[a] = set()
        if b not in graph:
            graph[b] = set()
        graph[a].add(b)
        graph[b].add(a)

    n_paths = search(graph, 'start','end',{'start':1},p=['start'])
    print('#paths: ', n_paths)

if __name__ == '__main__':
    main()


import csv
import os
import string
from nltk import word_tokenize, ngrams
import numpy as np

punctuations = string.punctuation

def _read_tsv(input_file, quoting=csv.QUOTE_MINIMAL):
    """Reads a tab separated value file."""
    with open(input_file, "r", encoding="utf-8-sig") as f:
        reader = csv.DictReader(f, delimiter="\t", quoting=quoting)
        lines = []
        for line in reader:
            lines.append(line)
        return lines

# if token is seen, head1, if token is unseen head0
def get_gate_type0(article, summary):
    gate = []
    article_tokens = set(article.lower().split(' '))
    summary_tokens = summary.lower().split(' ')
    for token in summary_tokens:
        if token in punctuations:
            gate.append('-1')
        elif token in article_tokens:
            gate.append('1')
        else:
            gate.append('0')

    assert len(gate) == len(summary_tokens), 'gate length does not match summary length'
    return gate

# if token is unseen and context is unseen, head1, if token is seen and context is seen head0
def get_gate_type2(article, summary):
    gate = []
    article_tokens = set(article.split(' '))
    summary_tokens = summary.split(' ')
    for idx, token in enumerate(summary_tokens):
        if idx < 3:
            gate.append('-1')
        elif token in punctuations:
            gate.append('-1')
        else:
            prefix_wtoken = ' '.join(summary_tokens[idx-2: idx + 1])
            prefix = ' '.join(summary_tokens[idx-2: idx])
            if prefix_wtoken in article:
                gate.append('1')
            elif prefix not in article and token not in article:
                gate.append('0')
            else:
                gate.append('-1')

    assert len(gate) == len(summary_tokens), 'gate length does not match summary length'
    return gate

def get_overlap(inp, out, ngram=2):
    grams_inp = set(ngrams(word_tokenize(inp.lower()), ngram))
    grams_out = set(ngrams(word_tokenize(out.lower()), ngram))

    total = len(grams_out)
    common = len(grams_inp.intersection(grams_out))
    if total == 0:
        return 0
    else:
        return float(common) / float(total)

def get_gate_type3(article, summary, mean_overlap):
    overlap = get_overlap(article, summary)
    if overlap < mean_overlap:
        return 0
    else:
        return 1

if __name__=='__main__':
    input_folder = '../../data/xsum/original_data'
    output_folder = '../../data/xsum/'

    split = 'train.tsv'
    input_file = os.path.join(input_folder, split)
    data = _read_tsv(input_file)

    outfile = open(os.path.join(output_folder, split), 'w')
    fieldnames = list(data[0].keys()) + ['gate', 'gate_sent']
    writer = csv.DictWriter(outfile, fieldnames=fieldnames, quoting=csv.QUOTE_MINIMAL, delimiter='\t')
    writer.writeheader()

    if split == 'train.tsv':
        mean_overlap = []
        for ex in data:
            article = ex['article']
            summary = ex['summary']
            overlap = get_overlap(article.lower(), summary.lower())
            mean_overlap.append(overlap)
        mean_overlap = np.mean(overlap)
    else:
        print('Please provide mean overlap based on the train file.')

    print(mean_overlap)

    num_0s = 0.
    num_1s = 0.
    num_blank = 0.
    num_1sent = 0.

    for ex in data:
        article = ex['article']
        summary = ex['summary']

        gate = get_gate_type0(article.lower(), summary.lower())
        gate_sent = get_gate_type3(article.lower(), summary.lower(), mean_overlap)

        ex['gate'] = ' '.join(gate)
        ex['gate_sent'] = str(gate_sent)
        num_0s += gate.count('0')
        num_1s += gate.count('1')
        num_blank += gate.count('-1')
        num_1sent += gate_sent

        writer.writerow(ex)

print(num_1s/(num_0s + num_1s + num_blank))
print(num_0s/(num_0s + num_1s + num_blank))
print(num_blank/(num_0s + num_1s + num_blank))

print(num_1sent/len(data))


# %%
import mesh2hrtf as m2h
import pyfar as pf
# import matplotlib.pyplot as plt
import os

parameters = (
    [None, "minimum", "equal"],                 # default parameter
    [1, "minimum", "equal"],                    # test smoothing
    [None, "linear", "equal"],                  # test linear phase filters
    [None, "zero", "equal"],                    # test zero phase filters
    [None, "minimum", 'voronoi'],               # test voronoi weights
    [None, "minimum", [1, 0, 0, 0, 0, 0]])      # test custom weights

for smooth_fractions, phase, weights in parameters:

    if isinstance(weights, list):
        name = f"compute_dtfs_{smooth_fractions}_{phase}_custom"
    else:
        name = f"compute_dtfs_{smooth_fractions}_{phase}_{weights}"

    print(name)

    dtf, dftf_inverse = m2h.compute_dtfs(
        os.path.join("..", "resources", "SOFA_files", "HRIR_6_points.sofa"),
        smooth_fractions, phase, weights)

    # plots used for manual verification of the results

    # pf.plot.time_freq(dftf_inverse)
    # plt.savefig(name + "_filter.png")
    # plt.close("all")

    # pf.plot.time_freq(pf.io.convert_sofa(dtf)[0])
    # plt.savefig(name + "_dtfs.png")
    # plt.close("all")

    pf.io.write(
        name, compress=True, dtf=dtf.Data_IR, dftf_inverse=dftf_inverse.time)


import model
import preprocess
import os.path
from torch import load, from_numpy
import sys

try:
    file_path = sys.argv[1]
except IndexError:
    raise BaseException('No file to predict')

net = model.NeuralNetwork(400, 350, 1)
# Load pretrained model
if os.path.isfile(model.net_path):
    net.load_state_dict(load(model.net_path))
else:
    raise BaseException('Neural Network not fitted!')
net = net.to(model.device)

# Preprocess to net
data_to_predict = from_numpy(preprocess.load_and_convert(file_path))
# Print prediction
print(int(net.predict(data_to_predict).item()))


from sqlalchemy import MetaData, Table, Column, Text
from sqlalchemy import Boolean, Integer, String, DateTime, Float

meta = MetaData()

alarm = Table(
    'alarm', meta,
    Column('id', String(255), primary_key=True, index=True),
    Column('enabled', Boolean),
    Column('name', Text()),
    Column('description', Text()),
    Column('timestamp', DateTime(timezone=False)),
    Column('counter_name', String(255), index=True),
    Column('user_id', String(255), index=True),
    Column('project_id', String(255), index=True),
    Column('comparison_operator', String(2)),
    Column('threshold', Float),
    Column('statistic', String(255)),
    Column('evaluation_periods', Integer),
    Column('period', Integer),
    Column('state', String(255)),
    Column('state_timestamp', DateTime(timezone=False)),
    Column('ok_actions', Text()),
    Column('alarm_actions', Text()),
    Column('insufficient_data_actions', Text()),
    Column('matching_metadata', Text()))

def upgrade(migrate_engine):
    meta.bind = migrate_engine
    alarm.create()

def downgrade(migrate_engine):
    meta.bind = migrate_engine
    alarm.drop()


#_*_coding:utf-8_*_

iphone7 = ['alex', 'rain', 'jack', 'old_driver']
iphone8 = ['alex', 'shanshan', 'jack', 'old_boy']

both_list = []

for name in iphone8:
    if name in iphone7:
        both_list.append(name)

print(both_list)

s = {1,3,4,5,6}
s.discard()