<commit_before>"""Module with main classes related to Interfaces."""
import json
import logging
from enum import IntEnum

from pyof.v0x01.common.phy_port import Port as PortNo01
from pyof.v0x01.common.phy_port import PortFeatures as PortFeatures01
from pyof.v0x04.common.port import PortFeatures as PortFeatures04
from pyof.v0x04.common.port import PortNo as PortNo04

from kytos.core.common import GenericEntity
from kytos.core.helpers import now

__all__ = ('Interface',)

LOG = logging.getLogger(__name__)

class TAGType(IntEnum):
    """Class that represents a TAG Type."""

    VLAN = 1
    VLAN_QINQ = 2
    MPLS = 3

class TAG:
    """Class that represents a TAG."""

    def __init__(self, tag_type, value):
        self.tag_type = tag_type
        self.value = value

    def __eq__(self, other):
        if not other:
            return False
        return self.tag_type == other.tag_type and self.value == other.value

    def as_dict(self):
        """Return a dictionary representating a tag object."""
        return {'tag_type': self.tag_type, 'value': self.value}

    @classmethod
    def from_dict(cls, tag_dict):
        """Return a TAG instance from python dictionary."""
        return cls(tag_dict.get('tag_type'), tag_dict.get('value'))

    @classmethod
    def from_json(cls, tag_json):
        """Return a TAG instance from json."""
        return cls.from_dict(json.loads(tag_json))

    def as_json(self):
        """Return a json representating a tag object."""
        return json.dumps(self.as_dict())

    def __repr__(self):
        return f"TAG({self.tag_type!r}, {self.value!r})"

class Interface(GenericEntity):  # pylint: disable=too-many-instance-attributes
    """Interface Class used to abstract the network interfaces."""

    # pylint: disable=too-many-arguments
    def __init__(self, name, port_number, switch, address=None, state=None,
                 features=None, speed=None):
        """Assign the parameters to instance attributes.

        Args:
            name (string): name from this interface.
            port_number (int): port number from this interface.
            switch (:class:`~.core.switch.Switch`): Switch with this interface.
            address (|hw_address|): Port address from this interface.
            state (|port_stats|): Port Stat from interface. It will be
            deprecated.
            features (|port_features|): Port feature used to calculate link
                utilization from this interface. It will be deprecated.
            speed (int, float): Interface speed in bytes per second. Defaults
                to what is informed by the switch. Return ``None`` if not set
                and switch does not inform the speed.
        """
        self.name = name
        self.port_number = int(port_number)
        self.switch = switch
        self.address = address
        self.state = state
        self.features = features
        self.nni = False
        self.endpoints = []
        self.stats = None
        self.link = None
        self.lldp = True
        self._custom_speed = speed
        self.set_available_tags(range(1, 4096))

        super().__init__()

    def __repr__(self):
        return f"Interface('{self.name}', {self.port_number}, {self.switch!r})"

    def __eq__(self, other):
        """Compare Interface class with another instance."""
        if isinstance(other, str):
            return self.address == other
        if isinstance(other, Interface):
            return self.port_number == other.port_number and \
                self.switch.dpid == other.switch.dpid
        return False

    @property
    def id(self):  # pylint: disable=invalid-name
        """Return id from Interface instance.

        Returns:
            string: Interface id.

        """
        return "{}:{}".format(self.switch.dpid, self.port_number)

    @property
    def uni(self):
        """Return if an interface is a user-to-network Interface."""
        return not self.nni

    def set_available_tags(self, iterable):
        """Set a range of VLAN tags to be used by this Interface.

        Args:
            iterable ([int]): range of VLANs.
        """
        self.available_tags = []

        for i in iterable:
            vlan = TAGType.VLAN
            tag = TAG(vlan, i)
            self.available_tags.append(tag)

    def enable(self):
        """Enable this interface instance.

        Also enable the switch instance this interface is attached to.
        """
        self.switch.enable()
        self._enabled = True

    def use_tag(self, tag):
        """Remove a specific tag from available_tags if it is there.

        Return False in case the tag is already removed.
        """
        try:
            self.available_tags.remove(tag)
        except ValueError:
            return False
        return True

    def is_tag_available(self, tag):
        """Check if a tag is available."""
        return tag in self.available_tags

    def get_next_available_tag(self):
        """Get the next available tag from the interface.

        Return the next available tag if exists and remove from the
        available tags.
        If no tag is available return False.
        """
        try:
            return self.available_tags.pop()
        except IndexError:
            return False

    def make_tag_available(self, tag):
        """Add a specific tag in available_tags."""
        if not self.is_tag_available(tag):
            self.available_tags.append(tag)
            return True
        return False

    def get_endpoint(self, endpoint):
        """Return a tuple with existent endpoint, None otherwise.

        Args:
            endpoint(|hw_address|, :class:`.Interface`): endpoint instance.

        Returns:
            tuple: A tuple with endpoint and time of last update.

        """
        for item in self.endpoints:
            if endpoint == item[0]:
                return item
        return None

    def add_endpoint(self, endpoint):
        """Create a new endpoint to Interface instance.

        Args:
            endpoint(|hw_address|, :class:`.Interface`): A target endpoint.
        """
        exists = self.get_endpoint(endpoint)
        if not exists:
            self.endpoints.append((endpoint, now()))

    def delete_endpoint(self, endpoint):
        """Delete a existent endpoint in Interface instance.

        Args:
            endpoint (|hw_address|, :class:`.Interface`): A target endpoint.
        """
        exists = self.get_endpoint(endpoint)
        if exists:
            self.endpoints.remove(exists)

    def update_endpoint(self, endpoint):
        """Update or create new endpoint to Interface instance.

        Args:
            endpoint(|hw_address|, :class:`.Interface`): A target endpoint.
        """
        exists = self.get_endpoint(endpoint)
        if exists:
            self.delete_endpoint(endpoint)
        self.add_endpoint(endpoint)

    def update_link(self, link):
        """Update link for this interface in a consistent way.

        Verify of the other endpoint of the link has the same Link information
        attached to it, and change it if necessary.

        Warning: This method can potentially change information of other
        Interface instances. Use it with caution.
        """
        if self not in (link.endpoint_a, link.endpoint_b):
            return False

        if self.link is None or self.link != link:
            self.link = link

        if link.endpoint_a == self:
            endpoint = link.endpoint_b
        else:
            endpoint = link.endpoint_a

        if endpoint.link is None or endpoint.link != link:
            endpoint.link = link

        return True

    @property
    def speed(self):
        """Return the link speed in bytes per second, None otherwise.

        If the switch was disconnected, we have :attr:`features` and speed is
        still returned for common values between v0x01 and v0x04. For specific
        v0x04 values (40 Gbps, 100 Gbps and 1 Tbps), the connection must be
        active so we can make sure the protocol version is v0x04.

        Returns:
            int, None: Link speed in bytes per second or ``None``.

        """
        if self._custom_speed is not None:
            return self._custom_speed

        if self._is_v0x04() and self.port_number == PortNo04.OFPP_LOCAL:
            return 0

        if not self._is_v0x04() and self.port_number == PortNo01.OFPP_LOCAL:
            return 0

        return self.get_of_features_speed()

    def set_custom_speed(self, bytes_per_second):
        """Set a speed that overrides switch OpenFlow information.

        If ``None`` is given, :attr:`speed` becomes the one given by the
        switch.
        """
        self._custom_speed = bytes_per_second

    def get_custom_speed(self):
        """Return custom speed or ``None`` if not set."""
        return self._custom_speed

    def get_of_features_speed(self):
        """Return the link speed in bytes per second, None otherwise.

        If the switch was disconnected, we have :attr:`features` and speed is
        still returned for common values between v0x01 and v0x04. For specific
        v0x04 values (40 Gbps, 100 Gbps and 1 Tbps), the connection must be
        active so we can make sure the protocol version is v0x04.

        Returns:
            int, None: Link speed in bytes per second or ``None``.

        """
        speed = self._get_v0x01_v0x04_speed()
        # Don't use switch.is_connected() because we can have the protocol
        if speed is None and self._is_v0x04():
            speed = self._get_v0x04_speed()
        if speed is not None:
            return speed
        # Warn unknown speed
        # Use shorter switch ID with its beginning and end
        if isinstance(self.switch.id, str) and len(self.switch.id) > 20:
            switch_id = self.switch.id[:3] + '...' + self.switch.id[-3:]
        else:
            switch_id = self.switch.id
        LOG.warning("Couldn't get port %s speed, sw %s, feats %s",
                    self.port_number, switch_id, self.features)
        return None

    def _is_v0x04(self):
        """Whether the switch is connected using OpenFlow 1.3."""
        return self.switch.is_connected() and \
            self.switch.connection.protocol.version == 0x04

    def _get_v0x01_v0x04_speed(self):
        """Check against all values of v0x01. They're part of v0x04."""
        fts = self.features
        pfts = PortFeatures01
        if fts and fts & pfts.OFPPF_10GB_FD:
            return 10 * 10**9 / 8
        if fts and fts & (pfts.OFPPF_1GB_HD | pfts.OFPPF_1GB_FD):
            return 10**9 / 8
        if fts and fts & (pfts.OFPPF_100MB_HD | pfts.OFPPF_100MB_FD):
            return 100 * 10**6 / 8
        if fts and fts & (pfts.OFPPF_10MB_HD | pfts.OFPPF_10MB_FD):
            return 10 * 10**6 / 8
        return None

    def _get_v0x04_speed(self):
        """Check against higher enums of v0x04.

        Must be called after :meth:`get_v0x01_speed` returns ``None``.
        """
        fts = self.features
        pfts = PortFeatures04
        if fts and fts & pfts.OFPPF_1TB_FD:
            return 10**12 / 8
        if fts and fts & pfts.OFPPF_100GB_FD:
            return 100 * 10**9 / 8
        if fts and fts & pfts.OFPPF_40GB_FD:
            return 40 * 10**9 / 8
        return None

    def get_hr_speed(self):
        """Return Human-Readable string for link speed.

        Returns:
            string: String with link speed. e.g: '350 Gbps' or '350 Mbps'.

        """
        speed = self.speed
        if speed is None:
            return ''
        speed *= 8
        if speed == 10**12:
            return '1 Tbps'
        if speed >= 10**9:
            return '{} Gbps'.format(round(speed / 10**9))
        return '{} Mbps'.format(round(speed / 10**6))

    def as_dict(self):
        """Return a dictionary with Interface attributes.

        Speed is in bytes/sec. Example of output (100 Gbps):

        .. code-block:: python3

            {'id': '00:00:00:00:00:00:00:01:2',
             'name': 'eth01',
             'port_number': 2,
             'mac': '00:7e:04:3b:c2:a6',
             'switch': '00:00:00:00:00:00:00:01',
             'type': 'interface',
             'nni': False,
             'uni': True,
             'speed': 12500000000,
             'metadata': {},
             'lldp': True,
             'active': True,
             'enabled': False,
             'link': ""
            }

        Returns:
            dict: Dictionary filled with interface attributes.

        """
        iface_dict = {'id': self.id,
                      'name': self.name,
                      'port_number': self.port_number,
                      'mac': self.address,
                      'switch': self.switch.dpid,
                      'type': 'interface',
                      'nni': self.nni,
                      'uni': self.uni,
                      'speed': self.speed,
                      'metadata': self.metadata,
                      'lldp': self.lldp,
                      'active': self.is_active(),
                      'enabled': self.is_enabled(),
                      'link': self.link.id if self.link else ""}
        if self.stats:
            iface_dict['stats'] = self.stats.as_dict()
        return iface_dict

    def as_json(self):
        """Return a json with Interfaces attributes.

        Example of output:

        .. code-block:: json

            {"mac": "00:7e:04:3b:c2:a6",
             "switch": "00:00:00:00:00:00:00:01",
             "type": "interface",
             "name": "eth01",
             "id": "00:00:00:00:00:00:00:01:2",
             "port_number": 2,
             "speed": "350 Mbps"}

        Returns:
            string: Json filled with interface attributes.

        """
        return json.dumps(self.as_dict())

class UNI:
    """Class that represents an User-to-Network Interface."""

    def __init__(self, interface, user_tag=None):
        self.user_tag = user_tag
        self.interface = interface

    def __eq__(self, other):
        """Override the default implementation."""
        return (self.user_tag == other.user_tag and
                self.interface == other.interface)

    def is_valid(self):
        """Check if TAG is possible for this interface TAG pool."""
        if self.user_tag:
            return self.interface.is_tag_available(self.user_tag)
        return True

    def as_dict(self):
        """Return a dict representating a UNI object."""
        uni_dict = {'interface_id': self.interface.id}
        if self.user_tag:
            uni_dict.update({'tag': self.user_tag.as_dict()})
        return uni_dict

    def as_json(self):
        """Return a json representating a UNI object."""
        return json.dumps(self.as_dict())

class NNI:
    """Class that represents an Network-to-Network Interface."""

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

class VNNI(NNI):
    """Class that represents an Virtual Network-to-Network Interface."""

    def __init__(self, service_tag, *args, **kwargs):
        self.service_tag = service_tag

        super().__init__(*args, **kwargs)
<commit_msg>Always return 'tag' key in UNI's as_dict.<commit_after>"""Module with main classes related to Interfaces."""
import json
import logging
from enum import IntEnum

from pyof.v0x01.common.phy_port import Port as PortNo01
from pyof.v0x01.common.phy_port import PortFeatures as PortFeatures01
from pyof.v0x04.common.port import PortFeatures as PortFeatures04
from pyof.v0x04.common.port import PortNo as PortNo04

from kytos.core.common import GenericEntity
from kytos.core.helpers import now

__all__ = ('Interface',)

LOG = logging.getLogger(__name__)

class TAGType(IntEnum):
    """Class that represents a TAG Type."""

    VLAN = 1
    VLAN_QINQ = 2
    MPLS = 3

class TAG:
    """Class that represents a TAG."""

    def __init__(self, tag_type, value):
        self.tag_type = tag_type
        self.value = value

    def __eq__(self, other):
        if not other:
            return False
        return self.tag_type == other.tag_type and self.value == other.value

    def as_dict(self):
        """Return a dictionary representating a tag object."""
        return {'tag_type': self.tag_type, 'value': self.value}

    @classmethod
    def from_dict(cls, tag_dict):
        """Return a TAG instance from python dictionary."""
        return cls(tag_dict.get('tag_type'), tag_dict.get('value'))

    @classmethod
    def from_json(cls, tag_json):
        """Return a TAG instance from json."""
        return cls.from_dict(json.loads(tag_json))

    def as_json(self):
        """Return a json representating a tag object."""
        return json.dumps(self.as_dict())

    def __repr__(self):
        return f"TAG({self.tag_type!r}, {self.value!r})"

class Interface(GenericEntity):  # pylint: disable=too-many-instance-attributes
    """Interface Class used to abstract the network interfaces."""

    # pylint: disable=too-many-arguments
    def __init__(self, name, port_number, switch, address=None, state=None,
                 features=None, speed=None):
        """Assign the parameters to instance attributes.

        Args:
            name (string): name from this interface.
            port_number (int): port number from this interface.
            switch (:class:`~.core.switch.Switch`): Switch with this interface.
            address (|hw_address|): Port address from this interface.
            state (|port_stats|): Port Stat from interface. It will be
            deprecated.
            features (|port_features|): Port feature used to calculate link
                utilization from this interface. It will be deprecated.
            speed (int, float): Interface speed in bytes per second. Defaults
                to what is informed by the switch. Return ``None`` if not set
                and switch does not inform the speed.
        """
        self.name = name
        self.port_number = int(port_number)
        self.switch = switch
        self.address = address
        self.state = state
        self.features = features
        self.nni = False
        self.endpoints = []
        self.stats = None
        self.link = None
        self.lldp = True
        self._custom_speed = speed
        self.set_available_tags(range(1, 4096))

        super().__init__()

    def __repr__(self):
        return f"Interface('{self.name}', {self.port_number}, {self.switch!r})"

    def __eq__(self, other):
        """Compare Interface class with another instance."""
        if isinstance(other, str):
            return self.address == other
        if isinstance(other, Interface):
            return self.port_number == other.port_number and \
                self.switch.dpid == other.switch.dpid
        return False

    @property
    def id(self):  # pylint: disable=invalid-name
        """Return id from Interface instance.

        Returns:
            string: Interface id.

        """
        return "{}:{}".format(self.switch.dpid, self.port_number)

    @property
    def uni(self):
        """Return if an interface is a user-to-network Interface."""
        return not self.nni

    def set_available_tags(self, iterable):
        """Set a range of VLAN tags to be used by this Interface.

        Args:
            iterable ([int]): range of VLANs.
        """
        self.available_tags = []

        for i in iterable:
            vlan = TAGType.VLAN
            tag = TAG(vlan, i)
            self.available_tags.append(tag)

    def enable(self):
        """Enable this interface instance.

        Also enable the switch instance this interface is attached to.
        """
        self.switch.enable()
        self._enabled = True

    def use_tag(self, tag):
        """Remove a specific tag from available_tags if it is there.

        Return False in case the tag is already removed.
        """
        try:
            self.available_tags.remove(tag)
        except ValueError:
            return False
        return True

    def is_tag_available(self, tag):
        """Check if a tag is available."""
        return tag in self.available_tags

    def get_next_available_tag(self):
        """Get the next available tag from the interface.

        Return the next available tag if exists and remove from the
        available tags.
        If no tag is available return False.
        """
        try:
            return self.available_tags.pop()
        except IndexError:
            return False

    def make_tag_available(self, tag):
        """Add a specific tag in available_tags."""
        if not self.is_tag_available(tag):
            self.available_tags.append(tag)
            return True
        return False

    def get_endpoint(self, endpoint):
        """Return a tuple with existent endpoint, None otherwise.

        Args:
            endpoint(|hw_address|, :class:`.Interface`): endpoint instance.

        Returns:
            tuple: A tuple with endpoint and time of last update.

        """
        for item in self.endpoints:
            if endpoint == item[0]:
                return item
        return None

    def add_endpoint(self, endpoint):
        """Create a new endpoint to Interface instance.

        Args:
            endpoint(|hw_address|, :class:`.Interface`): A target endpoint.
        """
        exists = self.get_endpoint(endpoint)
        if not exists:
            self.endpoints.append((endpoint, now()))

    def delete_endpoint(self, endpoint):
        """Delete a existent endpoint in Interface instance.

        Args:
            endpoint (|hw_address|, :class:`.Interface`): A target endpoint.
        """
        exists = self.get_endpoint(endpoint)
        if exists:
            self.endpoints.remove(exists)

    def update_endpoint(self, endpoint):
        """Update or create new endpoint to Interface instance.

        Args:
            endpoint(|hw_address|, :class:`.Interface`): A target endpoint.
        """
        exists = self.get_endpoint(endpoint)
        if exists:
            self.delete_endpoint(endpoint)
        self.add_endpoint(endpoint)

    def update_link(self, link):
        """Update link for this interface in a consistent way.

        Verify of the other endpoint of the link has the same Link information
        attached to it, and change it if necessary.

        Warning: This method can potentially change information of other
        Interface instances. Use it with caution.
        """
        if self not in (link.endpoint_a, link.endpoint_b):
            return False

        if self.link is None or self.link != link:
            self.link = link

        if link.endpoint_a == self:
            endpoint = link.endpoint_b
        else:
            endpoint = link.endpoint_a

        if endpoint.link is None or endpoint.link != link:
            endpoint.link = link

        return True

    @property
    def speed(self):
        """Return the link speed in bytes per second, None otherwise.

        If the switch was disconnected, we have :attr:`features` and speed is
        still returned for common values between v0x01 and v0x04. For specific
        v0x04 values (40 Gbps, 100 Gbps and 1 Tbps), the connection must be
        active so we can make sure the protocol version is v0x04.

        Returns:
            int, None: Link speed in bytes per second or ``None``.

        """
        if self._custom_speed is not None:
            return self._custom_speed

        if self._is_v0x04() and self.port_number == PortNo04.OFPP_LOCAL:
            return 0

        if not self._is_v0x04() and self.port_number == PortNo01.OFPP_LOCAL:
            return 0

        return self.get_of_features_speed()

    def set_custom_speed(self, bytes_per_second):
        """Set a speed that overrides switch OpenFlow information.

        If ``None`` is given, :attr:`speed` becomes the one given by the
        switch.
        """
        self._custom_speed = bytes_per_second

    def get_custom_speed(self):
        """Return custom speed or ``None`` if not set."""
        return self._custom_speed

    def get_of_features_speed(self):
        """Return the link speed in bytes per second, None otherwise.

        If the switch was disconnected, we have :attr:`features` and speed is
        still returned for common values between v0x01 and v0x04. For specific
        v0x04 values (40 Gbps, 100 Gbps and 1 Tbps), the connection must be
        active so we can make sure the protocol version is v0x04.

        Returns:
            int, None: Link speed in bytes per second or ``None``.

        """
        speed = self._get_v0x01_v0x04_speed()
        # Don't use switch.is_connected() because we can have the protocol
        if speed is None and self._is_v0x04():
            speed = self._get_v0x04_speed()
        if speed is not None:
            return speed
        # Warn unknown speed
        # Use shorter switch ID with its beginning and end
        if isinstance(self.switch.id, str) and len(self.switch.id) > 20:
            switch_id = self.switch.id[:3] + '...' + self.switch.id[-3:]
        else:
            switch_id = self.switch.id
        LOG.warning("Couldn't get port %s speed, sw %s, feats %s",
                    self.port_number, switch_id, self.features)
        return None

    def _is_v0x04(self):
        """Whether the switch is connected using OpenFlow 1.3."""
        return self.switch.is_connected() and \
            self.switch.connection.protocol.version == 0x04

    def _get_v0x01_v0x04_speed(self):
        """Check against all values of v0x01. They're part of v0x04."""
        fts = self.features
        pfts = PortFeatures01
        if fts and fts & pfts.OFPPF_10GB_FD:
            return 10 * 10**9 / 8
        if fts and fts & (pfts.OFPPF_1GB_HD | pfts.OFPPF_1GB_FD):
            return 10**9 / 8
        if fts and fts & (pfts.OFPPF_100MB_HD | pfts.OFPPF_100MB_FD):
            return 100 * 10**6 / 8
        if fts and fts & (pfts.OFPPF_10MB_HD | pfts.OFPPF_10MB_FD):
            return 10 * 10**6 / 8
        return None

    def _get_v0x04_speed(self):
        """Check against higher enums of v0x04.

        Must be called after :meth:`get_v0x01_speed` returns ``None``.
        """
        fts = self.features
        pfts = PortFeatures04
        if fts and fts & pfts.OFPPF_1TB_FD:
            return 10**12 / 8
        if fts and fts & pfts.OFPPF_100GB_FD:
            return 100 * 10**9 / 8
        if fts and fts & pfts.OFPPF_40GB_FD:
            return 40 * 10**9 / 8
        return None

    def get_hr_speed(self):
        """Return Human-Readable string for link speed.

        Returns:
            string: String with link speed. e.g: '350 Gbps' or '350 Mbps'.

        """
        speed = self.speed
        if speed is None:
            return ''
        speed *= 8
        if speed == 10**12:
            return '1 Tbps'
        if speed >= 10**9:
            return '{} Gbps'.format(round(speed / 10**9))
        return '{} Mbps'.format(round(speed / 10**6))

    def as_dict(self):
        """Return a dictionary with Interface attributes.

        Speed is in bytes/sec. Example of output (100 Gbps):

        .. code-block:: python3

            {'id': '00:00:00:00:00:00:00:01:2',
             'name': 'eth01',
             'port_number': 2,
             'mac': '00:7e:04:3b:c2:a6',
             'switch': '00:00:00:00:00:00:00:01',
             'type': 'interface',
             'nni': False,
             'uni': True,
             'speed': 12500000000,
             'metadata': {},
             'lldp': True,
             'active': True,
             'enabled': False,
             'link': ""
            }

        Returns:
            dict: Dictionary filled with interface attributes.

        """
        iface_dict = {'id': self.id,
                      'name': self.name,
                      'port_number': self.port_number,
                      'mac': self.address,
                      'switch': self.switch.dpid,
                      'type': 'interface',
                      'nni': self.nni,
                      'uni': self.uni,
                      'speed': self.speed,
                      'metadata': self.metadata,
                      'lldp': self.lldp,
                      'active': self.is_active(),
                      'enabled': self.is_enabled(),
                      'link': self.link.id if self.link else ""}
        if self.stats:
            iface_dict['stats'] = self.stats.as_dict()
        return iface_dict

    def as_json(self):
        """Return a json with Interfaces attributes.

        Example of output:

        .. code-block:: json

            {"mac": "00:7e:04:3b:c2:a6",
             "switch": "00:00:00:00:00:00:00:01",
             "type": "interface",
             "name": "eth01",
             "id": "00:00:00:00:00:00:00:01:2",
             "port_number": 2,
             "speed": "350 Mbps"}

        Returns:
            string: Json filled with interface attributes.

        """
        return json.dumps(self.as_dict())

class UNI:
    """Class that represents an User-to-Network Interface."""

    def __init__(self, interface, user_tag=None):
        self.user_tag = user_tag
        self.interface = interface

    def __eq__(self, other):
        """Override the default implementation."""
        return (self.user_tag == other.user_tag and
                self.interface == other.interface)

    def is_valid(self):
        """Check if TAG is possible for this interface TAG pool."""
        if self.user_tag:
            return self.interface.is_tag_available(self.user_tag)
        return True

    def as_dict(self):
        """Return a dict representating a UNI object."""
        return {
            'interface_id': self.interface.id,
            'tag': self.user_tag.as_dict() if self.user_tag else None
            }

    def as_json(self):
        """Return a json representating a UNI object."""
        return json.dumps(self.as_dict())

class NNI:
    """Class that represents an Network-to-Network Interface."""

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

class VNNI(NNI):
    """Class that represents an Virtual Network-to-Network Interface."""

    def __init__(self, service_tag, *args, **kwargs):
        self.service_tag = service_tag

        super().__init__(*args, **kwargs)
<|endoftext|>


<commit_before>"""Module with main classes related to Interfaces."""
import json
import logging
from enum import IntEnum

from pyof.v0x01.common.phy_port import Port as PortNo01
from pyof.v0x01.common.phy_port import PortFeatures as PortFeatures01
from pyof.v0x04.common.port import PortFeatures as PortFeatures04
from pyof.v0x04.common.port import PortNo as PortNo04

from kytos.core.common import GenericEntity
from kytos.core.helpers import now

__all__ = ('Interface',)

LOG = logging.getLogger(__name__)

class TAGType(IntEnum):
    """Class that represents a TAG Type."""

    VLAN = 1
    VLAN_QINQ = 2
    MPLS = 3

class TAG:
    """Class that represents a TAG."""

    def __init__(self, tag_type, value):
        self.tag_type = tag_type
        self.value = value

    def __eq__(self, other):
        return self.tag_type == other.tag_type and self.value == other.value

    def as_dict(self):
        """Return a dictionary representating a tag object."""
        return {'tag_type': self.tag_type, 'value': self.value}

    @classmethod
    def from_dict(cls, tag_dict):
        """Return a TAG instance from python dictionary."""
        return cls(tag_dict.get('tag_type'), tag_dict.get('value'))

    @classmethod
    def from_json(cls, tag_json):
        """Return a TAG instance from json."""
        return cls.from_dict(json.loads(tag_json))

    def as_json(self):
        """Return a json representating a tag object."""
        return json.dumps(self.as_dict())

    def __repr__(self):
        return f"TAG({self.tag_type!r}, {self.value!r})"

class Interface(GenericEntity):  # pylint: disable=too-many-instance-attributes
    """Interface Class used to abstract the network interfaces."""

    # pylint: disable=too-many-arguments
    def __init__(self, name, port_number, switch, address=None, state=None,
                 features=None, speed=None):
        """Assign the parameters to instance attributes.

        Args:
            name (string): name from this interface.
            port_number (int): port number from this interface.
            switch (:class:`~.core.switch.Switch`): Switch with this interface.
            address (|hw_address|): Port address from this interface.
            state (|port_stats|): Port Stat from interface. It will be
            deprecated.
            features (|port_features|): Port feature used to calculate link
                utilization from this interface. It will be deprecated.
            speed (int, float): Interface speed in bytes per second. Defaults
                to what is informed by the switch. Return ``None`` if not set
                and switch does not inform the speed.
        """
        self.name = name
        self.port_number = int(port_number)
        self.switch = switch
        self.address = address
        self.state = state
        self.features = features
        self.nni = False
        self.endpoints = []
        self.stats = None
        self.link = None
        self.lldp = True
        self._custom_speed = speed
        self.set_available_tags(range(1, 4096))

        super().__init__()

    def __repr__(self):
        return f"Interface('{self.name}', {self.port_number}, {self.switch!r})"

    def __eq__(self, other):
        """Compare Interface class with another instance."""
        if isinstance(other, str):
            return self.address == other
        if isinstance(other, Interface):
            return self.port_number == other.port_number and \
                self.switch.dpid == other.switch.dpid
        return False

    @property
    def id(self):  # pylint: disable=invalid-name
        """Return id from Interface instance.

        Returns:
            string: Interface id.

        """
        return "{}:{}".format(self.switch.dpid, self.port_number)

    @property
    def uni(self):
        """Return if an interface is a user-to-network Interface."""
        return not self.nni

    def set_available_tags(self, iterable):
        """Set a range of VLAN tags to be used by this Interface.

        Args:
            iterable ([int]): range of VLANs.
        """
        self.available_tags = []

        for i in iterable:
            vlan = TAGType.VLAN
            tag = TAG(vlan, i)
            self.available_tags.append(tag)

    def enable(self):
        """Enable this interface instance.

        Also enable the switch instance this interface is attached to.
        """
        self.switch.enable()
        self._enabled = True

    def use_tag(self, tag):
        """Remove a specific tag from available_tags if it is there.

        Return False in case the tag is already removed.
        """
        if tag in self.available_tags:
            self.available_tags.remove(tag)
            return True
        return False

    def is_tag_available(self, tag):
        """Check if a tag is available."""
        return tag in self.available_tags

    def get_next_available_tag(self):
        """Get the next available tag from the interface.

        Return the next available tag if exists and remove from the
        available tags.
        If no tag is available return False.
        """
        try:
            return self.available_tags.pop()
        except IndexError:
            return False

    def make_tag_available(self, tag):
        """Add a specific tag in available_tags."""
        if not self.is_tag_available(tag):
            self.available_tags.append(tag)
            return True
        return False

    def get_endpoint(self, endpoint):
        """Return a tuple with existent endpoint, None otherwise.

        Args:
            endpoint(|hw_address|, :class:`.Interface`): endpoint instance.

        Returns:
            tuple: A tuple with endpoint and time of last update.

        """
        for item in self.endpoints:
            if endpoint == item[0]:
                return item
        return None

    def add_endpoint(self, endpoint):
        """Create a new endpoint to Interface instance.

        Args:
            endpoint(|hw_address|, :class:`.Interface`): A target endpoint.
        """
        exists = self.get_endpoint(endpoint)
        if not exists:
            self.endpoints.append((endpoint, now()))

    def delete_endpoint(self, endpoint):
        """Delete a existent endpoint in Interface instance.

        Args:
            endpoint (|hw_address|, :class:`.Interface`): A target endpoint.
        """
        exists = self.get_endpoint(endpoint)
        if exists:
            self.endpoints.remove(exists)

    def update_endpoint(self, endpoint):
        """Update or create new endpoint to Interface instance.

        Args:
            endpoint(|hw_address|, :class:`.Interface`): A target endpoint.
        """
        exists = self.get_endpoint(endpoint)
        if exists:
            self.delete_endpoint(endpoint)
        self.add_endpoint(endpoint)

    def update_link(self, link):
        """Update link for this interface in a consistent way.

        Verify of the other endpoint of the link has the same Link information
        attached to it, and change it if necessary.

        Warning: This method can potentially change information of other
        Interface instances. Use it with caution.
        """
        if self not in (link.endpoint_a, link.endpoint_b):
            return False

        if self.link is None or self.link != link:
            self.link = link

        if link.endpoint_a == self:
            endpoint = link.endpoint_b
        else:
            endpoint = link.endpoint_a

        if endpoint.link is None or endpoint.link != link:
            endpoint.link = link

        return True

    @property
    def speed(self):
        """Return the link speed in bytes per second, None otherwise.

        If the switch was disconnected, we have :attr:`features` and speed is
        still returned for common values between v0x01 and v0x04. For specific
        v0x04 values (40 Gbps, 100 Gbps and 1 Tbps), the connection must be
        active so we can make sure the protocol version is v0x04.

        Returns:
            int, None: Link speed in bytes per second or ``None``.

        """
        if self._custom_speed is not None:
            return self._custom_speed

        if self._is_v0x04() and self.port_number == PortNo04.OFPP_LOCAL:
            return 0

        if not self._is_v0x04() and self.port_number == PortNo01.OFPP_LOCAL:
            return 0

        return self.get_of_features_speed()

    def set_custom_speed(self, bytes_per_second):
        """Set a speed that overrides switch OpenFlow information.

        If ``None`` is given, :attr:`speed` becomes the one given by the
        switch.
        """
        self._custom_speed = bytes_per_second

    def get_custom_speed(self):
        """Return custom speed or ``None`` if not set."""
        return self._custom_speed

    def get_of_features_speed(self):
        """Return the link speed in bytes per second, None otherwise.

        If the switch was disconnected, we have :attr:`features` and speed is
        still returned for common values between v0x01 and v0x04. For specific
        v0x04 values (40 Gbps, 100 Gbps and 1 Tbps), the connection must be
        active so we can make sure the protocol version is v0x04.

        Returns:
            int, None: Link speed in bytes per second or ``None``.

        """
        speed = self._get_v0x01_v0x04_speed()
        # Don't use switch.is_connected() because we can have the protocol
        if speed is None and self._is_v0x04():
            speed = self._get_v0x04_speed()
        if speed is not None:
            return speed
        # Warn unknown speed
        # Use shorter switch ID with its beginning and end
        if isinstance(self.switch.id, str) and len(self.switch.id) > 20:
            switch_id = self.switch.id[:3] + '...' + self.switch.id[-3:]
        else:
            switch_id = self.switch.id
        LOG.warning("Couldn't get port %s speed, sw %s, feats %s",
                    self.port_number, switch_id, self.features)
        return None

    def _is_v0x04(self):
        """Whether the switch is connected using OpenFlow 1.3."""
        return self.switch.is_connected() and \
            self.switch.connection.protocol.version == 0x04

    def _get_v0x01_v0x04_speed(self):
        """Check against all values of v0x01. They're part of v0x04."""
        fts = self.features
        pfts = PortFeatures01
        if fts and fts & pfts.OFPPF_10GB_FD:
            return 10 * 10**9 / 8
        if fts and fts & (pfts.OFPPF_1GB_HD | pfts.OFPPF_1GB_FD):
            return 10**9 / 8
        if fts and fts & (pfts.OFPPF_100MB_HD | pfts.OFPPF_100MB_FD):
            return 100 * 10**6 / 8
        if fts and fts & (pfts.OFPPF_10MB_HD | pfts.OFPPF_10MB_FD):
            return 10 * 10**6 / 8
        return None

    def _get_v0x04_speed(self):
        """Check against higher enums of v0x04.

        Must be called after :meth:`get_v0x01_speed` returns ``None``.
        """
        fts = self.features
        pfts = PortFeatures04
        if fts and fts & pfts.OFPPF_1TB_FD:
            return 10**12 / 8
        if fts and fts & pfts.OFPPF_100GB_FD:
            return 100 * 10**9 / 8
        if fts and fts & pfts.OFPPF_40GB_FD:
            return 40 * 10**9 / 8
        return None

    def get_hr_speed(self):
        """Return Human-Readable string for link speed.

        Returns:
            string: String with link speed. e.g: '350 Gbps' or '350 Mbps'.

        """
        speed = self.speed
        if speed is None:
            return ''
        speed *= 8
        if speed == 10**12:
            return '1 Tbps'
        if speed >= 10**9:
            return '{} Gbps'.format(round(speed / 10**9))
        return '{} Mbps'.format(round(speed / 10**6))

    def as_dict(self):
        """Return a dictionary with Interface attributes.

        Speed is in bytes/sec. Example of output (100 Gbps):

        .. code-block:: python3

            {'id': '00:00:00:00:00:00:00:01:2',
             'name': 'eth01',
             'port_number': 2,
             'mac': '00:7e:04:3b:c2:a6',
             'switch': '00:00:00:00:00:00:00:01',
             'type': 'interface',
             'nni': False,
             'uni': True,
             'speed': 12500000000,
             'metadata': {},
             'lldp': True,
             'active': True,
             'enabled': False,
             'link': ""
            }

        Returns:
            dict: Dictionary filled with interface attributes.

        """
        iface_dict = {'id': self.id,
                      'name': self.name,
                      'port_number': self.port_number,
                      'mac': self.address,
                      'switch': self.switch.dpid,
                      'type': 'interface',
                      'nni': self.nni,
                      'uni': self.uni,
                      'speed': self.speed,
                      'metadata': self.metadata,
                      'lldp': self.lldp,
                      'active': self.is_active(),
                      'enabled': self.is_enabled(),
                      'link': self.link.id if self.link else ""}
        if self.stats:
            iface_dict['stats'] = self.stats.as_dict()
        return iface_dict

    def as_json(self):
        """Return a json with Interfaces attributes.

        Example of output:

        .. code-block:: json

            {"mac": "00:7e:04:3b:c2:a6",
             "switch": "00:00:00:00:00:00:00:01",
             "type": "interface",
             "name": "eth01",
             "id": "00:00:00:00:00:00:00:01:2",
             "port_number": 2,
             "speed": "350 Mbps"}

        Returns:
            string: Json filled with interface attributes.

        """
        return json.dumps(self.as_dict())

class UNI:
    """Class that represents an User-to-Network Interface."""

    def __init__(self, interface, user_tag=None):
        self.user_tag = user_tag
        self.interface = interface

    def __eq__(self, other):
        """Override the default implementation."""
        return (self.user_tag == other.user_tag and
                self.interface == other.interface)

    def is_valid(self):
        """Check if TAG is possible for this interface TAG pool."""
        if self.user_tag:
            return self.interface.is_tag_available(self.user_tag)
        return True

    def as_dict(self):
        """Return a dict representating a UNI object."""
        return {'interface_id': self.interface.id,
                'tag': self.user_tag.as_dict()}

    def as_json(self):
        """Return a json representating a UNI object."""
        return json.dumps(self.as_dict())

class NNI:
    """Class that represents an Network-to-Network Interface."""

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

class VNNI(NNI):
    """Class that represents an Virtual Network-to-Network Interface."""

    def __init__(self, service_tag, *args, **kwargs):
        self.service_tag = service_tag

        super().__init__(*args, **kwargs)
<commit_msg>Replace if .. in with a try .. except.<commit_after>"""Module with main classes related to Interfaces."""
import json
import logging
from enum import IntEnum

from pyof.v0x01.common.phy_port import Port as PortNo01
from pyof.v0x01.common.phy_port import PortFeatures as PortFeatures01
from pyof.v0x04.common.port import PortFeatures as PortFeatures04
from pyof.v0x04.common.port import PortNo as PortNo04

from kytos.core.common import GenericEntity
from kytos.core.helpers import now

__all__ = ('Interface',)

LOG = logging.getLogger(__name__)

class TAGType(IntEnum):
    """Class that represents a TAG Type."""

    VLAN = 1
    VLAN_QINQ = 2
    MPLS = 3

class TAG:
    """Class that represents a TAG."""

    def __init__(self, tag_type, value):
        self.tag_type = tag_type
        self.value = value

    def __eq__(self, other):
        return self.tag_type == other.tag_type and self.value == other.value

    def as_dict(self):
        """Return a dictionary representating a tag object."""
        return {'tag_type': self.tag_type, 'value': self.value}

    @classmethod
    def from_dict(cls, tag_dict):
        """Return a TAG instance from python dictionary."""
        return cls(tag_dict.get('tag_type'), tag_dict.get('value'))

    @classmethod
    def from_json(cls, tag_json):
        """Return a TAG instance from json."""
        return cls.from_dict(json.loads(tag_json))

    def as_json(self):
        """Return a json representating a tag object."""
        return json.dumps(self.as_dict())

    def __repr__(self):
        return f"TAG({self.tag_type!r}, {self.value!r})"

class Interface(GenericEntity):  # pylint: disable=too-many-instance-attributes
    """Interface Class used to abstract the network interfaces."""

    # pylint: disable=too-many-arguments
    def __init__(self, name, port_number, switch, address=None, state=None,
                 features=None, speed=None):
        """Assign the parameters to instance attributes.

        Args:
            name (string): name from this interface.
            port_number (int): port number from this interface.
            switch (:class:`~.core.switch.Switch`): Switch with this interface.
            address (|hw_address|): Port address from this interface.
            state (|port_stats|): Port Stat from interface. It will be
            deprecated.
            features (|port_features|): Port feature used to calculate link
                utilization from this interface. It will be deprecated.
            speed (int, float): Interface speed in bytes per second. Defaults
                to what is informed by the switch. Return ``None`` if not set
                and switch does not inform the speed.
        """
        self.name = name
        self.port_number = int(port_number)
        self.switch = switch
        self.address = address
        self.state = state
        self.features = features
        self.nni = False
        self.endpoints = []
        self.stats = None
        self.link = None
        self.lldp = True
        self._custom_speed = speed
        self.set_available_tags(range(1, 4096))

        super().__init__()

    def __repr__(self):
        return f"Interface('{self.name}', {self.port_number}, {self.switch!r})"

    def __eq__(self, other):
        """Compare Interface class with another instance."""
        if isinstance(other, str):
            return self.address == other
        if isinstance(other, Interface):
            return self.port_number == other.port_number and \
                self.switch.dpid == other.switch.dpid
        return False

    @property
    def id(self):  # pylint: disable=invalid-name
        """Return id from Interface instance.

        Returns:
            string: Interface id.

        """
        return "{}:{}".format(self.switch.dpid, self.port_number)

    @property
    def uni(self):
        """Return if an interface is a user-to-network Interface."""
        return not self.nni

    def set_available_tags(self, iterable):
        """Set a range of VLAN tags to be used by this Interface.

        Args:
            iterable ([int]): range of VLANs.
        """
        self.available_tags = []

        for i in iterable:
            vlan = TAGType.VLAN
            tag = TAG(vlan, i)
            self.available_tags.append(tag)

    def enable(self):
        """Enable this interface instance.

        Also enable the switch instance this interface is attached to.
        """
        self.switch.enable()
        self._enabled = True

    def use_tag(self, tag):
        """Remove a specific tag from available_tags if it is there.

        Return False in case the tag is already removed.
        """
        try:
            self.available_tags.remove(tag)
        except ValueError:
            return False
        return True

    def is_tag_available(self, tag):
        """Check if a tag is available."""
        return tag in self.available_tags

    def get_next_available_tag(self):
        """Get the next available tag from the interface.

        Return the next available tag if exists and remove from the
        available tags.
        If no tag is available return False.
        """
        try:
            return self.available_tags.pop()
        except IndexError:
            return False

    def make_tag_available(self, tag):
        """Add a specific tag in available_tags."""
        if not self.is_tag_available(tag):
            self.available_tags.append(tag)
            return True
        return False

    def get_endpoint(self, endpoint):
        """Return a tuple with existent endpoint, None otherwise.

        Args:
            endpoint(|hw_address|, :class:`.Interface`): endpoint instance.

        Returns:
            tuple: A tuple with endpoint and time of last update.

        """
        for item in self.endpoints:
            if endpoint == item[0]:
                return item
        return None

    def add_endpoint(self, endpoint):
        """Create a new endpoint to Interface instance.

        Args:
            endpoint(|hw_address|, :class:`.Interface`): A target endpoint.
        """
        exists = self.get_endpoint(endpoint)
        if not exists:
            self.endpoints.append((endpoint, now()))

    def delete_endpoint(self, endpoint):
        """Delete a existent endpoint in Interface instance.

        Args:
            endpoint (|hw_address|, :class:`.Interface`): A target endpoint.
        """
        exists = self.get_endpoint(endpoint)
        if exists:
            self.endpoints.remove(exists)

    def update_endpoint(self, endpoint):
        """Update or create new endpoint to Interface instance.

        Args:
            endpoint(|hw_address|, :class:`.Interface`): A target endpoint.
        """
        exists = self.get_endpoint(endpoint)
        if exists:
            self.delete_endpoint(endpoint)
        self.add_endpoint(endpoint)

    def update_link(self, link):
        """Update link for this interface in a consistent way.

        Verify of the other endpoint of the link has the same Link information
        attached to it, and change it if necessary.

        Warning: This method can potentially change information of other
        Interface instances. Use it with caution.
        """
        if self not in (link.endpoint_a, link.endpoint_b):
            return False

        if self.link is None or self.link != link:
            self.link = link

        if link.endpoint_a == self:
            endpoint = link.endpoint_b
        else:
            endpoint = link.endpoint_a

        if endpoint.link is None or endpoint.link != link:
            endpoint.link = link

        return True

    @property
    def speed(self):
        """Return the link speed in bytes per second, None otherwise.

        If the switch was disconnected, we have :attr:`features` and speed is
        still returned for common values between v0x01 and v0x04. For specific
        v0x04 values (40 Gbps, 100 Gbps and 1 Tbps), the connection must be
        active so we can make sure the protocol version is v0x04.

        Returns:
            int, None: Link speed in bytes per second or ``None``.

        """
        if self._custom_speed is not None:
            return self._custom_speed

        if self._is_v0x04() and self.port_number == PortNo04.OFPP_LOCAL:
            return 0

        if not self._is_v0x04() and self.port_number == PortNo01.OFPP_LOCAL:
            return 0

        return self.get_of_features_speed()

    def set_custom_speed(self, bytes_per_second):
        """Set a speed that overrides switch OpenFlow information.

        If ``None`` is given, :attr:`speed` becomes the one given by the
        switch.
        """
        self._custom_speed = bytes_per_second

    def get_custom_speed(self):
        """Return custom speed or ``None`` if not set."""
        return self._custom_speed

    def get_of_features_speed(self):
        """Return the link speed in bytes per second, None otherwise.

        If the switch was disconnected, we have :attr:`features` and speed is
        still returned for common values between v0x01 and v0x04. For specific
        v0x04 values (40 Gbps, 100 Gbps and 1 Tbps), the connection must be
        active so we can make sure the protocol version is v0x04.

        Returns:
            int, None: Link speed in bytes per second or ``None``.

        """
        speed = self._get_v0x01_v0x04_speed()
        # Don't use switch.is_connected() because we can have the protocol
        if speed is None and self._is_v0x04():
            speed = self._get_v0x04_speed()
        if speed is not None:
            return speed
        # Warn unknown speed
        # Use shorter switch ID with its beginning and end
        if isinstance(self.switch.id, str) and len(self.switch.id) > 20:
            switch_id = self.switch.id[:3] + '...' + self.switch.id[-3:]
        else:
            switch_id = self.switch.id
        LOG.warning("Couldn't get port %s speed, sw %s, feats %s",
                    self.port_number, switch_id, self.features)
        return None

    def _is_v0x04(self):
        """Whether the switch is connected using OpenFlow 1.3."""
        return self.switch.is_connected() and \
            self.switch.connection.protocol.version == 0x04

    def _get_v0x01_v0x04_speed(self):
        """Check against all values of v0x01. They're part of v0x04."""
        fts = self.features
        pfts = PortFeatures01
        if fts and fts & pfts.OFPPF_10GB_FD:
            return 10 * 10**9 / 8
        if fts and fts & (pfts.OFPPF_1GB_HD | pfts.OFPPF_1GB_FD):
            return 10**9 / 8
        if fts and fts & (pfts.OFPPF_100MB_HD | pfts.OFPPF_100MB_FD):
            return 100 * 10**6 / 8
        if fts and fts & (pfts.OFPPF_10MB_HD | pfts.OFPPF_10MB_FD):
            return 10 * 10**6 / 8
        return None

    def _get_v0x04_speed(self):
        """Check against higher enums of v0x04.

        Must be called after :meth:`get_v0x01_speed` returns ``None``.
        """
        fts = self.features
        pfts = PortFeatures04
        if fts and fts & pfts.OFPPF_1TB_FD:
            return 10**12 / 8
        if fts and fts & pfts.OFPPF_100GB_FD:
            return 100 * 10**9 / 8
        if fts and fts & pfts.OFPPF_40GB_FD:
            return 40 * 10**9 / 8
        return None

    def get_hr_speed(self):
        """Return Human-Readable string for link speed.

        Returns:
            string: String with link speed. e.g: '350 Gbps' or '350 Mbps'.

        """
        speed = self.speed
        if speed is None:
            return ''
        speed *= 8
        if speed == 10**12:
            return '1 Tbps'
        if speed >= 10**9:
            return '{} Gbps'.format(round(speed / 10**9))
        return '{} Mbps'.format(round(speed / 10**6))

    def as_dict(self):
        """Return a dictionary with Interface attributes.

        Speed is in bytes/sec. Example of output (100 Gbps):

        .. code-block:: python3

            {'id': '00:00:00:00:00:00:00:01:2',
             'name': 'eth01',
             'port_number': 2,
             'mac': '00:7e:04:3b:c2:a6',
             'switch': '00:00:00:00:00:00:00:01',
             'type': 'interface',
             'nni': False,
             'uni': True,
             'speed': 12500000000,
             'metadata': {},
             'lldp': True,
             'active': True,
             'enabled': False,
             'link': ""
            }

        Returns:
            dict: Dictionary filled with interface attributes.

        """
        iface_dict = {'id': self.id,
                      'name': self.name,
                      'port_number': self.port_number,
                      'mac': self.address,
                      'switch': self.switch.dpid,
                      'type': 'interface',
                      'nni': self.nni,
                      'uni': self.uni,
                      'speed': self.speed,
                      'metadata': self.metadata,
                      'lldp': self.lldp,
                      'active': self.is_active(),
                      'enabled': self.is_enabled(),
                      'link': self.link.id if self.link else ""}
        if self.stats:
            iface_dict['stats'] = self.stats.as_dict()
        return iface_dict

    def as_json(self):
        """Return a json with Interfaces attributes.

        Example of output:

        .. code-block:: json

            {"mac": "00:7e:04:3b:c2:a6",
             "switch": "00:00:00:00:00:00:00:01",
             "type": "interface",
             "name": "eth01",
             "id": "00:00:00:00:00:00:00:01:2",
             "port_number": 2,
             "speed": "350 Mbps"}

        Returns:
            string: Json filled with interface attributes.

        """
        return json.dumps(self.as_dict())

class UNI:
    """Class that represents an User-to-Network Interface."""

    def __init__(self, interface, user_tag=None):
        self.user_tag = user_tag
        self.interface = interface

    def __eq__(self, other):
        """Override the default implementation."""
        return (self.user_tag == other.user_tag and
                self.interface == other.interface)

    def is_valid(self):
        """Check if TAG is possible for this interface TAG pool."""
        if self.user_tag:
            return self.interface.is_tag_available(self.user_tag)
        return True

    def as_dict(self):
        """Return a dict representating a UNI object."""
        return {'interface_id': self.interface.id,
                'tag': self.user_tag.as_dict()}

    def as_json(self):
        """Return a json representating a UNI object."""
        return json.dumps(self.as_dict())

class NNI:
    """Class that represents an Network-to-Network Interface."""

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

class VNNI(NNI):
    """Class that represents an Virtual Network-to-Network Interface."""

    def __init__(self, service_tag, *args, **kwargs):
        self.service_tag = service_tag

        super().__init__(*args, **kwargs)
<|endoftext|>


<commit_before>import cv2
import numpy as np
import play

neck_len=500
neck_top=0
neck_bottom=500

"""
strums = [
	['A',  'E4 C#3 A4 E3 A2 E2'],
	['C', 'E4 C3 G3 E3 C2 E2'],
	['D', 'F#4 D3 A4 D3 A2 E2'],
	['E', 'E4 B3 G#3 E3 B2 E2'],
	['G', 'G4 B3 G3 D3 B2 G2']]
"""

#(post_fix_num,note)

strums = {
	'A': [(3,7),(2,4),(3,0),(2,7),(1,0),(1,7)],
	'C': [(3,7),(2,3),(10,2),(2,7),(1,3),(1,7)],
	'D': [(3,9),(2,5),(3,0),(2,5),(1,0),(1,7)],
	'E': [(3,7),(2,2),(2,11),(2,7),(1,2),(1,7)],
	'G': [(3,10),(2,2),(2,10),(2,5),(1,2),(1,10)]}

notes_ar = [['A1','A#1','B1','C1','C#1','D1','D#1','E1','F1','F#1','G1','G#1'],
		['A2','A#2','B2','C2','C#2','D2','D#2','E2','F2','F#2','G2','G#2'],
		['A3','A#3','B3','C3','C#3','D3','D#3','E3','F3','F#3','G3','G#3'],
		['A4','A#4','B4','C4','C#4','D4','D#4','E4','F4','F#4','G4','G#4'],
		['A5','A#5','B5','C5','C#5','D5','D#5','E5','F5','F#5','G5','G#5']]

# HSV color ranges
ranges = [[(160, 179),(106, 255),(0, 255)], # Red
		[(60, 90),(81, 255),(0, 255)],	# Green
		[(100, 119),(136, 255),(0, 255)]] # Blue

def clearNoise(img):
	kernel = np.ones((10, 10), np.uint8)
	clrimg = cv2.morphologyEx(img, cv2.MORPH_CLOSE, kernel)
	clrimg = cv2.erode(img, kernel, iterations=1)
	kernel = np.ones((15, 15), np.uint8)
	clrimg = cv2.dilate(img, kernel, iterations=2)
	# Clears noise from image
	return clrimg

def get_strum(mode, direction):
	if direction == 'up':
		for row in strums:
			if row[0] == mode:
				return row[1]
	else:
		for row in strums:
			if row[0] == mode:
				rev = ""
				_array = row[1].split(" ")
				for row2 in reversed(_array):
					rev += row2 + " "
				return rev[:-1]

def filterFingers(img):

	height = img.shape[0]
	width = img.shape[1]

	imgcropped = img[0:(height/2), (width/2):(width-1)] # Cropping the ROI

	fingerArray = []

	min = np.array([ranges[0][0][0], ranges[0][1][0], ranges[0][2][0]], np.uint8)
	max = np.array([ranges[0][0][1], ranges[0][1][1], ranges[0][2][1]], np.uint8)
	red = cv2.inRange(imgcropped, min, max)
	red = clearNoise(red)
	fingerArray.append(red)

	min = np.array([ranges[1][0][0], ranges[1][1][0], ranges[1][2][0]], np.uint8)
	max = np.array([ranges[1][0][1], ranges[1][1][1], ranges[1][2][1]], np.uint8)
	green = cv2.inRange(imgcropped, min, max)
	green = clearNoise(green)
	fingerArray.append(green)

	min = np.array([ranges[2][0][0], ranges[2][1][0], ranges[2][2][0]], np.uint8)
	max = np.array([ranges[2][0][1], ranges[2][1][1], ranges[2][2][1]], np.uint8)
	blue = cv2.inRange(imgcropped, min, max)
	blue = clearNoise(blue)
	fingerArray.append(blue)

	# Returns an array of three filtered fingers images
	return fingerArray

def getPositions(imgArray):
	# Returns the topmost points of filtered blobs from given imgs

	positions=[]

	for img in imgArray:

		contours, hierarchy = cv2.findContours(img,cv2.RETR_TREE,cv2.CHAIN_APPROX_SIMPLE)

		if len(contours)<1:
			positions.append((0,0))
			continue

		cnt=contours[0]

		topmost = tuple(cnt[cnt[:,:,1].argmin()][0])

		positions.append(topmost)

	return positions

def getMode(positionArray):
	# Finds mode using by filtering three color strips and finding positions

	#Testing with Y coordinate Only

	R=positionArray[0][1]
	G=positionArray[1][1]
	B=positionArray[2][1]

	if R<G<B:
		mode='A'

	elif B<G<R:
		mode='C'

	elif R<B<G:
		mode='D'

	elif G<B<R:
		mode='E'

	else:
		mode='G'

	return mode

def getLowerBlob(img):
	# Filters lower blob and returns position
	position = (0, 0)
	height = img.shape[0]
	width = img.shape[1]

	imgcropped = img[(height/2):(height-1), 0:(width/2)] # Cropping the ROI

	min = np.array([ranges[0][0][0], ranges[0][1][0], ranges[0][2][0]], np.uint8)
	max = np.array([ranges[0][0][1], ranges[0][1][1], ranges[0][2][1]], np.uint8)
	lowerhand = [cv2.inRange(imgcropped, min, max)]
	lowerhand[0] = clearNoise(lowerhand[0])
	pos = getPositions(lowerhand)

	if pos != (0, 0):
		position = pos

	return position

def getPattern(mode,dist,direction):
	pattern=''
	for note in strums[mode]:
		if note[1]+dist>11:
			pattern=pattern+notes_ar[note[0]+1][(note[1]+dist)%12]+' '
		else:
			pattern=pattern+notes_ar[note[0]][(note[1]+dist)%12]+' '
	if direction=='up':
		return pattern
	else:
		rev = ""
		_array = pattern.split(" ")
		for row2 in reversed(_array):
			rev += row2 + " "
		return rev[:-1]

def initNeck(top,bottom):
	#Initialize the neck length
	neck_len=(bottom[0][0]-top[0])/2
	neck_bottom=bottom[0][0]
	neck_top=top[0]

def getDistance(positions):
	#Get Distance on the neck
	mean=0

	for pos in positions:
		mean+=pos[0]

	mean=mean/3-neck_top

	if mean > 3*neck_len/4:
		return 3
	elif mean > 2*neck_len/4:
		return 2
	elif mean > neck_len/4:
		return 1
	else:
		return 0
<commit_msg>Changes to Finding Distance.<commit_after>import cv2
import numpy as np
import play

neck_len=500
neck_top=0
neck_bottom=500

"""
strums = [
	['A',  'E4 C#3 A4 E3 A2 E2'],
	['C', 'E4 C3 G3 E3 C2 E2'],
	['D', 'F#4 D3 A4 D3 A2 E2'],
	['E', 'E4 B3 G#3 E3 B2 E2'],
	['G', 'G4 B3 G3 D3 B2 G2']]
"""

#(post_fix_num,note)

strums = {
	'A': [(3,7),(2,4),(3,0),(2,7),(1,0),(1,7)],
	'C': [(3,7),(2,3),(10,2),(2,7),(1,3),(1,7)],
	'D': [(3,9),(2,5),(3,0),(2,5),(1,0),(1,7)],
	'E': [(3,7),(2,2),(2,11),(2,7),(1,2),(1,7)],
	'G': [(3,10),(2,2),(2,10),(2,5),(1,2),(1,10)]}

notes_ar = [['A1','A#1','B1','C1','C#1','D1','D#1','E1','F1','F#1','G1','G#1'],
		['A2','A#2','B2','C2','C#2','D2','D#2','E2','F2','F#2','G2','G#2'],
		['A3','A#3','B3','C3','C#3','D3','D#3','E3','F3','F#3','G3','G#3'],
		['A4','A#4','B4','C4','C#4','D4','D#4','E4','F4','F#4','G4','G#4'],
		['A5','A#5','B5','C5','C#5','D5','D#5','E5','F5','F#5','G5','G#5']]

# HSV color ranges
ranges = [[(160, 179),(106, 255),(0, 255)], # Red
		[(60, 90),(81, 255),(0, 255)],	# Green
		[(100, 119),(136, 255),(0, 255)]] # Blue

def clearNoise(img):
	kernel = np.ones((10, 10), np.uint8)
	clrimg = cv2.morphologyEx(img, cv2.MORPH_CLOSE, kernel)
	clrimg = cv2.erode(img, kernel, iterations=1)
	kernel = np.ones((15, 15), np.uint8)
	clrimg = cv2.dilate(img, kernel, iterations=2)
	# Clears noise from image
	return clrimg

def get_strum(mode, direction):
	if direction == 'up':
		for row in strums:
			if row[0] == mode:
				return row[1]
	else:
		for row in strums:
			if row[0] == mode:
				rev = ""
				_array = row[1].split(" ")
				for row2 in reversed(_array):
					rev += row2 + " "
				return rev[:-1]

def filterFingers(img):

	height = img.shape[0]
	width = img.shape[1]

	imgcropped = img[0:(height/2), (width/2):(width-1)] # Cropping the ROI

	fingerArray = []

	min = np.array([ranges[0][0][0], ranges[0][1][0], ranges[0][2][0]], np.uint8)
	max = np.array([ranges[0][0][1], ranges[0][1][1], ranges[0][2][1]], np.uint8)
	red = cv2.inRange(imgcropped, min, max)
	red = clearNoise(red)
	fingerArray.append(red)

	min = np.array([ranges[1][0][0], ranges[1][1][0], ranges[1][2][0]], np.uint8)
	max = np.array([ranges[1][0][1], ranges[1][1][1], ranges[1][2][1]], np.uint8)
	green = cv2.inRange(imgcropped, min, max)
	green = clearNoise(green)
	fingerArray.append(green)

	min = np.array([ranges[2][0][0], ranges[2][1][0], ranges[2][2][0]], np.uint8)
	max = np.array([ranges[2][0][1], ranges[2][1][1], ranges[2][2][1]], np.uint8)
	blue = cv2.inRange(imgcropped, min, max)
	blue = clearNoise(blue)
	fingerArray.append(blue)

	# Returns an array of three filtered fingers images
	return fingerArray

def getPositions(imgArray):
	# Returns the topmost points of filtered blobs from given imgs

	positions=[]

	for img in imgArray:

		contours, hierarchy = cv2.findContours(img,cv2.RETR_TREE,cv2.CHAIN_APPROX_SIMPLE)

		if len(contours)<1:
			positions.append((0,0))
			continue

		cnt=contours[0]

		topmost = tuple(cnt[cnt[:,:,1].argmin()][0])

		positions.append(topmost)

	return positions

def getMode(positionArray):
	# Finds mode using by filtering three color strips and finding positions

	#Testing with Y coordinate Only

	R=positionArray[0][1]
	G=positionArray[1][1]
	B=positionArray[2][1]

	if R<G<B:
		mode='A'

	elif B<G<R:
		mode='C'

	elif R<B<G:
		mode='D'

	elif G<B<R:
		mode='E'

	else:
		mode='G'

	return mode

def getLowerBlob(img):
	# Filters lower blob and returns position
	position = (0, 0)
	height = img.shape[0]
	width = img.shape[1]

	imgcropped = img[(height/2):(height-1), 0:(width/2)] # Cropping the ROI

	min = np.array([ranges[0][0][0], ranges[0][1][0], ranges[0][2][0]], np.uint8)
	max = np.array([ranges[0][0][1], ranges[0][1][1], ranges[0][2][1]], np.uint8)
	lowerhand = [cv2.inRange(imgcropped, min, max)]
	lowerhand[0] = clearNoise(lowerhand[0])
	pos = getPositions(lowerhand)

	if pos != (0, 0):
		position = pos

	return position

def getPattern(mode,dist,direction):
	pattern=''
	for note in strums[mode]:
		if note[1]+dist>11:
			pattern=pattern+notes_ar[note[0]+1][(note[1]+dist)%12]+' '
		else:
			pattern=pattern+notes_ar[note[0]][(note[1]+dist)%12]+' '
	if direction=='up':
		return pattern
	else:
		rev = ""
		_array = pattern.split(" ")
		for row2 in reversed(_array):
			rev += row2 + " "
		return rev[:-1]

def initNeck(top,bottom):
	#Initialize the neck length
	neck_len=(top[0]-bottom[0][0])/2
	neck_bottom=bottom[0][0]
	neck_top=top[0]

def getDistance(positions):
	#Get Distance on the neck
	mean=0

	for pos in positions:
		mean+=pos[0]

	mean=neck_top-mean/3

	if mean > 3*neck_len/4:
		return 3
	elif mean > 2*neck_len/4:
		return 2
	elif mean > neck_len/4:
		return 1
	else:
		return 0
<|endoftext|>


<commit_before>#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""Machine learning model which uses Wikipedia data to predicts ILI levels.

Usage:
  model.py <year_start> <year_end> <dataset_path> <incidence_path> <keywords_file> [--p] [--f] [--v] [--d=<directory>] [--no-future]

  <year_start>      The first influenza season we want to predict.
  <year_end>        The last influenza season we want to predict.
  <dataset_path>    The path to dataset files (the ones which contain Wikipedia pageviews).
  <incidence_path>  The path to the files which specify the ILI levels for each week.
  <keywords_file>   The path to the file which contains the Wikipedia pages used.
  -p, --poisson     Use the Poisson model + LASSO instead of the linear one.
  -f, --file        Write informations to file
  -v, --verbose     Output more informations
  -d, --directory   Select a directory in which save your files
  -n, --no-future   Use a different method to train the model (avoid using seasonal influenza which are
                    later than the one we want to predict)
  -h, --help        Print this help message
"""

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import csv
from sklearn.linear_model import LassoCV
from sklearn.metrics import mean_squared_error
from tabulate import tabulate
from docopt import docopt

from models_utils import generate, generate_labels, generate_labels_one_year, generate_one_year, standardize_data, generate_keywords, stz, get_important_pages, correlation_matrix

from cvglmnetCoef import cvglmnetCoef
from cvglmnetPredict import cvglmnetPredict
from cvglmnet import cvglmnet

#########################
#### INITIALIZATION #####
#########################

# Parse the command line
arguments = docopt(__doc__)

# Directory where to save files
save_directory="./"
if arguments["--d"]:
    save_directory = arguments["--d"]

# Year selected
year_start = int(arguments["<year_start>"])
year_end = int(arguments["<year_end>"])+1

# Model type
model_type = "Linear Model"
using_poisson = False
if arguments["--p"]:
    using_poisson = True
    model_type="Poisson Model"

# Feature path and labels
path_features = arguments["<dataset_path>"]
path_labels = arguments["<incidence_path>"]
keywords = arguments["<keywords_file>"]

# Selected columns that will be extracted from the dataframes
selected_columns = generate_keywords(keywords)

all_features_values = pd.DataFrame()
all_predicted_values = []
all_true_labels = pd.DataFrame()
total_weeks = []
all_weighted_feature=dict()

# The pair start-end seasons
year_sel = (year_start, year_end)

######################
##### ALGORITHM ######
######################

# For each season, train a model with the complete dataset (without the season)
# and predict the ILI incidence.
for year_selected in range(year_sel[0], year_sel[1]):

    print("------------")
    print("[*] ", year_selected )

    # Generate the list of excluded years
    excluded=[]
    if arguments["--no-future"]:
        for i in range(2007, 2018):
            if i>year_selected:
                excluded.append(i)

    # Data generation from data files
    dataset = generate(year_selected, excluded, path_features)[selected_columns]
    labels = generate_labels(year_selected, excluded, path_labels)
    data = generate_one_year(year_selected, path_features)[selected_columns]
    labels_test = generate_labels_one_year(year_selected, path_labels)
    weeks = labels_test["week"]

    # Create a copy of the dataset where we fills NaN values
    # with 0.
    dataset_zero = dataset.fillna(0)
    data_zero = data.fillna(0)

    # Standardize data
    train, test = standardize_data(dataset_zero, data_zero)

    # Create a Lasso Cross-Validation instance which will be
    # trained on the dataset in which NaN values are replaced
    # with the column mean.
    if not using_poisson:
        model = LassoCV(max_iter=100000, n_jobs=-1, n_alphas=1000, random_state=1)
        model.fit(train, labels["incidence"].fillna(0))
        result = model.predict(test)
    else:
        model = cvglmnet(x=train.as_matrix().copy(), y=labels["incidence"].fillna(0).copy().as_matrix(), family='poisson', alpha=1.0,
                       ptype="mse", parallel=True, nfolds=10)
        result = cvglmnetPredict(model, test.as_matrix(), ptype='response', s="lambda_min")
        coeff = cvglmnetCoef(model, s="lambda_min")

    # Get the feature coefficients
    if not using_poisson:
        coeff = model.coef_
    else:
        coeff = cvglmnetCoef(model, s="lambda_min")

    # Add the pair (value, coeff) to a dictionary
    for i in list(zip(coeff, selected_columns)):
        if (i[0] != 0):
            if all_weighted_feature.get(i[1], None) is None:
                all_weighted_feature[i[1]] = [i[0]]
            else:
                all_weighted_feature[i[1]].append(i[0])

    # Add the value to global variables
    all_features_values= all_features_values.append(data.fillna(0))
    all_predicted_values.extend(result)
    all_true_labels = all_true_labels.append(labels_test)
    total_weeks.extend(weeks)

    # If we are verbose print more informations
    if arguments["--v"]:
        index_peak = labels_test['incidence'].fillna(0).idxmax(axis=1)
        index_peak_m = np.argmax(result)
        print("[*] Influenza Season Peak (Week):", labels_test.iloc[index_peak]["week"])
        print("[*] Influenza Seasons Predicted Peak (Week): ", labels_test.iloc[index_peak_m]["week"])
        print("[*] Influenza Season Peak Value: ", labels_test.iloc[index_peak]["incidence"])
        print("[*] Influenza Season Predicted Peak Value: ", result[index_peak_m])
        print("------------")

# Get important pages to generate the plot we need
important_pages = get_important_pages(all_weighted_feature, 5)

# Print MSE of the two models
mse = mean_squared_error(all_true_labels["incidence"].fillna(0), all_predicted_values)
print("------------")
print("MSE: ", mse)
print("------------")

# Print Pearson Coefficient
pcc = np.corrcoef(all_predicted_values, all_true_labels["incidence"].fillna(0), rowvar=False)[0][1]
print("------------")
print("Pearson Correlation Coeff: ", pcc)
print("------------")

# Print important pages
print(tabulate(important_pages, headers=["Page name", "Weight"]))

###########################
#### GRAPHS GENERATION ####
###########################

# Plot some informations
fig = plt.figure(3, figsize=(15, 6))
ax = fig.add_subplot(111)

# Set up the graph title, x and y axis titles
plt.title("Influenza Seasons "+str(year_sel[0]-1)+" - "+str(year_sel[1]-1), fontsize=18)
plt.ylabel("Incidence on 100000 people", fontsize=17)
plt.xlabel("Year-Week", fontsize=17)

# Generate the x axis labels in such way to
# have the year-week pair only every two ticks.
weeks_used = []
for k, v in enumerate(total_weeks):
    if k%2==0:
        weeks_used.append(v)
    else:
        weeks_used.append(" ")

# Set up the axes ticks
plt.xticks(range(0, len(weeks_used)), weeks_used, rotation="vertical", fontsize=15)
plt.yticks(fontsize=15)

# Plot the model result and the incidence
plt.plot(range(0, len(all_predicted_values)), all_predicted_values, 'r-', label=model_type, linewidth=3)
plt.plot(range(0, len(all_true_labels["incidence"])), all_true_labels["incidence"].fillna(0), 'k-', label="Incidence", linewidth=3)

# Add dotted line into the graph to delimit the influenza seasons.
for i in range(1, (year_sel[1]-year_sel[0])):
    plt.axvline(26*i, linestyle="dashed", color="k", linewidth=3)

# Update some graphical options
plt.grid()
plt.legend(fontsize=16, loc="upper right")
plt.tight_layout()

# Save the graph on file
plt.savefig(save_directory+"season_"+str(year_sel[0]-1)+"-"+str(year_sel[1]-1)+".png", dpi=150)

# Feature plot
figf = plt.figure(4, figsize=(15, 6))
axf = fig.add_subplot(111)

# Generate the graph showing the pageview variation of the top 5 Wikipedia's pages
# choosen by the model.
plt.title("Pageview variation "+str(year_sel[0]-1)+" - "+str(year_sel[1]-1), fontsize=18)
plt.ylabel("Pageview value", fontsize=17)
plt.xlabel("Year-Week", fontsize=17)

# Plot the axes labels
plt.xticks(range(0, len(weeks_used)), weeks_used, rotation="vertical", fontsize=15)
plt.yticks(fontsize=15)

# Plot all the pageview data and the ILI incidence
std_all_features_values = stz(all_features_values)
for key, value in important_pages:
 plt.plot(range(0, len(std_all_features_values[key])), std_all_features_values[key], '-', label=key, linewidth=3)
plt.plot(range(0, len(all_true_labels["incidence"])), stz(all_true_labels["incidence"]), 'k-', label="Incidence", linewidth=3)

# Add dotted line into the graph to delimit the influenza seasons.
for i in range(1, (year_sel[1]-year_sel[0])):
    plt.axvline(26*i, linestyle="dashed", color="k", linewidth=3)

# Update some graphical options
plt.grid()
plt.legend(fontsize=16, loc="best")
plt.tight_layout()

# Save again the graph on file
plt.savefig(save_directory+"season_"+str(year_sel[0]-1)+"-"+str(year_sel[1]-1)+"_features.png", dpi=150)

# If --f then we write some information to file
if arguments["--f"]:
    with open(save_directory+str(year_sel[0]-1)+"-"+str(year_sel[1]-1)+"_information.csv", 'w', newline='') as csvfile:
        spamwriter = csv.writer(csvfile, delimiter=',',
                                quotechar='|', quoting=csv.QUOTE_MINIMAL)
        spamwriter.writerow(['mse', mean_squared_error(all_true_labels["incidence"].fillna(0), all_predicted_values)])
        for p in important_pages:
            spamwriter.writerow([p[0], float(p[1])])

# Correlation Matrix
new_pages=[]
for p in important_pages:
    new_pages.append(p[0])
first = pd.DataFrame(all_features_values[new_pages])
second = pd.DataFrame(all_true_labels["incidence"]).set_index(first.index.values)
first['incidence'] = second
first = first.loc[:, (first != 0).any(axis=0)]
new_pages = list(first.columns.values)
print(first)
correlation_matrix(first, "Correlation Matrix "+str(year_sel[0]-1)+"-"+str(year_sel[1]-1), new_pages, save_directory+"cmatrix_"+str(year_sel[0]-1)+"-"+str(year_sel[1]-1)+".png")<commit_msg>[model] Update the model to write information to file (not on screen).<commit_after>#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""Machine learning model which uses Wikipedia data to predicts ILI levels.

Usage:
  model.py <year_start> <year_end> <dataset_path> <incidence_path> <keywords_file> [--p] [--f] [--v] [--d=<directory>] [--no-future]

  <year_start>      The first influenza season we want to predict.
  <year_end>        The last influenza season we want to predict.
  <dataset_path>    The path to dataset files (the ones which contain Wikipedia pageviews).
  <incidence_path>  The path to the files which specify the ILI levels for each week.
  <keywords_file>   The path to the file which contains the Wikipedia pages used.
  -p, --poisson     Use the Poisson model + LASSO instead of the linear one.
  -f, --file        Write informations to file
  -v, --verbose     Output more informations
  -d, --directory   Select a directory in which save your files
  -n, --no-future   Use a different method to train the model (avoid using seasonal influenza which are
                    later than the one we want to predict)
  -h, --help        Print this help message
"""

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import csv
from sklearn.linear_model import LassoCV
from sklearn.metrics import mean_squared_error
from tabulate import tabulate
from docopt import docopt

from models_utils import generate, generate_labels, generate_labels_one_year, generate_one_year, standardize_data, generate_keywords, stz, get_important_pages, correlation_matrix

from cvglmnetCoef import cvglmnetCoef
from cvglmnetPredict import cvglmnetPredict
from cvglmnet import cvglmnet

#########################
#### INITIALIZATION #####
#########################

# Parse the command line
arguments = docopt(__doc__)

# Directory where to save files
save_directory="./"
if arguments["--d"]:
    save_directory = arguments["--d"]

# Year selected
year_start = int(arguments["<year_start>"])
year_end = int(arguments["<year_end>"])+1

# Model type
model_type = "Linear Model"
using_poisson = False
if arguments["--p"]:
    using_poisson = True
    model_type="Poisson Model"

# Feature path and labels
path_features = arguments["<dataset_path>"]
path_labels = arguments["<incidence_path>"]
keywords = arguments["<keywords_file>"]

# Selected columns that will be extracted from the dataframes
selected_columns = generate_keywords(keywords)

all_features_values = pd.DataFrame()
all_predicted_values = []
all_true_labels = pd.DataFrame()
total_weeks = []
all_weighted_feature=dict()

# The pair start-end seasons
year_sel = (year_start, year_end)

# Useful information which can be written to file
peaks = []
predicted_peaks = []
peaks_value = []
predicted_peaks_value = []
mse_seasons = []

######################
##### ALGORITHM ######
######################

# For each season, train a model with the complete dataset (without the season)
# and predict the ILI incidence.
for year_selected in range(year_sel[0], year_sel[1]):

    print("------------")
    print("[*] ", year_selected )

    # Generate the list of excluded years
    excluded=[]
    if arguments["--no-future"]:
        for i in range(2007, 2018):
            if i>year_selected:
                excluded.append(i)

    # Data generation from data files
    dataset = generate(year_selected, excluded, path_features)[selected_columns]
    labels = generate_labels(year_selected, excluded, path_labels)
    data = generate_one_year(year_selected, path_features)[selected_columns]
    labels_test = generate_labels_one_year(year_selected, path_labels)
    weeks = labels_test["week"]

    # Create a copy of the dataset where we fills NaN values
    # with 0.
    dataset_zero = dataset.fillna(0)
    data_zero = data.fillna(0)

    # Standardize data
    train, test = standardize_data(dataset_zero, data_zero)

    # Create a Lasso Cross-Validation instance which will be
    # trained on the dataset in which NaN values are replaced
    # with the column mean.
    if not using_poisson:
        model = LassoCV(max_iter=100000, n_jobs=-1, n_alphas=1000, random_state=1)
        model.fit(train, labels["incidence"].fillna(0))
        result = model.predict(test)
    else:
        model = cvglmnet(x=train.as_matrix().copy(), y=labels["incidence"].fillna(0).copy().as_matrix(), family='poisson', alpha=1.0,
                       ptype="mse", parallel=True, nfolds=10)
        result = cvglmnetPredict(model, test.as_matrix(), ptype='response', s="lambda_min")
        coeff = cvglmnetCoef(model, s="lambda_min")

    # Get the feature coefficients
    if not using_poisson:
        coeff = model.coef_
    else:
        coeff = cvglmnetCoef(model, s="lambda_min")

    # Add the pair (value, coeff) to a dictionary
    for i in list(zip(coeff, selected_columns)):
        if (i[0] != 0):
            if all_weighted_feature.get(i[1], None) is None:
                all_weighted_feature[i[1]] = [i[0]]
            else:
                all_weighted_feature[i[1]].append(i[0])

    # Add the value to global variables
    all_features_values= all_features_values.append(data.fillna(0))
    all_predicted_values.extend(result)
    all_true_labels = all_true_labels.append(labels_test)
    total_weeks.extend(weeks)

    mse_seasons.append(mean_squared_error(labels_test['incidence'].fillna(0), result))

    # If we are verbose print more informations
    index_peak = labels_test['incidence'].fillna(0).idxmax(axis=1)
    index_peak_m = np.argmax(result)
    peaks.append(labels_test.iloc[index_peak]["week"])
    predicted_peaks.append(labels_test.iloc[index_peak_m]["week"])
    peaks_value.append(labels_test.iloc[index_peak]["incidence"])
    predicted_peaks_value.append(result[index_peak_m])
    if arguments["--v"]:
        print("[*] Influenza Season Peak (Week):", labels_test.iloc[index_peak]["week"])
        print("[*] Influenza Seasons Predicted Peak (Week): ", labels_test.iloc[index_peak_m]["week"])
        print("[*] Influenza Season Peak Value: ", labels_test.iloc[index_peak]["incidence"])
        print("[*] Influenza Season Predicted Peak Value: ", result[index_peak_m])
        print("------------")

# Get important pages to generate the plot we need
important_pages = get_important_pages(all_weighted_feature, 5)

# Print MSE of the two models
mse = mean_squared_error(all_true_labels["incidence"].fillna(0), all_predicted_values)
print("------------")
print("MSE: ", mse)
print("------------")

# Print Pearson Coefficient
pcc = np.corrcoef(all_predicted_values, all_true_labels["incidence"].fillna(0), rowvar=False)[0][1]
print("------------")
print("Pearson Correlation Coeff: ", pcc)
print("------------")

# Print important pages
print(tabulate(important_pages, headers=["Page name", "Weight"]))

###########################
#### GRAPHS GENERATION ####
###########################

# Plot some informations
fig = plt.figure(3, figsize=(15, 6))
ax = fig.add_subplot(111)

# Set up the graph title, x and y axis titles
plt.title("Influenza Seasons "+str(year_sel[0]-1)+" - "+str(year_sel[1]-1), fontsize=18)
plt.ylabel("Incidence on 100000 people", fontsize=17)
plt.xlabel("Year-Week", fontsize=17)

# Generate the x axis labels in such way to
# have the year-week pair only every two ticks.
weeks_used = []
for k, v in enumerate(total_weeks):
    if k%2==0:
        weeks_used.append(v)
    else:
        weeks_used.append(" ")

# Set up the axes ticks
plt.xticks(range(0, len(weeks_used)), weeks_used, rotation="vertical", fontsize=15)
plt.yticks(fontsize=15)

# Plot the model result and the incidence
plt.plot(range(0, len(all_predicted_values)), all_predicted_values, 'r-', label=model_type, linewidth=3)
plt.plot(range(0, len(all_true_labels["incidence"])), all_true_labels["incidence"].fillna(0), 'k-', label="Incidence", linewidth=3)

# Add dotted line into the graph to delimit the influenza seasons.
for i in range(1, (year_sel[1]-year_sel[0])):
    plt.axvline(26*i, linestyle="dashed", color="k", linewidth=3)

# Update some graphical options
plt.grid()
plt.legend(fontsize=16, loc="upper right")
plt.tight_layout()

# Save the graph on file
plt.savefig(save_directory+"season_"+str(year_sel[0]-1)+"-"+str(year_sel[1]-1)+".png", dpi=150)

# Feature plot
figf = plt.figure(4, figsize=(15, 6))
axf = fig.add_subplot(111)

# Generate the graph showing the pageview variation of the top 5 Wikipedia's pages
# choosen by the model.
plt.title("Pageview variation "+str(year_sel[0]-1)+" - "+str(year_sel[1]-1), fontsize=18)
plt.ylabel("Pageview value", fontsize=17)
plt.xlabel("Year-Week", fontsize=17)

# Plot the axes labels
plt.xticks(range(0, len(weeks_used)), weeks_used, rotation="vertical", fontsize=15)
plt.yticks(fontsize=15)

# Plot all the pageview data and the ILI incidence
std_all_features_values = stz(all_features_values)
for key, value in important_pages:
 plt.plot(range(0, len(std_all_features_values[key])), std_all_features_values[key], '-', label=key, linewidth=3)
plt.plot(range(0, len(all_true_labels["incidence"])), stz(all_true_labels["incidence"]), 'k-', label="Incidence", linewidth=3)

# Add dotted line into the graph to delimit the influenza seasons.
for i in range(1, (year_sel[1]-year_sel[0])):
    plt.axvline(26*i, linestyle="dashed", color="k", linewidth=3)

# Update some graphical options
plt.grid()
plt.legend(fontsize=16, loc="best")
plt.tight_layout()

# Save again the graph on file
plt.savefig(save_directory+"season_"+str(year_sel[0]-1)+"-"+str(year_sel[1]-1)+"_features.png", dpi=150)

# If --f then we write some information to file
if arguments["--f"]:
    with open(save_directory+str(year_sel[0]-1)+"-"+str(year_sel[1]-1)+"_information.csv", 'w', newline='') as csvfile:
        spamwriter = csv.writer(csvfile, delimiter=',',
                                quotechar='|', quoting=csv.QUOTE_MINIMAL)
        spamwriter.writerow(['real_influenza_peak', peaks])
        spamwriter.writerow(['predicted_influenza_peak', predicted_peaks])
        spamwriter.writerow(['real_influenza_peak_value', peaks_value])
        spamwriter.writerow(['predicted_influenza_peak_value', predicted_peaks_value])
        spamwriter.writerow(['mse_all_seasons', mse_seasons])
        spamwriter.writerow(['mse', mean_squared_error(all_true_labels["incidence"].fillna(0), all_predicted_values)])
        spamwriter.writerow(['pearson', pcc])
        for p in important_pages:
            spamwriter.writerow([p[0], float(p[1])])

# Correlation Matrix
new_pages=[]
for p in important_pages:
    new_pages.append(p[0])
first = pd.DataFrame(all_features_values[new_pages])
second = pd.DataFrame(all_true_labels["incidence"]).set_index(first.index.values)
first['incidence'] = second
first = first.loc[:, (first != 0).any(axis=0)]
new_pages = list(first.columns.values)
correlation_matrix(first, "Correlation Matrix "+str(year_sel[0]-1)+"-"+str(year_sel[1]-1), new_pages, save_directory+"cmatrix_"+str(year_sel[0]-1)+"-"+str(year_sel[1]-1)+".png")<|endoftext|>