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

import tempfile

from tensorflow.contrib import framework as contrib_framework
from tensorflow.contrib import layers
from tensorflow.contrib.framework.python.ops import variables
from tensorflow.contrib.layers.python.layers import initializers
from tensorflow.contrib.layers.python.layers import optimizers
from tensorflow.contrib.learn.python.learn import evaluable
from tensorflow.contrib.learn.python.learn import metric_spec
from tensorflow.contrib.learn.python.learn import trainable
from tensorflow.contrib.learn.python.learn.estimators import estimator
from tensorflow.contrib.metrics.python.ops import metric_ops
from tensorflow.python.framework import ops
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import control_flow_ops
from tensorflow.python.ops import init_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import nn
from tensorflow.python.ops import partitioned_variables
from tensorflow.python.ops import standard_ops
from tensorflow.python.ops import variable_scope
from tensorflow.python.training import training as train

_CLASSES = "classes"
_TOP_K = "top_k"
_PROBABILITIES = "probabilities"
_DEFAULT_LEARNING_RATE = 0.01

def _as_iterable(preds, output):
  for pred in preds:
    yield pred[output]

def _get_optimizer(optimizer):
  if callable(optimizer):
    return optimizer()
  else:
    return optimizer

def _get_default_optimizer():
  """Default optimizer for DNN models."""
  return train.AdagradOptimizer(_DEFAULT_LEARNING_RATE)

def _get_feature_dict(features):
  if isinstance(features, dict):
    return features
  return {"": features}

def _dnn_sampled_softmax_classifier_model_fn(features, targets, mode, params):
  """model_fn that uses candidate sampling.

  Args:
    features: Single Tensor or dict of Tensor (depends on data passed to `fit`)
    targets: A single Tensor of shape [batch_size, n_labels] containing
      the target indices.
    mode: Represents if this training, evaluation or prediction. See `ModeKeys`.
    params: A dict of hyperparameters that are listed below.
      hidden_units- List of hidden units per layer. All layers are fully
        connected. Ex. `[64, 32]` means first layer has 64 nodes and second one
        has 32.
      feature_columns- An iterable containing all the feature columns used by
        the model. All items in the set should be instances of classes derived
        from `FeatureColumn`.
      n_classes- number of target classes. It must be greater than 2.
      n_samples- number of sample target classes. Needs to be tuned - A good
        starting point could be 2% of n_classes.
      n_labels- number of labels in each example.
      top_k- The number of classes to predict.
      optimizer- An instance of `tf.Optimizer` used to train the model. If
        `None`, will use an Adagrad optimizer.
      dropout- When not `None`, the probability we will drop out a given
        coordinate.
      gradient_clip_norm- A float > 0. If provided, gradients are
        clipped to their global norm with this clipping ratio. See
        tf.clip_by_global_norm for more details.
      num_ps_replicas- The number of parameter server replicas.

  Returns:
    predictions: A single Tensor or a dict of Tensors.
    loss: A scalar containing the loss of the step.
    train_op: The op for training.
  """

  hidden_units = params["hidden_units"]
  feature_columns = params["feature_columns"]
  n_classes = params["n_classes"]
  n_samples = params["n_samples"]
  n_labels = params["n_labels"]
  top_k = params["top_k"]
  optimizer = params["optimizer"]
  dropout = params["dropout"]
  gradient_clip_norm = params["gradient_clip_norm"]
  num_ps_replicas = params["num_ps_replicas"]

  parent_scope = "dnn_ss"

  features = _get_feature_dict(features)
  targets = _reshape_targets(targets)

  # Setup the input layer partitioner.
  input_layer_partitioner = (
      partitioned_variables.min_max_variable_partitioner(
          max_partitions=num_ps_replicas,
          min_slice_size=64 << 20))

  # Create the input layer.
  with variable_scope.variable_scope(
      parent_scope + "/input_from_feature_columns",
      features.values(),
      partitioner=input_layer_partitioner) as scope:
    net = layers.input_from_feature_columns(
        features,
        feature_columns,
        weight_collections=[parent_scope],
        scope=scope)

  # Setup the hidden layer partitioner.
  hidden_layer_partitioner = (
      partitioned_variables.min_max_variable_partitioner(
          max_partitions=num_ps_replicas))

  final_hidden_layer_dim = None
  # Create hidden layers using fully_connected.
  for layer_id, num_hidden_units in enumerate(hidden_units):
    with variable_scope.variable_scope(
        parent_scope + "/hiddenlayer_%d" % layer_id, [net],
        partitioner=hidden_layer_partitioner) as scope:
      net = layers.fully_connected(net,
                                   num_hidden_units,
                                   variables_collections=[parent_scope],
                                   scope=scope)
      final_hidden_layer_dim = num_hidden_units
      # Add dropout if it is enabled.
      if dropout is not None and mode == estimator.ModeKeys.TRAIN:
        net = layers.dropout(net, keep_prob=(1.0 - dropout))

  # Create the weights and biases for the logit layer.
  with variable_scope.variable_scope(
      parent_scope + "/logits", [net],
      partitioner=hidden_layer_partitioner) as scope:
    dtype = net.dtype.base_dtype
    weights_shape = [n_classes, final_hidden_layer_dim]
    weights = variables.model_variable(
        "weights",
        shape=weights_shape,
        dtype=dtype,
        initializer=initializers.xavier_initializer(),
        trainable=True,
        collections=[parent_scope])
    biases = variables.model_variable(
        "biases",
        shape=[n_classes,],
        dtype=dtype,
        initializer=init_ops.zeros_initializer,
        trainable=True,
        collections=[parent_scope])

  if mode == estimator.ModeKeys.TRAIN:
    # Call the candidate sampling APIs and calculate the loss.
    sampled_values = nn.learned_unigram_candidate_sampler(
        true_classes=math_ops.to_int64(targets),
        num_true=n_labels,
        num_sampled=n_samples,
        unique=True,
        range_max=n_classes)

    sampled_softmax_loss = nn.sampled_softmax_loss(
        weights=weights,
        biases=biases,
        inputs=net,
        labels=math_ops.to_int64(targets),
        num_sampled=n_samples,
        num_classes=n_classes,
        num_true=n_labels,
        sampled_values=sampled_values)

    loss = math_ops.reduce_mean(sampled_softmax_loss, name="loss")

    train_op = optimizers.optimize_loss(
        loss=loss, global_step=contrib_framework.get_global_step(),
        learning_rate=_DEFAULT_LEARNING_RATE,
        optimizer=_get_optimizer(optimizer), clip_gradients=gradient_clip_norm,
        name=parent_scope)
    return None, loss, train_op

  elif mode == estimator.ModeKeys.EVAL:
    logits = nn.bias_add(standard_ops.matmul(net, array_ops.transpose(weights)),
                         biases)
    predictions = {}
    predictions[_PROBABILITIES] = nn.softmax(logits)
    predictions[_CLASSES] = math_ops.argmax(logits, 1)
    _, predictions[_TOP_K] = nn.top_k(logits, top_k)

    # Since the targets have multiple labels, setup the target probabilities
    # as 1.0/n_labels for each of the labels.
    target_one_hot = array_ops.one_hot(
        indices=targets, depth=n_classes, on_value=1.0 / n_labels)
    target_one_hot = math_ops.reduce_sum(
        input_tensor=target_one_hot,
        reduction_indices=[1])

    loss = math_ops.reduce_mean(
        nn.softmax_cross_entropy_with_logits(logits, target_one_hot))

    return predictions, loss, None

  elif mode == estimator.ModeKeys.INFER:
    logits = nn.bias_add(standard_ops.matmul(net, array_ops.transpose(weights)),
                         biases)
    predictions = {}
    predictions[_PROBABILITIES] = nn.softmax(logits)
    predictions[_CLASSES] = math_ops.argmax(logits, 1)
    _, predictions[_TOP_K] = nn.top_k(logits, top_k)

    return predictions, None, None

def _reshape_targets(targets):
  if targets is None:
    return None
  check_shape_op = control_flow_ops.Assert(
      math_ops.less_equal(array_ops.rank(targets), 2),
      ["target's should be either [batch_size, n_labels] or [batch_size]"])
  with ops.control_dependencies([check_shape_op]):
    targets = array_ops.reshape(
        targets, shape=[array_ops.shape(targets)[0], -1])
  return targets

def _top_k_fn_wrapper(metric_fn, k):

  def wrap_func(predictions, labels):
    return metric_fn(predictions, _reshape_targets(labels), k=k)

  wrap_func.__name__ = metric_fn.__name__
  return wrap_func

class _DNNSampledSoftmaxClassifier(trainable.Trainable, evaluable.Evaluable):
  """A classifier for TensorFlow DNN models.

  Example:

  ```python
  legos = sparse_column_with_hash_bucket(column_name="legos",
                                         hash_bucket_size=1000)
  watched_videos = sparse_column_with_hash_bucket(
                     column_name="watched_videos",
                     hash_bucket_size=20000)

  legos_emb = embedding_column(sparse_id_column=legos, dimension=16,
                               combiner="sum")
  watched_videos_emb = embedding_column(sparse_id_column=watched_videos,
                                        dimension=256,
                                        combiner="sum")

  estimator = DNNSampledSoftmaxClassifier(
      n_classes=500000, n_samples=10000, n_labels=5,
      feature_columns=[legos_emb, watched_videos_emb],
      hidden_units=[1024, 512, 256])

  # Or estimator using the Adam optimizer with dropout.
  estimator = DNNSampledSoftmaxClassifier(
      feature_columns=[education_emb, occupation_emb],
      hidden_units=[1024, 512, 256],
      optimizer=tf.train.ProximalAdagradOptimizer(
        learning_rate=0.1),
      dropout=0.1)

  # Input builders
  def input_fn_train: # returns x, Y
    pass
  estimator.fit(input_fn=input_fn_train)

  def input_fn_eval: # returns x, Y
    pass
  estimator.evaluate(input_fn=input_fn_eval)
  estimator.predict(x=x)
  ```

  Input of `fit` and `evaluate` should have following features,
    otherwise there will be a `KeyError`:

  * for each `column` in `feature_columns`:
    - if `column` is a `SparseColumn`, a feature with `key=column.name`
      whose `value` is a `SparseTensor`.
    - if `column` is a `EmbeddingColumn`, a feature with `key=column.name`
      whose `value` is a `SparseTensor`.
    - if `column` is a `WeightedSparseColumn`, two features: the first with
      `key` the id column name, the second with `key` the weight column name.
      Both features' `value` must be a `SparseTensor`.
    - if `column` is a `RealValuedColumn`, a feature with `key=column.name`
      whose `value` is a `Tensor`.
  """

  def __init__(self,
               hidden_units,
               feature_columns,
               n_classes,
               n_samples,
               n_labels=1,
               top_k=1,
               model_dir=None,
               optimizer=None,
               dropout=None,
               gradient_clip_norm=None,
               config=None,
               feature_engineering_fn=None):
    """Initializes a DNNSampledSoftmaxClassifier instance.

    Args:
      hidden_units: List of hidden units per layer. All layers are fully
        connected. Ex. `[64, 32]` means first layer has 64 nodes and second one
        has 32.
      feature_columns: An iterable containing all the feature columns used by
        the model. All items in the set should be instances of classes derived
        from `FeatureColumn`.
      n_classes: number of target classes. It must be greater than 2.
      n_samples: number of sample target classes. Needs to be tuned - A good
        starting point could be 2% of n_classes.
      n_labels: number of labels in each example.
      top_k: The number of classes to predict.
      model_dir: Directory to save model parameters, graph and etc. This can
        also be used to load checkpoints from the directory into a estimator to
        continue training a previously saved model.
      optimizer: An instance of `tf.Optimizer` used to train the model. If
        `None`, will use an Adagrad optimizer.
      dropout: When not `None`, the probability we will drop out a given
        coordinate.
      gradient_clip_norm: A float > 0. If provided, gradients are
        clipped to their global norm with this clipping ratio. See
        tf.clip_by_global_norm for more details.
      config: `RunConfig` object to configure the runtime settings.
      feature_engineering_fn: Feature engineering function. Takes features and
                        targets which are the output of `input_fn` and
                        returns features and targets which will be fed
                        into the model.

    Returns:
      A `DNNSampledSoftmaxClassifier` estimator.

    Raises:
      ValueError: If n_classes <= 2.
      ValueError: If n_classes < n_samples.
      ValueError: If n_classes < n_labels.
    """
    # Validate all the inputs.
    if n_classes <= 2:
      raise ValueError("n_classes should be greater than 2. For n_classes <= 2,"
                       " use DNNClassifier.")
    if n_classes < n_samples:
      raise ValueError("n_classes (%d) should be greater than n_samples (%d)." %
                       (n_classes, n_samples))
    if n_classes < n_labels:
      raise ValueError("n_classes (%d) should be greater than n_labels"
                       " (%d)." % (n_classes, n_labels))

    self._top_k = top_k
    self._feature_columns = feature_columns
    assert self._feature_columns
    self._model_dir = model_dir or tempfile.mkdtemp()

    # Build the estimator with _dnn_sampled_softmax_classifier_model_fn.
    self._estimator = estimator.Estimator(
        model_fn=_dnn_sampled_softmax_classifier_model_fn,
        model_dir=self._model_dir,
        config=config,
        params={
            "hidden_units": hidden_units,
            "feature_columns": feature_columns,
            "n_classes": n_classes,
            "n_samples": n_samples,
            "n_labels": n_labels,
            "top_k": top_k,
            "optimizer": optimizer or _get_default_optimizer(),
            "dropout": dropout,
            "gradient_clip_norm": gradient_clip_norm,
            "num_ps_replicas": config.num_ps_replicas if config else 0
        },
        feature_engineering_fn=feature_engineering_fn)

  def get_estimator(self):
    return self._estimator

  def fit(self, x=None, y=None, input_fn=None, steps=None, batch_size=None,
          monitors=None, max_steps=None):
    """See trainable.Trainable."""
    return self._estimator.fit(x=x, y=y, input_fn=input_fn, steps=steps,
                               batch_size=batch_size, monitors=monitors,
                               max_steps=max_steps)

  def evaluate(self, x=None, y=None, input_fn=None, feed_fn=None,
               batch_size=None, steps=None, metrics=None, name=None,
               range_k=None):
    # pylint: disable=g-doc-args,g-doc-return-or-yield
    """See evaluable.Evaluable for a description of the Args.

    Calculates the following metrics by default:
      loss
      average_precision@top_k: see
        https://en.wikipedia.org/wiki/Information_retrieval#Average_precision
      for k in range_k:
        precision@k and recall@k

    range_k: A list of numbers where precision and recall have to be obtained.
      For eg. range_k=[1,5] will calculate precision@1, precision@5,
      recall@1 and recall@5. If None, defaults to [1, top_k].
    """
    if not metrics:
      metrics = {}
    metrics.update({
        "average_precision_at_%d" % self._top_k: metric_spec.MetricSpec(
            metric_fn=_top_k_fn_wrapper(
                metric_ops.streaming_sparse_average_precision_at_k,
                k=self._top_k),
            prediction_key=_PROBABILITIES)
    })
    if range_k is None:
      if self._top_k > 1:
        range_k = [1, self._top_k]
      else:
        range_k = [1]
    for k in range_k:
      metrics.update({
          "precision_at_%d" % k: metric_spec.MetricSpec(
              metric_fn=_top_k_fn_wrapper(
                  metric_ops.streaming_sparse_precision_at_k, k=k),
              prediction_key=_PROBABILITIES,)
      })
      metrics.update({
          "recall_at_%d" % k: metric_spec.MetricSpec(
              metric_fn=_top_k_fn_wrapper(
                  metric_ops.streaming_sparse_recall_at_k, k=k),
              prediction_key=_PROBABILITIES,)
      })

    return self._estimator.evaluate(x=x, y=y, input_fn=input_fn,
                                    feed_fn=feed_fn, batch_size=batch_size,
                                    steps=steps, metrics=metrics, name=name)

  def predict(self, x=None, input_fn=None, batch_size=None, as_iterable=False,
              get_top_k=False):
    """Returns predicted classes for given features.

    Args:
      x: features.
      input_fn: Input function. If set, x must be None.
      batch_size: Override default batch size.
      as_iterable: If True, return an iterable which keeps yielding predictions
        for each example until inputs are exhausted. Note: The inputs must
        terminate if you want the iterable to terminate (e.g. be sure to pass
        num_epochs=1 if you are using something like read_batch_features).
      get_top_k : if set to true returns the top k classes otherwise returns
        the top class.

    Returns:
      Numpy array of predicted classes (or an iterable of predicted classes if
      as_iterable is True).
    """
    if get_top_k:
      key = _TOP_K
    else:
      key = _CLASSES
    preds = self._estimator.predict(x=x, input_fn=input_fn,
                                    batch_size=batch_size, outputs=[key],
                                    as_iterable=as_iterable)
    if as_iterable:
      return _as_iterable(preds, output=key)
    return preds[key]

  def predict_proba(self, x=None, input_fn=None, batch_size=None,
                    as_iterable=False):
    """Returns prediction probabilities for given features.

    Args:
      x: features.
      input_fn: Input function. If set, x and y must be None.
      batch_size: Override default batch size.
      as_iterable: If True, return an iterable which keeps yielding predictions
        for each example until inputs are exhausted. Note: The inputs must
        terminate if you want the iterable to terminate (e.g. be sure to pass
        num_epochs=1 if you are using something like read_batch_features).

    Returns:
      Numpy array of predicted probabilities (or an iterable of predicted
      probabilities if as_iterable is True).
    """
    preds = self._estimator.predict(x=x, input_fn=input_fn,
                                    batch_size=batch_size,
                                    outputs=[_PROBABILITIES],
                                    as_iterable=as_iterable)
    if as_iterable:
      return _as_iterable(preds, output=_PROBABILITIES)
    return preds[_PROBABILITIES]

  def export(self, export_dir, signature_fn=None,
             input_fn=None, default_batch_size=1,
             exports_to_keep=None):
    """Exports inference graph into given dir.

    Args:
      export_dir: A string containing a directory to write the exported graph
        and checkpoints.
      signature_fn: Function that returns a default signature and a named
        signature map, given `Tensor` of `Example` strings, `dict` of `Tensor`s
        for features and `Tensor` or `dict` of `Tensor`s for predictions.
      input_fn: If `use_deprecated_input_fn` is true, then a function that given
        `Tensor` of `Example` strings, parses it into features that are then
        passed to the model. Otherwise, a function that takes no argument and
        returns a tuple of (features, targets), where features is a dict of
        string key to `Tensor` and targets is a `Tensor` that's currently not
        used (and so can be `None`).
      default_batch_size: Default batch size of the `Example` placeholder.
      exports_to_keep: Number of exports to keep.

    Returns:
      The string path to the exported directory. NB: this functionality was
      added ca. 2016/09/25; clients that depend on the return value may need
      to handle the case where this function returns None because subclasses
      are not returning a value.
    """
    def default_input_fn(unused_estimator, examples):
      return layers.parse_feature_columns_from_examples(
          examples, self._feature_columns)
    return self._estimator.export(export_dir=export_dir,
                                  signature_fn=signature_fn,
                                  input_fn=input_fn or default_input_fn,
                                  default_batch_size=default_batch_size,
                                  exports_to_keep=exports_to_keep)

  def get_variable_names(self):
    return self._estimator.get_variable_names()

  @property
  def model_dir(self):
    return self._model_dir


import PyQt5.QtCore as qtc
import PyQt5.QtGui as qtg
import PyQt5.QtWidgets as qtw

from application_gui.common_gui_functions import CLabel, CHorizontalSeparator, CLabelledLineEdit
from application_gui.settings_scale.functions import SetScaleFunctions

##-\-\-\-\-\-\-\-\-\-\-\-\-\-\
## WINDOW FOR READING METADATA
##-/-/-/-/-/-/-/-/-/-/-/-/-/-/

class SetScaleWindow(qtw.QMainWindow, SetScaleFunctions):
    def __init__(self, parent, image_class=None):
        super(SetScaleWindow, self).__init__(parent)

        # Initialise the subwindow
        self.parent = parent
        self.image_class = image_class
        #self.setWindowModality(qtc.Qt.ApplicationModal)

        # Generate the window
        self.mainWidget = qtw.QWidget()
        self.mainLayout = qtw.QVBoxLayout(self.mainWidget)
        self.setWindowTitle("Set Scale")

        # Populate the panel
        self.createScaleSettings(self.mainLayout)
        #self.mainLayout.addWidget( CHorizontalSeparator() )
        self.createUserActions(self.mainLayout)

        # Display the panel
        self.mainWidget.setLayout(self.mainLayout)
        self.setCentralWidget(self.mainWidget)
        self.show()
        self.setFixedSize(self.size())

    # ---------------------------------------------------
    # Reinitialise the display when the window is closed
    def closeEvent(self, event=None):
        event.accept()
        self.parent.subWindows['set_scale'] = None

    ##-\-\-\-\-\-\-\-\-\-\
    ## GENERATE THE DISPLAY
    ##-/-/-/-/-/-/-/-/-/-/

    # ------------------------------------------
    # Generate the display for the spatial scale
    def createScaleSettings(self, parentWidget):

        # Generate the widget
        self.scaleSettingsWidget = qtw.QWidget()
        self.scaleSettingsLayout = qtw.QGridLayout(self.scaleSettingsWidget)

        # Label for space scale
        current_row = 0
        self.scaleSettingsLayout.addWidget( CLabel("Space Scale:"), current_row, 0, 1, 2 )

        # Entry for the distance in pixel
        current_row += 1
        self.scaleSettingsLayout.addWidget( qtw.QLabel("Distance in pixels"), current_row, 0 )
        self.pixelDistanceEntry = qtw.QLineEdit()
        self.pixelDistanceEntry.setText( str(self.image_class.scale.space_scale) )
        self.scaleSettingsLayout.addWidget( self.pixelDistanceEntry, current_row, 1 )

        # Entry for the real distance
        current_row += 1
        self.scaleSettingsLayout.addWidget( qtw.QLabel("Known distance"), current_row, 0 )
        self.knownDistanceEntry = qtw.QLineEdit()
        self.knownDistanceEntry.setText( "1" )
        self.scaleSettingsLayout.addWidget( self.knownDistanceEntry, current_row, 1 )

        # Entry for the length unit
        current_row += 1
        self.scaleSettingsLayout.addWidget( qtw.QLabel("Unit of length"), current_row, 0 )
        self.lengthUnitEntry = qtw.QLineEdit()
        self.lengthUnitEntry.setText( self.image_class.scale.space_unit )
        self.scaleSettingsLayout.addWidget( self.lengthUnitEntry, current_row, 1 )

        current_row += 1
        self.scaleSettingsLayout.addWidget( CHorizontalSeparator(), current_row, 0, 1, 2 )

        # Label for time scale
        current_row += 1
        self.scaleSettingsLayout.addWidget( CLabel("Time Scale:"), current_row, 0, 1, 2 )

        # Entry for the length unit
        current_row += 1
        self.scaleSettingsLayout.addWidget( qtw.QLabel("Frame rate"), current_row, 0 )
        self.frameRateEntry = qtw.QLineEdit()
        self.frameRateEntry.setText( str(self.image_class.scale.frame_rate) )
        self.scaleSettingsLayout.addWidget( self.frameRateEntry, current_row, 1 )

        # Display the widget
        self.scaleSettingsWidget.setLayout(self.scaleSettingsLayout)
        parentWidget.addWidget(self.scaleSettingsWidget)

    # ----------------------------------
    # Generate the controls for the user
    def createUserActions(self, parentWidget):

        # Generate the widget
        self.userActionWidget = qtw.QWidget()
        self.userActionLayout = qtw.QVBoxLayout(self.userActionWidget)

        # Add the global checkbox
        self.globalScaleCheckBox = qtw.QCheckBox("Global")
        self.userActionLayout.addWidget(self.globalScaleCheckBox)

        # Generate the widget
        self.userButtonWidget = qtw.QWidget()
        self.userButtonLayout = qtw.QHBoxLayout(self.userButtonWidget)

        # Add the button to open a new file
        self.applyButton = qtw.QPushButton("Apply")
        self.applyButton.clicked.connect(self.applyScale)
        self.applyButton.setStatusTip("Apply the scale to the stack.")
        self.applyButton.setFixedWidth(125)
        self.userButtonLayout.addWidget(self.applyButton, alignment=qtc.Qt.AlignLeft)

        # Add the button to close
        self.closeButton = qtw.QPushButton("Cancel")
        self.closeButton.clicked.connect(self.close)
        self.closeButton.setStatusTip("Close the current window.")
        self.closeButton.setFixedWidth(125)
        self.userButtonLayout.addWidget(self.closeButton, alignment=qtc.Qt.AlignRight)

        # Display the widget
        self.userButtonWidget.setLayout(self.userButtonLayout)
        self.userActionLayout.addWidget(self.userButtonWidget)

        # Display the widget
        self.userActionWidget.setLayout(self.userActionLayout)
        parentWidget.addWidget(self.userActionWidget)


import construct as cs
import construct_typed as cst
import dataclasses
import typing as t
from . import GalleryItem

constr = cs.Struct(
    "signature" / cs.Bytes(23),
    "data_start_pos" / cs.Tell,
    "data_peek" / cs.Peek(cs.Array(5, cs.Int24ub)),
    cs.Seek(15, 1),
    "pos_after_seek" / cs.Tell,
    "data_pointer"
    / cs.Pointer(lambda ctx: ctx.data_start_pos, cs.Array(5, cs.Int24ub)),
)

gallery_item = GalleryItem(
    construct=constr,
    example_binarys={
        "Zeros": bytes(23 + 15),
        "1": b"TestPointerPeekSeekTell0123456789abcde",
    },
)


"""Platform for degree days sensors."""
from homeassistant.components.sensor import SensorEntity
from homeassistant.helpers import update_coordinator
from homeassistant.helpers.entity import StateType

from . import DegreeDaysData
from .const import DOMAIN, SENSOR_TYPES, DegreeDaysSensorEntityDescription

async def async_setup_entry(hass, entry, async_add_entities):
    """Add degree days entry."""
    coordinator = hass.data[DOMAIN][entry.entry_id]
    async_add_entities(
        DegreeDaysSensor(coordinator, description) for description in SENSOR_TYPES
    )

class DegreeDaysSensor(update_coordinator.CoordinatorEntity, SensorEntity):
    """Representation of a Sensor."""

    entity_description: DegreeDaysSensorEntityDescription

    def __init__(
        self,
        coordinator: DegreeDaysData,
        description: DegreeDaysSensorEntityDescription,
    ) -> None:
        """Initialize the sensor."""
        super().__init__(coordinator)
        self.entity_description = description
        self._attr_name = f"{description.name}"
        self._attr_unique_id = f"{coordinator.unique_id}_{description.key}"

    @property
    def native_value(self) -> StateType:
        """Return the native sensor value."""
        state = getattr(self.coordinator.data, self.entity_description.key)
        return state


import itertools
from computer import Computer
import time

guesses = [5, 6, 7, 8, 9]
source = list(itertools.permutations(guesses, 5))

maxN = 0
for i in source:
    print("\n\n\n")
    nums = list(i)
    print(nums)

    ampA = Computer(nums[0])
    ampB = Computer(nums[1])
    ampC = Computer(nums[2])
    ampD = Computer(nums[3])
    ampE = Computer(nums[4])

    amps = [ampA, ampB, ampC, ampD, ampE]
    atAmp = 0

    subTotal = 0
    running = True

    while(running):
        answer = amps[atAmp].getOutput(subTotal)
        if not isinstance(answer, bool):
            subTotal = answer
            atAmp = atAmp + 1 if atAmp < 4 else 0
        else:
            running = False

    if(subTotal > maxN):
        maxN = subTotal
        print(subTotal)

print(maxN)


import sys, os
sys.path.append(os.path.join(os.path.dirname(__file__), '..'))

import logging
logger = logging.getLogger(__name__)

from IR import ir
from optimizer import common

# 删除无效的reshape  输入形状和输出相同
class reshape_nop_eliminate():
    # 判断是否满足条件
    def match_conditions(self, node):
        if node.op_type == "Reshape":

            if len(node.input[0].dims) != 0 and len(node.output[0].dims) != 0 :
                if node.input[0].dims == node.output[0].dims:
                    return True
            else:
                logger.warn("can not get reshape input shape and output shape.")

        return False

    # 运行一次优化
    def run_pass(self, ir_graph):
        for node in ir_graph.node_list:
            if self.match_conditions(node):
                logger.info("---- eliminate node %s %s", node.op_type, node.output[0].name)
                ir_graph = common.eliminate_node(ir_graph, node.name)
                return True

        return False


import scipy.optimize as opt
import numpy as np
from .constructions import PrimitiveObject
import random

__all__ = ["Problem"]

class Problem:
    x0 = []
    constraints = []
    all_objects = {}
    # total number of degrees of freedom
    n = 0
    # object_names = ["B", "M", "circ"]
    object_names = []
    # object_dofs = [2, 2, 3]
    object_dofs = []
    # object_starts = [0, 2, 4]
    object_starts = [0]

    def _add_object(self, object_name, object_type, *vals):
        obj = PrimitiveObject(object_type, object_name, vals)
        self.all_objects[object_name] = obj
        return obj

    def add_fixed_object(self, object_name, object_type, *vals):
        assert len(vals) == object_type.degrees_of_freedom()
        return self._add_object(object_name, object_type, *vals)

    def add_param_object(self, object_name, object_type, *vals):
        vals = list(vals)
        dof = object_type.degrees_of_freedom()
        assert len(vals) <= dof
        for i in range(dof - len(vals)):
            vals.append(random.random())
        self.object_names.append(object_name)
        self.n += dof
        self.object_dofs.append(dof)
        last = self.object_starts[-1]
        self.object_starts.append(last + dof)
        self.x0 += vals
        return self._add_object(object_name, object_type, *vals)

    def add_constraint(self, constraint):
        self.constraints.append(constraint)

    def error(self):
        return sum(constraint.error(self.all_objects) \
            for constraint in self.constraints)

    def _error_vals(self, vals):
        for i, (start, length) in enumerate(zip(self.object_starts[:-1], self.object_dofs)):
            name = self.object_names[i]
            obj_vals = vals[start:start+length]
            self.all_objects[name].set_vals(obj_vals)
        return self.error()

    def _error_methods(self):
        constraints = []
        for constraint in self.constraints:
            constraints.append({
                "type": "ineq",
                "fun": lambda vals: constraint.error(self.all_objects)
            })
        return constraints

    def solve(self):
        x0 = np.reshape(self.x0, (self.n, 1))
        # tries to get `fun` as close to 0 as possible
        # while `constraints` are non-negative
        self.solution = opt.minimize(fun=self._error_vals, x0=x0, \
            method="COBYLA", constraints=self._error_methods()
        )

    def solution_points(self):
        if (self.solution is not None):
            return self.all_objects
        else:
            return None

if __name__ == "__main__":
    problem = Problem()
    problem.add_point("A")
    problem.add_point("B")
    problem.add_point("C")
    print(problem.solution())


import astropy
from astropy.io import fits
import glob
from glob import glob
import numpy
from numpy import *
import matplotlib.pyplot as plt
plt.ioff()
plt.close('all')
figdir = '/Users/rsimons/Desktop/clear/figures/continuum'

fields = ['GS4']

for field in fields:
    fls = glob('/Volumes/pegasus/clear/grizli_extractions/%s/*/Prep/*GrismFLT.fits'%field)

    for f, fl in enumerate(fls):

        print (f, fl)
        a = fits.open(fl)
        fig, axes = plt.subplots(2,2, figsize = (10,10))

        xmn = 500
        xmx = 1100
        vmn = 0.0001
        vmx = 0.08

        vmn_im = 0.
        vmx_im = 1.e-20
        im = a['DREF'].data
        #im-=im.ravel().min()
        axes[0,0].imshow(im, cmap = 'Greys_r', vmin = vmn_im, vmax = vmx_im)
        axes[0,1].imshow(a['GSCI'].data, cmap = 'viridis', vmin = vmn, vmax = vmx)
        axes[1,0].imshow(a['MODEL'].data, cmap = 'viridis',  vmin = vmn, vmax = vmx)
        axes[1,1].imshow(a['GSCI'].data - a['MODEL'].data, cmap = 'viridis',  vmin = vmn, vmax = vmx)

        for ax in axes.ravel():
            ax.axis('off')
            ax.set_xlim(xmn, xmx)
            ax.set_ylim(xmn, xmx)

        outfile = fl.split('/')[-1].replace('GrismFLT.fits', '%s.png'%field)
        fig.subplots_adjust(hspace = 0.0, wspace = 0.0, top = 1.0, bottom = 0.0, right = 1.0, left = 0.0)
        fig.savefig(figdir + '/' + outfile)


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sun Jan 31 10:28:31 2021

@author: sethschimmel
"""

import json
from collections import Counter

## open the first seven dict and the mesh dict
sevDict = json.load(open("/Users/sethschimmel/Documents/GitHub/CUNY-Capstone/data/sbir_2008to2018_VOCABINDEX.json","r"))

meshDict = json.load(open("/Users/sethschimmel/Documents/GitHub/CUNY-Capstone/data/sbir_2008to2018_MESHINDEX.json", "r"))

## merge the first seven dicts and the MeSH dict
for k,v in meshDict.items():
    sevDict[str(k)].update(v)

sevDict['9897']

## convert the list to a dictionary of raw frequencies/counts
for k,v in sevDict.items():
    for voc in list(v.keys()):
        sevDict[k][voc] = dict(Counter(sevDict[k][voc]))

sevDict['9897']

newDict = {}
# clean up the counted collection to have term and freq, to make it easier for viz
for k,v in sevDict.items():
    newDict[k]  = {}
    for voc in list(v.keys()):
        cleanVoc = []
        for term,freq in sevDict[k][voc].items():
            #print(str(term).count("_"))

            t = term.replace(voc,"")
            ## get rid of unigrams to save space...?
            #if t.count("_") == 0:
            #    pass
            #else:
            entry = {"t":t,"f":freq}
            cleanVoc.append(entry)
        vocName = voc.replace("*","").replace("_","")
        newDict[k][vocName] = cleanVoc

newDict['9897']

######## get the original award_id from before the word2vec pipline re-indexing
#### NOTE: this adds lots of suize to the output data json... isn't necessary really...
import pandas as pd

#origDf = pd.read_csv("/Users/sethschimmel/Documents/GitHub/w2v_pipeline_sbir/pipeline_src/datasets/sbir_2008to2018_textFields.csv")

#origDf.columns
#origDf.head()

#awardIds = list(origDf.award_id)

#for k,v in sevDict.items():
#    sevDict[k]["award_id"] = awardIds[int(k)]

#sevDict['9']

out_file = open("sbir_2008to2018_FULLINDEX_clean.json", "w")
json.dump(newDict, out_file)

out_file.close()


import common

data = common.read_file('2018/02/data.txt').splitlines()

twos = 0
threes = 0
for line in data:
    letters = dict()
    for l in line:
        if l in letters:
            letters[l] += 1
        else:
            letters[l] = 1

    if any(filter(lambda x: letters[x] == 2, letters)):
        twos += 1
    if any(filter(lambda x: letters[x] == 3, letters)):
        threes += 1
print(twos * threes)

for i1 in range(len(data)):
    for i2 in range(i1+1, len(data)):
        line1 = data[i1]
        line2 = data[i2]
        acc = ''
        for i in range(len(line1)):
            if line1[i] == line2[i]:
                acc += line1[i]
        if len(acc) + 1 == len(line1):
            print(acc)


"""A LightningModule that represents a task
where a model predicts a label for each frame
in a time series, e.g., each time bin in
a window from a spectrogram."""
from __future__ import annotations

import logging
from typing import Callable, ClassVar, Mapping

import torch

from .. import transforms
from ..common import labels
from . import base
from .definition import ModelDefinition
from .registry import model_family

logger = logging.getLogger(__name__)

@model_family
class FrameClassificationModel(base.Model):
    """Class that represents a family of neural network models
    that predicts a label for each frame in a time series,
    e.g., each time bin in a window from a spectrogram.

    The task of predicting a label for each frame in a series
    is one way of predicting annotations for a vocalization,
    where the annotations consist of a sequence
    of segments, each with an onset, offset, and label.
    The model maps the spectrogram window
    to a vector of labels for each frame, i.e., each time bin.

    To annotate a vocalization with such a model,
    the spectrogram is converted into a batch of
    consecutive non-overlapping windows,
    for which the model produces predictions.
    These predictions are then concatenated
    into a vector of labeled frames,
    from which the segments can be recovered.

    Post-processing can be applied to the vector
    to clean up noisy predictions
    before recovering the segments.

    Attributes
    ----------
    network : torch.nn.Module, dict
        An instance of a ``torch.nn.Module``
        that implements a neural network,
        or a ``dict`` that maps human-readable string names
        to a set of such instances.
    loss : torch.nn.Module, callable
        An instance of a ``torch.nn.Module``
        that implements a loss function,
        or a callable Python function that
        computes a scalar loss.
    optimizer : torch.optim.Optimizer
        An instance of a ``torch.optim.Optimizer`` class
        used with ``loss`` to optimize
        the parameters of ``network``.
    metrics : dict
        A ``dict`` that maps human-readable string names
        to ``Callable`` functions, used to measure
        performance of the model.
    post_tfm : callable
        Post-processing transform applied to predictions.
    labelmap : dict-like
        That maps human-readable labels to integers predicted by network.
    eval_labelmap : dict-like
        Mapping from labels to integers predicted by network
        that is used by ``validation_step``.
        This is used when mapping from network outputs back to labels
        to compute metrics that require strings, such as edit distance.
        If ``labelmap`` contains keys with multiple characters,
        this will be ``labelmap`` re-mapped so that all labels have
        single characters (except "unlabeled"), to avoid artificially
        changing the edit distance.
        See https://github.com/vocalpy/vak/issues/373 for more detail.
        If all keys (except "unlabeled") are single-character,
        then ``eval_labelmap`` will just be ``labelmap``.
    to_labels_eval : vak.transforms.frame_labels.ToLabels
        Instance of :class:`~vak.transforms.frame_labels.ToLabels`
        that uses ``eval_labelmap`` to convert labeled timebins
        to string labels inside of ``validation_step``,
        for computing edit distance.
    """

    definition: ClassVar[ModelDefinition]

    def __init__(
        self,
        labelmap: Mapping,
        network: torch.nn.Module | dict | None = None,
        loss: torch.nn.Module | Callable | None = None,
        optimizer: torch.optim.Optimizer | None = None,
        metrics: dict | None = None,
        post_tfm: Callable | None = None,
    ):
        """Initialize a new instance of a
        :class:`~vak.models.frame_classification_model.FrameClassificationModel`.

        Parameters
        ----------
        labelmap : dict-like
            That maps human-readable labels to integers predicted by network.
        network : torch.nn.Module, dict
            An instance of a :class:`torch.nn.Module`
            that implements a neural network,
            or a ``dict`` that maps human-readable string names
            to a set of such instances.
        loss : torch.nn.Module, callable
            An instance of a :class:`torch.nn.Module`
            that implements a loss function,
            or a callable Python function that
            computes a scalar loss.
        optimizer : torch.optim.Optimizer
            An instance of a ``torch.optim.Optimizer`` class
            used with ``loss`` to optimize
            the parameters of ``network``.
        metrics : dict
            A ``dict`` that maps human-readable string names
            to ``Callable`` functions, used to measure
            performance of the model.
        post_tfm : callable
            Post-processing transform applied to predictions.
        """
        super().__init__(
            network=network, loss=loss, optimizer=optimizer, metrics=metrics
        )

        self.labelmap = labelmap
        # replace any multiple character labels in mapping
        # with single-character labels
        # so that we do not affect edit distance computation
        # see https://github.com/NickleDave/vak/issues/373
        labelmap_keys = [lbl for lbl in labelmap.keys() if lbl != "unlabeled"]
        if any(
            [len(label) > 1 for label in labelmap_keys]
        ):  # only re-map if necessary
            # (to minimize chance of knock-on bugs)
            logger.info(
                "Detected that labelmap has keys with multiple characters:"
                f"\n{labelmap_keys}\n"
                "Re-mapping labelmap used with to_labels_eval transform, using "
                "function vak.labels.multi_char_labels_to_single_char"
            )
            self.eval_labelmap = labels.multi_char_labels_to_single_char(
                labelmap
            )
        else:
            self.eval_labelmap = labelmap

        self.to_labels_eval = transforms.frame_labels.ToLabels(
            self.eval_labelmap
        )
        self.post_tfm = post_tfm

    def configure_optimizers(self):
        """Returns the model's optimizer.

        Method required by ``lightning.LightningModule``.
        This method returns the ``optimizer`` instance passed into ``__init__``.
        If None was passed in, an instance that was created
        with default arguments will be returned.
        """
        return self.optimizer

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Run a forward pass through this model's network.

        Parameters
        ----------
        x : torch.Tensor
            Input to network, with shape that matches ``self.input_shape``.

        Returns
        -------
        out : torch.Tensor
            Output from network.
        """
        return self.network(x)

    def training_step(self, batch: tuple, batch_idx: int):
        """Perform one training step.

        Method required by ``lightning.LightningModule``.

        Parameters
        ----------
        batch : tuple
            A batch from a dataloader.
        batch_idx : int
            The index of this batch in the dataloader.

        Returns
        -------
        loss : torch.Tensor
            Scalar loss value computed by
            the loss function, ``self.loss``.
        """
        x, y = batch[0], batch[1]
        out = self.network(x)
        loss = self.loss(out, y)
        self.log("train_loss", loss)
        return loss

    def validation_step(self, batch: tuple, batch_idx: int):
        """Perform one validation step.

        Method required by ``lightning.LightningModule``.
        Logs metrics using ``self.log``

        Parameters
        ----------
        batch : tuple
            A batch from a dataloader.
        batch_idx : int
            The index of this batch in the dataloader.

        Returns
        -------
        None
        """
        x, y = batch["frames"], batch["frame_labels"]
        # remove "batch" dimension added by collate_fn to x
        # we keep for y because loss still expects the first dimension to be batch
        # TODO: fix this weirdness. Diff't collate_fn?
        if x.ndim in (5, 4):
            if x.shape[0] == 1:
                x = torch.squeeze(x, dim=0)
        else:
            raise ValueError(f"invalid shape for x: {x.shape}")

        out = self.network(x)
        # permute and flatten out
        # so that it has shape (1, number classes, number of time bins)
        # ** NOTICE ** just calling out.reshape(1, out.shape(1), -1) does not work, it will change the data
        out = out.permute(1, 0, 2)
        out = torch.flatten(out, start_dim=1)
        out = torch.unsqueeze(out, dim=0)
        # reduce to predictions, assuming class dimension is 1
        y_pred = torch.argmax(
            out, dim=1
        )  # y_pred has dims (batch size 1, predicted label per time bin)

        if "padding_mask" in batch:
            padding_mask = batch[
                "padding_mask"
            ]  # boolean: 1 where valid, 0 where padding
            # remove "batch" dimension added by collate_fn
            # because this extra dimension just makes it confusing to use the mask as indices
            if padding_mask.ndim == 2:
                if padding_mask.shape[0] == 1:
                    padding_mask = torch.squeeze(padding_mask, dim=0)
            else:
                raise ValueError(
                    f"invalid shape for padding mask: {padding_mask.shape}"
                )

            out = out[:, :, padding_mask]
            y_pred = y_pred[:, padding_mask]

        y_labels = self.to_labels_eval(y.cpu().numpy())
        y_pred_labels = self.to_labels_eval(y_pred.cpu().numpy())

        if self.post_tfm:
            y_pred_tfm = self.post_tfm(
                y_pred.cpu().numpy(),
            )
            y_pred_tfm_labels = self.to_labels_eval(y_pred_tfm)
            # convert back to tensor so we can compute accuracy
            y_pred_tfm = torch.from_numpy(y_pred_tfm).to(self.device)

        # TODO: figure out smarter way to do this
        for metric_name, metric_callable in self.metrics.items():
            if metric_name == "loss":
                self.log(
                    f"val_{metric_name}",
                    metric_callable(out, y),
                    batch_size=1,
                    on_step=True,
                )
            elif metric_name == "acc":
                self.log(
                    f"val_{metric_name}",
                    metric_callable(y_pred, y),
                    batch_size=1,
                )
                if self.post_tfm:
                    self.log(
                        f"val_{metric_name}_tfm",
                        metric_callable(y_pred_tfm, y),
                        batch_size=1,
                        on_step=True,
                    )
            elif (
                metric_name == "levenshtein"
                or metric_name == "segment_error_rate"
            ):
                self.log(
                    f"val_{metric_name}",
                    metric_callable(y_pred_labels, y_labels),
                    batch_size=1,
                )
                if self.post_tfm:
                    self.log(
                        f"val_{metric_name}_tfm",
                        metric_callable(y_pred_tfm_labels, y_labels),
                        batch_size=1,
                        on_step=True,
                    )

    def predict_step(self, batch: tuple, batch_idx: int):
        """Perform one prediction step.

        Method required by ``lightning.LightningModule``.

        Parameters
        ----------
        batch : tuple
            A batch from a dataloader.
        batch_idx : int
            The index of this batch in the dataloader.

        Returns
        -------
        y_pred : dict
            Where the key is "source_path" and the value
            is the output of the network;
            "source_path" is the path to the file
            containing the spectrogram
            for which a prediction was generated.
        """
        x, frames_path = batch["frames"].to(self.device), batch["frames_path"]
        if isinstance(frames_path, list) and len(frames_path) == 1:
            frames_path = frames_path[0]
        # TODO: fix this weirdness. Diff't collate_fn?
        if x.ndim in (5, 4):
            if x.shape[0] == 1:
                x = torch.squeeze(x, dim=0)
        else:
            raise ValueError(f"invalid shape for x: {x.shape}")
        y_pred = self.network(x)
        return {frames_path: y_pred}

    @classmethod
    def from_config(
        cls, config: dict, labelmap: Mapping, post_tfm: Callable | None = None
    ):
        """Return an initialized model instance from a config ``dict``

        Parameters
        ----------
        config : dict
            Returned by calling :func:`vak.config.models.map_from_path`
            or :func:`vak.config.models.map_from_config_dict`.
        post_tfm : callable
            Post-processing transformation.
            A callable applied to the network output.
            Default is None.

        Returns
        -------
        cls : vak.models.base.Model
            An instance of the model with its attributes
            initialized using parameters from ``config``.
        """
        network, loss, optimizer, metrics = cls.attributes_from_config(config)
        return cls(
            labelmap=labelmap,
            network=network,
            optimizer=optimizer,
            loss=loss,
            metrics=metrics,
            post_tfm=post_tfm,
        )


from django.core.management.base import BaseCommand
import json, bulk_sync
from discord_messages.models import Guild, Category, Channel, Author, Message
class Command(BaseCommand):

    help = "Import messages from a JSON file."

    def add_arguments(self, parser):
        parser.add_argument("filepath", type=str)

    def handle(self, *args, **options):
        data = json.load(open(options["filepath"],'r'))
        guild, _ = Guild.objects.get_or_create(
            discord_id=data["guild"]["id"],
            name=data["guild"]["name"],
            icon_url=data["guild"]["iconUrl"]
        )
        category = None
        channel_data = data["channel"]
        if "category" in channel_data.keys():
            category_data = data["channel"]
            category, _ = Category.objects.get_or_create(
                discord_id = category_data["categoryId"],
                name = category_data["category"],
                guild=guild,
            )
        channel, _  = Channel.objects.get_or_create(
            discord_id=data["channel"]["id"],
            name=data["channel"]["name"],
            topic=data["channel"]["topic"],
            guild=guild,
            category=category if category else None
        )
        authors = [m["author"] for m in data["messages"]]
        authors = list({v["id"]:v for v in authors}.values())
        Author.objects.bulk_create(Author(
                    discord_id=a["id"],
                    name=a["name"],
                    discriminator=a["discriminator"],
                    avatar_url=a["avatarUrl"],
                    color=a["color"],
                    is_bot=a["isBot"],
                ) for a in authors)
        Message.objects.bulk_create(
            Message(
                discord_id=m["id"],
                channel=channel,
                timestamp_sent= m["timestamp"],
                timestamp_edited=m["timestampEdited"],
                content=m["content"],
                author_id=m["author"]["id"]
                ) for m in data["messages"]
            )


import numpy as np
from gym_simple_gridworlds.envs.grid_env import GridEnv
from gym_simple_gridworlds.envs.grid_2dplot import plot_value_function, plot_policy
from collections import defaultdict
from copy import deepcopy
from matplotlib import pyplot as plt

def encode_policy(grid_env, policy_matrix=None):
    """
     Convert deterministic policy matrix into stochastic policy representation

     :param grid_env: MDP environment
     :param policy_matrix: Deterministic policy matrix (one action per state)

     :return: (dict of dict) Dictionary of dictionaries where each element corresponds to the
             probability of selection an action a at a given state s
     """

    height, width = grid_env.grid.shape

    if policy_matrix is None:

        policy_matrix = np.array([[3,      3,  3,  -1],
                                  [0, np.NaN,  0,  -1],
                                  [0,      2,  0,   2]])

    result_policy = defaultdict(lambda: defaultdict(float))

    for i in range(height):
        for j in range(width):
            s = grid_env.grid[i, j]
            if np.isnan(s) or grid_env.is_terminal_state(i, j):
                continue

            for a, _ in grid_env.ACTIONS.items():
                result_policy[int(s)][int(a)] = 0.0

            if policy_matrix[i, j] >= 0 or not np.isnan(policy_matrix[i, j]):
                result_policy[int(s)][int(policy_matrix[i, j])] = 1.0

    return result_policy

def define_random_policy(grid_env):
    """
    Define random deterministic policy for given environment

    :param grid_env: MDP environment
    :return: (matrix) Deterministic policy matrix
    """
    np.random.seed(grid_env.seed()[0])

    policy_matrix = np.array([np.random.choice(grid_env.get_actions(), 4).tolist(),
                              np.random.choice(grid_env.get_actions(), 4).tolist(),
                              np.random.choice(grid_env.get_actions(), 4).tolist()])

    for (x, y) in grid_env.terminal_states:
        policy_matrix[x, y] = -1

    for (x, y) in grid_env.obstacles:
        policy_matrix[x, y] = -1

    return policy_matrix

def policy_evaluation(env, policy):

    v = {s: 0.0 for s in env.get_states()}
    theta = 0.0001
    delta = 1000

    while delta > theta:
        delta = 0.0
        for s in v.keys():

            old_v = v[s]
            new_v = 0

            for action, probability in policy[s].items():
                state_sum = 0
                for s_next in env.get_states():
                    state_sum += env.state_transitions[s, action, s_next] * v[s_next]

                new_v += probability * (env.rewards[s, action] + env.gamma * state_sum)

            delta = max(delta, np.abs(old_v - new_v))
            v[s] = new_v
    return v

def decode_policy(grid_env, policy=None):
    """
     Convert stochastic policy representation (dict of dict) to deterministic policy matrix

     :param grid_env: MDP environment
     :param policy: stochastic policy (probability of each action at each state)

     :return: (matrix) Deterministic policy matrix (one action per state)
     """

    height, width = grid_env.grid.shape
    policy_matrix = np.full((height, width), -1)

    for s, actions in policy.items():
        x, y = np.argwhere(grid_env.grid == s)[0]

        if not grid_env.is_terminal_state(x,y):
          action_keys = list(actions.keys())
          policy_matrix[x, y] = action_keys[np.argmax(list(actions.values()))]

    return policy_matrix