__author__ = "evas"
__docformat__ = "reStructuredText"

import logging

import numpy as np

# http://stackoverflow.com/questions/12459811/how-to-embed-matplotib-in-pyqt-for-dummies
# see also: http://matplotlib.org/users/navigation_toolbar.html
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.backends.backend_qt5agg import NavigationToolbar2QT as NavigationToolbar
from matplotlib.colors import LogNorm
from matplotlib.figure import Figure
from PyQt5 import QtWidgets
from PyQt5.QtCore import QObject, Qt, pyqtSignal

# a useful constant
from uwsift.common import Info
from uwsift.model.layer_model import LayerModel
from uwsift.queue import TASK_DOING, TASK_PROGRESS

# Stuff for custom toolbars
try:
    import matplotlib.backends.qt_editor.figureoptions as figureoptions
except ImportError:
    figureoptions = None

LOG = logging.getLogger(__name__)
DEFAULT_POINT_PROBE = "default_probe_name"

class CustomNavigationToolbar(NavigationToolbar):
    """Custom matplotlib toolbar."""

    def __init__(self, *args, **kwargs):
        self.__include_colorbar = kwargs.get("include_colorbar", False)
        super(CustomNavigationToolbar, self).__init__(*args, **kwargs)

    def edit_parameters(self):
        allaxes = self.canvas.figure.get_axes()
        if not allaxes:
            QtWidgets.QMessageBox.warning(self.parent, "Error", "There are no axes to edit.")
            return
        elif len(allaxes) == 1:
            (axes,) = allaxes
        else:
            titles = []
            not_colorbar_idx = -1
            for idx, axes in enumerate(allaxes):
                if any(x.colorbar for x in axes.images):
                    not_colorbar_idx = idx
                name = (
                    axes.get_title()
                    or " - ".join(filter(None, [axes.get_xlabel(), axes.get_ylabel()]))
                    or "<anonymous {} (id: {:#x})>".format(type(axes).__name__, id(axes))
                )
                titles.append(name)

            if len(titles) == 2 and not_colorbar_idx != -1 and not self.__include_colorbar:
                axes = allaxes[not_colorbar_idx]
            else:
                item, ok = QtWidgets.QInputDialog.getItem(self.parent, "Customize", "Select axes:", titles, 0, False)
                if ok:
                    axes = allaxes[titles.index(str(item))]
                else:
                    return

        figureoptions.figure_edit(axes, self)

class ProbeGraphManager(QObject):
    """The ProbeGraphManager manages the many tabs of the Area Probe Graphs."""

    # signals
    didChangeTab = pyqtSignal(
        tuple,
    )  # list of probe areas to show
    didClonePolygon = pyqtSignal(str, str)
    drawChildGraph = pyqtSignal(
        str,
    )
    pointProbeChanged = pyqtSignal(str, bool, tuple)

    def __init__(self, tab_widget, auto_update_checkbox, update_button, workspace, layer_model: LayerModel, queue):
        """Set up our tab widget with an appropriate graph object in the first
        tab.

        FUTURE, once we are saving our graph configurations, load those instead
        of setting up this default.

        :param auto_update_checkbox: the QCheckBox defined in the pov_main.ui
               file. It's logic - to switch on/off automatic update of graphs
               is managed here

        :param update_button: the QButton defined in the pov_main.ui
               file to trigger manual graph updates in case auto_update_checkbox
               is off.
        """

        super(ProbeGraphManager, self).__init__(tab_widget)

        # hang on to the workspace
        self.workspace = workspace
        self.layer_model = layer_model
        self.queue = queue

        # hang on to the tab widget
        self.tab_widget_object = tab_widget
        self.new_tab_button = QtWidgets.QToolButton()
        self.new_tab_button.setText("+")
        self.tab_widget_object.setCornerWidget(self.new_tab_button, corner=Qt.TopLeftCorner)
        self.tab_widget_object.clear()  # Delete all tabs that may have been created in the Designer
        self.auto_update_checkbox = auto_update_checkbox
        self.update_button = update_button
        # hold on to point probe locations (point probes are shared across tabs)
        self.point_probes: dict = {}

        # set up the first tab
        self.graphs: list = []
        self.selected_graph_index = -1
        self.max_tab_letter = "A"

        # hook things up so we know when the selected tab changes
        self.tab_widget_object.currentChanged[int].connect(self._handle_tab_change)
        self.drawChildGraph.connect(self._draw_child)

        # hook up signals relating to changes in the number of tabs
        self.new_tab_button.clicked.connect(self._add_tab)

        # hook up auto update vs manual update changes
        self.update_button.clicked.connect(self.handleActiveProductDatasetsChanged)
        self.update_button.clicked.connect(self._update_default_point_probe_graph)
        self.auto_update_checkbox.stateChanged.connect(self._on_auto_update_checkbox_state_changed)
        self.auto_update_checkbox.setCheckState(Qt.Unchecked)

    def _draw_child(self, child_name):
        for child in self.graphs:
            if child.myName == child_name:
                child._draw()
                break

    def set_up_tab(self, tab_index, do_increment_tab_letter=True):
        """Create a new tab at tab_index and add it to the list of graphs"""

        # increment our tab label letter if desired
        if do_increment_tab_letter:
            self.max_tab_letter = chr(ord(self.max_tab_letter) + 1)  # this will get strange after Z!

        # create our tab
        temp_widget = QtWidgets.QWidget()
        self.tab_widget_object.insertTab(tab_index, temp_widget, self.max_tab_letter)

        # create the associated graph display object
        graph = ProbeGraphDisplay(self, temp_widget, self.workspace, self.layer_model, self.queue, self.max_tab_letter)
        self.graphs.append(graph)

        # load up the layers for this new tab
        graph.set_possible_layers()

        # clone the previous tab
        if self.selected_graph_index != tab_index:
            # if we aren't setting up the initial tab, clone the current tab
            current_graph = self.graphs[self.selected_graph_index]
            graph.set_default_layer_selections([current_graph.xSelectedUUID, current_graph.ySelectedUUID])
            # give it a copy of the current full_data_selection or polygon
            if current_graph.full_data_selection:
                graph.setRegion(select_full_data=graph.full_data_selection)
            else:
                graph.setRegion(polygon_points=current_graph.polygon[:] if current_graph.polygon is not None else None)
            graph.checked = current_graph.checked
            point_status, point_xy = self.point_probes.get(DEFAULT_POINT_PROBE, (None, None))
            point_xy = point_xy if point_status else None
            graph.setPoint(point_xy, rebuild=False)

        # Create the initial plot
        graph.rebuildPlot()

        # go to the tab we just created
        self.tab_widget_object.setCurrentIndex(tab_index)

    def handleActiveProductDatasetsChanged(self):
        """Used when the layer model signals that something about the layers
        has changed
        """

        # reload the layer list for the existing graphs
        for graphObj in self.graphs:
            doRebuild = graphObj is self.graphs[self.selected_graph_index]
            graphObj.set_possible_layers(do_rebuild_plot=doRebuild)  # FIXME

    def current_graph_set_region(self, polygon_points=None, select_full_data=False):
        """Update the current region in the selected graph and rebuild its plot

        :return: Name of the current probe graph ('A', 'B', ...)

        Probably outdated comment (TODO):
        FUTURE, once the polygon is a layer, this signal will be unnecessary
        """

        return self.graphs[self.selected_graph_index].setRegion(
            polygon_points=polygon_points, select_full_data=select_full_data
        )

    def current_graph_has_polygon(self) -> bool:
        return self.graphs[self.selected_graph_index].polygon is not None

    def update_point_probe(self, probe_name=DEFAULT_POINT_PROBE, xy_pos=None, state=None):
        if xy_pos is None and state is None:
            if probe_name not in self.point_probes:
                # nothing to do
                return
            # they didn't ask to change anything
            # but they may want to refresh stuff
            state, xy_pos = self.point_probes[probe_name]
        elif probe_name not in self.point_probes:
            # new point
            if xy_pos is None:
                raise ValueError("Point probe '{}' does not exist".format(probe_name))
            # if this is a new point probe, then it must be enabled
            state = True if state is None else state
        else:
            old_state, old_xy_pos = self.point_probes[probe_name]
            if xy_pos is None:
                # state is what is changing
                xy_pos = old_xy_pos
            elif state is None:
                # they are updating the position only
                # we have to turn the probe back on
                state = True

            if old_state == state and old_xy_pos == xy_pos:
                # nothing has changed so no need to tell anyone
                return
            if old_state != state:
                LOG.info("Changing point probe '{}' state to '{}'".format(probe_name, "on" if state else "off"))
            if old_xy_pos != xy_pos:
                LOG.info("Changing point probe '{}' position to '{}'".format(probe_name, xy_pos))

        self.point_probes[probe_name] = [state, xy_pos]
        self.pointProbeChanged.emit(probe_name, state, xy_pos)

    def _update_default_point_probe_graph(self):
        probe_name = DEFAULT_POINT_PROBE
        point_probe = self.point_probes.get(probe_name, [None, None])
        self._update_point_probe_graph(probe_name, *point_probe)

    def _update_point_probe_graph(self, probe_name, state, xy_pos):
        # need to set the point for all graphs because the point probe
        # is used across all plots
        for idx, graph in enumerate(self.graphs):
            rebuild = idx == self.selected_graph_index
            if state:
                graph.setPoint(xy_pos, rebuild=rebuild)
            elif state is not None:
                # if it is False/"off"
                graph.setPoint(None, rebuild=rebuild)

    def current_point_probe_status(self, probe_name):
        if probe_name not in self.point_probes:
            return False, None
        return self.point_probes[probe_name]

    def toggle_point_probe(self, probe_name, state=None):
        if probe_name not in self.point_probes:
            LOG.info("No point probe to toggle")
            return

        old_state = self.point_probes[probe_name][0]
        state = state if state is not None else not old_state
        self.update_point_probe(probe_name, state=state)

    def set_default_layer_selections(self, layer_uuids):
        """Set the UUIDs for the current graph if it doesn't have a polygon"""
        return self.graphs[self.selected_graph_index].set_default_layer_selections(layer_uuids)

    def on_region_probe_tool_selected(self):
        if len(self.graphs) > 0:
            return
        # There is no graph tab yet, we must create one
        self.set_up_tab(self.tab_widget_object.count(), do_increment_tab_letter=False)

        current_name = self.graphs[self.selected_graph_index].getName()
        self.didChangeTab.emit((current_name,))

    def _add_tab(self):
        LOG.info("Creating new area probe graph tab.")

        old_name = self.graphs[self.selected_graph_index].getName()
        self.set_up_tab(self.tab_widget_object.count())

        # notify everyone that we cloned a polygon (if we did)
        if self.graphs[self.selected_graph_index].polygon is not None:
            new_name = self.graphs[-1].getName()
            self.didClonePolygon.emit(old_name, new_name)

        current_name = self.graphs[self.selected_graph_index].getName()
        self.didChangeTab.emit((current_name,))

    def _handle_tab_change(self):
        """Deal with the fact that the tab changed in the tab widget"""

        new_tab_index = self.tab_widget_object.currentIndex()

        self.selected_graph_index = new_tab_index
        self.graphs[self.selected_graph_index].rebuildPlot()

        current_name = self.graphs[self.selected_graph_index].getName()
        self.didChangeTab.emit((current_name,))

    def _on_auto_update_checkbox_state_changed(self, state):
        if self.auto_update_checkbox.isChecked():
            self.update_button.setEnabled(False)
            self.layer_model.didFinishActivateProductDatasets.connect(self.handleActiveProductDatasetsChanged)
            self.pointProbeChanged.connect(self._update_point_probe_graph)
        else:
            self.layer_model.didFinishActivateProductDatasets.disconnect(self.handleActiveProductDatasetsChanged)
            self.pointProbeChanged.disconnect(self._update_point_probe_graph)
            self.update_button.setEnabled(True)

class ProbeGraphDisplay(object):
    """The ProbeGraphDisplay controls one tab of the Area Probe Graphs.
    The ProbeGraphDisplay handles generating a displaying a single graph.
    """

    # the most data we are willing to plot in a scatter plot
    # this limit was determined experimentally on Eva's laptop for glance, may need to revisit this
    MAX_SCATTER_PLOT_DATA = 1e7

    # the default number of bins for the histogram and density scatter plot
    DEFAULT_NUM_BINS = 100

    def __init__(self, manager, qt_parent, workspace, layer_model: LayerModel, queue, name_str):
        """build the graph tab controls
        :param layer_model:
        :return:
        """

        # hang on to our name
        self.myName = name_str

        # plotting related controls
        self.toolbar = None
        self.yCheckBox = None
        self.xDropDown = None
        self.yDropDown = None

        # internal objects to reference for info and data
        self.polygon = None
        self.point = None
        self.full_data_selection = False

        # save the workspace and queue for use later
        self.manager = manager
        self.workspace = workspace
        self.layer_model = layer_model
        self.queue = queue

        # internal values that control the behavior of plotting and controls
        self.xSelectedUUID = None
        self.ySelectedUUID = None

        self.xCurrentDatasetUUID = None
        self.yCurrentDatasetUUID = None

        self.uuidMap = None  # this is needed because the drop downs can't properly handle objects as ids
        self._stale = True  # whether or not the plot needs to be redrawn

        # a figure instance to plot on
        self.figure = Figure(figsize=(3, 3), dpi=72)
        # this is the Canvas Widget that displays the `figure`
        # it takes the `figure` instance as a parameter to __init__
        self.canvas = FigureCanvas(self.figure)
        self.canvas.setMinimumSize(100, 100)
        # make sure our figure is clear
        self.clearPlot()

        # make a matplotlib toolbar to attach to the graph
        self.toolbar = CustomNavigationToolbar(self.canvas, qt_parent)

        # create our selection controls

        # the label for the x selection
        xLabel = QtWidgets.QLabel("X layer:")

        # the check box that turns on and off comparing to a y layer
        self.yCheckBox = QtWidgets.QCheckBox("vs Y layer:")
        self.yCheckBox.setToolTip("Plot X layer data vs Y layer when this is checked.")
        self.yCheckBox.stateChanged.connect(self.vsChecked)

        # the drop down for selecting the x layer
        self.xDropDown = QtWidgets.QComboBox(qt_parent)
        self.xDropDown.setToolTip("The X layer data to use for plotting.")
        self.xDropDown.activated.connect(self.xSelected)

        # the drop down for selecting the y layer
        self.yDropDown = QtWidgets.QComboBox(qt_parent)
        self.yDropDown.setDisabled(True)
        self.yDropDown.setToolTip("The Y layer data to use for plotting.")
        self.yDropDown.activated.connect(self.ySelected)

        # set the layout
        # Note: add in a grid is (widget, row#, col#) or (widget, row#, col#, row_span, col_span)
        layout = QtWidgets.QGridLayout()
        layout.addWidget(self.toolbar, 1, 1, 1, 3)
        layout.addWidget(self.canvas, 2, 1, 1, 3)
        layout.addWidget(xLabel, 3, 1)
        layout.addWidget(self.xDropDown, 3, 2, 1, 2)
        layout.addWidget(self.yCheckBox, 4, 1)
        layout.addWidget(self.yDropDown, 4, 2, 1, 2)
        qt_parent.setLayout(layout)

    def set_possible_layers(self, do_rebuild_plot=False):
        """Given a list of layer UUIDs, set the names and UUIDs in the drop downs"""

        # make a uuid map because the mapping in a combo box doesn't work with objects
        self.uuidMap = {}

        # clear out the current lists
        self.xDropDown.clear()
        self.yDropDown.clear()

        # fill up our lists of layers
        for layer in self.layer_model.get_probeable_layers():
            uuid_string = str(layer.uuid)
            self.xDropDown.addItem(layer.descriptor, uuid_string)
            self.yDropDown.addItem(layer.descriptor, uuid_string)

            self.uuidMap[uuid_string] = layer.uuid

        # if possible, set the selections back to the way they were
        need_rebuild = False
        x_index = self.xDropDown.findData(str(self.xSelectedUUID))
        if x_index >= 0:
            # Selection didn't change
            self.xDropDown.setCurrentIndex(x_index)
        elif self.xDropDown.count() > 0:
            # Setting to a new layer
            need_rebuild = True
            self.xSelectedUUID = self.uuidMap[self.xDropDown.itemData(0)]
            self.xDropDown.setCurrentIndex(0)
        else:
            # we had something selected but now there is nothing new to select
            need_rebuild = need_rebuild or self.xSelectedUUID is not None
            self.xSelectedUUID = None
        y_index = self.yDropDown.findData(str(self.ySelectedUUID))
        if y_index >= 0:
            # Selection didn't change
            self.yDropDown.setCurrentIndex(y_index)
        elif self.yDropDown.count() > 0:
            # Setting to a new layer
            need_rebuild = need_rebuild or self.yCheckBox.isChecked()
            self.ySelectedUUID = self.uuidMap[self.yDropDown.itemData(0)]
            self.yDropDown.setCurrentIndex(0)
        else:
            # we had something selected but now there is nothing new to select
            need_rebuild = need_rebuild or self.ySelectedUUID is not None
            self.ySelectedUUID = None

        need_rebuild |= self._check_active_datasets_changed()

        # refresh the plot
        self._stale = need_rebuild
        if do_rebuild_plot:
            # Rebuild the plot (stale is used to determine if actual rebuild
            # is needed)
            self.rebuildPlot()

    def _check_active_datasets_changed(self):
        # check whether active datasets have changed. If so, update stored
        # dataset uuids and indicate that the graph needs to be rebuilt.
        need_rebuild = False
        x_layer = self.layer_model.get_layer_by_uuid(self.xSelectedUUID)
        x_active_product_dataset = None if not x_layer else x_layer.get_first_active_product_dataset()
        if not x_active_product_dataset:
            if self.xCurrentDatasetUUID is not None:
                need_rebuild = True
                self.xCurrentDatasetUUID = None
        elif x_active_product_dataset.uuid != self.xCurrentDatasetUUID:
            need_rebuild = True
            self.xCurrentDatasetUUID = x_active_product_dataset.uuid
        y_layer = self.layer_model.get_layer_by_uuid(self.ySelectedUUID)
        y_active_product_dataset = None if not y_layer else y_layer.get_first_active_product_dataset()
        if not y_active_product_dataset:
            if self.yCurrentDatasetUUID is not None:
                need_rebuild |= self.yCheckBox.isChecked()
                self.yCurrentDatasetUUID = None
        elif y_active_product_dataset.uuid != self.yCurrentDatasetUUID:
            need_rebuild |= self.yCheckBox.isChecked()
            self.yCurrentDatasetUUID = y_active_product_dataset.uuid
        return need_rebuild

    def set_default_layer_selections(self, layer_uuids):
        # only set the defaults if we don't have a polygon yet
        if self.polygon is not None:
            return

        if len(layer_uuids) >= 1:
            xIndex = self.xDropDown.findData(str(layer_uuids[0]))
            if xIndex >= 0:
                self.xDropDown.setCurrentIndex(xIndex)
                self.xSelectedUUID = layer_uuids[0]
            else:
                LOG.error("Tried to set probe graph to non-existent layer: %s", layer_uuids[0])

        if len(layer_uuids) >= 2:
            yIndex = self.xDropDown.findData(str(layer_uuids[1]))
            if yIndex >= 0:
                self.yDropDown.setCurrentIndex(yIndex)
                self.ySelectedUUID = layer_uuids[1]
            else:
                LOG.error("Tried to set probe graph to non-existent layer: %s", layer_uuids[1])

    @property
    def checked(self):
        return self.yCheckBox.isChecked()

    @checked.setter
    def checked(self, is_checked):
        return self.yCheckBox.setChecked(is_checked)

    def xSelected(self):
        """The user selected something in the X layer list."""
        oldXStr = str(self.xSelectedUUID)
        newXStr = self.xDropDown.itemData(self.xDropDown.currentIndex())
        self.xSelectedUUID = self.uuidMap[newXStr]

        # regenerate the plot
        if oldXStr != newXStr:
            self._stale = True
            self.rebuildPlot()

    def ySelected(self):
        """The user selected something in the Y layer list."""
        oldYStr = str(self.ySelectedUUID)
        newYStr = self.yDropDown.itemData(self.yDropDown.currentIndex())
        self.ySelectedUUID = self.uuidMap[newYStr]

        # regenerate the plot
        if (oldYStr != newYStr) and self.yCheckBox.isChecked():
            self._stale = True
            self.rebuildPlot()

    def vsChecked(self):
        """The vs check box was checked!"""
        # look at the state of the vs box and enable/disable the y drop down accordingly
        doPlotVS = self.yCheckBox.isChecked()
        self.yDropDown.setDisabled(not doPlotVS)

        # regenerate the plot
        self._stale = True
        self.rebuildPlot()

    def setRegion(self, polygon_points=None, select_full_data=False):
        """Set the region for this graph as polygon selection or full data."""

        assert polygon_points is None or not select_full_data, (  # nosec B101
            "Must not give both 'polygon_points' and True for 'select_full_data':"
            " Defining region by polygon and as full data are mutually exclusive."
        )

        # Even with assertions switched off we will get a valid state here: a
        # polygonal region wins over the full data selection: the first one will
        # have a visual echo, the second one not, so this is more likely to give
        # a consistent state.
        self.polygon = polygon_points
        self.full_data_selection = False if self.polygon is not None else select_full_data

        # regenerate the plot
        self._stale = True
        self.rebuildPlot()

        # return our name to be used for the polygon name
        return self.myName

    def setPoint(self, coordinates, rebuild=True):
        self.point = coordinates
        self._stale = True
        # sometimes we set the point to be redrawn later
        if rebuild:
            self.rebuildPlot()

    def getName(self):
        """Accessor method for the graph's name."""
        return self.myName

    def rebuildPlot(self):
        """Given what we currently know about the selection area and selected bands, rebuild our plot

        Note: This should be called only when the selections change in some way.
        """
        if not self._stale:
            LOG.debug("Plot doesn't need to be rebuilt")
            return

        # should be plotting vs Y?
        doPlotVS = self.yCheckBox.isChecked()
        task_name = "%s_%s_region_plotting" % (self.xSelectedUUID, self.ySelectedUUID)
        task_description = (
            "Creating plot for full data" if self.full_data_selection else "Creating plot for region probe data"
        )
        self.queue.add(
            task_name,
            self._rebuild_plot_task(
                self.xSelectedUUID,
                self.ySelectedUUID,
                self.polygon,
                self.point,
                plot_versus=doPlotVS,
                plot_full_data=self.full_data_selection,
            ),
            task_description,
            interactive=True,
        )
        # Assume that the task gets resolved otherwise we might try to draw multiple times
        self._stale = False

    def _rebuild_plot_task(  # noqa: C901
        self, x_layer_uuid, y_layer_uuid, polygon, point_xy, plot_versus=False, plot_full_data=True
    ):
        data_source_description = "full data" if plot_full_data else "polygon data"

        x_layer = self.layer_model.get_layer_by_uuid(x_layer_uuid)
        x_active_product_dataset = None if not x_layer else x_layer.get_first_active_product_dataset()
        x_uuid = None if not x_active_product_dataset else x_active_product_dataset.uuid

        y_layer = self.layer_model.get_layer_by_uuid(y_layer_uuid)
        y_active_product_dataset = None if not y_layer else y_layer.get_first_active_product_dataset()
        y_uuid = None if not y_active_product_dataset else y_active_product_dataset.uuid

        # if we are plotting only x and we have a selected x and a polygon
        have_x_layer = x_layer_uuid is not None
        have_y_layer = y_layer_uuid is not None
        should_plot = polygon is not None or plot_full_data
        if not plot_versus and have_x_layer and should_plot:
            yield {TASK_DOING: f"Probe Plot: Collecting {data_source_description}...", TASK_PROGRESS: 0.0}

            # get the data and info we need for this plot
            if x_active_product_dataset:
                if plot_full_data:
                    data_polygon = self.workspace.get_content(x_active_product_dataset.uuid)
                else:
                    data_polygon = self.workspace.get_content_polygon(x_active_product_dataset.uuid, polygon)
            else:
                data_polygon = np.array([])

            x_conv_func = x_layer.info[Info.UNIT_CONVERSION][1]
            data_polygon = x_conv_func(data_polygon)
            time = x_active_product_dataset.info[Info.DISPLAY_TIME]
            title = f"{time}"
            x_axis_label = x_layer.descriptor
            y_axis_label = "Count of data points"

            # get point probe value
            if x_active_product_dataset and point_xy:
                x_point = self.workspace.get_content_point(x_active_product_dataset.uuid, point_xy)
                x_point = x_conv_func(x_point)
            else:
                x_point = None

            # plot a histogram
            yield {TASK_DOING: "Probe Plot: Creating histogram plot", TASK_PROGRESS: 0.25}
            self.plotHistogram(data_polygon, title, x_point, x_axis_label, y_axis_label)

        # if we are plotting x vs y and have x, y, and a polygon
        elif plot_versus and have_x_layer and have_y_layer and should_plot:
            yield {TASK_DOING: f"Probe Plot: Collecting {data_source_description} (layer 1)...", TASK_PROGRESS: 0.0}

            name1 = x_layer.descriptor
            name2 = y_layer.descriptor
            if not x_active_product_dataset or not y_active_product_dataset:
                x_point = None
                y_point = None
                time1 = None
                time2 = None
                data1 = np.array([0])
                data2 = np.array([0])
            else:
                # get the data and info we need for this plot
                x_info = x_active_product_dataset.info
                y_info = y_active_product_dataset.info
                time1 = x_info[Info.DISPLAY_TIME]
                time2 = y_info[Info.DISPLAY_TIME]
                hires_uuid = self.workspace.lowest_resolution_uuid(x_uuid, y_uuid)
                # hires_coord_mask are the lat/lon coordinates of each of the
                # pixels in hires_data. The coordinates are (lat, lon) to resemble
                # the (Y, X) indexing of numpy arrays
                if plot_full_data:
                    hires_coord_mask = None
                    hires_data = self.workspace.get_content(hires_uuid)
                else:
                    hires_coord_mask, hires_data = self.workspace.get_coordinate_mask_polygon(hires_uuid, polygon)

                x_conv_func = x_layer.info[Info.UNIT_CONVERSION][1]
                y_conv_func = y_layer.info[Info.UNIT_CONVERSION][1]
                yield {
                    TASK_DOING: f"Probe Plot: Collecting {data_source_description} (layer 2)...",
                    TASK_PROGRESS: 0.15,
                }
                if hires_uuid == x_uuid:
                    # the hires data was from the X UUID
                    data1 = x_conv_func(hires_data)
                    if plot_full_data:
                        data2 = self.workspace.get_content(y_uuid)
                    else:
                        data2 = self.workspace.get_content_coordinate_mask(y_uuid, hires_coord_mask)
                    data2 = y_conv_func(data2)
                else:
                    # the hires data was from the Y UUID
                    data2 = y_conv_func(hires_data)
                    if plot_full_data:
                        data1 = self.workspace.get_content(x_uuid)
                    else:
                        data1 = self.workspace.get_content_coordinate_mask(x_uuid, hires_coord_mask)
                    data1 = x_conv_func(data1)
                yield {TASK_DOING: "Probe Plot: Creating scatter plot...", TASK_PROGRESS: 0.25}

                if point_xy:
                    x_point = self.workspace.get_content_point(x_uuid, point_xy)
                    x_point = x_conv_func(x_point)
                    y_point = self.workspace.get_content_point(y_uuid, point_xy)
                    y_point = y_conv_func(y_point)
                else:
                    x_point = None
                    y_point = None

                # plot a scatter plot
                good_mask = ~(np.isnan(data1) | np.isnan(data2))
                data1 = data1[good_mask]
                data2 = data2[good_mask]

            self.plotDensityScatterplot(data1, name1, time1, data2, name2, time2, x_point, y_point)

        # if we have some combination of selections we don't understand, clear the figure
        else:
            yield {TASK_DOING: "Probe Plot: Clearing plot figure...", TASK_PROGRESS: 0.0}
            self.clearPlot()

        yield {TASK_DOING: "Probe Plot: Drawing plot...", TASK_PROGRESS: 0.95}
        self.manager.drawChildGraph.emit(self.myName)
        yield {TASK_DOING: "Probe Plot: Done", TASK_PROGRESS: 1.0}

    def _draw(self):
        self.canvas.draw()

    def plotHistogram(self, data, title, x_point, x_label, y_label, numBins=100):
        """Make a histogram using the given data and label it with the given title"""
        self.figure.clf()
        axes = self.figure.add_subplot(111)
        bars = axes.hist(data[~np.isnan(data)], bins=self.DEFAULT_NUM_BINS)
        if x_point is not None:
            # go through each rectangle object and make the one that contains x_point 'red'
            # default color is blue so red should stand out
            for bar in bars[2][::-1]:
                if bar.xy[0] <= x_point:
                    bar.set_color("red")
                    break
        axes.set_title(title)
        axes.set_xlabel(x_label)
        axes.set_ylabel(y_label)

    def plotDensityScatterplot(self, dataX, nameX, timeX, dataY, nameY, timeY, pointX, pointY):
        """Make a density scatter plot for the given data
        :param timeX:
        :param timeY:
        """

        # clear the figure and make a new subplot
        self.figure.clf()
        axes = self.figure.add_subplot(111)

        # figure out the range of the data
        # you might not be comparing the same units
        xmin_value = np.min(dataX)
        xmax_value = np.max(dataX)
        ymin_value = np.min(dataY)
        ymax_value = np.max(dataY)
        # bounds should be defined in the form [[xmin, xmax], [ymin, ymax]]
        bounds = [[xmin_value, xmax_value], [ymin_value, ymax_value]]

        # make the binned density map for this data set
        density_map, _, _ = np.histogram2d(dataX, dataY, bins=self.DEFAULT_NUM_BINS, range=bounds)
        # mask out zero counts; flip because y goes the opposite direction in an imshow graph
        density_map = np.flipud(np.transpose(np.ma.masked_array(density_map, mask=density_map == 0)))

        # display the density map data
        img = axes.imshow(
            density_map,
            extent=[xmin_value, xmax_value, ymin_value, ymax_value],
            aspect="auto",
            interpolation="nearest",
            norm=LogNorm(),
        )
        if pointX is not None:
            axes.set_autoscale_on(False)
            axes.plot(
                pointX,
                pointY,
                marker="o",
                markerfacecolor="white",
                markeredgecolor="black",
                markersize=10,
                markeredgewidth=1.0,
            )
            axes.set_autoscale_on(True)

        colorbar = self.figure.colorbar(img)
        colorbar.set_label("log(count of data points)")

        # set the various text labels
        axes.set_xlabel(f"{nameX}")
        axes.set_ylabel(f"{nameY}")
        axes.set_title(timeX)

        # draw the x vs y line
        self._draw_xy_line(axes)

    def clearPlot(self):
        """Clear our plot"""
        self.full_data_selection = False
        self.figure.clf()

    def _draw_xy_line(self, axes):
        # get the bounds for our calculations and so we can reset the viewing window later
        x_bounds = axes.get_xbound()
        y_bounds = axes.get_ybound()

        # draw the x=y line
        perfect = [max(x_bounds[0], y_bounds[0]), min(x_bounds[1], y_bounds[1])]
        axes.plot(perfect, perfect, "--", color="k", label="X = Y")

        # reset the bounds
        axes.set_xbound(x_bounds)
        axes.set_ybound(y_bounds)


from sqlalchemy import (
    Boolean,
    Column,
    ForeignKey,
    Integer,
    String,
    Date,
    DateTime,
    UniqueConstraint,
)
from sqlalchemy.orm import relationship, backref
from sqlalchemy.sql import func

from app.db.base_class import Base

class Saving(Base):
    id = Column(Integer, primary_key=True, index=True)
    month = Column(String, nullable=False)
    year = Column(Integer, nullable=False)
    amount = Column(Integer, nullable=False)
    date_sent = Column(Date, nullable=False)
    is_late = Column(Boolean(), default=False)
    fine_id = Column(Integer, ForeignKey("fine.id"), nullable=True)
    member_id = Column(Integer, ForeignKey("user.id"), nullable=False)
    created_at = Column(DateTime(timezone=True), server_default=func.now())
    updated_at = Column(DateTime(timezone=True), onupdate=func.now())

    fine = relationship("Fine", backref=backref("saving", uselist=False))

    __table_args__ = (
        UniqueConstraint("month", "year", "member_id", name="month_year_user_unique"),
    )


import common.parse_util as cpu
from common.common_based import CommonBased

class AdaType(CommonBased):
    RECORD_TYPE = "Record"
    ENUM_TYPE = "Enum"
    STR_TYPE = "String"
    ARRAY_TYPE = "Array"
    DERIVED_TYPE = "Derived"
    SUBTYPE = "Subtype"
    INT_TYPE = "Integer"
    REAL_TYPE = "Real"
    FIELD_TYPE = "Field"
    VAR_TYPE = "Var"
    ENUM_ITEM_TYPE = 'Enum_Item'
    ARRAY_INDEX_TYPE = 'Array_Index'
    def __init__(self, name, ttype, package=None, ctx=None, is_based=False):
        super(AdaType, self).__init__()
        self.name = name.upper() if name else None
        self.package = package.upper() if package else None
        self.ttype = ttype
        self.is_based = is_based
        self.discriminant = {}
        self.size = None
        self.ctx = ctx
        #self.ctx.cur_type = self
        self.size_solved = False
        self.constraint_solved = False
        self.type_chain = []
        self.first = None
        self.last = None
        self.must_print = ['first', 'last', 'size']
        self.to_print = ['discriminant', 'size']
        self.leader_str = "'%s type [%s] in [%s]:' % (self.ttype, self.name, self.package)"
        self.to_solve = []
        self.to_check = []

    def add_discrim(self, fs):
        for f in fs:
            self.discriminant[f.name] = f

    def solve_type_chain(self):
        pass

    def full_name(self):
        return '.'.join([self.package, self.name])

    def solve_a_type_or_expr(self, attr_n, i_expr=None, solved_n=None, mandatory=True, is_type=True):
        if not solved_n:
            solved_n = "%s_solved" % attr_n
        if not getattr(self, solved_n):
            if not mandatory and not i_expr and not getattr(self, attr_n):
                setattr(self, solved_n, True)
            else:
                if i_expr:
                    setattr(self, attr_n, i_expr)
                if is_type:
                    attr_v, solved_v = cpu.solve_type(self.ctx, getattr(self, attr_n))
                else:
                    attr_v, solved_v = cpu.solve_expr(self.ctx, getattr(self, attr_n))
                setattr(self, attr_n, attr_v)
                setattr(self, solved_n, solved_v)

    def solve_size(self, i_size=None):
        self.solve_a_type_or_expr('size', i_size)

    def solve_constraint(self, const=None):
        if const is not None:
            setattr(self, 'constraint', const)
        const = getattr(self, 'constraint', None)
        if self.constraint_solved or not const:
            self.constraint_solved = True
            return
        if const['type'] == 'range':
            if const['range']['type'] == 'range':
                self.first, solved1 = cpu.solve_expr(self.ctx, const['range']['first'])
                self.last, solved2 = cpu.solve_expr(self.ctx, const['range']['last'])
                self.constraint_solved = solved1 and solved2
            elif const['range']['type'] == 'attr':
                solved = True
                if isinstance(const['range']['base'], str):
                    const['range']['base'], solved = cpu.solve_type(self.ctx, const['range']['base'])
                if solved:
                    self.first, solved1 = cpu.solve_expr(self.ctx, const['range']['base'].first)
                    self.last, solved2 = cpu.solve_expr(self.ctx, const['range']['base'].last)
                    self.constraint_solved = solved1 and solved2
        #if self.constraint_solved:
        #    setattr(self, 'constraint', None)

    def check_solved(self):
        for attr in self.to_solve:
            method_n = 'solve_%s' % attr
            solved_n = '%s_solved' % attr
            if not getattr(self, solved_n):
                getattr(self, method_n)()
        ret = True
        if not self.to_check:
            self.to_check = self.to_solve
        for attr in self.to_check:
            solved_n = '%s_solved' % attr
            ret = ret and getattr(self, solved_n)
        if ret:
            self.is_based = True
        return  ret


"""
Create a function where will be needed as an argument string_a and number_b. With using assert - be sure, that user
input arguments of this types. After it concatenate string_a with number_b
"""

def special_function(string_a, number_b):
    assert type(string_a) == str
    assert type(number_b) == int

special_function('dddd', 4)

print(isinstance(True, int))
print(type(True) == int)


#!/usr/bin/env python

from __future__ import division, print_function

import argparse
import glob
from array import array

import numpy as np
import scipy as sp

import fitsio
from picca import constants
from picca.data import delta
from picca.Pk1D import (compute_cor_reso, compute_Pk_noise, compute_Pk_raw,
                        fill_masked_pixels, rebin_diff_noise, split_forest)
from picca.utils import print

def make_tree(tree,nb_bin_max):

    zqso = array( 'f', [ 0. ] )
    mean_z = array( 'f', [ 0. ] )
    mean_reso = array( 'f', [ 0. ] )
    mean_SNR = array( 'f', [ 0. ] )
    nb_mask_pix = array( 'f', [ 0. ] )

    lambda_min = array( 'f', [ 0. ] )
    lambda_max= array( 'f', [ 0. ] )

    plate = array( 'i', [ 0 ] )
    mjd = array( 'i', [ 0 ] )
    fiber = array( 'i', [ 0 ] )

    nb_r = array( 'i', [ 0 ] )
    k_r = array( 'f', nb_bin_max*[ 0. ] )
    Pk_r = array( 'f', nb_bin_max*[ 0. ] )
    Pk_raw_r = array( 'f', nb_bin_max*[ 0. ] )
    Pk_noise_r = array( 'f', nb_bin_max*[ 0. ] )
    Pk_diff_r = array( 'f', nb_bin_max*[ 0. ] )
    cor_reso_r = array( 'f', nb_bin_max*[ 0. ] )

    tree.Branch("zqso",zqso,"zqso/F")
    tree.Branch("mean_z",mean_z,"mean_z/F")
    tree.Branch("mean_reso",mean_reso,"mean_reso/F")
    tree.Branch("mean_SNR",mean_SNR,"mean_SNR/F")
    tree.Branch("lambda_min",lambda_min,"lambda_min/F")
    tree.Branch("lambda_max",lambda_max,"lambda_max/F")
    tree.Branch("nb_masked_pixel",nb_mask_pix,"nb_mask_pixel/F")

    tree.Branch("plate",plate,"plate/I")
    tree.Branch("mjd",mjd,"mjd/I")
    tree.Branch("fiber",fiber,"fiber/I")

    tree.Branch( 'NbBin', nb_r, 'NbBin/I' )
    tree.Branch( 'k', k_r, 'k[NbBin]/F' )
    tree.Branch( 'Pk_raw', Pk_raw_r, 'Pk_raw[NbBin]/F' )
    tree.Branch( 'Pk_noise', Pk_noise_r, 'Pk_noise[NbBin]/F' )
    tree.Branch( 'Pk_diff', Pk_diff_r, 'Pk_diff[NbBin]/F' )
    tree.Branch( 'cor_reso', cor_reso_r, 'cor_reso[NbBin]/F' )
    tree.Branch( 'Pk', Pk_r, 'Pk[NbBin]/F' )

    return zqso,mean_z,mean_reso,mean_SNR,lambda_min,lambda_max,plate,mjd,fiber,\
    nb_mask_pix,nb_r,k_r,Pk_r,Pk_raw_r,Pk_noise_r,cor_reso_r,Pk_diff_r

def compute_mean_delta(ll,delta,iv,zqso):

    for i, _ in enumerate (ll):
        ll_obs = np.power(10., ll[i])
        ll_rf = ll_obs/(1.+zqso)
        hdelta.Fill(ll_obs, ll_rf, delta[i])
        hdelta_RF.Fill(ll_rf, delta[i])
        hdelta_OBS.Fill(ll_obs, delta[i])
        hivar.Fill(iv[i])
        snr_pixel = (delta[i]+1)*np.sqrt(iv[i])
        hsnr.Fill(snr_pixel)
        hivar.Fill(iv[i])
        if (iv[i] < 1000):
            hdelta_RF_we.Fill(ll_rf, delta[i], iv[i])
            hdelta_OBS_we.Fill(ll_obs, delta[i], iv[i])

    return

if __name__ == '__main__':

    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter,
        description='Compute the 1D power spectrum')

    parser.add_argument('--out-dir', type=str, default=None, required=True,
        help='Output directory')

    parser.add_argument('--out-format', type=str, default='fits', required=False,
        help='Output format: root or fits (if root call PyRoot)')

    parser.add_argument('--in-dir', type=str, default=None, required=True,
        help='Directory to delta files')

    parser.add_argument('--in-format', type=str, default='fits', required=False,
        help=' Input format used for input files: ascii or fits')

    parser.add_argument('--SNR-min',type=float,default=2.,required=False,
        help='Minimal mean SNR per pixel ')

    parser.add_argument('--reso-max',type=float,default=85.,required=False,
        help='Maximal resolution in km/s ')

    parser.add_argument('--lambda-obs-min',type=float,default=3600.,required=False,
        help='Lower limit on observed wavelength [Angstrom]' )

    parser.add_argument('--nb-part',type=int,default=3,required=False,
        help='Number of parts in forest')

    parser.add_argument('--nb-pixel-min',type=int,default=75,required=False,
        help='Minimal number of pixels in a part of forest')

    parser.add_argument('--nb-pixel-masked-max',type=int,default=40,required=False,
        help='Maximal number of masked pixels in a part of forest')

    parser.add_argument('--no-apply-filling', action='store_true', default=False, required=False,
        help='Dont fill masked pixels')

    parser.add_argument('--noise-estimate', type=str, default='mean_diff', required=False,
        help='Estimate of Pk_noise pipeline/diff/mean_diff/rebin_diff/mean_rebin_diff')

    parser.add_argument('--forest-type', type=str, default='Lya', required=False,
        help='Forest used: Lya, SiIV, CIV')

    parser.add_argument('--debug', action='store_true', default=False, required=False,
        help='Fill root histograms for debugging')

    args = parser.parse_args()

#   Create root file
    if (args.out_format=='root') :
        from ROOT import TH1D, TFile, TTree, TProfile2D, TProfile
        storeFile = TFile(args.out_dir+"/Testpicca.root","RECREATE","PK 1D studies studies");
        nb_bin_max = 700
        tree = TTree("Pk1D","SDSS 1D Power spectrum Ly-a");
        zqso,mean_z,mean_reso,mean_SNR,lambda_min,lambda_max,plate,mjd,fiber,\
        nb_mask_pix,nb_r,k_r,Pk_r,Pk_raw_r,Pk_noise_r,cor_reso_r,Pk_diff_r = make_tree(tree,nb_bin_max)

        # control histograms
        if (args.forest_type=='Lya'):
            forest_inf=1040.
            forest_sup=1200.
        elif (args.forest_type=='SiIV'):
            forest_inf=1270.
            forest_sup=1380.
        elif (args.forest_type=='CIV'):
            forest_inf=1410.
            forest_sup=1520.
        hdelta  = TProfile2D( 'hdelta', 'delta mean as a function of lambda-lambdaRF', 36, 3600., 7200., 16, forest_inf, forest_sup, -5.0, 5.0)
        hdelta_RF  = TProfile( 'hdelta_RF', 'delta mean as a function of lambdaRF', 320, forest_inf, forest_sup, -5.0, 5.0)
        hdelta_OBS  = TProfile( 'hdelta_OBS', 'delta mean as a function of lambdaOBS', 1800, 3600., 7200., -5.0, 5.0)
        hdelta_RF_we  = TProfile( 'hdelta_RF_we', 'delta mean weighted as a function of lambdaRF', 320, forest_inf, forest_sup, -5.0, 5.0)
        hdelta_OBS_we  = TProfile( 'hdelta_OBS_we', 'delta mean weighted as a function of lambdaOBS', 1800, 3600., 7200., -5.0, 5.0)
        hivar = TH1D('hivar','  ivar ',10000,0.0,10000.)
        hsnr = TH1D('hsnr','  snr per pixel ',100,0.0,100.)
        hdelta_RF_we.Sumw2()
        hdelta_OBS_we.Sumw2()

    # Read deltas
    if (args.in_format=='fits') :
        fi = glob.glob(args.in_dir+"/*.fits.gz")
    elif (args.in_format=='ascii') :
        fi = glob.glob(args.in_dir+"/*.txt")

    data = {}
    ndata = 0

    # initialize randoms
    sp.random.seed(4)

    # loop over input files
    for i,f in enumerate(fi):
        if i%1==0:
            print("\rread {} of {} {}".format(i,len(fi),ndata),end="")

        # read fits or ascii file
        if (args.in_format=='fits') :
            hdus = fitsio.FITS(f)
            dels = [delta.from_fitsio(h,Pk1D_type=True) for h in hdus[1:]]
        elif (args.in_format=='ascii') :
            ascii_file = open(f,'r')
            dels = [delta.from_ascii(line) for line in ascii_file]

        ndata+=len(dels)
        print ("\n ndata =  ",ndata)
        out = None

        # loop over deltas
        for d in dels:

            # Selection over the SNR and the resolution
            if (d.mean_SNR<=args.SNR_min or d.mean_reso>=args.reso_max) : continue

            # first pixel in forest
            for first_pixel,first_pixel_ll in enumerate(d.ll):
                if 10.**first_pixel_ll>args.lambda_obs_min : break

            # minimum number of pixel in forest
            nb_pixel_min = args.nb_pixel_min
            if ((len(d.ll)-first_pixel)<nb_pixel_min) : continue

            # Split in n parts the forest
            nb_part_max = (len(d.ll)-first_pixel)//nb_pixel_min
            nb_part = min(args.nb_part,nb_part_max)
            m_z_arr,ll_arr,de_arr,diff_arr,iv_arr = split_forest(nb_part,d.dll,d.ll,d.de,d.diff,d.iv,first_pixel)
            for f in range(nb_part):

                # rebin diff spectrum
                if (args.noise_estimate=='rebin_diff' or args.noise_estimate=='mean_rebin_diff'):
                    diff_arr[f]=rebin_diff_noise(d.dll,ll_arr[f],diff_arr[f])

                # Fill masked pixels with 0.
                ll_new,delta_new,diff_new,iv_new,nb_masked_pixel = fill_masked_pixels(d.dll,ll_arr[f],de_arr[f],diff_arr[f],iv_arr[f],args.no_apply_filling)
                if (nb_masked_pixel> args.nb_pixel_masked_max) : continue
                if (args.out_format=='root' and  args.debug): compute_mean_delta(ll_new,delta_new,iv_new,d.zqso)

                lam_lya = constants.absorber_IGM["LYA"]
                z_abs =  np.power(10.,ll_new)/lam_lya - 1.0
                mean_z_new = sum(z_abs)/float(len(z_abs))

                # Compute Pk_raw
                k,Pk_raw = compute_Pk_raw(d.dll,delta_new,ll_new)

                # Compute Pk_noise
                run_noise = False
                if (args.noise_estimate=='pipeline'): run_noise=True
                Pk_noise,Pk_diff = compute_Pk_noise(d.dll,iv_new,diff_new,ll_new,run_noise)

                # Compute resolution correction
                delta_pixel = d.dll*np.log(10.)*constants.speed_light/1000.
                cor_reso = compute_cor_reso(delta_pixel,d.mean_reso,k)

                # Compute 1D Pk
                if (args.noise_estimate=='pipeline'):
                    Pk = (Pk_raw - Pk_noise)/cor_reso
                elif (args.noise_estimate=='diff' or args.noise_estimate=='rebin_diff'):
                    Pk = (Pk_raw - Pk_diff)/cor_reso
                elif (args.noise_estimate=='mean_diff' or args.noise_estimate=='mean_rebin_diff'):
                    selection = (k>0) & (k<0.02)
                    if (args.noise_estimate=='mean_rebin_diff'):
                        selection = (k>0.003) & (k<0.02)
                    Pk_mean_diff = sum(Pk_diff[selection])/float(len(Pk_diff[selection]))
                    Pk = (Pk_raw - Pk_mean_diff)/cor_reso

                # save in root format
                if (args.out_format=='root'):
                    zqso[0] = d.zqso
                    mean_z[0] = m_z_arr[f]
                    mean_reso[0] = d.mean_reso
                    mean_SNR[0] = d.mean_SNR
                    lambda_min[0] =  np.power(10.,ll_new[0])
                    lambda_max[0] =  np.power(10.,ll_new[-1])
                    nb_mask_pix[0] = nb_masked_pixel

                    plate[0] = d.plate
                    mjd[0] = d.mjd
                    fiber[0] = d.fid

                    nb_r[0] = min(len(k),nb_bin_max)
                    for i in range(nb_r[0]) :
                        k_r[i] = k[i]
                        Pk_raw_r[i] = Pk_raw[i]
                        Pk_noise_r[i] = Pk_noise[i]
                        Pk_diff_r[i] = Pk_diff[i]
                        Pk_r[i] = Pk[i]
                        cor_reso_r[i] = cor_reso[i]

                    tree.Fill()

                # save in fits format

                if (args.out_format=='fits'):
                    hd = [ {'name':'RA','value':d.ra,'comment':"QSO's Right Ascension [degrees]"},
                        {'name':'DEC','value':d.dec,'comment':"QSO's Declination [degrees]"},
                        {'name':'Z','value':d.zqso,'comment':"QSO's redshift"},
                        {'name':'MEANZ','value':m_z_arr[f],'comment':"Absorbers mean redshift"},
                        {'name':'MEANRESO','value':d.mean_reso,'comment':'Mean resolution [km/s]'},
                        {'name':'MEANSNR','value':d.mean_SNR,'comment':'Mean signal to noise ratio'},
                        {'name':'NBMASKPIX','value':nb_masked_pixel,'comment':'Number of masked pixels in the section'},
                        {'name':'PLATE','value':d.plate,'comment':"Spectrum's plate id"},
                        {'name':'MJD','value':d.mjd,'comment':'Modified Julian Date,date the spectrum was taken'},
                        {'name':'FIBER','value':d.fid,'comment':"Spectrum's fiber number"}
                    ]

                    cols=[k,Pk_raw,Pk_noise,Pk_diff,cor_reso,Pk]
                    names=['k','Pk_raw','Pk_noise','Pk_diff','cor_reso','Pk']
                    comments=['Wavenumber', 'Raw power spectrum', "Noise's power spectrum", 'Noise coadd difference power spectrum',\
                              'Correction resolution function', 'Corrected power spectrum (resolution and noise)']
                    units=['(km/s)^-1', 'km/s', 'km/s', 'km/s', 'km/s', 'km/s']

                    try:
                        out.write(cols,names=names,header=hd,comments=comments,units=units)
                    except AttributeError:
                        out = fitsio.FITS(args.out_dir+'/Pk1D-'+str(i)+'.fits.gz','rw',clobber=True)
                        out.write(cols,names=names,header=hd,comment=comments,units=units)
        if (args.out_format=='fits' and out is not None):
            out.close()

# Store root file results
    if (args.out_format=='root'):
         storeFile.Write()

    print ("all done ")


from typing import Dict, List

from twin_runtime.twin_runtime_core import TwinRuntime
from twin_runtime.twin_runtime_core import LogLevel

class TwinBuilderSimulator():
    def __init__(self, twin_model_file, state_variable_names: List,
                 action_variable_names: List,
                 number_of_warm_up_steps, warm_up_action_variable_values: List):
        self.state = {}
        self.twin_runtime = None #assigned in reset
        self.done = False
        self.twin_model_file = twin_model_file
        self.state_variable_names = state_variable_names
        self.action_variable_names = action_variable_names
        self.number_of_warm_up_steps = number_of_warm_up_steps
        self.warm_up_action_variable_values = warm_up_action_variable_values
        self.step_size = 0.5
        self.time_index = 0
        self.reset(self.step_size)

    def reset(self, step_size):
        self.done = False
        runtime_log = self.twin_model_file.replace('.twin', '.log')
        self.step_size = step_size

        if self.twin_runtime != None:
            self.twin_runtime.twin_close()
            self.twin_runtime = None

        # Load Twin, set the parameters values, initialize (and generate snapshots, output)
        self.twin_runtime = TwinRuntime(self.twin_model_file, runtime_log, log_level=LogLevel.TWIN_LOG_ALL)
        self.twin_runtime.twin_instantiate()
        self.twin_runtime.twin_initialize()

        for state_variable_name in self.state_variable_names:
            self.state[state_variable_name] = 0

        self.time_index = 0
        self.state['time_index'] = self.time_index

        # Run initial steps to "warm up" the simulation
        if self.number_of_warm_up_steps > 0:
            action = dict(zip(self.action_variable_names, self.warm_up_action_variable_values))
            for i in range(self.number_of_warm_up_steps):
                self.episode_step(action)

    def get_state(self) -> Dict[str, float]:
        """Called to retreive the current state of the simulator. """
        print(f"returning state: {self.state}")
        return self.state

    def halted(self) -> bool:
        """
        Should return True if the simulator cannot continue for some reason
        """
        return self.done

    def episode_start(self, config: Dict = None) -> None:
        """ Called at the start of each episode """
        self.reset(config["step_size"] or 0.5)

    def episode_step(self, action: Dict):
        """ Called for each step of the episode """
        for f in action.keys():
            self.twin_runtime.twin_set_input_by_name(f, action[f])

        self.twin_runtime.twin_simulate(self.time_index)

        for state_variable_name in self.state_variable_names:
            value = self.twin_runtime.twin_get_output_by_name(state_variable_name).value
            print(value)
            self.state[state_variable_name] = value

        self.state['time_index'] = self.time_index
        #increase the index
        self.time_index = self.time_index + self.step_size

    def episode_finish(self):
        """ Called at the end of each episode """
        self.twin_runtime.twin_close()
        self.twin_runtime = None


## Loading the required libraries

import os
import re
import pandas as pd

## Must have a text file folder
## Setting the directory which hold our pdf files that were converted to .txt files

leaflets_dir = 'dump_profissional/'

## Setting the keywords we want to check on the drug leaflets we scraped from the brazilian drug agency ANVISA,
## Our maing goal is to find drug related information that could help us creating clinical decision support API's.

keywords = [
            'insuficiência hepática',
            'Insuficiência hepática',
            'disfunção',
            'disfunções',
            'hepática',
            'reajuste de dose',
            'ajuste de dose',
            'ajuste posológico',
            'doença hepática',
            'Insuficiência renal',
            'insuficiência renal',
            'renal',
            'doença renal'

]
# Creating a 'for loop' to iterate through the text files
rows = []
for root, dirs, files in os.walk(leaflets_dir):
    for f in files:
        keywords_in_leaflet = []
        keywords_not_in_leaflet = []

        leaflet_num = f.rpartition('_bula')[0].partition('bula')[-1]
        leaflet_filename = os.path.join(root, f)

        with open(leaflet_filename,'r',encoding='utf-8',errors='ignore') as l:
            leaflet = l.read()
            for k in keywords:
                if k in leaflet:
                    keywords_in_leaflet.append(k)
                else:
                    keywords_not_in_leaflet.append(k)

            rows.append([
                        leaflet_num,
                        ', '.join(sorted(keywords_in_leaflet)),
                        ', '.join(sorted(keywords_not_in_leaflet))
            ])

# Wrangling the data into a pandas dataframe

df = pd.DataFrame(
    data = rows,
    columns = [
               'leaflet_id',
               'keywords_in_leaflet',
               'keywords_not_in_leaflet'
    ]
)

# Testing unique values for keywords in leaflet
# df['keywords_in_leaflet'].unique()

# Writing the dataframe into an excel file.
df.to_excel('leaflets_keywords_hep_ren.xlsx', index = False)


"""
  Capstone Project.  Code written by Evan Cochrane.
  Fall term, 2018-2019.
"""

import rosebotics_new as rb
import time

def main():
    """ Runs YOUR specific part of the project. Uncomment tests as needed. """

    print('why not?')

    # test_go_straight_inches_method()
    test_spin_in_place_degrees_method()
    # test_turn_degrees_method()
    # test_move_in_polygon_function()
    # infrared_tester(9, 15)

def move_in_polygon(robot, n, length_of_each_side=5, speed=100):
    """ Causes the given rb.Robot object to move in a regular n-gon pattern, with optional length and speed options.

    :type robot: rb.Snatch3rRobot
    :type n: int
    :type length_of_each_side: float
    :type speed: int
    """
    for _ in range(n):
        robot.drive_system.go_straight_inches(length_of_each_side, speed)
        robot.drive_system.spin_in_place_degrees((360 / n), speed)

def test_go_straight_inches_method():
    """ Tests the go_straight_inches method """
    my_robot = rb.Snatch3rRobot()

    print('Test 1: Forward 5 Inches')
    my_robot.drive_system.go_straight_inches(5, 50)

    time.sleep(10)

    print('Test 2: Backward 5 Inches')
    my_robot.drive_system.go_straight_inches(-5, -50)

    time.sleep(10)

    print('Test 3: Forward 0.5 Inches')
    my_robot.drive_system.go_straight_inches(0.5, 20)

    time.sleep(10)

    print('Test 4: Forward 0 Inches')
    my_robot.drive_system.go_straight_inches(0, 100)

    time.sleep(10)

    print('Test 5: Forward 60 Inches')
    my_robot.drive_system.go_straight_inches(60, 100)

    time.sleep(10)

def test_spin_in_place_degrees_method():
    """ Tests the spin_in_place_degrees method """
    my_robot = rb.Snatch3rRobot()

    print('Test 1: Clockwise 90 Degrees')
    my_robot.drive_system.spin_in_place_degrees(90, 100)

    time.sleep(10)

    print('Test 2: Counter 90 Degrees')
    my_robot.drive_system.spin_in_place_degrees(-90, -100)

    time.sleep(10)

    print('Test 3: Clockwise 360 Degrees')
    my_robot.drive_system.spin_in_place_degrees(360, 40)

    time.sleep(10)

    print('Test 4: Clockwise 0 Degrees')
    my_robot.drive_system.spin_in_place_degrees(0, 50)

    time.sleep(10)

    print('Test 5: Counter 720 Degrees')
    my_robot.drive_system.spin_in_place_degrees(-720, -100)

    time.sleep(10)

def test_turn_degrees_method():
    """ Tests the spin_in_place_degrees method """
    my_robot = rb.Snatch3rRobot()

    print('Test 1: Clockwise 90 Degrees')
    my_robot.drive_system.turn_degrees(90, 100)

    time.sleep(10)

    print('Test 2: Counter 90 Degrees')
    my_robot.drive_system.turn_degrees(-90, -100)

    time.sleep(10)

    print('Test 3: Clockwise 360 Degrees')
    my_robot.drive_system.turn_degrees(360, 40)

    time.sleep(10)

    print('Test 4: Clockwise 0 Degrees')
    my_robot.drive_system.turn_degrees(0, 50)

    time.sleep(10)

    print('Test 5: Counter 720 Degrees')
    my_robot.drive_system.turn_degrees(-720, -100)

    time.sleep(10)

def test_move_in_polygon_function():
    """ Tests the move_in_polygon function. """
    my_robot = rb.Snatch3rRobot()

    print('Test 1: Octagon')
    move_in_polygon(my_robot, 8)

    time.sleep(10)

    print('Test 2: Small Square')
    move_in_polygon(my_robot, 4, 2)

    time.sleep(10)

    print('Test 3: Slow Triangle')
    move_in_polygon(my_robot, 3, 10, 50)

    time.sleep(10)

    print('Test 4: Decagon')
    move_in_polygon(my_robot, 10, 1)

    time.sleep(10)

    print('Test 5: OH GOD WHY')
    move_in_polygon(my_robot, 100, 0.1)

    time.sleep(10)

def infrared_tester(low, high, beep_volume=100):
    """Causes the robot to beep with volume beep_volume between 0 and 100 when an object is placed between the lower bound
    and upper bound of the function.

        :type low: float
        :type high: float
        :type beep_volume: int

    """
    my_robot = rb.Snatch3rRobot()
    my_robot.sound.set_wait_time(1)
    my_robot.sound.set_volume(beep_volume)
    while True:
        dist = my_robot.proximity_sensor.get_distance_to_nearest_object_in_inches()
        print(dist)
        if low <= dist <= high:
            my_robot.sound.play_beep()
        if my_robot.touch_sensor.is_pressed():
            break

def test_movement(robot, duty_cycle_percent, seconds):
    """ DEPRECATED - USED TO CALCULATE LINEAR REGRESSION """

    robot.drive_system.right_wheel.reset_degrees_spun()
    robot.drive_system.left_wheel.reset_degrees_spun()
    robot.drive_system.start_moving(duty_cycle_percent, duty_cycle_percent)
    time.sleep(seconds)
    robot.drive_system.stop_moving()
    x = robot.drive_system.right_wheel.get_degrees_spun()
    y = robot.drive_system.left_wheel.get_degrees_spun()
    print('Degrees spun in right wheel, Degrees spun in left wheel:', x, y)

def all_movement_tests():
    """ DEPRECATED - USED TO CALCULATE LINEAR REGRESSION """
    my_robot = rb.Snatch3rRobot()
    duty = [100, 50, 100, 25, 10, 50, -100, -20]
    seconds = [3, 6, 1, 4, 10, 10, 10, 1, 5]
    for k in range(len(duty)):
        test_movement(my_robot, duty[k], seconds[k])
        time.sleep(30)

def test_rotation(bottom, top, n):
    """DEPRECATED - USED TO GUESS-AND-CHECK"""
    robot = rb.Snatch3rRobot()
    for k in range(bottom, round(top * n)):
        print('Testing rotation with test value', k / n)
        robot.drive_system.spin_in_place_degrees(90, k / n)
        time.sleep(20)

main()


import sys

class MyGraph:
    def __init__(self,n):
        self.vertices = {}
        for i in range(n):
            self.vertices[i] = []

    def addConnection(self, i,j):
        if (i and i) in self.vertices.keys():
            self.vertices[i].append(j)
            self.vertices[j].append(i)
        else:
            print("VERTEX DOES NOT EXIST")

    def dijkstra(self,start):

        visited = {}
        previous = {}
        distances = {}

        for v in self._vertices.keys():
            visited[v] = False
            previous[v] = None
            distances[v] = sys.maxsize

        distances[start] = 0

        for n in range(len(self._vertices)):
            u = self.minDistance(distances, visited)
            visited[u] = True

            for adj in self._vertices[u]:
                i = adj._vertex
                w = adj._weight

                if visited[i] == False and distances[i] > distances[u] + w:
                    distances[i] = distances[u] + w
                    previous[i] = u

        return distances, previous

    def minDistance(self, distances, visited):
        min = sys.maxsize

        for vertex in self.vertices.keys():
            if distances[vertex] <= min and visited[vertex] == False:
                min = distances[vertex]  # update the new smallest
                min_vertex = vertex  # update the index of the smallest

        return min_vertex

    def minimumPath(self, start, end):

        print("Mininum path from ", start, " to ", end)

        if (start or end) not in self.vertices.keys():
            return None, None

        distances, previous = self.dijkstra(start)
        if distances[end] == sys.maxsize:
            print("There is not path from {} to {}".format(start, end))
            return None, sys.maxsize

        minimum_path = []
        prev = previous[end]
        while prev != None:
            minimum_path.insert(0, prev)
            prev = previous[prev]

        minimum_path.append(end)

        return minimum_path, distances[end]

    def possible_paths(self, start, end):
        if start not in self.vertices.keys():
            return None
        if end not in self.vertices.keys():
            return None

        visited = {}
        for i in self.vertices.keys():
            visited[i] = False

        l_total = []
        self._possible_paths(start, end, visited, [], l_total)
        return l_total

    def _possible_paths(self, v, end, visited, l, l_total):

        visited[v] = True
        l.append(v)

        if v == end:
            l_total.append(l.copy())
        for adj in self.vertices[v]:
            if visited[adj] is False:
                self._possible_paths(adj, end, visited, l, l_total)
        l.pop()
        visited[v] = False

g = MyGraph(5)
g.addConnection(0, 2)  # A:0,  B:1
g.addConnection(0, 3)  # A:0,  C:2
g.addConnection(2, 4)  # A:0,  E:5
g.addConnection(2, 3)  # B:1,  D:4
g.addConnection(2, 1)  # C:2,  B:1
g.addConnection(4,2)

print(g.possible_paths(0, 2))


import pandas as pd
import datetime
import gc

#selected_aircrafts = ['ERJ 170-200 LR', 'BOEING 737-8AS', 'A320 214', '172S', 'A320-214', '737-7H4', 'A320-232', 'A320 232', 'A321-231', 'CL-600-2D24', '737-823', 'A320 214SL', '737-8H4', 'PA-28-181', '737-924ER', 'A321 231SL', 'AIRBUS A319-111', '737-800', 'A319 112', 'A321 231', 'CL-600-2B19', 'A320 232SL', 'A321-211', '172N', 'A319-112', 'A319 111', 'CL-600-2C10', '208B', 'EMB-145LR', 'AIRBUS A320-232', '737-824', '737NG 8FE/W', 'BD-100-1A10', 'SR22', '737NG 8K2/W', 'AIRBUS A320-214', '737NG 838/W', 'BOEING 737-800', '737-900ER', 'DHC-8-402 (Dehavilland)', 'AIRBUS A319-131', 'A330 343E', '172R', 'A319-132', '172M', '737NG 800/W', 'A380 861', 'SR20', '737NG 823/W', 'A319-111']
### get models removed temporarily ###
def get_models(df):
    return df[df.model.isin(selected_aircrafts)]

def add_dateinfo(df):
    df['lastseen'] = pd.to_datetime(df['lastseen'])
    df['valor_semana'] = df['lastseen'].apply(lambda x: x.isocalendar()[1])

    # 0 -> segunda .... 7 -> domingo
    df['dia_semana'] = df['lastseen'].apply(lambda x: x.weekday())
    return df

### get models removed temporarily ###

#2018
df = pd.read_csv('data/filghts_2018_cleaned.csv')
#df  = get_models(df)
df = add_dateinfo(df)

#2019 p1
df_2 = pd.read_csv('data/filghts_2019_1_cleaned.csv')
#df_2  = get_models(df_2)
df_2 = add_dateinfo(df_2)
df = df.append(df_2)
#clean memory
del df_2
gc.collect()

#2019 p2
df_3 = pd.read_csv('data/filghts_2019_2_cleaned.csv')
#df_3  = get_models(df_3)
df_3 = add_dateinfo(df_3)
df = df.append(df_3)
#full dataframe
df = df.append(df_3)
#clean memory
del df_3
gc.collect()

#Grouping origin df
df_origem = df.groupby(['pais_origem', 'valor_semana', 'dia_semana']).size().reset_index()
df_origem = pd.DataFrame(df_origem)
df_origem.columns = ['pais', 'valor_semana','dia_semana', 'voos']

#Grouping destination df
df_destino =df.groupby(['pais_destino', 'valor_semana', 'dia_semana']).size().reset_index()
df_destino = pd.DataFrame(df_destino)
df_destino.columns = ['pais', 'valor_semana','dia_semana', 'voos']

#Merging origin and destination (we use sum of fligths in the country)
df_final = df_origem.merge(df_destino, 'outer', on = ['pais', 'valor_semana','dia_semana'])
df_final['voos_soma'] = df_final['voos_x'].fillna(0) + df_final['voos_y'].fillna(0)
df_final = df_final[['pais','valor_semana','dia_semana','voos_soma']]
df_final.to_csv('data/2018-2019linreg_allAircrafts.csv', index=False)


import blankly
from blankly import StrategyState

class Grid:
    @staticmethod
    def init(retry_timeout, base_order_size, grid_spacing_percent, num_grid_above_below):
        def closure(symbol, state: StrategyState):
            state.variables['retry_timeout'] = retry_timeout
            state.variables['base_order_size'] = base_order_size
            state.variables['grid_spacing_percent'] = grid_spacing_percent
            state.variables['num_grid_above_below'] = num_grid_above_below
            state.variables['last_try'] = state.time
        return closure

    @staticmethod
    def price_event(price, symbol, state: StrategyState):
        now = state.time

        orders = state.interface.get_open_orders(symbol)
        if len(orders) > 0:
            if now - state.variables['last_try'] < state.variables['retry_timeout']:
                return

            for order in orders:
                state.interface.cancel_order(symbol, order_id=order['id'])

        state.variables['last_try'] = now

        buy_grids, sell_grids = generate_grid(price, state.variables['num_grid_above_below'], state.variables['grid_spacing_percent'])

        for i in range(len(buy_grids)):
            try:
                buy_price = blankly.trunc(buy_grids[i], 2)
                state.interface.limit_order(symbol, 'buy', buy_price, state.variables['base_order_size'])
                print("new buy order", "size: ", state.variables['base_order_size'], "price", buy_price)
            except Exception as e:
                print('"failed to place order"', '"error"', e, "price", buy_price, buy_grids[i])

            try:
                sell_price = blankly.trunc(sell_grids[i], 2)
                state.interface.limit_order(symbol, 'sell', sell_price, state.variables['base_order_size'])
                print("new sell order", "size: ", state.variables['base_order_size'], "price", sell_price)
            except Exception as e:
                print('"failed to place order"', '"error"', e, "price", sell_price, sell_grids[i])

def generate_grid(price, num_grid_above_below, grid_spacing_percent):
    buy_grid = []
    sell_grid = []

    for i in range(num_grid_above_below):
        price_above = price * (1 + grid_spacing_percent / 100) ** (i + 1)
        price_below = price * (1 - grid_spacing_percent / 100) ** (i + 1)
        sell_grid.append(price_above)
        buy_grid.append(price_below)

    return buy_grid, sell_grid


import os
import logging
import numpy as np
import pandas as pd
from sklearn.neighbors import KNeighborsClassifier

from .helpers import ROOT_DIR, DATA_DIR, RESULTS_DIR
from .helpers import GENRES, SUBSETS
from .helpers import start_experiment_log
from .helpers import relpath
from .classification import train_model
from .classification import get_scores

np.random.seed(0)

# Globals
# —————————————————————————————————————————————————————————————————————————————

PITCH_FEATURES = [f'freq_MIDI_{i}' for i in range(53, 87)]

PITCH_CLASS_FEATURES = [f'freq_pitch_class_{i}' for i in range(12)]

REPETITION_FEATURES = [f'repetition_score_MIDI_{i}' for i in range(53, 87)]

PROFILES = {
    'pitch': PITCH_FEATURES,
    'pitch_class': PITCH_CLASS_FEATURES,
    'repetition': REPETITION_FEATURES
}

# Helpers
# —————————————————————————————————————————————————————————————————————————————

def get_conditions(genres='all', subsets='all', profiles='all'):
    """Get a list of experimental conditions

    Parameters
    ----------
    genres : str or list, optional
        Genres to include, by default 'all'
    subsets : str or list, optional
        Subsets to include, by default 'all'
    profile : str or list, optional
        Profiles to include, by default 'all'

    Returns
    -------
    (list, dict)
        A list with all conditions, and a dictionary with keys 'genres',
        'subsets' and 'profiles' containing those values.
    """
    subsets = SUBSETS if subsets == 'all' else subsets
    genres = GENRES if genres == 'all' else genres
    profiles = list(PROFILES.keys()) if profiles == 'all' else profiles

    conditions = []
    for genre in genres:
        for subset in subsets:
            for profile in profiles:
                conditions.append(
                    dict(genre=genre, subset=subset, profile=profile))

    parts = dict(subsets=subsets, genres=genres, profiles=profiles)
    return conditions, parts

def load_dataset(genre, subset, profile, split, data_dir=DATA_DIR):
    """Load a dataset for training the classifier. Returns a dataframe of
    features and an array of corresponding targets (modes)"""
    feature_names = PROFILES[profile]
    features_fn = os.path.join(data_dir, genre, subset, f'{split}-features.csv')
    data = pd.read_csv(features_fn, index_col=0)[feature_names]
    chants_fn = os.path.join(data_dir, genre, subset, f'{split}-chants.csv')
    targets = pd.read_csv(chants_fn, index_col=0)['mode']
    assert len(targets) == len(data)
    return data, targets

# Experiment
# —————————————————————————————————————————————————————————————————————————————

def run_condition(genre, subset, profile,
                  data_dir, results_dir,
                  n_iter, n_splits):
    """Runs a single experimental condition: trains the classifier, stores
    all the model, cross-validation results, and evaluation scores."""
    # Start experiment
    logging.info(f'Training model...')
    logging.info(f'* profile={profile}')
    logging.info(f'* genre={genre}')
    logging.info(f'* subset={subset}')

    # Set up directories
    output_dir = os.path.join(results_dir, genre, subset)
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    # Load training and testing data
    kwargs = dict(genre=genre, subset=subset, profile=profile, data_dir=data_dir)
    train_data, train_targets = load_dataset(split='train', **kwargs)
    test_data, test_targets = load_dataset(split='test', **kwargs)
    logging.info(f'* Training/test size: {len(train_data)}/{len(test_data)}')
    logging.info(f'* Num. features: {train_data.shape[1]}')

    # Model parameters and param grid for tuning
    fixed_params = {
        'n_jobs': -1,
        'p': 2,
        'metric': 'minkowski'
    }
    tuned_params = {
        'n_neighbors': np.arange(1, 50),
        'weights': ['uniform', 'distance'],
        'algorithm': ['ball_tree', 'kd_tree', 'brute'],
        'leaf_size': np.arange(10, 100)
    }

    # Tune and train model
    model = KNeighborsClassifier(**fixed_params)
    train_model(
        model = model,
        train_data=train_data,
        train_targets=train_targets,
        test_data=test_data,
        test_targets=test_targets,
        param_grid=tuned_params,
        n_splits=n_splits,
        n_iter=n_iter,
        basepath=os.path.join(output_dir, profile)
    )

def run(experiment_name, description=None,
        genres='all', subsets='all', profiles='all',
        n_iter=100, n_splits=5,
        data_dir = DATA_DIR, results_dir = RESULTS_DIR):
    """Run a 'profile' mode classification experiment using pitch, pitch_class
    and repetition profiles."""
    # Set up directories
    results_dir = os.path.join(results_dir, experiment_name)
    if not os.path.exists(results_dir):
        os.makedirs(results_dir)

    # Get all conditions
    conditions, parts = get_conditions(genres, subsets, profiles)

    # Start log and log experiment settings
    start_experiment_log(
        name=experiment_name,
        description=description,
        data_dir=data_dir,
        results_dir=results_dir,
        n_iter=n_iter,
        n_splits=n_splits,
        num_conditions=len(conditions),
        **parts)

    # Train all models
    for condition in conditions:
        run_condition(
            data_dir=data_dir, results_dir=results_dir,
            n_iter=n_iter, n_splits=n_splits, **condition)

def evaluate(experiment_name,
             genres='all', subsets='all', profiles='all',
             data_dir = DATA_DIR, results_dir = RESULTS_DIR, **kwargs):
    """Evaluate an experiment and store accuracy and retrieval scores in a
    single CSV file that can be used to e.g. generate tables and figures."""
    logging.info('Evaluating experiment...')
    results_dir = os.path.join(results_dir, experiment_name)
    scores = []
    conditions, _ = get_conditions(genres, subsets, profiles)
    for condition in conditions:
        profile = condition['profile']
        output_dir = os.path.join(
            results_dir, condition['genre'], condition['subset'])
        condition_scores = get_scores(
            test_pred_fn = os.path.join(output_dir, f'{profile}-test-pred.txt'),
            train_pred_fn = os.path.join(output_dir, f'{profile}-train-pred.txt'),
            genre=condition['genre'],
            subset=condition['subset'],
            data_dir=data_dir)
        condition_scores.update(condition)
        scores.append(condition_scores)

    scores_fn = os.path.join(results_dir, f'{experiment_name}-scores.csv')
    pd.DataFrame(scores).to_csv(scores_fn)
    logging.info(f'> Stored scores to {relpath(scores_fn)}')


class Data():

    def getDataList(filename):
        """
        read the file and get the I and V values and return it
        :param filename: complete path to the file
        :return: list containing a list of the I values and a list of the V values
        """
        with open(filename, 'r') as file:
            wholeFile = file.read()
            data = wholeFile.split('**Data**')[1]
            print(filename)
            splitted = data.split()
            v = []
            i = []

            # the IV parameters are in columns after '**Data**' in the file
            # in the .drk files, 3 columns are present in the file, only the first 2 columns are important
            # in the .lgt files, 4 columns are present in the file, only the first 2 columns are important
            # in other files, only 2 columns are present
            if str(filename).endswith('drk'):
                jump = 3
            elif str(filename).endswith('lgt'):
                jump = 4
            else:
                jump = 2

            for j in range(0, len(splitted),jump):
                if(j+1 < len(splitted)):
                    v.append(float(splitted[j]))
                    i.append(float(splitted[j+1]))

            iv = [i, v]
            return iv

    def getDataListPsc(filename):
        with open(filename, 'r') as file:
            data = file.read()
            splitted = data.split()
            v = []
            i = []

            jump = 2

            for j in range(0, len(splitted), jump):
                if(j+1<len(splitted)):
                    i.append(float(splitted[j]))
                    v.append(float(splitted[j+1]))

            iv = [i, v]

            return iv

    def getDataListSm(filename):
        with open(filename, 'r') as file:
            data = file.read()
            data = data.replace(",", ".")
            rows = data.split('\n')
            v = []
            i = []

            for j in rows:
                if j != '':
                    splitted = j.split(';')
                    i.append(float(splitted[1]))
                    v.append(float(splitted[0]))

            iv = [i, v]

            return iv