# ***************************************************************************************************
# Copyright 2015, 2019 National Technology & Engineering Solutions of Sandia, LLC (NTESS).
# Under the terms of Contract DE-NA0003525 with NTESS, the U.S. Government retains certain rights
# in this software.
# Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
# in compliance with the License.  You may obtain a copy of the License at
# http://www.apache.org/licenses/LICENSE-2.0 or in the LICENSE file in the root pyGSTi directory.
# ***************************************************************************************************
"""
Functions for creating RPE data sets
"""

from pygsti.data import datasetconstruction as _dsc


def make_rpe_data_set(model_or_dataset, string_list_d, num_samples, sample_error='binomial', seed=None):
    """
    Generate a fake RPE DataSet using the probabilities obtained from a model.

    Is a thin wrapper for pygsti.data.simulate_data, changing
    default behavior of sample_error, and taking a dictionary of circuits
    as input.

    Parameters
    ----------
    model_or_dataset : Model or DataSet object
        If a Model, the model whose probabilities generate the data.
        If a DataSet, the data set whose frequencies generate the data.

    string_list_d : Dictionary of list of (tuples or Circuits)
        Each tuple or Circuit contains operation labels and
        specifies a gate sequence whose counts are included
        in the returned DataSet.  The dictionary must have the key
        'totalStrList'; easiest if this dictionary is generated by
        make_rpe_string_list_d.

    num_samples : int or list of ints or None
        The simulated number of samples for each circuit.  This only
        has effect when  sample_error == "binomial" or "multinomial".  If
        an integer, all circuits have this number of total samples. If
        a list, integer elements specify the number of samples for the
        corresponding circuit.  If None, then model_or_dataset must be
        a DataSet, and total counts are taken from it (on a per-circuit
        basis).

    sample_error : string, optional
        What type of sample error is included in the counts.  Can be:

        - "none"  - no sample error:
          counts are floating point numbers such that the exact
          probability can be found by the ratio of count / total.
        - "round" - same as "none", except counts are rounded to the nearest
          integer.
        - "binomial" - the number of counts is taken from a binomial
          distribution. Distribution has parameters p = probability of the
          circuit and n = number of samples.  This can only be used when
          there are exactly two SPAM labels in model_or_dataset.
        - "multinomial" - counts are taken from a multinomial distribution.
          Distribution has parameters p_k = probability of the circuit
          using the k-th SPAM label and n = number of samples.  This should not
          be used for RPE.

    seed : int, optional
        If not None, a seed for numpy's random number generator, which
        is used to sample from the binomial or multinomial distribution.

    Returns
    -------
    DataSet
        A static data set filled with counts for the specified circuits.
    """
    return _dsc.simulate_data(model_or_dataset,
                              string_list_d['totalStrList'],
                              num_samples, sample_error=sample_error, seed=seed)
