#!/usr/bin/env python
"""
Module containing various utility functions
"""

# IMPORTS #####################################################################

import re

from enum import Enum, IntEnum
from instruments.units import ureg as u

# CONSTANTS ###################################################################

_IDX_REGEX = re.compile(r"([a-zA-Z_][a-zA-Z0-9_]*)\[(-?[0-9]*)\]")

# FUNCTIONS ###################################################################

# pylint: disable=too-many-arguments
def assume_units(value, units):
    """
    If units are not provided for ``value`` (that is, if it is a raw
    `float`), then returns a `~pint.Quantity` with magnitude
    given by ``value`` and units given by ``units``.

    :param value: A value that may or may not be unitful.
    :param units: Units to be assumed for ``value`` if it does not already
        have units.

    :return: A unitful quantity that has either the units of ``value`` or
        ``units``, depending on if ``value`` is unitful.
    :rtype: `Quantity`
    """
    if isinstance(value, u.Quantity):
        return value
    elif isinstance(value, str):
        value = u.Quantity(value)
        if value.dimensionless:
            return u.Quantity(value.magnitude, units)
        return value
    return u.Quantity(value, units)

def setattr_expression(target, name_expr, value):
    """
    Recursively calls getattr/setattr for attribute
    names that are miniature expressions with subscripting.
    For instance, of the form ``a[0].b``.
    """
    # Allow "." in attribute names so that we can set attributes
    # recursively.
    if "." in name_expr:
        # Recursion: We have to strip off a level of getattr.
        head, name_expr = name_expr.split(".", 1)
        match = _IDX_REGEX.match(head)
        if match:
            head_name, head_idx = match.groups()
            target = getattr(target, head_name)[int(head_idx)]
        else:
            target = getattr(target, head)

        setattr_expression(target, name_expr, value)
    else:
        # Base case: We're in the last part of a dot-expression.
        match = _IDX_REGEX.match(name_expr)
        if match:
            name, idx = match.groups()
            getattr(target, name)[int(idx)] = value
        else:
            setattr(target, name_expr, value)

def convert_temperature(temperature, base):
    """
    Obsolete with the transition to Pint from Quantities.

    :param temperature: A quantity with units of Kelvin, Celsius, or Fahrenheit
    :type temperature: `pint.Quantity`
    :param base: A temperature unit to convert to
    :type base: `pint.Quantity`

    :return: The converted temperature
    :rtype: `pint.Quantity`
    """
    newval = assume_units(temperature, u.degC)
    return newval.to(base)

def split_unit_str(s, default_units=u.dimensionless, lookup=None):
    """
    Given a string of the form "12 C" or "14.7 GHz", returns a tuple of the
    numeric part and the unit part, irrespective of how many (if any) whitespace
    characters appear between.

    By design, the tuple should be such that it can be unpacked into
    :func:`u.Quantity`::

        >>> u.Quantity(*split_unit_str("1 s"))
        array(1) * s

    For this reason, the second element of the tuple may be a unit or
    a string, depending, since the quantity constructor takes either.

    :param str s: Input string that will be split up
    :param default_units: If no units are specified, this argument is given
        as the units.
    :param callable lookup: If specified, this function is called on the
        units part of the input string. If `None`, no lookup is performed.
        Lookups are never performed on the default units.
    :rtype: `tuple` of a `float` and a `str` or `u.Quantity`
    """
    if lookup is None:
        lookup = lambda x: x

    # Borrowed from:
    # http://stackoverflow.com/questions/430079/how-to-split-strings-into-text-and-number
    # Reg exp tweaked on May 30, 2015 by scasagrande to match on input with
    # scientific notation. General flow borrowed from:
    # http://www.regular-expressions.info/floatingpoint.html
    regex = r"([-+]?[0-9]*\.?[0-9]+)([eE][-+]?[0-9]+)?\s*([a-z]+)?"
    match = re.match(regex, str(s).strip(), re.I)
    if match:
        if match.groups()[1] is None:
            val, _, units = match.groups()
        else:
            val = float(match.groups()[0]) * 10 ** float(match.groups()[1][1:])
            units = match.groups()[2]

        if units is None:
            return float(val), default_units

        return float(val), lookup(units)

    try:
        return float(s), default_units
    except ValueError:
        raise ValueError(f"Could not split '{repr(s)}' into value and units.")

def rproperty(fget=None, fset=None, doc=None, readonly=False, writeonly=False):
    """
    Creates and returns a new property based on the input parameters.

    :param function fget: Function to be called for the new property's getter
    :param function fset: Function to be called for the new property's setter
    :param str doc: Docstring for the new property
    :param bool readonly: If `True`, the returned property does not have a
        setter.
    :param bool writeonly: If `True`, the returned property does not have a
        getter. Both readonly and writeonly cannot both be `True`.
    """
    if readonly and writeonly:
        raise ValueError("Properties cannot be both read- and write-only.")
    if readonly:
        return property(fget=fget, fset=None, doc=doc)
    elif writeonly:
        return property(fget=None, fset=fset, doc=doc)

    return property(fget=fget, fset=fset, doc=doc)

def bool_property(
    command,
    set_cmd=None,
    inst_true="ON",
    inst_false="OFF",
    doc=None,
    readonly=False,
    writeonly=False,
    set_fmt="{} {}",
):
    """
    Called inside of SCPI classes to instantiate boolean properties
    of the device cleanly.
    For example:

    >>> my_property = bool_property(
    ...     "BEST:PROPERTY",
    ...     inst_true="ON",
    ...     inst_false="OFF"
    ... ) # doctest: +SKIP

    This will result in "BEST:PROPERTY ON" or "BEST:PROPERTY OFF" being sent
    when setting, and "BEST:PROPERTY?" being sent when getting.

    :param str command: Name of the SCPI command corresponding to this property.
        If parameter set_cmd is not specified, then this parameter is also used
        for both getting and setting.
    :param str set_cmd: If not `None`, this parameter sets the command string
        to be used when sending commands with no return values to the
        instrument. This allows for non-symmetric properties that have different
        strings for getting vs setting a property.
    :param str inst_true: String returned and accepted by the instrument for
        `True` values.
    :param str inst_false: String returned and accepted by the instrument for
        `False` values.
    :param str doc: Docstring to be associated with the new property.
    :param bool readonly: If `True`, the returned property does not have a
        setter.
    :param bool writeonly: If `True`, the returned property does not have a
        getter. Both readonly and writeonly cannot both be `True`.
    :param str set_fmt: Specify the string format to use when sending a
        non-query to the instrument. The default is "{} {}" which places a
        space between the SCPI command the associated parameter. By switching
        to "{}={}" an equals sign would instead be used as the separator.
    """

    def _getter(self):
        return self.query(command + "?").strip() == inst_true

    def _setter(self, newval):
        if not isinstance(newval, bool):
            raise TypeError("Bool properties must be specified with a " "boolean value")
        self.sendcmd(
            set_fmt.format(
                command if set_cmd is None else set_cmd,
                inst_true if newval else inst_false,
            )
        )

    return rproperty(
        fget=_getter, fset=_setter, doc=doc, readonly=readonly, writeonly=writeonly
    )

def enum_property(
    command,
    enum,
    set_cmd=None,
    doc=None,
    input_decoration=None,
    output_decoration=None,
    readonly=False,
    writeonly=False,
    set_fmt="{} {}",
):
    """
    Called inside of SCPI classes to instantiate Enum properties
    of the device cleanly.
    The decorations can be functions which modify the incoming and outgoing
    values for dumb instruments that do stuff like include superfluous quotes
    that you might not want in your enum.
    Example:
    my_property = bool_property("BEST:PROPERTY", enum_class)

    :param str command: Name of the SCPI command corresponding to this property.
        If parameter set_cmd is not specified, then this parameter is also used
        for both getting and setting.
    :param str set_cmd: If not `None`, this parameter sets the command string
        to be used when sending commands with no return values to the
        instrument. This allows for non-symmetric properties that have different
        strings for getting vs setting a property.
    :param type enum: Class derived from `Enum` representing valid values.
    :param callable input_decoration: Function called on responses from
        the instrument before passing to user code.
    :param callable output_decoration: Function called on commands to the
        instrument.
    :param str doc: Docstring to be associated with the new property.
    :param bool readonly: If `True`, the returned property does not have a
        setter.
    :param bool writeonly: If `True`, the returned property does not have a
        getter. Both readonly and writeonly cannot both be `True`.
    :param str set_fmt: Specify the string format to use when sending a
        non-query to the instrument. The default is "{} {}" which places a
        space between the SCPI command the associated parameter. By switching
        to "{}={}" an equals sign would instead be used as the separator.
    :param str get_cmd: If not `None`, this parameter sets the command string
        to be used when reading/querying from the instrument. If used, the name
        parameter is still used to set the command for pure-write commands to
        the instrument.
    """

    def _in_decor_fcn(val):
        if input_decoration is None:
            return val
        elif hasattr(input_decoration, "__get__"):
            return input_decoration.__get__(None, object)(val)
        return input_decoration(val)

    def _out_decor_fcn(val):
        if output_decoration is None:
            return val
        elif hasattr(output_decoration, "__get__"):
            return output_decoration.__get__(None, object)(val)
        return output_decoration(val)

    def _getter(self):
        return enum(_in_decor_fcn(self.query(f"{command}?").strip()))

    def _setter(self, newval):
        try:  # First assume newval is Enum.value
            newval = enum[newval]
        except KeyError:  # Check if newval is Enum.name instead
            try:
                newval = enum(newval)
            except ValueError:
                raise ValueError("Enum property new value not in enum.")
        self.sendcmd(
            set_fmt.format(
                command if set_cmd is None else set_cmd,
                _out_decor_fcn(enum(newval).value),
            )
        )

    return rproperty(
        fget=_getter, fset=_setter, doc=doc, readonly=readonly, writeonly=writeonly
    )

def unitless_property(
    command,
    set_cmd=None,
    format_code="{:e}",
    doc=None,
    readonly=False,
    writeonly=False,
    set_fmt="{} {}",
):
    """
    Called inside of SCPI classes to instantiate properties with unitless
    numeric values.

    :param str command: Name of the SCPI command corresponding to this property.
        If parameter set_cmd is not specified, then this parameter is also used
        for both getting and setting.
    :param str set_cmd: If not `None`, this parameter sets the command string
        to be used when sending commands with no return values to the
        instrument. This allows for non-symmetric properties that have different
        strings for getting vs setting a property.
    :param str format_code: Argument to `str.format` used in sending values
        to the instrument.
    :param str doc: Docstring to be associated with the new property.
    :param bool readonly: If `True`, the returned property does not have a
        setter.
    :param bool writeonly: If `True`, the returned property does not have a
        getter. Both readonly and writeonly cannot both be `True`.
    :param str set_fmt: Specify the string format to use when sending a
        non-query to the instrument. The default is "{} {}" which places a
        space between the SCPI command the associated parameter. By switching
        to "{}={}" an equals sign would instead be used as the separator.
    """

    def _getter(self):
        raw = self.query(f"{command}?")
        return float(raw)

    def _setter(self, newval):
        if isinstance(newval, u.Quantity):
            if newval.units == u.dimensionless:
                newval = float(newval.magnitude)
            else:
                raise ValueError
        strval = format_code.format(newval)
        self.sendcmd(set_fmt.format(command if set_cmd is None else set_cmd, strval))

    return rproperty(
        fget=_getter, fset=_setter, doc=doc, readonly=readonly, writeonly=writeonly
    )

def int_property(
    command,
    set_cmd=None,
    format_code="{:d}",
    doc=None,
    readonly=False,
    writeonly=False,
    valid_set=None,
    set_fmt="{} {}",
):
    """
    Called inside of SCPI classes to instantiate properties with unitless
    numeric values.

    :param str command: Name of the SCPI command corresponding to this property.
        If parameter set_cmd is not specified, then this parameter is also used
        for both getting and setting.
    :param str set_cmd: If not `None`, this parameter sets the command string
        to be used when sending commands with no return values to the
        instrument. This allows for non-symmetric properties that have different
        strings for getting vs setting a property.
    :param str format_code: Argument to `str.format` used in sending values
        to the instrument.
    :param str doc: Docstring to be associated with the new property.
    :param bool readonly: If `True`, the returned property does not have a
        setter.
    :param bool writeonly: If `True`, the returned property does not have a
        getter. Both readonly and writeonly cannot both be `True`.
    :param valid_set: Set of valid values for the property, or `None` if all
        `int` values are valid.
    :param str set_fmt: Specify the string format to use when sending a
        non-query to the instrument. The default is "{} {}" which places a
        space between the SCPI command the associated parameter. By switching
        to "{}={}" an equals sign would instead be used as the separator.
    """

    def _getter(self):
        raw = self.query(f"{command}?")
        return int(raw)

    if valid_set is None:

        def _setter(self, newval):
            strval = format_code.format(newval)
            self.sendcmd(
                set_fmt.format(command if set_cmd is None else set_cmd, strval)
            )

    else:

        def _setter(self, newval):
            if newval not in valid_set:
                raise ValueError(
                    "{} is not an allowed value for this property; "
                    "must be one of {}.".format(newval, valid_set)
                )
            strval = format_code.format(newval)
            self.sendcmd(
                set_fmt.format(command if set_cmd is None else set_cmd, strval)
            )

    return rproperty(
        fget=_getter, fset=_setter, doc=doc, readonly=readonly, writeonly=writeonly
    )

def unitful_property(
    command,
    units,
    set_cmd=None,
    format_code="{:e}",
    doc=None,
    input_decoration=None,
    output_decoration=None,
    readonly=False,
    writeonly=False,
    set_fmt="{} {}",
    valid_range=(None, None),
):
    """
    Called inside of SCPI classes to instantiate properties with unitful numeric
    values. This function assumes that the instrument only accepts
    and returns magnitudes without unit annotations, such that all unit
    information is provided by the ``units`` argument. This is not suitable
    for instruments where the units can change dynamically due to front-panel
    interaction or due to remote commands.

    :param str command: Name of the SCPI command corresponding to this property.
        If parameter set_cmd is not specified, then this parameter is also used
        for both getting and setting.
    :param str set_cmd: If not `None`, this parameter sets the command string
        to be used when sending commands with no return values to the
        instrument. This allows for non-symmetric properties that have different
        strings for getting vs setting a property.
    :param units: Units to assume in sending and receiving magnitudes to and
        from the instrument.
    :param str format_code: Argument to `str.format` used in sending the
        magnitude of values to the instrument.
    :param str doc: Docstring to be associated with the new property.
    :param callable input_decoration: Function called on responses from
        the instrument before passing to user code.
    :param callable output_decoration: Function called on commands to the
        instrument.
    :param bool readonly: If `True`, the returned property does not have a
        setter.
    :param bool writeonly: If `True`, the returned property does not have a
        getter. Both readonly and writeonly cannot both be `True`.
    :param str set_fmt: Specify the string format to use when sending a
        non-query to the instrument. The default is "{} {}" which places a
        space between the SCPI command the associated parameter. By switching
        to "{}={}" an equals sign would instead be used as the separator.
    :param valid_range: Tuple containing min & max values when setting
        the property. Index 0 is minimum value, index 1 is maximum value.
        Setting `None` in either disables bounds checking for that end of the
        range. The default of `(None, None)` has no min or max constraints.
        The valid set is inclusive of the values provided.
    :type valid_range: `tuple` or `list` of `int` or `float`
    """

    def _in_decor_fcn(val):
        if input_decoration is None:
            return val
        elif hasattr(input_decoration, "__get__"):
            return input_decoration.__get__(None, object)(val)
        return input_decoration(val)

    def _out_decor_fcn(val):
        if output_decoration is None:
            return val
        elif hasattr(output_decoration, "__get__"):
            return output_decoration.__get__(None, object)(val)
        return output_decoration(val)

    def _getter(self):
        raw = _in_decor_fcn(self.query(f"{command}?"))
        return u.Quantity(*split_unit_str(raw, units)).to(units)

    def _setter(self, newval):
        newval = assume_units(newval, units).to(units)
        min_value, max_value = valid_range
        if min_value is not None:
            if callable(min_value):
                min_value = min_value(self)  # pylint: disable=not-callable
            else:
                min_value = assume_units(min_value, units)
            if newval < min_value:
                raise ValueError(
                    f"Unitful quantity is too low. Got {newval}, "
                    f"minimum value is {min_value}"
                )
        if max_value is not None:
            if callable(max_value):
                max_value = max_value(self)  # pylint: disable=not-callable
            else:
                max_value = assume_units(max_value, units)
            if newval > max_value:
                raise ValueError(
                    f"Unitful quantity is too high. Got {newval}, "
                    f"maximum value is {max_value}"
                )
        # Rescale to the correct unit before printing. This will also
        # catch bad units.
        strval = format_code.format(newval.magnitude)
        self.sendcmd(
            set_fmt.format(
                command if set_cmd is None else set_cmd, _out_decor_fcn(strval)
            )
        )

    return rproperty(
        fget=_getter, fset=_setter, doc=doc, readonly=readonly, writeonly=writeonly
    )

def bounded_unitful_property(
    command,
    units,
    min_fmt_str="{}:MIN?",
    max_fmt_str="{}:MAX?",
    valid_range=("query", "query"),
    **kwargs,
):
    """
    Called inside of SCPI classes to instantiate properties with unitful numeric
    values which have upper and lower bounds. This function in turn calls
    `unitful_property` where all kwargs for this function are passed on to.
    See `unitful_property` documentation for information about additional
    parameters that will be passed on.

    Compared to `unitful_property`, this function will return 3 properties:
    the one created by `unitful_property`, one for the minimum value, and one
    for the maximum value.

    :param str command: Name of the SCPI command corresponding to this property.
        If parameter set_cmd is not specified, then this parameter is also used
        for both getting and setting.
    :param str set_cmd: If not `None`, this parameter sets the command string
        to be used when sending commands with no return values to the
        instrument. This allows for non-symmetric properties that have different
        strings for getting vs setting a property.
    :param units: Units to assume in sending and receiving magnitudes to and
        from the instrument.
    :param str min_fmt_str: Specify the string format to use when sending a
        minimum value query. The default is ``"{}:MIN?"`` which will place
        the property name in before the colon. Eg: ``"MOCK:MIN?"``
    :param str max_fmt_str: Specify the string format to use when sending a
        maximum value query. The default is ``"{}:MAX?"`` which will place
        the property name in before the colon. Eg: ``"MOCK:MAX?"``
    :param valid_range: Tuple containing min & max values when setting
        the property. Index 0 is minimum value, index 1 is maximum value.
        Setting `None` in either disables bounds checking for that end of the
        range. The default of ``("query", "query")`` will query the instrument
        for min and max parameter values. The valid set is inclusive of
        the values provided.
    :type valid_range: `list` or `tuple` of `int`, `float`, `None`, or the
        string ``"query"``.
    :param kwargs: All other keyword arguments are passed onto
        `unitful_property`
    :return: Returns a `tuple` of 3 properties: first is as returned by
        `unitful_property`, second is a property representing the minimum
        value, and third is a property representing the maximum value
    """

    def _min_getter(self):
        if valid_range[0] == "query":
            return u.Quantity(
                *split_unit_str(self.query(min_fmt_str.format(command)), units)
            )

        return assume_units(valid_range[0], units).to(units)

    def _max_getter(self):
        if valid_range[1] == "query":
            return u.Quantity(
                *split_unit_str(self.query(max_fmt_str.format(command)), units)
            )

        return assume_units(valid_range[1], units).to(units)

    new_range = (
        None if valid_range[0] is None else _min_getter,
        None if valid_range[1] is None else _max_getter,
    )

    return (
        unitful_property(command, units, valid_range=new_range, **kwargs),
        property(_min_getter) if valid_range[0] is not None else None,
        property(_max_getter) if valid_range[1] is not None else None,
    )

def string_property(
    command,
    set_cmd=None,
    bookmark_symbol='"',
    doc=None,
    readonly=False,
    writeonly=False,
    set_fmt="{} {}{}{}",
):
    """
    Called inside of SCPI classes to instantiate properties with a string value.

    :param str command: Name of the SCPI command corresponding to this property.
        If parameter set_cmd is not specified, then this parameter is also used
        for both getting and setting.
    :param str set_cmd: If not `None`, this parameter sets the command string
        to be used when sending commands with no return values to the
        instrument. This allows for non-symmetric properties that have different
        strings for getting vs setting a property.
    :param str doc: Docstring to be associated with the new property.
    :param bool readonly: If `True`, the returned property does not have a
        setter.
    :param bool writeonly: If `True`, the returned property does not have a
        getter. Both readonly and writeonly cannot both be `True`.
    :param str set_fmt: Specify the string format to use when sending a
        non-query to the instrument. The default is "{} {}{}{}" which places a
        space between the SCPI command the associated parameter, and places
        the bookmark symbols on either side of the parameter.
    :param str bookmark_symbol: The symbol that will flank both sides of the
        parameter to be sent to the instrument. By default this is ``"``.
    """
    bookmark_length = len(bookmark_symbol)

    def _getter(self):
        string = self.query(f"{command}?")
        string = (
            string[bookmark_length:-bookmark_length] if bookmark_length > 0 else string
        )
        return string

    def _setter(self, newval):
        self.sendcmd(
            set_fmt.format(
                command if set_cmd is None else set_cmd,
                bookmark_symbol,
                newval,
                bookmark_symbol,
            )
        )

    return rproperty(
        fget=_getter, fset=_setter, doc=doc, readonly=readonly, writeonly=writeonly
    )

# CLASSES #####################################################################

class ProxyList:
    """
    This is a special class used to generate lists of objects where the valid
    keys are defined by the `valid_set` init parameter. This allows an
    instrument to have a single property through which all of its various
    identical input/ouput channels can be accessed.

    Search the code base of existing examples of how this is used for plenty
    of different examples.

    :param parent: The "parent" or "owner" of the of the proxy classes. In
        dev work, this is typically ``self``.
    :param proxy_cls: The child class that will be returned when the returned
        object is iterated through. These are usually objects that represent
        an entire channel/sensor/input/output, of which an instrument might
        have more than one but each are individually addressed. An example is
        an oscilloscope channel.
    :param valid_set: The set of valid keys by which the proxy class objects
        are accessed. Typically this is something like `range`, but can be
        any generator, list, enum, etc.
    """

    def __init__(self, parent, proxy_cls, valid_set):
        self._parent = parent
        self._proxy_cls = proxy_cls
        self._valid_set = valid_set

        # FIXME: This only checks the next level up the chain!
        if hasattr(valid_set, "__bases__"):
            self._isenum = (Enum in valid_set.__bases__) or (
                IntEnum in valid_set.__bases__
            )
        else:
            self._isenum = False

    def __iter__(self):
        for idx in self._valid_set:
            yield self._proxy_cls(self._parent, idx)

    def __getitem__(self, idx):
        # If we have an enum, try to normalize by using getitem. This will
        # allow for things like 'x' to be used instead of enum.x.
        if self._isenum:
            try:
                idx = self._valid_set[idx]
            except KeyError:
                try:
                    idx = self._valid_set(idx)
                except ValueError:
                    pass
            if not isinstance(idx, self._valid_set):
                raise IndexError(
                    "Index out of range. Must be " "in {}.".format(self._valid_set)
                )
            idx = idx.value
        else:
            if idx not in self._valid_set:
                raise IndexError(
                    "Index out of range. Must be " "in {}.".format(self._valid_set)
                )
        return self._proxy_cls(self._parent, idx)

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


'''
Student Name: Liyao Jiang
Student ID: 1512445
Section: EB1
Student Name: Xiaolei Zhang
Student ID: 1515335
Section: B1
Cmput 275 Project:
Link Picture Pairs

sidepan.py
contains a class 'Sidepan'
which is simply a side pannel for displaying
time remaining, score, and # of remaining picts
Also some instructions of game will be displayed
'''

import pygame

class Sidepan:
	# class attributes constants
	text_color = (127,255,0)
	circle_color = (47,79,79)
	s_width, s_height = 800, 800
	num_row = 12
	num_col = 12
	pict_width = int(s_width / num_col)
	pict_height = int(s_height/ num_row)

	# on __init__, Sidepan object takes display surface and font
	def __init__(self,_display_surf,font):
		pygame.init()
		self._display_surf = _display_surf
		self.font = font
		self.__on_init = True

	# called when main.py wants to initiated the side panel
	def on_init(self,score,count,time):
		self.update_score(score)
		self.update_count(count)
		self.update_time(time)
		self.instruction()
		self.__on_init = False

	# score is passed in, which is the score of the current game
	def update_score(self,score):
		pygame.draw.circle(self._display_surf,self.circle_color,(900,250),80,0)
		self._display_surf.blit(self.font.render("Score:",True,self.text_color), (850,200))
		self._display_surf.blit(self.font.render(str(score),True,self.text_color), (870,230))
		pygame.display.flip()

	# the time remaining is passed in
	def update_time(self,time):
		pygame.draw.circle(self._display_surf,self.circle_color,(900,450),80,0)
		self._display_surf.blit(self.font.render("Time:",True,self.text_color), (850,400))
		self._display_surf.blit(self.font.render(str(time),True,self.text_color), (870,430))
		pygame.display.flip()

	# the number of remaining picts is passed in
	def update_count(self,count):
		pygame.draw.circle(self._display_surf,self.circle_color,(900,650),80,0)
		self._display_surf.blit(self.font.render("Remain:",True,self.text_color), (850,600))
		self._display_surf.blit(self.font.render(str(count),True,self.text_color), (870,630))
		pygame.display.flip()

	# instruction strings
	def instruction(self):
		self._display_surf.blit(self.font.render("[s] Shuffle:",True,self.text_color), (820,10))
		self._display_surf.blit(self.font.render("  cost 200",True,self.text_color), (820,40))
		self._display_surf.blit(self.font.render("[a] Add 10s:",True,self.text_color), (820,80))
		self._display_surf.blit(self.font.render("  cost 300",True,self.text_color), (820,110))
		pygame.display.flip()

	# redraw the 4 rectangular sides of the background
	def redraw_bg(self):
		image_menu = pygame.image.load("images/menu.jpg")
		image_menu = pygame.transform.scale(image_menu,(self.s_width,self.s_height))
		self._display_surf.blit(image_menu,(0,0),(0,0,800,self.pict_height))
		self._display_surf.blit(image_menu,(0,11*self.pict_height),\
			(0,11*self.pict_height,12*self.pict_width,self.pict_height))
		self._display_surf.blit(image_menu,(0,self.pict_height),\
			(0,self.pict_height,self.pict_width,10*self.pict_height))
		self._display_surf.blit(image_menu,(11*self.pict_width,self.pict_height),\
			(11*self.pict_width,self.pict_height,self.pict_width,10*self.pict_height))
		pygame.display.flip()


import torch

def cumprod_exclusive(tensor: torch.Tensor) -> torch.Tensor:
    r"""Mimick functionality of tf.math.cumprod(..., exclusive=True), as it isn't available in PyTorch.

    Args:
    tensor (torch.Tensor): Tensor whose cumprod (cumulative product, see `torch.cumprod`) along dim=-1
      is to be computed.

    Returns:
    cumprod (torch.Tensor): cumprod of Tensor along dim=-1, mimiciking the functionality of
      tf.math.cumprod(..., exclusive=True) (see `tf.math.cumprod` for details).
    """
    # TESTED
    # Only works for the last dimension (dim=-1)
    dim = -1

    # Compute regular cumprod first (this is equivalent to `tf.math.cumprod(..., exclusive=False)`).
    cumprod = torch.cumprod(tensor, dim)

    # "Roll" the elements along dimension 'dim' by 1 element.
    cumprod = torch.roll(cumprod, 1, dim)

    # Replace the first element by "1" as this is what tf.cumprod(..., exclusive=True) does.
    cumprod[..., 0] = 1.0

    return cumprod

class VolumeRenderer(torch.nn.Module):
    def __init__(
        self,
        train_radiance_field_noise_std=0.0,
        val_radiance_field_noise_std=0.0,
        white_background=False,
        attenuation_threshold=1e-3,
        device='cuda',
        **kwargs
    ):
        """
        Perform volume rendering by computing the output of NeRF/NMF model ('radiance_field')
        Inputs:
            train_radiance_field_noise_std: (0 < float < 1): factor to perturb the model prediction of sigma when training.
            val_radiance_field_noise_std:  (0 < float < 1): factor to perturb the model prediction of sigma when evaluating.
            white_background: (bool):
                if true, By calculating the complement (`1-acc_map`) of the weight accumulation of each light (acc_map),
                if the complement is larger, the light is more likely to pass through the free-space,
                add the value of r, g, and b of each ray to make the value close to 1 (white).
            attenuation_threshold: (float) penetrability threshold, transparency of sampled point along the ray transparency,
                (`T_i`) below this threshold will be regarded as not intersecting with the object (sampled point in free space).
        """
        super(VolumeRenderer, self).__init__()
        self.train_radiance_field_noise_std = train_radiance_field_noise_std
        self.val_radiance_field_noise_std = val_radiance_field_noise_std
        self.attenuation_threshold = attenuation_threshold
        self.use_white_bkg = white_background
        epsilon = torch.tensor([1e10]).to(device)
        epsilon.requires_grad = False
        self.register_buffer("epsilon", epsilon)

    def forward(self, radiance_field, depth_values, ray_directions):
        """
        Perform volume rendering by computing the output of NeRF/NMF model ('radiance_field')
        Inputs:
            radiance_field: (torch.tensor): (ray_count, num_samples, 4), rgb and sigma : output from NeRF MLP
            depth_values:  (torch.tensor): (ray_count, num_samples) : Depths of the sampled positions along the ray
            ray_directions: (torch.tensor): (ray_count, num_samples, 3)
        Outputs:
            rgb_map: [num_rays, 3]. Estimated RGB color of a ray.
            disp_map: [num_rays]. Disparity map. 1 / depth.
            weights: [num_rays, num_samples] Weights assigned to each sampled color.
            acc_map: [num_rays]. Accumulated opacity along each ray.
            depth_map: [num_rays]. Estimated distance to object.
            mask_weights: [num_rays, num_samples]: Record all weights of sampled point whose penetrability is higher than threshold
        """
        # set up noise standard deviation (std) for raidance field
        if self.training:
            radiance_field_noise_std = self.train_radiance_field_noise_std
        else:
            radiance_field_noise_std = self.val_radiance_field_noise_std

        # distance between adjacent sampled points. concat a very small number (epsilon) at last as zero
        deltas = torch.cat(
            (
                depth_values[..., 1:] - depth_values[..., :-1],
                self.epsilon.expand(depth_values[..., :1].shape),
            ),
            dim=-1,
        )  # (ray_count, num_samples)

        # Multiply each distance between adjacent sampled points by the norm (2-norm) of its corresponding direction ray
        # to convert to real world distance (accounts for non-unit directions).
        deltas = deltas * ray_directions[..., None, :].norm(p=2, dim=-1)  # (N_rays, N_samples_)

        # Extract RGB of each sample position along each ray.
        rgb = radiance_field[..., :3]

        # Add noise to model's predictions for density. Can be used to
        # regularize network during training (prevents floater artifacts).
        noise = 0.0
        if radiance_field_noise_std > 0.0:
            noise = (
                torch.randn(radiance_field[..., 3].shape, dtype=radiance_field.dtype, device=radiance_field.device) *
                radiance_field_noise_std
            )  #(ray_count, num_samples)

        sigmas = torch.nn.functional.relu(radiance_field[..., 3] + noise)

        # Predict density of each sample along each ray. Higher values imply
        # higher likelihood of being absorbed at this point.
        alpha = 1.0 - torch.exp(-sigmas * deltas)

        # T_i = cumulative-product(j=1,j=(i-1))(1 - alpha_j)
        T_i = cumprod_exclusive(1.0 - alpha + 1e-10)  #(ray_count, num_samples)

        # weight_i = T_i * alpha_i
        # Compute weight for RGB of each sample along each ray.  A cumprod() is
        # used to express the idea of the ray not having reflected up to this
        # sampled point yet.
        weights = alpha * T_i  #(ray_count, num_samples)

        # Weight for RGB less than threshold will be masked out (for later tree ray integration)
        mask_weights = (T_i > self.attenuation_threshold).float()  #(ray_count, num_samples)

        # Computed weighted color of each sample along each ray.
        rgb_map = weights[..., None] * rgb  #(ray_count, num_samples, 3)

        # Sum over rgb value of each sample points (sum_over R,sum_over G, sum_over B) along the ray
        rgb_map = rgb_map.sum(dim=-2)  #(ray_count, 3)

        # sum over weight of each ray. This value is in [0, 1] up to numerical error.
        acc_map = weights.sum(dim=-1)  #(ray_count, )

        depth_map = (weights * depth_values).sum(dim=-1)  # (ray_count, )

        disp_map = 1.0 / torch.max(1e-10 * torch.ones_like(depth_map), depth_map / acc_map)  # (ray_count, )

        # filter out nan depth value
        disp_map[torch.isnan(disp_map)] = 0

        if self.use_white_bkg:
            rgb_map = rgb_map + (1.0 - acc_map[..., None])

        ret = {
            'rgb_map': rgb_map,
            'depth_map': depth_map,
            'weights': weights,
            'mask_weights': mask_weights,
            'acc_map': acc_map,
            'disp_map': disp_map
        }
        return ret


import time
import torch
import tqdm
import torch.nn.functional as F

from utils.eval_utils import accuracy
from utils.logging import AverageMeter, ProgressMeter

__all__ = ["train", "validate", "modifier"]

# Set threshold for score value to enforce global pruning across network layers
def global_prune_threshold(model, args):
    # Loop over all model parameters and concat flattened score tensors
    # Initialize list for storing parameters with scores
    p_list = []
    # Loop over all model parameters to extract scores
    for n, m in model.named_modules():
      # Only add parameters that have scores as attributes
      if hasattr(m,'scores'):
        # Add flattened clone of scores.abs() to p_list
        p_list.append(m.scores.clone().abs().flatten())

    # Concatenate flattened clones
    z = torch.cat(p_list)
    # Sort z
    z,_ = z.sort()
    # Determine number of elements to prune
    p_idx = int((1-args.prune_rate) * z.numel())
    # Identify prune_threshold for values in bottom
    # (1-args.prune_rate) percent of scores
    prune_threshold = z[p_idx-1]
    #print("prune_threshold = ", prune_threshold)

    # Loop over all model parameters to update prune_threshold
    for n, m in model.named_modules():
      # Only add parameters that have scores as attributes
      if hasattr(m,'scores'):
        # Pass prune_threshold value to model parameters
        m.set_prune_threshold(prune_threshold)

    # Exit function
    return

def train(train_loader, model, criterion, optimizer, epoch, args, writer):
    batch_time = AverageMeter("Time", ":6.3f")
    data_time = AverageMeter("Data", ":6.3f")
    losses = AverageMeter("Loss", ":.3f")
    top1 = AverageMeter("Acc@1", ":6.2f")
    top5 = AverageMeter("Acc@5", ":6.2f")
    progress = ProgressMeter(
        len(train_loader),
        [batch_time, data_time, losses, top1, top5],
        prefix=f"Epoch: [{epoch}]",
    )

    # switch to train mode
    model.train()

    batch_size = train_loader.batch_size
    num_batches = len(train_loader)
    end = time.time()
    for i, (images, target) in tqdm.tqdm(
        enumerate(train_loader), ascii=True, total=len(train_loader)
    ):
        # measure data loading time
        data_time.update(time.time() - end)

        if args.gpu is not None:
          if args.jsd:
            images_all = torch.cat(images, 0).cuda(args.gpu, non_blocking=True)
          else:
            images = images.cuda(args.gpu, non_blocking=True)

          target = target.cuda(args.gpu, non_blocking=True)

        # Write scores and weights to tensorboard at beginning of every other epoch
        if args.histograms:
          if (i % (num_batches * batch_size) == 0) and (epoch % 2 == 0):
            for param_name in model.state_dict():
              #print(param_name)
              # Only write scores for now (not weights and batch norm parameters since the pytorch parms don't actually change)
              #if 'score' not in param_name:
              #if 'score' in param_name or 'weight' in param_name:
                #print(param_name, model.state_dict()[param_name])
                writer.add_histogram(param_name, model.state_dict()[param_name], epoch)

        # Check if global pruning is being used
        if args.conv_type == "GlobalSubnetConv":
          # Set prune_threshold for all layers in model
          global_prune_threshold(model, args)

        # compute loss without Jensen-Shannon divergence
        if args.jsd == False:
          output = model(images)

          loss = criterion(output, target)

          # measure accuracy and record loss
          acc1, acc5 = accuracy(output, target, topk=(1, 5))
          losses.update(loss.item(), images.size(0))
          top1.update(acc1.item(), images.size(0))
          top5.update(acc5.item(), images.size(0))

        else: # compute loss with Jensen-Shannon divergence
          logits_all = model(images_all)
          logits_clean, logits_aug1, logits_aug2 = torch.split(logits_all, images[0].size(0))
          output = logits_clean # This is just used for accuracy function

          # Cross-entropy is only computed on clean images
          loss = criterion(logits_clean, target)

          # Terms for Jensen-Shannon divergence
          p_clean, p_aug1, p_aug2 = F.softmax(logits_clean, dim=1), F.softmax(logits_aug1, dim=1), F.softmax(logits_aug2, dim=1)

          # Clamp mixture distribution to avoid exploding KL divergence
          p_mixture = torch.clamp((p_clean + p_aug1 + p_aug2) / 3., 1e-7, 1).log()
          loss += 12 * (F.kl_div(p_mixture, p_clean, reduction='batchmean') + F.kl_div(p_mixture, p_aug1, reduction='batchmean') + F.kl_div(p_mixture, p_aug2, reduction='batchmean')) / 3.

          # measure accuracy and record loss
          acc1, acc5 = accuracy(logits_clean, target, topk=(1, 5))
          losses.update(loss.item(), images[0].size(0))
          top1.update(acc1.item(), images[0].size(0))
          top5.update(acc5.item(), images[0].size(0))

        # compute gradient and do SGD step
        optimizer.zero_grad()
        #torch.nn.utils.clip_grad_norm_(model.parameters(),1)
        loss.backward()
        # EDITED
        #print(torch.norm(torch.cat([p.grad.view(-1) for p in model.parameters()])))
        if args.grad_clip: torch.nn.utils.clip_grad_value_(model.parameters(),1)
        #print(torch.norm(torch.cat([p.grad.view(-1) for p in model.parameters()])))
        #for param_name in model.state_dict(): print(param_name, str(model.state_dict()[param_name])[:50])
        #torch.nn.utils.clip_grad_norm_(model.parameters(),1)
        # end
        optimizer.step()

        # Clamp updated scores to [-1,1] only when using binarized/quantized activations
        #for param_name in model.state_dict():
        #  if 'score' in param_name:
        #    #print(param_name)
        #    scores = model.state_dict()[param_name]
        #    #scores = torch.clamp(scores,min=-1.0,max=1.0)
        #    scores.clamp_(min=-1.0,max=1.0)

        # measure elapsed time
        batch_time.update(time.time() - end)
        end = time.time()

        #print(model.state_dict()['module.linear.3.scores'].grad)
        #params = list(model.parameters())
        #print(params[1].grad)

        if i % args.print_freq == 0:
            t = (num_batches * epoch + i) * batch_size
            progress.display(i)

            #_, predicted = torch.max(output, 1)
            progress.write_to_tensorboard(writer, prefix="train", global_step=t)

        # Write score gradients to tensorboard at end of every other epoch
        if args.histograms:
            if (i % (num_batches * batch_size) == 0) and (epoch % 2 == 0):
            #if ((i+1) % (num_batches-1) == 0) and (epoch % 2 == 0):
              params = list(model.parameters())
              param_names = list(model.state_dict())
              for j in range(len(params)):
                  if params[j].grad is not None:
                 # if 'score' in param_names[j] or 'weight' in param_names[j]:
                 # if 'score' not in param_name and params[j].grad is not None:
                      #print(param_names[j])
                      #print(params[j].grad)
                      writer.add_histogram(param_names[j] + '.grad', params[j].grad, epoch)
                  else:
                      writer.add_histogram(param_names[j] + '.grad', 0, epoch)
              #for param_name in model.state_dict():
              #  if 'score' in param_name:
              #    writer.add_histogram(param_name + '.grad', model.state_dict()[param_name].grad, epoch)
              #params = list(model.parameters())
              #for j in range(len(params)):
              #  writer.add_histogram('Layer' + str(j) + 'grad', params[j].grad, epoch)

    # Write final scores and weights to tensorboard
    if args.histograms:
        for param_name in model.state_dict():
          #writer.add_histogram(param_name, model.state_dict()[param_name], epoch)
          # Only write scores for now (not weights and batch norm parameters since the pytorch parms don't actually change)
          #if 'score' not in param_name:
          #print(param_name, model.state_dict()[param_name])
              writer.add_histogram(param_name, model.state_dict()[param_name], epoch)

    return top1.avg, top5.avg

def validate(val_loader, model, criterion, args, writer, epoch):
    batch_time = AverageMeter("Time", ":6.3f", write_val=False)
    losses = AverageMeter("Loss", ":.3f", write_val=False)
    top1 = AverageMeter("Acc@1", ":6.2f", write_val=False)
    top5 = AverageMeter("Acc@5", ":6.2f", write_val=False)
    progress = ProgressMeter(
        len(val_loader), [batch_time, losses, top1, top5], prefix="Test: "
    )

    # switch to evaluate mode
    model.eval()

    with torch.no_grad():
        end = time.time()
        for i, (images, target) in tqdm.tqdm(
            enumerate(val_loader), ascii=True, total=len(val_loader)
        ):
            if args.gpu is not None:
                images = images.cuda(args.gpu, non_blocking=True)

                target = target.cuda(args.gpu, non_blocking=True)

            # Check if global pruning is being used
            if args.conv_type == "GlobalSubnetConv":
              # Set prune_threshold for all layers in model
              global_prune_threshold(model, args)

            # compute output
            output = model(images)

            loss = criterion(output, target)

            # measure accuracy and record loss
            acc1, acc5 = accuracy(output, target, topk=(1, 5))
            losses.update(loss.item(), images.size(0))
            top1.update(acc1.item(), images.size(0))
            top5.update(acc5.item(), images.size(0))

            # measure elapsed time
            batch_time.update(time.time() - end)
            end = time.time()

            if i % args.print_freq == 0:
                progress.display(i)
                #_, predicted = torch.max(output, 1)
                #print(predicted,target)

        progress.display(len(val_loader))

        if writer is not None:
            progress.write_to_tensorboard(writer, prefix="test", global_step=epoch)

    return top1.avg, top5.avg

def modifier(args, epoch, model):
    return


from setuptools import setup, find_packages

requirements = [
    'Flask>=1.0.2,<2',
]

setup(
    name='flask-pundit',
    version='1.2.2',
    license='MIT',
    url='https://github.com/anurag90x/flask-pundit',
    author='Anurag Chaudhury',
    author_email='anuragchaudhury@gmail.com',
    description='Simple library to manage resource authorization and scoping',
    packages=find_packages(exclude=['tests']),
    zip_safe=False,
    include_package_data=True,
    platforms='any',
    test_suite='nose.collector',
    install_requires=requirements,
    tests_require=['mock==1.3.0'],
    classifiers=[
        'Framework :: Flask',
        'Programming Language :: Python :: 2.6',
        'Programming Language :: Python :: 2.7',
        'Programming Language :: Python :: 3.7',
        'License :: OSI Approved :: MIT License'
    ]
)


#!/usr/bin/env python3
import sys
import logging
from timeit import default_timer as timer

import util

LOG_LEVEL = logging.INFO

def build_pattern(base, position):
    pattern = []
    for x in base:
        for y in range(position):
            pattern.append(x)
    return pattern

def FFT(input):
    base_pattern = [0, 1, 0, -1]

    output = [None] * len(input)

    for i in range(len(input)):
        pattern = build_pattern(base_pattern, i+1)

        # Offset pattern by 1
        pattern.append(pattern.pop(0))

        value = 0
        debug_string = ""
        for digit in input:

            pattern_value = pattern.pop(0)

            debug_string = debug_string + f"{digit}*{pattern_value} + "

            value += digit * pattern_value

            pattern.append(pattern_value)

        # Take only the ones position
        value = abs(value) % 10
        output[i] = value
        debug_string = debug_string + f" = {value}"
        logging.debug(debug_string)
    return output

def main(input_file):

    with open(input_file) as f:
        while True:
            line = f.readline().strip()
            if not line:
                break
            input = [int(x) for x in line]

    logging.debug("Input Pattern: " + "".join([str(x) for x in input]))

    run_count = 100
    for x in range(run_count):
        input = FFT(input)
        logging.debug(f"After {x+1} phase: " + "".join([str(x) for x in input]))

    print("First 8 sequence: " + "".join([str(x) for x in input])[:8])
if __name__ == "__main__":

    logging.basicConfig(format='%(levelname)s:%(module)s: %(message)s', level=LOG_LEVEL)

    if len(sys.argv) < 2:
        logging.critical("Missing input file")
    else:
        start = timer()
        main(sys.argv[1])
        end = timer()

        logging.info(f'Executed in {(end-start):.5f} seconds')


from tqdm import tqdm
import os
import sys
import numpy as np
import torch
import torchvision
import torchvision.transforms as transforms
from torch.autograd import Variable
import torch.utils.data

torch.manual_seed(0)

# Uncomment these lines if you are having problem downloading PyTorch datasets:
# import ssl
# ssl._create_default_https_context = ssl._create_unverified_context
from pathlib import Path

ROOT = str(os.path.dirname(os.path.dirname(__file__)))

sys.path.append("{}/datasets/".format(ROOT))
sys.path.append("{}/pre_trained/".format(ROOT))
sys.path.append("{}/models/".format(ROOT))

from utils.logger import logger
from models.densenet import DenseNetBC100
from models.resnet import ResNet34

# Train set statistics
transform_dict = dict()
transform_dict["CIFAR10"] = (
    (0.4913997, 0.4821584, 0.446531),
    (0.2470323, 0.243485, 0.2615876),
)
transform_dict["CIFAR100"] = (
    (0.5070752, 0.486549, 0.4409178),
    (0.2673342, 0.2564384, 0.2761506),
)
transform_dict["SVHN"] = (
    (0.4379793, 0.4439904, 0.4729508),
    (0.1981116, 0.2011045, 0.1970895),
)

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

def transform_statistics(dataset_name):
    dataset_name = dataset_name.upper()
    return transforms.Compose(
        [
            transforms.ToTensor(),
            transforms.Normalize(
                transform_dict[dataset_name][0], transform_dict[dataset_name][1]
            ),
        ]
    )

def transform_statistics_tensor(dataset_name):
    dataset_name = dataset_name.upper()
    return transforms.Normalize(
        transform_dict[dataset_name][0], transform_dict[dataset_name][1]
    )

def transform_tensor(*args, **kwargs):
    return transforms.ToTensor()

def transform_none(*args, **kwargs):
    return transforms.Lambda(lambda x: x)

def load_train_dataset(name, transform_name, transform=transform_statistics):
    transform_name = transform_name.upper()
    if transform is None or transform_name is None:
        transform = transform_tensor

    if name.upper() == "CIFAR10":
        return torchvision.datasets.CIFAR10(
            root="{}/datasets".format(ROOT),
            train=True,
            download=True,
            transform=transform(transform_name),
        )
    elif name.upper() == "CIFAR100":
        return torchvision.datasets.CIFAR100(
            root="{}/datasets".format(ROOT),
            train=True,
            download=True,
            transform=transform(transform_name),
        )
    elif name.upper() == "SVHN":
        return torchvision.datasets.SVHN(
            root="{}/datasets/SVHN".format(ROOT),
            split="train",
            download=True,
            transform=transform(transform_name),
        )
    else:
        return torchvision.datasets.ImageFolder(
            "{}/datasets/{}".format(ROOT, name),
            transform=transform(transform_name),
        )

def train_dataloader(
    name, transform_name="CIFAR10", transform=transform_statistics, *args, **kwargs
):
    trainset = load_train_dataset(name, transform_name, transform)
    logger.info("dataset {} found. Preparing DataLoader".format(name))
    batch_size = kwargs.get("batch_size", 1)
    trainloader = torch.utils.data.DataLoader(
        trainset, shuffle=False, num_workers=2, batch_size=batch_size
    )
    logger.info("dataset {} loaded with batch size {}".format(name, batch_size))
    return trainloader

class CustomTensorDataset(torch.utils.data.Dataset):
    """TensorDataset with support of transforms."""

    def __init__(self, tensors, transform=None):
        assert all(tensors[0].size(0) == tensor.size(0) for tensor in tensors)
        self.tensors = tensors
        self.transform = transform

    def __getitem__(self, index):
        x = self.tensors[0][index]

        if self.transform:
            x = self.transform(x)

        y = int(self.tensors[1][index])
        return x, y

    def __len__(self):
        return self.tensors[0].size(0)

def load_test_dataset(name, transform_dataset, transform=transform_statistics):
    if transform is None or transform_dataset is None:
        transform = transform_tensor
    elif name.upper() == "PLACES365":
        # original size: 3x256x256
        # original length 36500
        pre_transform = torchvision.transforms.Resize(32)

        return LimitDataset(
            torchvision.datasets.ImageFolder(
                "{}/datasets/Places365/".format(
                    ROOT,
                ),
                transform=torchvision.transforms.Compose(
                    [pre_transform, transform(transform_dataset)]
                ),
            ),
            10000,
        )
    elif name.upper() == "TEXTURES":
        # original size: 3x300x300
        # length 5640
        pre_transform = torchvision.transforms.Compose(
            [
                torchvision.transforms.Resize(32),
                torchvision.transforms.CenterCrop(32),
            ]
        )
        return torchvision.datasets.ImageFolder(
            "{}/datasets/Textures/dtd/images".format(
                ROOT,
            ),
            transform=torchvision.transforms.Compose(
                [
                    pre_transform,
                    transform(transform_dataset),
                ]
            ),
        )
    elif name.upper() == "CHARS74K":
        # original size: 3x42x34
        # length 7705
        pre_transform = torchvision.transforms.Compose(
            [
                torchvision.transforms.Resize(32),
                torchvision.transforms.CenterCrop(32),
            ]
        )
        return torchvision.datasets.ImageFolder(
            "{}/datasets/Chars74K/English/Img/GoodImg/Bmp/".format(
                ROOT,
            ),
            transform=torchvision.transforms.Compose(
                [
                    pre_transform,
                    transform(transform_dataset),
                ]
            ),
        )
    elif name.upper() == "CIFAR100":
        return torchvision.datasets.CIFAR100(
            root="{}/datasets".format(ROOT),
            train=False,
            download=True,
            transform=transform(transform_dataset),
        )
    elif name.upper() == "CIFAR10":
        return torchvision.datasets.CIFAR10(
            root="{}/datasets".format(ROOT),
            train=False,
            download=True,
            transform=transform(transform_dataset),
        )
    elif name.upper() == "SVHN":
        return torchvision.datasets.SVHN(
            root="{}/datasets/SVHN".format(ROOT),
            split="test",
            download=True,
            transform=transform(transform_dataset),
        )
    elif name.upper() == "GAUSSIAN_NOISE_DATASET":
        nb_samples = 10000
        return CustomTensorDataset(
            [
                torch.clamp(torch.randn(nb_samples, 3, 32, 32) + 0.5, 0, 1),
                torch.tensor([1.0] * nb_samples),
            ],
            transform=transform_statistics_tensor(transform_dataset),
        )
    elif name.upper() == "DENSENET10_ADV":
        return load_adv_dataset("densenet10")

    elif name.upper() == "RESNET_CIFAR10_ADV":
        return load_adv_dataset("resnet_cifar10")

    elif name.upper() == "DENSENET100_ADV":
        return load_adv_dataset("densenet100")

    elif name.upper() == "RESNET_CIFAR100_ADV":
        return load_adv_dataset("resnet_cifar100")

    elif name.upper() == "DENSENET_SVHN_ADV":
        return load_adv_dataset("densenet_svhn")

    elif name.upper() == "RESNET_SVHN_ADV":
        return load_adv_dataset("resnet_svhn")
    else:
        return torchvision.datasets.ImageFolder(
            "{}/datasets/{}".format(ROOT, name),
            transform=transform(transform_dataset),
        )

def dataset_channel_statistics(dataloader: torch.utils.data.DataLoader, decimal=4):
    """Calculate a dataset's empirical mean and standard deviation and round to `decimal` points"""
    channels_sum, channels_sq_sum, total_batches = 0, 0, 0
    for (data, _) in tqdm(dataloader):
        channels_sum += torch.mean(data, dim=[0, 2, 3])
        channels_sq_sum += torch.mean(data ** 2, dim=[0, 2, 3])
        total_batches += 1

    logger.debug(f"total samples {total_batches}")
    mean = channels_sum / total_batches
    meansq = channels_sq_sum / total_batches
    std = (meansq - mean ** 2) ** 0.5
    return list(np.around(mean.numpy(), decimal)), list(np.around(std.numpy(), decimal))

def test_dataloader(
    name, transform_name="CIFAR10", transform=transform_statistics, *args, **kwargs
):
    shuffle = False
    if name == "densenet10_adv":
        testset = load_adv_dataset("densenet10", *args, **kwargs)
    elif name == "densenet100_adv":
        testset = load_adv_dataset("densenet100", *args, **kwargs)
    elif name == "densenet_svhn_adv":
        testset = load_adv_dataset("densenet_svhn", *args, **kwargs)
    elif name == "resnet_cifar10_adv":
        testset = load_adv_dataset("resnet_cifar10", *args, **kwargs)
    elif name == "resnet_cifar100_adv":
        testset = load_adv_dataset("resnet_cifar100", *args, **kwargs)
    elif name == "resnet_svhn_adv":
        testset = load_adv_dataset("resnet_svhn", *args, **kwargs)
    elif "densenet" in name or "resnet" in name:
        testset = load_adv_dataset(name.split("ADV")[-1], *args, **kwargs)
    else:
        testset = load_test_dataset(name, transform_name, transform)

    batch_size = kwargs.get("batch_size", 1)
    testloader = torch.utils.data.DataLoader(
        testset, shuffle=shuffle, num_workers=0, batch_size=batch_size
    )
    logger.info("dataset {} loaded with batch size {}".format(name, batch_size))
    return testloader

def load_adv_dataset(nn_name, *args, **kwargs):
    in_dataset_name = get_in_dataset_name(nn_name)
    filename = "{}/tensors/{}/{}/adv_data_fgsm.pth".format(
        ROOT, nn_name, in_dataset_name
    )
    adv_data = load_tensor(filename)
    filename = "{}/tensors/{}/{}/label_fgsm.pth".format(ROOT, nn_name, in_dataset_name)
    adv_labels = load_tensor(filename)

    if adv_data is not None and adv_labels is not None:
        return torch.utils.data.TensorDataset(
            adv_data[:10000].cpu(), adv_labels[:10000].cpu()
        )

def gradient_trasform(in_dataset_name):
    """Transforms the standard deviation of the gradient tensor to adapt to the in-distribution data.

    Args:
        in_dataset_name (str): name of the in-distribution dataset.

    Returns:
        transforms.Normalize: callable torch based normalizer.
    """
    n_channels = get_number_channels(in_dataset_name)
    return transforms.Normalize(
        np.zeros(n_channels), transform_dict[in_dataset_name.upper()][1]
    )

def load_pre_trained_nn(nn_name, gpu=None):
    if type(gpu) == int:
        map_location = torch.device("cuda:{}".format(gpu))
    elif gpu is None:
        map_location = torch.device("cpu")
    else:
        map_location = gpu
    num_c = get_num_classes(get_in_dataset_name(nn_name))
    if "densenet" in nn_name:
        model = DenseNetBC100(num_c)
        if "svhn" not in nn_name and "cifar" not in nn_name:
            return load_nn(nn_name, map_location)
        return load_nn_from_state_dict(nn_name, model, map_location)
    elif "resnet" in nn_name:
        model = ResNet34(num_c)
        return load_nn_from_state_dict(nn_name, model, map_location)
    return load_nn(nn_name, map_location)

def load_nn(nn_name, map_location=torch.device("cpu")):
    logger.info("model {}/pre_trained/{}.pth loaded".format(ROOT, nn_name))
    model = torch.load("{}/pre_trained/{}.pth".format(ROOT, nn_name), map_location)
    return model

def load_nn_from_state_dict(nn_name, model, map_location=torch.device("cpu")):
    model.load_state_dict(load_nn(nn_name), strict=False)
    model.to(map_location)
    model.eval()
    return model

def save_model(model, PATH):
    torch.save(model.module.state_dict(), PATH)

def get_in_dataset_name(nn_name):
    if "1" not in nn_name:
        return "SVHN"
    l = nn_name.split("1")
    if len(l[1]) == 1:
        return "CIFAR10"
    else:
        return "CIFAR100"

def get_nn_name(architecture, in_dataset_name):
    if architecture.lower() == "densenet":
        if "CIFAR10" == in_dataset_name.upper():
            nn_name = "densenet10"
        elif "CIFAR100" == in_dataset_name.upper():
            nn_name = "densenet100"
        else:
            nn_name = "densenet_svhn"
    elif architecture.lower() == "resnet":
        nn_name = "{}_{}".format(architecture.lower(), in_dataset_name.lower())
    else:
        return
    return nn_name

def get_number_channels(dataset_name):
    if "MNIST" in dataset_name.upper():
        return 1
    return 3

def get_num_classes(dataset_name):
    dataset_name = dataset_name.upper()
    if "MNIST" in dataset_name or "SVHN" in dataset_name or "CIFAR10" == dataset_name:
        return 10
    else:
        return 100

def save_tensor(x, filepath):
    torch.save(x, "{}/tensors/{}".format(ROOT, filepath))

def load_tensor(filename):
    if os.path.isfile(filename):
        return torch.load(filename, map_location=DEVICE)
    logger.warning("file {} not found, returning None".format(filename))
    return None

def pred_loop(model, dataloader, gpu, *args, **kwargs):
    logits = []
    targets = []
    with torch.no_grad():
        for (data, target) in tqdm(dataloader):
            if gpu is not None:
                data = data.to(gpu)
            pred = model(data, *args, **kwargs)
            logits.append(pred.detach().cpu())
            targets.append(target.detach().cpu().reshape(-1, 1))

    logits = torch.vstack(logits)
    targets = torch.vstack(targets).reshape(-1)
    return logits, targets

def hidden_features_pred_loop(model, dataloader, gpu, *args, **kwargs):
    features = {}
    targets = []
    with torch.no_grad():
        for idx, (data, target) in enumerate(tqdm(dataloader)):
            if gpu is not None:
                data = data.to(gpu)
            logits, out_features = model.feature_list(data)
            for i, out_feature in enumerate(out_features):
                if idx == 0:
                    features[i] = []
                out_feature = out_feature.reshape(
                    out_feature.shape[0], out_feature.shape[1], -1
                )
                out_feature = torch.mean(out_feature, 2)
                features[i].append(out_feature)
            targets.append(target.detach().cpu().reshape(-1, 1))

    for k in features:
        features[k] = torch.vstack(features[k])
    targets = torch.vstack(targets).reshape(-1)
    return features, targets

def get_and_save_test_logits(nn_name, dataset_name, batch_size, gpu, *args, **kwargs):
    os.makedirs(
        "{}/tensors/{}/{}".format(ROOT, nn_name, dataset_name),
        exist_ok=True,
    )
    model = load_pre_trained_nn(nn_name, gpu)
    in_dataset_name = get_in_dataset_name(nn_name)
    dataloader = test_dataloader(dataset_name, in_dataset_name, batch_size=batch_size)
    logits, targets = pred_loop(model, dataloader, gpu, *args, **kwargs)
    torch.save(
        logits,
        "{}/tensors/{}/{}/logits_test.pt".format(ROOT, nn_name, dataset_name),
    )
    torch.save(
        targets,
        "{}/tensors/{}/{}/targets_test.pt".format(ROOT, nn_name, dataset_name),
    )
    logits = load_test_logits(nn_name, dataset_name)
    return logits

def get_and_save_test_features(nn_name, dataset_name, batch_size, gpu, *args, **kwargs):
    os.makedirs(
        "{}/tensors/{}/{}".format(ROOT, nn_name, dataset_name),
        exist_ok=True,
    )
    model = load_pre_trained_nn(nn_name, gpu)
    in_dataset_name = get_in_dataset_name(nn_name)
    dataloader = test_dataloader(dataset_name, in_dataset_name, batch_size=batch_size)
    features, targets = hidden_features_pred_loop(
        model, dataloader, gpu, *args, **kwargs
    )
    torch.save(
        features,
        "{}/tensors/{}/{}/features_test.pt".format(ROOT, nn_name, dataset_name),
    )
    features = load_test_features(nn_name, dataset_name)
    return features

def load_test_logits(nn_name, dataset_name):
    filename = "{}/tensors/{}/{}/logits_test.pt".format(ROOT, nn_name, dataset_name)
    return load_tensor(filename)

def load_test_features(nn_name, dataset_name):
    filename = "{}/tensors/{}/{}/features_test.pt".format(ROOT, nn_name, dataset_name)
    return load_tensor(filename)

def load_test_targets(nn_name, dataset_name):
    filename = "{}/tensors/{}/{}/targets_test.pt".format(ROOT, nn_name, dataset_name)
    return load_tensor(filename)

def load_test_penultimate_features(nn_name, dataset_name):
    filename = "{}/tensors/{}/{}/penultimate_features_test.pt".format(
        ROOT, nn_name, dataset_name
    )
    return load_tensor(filename)

def get_and_save_train_logits(nn_name, dataset_name, batch_size, gpu, *args, **kwargs):
    os.makedirs(
        "{}/tensors/{}/{}".format(ROOT, nn_name, dataset_name),
        exist_ok=True,
    )
    model = load_pre_trained_nn(nn_name, gpu)
    in_dataset_name = get_in_dataset_name(nn_name)
    dataloader = train_dataloader(
        dataset_name,
        in_dataset_name,
        batch_size=batch_size,
    )
    logits, targets = pred_loop(model, dataloader, gpu, *args, **kwargs)
    torch.save(
        logits,
        "{}/tensors/{}/{}/logits_train.pt".format(ROOT, nn_name, dataset_name),
    )
    torch.save(
        targets,
        "{}/tensors/{}/{}/targets_train.pt".format(ROOT, nn_name, dataset_name),
    )
    logits = load_train_logits(nn_name, dataset_name)
    return logits

def get_and_save_train_features(
    nn_name, dataset_name, batch_size, gpu, *args, **kwargs
):
    os.makedirs(
        "{}/tensors/{}/{}".format(ROOT, nn_name, dataset_name),
        exist_ok=True,
    )
    model = load_pre_trained_nn(nn_name, gpu)
    in_dataset_name = get_in_dataset_name(nn_name)
    dataloader = train_dataloader(
        dataset_name,
        in_dataset_name,
        batch_size=batch_size,
    )
    features, targets = hidden_features_pred_loop(
        model, dataloader, gpu, *args, **kwargs
    )
    torch.save(
        features,
        "{}/tensors/{}/{}/features_train.pt".format(ROOT, nn_name, dataset_name),
    )
    features = load_train_features(nn_name, dataset_name)
    return features

def load_train_logits(nn_name, dataset_name):
    filename = "{}/tensors/{}/{}/logits_train.pt".format(ROOT, nn_name, dataset_name)
    return load_tensor(filename)

def load_train_features(nn_name, dataset_name):
    filename = "{}/tensors/{}/{}/features_train.pt".format(ROOT, nn_name, dataset_name)
    return load_tensor(filename)

def load_train_targets(nn_name, dataset_name):
    filename = "{}/tensors/{}/{}/targets_train.pt".format(ROOT, nn_name, dataset_name)
    return load_tensor(filename)

def load_logits_centroid(nn_name, dataset_name, method_name=None, new=False, cap=1000):
    # check file existence
    filename = "{}/tensors/centroid_logits_{}_{}.pt".format(ROOT, nn_name, dataset_name)
    if new:
        filename = "{}/tensors/{}/{}/{}_centroid_logits.pt".format(
            ROOT, nn_name, dataset_name, method_name
        )
    if os.path.isfile(filename):
        centroid = torch.load(filename, map_location=DEVICE)
    else:
        logger.warning("file not found. Returning None")
        return None
    return centroid.detach()

def load_test_logits_centroid(nn_name, dataset_name, cap=1000):
    # check file existence
    filename = "{}/tensors/{}/{}/ige_centroid_logits_{}.pt".format(
        ROOT, nn_name, dataset_name, cap
    )
    if os.path.isfile(filename):
        centroid = torch.load(filename, map_location=DEVICE)
    else:
        logger.warning("file not found. Returning None")
        return None
    return centroid.detach()

def get_feature_list(model, gpu):
    if gpu is not None:
        temp_x = torch.rand(2, 3, 32, 32).cuda()
    else:
        temp_x = torch.rand(2, 3, 32, 32)
    temp_x = Variable(temp_x)
    temp_list = model.feature_list(temp_x)[1]
    num_output = len(temp_list)
    feature_list = np.empty(num_output)
    for count, out in enumerate(temp_list):
        feature_list[count] = out.size(1)
    return feature_list

def load_hidden_features_cov(
    nn_name, dataset_name, diag=False, cap=None, per_class=False
):
    mat_type = ""
    if diag:
        mat_type = "_diag"
    cap_name = "_{}".format(cap) if cap is not None else ""
    if per_class:
        cap_name += "_per_class"

    filename = "{}/tensors/{}/{}/hidden_features{}_cov_mat{}.pt".format(
        ROOT, nn_name, dataset_name, mat_type, cap_name
    )
    if os.path.isfile(filename):
        logger.info("loading cov tensors from file {}".format(filename.split("/")[-1]))
        return torch.load(filename, map_location=DEVICE)
    else:
        logger.warn("file {} not found, returning".format(filename.split("/")[-1]))
        return

def load_hidden_features_inv(
    nn_name, dataset_name, diag=False, cap=None, per_class=False
):
    cap_name = "_{}".format(cap) if cap is not None else ""
    if per_class:
        cap_name += "_per_class"
    mat_type = ""
    if diag:
        mat_type = "_diag"

    filename = "{}/tensors/{}/{}/hidden_features{}_invs_cov{}.pt".format(
        ROOT, nn_name, dataset_name, mat_type, cap_name
    )
    if os.path.isfile(filename):
        logger.info("loading inv tensors from file {}".format(filename.split("/")[-1]))
        return torch.load(filename, map_location=DEVICE)
    else:
        logger.warn("file {} not found, returning".format(filename.split("/")[-1]))
        return

def load_hidden_features_means(nn_name, dataset_name, cap=None):
    cap_name = "_{}".format(cap) if cap is not None else ""
    filename = "{}/tensors/{}/{}/hidden_features_means{}.pt".format(
        ROOT, nn_name, dataset_name, cap_name
    )
    if os.path.isfile(filename):
        logger.info(
            "loading means tensors from file {}".format(filename.split("/")[-1])
        )
        return torch.load(filename, map_location=DEVICE)
    else:
        logger.warn("file {} not found, returning".format(filename.split("/")[-1]))
        return

class LimitDataset(torch.utils.data.Dataset):
    def __init__(self, dataset, nb_samples):
        self.dataset = dataset
        self.nb_samples = min(nb_samples, len(dataset))

    def __getitem__(self, i):
        return self.dataset[i]

    def __len__(self):
        return self.nb_samples