import six


def _add_is_control_and_control_id(df, ctrl_id):
  """Adds Is_Control and Control_Id columns to df.

  Args:
    df: A dataframe similar to the one returned by metrics_types.as_dataframe().
    ctrl_id: The control experiment id(s). For single control case, it can be
      basically any type that can be used as an experiment key except dict. For
      multiple controls, it should be a dict, with keys being control ids,
      values being list of corresponding experiment id(s).

  Returns:
    A DataFrame. Is_Control and Control_Id are either both presented or neither.
  """
  if 'Experiment_Id' not in df:
    return df
  if ctrl_id is None:
    if 'Is_Control' not in df and 'Control_Id' in df and 'Experiment_Id' in df:
      df['Is_Control'] = df['Control_Id'] == df['Experiment_Id']
    return df

  if not isinstance(ctrl_id, dict):
    df['Is_Control'] = df['Experiment_Id'] == ctrl_id
    df['Control_Id'] = ctrl_id
    return df

  ctrl_id_lookup = {}
  for ctrl, exp_ids in six.iteritems(ctrl_id):
    if not isinstance(exp_ids, list):
      raise ValueError('The experiment id(s) {} is not a list.'.format(exp_ids))
    for e in exp_ids:
      ctrl_id_lookup[e] = ctrl
    ctrl_id_lookup[ctrl] = ctrl  # The control of a control is itself.
  df['Is_Control'] = df['Experiment_Id'].apply(lambda k: k in ctrl_id)
  df['Control_Id'] = df['Experiment_Id'].apply(ctrl_id_lookup.get)
  return df