<commit_before>import numpy as np
import utils
from skimage import morphology

def extract_noise_part(spectrogram):
    """ Extract the noise part of a spectrogram
    """
    mask = compute_noise_mask(spectrogram)
    noise_part = extract_masked_part(mask, spectrogram)
    return noise_part

def extract_signal_part(spectrogram):
    """ Extract the signal part of a spectrogram
    """
    mask = compute_signal_mask(spectrogram)
    signal_part = extract_masked_part(mask, spectrogram)
    return signal_part

def extract_masked_part(mask, spectrogram):
    """ Extract the masked part of the spectrogram
    """
    # the mask is binary and numpy arrays use boolean masks
    mask = mask==1
    return spectrogram[:,mask]

def compute_signal_mask(spectrogram):
    """ Computes a binary noise mask (convenience function)

    # Arguments
        spectrogram : a numpy array representation of a spectrogram (2-dim)

    # Returns
        binary_mask : the binary signal mask
    """
    threshold = 3
    mask = compute_binary_mask(spectrogram, threshold)
    return mask

def compute_noise_mask(spectrogram):
    """ Computes a binary noise mask (convenience function)

    # Arguments
        spectrogram : a numpy array representation of a spectrogram (2-dim)

    # Returns
        binary_mask : the binary noise mask
    """
    threshold = 2.5
    mask = compute_binary_mask(spectrogram, threshold)
    # invert mask
    return 1-mask

def compute_binary_mask(spectrogram, threshold):
    """ Computes a binary mask for the spectrogram
    # Arguments
        spectrogram : a numpy array representation of a spectrogram (2-dim)
        threshold   : a threshold for times larger than the median
    # Returns
        binary_mask : the binary mask
    """
    norm_spectrogram = normalize(spectrogram)
    binary_structure = mark_cells_times_larger_than_median(norm_spectrogram, threshold)

    n_hood = np.ones((4, 4))
    binary_structure = morphology.binary_erosion(binary_structure, n_hood)
    binary_structure = morphology.binary_dilation(binary_structure, n_hood)

    # TODO: transpose is O(n^2)
    mask = np.array([np.max(col) for col in np.transpose(binary_structure)])
    mask = smooth_mask(mask)

    return mask

def smooth_mask(mask):
    """ Smooths a binary mask using 4x4 dilation
        # Arguments
            mask : the binary mask
        # Returns
            mask : a smoother binary mask
    """
    n_hood = np.ones(4)
    mask = morphology.binary_dilation(mask, n_hood)
    mask = morphology.binary_dilation(mask, n_hood)

    # type casting is a bitch
    return 1*mask

def mark_cells_times_larger_than_median(spectrogram, number_times_larger):
    """ Compute binary image from spectrogram where cells are marked as 1 if
    number_times_larger than the row AND column median, otherwise 0
    """
    row_medians = np.median(spectrogram, axis=1)
    col_medians = np.median(spectrogram, axis=0)
    for row in range(spectrogram.shape[0]):
        for col in range(spectrogram.shape[1]):
            if spectrogram[row][col] > 3*row_medians[row] \
               and spectrogram[row][col] > 3*col_medians[col]:
                spectrogram[row][col] = 1
            else:
                spectrogram[row][col] = 0

    return spectrogram

def remove_low_and_high_frequency_bins(spectrogram, nb_low=5, nb_high=25):
    """ Removes low and high frequency bins from a spectrogram

    # Argmuents
        nb_low  : the number of lower frequency bins to remove
        nb_high : the number of higher frequency bins to remove
    # Returns
        spectrogram : the narrowed spectrogram
    """
    return spectrogram[nb_low:spectrogram.shape[0]-nb_high]

def normalize(X):
    """ Normalize numpy array to interval [0, 1]
    """
    mi = np.min(X)
    ma = np.max(X)

    X = (X-mi)/(ma-mi)
    return X
<commit_msg>Change to dilation, erosion, dilation, and remove objs<commit_after>import numpy as np
import utils
from skimage import morphology
import skimage.filters as filters

def extract_noise_part(spectrogram):
    """ Extract the noise part of a spectrogram
    """
    mask = compute_noise_mask(spectrogram)
    noise_part = extract_masked_part(mask, spectrogram)
    return noise_part

def extract_signal_part(spectrogram):
    """ Extract the signal part of a spectrogram
    """
    mask = compute_signal_mask(spectrogram)
    signal_part = extract_masked_part(mask, spectrogram)
    return signal_part

def extract_masked_part(mask, spectrogram):
    """ Extract the masked part of the spectrogram
    """
    # the mask is binary and numpy arrays use boolean masks
    mask = mask==1
    return spectrogram[:,mask]

def compute_signal_mask(spectrogram):
    """ Computes a binary noise mask (convenience function)

    # Arguments
        spectrogram : a numpy array representation of a spectrogram (2-dim)

    # Returns
        binary_mask : the binary signal mask
    """
    threshold = 3
    mask = compute_binary_mask(spectrogram, threshold)
    return mask

def compute_noise_mask(spectrogram):
    """ Computes a binary noise mask (convenience function)

    # Arguments
        spectrogram : a numpy array representation of a spectrogram (2-dim)

    # Returns
        binary_mask : the binary noise mask
    """
    threshold = 2.5
    mask = compute_binary_mask(spectrogram, threshold)
    # invert mask
    return 1-mask

def compute_binary_mask(spectrogram, threshold):
    """ Computes a binary mask for the spectrogram
    # Arguments
        spectrogram : a numpy array representation of a spectrogram (2-dim)
        threshold   : a threshold for times larger than the median
    # Returns
        binary_mask : the binary mask
    """
    norm_spectrogram = normalize(spectrogram)
    binary_image = mark_cells_times_larger_than_median(norm_spectrogram, threshold)

    utils.plot_matrix(binary_image, "Median Clipping")

    # closing binary image (dilation followed by erosion)
    binary_image = morphology.binary_closing(binary_image, selem=np.ones((4, 4)))
    # dialate binary image
    binary_image = morphology.binary_dilation(binary_image, selem=np.ones((4, 4)))
    utils.plot_matrix(binary_image, "Closing and Dilation")
    # apply median filter
    #binary_image = filters.median(binary_image, selem=np.ones((2, 2)))
    # remove small objects
    binary_image = morphology.remove_small_objects(binary_image, min_size=32,
                                                   connectivity=1)

    utils.plot_matrix(binary_image, "Median Filter and Small Objects Removed")

    # TODO: transpose is O(n^2)
    mask = np.array([np.max(col) for col in np.transpose(binary_image)])
    mask = smooth_mask(mask)

    return mask

def smooth_mask(mask):
    """ Smooths a binary mask using 4x4 dilation
        # Arguments
            mask : the binary mask
        # Returns
            mask : a smoother binary mask
    """
    n_hood = np.ones(4)
    mask = morphology.binary_dilation(mask, n_hood)
    mask = morphology.binary_dilation(mask, n_hood)

    # type casting is a bitch
    return 1*mask

def mark_cells_times_larger_than_median(spectrogram, number_times_larger):
    """ Compute binary image from spectrogram where cells are marked as 1 if
    number_times_larger than the row AND column median, otherwise 0
    """
    row_medians = np.median(spectrogram, axis=1)
    col_medians = np.median(spectrogram, axis=0)
    for row in range(spectrogram.shape[0]):
        for col in range(spectrogram.shape[1]):
            if spectrogram[row][col] > 3*row_medians[row] \
               and spectrogram[row][col] > 3*col_medians[col]:
                spectrogram[row][col] = 1
            else:
                spectrogram[row][col] = 0

    return spectrogram

def remove_low_and_high_frequency_bins(spectrogram, nb_low=5, nb_high=25):
    """ Removes low and high frequency bins from a spectrogram

    # Argmuents
        nb_low  : the number of lower frequency bins to remove
        nb_high : the number of higher frequency bins to remove
    # Returns
        spectrogram : the narrowed spectrogram
    """
    return spectrogram[nb_low:spectrogram.shape[0]-nb_high]

def normalize(X):
    """ Normalize numpy array to interval [0, 1]
    """
    mi = np.min(X)
    ma = np.max(X)

    X = (X-mi)/(ma-mi)
    return X
<|endoftext|>


<commit_before>#!/usr/bin/env python2
# IBM_PROLOG_BEGIN_TAG
# This is an automatically generated prolog.
#
# $Source: op-test-framework/testcases/OpTestOpenCAPI.py $
#
# OpenPOWER Automated Test Project
#
# Contributors Listed Below - COPYRIGHT 2018
# [+] International Business Machines Corp.
#
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
# implied. See the License for the specific language governing
# permissions and limitations under the License.
#
# IBM_PROLOG_END_TAG

'''
OpTestOpenCAPI
--------------

OpenCAPI tests for OpenPower testing.

This class will test the functionality of OpenCAPI

Prerequisites:

1. Host must have a OpenCAPI FPGA card
2. OpenCAPI card must have been flashed with memcpy3 AFU
'''

import time
import subprocess
import re

import unittest

import OpTestConfiguration
from common.OpTestSystem import OpSystemState
from common.OpTestError import OpTestError
from common.Exceptions import CommandFailed

import logging
import OpTestLogger
log = OpTestLogger.optest_logger_glob.get_logger(__name__)

class OpTestOpenCAPI(unittest.TestCase):
    def setUp(self):
        conf = OpTestConfiguration.conf
        self.cv_IPMI = conf.ipmi()
        self.cv_SYSTEM = conf.system()
        self.cv_HOST = conf.host()

    def set_up(self):
        self.cv_SYSTEM.goto_state(OpSystemState.OS)
        # Check that host has a OpenCAPI FPGA card
        if (self.cv_HOST.host_has_opencapi_fpga_card() != True):
            raise unittest.SkipTest("No OpenCAPI card available on host \
                                     Skipping the OpenCAPI tests")

        self.cv_HOST.host_check_command("git", "gcc")

        # Get OS level
        l_oslevel = self.cv_HOST.host_get_OS_Level()

        # Get kernel version
        l_kernel = self.cv_HOST.host_get_kernel_version()

        # Load module ocxl based on config option
        l_config = "CONFIG_OCXL"
        l_module = "ocxl"
        self.cv_HOST.host_load_module_based_on_config(l_kernel, l_config, \
                                                   l_module)

class OcxlDeviceFileTest(OpTestOpenCAPI, unittest.TestCase):
    '''
    If the system has a OpenCAPI FPGA card, then this test load the ocxl module
    if required and check that the ocxl device files afu0.0m and afu0.0s exist
    '''
    def setUp(self):
        super(OcxlDeviceFileTest, self).setUp()

    def runTest(self):
        self.set_up()
        # Check device files /dev/ocxl/IBM,MEMCPY3.* existence
        l_cmd = "ls -l /dev/ocxl/IBM,MEMCPY3.*"
        try:
            self.cv_HOST.host_run_command(l_cmd)
        except CommandFailed:
            self.assertTrue(False, "ocxl device file tests fail")

class MemCpy3AFUTest(OpTestOpenCAPI, unittest.TestCase):
    '''
    If the system has a OpenCAPI FPGA card, then this test load the ocxl module
    if required and test the memcpy3 AFU with ocxl_memcpy
    '''
    def setUp(self):
        super(MemCpy3AFUTest, self).setUp()

    def runTest(self):
        self.set_up()

        # Check that the afutests binary are available
        # If not, clone and build libocxl and afutests
        l_dir = "/tmp/libocxl"
        if (self.cv_HOST.host_check_binary(l_dir, "afuobj/ocxl_memcpy") != True):
            self.cv_HOST.host_clone_libocxl(l_dir)
            self.cv_HOST.host_build_libocxl(l_dir)

        # Run memcpy3 afu tests
        l_exec = "afuobj/ocxl_memcpy -p0 -l10000 >/tmp/ocxl_memcpy.log"
        cmd = "cd %s; ./%s" % (l_dir, l_exec)
        log.debug(cmd)
        try:
            self.cv_HOST.host_run_command(cmd)
            l_msg = "ocxl_memcpy tests pass"
            log.debug(l_msg)
        except CommandFailed:
            self.assertTrue(False, "ocxl_memcpy tests failed")

class MemCpy3AFUReallocTest(OpTestOpenCAPI, unittest.TestCase):
    '''
    If the system has a OpenCAPI FPGA card, then load the ocxl module
    if required and test the memcpy3 AFU with ocxl_memcpy -r
    '''
    def setUp(self):
        super(MemCpy3AFUReallocTest, self).setUp()

    def runTest(self):
        self.set_up()

        # Check that the afutests binary are available
        # If not, clone and build libocxl and afutests
        l_dir = "/tmp/libocxl"
        if (self.cv_HOST.host_check_binary(l_dir, "afuobj/ocxl_memcpy") != True):
            self.cv_HOST.host_clone_libocxl(l_dir)
            self.cv_HOST.host_build_libocxl(l_dir)

        # Run memcpy3 afu tests
        l_exec = "afuobj/ocxl_memcpy -p0 -r -l3000 >/tmp/ocxl_memcpy-r.log"
        cmd = "cd %s; ./%s" % (l_dir, l_exec)
        log.debug(cmd)
        try:
            self.cv_HOST.host_run_command(cmd)
            l_msg = "ocxl_memcpy -r tests pass"
            log.debug(l_msg)
        except CommandFailed:
            self.assertTrue(False, "ocxl_memcpy -r tests failed")

def opencapi_test_suite():
    s = unittest.TestSuite()
    s.addTest(OcxlDeviceFileTest())
    s.addTest(MemCpy3AFUTest())
    s.addTest(MemCpy3AFUReallocTest())
    return s
<commit_msg>OpTestOpenCAPI: new memcpy test MemCpy3AFUIrqTest<commit_after>#!/usr/bin/env python2
# IBM_PROLOG_BEGIN_TAG
# This is an automatically generated prolog.
#
# $Source: op-test-framework/testcases/OpTestOpenCAPI.py $
#
# OpenPOWER Automated Test Project
#
# Contributors Listed Below - COPYRIGHT 2018
# [+] International Business Machines Corp.
#
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
# implied. See the License for the specific language governing
# permissions and limitations under the License.
#
# IBM_PROLOG_END_TAG

'''
OpTestOpenCAPI
--------------

OpenCAPI tests for OpenPower testing.

This class will test the functionality of OpenCAPI

Prerequisites:

1. Host must have a OpenCAPI FPGA card
2. OpenCAPI card must have been flashed with memcpy3 AFU
'''

import time
import subprocess
import re

import unittest

import OpTestConfiguration
from common.OpTestSystem import OpSystemState
from common.OpTestError import OpTestError
from common.Exceptions import CommandFailed

import logging
import OpTestLogger
log = OpTestLogger.optest_logger_glob.get_logger(__name__)

class OpTestOpenCAPI(unittest.TestCase):
    def setUp(self):
        conf = OpTestConfiguration.conf
        self.cv_IPMI = conf.ipmi()
        self.cv_SYSTEM = conf.system()
        self.cv_HOST = conf.host()

    def set_up(self):
        self.cv_SYSTEM.goto_state(OpSystemState.OS)
        # Check that host has a OpenCAPI FPGA card
        if (self.cv_HOST.host_has_opencapi_fpga_card() != True):
            raise unittest.SkipTest("No OpenCAPI card available on host \
                                     Skipping the OpenCAPI tests")

        self.cv_HOST.host_check_command("git", "gcc")

        # Get OS level
        l_oslevel = self.cv_HOST.host_get_OS_Level()

        # Get kernel version
        l_kernel = self.cv_HOST.host_get_kernel_version()

        # Load module ocxl based on config option
        l_config = "CONFIG_OCXL"
        l_module = "ocxl"
        self.cv_HOST.host_load_module_based_on_config(l_kernel, l_config, \
                                                   l_module)

class OcxlDeviceFileTest(OpTestOpenCAPI, unittest.TestCase):
    '''
    If the system has a OpenCAPI FPGA card, then this test load the ocxl module
    if required and check that the ocxl device files afu0.0m and afu0.0s exist
    '''
    def setUp(self):
        super(OcxlDeviceFileTest, self).setUp()

    def runTest(self):
        self.set_up()
        # Check device files /dev/ocxl/IBM,MEMCPY3.* existence
        l_cmd = "ls -l /dev/ocxl/IBM,MEMCPY3.*"
        try:
            self.cv_HOST.host_run_command(l_cmd)
        except CommandFailed:
            self.assertTrue(False, "ocxl device file tests fail")

class MemCpy3AFUTest(OpTestOpenCAPI, unittest.TestCase):
    '''
    If the system has a OpenCAPI FPGA card, then this test load the ocxl module
    if required and test the memcpy3 AFU with ocxl_memcpy
    '''
    def setUp(self):
        super(MemCpy3AFUTest, self).setUp()

    def runTest(self):
        self.set_up()

        # Check that the afutests binary are available
        # If not, clone and build libocxl and afutests
        l_dir = "/tmp/libocxl"
        if (self.cv_HOST.host_check_binary(l_dir, "afuobj/ocxl_memcpy") != True):
            self.cv_HOST.host_clone_libocxl(l_dir)
            self.cv_HOST.host_build_libocxl(l_dir)

        # Run memcpy3 afu tests
        l_exec = "afuobj/ocxl_memcpy -p0 -l10000 >/tmp/ocxl_memcpy.log"
        cmd = "cd %s; ./%s" % (l_dir, l_exec)
        log.debug(cmd)
        try:
            self.cv_HOST.host_run_command(cmd)
            l_msg = "ocxl_memcpy tests pass"
            log.debug(l_msg)
        except CommandFailed:
            self.assertTrue(False, "ocxl_memcpy tests failed")

class MemCpy3AFUIrqTest(OpTestOpenCAPI, unittest.TestCase):
    '''
    If the system has a OpenCAPI FPGA card, then load the ocxl module
    if required and test the memcpy3 AFU with ocxl_memcpy -i
    '''
    def setUp(self):
        super(MemCpy3AFUIrqTest, self).setUp()

    def runTest(self):
        self.set_up()

        # Check that the afutests binary are available
        # If not, clone and build libocxl and afutests
        l_dir = "/tmp/libocxl"
        if (self.cv_HOST.host_check_binary(l_dir, "afuobj/ocxl_memcpy") != True):
            self.cv_HOST.host_clone_libocxl(l_dir)
            self.cv_HOST.host_build_libocxl(l_dir)

        # Run memcpy3 afu tests
        l_exec = "afuobj/ocxl_memcpy -p0 -i -l10000 >/tmp/ocxl_memcpy-i.log"
        cmd = "cd %s; ./%s" % (l_dir, l_exec)
        log.debug(cmd)
        try:
            self.cv_HOST.host_run_command(cmd)
            l_msg = "ocxl_memcpy -i tests pass"
            log.debug(l_msg)
        except CommandFailed:
            self.assertTrue(False, "ocxl_memcpy -i tests failed")

class MemCpy3AFUReallocTest(OpTestOpenCAPI, unittest.TestCase):
    '''
    If the system has a OpenCAPI FPGA card, then load the ocxl module
    if required and test the memcpy3 AFU with ocxl_memcpy -r
    '''
    def setUp(self):
        super(MemCpy3AFUReallocTest, self).setUp()

    def runTest(self):
        self.set_up()

        # Check that the afutests binary are available
        # If not, clone and build libocxl and afutests
        l_dir = "/tmp/libocxl"
        if (self.cv_HOST.host_check_binary(l_dir, "afuobj/ocxl_memcpy") != True):
            self.cv_HOST.host_clone_libocxl(l_dir)
            self.cv_HOST.host_build_libocxl(l_dir)

        # Run memcpy3 afu tests
        l_exec = "afuobj/ocxl_memcpy -p0 -r -l3000 >/tmp/ocxl_memcpy-r.log"
        cmd = "cd %s; ./%s" % (l_dir, l_exec)
        log.debug(cmd)
        try:
            self.cv_HOST.host_run_command(cmd)
            l_msg = "ocxl_memcpy -r tests pass"
            log.debug(l_msg)
        except CommandFailed:
            self.assertTrue(False, "ocxl_memcpy -r tests failed")

def opencapi_test_suite():
    s = unittest.TestSuite()
    s.addTest(OcxlDeviceFileTest())
    s.addTest(MemCpy3AFUTest())
    s.addTest(MemCpy3AFUIrqTest())
    s.addTest(MemCpy3AFUReallocTest())
    return s
<|endoftext|>


<commit_before>"""Support for Xiaomi Miio."""
from __future__ import annotations

from dataclasses import dataclass
from datetime import timedelta
import logging

import async_timeout
from miio import (
    AirFresh,
    AirFreshA1,
    AirFreshT2017,
    AirHumidifier,
    AirHumidifierMiot,
    AirHumidifierMjjsq,
    AirPurifier,
    AirPurifierMiot,
    CleaningDetails,
    CleaningSummary,
    ConsumableStatus,
    Device as MiioDevice,
    DeviceException,
    DNDStatus,
    Fan,
    Fan1C,
    FanMiot,
    FanP5,
    FanZA5,
    RoborockVacuum,
    Timer,
    VacuumStatus,
)
from miio.gateway.gateway import GatewayException

from homeassistant.config_entries import ConfigEntry
from homeassistant.const import CONF_HOST, CONF_MODEL, CONF_TOKEN, Platform
from homeassistant.core import HomeAssistant, callback
from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady
from homeassistant.helpers import device_registry as dr, entity_registry as er
from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed

from .const import (
    ATTR_AVAILABLE,
    CONF_DEVICE,
    CONF_FLOW_TYPE,
    CONF_GATEWAY,
    DOMAIN,
    KEY_COORDINATOR,
    KEY_DEVICE,
    MODEL_AIRFRESH_A1,
    MODEL_AIRFRESH_T2017,
    MODEL_FAN_1C,
    MODEL_FAN_P5,
    MODEL_FAN_P9,
    MODEL_FAN_P10,
    MODEL_FAN_P11,
    MODEL_FAN_ZA5,
    MODELS_AIR_MONITOR,
    MODELS_FAN,
    MODELS_FAN_MIIO,
    MODELS_HUMIDIFIER,
    MODELS_HUMIDIFIER_MIIO,
    MODELS_HUMIDIFIER_MIOT,
    MODELS_HUMIDIFIER_MJJSQ,
    MODELS_LIGHT,
    MODELS_PURIFIER_MIOT,
    MODELS_SWITCH,
    MODELS_VACUUM,
    ROBOROCK_GENERIC,
    ROCKROBO_GENERIC,
    AuthException,
    SetupException,
)
from .gateway import ConnectXiaomiGateway

_LOGGER = logging.getLogger(__name__)

POLLING_TIMEOUT_SEC = 10
UPDATE_INTERVAL = timedelta(seconds=15)

GATEWAY_PLATFORMS = [
    Platform.ALARM_CONTROL_PANEL,
    Platform.LIGHT,
    Platform.SENSOR,
    Platform.SWITCH,
]
SWITCH_PLATFORMS = [Platform.SWITCH]
FAN_PLATFORMS = [
    Platform.BINARY_SENSOR,
    Platform.BUTTON,
    Platform.FAN,
    Platform.NUMBER,
    Platform.SELECT,
    Platform.SENSOR,
    Platform.SWITCH,
]
HUMIDIFIER_PLATFORMS = [
    Platform.BINARY_SENSOR,
    Platform.HUMIDIFIER,
    Platform.NUMBER,
    Platform.SELECT,
    Platform.SENSOR,
    Platform.SWITCH,
]
LIGHT_PLATFORMS = [Platform.LIGHT]
VACUUM_PLATFORMS = [Platform.BINARY_SENSOR, Platform.SENSOR, Platform.VACUUM]
AIR_MONITOR_PLATFORMS = [Platform.AIR_QUALITY, Platform.SENSOR]

MODEL_TO_CLASS_MAP = {
    MODEL_FAN_1C: Fan1C,
    MODEL_FAN_P9: FanMiot,
    MODEL_FAN_P10: FanMiot,
    MODEL_FAN_P11: FanMiot,
    MODEL_FAN_P5: FanP5,
    MODEL_FAN_ZA5: FanZA5,
}

async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
    """Set up the Xiaomi Miio components from a config entry."""
    hass.data.setdefault(DOMAIN, {})
    if entry.data[CONF_FLOW_TYPE] == CONF_GATEWAY:
        await async_setup_gateway_entry(hass, entry)
        return True

    return bool(
        entry.data[CONF_FLOW_TYPE] != CONF_DEVICE
        or await async_setup_device_entry(hass, entry)
    )

@callback
def get_platforms(config_entry):
    """Return the platforms belonging to a config_entry."""
    model = config_entry.data[CONF_MODEL]
    flow_type = config_entry.data[CONF_FLOW_TYPE]

    if flow_type == CONF_GATEWAY:
        return GATEWAY_PLATFORMS
    if flow_type == CONF_DEVICE:
        if model in MODELS_SWITCH:
            return SWITCH_PLATFORMS
        if model in MODELS_HUMIDIFIER:
            return HUMIDIFIER_PLATFORMS
        if model in MODELS_FAN:
            return FAN_PLATFORMS
        if model in MODELS_LIGHT:
            return LIGHT_PLATFORMS
        for vacuum_model in MODELS_VACUUM:
            if model.startswith(vacuum_model):
                return VACUUM_PLATFORMS
        for air_monitor_model in MODELS_AIR_MONITOR:
            if model.startswith(air_monitor_model):
                return AIR_MONITOR_PLATFORMS
    _LOGGER.error(
        "Unsupported device found! Please create an issue at "
        "https://github.com/syssi/xiaomi_airpurifier/issues "
        "and provide the following data: %s",
        model,
    )
    return []

def _async_update_data_default(hass, device):
    async def update():
        """Fetch data from the device using async_add_executor_job."""

        async def _async_fetch_data():
            """Fetch data from the device."""
            async with async_timeout.timeout(POLLING_TIMEOUT_SEC):
                state = await hass.async_add_executor_job(device.status)
                _LOGGER.debug("Got new state: %s", state)
                return state

        try:
            return await _async_fetch_data()
        except DeviceException as ex:
            if getattr(ex, "code", None) != -9999:
                raise UpdateFailed(ex) from ex
            _LOGGER.info("Got exception while fetching the state, trying again: %s", ex)
        # Try to fetch the data a second time after error code -9999
        try:
            return await _async_fetch_data()
        except DeviceException as ex:
            raise UpdateFailed(ex) from ex

    return update

@dataclass(frozen=True)
class VacuumCoordinatorData:
    """A class that holds the vacuum data retrieved by the coordinator."""

    status: VacuumStatus
    dnd_status: DNDStatus
    last_clean_details: CleaningDetails
    consumable_status: ConsumableStatus
    clean_history_status: CleaningSummary
    timers: list[Timer]
    fan_speeds: dict[str, int]
    fan_speeds_reverse: dict[int, str]

@dataclass(init=False, frozen=True)
class VacuumCoordinatorDataAttributes:
    """
    A class that holds attribute names for VacuumCoordinatorData.

    These attributes can be used in methods like `getattr` when a generic solutions is
    needed.
    See homeassistant.components.xiaomi_miio.device.XiaomiCoordinatedMiioEntity
    ._extract_value_from_attribute for
    an example.
    """

    status: str = "status"
    dnd_status: str = "dnd_status"
    last_clean_details: str = "last_clean_details"
    consumable_status: str = "consumable_status"
    clean_history_status: str = "clean_history_status"
    timer: str = "timer"
    fan_speeds: str = "fan_speeds"
    fan_speeds_reverse: str = "fan_speeds_reverse"

def _async_update_data_vacuum(hass, device: RoborockVacuum):
    def update() -> VacuumCoordinatorData:
        timer = []

        # See https://github.com/home-assistant/core/issues/38285 for reason on
        # Why timers must be fetched separately.
        try:
            timer = device.timer()
        except DeviceException as ex:
            _LOGGER.debug(
                "Unable to fetch timers, this may happen on some devices: %s", ex
            )

        fan_speeds = device.fan_speed_presets()

        data = VacuumCoordinatorData(
            device.status(),
            device.dnd_status(),
            device.last_clean_details(),
            device.consumable_status(),
            device.clean_history(),
            timer,
            fan_speeds,
            {v: k for k, v in fan_speeds.items()},
        )

        return data

    async def update_async():
        """Fetch data from the device using async_add_executor_job."""

        async def execute_update():
            async with async_timeout.timeout(POLLING_TIMEOUT_SEC):
                state = await hass.async_add_executor_job(update)
                _LOGGER.debug("Got new vacuum state: %s", state)
                return state

        try:
            return await execute_update()
        except DeviceException as ex:
            if getattr(ex, "code", None) != -9999:
                raise UpdateFailed(ex) from ex
            _LOGGER.info("Got exception while fetching the state, trying again: %s", ex)

        # Try to fetch the data a second time after error code -9999
        try:
            return await execute_update()
        except DeviceException as ex:
            raise UpdateFailed(ex) from ex

    return update_async

async def async_create_miio_device_and_coordinator(
    hass: HomeAssistant, entry: ConfigEntry
) -> None:
    """Set up a data coordinator and one miio device to service multiple entities."""
    model: str = entry.data[CONF_MODEL]
    host = entry.data[CONF_HOST]
    token = entry.data[CONF_TOKEN]
    name = entry.title
    device: MiioDevice | None = None
    migrate = False
    update_method = _async_update_data_default
    coordinator_class: type[DataUpdateCoordinator] = DataUpdateCoordinator

    if (
        model not in MODELS_HUMIDIFIER
        and model not in MODELS_FAN
        and model not in MODELS_VACUUM
        and not model.startswith(ROBOROCK_GENERIC)
        and not model.startswith(ROCKROBO_GENERIC)
    ):
        return

    _LOGGER.debug("Initializing with host %s (token %s...)", host, token[:5])

    # Humidifiers
    if model in MODELS_HUMIDIFIER_MIOT:
        device = AirHumidifierMiot(host, token)
        migrate = True
    elif model in MODELS_HUMIDIFIER_MJJSQ:
        device = AirHumidifierMjjsq(host, token, model=model)
        migrate = True
    elif model in MODELS_HUMIDIFIER_MIIO:
        device = AirHumidifier(host, token, model=model)
        migrate = True
    # Airpurifiers and Airfresh
    elif model in MODELS_PURIFIER_MIOT:
        device = AirPurifierMiot(host, token)
    elif model.startswith("zhimi.airpurifier."):
        device = AirPurifier(host, token)
    elif model.startswith("zhimi.airfresh."):
        device = AirFresh(host, token)
    elif model == MODEL_AIRFRESH_A1:
        device = AirFreshA1(host, token)
    elif model == MODEL_AIRFRESH_T2017:
        device = AirFreshT2017(host, token)
    elif (
        model in MODELS_VACUUM
        or model.startswith(ROBOROCK_GENERIC)
        or model.startswith(ROCKROBO_GENERIC)
    ):
        device = RoborockVacuum(host, token)
        update_method = _async_update_data_vacuum
        coordinator_class = DataUpdateCoordinator[VacuumCoordinatorData]
    # Pedestal fans
    elif model in MODEL_TO_CLASS_MAP:
        device = MODEL_TO_CLASS_MAP[model](host, token)
    elif model in MODELS_FAN_MIIO:
        device = Fan(host, token, model=model)
    else:
        _LOGGER.error(
            "Unsupported device found! Please create an issue at "
            "https://github.com/syssi/xiaomi_airpurifier/issues "
            "and provide the following data: %s",
            model,
        )
        return

    if migrate:
        # Removing fan platform entity for humidifiers and migrate the name to the config entry for migration
        entity_registry = er.async_get(hass)
        assert entry.unique_id
        entity_id = entity_registry.async_get_entity_id("fan", DOMAIN, entry.unique_id)
        if entity_id:
            # This check is entities that have a platform migration only and should be removed in the future
            if (entity := entity_registry.async_get(entity_id)) and (
                migrate_entity_name := entity.name
            ):
                hass.config_entries.async_update_entry(entry, title=migrate_entity_name)
            entity_registry.async_remove(entity_id)

    # Create update miio device and coordinator
    coordinator = coordinator_class(
        hass,
        _LOGGER,
        name=name,
        update_method=update_method(hass, device),
        # Polling interval. Will only be polled if there are subscribers.
        update_interval=UPDATE_INTERVAL,
    )
    hass.data[DOMAIN][entry.entry_id] = {
        KEY_DEVICE: device,
        KEY_COORDINATOR: coordinator,
    }

    # Trigger first data fetch
    await coordinator.async_config_entry_first_refresh()

async def async_setup_gateway_entry(hass: HomeAssistant, entry: ConfigEntry) -> None:
    """Set up the Xiaomi Gateway component from a config entry."""
    host = entry.data[CONF_HOST]
    token = entry.data[CONF_TOKEN]
    name = entry.title
    gateway_id = entry.unique_id

    assert gateway_id

    # For backwards compat
    if gateway_id.endswith("-gateway"):
        hass.config_entries.async_update_entry(entry, unique_id=entry.data["mac"])

    entry.async_on_unload(entry.add_update_listener(update_listener))

    # Connect to gateway
    gateway = ConnectXiaomiGateway(hass, entry)
    try:
        await gateway.async_connect_gateway(host, token)
    except AuthException as error:
        raise ConfigEntryAuthFailed() from error
    except SetupException as error:
        raise ConfigEntryNotReady() from error
    gateway_info = gateway.gateway_info

    device_registry = dr.async_get(hass)
    device_registry.async_get_or_create(
        config_entry_id=entry.entry_id,
        connections={(dr.CONNECTION_NETWORK_MAC, gateway_info.mac_address)},
        identifiers={(DOMAIN, gateway_id)},
        manufacturer="Xiaomi",
        name=name,
        model=gateway_info.model,
        sw_version=gateway_info.firmware_version,
        hw_version=gateway_info.hardware_version,
    )

    def update_data_factory(sub_device):
        """Create update function for a subdevice."""

        async def async_update_data():
            """Fetch data from the subdevice."""
            try:
                await hass.async_add_executor_job(sub_device.update)
            except GatewayException as ex:
                _LOGGER.error("Got exception while fetching the state: %s", ex)
                return {ATTR_AVAILABLE: False}
            return {ATTR_AVAILABLE: True}

        return async_update_data

    coordinator_dict: dict[str, DataUpdateCoordinator] = {}
    for sub_device in gateway.gateway_device.devices.values():
        # Create update coordinator
        coordinator_dict[sub_device.sid] = DataUpdateCoordinator(
            hass,
            _LOGGER,
            name=name,
            update_method=update_data_factory(sub_device),
            # Polling interval. Will only be polled if there are subscribers.
            update_interval=UPDATE_INTERVAL,
        )

    hass.data[DOMAIN][entry.entry_id] = {
        CONF_GATEWAY: gateway.gateway_device,
        KEY_COORDINATOR: coordinator_dict,
    }

    await hass.config_entries.async_forward_entry_setups(entry, GATEWAY_PLATFORMS)

async def async_setup_device_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
    """Set up the Xiaomi Miio device component from a config entry."""
    platforms = get_platforms(entry)
    await async_create_miio_device_and_coordinator(hass, entry)

    if not platforms:
        return False

    entry.async_on_unload(entry.add_update_listener(update_listener))

    await hass.config_entries.async_forward_entry_setups(entry, platforms)

    return True

async def async_unload_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> bool:
    """Unload a config entry."""
    platforms = get_platforms(config_entry)

    unload_ok = await hass.config_entries.async_unload_platforms(
        config_entry, platforms
    )

    if unload_ok:
        hass.data[DOMAIN].pop(config_entry.entry_id)

    return unload_ok

async def update_listener(hass: HomeAssistant, config_entry: ConfigEntry) -> None:
    """Handle options update."""
    await hass.config_entries.async_reload(config_entry.entry_id)
<commit_msg>Register xiaomi_miio unload callbacks later in setup (#76714)<commit_after>"""Support for Xiaomi Miio."""
from __future__ import annotations

from dataclasses import dataclass
from datetime import timedelta
import logging

import async_timeout
from miio import (
    AirFresh,
    AirFreshA1,
    AirFreshT2017,
    AirHumidifier,
    AirHumidifierMiot,
    AirHumidifierMjjsq,
    AirPurifier,
    AirPurifierMiot,
    CleaningDetails,
    CleaningSummary,
    ConsumableStatus,
    Device as MiioDevice,
    DeviceException,
    DNDStatus,
    Fan,
    Fan1C,
    FanMiot,
    FanP5,
    FanZA5,
    RoborockVacuum,
    Timer,
    VacuumStatus,
)
from miio.gateway.gateway import GatewayException

from homeassistant.config_entries import ConfigEntry
from homeassistant.const import CONF_HOST, CONF_MODEL, CONF_TOKEN, Platform
from homeassistant.core import HomeAssistant, callback
from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady
from homeassistant.helpers import device_registry as dr, entity_registry as er
from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed

from .const import (
    ATTR_AVAILABLE,
    CONF_DEVICE,
    CONF_FLOW_TYPE,
    CONF_GATEWAY,
    DOMAIN,
    KEY_COORDINATOR,
    KEY_DEVICE,
    MODEL_AIRFRESH_A1,
    MODEL_AIRFRESH_T2017,
    MODEL_FAN_1C,
    MODEL_FAN_P5,
    MODEL_FAN_P9,
    MODEL_FAN_P10,
    MODEL_FAN_P11,
    MODEL_FAN_ZA5,
    MODELS_AIR_MONITOR,
    MODELS_FAN,
    MODELS_FAN_MIIO,
    MODELS_HUMIDIFIER,
    MODELS_HUMIDIFIER_MIIO,
    MODELS_HUMIDIFIER_MIOT,
    MODELS_HUMIDIFIER_MJJSQ,
    MODELS_LIGHT,
    MODELS_PURIFIER_MIOT,
    MODELS_SWITCH,
    MODELS_VACUUM,
    ROBOROCK_GENERIC,
    ROCKROBO_GENERIC,
    AuthException,
    SetupException,
)
from .gateway import ConnectXiaomiGateway

_LOGGER = logging.getLogger(__name__)

POLLING_TIMEOUT_SEC = 10
UPDATE_INTERVAL = timedelta(seconds=15)

GATEWAY_PLATFORMS = [
    Platform.ALARM_CONTROL_PANEL,
    Platform.LIGHT,
    Platform.SENSOR,
    Platform.SWITCH,
]
SWITCH_PLATFORMS = [Platform.SWITCH]
FAN_PLATFORMS = [
    Platform.BINARY_SENSOR,
    Platform.BUTTON,
    Platform.FAN,
    Platform.NUMBER,
    Platform.SELECT,
    Platform.SENSOR,
    Platform.SWITCH,
]
HUMIDIFIER_PLATFORMS = [
    Platform.BINARY_SENSOR,
    Platform.HUMIDIFIER,
    Platform.NUMBER,
    Platform.SELECT,
    Platform.SENSOR,
    Platform.SWITCH,
]
LIGHT_PLATFORMS = [Platform.LIGHT]
VACUUM_PLATFORMS = [Platform.BINARY_SENSOR, Platform.SENSOR, Platform.VACUUM]
AIR_MONITOR_PLATFORMS = [Platform.AIR_QUALITY, Platform.SENSOR]

MODEL_TO_CLASS_MAP = {
    MODEL_FAN_1C: Fan1C,
    MODEL_FAN_P9: FanMiot,
    MODEL_FAN_P10: FanMiot,
    MODEL_FAN_P11: FanMiot,
    MODEL_FAN_P5: FanP5,
    MODEL_FAN_ZA5: FanZA5,
}

async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
    """Set up the Xiaomi Miio components from a config entry."""
    hass.data.setdefault(DOMAIN, {})
    if entry.data[CONF_FLOW_TYPE] == CONF_GATEWAY:
        await async_setup_gateway_entry(hass, entry)
        return True

    return bool(
        entry.data[CONF_FLOW_TYPE] != CONF_DEVICE
        or await async_setup_device_entry(hass, entry)
    )

@callback
def get_platforms(config_entry):
    """Return the platforms belonging to a config_entry."""
    model = config_entry.data[CONF_MODEL]
    flow_type = config_entry.data[CONF_FLOW_TYPE]

    if flow_type == CONF_GATEWAY:
        return GATEWAY_PLATFORMS
    if flow_type == CONF_DEVICE:
        if model in MODELS_SWITCH:
            return SWITCH_PLATFORMS
        if model in MODELS_HUMIDIFIER:
            return HUMIDIFIER_PLATFORMS
        if model in MODELS_FAN:
            return FAN_PLATFORMS
        if model in MODELS_LIGHT:
            return LIGHT_PLATFORMS
        for vacuum_model in MODELS_VACUUM:
            if model.startswith(vacuum_model):
                return VACUUM_PLATFORMS
        for air_monitor_model in MODELS_AIR_MONITOR:
            if model.startswith(air_monitor_model):
                return AIR_MONITOR_PLATFORMS
    _LOGGER.error(
        "Unsupported device found! Please create an issue at "
        "https://github.com/syssi/xiaomi_airpurifier/issues "
        "and provide the following data: %s",
        model,
    )
    return []

def _async_update_data_default(hass, device):
    async def update():
        """Fetch data from the device using async_add_executor_job."""

        async def _async_fetch_data():
            """Fetch data from the device."""
            async with async_timeout.timeout(POLLING_TIMEOUT_SEC):
                state = await hass.async_add_executor_job(device.status)
                _LOGGER.debug("Got new state: %s", state)
                return state

        try:
            return await _async_fetch_data()
        except DeviceException as ex:
            if getattr(ex, "code", None) != -9999:
                raise UpdateFailed(ex) from ex
            _LOGGER.info("Got exception while fetching the state, trying again: %s", ex)
        # Try to fetch the data a second time after error code -9999
        try:
            return await _async_fetch_data()
        except DeviceException as ex:
            raise UpdateFailed(ex) from ex

    return update

@dataclass(frozen=True)
class VacuumCoordinatorData:
    """A class that holds the vacuum data retrieved by the coordinator."""

    status: VacuumStatus
    dnd_status: DNDStatus
    last_clean_details: CleaningDetails
    consumable_status: ConsumableStatus
    clean_history_status: CleaningSummary
    timers: list[Timer]
    fan_speeds: dict[str, int]
    fan_speeds_reverse: dict[int, str]

@dataclass(init=False, frozen=True)
class VacuumCoordinatorDataAttributes:
    """
    A class that holds attribute names for VacuumCoordinatorData.

    These attributes can be used in methods like `getattr` when a generic solutions is
    needed.
    See homeassistant.components.xiaomi_miio.device.XiaomiCoordinatedMiioEntity
    ._extract_value_from_attribute for
    an example.
    """

    status: str = "status"
    dnd_status: str = "dnd_status"
    last_clean_details: str = "last_clean_details"
    consumable_status: str = "consumable_status"
    clean_history_status: str = "clean_history_status"
    timer: str = "timer"
    fan_speeds: str = "fan_speeds"
    fan_speeds_reverse: str = "fan_speeds_reverse"

def _async_update_data_vacuum(hass, device: RoborockVacuum):
    def update() -> VacuumCoordinatorData:
        timer = []

        # See https://github.com/home-assistant/core/issues/38285 for reason on
        # Why timers must be fetched separately.
        try:
            timer = device.timer()
        except DeviceException as ex:
            _LOGGER.debug(
                "Unable to fetch timers, this may happen on some devices: %s", ex
            )

        fan_speeds = device.fan_speed_presets()

        data = VacuumCoordinatorData(
            device.status(),
            device.dnd_status(),
            device.last_clean_details(),
            device.consumable_status(),
            device.clean_history(),
            timer,
            fan_speeds,
            {v: k for k, v in fan_speeds.items()},
        )

        return data

    async def update_async():
        """Fetch data from the device using async_add_executor_job."""

        async def execute_update():
            async with async_timeout.timeout(POLLING_TIMEOUT_SEC):
                state = await hass.async_add_executor_job(update)
                _LOGGER.debug("Got new vacuum state: %s", state)
                return state

        try:
            return await execute_update()
        except DeviceException as ex:
            if getattr(ex, "code", None) != -9999:
                raise UpdateFailed(ex) from ex
            _LOGGER.info("Got exception while fetching the state, trying again: %s", ex)

        # Try to fetch the data a second time after error code -9999
        try:
            return await execute_update()
        except DeviceException as ex:
            raise UpdateFailed(ex) from ex

    return update_async

async def async_create_miio_device_and_coordinator(
    hass: HomeAssistant, entry: ConfigEntry
) -> None:
    """Set up a data coordinator and one miio device to service multiple entities."""
    model: str = entry.data[CONF_MODEL]
    host = entry.data[CONF_HOST]
    token = entry.data[CONF_TOKEN]
    name = entry.title
    device: MiioDevice | None = None
    migrate = False
    update_method = _async_update_data_default
    coordinator_class: type[DataUpdateCoordinator] = DataUpdateCoordinator

    if (
        model not in MODELS_HUMIDIFIER
        and model not in MODELS_FAN
        and model not in MODELS_VACUUM
        and not model.startswith(ROBOROCK_GENERIC)
        and not model.startswith(ROCKROBO_GENERIC)
    ):
        return

    _LOGGER.debug("Initializing with host %s (token %s...)", host, token[:5])

    # Humidifiers
    if model in MODELS_HUMIDIFIER_MIOT:
        device = AirHumidifierMiot(host, token)
        migrate = True
    elif model in MODELS_HUMIDIFIER_MJJSQ:
        device = AirHumidifierMjjsq(host, token, model=model)
        migrate = True
    elif model in MODELS_HUMIDIFIER_MIIO:
        device = AirHumidifier(host, token, model=model)
        migrate = True
    # Airpurifiers and Airfresh
    elif model in MODELS_PURIFIER_MIOT:
        device = AirPurifierMiot(host, token)
    elif model.startswith("zhimi.airpurifier."):
        device = AirPurifier(host, token)
    elif model.startswith("zhimi.airfresh."):
        device = AirFresh(host, token)
    elif model == MODEL_AIRFRESH_A1:
        device = AirFreshA1(host, token)
    elif model == MODEL_AIRFRESH_T2017:
        device = AirFreshT2017(host, token)
    elif (
        model in MODELS_VACUUM
        or model.startswith(ROBOROCK_GENERIC)
        or model.startswith(ROCKROBO_GENERIC)
    ):
        device = RoborockVacuum(host, token)
        update_method = _async_update_data_vacuum
        coordinator_class = DataUpdateCoordinator[VacuumCoordinatorData]
    # Pedestal fans
    elif model in MODEL_TO_CLASS_MAP:
        device = MODEL_TO_CLASS_MAP[model](host, token)
    elif model in MODELS_FAN_MIIO:
        device = Fan(host, token, model=model)
    else:
        _LOGGER.error(
            "Unsupported device found! Please create an issue at "
            "https://github.com/syssi/xiaomi_airpurifier/issues "
            "and provide the following data: %s",
            model,
        )
        return

    if migrate:
        # Removing fan platform entity for humidifiers and migrate the name to the config entry for migration
        entity_registry = er.async_get(hass)
        assert entry.unique_id
        entity_id = entity_registry.async_get_entity_id("fan", DOMAIN, entry.unique_id)
        if entity_id:
            # This check is entities that have a platform migration only and should be removed in the future
            if (entity := entity_registry.async_get(entity_id)) and (
                migrate_entity_name := entity.name
            ):
                hass.config_entries.async_update_entry(entry, title=migrate_entity_name)
            entity_registry.async_remove(entity_id)

    # Create update miio device and coordinator
    coordinator = coordinator_class(
        hass,
        _LOGGER,
        name=name,
        update_method=update_method(hass, device),
        # Polling interval. Will only be polled if there are subscribers.
        update_interval=UPDATE_INTERVAL,
    )
    hass.data[DOMAIN][entry.entry_id] = {
        KEY_DEVICE: device,
        KEY_COORDINATOR: coordinator,
    }

    # Trigger first data fetch
    await coordinator.async_config_entry_first_refresh()

async def async_setup_gateway_entry(hass: HomeAssistant, entry: ConfigEntry) -> None:
    """Set up the Xiaomi Gateway component from a config entry."""
    host = entry.data[CONF_HOST]
    token = entry.data[CONF_TOKEN]
    name = entry.title
    gateway_id = entry.unique_id

    assert gateway_id

    # For backwards compat
    if gateway_id.endswith("-gateway"):
        hass.config_entries.async_update_entry(entry, unique_id=entry.data["mac"])

    # Connect to gateway
    gateway = ConnectXiaomiGateway(hass, entry)
    try:
        await gateway.async_connect_gateway(host, token)
    except AuthException as error:
        raise ConfigEntryAuthFailed() from error
    except SetupException as error:
        raise ConfigEntryNotReady() from error
    gateway_info = gateway.gateway_info

    device_registry = dr.async_get(hass)
    device_registry.async_get_or_create(
        config_entry_id=entry.entry_id,
        connections={(dr.CONNECTION_NETWORK_MAC, gateway_info.mac_address)},
        identifiers={(DOMAIN, gateway_id)},
        manufacturer="Xiaomi",
        name=name,
        model=gateway_info.model,
        sw_version=gateway_info.firmware_version,
        hw_version=gateway_info.hardware_version,
    )

    def update_data_factory(sub_device):
        """Create update function for a subdevice."""

        async def async_update_data():
            """Fetch data from the subdevice."""
            try:
                await hass.async_add_executor_job(sub_device.update)
            except GatewayException as ex:
                _LOGGER.error("Got exception while fetching the state: %s", ex)
                return {ATTR_AVAILABLE: False}
            return {ATTR_AVAILABLE: True}

        return async_update_data

    coordinator_dict: dict[str, DataUpdateCoordinator] = {}
    for sub_device in gateway.gateway_device.devices.values():
        # Create update coordinator
        coordinator_dict[sub_device.sid] = DataUpdateCoordinator(
            hass,
            _LOGGER,
            name=name,
            update_method=update_data_factory(sub_device),
            # Polling interval. Will only be polled if there are subscribers.
            update_interval=UPDATE_INTERVAL,
        )

    hass.data[DOMAIN][entry.entry_id] = {
        CONF_GATEWAY: gateway.gateway_device,
        KEY_COORDINATOR: coordinator_dict,
    }

    await hass.config_entries.async_forward_entry_setups(entry, GATEWAY_PLATFORMS)

    entry.async_on_unload(entry.add_update_listener(update_listener))

async def async_setup_device_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
    """Set up the Xiaomi Miio device component from a config entry."""
    platforms = get_platforms(entry)
    await async_create_miio_device_and_coordinator(hass, entry)

    if not platforms:
        return False

    await hass.config_entries.async_forward_entry_setups(entry, platforms)

    entry.async_on_unload(entry.add_update_listener(update_listener))

    return True

async def async_unload_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> bool:
    """Unload a config entry."""
    platforms = get_platforms(config_entry)

    unload_ok = await hass.config_entries.async_unload_platforms(
        config_entry, platforms
    )

    if unload_ok:
        hass.data[DOMAIN].pop(config_entry.entry_id)

    return unload_ok

async def update_listener(hass: HomeAssistant, config_entry: ConfigEntry) -> None:
    """Handle options update."""
    await hass.config_entries.async_reload(config_entry.entry_id)
<|endoftext|>


<commit_before>import csv
import fileinput
import re
from time import strftime
import time
import os

from collections import deque
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns

weightAdjustment = 0.8 #this value will be added to every weight entered by user from here on out
#a negative value here will subtract weight from what the user enters, and a positive value will add it.
logFilename = "weight.log"
avgWindow = 10
height = 72 #height in inches for BMI calcs
debug = False #set this to True to avoid saving any data (graphs are still generated)
keys = ['date', 'time', 'timestamp', 'bmi', 'weightAvg', 'weight']

secsInDay   = 60 * 60 * 24 *  1
secsInWeek  = 60 * 60 * 24 *  7
secsInMonth = 60 * 60 * 24 * 30

#find average of all values in a queue
def queAvg(q):
	s = 0
	for item in q:
		s += item
	s /= len(q)
	return s

#calc BMI from weight
def calcBmi(weight):
	return weight * 703 / height / height  #http://www.cdc.gov/healthyweight/assessing/bmi/adult_bmi/

scriptPath = os.path.dirname(os.path.realpath(__file__)) + "/"
try: #open log and read data, or create an empty list
	with open(scriptPath + logFilename, 'rb') as inFile:
		reader = csv.DictReader(inFile, delimiter=';', quotechar='"')
		weightData = list(reader) #get the data into a list
except IOError:
	weightData = []

currentTime = strftime("%H:%M:%S")
currentDate = strftime("%Y-%m-%d")
unixTime = int(time.time())
earliestTimestamp = float("inf")
minWeight = float("inf")
maxWeight = -1

pointsInThisWeek = 0
pointsInThisMonth = 0
curWeekNum = 0
curMonthNum = 0
sumThisWeek = 0
sumThisMonth = 0
weekAvgs = []
monthAvgs = []
weightQueue = deque() #a FIFO queue

#find the max and min weights
for row in weightData:
	if float(row['timestamp']) < earliestTimestamp:
		earliestTimestamp = float(row['timestamp'])
	thisWeight = float(row['weight'])
	if thisWeight > maxWeight:
		maxWeight = thisWeight
	if thisWeight < minWeight:
		minWeight = thisWeight

#ensure that there's enough data for two windows of avgWidnow length
if avgWindow * 2 > len(weightData):
	avgWindow = int(len(weightData) / 2)
	if avgWindow == 0:
		avgWindow = 1 #prevent divide by zero later

#find avg of the most recent avgWindow number of entries
tempSum = 0
for row in weightData[-avgWindow:]:
	tempSum += float(row['weight'])
recentAvg = tempSum / avgWindow

#find avg of the previous avgWindow number of entries
tempSum = 0
for row in weightData[-avgWindow * 2:-avgWindow]:
	tempSum += float(row['weight'])
prevAvg = tempSum / avgWindow

#show user the averages and get the current weight
print "Previous average: " + str(round(prevAvg, 2))
print "Current average:  " + str(round(recentAvg, 2))
currentWeight = raw_input("Enter current weight: ")
currentWeight = float(currentWeight)
currentWeight = round(currentWeight + weightAdjustment, 1) #adjust user entered weight based on above offset
print ""
newRow = [{'date': currentDate, 'timestamp': unixTime, 'weight': currentWeight, 'time': currentTime, 'weightAvg': currentWeight}]
weightData += newRow

#go through weights now that we have current value, find weekly averages, find bmi, find weekly avgs
for row in weightData:
	thisWeight = float(row['weight'])
	thisWeekNum = int((unixTime - float(row['timestamp'])) / secsInWeek)
	thisMonthNum = int((unixTime - float(row['timestamp'])) / secsInMonth)

	#calc BMI:
	thisBmi = round(calcBmi(thisWeight), 1)
	row['bmi'] = thisBmi

	#calc running average:
	weightQueue.append(thisWeight)
	if len(weightQueue) > avgWindow:
		weightQueue.popleft()
	runningAvg = queAvg(weightQueue)
	row['weightAvg'] = round(runningAvg, 1)

	#calc weekly averages:
	if thisWeekNum == curWeekNum:
		pointsInThisWeek += 1
		sumThisWeek += thisWeight
	else: #a new week
		try:
			weekAvgs.append(sumThisWeek / pointsInThisWeek)
		except ZeroDivisionError:
			pass
		#the first point in this new week is this point
		curWeekNum = thisWeekNum
		sumThisWeek = thisWeight
		pointsInThisWeek = 1

	#calc monthly averages:
	if thisMonthNum == curMonthNum:
		pointsInThisMonth += 1
		sumThisMonth += thisWeight
	else: #a new month
		try:
			monthAvgs.append(sumThisMonth / pointsInThisMonth)
		except ZeroDivisionError:
			pass
		#the first point in this new month is this point
		curMonthNum = thisMonthNum
		sumThisMonth = thisWeight
		pointsInThisMonth = 1

weekAvgs.append(sumThisWeek / pointsInThisWeek)
monthAvgs.append(sumThisMonth / pointsInThisMonth)

#calc, and display weekly deltas, BMIs
for ii, thisWeek in enumerate(weekAvgs):
	weekBmi = calcBmi(thisWeek)
	try:
		delta = thisWeek - lastWeek
	except NameError:
		delta = 0
	lastWeek = thisWeek
	print("Week #{:2d} Average: {:06.2f}, BMI: {:05.2f}, delta: {:+5.2f}".format(ii, round(thisWeek, 2), round(weekBmi, 2), round(delta, 2)))
print ""

#calc, and display monthly deltas, BMIs
for ii, thisMonth in enumerate(monthAvgs):
	monthBmi = calcBmi(thisMonth)
	try:
		delta = thisMonth - lastMonth
	except NameError:
		delta = 0
	lastMonth = thisMonth
	print("Month #{:2d} Average: {:06.2f}, BMI: {:05.2f}, delta: {:+5.2f}".format(ii, round(thisMonth, 2), round(monthBmi, 2), round(delta, 2)))
print ""

#skip plotting if we don't have enough data points
if len(weightData) > 2:
	#create the pandas dataframe which is used for plotting
	weightFrame = pd.DataFrame(weightData)
	weightFrame.weight    = weightFrame.weight.astype(float)
	weightFrame.weightAvg = weightFrame.weightAvg.astype(float)
	weightFrame.timestamp = weightFrame.timestamp.astype(float)
	weightFrame["datetime"] = weightFrame.date + " " + weightFrame.time
	weightFrame.datetime = pd.to_datetime(weightFrame.datetime)

	#here we convert from unix timestamp to # of days since start of data:
	convertToDays = lambda x: (x - earliestTimestamp) / secsInDay
	weightFrame.timestamp = weightFrame.timestamp.apply(convertToDays)

	#plotting code
	plot = weightFrame.plot(x='datetime', y='weight', figsize=(16, 10)).get_figure()
	plot.savefig(scriptPath + "line.png")
	plot = sns.lmplot(x='timestamp', y='weight', data=weightFrame, size=8, aspect=2, lowess=True)
	plot.set_axis_labels("Day", "Weight")
	plot.set(xlim=(0, None))
	plot.savefig(scriptPath + "scatter.png")

	#plot running average
	from scipy.interpolate import spline
	smoothTime = np.linspace(weightFrame.timestamp.min(), weightFrame.timestamp.max(), 500)
	weightSmooth = spline(weightFrame.timestamp, weightFrame.weightAvg, smoothTime)
	plot = sns.lmplot(x='timestamp', y='weight', data=weightFrame, size=8, aspect=2, fit_reg=False)
	plot.set_axis_labels("Day", "Weight")
	plot.set(xlim=(0, None))
	plt.plot(smoothTime, weightSmooth, sns.xkcd_rgb["electric blue"], lw=1)
	plt.savefig(scriptPath + "smooth.png")

	if currentWeight < minWeight:
		print "Lowest weight ever!"
	if currentWeight > maxWeight:
		print "Highest weight ever!"

if debug:
	print "WARNING: Debug mode, data not saved!"
else: #write data to log
	with open(scriptPath + logFilename, 'w') as outFile:
		csvWriter = csv.DictWriter(outFile, keys, delimiter=';')
		csvWriter.writeheader()
		csvWriter.writerows(weightData)
	outFile.close()
<commit_msg>Change offset to 0<commit_after>import csv
import fileinput
import re
from time import strftime
import time
import os

from collections import deque
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns

weightAdjustment = 0.0 #this value will be added to every weight entered by user from here on out
#a negative value here will subtract weight from what the user enters, and a positive value will add it.
logFilename = "weight.log"
avgWindow = 10
height = 72 #height in inches for BMI calcs
debug = False #set this to True to avoid saving any data (graphs are still generated)
keys = ['date', 'time', 'timestamp', 'bmi', 'weightAvg', 'weight']

secsInDay   = 60 * 60 * 24 *  1
secsInWeek  = 60 * 60 * 24 *  7
secsInMonth = 60 * 60 * 24 * 30

#find average of all values in a queue
def queAvg(q):
	s = 0
	for item in q:
		s += item
	s /= len(q)
	return s

#calc BMI from weight
def calcBmi(weight):
	return weight * 703 / height / height  #http://www.cdc.gov/healthyweight/assessing/bmi/adult_bmi/

scriptPath = os.path.dirname(os.path.realpath(__file__)) + "/"
try: #open log and read data, or create an empty list
	with open(scriptPath + logFilename, 'rb') as inFile:
		reader = csv.DictReader(inFile, delimiter=';', quotechar='"')
		weightData = list(reader) #get the data into a list
except IOError:
	weightData = []

currentTime = strftime("%H:%M:%S")
currentDate = strftime("%Y-%m-%d")
unixTime = int(time.time())
earliestTimestamp = float("inf")
minWeight = float("inf")
maxWeight = -1

pointsInThisWeek = 0
pointsInThisMonth = 0
curWeekNum = 0
curMonthNum = 0
sumThisWeek = 0
sumThisMonth = 0
weekAvgs = []
monthAvgs = []
weightQueue = deque() #a FIFO queue

#find the max and min weights
for row in weightData:
	if float(row['timestamp']) < earliestTimestamp:
		earliestTimestamp = float(row['timestamp'])
	thisWeight = float(row['weight'])
	if thisWeight > maxWeight:
		maxWeight = thisWeight
	if thisWeight < minWeight:
		minWeight = thisWeight

#ensure that there's enough data for two windows of avgWidnow length
if avgWindow * 2 > len(weightData):
	avgWindow = int(len(weightData) / 2)
	if avgWindow == 0:
		avgWindow = 1 #prevent divide by zero later

#find avg of the most recent avgWindow number of entries
tempSum = 0
for row in weightData[-avgWindow:]:
	tempSum += float(row['weight'])
recentAvg = tempSum / avgWindow

#find avg of the previous avgWindow number of entries
tempSum = 0
for row in weightData[-avgWindow * 2:-avgWindow]:
	tempSum += float(row['weight'])
prevAvg = tempSum / avgWindow

#show user the averages and get the current weight
print "Previous average: " + str(round(prevAvg, 2))
print "Current average:  " + str(round(recentAvg, 2))
currentWeight = raw_input("Enter current weight: ")
currentWeight = float(currentWeight)
currentWeight = round(currentWeight + weightAdjustment, 1) #adjust user entered weight based on above offset
print ""
newRow = [{'date': currentDate, 'timestamp': unixTime, 'weight': currentWeight, 'time': currentTime, 'weightAvg': currentWeight}]
weightData += newRow

#go through weights now that we have current value, find weekly averages, find bmi, find weekly avgs
for row in weightData:
	thisWeight = float(row['weight'])
	thisWeekNum = int((unixTime - float(row['timestamp'])) / secsInWeek)
	thisMonthNum = int((unixTime - float(row['timestamp'])) / secsInMonth)

	#calc BMI:
	thisBmi = round(calcBmi(thisWeight), 1)
	row['bmi'] = thisBmi

	#calc running average:
	weightQueue.append(thisWeight)
	if len(weightQueue) > avgWindow:
		weightQueue.popleft()
	runningAvg = queAvg(weightQueue)
	row['weightAvg'] = round(runningAvg, 1)

	#calc weekly averages:
	if thisWeekNum == curWeekNum:
		pointsInThisWeek += 1
		sumThisWeek += thisWeight
	else: #a new week
		try:
			weekAvgs.append(sumThisWeek / pointsInThisWeek)
		except ZeroDivisionError:
			pass
		#the first point in this new week is this point
		curWeekNum = thisWeekNum
		sumThisWeek = thisWeight
		pointsInThisWeek = 1

	#calc monthly averages:
	if thisMonthNum == curMonthNum:
		pointsInThisMonth += 1
		sumThisMonth += thisWeight
	else: #a new month
		try:
			monthAvgs.append(sumThisMonth / pointsInThisMonth)
		except ZeroDivisionError:
			pass
		#the first point in this new month is this point
		curMonthNum = thisMonthNum
		sumThisMonth = thisWeight
		pointsInThisMonth = 1

weekAvgs.append(sumThisWeek / pointsInThisWeek)
monthAvgs.append(sumThisMonth / pointsInThisMonth)

#calc, and display weekly deltas, BMIs
for ii, thisWeek in enumerate(weekAvgs):
	weekBmi = calcBmi(thisWeek)
	try:
		delta = thisWeek - lastWeek
	except NameError:
		delta = 0
	lastWeek = thisWeek
	print("Week #{:2d} Average: {:06.2f}, BMI: {:05.2f}, delta: {:+5.2f}".format(ii, round(thisWeek, 2), round(weekBmi, 2), round(delta, 2)))
print ""

#calc, and display monthly deltas, BMIs
for ii, thisMonth in enumerate(monthAvgs):
	monthBmi = calcBmi(thisMonth)
	try:
		delta = thisMonth - lastMonth
	except NameError:
		delta = 0
	lastMonth = thisMonth
	print("Month #{:2d} Average: {:06.2f}, BMI: {:05.2f}, delta: {:+5.2f}".format(ii, round(thisMonth, 2), round(monthBmi, 2), round(delta, 2)))
print ""

#skip plotting if we don't have enough data points
if len(weightData) > 2:
	#create the pandas dataframe which is used for plotting
	weightFrame = pd.DataFrame(weightData)
	weightFrame.weight    = weightFrame.weight.astype(float)
	weightFrame.weightAvg = weightFrame.weightAvg.astype(float)
	weightFrame.timestamp = weightFrame.timestamp.astype(float)
	weightFrame["datetime"] = weightFrame.date + " " + weightFrame.time
	weightFrame.datetime = pd.to_datetime(weightFrame.datetime)

	#here we convert from unix timestamp to # of days since start of data:
	convertToDays = lambda x: (x - earliestTimestamp) / secsInDay
	weightFrame.timestamp = weightFrame.timestamp.apply(convertToDays)

	#plotting code
	plot = weightFrame.plot(x='datetime', y='weight', figsize=(16, 10)).get_figure()
	plot.savefig(scriptPath + "line.png")
	plot = sns.lmplot(x='timestamp', y='weight', data=weightFrame, size=8, aspect=2, lowess=True)
	plot.set_axis_labels("Day", "Weight")
	plot.set(xlim=(0, None))
	plot.savefig(scriptPath + "scatter.png")

	#plot running average
	from scipy.interpolate import spline
	smoothTime = np.linspace(weightFrame.timestamp.min(), weightFrame.timestamp.max(), 500)
	weightSmooth = spline(weightFrame.timestamp, weightFrame.weightAvg, smoothTime)
	plot = sns.lmplot(x='timestamp', y='weight', data=weightFrame, size=8, aspect=2, fit_reg=False)
	plot.set_axis_labels("Day", "Weight")
	plot.set(xlim=(0, None))
	plt.plot(smoothTime, weightSmooth, sns.xkcd_rgb["electric blue"], lw=1)
	plt.savefig(scriptPath + "smooth.png")

	if currentWeight < minWeight:
		print "Lowest weight ever!"
	if currentWeight > maxWeight:
		print "Highest weight ever!"

if debug:
	print "WARNING: Debug mode, data not saved!"
else: #write data to log
	with open(scriptPath + logFilename, 'w') as outFile:
		csvWriter = csv.DictWriter(outFile, keys, delimiter=';')
		csvWriter.writeheader()
		csvWriter.writerows(weightData)
	outFile.close()
<|endoftext|>


<commit_before>"""
Various types of graph views for the Bitcoin block chain.

"""

__author__ = ['Bernhard Haslhofer (bernhard.haslhofer@ait.ac.at)',
              'Aljosha Judmayer (judmayer@xylem-technologies.com)']
__copyright__ = 'Copyright 2015, Bernhard Haslhofer, Aljosha Judmayer'
__license__ = "MIT"

# tx_graph/et_graph specific csv header fields
TXID = "txid"
SRC = "src"
DST = "tgt"
BTC = "value"
TIMESTAMP = "timestamp"
BLOCKID = "block_height"

# etmap specific csv header fields
ENTITYID = "entity_id"
BTCADDRS = "btc_addrs"

# btcmap specific csv header field
BTCADDR = "btc_addr"

# entity graph output files
ETG     = "etg.csv"
ETMAP   = "etmap.csv"
BTCMAP  = "btcmap.csv"

# csv special chars
DELIMCHR = ';'
QUOTECHR = '|'

# ensure list is sorted, has an impact on performance
EDGESORT = True

import os
import operator
import csv
import logging
logger = logging.getLogger("graph")

from bitcoingraph.blockchain import BlockchainException

class GraphException(Exception):
    """
    Exception raised when interacting with blockchain graph view.
    """

    def __init__(self, msg, inner_exc=None):
        self.msg = msg
        self.inner_exc = inner_exc

    def __str__(self):
        return repr(self.msg)

class Graph(object):
    """
    A generic directed graph representation.

    TODO: wrap third party library
    TODO: implement required graph search algorithms
    """
    def __init__(self):
        self._edges = []

    def add_edge(self, edge):
        """
        Add edge to graph.
        """
        self._edges.append(edge)

    def count_edges(self):
        """
        Returns number of edges in graph.
        """
        return len(self._edges)

    def list_edges(self, sort_key=None):
        """
        Returns an optionally sorted (by sort_key) iterator overall edges.
        """
        edge_list = self._edges
        if sort_key is not None:
            edge_list = sorted(self._edges, key=lambda k: k[sort_key])
        for edge in edge_list:
            yield edge

    def sort_edges(self):
        """
        Sorts edges according to (BLOCKID, TXID)
        This is slightly faster than useing sorted() and prioritizes
        low BLOCKIDs which should reduce entity changes
        """
        self._edges.sort(key = operator.itemgetter(BLOCKID, TXID))

class TransactionGraph(Graph):
    """
    A directed graph view on block chain transactions.

    Vertex = public key address (wallets)
    Edge = bitcoin flow from source to target address
    """

    def __init__(self, blockchain=None):
        """
        Creates transaction graph view on blockchain.
        """
        if blockchain is not None:
            self._blockchain = blockchain
        super().__init__()

    def generate_from_blockchain(self, start_block=None, end_block=None):
        """
        Generates transaction graph by extracting edges from blockchain.
        """
        for edge in self._generate(start_block, end_block):
            self.add_edge(edge)

    def load_from_file(self, tx_graph_file):
        """
        Loads already generated transaction graph from file.
        """
        with open(tx_graph_file, newline='') as csvfile:
                csv_reader = csv.DictReader(csvfile, delimiter=DELIMCHR,
                                            quotechar=QUOTECHR,
                                            quoting=csv.QUOTE_MINIMAL)
                for edge in csv_reader:
                    self.add_edge(edge)

    def export_to_csv(self, start_block=None,
                      end_block=None, output_file=None, progress=None):
        """
        Exports transaction graph to CSV file directly from blockchain.
        """
        with open(output_file, 'w', newline='') as csvfile:
            fieldnames = [TXID, SRC, DST,
                          BTC, TIMESTAMP, BLOCKID]
            csv_writer = csv.DictWriter(csvfile, delimiter=DELIMCHR,
                                        quotechar=QUOTECHR,
                                        fieldnames=fieldnames,
                                        quoting=csv.QUOTE_MINIMAL)
            csv_writer.writeheader()
            for edge in self._generate(start_block, end_block):
                csv_writer.writerow(edge)
                if progress:
                    progress(edge[BLOCKID] / (end_block - start_block))

    def _generate(self, start_block=None, end_block=None):
        """
        Generates transaction graph by extracting edges from blockchain.
        """
        if self._blockchain is None:
            raise GraphException("Cannot generated transaction graph without \
                reference to underlying blockchain")

        if start_block is None:
            start_block = 1
        if end_block is None:
            end_block = self._blockchain.get_max_blockheight()

        for block in self._blockchain.get_blocks_in_range(start_block,
                                                          end_block):
            for tx in block.transactions:
                try:
                    for bc_flow in tx.bc_flows:
                        # Construction transaction graph edge
                        edge = {}
                        edge[SRC] = bc_flow['src']
                        edge[DST] = bc_flow['tgt']
                        # Addding named edge descriptior
                        edge[TXID] = tx.id
                        edge[BTC] = bc_flow['value']
                        edge[TIMESTAMP] = tx.time
                        edge[BLOCKID] = block.height
                        yield edge
                except BlockchainException as exc:
                    raise GraphException("Transaction graph generation failed",
                                         exc)

class EntityGraph(Graph):
    """
    A directed graph view on entites in the blockchain.

    Vertex = entity
    Edge = bitcoin flow between entities

    Generator for Entity Graphs from transaction graphs
    Memory intensive variant.
    """
    def __init__(self, tx_graph=None, customlogger=None):
        """
        Creates entity graph view based on transaction graph.
        """
        self._tx_graph = tx_graph
        self._etdict = dict()   # dict() with entities as key
        self._btcdict = dict()   # dict() with btc addresses as key
        if customlogger:
            self._logger = customlogger
        else:
            self._logger = logger

        super().__init__()

        return

    def _handle_tx_inputs(self, txstack):
        """
        Generates or updates entity mapping for given group of txs of same id.
        """
        self._logger.debug("Handle txstack with len: {} in memory".format(len(txstack)))

        entity_id = None   # the entity id of all btc src addresses in this tx
        entitylist = set()  # set of all entity mappings for btc src addresses in this tx
        btcaddrlist = set()  # set all btc src addresses in this tx

        for edge in txstack:
            if (len(txstack) > 1 and edge[SRC] == 'NULL'):
                # ignore coinbase transactions or 'NULL' inputs
                self._logger.debug("Found NULL/coinbase tx input. Ignoring edge: {}".fomat(edge))
                continue

            if (edge[SRC] in self._btcdict.keys()):
                entity_id = self._btcdict[edge[SRC]]
                entitylist.add(self._btcdict[edge[SRC]])

            # create set of all unique source addresses
            btcaddrlist.add(edge[SRC])

        if len(entitylist) > 0:
            # handle entity collision and add new btc src addresses
            entity_id = min(entitylist)
            for et in entitylist:
                if et != entity_id:
                    #TODO: do the change btc address entity step here
                    self._etdict[entity_id] = self._etdict[entity_id].union(self._etdict[et])
                    self._etdict[et] = None

            # change btc address entity mapping of entries already in btc->entity dict
            for btcaddr in self._etdict[entity_id]:
                self._btcdict[btcaddr] = entity_id

            # add btc addresses to entity->btc dict
            for btcaddr in btcaddrlist:
                self._etdict[entity_id].add(btcaddr)
        else:
            # generate new entity and add new btc src addresses
            entity_id = len(self._etdict)+1
            self._etdict[entity_id] = btcaddrlist

        # add Bitcoin source address to btc->entity dict
        for btcaddr in btcaddrlist:
            self._btcdict[btcaddr] = entity_id

    def _generate_entity_mapping(self):
        """
        Generate entiy mapping from transaction graph.
        """
        txstack = list()
        txid = None

        #for edge in self._tx_graph.list_edges(sort_key=TXID):
        if EDGESORT: self.sort_edges()
        for edge in self._tx_graph._edges:
            # check if still the same transaction
            if (edge[TXID] != txid and txid is not None):
                self._handle_tx_inputs(txstack)
                txstack.clear()
            txid = edge[TXID]
            txstack.append(edge)
        if (len(txstack) > 0):
            self._handle_tx_inputs(txstack)
            txstack.clear()

        self._logger.info("Finished entity graph generation.")
        return 0

    def _generate_edges(self):
        """
        Generates entity graph from tranaction graph
        """
        ret = self._generate_entity_mapping()
        if ret != 0:
            self._logger.error("Failed generating entity graph")
            raise GraphException("Entity graph generation failed", exc)

        for tx_edge in self._tx_graph._edges:
            if tx_edge[SRC] in self._btcdict.keys():
                entity_src = self._btcdict[tx_edge[SRC]]
            else:
                entity_src = "0" #undefined entity
            if tx_edge[DST] in self._btcdict.keys():
                entity_dst = self._btcdict[tx_edge[DST]]
            else:
                entity_dst = "0"
            et_edge = dict()
            et_edge[TXID] = tx_edge[TXID]
            et_edge[SRC] = str(entity_src)
            et_edge[DST] = str(entity_dst)
            et_edge[BTC] = tx_edge[BTC]
            et_edge[TIMESTAMP] = tx_edge[TIMESTAMP]
            et_edge[BLOCKID] = tx_edge[BLOCKID]
            yield et_edge

    def generate_from_tx_graph(self, tx_graph=None):
        """
        Populate entity graph edges
        """
        if tx_graph:
            self._tx_graph = tx_graph
        elif self._tx_graph is None:
            self._logger.error("No tx_graph given")
            raise GraphException("No tx_graph given", None)

        generator = self._generate_edges()
        for edge in generator:
            self.add_edge(edge)

        return 0

    def load_from_dir(self, et_graph_dir):
        """
        Generates entity graph by loading edges from CSV file.
        """
        if not os.path.isdir(et_graph_dir):
            self._logger.error("Input entity graph directory is not a directory")
            return 1

        # clear all data structures
        self._edges.clear()
        self._etdict.clear()
        self._btcdict.clear()

        with open(et_graph_dir + "/" + ETG,'r') as etgfp:
            etgreader = csv.DictReader(etgfp,delimiter=DELIMCHR, quotechar=QUOTECHR)
            for edge in etgreader:
                self.add_edge(edge)

        with open(et_graph_dir + "/" + ETMAP, 'r') as etmapfp:
            etmapreader = csv.DictReader(etmapfp,delimiter=DELIMCHR, quotechar=QUOTECHR)
            for etmap in etmapreader:
                if etmap[BTCADDRS] == "None":
                    self._etdict[int(etmap[ENTITYID])] = None
                else:
                    self._etdict[int(etmap[ENTITYID])] = etmap[BTCADDRS]

        with open(et_graph_dir + "/btcmap.csv", 'r') as btcmapfp:
            btcmapreader = csv.DictReader(btcmapfp,delimiter=DELIMCHR, quotechar=QUOTECHR)
            for btcmap in btcmapreader:
                self._btcdict[str(btcmap[BTCADDR])] = int(btcmap[ENTITYID])

        return 0

    def export_to_csv(self, output_dir):
        """
        Export entity graph data to ouptput directory
        """
        if (len(self._btcdict) == 0 or len(self._etdict) == 0 or len(self._edges) == 0):
            self._logger.error("EntityGraph not initialized")
            return 1

        try:
            os.makedirs(output_dir)
        except OSError as exc:
            if not os.path.isdir(output_dir):
                raise GraphException("Output direcotry is not a directory", exc)

        # write entity graph
        with open(output_dir + "/" + ETG,'w') as etgfp:
            fieldnames = [ TXID, SRC, DST, BTC, TIMESTAMP, BLOCKID ]
            csv_writer = csv.DictWriter(etgfp,
                                        delimiter=DELIMCHR,
                                        quotechar=QUOTECHR,
                                        fieldnames=fieldnames,
                                        quoting=csv.QUOTE_MINIMAL)
            csv_writer.writeheader()
            for edge in self._edges:
                csv_writer.writerow(edge)

        # write entity mapping
        with open(output_dir + "/" + ETMAP,'w') as etmapfp:
            print(ENTITYID + DELIMCHR + BTCADDRS,file=etmapfp)
            for entity in self._etdict.items():
                print(str(entity[0]) + DELIMCHR,file=etmapfp,end='')
                if (entity[1] == None):
                    print('None',file=etmapfp)
                else:
                    for btcaddr in entity[1]:
                        print(str(btcaddr) + " ",file=etmapfp,end='')
                    print('',file=etmapfp)

        # write btc mapping
        with open(output_dir + "/" + BTCMAP ,'w') as btcmapfp:
            print(BTCADDR + DELIMCHR + ENTITYID,file=btcmapfp)
            for btcaddr in self._btcdict.items():
                print(str(btcaddr[0]) + DELIMCHR + str(btcaddr[1]),file=btcmapfp)

        return 0

<commit_msg>fixed typo in logging invocation<commit_after>"""
Various types of graph views for the Bitcoin block chain.

"""

__author__ = ['Bernhard Haslhofer (bernhard.haslhofer@ait.ac.at)',
              'Aljosha Judmayer (judmayer@xylem-technologies.com)']
__copyright__ = 'Copyright 2015, Bernhard Haslhofer, Aljosha Judmayer'
__license__ = "MIT"

# tx_graph/et_graph specific csv header fields
TXID = "txid"
SRC = "src"
DST = "tgt"
BTC = "value"
TIMESTAMP = "timestamp"
BLOCKID = "block_height"

# etmap specific csv header fields
ENTITYID = "entity_id"
BTCADDRS = "btc_addrs"

# btcmap specific csv header field
BTCADDR = "btc_addr"

# entity graph output files
ETG     = "etg.csv"
ETMAP   = "etmap.csv"
BTCMAP  = "btcmap.csv"

# csv special chars
DELIMCHR = ';'
QUOTECHR = '|'

# ensure list is sorted, has an impact on performance
EDGESORT = True

import os
import operator
import csv
import logging
logger = logging.getLogger("graph")

from bitcoingraph.blockchain import BlockchainException

class GraphException(Exception):
    """
    Exception raised when interacting with blockchain graph view.
    """

    def __init__(self, msg, inner_exc=None):
        self.msg = msg
        self.inner_exc = inner_exc

    def __str__(self):
        return repr(self.msg)

class Graph(object):
    """
    A generic directed graph representation.

    TODO: wrap third party library
    TODO: implement required graph search algorithms
    """
    def __init__(self):
        self._edges = []

    def add_edge(self, edge):
        """
        Add edge to graph.
        """
        self._edges.append(edge)

    def count_edges(self):
        """
        Returns number of edges in graph.
        """
        return len(self._edges)

    def list_edges(self, sort_key=None):
        """
        Returns an optionally sorted (by sort_key) iterator overall edges.
        """
        edge_list = self._edges
        if sort_key is not None:
            edge_list = sorted(self._edges, key=lambda k: k[sort_key])
        for edge in edge_list:
            yield edge

    def sort_edges(self):
        """
        Sorts edges according to (BLOCKID, TXID)
        This is slightly faster than useing sorted() and prioritizes
        low BLOCKIDs which should reduce entity changes
        """
        self._edges.sort(key = operator.itemgetter(BLOCKID, TXID))

class TransactionGraph(Graph):
    """
    A directed graph view on block chain transactions.

    Vertex = public key address (wallets)
    Edge = bitcoin flow from source to target address
    """

    def __init__(self, blockchain=None):
        """
        Creates transaction graph view on blockchain.
        """
        if blockchain is not None:
            self._blockchain = blockchain
        super().__init__()

    def generate_from_blockchain(self, start_block=None, end_block=None):
        """
        Generates transaction graph by extracting edges from blockchain.
        """
        for edge in self._generate(start_block, end_block):
            self.add_edge(edge)

    def load_from_file(self, tx_graph_file):
        """
        Loads already generated transaction graph from file.
        """
        with open(tx_graph_file, newline='') as csvfile:
                csv_reader = csv.DictReader(csvfile, delimiter=DELIMCHR,
                                            quotechar=QUOTECHR,
                                            quoting=csv.QUOTE_MINIMAL)
                for edge in csv_reader:
                    self.add_edge(edge)

    def export_to_csv(self, start_block=None,
                      end_block=None, output_file=None, progress=None):
        """
        Exports transaction graph to CSV file directly from blockchain.
        """
        with open(output_file, 'w', newline='') as csvfile:
            fieldnames = [TXID, SRC, DST,
                          BTC, TIMESTAMP, BLOCKID]
            csv_writer = csv.DictWriter(csvfile, delimiter=DELIMCHR,
                                        quotechar=QUOTECHR,
                                        fieldnames=fieldnames,
                                        quoting=csv.QUOTE_MINIMAL)
            csv_writer.writeheader()
            for edge in self._generate(start_block, end_block):
                csv_writer.writerow(edge)
                if progress:
                    progress(edge[BLOCKID] / (end_block - start_block))

    def _generate(self, start_block=None, end_block=None):
        """
        Generates transaction graph by extracting edges from blockchain.
        """
        if self._blockchain is None:
            raise GraphException("Cannot generated transaction graph without \
                reference to underlying blockchain")

        if start_block is None:
            start_block = 1
        if end_block is None:
            end_block = self._blockchain.get_max_blockheight()

        for block in self._blockchain.get_blocks_in_range(start_block,
                                                          end_block):
            for tx in block.transactions:
                try:
                    for bc_flow in tx.bc_flows:
                        # Construction transaction graph edge
                        edge = {}
                        edge[SRC] = bc_flow['src']
                        edge[DST] = bc_flow['tgt']
                        # Addding named edge descriptior
                        edge[TXID] = tx.id
                        edge[BTC] = bc_flow['value']
                        edge[TIMESTAMP] = tx.time
                        edge[BLOCKID] = block.height
                        yield edge
                except BlockchainException as exc:
                    raise GraphException("Transaction graph generation failed",
                                         exc)

class EntityGraph(Graph):
    """
    A directed graph view on entites in the blockchain.

    Vertex = entity
    Edge = bitcoin flow between entities

    Generator for Entity Graphs from transaction graphs
    Memory intensive variant.
    """
    def __init__(self, tx_graph=None, customlogger=None):
        """
        Creates entity graph view based on transaction graph.
        """
        self._tx_graph = tx_graph
        self._etdict = dict()   # dict() with entities as key
        self._btcdict = dict()   # dict() with btc addresses as key
        if customlogger:
            self._logger = customlogger
        else:
            self._logger = logger

        super().__init__()

        return

    def _handle_tx_inputs(self, txstack):
        """
        Generates or updates entity mapping for given group of txs of same id.
        """
        self._logger.debug("Handle txstack with len: {} in memory".format(len(txstack)))

        entity_id = None   # the entity id of all btc src addresses in this tx
        entitylist = set()  # set of all entity mappings for btc src addresses in this tx
        btcaddrlist = set()  # set all btc src addresses in this tx

        for edge in txstack:
            if (len(txstack) > 1 and edge[SRC] == 'NULL'):
                # ignore coinbase transactions or 'NULL' inputs
                self._logger.debug("Found NULL/coinbase tx input. Ignoring edge: {}".format(edge))
                continue

            if (edge[SRC] in self._btcdict.keys()):
                entity_id = self._btcdict[edge[SRC]]
                entitylist.add(self._btcdict[edge[SRC]])

            # create set of all unique source addresses
            btcaddrlist.add(edge[SRC])

        if len(entitylist) > 0:
            # handle entity collision and add new btc src addresses
            entity_id = min(entitylist)
            for et in entitylist:
                if et != entity_id:
                    #TODO: do the change btc address entity step here
                    self._etdict[entity_id] = self._etdict[entity_id].union(self._etdict[et])
                    self._etdict[et] = None

            # change btc address entity mapping of entries already in btc->entity dict
            for btcaddr in self._etdict[entity_id]:
                self._btcdict[btcaddr] = entity_id

            # add btc addresses to entity->btc dict
            for btcaddr in btcaddrlist:
                self._etdict[entity_id].add(btcaddr)
        else:
            # generate new entity and add new btc src addresses
            entity_id = len(self._etdict)+1
            self._etdict[entity_id] = btcaddrlist

        # add Bitcoin source address to btc->entity dict
        for btcaddr in btcaddrlist:
            self._btcdict[btcaddr] = entity_id

    def _generate_entity_mapping(self):
        """
        Generate entiy mapping from transaction graph.
        """
        txstack = list()
        txid = None

        #for edge in self._tx_graph.list_edges(sort_key=TXID):
        if EDGESORT: self.sort_edges()
        for edge in self._tx_graph._edges:
            # check if still the same transaction
            if (edge[TXID] != txid and txid is not None):
                self._handle_tx_inputs(txstack)
                txstack.clear()
            txid = edge[TXID]
            txstack.append(edge)
        if (len(txstack) > 0):
            self._handle_tx_inputs(txstack)
            txstack.clear()

        self._logger.info("Finished entity graph generation.")
        return 0

    def _generate_edges(self):
        """
        Generates entity graph from tranaction graph
        """
        ret = self._generate_entity_mapping()
        if ret != 0:
            self._logger.error("Failed generating entity graph")
            raise GraphException("Entity graph generation failed", exc)

        for tx_edge in self._tx_graph._edges:
            if tx_edge[SRC] in self._btcdict.keys():
                entity_src = self._btcdict[tx_edge[SRC]]
            else:
                entity_src = "0" #undefined entity
            if tx_edge[DST] in self._btcdict.keys():
                entity_dst = self._btcdict[tx_edge[DST]]
            else:
                entity_dst = "0"
            et_edge = dict()
            et_edge[TXID] = tx_edge[TXID]
            et_edge[SRC] = str(entity_src)
            et_edge[DST] = str(entity_dst)
            et_edge[BTC] = tx_edge[BTC]
            et_edge[TIMESTAMP] = tx_edge[TIMESTAMP]
            et_edge[BLOCKID] = tx_edge[BLOCKID]
            yield et_edge

    def generate_from_tx_graph(self, tx_graph=None):
        """
        Populate entity graph edges
        """
        if tx_graph:
            self._tx_graph = tx_graph
        elif self._tx_graph is None:
            self._logger.error("No tx_graph given")
            raise GraphException("No tx_graph given", None)

        generator = self._generate_edges()
        for edge in generator:
            self.add_edge(edge)

        return 0

    def load_from_dir(self, et_graph_dir):
        """
        Generates entity graph by loading edges from CSV file.
        """
        if not os.path.isdir(et_graph_dir):
            self._logger.error("Input entity graph directory is not a directory")
            return 1

        # clear all data structures
        self._edges.clear()
        self._etdict.clear()
        self._btcdict.clear()

        with open(et_graph_dir + "/" + ETG,'r') as etgfp:
            etgreader = csv.DictReader(etgfp,delimiter=DELIMCHR, quotechar=QUOTECHR)
            for edge in etgreader:
                self.add_edge(edge)

        with open(et_graph_dir + "/" + ETMAP, 'r') as etmapfp:
            etmapreader = csv.DictReader(etmapfp,delimiter=DELIMCHR, quotechar=QUOTECHR)
            for etmap in etmapreader:
                if etmap[BTCADDRS] == "None":
                    self._etdict[int(etmap[ENTITYID])] = None
                else:
                    self._etdict[int(etmap[ENTITYID])] = etmap[BTCADDRS]

        with open(et_graph_dir + "/btcmap.csv", 'r') as btcmapfp:
            btcmapreader = csv.DictReader(btcmapfp,delimiter=DELIMCHR, quotechar=QUOTECHR)
            for btcmap in btcmapreader:
                self._btcdict[str(btcmap[BTCADDR])] = int(btcmap[ENTITYID])

        return 0

    def export_to_csv(self, output_dir):
        """
        Export entity graph data to ouptput directory
        """
        if (len(self._btcdict) == 0 or len(self._etdict) == 0 or len(self._edges) == 0):
            self._logger.error("EntityGraph not initialized")
            return 1

        try:
            os.makedirs(output_dir)
        except OSError as exc:
            if not os.path.isdir(output_dir):
                raise GraphException("Output direcotry is not a directory", exc)

        # write entity graph
        with open(output_dir + "/" + ETG,'w') as etgfp:
            fieldnames = [ TXID, SRC, DST, BTC, TIMESTAMP, BLOCKID ]
            csv_writer = csv.DictWriter(etgfp,
                                        delimiter=DELIMCHR,
                                        quotechar=QUOTECHR,
                                        fieldnames=fieldnames,
                                        quoting=csv.QUOTE_MINIMAL)
            csv_writer.writeheader()
            for edge in self._edges:
                csv_writer.writerow(edge)

        # write entity mapping
        with open(output_dir + "/" + ETMAP,'w') as etmapfp:
            print(ENTITYID + DELIMCHR + BTCADDRS,file=etmapfp)
            for entity in self._etdict.items():
                print(str(entity[0]) + DELIMCHR,file=etmapfp,end='')
                if (entity[1] == None):
                    print('None',file=etmapfp)
                else:
                    for btcaddr in entity[1]:
                        print(str(btcaddr) + " ",file=etmapfp,end='')
                    print('',file=etmapfp)

        # write btc mapping
        with open(output_dir + "/" + BTCMAP ,'w') as btcmapfp:
            print(BTCADDR + DELIMCHR + ENTITYID,file=btcmapfp)
            for btcaddr in self._btcdict.items():
                print(str(btcaddr[0]) + DELIMCHR + str(btcaddr[1]),file=btcmapfp)

        return 0

<|endoftext|>