import six


def get_probs_for_labels(labels, prediction_results):
  """ Given ML Workbench prediction results, get probs of each label for each instance.

  The prediction results are like:
  [
    {'predicted': 'daisy', 'probability': 0.8, 'predicted_2': 'rose', 'probability_2': 0.1},
    {'predicted': 'sunflower', 'probability': 0.9, 'predicted_2': 'daisy', 'probability_2': 0.01},
    ...
  ]

  Each instance is ordered by prob. But in some cases probs are needed for fixed
  order of labels. For example, given labels = ['daisy', 'rose', 'sunflower'], the
  results of above is expected to be:
  [
    [0.8, 0.1, 0.0],
    [0.01, 0.0, 0.9],
    ...
  ]
  Note that the sum of each instance may not be always 1. If model's top_n is set to
  none-zero, and is less than number of labels, then prediction results may not contain
  probs for all labels.

  Args:
    labels: a list of labels specifying the order of the labels.
    prediction_results: a pandas DataFrame containing prediction results, usually returned
        by get_prediction_results() call.

  Returns:
    A list of list of probs for each class.
  """

  probs = []
  if 'probability' in prediction_results:
    # 'probability' exists so top-n is set to none zero, and results are like
    # "predicted, predicted_2,...,probability,probability_2,...
    for i, r in prediction_results.iterrows():
      probs_one = [0.0] * len(labels)
      for k, v in six.iteritems(r):
        if v in labels and k.startswith('predicted'):
          if k == 'predict':
            prob_name = 'probability'
          else:
            prob_name = 'probability' + k[9:]
          probs_one[labels.index(v)] = r[prob_name]
      probs.append(probs_one)
    return probs
  else:
    # 'probability' does not exist, so top-n is set to zero. Results are like
    # "predicted, class_name1, class_name2,...
    for i, r in prediction_results.iterrows():
      probs_one = [0.0] * len(labels)
      for k, v in six.iteritems(r):
        if k in labels:
          probs_one[labels.index(k)] = v
      probs.append(probs_one)
    return probs