# -*- coding: utf-8 -*-
"""A model of Infrastructure Virtual Machines area of CFME.  This includes the VMs explorer tree,
quadicon lists, and VM details page.
"""
import fauxfactory
import re
from cfme import js
from cfme.exceptions import CandidateNotFound, VmNotFound, OptionNotAvailable, TemplateNotFound
from cfme.fixtures import pytest_selenium as sel
from cfme.services import requests
from cfme.web_ui import (
    CheckboxTree, Form, Region, Quadicon, Tree, accordion, fill, flash, form_buttons, paginator,
    toolbar, Calendar, Select, Input, mixins
)
from cfme.web_ui.menu import nav
from datetime import date
from contextlib import contextmanager
from functools import partial
from selenium.common.exceptions import NoSuchElementException
from utils.conf import cfme_data
from utils.log import logger
from utils.pretty import Pretty
from utils.timeutil import parsetime
from utils.virtual_machines import deploy_template
from utils.wait import wait_for, TimedOutError
from utils.mgmt_system import ActionNotSupported, VMInstanceNotFound
from utils import version

QUADICON_TITLE_LOCATOR = ("//div[@id='quadicon']/../../../tr/td/a[contains(@href,'vm_infra/x_show')"
                         " or contains(@href, '/show/')]")  # for provider specific vm/template page

details_page = Region(infoblock_type='detail')

cfg_btn = partial(toolbar.select, 'Configuration')
pol_btn = partial(toolbar.select, 'Policy')
lcl_btn = partial(toolbar.select, 'Lifecycle')
mon_btn = partial(toolbar.select, 'Monitoring')
pwr_btn = partial(toolbar.select, 'Power')

create_button = form_buttons.FormButton("Create")

visible_tree = Tree("//div[@class='dhxcont_global_content_area']"
                    "[not(contains(@style, 'display: none'))]/div/div/div"
                    "/ul[@class='dynatree-container']")

manage_policies_tree = CheckboxTree(
    {
        version.LOWEST: "//div[@id='treebox']/div/table",
        "5.3": "//div[@id='protect_treebox']/ul"
    }
)

manage_policies_page = Region(
    locators={
        'save_button': form_buttons.save,
    })

snapshot_tree = Tree(
    {
        version.LOWEST: "//div[@id='treebox']/div/table",
        "5.3": "//div[@id='snapshots_treebox']/ul"
    }
)

snapshot_form = Form(
    fields=[
        ('name', Input('name')),
        ('descrition', Input('description')),
        ('snapshot_memory', Input('snap_memory')),
        ('create_button', create_button),
        ('cancel_button', form_buttons.cancel)
    ])

def date_retire_element(fill_data):
    """We need to call this function that will mimic clicking the calendar, picking the date and
    the subsequent callbacks from the server"""
    # TODO: Move the code in the Calendar itself? I did not check other calendars
    if isinstance(fill_data, date):
        date_str = '%s/%s/%s' % (fill_data.month, fill_data.day, fill_data.year)
    else:
        date_str = str(fill_data)
    sel.execute_script(
        js.update_retirement_date_function_script +
        "updateDate(arguments[0]);",
        date_str
    )

retire_form = Form(fields=[
    ('date_retire', date_retire_element),
    ('warn', sel.Select("select#retirement_warn"))
])

retirement_date_form = Form(fields=[
    ('retirement_date_text', Calendar("miq_date_1")),
    ('retirement_warning_select', Select("//select[@id='retirement_warn']"))
])

retire_remove_button = "//span[@id='remove_button']/a/img"

set_ownership_form = Form(fields=[
    ('user_name', Select("//select[@id='user_name']")),
    ('group_name', Select("//select[@id='group_name']")),
    ('create_button', form_buttons.save),
    ('reset_button', form_buttons.reset),
    ('cancel_button', form_buttons.cancel)
])

nav.add_branch(
    'infrastructure_virtual_machines',
    {
        "infra_vm_and_templates":
        [
            lambda _: accordion.tree("VMs & Templates", "All VMs & Templates"),
            {
                "vm_templates_provider_branch":
                [
                    lambda ctx: visible_tree.click_path(ctx["provider_name"]),
                    {
                        "datacenter_branch":
                        [
                            lambda ctx: visible_tree.click_path(ctx["datacenter_name"]),
                            {
                                "infra_vm_obj": lambda ctx: visible_tree.click_path(ctx["vm_name"]),
                            }
                        ],
                    }
                ],

                "vm_templates_archived_branch":
                [
                    lambda ctx: visible_tree.click_path("<Archived>"),
                    {
                        "infra_archive_obj":
                        lambda ctx: visible_tree.click_path(ctx["archive_name"]),
                    }
                ],

                "vm_templates_orphaned_branch":
                [
                    lambda ctx: visible_tree.click_path('<Orphaned>'),
                    {
                        "infra_orphan_obj": lambda ctx: visible_tree.click_path(ctx["orphan_name"]),
                    }
                ],
            }
        ],

        "infra_vms":
        [
            lambda _: accordion.tree("VMs", "All VMs"),
            {
                "infra_vms_filter_folder":
                [
                    lambda ctx: visible_tree.click_path(ctx["folder_name"]),
                    {
                        "infra_vms_filter": lambda ctx: visible_tree.click_path(ctx["filter_name"]),
                    }
                ],

                "infra_vm_by_name": lambda ctx: sel.click(ctx['vm'].find_quadicon(
                    do_not_navigate=True))
            }
        ],

        "infra_templates":
        [
            lambda _: (accordion.tree("Templates", "All Templates"), toolbar.set_vms_grid_view()),
            {
                "infra_templates_filter_folder":
                [
                    lambda ctx: visible_tree.click_path(ctx["folder_name"]),
                    {
                        "infra_templates_filter":
                        lambda ctx: visible_tree.click_path(ctx["filter_name"]),
                    }
                ],
            }
        ],
    }
)

class Common(Pretty):

    pretty_attrs = ["name", "provider_crud", "template_name"]

    def __init__(self):
        self._assigned_pp = set([])

    def _load_details(self, refresh=False, is_vm=True):
        """Navigates to a VM's details page.

        Args:
            refresh: Refreshes the vm page if already there

        Raises:
            VmNotFound:
                When unable to find the VM passed
        """
        if not self.on_details(is_vm=is_vm):
            logger.debug("load_vm_details: not on details already")
            sel.click(self._find_quadicon(is_vm=is_vm))
        else:
            if refresh:
                toolbar.refresh()

    def on_details(self, force=False, is_vm=True):
        """A function to determine if the browser is already on the proper vm details page.
        """
        locator = ("//div[@class='dhtmlxInfoBarLabel' and contains(. , 'VM and Instance \"%s\"')]" %
                   self.name)
        if not sel.is_displayed(locator):
            if not force:
                return False
            else:
                self.load_details(is_vm=is_vm)
                return True

        text = sel.text(locator).encode("utf-8")
        pattern = r'("[A-Za-z0-9_\./\\-]*")'
        m = re.search(pattern, text)

        if not force:
            return self.name == m.group().replace('"', '')
        else:
            if self.name != m.group().replace('"', ''):
                self._load_details(is_vm=is_vm)
            else:
                return False

    def _find_quadicon(
            self, is_vm=True, do_not_navigate=False, mark=False, refresh=True,
            from_any_provider=False):
        """Find and return a quadicon belonging to a specific vm

        Args:
            from_any_provider: Whether to look for it anywhere (root of the tree). Useful when
                looking up archived or orphaned VMs

        Returns: :py:class:`cfme.web_ui.Quadicon` instance
        Raises: VmNotFound
        """
        quadicon = Quadicon(self.name, "vm")
        if not do_not_navigate:
            if from_any_provider:
                if is_vm:
                    sel.force_navigate("infra_vms")
                else:
                    sel.force_navigate("infra_templates")
            elif is_vm:
                self.provider_crud.load_all_provider_vms()
            else:
                self.provider_crud.load_all_provider_templates()
            toolbar.set_vms_grid_view()
        elif refresh:
            sel.refresh()
        if not paginator.page_controls_exist():
            if is_vm:
                raise VmNotFound("VM '{}' not found in UI!".format(self.name))
            else:
                raise TemplateNotFound("Template '{}' not found in UI!".format(self.name))

        paginator.results_per_page(1000)
        for page in paginator.pages():
            if sel.is_displayed(quadicon, move_to=True):
                if mark:
                    sel.check(quadicon.checkbox())
                return quadicon
        else:
            raise VmNotFound("VM '{}' not found in UI!".format(self.name))

    def does_vm_exist_on_provider(self):
        """Check if VM exists on provider itself"""
        return self.provider_crud.get_mgmt_system().does_vm_exist(self.name)

    def _method_helper(self, from_details=False):
        if from_details:
            self.load_details()
        else:
            self.find_quadicon(mark=True)

    def _remove_from_cfme(self, is_vm=True, cancel=True, from_details=False):
        """Removes a VM from CFME VMDB

        Args:
            cancel: Whether to cancel the deletion, defaults to True
            from_details: whether to delete from the details page
        """
        self._method_helper(from_details)
        if from_details:
            cfg_btn('Remove from the VMDB', invokes_alert=True)
        elif not is_vm:
            cfg_btn('Remove Templates from the VMDB', invokes_alert=True)
        else:
            cfg_btn('Remove selected items from the VMDB', invokes_alert=True)
        sel.handle_alert(cancel=cancel)

    def delete_from_provider(self):
        provider_mgmt = self.provider_crud.get_mgmt_system()
        if provider_mgmt.does_vm_exist(self.name):
            try:
                if provider_mgmt.is_vm_suspended(self.name):
                    logger.debug("Powering up VM %s to shut it down correctly on %s." %
                                 (self.name, self.provider_crud.key))
                    provider_mgmt.start_vm(self.name)
            except ActionNotSupported:
                # Action is not supported on mgmt system. Simply continue
                pass
            # One more check (for the suspended one)
            if provider_mgmt.does_vm_exist(self.name):
                try:
                    return provider_mgmt.delete_vm(self.name)
                except VMInstanceNotFound:
                    # Does not exist already
                    return True
        else:
            return True

    def get_detail(self, properties=None, icon_href=False):
        """Gets details from the details infoblock

        The function first ensures that we are on the detail page for the specific vm.

        Args:
            properties: An InfoBlock title, followed by the Key name, e.g. "Relationships", "Images"

        Returns:
            A string representing the contents of the InfoBlock's value.
        """
        self.load_details(refresh=True)
        if icon_href:
            return details_page.infoblock.icon_href(*properties)
        else:
            return details_page.infoblock.text(*properties)

    @property
    def compliance_status(self):
        """Returns the title of the compliance infoblock. The title contains datetime so it can be
        compared.

        Returns:
            :py:class:`NoneType` if no title is present (no compliance checks before), otherwise str
        """
        self.load_details(refresh=True)
        return details_page.infoblock("Compliance", "Status").title

    @property
    def compliant(self):
        """Check if the VM is compliant

        Returns:
            :py:class:`NoneType` if the VM was never verified, otherwise :py:class:`bool`
        """
        text = self.get_detail(properties=("Compliance", "Status")).strip().lower()
        if text == "never verified":
            return None
        elif text.startswith("non-compliant"):
            return False
        elif text.startswith("compliant"):
            return True
        else:
            raise ValueError("{} is not a known state for compliance".format(text))

    @contextmanager
    def check_compliance_wrapper(self, timeout=240):
        """This wrapper takes care of waiting for the compliance status to change

        Args:
            timeout: Wait timeout in seconds.
        """
        self.load_details(refresh=True)
        original_state = self.compliance_status
        yield
        wait_for(
            lambda: self.compliance_status != original_state,
            num_sec=timeout, delay=5, message="compliance of {} checked".format(self.name),
            fail_func=lambda: toolbar.select("Reload"))

    def check_compliance_and_wait(self, timeout=240):
        with self.check_compliance_wrapper(timeout=timeout):
            self.check_compliance()
        return self.compliant

    def rediscover(self):
        """Deletes the VM from the provider and lets it discover again"""
        self.remove_from_cfme(cancel=False, from_details=True)
        wait_for(self.does_vm_exist_in_cfme, fail_condition=True,
                 num_sec=300, delay=15, fail_func=sel.refresh)
        self.provider_crud.refresh_provider_relationships()
        self.wait_to_appear()

    def rediscover_if_analysis_data_present(self):
        if self.get_detail(properties=('Lifecycle', 'Last Analyzed')).lower() != 'never':
            self.rediscover()
            return True
        return False

    def get_tags(self, tag="My Company Tags"):
        """Returns all tags that are associated with this VM"""
        self.load_details(refresh=True)
        return mixins.get_tags(tag)

    def refresh_relationships(self, from_details=False, cancel=False):
        """Executes a refresh relationships action against a list of VMs.

        Args:
            from_details: Whether or not to perform action from vm details page
            cancel: Whether or not to cancel the refresh relationships action
        """
        if from_details:
            self.load_details()
        else:
            self.find_quadicon(mark=True)
        cfg_btn('Refresh Relationships and Power States', invokes_alert=True)
        sel.handle_alert(cancel=cancel)

    def smartstate_scan(self, cancel=True, from_details=False):
        self._method_helper(from_details=from_details)
        cfg_btn('Perform SmartState Analysis', invokes_alert=True)
        sel.handle_alert(cancel=cancel)

    def wait_to_appear(self, timeout=600, load_details=True):
        """Wait for a VM to appear within CFME

        Args:
            timeout: time (in seconds) to wait for it to appear
            from_details: when found, should it load the vm details
        """
        wait_for(self.does_vm_exist_in_cfme, num_sec=timeout, delay=30, fail_func=sel.refresh)
        if load_details:
            self.load_details()

    def wait_for_delete(self, timeout=600):
        wait_for(self.does_vm_exist_in_cfme, num_sec=timeout, delay=30, fail_func=sel.refresh,
                 fail_condition=True)

    def does_vm_exist_in_cfme(self):
        """A function to tell you if a VM exists or not.
        """
        try:
            self.find_quadicon()
            return True
        except VmNotFound:
            return False

    @property
    def genealogy(self):
        return Genealogy(self)

    def check_compliance(self):
        self.load_details(refresh=True)
        pol_btn("Check Compliance of Last Known Configuration", invokes_alert=True)
        sel.handle_alert()
        flash.assert_no_errors()

    def assign_policy_profiles(self, *policy_profile_names):
        """ Assign Policy Profiles to this VM.

        Args:
            policy_profile_names: :py:class:`str` with Policy Profile names. After Control/Explorer
                coverage goes in, PolicyProfile objects will be also passable.
        """
        map(self._assigned_pp.add, policy_profile_names)
        self._assign_unassign_policy_profiles(True, *policy_profile_names)

    def unassign_policy_profiles(self, *policy_profile_names):
        """ Unssign Policy Profiles to this VM.

        Args:
            policy_profile_names: :py:class:`str` with Policy Profile names. After Control/Explorer
                coverage goes in, PolicyProfile objects will be also passable.
        """
        for pp_name in policy_profile_names:
            try:
                self._assigned_pp.remove(pp_name)
            except KeyError:
                pass
        self._assign_unassign_policy_profiles(False, *policy_profile_names)

    def _assign_unassign_policy_profiles(self, assign, *policy_profile_names):
        """DRY function for managing policy profiles.

        See :py:func:`assign_policy_profiles` and :py:func:`assign_policy_profiles`

        Args:
            assign: Wheter to assign or unassign.
            policy_profile_names: :py:class:`str` with Policy Profile names.
        """
        self.load_details(refresh=True)
        pol_btn("Manage Policies")
        for policy_profile in policy_profile_names:
            if assign:
                manage_policies_tree.check_node(policy_profile)
            else:
                manage_policies_tree.uncheck_node(policy_profile)
        sel.move_to_element('#tP')
        form_buttons.save()

    def set_ownership(self, user=None, group=None, click_cancel=False, click_reset=False):
        sel.click(self.find_quadicon(False, False, False))
        cfg_btn('Set Ownership')
        if click_reset:
            action = form_buttons.reset
            msg_assert = lambda: flash.assert_message_match(
                'All changes have been reset')
        elif click_cancel:
            action = form_buttons.cancel
            msg_assert = lambda: flash.assert_success_message(
                'Set Ownership was cancelled by the user')
        else:
            action = form_buttons.save
            msg_assert = lambda: flash.assert_success_message(
                'Ownership saved for selected Virtual Machine')
        fill(set_ownership_form, {'user_name': user, 'group_name': group},
             action=action)
        msg_assert()

    def unset_ownership(self):
        # choose the vm code comes here
        sel.click(self.find_quadicon(False, False, False))
        cfg_btn('Set Ownership')
        fill(set_ownership_form, {'user_name': '<No Owner>', 'group_name': '<No Group>'},
            action=form_buttons.save)
        flash.assert_success_message('Ownership saved for selected Virtual Machine')

class Vm(Common):
    """Represents a VM in CFME

    Args:
        name: Name of the VM
        provider_crud: :py:class:`cfme.cloud.provider.Provider` object
        template_name: Name of the template to use for provisioning
    """

    class Snapshot(object):
        def __init__(self, name=None, description=None, memory=None, parent_vm=None):
            super(Vm.Snapshot, self).__init__()
            self.name = name
            self.description = description
            self.memory = memory
            self.vm = parent_vm

        def _nav_to_snapshot_mgmt(self):
            locator = ("//div[@class='dhtmlxInfoBarLabel' and " +
                       "contains(. , '\"Snapshots\" for Virtual Machine \"%s\"') ]" % self.name)
            if not sel.is_displayed(locator):
                self.vm.load_details()
                sel.click(details_page.infoblock.element("Properties", "Snapshots"))

        def does_snapshot_exist(self):
            self._nav_to_snapshot_mgmt()
            try:
                snapshot_tree.find_path_to(re.compile(r".*?\(Active\)$"))
                return True
            except CandidateNotFound:
                try:
                    snapshot_tree.find_path_to(re.compile(self.name))
                    return True
                except CandidateNotFound:
                    return False
            except NoSuchElementException:
                return False

        def wait_for_snapshot_active(self):
            self._nav_to_snapshot_mgmt()
            try:
                snapshot_tree.click_path(*snapshot_tree.find_path_to(re.compile(self.name)))
                if sel.is_displayed_text(self.name + " (Active)"):
                    return True
            except CandidateNotFound:
                return False

        def create(self):
            self._nav_to_snapshot_mgmt()
            toolbar.select('Create a new snapshot for this VM')
            fill(snapshot_form, {'name': self.name,
                                 'description': self.description,
                                 'snapshot_memory': self.memory
                                 },
                 action=snapshot_form.create_button)
            wait_for(self.does_snapshot_exist, num_sec=300, delay=20, fail_func=sel.refresh)

        def delete(self, cancel=False):
            self._nav_to_snapshot_mgmt()
            toolbar.select('Delete Snapshots', 'Delete Selected Snapshot', invokes_alert=True)
            sel.handle_alert(cancel=cancel)
            if not cancel:
                flash.assert_message_match(
                    'Delete Snapshot initiated for 1 VM and Instance from the CFME Database')

        def delete_all(self, cancel=False):
            self._nav_to_snapshot_mgmt()
            toolbar.select('Delete Snapshots', 'Delete All Existing Snapshots', invokes_alert=True)
            sel.handle_alert(cancel=cancel)
            if not cancel:
                flash.assert_message_match(
                    'Delete All Snapshots initiated for 1 VM and Instance from the CFME Database')

        def revert_to(self, cancel=False):
            self._nav_to_snapshot_mgmt()
            snapshot_tree.click_path(*snapshot_tree.find_path_to(re.compile(self.name)))
            toolbar.select('Revert to selected snapshot', invokes_alert=True)
            sel.handle_alert(cancel=cancel)
            flash.assert_message_match(
                'Revert To A Snapshot initiated for 1 VM and Instance from the CFME Database')

    # POWER CONTROL OPTIONS
    SUSPEND = "Suspend"
    POWER_ON = "Power On"
    POWER_OFF = "Power Off"
    GUEST_RESTART = "Restart Guest"
    GUEST_SHUTDOWN = "Shutdown Guest"
    RESET = "Reset"
    # POWER STATE
    STATE_ON = "on"
    STATE_OFF = "off"
    STATE_SUSPENDED = "suspended"

    def __init__(self, name, provider_crud, template_name=None):
        super(Vm, self).__init__()
        self.name = name
        self.template_name = template_name
        self.provider_crud = provider_crud

    def create_on_provider(self, timeout=900, find_in_cfme=False, **kwargs):
        """Create the VM on the provider

        Args:
            timeout: Number of seconds to wait for the VM to appear in CFME
                     Will not wait at all, if set to 0 (Defaults to ``900``)
        """
        deploy_template(self.provider_crud.key, self.name, self.template_name, **kwargs)
        if find_in_cfme:
            self.provider_crud.refresh_provider_relationships()
            self.wait_to_appear(timeout=timeout, load_details=False)

    def load_details(self, refresh=False):
        self._load_details(is_vm=True, refresh=refresh)

    def open_edit(self):
        self.load_details(refresh=True)
        toolbar.select("Configuration", "Edit this VM")

    def find_quadicon(self, do_not_navigate=False, mark=False, refresh=True, **kwargs):
        return self._find_quadicon(
            is_vm=True, do_not_navigate=do_not_navigate, mark=mark, refresh=refresh, **kwargs)

    def remove_from_cfme(self, cancel=False, from_details=False):
        """Removes a VM from CFME VMDB"""
        self._remove_from_cfme(is_vm=True, cancel=cancel, from_details=from_details)

    def power_control_from_provider(self, option):
        """Power control a vm from the provider

        Args:
            option: power control action to take against vm

        Raises:
            OptionNotAvailable: option parm must have proper value
        """
        if option == Vm.POWER_ON:
            self.provider_crud.get_mgmt_system().start_vm(self.name)
        elif option == Vm.POWER_OFF:
            self.provider_crud.get_mgmt_system().stop_vm(self.name)
        elif option == Vm.SUSPEND:
            self.provider_crud.get_mgmt_system().suspend_vm(self.name)
        # elif reset:
        # elif shutdown:
        else:
            raise OptionNotAvailable(option + " is not a supported action")

    def power_control_from_cfme(self, option, cancel=True, from_details=False):
        """Power controls a VM from within CFME

        Args:
            option: corresponds to option values under the power button
            cancel: Whether or not to cancel the power operation on confirmation
            from_details: Whether or not to perform action from vm details page
        """
        if (self.is_pwr_option_available_in_cfme(option=option, from_details=from_details)):
                pwr_btn(option, invokes_alert=True)
                sel.handle_alert(cancel=cancel)
                logger.info(
                    "Power control action of vm %s, option %s, cancel %s executed" %
                    (self.name, option, str(cancel)))
        else:
            raise OptionNotAvailable(option + " is not visible or enabled")

    def is_pwr_option_available_in_cfme(self, option, from_details=False):
        """Checks to see if a power option is available on the VM

        Args:
            option: corresponds to option values under the power button, preferred approach
                is to use Vm option constansts
            from_details: Whether or not to perform action from vm details page
        """
        self._method_helper(from_details=from_details)
        try:
            return not toolbar.is_greyed('Power', option)
        except NoSuchElementException:
            return False

    def retire(self):
        sel.force_navigate("infra_vm_by_name", context={'vm': self})
        lcl_btn("Retire this VM", invokes_alert=True)
        sel.handle_alert()
        flash.assert_success_message(
            "Retire initiated for 1 VM and Instance from the CFME Database")

    def migrate_vm(self, email=None, first_name=None, last_name=None,
                   host_name=None, datastore_name=None):
        sel.force_navigate("infra_vm_by_name", context={'vm': self})
        lcl_btn("Migrate this VM")
        first_name = first_name or fauxfactory.gen_alphanumeric()
        last_name = last_name or fauxfactory.gen_alphanumeric()
        email = email or "{}@{}.test".format(first_name, last_name)
        try:
            prov_data = cfme_data["management_systems"][self.provider_crud.key]["provisioning"]
        except (KeyError, IndexError):
            raise ValueError("You have to specify the correct options in cfme_data.yaml")
        provisioning_data = {
            "first_name": first_name,
            "last_name": last_name,
            "email": email,
            "host_name": {"name": prov_data.get("host")},
            "datastore_name": {"name": prov_data.get("datastore")}
        }
        from cfme.provisioning import provisioning_form
        fill(provisioning_form, provisioning_data, action=provisioning_form.submit_button)

    def clone_vm(self, email=None, first_name=None, last_name=None,
                 vm_name=None, provision_type=None):
        sel.force_navigate("infra_vm_by_name", context={'vm': self})
        lcl_btn("Clone this VM")
        first_name = first_name or fauxfactory.gen_alphanumeric()
        last_name = last_name or fauxfactory.gen_alphanumeric()
        email = email or "{}@{}.test".format(first_name, last_name)
        try:
            prov_data = cfme_data["management_systems"][self.provider_crud.key]["provisioning"]
        except (KeyError, IndexError):
            raise ValueError("You have to specify the correct options in cfme_data.yaml")
        provisioning_data = {
            "first_name": first_name,
            "last_name": last_name,
            "email": email,
            "provision_type": provision_type,
            "vm_name": vm_name,
            "host_name": {"name": prov_data.get("host")},
            "datastore_name": {"name": prov_data.get("datastore")},
            "vlan": prov_data.get("vlan")
        }
        from cfme.provisioning import provisioning_form
        fill(provisioning_form, provisioning_data, action=provisioning_form.submit_button)

    @property
    def retirement_date(self):
        """Returns the retirement date of the selected machine.

        Returns:
            :py:class:`NoneType` if there is none, or :py:class:`utils.timeutil.parsetime`
        """
        date_str = self.get_detail(properties=("Lifecycle", "Retirement Date")).strip()
        if date_str.lower() == "never":
            return None
        return parsetime.from_american_date_only(date_str)

    def set_retirement_date(self, when, warn=None):
        """Sets the retirement date for this Vm object.

        It incorporates some magic to make it work reliably since the retirement form is not very
        pretty and it can't be just "done".

        Args:
            when: When to retire. :py:class:`str` in format mm/dd/yy of
                :py:class:`datetime.datetime` or :py:class:`utils.timeutil.parsetime`.
            warn: When to warn, fills the select in the form in case the ``when`` is specified.
        """
        self.load_details()
        lcl_btn("Set Retirement Date")
        sel.wait_for_element("#miq_date_1")
        if when is None:
            try:
                wait_for(lambda: sel.is_displayed(retire_remove_button), num_sec=5, delay=0.2)
                sel.click(retire_remove_button)
                wait_for(lambda: not sel.is_displayed(retire_remove_button), num_sec=10, delay=0.2)
                sel.click(form_buttons.save)
            except TimedOutError:
                pass
        else:
            if sel.is_displayed(retire_remove_button):
                sel.click(retire_remove_button)
                wait_for(lambda: not sel.is_displayed(retire_remove_button), num_sec=15, delay=0.2)
            fill(retire_form.date_retire, when)
            wait_for(lambda: sel.is_displayed(retire_remove_button), num_sec=15, delay=0.2)
            if warn is not None:
                fill(retire_form.warn, warn)
            sel.click(form_buttons.save)

    def publish_to_template(self, template_name, email=None, first_name=None, last_name=None):
        self.load_details()
        lcl_btn("Publish this VM to a Template")
        first_name = first_name or fauxfactory.gen_alphanumeric()
        last_name = last_name or fauxfactory.gen_alphanumeric()
        email = email or "{}@{}.test".format(first_name, last_name)
        try:
            prov_data = cfme_data["management_systems"][self.provider_crud.key]["provisioning"]
        except (KeyError, IndexError):
            raise ValueError("You have to specify the correct options in cfme_data.yaml")

        provisioning_data = {
            "first_name": first_name,
            "last_name": last_name,
            "email": email,
            "vm_name": template_name,
            "host_name": {"name": prov_data.get("host")},
            "datastore_name": {"name": prov_data.get("datastore")},
        }
        from cfme.provisioning import provisioning_form
        fill(provisioning_form, provisioning_data, action=provisioning_form.submit_button)
        cells = {'Description': 'Publish from [%s] to [%s]' % (self.name, template_name)}
        row, __ = wait_for(
            requests.wait_for_request, [cells], fail_func=requests.reload, num_sec=900, delay=20)
        return Template(template_name, self.provider_crud)

    def wait_for_vm_state_change(
            self, desired_state=None, timeout=300, from_details=False):
        """Wait for VM to come to desired state.

        This function waits just the needed amount of time thanks to wait_for.

        Args:
            desired_state: on, off, suspended... corresponds to values in cfme, preferred approach
                is to use Vm.STATE_* constansts
            timeout: Specify amount of time (in seconds) to wait
        Raises:
            TimedOutError:
                When VM does not come up to desired state in specified period of time.
            NoVmFound:
                When unable to find the VM passed
        """
        def _looking_for_state_change():
            if from_details:
                self.load_details(refresh=True)
                detail_t = ("Power Management", "Power State")
                return self.get_detail(properties=detail_t) == desired_state
            else:
                return self.find_quadicon().state == 'currentstate-' + desired_state
        return wait_for(_looking_for_state_change, num_sec=timeout, delay=30)

    class CfmeRelationship(object):

        relationship_form = Form(
            fields=[
                ('server_select', Select("//*[@id='server_id']")),
                ('save_button', form_buttons.save),
                ('reset_button', form_buttons.reset),
                ('cancel_button', form_buttons.cancel)
            ])

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

        def navigate(self):
            self.o.load_details()
            cfg_btn('Edit Management Engine Relationship')

        def is_relationship_set(self):
            return "<Not a Server>" not in self.get_relationship()

        def get_relationship(self):
            self.navigate()
            rel = str(self.relationship_form.server_select.all_selected_options[0].text)
            form_buttons.cancel()
            return rel

        def set_relationship(self, server_name, server_id, click_cancel=False):
            self.navigate()
            option = "%s (%d)" % (server_name, server_id)

            if click_cancel:
                fill(self.relationship_form, {'server_select': option},
                     action=self.relationship_form.cancel_button)
            else:
                fill(self.relationship_form, {'server_select': option},
                     action=self.relationship_form.save_button)
                # something weird going on where changing the select doesn't POST to undim save
                sel.wait_for_ajax()
                if self.relationship_form.save_button.is_dimmed:
                    logger.warning("Worked around dimmed save button")
                    sel.browser().execute_script(
                        "$j.ajax({type: 'POST', url: '/vm_infra/evm_relationship_field_changed',"
                        " data: {'server_id':'%s'}})" % (server_id))
                    sel.click(form_buttons.FormButton(
                        "Save Changes", dimmed_alt="Save", force_click=True))
                flash.assert_success_message("Management Engine Relationship saved")

class Template(Common):

    def __init__(self, name, provider_crud):
        super(Template, self).__init__()
        self.name = name
        self.template_name = name
        self.provider_crud = provider_crud

    def load_details(self, refresh=False):
        self._load_details(refresh=refresh, is_vm=False)

    def find_quadicon(self, do_not_navigate=False, mark=False, refresh=True, **kwargs):
        return self._find_quadicon(
            is_vm=False, do_not_navigate=do_not_navigate, mark=mark, refresh=refresh, **kwargs)

    def remove_from_cfme(self, cancel=False, from_details=False):
        """Removes a VM from CFME VMDB"""
        self._remove_from_cfme(is_vm=False, cancel=cancel, from_details=from_details)

class Genealogy(object):
    """Class, representing genealogy of an infra object with possibility of data retrieval
    and comparison.

    Args:
        o: The :py:class:`Vm` or :py:class:`Template` object.
    """
    genealogy_tree = CheckboxTree({
        version.LOWEST: "//div[@id='treebox']/div/table",
        "5.3": "//div[@id='genealogy_treebox']/ul",
    })

    section_comparison_tree = CheckboxTree("//div[@id='all_sections_treebox']/div/table")
    apply_button = form_buttons.FormButton("Apply sections")

    mode_mapping = {
        "exists": "Exists Mode",
        "details": "Details Mode",
    }

    attr_mapping = {
        "all": "All Attributes",
        "different": "Attributes with different values",
        "same": "Attributes with same values",
    }

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

    def navigate(self):
        self.o.load_details()
        sel.click(details_page.infoblock.element("Relationships", "Genealogy"))

    def compare(self, *objects, **kwargs):
        """Compares two or more objects in the genealogy.

        Args:
            *objects: :py:class:`Vm` or :py:class:`Template` or :py:class:`str` with name.

        Keywords:
            sections: Which sections to compare.
            attributes: `all`, `different` or `same`. Default: `all`.
            mode: `exists` or `details`. Default: `exists`."""
        sections = kwargs.get("sections", None)
        attributes = kwargs.get("attributes", "all").lower()
        mode = kwargs.get("mode", "exists").lower()
        assert len(objects) >= 2, "You must specify at least two objects"
        objects = map(lambda o: o.name if isinstance(o, (Vm, Template)) else o, objects)
        self.navigate()
        for obj in objects:
            if not isinstance(obj, list):
                path = self.genealogy_tree.find_path_to(obj)
            self.genealogy_tree.check_node(*path)
        toolbar.select("Compare selected VMs")
        # COMPARE PAGE
        flash.assert_no_errors()
        if sections is not None:
            map(lambda path: self.section_comparison_tree.check_node(*path), sections)
            sel.click(self.apply_button)
            flash.assert_no_errors()
        # Set requested attributes sets
        toolbar.select(self.attr_mapping[attributes])
        # Set the requested mode
        toolbar.select(self.mode_mapping[mode])

    @property
    def tree(self):
        """Returns contents of the tree with genealogy"""
        self.navigate()
        return self.genealogy_tree.read_contents()

    @property
    def ancestors(self):
        """Returns list of ancestors of the represented object."""
        self.navigate()
        path = self.genealogy_tree.find_path_to(re.compile(r"^.*?\(Selected\)$"))
        if not path:
            raise ValueError("Something wrong happened, path not found!")
        processed_path = []
        for step in path[:-1]:
            # We will remove the (parent) and (Selected) suffixes
            processed_path.append(re.sub(r"\s*(?:\(Current\)|\(Parent\))$", "", step))
        return processed_path

def _method_setup(vm_names, provider_crud=None):
    """ Reduces some redundant code shared between methods """
    if isinstance(vm_names, basestring):
        vm_names = [vm_names]

    if provider_crud:
        provider_crud.load_all_provider_vms()
    else:
        sel.force_navigate('infra_vms')
    if paginator.page_controls_exist():
        paginator.results_per_page(1000)
    for vm_name in vm_names:
        sel.check(Quadicon(vm_name, 'vm').checkbox())

def find_quadicon(vm_name, do_not_navigate=False):
    """Find and return a quadicon belonging to a specific vm

    Args:
        vm: vm name as displayed at the quadicon
    Returns: :py:class:`cfme.web_ui.Quadicon` instance
    """
    if not do_not_navigate:
        sel.force_navigate('infra_vms')
    if not paginator.page_controls_exist():
        raise VmNotFound("VM '{}' not found in UI!".format(vm_name))

    paginator.results_per_page(1000)
    for page in paginator.pages():
        quadicon = Quadicon(vm_name, "vm")
        if sel.is_displayed(quadicon):
            return quadicon
    else:
        raise VmNotFound("VM '{}' not found in UI!".format(vm_name))

def remove(vm_names, cancel=True, provider_crud=None):
    """Removes multiple VMs from CFME VMDB

    Args:
        vm_names: List of VMs to interact with
        cancel: Whether to cancel the deletion, defaults to True
        provider_crud: provider object where vm resides on (optional)
    """
    _method_setup(vm_names, provider_crud)
    cfg_btn('Remove selected items from the VMDB', invokes_alert=True)
    sel.handle_alert(cancel=cancel)

def wait_for_vm_state_change(vm_name, desired_state, timeout=300, provider_crud=None):
    """Wait for VM to come to desired state.

    This function waits just the needed amount of time thanks to wait_for.

    Args:
        vm_name: Displayed name of the VM
        desired_state: 'on' or 'off'
        timeout: Specify amount of time (in seconds) to wait until TimedOutError is raised
        provider_crud: provider object where vm resides on (optional)
    """
    def _looking_for_state_change():
        toolbar.refresh()
        find_quadicon(vm_name, do_not_navigate=False).state == 'currentstate-' + desired_state

    _method_setup(vm_name, provider_crud)
    return wait_for(_looking_for_state_change, num_sec=timeout)

def is_pwr_option_visible(vm_names, option, provider_crud=None):
    """Returns whether a particular power option is visible.

    Args:
        vm_names: List of VMs to interact with, if from_details=True is passed, only one VM can
            be passed in the list.
        option: Power option param.
        provider_crud: provider object where vm resides on (optional)
    """
    _method_setup(vm_names, provider_crud)
    try:
        toolbar.is_greyed('Power', option)
        return True
    except NoSuchElementException:
        return False

def is_pwr_option_enabled(vm_names, option, provider_crud=None):
    """Returns whether a particular power option is enabled.

    Args:
        vm_names: List of VMs to interact with
        provider_crud: provider object where vm resides on (optional)
        option: Power option param.

    Raises:
        NoOptionAvailable:
            When unable to find the power option passed
    """
    _method_setup(vm_names, provider_crud)
    try:
        return not toolbar.is_greyed('Power', option)
    except NoSuchElementException:
        raise OptionNotAvailable("No such power option (" + str(option) + ") is available")

def do_power_control(vm_names, option, provider_crud=None, cancel=True):
    """Executes a power option against a list of VMs.

    Args:
        vm_names: List of VMs to interact with
        option: Power option param.
        provider_crud: provider object where vm resides on (optional)
        cancel: Whether or not to cancel the power control action
    """
    _method_setup(vm_names, provider_crud)

    if (is_pwr_option_visible(vm_names, provider_crud=provider_crud, option=option) and
            is_pwr_option_enabled(vm_names, provider_crud=provider_crud, option=option)):
                pwr_btn(option, invokes_alert=True)
                sel.handle_alert(cancel=cancel)

def refresh_relationships(vm_names, provider_crud=None, cancel=True):
    """Executes a refresh relationships action against a list of VMs.

    Args:
        vm_names: List of VMs to interact with
        provider_crud: provider object where vm resides on (optional)
        cancel: Whether or not to cancel the refresh relationships action
    """
    _method_setup(vm_names, provider_crud)
    cfg_btn('Refresh Relationships and Power States', invokes_alert=True)
    sel.handle_alert(cancel=cancel)

def perform_smartstate_analysis(vm_names, provider_crud=None, cancel=True):
    """Executes a refresh relationships action against a list of VMs.

    Args:
        vm_names: List of VMs to interact with
        provider_crud: provider object where vm resides on (optional)
        cancel: Whether or not to cancel the refresh relationships action
    """
    _method_setup(vm_names, provider_crud)
    cfg_btn('Perform SmartState Analysis', invokes_alert=True)
    sel.handle_alert(cancel=cancel)

def get_first_vm_title(do_not_navigate=False):
    if not do_not_navigate:
        sel.force_navigate('infra_vms')
    return Quadicon.get_first_quad_title()

def get_all_vms(do_not_navigate=False):
    """Returns list of all vms"""
    if not do_not_navigate:
        sel.force_navigate('infra_vms')
    vms = set([])
    if not paginator.page_controls_exist():
        return vms

    paginator.results_per_page(1000)
    for page in paginator.pages():
        try:
            for page in paginator.pages():
                for title in sel.elements(QUADICON_TITLE_LOCATOR):
                    vms.add(sel.get_attribute(title, "title"))
        except sel.NoSuchElementException:
            pass
    return vms

def get_number_of_vms(do_not_navigate=False):
    """
    Returns the total number of VMs visible to the user,
    including those archived or orphaned
    """
    logger.info("Getting number of vms")
    if not do_not_navigate:
        sel.force_navigate('infra_vms')
    if not paginator.page_controls_exist():
        logger.debug("No page controls")
        return 0
    total = paginator.rec_total()
    logger.debug("Number of VMs: {}".format(total))
    return int(total)

def _assign_unassign_policy_profiles(vm_name, assign, *policy_profile_names, **kwargs):
    """DRY function for managing policy profiles.

    See :py:func:`assign_policy_profiles` and :py:func:`assign_policy_profiles`

    Args:
        vm_name: Name of the VM.
        assign: Wheter to assign or unassign.
        policy_profile_names: :py:class:`str` with Policy Profile names.
    """
    if kwargs.get("via_details", True):
        sel.click(find_quadicon(vm_name))
    else:
        _method_setup(vm_name, **kwargs)
    toolbar.select("Policy", "Manage Policies")
    for policy_profile in policy_profile_names:
        if assign:
            manage_policies_tree.check_node(policy_profile)
        else:
            manage_policies_tree.uncheck_node(policy_profile)
    sel.click(form_buttons.save)

def assign_policy_profiles(vm_name, *policy_profile_names, **kwargs):
    """Assign Policy Profiles to specified VM.

    Args:
        vm_name: Name of the VM.
        policy_profile_names: :py:class:`str` with Policy Profile names.
    """
    return _assign_unassign_policy_profiles(vm_name, True, *policy_profile_names, **kwargs)

def unassign_policy_profiles(vm_name, *policy_profile_names, **kwargs):
    """Unassign Policy Profiles to specified VM.

    Args:
        vm_name: Name of the VM.
        policy_profile_names: :py:class:`str` with Policy Profile names.
    """
    return _assign_unassign_policy_profiles(vm_name, False, *policy_profile_names, **kwargs)


# Copyright cozybit, Inc. - 2014
# All rights reserved

"""
An adapter used to map the wpa_supplicant tests cases for mesh onto real
hardware using wtf.
"""

import importlib
import os
import Queue as queue
import re
import subprocess
import sys
import threading
import time
import unittest
import wtf

WTFCONFIG = wtf.conf
STATIONS = WTFCONFIG.mps
ADAPTERS = None
THIS_FILE = os.path.abspath(__file__)
THIS_DIR = os.path.dirname(THIS_FILE)

wpa_cli_error = ("Expected exit status of 0 from wpa_cli, got %d instead "
                 "(for command: '%s')")

def setUp(self):
    global ADAPTERS
    for station in STATIONS:
        station.comm.send_cmd("stop wpa_supplicant")
        station.comm.send_cmd("svc wifi disable")
    time.sleep(2)
    ADAPTERS = [WpasTestAdapter(sta) for sta in STATIONS]

def tearDown(self):
    for n in WTFCONFIG.nodes:
        n.stop()

class AdbCmdGen(object):
    @staticmethod
    def clear_logcat(adb_id):
        return ["adb", "-s", adb_id, "logcat", "-c"]

    @staticmethod
    def start_logcat(adb_id):
        return ["adb", "-s", adb_id, "logcat"]

    @staticmethod
    def wpa_cli(adb_id):
        return ["adb", "-s", adb_id, "shell", "wpa_cli",
                "-p/data/misc/wifi", "-iwlan0"]

    @staticmethod
    def dmesg(adb_id):
        return ["adb", "-s", adb_id, "shell", "cat", "/proc/kmsg"]

class WpaCliMonitor (object):

    EVENT_TRIM = re.compile(r'^[\s>]+')

    def __init__(self, q, adb_id, device_id):
        self._q = q
        self._adb_id = adb_id
        self._device_id = device_id

        self._wpacli_process = None
        self._logcat_process = None
        self._dmesg_process = None

        self._wpacli_logfp = None
        self._logcat_logfp = None
        self._dmesg_logfp = None

        self._wpacli_thread = None
        self._wpacli_stderr_thread = None
        self._logcat_out_thread = None
        self._logcat_err_thread = None
        self._dmesg_out_thread = None
        self._dmesg_err_thread = None

    def _gen(self):
        while True:
            line = self._wpacli_process.stdout.readline()
            if not line:
                return
            self._wpacli_logfp.write(line)
            line = line.strip()
            #print 'LaunchMonitorProcess:', repr(line)
            split = self.EVENT_TRIM.split(line)
            if len(split) > 1:
                line = split[1]
            if line.startswith("<3>"):
                yield line[3:]
            else:
                continue

    @staticmethod
    def _dump_fp_to_log(fp, logfp):
        while True:
            line = fp.readline()
            if not line:
                break
            logfp.write(line)

    def _wpa_cli_run(self, test_name):
        self._wpacli_process = subprocess.Popen(
            AdbCmdGen.wpa_cli(self._adb_id),
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE)

        subprocess.call(AdbCmdGen.clear_logcat(self._adb_id))

        self._logcat_process = (
            subprocess.Popen(AdbCmdGen.start_logcat(self._adb_id),
                             stdout=subprocess.PIPE,
                             stderr=subprocess.PIPE))

        self._dmesg_process = (
            subprocess.Popen(AdbCmdGen.dmesg(self._adb_id),
                             stdout=subprocess.PIPE,
                             stderr=subprocess.PIPE))

        self._wpacli_logfp = open("wpa_cli.%s.%s.log" %
                                  (test_name, self._device_id), "w")

        self._logcat_logfp = open("logcat.%s.%s.log" %
                                  (test_name, self._device_id), "w")

        self._dmesg_logfp = open("dmesg.%s.%s.log" %
                                 (test_name, self._device_id), "w")

        self._wpacli_stderr_thread = threading.Thread(
            target=self._dump_fp_to_log,
            args=(self._wpacli_process.stderr,
                  self._wpacli_logfp))

        self._wpacli_stderr_thread.start()

        self._logcat_out_thread = threading.Thread(
            target=self._dump_fp_to_log,
            args=(self._logcat_process.stdout,
                  self._logcat_logfp))

        self._logcat_out_thread.start()

        self._logcat_err_thread = threading.Thread(
            target=self._dump_fp_to_log,
            args=(self._logcat_process.stderr,
                  self._logcat_logfp))

        self._logcat_err_thread.start()

        self._dmesg_out_thread = threading.Thread(
            target=self._dump_fp_to_log,
            args=(self._dmesg_process.stdout,
                  self._dmesg_logfp))

        self._dmesg_out_thread.start()

        self._dmesg_err_thread = threading.Thread(
            target=self._dump_fp_to_log,
            args=(self._dmesg_process.stderr,
                  self._dmesg_logfp))

        self._dmesg_err_thread.start()

        for line in self._gen():
            self._q.put(line)

    def start(self, test_name):
        self._wpacli_thread = threading.Thread(target=self._wpa_cli_run,
                                               args=(test_name,))
        self._wpacli_thread.start()

    def stopMonitor(self):
        self._wpacli_process.kill()
        self._logcat_process.kill()
        self._dmesg_process.kill()
        self._wpacli_thread.join()
        self._wpacli_stderr_thread.join()
        self._logcat_out_thread.join()
        self._logcat_err_thread.join()
        self._dmesg_out_thread.join()
        self._dmesg_err_thread.join()
        self._wpacli_logfp.close()
        self._logcat_logfp.close()
        self._dmesg_logfp.close()

class EventMonitor (object):

    def __init__(self, sta):
        self._sta = sta
        self._q = queue.Queue()
        self._monitor = WpaCliMonitor(self._q,
                                      self._sta.comm.get_adb_id(),
                                      self._sta.comm.get_device_id())

    def waitForEvent(self, eventName):
        while True:
            line = self._q.get(timeout=10)
            #print 'waitForEvent:', line
            if line.startswith(eventName):
                return line

    def startMonitor(self, test_name):
        self._monitor.start(test_name=test_name)

    def stopMonitor(self):
        self._monitor.stopMonitor()

class WpasTestAdapter (object):

    def __init__(self, sta):
        self._sta = sta
        self._comm = sta.comm
        self._monitor = EventMonitor(sta)

    def startWpaCliMonitor(self, test_name):
        self._monitor.startMonitor(test_name)

    def stopWpaCliMonitor(self):
        self._monitor.stopMonitor()

    def getSTA(self):
        return self._sta

    def _send_to_wpa_cli(self, cmd):
        cmd = "wpa_cli -p/data/misc/wifi -iwlan0 " + cmd
        #print cmd
        code, resp = self._comm.send_cmd(cmd)
        if code != 0:
            raise ValueError(wpa_cli_error % (code, cmd))
        return str.join('', resp)

    def ping_wpa_supplicant(self):
        cmd = "wpa_cli -p/data/misc/wifi -iwlan0 ping"
        code, resp = self._comm.send_cmd(cmd)
        return code == 0 and str.join('', resp) == "PONG"

    def _send_cmd_bool_resp(self, cmd):
        resp = self._send_to_wpa_cli(cmd).strip()
        if resp != "OK":
            raise ValueError(resp)
        return resp

    def _send_cmd_int_resp(self, cmd):
        resp = self._send_to_wpa_cli(cmd).strip()
        if not resp.isdigit():
            raise ValueError(resp)
        return resp

    def dump_monitor(self):
        pass

    def request(self, req):
        return self._send_to_wpa_cli(req)

    def add_network(self):
        return self._send_cmd_int_resp("add_network")

    def set_network(self, networkId, prop, val):
        return self._send_cmd_bool_resp(
            "set_network " + networkId + " " + prop + " " + val)

    def set_network_quoted(self, networkId, prop, val):
        return self.set_network(networkId, prop, '\'"' + val + '"\'')

    def mesh_group_add(self, meshId):
        return self._send_cmd_bool_resp("mesh_group_add " + meshId)

    def mesh_group_remove(self):
        return self._send_cmd_bool_resp("mesh_group_remove wlan0")

    def remove_network(self, networkId):
        return self._send_cmd_bool_resp("remove_network " + networkId)

    def wait_event(self, eventList):
        for event in eventList:
            ev = self._monitor.waitForEvent(event)
            if ev is not None:
                return ev

def adb_ping_check(dev1, dev2):

    dev1.getSTA().comm.send_cmd("ip addr flush dev wlan0")
    dev1.getSTA().comm.send_cmd("ip addr add 10.0.0.2/24 dev wlan0")

    dev2.getSTA().comm.send_cmd("ip addr flush dev wlan0")
    dev2.getSTA().comm.send_cmd("ip addr add 10.0.0.3/24 dev wlan0")

    code, outp = dev1.getSTA().comm.send_cmd("ping -c5 10.0.0.3")
    if code != 0:
        raise Exception("No connectivity from dev1 to dev2")

    return

class WpasMeshTest(unittest.TestCase):

    def __init__(self, *args, **kw):
        unittest.TestCase.__init__(self, *args, **kw)
        fp = os.path.join(THIS_DIR, "..", "submodules", "wpa_supplicant",
                          "tests", "hwsim")
        sys.path.append(fp)
        module = importlib.import_module("test_wpas_mesh")
        self._tests = {}
        for name in module.__dict__.keys():
            if name.startswith("test_") or name.startswith("_test"):
                self._tests[name] = module.__dict__[name]

    def setUp(self):
        for adapter in ADAPTERS:
            adapter.getSTA().comm.send_cmd("start wpa_supplicant")
            while not adapter.ping_wpa_supplicant():
                time.sleep(0.5)
            adapter.startWpaCliMonitor(self.id())

    def tearDown(self):
        for adapter in ADAPTERS:
            adapter.stopWpaCliMonitor()
            adapter.getSTA().comm.send_cmd("stop wpa_supplicant")

    def test_wpas_add_set_remove_support(self):
        self._tests['test_wpas_add_set_remove_support'](ADAPTERS)

    def test_wpas_mesh_group_added(self):
        self._tests['test_wpas_mesh_group_added'](ADAPTERS)

    def test_wpas_mesh_group_remove(self):
        self._tests['test_wpas_mesh_group_remove'](ADAPTERS)

    def test_wpas_mesh_peer_connected(self):
        self._tests['test_wpas_mesh_peer_connected'](ADAPTERS)

    def test_wpas_mesh_peer_disconnected(self):
        self._tests['test_wpas_mesh_peer_disconnected'](ADAPTERS)

    def test_wpas_mesh_mode_scan(self):
        self._tests['test_wpas_mesh_mode_scan'](ADAPTERS)

    def test_wpas_mesh_open(self):
        self._tests['_test_wpas_mesh_open'](ADAPTERS, None, adb_ping_check)

    def test_wpas_mesh_open_no_auto(self):
        self._tests['_test_wpas_mesh_open_no_auto'](ADAPTERS, None,
                                                    adb_ping_check)

    def test_wpas_mesh_secure(self):
        self._tests['_test_wpas_mesh_secure'](ADAPTERS, None, adb_ping_check)

    def test_wpas_mesh_secure_no_auto(self):
        self._tests['_test_wpas_mesh_secure_no_auto'](ADAPTERS, None,
                                                      adb_ping_check)


#
# Copyright 2013 Quantopian, Inc.
#
# 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.
from datetime import timedelta

import pytz
import math

from functools import partial

import numpy as np

from logbook import Processor

def transact_stub(slippage, commission, event, open_orders):
    """
    This is intended to be wrapped in a partial, so that the
    slippage and commission models can be enclosed.
    """
    def inject_algo_dt(record):
        if not 'algo_dt' in record.extra:
            record.extra['algo_dt'] = event['dt']

    with Processor(inject_algo_dt).threadbound():

        transaction = slippage.simulate(event, open_orders)
        if transaction and not np.allclose(transaction.amount, 0):
            direction = math.copysign(1, transaction.amount)
            per_share, total_commission = commission.calculate(transaction)
            transaction.price = transaction.price + (per_share * direction)
            transaction.commission = total_commission
        return transaction

def transact_partial(slippage, commission):
    return partial(transact_stub, slippage, commission)

class Transaction(object):

    def __init__(self, sid, amount, dt, price, commission=None):
        self.sid = sid
        self.amount = amount
        self.dt = dt
        self.price = price
        self.commission = commission

    def __getitem__(self, name):
        return self.__dict__[name]

def create_transaction(sid, amount, price, dt):

    txn = {
        'sid': sid,
        'amount': int(amount),
        'dt': dt,
        'price': price,
    }

    transaction = Transaction(**txn)
    return transaction

class VolumeShareSlippage(object):

    def __init__(self,
                 volume_limit=.25,
                 price_impact=0.1,
                 delay=timedelta(minutes=1)):

        self.volume_limit = volume_limit
        self.price_impact = price_impact
        self.delay = delay

    def simulate(self, event, open_orders):

        if np.allclose(event.volume, 0):
            #there are zero volume events bc some stocks trade
            #less frequently than once per minute.
            return None

        if event.sid in open_orders:
            orders = open_orders[event.sid]
            orders = sorted(orders, key=lambda o: o.dt)
            # Only use orders for the current day or before
            current_orders = filter(
                lambda o: o.dt + self.delay <= event.dt,
                orders)
        else:
            return None

        dt = event.dt
        total_order = 0
        simulated_amount = 0
        simulated_impact = 0.0

        for order in current_orders:

            open_amount = order.amount - order.filled

            if np.allclose(open_amount, 0):
                continue

            direction = math.copysign(1, open_amount)

            # if the stop price is reached, simply set stop to None
            # othrewise we skip this order with a continue
            if order.stop is not None:
                if (direction * (event.price - order.stop) < 0):
                    # convert stop -> limit or market
                    order.stop = None
                else:
                    continue

            # if the limit price is reached, we execute this order at
            # (event.price + simulated_impact)
            # we skip this order with a continue when the limit is not reached
            if order.limit is not None:
                # if limit conditions not met, then continue
                if (direction * (event.price - order.limit) > 0):
                    continue

            desired_order = total_order + open_amount

            volume_share = min(direction * (desired_order) / event.volume,
                               self.volume_limit)

            if np.allclose(volume_share, self.volume_limit):
                simulated_amount = \
                    int(self.volume_limit * event.volume * direction)
            else:
                # we can fill the entire desired order
                # let's not deal with floating-point errors
                simulated_amount = desired_order

            simulated_impact = (volume_share) ** 2 \
                * self.price_impact * direction * event.price

            order.filled += (simulated_amount - total_order)
            total_order = simulated_amount

            # we cap the volume share at configured % of a trade
            if np.allclose(volume_share, self.volume_limit):
                break

        filled_orders = [x for x in orders
                         if abs(x.amount - x.filled) > 0
                         and x.dt.day >= event.dt.day]

        open_orders[event.sid] = filled_orders

        if simulated_amount != 0:
            return create_transaction(
                event.sid,
                simulated_amount,
                # In the future, we may want to change the next line
                # for limit pricing
                event.price + simulated_impact,
                dt.replace(tzinfo=pytz.utc),
            )

class FixedSlippage(object):

    def __init__(self, spread=0.0):
        """
        Use the fixed slippage model, which will just add/subtract
        a specified spread spread/2 will be added on buys and subtracted
        on sells per share
        """
        self.spread = spread

    def simulate(self, event, open_orders):
        if event.sid in open_orders:
            orders = open_orders[event.sid]
            orders = sorted(orders, key=lambda o: o.dt)
        else:
            return None

        amount = 0
        for order in orders:
            # what if we have 2 orders, one for 100 shares long,
            # and one for 100 shares short
            # such as in a hedging scenario?
            amount += order.amount
            direction = math.copysign(1, amount)

            # if the stop price is reached, simply set stop to None
            # othrewise we skip this order with a continue
            if order.stop is not None:
                if (direction * (event.price - order.stop) < 0):
                    # convert stop -> limit or market
                    order.stop = None
                else:
                    continue

            # if the limit price is reached, we execute this order at
            # (event.price + simulated_impact)
            # we skip this order with a continue when the limit is not reached
            if order.limit is not None:
                # if limit conditions not met, then continue
                if (direction * (event.price - order.limit) > 0):
                    continue

        if np.allclose(amount, 0):
            return

        txn = create_transaction(
            event.sid,
            amount,
            event.price + (self.spread / 2.0 * direction),
            event.dt
        )

        open_orders[event.sid] = []

        return txn


#!/usr/bin/env python3
import os
import rospy

import cv2
import time
import copy
import pickle
import traceback
import numpy as np

import tensorflow
from object_detection.utils import label_map_util
from object_detection.utils import visualization_utils as viz_utils

from dynamic_reconfigure.server import Server
from db_tensorflow.cfg import DetectConfig

import message_filters
from sensor_msgs.msg import Image, CameraInfo

from image_geometry import PinholeCameraModel

from geometry_msgs.msg import PoseStamped
from geometry_msgs.msg import PoseWithCovariance
from geometry_msgs.msg import Vector3

from vision_msgs.msg import Detection2D
from vision_msgs.msg import Detection2DArray
from vision_msgs.msg import ObjectHypothesisWithPose

from visualization_msgs.msg import MarkerArray
from visualization_msgs.msg import Marker

from std_msgs.msg import ColorRGBA

import tf2_ros
import tf2_geometry_msgs

import ctypes
# a thread gets killed improperly within CvBridge without this causing segfaults
libgcc_s = ctypes.CDLL('libgcc_s.so.1')

from cv_bridge import CvBridge, CvBridgeError

class RollingAvgRateLogger:
    def __init__(self, rolling_avg_window):
        self.rolling_avg_window = rolling_avg_window
        self.data = np.zeros(self.rolling_avg_window)
        self.index = 0

    def append(self, value):
        self.data[self.index] = value
        self.index += 1
        if self.index >= len(self.data):
            self.index = 0

    def rate(self):
        rate = 1.0 / np.mean(self.data)
        return rate

    def avg(self):
        return np.mean(self.data)

class BoxDescription:
    def __init__(self):
        self.class_index = 0
        self.class_name = ""
        self.score = 0.0
        self.box = (0.0, 0.0, 0.0, 0.0)
        self.xcenter = 0.0
        self.ycenter = 0.0
        self.xsize = 0.0
        self.ysize = 0.0
        self.obj_radius = 0.0
        self.detect_radius = 0.0

        self.pose_stamped = PoseStamped()
        self.xedge = 0.0
        self.yedge = 0.0
        self.zedge = 0.0

class DodobotTensorflow:
    def __init__(self):
        self.node_name = "db_tensorflow"
        rospy.init_node(
            self.node_name
            # disable_signals=True
            # log_level=rospy.DEBUG
        )
        # rospy.on_shutdown(self.shutdown)

        self.image_topic_name = rospy.get_param("~image_topic", "color/image_raw")
        self.depth_topic_name = rospy.get_param("~depth_topic", "depth/image_raw")
        self.info_topic_name = rospy.get_param("~camera_info_topic", "color/camera_info")

        self.labels_path = rospy.get_param("~labels_path", "annotations/label_map.pbtxt")
        self.model_path = rospy.get_param("~model_path", "models/dodobot_objects_ssd_resnet50_v1_fpn")

        self.min_score_threshold = rospy.get_param("~min_score_threshold", 0.3)
        self.max_boxes_to_draw = rospy.get_param("~max_boxes_to_draw", 20)
        self.bounding_box_border_px = rospy.get_param("~bounding_box_border_px", 30)
        self.z_size_estimations = rospy.get_param("~z_size_estimations", None)
        marker_colors = rospy.get_param("~marker_colors", None)
        self.publish_in_robot_frame = rospy.get_param("~publish_in_robot_frame", True)
        self.robot_frame = rospy.get_param("~robot_frame", "base_link")
        self.marker_persistance = rospy.get_param("~marker_persistance_s", 0.5)
        self.default_z_size = rospy.get_param("~default_z_size", 0.05)
        self.min_valid_z = rospy.get_param("~min_valid_z", 0.03)
        self.max_valid_z = rospy.get_param("~max_valid_z", 3.0)

        self.marker_persistance = rospy.Duration(self.marker_persistance)
        self.marker_colors = {name: ColorRGBA(*args) for name, args in marker_colors.items()}

        self.image_sub = message_filters.Subscriber(self.image_topic_name, Image)
        self.depth_sub = message_filters.Subscriber(self.depth_topic_name, Image)

        self.detect_fn = None
        self.category_index = None

        self.dyn_server = Server(DetectConfig, self.dyn_callback)

        self.time_sync_sub = message_filters.TimeSynchronizer([self.image_sub, self.depth_sub], 1)
        # self.time_sync_sub = message_filters.ApproximateTimeSynchronizer([self.image_sub, self.depth_sub], 10, 0.005)

        self.visualization_image_pub = rospy.Publisher("detect_overlay", Image, queue_size=1)

        self.detect_pub = rospy.Publisher("detections", Detection2DArray, queue_size=10)
        self.marker_pub = rospy.Publisher("obj_markers", MarkerArray, queue_size=25)

        self.bridge = CvBridge()
        self.camera_model = PinholeCameraModel()

        self.detect_rate_logger = RollingAvgRateLogger(10)
        self.update_rate_logger = RollingAvgRateLogger(10)
        self.camera_rate_logger = RollingAvgRateLogger(10)
        self.depth_rate_logger = RollingAvgRateLogger(10)

        self.tf_buffer = tf2_ros.Buffer()
        self.tf_listener = tf2_ros.TransformListener(self.tf_buffer)

        # with open("/home/ben/detection.pkl", 'rb') as file:
        #     self.dummy_detections = pickle.load(file)

        self.generate_detect_fn()

        self.time_sync_sub.registerCallback(self.rgbd_callback)
        self.info_sub = rospy.Subscriber(self.info_topic_name, CameraInfo, self.camera_info_callback, queue_size=5)

        rospy.loginfo("%s init done" % self.node_name)

    def generate_detect_fn(self):
        if not os.path.isfile(self.labels_path):
            raise FileNotFoundError("Labels path doesn't exist: %s" % self.labels_path)

        if not os.path.isdir(self.model_path):
            raise FileNotFoundError("Model path doesn't exist: %s" % self.model_path)

        category_index = label_map_util.create_category_index_from_labelmap(self.labels_path, use_display_name=True)

        rospy.loginfo("Loading model...")
        start_time = time.time()
        # Load saved model and build the detection function
        detect_fn = tensorflow.saved_model.load(self.model_path)
        # detect_fn = None

        end_time = time.time()
        elapsed_time = end_time - start_time
        rospy.loginfo("Model took %0.2f seconds to load" % (elapsed_time))

        self.detect_fn = detect_fn
        self.category_index = category_index

        rospy.loginfo("Priming the pump...")
        color_image_np = np.zeros((540, 960, 3), dtype=np.uint8)
        start_time = time.time()
        self.run_detection(color_image_np)
        end_time = time.time()
        elapsed_time = end_time - start_time

        rospy.loginfo("Priming took %0.2f seconds" % (elapsed_time))

    def rgbd_callback(self, color_image, depth_image):
        t0 = rospy.Time.now()
        detections = self.detection_pipeline(color_image)

        detect_array_msg = Detection2DArray()
        markers_msg = MarkerArray()

        for index in range(detections["num_detections"]):
            score = detections["detection_scores"][index]
            if score < self.min_score_threshold:
                continue
            box = detections["detection_boxes"][index]
            class_index = detections["detection_classes"][index]
            desc = self.depth_pipeline(box, score, class_index, depth_image)

            if desc.pose_stamped is None or not (self.min_valid_z <= desc.pose_stamped.pose.position.z <= self.max_valid_z):
                continue

            detect_msg = self.get_detect_msg(desc)

            if self.publish_in_robot_frame:
                new_pose = self.object_to_robot_frame(desc.pose_stamped, color_image.header)
                if new_pose is None:
                    continue
                else:
                    desc.pose_stamped = new_pose

            detection_pose = PoseWithCovariance()
            detection_pose.pose = desc.pose_stamped.pose
            detect_msg.results[0].pose = detection_pose
            detect_msg.header = desc.pose_stamped.header

            markers = self.make_markers(desc)

            detect_array_msg.detections.append(detect_msg)
            markers_msg.markers.extend(markers)

        detect_array_msg.header.stamp = color_image.header.stamp

        self.detect_pub.publish(detect_array_msg)
        self.marker_pub.publish(markers_msg)

        t1 = rospy.Time.now()
        self.update_rate_logger.append((t1 - t0).to_sec())
        self.camera_rate_logger.append((t0 - color_image.header.stamp).to_sec())
        rospy.loginfo_throttle(10, "Detection rate avg: %0.3f Hz" % self.update_rate_logger.rate())
        rospy.loginfo_throttle(10, "Camera delay avg: %0.3fs" % self.camera_rate_logger.avg())

    def camera_info_callback(self, camera_info):
        self.camera_model.fromCameraInfo(camera_info)

    def run_detection(self, color_image_np):
        input_tensor = tensorflow.convert_to_tensor(color_image_np)
        input_tensor = input_tensor[tensorflow.newaxis, ...]

        t0 = rospy.Time.now()
        # if self.detect_fn is None:
        #     detections = copy.deepcopy(self.dummy_detections)
        # else:
        detections = self.detect_fn(input_tensor)
            # with open("/home/ben/detection.pkl", 'wb') as file:
            #     pickle.dump(detections, file)
        t1 = rospy.Time.now()
        self.detect_rate_logger.append((t1 - t0).to_sec())

        rospy.loginfo_throttle(10, "Detection rate avg: %0.3f Hz" % self.detect_rate_logger.rate())

        return detections

    def detection_pipeline(self, color_image):
        try:
            # Convert ROS Image message to numpy array
            color_image_np = self.bridge.imgmsg_to_cv2(color_image, "rgb8")
        except CvBridgeError as e:
            rospy.logerr(e)
            return

        detections = self.run_detection(color_image_np)

        num_detections = int(detections.pop("num_detections"))
        detections = {key: value[0, :num_detections].numpy() for key, value in detections.items()}
        detections["num_detections"] = num_detections

        # detection_classes should be ints.
        detections["detection_classes"] = detections["detection_classes"].astype(np.int64)

        image_with_detections = color_image_np.copy()

        self.publish_visualization(detections, image_with_detections)

        return detections

    def depth_pipeline(self, box, score, class_index, depth_image):
        t0 = time.time()
        class_name = self.category_index[class_index]["name"]
        ymin, xmin, ymax, xmax = box

        # depth image dimensions should be the same as the color image because
        # the depth image has been registered/shifted to be overlaid on top of the color image
        ymin *= depth_image.height
        ymax *= depth_image.height
        xmin *= depth_image.width
        xmax *= depth_image.width

        xcenter = (xmax + xmin) / 2.0
        ycenter = (ymax + ymin) / 2.0
        xsize = xmax - xmin
        ysize = ymax - ymin
        obj_radius = min(xsize, ysize) / 2.0
        detect_radius = max(obj_radius - self.bounding_box_border_px, 1)

        desc = BoxDescription()
        desc.class_index = class_index
        desc.class_name = class_name
        desc.score = score
        desc.box = (ymin, xmin, ymax, xmax)
        desc.xcenter = xcenter
        desc.ycenter = ycenter
        desc.xsize = xsize
        desc.ysize = ysize
        desc.obj_radius = obj_radius
        desc.detect_radius = detect_radius

        if self.z_size_estimations is not None:
            zsize = self.z_size_estimations.get(class_name, self.default_z_size)
        else:
            zsize = self.default_z_size
        z_depth = self.z_depth_from_image(depth_image, desc)
        if z_depth is None:
            return None
        z_depth += zsize / 2.0  # move to the rough center of the object
        ray = self.camera_model.projectPixelTo3dRay((desc.xcenter, desc.ycenter))

        x_pos = ray[0] * z_depth
        y_pos = ray[1] * z_depth

        pose = PoseStamped()
        pose.header.frame_id = self.camera_model.tfFrame()
        pose.header.stamp = depth_image.header.stamp
        pose.pose.position.x = x_pos
        pose.pose.position.y = y_pos
        pose.pose.position.z = z_depth
        pose.pose.orientation.x = 0.0
        pose.pose.orientation.y = 0.0
        pose.pose.orientation.z = 0.0
        pose.pose.orientation.w = 1.0

        edge_point_x = desc.xcenter - xsize / 2.0
        edge_point_y = desc.ycenter - ysize / 2.0
        ray = self.camera_model.projectPixelTo3dRay((edge_point_x, edge_point_y))

        desc.pose_stamped = pose
        desc.xedge = abs(ray[0] * z_depth - x_pos) * 2.0
        desc.yedge = abs(ray[1] * z_depth - y_pos) * 2.0
        desc.zedge = zsize

        t1 = time.time()
        self.depth_rate_logger.append(t1 - t0)
        rospy.loginfo_throttle(10, "Depth pipeline rate avg: %0.3f Hz" % self.depth_rate_logger.rate())

        return desc

    def z_depth_from_image(self, depth_image, desc: BoxDescription):
        try:
            # Convert ROS Image message to numpy array
            depth_image_np = self.bridge.imgmsg_to_cv2(depth_image, desired_encoding="passthrough")
        except CvBridgeError as e:
            rospy.logerr(e)
            return None
        depth_image_np = depth_image_np.astype(np.float32)
        circle_mask = np.zeros(depth_image_np.shape, np.uint8)
        cv2.circle(circle_mask, (int(desc.xcenter), int(desc.ycenter)), int(desc.detect_radius), (255, 255, 255), cv2.FILLED)
        nonzero_mask = (depth_image_np > 0.0).astype(np.uint8)
        target_mask = cv2.bitwise_and(circle_mask, nonzero_mask)
        z_depth = cv2.mean(depth_image_np, target_mask)[0]
        z_depth /= 1000.0
        return z_depth

    def publish_visualization(self, detections, image_with_detections):
        if self.visualization_image_pub.get_num_connections() == 0:
            return
        viz_utils.visualize_boxes_and_labels_on_image_array(
            image_with_detections,
            detections["detection_boxes"],
            detections["detection_classes"],
            detections["detection_scores"],
            self.category_index,
            use_normalized_coordinates=True,
            max_boxes_to_draw=self.max_boxes_to_draw,
            min_score_thresh=self.min_score_threshold,
            agnostic_mode=False
        )

        try:
            visualize_image_msg = self.bridge.cv2_to_imgmsg(image_with_detections, "rgb8")
        except CvBridgeError as e:
            rospy.logerr(e)
            return
        self.visualization_image_pub.publish(visualize_image_msg)

    def get_detect_msg(self, desc: BoxDescription):
        detect_msg = Detection2D()
        detect_msg.bbox.size_x = desc.xsize
        detect_msg.bbox.size_y = desc.ysize
        detect_msg.bbox.center.x = desc.xcenter
        detect_msg.bbox.center.y = desc.ycenter
        hyp = ObjectHypothesisWithPose()

        hyp.id = desc.class_index
        hyp.score = desc.score
        detect_msg.results.append(hyp)
        return detect_msg

    def object_to_robot_frame(self, pose_stamped, image_header):
        try:
            transform = self.tf_buffer.lookup_transform(
                self.robot_frame,
                image_header.frame_id,
                image_header.stamp,
                rospy.Duration(1.0)
            )
        except (tf2_ros.LookupException, tf2_ros.ConnectivityException, tf2_ros.ExtrapolationException) as e:
            rospy.logwarn("Failed to look up %s to %s. %s" % (self.robot_frame, image_header.frame_id, e))
            return None
        pose_robot_frame = tf2_geometry_msgs.do_transform_pose(pose_stamped, transform)
        return pose_robot_frame

    def dyn_callback(self, config, level):
        rospy.loginfo("Reconfigure Request: %s. Level: %s" % (str(config), level))
        if level >= 0:
            rospy.loginfo("Reconfiguring")
            self.min_score_threshold = config["min_score_threshold"]
            self.max_boxes_to_draw = config["max_boxes_to_draw"]
            self.bounding_box_border_px = config["bounding_box_border_px"]
            self.min_valid_z = config["min_valid_z"]
            self.max_valid_z = config["max_valid_z"]
        return config

    def make_markers(self, desc: BoxDescription):
        sphere_marker = self.make_marker_base(desc)
        text_marker = self.make_marker_base(desc)

        sphere_marker.type = Marker.SPHERE
        sphere_marker.ns = "sphere_" + sphere_marker.ns

        text_marker.type = Marker.TEXT_VIEW_FACING
        text_marker.ns = "text_" + text_marker.ns
        text_marker.text = desc.class_name + "_" + str(desc.class_index)
        text_marker.scale.x = 0.0
        text_marker.scale.y = 0.0

        return [sphere_marker, text_marker]

    def make_marker_base(self, desc: BoxDescription):
        marker = Marker()
        marker.action = Marker.ADD
        marker.pose = desc.pose_stamped.pose
        marker.header = desc.pose_stamped.header
        marker.lifetime = self.marker_persistance
        marker.ns = desc.class_name
        marker.id = desc.class_index

        scale_vector = Vector3()
        scale_vector.x = desc.xedge
        scale_vector.y = desc.yedge
        scale_vector.z = desc.zedge
        marker.scale = scale_vector
        if self.marker_colors is not None and desc.class_name in self.marker_colors:
            marker.color = self.marker_colors[desc.class_name]
        else:
            marker.color = ColorRGBA(1.0, 1.0, 1.0, 1.0)

        return marker

    def run(self):
        rospy.spin()

if __name__ == "__main__":
    try:
        node = DodobotTensorflow()
        node.run()
    except rospy.ROSInterruptException:
        pass
    except BaseException as e:
        rospy.logerr("%s: %s\n%s" % (e.__class__.__name__, str(e), traceback.format_exc()))
    finally:
        rospy.loginfo("Exiting db_tensorflow node")


import functools
import unittest
import numpy as np

def read_lines(filename):
    with open(filename) as f:
        lines = f.read().splitlines()
    return np.array([list(l) for l in lines])

def part1(file_name):
    lines = read_lines(file_name).transpose()
    # for each column, shift and accumulate 1 or 0 depending on count of 1's or 0's
    gamma = functools.reduce(lambda accum, line: (accum << 1) + (line[line == '1'].size > len(line)/2), lines, 0)
    # epsilon is the "negated" value of gamma, but needs to be masked to the length of the number
    mask = ~(~0 << len(lines))
    epsilon = ~gamma & mask
    print(f'Gamma: {gamma}, Epsilon: {epsilon}')
    return gamma * epsilon

def part2(file_name):
    return calc_oxygen(file_name) * calc_co2(file_name)

def calc_oxygen(file_name):
    all_lines = read_lines(file_name)
    col = 0
    while len(all_lines) > 1:
        line = all_lines[:, col]
        ones = line[line == '1'].size
        zeroes = line[line == '0'].size
        if ones >= zeroes:
            all_lines = all_lines[all_lines[:, col] == '1']
        else:
            all_lines = all_lines[all_lines[:, col] == '0']
        col += 1
    res = ''.join(all_lines[0])
    print(f'num: {res}')
    return int(res, 2)

def calc_co2(file_name):
    all_lines = read_lines(file_name)
    col = 0
    while len(all_lines) > 1:
        line = all_lines[:, col]
        ones = line[line == '1'].size
        zeroes = line[line == '0'].size
        if zeroes <= ones:
            all_lines = all_lines[all_lines[:, col] == '0']
        else:
            all_lines = all_lines[all_lines[:, col] == '1']
        col += 1
    res = ''.join(all_lines[0])
    print(f'num: {res}')
    return int(res, 2)

class TestPart1(unittest.TestCase):
    def test10(self):
        self.assertEqual(198, part1('input0.txt'))

    def test1(self):
        self.assertEqual(3429254, part1('input.txt'))

class TestPart2(unittest.TestCase):
    def test20(self):
        self.assertEqual(230, part2('input0.txt'))

    def test21(self):
        self.assertEqual(5410338, part2('input.txt'))


import tensorflow as tf

import pyhopper
import numpy as np

def train_1ep_mnist(hparams):
    (x_train, y_train), _ = tf.keras.datasets.mnist.load_data()
    x_train = x_train.reshape((-1, 28, 28, 1)) / 255.0
    model = tf.keras.models.Sequential(
        [
            tf.keras.Input((28, 28, 1), dtype=tf.uint8),
            tf.keras.layers.Conv2D(16, 5, activation="relu"),
            tf.keras.layers.MaxPool2D(),
            tf.keras.layers.Conv2D(32, 3, activation="relu"),
            tf.keras.layers.MaxPool2D(),
            tf.keras.layers.Conv2D(64, 3, activation="relu"),
            tf.keras.layers.MaxPool2D(),
            tf.keras.layers.Flatten(),
            tf.keras.layers.Dropout(0.1),
            tf.keras.layers.Dense(128, "relu"),
            tf.keras.layers.Dropout(0.1),
            tf.keras.layers.Dense(10, "softmax"),
        ]
    )
    model.compile(
        loss=tf.keras.losses.sparse_categorical_crossentropy,
        metrics=[tf.keras.metrics.sparse_categorical_accuracy],
        optimizer=tf.keras.optimizers.Adam(0.0005),
    )
    hist = model.fit(
        x_train,
        y_train,
        batch_size=128,
        epochs=8,
        validation_split=0.05,
    )
    return hist.history["val_sparse_categorical_accuracy"][-1]

search = pyhopper.Search(
    {
        "lr": pyhopper.float(1e-4, 1e-2),
    }
)
policy = search.run(
    pyhopper.wrap_n_times(train_1ep_mnist, n=5, yield_after=4),
    direction="max",
    seeding_steps=10,
    timeout="5s",
    n_jobs=4,
    kwargs={},
)

import matplotlib.pyplot as plt

fig, ax = plt.subplots()
ax.scatter(search.history.seconds, search.history.fs)
ax.plot(
    search.history.seconds,
    search.history.best_fs,
    label="Best",
    color="red",
)
fig.savefig("test.png")


import public_fun
import sqlite3
import os
import traceback
from bs4 import BeautifulSoup
import re
import time

def init_sqlite_db(dbPath):
    '''
    返回一个sqlite数据库的实例
    '''
    conn = sqlite3.connect(dbPath)
    return conn

def file_list(filepath):
    '''
    遍历文件夹中的文件
    '''
    for root, dirs, files in os.walk(filepath):
        for file in files:
            yield os.path.join(root, file)

def get_title(pTagList):
    '''
    获取文献标题(中英文)
    '''
    aTag = pTagList[0].find('a')
    titleOne = aTag.string
    aTag = pTagList[1].find('a')
    titleTwo = aTag.string
    if titleTwo is None:
        return titleOne, ''
    regepx = re.compile(r'[\u4e00-\u9fa5]')
    match = regepx.search(titleOne)
    if match:
        title_c = titleOne
        title_e = titleTwo
    else:
        title_c = titleTwo
        title_e = titleOne

    return title_c, title_e

def get_creator(pTagList):
    '''
    获取文章作者
    '''
    creators_c = []
    creators_e = []
    aTagList = pTagList[2].find_all('a')
    for aTag in aTagList:
        creator = ''.join(aTag.stripped_strings)
        regepx = re.compile(r'\(')
        match = re.findall(regepx, creator)
        if len(match) > 0:
            tmp = creator.split('(')
            creators_c.append(tmp[0])
            creators_e.append(tmp[1][:-1])
        else:
            creators_c.append(creator)
    return ';'.join(creators_c), ';'.join(creators_e)

def get_institution_date_page(pTagList):
    """
    获取单位，发表日期，还有文章所在的页码
    """
    publisher_date_issues_page = (
        ''.join((pTagList[3].string).split())).split('；')
    creator_institution = publisher_date_issues_page[0][:-4]
    date_issue_page = publisher_date_issues_page[1].split('，')
    page = date_issue_page[1]
    date_issue = date_issue_page[0].split('(')
    date_created = ''.join((date_issue[1][:-1]).split('/'))
    return creator_institution, date_created, page

def get_lang_doi(pTagList):
    """
    获取语言与doi
    """
    lang_doi = list()
    for string in pTagList[4].stripped_strings:
        lang_doi.append(string)
    regepx = re.compile(r'(中文)')
    match = re.findall(regepx, lang_doi[0])
    if match:
        lang = "ZH"
    else:
        lang = "EN"
    try:
        doi = lang_doi[2]
    except IndexError:
        doi = ''
    return lang, doi

def get_keyword(pTagList):
    '''
    获取文献关键字
    '''
    subject = list()
    subject_en = list()
    aTagL = pTagList[5].find_all('a')
    for aTag in aTagL:
        keyword = ''.join(aTag.stripped_strings)
        if re.search(r'[\u4e00-\u9fa5]', keyword):
            subject.append(keyword)
        else:
            subject_en.append(keyword)
    return ';'.join(subject), ';'.join(subject_en)

def get_summy(soup):
    '''
    获取文献的摘要（中英）
    '''
    summy_e = soup.select('.TextArea2')
    # print(xxxx[0].string)
    summy_c = soup.select('.TextArea1')
    # print(xxx[0].string)
    return summy_c[0].string, summy_e[0].string

def get_university(rawid):
    uid = rawid.split('-')[0]
    conn = public_fun.init_db()
    cursor = conn.cursor()
    cursor.execute(
        'select university from universities where FirstID=%s', [uid])
    university = cursor.fetchone()
    cursor.close()
    conn.close()
    return university

def ProcOneFile(file):
    rawid = (file.split('\\')[1])[:-5]
    html = public_fun.read_html(file)
    soup = BeautifulSoup(html, 'lxml')
    pTagList = soup.select('.booksInformation > .detail > p')
    language, identifier_doi = get_lang_doi(pTagList)
    subject, subject_en = get_keyword(pTagList)
    lngid = 'AIRITILIBRARY_BS_' + rawid
    creator_discipline, date_created, page = get_institution_date_page(
        pTagList)
    title, title_alternative = get_title(pTagList)
    description, description_en = get_summy(soup)
    country = "CN"
    provider = 'airitithesis'
    provider_url = provider + \
        '@http://www.airitilibrary.cn/Publication/alDetailedMesh?DocID=' + rawid
    provider_id = provider + '@' + rawid
    university = get_university(rawid)
    creator, creator_en = get_creator(pTagList)
    tup = (
        lngid,
        rawid,
        identifier_doi,
        title,
        title_alternative,
        creator,
        creator_en,
        university[0],  # 论文发表者所在大学
        creator_discipline[len(university[0]):],  # 论文发表者的专业
        # creator_degree,  # 论文发表者的学位
        # contributor,  # 指导老师
        university[0],  # 所在单位
        page,
        description,
        description_en,
        subject,
        subject_en,
        date_created[:4],
        date_created,
        language,
        country,
        provider,
        provider_url,
        provider_id,
        4,
        2,
        time.strftime('%Y%m%d', time.localtime(time.time())) + '00'
    )
    return tup

if __name__ == "__main__":
    filePath = './down_by_doc_id/'
    conn = init_sqlite_db('./sqlite.db3')
    cursor = conn.cursor()
    sql = '''insert OR IGNORE into modify_title_info_zt
    (lngid,rawid,identifier_doi,title,title_alternative,creator,creator_en,creator_institution,creator_discipline,source_institution,page,
    description,description_en,subject,subject_en,date,date_created,
    language,country,provider,provider_url,provider_id,type,medium,batch) values(?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)'''

    sqlList = list()
    for file in file_list(filePath):
        try:
            tup = ProcOneFile(file)
            sqlList.append(tup)
        except:
            exMsg = 'err file:' + file + '\n'
            exMsg += traceback.format_exc()
            print(exMsg)
            with open('./bug.log', 'a') as fp:
                fp.write(file)
                fp.write('\n')
                fp.write(exMsg)
                fp.write('\n')
            # sys.exit(-1)

        if len(sqlList) > 10:
            cursor.executemany(sql, sqlList)
            conn.commit()
            sqlList.clear()
        print('Parser ' + file + ' Successfully!!!')
    if len(sqlList) > 0:
        cursor.executemany(sql, sqlList)
        conn.commit()
        sqlList.clear()

    cursor.close()
    conn.close()
    print('Done!!!')