import collections
import functools
import json
import typing
from pathlib import Path
from random import Random
from typing import Optional, Dict, Set, List, Tuple, Iterator, Union

import matplotlib.pyplot as plt
import networkx
from PySide2 import QtWidgets
from PySide2.QtCore import Qt
from PySide2.QtWidgets import QMainWindow, QTreeWidgetItem, QCheckBox, QLabel, QGridLayout, QWidget, QMessageBox
from matplotlib.axes import Axes
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.backends.backend_qt5agg import NavigationToolbar2QT as NavigationToolbar
from matplotlib.figure import Figure

from randovania.game_description.area_location import AreaLocation
from randovania.game_description.game_description import GameDescription
from randovania.game_description.item.item_category import ItemCategory
from randovania.game_description.node import Node, ResourceNode, TranslatorGateNode, TeleporterNode, DockNode
from randovania.game_description.resources.item_resource_info import ItemResourceInfo
from randovania.game_description.resources.pickup_entry import PickupEntry
from randovania.game_description.resources.resource_info import add_resource_gain_to_current_resources
from randovania.game_description.resources.translator_gate import TranslatorGate
from randovania.game_description.world import World
from randovania.games.game import RandovaniaGame
from randovania.games.prime import patcher_file
from randovania.generator import generator
from randovania.gui.generated.tracker_window_ui import Ui_TrackerWindow
from randovania.gui.lib.common_qt_lib import set_default_window_icon
from randovania.gui.lib.custom_spin_box import CustomSpinBox
from randovania.layout import translator_configuration
from randovania.layout.echoes_configuration import EchoesConfiguration
from randovania.layout.teleporters import TeleporterShuffleMode
from randovania.layout.translator_configuration import LayoutTranslatorRequirement
from randovania.resolver.bootstrap import logic_bootstrap
from randovania.resolver.logic import Logic
from randovania.resolver.resolver_reach import ResolverReach
from randovania.resolver.state import State, add_pickup_to_state

class InvalidLayoutForTracker(Exception):
    pass

def _load_previous_state(persistence_path: Path,
                         layout_configuration: EchoesConfiguration,
                         ) -> Optional[dict]:
    previous_layout_path = persistence_path.joinpath("layout_configuration.json")
    try:
        with previous_layout_path.open() as previous_layout_file:
            previous_layout = EchoesConfiguration.from_json(json.load(previous_layout_file))
    except (FileNotFoundError, TypeError, KeyError, ValueError, json.JSONDecodeError):
        return None

    if previous_layout != layout_configuration:
        return None

    previous_state_path = persistence_path.joinpath("state.json")
    try:
        with previous_state_path.open() as previous_state_file:
            return json.load(previous_state_file)
    except (FileNotFoundError, json.JSONDecodeError):
        return None

class MatplotlibWidget(QtWidgets.QWidget):
    ax: Axes

    def __init__(self, parent=None):
        super().__init__(parent)

        fig = Figure(figsize=(7, 5), dpi=65, facecolor=(1, 1, 1), edgecolor=(0, 0, 0))
        self.canvas = FigureCanvas(fig)
        self.toolbar = NavigationToolbar(self.canvas, self)
        lay = QtWidgets.QVBoxLayout(self)
        lay.addWidget(self.toolbar)
        lay.addWidget(self.canvas)

        self.ax = fig.add_subplot(111)
        self.line, *_ = self.ax.plot([])

class TrackerWindow(QMainWindow, Ui_TrackerWindow):
    # Tracker state
    _collected_pickups: Dict[PickupEntry, int]
    _actions: List[Node]

    # Tracker configuration
    logic: Logic
    game_description: GameDescription
    layout_configuration: EchoesConfiguration
    persistence_path: Path
    _initial_state: State
    _elevator_id_to_combo: Dict[int, QtWidgets.QComboBox]
    _translator_gate_to_combo: Dict[TranslatorGate, QtWidgets.QComboBox]
    _starting_nodes: Set[ResourceNode]
    _undefined_item = ItemResourceInfo(-1, "Undefined", "Undefined", 0, None)

    # UI tools
    _asset_id_to_item: Dict[int, QTreeWidgetItem]
    _node_to_item: Dict[Node, QTreeWidgetItem]
    _widget_for_pickup: Dict[PickupEntry, Union[QCheckBox, CustomSpinBox]]
    _during_setup = False

    def __init__(self, persistence_path: Path, layout_configuration: EchoesConfiguration):
        super().__init__()
        self.setupUi(self)
        set_default_window_icon(self)

        self._collected_pickups = {}
        self._widget_for_pickup = {}
        self._actions = []
        self._asset_id_to_item = {}
        self._node_to_item = {}
        self.layout_configuration = layout_configuration
        self.persistence_path = persistence_path

        player_pool = generator.create_player_pool(Random(0), self.layout_configuration, 0, 1)
        pool_patches = player_pool.patches
        self.game_description, self._initial_state = logic_bootstrap(layout_configuration,
                                                                     player_pool.game,
                                                                     pool_patches)
        self.logic = Logic(self.game_description, layout_configuration)

        self._initial_state.resources["add_self_as_requirement_to_resources"] = 1

        self.menu_reset_action.triggered.connect(self._confirm_reset)
        self.resource_filter_check.stateChanged.connect(self.update_locations_tree_for_reachable_nodes)
        self.hide_collected_resources_check.stateChanged.connect(self.update_locations_tree_for_reachable_nodes)
        self.undo_last_action_button.clicked.connect(self._undo_last_action)

        self.configuration_label.setText("Trick Level: {}; Starts with:\n{}".format(
            layout_configuration.trick_level.pretty_description,
            ", ".join(
                resource.short_name
                for resource in pool_patches.starting_items.keys()
            )
        ))

        self.setup_pickups_box(player_pool.pickups)
        self.setup_possible_locations_tree()
        self.setup_elevators()
        self.setup_translator_gates()

        self.matplot_widget = MatplotlibWidget(self.tab_graph_map)
        self.tab_graph_map_layout.addWidget(self.matplot_widget)
        self._world_to_node_positions = {}
        self.map_tab_widget.currentChanged.connect(self._on_tab_changed)

        for world in self.game_description.world_list.worlds:
            self.graph_map_world_combo.addItem(world.name, world)
        self.graph_map_world_combo.currentIndexChanged.connect(self.on_graph_map_world_combo)

        persistence_path.mkdir(parents=True, exist_ok=True)
        previous_state = _load_previous_state(persistence_path, layout_configuration)

        if not self.apply_previous_state(previous_state):
            self.setup_starting_location(None)

            with persistence_path.joinpath("layout_configuration.json").open("w") as layout_file:
                json.dump(layout_configuration.as_json, layout_file)
            self._add_new_action(self._initial_state.node)

    def apply_previous_state(self, previous_state: Optional[dict]) -> bool:
        if previous_state is None:
            return False

        starting_location = None
        needs_starting_location = len(self.layout_configuration.starting_location.locations) > 1
        resource_db = self.game_description.resource_database
        translator_gates = {}

        try:
            pickup_name_to_pickup = {pickup.name: pickup for pickup in self._collected_pickups.keys()}
            quantity_to_change = {
                pickup_name_to_pickup[pickup_name]: quantity
                for pickup_name, quantity in previous_state["collected_pickups"].items()
            }
            previous_actions = [
                self.game_description.world_list.all_nodes[index]
                for index in previous_state["actions"]
            ]
            if needs_starting_location:
                starting_location = AreaLocation.from_json(previous_state["starting_location"])

            elevators = {
                int(elevator_id): AreaLocation.from_json(location) if location is not None else None
                for elevator_id, location in previous_state["elevators"].items()
            }
            if self.layout_configuration.game == RandovaniaGame.PRIME2:
                translator_gates = {
                    TranslatorGate(int(gate)): (resource_db.get_item(item)
                                                if item is not None
                                                else self._undefined_item)
                    for gate, item in previous_state["translator_gates"].items()
                }
        except KeyError:
            return False

        self.setup_starting_location(starting_location)

        for elevator_id, area_location in elevators.items():
            combo = self._elevator_id_to_combo[elevator_id]
            if area_location is None:
                combo.setCurrentIndex(0)
                continue
            for i in range(combo.count()):
                if area_location == combo.itemData(i):
                    combo.setCurrentIndex(i)
                    break

        for gate, item in translator_gates.items():
            combo = self._translator_gate_to_combo[gate]
            for i in range(combo.count()):
                if item == combo.itemData(i):
                    combo.setCurrentIndex(i)
                    break

        self.bulk_change_quantity(quantity_to_change)
        self._add_new_actions(previous_actions)
        return True

    def reset(self):
        self.bulk_change_quantity({
            pickup: 0
            for pickup in self._collected_pickups.keys()
        })

        while len(self._actions) > 1:
            self._actions.pop()
            self.actions_list.takeItem(len(self._actions))

        for elevator in self._elevator_id_to_combo.values():
            elevator.setCurrentIndex(0)
        for elevator in self._translator_gate_to_combo.values():
            elevator.setCurrentIndex(0)

        self._refresh_for_new_action()

    def _confirm_reset(self):
        reply = QMessageBox.question(self, "Reset Tracker?", "Do you want to reset the tracker progression?",
                                     QMessageBox.Yes | QMessageBox.No, QMessageBox.No)
        if reply == QMessageBox.Yes:
            self.reset()

    @property
    def _show_only_resource_nodes(self) -> bool:
        return self.resource_filter_check.isChecked()

    @property
    def _hide_collected_resources(self) -> bool:
        return self.hide_collected_resources_check.isChecked()

    @property
    def _collected_nodes(self) -> Set[ResourceNode]:
        return self._starting_nodes | set(action for action in self._actions if action.is_resource_node)

    def _pretty_node_name(self, node: Node) -> str:
        world_list = self.game_description.world_list
        return "{} / {}".format(world_list.area_name(world_list.nodes_to_area(node)), node.name)

    def _refresh_for_new_action(self):
        self.undo_last_action_button.setEnabled(len(self._actions) > 1)
        self.current_location_label.setText("Current location: {}".format(self._pretty_node_name(self._actions[-1])))
        self.update_locations_tree_for_reachable_nodes()

    def _add_new_action(self, node: Node):
        self._add_new_actions([node])

    def _add_new_actions(self, nodes: Iterator[Node]):
        for node in nodes:
            self.actions_list.addItem(self._pretty_node_name(node))
            self._actions.append(node)
        self._refresh_for_new_action()

    def _undo_last_action(self):
        self._actions.pop()
        self.actions_list.takeItem(len(self._actions))
        self._refresh_for_new_action()

    def _on_tree_node_double_clicked(self, item: QTreeWidgetItem, _):
        node: Optional[Node] = getattr(item, "node", None)

        if not item.isDisabled() and node is not None and node != self._actions[-1]:
            self._add_new_action(node)

    def _positions_for_world(self, world: World):
        g = networkx.DiGraph()
        world_list = self.game_description.world_list
        state = self.state_for_current_configuration()

        for area in world.areas:
            g.add_node(area)

        for area in world.areas:
            nearby_areas = set()
            for node in area.nodes:
                if isinstance(node, DockNode):
                    try:
                        target_node = world_list.resolve_dock_node(node, state.patches)
                        nearby_areas.add(world_list.nodes_to_area(target_node))
                    except IndexError as e:
                        print(f"For {node.name} in {area.name}, received {e}")
                        continue
            for other_area in nearby_areas:
                g.add_edge(area, other_area)

        return networkx.drawing.spring_layout(g)

    def update_matplot_widget(self, nodes_in_reach: Set[Node]):
        g = networkx.DiGraph()
        world_list = self.game_description.world_list
        state = self.state_for_current_configuration()

        world = self.graph_map_world_combo.currentData()
        for area in world.areas:
            g.add_node(area)

        for area in world.areas:
            nearby_areas = set()
            for node in area.nodes:
                if node not in nodes_in_reach:
                    continue

                if isinstance(node, DockNode):
                    # TODO: respect is_blast_shield: if already opened once, no requirement needed.
                    # Includes opening form behind with different criteria
                    try:
                        target_node = world_list.resolve_dock_node(node, state.patches)
                        dock_weakness = state.patches.dock_weakness.get((area.area_asset_id, node.dock_index),
                                                                        node.default_dock_weakness)
                        if dock_weakness.requirement.satisfied(state.resources, state.energy):
                            nearby_areas.add(world_list.nodes_to_area(target_node))
                    except IndexError as e:
                        print(f"For {node.name} in {area.name}, received {e}")
                        continue

            for other_area in nearby_areas:
                g.add_edge(area, other_area)

        self.matplot_widget.ax.clear()

        cf = self.matplot_widget.ax.get_figure()
        cf.set_facecolor("w")

        if world.world_asset_id not in self._world_to_node_positions:
            self._world_to_node_positions[world.world_asset_id] = self._positions_for_world(world)
        pos = self._world_to_node_positions[world.world_asset_id]

        networkx.draw_networkx_nodes(g, pos, ax=self.matplot_widget.ax)
        networkx.draw_networkx_edges(g, pos, arrows=True, ax=self.matplot_widget.ax)
        networkx.draw_networkx_labels(g, pos, ax=self.matplot_widget.ax,
                                      labels={area: area.name for area in world.areas},
                                      verticalalignment='top')

        self.matplot_widget.ax.set_axis_off()

        plt.draw_if_interactive()
        self.matplot_widget.canvas.draw()

    def on_graph_map_world_combo(self):
        nodes_in_reach = self.current_nodes_in_reach(self.state_for_current_configuration())
        self.update_matplot_widget(nodes_in_reach)

    def current_nodes_in_reach(self, state):
        if state is None:
            nodes_in_reach = set()
        else:
            reach = ResolverReach.calculate_reach(self.logic, state)
            nodes_in_reach = set(reach.nodes)
            nodes_in_reach.add(state.node)
        return nodes_in_reach

    def _on_tab_changed(self):
        if self.map_tab_widget.currentWidget() == self.tab_graph_map:
            self.on_graph_map_world_combo()

    def update_locations_tree_for_reachable_nodes(self):
        state = self.state_for_current_configuration()
        nodes_in_reach = self.current_nodes_in_reach(state)
        if self.map_tab_widget.currentWidget() == self.tab_graph_map:
            self.update_matplot_widget(nodes_in_reach)

        all_nodes = self.game_description.world_list.all_nodes
        for world in self.game_description.world_list.worlds:
            for area in world.areas:
                area_is_visible = False
                for node in area.nodes:
                    is_collected = node in self._collected_nodes
                    is_visible = node in nodes_in_reach and not (self._hide_collected_resources
                                                                 and is_collected)

                    if self._show_only_resource_nodes:
                        is_visible = is_visible and node.is_resource_node

                    node_item = self._node_to_item[node]
                    node_item.setHidden(not is_visible)
                    if node.is_resource_node:
                        resource_node = typing.cast(ResourceNode, node)
                        node_item.setDisabled(not resource_node.can_collect(state.patches, state.resources, all_nodes))
                        node_item.setCheckState(0, Qt.Checked if is_collected else Qt.Unchecked)

                    area_is_visible = area_is_visible or is_visible
                self._asset_id_to_item[area.area_asset_id].setHidden(not area_is_visible)

        # Persist the current state
        self.persist_current_state()

    def persist_current_state(self):
        world_list = self.game_description.world_list
        with self.persistence_path.joinpath("state.json").open("w") as state_file:
            json.dump(
                {
                    "actions": [
                        node.index
                        for node in self._actions
                    ],
                    "collected_pickups": {
                        pickup.name: quantity
                        for pickup, quantity in self._collected_pickups.items()
                    },
                    "elevators": {
                        str(elevator_id): combo.currentData().as_json if combo.currentIndex() > 0 else None
                        for elevator_id, combo in self._elevator_id_to_combo.items()
                    },
                    "translator_gates": {
                        str(gate.index): combo.currentData().index if combo.currentIndex() > 0 else None
                        for gate, combo in self._translator_gate_to_combo.items()
                    },
                    "starting_location": world_list.node_to_area_location(self._initial_state.node).as_json,
                },
                state_file
            )

    def setup_possible_locations_tree(self):
        """
        Creates the possible_locations_tree with all worlds, areas and nodes.
        """
        self.possible_locations_tree.itemDoubleClicked.connect(self._on_tree_node_double_clicked)

        # TODO: Dark World names
        for world in self.game_description.world_list.worlds:
            world_item = QTreeWidgetItem(self.possible_locations_tree)
            world_item.setText(0, world.name)
            world_item.setExpanded(True)
            self._asset_id_to_item[world.world_asset_id] = world_item

            for area in world.areas:
                area_item = QTreeWidgetItem(world_item)
                area_item.area = area
                area_item.setText(0, area.name)
                area_item.setHidden(True)
                self._asset_id_to_item[area.area_asset_id] = area_item

                for node in area.nodes:
                    node_item = QTreeWidgetItem(area_item)
                    if isinstance(node, TranslatorGateNode):
                        node_item.setText(0, "{} ({})".format(node.name, node.gate))
                    else:
                        node_item.setText(0, node.name)
                    node_item.node = node
                    if node.is_resource_node:
                        node_item.setFlags(node_item.flags() & ~Qt.ItemIsUserCheckable)
                    self._node_to_item[node] = node_item

    def setup_elevators(self):
        world_list = self.game_description.world_list
        nodes_by_world: Dict[str, List[TeleporterNode]] = collections.defaultdict(list)
        self._elevator_id_to_combo = {}

        areas_to_not_change = {
            2278776548,  # Sky Temple Gateway
            2068511343,  # Sky Temple Energy Controller
            3136899603,  # Aerie Transport Station
            1564082177,  # Aerie
        }
        targets = {}
        for world, area, node in world_list.all_worlds_areas_nodes:
            if isinstance(node, TeleporterNode) and node.editable and area.area_asset_id not in areas_to_not_change:
                name = world.correct_name(area.in_dark_aether)
                nodes_by_world[name].append(node)

                location = AreaLocation(world.world_asset_id, area.area_asset_id)
                targets[patcher_file.elevator_area_name(world_list, location, True)] = location

        if self.layout_configuration.elevators.mode == TeleporterShuffleMode.ONE_WAY_ANYTHING:
            targets = {}
            for world in world_list.worlds:
                for area in world.areas:
                    name = world.correct_name(area.in_dark_aether)
                    targets[f"{name} - {area.name}"] = AreaLocation(world.world_asset_id, area.area_asset_id)

        combo_targets = sorted(targets.items(), key=lambda it: it[0])

        for world_name in sorted(nodes_by_world.keys()):
            nodes = nodes_by_world[world_name]
            nodes_locations = [AreaLocation(world_list.nodes_to_world(node).world_asset_id,
                                            world_list.nodes_to_area(node).area_asset_id)
                               for node in nodes]
            nodes_names = [patcher_file.elevator_area_name(world_list, location, True)
                           for location in nodes_locations]

            nodes = sorted(nodes_by_world[world_name], key=lambda it: world_list.nodes_to_area(it).name)

            group = QtWidgets.QGroupBox(self.elevators_scroll_contents)
            group.setTitle(world_name)
            self.elevators_scroll_layout.addWidget(group)
            layout = QtWidgets.QGridLayout(group)

            for i, (node, location, name) in enumerate(sorted(zip(nodes, nodes_locations, nodes_names),
                                                              key=lambda it: it[2])):
                node_name = QtWidgets.QLabel(group)
                node_name.setText(name)
                layout.addWidget(node_name, i, 0)

                combo = QtWidgets.QComboBox(group)
                if self.layout_configuration.elevators.is_vanilla:
                    combo.addItem("Vanilla", node.default_connection)
                    combo.setEnabled(False)
                else:
                    combo.addItem("Undefined", location)
                    for target_name, connection in combo_targets:
                        combo.addItem(target_name, connection)

                combo.currentIndexChanged.connect(self.update_locations_tree_for_reachable_nodes)
                self._elevator_id_to_combo[node.teleporter_instance_id] = combo
                layout.addWidget(combo, i, 1)

    def setup_translator_gates(self):
        world_list = self.game_description.world_list
        resource_db = self.game_description.resource_database
        self._translator_gate_to_combo = {}

        if self.layout_configuration.game != RandovaniaGame.PRIME2:
            return

        gates = {
            f"{area.name} ({node.gate.index})": node.gate
            for world, area, node in world_list.all_worlds_areas_nodes
            if isinstance(node, TranslatorGateNode)
        }
        translator_requirement = self.layout_configuration.translator_configuration.translator_requirement

        for i, (gate_name, gate) in enumerate(sorted(gates.items(), key=lambda it: it[0])):
            node_name = QtWidgets.QLabel(self.translator_gate_scroll_contents)
            node_name.setText(gate_name)
            self.translator_gate_scroll_layout.addWidget(node_name, i, 0)

            combo = QtWidgets.QComboBox(self.translator_gate_scroll_contents)
            gate_requirement = translator_requirement[gate]

            if gate_requirement in (LayoutTranslatorRequirement.RANDOM,
                                    LayoutTranslatorRequirement.RANDOM_WITH_REMOVED):
                combo.addItem("Undefined", self._undefined_item)
                for translator, index in translator_configuration.ITEM_INDICES.items():
                    combo.addItem(translator.long_name, resource_db.get_item(index))
            else:
                combo.addItem(gate_requirement.long_name, resource_db.get_item(gate_requirement.item_index))
                combo.setEnabled(False)

            combo.currentIndexChanged.connect(self.update_locations_tree_for_reachable_nodes)
            self._translator_gate_to_combo[gate] = combo
            self.translator_gate_scroll_layout.addWidget(combo, i, 1)

    def setup_starting_location(self, area_location: Optional[AreaLocation]):
        world_list = self.game_description.world_list

        if len(self.layout_configuration.starting_location.locations) > 1:
            if area_location is None:
                area_locations = sorted(self.layout_configuration.starting_location.locations,
                                        key=lambda it: world_list.area_name(world_list.area_by_area_location(it)))

                location_names = [world_list.area_name(world_list.area_by_area_location(it))
                                  for it in area_locations]
                selected_name = QtWidgets.QInputDialog.getItem(self, "Starting Location", "Select starting location",
                                                               location_names, 0, False)
                area_location = area_locations[location_names.index(selected_name[0])]

            self._initial_state.node = world_list.resolve_teleporter_connection(area_location)

        self._starting_nodes = {
            node
            for node in world_list.all_nodes
            if node.is_resource_node and node.resource() in self._initial_state.resources
        }

    def _change_item_quantity(self, pickup: PickupEntry, use_quantity_as_bool: bool, quantity: int):
        if use_quantity_as_bool:
            if bool(quantity):
                quantity = 1
            else:
                quantity = 0

        self._collected_pickups[pickup] = quantity
        if not self._during_setup:
            self.update_locations_tree_for_reachable_nodes()

    def bulk_change_quantity(self, new_quantity: Dict[PickupEntry, int]):
        self._during_setup = True
        for pickup, quantity in new_quantity.items():
            widget = self._widget_for_pickup[pickup]
            if isinstance(widget, QCheckBox):
                widget.setChecked(quantity > 0)
            else:
                widget.setValue(quantity)
        self._during_setup = False

    def _create_widgets_with_quantity(self,
                                      pickup: PickupEntry,
                                      parent_widget: QWidget,
                                      parent_layout: QGridLayout,
                                      row: int,
                                      quantity: int,
                                      ):
        label = QLabel(parent_widget)
        label.setText(pickup.name)
        parent_layout.addWidget(label, row, 0)

        spin_bix = CustomSpinBox(parent_widget)
        spin_bix.setMaximumWidth(50)
        spin_bix.setMaximum(quantity)
        spin_bix.valueChanged.connect(functools.partial(self._change_item_quantity, pickup, False))
        self._widget_for_pickup[pickup] = spin_bix
        parent_layout.addWidget(spin_bix, row, 1)

    def setup_pickups_box(self, item_pool: List[PickupEntry]):

        parent_widgets: Dict[ItemCategory, Tuple[QWidget, QGridLayout]] = {
            ItemCategory.EXPANSION: (self.expansions_box, self.expansions_layout),
            ItemCategory.ENERGY_TANK: (self.expansions_box, self.expansions_layout),
            ItemCategory.TRANSLATOR: (self.translators_box, self.translators_layout),
            ItemCategory.TEMPLE_KEY: (self.keys_box, self.keys_layout),
            ItemCategory.SKY_TEMPLE_KEY: (self.keys_box, self.keys_layout),
        }
        major_pickup_parent_widgets = (self.upgrades_box, self.upgrades_layout)

        row_for_parent = {
            self.expansions_box: 0,
            self.translators_box: 0,
            self.upgrades_box: 0,
            self.keys_box: 0,
        }
        column_for_parent = {
            self.translators_box: 0,
            self.upgrades_box: 0,
            self.keys_box: 0,
        }
        k_column_count = 2

        pickup_by_name = {}
        pickup_with_quantity = {}

        for pickup in item_pool:
            if pickup.name in pickup_by_name:
                pickup_with_quantity[pickup_by_name[pickup.name]] += 1
            else:
                pickup_by_name[pickup.name] = pickup
                pickup_with_quantity[pickup] = 1

        non_expansions_with_quantity = []

        for pickup, quantity in pickup_with_quantity.items():
            self._collected_pickups[pickup] = 0
            parent_widget, parent_layout = parent_widgets.get(pickup.item_category, major_pickup_parent_widgets)

            row = row_for_parent[parent_widget]

            if parent_widget is self.expansions_box:
                self._create_widgets_with_quantity(pickup, parent_widget, parent_layout, row, quantity)
                row_for_parent[parent_widget] += 1
            else:
                if quantity > 1:
                    non_expansions_with_quantity.append((parent_widget, parent_layout, pickup, quantity))
                else:
                    check_box = QCheckBox(parent_widget)
                    check_box.setText(pickup.name)
                    check_box.stateChanged.connect(functools.partial(self._change_item_quantity, pickup, True))
                    self._widget_for_pickup[pickup] = check_box

                    column = column_for_parent[parent_widget]
                    parent_layout.addWidget(check_box, row, column)
                    column += 1

                    if column >= k_column_count:
                        column = 0
                        row += 1

                    row_for_parent[parent_widget] = row
                    column_for_parent[parent_widget] = column

        for parent_widget, parent_layout, pickup, quantity in non_expansions_with_quantity:
            if column_for_parent[parent_widget] != 0:
                column_for_parent[parent_widget] = 0
                row_for_parent[parent_widget] += 1

            self._create_widgets_with_quantity(pickup, parent_widget, parent_layout,
                                               row_for_parent[parent_widget],
                                               quantity)
            row_for_parent[parent_widget] += 1

    def state_for_current_configuration(self) -> Optional[State]:
        all_nodes = self.game_description.world_list.all_nodes

        state = self._initial_state.copy()
        if self._actions:
            state.node = self._actions[-1]

        for teleporter, combo in self._elevator_id_to_combo.items():
            assert combo.currentData() is not None
            state.patches.elevator_connection[teleporter] = combo.currentData()

        for gate, item in self._translator_gate_to_combo.items():
            state.patches.translator_gates[gate] = item.currentData()

        for pickup, quantity in self._collected_pickups.items():
            for _ in range(quantity):
                add_pickup_to_state(state, pickup)

        for node in self._collected_nodes:
            add_resource_gain_to_current_resources(node.resource_gain_on_collect(state.patches, state.resources,
                                                                                all_nodes),
                                                   state.resources)

        return state


"""
TimeFrameResults steals methods from oemof.outputlib adapted to the structure
applied here. Most relevant difference is results data stored in self.data

"""

from oemof.outputlib import DataFramePlot, ResultsDataFrame

import pickle
from matplotlib import pyplot as plt
import logging
import pandas as pd

class TimeFrameResults:
    """
    Container for results of one time frame (i.e. one year)

    Attributes
    ----------
    data : DataFrame
        Structure multi-indexed result data

    """

    def __init__(self, **kwargs):
        """
        Initializes data object based on oemof results class
        """

        results_file = kwargs.get('results_file', None)
        self.subset = kwargs.get('subset', None)
        self.ax = kwargs.get('ax')

        if results_file is None:
            # self.data = DataFramePlot(energy_system=kwargs.get('energy_system'))
            self.data = ResultsDataFrame(energy_system=kwargs.get('energy_system'))
        else:
            self.data = pickle.load(open(results_file, 'rb'))

        self.reformat_data()

    def preview(self):
        """
        Print short preview of data
        """
        return self.data.head()

    def reformat_data(self):
        """
        Extract region information from bus label put into separate index label
        """
        # TODO: get regions list from elsewhere
        regions = ['deu', 'xfra', 'xbnl']
        regions_leading_underscore = ['_' + x for x in regions]

        # put bus_label to column (required to work on)
        self.data.reset_index(level='bus_label', inplace=True)
        self.data.reset_index(level='obj_label', inplace=True)

        # extra region from bus label and write to new column
        self.data['region'] = self.data['bus_label'].str.extract(
            r"(?=(" + '|'.join(regions) + r"))", expand=True)
        self.data['region'].fillna('global', inplace=True)

        # remove region from bus_label and obj_label
        self.data['bus_label'] = self.data['bus_label'].str.replace(
            r"(" + '|'.join(regions_leading_underscore) + r")", '')
        self.data['obj_label'] = self.data['obj_label'].str.replace(
            r"(" + '|'.join(regions_leading_underscore) + r")", '')

        # put bus_label back to index
        self.data = self.data.set_index(['bus_label', 'region', 'obj_label'],
                                        append=True)

        # reorder and resort levels
        level_order = ['bus_label', 'type', 'obj_label', 'region', 'datetime']
        self.data = self.data.reorder_levels(level_order)

    def slice_by(self, **kwargs):
        r""" Method for slicing the ResultsDataFrame. A subset is returned.

        Parameters
        ----------
        bus_label : string
        type : string (to_bus/from_bus/other)
        obj_label: string
        date_from : string
            Start date selection e.g. "2016-01-01 00:00:00". If not set, the
            whole time range will be plotted.
        date_to : string
            End date selection e.g. "2016-03-01 00:00:00". If not set, the
            whole time range will be plotted.

        """
        kwargs.setdefault('bus_label', slice(None))
        kwargs.setdefault('type', slice(None))
        kwargs.setdefault('obj_label', slice(None))
        kwargs.setdefault(
            'date_from', self.data.index.get_level_values('datetime')[0])
        kwargs.setdefault(
            'date_to', self.data.index.get_level_values('datetime')[-1])

        # slicing
        idx = pd.IndexSlice

        subset = self.data.loc[idx[
            kwargs['bus_label'],
            kwargs['type'],
            kwargs['obj_label'],
            slice(pd.Timestamp(kwargs['date_from']),
                  pd.Timestamp(kwargs['date_to']))], :]

        return subset

    def slice_unstacked(self, unstacklevel='obj_label',
                        formatted=False, **kwargs):
        r"""Method for slicing the ResultsDataFrame. An unstacked
        subset is returned.

        Parameters
        ----------
        unstacklevel : string (default: 'obj_label')
            Level to unstack the subset of the DataFrame.
        formatted : boolean
            missing...
        """
        subset = self.slice_by(**kwargs)
        subset = subset.unstack(level=unstacklevel)
        if formatted is True:
            subset.reset_index(level=['bus_label', 'type'], drop=True,
                               inplace=True)
        # user standard insteadt of multi-indexed columns
        subset.columns = subset.columns.get_level_values(1).unique()
        # return subset
        self.subset = subset

    def plot(self, **kwargs):
        r""" Passing the data attribute to the pandas plotting method. All
        parameters will be directly passed to pandas.DataFrame.plot(). See
        http://pandas.pydata.org/pandas-docs/stable/generated/pandas.DataFrame.plot.html
        for more information.

        Returns
        -------
        self
        """
        self.ax = self.subset.plot(**kwargs)
        return self

    def io_plot(self, bus_label, cdict, line_kwa=None, lineorder=None,
                bar_kwa=None, barorder=None, **kwargs):
        r""" Plotting a combined bar and line plot to see the fitting of in-
        and outcomming flows of a bus balance.

        Parameters
        ----------
        bus_label : string
            Uid of the bus to plot the balance.
        cdict : dictionary
            A dictionary that has all possible components as keys and its
            colors as items.
        line_kwa : dictionary
            Keyword arguments to be passed to the pandas line plot.
        bar_kwa : dictionary
            Keyword arguments to be passed to the pandas bar plot.
        lineorder : list
            Order of columns to plot the line plot
        barorder : list
            Order of columns to plot the bar plot

        Note
        ----
        Further keyword arguments will be passed to the
        :class:`slice_unstacked method <DataFramePlot.slice_unstacked>`.

        Returns
        -------
        handles, labels
            Manipulated labels to correct the unsual construction of the
            stack line plot. You can use them for further maipulations.
        """
        self.ax = kwargs.get('ax', self.ax)

        if bar_kwa is None:
            bar_kwa = dict()
        if line_kwa is None:
            line_kwa = dict()

        if self.ax is None:
            fig = plt.figure()
            self.ax = fig.add_subplot(1, 1, 1)

        # Create a bar plot for all input flows
        self.slice_unstacked(bus_label=bus_label, type='to_bus', **kwargs)

        if barorder is not None:
            self.rearrange_subset(barorder)

        self.subset.plot(kind='bar', linewidth=0, stacked=True, width=1,
                         ax=self.ax, color=self.color_from_dict(cdict),
                         **bar_kwa)

        # Create a line plot for all output flows
        self.slice_unstacked(bus_label=bus_label, type='from_bus', **kwargs)
        if lineorder is not None:
            self.rearrange_subset(lineorder)
        # The following changes are made to have the bottom line on top layer
        # of all lines. Normally the bottom line is the first line that is
        # plotted and will be on the lowest layer. This is difficult to read.
        new_df = pd.DataFrame(index=self.subset.index)
        n = 0
        tmp = 0
        for col in self.subset.columns:
            if n < 1:
                new_df[col] = self.subset[col]
            else:
                new_df[col] = self.subset[col] + tmp
            tmp = new_df[col]
            n += 1
        if lineorder is None:
            new_df.sort_index(axis=1, ascending=False, inplace=True)
        else:
            lineorder.reverse()
            new_df = new_df[lineorder]
        colorlist = self.color_from_dict(cdict)
        if isinstance(colorlist, list):
            colorlist.reverse()
        separator = len(colorlist)
        new_df.plot(kind='line', ax=self.ax, color=colorlist,
                    drawstyle='steps-mid', **line_kwa)

        # Adapt the legend to the new oder
        handles, labels = self.ax.get_legend_handles_labels()
        tmp_lab = [x for x in reversed(labels[0:separator])]
        tmp_hand = [x for x in reversed(handles[0:separator])]
        handles = tmp_hand + handles[separator:]
        labels = tmp_lab + labels[separator:]
        labels.reverse()
        handles.reverse()

        self.ax.legend(handles, labels)
        return handles, labels

    def rearrange_subset(self, order):
        r"""
        Change the order of the subset DataFrame

        Parameters
        ----------
        order : list
            New order of columns

        Returns
        -------
        self
        """
        cols = list(self.subset.columns.values)
        neworder = [x for x in list(order) if x in set(cols)]
        missing = [x for x in list(cols) if x not in set(order)]
        if len(missing) > 0:
            logging.warning(
                "Columns that are not part of the order list are removed: " +
                str(missing))
        self.subset = self.subset[neworder]

    def color_from_dict(self, colordict):
        r""" Method to convert a dictionary containing the components and its
        colors to a color list that can be directly useed with the color
        parameter of the pandas plotting method.

        Parameters
        ----------
        colordict : dictionary
            A dictionary that has all possible components as keys and its
            colors as items.

        Returns
        -------
        list
            Containing the colors of all components of the subset attribute
        """
        tmplist = list(
            map(colordict.get, list(self.subset.columns)))
        tmplist = ['#00FFFF' if v is None else v for v in tmplist]
        if len(tmplist) == 1:
            colorlist = tmplist[0]
        else:
            colorlist = tmplist
        return colorlist

    def set_datetime_ticks(self, tick_distance=None, number_autoticks=3,
                           date_format='%d-%m-%Y %H:%M'):
        r""" Set configurable ticks for the time axis. One can choose the
        number of ticks or the distance between ticks and the format.

        Parameters
        ----------
        tick_distance : real
            The disctance between to ticks in hours. If not set autoticks are
            set (see number_autoticks).
        number_autoticks : int (default: 3)
            The number of ticks on the time axis, independent of the time
            range. The higher the number of ticks is, the shorter should be the
            date_format string.
        date_format : string (default: '%d-%m-%Y %H:%M')
            The string to define the format of the date and time. See
            https://docs.python.org/3/library/datetime.html#strftime-and-strptime-behavior
            for more information.
        """
        dates = self.subset.index.get_level_values('datetime').unique()
        if tick_distance is None:
            tick_distance = int(len(dates) / number_autoticks) - 1
        self.ax.set_xticks(range(0, len(dates), tick_distance),
                           minor=False)
        self.ax.set_xticklabels(
            [item.strftime(date_format)
             for item in dates.tolist()[0::tick_distance]],
            rotation=0, minor=False)

    def outside_legend(self, reverse=False, plotshare=0.9, **kwargs):
        r""" Move the legend outside the plot. Bases on the ideas of Joe
        Kington. See
        http://stackoverflow.com/questions/4700614/how-to-put-the-legend-out-of-the-plot
        for more information.

        Parameters
        ----------
        reverse : boolean (default: False)
            Print out the legend in reverse order. This is interesting for
            stack-plots to have the legend in the same order as the stacks.
        plotshare : real (default: 0.9)
            Share of the plot area to create space for the legend (0 to 1).
        loc : string (default: 'center left')
            Location of the plot.
        bbox_to_anchor : tuple (default: (1, 0.5))
            Set the anchor for the legend.
        ncol : integer (default: 1)
            Number of columns of the legend.
        handles : list of handles
            A list of handels if they are already modified by another function
            or method. Normally these handles will be automatically taken from
            the artis object.
        lables : list of labels
            A list of labels if they are already modified by another function
            or method. Normally these handles will be automatically taken from
            the artis object.
        Note
        ----
        All keyword arguments (kwargs) will be directly passed to the
        matplotlib legend class. See
        http://matplotlib.org/api/legend_api.html#matplotlib.legend.Legend
        for more parameters.
        """
        kwargs.setdefault('loc', 'center left')
        kwargs.setdefault('bbox_to_anchor', (1, 0.5))
        kwargs.setdefault('ncol', 1)
        handles = kwargs.pop('handles', self.ax.get_legend_handles_labels()[0])
        labels = kwargs.pop('labels', self.ax.get_legend_handles_labels()[1])

        if reverse:
            handles.reverse()
            labels.reverse()

        box = self.ax.get_position()
        self.ax.set_position([box.x0, box.y0, box.width * plotshare,
                              box.height])

        self.ax.legend(handles, labels, **kwargs)

if __name__ == '__main__':
    pass


# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
#
# 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.
# ==============================================================================
"""Tests for dnn.py."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import shutil
import tempfile

import numpy as np
import six

from tensorflow.core.example import example_pb2
from tensorflow.core.example import feature_pb2
from tensorflow.python.estimator.canned import dnn
from tensorflow.python.estimator.canned import dnn_testing_utils
from tensorflow.python.estimator.canned import prediction_keys
from tensorflow.python.estimator.export import export
from tensorflow.python.estimator.inputs import numpy_io
from tensorflow.python.estimator.inputs import pandas_io
from tensorflow.python.feature_column import feature_column
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.ops import data_flow_ops
from tensorflow.python.ops import parsing_ops
from tensorflow.python.platform import gfile
from tensorflow.python.platform import test
from tensorflow.python.summary.writer import writer_cache
from tensorflow.python.training import input as input_lib
from tensorflow.python.training import queue_runner

try:
  # pylint: disable=g-import-not-at-top
  import pandas as pd
  HAS_PANDAS = True
except IOError:
  # Pandas writes a temporary file during import. If it fails, don't use pandas.
  HAS_PANDAS = False
except ImportError:
  HAS_PANDAS = False

def _dnn_classifier_fn(*args, **kwargs):
  return dnn.DNNClassifier(*args, **kwargs)

class DNNModelFnTest(dnn_testing_utils.BaseDNNModelFnTest, test.TestCase):

  def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
    test.TestCase.__init__(self, methodName)
    dnn_testing_utils.BaseDNNModelFnTest.__init__(self, dnn._dnn_model_fn)

class DNNLogitFnTest(dnn_testing_utils.BaseDNNLogitFnTest, test.TestCase):

  def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
    test.TestCase.__init__(self, methodName)
    dnn_testing_utils.BaseDNNLogitFnTest.__init__(self,
                                                  dnn._dnn_logit_fn_builder)

class DNNWarmStartingTest(dnn_testing_utils.BaseDNNWarmStartingTest,
                          test.TestCase):

  def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
    test.TestCase.__init__(self, methodName)
    dnn_testing_utils.BaseDNNWarmStartingTest.__init__(self, _dnn_classifier_fn,
                                                       _dnn_regressor_fn)

class DNNClassifierEvaluateTest(
    dnn_testing_utils.BaseDNNClassifierEvaluateTest, test.TestCase):

  def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
    test.TestCase.__init__(self, methodName)
    dnn_testing_utils.BaseDNNClassifierEvaluateTest.__init__(
        self, _dnn_classifier_fn)

class DNNClassifierPredictTest(
    dnn_testing_utils.BaseDNNClassifierPredictTest, test.TestCase):

  def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
    test.TestCase.__init__(self, methodName)
    dnn_testing_utils.BaseDNNClassifierPredictTest.__init__(
        self, _dnn_classifier_fn)

class DNNClassifierTrainTest(
    dnn_testing_utils.BaseDNNClassifierTrainTest, test.TestCase):

  def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
    test.TestCase.__init__(self, methodName)
    dnn_testing_utils.BaseDNNClassifierTrainTest.__init__(
        self, _dnn_classifier_fn)

def _dnn_regressor_fn(*args, **kwargs):
  return dnn.DNNRegressor(*args, **kwargs)

class DNNRegressorEvaluateTest(
    dnn_testing_utils.BaseDNNRegressorEvaluateTest, test.TestCase):

  def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
    test.TestCase.__init__(self, methodName)
    dnn_testing_utils.BaseDNNRegressorEvaluateTest.__init__(
        self, _dnn_regressor_fn)

class DNNRegressorPredictTest(
    dnn_testing_utils.BaseDNNRegressorPredictTest, test.TestCase):

  def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
    test.TestCase.__init__(self, methodName)
    dnn_testing_utils.BaseDNNRegressorPredictTest.__init__(
        self, _dnn_regressor_fn)

class DNNRegressorTrainTest(
    dnn_testing_utils.BaseDNNRegressorTrainTest, test.TestCase):

  def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
    test.TestCase.__init__(self, methodName)
    dnn_testing_utils.BaseDNNRegressorTrainTest.__init__(
        self, _dnn_regressor_fn)

def _queue_parsed_features(feature_map):
  tensors_to_enqueue = []
  keys = []
  for key, tensor in six.iteritems(feature_map):
    keys.append(key)
    tensors_to_enqueue.append(tensor)
  queue_dtypes = [x.dtype for x in tensors_to_enqueue]
  input_queue = data_flow_ops.FIFOQueue(capacity=100, dtypes=queue_dtypes)
  queue_runner.add_queue_runner(
      queue_runner.QueueRunner(
          input_queue,
          [input_queue.enqueue(tensors_to_enqueue)]))
  dequeued_tensors = input_queue.dequeue()
  return {keys[i]: dequeued_tensors[i] for i in range(len(dequeued_tensors))}

class DNNRegressorIntegrationTest(test.TestCase):

  def setUp(self):
    self._model_dir = tempfile.mkdtemp()

  def tearDown(self):
    if self._model_dir:
      writer_cache.FileWriterCache.clear()
      shutil.rmtree(self._model_dir)

  def _test_complete_flow(
      self, train_input_fn, eval_input_fn, predict_input_fn, input_dimension,
      label_dimension, batch_size):
    feature_columns = [
        feature_column.numeric_column('x', shape=(input_dimension,))]
    est = dnn.DNNRegressor(
        hidden_units=(2, 2),
        feature_columns=feature_columns,
        label_dimension=label_dimension,
        model_dir=self._model_dir)

    # TRAIN
    num_steps = 10
    est.train(train_input_fn, steps=num_steps)

    # EVALUTE
    scores = est.evaluate(eval_input_fn)
    self.assertEqual(num_steps, scores[ops.GraphKeys.GLOBAL_STEP])
    self.assertIn('loss', six.iterkeys(scores))

    # PREDICT
    predictions = np.array([
        x[prediction_keys.PredictionKeys.PREDICTIONS]
        for x in est.predict(predict_input_fn)
    ])
    self.assertAllEqual((batch_size, label_dimension), predictions.shape)

    # EXPORT
    feature_spec = feature_column.make_parse_example_spec(feature_columns)
    serving_input_receiver_fn = export.build_parsing_serving_input_receiver_fn(
        feature_spec)
    export_dir = est.export_savedmodel(tempfile.mkdtemp(),
                                       serving_input_receiver_fn)
    self.assertTrue(gfile.Exists(export_dir))

  def test_numpy_input_fn(self):
    """Tests complete flow with numpy_input_fn."""
    label_dimension = 2
    batch_size = 10
    data = np.linspace(0., 2., batch_size * label_dimension, dtype=np.float32)
    data = data.reshape(batch_size, label_dimension)
    # learn y = x
    train_input_fn = numpy_io.numpy_input_fn(
        x={'x': data},
        y=data,
        batch_size=batch_size,
        num_epochs=None,
        shuffle=True)
    eval_input_fn = numpy_io.numpy_input_fn(
        x={'x': data},
        y=data,
        batch_size=batch_size,
        shuffle=False)
    predict_input_fn = numpy_io.numpy_input_fn(
        x={'x': data},
        batch_size=batch_size,
        shuffle=False)

    self._test_complete_flow(
        train_input_fn=train_input_fn,
        eval_input_fn=eval_input_fn,
        predict_input_fn=predict_input_fn,
        input_dimension=label_dimension,
        label_dimension=label_dimension,
        batch_size=batch_size)

  def test_pandas_input_fn(self):
    """Tests complete flow with pandas_input_fn."""
    if not HAS_PANDAS:
      return
    label_dimension = 1
    batch_size = 10
    data = np.linspace(0., 2., batch_size, dtype=np.float32)
    x = pd.DataFrame({'x': data})
    y = pd.Series(data)
    train_input_fn = pandas_io.pandas_input_fn(
        x=x,
        y=y,
        batch_size=batch_size,
        num_epochs=None,
        shuffle=True)
    eval_input_fn = pandas_io.pandas_input_fn(
        x=x,
        y=y,
        batch_size=batch_size,
        shuffle=False)
    predict_input_fn = pandas_io.pandas_input_fn(
        x=x,
        batch_size=batch_size,
        shuffle=False)

    self._test_complete_flow(
        train_input_fn=train_input_fn,
        eval_input_fn=eval_input_fn,
        predict_input_fn=predict_input_fn,
        input_dimension=label_dimension,
        label_dimension=label_dimension,
        batch_size=batch_size)

  def test_input_fn_from_parse_example(self):
    """Tests complete flow with input_fn constructed from parse_example."""
    label_dimension = 2
    batch_size = 10
    data = np.linspace(0., 2., batch_size * label_dimension, dtype=np.float32)
    data = data.reshape(batch_size, label_dimension)

    serialized_examples = []
    for datum in data:
      example = example_pb2.Example(features=feature_pb2.Features(
          feature={
              'x': feature_pb2.Feature(
                  float_list=feature_pb2.FloatList(value=datum)),
              'y': feature_pb2.Feature(
                  float_list=feature_pb2.FloatList(value=datum)),
          }))
      serialized_examples.append(example.SerializeToString())

    feature_spec = {
        'x': parsing_ops.FixedLenFeature([label_dimension], dtypes.float32),
        'y': parsing_ops.FixedLenFeature([label_dimension], dtypes.float32),
    }
    def _train_input_fn():
      feature_map = parsing_ops.parse_example(serialized_examples, feature_spec)
      features = _queue_parsed_features(feature_map)
      labels = features.pop('y')
      return features, labels
    def _eval_input_fn():
      feature_map = parsing_ops.parse_example(
          input_lib.limit_epochs(serialized_examples, num_epochs=1),
          feature_spec)
      features = _queue_parsed_features(feature_map)
      labels = features.pop('y')
      return features, labels
    def _predict_input_fn():
      feature_map = parsing_ops.parse_example(
          input_lib.limit_epochs(serialized_examples, num_epochs=1),
          feature_spec)
      features = _queue_parsed_features(feature_map)
      features.pop('y')
      return features, None

    self._test_complete_flow(
        train_input_fn=_train_input_fn,
        eval_input_fn=_eval_input_fn,
        predict_input_fn=_predict_input_fn,
        input_dimension=label_dimension,
        label_dimension=label_dimension,
        batch_size=batch_size)

class DNNClassifierIntegrationTest(test.TestCase):

  def setUp(self):
    self._model_dir = tempfile.mkdtemp()

  def tearDown(self):
    if self._model_dir:
      writer_cache.FileWriterCache.clear()
      shutil.rmtree(self._model_dir)

  def _as_label(self, data_in_float):
    return np.rint(data_in_float).astype(np.int64)

  def _test_complete_flow(
      self, train_input_fn, eval_input_fn, predict_input_fn, input_dimension,
      n_classes, batch_size):
    feature_columns = [
        feature_column.numeric_column('x', shape=(input_dimension,))]
    est = dnn.DNNClassifier(
        hidden_units=(2, 2),
        feature_columns=feature_columns,
        n_classes=n_classes,
        model_dir=self._model_dir)

    # TRAIN
    num_steps = 10
    est.train(train_input_fn, steps=num_steps)

    # EVALUTE
    scores = est.evaluate(eval_input_fn)
    self.assertEqual(num_steps, scores[ops.GraphKeys.GLOBAL_STEP])
    self.assertIn('loss', six.iterkeys(scores))

    # PREDICT
    predicted_proba = np.array([
        x[prediction_keys.PredictionKeys.PROBABILITIES]
        for x in est.predict(predict_input_fn)
    ])
    self.assertAllEqual((batch_size, n_classes), predicted_proba.shape)

    # EXPORT
    feature_spec = feature_column.make_parse_example_spec(feature_columns)
    serving_input_receiver_fn = export.build_parsing_serving_input_receiver_fn(
        feature_spec)
    export_dir = est.export_savedmodel(tempfile.mkdtemp(),
                                       serving_input_receiver_fn)
    self.assertTrue(gfile.Exists(export_dir))

  def test_numpy_input_fn(self):
    """Tests complete flow with numpy_input_fn."""
    n_classes = 3
    input_dimension = 2
    batch_size = 10
    data = np.linspace(
        0., n_classes - 1., batch_size * input_dimension, dtype=np.float32)
    x_data = data.reshape(batch_size, input_dimension)
    y_data = np.reshape(self._as_label(data[:batch_size]), (batch_size, 1))
    # learn y = x
    train_input_fn = numpy_io.numpy_input_fn(
        x={'x': x_data},
        y=y_data,
        batch_size=batch_size,
        num_epochs=None,
        shuffle=True)
    eval_input_fn = numpy_io.numpy_input_fn(
        x={'x': x_data},
        y=y_data,
        batch_size=batch_size,
        shuffle=False)
    predict_input_fn = numpy_io.numpy_input_fn(
        x={'x': x_data},
        batch_size=batch_size,
        shuffle=False)

    self._test_complete_flow(
        train_input_fn=train_input_fn,
        eval_input_fn=eval_input_fn,
        predict_input_fn=predict_input_fn,
        input_dimension=input_dimension,
        n_classes=n_classes,
        batch_size=batch_size)

  def test_pandas_input_fn(self):
    """Tests complete flow with pandas_input_fn."""
    if not HAS_PANDAS:
      return
    input_dimension = 1
    n_classes = 3
    batch_size = 10
    data = np.linspace(0., n_classes - 1., batch_size, dtype=np.float32)
    x = pd.DataFrame({'x': data})
    y = pd.Series(self._as_label(data))
    train_input_fn = pandas_io.pandas_input_fn(
        x=x,
        y=y,
        batch_size=batch_size,
        num_epochs=None,
        shuffle=True)
    eval_input_fn = pandas_io.pandas_input_fn(
        x=x,
        y=y,
        batch_size=batch_size,
        shuffle=False)
    predict_input_fn = pandas_io.pandas_input_fn(
        x=x,
        batch_size=batch_size,
        shuffle=False)

    self._test_complete_flow(
        train_input_fn=train_input_fn,
        eval_input_fn=eval_input_fn,
        predict_input_fn=predict_input_fn,
        input_dimension=input_dimension,
        n_classes=n_classes,
        batch_size=batch_size)

  def test_input_fn_from_parse_example(self):
    """Tests complete flow with input_fn constructed from parse_example."""
    input_dimension = 2
    n_classes = 3
    batch_size = 10
    data = np.linspace(
        0., n_classes - 1., batch_size * input_dimension, dtype=np.float32)
    data = data.reshape(batch_size, input_dimension)

    serialized_examples = []
    for datum in data:
      example = example_pb2.Example(features=feature_pb2.Features(
          feature={
              'x':
                  feature_pb2.Feature(float_list=feature_pb2.FloatList(
                      value=datum)),
              'y':
                  feature_pb2.Feature(int64_list=feature_pb2.Int64List(
                      value=self._as_label(datum[:1]))),
          }))
      serialized_examples.append(example.SerializeToString())

    feature_spec = {
        'x': parsing_ops.FixedLenFeature([input_dimension], dtypes.float32),
        'y': parsing_ops.FixedLenFeature([1], dtypes.int64),
    }
    def _train_input_fn():
      feature_map = parsing_ops.parse_example(serialized_examples, feature_spec)
      features = _queue_parsed_features(feature_map)
      labels = features.pop('y')
      return features, labels
    def _eval_input_fn():
      feature_map = parsing_ops.parse_example(
          input_lib.limit_epochs(serialized_examples, num_epochs=1),
          feature_spec)
      features = _queue_parsed_features(feature_map)
      labels = features.pop('y')
      return features, labels
    def _predict_input_fn():
      feature_map = parsing_ops.parse_example(
          input_lib.limit_epochs(serialized_examples, num_epochs=1),
          feature_spec)
      features = _queue_parsed_features(feature_map)
      features.pop('y')
      return features, None

    self._test_complete_flow(
        train_input_fn=_train_input_fn,
        eval_input_fn=_eval_input_fn,
        predict_input_fn=_predict_input_fn,
        input_dimension=input_dimension,
        n_classes=n_classes,
        batch_size=batch_size)

if __name__ == '__main__':
  test.main()


"""
==============================
Probability Calibration curves
==============================

When performing classification one often wants to predict not only the class
label, but also the associated probability. This probability gives some
kind of confidence on the prediction. This example demonstrates how to display
how well calibrated the predicted probabilities are and how to calibrate an
uncalibrated classifier.

The experiment is performed on an artificial dataset for binary classification
with 100.000 samples (1.000 of them are used for model fitting) with 20
features. Of the 20 features, only 2 are informative and 10 are redundant. The
first figure shows the estimated probabilities obtained with logistic
regression, Gaussian naive Bayes, and Gaussian naive Bayes with both isotonic
calibration and sigmoid calibration. The calibration performance is evaluated
with Brier score, reported in the legend (the smaller the better). One can
observe here that logistic regression is well calibrated while raw Gaussian
naive Bayes performs very badly. This is because of the redundant features
which violate the assumption of feature-independence and result in an overly
confident classifier, which is indicated by the typical transposed-sigmoid
curve.

Calibration of the probabilities of Gaussian naive Bayes with isotonic
regression can fix this issue as can be seen from the nearly diagonal
calibration curve. Sigmoid calibration also improves the brier score slightly,
albeit not as strongly as the non-parametric isotonic regression. This can be
attributed to the fact that we have plenty of calibration data such that the
greater flexibility of the non-parametric model can be exploited.

The second figure shows the calibration curve of a linear support-vector
classifier (LinearSVC). LinearSVC shows the opposite behavior as Gaussian
naive Bayes: the calibration curve has a sigmoid curve, which is typical for
an under-confident classifier. In the case of LinearSVC, this is caused by the
margin property of the hinge loss, which lets the model focus on hard samples
that are close to the decision boundary (the support vectors).

Both kinds of calibration can fix this issue and yield nearly identical
results. This shows that sigmoid calibration can deal with situations where
the calibration curve of the base classifier is sigmoid (e.g., for LinearSVC)
but not where it is transposed-sigmoid (e.g., Gaussian naive Bayes).
"""
print(__doc__)

# Author: Alexandre Gramfort <alexandre.gramfort@telecom-paristech.fr>
#         Jan Hendrik Metzen <jhm@informatik.uni-bremen.de>
# License: BSD Style.

import matplotlib.pyplot as plt

from sklearn import datasets
from sklearn.naive_bayes import GaussianNB
from sklearn.svm import LinearSVC
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import (brier_score_loss, precision_score, recall_score,
                             f1_score)
from sklearn.calibration import CalibratedClassifierCV, calibration_curve
from sklearn.cross_validation import train_test_split

# Create dataset of classification task with many redundant and few
# informative features
X, y = datasets.make_classification(n_samples=100000, n_features=20,
                                    n_informative=2, n_redundant=10,
                                    random_state=42)

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.99,
                                                    random_state=42)

def plot_calibration_curve(est, name, fig_index):
    """Plot calibration curve for est w/o and with calibration. """
    # Calibrated with isotonic calibration
    isotonic = CalibratedClassifierCV(est, cv=2, method='isotonic')

    # Calibrated with sigmoid calibration
    sigmoid = CalibratedClassifierCV(est, cv=2, method='sigmoid')

    # Logistic regression with no calibration as baseline
    lr = LogisticRegression(C=1., solver='lbfgs')

    fig = plt.figure(fig_index, figsize=(10, 10))
    ax1 = plt.subplot2grid((3, 1), (0, 0), rowspan=2)
    ax2 = plt.subplot2grid((3, 1), (2, 0))

    ax1.plot([0, 1], [0, 1], "k:", label="Perfectly calibrated")
    for clf, name in [(lr, 'Logistic'),
                      (est, name),
                      (isotonic, name + ' + Isotonic'),
                      (sigmoid, name + ' + Sigmoid')]:
        clf.fit(X_train, y_train)
        y_pred = clf.predict(X_test)
        if hasattr(clf, "predict_proba"):
            prob_pos = clf.predict_proba(X_test)[:, 1]
        else:  # use decision function
            prob_pos = clf.decision_function(X_test)
            prob_pos = \
                (prob_pos - prob_pos.min()) / (prob_pos.max() - prob_pos.min())

        clf_score = brier_score_loss(y_test, prob_pos, pos_label=y.max())
        print("%s:" % name)
        print("\tBrier: %1.3f" % (clf_score))
        print("\tPrecision: %1.3f" % precision_score(y_test, y_pred))
        print("\tRecall: %1.3f" % recall_score(y_test, y_pred))
        print("\tF1: %1.3f\n" % f1_score(y_test, y_pred))

        fraction_of_positives, mean_predicted_value = \
            calibration_curve(y_test, prob_pos, n_bins=10)

        ax1.plot(mean_predicted_value, fraction_of_positives, "s-",
                 label="%s (%1.3f)" % (name, clf_score))

        ax2.hist(prob_pos, range=(0, 1), bins=10, label=name,
                 histtype="step", lw=2)

    ax1.set_ylabel("Fraction of positives")
    ax1.set_ylim([-0.05, 1.05])
    ax1.legend(loc="lower right")
    ax1.set_title('Calibration plots  (reliability curve)')

    ax2.set_xlabel("Mean predicted value")
    ax2.set_ylabel("Count")
    ax2.legend(loc="upper center", ncol=2)

    plt.tight_layout()

# Plot calibration cuve for Gaussian Naive Bayes
plot_calibration_curve(GaussianNB(), "Naive Bayes", 1)

# Plot calibration cuve for Linear SVC
plot_calibration_curve(LinearSVC(), "SVC", 2)

plt.show()


# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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.
#!/usr/bin/env python

# ===================================
# file: metrics.py
# ===================================

import inspect
import json
import logging
import os
import signal
import subprocess
import sys
import traceback

import csv
import time
import matplotlib as mpl
mpl.use('Agg')
import matplotlib.pyplot as plt
from collections import namedtuple
import numpy

from pyh import *
import kafka_system_test_utils
import system_test_utils

logger     = logging.getLogger("namedLogger")
thisClassName = '(metrics)'
d = {'name_of_class': thisClassName}

attributeNameToNameInReportedFileMap = {
    'Min': 'min',
    'Max': 'max',
    'Mean': 'mean',
    '50thPercentile': 'median',
    'StdDev': 'stddev',
    '95thPercentile': '95%',
    '99thPercentile': '99%',
    '999thPercentile': '99.9%',
    'Count': 'count',
    'OneMinuteRate': '1 min rate',
    'MeanRate': 'mean rate',
    'FiveMinuteRate': '5 min rate',
    'FifteenMinuteRate': '15 min rate',
    'Value': 'value'
}

def getCSVFileNameFromMetricsMbeanName(mbeanName):
    return mbeanName.replace(":type=", ".").replace(",name=", ".") + ".csv"

def read_metrics_definition(metricsFile):
    metricsFileData = open(metricsFile, "r").read()
    metricsJsonData = json.loads(metricsFileData)
    allDashboards = metricsJsonData['dashboards']
    allGraphs = []
    for dashboard in allDashboards:
        dashboardName = dashboard['name']
        graphs = dashboard['graphs']
        for graph in graphs:
            bean = graph['bean_name']
            allGraphs.append(graph)
            attributes = graph['attributes']
            #print "Filtering on attributes " + attributes
    return allGraphs

def get_dashboard_definition(metricsFile, role):
    metricsFileData = open(metricsFile, "r").read()
    metricsJsonData = json.loads(metricsFileData)
    allDashboards = metricsJsonData['dashboards']
    dashboardsForRole = []
    for dashboard in allDashboards:
        if dashboard['role'] == role:
            dashboardsForRole.append(dashboard)
    return dashboardsForRole

def ensure_valid_headers(headers, attributes):
    if headers[0] != "# time":
        raise Exception("First column should be time")
    for header in headers:
        logger.debug(header, extra=d)
    # there should be exactly one column with a name that matches attributes
    try:
        attributeColumnIndex = headers.index(attributes)
        return attributeColumnIndex
    except ValueError as ve:
        #print "#### attributes : ", attributes
        #print "#### headers    : ", headers
        raise Exception("There should be exactly one column that matches attribute: {0} in".format(attributes) +
                        " headers: {0}".format(",".join(headers)))

def plot_graphs(inputCsvFiles, labels, title, xLabel, yLabel, attribute, outputGraphFile):
    # create empty plot
    fig=plt.figure()
    fig.subplots_adjust(bottom=0.2)
    ax=fig.add_subplot(111)
    labelx = -0.3  # axes coords
    ax.set_xlabel(xLabel)
    ax.set_ylabel(yLabel)
    ax.grid()
    #ax.yaxis.set_label_coords(labelx, 0.5)
    Coordinates = namedtuple("Coordinates", 'x y')
    plots = []
    coordinates = []
    # read data for all files, organize by label in a dict
    for fileAndLabel in zip(inputCsvFiles, labels):
        inputCsvFile = fileAndLabel[0]
        label = fileAndLabel[1]
        csv_reader = list(csv.reader(open(inputCsvFile, "rb")))
        x,y = [],[]
        xticks_labels = []
        try:
            # read first line as the headers
            headers = csv_reader.pop(0)
            attributeColumnIndex = ensure_valid_headers(headers, attributeNameToNameInReportedFileMap[attribute])
            logger.debug("Column index for attribute {0} is {1}".format(attribute, attributeColumnIndex), extra=d)
            start_time = (int)(os.path.getctime(inputCsvFile) * 1000)
            int(csv_reader[0][0])
            for line in csv_reader:
                if(len(line) == 0):
                    continue
                yVal = float(line[attributeColumnIndex])
                xVal = int(line[0])
                y.append(yVal)
                epoch= start_time + int(line[0])
                x.append(xVal)
                xticks_labels.append(time.strftime("%H:%M:%S", time.localtime(epoch)))
                coordinates.append(Coordinates(xVal, yVal))
            p1 = ax.plot(x,y)
            plots.append(p1)
        except Exception as e:
            logger.error("ERROR while plotting data for {0}: {1}".format(inputCsvFile, e), extra=d)
            traceback.print_exc()
    # find xmin, xmax, ymin, ymax from all csv files
    xmin = min(map(lambda coord: coord.x, coordinates))
    xmax = max(map(lambda coord: coord.x, coordinates))
    ymin = min(map(lambda coord: coord.y, coordinates))
    ymax = max(map(lambda coord: coord.y, coordinates))
    # set x and y axes limits
    plt.xlim(xmin, xmax)
    plt.ylim(ymin, ymax)
    # set ticks accordingly
    xticks = numpy.arange(xmin, xmax, 0.2*xmax)
#    yticks = numpy.arange(ymin, ymax)
    plt.xticks(xticks,xticks_labels,rotation=17)
#    plt.yticks(yticks)
    plt.legend(plots,labels, loc=2)
    plt.title(title)
    plt.savefig(outputGraphFile)

def draw_all_graphs(metricsDescriptionFile, testcaseEnv, clusterConfig):
    # go through each role and plot graphs for the role's metrics
    roles = set(map(lambda config: config['role'], clusterConfig))
    for role in roles:
        dashboards = get_dashboard_definition(metricsDescriptionFile, role)
        entities = kafka_system_test_utils.get_entities_for_role(clusterConfig, role)
        for dashboard in dashboards:
            graphs = dashboard['graphs']
            # draw each graph for all entities
            draw_graph_for_role(graphs, entities, role, testcaseEnv)

def draw_graph_for_role(graphs, entities, role, testcaseEnv):
    for graph in graphs:
        graphName = graph['graph_name']
        yLabel = graph['y_label']
        inputCsvFiles = []
        graphLegendLabels = []
        for entity in entities:
            entityMetricsDir = kafka_system_test_utils.get_testcase_config_log_dir_pathname(testcaseEnv, role, entity['entity_id'], "metrics")
            entityMetricCsvFile = entityMetricsDir + "/" + getCSVFileNameFromMetricsMbeanName(graph['bean_name'])
            if(not os.path.exists(entityMetricCsvFile)):
                logger.warn("The file {0} does not exist for plotting".format(entityMetricCsvFile), extra=d)
            else:
                inputCsvFiles.append(entityMetricCsvFile)
                graphLegendLabels.append(role + "-" + entity['entity_id'])
#            print "Plotting graph for metric {0} on entity {1}".format(graph['graph_name'], entity['entity_id'])
        try:
            # plot one graph per mbean attribute
            labels = graph['y_label'].split(',')
            fullyQualifiedAttributeNames = map(lambda attribute: graph['bean_name'] + ':' + attribute,
                                           graph['attributes'].split(','))
            attributes = graph['attributes'].split(',')
            for labelAndAttribute in zip(labels, fullyQualifiedAttributeNames, attributes):
                outputGraphFile = testcaseEnv.testCaseDashboardsDir + "/" + role + "/" + labelAndAttribute[1] + ".svg"
                plot_graphs(inputCsvFiles, graphLegendLabels, graph['graph_name'] + '-' + labelAndAttribute[2],
                            "time", labelAndAttribute[0], labelAndAttribute[2], outputGraphFile)
#            print "Finished plotting graph for metric {0} on entity {1}".format(graph['graph_name'], entity['entity_id'])
        except Exception as e:
            logger.error("ERROR while plotting graph {0}: {1}".format(outputGraphFile, e), extra=d)
            traceback.print_exc()

def build_all_dashboards(metricsDefinitionFile, testcaseDashboardsDir, clusterConfig):
    metricsHtmlFile = testcaseDashboardsDir + "/metrics.html"
    centralDashboard = PyH('Kafka Metrics Dashboard')
    centralDashboard << h1('Kafka Metrics Dashboard', cl='center')
    roles = set(map(lambda config: config['role'], clusterConfig))
    for role in roles:
        entities = kafka_system_test_utils.get_entities_for_role(clusterConfig, role)
        dashboardPagePath = build_dashboard_for_role(metricsDefinitionFile, role,
                                                     entities, testcaseDashboardsDir)
        centralDashboard << a(role, href = dashboardPagePath)
        centralDashboard << br()

    centralDashboard.printOut(metricsHtmlFile)

def build_dashboard_for_role(metricsDefinitionFile, role, entities, testcaseDashboardsDir):
    # build all dashboards for the input entity's based on its role. It can be one of kafka, zookeeper, producer
    # consumer
    dashboards = get_dashboard_definition(metricsDefinitionFile, role)
    entityDashboard = PyH('Kafka Metrics Dashboard for ' + role)
    entityDashboard << h1('Kafka Metrics Dashboard for ' + role, cl='center')
    entityDashboardHtml = testcaseDashboardsDir + "/" + role + "-dashboards.html"
    for dashboard in dashboards:
        # place the graph svg files in this dashboard
        allGraphs = dashboard['graphs']
        for graph in allGraphs:
            attributes = map(lambda attribute: graph['bean_name'] + ':' + attribute,
                                           graph['attributes'].split(','))
            for attribute in attributes:
                graphFileLocation = testcaseDashboardsDir + "/" + role + "/" + attribute + ".svg"
                entityDashboard << embed(src = graphFileLocation, type = "image/svg+xml")
    entityDashboard.printOut(entityDashboardHtml)
    return entityDashboardHtml

def start_metrics_collection(jmxHost, jmxPort, role, entityId, systemTestEnv, testcaseEnv):
    logger.info("starting metrics collection on jmx port : " + jmxPort, extra=d)
    jmxUrl = "service:jmx:rmi:///jndi/rmi://" + jmxHost + ":" + jmxPort + "/jmxrmi"
    clusterConfig = systemTestEnv.clusterEntityConfigDictList
    metricsDefinitionFile = systemTestEnv.METRICS_PATHNAME
    entityMetricsDir = kafka_system_test_utils.get_testcase_config_log_dir_pathname(testcaseEnv, role, entityId, "metrics")
    dashboardsForRole = get_dashboard_definition(metricsDefinitionFile, role)
    mbeansForRole = get_mbeans_for_role(dashboardsForRole)

    kafkaHome = system_test_utils.get_data_by_lookup_keyval(clusterConfig, "entity_id", entityId, "kafka_home")
    javaHome  = system_test_utils.get_data_by_lookup_keyval(clusterConfig, "entity_id", entityId, "java_home")

    for mbean in mbeansForRole:
        outputCsvFile = entityMetricsDir + "/" + mbean + ".csv"
        startMetricsCmdList = ["ssh " + jmxHost,
                               "'JAVA_HOME=" + javaHome,
                               "JMX_PORT= " + kafkaHome + "/bin/kafka-run-class.sh kafka.tools.JmxTool",
                               "--jmx-url " + jmxUrl,
                               "--object-name " + mbean + " 1> ",
                                outputCsvFile + " & echo pid:$! > ",
                                entityMetricsDir + "/entity_pid'"]

        startMetricsCommand = " ".join(startMetricsCmdList)
        logger.debug("executing command: [" + startMetricsCommand + "]", extra=d)
        system_test_utils.async_sys_call(startMetricsCommand)
        time.sleep(1)

        pidCmdStr = "ssh " + jmxHost + " 'cat " + entityMetricsDir + "/entity_pid' 2> /dev/null"
        logger.debug("executing command: [" + pidCmdStr + "]", extra=d)
        subproc = system_test_utils.sys_call_return_subproc(pidCmdStr)

        # keep track of JMX ppid in a dictionary of entity_id to list of JMX ppid
        # testcaseEnv.entityJmxParentPidDict:
        #   key: entity_id
        #   val: list of JMX ppid associated to that entity_id
        #   { 1: [1234, 1235, 1236], 2: [2234, 2235, 2236], ... }
        for line in subproc.stdout.readlines():
            line = line.rstrip('\n')
            logger.debug("line: [" + line + "]", extra=d)
            if line.startswith("pid"):
                logger.debug("found pid line: [" + line + "]", extra=d)
                tokens  = line.split(':')
                thisPid = tokens[1]
                if entityId not in testcaseEnv.entityJmxParentPidDict:
                    testcaseEnv.entityJmxParentPidDict[entityId] = []
                testcaseEnv.entityJmxParentPidDict[entityId].append(thisPid)
                #print "\n#### testcaseEnv.entityJmxParentPidDict ", testcaseEnv.entityJmxParentPidDict, "\n"

def stop_metrics_collection(jmxHost, jmxPort):
    logger.info("stopping metrics collection on " + jmxHost + ":" + jmxPort, extra=d)
    system_test_utils.sys_call("ps -ef | grep JmxTool | grep -v grep | grep " + jmxPort + " | awk '{print $2}' | xargs kill -9")

def get_mbeans_for_role(dashboardsForRole):
    graphs = reduce(lambda x,y: x+y, map(lambda dashboard: dashboard['graphs'], dashboardsForRole))
    return set(map(lambda metric: metric['bean_name'], graphs))