# Copyright (c) 2015-2019 Shawn Chin

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
"""
Example usage:

        # Crete a mesh (collection of connected components)
        m = Mesh()

        # Add components (hexagons) to the mesh, each with a set of 'needs' ports (output ports) and
        # and 'provides' ports (input ports)
        m.add_component('A', needs_ports=['n1', 'n2'])
        m.add_component('B', provides_ports=['p1', 'p2'])
        m.add_component('C', needs_ports=['nX'])
        m.add_component('D', needs_ports=['data'], provides_ports=['pX'])

        # Add pass-through adapters to connect ports between components
        m.add_connection('A', 'n1', 'B', 'p1')
        m.add_connection('A', 'n2', 'D', 'pX')
        m.add_connection('C', 'nX', 'B', 'p2')

        # Define an external adapter and attach it to a 'needs' port
        m.add_resource('Resource X')
        m.add_connection_to_resource('D', 'data', 'Resource X')

        # Render as a Graphviz dot file
        print render_mesh_as_dot(m)

When visualised using Graphviz, this should produce an output that looks like:

            _____________              _____________
           |      A      |            |      B      |
           |-------------|            |-------------|
           |      |  n1  |----------->|  p1  |      |
           |      |______|            |______|      |
           |      |  n2  |-----+  +-->|  p2  |      |
           |______|______|     |  |   |______|______|
                            +- | -+
                            |  |
            _____________   |  |      _____________
           |      C      |  |  |     |      D      |
           |-------------|  |  |     |-------------|
           |      |  nX  |--+  +---->|  pX  | data |------->[  Resource X  ]
           |______|______|           |______|______|

"""
__author__ = "Shawn Chin"
__version__ = "v1.2.1"

import re
import hashlib
import textwrap
from jinja2 import Environment
try:
    from collections import OrderedDict
except ImportError:
    from ordereddict import OrderedDict

class DuplicateEntry(Exception):
    """Raised when a duplicate element is added to the mesh.
    """

class InvalidComponent(Exception):
    """Raised when a referencing a component that does not exist in the mesh.
    """

class InvalidDomain(Exception):
    """Raised when a referencing a domain that does not exist in the mesh.
    """

class InvalidPort(Exception):
    """Raised when a referencing an invalid port on a component in the mesh.
    """

class InvalidConnection(Exception):
    """Raised when there is an attempt to create an invalid connection in the mesh.
    """

class InvalidResource(Exception):
    """Raised when there is an attempt to create an invalid resource in the mesh.
    """

class ConnectionNode(object):
    """Internal representation of a connection between components within the mesh.
    """
    def __init__(self, consumer, producer):
        """Instantiates a new connection between the given consumer and producer.

        :param tuple consumer: consumer of the connection defined as (consumer_component, consumer_port)
        :param tuple producer: producer of the connection defined as (producer_component, producer_port)
        """
        self.consumer = consumer
        self.producer = producer
        self.highlighted = False

    def as_dict(self):
        """Returns a dict representation of the connection.
        """
        consumer_component, consumer_port = self.consumer
        producer_component, producer_port = self.producer
        d = {
            'consumer_component': consumer_component,
            'consumer_port': consumer_port,
            'producer_component': producer_component,
            'producer_port': producer_port,
        }

        if self.highlighted:
            d['highlighted'] = True

        return d

class DomainNeedsConnectionNode(ConnectionNode):
    """Internal representation of component's port being exposed as that of its parent domain."""

    def as_dict(self):
        d = super(DomainNeedsConnectionNode, self).as_dict()
        d['domain_export'] = "needs"
        return d

class DomainProvidesConnectionNode(ConnectionNode):
    """Internal representation of component's port being exposed as that of its parent domain."""

    def as_dict(self):
        d = super(DomainProvidesConnectionNode, self).as_dict()
        d['domain_export'] = "provides"
        return d

class ResourceConnectionNode(ConnectionNode):
    """Internal representation of a connection between components and resources within the mesh.
    """
    def as_dict(self):
        """Returns a dict representation of the connection.
        """
        consumer_component, consumer_port = self.consumer
        d = {
            'consumer_component': consumer_component,
            'consumer_port': consumer_port,
            'resource': self.producer,
        }

        if self.highlighted:
            d['highlighted'] = True

        return d

class Mesh(object):
    """Light-weight representation of a hexagonal mesh that can drive template renderers (e.g. jinja2) to
    produce a textual representation of the mesh.
    """

    GRAPH_TEMPLATE = ''

    def __init__(self):
        self.components = OrderedDict()
        self.connections = OrderedDict()
        self.resources = []
        self._highlighted_resource = set()
        self.connected_consumers = set()

    def add_component(self, component_name, needs_ports=None, provides_ports=None):
        """Adds a component to the mesh.

        :param str component_name: Name of component to add
        :param list needs_ports: List of needs ports for this component
        :param list provides_ports: List of provides ports for this component
        :raises: DuplicateEntry if component of that name already exists
        """
        if component_name in self.components:
            raise DuplicateEntry('Component or Domain with name {0} already exists'.format(component_name))

        self.components[component_name] = ComponentNode(component_name)

        needs = needs_ports or tuple()
        for port_name in needs:
            self.add_needs_port(component_name, port_name)

        provides = provides_ports or tuple()
        for port_name in provides:
            self.add_provides_port(component_name, port_name)

    def add_resource(self, resource_name):
        """Adds a resource (adapter to external data) to the mesh.

        :param str resource_name: Name of resource to add
        :raises: DuplicateEntry if component of that name already exists
        """
        if resource_name in self.resources:
            raise DuplicateEntry('Resource with name {0} already exists'.format(resource_name))

        self.resources.append(resource_name)

    def add_needs_port(self, component_name, port_name):
        """Assigns an additional needs port to an existing component.

        :param str component_name: name of component to add the port to
        :param str port_name: name of needs port to add
        :raises: DuplicateEntry if that needs port already exists on the component
        """
        self.components[component_name].add_needs_port(port_name)

    def add_provides_port(self, component_name, port_name):
        """Assigns an additional provides port to an existing component.

        :param str component_name: name of component to add the port to
        :param str port_name: name of provides port to add
        :raises: DuplicateEntry if that provides port already exists on the component
        """
        self.components[component_name].add_provides_port(port_name)

    def add_connection(self, consumer_component, consumer_port, producer_component, producer_port):
        """Adds a connection between a needs port from a consumer component to the provides port of a producer.

        :param str consumer_component: name of the consumer component
        :param str consumer_port: name of the needs port of the consumer
        :param str producer_component: name of the producer component
        :param str producer_port: name of the provides port of the producer
        """
        try:
            consumer = self.components[consumer_component]
            consumer.assert_is_valid_needs_port(consumer_port)
        except KeyError:
            raise InvalidComponent('{0} component does not exist in the mesh'.format(consumer_component))

        try:
            producer = self.components[producer_component]
            producer.assert_is_valid_provides_port(producer_port)
        except KeyError:
            raise InvalidComponent('{0} component does not exist in the mesh'.format(producer_component))

        # FIXME: using label_for_* here leaks the requirements of a specific viz template (dot) to the conceptual model
        #        of the mesh. Not great. This is essentially a hack to get domains working for now. Need to rethink this
        consumer = consumer.label_for_needs, consumer_port
        producer = producer.label_for_provides, producer_port
        self._add_connection_between_consumer_and_producer(consumer, producer)

    def add_connection_to_resource(self, consumer_component, consumer_port, resource):
        """Adds a connection between a needs port from a consumer component to a resource.

        :param str consumer_component: name of the consumer component
        :param str consumer_port: name of the needs port of the consumer
        :param str resource: name of resource
        """
        try:
            consumer = self.components[consumer_component]
            consumer.assert_is_valid_needs_port(consumer_port)
        except KeyError:
            raise InvalidComponent('{0} component does not exist in the mesh'.format(consumer_component))

        if resource not in self.resources:
            raise InvalidResource('{0} resource does not exist in the mesh'.format(resource))

        consumer = consumer.label_for_needs, consumer_port
        self._add_connection_between_consumer_and_producer(consumer, resource, connectionClass=ResourceConnectionNode)

    def _add_connection_between_consumer_and_producer(self, consumer, producer, connectionClass=ConnectionNode):
        if consumer in self.connected_consumers:
            raise InvalidConnection('{0} already connected'.format(consumer))

        self.connections[consumer, producer] = (connectionClass(consumer, producer))
        self.connected_consumers.add(consumer)

    def add_domain(self, domain_name):
        """Creates a domain which groups together components as a single entity.

        :param str domain_name: name of domain to add
        :raises: DuplicateEntry if component or domain of that name already exists
        """
        if domain_name in self.components:
            raise DuplicateEntry('Component or Domain with name {0} already exists'.format(domain_name))

        self.components[domain_name] = DomainNode(domain_name)

    def add_component_to_domain(self, component_name, domain_name):
        """Adds a component into the given domain.

        :param str component_name: Name of the component
        :param str domain_name: Name of the domain
        """
        try:
            component = self.components[component_name]
        except KeyError:
            raise InvalidComponent('{0} component does not exist in the mesh'.format(component_name))

        # No reason why we cannot eventually support nested domains. We just need to test it out and potentially
        # generalise the way we set up connections between ports.
        if isinstance(component, DomainNode):
            raise InvalidComponent(
                'Cannot add {0} to {1}. Nested domains not yet supported'.format(component_name, domain_name))

        if component.parent:
            raise DuplicateEntry('Component {0} is already part of domain {1}'.format(component.name, component.parent))

        try:
            domain = self.components[domain_name]
        except KeyError:
            raise InvalidDomain('{0} domain does not exist in the mesh'.format(domain_name))

        component.parent = domain_name
        domain.add_child_component(component_name)

    def expose_component_needs_port(self, component_name, port_name):
        """Associated a component's need port to that of its parent domain.

        :param str component_name: Name of component
        :param str port_name: Name of needs port
        """
        try:
            component = self.components[component_name]
        except KeyError:
            raise InvalidComponent('{0} component does not exist in the mesh'.format(component_name))

        component.assert_is_valid_needs_port(port_name)

        try:
            domain = self.components[component.parent]
        except KeyError:
            raise InvalidDomain('Parent domain for {0} component is unspecified or invalid'.format(component_name))

        if port_name not in domain.needs_ports:
            domain.add_needs_port(port_name)

        consumer = component_name, port_name
        producer = domain.label_for_needs, port_name
        self._add_connection_between_consumer_and_producer(consumer, producer, connectionClass=DomainNeedsConnectionNode)

    def expose_component_provides_port(self, component_name, port_name):
        """Associated a component's provides port to that of its parent domain.

        :param str component_name: Name of component
        :param str port_name: Name of provides port
        """
        try:
            component = self.components[component_name]
        except KeyError:
            raise InvalidComponent('{0} component does not exist in the mesh'.format(component_name))

        component.assert_is_valid_provides_port(port_name)

        try:
            domain = self.components[component.parent]
        except KeyError:
            raise InvalidDomain('Parent domain for {0} component is unspecified or invalid'.format(component_name))

        if port_name in domain.provides_ports:
            raise DuplicateEntry('{0} domain already has exposed provides port for {1}'.format(domain.name, port_name))
        else:
            domain.add_provides_port(port_name)

        consumer = domain.label_for_provides, port_name
        producer = component_name, port_name
        self._add_connection_between_consumer_and_producer(consumer, producer, connectionClass=DomainProvidesConnectionNode)

    def highlight_component(self, component_name):
        """Highlights a component in the mesh.

        :param str component_name: name of component to highlight
        """
        try:
            self.components[component_name].highlighted = True
        except KeyError:
            raise InvalidComponent('{0} component does not exist in the mesh'.format(component_name))

    def highlight_connection(self, consumer_component, consumer_port, producer_component, producer_port):
        """Highlights a connection between a needs port from a consumer component to the provides port of a producer.

        :param str consumer_component: name of the consumer component
        :param str consumer_port: name of the needs port of the consumer
        :param str producer_component: name of the producer component
        :param str producer_port: name of the provides port of the producer
        """
        consumer = consumer_component, consumer_port
        producer = producer_component, producer_port
        try:
            self.connections[consumer, producer].highlighted = True
        except KeyError:
            raise InvalidConnection('Invalid Connection: {0} -> {1}'.format(consumer, producer))

    def highlight_connection_to_resource(self, consumer_component, consumer_port, resource):
        """Highlights a connection between a needs port and a resource.

        :param str consumer_component: name of the consumer component
        :param str consumer_port: name of the needs port of the consumer
        :param str resource: name of resource
        """
        consumer = consumer_component, consumer_port
        try:
            self.connections[consumer, resource].highlighted = True
        except KeyError:
            raise InvalidConnection('Invalid Connection: {0} -> {1}'.format(consumer, resource))

    def highlight_resource(self, resource):
        """Highlights a resource in the mesh.

        :param str resource: name of resource to highlight
        """
        if resource not in self.resources:
            raise InvalidResource('{0} resource does not exist in the mesh'.format(resource))

        self._highlighted_resource.add(resource)

    def as_dict(self):
        """Returns a dict representation of the mesh.
        """
        d = {
            'components': [c.as_dict() for c in self.components.values() if not isinstance(c, DomainNode)],
            'domains': [c.as_dict() for c in self.components.values() if isinstance(c, DomainNode)],
            'connections': [c.as_dict() for c in self.connections.values()],
            'resources': [r for r in self.resources],
        }

        if self._highlighted_resource:
            d['highlighted_resources'] = list(self._highlighted_resource)

        return d

class ComponentNode(object):
    """Internal representation of a Component within the mesh.
    """

    def __init__(self, name):
        """Instantiates the component node with a given name

        :param str name: component name
        """
        self.name = name
        self.needs_ports = []
        self.provides_ports = []
        self.highlighted = False
        self.parent = None

    def add_needs_port(self, port_name):
        """Assigns a needs port.

        :param str port_name: name of needs port to add
        :raises: DuplicateEntry if that needs port already exists
        """
        if port_name in self.needs_ports:
            raise DuplicateEntry('Needs port with name {0} already exists for {1}'.format(port_name, self.name))

        self.needs_ports.append(port_name)

    def add_provides_port(self, port_name):
        """Assigns a provides port.

        :param str port_name: name of provides port to add
        :raises: DuplicateEntry if that provides port already exists
        """
        if port_name in self.provides_ports:
            raise DuplicateEntry('Provides port with name {0} already exists for {1}'.format(port_name, self.name))

        self.provides_ports.append(port_name)

    def as_dict(self):
        """Returns a dict representation of the component.
        """
        d = {
            'name': self.name,
            'needs_ports': list(self.needs_ports),
            'provides_ports': list(self.provides_ports),
        }

        if self.highlighted:
            d['highlighted'] = True

        return d

    @property
    def label_for_needs(self):
        """Label used by connections to differentiate between needs and provides placeholders of components."""
        return self.name

    @property
    def label_for_provides(self):
        """Label used by connections to differentiate between needs and provides placeholders of components."""
        return self.name

    def assert_is_valid_needs_port(self, port_name):
        """Raises InvalidPort if given port is not a valid needs port.
        """
        if port_name not in self.needs_ports:
            raise InvalidPort('{0} is not a valid valid needs port for component {1}'.format(port_name, self.name))

    def assert_is_valid_provides_port(self, port_name):
        """Raises InvalidPort if given port is not a valid provides port.
        """
        if port_name not in self.provides_ports:
            raise InvalidPort('{0} is not a valid valid provides port for component {1}'.format(port_name, self.name))

class DomainNode(ComponentNode):
    def __init__(self, name):
        """Instantiates the domain node with a given name

        :param str name: domain name
        """
        self.name = name
        self.needs_ports = []
        self.provides_ports = []
        self.children = set()

    def add_child_component(self, component_name):
        """Adds a component as a child to this domain.

        :param str component_name: Name of component to add
        :raises: DuplicateEntry if compoent already a child of this domain.
        """
        if component_name in self.children:
            raise DuplicateEntry('{0} component is already a child of {1} domain'.format(component_name, self.name))

        self.children.add(component_name)

    def as_dict(self):
        """Returns a dict representation of the component.
        """
        d = {
            'name': self.name,
            'needs_ports': list(self.needs_ports),
            'provides_ports': list(self.provides_ports),
            'children': list(self.children),
            'label_for_needs': self.label_for_needs,
            'label_for_provides': self.label_for_provides,
        }
        return d

    @property
    def label_for_needs(self):
        """Label used by connections to differentiate between needs and provides placeholders of components."""
        return self.name + "__needs"

    @property
    def label_for_provides(self):
        """Label used by connections to differentiate between needs and provides placeholders of components."""
        return self.name + "__provides"

    def assert_is_valid_needs_port(self, port_name):
        """Raises InvalidPort if given port is not a valid needs port.
        """
        if port_name not in self.needs_ports:
            raise InvalidPort('{0} is not a valid needs port for component {1}'.format(port_name, self.name))

    def assert_is_valid_provides_port(self, port_name):
        """Raises InvalidPort if given port is not a valid provides port.
        """
        if port_name not in self.provides_ports:
            raise InvalidPort('{0} is not a valid provides port for component {1}'.format(port_name, self.name))

def render(mesh, template, custom_filters=None):
    """Renders the given mesh using the template text provided.

    The template and filters should be compatible with jinja2

    :param str template: the template text to be used for rendering
    :param dict custom_filters: dict of custom filters where the key is the filter tag and the value is the callables
    :returns: rendered textual representation of the mesh
    """
    jinja_env = Environment()
    if custom_filters:
        jinja_env.filters.update(custom_filters)

    jinja_template = jinja_env.from_string(template)

    return jinja_template.render(mesh.as_dict())

DOT_TEMPLATE = textwrap.dedent('''
    digraph G {

        rankdir=LR;
        node [shape=plaintext];

        {% for component in components %}
        {{ component.name|hash }} [label=<
        <TABLE BORDER="0" CELLBORDER="1" CELLSPACING="0" CELLPADDING="4"{% if component.highlighted %} BGCOLOR="yellow"{% endif %}>
        <TR>
            <TD COLSPAN="2"> {{ component.name|escape }}</TD>
        </TR>
        <TR>
            <TD>
                {% if component.provides_ports %}
                <TABLE BORDER="0" CELLBORDER="1" CELLSPACING="0" CELLPADDING="4">
                {% for port in component.provides_ports %}
                <TR><TD PORT="{{ port|hash_p }}" BGCOLOR="grey">{{ port|escape }}</TD></TR>
                {% endfor %}
                </TABLE>
                {% else %} {% endif %}
            </TD>
            <TD>
                {% if component.needs_ports %}
                <TABLE BORDER="0" CELLBORDER="1" CELLSPACING="0" CELLPADDING="4">
                {% for port in component.needs_ports %}
                <TR><TD PORT="{{ port|hash }}" BGCOLOR="grey">{{ port|escape }}</TD></TR>
                {% endfor %}
                </TABLE>
                {% else %} {% endif %}
            </TD>
        </TR>

        </TABLE>>];
        {% endfor %}

        {% for domain in domains %}
        subgraph cluster_domain_{{ domain.name|hash }} {
            label="{{ domain.name }}";
            style="rounded";
            rank=same;

            {% if domain.provides_ports %}
            {{ domain.label_for_provides|hash }} [label=<<TABLE BORDER="0" CELLBORDER="1" CELLSPACING="4" CELLPADDING="4">
                {% for port in domain.provides_ports %}
                <TR><TD PORT="{{ port|hash_p }}" BGCOLOR="lightgrey">{{ port|escape }}</TD></TR>
                {% endfor %}
            </TABLE>>];
            {% endif %}

            subgraph cluster_domain_{{ domain.name|hash }}_services {
                style="rounded,dashed";
                label="";

                {% for child in domain.children %}
                {{ child|hash }};
                {% endfor %}
            }

            {% if domain.needs_ports %}
            {{ domain.label_for_needs|hash }} [label=<<TABLE BORDER="0" CELLBORDER="1" CELLSPACING="4" CELLPADDING="4">
                {% for port in domain.needs_ports %}
                <TR><TD PORT="{{ port|hash }}" BGCOLOR="lightgrey">{{ port|escape }}</TD></TR>
                {% endfor %}
            </TABLE>>];
            {% endif %}
        }
        {% endfor %}

        {% for resource in resources %}
        {{ resource|hash_p }} [shape="rect";label="{{ resource }}", style="dashed"{% if resource in highlighted_resources %}, color="red"{% endif %}];
        {% endfor %}

        {% for conn in connections %}
        {% if "resource" in conn %}
        {{ conn.consumer_component|hash }}:{{ conn.consumer_port|hash }} -> {{ conn.resource|hash_p }} [style="dashed"{% if conn.highlighted %}, color="red"{% endif %}];
        {% elif "domain_export" in conn %}
        {% if conn.domain_export == "needs" %}
        {{ conn.consumer_component|hash }}:{{ conn.consumer_port|hash }} -> {{ conn.producer_component|hash }}:{{ conn.producer_port|hash }} [color="grey",arrowhead="dot"];
        {% else %}
        {{ conn.consumer_component|hash }}:{{ conn.consumer_port|hash_p }} -> {{ conn.producer_component|hash }}:{{ conn.producer_port|hash_p }} [color="grey",dir="back",arrowtail="dot"];
        {% endif %}
        {% else %}
        {{ conn.consumer_component|hash }}:{{ conn.consumer_port|hash }} -> {{ conn.producer_component|hash }}:{{ conn.producer_port|hash_p }}{% if conn.highlighted %}[color="red"]{% endif %};
        {% endif %}
        {% endfor %}
    }
''').lstrip()

def render_mesh_as_dot(mesh, template=DOT_TEMPLATE):
    """Renders the given mesh in the Graphviz dot format.

    :param Mesh mesh: the mesh to be rendered
    :param str template: alternative template to use
    :returns: textual dot representation of the mesh
    """

    custom_filters = {
        'hash': lambda s: "id" + hashlib.md5(s).hexdigest()[:6],
        # alternative hash for provides ports to avoid conflicts with needs ports with same name
        'hash_p': lambda s: "idp" + hashlib.md5(s).hexdigest()[:6],
        'escape': lambda s: re.sub(r'([{}|"<>])', r'\\\1', s),
    }

    return render(mesh, template, custom_filters=custom_filters)


#!/usr/bin/python
"""
WSGI application that renders du-diff output
"""

import os
import re
import textwrap

import mako.template

from .utils import mako_error_handler
from .du_diff import du_diff, format_du_diff

__author__ = 'Marius Gedminas <marius@gedmin.as>'
__version__ = '0.5'
__date__ = '2018-01-19'

STATIC_ASSETS = os.path.join(os.path.dirname(__file__), 'static')
if not os.path.exists(STATIC_ASSETS):  # nocover: testing installed package
    STATIC_ASSETS = '/usr/share/pov-server-page/static'

DATE_RANGE_RX = re.compile(r'^(\d\d\d\d-\d\d-\d\d)\.\.(\d\d\d\d-\d\d-\d\d)(\.txt)?$')

def get_directory(environ):
    return environ.get('DIRECTORY') or os.getenv('DIRECTORY') or '.'

def get_prefix(environ):
    script_name = environ['SCRIPT_NAME']
    return script_name.rstrip('/')

def fmt(delta):
    # du reports sizes in kibibytes, let's convert to bytes then humanize
    size = delta * 1024
    for unit in 'kB', 'MB', 'GB', 'TB':
        size /= 1000.0
        if abs(size) < 1000:
            break
    return '{size:+,.1f} {unit}'.format(size=size, unit=unit)

class Response(object):

    def __init__(self, body='', content_type='text/html; charset=UTF-8',
                 status='200 OK', headers={}):
        self.body = body
        self.status = status
        self.headers = {'Content-Type': content_type}
        self.headers.update(headers)

def not_found():
    return Response('
404 Not Found
', status='404 Not Found')

def Template(*args, **kw):
    return mako.template.Template(
        error_handler=mako_error_handler,
        strict_undefined=True,
        default_filters=['unicode', 'h'],
        *args, **kw)

STYLESHEET = textwrap.dedent('''
    body {
        margin: 1em;
    }

    .du-diff th:first-child {
        width: 8em;
    }
    .du-diff td:first-child {
        text-align: right;
    }

    .du-diff .sorting {
        color: #888;
        cursor: progress;
    }
    .du-diff th {
        cursor: pointer;
    }
    .du-diff th:hover {
        background: #f5f5f5;
    }
'''.lstrip('\n'))

def stylesheet():
    return Response(STYLESHEET, content_type='text/css')

def bootstrap_stylesheet():
    filename = os.path.join(STATIC_ASSETS, 'css', 'bootstrap.min.css')
    with open(filename, 'rb') as f:
        return Response(f.read(), content_type='text/css')

dudiff_template = Template(textwrap.dedent('''
    <!DOCTYPE html>

        <meta charset="UTF-8">
        <meta http-equiv="X-UA-Compatible" content="IE=edge">
        <meta name="viewport" content="width=device-width, initial-scale=1">

        <title>du-diff for ${location}: ${old}..${new}</title>

        <link rel="stylesheet" href="/static/css/bootstrap.min.css">
        <link rel="stylesheet" href="${prefix}/style.css">

    <%
        if dudiff:
            max_depth = max(row.path.count('/') for row in dudiff)
        else:
            max_depth = 0
    %>

du-diff for ${location} ${old} to ${new}

    % if max_depth > 1:

      Limit to depth

    %     for n in range(1, max_depth + 1):
        <button class="btn btn-default" id="depth-btn-${n}" role="group" aria-label="${n}">${n}
    %     endfor

    % endif

              Delta
              Location

    % for row in dudiff:

              ${fmt(row.delta)}
              ${row.path}

    % endfor

'''))

def render_du_diff(environ, location, old, new, format=None):
    if '.' in location or '/' in location:
        return not_found()
    directory = os.path.join(get_directory(environ), location)
    if not os.path.isdir(directory):
        return not_found()
    old_file = os.path.join(directory, 'du-%s.gz' % old)
    new_file = os.path.join(directory, 'du-%s.gz' % new)
    if not os.path.exists(old_file):
        return not_found()
    if not os.path.exists(new_file):
        return not_found()
    diff = du_diff(old_file, new_file)
    if format == '.txt':
        return Response(format_du_diff(diff).encode('UTF-8'), content_type='text/plain; charset=UTF-8')
    html = dudiff_template.render_unicode(
        location=location, old=old, new=new,
        dudiff=diff, fmt=fmt,
        prefix=get_prefix(environ))
    return Response(html)

def dispatch(environ):
    path_info = environ['PATH_INFO'] or '/'
    if path_info == '/style.css':
        return stylesheet, ()
    if path_info == '/static/css/bootstrap.min.css':
        # only used for debugging
        return bootstrap_stylesheet, ()
    components = path_info.strip('/').split('/')
    if len(components) == 2:
        location, dates = components
        m = DATE_RANGE_RX.match(dates)
        if m:
            old, new, format = m.groups()
            return render_du_diff, (environ, location, old, new, format)
    return not_found, ()

def wsgi_app(environ, start_response):
    view, args = dispatch(environ)
    response = view(*args)
    start_response(response.status, sorted(response.headers.items()))
    body = response.body
    if not isinstance(body, bytes):
        body = body.encode('UTF-8')  # pragma: PY3
    return [body]

application = wsgi_app  # for mod_wsgi

try:
    from importlib import reload
except ImportError:
    pass

def reloading_wsgi_app(environ, start_response):
    # Horrible hack that gives me a fast development loop: reload the code on
    # every request!
    from . import dudiff2html
    reload(dudiff2html)
    return dudiff2html.wsgi_app(environ, start_response)

def main():
    import datetime
    import optparse
    from wsgiref.simple_server import make_server
    parser = optparse.OptionParser(
        'usage: %prog [options] [directory]',
        description="Launch web server showing du-diff output")
    parser.add_option('-p', '--port', type='int', default=8080)
    parser.add_option('--public', default=False,
                      help='accept non-localhost connections')
    opts, args = parser.parse_args()
    if len(args) > 1:
        parser.error("too many arguments")
    if args:
        os.environ['DIRECTORY'] = args[0]
    host = '0.0.0.0' if opts.public else 'localhost'
    httpd = make_server(host, opts.port, reloading_wsgi_app)
    print("Looking for files under subdirectories under %s" % get_directory({}))
    print("Serving http://%s:%d/<subdir>/<date1>..<date2>[.txt]" % (host, opts.port))
    print("Try http://%s:%d/root/%s..%s"
          % (host, opts.port, datetime.date.today() - datetime.timedelta(1), datetime.date.today()))
    try:
        httpd.serve_forever()
    except KeyboardInterrupt:
        pass

if __name__ == '__main__':
    main()


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

from flask import Flask
import proxyproc
import secrets

app_relay = Flask(__name__)
app_relay.config["JSON_AS_ASCII"] = False

secret = secrets.token_urlsafe(32)
app_relay.secret_key = secret

@app_relay.route("/proxyproc", methods=["GET", "POST"])
def get():
    return proxyproc.proxyproc()

if __name__ == "__main__":
    app_relay.run()


# f = open("measurements1.txt")
# print("Opened Measurements")
# registers = {}
# for line in f:
#     vals = line.split("-->")
#     if vals[0] not in registers:
#         registers[vals[0]] = [vals[1]]
#     else:
#         registers[vals[0]].append(vals[1])
#
# f.close()
# print(len(registers["0x01"]))

# Cable 1 (green marker small and fixed) loss:
# 1.07dB (at 865.9MHz)
# 1.07dB (at 865.9MHz)
# 1.07dB (at 865.9MHz)
# 1.07dB (at 865.9MHz)
# 1.08dB (at 865.9MHz)
# 1.08dB (at 865.7MHz)
# 1.08dB (at 865.7MHz)
# 1.08dB (at 865.7MHz)
# 1.08dB (at 867.5MHz)
# 1.08dB (at 867.5MHz)
# 1.08dB (at 867.5MHz)
# 1.08dB (at 867.5MHz)
# 1.076666667 average
# Cable 2 (green marker big and not fixed) loss:
# 1.16dB (at 867.5MHz)
# 1.16dB (at 867.5MHz)
# 1.15dB (at 867.5MHz)
# 1.15dB (at 867.5MHz)
# 1.15dB (at 867.5MHz)
# 1.17dB (at 867.5MHz)
# 1.16dB (at 867.5MHz)
# 1.14dB (at 865.9MHz)
# 1.14dB (at 865.9MHz)
# 1.14dB (at 865.9MHz)
# 1.15dB (at 865.9MHz)
# 1.151818182 average

# 0.0751dB difference
# setup

#       WARP ---- cable 1 ------ splitter ------- cable 2 -------- spectrum analyser
#                                   ||
#                                   ||
#                                   sig gen

# Time: 14:21
# Measured frequency: 865.9MHz, started in 1dBm increments, then moved to 5dBm increments
# so that you can see it more easily when plotted:
# -10
# -15
# -20
# -25
# -30
# squiggle then -35
# -40
# -45
# -50
# -55

# 14:28
# Same freq settings but this time measurement of power from spectrum analyser
# Multiple readings are because the power was varying heavily
# -5.42
# -10.43
# -10.42
# -15.46
# -15.46
# -20.38
# -20.38
# -25.40
# -30.40
# -35.39
# -35.40
# -40.40
# -40.38
# -45.42
# -45.39
# -50.38
# -50.40
# -50.35
# -55.42
# -55.38
# -55.36
# -60.42
# -60.34
# -60.29
# -65.22
# -65.30

import matplotlib.pyplot as plt
import numpy as np
from copy import deepcopy
from scipy.optimize import curve_fit
import scipy.stats as stats

f = open("Mon13Jan1429_test_865.9MHz.log")
print("Opened Measurements")
real_measurements = []
errors = 0
for line in f:

    try:
        vals = line.split(",")
        reading_count = int(vals[0])
        # Start and end readings are within limits for file
        assert reading_count > 359022
        assert reading_count < 505980
        # TSR values match
        assert vals[1] == vals[7]
        # No config errors
        assert int(vals[9]) == 0
        # 0,1,2
        # 3,4,5 are x,y,z orientations
        adc_value = int(vals[6])
        # Must be under 2^12 and over 0
        assert adc_value > 0
        assert adc_value < 2 ** 12 + 1
        real_measurements.append(adc_value)

    except:
        errors += 1
        pass

f.close()

print("{} Errors reading from the file".format(errors))
print("{} Values read successfully".format(len(real_measurements)))

real_measurements = np.array(real_measurements)
grad_measurements = np.diff(real_measurements)
plt.figure(1)
plt.plot(range(len(real_measurements)), real_measurements)
plt.title("Raw Measurements Taken")
plt.xlabel("Measurement Number")
plt.ylabel("ADC Value")
# plt.plot(range(len(grad_measurements)), grad_measurements)

signal_generator_output_levels = [
    -5.42,
    -10.43,
    -15.46,
    -20.38,
    -25.40,
    -30.40,
    -35.40,
    -40.39,
    -45.41,
    -50.39,
    -55.39,
    -60.35,
    -65.26,
]

single_power_measurement_vals = []
a = []

data_to_be_fit_x = []
data_to_be_fit_y = []

# real_measurements[39456]
auto_index_finish = 0
# Separating the different power measurements using threshold on gradient
measurement_index = 0
for i, grad in enumerate(grad_measurements):
    a.append(real_measurements[i])
    if grad > 140:
        b = np.array(deepcopy(a))
        a = []

        if measurement_index < 10:
            single_power_measurement_vals.append(b)
            print(
                "Saving at index: {:5} = {: 4.2f}dBm Mean ADC Value: {:04.0f}".format(
                    i, signal_generator_output_levels[measurement_index], np.mean(b)
                )
            )
            plt.figure(1)
            plt.plot(i, real_measurements[i], "or")
            # plt.annotate("{:.0f}".format(np.mean(b)), (i, real_measurements[i]))
            plt.annotate(
                "{}dBm".format(signal_generator_output_levels[measurement_index]),
                (i, real_measurements[i]),
            )
            measurement_index += 1
            auto_index_finish = i
        else:
            break

# Manually splice last two values since they are very noisy
end_55 = 39455
end_60 = 45090

# 55
plt.plot(auto_index_finish + 1, real_measurements[auto_index_finish + 1], "og")
plt.plot(end_55, real_measurements[end_55], "og")
b = np.array(real_measurements[auto_index_finish + 1 : end_55 + 1])

assert b[-1] == real_measurements[end_55]
single_power_measurement_vals.append(b)

# 60
plt.plot(end_55 + 1, real_measurements[end_55 + 1], "og")
plt.plot(end_60, real_measurements[end_60], "og")

b = np.array(real_measurements[end_55 + 1 : end_60 + 1])
single_power_measurement_vals.append(b)
# 65
plt.plot(end_60 + 1, real_measurements[end_60 + 1], "og")
plt.plot(len(real_measurements), real_measurements[-1], "og")
b = np.array(real_measurements[end_60 + 1 :])

single_power_measurement_vals.append(b)

plt.figure(2)
plt.title("Average Measured ADC Values vs True Signal Power")
plt.ylabel("Raw ADC Reading [N/A]")
plt.xlabel("True Signal Power [dBm]")

y_err = np.empty((2, len(single_power_measurement_vals)))

for i, vals in enumerate(single_power_measurement_vals):
    data_to_be_fit_x.append(signal_generator_output_levels[i])
    data_to_be_fit_y.append(np.mean(vals))

    y_err[0, i] = np.mean(vals) - np.amin(vals)
    y_err[1, i] = np.amax(vals) - np.mean(vals)
    print(
        "Reading at {: 4.2f}dBm, Minimum data point at index {}, length {}".format(
            signal_generator_output_levels[i],
            np.where(vals == np.amin(vals)),
            len(vals),
        )
    )

data_to_be_fit_x = np.array(data_to_be_fit_x)
data_to_be_fit_y = np.array(data_to_be_fit_y)
plt.errorbar(data_to_be_fit_x, data_to_be_fit_y, yerr=y_err, fmt="r+")

def lin(x, m, c):
    return m * x + c

# assuming ydata = f(xdata, *params)
popt, pcov = curve_fit(lin, data_to_be_fit_x[:-3], data_to_be_fit_y[:-3])
# print(data_to_be_fit_x[:-3])
print("Gradient: {}, Intercept: {}".format(*popt))
plt.plot(
    data_to_be_fit_x,
    lin(data_to_be_fit_x, *popt),
    "-b",
    label="ADC Reading = {:.1f} (True Power) + {:.1f}".format(*popt),
)
plt.legend()

plt.figure(3)

plt.hist(
    single_power_measurement_vals[-3:],
    bins=range(
        min(single_power_measurement_vals[-3]), max(single_power_measurement_vals[-1])
    ),
    width=1,
)
plt.title("Histogram of Measurements for -55.39 ,-60.35 -65.26dBm")
plt.xlabel("ADC Value")
plt.ylabel("Occurence")
plt.figure(4)

# 5 % Hypothesis Test
alpha = 0.05
print(len(signal_generator_output_levels))
# enumerate(single_power_measurement_vals[:-3]):
for j, good_measurements in enumerate(single_power_measurement_vals[-3:]):
    i = j + 10
    stat, pvalue = stats.normaltest(good_measurements)
    mean = np.mean(good_measurements)
    std = np.std(good_measurements)
    name = "{: 4.2f}dBm".format(signal_generator_output_levels[i])
    normal = pvalue < alpha
    print(
        "For {:>9} measurements: n={:5d} mean={:4.0f} std={:4.2f}, Normal Fit: {} stat={} pvalue={} ".format(
            name, len(good_measurements), mean, std, normal, stat, pvalue
        )
    )
    plt.hist(
        good_measurements,
        bins=range(int(mean) - 100, int(mean) + 100),
        width=1,
        label="{}dBm".format(signal_generator_output_levels[i]),
        density=True,
    )
    if normal:
        x = np.linspace(mean - 3 * std, mean + 3 * std, 100)
        y = stats.norm.pdf(x, loc=mean, scale=std)
        plt.plot(
            x,
            y,
            "r",
            label="{}dBm Normal Fit".format(signal_generator_output_levels[i]),
        )
    # zero_mean = good_measurements - np.mean(good_measurements)
    # # n = len(good_measurements)
    # # S = np.std(good_measurements, ddof=1)
    # # t_dist = zero_mean / (S / np.sqrt(n))
    # s = np.std(good_measurements)
    # dist = zero_mean / s
    #
    # plt.hist(
    #     dist,
    #     bins=np.linspace(-3, 3, 40),
    #     # bins=range(int(np.floor(min(t_dist))), int(np.ceil(max(t_dist)))),
    #     label="{}dBm".format(signal_generator_output_levels[i]),
    # )
plt.title("Histogram of Noisiest Measurements")
plt.xlabel("ADC Value")
plt.ylabel("Occurence")
plt.legend()
# plt.figure(4)
# plt.hist(single_power_measurement_vals[-2], bins=50)
# plt.figure(5)
# plt.hist(single_power_measurement_vals[-1], bins=50)
# single_power_measurement_vals = np.array(single_power_measurement_vals)

# plt.hist(
#     single_power_measurement_vals,
#     rwidth=1,
#     bins=range(
#         min(single_power_measurement_vals[0]), max(single_power_measurement_vals[-1])
#     ),
# )

# for index, measured_adc_vals in enumerate(single_power_measurement_vals):
#     real_power = signal_generator_output_levels[index]
#     plt.figure(3+index)
#     plt.title("Distribution of Measured Power Values")
#

plt.show()


import torch
import numpy as np

import random
from torch._six import nan
from itertools import permutations

from torch.testing._internal.common_utils import \
    (TestCase, run_tests, make_tensor, slowTest)
from torch.testing._internal.common_device_type import \
    (instantiate_device_type_tests, dtypes, onlyOnCPUAndCUDA,
     skipCUDAIfRocm, onlyCUDA, dtypesIfCUDA, onlyCPU, largeTensorTest)

# TODO: remove this
SIZE = 100

class TestSortAndSelect(TestCase):

    def assertIsOrdered(self, order, x, mxx, ixx, task):
        SIZE = x.size(1)
        if order == 'descending':
            def check_order(a, b):
                # `a != a` because we put NaNs
                # at the end of ascending sorted lists,
                # and the beginning of descending ones.
                return ((a != a) | (a >= b)).all().item()
        elif order == 'ascending':
            def check_order(a, b):
                # see above
                return ((b != b) | (a <= b)).all().item()
        else:
            error('unknown order "{}", must be "ascending" or "descending"'.format(order))

        are_ordered = True
        for k in range(1, SIZE):
            self.assertTrue(check_order(mxx[:, k - 1], mxx[:, k]),
                            'torch.sort ({}) values unordered for {}'.format(order, task))

        seen = set()
        indicesCorrect = True
        size0 = x.size(0)
        size = x.size(x.dim() - 1)
        x = x.tolist()
        mxx = mxx.tolist()
        ixx = ixx.tolist()
        for k in range(size0):
            seen.clear()
            for j in range(size):
                self.assertEqual(x[k][ixx[k][j]], mxx[k][j],
                                 msg='torch.sort ({}) indices wrong for {}'.format(order, task))
                seen.add(ixx[k][j])
            self.assertEqual(len(seen), size)

    def test_sort(self, device):
        # on CUDA 2048 vs >2048 have different code path for the dim being sorted
        for SIZE in (4, 2049):
            x = torch.rand(4, SIZE, device=device)
            res1val, res1ind = torch.sort(x)

            # Test use of result tensor
            res2val = torch.tensor((), device=device)
            res2ind = torch.tensor((), device=device, dtype=torch.long)
            torch.sort(x, out=(res2val, res2ind))
            self.assertEqual(res1val, res2val, atol=0, rtol=0)
            self.assertEqual(res1ind, res2ind, atol=0, rtol=0)
            self.assertEqual(torch.argsort(x), res1ind)
            self.assertEqual(x.argsort(), res1ind)

            # Test sorting of random numbers
            self.assertIsOrdered('ascending', x, res2val, res2ind, 'random')

            # Test simple sort
            self.assertEqual(
                torch.sort(torch.tensor((50, 40, 30, 20, 10), device=device))[0],
                torch.tensor((10, 20, 30, 40, 50), device=device),
                atol=0, rtol=0
            )

            # Test that we still have proper sorting with duplicate keys
            x = torch.floor(torch.rand(4, SIZE, device=device) * 10)
            torch.sort(x, out=(res2val, res2ind))
            self.assertIsOrdered('ascending', x, res2val, res2ind, 'random with duplicate keys')

            # DESCENDING SORT
            x = torch.rand(4, SIZE, device=device)
            res1val, res1ind = torch.sort(x, x.dim() - 1, True)

            # Test use of result tensor
            res2val = torch.tensor((), device=device)
            res2ind = torch.tensor((), device=device, dtype=torch.long)
            torch.sort(x, x.dim() - 1, True, out=(res2val, res2ind))
            self.assertEqual(res1val, res2val, atol=0, rtol=0)
            self.assertEqual(res1ind, res2ind, atol=0, rtol=0)
            self.assertEqual(torch.argsort(x, x.dim() - 1, True), res1ind)
            self.assertEqual(x.argsort(x.dim() - 1, True), res1ind)

            # Test sorting of random numbers
            self.assertIsOrdered('descending', x, res2val, res2ind, 'random')

            # Test simple sort task
            self.assertEqual(
                torch.sort(torch.tensor((10, 20, 30, 40, 50), device=device), 0, True)[0],
                torch.tensor((50, 40, 30, 20, 10), device=device),
                atol=0, rtol=0
            )

            # Test that we still have proper sorting with duplicate keys
            self.assertIsOrdered('descending', x, res2val, res2ind, 'random with duplicate keys')

            # Test sorting with NaNs
            x = torch.rand(4, SIZE, device=device)
            x[1][2] = float('NaN')
            x[3][0] = float('NaN')
            torch.sort(x, out=(res2val, res2ind))
            self.assertIsOrdered('ascending', x, res2val, res2ind,
                                 'random with NaNs')
            torch.sort(x, out=(res2val, res2ind), descending=True)
            self.assertIsOrdered('descending', x, res2val, res2ind,
                                 'random with NaNs')

    # FIXME: remove torch.bool from unsupported types once support is added for cub sort
    @dtypes(*set(torch.testing.get_all_dtypes()) - {torch.bool, torch.bfloat16, torch.complex64, torch.complex128})
    def test_stable_sort(self, device, dtype):
        if self.device_type == 'cpu':
            sizes = (100, 1000, 10000)
        elif self.device_type == 'cuda':
            # On CUDA, stable sort is supported only when the size of
            # the sorted dim is greater than 2048
            sizes = (1025, 10000)
        else:
            return
        for ncopies in sizes:
            x = torch.tensor([0, 1] * ncopies, dtype=dtype, device=device)
            _, idx = x.sort(stable=True)
            self.assertEqual(
                idx[:ncopies],
                torch.arange(start=0, end=2 * ncopies, step=2, device=device)
            )
            self.assertEqual(
                idx[ncopies:],
                torch.arange(start=1, end=2 * ncopies, step=2, device=device)
            )

    @onlyCUDA
    @dtypes(torch.uint8)
    @largeTensorTest('200GB')  # Unfortunately 80GB A100 is not large enough
    def test_sort_large(self, device, dtype):
        t0 = torch.randperm(8192, device=device).to(dtype)
        t = t0.view(1, 8192).expand(2 ** 18 + 1, -1).contiguous()
        v, i = t.sort()
        del t
        iv, im = i.var_mean(dim=0)
        del i
        vv, vm = v.var_mean(dim=0)
        del v
        self.assertEqual(vv, torch.zeros_like(vv))
        self.assertEqual(iv, torch.zeros_like(iv))
        self.assertEqual(vm, torch.arange(255, dtype=dtype, device=device))
        self.assertEqual(im, t0.sort().indices)

    def _test_sort_discontiguous(self, device, dtype):
        # on CUDA 2048 vs >2048 have different code path for the dim being sorted
        sizes = (5, 7, 2049)
        for shape in permutations(sizes):
            for perm in permutations((0, 1, 2)):
                for dim in range(3):
                    t = torch.randn(shape, device=device, dtype=dtype).permute(perm)
                    r1 = t.sort(dim=dim)
                    r2 = t.contiguous().sort(dim=dim)
                    self.assertEqual(r1, r2)
                    n = t.size(dim)

                    # assert ordered
                    self.assertTrue((r1.values.narrow(dim, 1, n - 1) >= r1.values.narrow(dim, 0, n - 1)).all())

                    # assert that different segments does not mix, which can easily happen
                    # if the stride is not handled correctly
                    self.assertTrue((t.unsqueeze(-1).transpose(dim, -1) == r1.values.unsqueeze(-1)).any(dim=dim).any(dim=-1).all())

                    # assert stride is preserved
                    if self.device_type == 'cuda':
                        # FIXME: this behavior should be true for all cases, not
                        # just the one specified in if condition
                        self.assertEqual(r1.values.stride(), t.stride())
                        self.assertEqual(r1.indices.stride(), t.stride())

    @onlyCUDA
    @dtypes(torch.float32)
    def test_sort_discontiguous(self, device, dtype):
        self._test_sort_discontiguous(device, dtype)

    @slowTest  # this test is slow on CPU, but not on CUDA
    @onlyCPU
    @dtypes(torch.float32)
    def test_sort_discontiguous_slow(self, device, dtype):
        self._test_sort_discontiguous(device, dtype)

    # FIXME: remove torch.bool from unsupported types once support is added for cub sort
    @dtypes(*set(torch.testing.get_all_dtypes()) - {torch.bool, torch.bfloat16, torch.complex64, torch.complex128})
    def test_stable_sort_against_numpy(self, device, dtype):
        if dtype in torch.testing.floating_types_and(torch.float16):
            inf = float('inf')
            neg_inf = -float('inf')
            nan = float('nan')
        else:
            if dtype != torch.bool:
                # no torch.iinfo support for torch.bool
                inf = torch.iinfo(dtype).max
                neg_inf = torch.iinfo(dtype).min
            else:
                inf = True
                neg_inf = ~inf
            # no nan for integral types, we use inf instead for simplicity
            nan = inf

        def generate_samples():
            from itertools import chain, combinations

            for sizes in [(1025,), (10000,)]:
                size = sizes[0]
                # binary strings
                yield (torch.tensor([0, 1] * size, dtype=dtype, device=device), 0)

            if self.device_type == 'cuda':
                return

            yield (torch.tensor([0, 1] * 100, dtype=dtype, device=device), 0)

            def repeated_index_fill(t, dim, idxs, vals):
                res = t
                for idx, val in zip(idxs, vals):
                    res = res.index_fill(dim, idx, val)
                return res

            for sizes in [(1, 10), (10, 1), (10, 10), (10, 10, 10)]:
                size = min(*sizes)
                x = (torch.randn(*sizes, device=device) * size).to(dtype)
                yield (x, 0)

                # Generate tensors which are being filled at random locations
                # with values from the non-empty subsets of the set (inf, neg_inf, nan)
                # for each dimension.
                n_fill_vals = 3  # cardinality of (inf, neg_inf, nan)
                for dim in range(len(sizes)):
                    idxs = (torch.randint(high=size, size=(size // 10,)) for i in range(n_fill_vals))
                    vals = (inf, neg_inf, nan)
                    subsets = chain.from_iterable(combinations(list(zip(idxs, vals)), r)
                                                  for r in range(1, n_fill_vals + 1))
                    for subset in subsets:
                        idxs_subset, vals_subset = zip(*subset)
                        yield (repeated_index_fill(x, dim, idxs_subset, vals_subset), dim)

        for sample, dim in generate_samples():
            _, idx_torch = sample.sort(dim=dim, stable=True)
            sample_numpy = sample.cpu().numpy()
            idx_numpy = np.argsort(sample_numpy, axis=dim, kind='stable')
            self.assertEqual(idx_torch, idx_numpy)

    @dtypes(*(torch.testing.get_all_int_dtypes() + torch.testing.get_all_fp_dtypes(include_bfloat16=False)))
    def test_msort(self, device, dtype):
        def test(shape):
            tensor = make_tensor(shape, device, dtype, low=-9, high=9)
            if tensor.size() != torch.Size([]):
                expected = torch.from_numpy(np.msort(tensor.cpu().numpy()))
            else:
                expected = tensor  # numpy.msort() does not support empty shapes tensor

            result = torch.msort(tensor)
            self.assertEqual(result, expected)

            out = torch.empty_like(result)
            torch.msort(tensor, out=out)
            self.assertEqual(out, expected)

        shapes = (
            [],
            [0, ],
            [20, ],
            [1, 20],
            [30, 30],
            [10, 20, 30]
        )
        for shape in shapes:
            test(shape)

    def test_topk(self, device):
        def topKViaSort(t, k, dim, dir):
            sorted, indices = t.sort(dim, dir)
            return sorted.narrow(dim, 0, k), indices.narrow(dim, 0, k)

        def compareTensors(t, res1, ind1, res2, ind2, dim):
            # Values should be exactly equivalent
            self.assertEqual(res1, res2, atol=0, rtol=0)

            # Indices might differ based on the implementation, since there is
            # no guarantee of the relative order of selection
            if not ind1.eq(ind2).all():
                # To verify that the indices represent equivalent elements,
                # gather from the input using the topk indices and compare against
                # the sort indices
                vals = t.gather(dim, ind2)
                self.assertEqual(res1, vals, atol=0, rtol=0)

        def compare(t, k, dim, dir):
            topKVal, topKInd = t.topk(k, dim, dir, True)
            sortKVal, sortKInd = topKViaSort(t, k, dim, dir)
            compareTensors(t, sortKVal, sortKInd, topKVal, topKInd, dim)

        t = torch.rand(random.randint(1, SIZE),
                       random.randint(1, SIZE),
                       random.randint(1, SIZE), device=device)

        for _kTries in range(3):
            for _dimTries in range(3):
                for transpose in (True, False):
                    for dir in (True, False):
                        testTensor = t
                        if transpose:
                            dim1 = random.randrange(t.ndimension())
                            dim2 = dim1
                            while dim1 == dim2:
                                dim2 = random.randrange(t.ndimension())

                            testTensor = t.transpose(dim1, dim2)

                        dim = random.randrange(testTensor.ndimension())
                        k = random.randint(1, testTensor.size(dim))
                        compare(testTensor, k, dim, dir)

    def test_topk_arguments(self, device):
        q = torch.randn(10, 2, 10, device=device)
        # Make sure True isn't mistakenly taken as the 2nd dimension (interpreted as 1)
        self.assertRaises(TypeError, lambda: q.topk(4, True))

    @skipCUDAIfRocm
    def test_unique_dim(self, device):
        self.assertFalse(hasattr(torch, 'unique_dim'))

        def run_test(device, dtype):
            x = torch.tensor([[[1., 1.],
                               [0., 1.],
                               [2., 1.],
                               [0., 1.]],
                              [[1., 1.],
                               [0., 1.],
                               [2., 1.],
                               [0., 1.]]],
                             dtype=dtype,
                             device=device)
            x_empty = torch.empty(5, 0, dtype=dtype, device=device)
            x_ill_formed_empty = torch.empty(5, 0, 0, dtype=dtype, device=device)
            x_ill_formed_empty_another = torch.empty(5, 0, 5, dtype=dtype, device=device)
            expected_unique_dim0 = torch.tensor([[[1., 1.],
                                                  [0., 1.],
                                                  [2., 1.],
                                                  [0., 1.]]],
                                                dtype=dtype,
                                                device=device)
            expected_inverse_dim0 = torch.tensor([0, 0])
            expected_counts_dim0 = torch.tensor([2])
            expected_unique_dim1 = torch.tensor([[[0., 1.],
                                                  [1., 1.],
                                                  [2., 1.]],
                                                 [[0., 1.],
                                                  [1., 1.],
                                                  [2., 1.]]],
                                                dtype=dtype,
                                                device=device)
            expected_unique_dim1_bool = torch.tensor([[[False, True], [True, True]],
                                                      [[False, True], [True, True]]],
                                                     dtype=torch.bool,
                                                     device=device)
            expected_inverse_dim1 = torch.tensor([1, 0, 2, 0])
            expected_inverse_dim1_bool = torch.tensor([1, 0, 1, 0])
            expected_counts_dim1 = torch.tensor([2, 1, 1])
            expected_counts_dim1_bool = torch.tensor([2, 2])
            expected_unique_dim2 = torch.tensor([[[1., 1.],
                                                  [0., 1.],
                                                  [2., 1.],
                                                  [0., 1.]],
                                                 [[1., 1.],
                                                  [0., 1.],
                                                  [2., 1.],
                                                  [0., 1.]]],
                                                dtype=dtype,
                                                device=device)
            expected_inverse_dim2 = torch.tensor([0, 1])
            expected_counts_dim2 = torch.tensor([1, 1])
            expected_unique_empty = torch.tensor([], dtype=dtype, device=device)
            expected_inverse_empty = torch.tensor([], dtype=torch.long, device=device)
            expected_counts_empty = torch.tensor([], dtype=torch.long, device=device)
            # dim0
            x_unique = torch.unique(x, dim=0)
            self.assertEqual(expected_unique_dim0, x_unique)

            x_unique, x_inverse = torch.unique(
                x,
                return_inverse=True,
                dim=0)
            self.assertEqual(expected_unique_dim0, x_unique)
            self.assertEqual(expected_inverse_dim0, x_inverse)

            x_unique, x_counts = torch.unique(
                x,
                return_inverse=False,
                return_counts=True,
                dim=0)
            self.assertEqual(expected_unique_dim0, x_unique)
            self.assertEqual(expected_counts_dim0, x_counts)

            x_unique, x_inverse, x_counts = torch.unique(
                x,
                return_inverse=True,
                return_counts=True,
                dim=0)
            self.assertEqual(expected_unique_dim0, x_unique)
            self.assertEqual(expected_inverse_dim0, x_inverse)
            self.assertEqual(expected_counts_dim0, x_counts)

            # dim1
            x_unique = torch.unique(x, dim=1)
            if x.dtype == torch.bool:
                self.assertEqual(expected_unique_dim1_bool, x_unique)
            else:
                self.assertEqual(expected_unique_dim1, x_unique)

            x_unique, x_inverse = torch.unique(
                x,
                return_inverse=True,
                dim=1)
            if x.dtype == torch.bool:
                self.assertEqual(expected_unique_dim1_bool, x_unique)
                self.assertEqual(expected_inverse_dim1_bool, x_inverse)
            else:
                self.assertEqual(expected_unique_dim1, x_unique)
                self.assertEqual(expected_inverse_dim1, x_inverse)

            x_unique, x_counts = torch.unique(
                x,
                return_inverse=False,
                return_counts=True,
                dim=1)
            if x.dtype == torch.bool:
                self.assertEqual(expected_unique_dim1_bool, x_unique)
                self.assertEqual(expected_counts_dim1_bool, x_counts)
            else:
                self.assertEqual(expected_unique_dim1, x_unique)
                self.assertEqual(expected_counts_dim1, x_counts)

            x_unique, x_inverse, x_counts = torch.unique(
                x,
                return_inverse=True,
                return_counts=True,
                dim=1)
            if x.dtype == torch.bool:
                self.assertEqual(expected_unique_dim1_bool, x_unique)
                self.assertEqual(expected_inverse_dim1_bool, x_inverse)
                self.assertEqual(expected_counts_dim1_bool, x_counts)
            else:
                self.assertEqual(expected_unique_dim1, x_unique)
                self.assertEqual(expected_inverse_dim1, x_inverse)
                self.assertEqual(expected_counts_dim1, x_counts)

            # dim2
            x_unique = torch.unique(x, dim=2)
            self.assertEqual(expected_unique_dim2, x_unique)

            x_unique, x_inverse = torch.unique(
                x,
                return_inverse=True,
                dim=2)
            self.assertEqual(expected_unique_dim2, x_unique)
            self.assertEqual(expected_inverse_dim2, x_inverse)

            x_unique, x_counts = torch.unique(
                x,
                return_inverse=False,
                return_counts=True,
                dim=2)
            self.assertEqual(expected_unique_dim2, x_unique)
            self.assertEqual(expected_counts_dim2, x_counts)

            x_unique, x_inverse, x_counts = torch.unique(
                x,
                return_inverse=True,
                return_counts=True,
                dim=2)
            self.assertEqual(expected_unique_dim2, x_unique)
            self.assertEqual(expected_inverse_dim2, x_inverse)
            self.assertEqual(expected_counts_dim2, x_counts)

            # test empty tensor
            x_unique, x_inverse, x_counts = torch.unique(
                x_empty,
                return_inverse=True,
                return_counts=True,
                dim=1)
            self.assertEqual(expected_unique_empty, x_unique)
            self.assertEqual(expected_inverse_empty, x_inverse)
            self.assertEqual(expected_counts_empty, x_counts)

            # test not a well formed tensor
            # Checking for runtime error, as this is the expected behaviour
            with self.assertRaises(RuntimeError):
                torch.unique(
                    x_ill_formed_empty,
                    return_inverse=True,
                    return_counts=True,
                    dim=1)

            # test along dim2
            with self.assertRaises(RuntimeError):
                torch.unique(
                    x_ill_formed_empty_another,
                    return_inverse=True,
                    return_counts=True,
                    dim=2)

            # test consecutive version
            y = torch.tensor(
                [[0, 1],
                 [0, 1],
                 [0, 1],
                 [1, 2],
                 [1, 2],
                 [3, 4],
                 [0, 1],
                 [0, 1],
                 [3, 4],
                 [1, 2]],
                dtype=dtype,
                device=device
            )
            expected_y_unique = torch.tensor(
                [[0, 1],
                 [1, 2],
                 [3, 4],
                 [0, 1],
                 [3, 4],
                 [1, 2]],
                dtype=dtype,
                device=device
            )
            expected_y_inverse = torch.tensor([0, 0, 0, 1, 1, 2, 3, 3, 4, 5], dtype=torch.int64, device=device)
            expected_y_counts = torch.tensor([3, 2, 1, 2, 1, 1], dtype=torch.int64, device=device)
            expected_y_inverse_bool = torch.tensor([0, 0, 0, 1, 1, 1, 2, 2, 3, 3], dtype=torch.int64, device=device)
            expected_y_counts_bool = torch.tensor([3, 3, 2, 2], dtype=torch.int64, device=device)
            y_unique, y_inverse, y_counts = torch.unique_consecutive(y, return_inverse=True, return_counts=True, dim=0)
            if x.dtype == torch.bool:
                self.assertEqual(expected_y_inverse_bool, y_inverse)
                self.assertEqual(expected_y_counts_bool, y_counts)
            else:
                self.assertEqual(expected_y_inverse, y_inverse)
                self.assertEqual(expected_y_counts, y_counts)

        run_test(device, torch.float)
        run_test(device, torch.double)
        run_test(device, torch.long)
        run_test(device, torch.uint8)
        run_test(device, torch.bool)

    @onlyCUDA
    def test_topk_noncontiguous_gpu(self, device):
        t = torch.randn(20, device=device)[::2]
        top1, idx1 = t.topk(5)
        top2, idx2 = t.contiguous().topk(5)
        self.assertEqual(top1, top2)
        self.assertEqual(idx1, idx2)

    @dtypes(torch.int8, torch.uint8, torch.int16, torch.int32, torch.int64)
    def test_topk_integral(self, device, dtype):
        a = torch.randint(torch.iinfo(dtype).min, torch.iinfo(dtype).max, size=(10,),
                          dtype=dtype, device=device)
        sort_topk = a.sort()[0][-5:].flip(0)
        topk = a.topk(5)
        self.assertEqual(sort_topk, topk[0])      # check values
        self.assertEqual(sort_topk, a[topk[1]])   # check indices

    @dtypesIfCUDA(*torch.testing.get_all_fp_dtypes())
    @dtypes(torch.float, torch.double)
    def test_topk_nonfinite(self, device, dtype):
        x = torch.tensor([float('nan'), float('inf'), 1e4, 0, -1e4, -float('inf')], device=device, dtype=dtype)
        val, idx = x.topk(4)
        expect = torch.tensor([float('nan'), float('inf'), 1e4, 0], device=device, dtype=dtype)
        self.assertEqual(val, expect)
        self.assertEqual(idx, [0, 1, 2, 3])

        val, idx = x.topk(4, largest=False)
        expect = torch.tensor([-float('inf'), -1e4, 0, 1e4], device=device, dtype=dtype)
        self.assertEqual(val, expect)
        self.assertEqual(idx, [5, 4, 3, 2])

    def test_topk_4d(self, device):
        x = torch.ones(2, 3072, 2, 2, device=device)
        x[:, 1, :, :] *= 2.
        x[:, 10, :, :] *= 1.5
        val, ind = torch.topk(x, k=2, dim=1)
        expected_ind = torch.ones(2, 2, 2, 2, dtype=torch.long, device=device)
        expected_ind[:, 1, :, :] = 10
        expected_val = torch.ones(2, 2, 2, 2, device=device)
        expected_val[:, 0, :, :] *= 2.
        expected_val[:, 1, :, :] *= 1.5
        self.assertEqual(val, expected_val, atol=0, rtol=0)
        self.assertEqual(ind, expected_ind, atol=0, rtol=0)

    def _test_unique_scalar_empty(self, dtype, device, f):
        # test scalar
        x = torch.tensor(0, dtype=dtype, device=device)
        unique, inverse, counts = f(x, return_inverse=True, return_counts=True)
        expected_unique = torch.tensor([0], dtype=dtype, device=device)
        expected_inverse = torch.tensor(0, device=device)
        expected_counts = torch.tensor([1], device=device)
        self.assertEqual(unique, expected_unique)
        self.assertEqual(inverse, expected_inverse)
        self.assertEqual(counts, expected_counts)

        # test zero sized tensor
        x = torch.zeros((0, 0, 3), dtype=dtype, device=device)
        unique, inverse, counts = f(x, return_inverse=True, return_counts=True)
        expected_unique = torch.tensor([], dtype=dtype, device=device)
        expected_inverse = torch.empty((0, 0, 3), dtype=torch.long, device=device)
        expected_counts = torch.tensor([], dtype=torch.long, device=device)
        self.assertEqual(unique, expected_unique)
        self.assertEqual(inverse, expected_inverse)
        self.assertEqual(counts, expected_counts)

    def _test_unique_with_expects(self, device, dtype, f, x, expected_unique, expected_inverse, expected_counts, additional_shape):
        def ensure_tuple(x):
            if isinstance(x, torch.Tensor):
                return (x,)
            return x

        for return_inverse in [True, False]:
            for return_counts in [True, False]:
                # test with expected
                ret = ensure_tuple(f(x, return_inverse=return_inverse, return_counts=return_counts))
                self.assertEqual(len(ret), 1 + int(return_inverse) + int(return_counts))
                self.assertEqual(expected_unique, ret[0])
                if return_inverse:
                    self.assertEqual(expected_inverse, ret[1])
                if return_counts:
                    count_index = 1 + int(return_inverse)
                    self.assertEqual(expected_counts, ret[count_index])

                # tests per-element unique on a higher rank tensor.
                y = x.view(additional_shape)
                y_unique, y_inverse, y_counts = f(y, return_inverse=True, return_counts=True)
                self.assertEqual(expected_unique, y_unique)
                self.assertEqual(expected_inverse.view(additional_shape), y_inverse)
                self.assertEqual(expected_counts, y_counts)

    @dtypes(*set(torch.testing.get_all_dtypes()) - {torch.bfloat16, torch.complex64, torch.complex128})
    def test_unique(self, device, dtype):
        if dtype is torch.half and self.device_type == 'cpu':
            return  # CPU does not have half support

        def ensure_tuple(x):
            if isinstance(x, torch.Tensor):
                return (x,)
            return x

        if dtype is torch.bool:
            x = torch.tensor([True, False, False, False, True, False, True, False], dtype=torch.bool, device=device)
            expected_unique = torch.tensor([False, True], dtype=torch.bool, device=device)
            expected_inverse = torch.tensor([1, 0, 0, 0, 1, 0, 1, 0], dtype=torch.long, device=device)
            expected_counts = torch.tensor([5, 3], dtype=torch.long, device=device)
        else:
            x = torch.tensor([1, 2, 3, 2, 8, 5, 2, 3], dtype=dtype, device=device)
            expected_unique = torch.tensor([1, 2, 3, 5, 8], dtype=dtype, device=device)
            expected_inverse = torch.tensor([0, 1, 2, 1, 4, 3, 1, 2], device=device)
            expected_counts = torch.tensor([1, 3, 2, 1, 1], device=device)

        # test sorted unique
        fs = [
            lambda x, **kwargs: torch.unique(x, sorted=True, **kwargs),
            lambda x, **kwargs: x.unique(sorted=True, **kwargs),
        ]
        for f in fs:
            self._test_unique_with_expects(device, dtype, f, x, expected_unique, expected_inverse, expected_counts, (2, 2, 2))
            self._test_unique_scalar_empty(dtype, device, f)

        # test unsorted unique
        fs = [
            lambda x, **kwargs: torch.unique(x, sorted=False, **kwargs),
            lambda x, **kwargs: x.unique(sorted=False, **kwargs)
        ]
        for f in fs:
            self._test_unique_scalar_empty(dtype, device, f)
            for return_inverse in [True, False]:
                for return_counts in [True, False]:
                    ret = ensure_tuple(f(x, return_inverse=return_inverse, return_counts=return_counts))
                    self.assertEqual(len(ret), 1 + int(return_inverse) + int(return_counts))
                    x_list = x.tolist()
                    x_unique_list = ret[0].tolist()
                    self.assertEqual(expected_unique.tolist(), sorted(x_unique_list))
                    if return_inverse:
                        x_inverse_list = ret[1].tolist()
                        for i, j in enumerate(x_inverse_list):
                            self.assertEqual(x_list[i], x_unique_list[j])
                    if return_counts:
                        count_index = 1 + int(return_inverse)
                        x_counts_list = ret[count_index].tolist()
                        for i, j in zip(x_unique_list, x_counts_list):
                            count = 0
                            for k in x_list:
                                if k == i:
                                    count += 1
                            self.assertEqual(j, count)

    @dtypes(*set(torch.testing.get_all_dtypes()) - {torch.bfloat16, torch.complex64, torch.complex128})
    def test_unique_consecutive(self, device, dtype):
        if dtype is torch.half and self.device_type == 'cpu':
            return  # CPU does not have half support

        if dtype is torch.bool:
            x = torch.tensor([True, False, False, False, True, True, False, False, False], dtype=torch.bool, device=device)
            expected_unique = torch.tensor([True, False, True, False], dtype=torch.bool, device=device)
            expected_inverse = torch.tensor([0, 1, 1, 1, 2, 2, 3, 3, 3], dtype=torch.long, device=device)
            expected_counts = torch.tensor([1, 3, 2, 3], dtype=torch.long, device=device)
        else:
            x = torch.tensor([1, 2, 2, 2, 5, 5, 2, 2, 3], dtype=dtype, device=device)
            expected_unique = torch.tensor([1, 2, 5, 2, 3], dtype=dtype, device=device)
            expected_inverse = torch.tensor([0, 1, 1, 1, 2, 2, 3, 3, 4], device=device)
            expected_counts = torch.tensor([1, 3, 2, 2, 1], device=device)

        for f in [torch.unique_consecutive, lambda x, **kwargs: x.unique_consecutive(**kwargs)]:
            self._test_unique_with_expects(device, dtype, f, x, expected_unique, expected_inverse, expected_counts, (3, 3))
            self._test_unique_scalar_empty(dtype, device, f)

    @dtypes(torch.double)
    def test_kthvalue(self, device, dtype):
        SIZE = 50
        x = torch.rand(SIZE, SIZE, SIZE, dtype=dtype, device=device)
        x0 = x.clone()

        k = random.randint(1, SIZE)
        res1val, res1ind = torch.kthvalue(x, k, keepdim=False)
        res2val, res2ind = torch.sort(x)

        self.assertEqual(res1val[:, :], res2val[:, :, k - 1], atol=0, rtol=0)
        self.assertEqual(res1ind[:, :], res2ind[:, :, k - 1], atol=0, rtol=0)
        # test use of result tensors
        k = random.randint(1, SIZE)
        res1val = torch.tensor([], dtype=dtype, device=device)
        res1ind = torch.tensor([], dtype=torch.long, device=device)
        torch.kthvalue(x, k, keepdim=False, out=(res1val, res1ind))
        res2val, res2ind = torch.sort(x)
        self.assertEqual(res1val[:, :], res2val[:, :, k - 1], atol=0, rtol=0)
        self.assertEqual(res1ind[:, :], res2ind[:, :, k - 1], atol=0, rtol=0)

        # test non-default dim
        k = random.randint(1, SIZE)
        res1val, res1ind = torch.kthvalue(x, k, 0, keepdim=False)
        res2val, res2ind = torch.sort(x, 0)
        self.assertEqual(res1val, res2val[k - 1], atol=0, rtol=0)
        self.assertEqual(res1ind, res2ind[k - 1], atol=0, rtol=0)

        # non-contiguous
        y = x.narrow(1, 0, 1)
        y0 = y.contiguous()
        k = random.randint(1, SIZE)
        res1val, res1ind = torch.kthvalue(y, k)
        res2val, res2ind = torch.kthvalue(y0, k)
        self.assertEqual(res1val, res2val, atol=0, rtol=0)
        self.assertEqual(res1ind, res2ind, atol=0, rtol=0)

        # non-contiguous [Reference: https://github.com/pytorch/pytorch/issues/45721]
        non_contig_t = torch.tensor([0, -1, 1, -2, 2], dtype=dtype, device=device)[::2]
        expected_val, expected_ind = non_contig_t.contiguous().kthvalue(2)
        non_contig_cpu_t = non_contig_t.cpu()
        expected_val_cpu, expected_ind_cpu = non_contig_cpu_t.kthvalue(2)

        out_val, out_ind = non_contig_t.kthvalue(2)
        self.assertEqual(expected_val, out_val, atol=0, rtol=0)
        self.assertEqual(expected_ind, out_ind, atol=0, rtol=0)
        self.assertEqual(expected_val_cpu, out_val, atol=0, rtol=0)
        self.assertEqual(expected_ind_cpu, out_ind, atol=0, rtol=0)

        # check that the input wasn't modified
        self.assertEqual(x, x0, atol=0, rtol=0)

        # simple test case (with repetitions)
        y = torch.tensor((3., 5, 4, 1, 1, 5), dtype=dtype, device=device)
        self.assertEqual(torch.kthvalue(y, 3)[0], 3, atol=0, rtol=0)
        self.assertEqual(torch.kthvalue(y, 2)[0], 1, atol=0, rtol=0)

        # simple test case (with NaN)
        SIZE = 50
        x = torch.rand(SIZE, SIZE, SIZE, dtype=dtype, device=device)
        x[torch.arange(SIZE), :, torch.randint(50, (50,))] = nan
        ks = [random.randint(1, SIZE), 1, SIZE, SIZE - 1]
        res2val, res2ind = torch.sort(x)
        for k in ks:
            res1val, res1ind = torch.kthvalue(x, k, keepdim=False)
            self.assertEqual(res1val[:, :], res2val[:, :, k - 1], atol=0, rtol=0)
            self.assertEqual(res1ind[:, :], res2ind[:, :, k - 1], atol=0, rtol=0)

    # test overlapping output
    @dtypes(torch.double)
    @onlyOnCPUAndCUDA   # Fails on XLA
    def test_kthvalue_overlap(self, device, dtype):
        S = 10
        k = 5
        a = torch.randn(S)
        indices = torch.empty((), device=device, dtype=torch.long)
        with self.assertRaisesRegex(RuntimeError, "unsupported operation:"):
            torch.kthvalue(a, k, out=(a, indices))

    @dtypes(torch.float)
    @onlyOnCPUAndCUDA   # Fails on XLA
    def test_kthvalue_scalar(self, device, dtype):
        # Test scalar input (test case from https://github.com/pytorch/pytorch/issues/30818)
        # Tests that passing a scalar tensor or 1D tensor with 1 element work either way
        res = torch.tensor(2, device=device, dtype=dtype).kthvalue(1)
        ref = torch.tensor([2], device=device, dtype=dtype).kthvalue(1)
        self.assertEqual(res[0], ref[0].squeeze())
        self.assertEqual(res[1], ref[1].squeeze())

instantiate_device_type_tests(TestSortAndSelect, globals())

if __name__ == '__main__':
    run_tests()


from collections import Counter, namedtuple, defaultdict
import csv
from datetime import datetime

fields = ("MSISDN","Network","Date","Product","Amount")

class LoansRecord(namedtuple('LoansRecord_', fields)):

    @classmethod
    def parse(klass, row):
        row = list(row)
        row[2] = datetime.strptime(row[2], "%d-%b-%y").strftime('%b-%y') # Parse the "Date" to month-year format
        row[4] = int(row[4]) # Convert "Amount" to an integer
        return klass(*row)

class LoansReader(object):

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

        self._length  = None
        self._counter = None

    def __iter__(self):
        self._length  = 0
        self._counter = Counter()

        with open(self.path, 'rU') as data:
            data.readline()  # Skips the header
            reader = csv.reader(data)  # Create a regular tuple reader
            for row in map(LoansRecord.parse, reader):

                self._length += 1
                self._counter[row[0]] += 1
                yield row

    def __len__(self):
        if self._length is None:
            for row in self: continue # Read the data for length and counter
        return self._length

    @property
    def counter(self):
        if self._counter is None:
            for row in self: continue # Read the data for length and counter
        return self._counter

    @property
    def msisdns(self):
        return self.counter.keys()

    def reset(self):
        """
        In case of partial seeks (e.g. breaking in the middle of the read)
        """
        self._length  = None
        self._counter = None

    @property
    def aggregate_loans(self):
        """
        Aggregate loans by (Network, Product, Month)
        """
        result = []
        aggregated = defaultdict(list)
        for row in self:
            keys = (row.Network, row.Product, row.Date)
            aggregated[keys].append(row.Amount)
        for key, values in aggregated.items():
            result.append(list(key) + [sum(values), len(values)])
        return result

class CSVOutput:

    @classmethod
    def csv_out(klass, header, rows, path):
        with open(path, 'w') as file:
            writer = csv.writer(file)
            writer.writerow(header)
            for row in rows:
                writer.writerow(row)


import pytesseract
from PIL import Image
import cv2

# Simple image to string

#img = Image.open('test.png')

template_img = cv2.imread('./state-frames/two_life_left_cross.png')
cropped_template = template_img[0:40, 0:150]
#dimensions = img.shape
#cv2.imshow('cropped', cropped)
#cv2.waitKey(0)

img = cv2.imread('./state-frames/one_life_left_test.png')
cropped_image = img[0:40, 0:150]
res = cv2.matchTemplate(cropped_image, cropped_template, cv2.TM_CCORR_NORMED)
print(res)
#cv2.imshow('cropped', res)
#cv2.waitKey(0)
#print(pytesseract.image_to_string(img))

def draw_boxes_and_show(cv_image):
    cv2.rectangle(cv_image, (250, 10), (600, 60), (0, 0, 255), 2)
    cv2.imshow('window',cv_image)
    cv2.waitKey(0)
    return cv_image

#print(pytesseract.image_to_string(draw_boxes_and_show(img)))


'''
Script to facilitate making seeds for solving the full problem. This script finds games which must be zero or one based on
(1) the given constraints and
(2) running the model quickly to determine infesible games
'''
import gurobipy as grb
import time
import pandas as pd

def set_constrained_variables():
    '''
    This function sets the games which must be zero based on the constraints to zero (i.e. zero upper and lower bound)
    '''

    myConstrs = NFLmodel.getConstrs()
    for c in myConstrs:
        SoftlinkingConstr = False
        if c.sense =='<'and c.RHS == 0:
            row = NFLmodel.getRow(c)
            #check if the row contains a linking term or penalty term
            for r in range(row.size()):
                if 'GO_' not in row.getVar(r).varName:
                    SoftlinkingConstr = True
                    break
            # if it doesnt contain one of those terms then set all the variables in that constraint to zero
            if not SoftlinkingConstr:
                for r in range(row.size()):
                    row.getVar(r).lb = 0
                    row.getVar(r).ub = 0
                    NFLmodel.update()
    return

def get_variables():
    '''
    makes two dictionarys keyed over cleaned up tuples of the game variables:
    (1) freevars - are the gurobi objects of the variables which are not yet assigned
    (2) var_status contains a tuple indicating the bounds of the variables
    '''
    free_vars = {}
    var_status = {}
    temp=[]
    myVars = NFLmodel.getVars()
    for v in myVars:
        if 'GO' == v.varName[:2]:
            temp = v.varName.split('_')
            if 'PRIME' in temp:
                free_vars[tuple(temp[1:])] = v
                var_status[tuple(temp[1:])]=(v.lb,v.ub)
    free_vars = cleanfreevars(free_vars,var_status)
    print(len(var_status))
    print(len(free_vars) )
    return free_vars, var_status

def cleanfreevars(free_vars,var_status):
    '''
    this kills any variables in freevars which now have fixed bounds as recorded in var_status
    '''
    for v in var_status:
        if var_status[v][0]==var_status[v][1]:
            if v in free_vars:
                free_vars.pop(v,None)
    return free_vars

def fix_impossible_games(free_vars,var_status):
    start=time.time()
    NFLmodel.setParam('OutputFlag', False )
    NFLmodel.setParam('TimeLimit',10)

    print('Starting...')
    stop = False
    while not stop:
        stop = True
        end = time.time()
        print(str(len(free_vars))+ ' free variables remaining. Runtime = ' + str(round((end-start)/60,2)) + ' minutes.')
        count = 0
        for v in free_vars:
            free_vars[v].lb = 1
            NFLmodel.update()
            NFLmodel.optimize()
            if NFLmodel.Status == grb.GRB.INFEASIBLE:
                var_status[v] = (0,0)
                free_vars[v].lb = 0
                free_vars[v].ub = 0
                NFLmodel.update()
                stop = False
                print(str(v) + ' is fixed to zero ' + str(count)+ '/'+str(len(free_vars)))
            else:
                free_vars[v].lb = 0
                NFLmodel.update()
                print(str(v) + ' is free '+ str(count)+ '/'+str(len(free_vars)))
            NFLmodel.update()
            count+=1
        free_vars = cleanfreevars(free_vars,var_status)
        NFLmodel.write('updated.lp')
        print(str(len(free_vars))+ ' free variables remaining. Runtime =  ', str(round((end-start)/60,2))+ ' minutes.')

if __name__ == "__main__":
    #load model
    NFLmodel=grb.read('OR604 Model File v2.lp')
    #load constraints
    set_constrained_variables()
    free_vars, var_status = get_variables()
    fix_impossible_games(free_vars,var_status)
    write = pd.DataFrame.from_dict(var_status,orient="index")
    write.to_csv("GameBounds.csv")
    NFLmodel.write('updated.lp')


from flask import Flask, session, render_template, request, redirect

from mysqlconnection import connectToMySQL

app = Flask(__name__)
app.secret_key = "ThisKeyIsSecret"

mysql = connectToMySQL('friendsdb')
all_friends = mysql.query_db("SELECT * FROM friends")
print("Fetched all friends", all_friends)

# @app.route('/')
# def index():
#     mysql = connectToMySQL('friendsdb')
#     all_friends = mysql.query_db("SELECT * FROM friends")
#     print("Fetched all friends", all_friends)
#     return render_template('index.html', friends = all_friends)

if __name__ == "__main__":
    app.run(debug=True)


import numpy as np
import os
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import KFold
from sklearn.utils import resample
from sklearn import linear_model
import matplotlib.pyplot as plt
import seaborn as sns

import pickle

def type_of_script():
  try:
    ipy_str = str(type(get_ipython()))
    if 'zmqshell' in ipy_str:
      return 'jupyter'
    if 'terminal' in ipy_str:
      return 'ipython'
  except:
    return 'terminal'

if type_of_script() == 'jupyter':
  from tqdm import tqdm_notebook as tqdm
else:
  from tqdm import tqdm

def readPickle(filepath):
  assert os.path.exists(filepath)
  with open(filepath,'rb') as f:
    data = pickle.load(f)

  nfiles = len(data)

  #views
  localviews = np.zeros((nfiles,3,61))
  globalviews = np.zeros((nfiles,3,201))

  std_depths     = np.zeros((nfiles,2))
  true_depths    = np.zeros((nfiles,6))
  depth_errors   = np.zeros((nfiles,6))
  astronet_score = np.zeros((nfiles,1))
  stellar_radii  = np.zeros((nfiles,1))
  magnitudes     = np.zeros((nfiles, 1))

  # Labels
  labels = np.zeros((nfiles,))

  for i,file in enumerate(data):
    stdcut = (np.std(file['LocalDepths']) < 1) and (np.std(file['GlobalDepths']) < 1)
    non_zero_depths = ((np.sum(file['LocalDepths'] == 0) + np.sum(file['GlobalDepths'] == 0)) == 0)
    no_nans = np.sum(np.isnan(file['LocalView'])) + np.sum(np.isnan(file['GlobalView'])) == 0

    if no_nans and non_zero_depths and stdcut:

      localviews[i] = file['LocalView']
      globalviews[i] = file['GlobalView']
      labels[i] = file['label']

      astronet_score[i] = np.array([file['AstroNetScore']])
      std_depths[i]     = np.array([np.std(file['LocalDepths']),np.std(file['GlobalDepths'])])
      true_depths[i]    = np.hstack([file['LocalDepths'],file['GlobalDepths']])
      depth_errors[i]    = np.hstack([file['LocalDepthErrors'],file['GlobalDepthErrors']])
      stellar_radii[i]  = np.array([file['StellarParams']['rad']])
      magnitudes[i]     = np.array(file['StellarParams']['tmag'])

  return labels, localviews, globalviews, std_depths, true_depths, depth_errors, astronet_score, stellar_radii, magnitudes

def loadSortedData(pickleDir, sectors, verbose=True):
  std_depths = np.zeros((0,2))
  true_depths = np.zeros((0,6))
  depth_errors = np.zeros((0,6))
  astronets = np.zeros((0,1))
  stellar_radii = np.zeros((0,1))
  all_mag    = np.zeros((0,1))

  all_localviews = np.zeros((0,3,61))
  all_globalviews = np.zeros((0,3,201))
  all_labels = np.zeros((0,))

  # append with data from each sector
  for sector in tqdm(sectors):
    filepath = pickleDir + sector + '.pickle'
    labels, localviews, globalviews, std_depth, true_depth, depth_error, astronet_score, stellar_rad, magnitude = readPickle(filepath)

    std_depths    = np.append(std_depths, std_depth,axis=0)
    true_depths   = np.append(true_depths,true_depth,axis=0)
    depth_errors  = np.append(depth_errors, depth_error, axis=0)
    astronets     = np.append(astronets, astronet_score, axis = 0)
    stellar_radii = np.append(stellar_radii, stellar_rad , axis = 0)
    all_mag       = np.append(all_mag, magnitude, axis=0)

    all_localviews = np.append(all_localviews, localviews,axis=0)
    all_globalviews = np.append(all_globalviews, globalviews,axis=0)
    all_labels = np.append(all_labels,labels)

  if verbose:
    print("X_global.shape = ", all_globalviews.shape)
    print("X_local.shape  = ", all_localviews.shape)
    print("X_depths.shape = ", true_depths.shape)
    print("y.shape        = ", all_labels.shape)

  return all_labels, all_localviews, all_globalviews, true_depths, depth_errors, std_depths, astronets, stellar_radii, all_mag

def addToDataMatrix(X, addon):
  n = np.shape(X)[0]
  ncols = np.prod(np.shape(addon)[1:])
  X = np.append(X, addon.reshape(n,ncols), axis=1)
  return X

def genDataMatrix(y, *features):
  n = len(y)
  X = np.zeros((n,0))
  for feat in features:
    X = addToDataMatrix(X, feat)
  return X, y

##########################################

def getPCs(X,y):
  pcs = X[np.where(y==1)]
  npcs = X[np.where(y==0)]
  return pcs, npcs

def magnitudeCut(X, y, magnitudes, cut):
  pcs, npcs = getPCs(X,y)
  npcIndex = np.where(y==0)
  npcMags = magnitudes[npcIndex]

  indicies = np.where(npcMags < cut)[0]
  npcs = npcs[indicies]

  X = np.vstack((pcs,npcs))
  y = np.concatenate((np.ones(len(pcs)), np.zeros(len(npcs))))
  return X,y

def overSamplePos(x,y,random_state=420):
  data = np.hstack((x,y[:,np.newaxis]))

  pos = data[data[:,-1] == 1]
  neg = data[data[:,-1] == 0]

  pos_oversample = resample(pos, replace=True, n_samples=len(neg),
                            random_state=random_state)

  data = np.append(neg, pos_oversample, axis=0)
  x = data[:,:-1]
  y = data[:,-1]
  return x,y

# -- Models
def RFModel(x_train,y_train,x_test,y_test,
            depth, n_jobs, random_state=420):
  print(n_jobs)
  clf = RandomForestClassifier(n_estimators = 100,
                               max_depth=depth,
                               n_jobs=n_jobs,
                               random_state=random_state)
  clf.fit(x_train,y_train)

  predictions = clf.predict(x_test)
  p_planet = clf.predict_proba(x_test)[:,1]
  return p_planet, clf.feature_importances_

def LRModel(x_train,y_train, x_test,y_test,
            random_state=420):
  reg = linear_model.LinearRegression()
  reg.fit(x_train, y_train)

  score = np.dot(x_test, reg.coef_)
  return score, [reg.coef_]
###

#-- Model Running
def testModel(x,y, model, *params, k=5, overSample=True, random_state=420):
  kf = KFold(n_splits=k, shuffle=True, random_state=random_state)
  results = []
  model_info = []
  for i, (train_index, test_index) in enumerate(kf.split(x)):
    print(f'{i+1} / {k}',end=' ')
    x_train, x_test = x[train_index], x[test_index]
    y_train, y_test = y[train_index], y[test_index]

    if overSample:
      x_train, y_train = overSamplePos(x_train, y_train, random_state=random_state)

    scores, info = model(x_train,y_train,x_test,y_test,
                             *params, random_state=random_state)

    results.append(np.array([scores,y_test]))
    model_info.append(info)
  return results, np.array(model_info)

def aggregateScores(modelOutput):
  return np.concatenate(modelOutput,1)

def getConfusion(scores, labels, threshold):
  pos      = (scores > threshold)
  real     = (labels == 1)

  true_neg  = np.sum(np.logical_and(~pos, ~real))
  true_pos  = np.sum(np.logical_and( pos,  real))
  false_neg = np.sum(np.logical_and(~pos,  real))
  false_pos = np.sum(np.logical_and( pos, ~real))

  return np.array([true_neg, true_pos, false_neg, false_pos])

def calcPrecision(confusion):
  true_neg, true_pos, false_neg, false_pos = confusion
  return true_pos / (true_pos + false_pos)

def calcRecall(confusion):
  true_neg, true_pos, false_neg, false_pos = confusion
  return true_pos / (true_pos + false_neg)

def getPRCurve(scores, labels, thresholds):
  precision = []
  recall    = []
  for threshold in thresholds:
    confusion = getConfusion(scores, labels, threshold)
    precision.append(calcPrecision(confusion))
    recall.append(calcRecall(confusion))

  return np.array(recall), np.array(precision)

def interpretModel(y, *data):
  nPos = int(np.sum(y))
  nNeg = int(np.sum(1-y))
  print(f'{nPos} - True Positive, {nNeg} - True Negative')

  print('Feature Importances')
  for each in data:
    info = each[1]
    print(np.median(info,0))

  sns.set()
  plt.figure()
  for i,each in enumerate(data):
    stack = np.hstack(each[0])
    plt.plot(*getPRCurve(*stack,np.arange(0,1,.01)), label=i)

  guessLevel = nPos / (nNeg + nPos)
  plt.plot((0,1),(guessLevel, guessLevel))
  plt.legend()
###


#!/usr/bin/python3

import logging
logger = logging.getLogger()
from . import tenv
from . import apicall
from .error import *

class PerfTest:
    def __init__(self, suite, cases, tenv, plan):
        self.suite = suite
        self.tenv = tenv
        self.cases = cases
        self.distro = tenv.distro
        self.plan = plan
        self.trpairs = []

    def compare1(self, dynapi, case, trpair):
        endpoint = "/api/report/v1/comparison/run"
        pargs = {}
        pargs["suite"] = self.suite
        pargs["case"] = case
        pargs["q_tenv_id"] = self.tenv.get_id(dynapi)
        pargs["r_tenv_id"] = trpair["tenv_id"]
        pargs["q_run_id"] = self.plan.run_id
        pargs["r_run_id"] = trpair["run_id"]
        return dynapi.get(endpoint, pargs, []).json()["conclusion"]

    def update_conclusion(self, dynapi, case, trpair, conclusion):
        endpoint = "/api/report/v1/status/run"
        pargs = {}
        pargs["suite"] = self.suite
        pargs["case"] = case
        pargs["q_tenv_id"] = self.tenv.get_id(dynapi)
        pargs["r_tenv_id"] = trpair["tenv_id"]
        pargs["q_run_id"] = self.plan.run_id
        pargs["r_run_id"] = trpair["run_id"]
        pargs["status"] = conclusion
        return dynapi.post_no_body(endpoint, pargs, [], [])

    def complete(self, dynapi):
        if self.trpairs : return
        self.run_id = self.plan.run_id
        self.plan.references(dynapi)
        for drpair in self.plan.drpairs:
            distro = drpair["distro"]
            run_id = drpair["run_id"]
            try:
                r_tenv = tenv.Tenv(distro, self.tenv.component,
                                   self.tenv.inst, self.tenv.machine,
                                   self.tenv.extra)
                r_tenv_id = r_tenv.get_id(dynapi)
            except HTTPError as e:
                logger.error(str(e))
                continue
            trpair = {"tenv_id":r_tenv_id, "run_id":run_id}
            self.trpairs.append(trpair)

    def compare(self, dynapi):
        self.complete(dynapi)
        for case in self.cases:
            for trpair in self.trpairs:
                try:
                    conclusion = self.compare1(dynapi, case, trpair)
                    self.update_conclusion(dynapi, case, trpair, conclusion)
                except HTTPError as e:
                    logger.error(str(e))